diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 7b69f67d7..16f3d2320 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -9,7 +9,6 @@ import ( "fmt" "io" "net/http" - "regexp" "sort" "strings" "time" @@ -41,7 +40,6 @@ const ( ) var dataTag = []byte("data:") -var codexClaudeCodeSessionPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`) // Streamed Codex responses may emit response.output_item.done events while leaving // response.completed.response.output empty. Keep the stream path aligned with the @@ -294,57 +292,14 @@ func sourceFormatEqual(from, want sdktranslator.Format) bool { return strings.EqualFold(strings.TrimSpace(from.String()), want.String()) } -func codexClaudeCodeReplaySessionKey(payload []byte) string { - sessionID := extractClaudeCodeSessionIDForCodexReplay(payload) +func codexClaudeCodeReplaySessionKey(ctx context.Context, payload []byte, headers http.Header) string { + sessionID := helps.ExtractClaudeCodeSessionID(ctx, payload, headers) if sessionID == "" { return "" } return "claude:" + sessionID } -func codexClaudeCodePromptCacheStorageKey(req cliproxyexecutor.Request) string { - sessionID := extractClaudeCodeSessionIDForCodexReplay(req.Payload) - if sessionID == "" { - return "" - } - return helps.CodexPromptCacheKey(req.Model, "claude:"+sessionID) -} - -func codexClaudeCodePromptCache(ctx context.Context, req cliproxyexecutor.Request) (helps.CodexCache, bool, error) { - key := codexClaudeCodePromptCacheStorageKey(req) - if key == "" { - return helps.CodexCache{}, false, nil - } - if cache, ok, errCache := helps.GetCodexCacheRequired(ctx, key); errCache != nil || ok { - return cache, ok, errCache - } - cache := helps.CodexCache{ - ID: uuid.New().String(), - Expire: time.Now().Add(1 * time.Hour), - } - if errSet := helps.SetCodexCacheRequired(ctx, key, cache); errSet != nil { - return helps.CodexCache{}, false, errSet - } - return cache, true, nil -} - -func extractClaudeCodeSessionIDForCodexReplay(payload []byte) string { - if len(payload) == 0 { - return "" - } - userID := gjson.GetBytes(payload, "metadata.user_id").String() - if userID == "" { - return "" - } - if matches := codexClaudeCodeSessionPattern.FindStringSubmatch(userID); len(matches) >= 2 { - return matches[1] - } - if len(userID) > 0 && userID[0] == '{' { - return gjson.Get(userID, "session_id").String() - } - return "" -} - func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) string { if ctx == nil { ctx = context.Background() @@ -370,7 +325,7 @@ func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Form } } if sourceFormatEqual(from, sdktranslator.FormatClaude) { - return codexClaudeCodeReplaySessionKey(req.Payload) + return codexClaudeCodeReplaySessionKey(ctx, req.Payload, opts.Headers) } if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { @@ -1464,7 +1419,7 @@ type codexIdentityReplacement struct { 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) { var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - cached, ok, errCache := codexClaudeCodePromptCache(ctx, req) + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, nil) 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 d33d7fc64..8e28340f4 100644 --- a/internal/runtime/executor/codex_executor_cache_test.go +++ b/internal/runtime/executor/codex_executor_cache_test.go @@ -3,12 +3,14 @@ package executor import ( "context" "io" + "net/http" "net/http/httptest" "testing" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" @@ -259,3 +261,49 @@ func TestCodexIdentityConfuseKeepsClientBodySeparateFromUpstreamBody(t *testing. t.Fatalf("client prompt_cache_key = %q, want cache-1", gotKey) } } + +func TestCodexExecutorCacheHelper_ClaudeUsesSessionHeader(t *testing.T) { + executor := &CodexExecutor{} + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + ginCtx.Request.Header.Set(helps.ClaudeCodeSessionHeader, "cache-session-header") + ctx := context.WithValue(context.Background(), "gin", ginCtx) + + firstReq := cliproxyexecutor.Request{ + Model: "gpt-5.4-claude-cache-header", + Payload: []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":[{"type":"text","text":"first"}]}]}`), + } + secondReq := cliproxyexecutor.Request{ + Model: "gpt-5.4-claude-cache-header", + Payload: []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":[{"type":"text","text":"next"}]}]}`), + } + rawJSON := []byte(`{"model":"gpt-5.4","stream":true}`) + url := "https://example.com/responses" + + firstHTTPReq, _, _, err := executor.cacheHelper(ctx, sdktranslator.FromString("claude"), url, nil, firstReq, firstReq.Payload, rawJSON) + if err != nil { + t.Fatalf("cacheHelper first error: %v", err) + } + secondHTTPReq, _, _, err := executor.cacheHelper(ctx, sdktranslator.FromString("claude"), url, nil, secondReq, secondReq.Payload, rawJSON) + if err != nil { + t.Fatalf("cacheHelper second error: %v", err) + } + + firstBody, errRead := io.ReadAll(firstHTTPReq.Body) + if errRead != nil { + t.Fatalf("read first request body: %v", errRead) + } + secondBody, errRead := io.ReadAll(secondHTTPReq.Body) + if errRead != nil { + t.Fatalf("read second request body: %v", errRead) + } + firstKey := gjson.GetBytes(firstBody, "prompt_cache_key").String() + secondKey := gjson.GetBytes(secondBody, "prompt_cache_key").String() + if firstKey == "" { + t.Fatalf("first prompt_cache_key is empty; body=%s", string(firstBody)) + } + if secondKey != firstKey { + t.Fatalf("same Claude Code session header produced different prompt_cache_key: first=%q second=%q", firstKey, secondKey) + } +} diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index d96abfee1..7e74d953d 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -870,7 +870,7 @@ func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktransl var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - cached, ok, errCache := codexClaudeCodePromptCache(ctx, req) + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, nil) if errCache != nil { return nil, nil, errCache } diff --git a/internal/runtime/executor/helps/claude_code_session.go b/internal/runtime/executor/helps/claude_code_session.go new file mode 100644 index 000000000..cd986302d --- /dev/null +++ b/internal/runtime/executor/helps/claude_code_session.go @@ -0,0 +1,71 @@ +package helps + +import ( + "context" + "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" + +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 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 extractClaudeCodeSessionIDFromPayload(payload) +} + +func extractClaudeCodeSessionIDFromPayload(payload []byte) string { + if len(payload) == 0 { + return "" + } + userID := gjson.GetBytes(payload, "metadata.user_id").String() + if userID == "" { + return "" + } + if matches := claudeCodeSessionSuffixPattern.FindStringSubmatch(userID); len(matches) >= 2 { + return matches[1] + } + if len(userID) > 0 && userID[0] == '{' { + return strings.TrimSpace(gjson.Get(userID, "session_id").String()) + } + return "" +} + +// ClaudeCodePromptCache maps a Claude Code session to a stable upstream prompt_cache_key. +func ClaudeCodePromptCache(ctx context.Context, modelName string, payload []byte, headers http.Header) (CodexCache, bool, error) { + sessionID := ExtractClaudeCodeSessionID(ctx, payload, headers) + if sessionID == "" { + 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 +} diff --git a/internal/runtime/executor/helps/claude_code_session_test.go b/internal/runtime/executor/helps/claude_code_session_test.go new file mode 100644 index 000000000..4d1b76569 --- /dev/null +++ b/internal/runtime/executor/helps/claude_code_session_test.go @@ -0,0 +1,61 @@ +package helps + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestExtractClaudeCodeSessionIDFromPayloadJSON(t *testing.T) { + payload := []byte(`{"metadata":{"user_id":"{\"device_id\":\"d\",\"session_id\":\"cache-session-1\"}"}}`) + got := ExtractClaudeCodeSessionID(context.Background(), payload, nil) + if got != "cache-session-1" { + t.Fatalf("ExtractClaudeCodeSessionID() = %q, want cache-session-1", got) + } +} + +func TestExtractClaudeCodeSessionIDFromHeader(t *testing.T) { + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + ginCtx.Request.Header.Set(ClaudeCodeSessionHeader, "header-session-1") + ctx := context.WithValue(context.Background(), "gin", ginCtx) + + got := ExtractClaudeCodeSessionID(ctx, []byte(`{"model":"gpt-5.4"}`), nil) + if got != "header-session-1" { + t.Fatalf("ExtractClaudeCodeSessionID() = %q, want header-session-1", got) + } +} + +func TestClaudeCodePromptCacheStableAcrossRequests(t *testing.T) { + ctx := context.Background() + payload := []byte(`{"metadata":{"user_id":"{\"session_id\":\"cache-session-2\"}"}}`) + first, ok, err := ClaudeCodePromptCache(ctx, "grok-composer-2.5-fast", payload, nil) + if err != nil { + t.Fatalf("ClaudeCodePromptCache first error: %v", err) + } + if !ok || first.ID == "" { + t.Fatalf("ClaudeCodePromptCache first = %#v, ok=%v, want cached id", first, ok) + } + second, ok, err := ClaudeCodePromptCache(ctx, "grok-composer-2.5-fast", payload, nil) + if err != nil { + t.Fatalf("ClaudeCodePromptCache second error: %v", err) + } + if !ok || second.ID != first.ID { + t.Fatalf("second cache id = %q, want %q", second.ID, first.ID) + } +} + +func TestExtractClaudeCodeSessionIDPrefersHeaderOverPayload(t *testing.T) { + payload := []byte(`{"metadata":{"user_id":"{"session_id":"payload-session"}"}}`) + headers := http.Header{} + headers.Set(ClaudeCodeSessionHeader, "header-session") + + got := ExtractClaudeCodeSessionID(context.Background(), payload, headers) + if got != "header-session" { + t.Fatalf("ExtractClaudeCodeSessionID() = %q, want header-session", got) + } +} diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index c6795ef98..d488ae104 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -838,9 +838,9 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox body = normalizeCodexInstructions(body) body = sanitizeXAIResponsesBody(body, baseModel) - sessionID := xaiExecutionSessionID(req, opts) - if sessionID == "" && xaiRequiresIsolatedConversation(baseModel) { - sessionID = uuid.NewString() + sessionID, errSession := xaiResolveComposerSessionID(ctx, req, opts, baseModel) + if errSession != nil { + return nil, errSession } if sessionID != "" { body, _ = sjson.SetBytes(body, "prompt_cache_key", sessionID) @@ -917,6 +917,23 @@ func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, str util.ApplyCustomHeadersFromAttrs(r, attrs) } +func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) { + if sessionID := xaiExecutionSessionID(req, opts); sessionID != "" { + return sessionID, nil + } + if !xaiRequiresIsolatedConversation(baseModel) { + return "", nil + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, req.Model, req.Payload, opts.Headers) + if errCache != nil { + return "", errCache + } + if ok { + return cached.ID, nil + } + return uuid.NewString(), nil +} + func xaiExecutionSessionID(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) string { if value := xaiMetadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { return value diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 5e7b371a2..e674b5cf1 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -987,3 +987,41 @@ func TestNormalizeXAIToolChoiceForTools_NoOpWhenBothAbsent(t *testing.T) { t.Fatalf("tool_choice should not appear: %s", string(out)) } } + +func TestXAIExecutorComposerReusesClaudeCodeSession(t *testing.T) { + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Metadata: map[string]any{"access_token": "xai-token"}, + } + payload := []byte(`{"model":"grok-composer-2.5-fast","metadata":{"user_id":"{\"session_id\":\"cache-session-1\"}"},"input":"hello"}`) + req := cliproxyexecutor.Request{Model: "grok-composer-2.5-fast", Payload: payload} + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, Stream: true} + + first, err := exec.prepareResponsesRequest(context.Background(), req, opts, true) + if err != nil { + t.Fatalf("prepareResponsesRequest first error: %v", err) + } + second, err := exec.prepareResponsesRequest(context.Background(), req, opts, true) + if err != nil { + t.Fatalf("prepareResponsesRequest second error: %v", err) + } + + firstKey := gjson.GetBytes(first.body, "prompt_cache_key").String() + secondKey := gjson.GetBytes(second.body, "prompt_cache_key").String() + if firstKey == "" { + t.Fatalf("first prompt_cache_key is empty; body=%s", string(first.body)) + } + if secondKey != firstKey { + t.Fatalf("same Claude Code session produced different prompt_cache_key: first=%q second=%q", firstKey, secondKey) + } + + httpReq, errRequest := http.NewRequest(http.MethodPost, "https://example.test/responses", bytes.NewReader(first.body)) + if errRequest != nil { + t.Fatalf("NewRequest() error = %v", errRequest) + } + applyXAIHeaders(httpReq, auth, "xai-token", true, first.sessionID) + if got := httpReq.Header.Get("x-grok-conv-id"); got != firstKey { + t.Fatalf("x-grok-conv-id = %q, want %q", got, firstKey) + } +}