From baffbe2cc703fa694c0c72baad13715b9bd28e57 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 18 Jul 2026 02:18:09 +0800 Subject: [PATCH] feat(executor): implement agent-scoped cache key isolation and enhance replay handling - Introduced agent-specific cache key isolation to ensure distinct prompt caching across agents. - Improved reasoning replay handling for Claude models, including cumulative tool turn restoration and validation. - Enhanced HTTP and WebSocket session consistency during execution scope and cache replay synchronization. - Added new unit tests and extended coverage for prompt caching, agent isolation, and replay cache functionality. Closes: #4352 --- .../cache/codex_reasoning_replay_cache.go | 162 +++++++- .../codex_reasoning_replay_cache_test.go | 117 ++++++ internal/home/client.go | 34 ++ internal/home/client_test.go | 20 + internal/runtime/executor/codex_executor.go | 358 ++++++++++++++++-- .../executor/codex_executor_cache_test.go | 59 +++ ...ex_executor_reasoning_replay_cache_test.go | 277 +++++++++++++- .../executor/codex_websockets_executor.go | 119 ++++-- .../codex_websockets_executor_test.go | 92 +++++ .../executor/helps/claude_code_session.go | 82 ++-- .../helps/claude_code_session_test.go | 61 +++ internal/runtime/executor/xai_executor.go | 2 +- 12 files changed, 1285 insertions(+), 98 deletions(-) diff --git a/internal/cache/codex_reasoning_replay_cache.go b/internal/cache/codex_reasoning_replay_cache.go index 274d131b8..bf76372fc 100644 --- a/internal/cache/codex_reasoning_replay_cache.go +++ b/internal/cache/codex_reasoning_replay_cache.go @@ -16,6 +16,9 @@ import ( ) const ( + // CodexReasoningReplayTurnType identifies an internal turn-boundary marker. + CodexReasoningReplayTurnType = "cpa_codex_replay_turn" + // CodexReasoningReplayCacheTTL limits how long encrypted reasoning replay // items stay in process memory. CodexReasoningReplayCacheTTL = 1 * time.Hour @@ -24,6 +27,12 @@ const ( // continuity. Oldest entries are evicted first. CodexReasoningReplayCacheMaxEntries = 10240 + // CodexReasoningReplayCacheMaxTurnsPerEntry bounds cumulative state for one agent. + CodexReasoningReplayCacheMaxTurnsPerEntry = 256 + + // CodexReasoningReplayCacheMaxBytesPerEntry bounds cumulative serialized items for one agent. + CodexReasoningReplayCacheMaxBytesPerEntry = 16 << 20 + // CodexReasoningReplayCacheEvictBatchSize leaves headroom after the cache // reaches capacity so high write volume does not rescan the map every turn. CodexReasoningReplayCacheEvictBatchSize = 128 @@ -42,6 +51,7 @@ var ( type codexReasoningReplayKVClient interface { KVGet(ctx context.Context, key string) ([]byte, bool, error) KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVCompareAndSwap(ctx context.Context, key string, expected []byte, expectedExists bool, value []byte, ttl time.Duration) (bool, error) KVDel(ctx context.Context, keys ...string) (int64, error) KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) } @@ -105,13 +115,132 @@ func CacheCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, se return true } -// GetCodexReasoningReplayItem retrieves a normalized reasoning replay item. +// AppendCodexReasoningReplayItemsBestEffort appends one completed turn to existing replay state. +func AppendCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool { + if ctx == nil { + ctx = context.Background() + } + key := codexReasoningReplayCacheKey(modelName, sessionKey) + if key == "" { + return false + } + normalized, ok := normalizeCodexReasoningReplayItems(items) + if !ok { + return false + } + if client, homeMode, errClient := currentCodexReasoningReplayKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errClient) + return false + } + kvKey := codexReasoningReplayKVKey(modelName, sessionKey) + const maxCASAttempts = 32 + for attempt := 0; attempt < maxCASAttempts; attempt++ { + if errContext := ctx.Err(); errContext != nil { + return false + } + existingRaw, found, errGet := client.KVGet(ctx, kvKey) + if errGet != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errGet) + return false + } + var existing [][]byte + if found { + if errUnmarshal := json.Unmarshal(existingRaw, &existing); errUnmarshal != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errUnmarshal) + return false + } + } + combined := appendCodexReasoningReplayTurn(existing, normalized) + raw, errMarshal := json.Marshal(combined) + if errMarshal != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errMarshal) + return false + } + written, errCAS := client.KVCompareAndSwap(ctx, kvKey, existingRaw, found, raw, CodexReasoningReplayCacheTTL) + if errCAS != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errCAS) + return false + } + if written { + return true + } + } + log.Warn("home kv best-effort codex reasoning replay append exhausted compare-and-swap attempts") + return false + } + + cacheCleanupOnce.Do(startCacheCleanup) + now := time.Now() + codexReasoningReplayMu.Lock() + entry := codexReasoningReplayEntries[key] + if now.Sub(entry.Timestamp) > CodexReasoningReplayCacheTTL { + entry.Items = nil + } + entry.Items = appendCodexReasoningReplayTurn(entry.Items, normalized) + entry.Timestamp = now + codexReasoningReplayEntries[key] = entry + if len(codexReasoningReplayEntries) > CodexReasoningReplayCacheMaxEntries { + evictOldestCodexReasoningReplayEntries(CodexReasoningReplayCacheEvictBatchSize) + } + codexReasoningReplayMu.Unlock() + return true +} + +func appendCodexReasoningReplayTurn(existing, turn [][]byte) [][]byte { + if len(existing) > 0 && strings.TrimSpace(gjson.GetBytes(existing[0], "type").String()) != CodexReasoningReplayTurnType { + existing = nil + } + turnID := "" + if len(turn) > 0 && strings.TrimSpace(gjson.GetBytes(turn[0], "type").String()) == CodexReasoningReplayTurnType { + turnID = strings.TrimSpace(gjson.GetBytes(turn[0], "id").String()) + } + if turnID != "" { + for _, item := range existing { + if strings.TrimSpace(gjson.GetBytes(item, "type").String()) == CodexReasoningReplayTurnType && + strings.TrimSpace(gjson.GetBytes(item, "id").String()) == turnID { + return trimCodexReasoningReplayItems(cloneCodexReasoningReplayItems(existing)) + } + } + } + combined := make([][]byte, 0, len(existing)+len(turn)) + combined = append(combined, cloneCodexReasoningReplayItems(existing)...) + combined = append(combined, cloneCodexReasoningReplayItems(turn)...) + return trimCodexReasoningReplayItems(combined) +} + +func trimCodexReasoningReplayItems(items [][]byte) [][]byte { + for { + turnStarts := []int{0} + totalBytes := 0 + for index, item := range items { + totalBytes += len(item) + if index > 0 && strings.TrimSpace(gjson.GetBytes(item, "type").String()) == CodexReasoningReplayTurnType { + turnStarts = append(turnStarts, index) + } + } + if len(turnStarts) <= CodexReasoningReplayCacheMaxTurnsPerEntry && totalBytes <= CodexReasoningReplayCacheMaxBytesPerEntry { + return items + } + if len(turnStarts) <= 1 { + return nil + } + items = items[turnStarts[1]:] + } +} + +// GetCodexReasoningReplayItem retrieves the first normalized upstream replay item. func GetCodexReasoningReplayItem(modelName, sessionKey string) ([]byte, bool) { items, ok := GetCodexReasoningReplayItems(modelName, sessionKey) - if !ok || len(items) == 0 { + if !ok { return nil, false } - return items[0], true + for _, item := range items { + if strings.TrimSpace(gjson.GetBytes(item, "type").String()) != CodexReasoningReplayTurnType { + return item, true + } + } + return nil, false } // GetCodexReasoningReplayItems retrieves normalized assistant output items. @@ -223,12 +352,15 @@ func normalizeCodexReasoningReplayItems(items [][]byte) ([][]byte, bool) { normalized = append(normalized, normalizedItem) } } + normalized = trimCodexReasoningReplayItems(normalized) return normalized, len(normalized) > 0 } func normalizeCodexReasoningReplayItem(item []byte) ([]byte, bool) { itemResult := gjson.ParseBytes(item) switch strings.TrimSpace(itemResult.Get("type").String()) { + case CodexReasoningReplayTurnType: + return normalizeCodexReasoningReplayTurn(itemResult) case "reasoning": return normalizeCodexReasoningReplayReasoningItem(itemResult) case "function_call": @@ -240,6 +372,30 @@ func normalizeCodexReasoningReplayItem(item []byte) ([]byte, bool) { } } +func normalizeCodexReasoningReplayTurn(itemResult gjson.Result) ([]byte, bool) { + turnID := strings.TrimSpace(itemResult.Get("id").String()) + if turnID == "" { + return nil, false + } + normalized := []byte(`{"type":"` + CodexReasoningReplayTurnType + `"}`) + normalized, _ = sjson.SetBytes(normalized, "id", turnID) + if fingerprint := strings.TrimSpace(itemResult.Get("assistant_fingerprint").String()); fingerprint != "" { + normalized, _ = sjson.SetBytes(normalized, "assistant_fingerprint", fingerprint) + } + if fingerprint := strings.TrimSpace(itemResult.Get("request_fingerprint").String()); fingerprint != "" { + normalized, _ = sjson.SetBytes(normalized, "request_fingerprint", fingerprint) + } + callIDs := itemResult.Get("call_ids") + if callIDs.IsArray() { + for _, callIDResult := range callIDs.Array() { + if callID := strings.TrimSpace(callIDResult.String()); callID != "" { + normalized, _ = sjson.SetBytes(normalized, "call_ids.-1", callID) + } + } + } + return normalized, true +} + func normalizeCodexReasoningReplayReasoningItem(itemResult gjson.Result) ([]byte, bool) { encryptedContentResult := itemResult.Get("encrypted_content") if encryptedContentResult.Type != gjson.String { diff --git a/internal/cache/codex_reasoning_replay_cache_test.go b/internal/cache/codex_reasoning_replay_cache_test.go index 8bfe494f8..f5d05d739 100644 --- a/internal/cache/codex_reasoning_replay_cache_test.go +++ b/internal/cache/codex_reasoning_replay_cache_test.go @@ -1,18 +1,22 @@ package cache import ( + "bytes" "context" "encoding/base64" "encoding/json" "errors" "fmt" + "sync" "testing" "time" homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/tidwall/gjson" ) type fakeCodexReasoningReplayKVClient struct { + mu sync.Mutex values map[string][]byte getErr error setErr error @@ -31,6 +35,8 @@ func newFakeCodexReasoningReplayKVClient() *fakeCodexReasoningReplayKVClient { } func (c *fakeCodexReasoningReplayKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.mu.Lock() + defer c.mu.Unlock() c.getCount++ if c.getErr != nil { return nil, false, c.getErr @@ -43,6 +49,8 @@ func (c *fakeCodexReasoningReplayKVClient) KVGet(_ context.Context, key string) } func (c *fakeCodexReasoningReplayKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() c.setCount++ c.lastSetTTL = opts.EX if c.setErr != nil { @@ -52,7 +60,25 @@ func (c *fakeCodexReasoningReplayKVClient) KVSet(_ context.Context, key string, return true, nil } +func (c *fakeCodexReasoningReplayKVClient) KVCompareAndSwap(_ context.Context, key string, expected []byte, expectedExists bool, value []byte, ttl time.Duration) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.setCount++ + c.lastSetTTL = ttl + if c.setErr != nil { + return false, c.setErr + } + current, exists := c.values[key] + if exists != expectedExists || (exists && !bytes.Equal(current, expected)) { + return false, nil + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + func (c *fakeCodexReasoningReplayKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { + c.mu.Lock() + defer c.mu.Unlock() c.delCount++ if c.delErr != nil { return 0, c.delErr @@ -68,6 +94,8 @@ func (c *fakeCodexReasoningReplayKVClient) KVDel(_ context.Context, keys ...stri } func (c *fakeCodexReasoningReplayKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() c.expireCount++ c.lastExpireTTL = ttl if c.expireErr != nil { @@ -185,6 +213,95 @@ func TestCodexReasoningReplayBestEffortHomeWriteFailureDoesNotUseLocalCache(t *t } } +func TestCodexReasoningReplayAppendPreservesCumulativeTurnsInHome(t *testing.T) { + ClearCodexReasoningReplayCache() + t.Cleanup(ClearCodexReasoningReplayCache) + client := newFakeCodexReasoningReplayKVClient() + useFakeCodexReasoningReplayKVClient(t, client, true, nil) + + first := [][]byte{ + []byte(`{"type":"` + CodexReasoningReplayTurnType + `","id":"turn-1","assistant_fingerprint":"answer-1"}`), + validCodexReasoningReplayItemForTest(11), + } + second := [][]byte{ + []byte(`{"type":"` + CodexReasoningReplayTurnType + `","id":"turn-2","call_ids":["call-2"]}`), + validCodexReasoningReplayItemForTest(12), + } + if !AppendCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "session-home-append", first) { + t.Fatal("first append failed") + } + if !AppendCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "session-home-append", second) { + t.Fatal("second append failed") + } + if !AppendCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "session-home-append", second) { + t.Fatal("duplicate append failed") + } + + items, found, errGet := GetCodexReasoningReplayItemsRequired(context.Background(), "gpt-5.4", "session-home-append") + if errGet != nil || !found { + t.Fatalf("get cumulative turns = found %v err %v", found, errGet) + } + if len(items) != 4 { + t.Fatalf("cumulative item count = %d, want 4: %q", len(items), items) + } + if got := gjson.GetBytes(items[0], "id").String(); got != "turn-1" { + t.Fatalf("first turn id = %q, want turn-1", got) + } + if got := gjson.GetBytes(items[2], "id").String(); got != "turn-2" { + t.Fatalf("second turn id = %q, want turn-2", got) + } +} + +func TestCodexReasoningReplayAppendHomeCASPreservesConcurrentTurns(t *testing.T) { + ClearCodexReasoningReplayCache() + t.Cleanup(ClearCodexReasoningReplayCache) + client := newFakeCodexReasoningReplayKVClient() + useFakeCodexReasoningReplayKVClient(t, client, true, nil) + + const turnCount = 16 + var waitGroup sync.WaitGroup + for turn := 0; turn < turnCount; turn++ { + waitGroup.Add(1) + go func(turnID int) { + defer waitGroup.Done() + items := [][]byte{ + []byte(fmt.Sprintf(`{"type":"%s","id":"turn-%d"}`, CodexReasoningReplayTurnType, turnID)), + validCodexReasoningReplayItemForTest(byte(30 + turnID)), + } + if !AppendCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "session-home-concurrent", items) { + t.Errorf("append turn %d failed", turnID) + } + }(turn) + } + waitGroup.Wait() + + items, found, errGet := GetCodexReasoningReplayItemsRequired(context.Background(), "gpt-5.4", "session-home-concurrent") + if errGet != nil || !found { + t.Fatalf("get concurrent turns = found %v err %v", found, errGet) + } + if len(items) != turnCount*2 { + t.Fatalf("concurrent cumulative item count = %d, want %d", len(items), turnCount*2) + } +} + +func TestCodexReasoningReplayAppendBoundsTurnsPerEntry(t *testing.T) { + items := make([][]byte, 0, (CodexReasoningReplayCacheMaxTurnsPerEntry+1)*2) + for turn := 0; turn <= CodexReasoningReplayCacheMaxTurnsPerEntry; turn++ { + items = append(items, + []byte(fmt.Sprintf(`{"type":"%s","id":"turn-%d"}`, CodexReasoningReplayTurnType, turn)), + validCodexReasoningReplayItemForTest(byte(50+turn)), + ) + } + + trimmed := trimCodexReasoningReplayItems(items) + if len(trimmed) != CodexReasoningReplayCacheMaxTurnsPerEntry*2 { + t.Fatalf("trimmed item count = %d, want %d", len(trimmed), CodexReasoningReplayCacheMaxTurnsPerEntry*2) + } + if firstID := gjson.GetBytes(trimmed[0], "id").String(); firstID != "turn-1" { + t.Fatalf("first retained turn = %q, want turn-1", firstID) + } +} + func TestCodexReasoningReplayHomeRejectsEmptyScopeWithoutKV(t *testing.T) { client := newFakeCodexReasoningReplayKVClient() useFakeCodexReasoningReplayKVClient(t, client, true, nil) diff --git a/internal/home/client.go b/internal/home/client.go index e24878566..279f5e832 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -647,6 +647,40 @@ func (c *Client) KVSetNX(ctx context.Context, key string, value []byte, ttl time return c.KVSet(ctx, key, value, opts) } +// KVCompareAndSwap atomically replaces a value only when its current state matches the expected state. +func (c *Client) KVCompareAndSwap(ctx context.Context, key string, expected []byte, expectedExists bool, value []byte, ttl time.Duration) (bool, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return false, errClient + } + const script = ` +local current = redis.call("GET", KEYS[1]) +if ARGV[1] == "1" then + if not current or current ~= ARGV[2] then + return 0 + end +elseif current then + return 0 +end +local ttl = tonumber(ARGV[4]) +if ttl and ttl > 0 then + redis.call("SET", KEYS[1], ARGV[3], "PX", ttl) +else + redis.call("SET", KEYS[1], ARGV[3]) +end +return 1 +` + expectedFlag := "0" + if expectedExists { + expectedFlag = "1" + } + result, errEval := cmd.Eval(ctx, script, []string{key}, expectedFlag, expected, value, durationCeil(ttl, time.Millisecond)).Int64() + if errEval != nil { + return false, errEval + } + return result == 1, nil +} + func (c *Client) KVDel(ctx context.Context, keys ...string) (int64, error) { if len(keys) == 0 { return 0, nil diff --git a/internal/home/client_test.go b/internal/home/client_test.go index bf7568cb5..b9480c080 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -254,6 +254,26 @@ func TestKVSetConditionUnmetReturnsFalse(t *testing.T) { } } +func TestKVCompareAndSwapReturnsScriptResult(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "EVAL") { + return ":1\r\n" + } + return "-ERR unexpected command\r\n" + }) + + swapped, errCAS := client.KVCompareAndSwap(context.Background(), "key", []byte("old"), true, []byte("new"), 1500*time.Millisecond) + if errCAS != nil { + t.Fatalf("KVCompareAndSwap() error = %v", errCAS) + } + if !swapped { + t.Fatal("KVCompareAndSwap() swapped = false, want true") + } + if lastCommand := commands.Last(); len(lastCommand) < 2 || !strings.EqualFold(lastCommand[0], "EVAL") { + t.Fatalf("last command = %#v, want EVAL", lastCommand) + } +} + func TestKVMSetUsesStableKeyOrder(t *testing.T) { client, commands := newRedisCommandTestClient(t, func(args []string) string { if len(args) > 0 && strings.EqualFold(args[0], "MSET") { diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 004253779..7ed57b315 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -302,8 +302,9 @@ func translateCodexRequestPair(from, to sdktranslator.Format, model string, orig } type codexReasoningReplayScope struct { - modelName string - sessionKey string + modelName string + sessionKey string + requestFingerprint string } func (s codexReasoningReplayScope) valid() bool { @@ -324,11 +325,7 @@ func applyCodexReasoningReplayCacheRequired(ctx context.Context, from sdktransla if errReplay != nil || !ok { return body, scope, errReplay } - items = filterCodexReasoningReplayItemsForInput(body, items) - if len(items) == 0 { - return body, scope, nil - } - updated, ok := insertCodexReasoningReplayItems(body, items) + updated, ok := insertCodexReasoningReplayTurns(body, items) if !ok { return body, scope, nil } @@ -339,9 +336,15 @@ func codexReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslato if !codexReasoningReplayEnabledForSource(from) { return codexReasoningReplayScope{} } + modelName := strings.TrimSpace(gjson.GetBytes(body, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + inputItems := gjson.GetBytes(body, "input").Array() return codexReasoningReplayScope{ - modelName: thinking.ParseSuffix(req.Model).ModelName, - sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body), + modelName: modelName, + sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body), + requestFingerprint: codexReplayInputPrefixFingerprint(inputItems, len(inputItems)), } } @@ -354,17 +357,19 @@ func sourceFormatEqual(from, want sdktranslator.Format) bool { } func codexClaudeCodeReplaySessionKey(ctx context.Context, payload []byte, headers http.Header) string { - sessionID := helps.ExtractClaudeCodeSessionID(ctx, payload, headers) - if sessionID == "" { - return "" - } - return "claude:" + sessionID + sessionKey, _ := helps.ClaudeCodeExecutionScope(ctx, payload, headers) + return sessionKey } func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) string { if ctx == nil { ctx = context.Background() } + if sourceFormatEqual(from, sdktranslator.FormatClaude) { + if sessionKey := codexClaudeCodeReplaySessionKey(ctx, req.Payload, opts.Headers); sessionKey != "" { + return sessionKey + } + } if value := metadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { return "execution:" + value } @@ -385,9 +390,6 @@ func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Form return value } } - if sourceFormatEqual(from, sdktranslator.FormatClaude) { - return codexClaudeCodeReplaySessionKey(ctx, req.Payload, opts.Headers) - } if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { return "prompt-cache:" + uuid.NewSHA1(uuid.NameSpaceOID, []byte("cli-proxy-api:codex:prompt-cache:"+apiKey)).String() @@ -483,6 +485,261 @@ func codexInputHasValidReasoningEncryptedContent(body []byte) bool { return false } +type codexReasoningReplayTurn struct { + marked bool + assistantFingerprint string + requestFingerprint string + callIDs []string + items [][]byte +} + +func insertCodexReasoningReplayTurns(body []byte, replayItems [][]byte) ([]byte, bool) { + input := gjson.GetBytes(body, "input") + if !input.IsArray() || len(replayItems) == 0 { + return body, false + } + inputItems := input.Array() + turns := splitCodexReasoningReplayTurns(replayItems) + insertions := make(map[int][][]byte) + usedAnchorIndexes := make(map[int]bool) + fallbackAnchorEnd := len(inputItems) - 1 + inserted := false + for turnIndex := len(turns) - 1; turnIndex >= 0; turnIndex-- { + turn := turns[turnIndex] + if len(turn.items) == 0 { + continue + } + if !turn.marked { + items := filterCodexReasoningReplayItemsForInput(body, turn.items) + if len(items) == 0 { + continue + } + index := codexReasoningReplayInsertIndex(inputItems, items) + items = codexAlignReasoningReplayToolCallIDs(inputItems, items) + insertions[index] = append(items, insertions[index]...) + inserted = true + continue + } + + anchorIndex, matched := codexReasoningReplayTurnAnchorIndex(inputItems, turn, fallbackAnchorEnd, usedAnchorIndexes) + if !matched { + continue + } + usedAnchorIndexes[anchorIndex] = true + if turn.requestFingerprint == "" { + fallbackAnchorEnd = anchorIndex - 1 + } + items := filterCodexReasoningReplayTurnItems(inputItems, turn.items) + if len(items) == 0 { + continue + } + items = codexAlignReasoningReplayToolCallIDs(inputItems, items) + insertions[anchorIndex] = append(items, insertions[anchorIndex]...) + inserted = true + } + if !inserted { + return body, false + } + + items := make([]string, 0, len(inputItems)+len(replayItems)) + for index, inputItem := range inputItems { + for _, replayItem := range insertions[index] { + items = append(items, string(replayItem)) + } + items = append(items, inputItem.Raw) + } + for _, replayItem := range insertions[len(inputItems)] { + items = append(items, string(replayItem)) + } + updated, err := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(items, ",")+"]")) + if err != nil { + return body, false + } + return updated, true +} + +func splitCodexReasoningReplayTurns(items [][]byte) []codexReasoningReplayTurn { + turns := make([]codexReasoningReplayTurn, 0) + current := codexReasoningReplayTurn{} + appendCurrent := func() { + if len(current.items) > 0 { + turns = append(turns, current) + } + } + for _, item := range items { + itemResult := gjson.ParseBytes(item) + if strings.TrimSpace(itemResult.Get("type").String()) == internalcache.CodexReasoningReplayTurnType { + appendCurrent() + current = codexReasoningReplayTurn{ + marked: true, + assistantFingerprint: strings.TrimSpace(itemResult.Get("assistant_fingerprint").String()), + requestFingerprint: strings.TrimSpace(itemResult.Get("request_fingerprint").String()), + } + if callIDs := itemResult.Get("call_ids"); callIDs.IsArray() { + for _, callIDResult := range callIDs.Array() { + if callID := strings.TrimSpace(callIDResult.String()); callID != "" { + current.callIDs = append(current.callIDs, callID) + } + } + } + continue + } + current.items = append(current.items, item) + } + appendCurrent() + return turns +} + +func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexReasoningReplayTurn, fallbackEnd int, used map[int]bool) (int, bool) { + searchEnd := fallbackEnd + if turn.requestFingerprint != "" { + searchEnd = len(inputItems) - 1 + } + if searchEnd >= len(inputItems) { + searchEnd = len(inputItems) - 1 + } + matchesRequestPrefix := func(index int) bool { + return turn.requestFingerprint == "" || codexReplayInputPrefixFingerprint(inputItems, index) == turn.requestFingerprint + } + if len(turn.callIDs) > 0 { + callIDs := make(map[string]bool) + for _, callID := range turn.callIDs { + for _, candidate := range codexReplayComparableCallIDs(callID) { + callIDs[candidate] = true + } + } + for index := searchEnd; index >= 0; index-- { + if used[index] || !matchesRequestPrefix(index) { + continue + } + itemType := strings.TrimSpace(inputItems[index].Get("type").String()) + if itemType != "function_call" && itemType != "custom_tool_call" && itemType != "function_call_output" && itemType != "custom_tool_call_output" { + continue + } + for _, candidate := range codexReplayComparableCallIDs(inputItems[index].Get("call_id").String()) { + if callIDs[candidate] { + return index, true + } + } + } + } + if turn.assistantFingerprint != "" { + for index := searchEnd; index >= 0; index-- { + if used[index] || !matchesRequestPrefix(index) { + continue + } + if codexReplayAssistantMessageFingerprint(inputItems[index]) == turn.assistantFingerprint { + return index, true + } + } + } + if len(turn.callIDs) == 0 && turn.assistantFingerprint == "" { + return codexReasoningReplayInsertIndex(inputItems, turn.items), true + } + return 0, false +} + +func filterCodexReasoningReplayTurnItems(inputItems []gjson.Result, items [][]byte) [][]byte { + existingReasoning := make(map[string]bool) + existingCalls := make(map[string]bool) + existingOutputs := make(map[string]bool) + for _, inputItem := range inputItems { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + switch itemType { + case "reasoning": + if encryptedContent := strings.TrimSpace(inputItem.Get("encrypted_content").String()); encryptedContent != "" { + existingReasoning[encryptedContent] = true + } + case "function_call_output", "custom_tool_call_output": + for _, candidate := range codexReplayComparableCallIDs(inputItem.Get("call_id").String()) { + existingOutputs[candidate] = true + } + } + for _, key := range codexReplayToolCallKeys(inputItem) { + existingCalls[key] = true + } + } + + filtered := make([][]byte, 0, len(items)) + for _, item := range items { + itemResult := gjson.ParseBytes(item) + switch strings.TrimSpace(itemResult.Get("type").String()) { + case "reasoning": + if existingReasoning[strings.TrimSpace(itemResult.Get("encrypted_content").String())] { + continue + } + case "function_call", "custom_tool_call": + keys := codexReplayToolCallKeys(itemResult) + if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { + continue + } + hasMatchingOutput := false + for _, candidate := range codexReplayComparableCallIDs(itemResult.Get("call_id").String()) { + if existingOutputs[candidate] { + hasMatchingOutput = true + break + } + } + if !hasMatchingOutput { + continue + } + for _, key := range keys { + existingCalls[key] = true + } + default: + continue + } + filtered = append(filtered, item) + } + return filtered +} + +func codexReplayAssistantMessageFingerprint(item gjson.Result) string { + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType != "" && itemType != "message" { + return "" + } + if !strings.EqualFold(strings.TrimSpace(item.Get("role").String()), "assistant") { + return "" + } + content := item.Get("content") + var builder strings.Builder + if content.Type == gjson.String { + builder.WriteString(content.String()) + } else if content.IsArray() { + for _, part := range content.Array() { + switch strings.TrimSpace(part.Get("type").String()) { + case "input_text", "output_text": + builder.WriteString(part.Get("text").String()) + case "refusal": + builder.WriteString("\x00refusal\x00") + builder.WriteString(part.Get("refusal").String()) + default: + return "" + } + } + } else { + return "" + } + if builder.Len() == 0 { + return "" + } + sum := sha256.Sum256([]byte(builder.String())) + return hex.EncodeToString(sum[:]) +} + +func codexReplayInputPrefixFingerprint(inputItems []gjson.Result, end int) string { + if end < 0 || end > len(inputItems) { + return "" + } + hasher := sha256.New() + for index := 0; index < end; index++ { + _, _ = hasher.Write([]byte("\x00item\x00")) + _, _ = hasher.Write([]byte(inputItems[index].Raw)) + } + return hex.EncodeToString(hasher.Sum(nil)) +} + func filterCodexReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte { input := gjson.GetBytes(body, "input") if !input.IsArray() { @@ -751,18 +1008,53 @@ func cacheCodexReasoningReplayFromCompleted(scope codexReasoningReplayScope, com if !output.IsArray() { return } - items := make([][]byte, 0, len(output.Array())) + replayItems := make([][]byte, 0, len(output.Array())) + callIDs := make([]string, 0) + assistantFingerprint := "" for _, item := range output.Array() { switch strings.TrimSpace(item.Get("type").String()) { - case "reasoning", "function_call", "custom_tool_call": - items = append(items, []byte(item.Raw)) - default: - continue + case "reasoning": + replayItems = append(replayItems, []byte(item.Raw)) + case "function_call", "custom_tool_call": + replayItems = append(replayItems, []byte(item.Raw)) + if callID := strings.TrimSpace(item.Get("call_id").String()); callID != "" { + callIDs = append(callIDs, callID) + } + case "message": + if fingerprint := codexReplayAssistantMessageFingerprint(item); fingerprint != "" { + assistantFingerprint = fingerprint + } } } - if !internalcache.CacheCodexReasoningReplayItemsBestEffort(context.Background(), scope.modelName, scope.sessionKey, items) { - internalcache.DeleteCodexReasoningReplayItem(scope.modelName, scope.sessionKey) + if len(replayItems) == 0 { + return } + + hasher := sha256.New() + _, _ = hasher.Write([]byte(scope.requestFingerprint)) + _, _ = hasher.Write([]byte("\x00assistant\x00" + assistantFingerprint)) + for _, callID := range callIDs { + _, _ = hasher.Write([]byte("\x00call\x00" + callID)) + } + for _, item := range replayItems { + _, _ = hasher.Write([]byte("\x00item\x00")) + _, _ = hasher.Write(item) + } + marker := []byte(`{"type":"` + internalcache.CodexReasoningReplayTurnType + `"}`) + marker, _ = sjson.SetBytes(marker, "id", hex.EncodeToString(hasher.Sum(nil))) + if assistantFingerprint != "" { + marker, _ = sjson.SetBytes(marker, "assistant_fingerprint", assistantFingerprint) + } + if scope.requestFingerprint != "" { + marker, _ = sjson.SetBytes(marker, "request_fingerprint", scope.requestFingerprint) + } + for _, callID := range callIDs { + marker, _ = sjson.SetBytes(marker, "call_ids.-1", callID) + } + items := make([][]byte, 0, len(replayItems)+1) + items = append(items, marker) + items = append(items, replayItems...) + internalcache.AppendCodexReasoningReplayItemsBestEffort(context.Background(), scope.modelName, scope.sessionKey, items) } func clearCodexReasoningReplayOnInvalidSignature(ctx context.Context, scope codexReasoningReplayScope, statusCode int, body []byte) error { @@ -864,7 +1156,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re url := strings.TrimSuffix(baseURL, "/") + "/responses" var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body) + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) if err != nil { return resp, err } @@ -1040,7 +1332,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A url := strings.TrimSuffix(baseURL, "/") + "/responses/compact" var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body) + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) if err != nil { return resp, err } @@ -1155,7 +1447,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au url := strings.TrimSuffix(baseURL, "/") + "/responses" var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body) + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) if err != nil { return nil, err } @@ -1506,10 +1798,18 @@ type codexIdentityReplacement struct { confused string } -func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Format, url string, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, userPayload []byte, rawJSON []byte) (*http.Request, []byte, codexIdentityConfuseState, error) { +func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Format, url string, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, userPayload []byte, rawJSON []byte, headerSets ...http.Header) (*http.Request, []byte, codexIdentityConfuseState, error) { + var headers http.Header + if len(headerSets) > 0 { + headers = headerSets[0] + } var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, nil) + modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, headers) if errCache != nil { return nil, nil, codexIdentityConfuseState{}, errCache } diff --git a/internal/runtime/executor/codex_executor_cache_test.go b/internal/runtime/executor/codex_executor_cache_test.go index 8e28340f4..c0b7523b4 100644 --- a/internal/runtime/executor/codex_executor_cache_test.go +++ b/internal/runtime/executor/codex_executor_cache_test.go @@ -307,3 +307,62 @@ func TestCodexExecutorCacheHelper_ClaudeUsesSessionHeader(t *testing.T) { t.Fatalf("same Claude Code session header produced different prompt_cache_key: first=%q second=%q", firstKey, secondKey) } } + +func TestCodexExecutorCacheHelper_ClaudeAgentScopeUsesResolvedModelAcrossHTTPAndWebsocket(t *testing.T) { + executor := &CodexExecutor{} + url := "https://example.com/responses" + req := cliproxyexecutor.Request{ + Model: "requested-alias-high", + Payload: []byte(`{"model":"requested-alias","messages":[{"role":"user","content":"hello"}]}`), + } + rootHeaders := http.Header{} + rootHeaders.Set(helps.ClaudeCodeSessionHeader, "resolved-model-session") + childHeaders := rootHeaders.Clone() + childHeaders.Set(helps.ClaudeCodeAgentHeader, "agent-a") + rawJSON := []byte(`{"model":"gpt-5.4","stream":true}`) + + rootRequest, _, _, errRoot := executor.cacheHelper(context.Background(), sdktranslator.FromString("claude"), url, nil, req, req.Payload, rawJSON, rootHeaders) + if errRoot != nil { + t.Fatalf("root cacheHelper error: %v", errRoot) + } + rootBody, errReadRoot := io.ReadAll(rootRequest.Body) + if errReadRoot != nil { + t.Fatalf("read root body: %v", errReadRoot) + } + rootKey := gjson.GetBytes(rootBody, "prompt_cache_key").String() + + childRequest, _, _, errChild := executor.cacheHelper(context.Background(), sdktranslator.FromString("claude"), url, nil, req, req.Payload, rawJSON, childHeaders) + if errChild != nil { + t.Fatalf("child cacheHelper error: %v", errChild) + } + childBody, errReadChild := io.ReadAll(childRequest.Body) + if errReadChild != nil { + t.Fatalf("read child body: %v", errReadChild) + } + childKey := gjson.GetBytes(childBody, "prompt_cache_key").String() + if rootKey == "" || childKey == "" || rootKey == childKey { + t.Fatalf("agent prompt keys are not isolated: root=%q child=%q", rootKey, childKey) + } + + aliasReq := req + aliasReq.Model = "another-local-alias-low" + aliasRequest, _, _, errAlias := executor.cacheHelper(context.Background(), sdktranslator.FromString("claude"), url, nil, aliasReq, aliasReq.Payload, rawJSON, childHeaders) + if errAlias != nil { + t.Fatalf("alias cacheHelper error: %v", errAlias) + } + aliasBody, errReadAlias := io.ReadAll(aliasRequest.Body) + if errReadAlias != nil { + t.Fatalf("read alias body: %v", errReadAlias) + } + if aliasKey := gjson.GetBytes(aliasBody, "prompt_cache_key").String(); aliasKey != childKey { + t.Fatalf("resolved model key fragmented by request alias: first=%q alias=%q", childKey, aliasKey) + } + + websocketBody, _, errWebsocket := applyCodexPromptCacheHeadersWithContext(context.Background(), sdktranslator.FromString("claude"), aliasReq, rawJSON, childHeaders) + if errWebsocket != nil { + t.Fatalf("websocket prompt cache error: %v", errWebsocket) + } + if websocketKey := gjson.GetBytes(websocketBody, "prompt_cache_key").String(); websocketKey != childKey { + t.Fatalf("HTTP/WebSocket prompt keys differ: http=%q websocket=%q", childKey, websocketKey) + } +} diff --git a/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go b/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go index 8c94b146b..cd1b6a785 100644 --- a/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go +++ b/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go @@ -154,8 +154,34 @@ func TestCodexExecutorReasoningReplaySessionKeyUsesClaudeCodeJSONSessionID(t *te body := []byte(`{"model":"gpt-5.4","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"next"}]}]}`) got := codexReasoningReplaySessionKey(context.Background(), from, req, cliproxyexecutor.Options{SourceFormat: from}, body) - if got != "claude:session-json-1" { - t.Fatalf("codexReasoningReplaySessionKey() = %q, want claude:session-json-1", got) + if got != "claude:session-json-1:agent:main" { + t.Fatalf("codexReasoningReplaySessionKey() = %q, want claude:session-json-1:agent:main", got) + } +} + +func TestCodexExecutorReasoningReplaySessionKeyIsolatesClaudeCodeAgents(t *testing.T) { + from := sdktranslator.FromString("claude") + req := cliproxyexecutor.Request{ + Model: "local-alias-high", + Payload: []byte(`{"model":"local-alias","messages":[{"role":"user","content":"next"}]}`), + } + body := []byte(`{"model":"gpt-5.4","prompt_cache_key":"shared-client-key","input":[{"type":"message","role":"user","content":"next"}]}`) + rootHeaders := http.Header{} + rootHeaders.Set("X-Claude-Code-Session-Id", "session-agents") + childAHeaders := rootHeaders.Clone() + childAHeaders.Set("X-Claude-Code-Agent-Id", "agent-a") + childBHeaders := rootHeaders.Clone() + childBHeaders.Set("X-Claude-Code-Agent-Id", "agent-b") + + metadata := map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "shared-execution-session"} + root := codexReasoningReplayScopeFromRequest(context.Background(), from, req, cliproxyexecutor.Options{SourceFormat: from, Headers: rootHeaders, Metadata: metadata}, body) + childA := codexReasoningReplayScopeFromRequest(context.Background(), from, req, cliproxyexecutor.Options{SourceFormat: from, Headers: childAHeaders, Metadata: metadata}, body) + childB := codexReasoningReplayScopeFromRequest(context.Background(), from, req, cliproxyexecutor.Options{SourceFormat: from, Headers: childBHeaders, Metadata: metadata}, body) + if root.modelName != "gpt-5.4" || childA.modelName != "gpt-5.4" || childB.modelName != "gpt-5.4" { + t.Fatalf("replay scopes did not use resolved model: root=%#v a=%#v b=%#v", root, childA, childB) + } + if root.sessionKey == childA.sessionKey || childA.sessionKey == childB.sessionKey || root.sessionKey == childB.sessionKey { + t.Fatalf("agent replay scopes are not isolated: root=%#v a=%#v b=%#v", root, childA, childB) } } @@ -367,7 +393,7 @@ func TestCodexExecutorReasoningReplayCacheDoesNotDuplicateClaudeClientReasoning( cachedEncryptedContent := validCodexReasoningEncryptedContentForTestSeed(5) clientEncryptedContent := validCodexReasoningEncryptedContentForTestSeed(6) - internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-2", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) + internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-2:agent:main", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -418,7 +444,7 @@ func TestCodexExecutorReasoningReplayCacheInsertsReasoningBeforeAssistantOutputI t.Cleanup(internalcache.ClearCodexReasoningReplayCache) cachedEncryptedContent := validCodexReasoningEncryptedContentForTestSeed(7) - internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-history", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) + internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-history:agent:main", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -546,7 +572,7 @@ func TestCodexExecutorReasoningReplayCacheClearsOnNonStreamResponseFailedInvalid t.Cleanup(internalcache.ClearCodexReasoningReplayCache) cachedEncryptedContent := validCodexReasoningEncryptedContentForTestSeed(9) - internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-nonstream", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) + internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-nonstream:agent:main", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.ReadAll(r.Body) @@ -572,7 +598,7 @@ func TestCodexExecutorReasoningReplayCacheClearsOnNonStreamResponseFailedInvalid if err == nil { t.Fatal("expected invalid signature error") } - if _, ok := internalcache.GetCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-nonstream"); ok { + if _, ok := internalcache.GetCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-nonstream:agent:main"); ok { t.Fatal("invalid signature response.failed should clear cached replay item") } } @@ -582,7 +608,7 @@ func TestCodexExecutorReasoningReplayCacheClearsOnStreamResponseFailedInvalidSig t.Cleanup(internalcache.ClearCodexReasoningReplayCache) cachedEncryptedContent := validCodexReasoningEncryptedContentForTestSeed(10) - internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-stream", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) + internalcache.CacheCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-stream:agent:main", []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+cachedEncryptedContent+`"}`)) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.ReadAll(r.Body) @@ -618,7 +644,7 @@ func TestCodexExecutorReasoningReplayCacheClearsOnStreamResponseFailedInvalidSig if !gotChunkErr { t.Fatal("expected stream chunk error for invalid signature response.failed") } - if _, ok := internalcache.GetCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-stream"); ok { + if _, ok := internalcache.GetCodexReasoningReplayItem("gpt-5.4", "claude:session-invalid-stream:agent:main"); ok { t.Fatal("invalid signature response.failed should clear cached replay item") } } @@ -710,6 +736,224 @@ func TestCodexExecutorReasoningReplayCacheReplaysFunctionCallForClaudeToolResult } } +func TestCodexExecutorReasoningReplayCacheRestoresCumulativeToolTurns(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + scope := codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:session-cumulative-tools:agent:main", + } + firstEncrypted := validCodexReasoningEncryptedContentForTestSeed(21) + secondEncrypted := validCodexReasoningEncryptedContentForTestSeed(22) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+firstEncrypted+`"},`+ + `{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"first\"}"}`+ + `]}}`)) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+secondEncrypted+`"},`+ + `{"type":"function_call","call_id":"call_2","name":"lookup","arguments":"{\"q\":\"second\"}"}`+ + `]}}`)) + + body := []byte(`{"model":"gpt-5.4","input":[` + + `{"type":"message","role":"user","content":"first"},` + + `{"type":"function_call_output","call_id":"call_1","output":"one"},` + + `{"type":"message","role":"user","content":"second"},` + + `{"type":"function_call_output","call_id":"call_2","output":"two"},` + + `{"type":"message","role":"user","content":"third"}` + + `]}`) + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"metadata":{"user_id":"{\"session_id\":\"session-cumulative-tools\"}"}}`), + } + updated, gotScope := applyCodexReasoningReplayCache(context.Background(), sdktranslator.FromString("claude"), req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body) + if gotScope.modelName != scope.modelName || gotScope.sessionKey != scope.sessionKey { + t.Fatalf("replay scope = %#v, want model/session %#v", gotScope, scope) + } + wantTypes := []string{"message", "reasoning", "function_call", "function_call_output", "message", "reasoning", "function_call", "function_call_output", "message"} + gotItems := gjson.GetBytes(updated, "input").Array() + if len(gotItems) != len(wantTypes) { + t.Fatalf("input length = %d, want %d; body=%s", len(gotItems), len(wantTypes), updated) + } + for index, wantType := range wantTypes { + if gotType := gotItems[index].Get("type").String(); gotType != wantType { + t.Fatalf("input.%d.type = %q, want %q; body=%s", index, gotType, wantType, updated) + } + } + if gotItems[1].Get("encrypted_content").String() != firstEncrypted || gotItems[5].Get("encrypted_content").String() != secondEncrypted { + t.Fatalf("cumulative reasoning was not restored in turn order: %s", updated) + } +} + +func TestCodexExecutorReasoningReplayCacheRestoresCumulativeAssistantTurns(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + scope := codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:session-cumulative-messages:agent:main", + } + firstEncrypted := validCodexReasoningEncryptedContentForTestSeed(23) + secondEncrypted := validCodexReasoningEncryptedContentForTestSeed(24) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+firstEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]}`+ + `]}}`)) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+secondEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"second answer"}]}`+ + `]}}`)) + + body := []byte(`{"model":"gpt-5.4","input":[` + + `{"type":"message","role":"user","content":"first"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"first answer"}]},` + + `{"type":"message","role":"user","content":"second"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"second answer"}]},` + + `{"type":"message","role":"user","content":"third"}` + + `]}`) + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"metadata":{"user_id":"{\"session_id\":\"session-cumulative-messages\"}"}}`), + } + updated, gotScope := applyCodexReasoningReplayCache(context.Background(), sdktranslator.FromString("claude"), req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body) + if gotScope.modelName != scope.modelName || gotScope.sessionKey != scope.sessionKey { + t.Fatalf("replay scope = %#v, want model/session %#v", gotScope, scope) + } + wantTypes := []string{"message", "reasoning", "message", "message", "reasoning", "message", "message"} + gotItems := gjson.GetBytes(updated, "input").Array() + if len(gotItems) != len(wantTypes) { + t.Fatalf("input length = %d, want %d; body=%s", len(gotItems), len(wantTypes), updated) + } + for index, wantType := range wantTypes { + if gotType := gotItems[index].Get("type").String(); gotType != wantType { + t.Fatalf("input.%d.type = %q, want %q; body=%s", index, gotType, wantType, updated) + } + } + if gotItems[1].Get("encrypted_content").String() != firstEncrypted || gotItems[4].Get("encrypted_content").String() != secondEncrypted { + t.Fatalf("assistant reasoning was not restored at its original turns: %s", updated) + } +} + +func TestCodexExecutorReasoningReplayCacheSkipsDetachedTurnAfterCompaction(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + scope := codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:session-compacted:agent:main", + } + detachedEncrypted := validCodexReasoningEncryptedContentForTestSeed(25) + retainedEncrypted := validCodexReasoningEncryptedContentForTestSeed(26) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+detachedEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"removed answer"}]}`+ + `]}}`)) + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+retainedEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"retained answer"}]}`+ + `]}}`)) + + body := []byte(`{"model":"gpt-5.4","input":[` + + `{"type":"message","role":"user","content":"compacted summary"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"retained answer"}]},` + + `{"type":"message","role":"user","content":"continue"}` + + `]}`) + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"metadata":{"user_id":"{\"session_id\":\"session-compacted\"}"}}`), + } + updated, _ := applyCodexReasoningReplayCache(context.Background(), sdktranslator.FromString("claude"), req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body) + gotItems := gjson.GetBytes(updated, "input").Array() + if len(gotItems) != 4 || gotItems[1].Get("encrypted_content").String() != retainedEncrypted { + t.Fatalf("retained turn reasoning was not restored: %s", updated) + } + for _, item := range gotItems { + if item.Get("encrypted_content").String() == detachedEncrypted { + t.Fatalf("detached reasoning moved into compacted history: %s", updated) + } + } +} + +func TestCodexExecutorReasoningReplayCacheMatchesNewestDuplicateAssistantAfterCompaction(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + scope := codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:session-duplicate-compaction:agent:main", + } + oldEncrypted := validCodexReasoningEncryptedContentForTestSeed(27) + newEncrypted := validCodexReasoningEncryptedContentForTestSeed(28) + for _, encryptedContent := range []string{oldEncrypted, newEncrypted} { + cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]}`+ + `]}}`)) + } + + body := []byte(`{"model":"gpt-5.4","input":[` + + `{"type":"message","role":"user","content":"compacted summary"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]},` + + `{"type":"message","role":"user","content":"continue"}` + + `]}`) + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"metadata":{"user_id":"{\"session_id\":\"session-duplicate-compaction\"}"}}`), + } + updated, _ := applyCodexReasoningReplayCache(context.Background(), sdktranslator.FromString("claude"), req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body) + gotItems := gjson.GetBytes(updated, "input").Array() + if len(gotItems) != 4 || gotItems[1].Get("encrypted_content").String() != newEncrypted { + t.Fatalf("newest duplicate assistant turn was not retained: %s", updated) + } + for _, item := range gotItems { + if item.Get("encrypted_content").String() == oldEncrypted { + t.Fatalf("detached duplicate assistant reasoning was restored: %s", updated) + } + } +} + +func TestCodexExecutorReasoningReplayCacheUsesRequestPrefixForDuplicateOutOfOrderTurns(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + body := []byte(`{"model":"gpt-5.4","input":[` + + `{"type":"message","role":"user","content":"first"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]},` + + `{"type":"message","role":"user","content":"second"},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]},` + + `{"type":"message","role":"user","content":"third"}` + + `]}`) + inputItems := gjson.GetBytes(body, "input").Array() + baseScope := codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:session-duplicate-prefix:agent:main", + } + oldEncrypted := validCodexReasoningEncryptedContentForTestSeed(29) + newEncrypted := validCodexReasoningEncryptedContentForTestSeed(30) + newScope := baseScope + newScope.requestFingerprint = codexReplayInputPrefixFingerprint(inputItems, 3) + cacheCodexReasoningReplayFromCompleted(newScope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+newEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]}`+ + `]}}`)) + oldScope := baseScope + oldScope.requestFingerprint = codexReplayInputPrefixFingerprint(inputItems, 1) + cacheCodexReasoningReplayFromCompleted(oldScope, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+oldEncrypted+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Done"}]}`+ + `]}}`)) + + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"metadata":{"user_id":"{\"session_id\":\"session-duplicate-prefix\"}"}}`), + } + updated, _ := applyCodexReasoningReplayCache(context.Background(), sdktranslator.FromString("claude"), req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body) + gotItems := gjson.GetBytes(updated, "input").Array() + if len(gotItems) != 7 || gotItems[1].Get("encrypted_content").String() != oldEncrypted || gotItems[4].Get("encrypted_content").String() != newEncrypted { + t.Fatalf("duplicate out-of-order turns were not matched by request prefix: %s", updated) + } +} + func TestCodexExecutorReasoningReplayCacheDropsFunctionCallWithoutMatchingOutput(t *testing.T) { internalcache.ClearCodexReasoningReplayCache() t.Cleanup(internalcache.ClearCodexReasoningReplayCache) @@ -717,7 +961,7 @@ func TestCodexExecutorReasoningReplayCacheDropsFunctionCallWithoutMatchingOutput encryptedContent := validCodexReasoningEncryptedContentForTestSeed(14) scope := codexReasoningReplayScope{ modelName: "gpt-5.4", - sessionKey: "claude:session-dropped-tool", + sessionKey: "claude:session-dropped-tool:agent:main", } cacheCodexReasoningReplayFromCompleted(scope, []byte(`{"response":{"output":[`+ `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"},`+ @@ -741,21 +985,18 @@ func TestCodexExecutorReasoningReplayCacheDropsFunctionCallWithoutMatchingOutput cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}, body, ) - if replayScope != scope { - t.Fatalf("replay scope = %#v, want %#v", replayScope, scope) + if replayScope.modelName != scope.modelName || replayScope.sessionKey != scope.sessionKey { + t.Fatalf("replay scope = %#v, want model/session %#v", replayScope, scope) } - if got := gjson.GetBytes(updated, "input.0.type").String(); got != "reasoning" { - t.Fatalf("input.0.type = %q, want reasoning; body=%s", got, string(updated)) + if got := gjson.GetBytes(updated, "input.0.role").String(); got != "user" { + t.Fatalf("input.0.role = %q, want detached turn to be skipped; body=%s", got, string(updated)) } - if got := gjson.GetBytes(updated, "input.0.encrypted_content").String(); got != encryptedContent { - t.Fatalf("input.0.encrypted_content = %q, want cached reasoning; body=%s", got, string(updated)) + if gjson.GetBytes(updated, `input.#(type=="reasoning")`).Exists() { + t.Fatalf("detached turn reasoning should not move to the front; body=%s", string(updated)) } if gjson.GetBytes(updated, `input.#(call_id=="call_dropped")`).Exists() { t.Fatalf("cached function_call without matching output should not be replayed; body=%s", string(updated)) } - if got := gjson.GetBytes(updated, "input.1.role").String(); got != "user" { - t.Fatalf("input.1.role = %q, want user; body=%s", got, string(updated)) - } } func TestCodexExecutorReasoningReplayCacheMatchesShortenedClaudeToolResultCallID(t *testing.T) { diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index af0e415c2..64eff7452 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -275,6 +275,10 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return resp, errReplay + } httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildCodexResponsesWebsocketURL(httpURL) @@ -282,7 +286,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut return resp, err } - body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body) + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) if errPromptCache != nil { return resp, errPromptCache } @@ -405,6 +409,8 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } } + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte for { if ctx != nil && ctx.Err() != nil { return resp, ctx.Err() @@ -439,13 +445,27 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if sess != nil { e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) } + if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { + return resp, errClearReplay + } helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) return resp, wsErr } + if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + return resp, errClearReplay + } + return resp, streamErr + } payload = normalizeCodexWebsocketCompletion(payload) eventType := gjson.GetBytes(payload, "type").String() - if eventType == "response.completed" { + switch eventType { + case "response.output_item.done": + collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) + case "response.completed": + payload = patchCodexCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) + cacheCodexReasoningReplayFromCompleted(replayScope, payload) if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) } @@ -479,11 +499,12 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr from := opts.SourceFormat responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) to := sdktranslator.FromString("codex") - body := req.Payload - userPayload := req.Payload + originalPayloadSource := req.Payload if len(opts.OriginalRequest) > 0 { - userPayload = opts.OriginalRequest + originalPayloadSource = opts.OriginalRequest } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { @@ -492,13 +513,17 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, body, requestedModel, requestPath, opts.Headers) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) body, _ = sjson.SetBytes(body, "model", baseModel) body = normalizeCodexInstructions(body) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return nil, errReplay + } httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildCodexResponsesWebsocketURL(httpURL) @@ -506,13 +531,13 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return nil, err } - body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body) + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) if errPromptCache != nil { return nil, errPromptCache } clientBody := body var identityState codexIdentityConfuseState - upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, userPayload, body) + upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) reporter.SetTranslatedReasoningEffort(clientBody, to.String()) wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) applyModelHeaderOverrides(wsHeaders, baseModel) @@ -659,6 +684,8 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte for { if ctx != nil && ctx.Err() != nil { terminateReason = "context_done" @@ -709,24 +736,54 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if wsErr, ok := parseCodexWebsocketError(payload); ok { terminateReason = "upstream_error" terminateErr = wsErr - helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) - reporter.PublishFailure(ctx, wsErr) if sess != nil { e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) } + if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { + terminateErr = errClearReplay + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + return + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) + reporter.PublishFailure(ctx, wsErr) _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) return } + if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + terminateReason = "upstream_error" + terminateErr = streamErr + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + terminateErr = errClearReplay + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + return + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr) + reporter.PublishFailure(ctx, streamErr) + _ = send(cliproxyexecutor.StreamChunk{Err: streamErr}) + return + } eventType := gjson.GetBytes(payload, "type").String() isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" + if eventType == "response.output_item.done" { + collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) + } + completedPayload := payload + if eventType == "response.completed" || eventType == "response.done" { + completedPayload = normalizeCodexWebsocketCompletion(completedPayload) + completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback) + cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) + if detail, ok := helps.ParseCodexUsage(completedPayload); ok { + reporter.Publish(ctx, detail) + } + } + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) if cliproxyexecutor.DownstreamWebsocket(ctx) { - if eventType == "response.completed" || eventType == "response.done" { - if detail, ok := helps.ParseCodexUsage(payload); ok { - reporter.Publish(ctx, detail) - } - } if !send(cliproxyexecutor.StreamChunk{Payload: clientPayload}) { terminateReason = "context_done" terminateErr = ctx.Err() @@ -739,16 +796,13 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } payload = normalizeCodexWebsocketCompletion(payload) - eventType = gjson.GetBytes(payload, "type").String() if eventType == "response.completed" || eventType == "response.done" { - if detail, ok := helps.ParseCodexUsage(payload); ok { - reporter.Publish(ctx, detail) - } + payload = completedPayload } - + eventType = gjson.GetBytes(payload, "type").String() clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState) line := encodeCodexWebsocketAsSSE(clientPayload) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, clientBody, clientBody, line, ¶m) + chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m) for i := range chunks { if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { terminateReason = "context_done" @@ -941,15 +995,23 @@ func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecuto return body, headers } -func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header, error) { +func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte, headerSets ...http.Header) ([]byte, http.Header, error) { headers := http.Header{} if len(rawJSON) == 0 { return rawJSON, headers, nil } + var requestHeaders http.Header + if len(headerSets) > 0 { + requestHeaders = headerSets[0] + } var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, nil) + modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, requestHeaders) if errCache != nil { return nil, nil, errCache } @@ -1271,6 +1333,17 @@ func parseCodexWebsocketError(payload []byte) (error, bool) { }, true } +func clearCodexReasoningReplayOnWebsocketError(ctx context.Context, scope codexReasoningReplayScope, payload []byte) error { + status := int(gjson.GetBytes(payload, "status").Int()) + if status == 0 { + status = int(gjson.GetBytes(payload, "status_code").Int()) + } + if status <= 0 { + return nil + } + return clearCodexReasoningReplayOnInvalidSignature(ctx, scope, status, buildCodexWebsocketErrorPayload(payload, status)) +} + func buildCodexWebsocketErrorPayload(payload []byte, status int) []byte { out := []byte(`{}`) out, _ = sjson.SetBytes(out, "status", status) diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 42ab3b59a..2b3c10dd5 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -12,6 +12,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -73,6 +74,97 @@ func TestBuildCodexWebsocketRequestBodyShortensOverlongInputItemIDs(t *testing.T } } +func TestCodexWebsocketsExecuteRestoresClaudeAgentReasoningReplay(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + encryptedContent := validCodexReasoningEncryptedContentForTestSeed(31) + cacheCodexReasoningReplayFromCompleted(codexReasoningReplayScope{ + modelName: "gpt-5.4", + sessionKey: "claude:ws-replay-session:agent:agent-a", + }, []byte(`{"response":{"output":[`+ + `{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"},`+ + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"previous answer"}]}`+ + `]}}`)) + + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + capturedPayload := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Fatalf("upgrade websocket: %v", errUpgrade) + } + defer func() { _ = conn.Close() }() + + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read upstream websocket message: %v", errRead) + } + capturedPayload <- bytes.Clone(payload) + completed := []byte(`{"type":"response.completed","response":{"id":"resp-ws-replay","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"next answer"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Fatalf("write completed websocket message: %v", errWrite) + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{ + "model":"gpt-5.4", + "messages":[ + {"role":"user","content":"first"}, + {"role":"assistant","content":"previous answer"}, + {"role":"user","content":"next"} + ] + }`), + } + headers := http.Header{} + headers.Set("X-Claude-Code-Session-Id", "ws-replay-session") + headers.Set("X-Claude-Code-Agent-Id", "agent-a") + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude"), Headers: headers} + + if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + + select { + case payload := <-capturedPayload: + input := gjson.GetBytes(payload, "input").Array() + if len(input) != 4 { + t.Fatalf("upstream input length = %d, want 4; payload=%s", len(input), payload) + } + if input[1].Get("type").String() != "reasoning" || input[1].Get("encrypted_content").String() != encryptedContent { + t.Fatalf("websocket reasoning replay missing before assistant message: %s", payload) + } + if input[2].Get("role").String() != "assistant" { + t.Fatalf("input.2.role = %q, want assistant; payload=%s", input[2].Get("role").String(), payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for upstream websocket payload") + } +} + +func TestClearCodexReasoningReplayOnWebsocketInvalidSignature(t *testing.T) { + internalcache.ClearCodexReasoningReplayCache() + t.Cleanup(internalcache.ClearCodexReasoningReplayCache) + + scope := codexReasoningReplayScope{modelName: "gpt-5.4", sessionKey: "claude:ws-invalid:agent:main"} + encryptedContent := validCodexReasoningEncryptedContentForTestSeed(32) + if !internalcache.CacheCodexReasoningReplayItem(scope.modelName, scope.sessionKey, []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"`+encryptedContent+`"}`)) { + t.Fatal("failed to seed websocket replay cache") + } + payload := []byte(`{"type":"error","status":400,"body":{"error":{"message":"Invalid signature in thinking block","type":"invalid_request_error","code":"invalid_request_error"}}}`) + if errClear := clearCodexReasoningReplayOnWebsocketError(context.Background(), scope, payload); errClear != nil { + t.Fatalf("clear websocket replay error: %v", errClear) + } + if _, ok := internalcache.GetCodexReasoningReplayItem(scope.modelName, scope.sessionKey); ok { + t.Fatal("websocket invalid signature did not clear replay state") + } +} + func TestCodexWebsocketsExecuteResponsesLiteDoesNotInjectImageGenerationTool(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) diff --git a/internal/runtime/executor/helps/claude_code_session.go b/internal/runtime/executor/helps/claude_code_session.go index cd986302d..d63584692 100644 --- a/internal/runtime/executor/helps/claude_code_session.go +++ b/internal/runtime/executor/helps/claude_code_session.go @@ -5,32 +5,75 @@ import ( "net/http" "regexp" "strings" - "time" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/tidwall/gjson" ) -const ClaudeCodeSessionHeader = "X-Claude-Code-Session-Id" +const ( + ClaudeCodeSessionHeader = "X-Claude-Code-Session-Id" + ClaudeCodeAgentHeader = "X-Claude-Code-Agent-Id" + ClaudeCodeMainAgentID = "main" +) var claudeCodeSessionSuffixPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`) // ExtractClaudeCodeSessionID resolves a Claude Code session ID, preferring X-Claude-Code-Session-Id over payload metadata. func ExtractClaudeCodeSessionID(ctx context.Context, payload []byte, headers http.Header) string { - if headers != nil { - if sessionID := strings.TrimSpace(headers.Get(ClaudeCodeSessionHeader)); sessionID != "" { - return sessionID - } + if sessionID := claudeCodeHeader(ctx, headers, ClaudeCodeSessionHeader); sessionID != "" { + return sessionID + } + return extractClaudeCodeSessionIDFromPayload(payload) +} + +// ExtractClaudeCodeAgentID resolves the Claude Code agent ID and uses a stable sentinel for the root agent. +func ExtractClaudeCodeAgentID(ctx context.Context, headers http.Header) string { + if agentID := claudeCodeHeader(ctx, headers, ClaudeCodeAgentHeader); agentID != "" { + return agentID + } + return ClaudeCodeMainAgentID +} + +// ClaudeCodeExecutionScope returns the stable root-session and agent identity used by Codex execution state. +func ClaudeCodeExecutionScope(ctx context.Context, payload []byte, headers http.Header) (string, bool) { + sessionID := ExtractClaudeCodeSessionID(ctx, payload, headers) + if sessionID == "" { + return "", false + } + return "claude:" + sessionID + ":agent:" + ExtractClaudeCodeAgentID(ctx, headers), true +} + +func claudeCodeHeader(ctx context.Context, headers http.Header, name string) string { + if value := headerValueCaseInsensitive(headers, name); value != "" { + return value } if ctx != nil { if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - if sessionID := strings.TrimSpace(ginCtx.Request.Header.Get(ClaudeCodeSessionHeader)); sessionID != "" { - return sessionID + return headerValueCaseInsensitive(ginCtx.Request.Header, name) + } + } + return "" +} + +func headerValueCaseInsensitive(headers http.Header, name string) string { + if headers == nil { + return "" + } + if value := strings.TrimSpace(headers.Get(name)); value != "" { + return value + } + for key, values := range headers { + if !strings.EqualFold(key, name) { + continue + } + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value } } } - return extractClaudeCodeSessionIDFromPayload(payload) + return "" } func extractClaudeCodeSessionIDFromPayload(payload []byte) string { @@ -50,22 +93,13 @@ func extractClaudeCodeSessionIDFromPayload(payload []byte) string { return "" } -// ClaudeCodePromptCache maps a Claude Code session to a stable upstream prompt_cache_key. +// ClaudeCodePromptCache derives a deterministic upstream prompt_cache_key for one Claude Code agent. func ClaudeCodePromptCache(ctx context.Context, modelName string, payload []byte, headers http.Header) (CodexCache, bool, error) { - sessionID := ExtractClaudeCodeSessionID(ctx, payload, headers) - if sessionID == "" { + modelName = strings.TrimSpace(modelName) + executionScope, ok := ClaudeCodeExecutionScope(ctx, payload, headers) + if modelName == "" || !ok { return CodexCache{}, false, nil } - key := CodexPromptCacheKey(modelName, "claude:"+sessionID) - if cache, ok, errCache := GetCodexCacheRequired(ctx, key); errCache != nil || ok { - return cache, ok, errCache - } - cache := CodexCache{ - ID: uuid.New().String(), - Expire: time.Now().Add(1 * time.Hour), - } - if errSet := SetCodexCacheRequired(ctx, key, cache); errSet != nil { - return CodexCache{}, false, errSet - } - return cache, true, nil + identity := strings.Join([]string{"cli-proxy-api:codex:claude-code", modelName, executionScope}, "\x00") + return CodexCache{ID: uuid.NewSHA1(uuid.NameSpaceOID, []byte(identity)).String()}, true, nil } diff --git a/internal/runtime/executor/helps/claude_code_session_test.go b/internal/runtime/executor/helps/claude_code_session_test.go index 4d1b76569..df6e48d59 100644 --- a/internal/runtime/executor/helps/claude_code_session_test.go +++ b/internal/runtime/executor/helps/claude_code_session_test.go @@ -59,3 +59,64 @@ func TestExtractClaudeCodeSessionIDPrefersHeaderOverPayload(t *testing.T) { t.Fatalf("ExtractClaudeCodeSessionID() = %q, want header-session", got) } } + +func TestClaudeCodeExecutionScopeAcceptsLowercaseHeaderMapKeys(t *testing.T) { + headers := http.Header{ + "x-claude-code-session-id": []string{"lower-session"}, + "x-claude-code-agent-id": []string{"lower-agent"}, + } + + scope, ok := ClaudeCodeExecutionScope(context.Background(), nil, headers) + if !ok || scope != "claude:lower-session:agent:lower-agent" { + t.Fatalf("lowercase header scope = %q, %v", scope, ok) + } +} + +func TestClaudeCodeExecutionScopeIsolatesAgents(t *testing.T) { + rootHeaders := http.Header{} + rootHeaders.Set(ClaudeCodeSessionHeader, "session-agents") + childAHeaders := rootHeaders.Clone() + childAHeaders.Set(ClaudeCodeAgentHeader, "agent-a") + childBHeaders := rootHeaders.Clone() + childBHeaders.Set(ClaudeCodeAgentHeader, "agent-b") + + rootScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, rootHeaders) + if !ok || rootScope != "claude:session-agents:agent:main" { + t.Fatalf("root scope = %q, %v", rootScope, ok) + } + childAScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, childAHeaders) + if !ok || childAScope != "claude:session-agents:agent:agent-a" { + t.Fatalf("child A scope = %q, %v", childAScope, ok) + } + childBScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, childBHeaders) + if !ok || childBScope != "claude:session-agents:agent:agent-b" { + t.Fatalf("child B scope = %q, %v", childBScope, ok) + } + if rootScope == childAScope || childAScope == childBScope || rootScope == childBScope { + t.Fatalf("agent scopes are not isolated: root=%q a=%q b=%q", rootScope, childAScope, childBScope) + } +} + +func TestClaudeCodePromptCacheDeterministicAndAgentScoped(t *testing.T) { + rootHeaders := http.Header{} + rootHeaders.Set(ClaudeCodeSessionHeader, "session-cache-agents") + childHeaders := rootHeaders.Clone() + childHeaders.Set(ClaudeCodeAgentHeader, "agent-a") + + rootFirst, ok, errFirst := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, rootHeaders) + if errFirst != nil || !ok { + t.Fatalf("root first cache = %#v, %v, %v", rootFirst, ok, errFirst) + } + rootSecond, ok, errSecond := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, rootHeaders) + if errSecond != nil || !ok || rootSecond.ID != rootFirst.ID { + t.Fatalf("root second cache = %#v, %v, %v; want ID %q", rootSecond, ok, errSecond, rootFirst.ID) + } + child, ok, errChild := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, childHeaders) + if errChild != nil || !ok || child.ID == rootFirst.ID { + t.Fatalf("child cache = %#v, %v, %v; root ID %q", child, ok, errChild, rootFirst.ID) + } + otherModel, ok, errModel := ClaudeCodePromptCache(context.Background(), "gpt-5.5", nil, rootHeaders) + if errModel != nil || !ok || otherModel.ID == rootFirst.ID { + t.Fatalf("other model cache = %#v, %v, %v; root ID %q", otherModel, ok, errModel, rootFirst.ID) + } +} diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 73d265c9e..d2fbc9a49 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -1150,7 +1150,7 @@ func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Reque if !xaiRequiresIsolatedConversation(baseModel) { return "", nil } - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, opts.Headers) + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, baseModel, req.Payload, opts.Headers) if errCache != nil { return "", errCache }