You can close this window.
This window will close automatically in 5 seconds.
` var corsExposedResponseHeaders = []string{ + logging.CPATraceIDHeader, "X-CPA-VERSION", "X-CPA-COMMIT", "X-CPA-BUILD-DATE", @@ -292,6 +293,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk // Add middleware engine.Use(logging.GinLogrusLogger()) engine.Use(logging.GinLogrusRecovery()) + engine.Use(logging.CPATraceIDMiddleware()) for _, mw := range optionState.extraMiddleware { engine.Use(mw) } @@ -689,6 +691,7 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { c.JSON(status, gin.H{"error": err.Error()}) return } + logging.SetGinCPATraceID(c, selected.EnsureIndex()) headers := make(http.Header) headers.Set("Content-Type", "application/json") diff --git a/internal/api/server_test.go b/internal/api/server_test.go index 1f9c6a3c..0a83c852 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -212,6 +212,14 @@ func TestCodexAlphaSearchForwardsRequest(t *testing.T) { if got := rr.Header().Get("Content-Type"); got != "application/json" { t.Fatalf("response Content-Type = %q", got) } + traceID := rr.Header().Get(internallogging.CPATraceIDHeader) + parts := strings.Split(traceID, "-") + if len(parts) != 3 || parts[1] != credential.Index || len(parts[2]) != 8 { + t.Fatalf("trace ID = %q, want timestamp-%s-requestID", traceID, credential.Index) + } + if _, errParse := time.Parse("20060102150405", parts[0]); errParse != nil { + t.Fatalf("trace timestamp = %q: %v", parts[0], errParse) + } } func TestCodexAlphaSearchSanitizesResponsesOnlyFields(t *testing.T) { @@ -808,6 +816,9 @@ func TestExampleAPIKeySafeModeShowsWarningAndKeepsManagement(t *testing.T) { if !strings.Contains(rr.Body.String(), "/management.html?safe-mode=configure") { t.Fatalf("body missing management link in message: %s", rr.Body.String()) } + if got := rr.Header().Get(internallogging.CPATraceIDHeader); got != "" { + t.Fatalf("trace ID = %q, want empty before auth selection", got) + } }) t.Run("management endpoints still work", func(t *testing.T) { @@ -818,6 +829,9 @@ func TestExampleAPIKeySafeModeShowsWarningAndKeepsManagement(t *testing.T) { if rr.Code != http.StatusOK { t.Fatalf("status = %d, want %d body=%s", rr.Code, http.StatusOK, rr.Body.String()) } + if got := rr.Header().Get(internallogging.CPATraceIDHeader); got != "" { + t.Fatalf("management trace ID = %q, want empty", got) + } }) t.Run("safe mode clears after key update", func(t *testing.T) { diff --git a/internal/logging/cpa_trace.go b/internal/logging/cpa_trace.go new file mode 100644 index 00000000..bc1d243a --- /dev/null +++ b/internal/logging/cpa_trace.go @@ -0,0 +1,151 @@ +package logging + +import ( + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +// CPATraceIDHeader is the downstream response header used to correlate requests with selected credentials. +const CPATraceIDHeader = "X-CPA-TRACE-ID" + +const ginCPATraceStateKey = "__cpa_trace_state__" + +// FormatCPATraceID builds a CPA trace ID from the selection time, auth index, and request ID. +func FormatCPATraceID(selectedAt time.Time, authIndex, requestID string) string { + authIndex = strings.TrimSpace(authIndex) + requestID = strings.TrimSpace(requestID) + if selectedAt.IsZero() || authIndex == "" || requestID == "" { + return "" + } + return selectedAt.Format("20060102150405") + "-" + authIndex + "-" + requestID +} + +type cpaTraceState struct { + mu sync.RWMutex + traceID string +} + +func (s *cpaTraceState) set(traceID string) { + if s == nil { + return + } + s.mu.Lock() + s.traceID = strings.TrimSpace(traceID) + s.mu.Unlock() +} + +func (s *cpaTraceState) get() string { + if s == nil { + return "" + } + s.mu.RLock() + traceID := s.traceID + s.mu.RUnlock() + return traceID +} + +func ginCPATraceState(c *gin.Context) *cpaTraceState { + if c == nil { + return nil + } + if value, exists := c.Get(ginCPATraceStateKey); exists { + if state, ok := value.(*cpaTraceState); ok && state != nil { + return state + } + } + state := &cpaTraceState{} + c.Set(ginCPATraceStateKey, state) + return state +} + +// GinCPATraceIDCallback returns a callback that is safe to invoke after the Gin context is released. +func GinCPATraceIDCallback(c *gin.Context) func(string) { + state := ginCPATraceState(c) + if state == nil { + return nil + } + requestID := GetGinRequestID(c) + if requestID == "" && c.Request != nil { + requestID = GetRequestID(c.Request.Context()) + } + requestID = strings.TrimSpace(requestID) + if requestID == "" { + return nil + } + return func(authIndex string) { + if traceID := FormatCPATraceID(time.Now(), authIndex, requestID); traceID != "" { + state.set(traceID) + } + } +} + +// SetGinCPATraceID stores the trace ID until the downstream response headers are committed. +func SetGinCPATraceID(c *gin.Context, authIndex string) { + if callback := GinCPATraceIDCallback(c); callback != nil { + callback(authIndex) + } +} + +// GetGinCPATraceID returns the trace ID stored for the current request. +func GetGinCPATraceID(c *gin.Context) string { + if c == nil { + return "" + } + value, exists := c.Get(ginCPATraceStateKey) + if !exists { + return "" + } + state, _ := value.(*cpaTraceState) + return state.get() +} + +// CPATraceIDMiddleware injects a stored trace ID immediately before response headers are committed. +func CPATraceIDMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + state := ginCPATraceState(c) + c.Writer = &cpaTraceResponseWriter{ResponseWriter: c.Writer, state: state} + c.Next() + } +} + +type cpaTraceResponseWriter struct { + gin.ResponseWriter + state *cpaTraceState +} + +func (w *cpaTraceResponseWriter) WriteHeader(statusCode int) { + w.applyTraceHeader() + w.ResponseWriter.WriteHeader(statusCode) +} + +func (w *cpaTraceResponseWriter) WriteHeaderNow() { + w.applyTraceHeader() + w.ResponseWriter.WriteHeaderNow() +} + +func (w *cpaTraceResponseWriter) Write(data []byte) (int, error) { + w.applyTraceHeader() + return w.ResponseWriter.Write(data) +} + +func (w *cpaTraceResponseWriter) WriteString(data string) (int, error) { + w.applyTraceHeader() + return w.ResponseWriter.WriteString(data) +} + +func (w *cpaTraceResponseWriter) Flush() { + w.applyTraceHeader() + w.ResponseWriter.Flush() +} + +func (w *cpaTraceResponseWriter) applyTraceHeader() { + if w == nil || w.ResponseWriter == nil || w.ResponseWriter.Written() { + return + } + if traceID := w.state.get(); traceID != "" { + w.ResponseWriter.Header().Set(CPATraceIDHeader, traceID) + } +} diff --git a/internal/logging/cpa_trace_test.go b/internal/logging/cpa_trace_test.go new file mode 100644 index 00000000..2202e509 --- /dev/null +++ b/internal/logging/cpa_trace_test.go @@ -0,0 +1,115 @@ +package logging + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestFormatCPATraceID(t *testing.T) { + selectedAt := time.Date(2026, time.July, 17, 21, 58, 49, 0, time.UTC) + got := FormatCPATraceID(selectedAt, "auth-index", "request1") + if want := "20260717215849-auth-index-request1"; got != want { + t.Fatalf("FormatCPATraceID() = %q, want %q", got, want) + } + + for _, test := range []struct { + name string + selectedAt time.Time + authIndex string + requestID string + }{ + {name: "zero time", authIndex: "auth-index", requestID: "request1"}, + {name: "empty auth index", selectedAt: selectedAt, requestID: "request1"}, + {name: "empty request ID", selectedAt: selectedAt, authIndex: "auth-index"}, + } { + t.Run(test.name, func(t *testing.T) { + if gotEmpty := FormatCPATraceID(test.selectedAt, test.authIndex, test.requestID); gotEmpty != "" { + t.Fatalf("FormatCPATraceID() = %q, want empty", gotEmpty) + } + }) + } +} + +func TestCPATraceIDMiddlewareRequiresAuthIndexBeforeResponseCommit(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Use(CPATraceIDMiddleware()) + engine.GET("/selected", func(c *gin.Context) { + SetGinRequestID(c, "1234abcd") + SetGinCPATraceID(c, "auth-index") + c.Status(http.StatusOK) + }) + engine.GET("/unselected", func(c *gin.Context) { + SetGinRequestID(c, "1234abcd") + SetGinCPATraceID(c, "") + c.Status(http.StatusOK) + }) + engine.GET("/committed", func(c *gin.Context) { + SetGinRequestID(c, "1234abcd") + c.Writer.WriteHeaderNow() + SetGinCPATraceID(c, "auth-index") + }) + + t.Run("writes selected auth trace", func(t *testing.T) { + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/selected", nil)) + + traceID := recorder.Header().Get(CPATraceIDHeader) + if len(traceID) != len("20060102150405-auth-index-1234abcd") { + t.Fatalf("trace ID = %q, unexpected length", traceID) + } + if got := traceID[15:]; got != "auth-index-1234abcd" { + t.Fatalf("trace suffix = %q, want %q", got, "auth-index-1234abcd") + } + }) + + t.Run("skips empty auth index", func(t *testing.T) { + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/unselected", nil)) + + if got := recorder.Header().Get(CPATraceIDHeader); got != "" { + t.Fatalf("trace ID = %q, want empty", got) + } + }) + + t.Run("skips committed response", func(t *testing.T) { + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/committed", nil)) + + if got := recorder.Header().Get(CPATraceIDHeader); got != "" { + t.Fatalf("trace ID = %q, want empty", got) + } + }) +} + +func TestCPATraceIDConcurrentSelectionAndResponseCommit(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Use(CPATraceIDMiddleware()) + engine.GET("/race", func(c *gin.Context) { + SetGinRequestID(c, "1234abcd") + traceCallback := GinCPATraceIDCallback(c) + start := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + <-start + traceCallback("auth-index") + }() + close(start) + _, _ = c.Writer.Write([]byte("\n")) + <-done + }) + + for range 100 { + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/race", nil)) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusOK) + } + } +} diff --git a/internal/logging/gin_logger.go b/internal/logging/gin_logger.go index 446c97fb..ad905fcb 100644 --- a/internal/logging/gin_logger.go +++ b/internal/logging/gin_logger.go @@ -18,15 +18,10 @@ import ( // aiAPIPrefixes defines path prefixes for AI API requests that should have request ID tracking. var aiAPIPrefixes = []string{ - "/v1/chat/completions", - "/v1/completions", - "/v1/images", - "/v1/videos", - "/v1/messages", - "/v1/responses", - "/openai/v1/videos", - "/v1beta/models/", - "/backend-api/codex/", + "/v1", + "/v1beta", + "/openai/v1", + "/backend-api/codex", } const ( @@ -107,7 +102,7 @@ func GinLogrusLogger() gin.HandlerFunc { // isAIAPIPath checks if the given path is an AI API endpoint that should have request ID tracking. func isAIAPIPath(path string) bool { for _, prefix := range aiAPIPrefixes { - if strings.HasPrefix(path, prefix) { + if path == prefix || strings.HasPrefix(path, prefix+"/") { return true } } diff --git a/internal/logging/gin_logger_test.go b/internal/logging/gin_logger_test.go index a3c203ae..20ade057 100644 --- a/internal/logging/gin_logger_test.go +++ b/internal/logging/gin_logger_test.go @@ -59,6 +59,31 @@ func TestGinLogrusRecoveryHandlesRegularPanic(t *testing.T) { } } +func TestIsAIAPIPathIncludesPublicAPIGroups(t *testing.T) { + for _, path := range []string{ + "/v1", + "/v1/models", + "/v1/alpha/search", + "/v1beta/interactions", + "/openai/v1/videos", + "/backend-api/codex/responses", + } { + if !isAIAPIPath(path) { + t.Fatalf("expected %s to be treated as AI API path", path) + } + } + for _, path := range []string{ + "/v0/management/config", + "/v10/models", + "/openai/v10/videos", + "/backend-api/codex-status", + } { + if isAIAPIPath(path) { + t.Fatalf("expected %s not to be treated as AI API path", path) + } + } +} + func TestIsAIAPIPathIncludesImages(t *testing.T) { if !isAIAPIPath("/v1/images/generations") { t.Fatalf("expected /v1/images/generations to be treated as AI API path") diff --git a/sdk/api/handlers/claude/code_handlers.go b/sdk/api/handlers/claude/code_handlers.go index e9bdd600..16a75d83 100644 --- a/sdk/api/handlers/claude/code_handlers.go +++ b/sdk/api/handlers/claude/code_handlers.go @@ -418,7 +418,7 @@ func (h *ClaudeCodeAPIHandler) WriteErrorResponse(c *gin.Context, msg *interface } if msg != nil && msg.Addon != nil && handlers.PassthroughHeadersEnabled(h.Cfg) { for key, values := range msg.Addon { - if len(values) == 0 { + if len(values) == 0 || handlers.IsCPAReservedResponseHeader(key) { continue } c.Writer.Header().Del(key) diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index f3ab7178..a05b1e41 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -16,6 +16,7 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" . "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" @@ -260,8 +261,10 @@ func requestExecutionMetadata(ctx context.Context) map[string]any { // Only include it if the client explicitly provides it. key := "" requestPath := "" + var ginCtx *gin.Context if ctx != nil { - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if requestGinCtx, ok := ctx.Value("gin").(*gin.Context); ok && requestGinCtx != nil && requestGinCtx.Request != nil { + ginCtx = requestGinCtx key = strings.TrimSpace(ginCtx.GetHeader("Idempotency-Key")) requestPath = strings.TrimSpace(ginCtx.FullPath()) if requestPath == "" && ginCtx.Request.URL != nil { @@ -283,6 +286,11 @@ func requestExecutionMetadata(ctx context.Context) map[string]any { if selectedCallback := selectedAuthIDCallbackFromContext(ctx); selectedCallback != nil { meta[coreexecutor.SelectedAuthCallbackMetadataKey] = selectedCallback } + if ginCtx != nil && !websocket.IsWebSocketUpgrade(ginCtx.Request) { + if traceCallback := logging.GinCPATraceIDCallback(ginCtx); traceCallback != nil { + meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey] = traceCallback + } + } if executionSessionID := executionSessionIDFromContext(ctx); executionSessionID != "" { meta[coreexecutor.ExecutionSessionMetadataKey] = executionSessionID } @@ -2200,7 +2208,7 @@ func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.Erro } if msg != nil && msg.Addon != nil && PassthroughHeadersEnabled(h.Cfg) { for key, values := range msg.Addon { - if len(values) == 0 { + if len(values) == 0 || IsCPAReservedResponseHeader(key) { continue } c.Writer.Header().Del(key) diff --git a/sdk/api/handlers/handlers_error_response_test.go b/sdk/api/handlers/handlers_error_response_test.go index 0c206e38..0f2ad88e 100644 --- a/sdk/api/handlers/handlers_error_response_test.go +++ b/sdk/api/handlers/handlers_error_response_test.go @@ -47,14 +47,18 @@ func TestWriteErrorResponse_AddonHeadersEnabled(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodGet, "/", nil) c.Writer.Header().Set("X-Request-Id", "old-value") + c.Writer.Header().Set("x-cpa-trace-id", "local-trace") + c.Writer.Header().Set("Access-Control-Expose-Headers", "x-cpa-trace-id") handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true}, nil) handler.WriteErrorResponse(c, &interfaces.ErrorMessage{ StatusCode: http.StatusTooManyRequests, Error: errors.New("rate limit"), Addon: http.Header{ - "Retry-After": {"30"}, - "X-Request-Id": {"new-1", "new-2"}, + "Retry-After": {"30"}, + "X-Request-Id": {"new-1", "new-2"}, + "x-cpa-trace-id": {"upstream-trace"}, + "Access-Control-Expose-Headers": {"upstream-header"}, }, }) @@ -67,6 +71,12 @@ func TestWriteErrorResponse_AddonHeadersEnabled(t *testing.T) { if got := recorder.Header().Values("X-Request-Id"); !reflect.DeepEqual(got, []string{"new-1", "new-2"}) { t.Fatalf("X-Request-Id = %#v, want %#v", got, []string{"new-1", "new-2"}) } + if got := recorder.Header().Get("x-cpa-trace-id"); got != "local-trace" { + t.Fatalf("x-cpa-trace-id = %q, want local trace", got) + } + if got := recorder.Header().Get("Access-Control-Expose-Headers"); got != "x-cpa-trace-id" { + t.Fatalf("Access-Control-Expose-Headers = %q, want CPA value", got) + } } func TestEnrichAuthSelectionError_DefaultsTo503WithContext(t *testing.T) { diff --git a/sdk/api/handlers/handlers_metadata_test.go b/sdk/api/handlers/handlers_metadata_test.go index 60022b56..30b3984d 100644 --- a/sdk/api/handlers/handlers_metadata_test.go +++ b/sdk/api/handlers/handlers_metadata_test.go @@ -1,8 +1,12 @@ package handlers import ( + "net/http" + "net/http/httptest" "testing" + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "golang.org/x/net/context" ) @@ -19,6 +23,39 @@ func TestRequestExecutionMetadataIncludesExecutionSessionWithoutIdempotencyKey(t } } +func TestRequestExecutionMetadataTraceCallbackWebsocketDetection(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("skips websocket upgrade", func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + ginCtx.Request.Header.Set("Connection", "Upgrade") + ginCtx.Request.Header.Set("Upgrade", "websocket") + logging.SetGinRequestID(ginCtx, "1234abcd") + ctx := context.WithValue(context.Background(), "gin", ginCtx) + + meta := requestExecutionMetadata(ctx) + + if _, exists := meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey]; exists { + t.Fatal("unexpected selected auth index callback for websocket upgrade") + } + }) + + t.Run("keeps callback for incomplete upgrade headers", func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + ginCtx.Request.Header.Set("Upgrade", "websocket") + logging.SetGinRequestID(ginCtx, "1234abcd") + ctx := context.WithValue(context.Background(), "gin", ginCtx) + + meta := requestExecutionMetadata(ctx) + + if _, exists := meta[coreexecutor.SelectedAuthIndexCallbackMetadataKey]; !exists { + t.Fatal("missing selected auth index callback for ordinary HTTP request") + } + }) +} + func TestSetReasoningEffortMetadataUsesSuffixOverBody(t *testing.T) { meta := make(map[string]any) diff --git a/sdk/api/handlers/handlers_stream_bootstrap_test.go b/sdk/api/handlers/handlers_stream_bootstrap_test.go index 551baac3..9e3789d2 100644 --- a/sdk/api/handlers/handlers_stream_bootstrap_test.go +++ b/sdk/api/handlers/handlers_stream_bootstrap_test.go @@ -4,11 +4,15 @@ import ( "context" "errors" "net/http" + "net/http/httptest" "strings" "sync" "testing" + "time" + "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -634,8 +638,15 @@ func TestExecuteStreamWithAuthManager_SelectedAuthCallbackReceivesAuthID(t *test }, }, manager) + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + logging.SetGinRequestID(ginCtx, "1234abcd") + selectedAuthID := "" - ctx := WithSelectedAuthIDCallback(context.Background(), func(authID string) { + ctx := context.WithValue(context.Background(), "gin", ginCtx) + ctx = WithSelectedAuthIDCallback(ctx, func(authID string) { selectedAuthID = authID }) dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", "test-model", []byte(`{"model":"test-model"}`), "") @@ -659,6 +670,14 @@ func TestExecuteStreamWithAuthManager_SelectedAuthCallbackReceivesAuthID(t *test if selectedAuthID != "auth2" { t.Fatalf("selectedAuthID = %q, want %q", selectedAuthID, "auth2") } + traceID := logging.GetGinCPATraceID(ginCtx) + parts := strings.Split(traceID, "-") + if len(parts) != 3 || parts[1] != auth2.Index || parts[2] != "1234abcd" { + t.Fatalf("trace ID = %q, want timestamp-%s-1234abcd", traceID, auth2.Index) + } + if _, errParse := time.Parse("20060102150405", parts[0]); errParse != nil { + t.Fatalf("trace timestamp = %q: %v", parts[0], errParse) + } } func TestExecuteStreamWithAuthManager_ValidatesOpenAIResponsesStreamDataJSON(t *testing.T) { diff --git a/sdk/api/handlers/header_filter.go b/sdk/api/handlers/header_filter.go index 73626d38..4e2617ec 100644 --- a/sdk/api/handlers/header_filter.go +++ b/sdk/api/handlers/header_filter.go @@ -36,6 +36,17 @@ var hopByHopHeaders = map[string]struct{}{ "Content-Encoding": {}, } +var cpaReservedResponseHeaders = map[string]struct{}{ + "Access-Control-Expose-Headers": {}, + "X-Cpa-Trace-Id": {}, +} + +// IsCPAReservedResponseHeader reports whether a downstream response header is managed by CPA. +func IsCPAReservedResponseHeader(name string) bool { + _, reserved := cpaReservedResponseHeaders[http.CanonicalHeaderKey(name)] + return reserved +} + // FilterUpstreamHeaders returns a copy of src with hop-by-hop and security-sensitive // headers removed. Returns nil if src is nil or empty after filtering. func FilterUpstreamHeaders(src http.Header) http.Header { @@ -49,6 +60,9 @@ func FilterUpstreamHeaders(src http.Header) http.Header { if _, blocked := hopByHopHeaders[canonicalKey]; blocked { continue } + if _, reserved := cpaReservedResponseHeaders[canonicalKey]; reserved { + continue + } if _, scoped := connectionScoped[canonicalKey]; scoped { continue } diff --git a/sdk/api/handlers/header_filter_test.go b/sdk/api/handlers/header_filter_test.go index a87e65a1..38ed9aa3 100644 --- a/sdk/api/handlers/header_filter_test.go +++ b/sdk/api/handlers/header_filter_test.go @@ -15,6 +15,8 @@ func TestFilterUpstreamHeaders_RemovesConnectionScopedHeaders(t *testing.T) { src.Set("X-Hop-C", "c") src.Set("X-Request-Id", "req-1") src.Set("Set-Cookie", "session=secret") + src.Set("x-cpa-trace-id", "upstream-trace") + src.Set("Access-Control-Expose-Headers", "upstream-header") filtered := FilterUpstreamHeaders(src) if filtered == nil { @@ -33,6 +35,8 @@ func TestFilterUpstreamHeaders_RemovesConnectionScopedHeaders(t *testing.T) { "X-Hop-B", "X-Hop-C", "Set-Cookie", + "x-cpa-trace-id", + "Access-Control-Expose-Headers", } for _, key := range blockedHeaderKeys { value := filtered.Get(key) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index fedf4951..a91a12bf 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -2564,7 +2564,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, provider, routeModel) - publishSelectedAuthMetadata(opts.Metadata, auth.ID) + publishSelectedAuthMetadata(opts.Metadata, auth) tried[auth.ID] = struct{}{} execCtx := ctx @@ -2677,7 +2677,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, provider, routeModel) - publishSelectedAuthMetadata(opts.Metadata, auth.ID) + publishSelectedAuthMetadata(opts.Metadata, auth) tried[auth.ID] = struct{}{} execCtx := ctx @@ -2790,7 +2790,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, provider, routeModel) - publishSelectedAuthMetadata(opts.Metadata, auth.ID) + publishSelectedAuthMetadata(opts.Metadata, auth) tried[auth.ID] = struct{}{} execCtx := ctx @@ -3157,17 +3157,21 @@ func isFreeCodexAuth(auth *Auth) bool { return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") } -func publishSelectedAuthMetadata(meta map[string]any, authID string) { - if len(meta) == 0 { +func publishSelectedAuthMetadata(meta map[string]any, auth *Auth) { + if len(meta) == 0 || auth == nil { return } - authID = strings.TrimSpace(authID) - if authID == "" { - return + if authID := strings.TrimSpace(auth.ID); authID != "" { + meta[cliproxyexecutor.SelectedAuthMetadataKey] = authID + if callback, ok := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string)); ok && callback != nil { + callback(authID) + } } - meta[cliproxyexecutor.SelectedAuthMetadataKey] = authID - if callback, ok := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string)); ok && callback != nil { - callback(authID) + if authIndex := strings.TrimSpace(auth.EnsureIndex()); authIndex != "" { + meta[cliproxyexecutor.SelectedAuthIndexMetadataKey] = authIndex + if callback, ok := meta[cliproxyexecutor.SelectedAuthIndexCallbackMetadataKey].(func(string)); ok && callback != nil { + callback(authIndex) + } } } @@ -5486,7 +5490,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy continue } c.auth = preparedAuth - publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID) + publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) if len(models) == 0 { continue @@ -5534,7 +5538,7 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl continue } c.auth = preparedAuth - publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID) + publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) if len(models) == 0 { continue diff --git a/sdk/cliproxy/auth/selected_auth_metadata_test.go b/sdk/cliproxy/auth/selected_auth_metadata_test.go new file mode 100644 index 00000000..2a7433e4 --- /dev/null +++ b/sdk/cliproxy/auth/selected_auth_metadata_test.go @@ -0,0 +1,40 @@ +package auth + +import ( + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestPublishSelectedAuthMetadataIncludesStableIndex(t *testing.T) { + auth := &Auth{ + ID: "auth-1", + Provider: "codex", + FileName: "auth-1.json", + } + selectedAuthID := "" + selectedAuthIndex := "" + meta := map[string]any{ + cliproxyexecutor.SelectedAuthCallbackMetadataKey: func(authID string) { + selectedAuthID = authID + }, + cliproxyexecutor.SelectedAuthIndexCallbackMetadataKey: func(authIndex string) { + selectedAuthIndex = authIndex + }, + } + + publishSelectedAuthMetadata(meta, auth) + + if selectedAuthID != auth.ID { + t.Fatalf("selected auth ID = %q, want %q", selectedAuthID, auth.ID) + } + if selectedAuthIndex == "" || selectedAuthIndex != auth.Index { + t.Fatalf("selected auth index = %q, want %q", selectedAuthIndex, auth.Index) + } + if got := meta[cliproxyexecutor.SelectedAuthMetadataKey]; got != auth.ID { + t.Fatalf("selected auth metadata = %#v, want %q", got, auth.ID) + } + if got := meta[cliproxyexecutor.SelectedAuthIndexMetadataKey]; got != auth.Index { + t.Fatalf("selected auth index metadata = %#v, want %q", got, auth.Index) + } +} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index d114e33d..e779f30d 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -38,6 +38,10 @@ const ( SelectedAuthMetadataKey = "selected_auth_id" // SelectedAuthCallbackMetadataKey carries an optional callback invoked with the selected auth ID. SelectedAuthCallbackMetadataKey = "selected_auth_callback" + // SelectedAuthIndexMetadataKey stores the stable index of the auth selected by the scheduler. + SelectedAuthIndexMetadataKey = "selected_auth_index" + // SelectedAuthIndexCallbackMetadataKey carries an optional callback invoked with the selected auth index. + SelectedAuthIndexCallbackMetadataKey = "selected_auth_index_callback" // ExecutionSessionMetadataKey identifies a long-lived downstream execution session. ExecutionSessionMetadataKey = "execution_session_id" )