diff --git a/config.example.yaml b/config.example.yaml index 7c0a5180b..f4e771470 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -312,6 +312,9 @@ nonstream-keepalive-interval: 0 # base-url: "https://generativelanguage.googleapis.com" # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # proxy-url: "socks5://proxy.example.com:1080" # # proxy-url: "direct" # optional: explicit direct connect for this credential # models: @@ -346,6 +349,9 @@ nonstream-keepalive-interval: 0 # base-url: "https://generativelanguage.googleapis.com" # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # proxy-url: "socks5://proxy.example.com:1080" # # proxy-url: "direct" # optional: explicit direct connect for this credential # models: @@ -374,6 +380,9 @@ nonstream-keepalive-interval: 0 # alpha-search: false # optional: allow this key to serve /v1/alpha/search via base-url + /alpha/search # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential # models: @@ -413,6 +422,9 @@ nonstream-keepalive-interval: 0 # websockets: true # optional: use the xAI upstream websocket transport for downstream websocket requests # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential # models: @@ -444,6 +456,9 @@ nonstream-keepalive-interval: 0 # base-url: "https://www.example.com" # use the custom claude API endpoint # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential # models: @@ -589,6 +604,9 @@ nonstream-keepalive-interval: 0 # action: "stop" # "stop", "stop-and-cooldown", "continue", "continue-and-cooldown" # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # api-key-entries: # - api-key: "sk-or-v1-...b780" # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 @@ -630,6 +648,9 @@ nonstream-keepalive-interval: 0 # # proxy-url: "direct" # optional: explicit direct connect for this credential # headers: # X-Custom-Header: "custom-value" +# # Values starting with "$" dynamically copy the header value from downstream client requests. +# # If the client did not send the specified header, the header is omitted. +# # X-Claude-Code-Session-Id: "$ABC" # copies client's "ABC" header # models: # optional: map aliases to upstream model names # - name: "gemini-2.5-flash" # upstream model name # alias: "vertex-flash" # client-visible alias diff --git a/internal/runtime/executor/claude_executor_request.go b/internal/runtime/executor/claude_executor_request.go index 17ddc80e6..39e0f4ad9 100644 --- a/internal/runtime/executor/claude_executor_request.go +++ b/internal/runtime/executor/claude_executor_request.go @@ -927,7 +927,7 @@ func applyClaudeHeadersWithNativeProfile( if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(r, attrs) + util.ApplyCustomHeadersFromAttrs(r, attrs, incomingHeaders) // Custom credential headers are a configuration escape hatch for third-party // gateways, so they keep the last word there. On api.anthropic.com they must // not rewrite the reconstructed identity: an overridden Anthropic-Beta yields a diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go index 5d774908c..d7c6dbd7f 100644 --- a/internal/runtime/executor/codex_executor_execute.go +++ b/internal/runtime/executor/codex_executor_execute.go @@ -79,7 +79,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if err != nil { return resp, err } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg, opts.Headers) applyModelHeaderOverrides(httpReq.Header, baseModel) applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) var authID, authLabel, authType, authValue string @@ -239,7 +239,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A if err != nil { return resp, err } - applyCodexHeaders(httpReq, auth, apiKey, false, e.cfg) + applyCodexHeaders(httpReq, auth, apiKey, false, e.cfg, opts.Headers) applyModelHeaderOverrides(httpReq.Header, baseModel) applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) var authID, authLabel, authType, authValue string diff --git a/internal/runtime/executor/codex_executor_request.go b/internal/runtime/executor/codex_executor_request.go index ea65cafc6..205a40235 100644 --- a/internal/runtime/executor/codex_executor_request.go +++ b/internal/runtime/executor/codex_executor_request.go @@ -276,9 +276,11 @@ func codexIdentityConfuseUUID(authID string, kind string, value string) string { return uuid.NewSHA1(uuid.NameSpaceOID, []byte(name)).String() } -func applyCodexHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { +func applyCodexHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config, clientHeaders ...http.Header) { var ginHeaders http.Header - if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if len(clientHeaders) > 0 && clientHeaders[0] != nil { + ginHeaders = clientHeaders[0] + } else if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { ginHeaders = ginCtx.Request.Header } applyCodexHeadersFromSources(r, auth, token, stream, cfg, ginHeaders) @@ -303,9 +305,12 @@ func applyModelHeaderOverrides(headers http.Header, modelName string) { // applyCodexDirectImageHeaders sets Codex upstream headers for direct /images/* calls. // Downstream client User-Agent values are not forwarded to reduce Cloudflare 1010 blocks. -func applyCodexDirectImageHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { +func applyCodexDirectImageHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config, clientHeaders ...http.Header) { var ginHeaders http.Header - if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if len(clientHeaders) > 0 && clientHeaders[0] != nil { + ginHeaders = clientHeaders[0].Clone() + ginHeaders.Del("User-Agent") + } else if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { ginHeaders = ginCtx.Request.Header.Clone() ginHeaders.Del("User-Agent") } @@ -359,7 +364,7 @@ func applyCodexHeadersFromSources(r *http.Request, auth *cliproxyauth.Auth, toke if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(r, attrs) + util.ApplyCustomHeadersFromAttrs(r, attrs, ginHeaders) applyCodexCloakingHeaders(r.Header, cfg) } diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index d6e622646..00dc46ad4 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -83,7 +83,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if err != nil { return nil, err } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg, opts.Headers) applyModelHeaderOverrides(httpReq.Header, baseModel) applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) var authID, authLabel, authType, authValue string diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 18ef44186..5492f37ef 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -112,7 +112,7 @@ func (e *CodexExecutor) executeOpenAIImage(ctx context.Context, auth *cliproxyau if errCache != nil { return resp, errCache } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg, opts.Headers) applyModelHeaderOverrides(httpReq.Header, mainModel) applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) recordCodexOpenAIImageRequest(ctx, e.cfg, e.Identifier(), auth, url, httpReq.Header.Clone(), body) @@ -209,7 +209,7 @@ func (e *CodexExecutor) executeOpenAIImageStream(ctx context.Context, auth *clip if errCache != nil { return nil, errCache } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg, opts.Headers) applyModelHeaderOverrides(httpReq.Header, mainModel) applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) recordCodexOpenAIImageRequest(ctx, e.cfg, e.Identifier(), auth, url, httpReq.Header.Clone(), body) diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index 136c51273..72bace8cf 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -85,7 +85,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut var identityState codexIdentityConfuseState upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) reporter.SetTranslatedReasoningEffort(clientBody, to.String()) - wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) + wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg, opts.Headers) applyModelHeaderOverrides(wsHeaders, baseModel) applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) diff --git a/internal/runtime/executor/codex_websockets_request.go b/internal/runtime/executor/codex_websockets_request.go index 7d1cf92dd..977102fe1 100644 --- a/internal/runtime/executor/codex_websockets_request.go +++ b/internal/runtime/executor/codex_websockets_request.go @@ -64,7 +64,7 @@ func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktransl return rawJSON, headers, nil } -func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config) http.Header { +func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config, clientHeaders ...http.Header) http.Header { if headers == nil { headers = http.Header{} } @@ -73,7 +73,9 @@ func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth * } var ginHeaders http.Header - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if len(clientHeaders) > 0 && clientHeaders[0] != nil { + ginHeaders = clientHeaders[0].Clone() + } else if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { ginHeaders = ginCtx.Request.Header.Clone() } @@ -123,7 +125,7 @@ func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth * if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs) + util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs, ginHeaders) applyCodexCloakingHeaders(headers, cfg) return headers diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 1aec1c3f9..d094a894b 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -82,7 +82,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr var identityState codexIdentityConfuseState upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) reporter.SetTranslatedReasoningEffort(clientBody, to.String()) - wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) + wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg, opts.Headers) applyModelHeaderOverrides(wsHeaders, baseModel) applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) diff --git a/internal/runtime/executor/custom_magic_headers_test.go b/internal/runtime/executor/custom_magic_headers_test.go new file mode 100644 index 000000000..5fc0ad86c --- /dev/null +++ b/internal/runtime/executor/custom_magic_headers_test.go @@ -0,0 +1,408 @@ +package executor + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + 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" +) + +func TestCustomMagicHeaders_OpenAICompat(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"chatcmpl-1","choices":[{"message":{"role":"assistant","content":"ok"}}]}`)) + })) + defer server.Close() + + executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{ + OpenAICompatibility: []config.OpenAICompatibility{{ + Name: "compat", + }}, + }) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Forwarded-Session": "$X-Client-Session", + "header:X-Missing": "$NONEXISTENT", + "header:X-Static": "static-value", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gpt-4o", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Headers: http.Header{ + "Abc": []string{"session-abc-value"}, + "X-Client-Session": []string{"client-session-uuid-123"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "session-abc-value" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "session-abc-value") + } + if got := gotHeaders.Get("X-Forwarded-Session"); got != "client-session-uuid-123" { + t.Errorf("X-Forwarded-Session = %q, want %q", got, "client-session-uuid-123") + } + if got := gotHeaders.Get("X-Static"); got != "static-value" { + t.Errorf("X-Static = %q, want %q", got, "static-value") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_Gemini(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}`)) + })) + defer server.Close() + + executor := NewGeminiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "gemini", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "gemini-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + "header:X-Static": "gemini-static", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"contents":[{"parts":[{"text":"hi"}]}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatGemini, + Headers: http.Header{ + "Abc": []string{"gemini-session-abc"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "gemini-session-abc" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "gemini-session-abc") + } + if got := gotHeaders.Get("X-Static"); got != "gemini-static" { + t.Errorf("X-Static = %q, want %q", got, "gemini-static") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_GeminiInteractions(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"interaction_1","status":"completed","outputs":[{"text":"ok"}]}`)) + })) + defer server.Close() + + executor := NewGeminiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "gemini-interactions", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "interactions-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + "header:X-Static": "interactions-static", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gemini-3.1-flash-lite", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Headers: http.Header{ + "Abc": []string{"interactions-session-123"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "interactions-session-123" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "interactions-session-123") + } + if got := gotHeaders.Get("X-Static"); got != "interactions-static" { + t.Errorf("X-Static = %q, want %q", got, "interactions-static") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_GeminiVertex(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"candidates":[{"content":{"parts":[{"text":"vertex-response"}]}}]}`)) + })) + defer server.Close() + + executor := NewGeminiVertexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "vertex", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "vertex-api-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"contents":[{"parts":[{"text":"hi"}]}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatGemini, + Headers: http.Header{ + "Abc": []string{"vertex-session-123"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "vertex-session-123" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "vertex-session-123") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_XAI(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":0,\"status\":\"completed\",\"background\":false,\"error\":null,\"output\":[]}}\n\n")) + })) + defer server.Close() + + executor := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "xai-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + }, + } + + req := cliproxyexecutor.Request{ + Model: "grok-2", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Headers: http.Header{ + "ABC": []string{"xai-session-value"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "xai-session-value" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "xai-session-value") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_Claude(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}]}`)) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "claude", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "sk-ant-test", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + }, + } + + req := cliproxyexecutor.Request{ + Model: "claude-3-7-sonnet-20250219", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Headers: http.Header{ + "Abc": []string{"claude-session-value"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "claude-session-value" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "claude-session-value") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} + +func TestCustomMagicHeaders_OpenAICompat_Stream(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n")) + })) + defer server.Close() + + executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{ + OpenAICompatibility: []config.OpenAICompatibility{{ + Name: "compat", + }}, + }) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Empty-Var": "$ ", + "header:X-Only-Dollar": "$", + "header:X-Missing": "$NONEXISTENT", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gpt-4o", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Stream: true, + Headers: http.Header{ + "Abc": []string{"stream-session-abc"}, + }, + } + + result, err := executor.ExecuteStream(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + for range result.Chunks { + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "stream-session-abc" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "stream-session-abc") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } + if _, exists := gotHeaders["X-Empty-Var"]; exists { + t.Errorf("expected X-Empty-Var to be omitted, got %q", gotHeaders.Get("X-Empty-Var")) + } + if _, exists := gotHeaders["X-Only-Dollar"]; exists { + t.Errorf("expected X-Only-Dollar to be omitted, got %q", gotHeaders.Get("X-Only-Dollar")) + } +} + +func TestCustomMagicHeaders_Codex(t *testing.T) { + var gotHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeaders = r.Header.Clone() + body, _ := io.ReadAll(r.Body) + _ = body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":0,\"status\":\"completed\",\"background\":false,\"error\":null,\"output\":[]}}\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{ + Codex: config.CodexConfig{ + DisableCodexCloaking: true, + }, + }) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "codex-key", + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Missing": "$NONEXISTENT", + }, + } + + req := cliproxyexecutor.Request{ + Model: "gpt-5-codex", + Payload: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatCodex, + Headers: http.Header{ + "Abc": []string{"codex-session-value"}, + }, + } + + _, err := executor.Execute(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if got := gotHeaders.Get("X-Claude-Code-Session-Id"); got != "codex-session-value" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "codex-session-value") + } + if _, exists := gotHeaders["X-Missing"]; exists { + t.Errorf("expected X-Missing to be omitted, got %q", gotHeaders.Get("X-Missing")) + } +} diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 4b878ba53..70e3dcc7f 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -183,7 +183,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -295,7 +295,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -415,7 +415,7 @@ func (e *GeminiExecutor) executeInteractions(ctx context.Context, auth *cliproxy if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) applyGeminiInteractionsRequestHeaders(httpReq, opts.Headers) applyGeminiInteractionsRevisionHeader(httpReq) @@ -495,7 +495,7 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) applyGeminiInteractionsRequestHeaders(httpReq, opts.Headers) applyGeminiInteractionsRevisionHeader(httpReq) @@ -654,7 +654,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -896,12 +896,12 @@ func geminiAuthLogFields(auth *cliproxyauth.Auth) (string, string, string, strin return auth.ID, auth.Label, authType, authValue } -func applyGeminiHeaders(req *http.Request, auth *cliproxyauth.Auth) { +func applyGeminiHeaders(req *http.Request, auth *cliproxyauth.Auth, clientHeaders ...http.Header) { var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(req, attrs) + util.ApplyCustomHeadersFromAttrs(req, attrs, clientHeaders...) } func capGeminiMaxOutputTokens(body []byte, modelName string) []byte { diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index 395ddf197..6d6eb8ffb 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -369,12 +369,12 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au log.Errorf("vertex executor: access token error: %v", errTok) return resp, statusErr{code: 500, msg: "internal server error"} } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { @@ -497,12 +497,12 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { @@ -614,12 +614,12 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte log.Errorf("vertex executor: access token error: %v", errTok) return nil, statusErr{code: 500, msg: "internal server error"} } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { @@ -761,12 +761,12 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { @@ -889,12 +889,12 @@ func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context log.Errorf("vertex executor: access token error: %v", errTok) return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"} } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { @@ -980,12 +980,12 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth * if apiKey != "" { httpReq.Header.Set("x-goog-api-key", apiKey) } - applyGeminiHeaders(httpReq, auth) + applyGeminiHeaders(httpReq, auth, opts.Headers) var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index 9429568d1..ee679d6d8 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -157,7 +157,7 @@ func (e *OpenAICompatExecutor) Execute(ctx context.Context, auth *cliproxyauth.A if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -251,7 +251,7 @@ func (e *OpenAICompatExecutor) executeImages(ctx context.Context, auth *cliproxy if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -369,7 +369,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) httpReq.Header.Set("Accept", "text/event-stream") httpReq.Header.Set("Cache-Control", "no-cache") var authID, authLabel, authType, authValue string @@ -608,7 +608,7 @@ func (e *OpenAICompatExecutor) executeImagesStream(ctx context.Context, auth *cl if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + util.ApplyCustomHeadersFromAttrs(httpReq, attrs, opts.Headers) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go index a7b59c8c8..72f769729 100644 --- a/internal/runtime/executor/xai_executor_execute.go +++ b/internal/runtime/executor/xai_executor_execute.go @@ -23,7 +23,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req return e.executeCompact(ctx, auth, req, opts) } if endpointPath := xaiImageEndpointPath(opts); endpointPath != "" { - return e.executeImages(ctx, auth, req, endpointPath) + return e.executeImages(ctx, auth, req, opts, endpointPath) } if xaiIsVideoRequest(opts) { return e.executeVideos(ctx, auth, req, opts) @@ -47,7 +47,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req if err != nil { return resp, err } - applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) + applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID, opts.Headers) e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) @@ -159,7 +159,7 @@ func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxya } // Official API / custom compact endpoints use standard API headers, not CLI // chat-proxy identity headers (which applyXAIChatHeaders may still attach for OAuth chat). - applyXAIHeaders(httpReq, auth, token, false, prepared.sessionID) + applyXAIHeaders(httpReq, auth, token, false, prepared.sessionID, opts.Headers) e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), prepared.body) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) diff --git a/internal/runtime/executor/xai_executor_media.go b/internal/runtime/executor/xai_executor_media.go index 29128c8ab..f5df302d5 100644 --- a/internal/runtime/executor/xai_executor_media.go +++ b/internal/runtime/executor/xai_executor_media.go @@ -16,7 +16,7 @@ import ( "github.com/tidwall/gjson" ) -func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, endpointPath string) (resp cliproxyexecutor.Response, err error) { +func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, endpointPath string) (resp cliproxyexecutor.Response, err error) { model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String()) if model == "" { model = strings.TrimSpace(req.Model) @@ -39,7 +39,7 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth if err != nil { return resp, err } - applyXAIHeaders(httpReq, auth, token, false, "") + applyXAIHeaders(httpReq, auth, token, false, "", opts.Headers) e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), payload) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) @@ -107,7 +107,7 @@ func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth if err != nil { return resp, err } - applyXAIHeaders(httpReq, auth, token, false, "") + applyXAIHeaders(httpReq, auth, token, false, "", opts.Headers) if method == http.MethodPost { key := xaiMetadataString(opts.Metadata, xaiIdempotencyKeyMetaKey) if key == "" && opts.Headers != nil { diff --git a/internal/runtime/executor/xai_executor_request.go b/internal/runtime/executor/xai_executor_request.go index 6ee90d66f..04810dcbc 100644 --- a/internal/runtime/executor/xai_executor_request.go +++ b/internal/runtime/executor/xai_executor_request.go @@ -276,9 +276,9 @@ func logXAIResolvedBaseURL(ctx context.Context, baseURL string) { helps.LogWithRequestID(ctx).Infof("xai: using base_url=%s source=%s", baseURL, xaiBaseURLSource(baseURL)) } -func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { +func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string, clientHeaders ...http.Header) { applyXAIDefaultHeaders(r, token, stream, sessionID) - applyXAICustomHeaders(r, auth) + applyXAICustomHeaders(r, auth, clientHeaders...) } func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionID string) { @@ -297,12 +297,12 @@ func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionI } } -func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) { +func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth, clientHeaders ...http.Header) { var attrs map[string]string if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(r, attrs) + util.ApplyCustomHeadersFromAttrs(r, attrs, clientHeaders...) } // applyXAIChatHeaders applies standard xAI headers for non-image/video chat @@ -310,9 +310,9 @@ func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) { // applyXAIHeaders behavior. CLI chat-proxy identity headers are only attached // when using_api is false and the resolved chat base URL is the official CLI // chat-proxy endpoint. -func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { +func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string, clientHeaders ...http.Header) { if xaiUsingAPI(auth) { - applyXAIHeaders(r, auth, token, stream, sessionID) + applyXAIHeaders(r, auth, token, stream, sessionID, clientHeaders...) return } applyXAIDefaultHeaders(r, token, stream, sessionID) @@ -323,7 +323,7 @@ func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, r.Header.Set(xaiClientIdentifierHeader, xaiClientIdentifierValue) r.Header.Set(xaiAuthenticateResponseHeader, xaiAuthenticateResponseValue) } - applyXAICustomHeaders(r, auth) + applyXAICustomHeaders(r, auth, clientHeaders...) } func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) { diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go index 5ccbe292c..4dad2c51f 100644 --- a/internal/runtime/executor/xai_executor_stream.go +++ b/internal/runtime/executor/xai_executor_stream.go @@ -41,7 +41,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth if err != nil { return nil, err } - applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) + applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID, opts.Headers) e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index f0d39b047..956046812 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -507,7 +507,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) - wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID) + wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID, opts.Headers) wsReqBody := buildXAIWebsocketRequestBody(prepared.body) requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String()) transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" && @@ -1442,7 +1442,7 @@ func buildXAIResponsesWebsocketURL(httpURL string) (string, error) { return parsed.String(), nil } -func applyXAIWebsocketHeaders(headers http.Header, auth *cliproxyauth.Auth, token string, sessionID string) http.Header { +func applyXAIWebsocketHeaders(headers http.Header, auth *cliproxyauth.Auth, token string, sessionID string, clientHeaders ...http.Header) http.Header { if headers == nil { headers = http.Header{} } @@ -1457,7 +1457,7 @@ func applyXAIWebsocketHeaders(headers http.Header, auth *cliproxyauth.Auth, toke if auth != nil { attrs = auth.Attributes } - util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs) + util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs, clientHeaders...) return headers } diff --git a/internal/util/header_helpers.go b/internal/util/header_helpers.go index 0b8d72bcb..f100fab80 100644 --- a/internal/util/header_helpers.go +++ b/internal/util/header_helpers.go @@ -3,18 +3,34 @@ package util import ( "net/http" "strings" + + "github.com/gin-gonic/gin" ) // ApplyCustomHeadersFromAttrs applies user-defined headers stored in the provided attributes map. // Custom headers override built-in defaults when conflicts occur. -func ApplyCustomHeadersFromAttrs(r *http.Request, attrs map[string]string) { +// If clientHeaders is provided (or if the request context carries a Gin context), any custom header +// whose value starts with "$" (e.g. "$ABC" or "$X-Claude-Code-Session-Id") is dynamically +// resolved from the client's request headers. If the client did not provide that header, +// the custom header is omitted from the outgoing request. +func ApplyCustomHeadersFromAttrs(r *http.Request, attrs map[string]string, clientHeaders ...http.Header) { if r == nil { return } - applyCustomHeaders(r, extractCustomHeaders(attrs)) + var ch http.Header + if len(clientHeaders) > 0 && clientHeaders[0] != nil { + ch = clientHeaders[0] + } else if r.Context() != nil { + if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + ch = ginCtx.Request.Header + } else if ginCtx, ok := r.Context().(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + ch = ginCtx.Request.Header + } + } + applyCustomHeaders(r, extractCustomHeaders(attrs, ch)) } -func extractCustomHeaders(attrs map[string]string) map[string]string { +func extractCustomHeaders(attrs map[string]string, clientHeaders http.Header) map[string]string { if len(attrs) == 0 { return nil } @@ -31,6 +47,25 @@ func extractCustomHeaders(attrs map[string]string) map[string]string { if val == "" { continue } + if strings.HasPrefix(val, "$") { + varName := strings.TrimSpace(strings.TrimPrefix(val, "$")) + if varName == "" || clientHeaders == nil { + continue + } + clientVal := clientHeaders.Get(varName) + if clientVal == "" { + for ck, cv := range clientHeaders { + if strings.EqualFold(ck, varName) && len(cv) > 0 && cv[0] != "" { + clientVal = cv[0] + break + } + } + } + if clientVal == "" { + continue + } + val = clientVal + } headers[name] = val } if len(headers) == 0 { diff --git a/internal/util/header_helpers_test.go b/internal/util/header_helpers_test.go new file mode 100644 index 000000000..1f9d29ada --- /dev/null +++ b/internal/util/header_helpers_test.go @@ -0,0 +1,116 @@ +package util + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestApplyCustomHeadersFromAttrs_StaticHeaders(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "https://api.example.com", nil) + attrs := map[string]string{ + "header:X-Custom-Static": "static-value", + "header:Host": "custom.host.com", + } + + ApplyCustomHeadersFromAttrs(req, attrs) + + if got := req.Header.Get("X-Custom-Static"); got != "static-value" { + t.Errorf("X-Custom-Static = %q, want %q", got, "static-value") + } + if got := req.Host; got != "custom.host.com" { + t.Errorf("req.Host = %q, want %q", got, "custom.host.com") + } +} + +func TestApplyCustomHeadersFromAttrs_MagicVariable(t *testing.T) { + t.Run("present in clientHeaders sets header", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "https://api.example.com", nil) + attrs := map[string]string{ + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Target-Session": "$X-Claude-Code-Session-Id", + "header:Static-Header": "static-123", + } + clientHeaders := http.Header{ + "Abc": []string{"session-abc-456"}, + "X-Claude-Code-Session-Id": []string{"claude-code-uuid-789"}, + } + + ApplyCustomHeadersFromAttrs(req, attrs, clientHeaders) + + if got := req.Header.Get("X-Claude-Code-Session-Id"); got != "session-abc-456" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "session-abc-456") + } + if got := req.Header.Get("X-Target-Session"); got != "claude-code-uuid-789" { + t.Errorf("X-Target-Session = %q, want %q", got, "claude-code-uuid-789") + } + if got := req.Header.Get("Static-Header"); got != "static-123" { + t.Errorf("Static-Header = %q, want %q", got, "static-123") + } + }) + + t.Run("absent in clientHeaders does not set header", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "https://api.example.com", nil) + attrs := map[string]string{ + "header:X-Claude-Code-Session-Id": "$ABC", + "header:X-Other": "$NONEXISTENT", + "header:Static-Header": "static-123", + } + clientHeaders := http.Header{ + "Other-Header": []string{"some-value"}, + } + + ApplyCustomHeadersFromAttrs(req, attrs, clientHeaders) + + if _, exists := req.Header["X-Claude-Code-Session-Id"]; exists { + t.Errorf("expected X-Claude-Code-Session-Id to be omitted when $ABC is absent in clientHeaders, got %q", req.Header.Get("X-Claude-Code-Session-Id")) + } + if _, exists := req.Header["X-Other"]; exists { + t.Errorf("expected X-Other to be omitted when $NONEXISTENT is absent in clientHeaders, got %q", req.Header.Get("X-Other")) + } + if got := req.Header.Get("Static-Header"); got != "static-123" { + t.Errorf("Static-Header = %q, want %q", got, "static-123") + } + }) + + t.Run("nil clientHeaders does not set variable headers", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "https://api.example.com", nil) + attrs := map[string]string{ + "header:X-Claude-Code-Session-Id": "$ABC", + "header:Static-Header": "static-123", + } + + ApplyCustomHeadersFromAttrs(req, attrs) + + if _, exists := req.Header["X-Claude-Code-Session-Id"]; exists { + t.Errorf("expected X-Claude-Code-Session-Id to be omitted with nil clientHeaders, got %q", req.Header.Get("X-Claude-Code-Session-Id")) + } + if got := req.Header.Get("Static-Header"); got != "static-123" { + t.Errorf("Static-Header = %q, want %q", got, "static-123") + } + }) + + t.Run("fallback to gin context in request context", func(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(w) + ginReq := httptest.NewRequest(http.MethodPost, "/", nil) + ginReq.Header.Set("ABC", "from-gin-ctx-123") + ginCtx.Request = ginReq + + req := httptest.NewRequest(http.MethodPost, "https://api.example.com", nil) + req = req.WithContext(ginCtx) + + attrs := map[string]string{ + "header:X-Claude-Code-Session-Id": "$ABC", + } + + ApplyCustomHeadersFromAttrs(req, attrs) + + if got := req.Header.Get("X-Claude-Code-Session-Id"); got != "from-gin-ctx-123" { + t.Errorf("X-Claude-Code-Session-Id = %q, want %q", got, "from-gin-ctx-123") + } + }) +}