diff --git a/config.example.yaml b/config.example.yaml index 2acab753c82..baaafcaca17 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -185,9 +185,13 @@ routing: # CPA preserves the retained credential and concurrency slot instead of redispatching each turn. # For OAuth/file credentials, add a top-level numeric "weight" field to the auth JSON. # Enable universal session-sticky routing for all clients. - # Explicit Claude Code, Codex, OpenCode, and pi session headers are preferred, - # followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, - # execution or derived session identity, and the existing first-message hash fallback. + # Explicit Claude Code, Codex, OpenCode, pi, and generic (X-Session-ID) session headers + # are preferred, followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, + # execution session identity, Merkle LCP prefix/compaction matching, derived session + # identity, and the existing first-message hash fallback. + # LCP compaction matching requires at least 2 shared trailing turns; clients whose + # compaction replaces all earlier turns with a summary should send X-Session-ID (or + # another explicit session ID), or affinity starts a new binding at every compaction. # Automatic failover is always enabled when bound auth becomes unavailable. # An established binding outranks credential priority: once a session is bound, that # credential is kept even if a higher-priority credential recovers. Credential priority @@ -295,7 +299,7 @@ requests: # default: # Default rules only set parameters when they are missing in the payload. # - models: # - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*") - # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity + # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity, meta # from-protocol: "responses" # restricts the rule to the source protocol, options: openai, responses, gemini, claude # headers: # all configured request headers must match; values support "*" wildcards # X-Client-Tier: "tenant-*-region-*" @@ -309,10 +313,15 @@ requests: # - "metadata.disable_payload" # params: # JSON path (gjson/sjson syntax) -> value # "generationConfig.thinkingConfig.thinkingBudget": 32768 + # - models: + # - name: "muse-*" + # protocol: "meta" + # params: + # prompt_cache_retention: "24h" # opt into Meta's 24h prompt cache when the client omits prompt_cache_retention # default-raw: # Default raw rules set parameters using raw JSON when missing (must be valid JSON). # - models: # - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*") - # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity + # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity, meta # params: # JSON path (gjson/sjson syntax) -> raw JSON value (strings are used as-is, must be valid JSON) # "generationConfig.responseJsonSchema": "{\"type\":\"object\",\"properties\":{\"answer\":{\"type\":\"string\"}}}" # override: # Override rules always set parameters, overwriting any existing values. @@ -325,7 +334,7 @@ requests: # service_tier: priority # - models: # - name: "gpt-*" # Supports wildcards (e.g., "gpt-*") - # protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity + # protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity, meta # params: # JSON path (gjson/sjson syntax) -> value # "reasoning.effort": "high" # - models: @@ -336,13 +345,13 @@ requests: # override-raw: # Override raw rules always set parameters using raw JSON (must be valid JSON). # - models: # - name: "gpt-*" # Supports wildcards (e.g., "gpt-*") - # protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity + # protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity, meta # params: # JSON path (gjson/sjson syntax) -> raw JSON value (strings are used as-is, must be valid JSON) # "response_format": "{\"type\":\"json_schema\",\"json_schema\":{\"name\":\"answer\",\"schema\":{\"type\":\"object\"}}}" # filter: # Filter rules remove specified parameters from the payload. # - models: # - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*") - # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity + # protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity, meta # params: # JSON paths (gjson/sjson syntax) to remove from the payload # - "generationConfig.thinkingConfig.thinkingBudget" # - "generationConfig.responseJsonSchema" @@ -488,6 +497,16 @@ upstream: # Default false keeps the translation and catalog behavior that predate tool_search support. tool-search: false + # When true, inject Anthropic server-side refusal fallbacks on cloaked Claude requests + # when the request body does not already specify fallbacks (claude-opus-5-5 -> claude-opus-4-8, + # claude-fable-5-1 -> claude-opus-5). Default false never opts requests into cross-model + # fallback unless the caller or requests.payload rules explicitly provide fallbacks. + # When fallback occurs, translated OpenAI Chat Completions and Responses payloads report + # the actual serving model and attach structured fallback metadata (choices[0].delta.fallback / + # choices[0].message.fallback and fallbacks for Chat; response.fallback and response.fallbacks + # for Responses). + server-side-fallback: false + # Days to retain the cloaked currentDate pin per session in in-memory mode # (server.kv-store.disabled: true without Home), so an idle gap or a credential # switch cannot re-anchor the date and invalidate earlier thinking signatures and the @@ -744,6 +763,7 @@ upstream: # max-context-length: 1048576 # optional: override Codex client context window metadata # force-mapping: true # optional: rewrite response model fields back to the alias # # Preserve thinking blocks with empty signatures for compatible upstreams. +# # Also handle compaction_trigger requests with local summary compaction instead of forwarding them upstream. # # client.codex.optimize-multi-agent-v2 applies to both OAuth and API-key credentials. # is-compat: false # # Optional: enable configuration_update for this API-key model; omitted defaults to false. @@ -807,7 +827,7 @@ upstream: # rebuild-mid-system-message: false # optional: default is false; when true, move messages with role "system" into the top-level Claude system field # native-tool-search: false # optional: default is false; when true, trust this key's custom base-url with Anthropic native tool search (defer_loading + tool_reference) for Codex tool_search; requires upstream.claude.tool-search: true # cloak: # optional: explicitly enable request cloaking for non-Claude-Code clients -# mode: "auto" # "auto" (default inside this block): cloak only when client is not Claude Code +# mode: "auto" # omitted inherits shared default; explicit "auto": cloak only when client is not Claude Code # # "always": cloak every unconfirmed client; confirmed native Claude Code still passes through # # "never" / "uncloak": never apply cloaking; on non-CLI API keys, also strip incoming Claude Code identity/billing attribution # # This "cloak" block applies to this api-keys.claude key entry only. For Claude OAuth @@ -820,6 +840,8 @@ upstream: # sensitive-words: # optional: words to obfuscate with zero-width characters # - "API" # - "proxy" +# sensitive-words-replace: # optional: literal sensitive-word replacements ("target": "replacement") +# "Claude Code": "Claude" # cache-user-id: true # optional: default is false; set true to reuse cached user_id per API key instead of generating a random one each request # # Every custom tool on a cloaked OAuth request automatically uses a caller-stable opaque mcp____ alias. # @@ -1192,9 +1214,11 @@ oauth: # Antigravity provider behavior. antigravity: - # sensitive-words: # optional: words to obfuscate with zero-width characters in system instructions + # sensitive-words: # optional: words to obfuscate with combining grapheme joiners (U+034F) in system instructions # - "API" # - "proxy" + # sensitive-words-replace: # optional: literal sensitive-word replacements ("target": "replacement") in system instructions + # "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant." # connection-pool: # optional: upstream HTTP connection pool settings # enabled: false # whether to enable connection pooling (default: false) # idle-conn-timeout: "30s" # idle keep-alive timeout (default: 30s, capped at 210s) @@ -1214,9 +1238,11 @@ oauth: # Devin provider behavior. # devin: - # sensitive-words: # optional: words to obfuscate with zero-width characters in system instructions and prompts + # sensitive-words: # optional: words to remove matching lines in system prompts # - "API" # - "proxy" + # sensitive-words-replace: # optional: literal sensitive-word replacements ("target": "replacement") in system prompts + # "Claude Code": "Claude" # Image and video behavior. multimedia: diff --git a/internal/api/handlers/management/auth_files_cooldown_test.go b/internal/api/handlers/management/auth_files_cooldown_test.go index b8f15ecb1e7..da85dce0eb4 100644 --- a/internal/api/handlers/management/auth_files_cooldown_test.go +++ b/internal/api/handlers/management/auth_files_cooldown_test.go @@ -1237,8 +1237,14 @@ func (s *targetedRecoveryTestCodexOAuthService) CreateTokenStorage(bundle *codex func pollOAuthStatusUntilDone(t *testing.T, router http.Handler, state string) map[string]any { t.Helper() - deadline := time.Now().Add(5 * time.Second) - for time.Now().Before(deadline) { + // The OAuth flow polls for its callback file every 500ms, so a slow CI host + // needs a generous deadline. Pause between polls so this loop does not + // starve the flow goroutine it is waiting on. + deadline := time.Now().Add(15 * time.Second) + for first := true; time.Now().Before(deadline); first = false { + if !first { + time.Sleep(10 * time.Millisecond) + } w := httptest.NewRecorder() router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v8/management/oauth/status?state="+state, nil)) if w.Code != http.StatusOK { diff --git a/internal/api/handlers/management/config_v8_test.go b/internal/api/handlers/management/config_v8_test.go index 95741ac20b0..9943e213db0 100644 --- a/internal/api/handlers/management/config_v8_test.go +++ b/internal/api/handlers/management/config_v8_test.go @@ -975,3 +975,52 @@ api-keys: } } } + +func TestConfigV8_SensitiveWordsReplace(t *testing.T) { + gin.SetMode(gin.TestMode) + path := filepath.Join(t.TempDir(), "config.yaml") + if err := os.WriteFile(path, []byte("server: {port: 8317}\n"), 0600); err != nil { + t.Fatal(err) + } + cfg, err := config.LoadConfig(path) + if err != nil { + t.Fatal(err) + } + h := &Handler{cfg: cfg, configFilePath: path} + router := gin.New() + router.GET("/v8/management/config/*path", h.ConfigV8) + router.PUT("/v8/management/config/*path", h.ConfigV8) + router.PATCH("/v8/management/config/*path", h.ConfigV8) + router.DELETE("/v8/management/config/*path", h.ConfigV8) + + for _, endpoint := range []string{ + "/v8/management/config/oauth/providers/antigravity/sensitive-words-replace", + "/v8/management/config/oauth/providers/devin/sensitive-words-replace", + } { + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, endpoint, strings.NewReader(`{"You are Codex, an agent based on GPT-5.":"You are a helpful coding assistant."}`))) + if w.Code != http.StatusOK { + t.Fatalf("PUT %s: status=%d body=%s", endpoint, w.Code, w.Body.String()) + } + getRec := httptest.NewRecorder() + router.ServeHTTP(getRec, httptest.NewRequest(http.MethodGet, endpoint, nil)) + if getRec.Code != http.StatusOK || !strings.Contains(getRec.Body.String(), "You are a helpful coding assistant.") { + t.Fatalf("GET %s: status=%d body=%s", endpoint, getRec.Code, getRec.Body.String()) + } + delRec := httptest.NewRecorder() + router.ServeHTTP(delRec, httptest.NewRequest(http.MethodDelete, endpoint, nil)) + if delRec.Code != http.StatusOK { + t.Fatalf("DELETE %s: status=%d body=%s", endpoint, delRec.Code, delRec.Body.String()) + } + } + + claudePut := httptest.NewRecorder() + router.ServeHTTP(claudePut, httptest.NewRequest( + http.MethodPut, + "/v8/management/config/api-keys/claude", + strings.NewReader(`[{"name":"claude-1","keys":[{"api-key":"sk-ant-1","cloak":{"sensitive-words-replace":{"Claude Code":"Claude"}}}] }]`), + )) + if claudePut.Code != http.StatusOK { + t.Fatalf("PUT api-keys/claude: status=%d body=%s", claudePut.Code, claudePut.Body.String()) + } +} diff --git a/internal/api/middleware/request_logging.go b/internal/api/middleware/request_logging.go index 33c01abd8e8..9467000d0d2 100644 --- a/internal/api/middleware/request_logging.go +++ b/internal/api/middleware/request_logging.go @@ -66,6 +66,10 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc { wrapper.ginContext = c if !loggerEnabled { wrapper.logOnErrorOnly = true + deferredAPICapture := logging.NewDeferredAPICapture() + wrapper.deferredAPICapture = deferredAPICapture + c.Request = c.Request.WithContext(logging.WithDeferredAPICapture(c.Request.Context(), deferredAPICapture)) + c.Set(logging.DeferredAPIResponseContextKey, deferredAPICapture) } c.Writer = wrapper attachRequestLogSources(c, logger, loggerEnabled && hasMatchedRoute) diff --git a/internal/api/middleware/request_logging_test.go b/internal/api/middleware/request_logging_test.go index a0d20aeba76..f328c3c1a81 100644 --- a/internal/api/middleware/request_logging_test.go +++ b/internal/api/middleware/request_logging_test.go @@ -13,6 +13,7 @@ import ( "path/filepath" "runtime" "strings" + "sync" "testing" "time" @@ -24,6 +25,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" ) func TestShouldSkipMethodForRequestLogging(t *testing.T) { @@ -1175,3 +1177,446 @@ func TestRequestLoggingMiddleware_Unrouted404DoesNotEvictErrorDumpRing(t *testin } }) } + +func TestRequestLoggingMiddlewareErrorLogKeepsUpstreamResponseWhenRequestLogDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + + readSingleErrorLog := func(t *testing.T, logsDir string) string { + t.Helper() + entries, errReadDir := os.ReadDir(logsDir) + if errReadDir != nil { + t.Fatalf("read logs dir: %v", errReadDir) + } + var errorLogPaths []string + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") { + errorLogPaths = append(errorLogPaths, filepath.Join(logsDir, entry.Name())) + } + } + if len(errorLogPaths) != 1 { + t.Fatalf("error log count = %d, want 1 (entries=%v)", len(errorLogPaths), entries) + } + content, errReadFile := os.ReadFile(errorLogPaths[0]) + if errReadFile != nil { + t.Fatalf("read error log: %v", errReadFile) + } + return string(content) + } + + t.Run("non-stream WriteErrorResponse keeps per-attempt upstream status masked headers body and error", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + baseHandler := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, nil) + largeAttemptTwoBody := "data: {\"type\":\"response.created\",\"id\":\"resp_2\"}\n\n" + strings.Repeat("x", 70*1024) + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/chat/completions", func(c *gin.Context) { + executorCtx := context.WithValue(context.Background(), "gin", c) + + helps.RecordAPIRequest(executorCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://cli-chat-proxy.grok.com/v1/responses", + Method: http.MethodPost, + Headers: http.Header{"Authorization": []string{"Bearer super-secret-request-token"}}, + Body: []byte(`{"model":"grok-4.7","stream":true}`), + }) + helps.RecordAPIResponseMetadata(executorCtx, cfg, http.StatusBadGateway, http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Authorization": []string{"Bearer super-secret-response-token"}, + "X-Request-Id": []string{"xai-req-1"}, + }) + helps.AppendAPIResponseChunk(executorCtx, cfg, []byte("event: response.created")) + helps.AppendAPIResponseChunk(executorCtx, cfg, []byte(`data: {"type":"response.created","id":"resp_1"}`)) + helps.RecordAPIResponseError(executorCtx, cfg, errors.New("unexpected EOF")) + + helps.RecordAPIRequest(executorCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://cli-chat-proxy.grok.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"model":"grok-4.7","stream":true,"attempt":2}`), + }) + helps.RecordAPIResponseMetadata(executorCtx, cfg, http.StatusOK, http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Set-Cookie": []string{"session=super-secret-cookie-value"}, + }) + helps.AppendAPIResponseChunk(executorCtx, cfg, []byte(largeAttemptTwoBody)) + + baseHandler.WriteErrorResponse(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusRequestTimeout, + Error: errors.New("xai stream error: stream disconnected before response.completed or response.incomplete"), + }) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"grok-4.7","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + if resp.Code != http.StatusRequestTimeout { + t.Fatalf("response status = %d, want %d", resp.Code, http.StatusRequestTimeout) + } + + logText := readSingleErrorLog(t, logsDir) + apiRequestIdx := strings.Index(logText, "=== API REQUEST 1 ===") + apiResponseIdx := strings.Index(logText, "=== API RESPONSE 1 ===") + downstreamResponseIdx := strings.Index(logText, "=== RESPONSE ===") + if apiRequestIdx == -1 || apiResponseIdx == -1 || downstreamResponseIdx == -1 { + t.Fatalf("missing expected log sections (req=%d, apiResp=%d, resp=%d):\n%s", apiRequestIdx, apiResponseIdx, downstreamResponseIdx, logText) + } + if apiRequestIdx >= apiResponseIdx || apiResponseIdx >= downstreamResponseIdx { + t.Fatalf("unexpected section order (req=%d, apiResp=%d, resp=%d):\n%s", apiRequestIdx, apiResponseIdx, downstreamResponseIdx, logText) + } + + apiResponseSection := logText[apiResponseIdx:downstreamResponseIdx] + downstreamSection := logText[downstreamResponseIdx:] + for _, want := range []string{ + "=== API RESPONSE 1 ===", + "Status: 502", + "Authorization: Bearer supe...oken", + "X-Request-Id: xai-req-1", + "event: response.created\ndata: {\"type\":\"response.created\",\"id\":\"resp_1\"}", + "Error: unexpected EOF", + "=== API RESPONSE 2 ===", + "Status: 200", + "data: {\"type\":\"response.created\",\"id\":\"resp_2\"}", + "[API RESPONSE BODY TRUNCATED: captured first ", + } { + if !strings.Contains(apiResponseSection, want) { + t.Fatalf("API RESPONSE section missing %q:\n%s", want, apiResponseSection) + } + } + for _, secret := range []string{"super-secret-response-token", "super-secret-cookie-value"} { + if strings.Contains(logText, secret) { + t.Fatalf("error log leaked sensitive header value %q:\n%s", secret, logText) + } + } + if strings.Contains(apiResponseSection, `"code":"request_timeout"`) { + t.Fatalf("API RESPONSE section should contain upstream attempts rather than proxy synthesized error JSON:\n%s", apiResponseSection) + } + if !strings.Contains(downstreamSection, "Status: 408") || !strings.Contains(downstreamSection, "xai stream error: stream disconnected before response.completed or response.incomplete") { + t.Fatalf("downstream RESPONSE section missing final proxy error:\n%s", downstreamSection) + } + }) + + t.Run("mid-stream SSE error keeps upstream response in error log", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + executorCtx := context.WithValue(context.Background(), "gin", c) + helps.RecordAPIRequest(executorCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://cli-chat-proxy.grok.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"model":"grok-4.7","stream":true}`), + }) + helps.RecordAPIResponseMetadata(executorCtx, cfg, http.StatusOK, http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"xai-stream-1"}, + }) + helps.AppendAPIResponseChunk(executorCtx, cfg, []byte(`data: {"type":"response.failed","response":{"error":{"message":"upstream quota exceeded"}}}`)) + helps.RecordAPIResponseError(executorCtx, cfg, errors.New("upstream quota exceeded")) + + c.Writer.WriteHeader(http.StatusOK) + _, _ = c.Writer.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")) + c.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{ + { + StatusCode: http.StatusTooManyRequests, + Error: errors.New("upstream quota exceeded"), + }, + }) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"grok-4.7","stream":true}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + logText := readSingleErrorLog(t, logsDir) + for _, want := range []string{ + "=== API REQUEST 1 ===", + "=== API RESPONSE 1 ===", + "Status: 200", + "X-Request-Id: xai-stream-1", + `data: {"type":"response.failed","response":{"error":{"message":"upstream quota exceeded"}}}`, + "Error: upstream quota exceeded", + } { + if !strings.Contains(logText, want) { + t.Fatalf("mid-stream error log missing %q:\n%s", want, logText) + } + } + }) + + t.Run("rejected upstream websocket upgrade keeps handshake request and response in error log", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/responses", func(c *gin.Context) { + executorCtx := context.WithValue(context.Background(), "gin", c) + helps.RecordAPIWebsocketUpgradeRejection( + executorCtx, + cfg, + helps.UpstreamRequestLog{ + URL: "https://chatgpt.com/backend-api/codex/responses", + Method: http.MethodGet, + Headers: http.Header{"Authorization": []string{"Bearer super-secret-ws-token"}}, + }, + http.StatusForbidden, + http.Header{ + "Content-Type": []string{"application/json"}, + "Set-Cookie": []string{"cf_clearance=super-secret-cf-cookie"}, + }, + []byte(`{"error":{"message":"websocket handshake rejected"}}`), + ) + c.JSON(http.StatusForbidden, gin.H{"error": "websocket handshake rejected"}) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5-codex","input":"hi"}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + logText := readSingleErrorLog(t, logsDir) + for _, want := range []string{ + "=== API REQUEST 1 ===", + "Upstream URL: https://chatgpt.com/backend-api/codex/responses", + "=== API RESPONSE 1 ===", + "Status: 403", + `{"error":{"message":"websocket handshake rejected"}}`, + } { + if !strings.Contains(logText, want) { + t.Fatalf("websocket rejection error log missing %q:\n%s", want, logText) + } + } + for _, secret := range []string{"super-secret-ws-token", "super-secret-cf-cookie"} { + if strings.Contains(logText, secret) { + t.Fatalf("websocket rejection error log leaked secret %q:\n%s", secret, logText) + } + } + }) + + t.Run("successful request with request-log off writes no log file", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/chat/completions", func(c *gin.Context) { + executorCtx := context.WithValue(context.Background(), "gin", c) + helps.RecordAPIRequest(executorCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://cli-chat-proxy.grok.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"model":"grok-4.7"}`), + }) + helps.RecordAPIResponseMetadata(executorCtx, cfg, http.StatusOK, http.Header{ + "Content-Type": []string{"text/event-stream"}, + }) + helps.AppendAPIResponseChunk(executorCtx, cfg, []byte(`data: {"type":"response.completed"}`)) + c.JSON(http.StatusOK, gin.H{"id": "chatcmpl_1"}) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"grok-4.7"}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + + if resp.Code != http.StatusOK { + t.Fatalf("response status = %d, want %d", resp.Code, http.StatusOK) + } + entries, errReadDir := os.ReadDir(logsDir) + if errReadDir != nil && !os.IsNotExist(errReadDir) { + t.Fatalf("read logs dir: %v", errReadDir) + } + if len(entries) != 0 { + t.Fatalf("expected 0 log files for successful request with request-log=false, got: %v", entries) + } + }) +} + +func TestRequestLoggingMiddlewareIsolatesDeferredResponseAcrossStreamCancelAndGinContextReuse(t *testing.T) { + gin.SetMode(gin.TestMode) + + readSingleErrorLog := func(t *testing.T, logsDir string) string { + t.Helper() + entries, errReadDir := os.ReadDir(logsDir) + if errReadDir != nil { + t.Fatalf("read logs dir: %v", errReadDir) + } + var errorLogPaths []string + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") { + errorLogPaths = append(errorLogPaths, filepath.Join(logsDir, entry.Name())) + } + } + if len(errorLogPaths) != 1 { + t.Fatalf("error log count = %d, want 1 (entries=%v)", len(errorLogPaths), entries) + } + content, errReadFile := os.ReadFile(errorLogPaths[0]) + if errReadFile != nil { + t.Fatalf("read error log: %v", errReadFile) + } + return string(content) + } + + t.Run("late stream producer after cancel does not pollute reused Gin context", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + baseHandler := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, nil) + + const cancelledTurns = 8 + turn2Active := make(chan struct{}) + var producers sync.WaitGroup + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/chat/completions", func(c *gin.Context) { + if c.GetHeader("X-Turn") == "2" { + cliCtx, cliCancel := baseHandler.GetContextWithCancel(nil, c, context.Background()) + defer cliCancel() + helps.RecordAPIRequest(cliCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://api.example.com/turn2", + Method: http.MethodPost, + Body: []byte(`{"turn":2}`), + }) + helps.RecordAPIResponseMetadata(cliCtx, cfg, http.StatusBadGateway, http.Header{ + "Content-Type": []string{"application/json"}, + }) + helps.AppendAPIResponseChunk(cliCtx, cfg, []byte(`{"error":"turn2-upstream-failure"}`)) + close(turn2Active) + producers.Wait() + baseHandler.WriteErrorResponse(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadGateway, + Error: errors.New("turn2 failed"), + }) + return + } + + cliCtx, cliCancel := baseHandler.GetContextWithCancel(nil, c, context.Background()) + helps.RecordAPIRequest(cliCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://api.example.com/turn1", + Method: http.MethodPost, + Body: []byte(`{"turn":1}`), + }) + helps.RecordAPIResponseMetadata(cliCtx, cfg, http.StatusOK, http.Header{ + "Content-Type": []string{"text/event-stream"}, + }) + helps.AppendAPIResponseChunk(cliCtx, cfg, []byte(`data: {"turn":1,"early":true}`)) + producerStarted := make(chan struct{}) + producers.Add(1) + go func(streamCtx context.Context) { + defer producers.Done() + close(producerStarted) + <-turn2Active + helps.RecordAPIRequest(streamCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://api.example.com/turn1-late-req", + Method: http.MethodPost, + Body: []byte(`{"turn":1,"late":true}`), + }) + helps.RecordAPIResponseMetadata(streamCtx, cfg, http.StatusInternalServerError, http.Header{ + "X-Turn1-Late": []string{"leaked"}, + }) + helps.AppendAPIResponseChunk(streamCtx, cfg, []byte("turn1-late-stream-chunk")) + helps.RecordAPIResponseError(streamCtx, cfg, errors.New("turn1-late-stream-error")) + }(cliCtx) + <-producerStarted + cliCancel() + c.Status(clienterror.StatusClientClosedRequest) + }) + + for range cancelledTurns { + req1 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"stream":true}`)) + req1.Header.Set("Content-Type", "application/json") + req1.Header.Set("X-Turn", "1") + router.ServeHTTP(httptest.NewRecorder(), req1) + } + + req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"stream":false}`)) + req2.Header.Set("Content-Type", "application/json") + req2.Header.Set("X-Turn", "2") + resp2 := httptest.NewRecorder() + router.ServeHTTP(resp2, req2) + if resp2.Code != http.StatusBadGateway { + t.Fatalf("turn 2 status = %d, want %d", resp2.Code, http.StatusBadGateway) + } + + logText := readSingleErrorLog(t, logsDir) + for _, want := range []string{ + "https://api.example.com/turn2", + `{"error":"turn2-upstream-failure"}`, + } { + if !strings.Contains(logText, want) { + t.Fatalf("turn 2 error log missing %q:\n%s", want, logText) + } + } + for _, leaked := range []string{ + "turn1-late-req", + "X-Turn1-Late", + "turn1-late-stream-chunk", + "turn1-late-stream-error", + } { + if strings.Contains(logText, leaked) { + t.Fatalf("turn 2 error log leaked stale turn 1 data %q:\n%s", leaked, logText) + } + } + }) + + t.Run("concurrent stream producer during Finalize is race-free", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 10) + cfg := &config.Config{} + baseHandler := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, nil) + + producerStarted := make(chan struct{}) + producerDone := make(chan struct{}) + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/chat/completions", func(c *gin.Context) { + cliCtx, cliCancel := baseHandler.GetContextWithCancel(nil, c, context.Background()) + defer cliCancel() + helps.RecordAPIRequest(cliCtx, cfg, helps.UpstreamRequestLog{ + URL: "https://api.example.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"stream":true}`), + }) + helps.RecordAPIResponseMetadata(cliCtx, cfg, http.StatusOK, http.Header{ + "Content-Type": []string{"text/event-stream"}, + }) + go func(streamCtx context.Context) { + defer close(producerDone) + close(producerStarted) + for range 32 { + helps.AppendAPIResponseChunk(streamCtx, cfg, []byte(`data: {"delta":"chunk"}`)) + } + helps.RecordAPIResponseError(streamCtx, cfg, errors.New("stream aborted")) + }(cliCtx) + <-producerStarted + c.Header("Content-Type", "text/event-stream") + c.Writer.WriteHeader(http.StatusOK) + c.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{ + { + StatusCode: http.StatusBadGateway, + Error: errors.New("mid-stream abort"), + }, + }) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"stream":true}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(httptest.NewRecorder(), req) + <-producerDone + + logText := readSingleErrorLog(t, logsDir) + if !strings.Contains(logText, "=== API RESPONSE 1 ===") { + t.Fatalf("concurrent error log missing API RESPONSE 1:\n%s", logText) + } + }) +} diff --git a/internal/api/middleware/response_writer.go b/internal/api/middleware/response_writer.go index 57323930e1e..c540cd9e011 100644 --- a/internal/api/middleware/response_writer.go +++ b/internal/api/middleware/response_writer.go @@ -40,21 +40,22 @@ type RequestInfo struct { // It is designed to handle both standard and streaming responses, ensuring that logging operations do not block the client response. type ResponseWriterWrapper struct { gin.ResponseWriter - body *bytes.Buffer // body is a buffer to store the response body for non-streaming responses. - isStreaming bool // isStreaming indicates whether the response is a streaming type (e.g., text/event-stream). - streamWriter logging.StreamingLogWriter // streamWriter is a writer for handling streaming log entries. - chunkChannel chan []byte // chunkChannel is a channel for asynchronously passing response chunks to the logger. - streamDone chan struct{} // streamDone signals when the streaming goroutine completes. - logger logging.RequestLogger // logger is the instance of the request logger service. - requestInfo *RequestInfo // requestInfo holds the details of the original request. - ginContext *gin.Context // ginContext holds the Gin request context for mid-stream error inspection. - statusCode int // statusCode stores the HTTP status code of the response. - headers map[string][]string // headers stores the response headers. - logOnErrorOnly bool // logOnErrorOnly enables logging only when an error response is detected. - hijacked bool // hijacked reports whether the underlying connection was hijacked. - panicked bool // panicked records whether a handler panicked (even after committing a 2xx response). - finalized bool // finalized prevents post-Finalize writes (e.g. outer recovery) from reopening stream writers. - firstChunkTimestamp time.Time // firstChunkTimestamp captures TTFB for streaming responses. + body *bytes.Buffer // body is a buffer to store the response body for non-streaming responses. + isStreaming bool // isStreaming indicates whether the response is a streaming type (e.g., text/event-stream). + streamWriter logging.StreamingLogWriter // streamWriter is a writer for handling streaming log entries. + chunkChannel chan []byte // chunkChannel is a channel for asynchronously passing response chunks to the logger. + streamDone chan struct{} // streamDone signals when the streaming goroutine completes. + logger logging.RequestLogger // logger is the instance of the request logger service. + requestInfo *RequestInfo // requestInfo holds the details of the original request. + ginContext *gin.Context // ginContext holds the Gin request context for mid-stream error inspection. + statusCode int // statusCode stores the HTTP status code of the response. + headers map[string][]string // headers stores the response headers. + logOnErrorOnly bool // logOnErrorOnly enables logging only when an error response is detected. + deferredAPICapture *logging.DeferredAPICapture // deferredAPICapture retains bounded upstream request/response attempts for error logs. + hijacked bool // hijacked reports whether the underlying connection was hijacked. + panicked bool // panicked records whether a handler panicked (even after committing a 2xx response). + finalized bool // finalized prevents post-Finalize writes (e.g. outer recovery) from reopening stream writers. + firstChunkTimestamp time.Time // firstChunkTimestamp captures TTFB for streaming responses. } // NewResponseWriterWrapper creates and initializes a new ResponseWriterWrapper. @@ -373,6 +374,7 @@ func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error { if w.requestInfo != nil && w.requestInfo.deferredBodyCapture != nil { defer w.requestInfo.deferredBodyCapture.Cleanup() } + defer w.closeDeferredAPICapture(c) if w.logger == nil { return nil } @@ -492,7 +494,13 @@ func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error { if forceLog && len(apiRequest) == 0 { apiRequest = w.extractDeferredAPIRequest(c) } - return w.logRequest(w.extractRequestBody(c), finalStatusCode, w.cloneHeaders(), w.extractResponseBody(c), w.extractWebsocketTimeline(c), websocketTimelineSource, apiRequest, apiRequestSource, w.extractAPIResponse(c), apiResponseSource, w.extractAPIWebsocketTimeline(c), apiWebsocketTimelineSource, w.extractAPIResponseTimestamp(c), slicesAPIResponseError, forceLog) + apiResponse := w.extractAPIResponse(c) + if forceLog { + if deferredResponse := w.extractDeferredAPIResponse(c); len(deferredResponse) > 0 { + apiResponse = deferredResponse + } + } + return w.logRequest(w.extractRequestBody(c), finalStatusCode, w.cloneHeaders(), w.extractResponseBody(c), w.extractWebsocketTimeline(c), websocketTimelineSource, apiRequest, apiRequestSource, apiResponse, apiResponseSource, w.extractAPIWebsocketTimeline(c), apiWebsocketTimelineSource, w.extractAPIResponseTimestamp(c), slicesAPIResponseError, forceLog) } func (w *ResponseWriterWrapper) cloneHeaders() map[string][]string { @@ -521,9 +529,17 @@ func (w *ResponseWriterWrapper) extractAPIRequest(c *gin.Context) []byte { } func (w *ResponseWriterWrapper) extractDeferredAPIRequest(c *gin.Context) []byte { + if w != nil && w.deferredAPICapture != nil && w.deferredAPICapture.HasRequests() { + return w.deferredAPICapture.RenderAPIRequests() + } if c == nil { return nil } + if value, exists := c.Get(logging.DeferredAPIResponseContextKey); exists { + if capture, ok := value.(*logging.DeferredAPICapture); ok && capture != nil && capture.HasRequests() { + return capture.RenderAPIRequests() + } + } value, exists := c.Get(logging.DeferredAPIRequestContextKey) if !exists { return nil @@ -542,6 +558,41 @@ func (w *ResponseWriterWrapper) extractDeferredAPIRequest(c *gin.Context) []byte return body.Bytes() } +func (w *ResponseWriterWrapper) extractDeferredAPIResponse(c *gin.Context) []byte { + if w != nil && w.deferredAPICapture != nil { + if !w.deferredAPICapture.HasPayload() { + return nil + } + return w.deferredAPICapture.RenderAPIResponse() + } + if c == nil { + return nil + } + value, exists := c.Get(logging.DeferredAPIResponseContextKey) + if !exists { + return nil + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil || !deferred.HasPayload() { + return nil + } + return deferred.RenderAPIResponse() +} + +func (w *ResponseWriterWrapper) closeDeferredAPICapture(c *gin.Context) { + if w != nil && w.deferredAPICapture != nil { + w.deferredAPICapture.Close() + } + if c == nil { + return + } + if value, exists := c.Get(logging.DeferredAPIResponseContextKey); exists { + if closer, ok := value.(interface{ Close() }); ok && closer != nil { + closer.Close() + } + } +} + func (w *ResponseWriterWrapper) extractAPIResponse(c *gin.Context) []byte { apiResponse, isExist := c.Get("API_RESPONSE") if !isExist { @@ -947,6 +998,14 @@ func hasRecordedUpstreamActivity(c *gin.Context) bool { return true } } + if value, exists := c.Get(logging.DeferredAPIResponseContextKey); exists { + if capture, ok := value.(*logging.DeferredAPICapture); ok && capture != nil && capture.HasRequests() { + return true + } + if response, ok := value.(logging.DeferredAPIResponse); ok && response != nil && response.HasPayload() { + return true + } + } for _, key := range []string{ logging.APIRequestSourceContextKey, logging.APIResponseSourceContextKey, diff --git a/internal/api/middleware/response_writer_test.go b/internal/api/middleware/response_writer_test.go index 5565a69ee84..0e987696bc5 100644 --- a/internal/api/middleware/response_writer_test.go +++ b/internal/api/middleware/response_writer_test.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -316,12 +317,14 @@ func TestHasActionableError(t *testing.T) { } type recordingRequestLogger struct { - loggedCalls []int - enabled bool + loggedCalls []int + apiResponses [][]byte + enabled bool } func (l *recordingRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { l.loggedCalls = append(l.loggedCalls, statusCode) + l.apiResponses = append(l.apiResponses, bytes.Clone(apiResponse)) return nil } @@ -336,10 +339,113 @@ func (l *recordingRequestLogger) IsEnabled() bool { func (l *recordingRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { if force || l.enabled { l.loggedCalls = append(l.loggedCalls, statusCode) + l.apiResponses = append(l.apiResponses, bytes.Clone(apiResponse)) } return nil } +type stubDeferredAPIResponse struct { + section []byte + renders int +} + +func (s *stubDeferredAPIResponse) HasPayload() bool { + return s != nil && len(s.section) > 0 +} + +func (s *stubDeferredAPIResponse) RenderAPIResponse() []byte { + if s == nil { + return nil + } + s.renders++ + return s.section +} + +func TestFinalizeForceLogUsesDeferredAPIResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("error log uses deferred upstream response instead of fallback proxy error", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + stub := &stubDeferredAPIResponse{ + section: []byte("=== API RESPONSE 1 ===\nTimestamp: 2026-11-11T00:45:05Z\n\nStatus: 502\n\nBody:\nupstream-error-body\n"), + } + c.Set(logging.DeferredAPIResponseContextKey, stub) + c.Set("API_RESPONSE", []byte(`{"error":{"message":"proxy synthesized error"}}`)) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + statusCode: http.StatusBadGateway, + requestInfo: &RequestInfo{URL: "/v1/chat/completions", Method: http.MethodPost, Timestamp: time.Now()}, + } + if errFinalize := wrapper.Finalize(c); errFinalize != nil { + t.Fatalf("Finalize error: %v", errFinalize) + } + if len(logger.apiResponses) != 1 { + t.Fatalf("logged apiResponses = %d, want 1", len(logger.apiResponses)) + } + got := string(logger.apiResponses[0]) + if !strings.Contains(got, "upstream-error-body") { + t.Fatalf("logged API response missing upstream body:\n%s", got) + } + if strings.Contains(got, "proxy synthesized error") { + t.Fatalf("logged API response should not include proxy synthesized fallback JSON:\n%s", got) + } + }) + + t.Run("local error without upstream attempt preserves fallback API_RESPONSE", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + c.Set(logging.DeferredAPIResponseContextKey, &stubDeferredAPIResponse{}) + c.Set("API_RESPONSE", []byte(`{"error":{"message":"local cooldown"}}`)) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + statusCode: http.StatusTooManyRequests, + requestInfo: &RequestInfo{URL: "/v1/chat/completions", Method: http.MethodPost, Timestamp: time.Now()}, + } + if errFinalize := wrapper.Finalize(c); errFinalize != nil { + t.Fatalf("Finalize error: %v", errFinalize) + } + if len(logger.apiResponses) != 1 || !strings.Contains(string(logger.apiResponses[0]), "local cooldown") { + t.Fatalf("logged API response = %q, want fallback local error", logger.apiResponses) + } + }) + + t.Run("successful request neither logs nor renders deferred response", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + stub := &stubDeferredAPIResponse{ + section: []byte("=== API RESPONSE 1 ===\n\nBody:\nupstream-ok\n"), + } + c.Set(logging.DeferredAPIResponseContextKey, stub) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + statusCode: http.StatusOK, + requestInfo: &RequestInfo{URL: "/v1/chat/completions", Method: http.MethodPost, Timestamp: time.Now()}, + } + if errFinalize := wrapper.Finalize(c); errFinalize != nil { + t.Fatalf("Finalize error: %v", errFinalize) + } + if len(logger.loggedCalls) != 0 || stub.renders != 0 { + t.Fatalf("loggedCalls=%d renders=%d, want 0 and 0", len(logger.loggedCalls), stub.renders) + } + }) +} + func TestFinalizeExcludes499FromForceLog(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -451,6 +557,30 @@ func TestFinalizeLogsUnrouted404WhenUpstreamEvidenceOrPanicExists(t *testing.T) } }) + t.Run("deferred api response recorded", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/files", nil) + c.Set(logging.DeferredAPIResponseContextKey, &stubDeferredAPIResponse{ + section: []byte("=== API RESPONSE 1 ===\n"), + }) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + statusCode: http.StatusNotFound, + requestInfo: &RequestInfo{URL: "/v1/files", Method: http.MethodPost, Timestamp: time.Now()}, + } + if errFinalize := wrapper.Finalize(c); errFinalize != nil { + t.Fatalf("Finalize error: %v", errFinalize) + } + if len(logger.loggedCalls) != 1 || logger.loggedCalls[0] != http.StatusNotFound { + t.Fatalf("expected 1 logged call when deferred API response exists, got: %v", logger.loggedCalls) + } + }) + t.Run("file body source with payload recorded", func(t *testing.T) { recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) diff --git a/internal/api/protocol_multiplexer_test.go b/internal/api/protocol_multiplexer_test.go index 6974c78e442..110d0946d64 100644 --- a/internal/api/protocol_multiplexer_test.go +++ b/internal/api/protocol_multiplexer_test.go @@ -316,7 +316,11 @@ func TestServerRejectsIncompleteUnauthenticatedRequestBodiesPromptly(t *testing. _ = conn.Close() }() - if errDeadline := conn.SetDeadline(time.Now().Add(500 * time.Millisecond)); errDeadline != nil { + // The server writes its request log before responding. Under heavy + // parallel test I/O, Windows can stall opening the request body temp + // file for several seconds, so allow a generous bound; a server that + // never closes the connection still fails. + if errDeadline := conn.SetDeadline(time.Now().Add(15 * time.Second)); errDeadline != nil { t.Fatalf("set deadline: %v", errDeadline) } if _, errWrite := conn.Write([]byte(tc.rawRequest)); errWrite != nil { diff --git a/internal/auth/claude/identity.go b/internal/auth/claude/identity.go index 3e4bde729c0..7cbf917d3a4 100644 --- a/internal/auth/claude/identity.go +++ b/internal/auth/claude/identity.go @@ -189,6 +189,173 @@ func ReadMetadataString(metadata *map[string]any, key string) string { return value } +// ReadMetadataStringMap reads a map[string]string metadata entry under the metadata +// lock, copying out string pairs from either map[string]string or map[string]any. +func ReadMetadataStringMap(metadata *map[string]any, key string) map[string]string { + if metadata == nil { + return nil + } + claudeDevicePoolMu.Lock() + defer claudeDevicePoolMu.Unlock() + + if *metadata == nil { + return nil + } + switch typed := (*metadata)[key].(type) { + case map[string]string: + if len(typed) == 0 { + return nil + } + out := make(map[string]string, len(typed)) + for k, v := range typed { + out[k] = v + } + return out + case map[string]any: + if len(typed) == 0 { + return nil + } + out := make(map[string]string, len(typed)) + for k, rawVal := range typed { + if strVal, ok := rawVal.(string); ok { + out[k] = strVal + } + } + if len(out) == 0 { + return nil + } + return out + default: + return nil + } +} + +func copyStringSliceFromAny(raw any) []string { + switch typed := raw.(type) { + case []string: + if len(typed) == 0 { + return nil + } + out := make([]string, 0, len(typed)) + for _, s := range typed { + if trimmed := strings.TrimSpace(s); trimmed != "" { + out = append(out, trimmed) + } + } + if len(out) == 0 { + return nil + } + return out + case []any: + if len(typed) == 0 { + return nil + } + out := make([]string, 0, len(typed)) + for _, elem := range typed { + if s, ok := elem.(string); ok { + if trimmed := strings.TrimSpace(s); trimmed != "" { + out = append(out, trimmed) + } + } + } + if len(out) == 0 { + return nil + } + return out + default: + return nil + } +} + +func copyStringMapFromAny(raw any) map[string]string { + switch typed := raw.(type) { + case map[string]string: + if len(typed) == 0 { + return nil + } + out := make(map[string]string, len(typed)) + for k, v := range typed { + out[k] = v + } + return out + case map[string]any: + if len(typed) == 0 { + return nil + } + out := make(map[string]string, len(typed)) + for k, rawVal := range typed { + if strVal, ok := rawVal.(string); ok { + out[k] = strVal + } + } + if len(out) == 0 { + return nil + } + return out + default: + return nil + } +} + +// ReadMetadataStringSlice reads a []string metadata entry under the metadata lock, +// copying string elements from either []string or []any. +func ReadMetadataStringSlice(metadata *map[string]any, key string) []string { + if metadata == nil { + return nil + } + claudeDevicePoolMu.Lock() + defer claudeDevicePoolMu.Unlock() + + if *metadata == nil { + return nil + } + return copyStringSliceFromAny((*metadata)[key]) +} + +// ReadMetadataHomeCloak extracts credential_options.cloak from metadata under +// the metadata lock without JSON serialization. +func ReadMetadataHomeCloak(metadata *map[string]any) ( + mode string, + strictMode bool, + sensitiveWords []string, + sensitiveWordsReplace map[string]string, + cacheUserID *bool, + present bool, +) { + if metadata == nil { + return "", false, nil, nil, nil, false + } + claudeDevicePoolMu.Lock() + defer claudeDevicePoolMu.Unlock() + + if *metadata == nil { + return "", false, nil, nil, nil, false + } + opts, ok := (*metadata)["credential_options"].(map[string]any) + if !ok || opts == nil { + return "", false, nil, nil, nil, false + } + rawCloak, exists := opts["cloak"] + if !exists || rawCloak == nil { + return "", false, nil, nil, nil, false + } + cloakMap, ok := rawCloak.(map[string]any) + if !ok { + return "", false, nil, nil, nil, false + } + mode, _ = cloakMap["mode"].(string) + strictMode, _ = cloakMap["strict-mode"].(bool) + sensitiveWords = copyStringSliceFromAny(cloakMap["sensitive-words"]) + sensitiveWordsReplace = copyStringMapFromAny(cloakMap["sensitive-words-replace"]) + if rawCache, hasCache := cloakMap["cache-user-id"]; hasCache { + if b, okBool := rawCache.(bool); okBool { + cacheCopy := b + cacheUserID = &cacheCopy + } + } + return mode, strictMode, sensitiveWords, sensitiveWordsReplace, cacheUserID, true +} + // StoreMetadataString writes a string-valued metadata entry under the metadata // lock, initializing the map when needed. Empty values are skipped so callers can // forward optional fields without erasing a previously resolved value. diff --git a/internal/config/claude_code_test.go b/internal/config/claude_code_test.go index ff40947028f..5b8e901609e 100644 --- a/internal/config/claude_code_test.go +++ b/internal/config/claude_code_test.go @@ -182,3 +182,28 @@ func TestParseConfigBytesClaudeToolSearch(t *testing.T) { }) } } + +func TestParseConfigBytesClaudeServerSideFallback(t *testing.T) { + tests := []struct { + name string + yaml string + want bool + }{ + {name: "defaults to false", yaml: "port: 8317\n", want: false}, + {name: "legacy layout", yaml: "claude:\n server-side-fallback: true\n", want: true}, + {name: "v8 layout", yaml: "config-version: 8\nupstream:\n claude:\n server-side-fallback: true\n", want: true}, + {name: "v8 explicit false", yaml: "config-version: 8\nupstream:\n claude:\n server-side-fallback: false\n", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg, errParse := ParseConfigBytes([]byte(tt.yaml)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + if got := cfg.Claude.ServerSideFallback; got != tt.want { + t.Fatalf("cfg.Claude.ServerSideFallback = %t, want %t", got, tt.want) + } + }) + } +} diff --git a/internal/config/cloak_save_test.go b/internal/config/cloak_save_test.go index 3f97d060b19..545ff9b00f8 100644 --- a/internal/config/cloak_save_test.go +++ b/internal/config/cloak_save_test.go @@ -295,3 +295,79 @@ claude-api-key: t.Errorf("saved YAML lost 'allow-remote: true'; got:\n%s", savedText) } } + +func TestSaveConfigPreserveComments_SensitiveWordsReplaceEmptyValueAndPruning(t *testing.T) { + for _, v8Layout := range []bool{false, true} { + name := "legacy" + initialYAML := "port: 8317\nclaude-api-key:\n - api-key: sk-ant-test\n" + if v8Layout { + name = "v8" + initialYAML = "config-version: 8\nserver:\n port: 8317\napi-keys:\n claude:\n - name: claude-1\n keys:\n - api-key: sk-ant-test\n" + } + t.Run(name, func(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + if errWrite := os.WriteFile(configPath, []byte(initialYAML), 0o600); errWrite != nil { + t.Fatalf("os.WriteFile() error = %v", errWrite) + } + cfg, errLoad := LoadConfigOptional(configPath, false) + if errLoad != nil { + t.Fatalf("initial LoadConfigOptional() error = %v", errLoad) + } + + // Step 1: Save empty-string deletion entries in all three scopes. + cfg.Antigravity.SensitiveWordsReplace = map[string]string{ + "You are Codex, an agent based on GPT-5.": "", + "TempAntigrav": "Keep", + } + cfg.Devin.SensitiveWordsReplace = map[string]string{ + "DropDevinPhrase": "", + } + cfg.ClaudeKey[0].Cloak = &CloakConfig{ + SensitiveWordsReplace: map[string]string{ + "DropClaudePhrase": "", + "OldClaudeTarget": "OldRepl", + }, + } + if errSave := SaveConfigPreserveComments(configPath, cfg); errSave != nil { + t.Fatalf("SaveConfigPreserveComments() step 1 error = %v", errSave) + } + reloaded, errReload := LoadConfigOptional(configPath, false) + if errReload != nil { + t.Fatalf("reload step 1 error = %v", errReload) + } + if val, ok := reloaded.Antigravity.SensitiveWordsReplace["You are Codex, an agent based on GPT-5."]; !ok || val != "" { + t.Fatalf("Antigravity empty replacement lost: %#v", reloaded.Antigravity.SensitiveWordsReplace) + } + if val, ok := reloaded.Devin.SensitiveWordsReplace["DropDevinPhrase"]; !ok || val != "" { + t.Fatalf("Devin empty replacement lost: %#v", reloaded.Devin.SensitiveWordsReplace) + } + if reloaded.ClaudeKey[0].Cloak == nil { + t.Fatal("Claude Cloak is nil after step 1") + } + if val, ok := reloaded.ClaudeKey[0].Cloak.SensitiveWordsReplace["DropClaudePhrase"]; !ok || val != "" { + t.Fatalf("Claude empty replacement lost: %#v", reloaded.ClaudeKey[0].Cloak.SensitiveWordsReplace) + } + + // Step 2: Prune individual keys and clear Devin replacements completely. + delete(reloaded.Antigravity.SensitiveWordsReplace, "TempAntigrav") + reloaded.Devin.SensitiveWordsReplace = nil + delete(reloaded.ClaudeKey[0].Cloak.SensitiveWordsReplace, "OldClaudeTarget") + if errSave2 := SaveConfigPreserveComments(configPath, reloaded); errSave2 != nil { + t.Fatalf("SaveConfigPreserveComments() step 2 error = %v", errSave2) + } + reloaded2, errReload2 := LoadConfigOptional(configPath, false) + if errReload2 != nil { + t.Fatalf("reload step 2 error = %v", errReload2) + } + if _, exists := reloaded2.Antigravity.SensitiveWordsReplace["TempAntigrav"]; exists { + t.Fatalf("deleted Antigravity key persisted: %#v", reloaded2.Antigravity.SensitiveWordsReplace) + } + if len(reloaded2.Devin.SensitiveWordsReplace) != 0 { + t.Fatalf("cleared Devin replacements persisted: %#v", reloaded2.Devin.SensitiveWordsReplace) + } + if _, exists := reloaded2.ClaudeKey[0].Cloak.SensitiveWordsReplace["OldClaudeTarget"]; exists { + t.Fatalf("deleted Claude replacement key persisted: %#v", reloaded2.ClaudeKey[0].Cloak.SensitiveWordsReplace) + } + }) + } +} diff --git a/internal/config/clone_test.go b/internal/config/clone_test.go index 45a94a72fea..9740c019c53 100644 --- a/internal/config/clone_test.go +++ b/internal/config/clone_test.go @@ -164,6 +164,14 @@ func sampleCloneRuntimeConfig() *Config { }, AntigravitySignatureCacheEnabled: &cacheStrict, AntigravitySignatureBypassStrict: &bypassStrict, + Antigravity: AntigravityConfig{ + SensitiveWords: []string{"Codex"}, + SensitiveWordsReplace: map[string]string{"You are Codex": "You are an assistant"}, + }, + Devin: DevinConfig{ + SensitiveWords: []string{"DevinWord"}, + SensitiveWordsReplace: map[string]string{"Claude Code": "Devin"}, + }, GeminiKey: []GeminiKey{{ APIKey: "gemini-key", Models: []GeminiModel{{Name: "gemini-upstream", Alias: "gemini-upstream-alias"}}, @@ -182,8 +190,9 @@ func sampleCloneRuntimeConfig() *Config { Headers: map[string]string{"X-Claude": "one"}, ExcludedModels: []string{"claude-hidden"}, Cloak: &CloakConfig{ - SensitiveWords: []string{"secret"}, - CacheUserID: &cacheUserID, + SensitiveWords: []string{"secret"}, + SensitiveWordsReplace: map[string]string{"Claude Code": "Claude"}, + CacheUserID: &cacheUserID, }, }}, OpenAICompatibility: []OpenAICompatibility{{ @@ -366,3 +375,46 @@ func TestParseConfigBytes_ProxyConnectionPool(t *testing.T) { t.Fatal("clone shares enabled pointer") } } + +func TestV8Validate_SensitiveWordsReplace(t *testing.T) { + v8YAML := []byte(`config-version: 8 +oauth: + providers: + antigravity: + sensitive-words: + - "Codex" + sensitive-words-replace: + "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant." + devin: + sensitive-words: + - "API" + sensitive-words-replace: + "Claude Code": "Devin Assistant" +api-keys: + claude: + - name: claude-1 + keys: + - api-key: sk-ant-test + cloak: + sensitive-words: + - "proxy" + sensitive-words-replace: + "Claude Code": "Claude" +`) + if err := ValidateV8Config(v8YAML); err != nil { + t.Fatalf("ValidateV8Config() error = %v", err) + } + cfg, errParse := ParseConfigBytes(v8YAML) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + if got := cfg.Antigravity.SensitiveWordsReplace["You are Codex, an agent based on GPT-5."]; got != "You are a helpful coding assistant." { + t.Fatalf("Antigravity.SensitiveWordsReplace = %#v", cfg.Antigravity.SensitiveWordsReplace) + } + if got := cfg.Devin.SensitiveWordsReplace["Claude Code"]; got != "Devin Assistant" { + t.Fatalf("Devin.SensitiveWordsReplace = %#v", cfg.Devin.SensitiveWordsReplace) + } + if len(cfg.ClaudeKey) != 1 || cfg.ClaudeKey[0].Cloak == nil || cfg.ClaudeKey[0].Cloak.SensitiveWordsReplace["Claude Code"] != "Claude" { + t.Fatalf("ClaudeKey[0].Cloak = %#v", cfg.ClaudeKey) + } +} diff --git a/internal/config/config_load.go b/internal/config/config_load.go index c4ea79e5a35..4f9daa2e697 100644 --- a/internal/config/config_load.go +++ b/internal/config/config_load.go @@ -318,6 +318,10 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { // Sanitize Claude key headers cfg.SanitizeClaudeKeys() + // Sanitize Antigravity and Devin sensitive-word replacement maps. + cfg.SanitizeAntigravity() + cfg.SanitizeDevin() + // Sanitize OpenAI compatibility providers: drop entries without base-url cfg.SanitizeOpenAICompatibility() diff --git a/internal/config/config_normalization.go b/internal/config/config_normalization.go index 163f300a112..40addc36245 100644 --- a/internal/config/config_normalization.go +++ b/internal/config/config_normalization.go @@ -379,6 +379,50 @@ func sanitizeCodexKeyEntries(entries []CodexKey) []CodexKey { return out } +// NormalizeSensitiveWordsReplace trims replacement target keys, drops blank keys, +// and preserves replacement values (including empty strings for phrase deletion). +func NormalizeSensitiveWordsReplace(in map[string]string) map[string]string { + if len(in) == 0 { + return nil + } + keys := make([]string, 0, len(in)) + for k := range in { + keys = append(keys, k) + } + sort.Strings(keys) + out := make(map[string]string, len(in)) + for _, rawKey := range keys { + trimmedKey := strings.TrimSpace(rawKey) + if trimmedKey == "" { + continue + } + if _, exists := out[trimmedKey]; exists && rawKey != trimmedKey { + continue + } + out[trimmedKey] = in[rawKey] + } + if len(out) == 0 { + return nil + } + return out +} + +// SanitizeAntigravity normalizes Antigravity sensitive-word settings. +func (cfg *Config) SanitizeAntigravity() { + if cfg == nil { + return + } + cfg.Antigravity.SensitiveWordsReplace = NormalizeSensitiveWordsReplace(cfg.Antigravity.SensitiveWordsReplace) +} + +// SanitizeDevin normalizes Devin sensitive-word settings. +func (cfg *Config) SanitizeDevin() { + if cfg == nil { + return + } + cfg.Devin.SensitiveWordsReplace = NormalizeSensitiveWordsReplace(cfg.Devin.SensitiveWordsReplace) +} + // NormalizeCloakConfig trims strings and removes blank sensitive words. func NormalizeCloakConfig(cloak *CloakConfig) *CloakConfig { if cloak == nil { @@ -400,6 +444,7 @@ func NormalizeCloakConfig(cloak *CloakConfig) *CloakConfig { } else { cloak.SensitiveWords = nil } + cloak.SensitiveWordsReplace = NormalizeSensitiveWordsReplace(cloak.SensitiveWordsReplace) return cloak } diff --git a/internal/config/config_types.go b/internal/config/config_types.go index c186d21ab18..48dc6d28900 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -122,6 +122,12 @@ type ClaudeConfig struct { // translation elsewhere, and per-credential search advertising in the Codex // catalog. Default false keeps the translation that predates tool_search support. ToolSearch bool `yaml:"tool-search" json:"tool-search"` + // ServerSideFallback opts cloaked Claude requests into Anthropic server-side + // refusal fallback injection (claude-opus-5-5 -> claude-opus-4-8 and + // claude-fable-5-1 -> claude-opus-5) when the request body does not already + // specify fallbacks. Default false leaves cross-model fallbacks omitted unless + // explicitly configured by the caller or payload rules. + ServerSideFallback bool `yaml:"server-side-fallback" json:"server-side-fallback"` // SessionDatePinDays retains the cloaked currentDate pin per session for // this many days in in-memory mode (server.kv-store.disabled: true without // Home), independently of continuity expiry and credential switches. A nil @@ -206,6 +212,9 @@ type XAIConfig struct { type DevinConfig struct { // SensitiveWords is a list of words to obfuscate with zero-width characters in system prompts and messages. SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"` + + // SensitiveWordsReplace maps sensitive words or phrases to literal replacements in system prompts. + SensitiveWordsReplace map[string]string `yaml:"sensitive-words-replace,omitempty" json:"sensitive-words-replace,omitempty"` } // ProxyConnectionPoolConfig configures upstream connection reuse for one api-key @@ -237,9 +246,12 @@ type ProxyConnectionPoolConfig struct { // AntigravityConfig configures provider-wide Antigravity request behavior. type AntigravityConfig struct { - // SensitiveWords is a list of words to obfuscate with zero-width characters in system instructions. + // SensitiveWords is a list of words to obfuscate with invisible combining grapheme joiners in system instructions. SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"` + // SensitiveWordsReplace maps sensitive words or phrases to literal replacements in system instructions. + SensitiveWordsReplace map[string]string `yaml:"sensitive-words-replace,omitempty" json:"sensitive-words-replace,omitempty"` + // ConnectionPool configures upstream HTTP connection pooling behavior for Antigravity. ConnectionPool AntigravityConnectionPoolConfig `yaml:"connection-pool,omitempty" json:"connection-pool,omitempty"` } @@ -535,9 +547,13 @@ type RoutingConfig struct { Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` // SessionAffinity enables universal session-sticky routing for all clients. - // Explicit Claude Code, Codex, OpenCode, and pi session headers are preferred, - // followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, - // execution or derived session identity, and the existing message-content hash fallback. + // Explicit Claude Code, Codex, OpenCode, pi, and generic (X-Session-ID) session headers + // are preferred, followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, + // execution session identity, Merkle LCP prefix/compaction matching, derived session + // identity, and the existing message-content hash fallback. + // LCP compaction matching requires at least 2 shared trailing turns; clients whose + // compaction replaces all earlier turns with a summary should send X-Session-ID (or + // another explicit session ID), or affinity starts a new binding at every compaction. // Automatic failover is always enabled when bound auth becomes unavailable. SessionAffinity bool `yaml:"session-affinity,omitempty" json:"session-affinity,omitempty"` @@ -703,6 +719,9 @@ type CloakConfig struct { // This can help bypass certain content filters. SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"` + // SensitiveWordsReplace maps sensitive words or phrases to literal replacements. + SensitiveWordsReplace map[string]string `yaml:"sensitive-words-replace,omitempty" json:"sensitive-words-replace,omitempty"` + // CacheUserID controls whether Claude user_id values are cached per API key. // When false, a fresh random user_id is generated for every request. CacheUserID *bool `yaml:"cache-user-id,omitempty" json:"cache-user-id,omitempty"` @@ -923,6 +942,7 @@ type CodexModel struct { // Use this for third-party Responses-compatible endpoints that do not accept // native agent_message items or empty-signature thinking blocks. Default false // keeps the native behavior unchanged. + // When enabled, compaction_trigger items also use local summary compaction. IsCompat bool `yaml:"is-compat,omitempty" json:"is-compat,omitempty"` // SupportConfigurationUpdate enables configuration_update for this API-key model. diff --git a/internal/config/config_v8_test.go b/internal/config/config_v8_test.go index e8c66625c35..1c04a1510bd 100644 --- a/internal/config/config_v8_test.go +++ b/internal/config/config_v8_test.go @@ -933,3 +933,33 @@ func TestV8SecretHashResolvesReferences(t *testing.T) { }) } } + +func TestV8ExampleDocumentsSessionAffinityCompactionGuidance(t *testing.T) { + example, errRead := os.ReadFile(filepath.Join("..", "..", "config.example.yaml")) + if errRead != nil { + t.Fatal(errRead) + } + var doc yaml.Node + if errUnmarshal := yaml.Unmarshal(example, &doc); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + routing := yamlPath(doc.Content[0], "routing") + if routing == nil { + t.Fatal("config.example.yaml is missing routing block") + } + idx := findMapKeyIndex(routing, "session-affinity") + if idx < 0 { + t.Fatal("config.example.yaml is missing routing.session-affinity") + } + comment := routing.Content[idx].HeadComment + for _, required := range []string{ + "Merkle LCP prefix/compaction matching", + "2 shared trailing turns", + "replaces all earlier turns with a summary", + "send X-Session-ID", + } { + if !strings.Contains(comment, required) { + t.Fatalf("routing.session-affinity comment in config.example.yaml is missing %q; got:\n%s", required, comment) + } + } +} diff --git a/internal/config/config_yaml.go b/internal/config/config_yaml.go index 2ee6c1db287..9f34ef2d433 100644 --- a/internal/config/config_yaml.go +++ b/internal/config/config_yaml.go @@ -70,6 +70,8 @@ func SaveConfigPreserveComments(configFile string, cfg *Config, migrateV8 ...boo pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-model-alias") pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-request-scoped-errors") pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-settings") + pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "antigravity", "sensitive-words-replace") + pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "devin", "sensitive-words-replace") replacePluginConfigsSubtree(original.Content[0], generated.Content[0]) // Merge generated into original in-place, preserving comments/order of existing nodes. @@ -382,8 +384,49 @@ func appendPath(path []string, key string) []string { // an SDK embedder adds through sdk/config.SaveConfigPreserveComments to a file that // never carried it: pruning drops those zero-valued sub-keys. No management // endpoint mutates ProxyConnectionPool, so that path is reachable only from code. +func isSensitiveWordsReplaceEntryPath(path []string) bool { + for i, part := range path { + if part == "sensitive-words-replace" && i < len(path)-1 { + return true + } + } + return false +} + +func hasNonEmptySensitiveWordsReplace(node *yaml.Node) bool { + if node == nil { + return false + } + switch node.Kind { + case yaml.MappingNode: + for i := 0; i+1 < len(node.Content); i += 2 { + k := node.Content[i] + v := node.Content[i+1] + if k != nil && k.Value == "sensitive-words-replace" && v != nil && v.Kind == yaml.MappingNode && len(v.Content) > 0 { + return true + } + if hasNonEmptySensitiveWordsReplace(v) { + return true + } + } + case yaml.SequenceNode: + for _, child := range node.Content { + if hasNonEmptySensitiveWordsReplace(child) { + return true + } + } + } + return false +} + func isKnownDefaultValue(path []string, node *yaml.Node) bool { - if isPluginConfigsSubtreePath(path) { + if isPluginConfigsSubtreePath(path) || isSensitiveWordsReplaceEntryPath(path) { + return false + } + if len(path) > 0 && path[len(path)-1] == "sensitive-words-replace" && node != nil && node.Kind == yaml.MappingNode && len(node.Content) > 0 { + return false + } + if hasNonEmptySensitiveWordsReplace(node) { return false } if len(path) == 1 && path[0] == "plugins" && node != nil && node.Kind == yaml.MappingNode { @@ -462,7 +505,10 @@ func isKnownDefaultValue(path []string, node *yaml.Node) bool { // pruneKnownDefaultsInNewNode removes default-valued descendants from a new node // before it is appended into the destination YAML tree. func pruneKnownDefaultsInNewNode(path []string, node *yaml.Node) { - if node == nil || isPluginConfigsSubtreePath(path) { + if node == nil || isPluginConfigsSubtreePath(path) || isSensitiveWordsReplaceEntryPath(path) { + return + } + if len(path) > 0 && path[len(path)-1] == "sensitive-words-replace" { return } @@ -869,6 +915,9 @@ func shouldPruneNestedMappingKeys(path []string) bool { } parent := path[len(path)-2] last := path[len(path)-1] + if last == "sensitive-words-replace" && (parent == "antigravity" || parent == "devin" || parent == "cloak") { + return true + } switch parent { case "claude-api-key": return last == "cloak" || last == "headers" diff --git a/internal/config/parse.go b/internal/config/parse.go index bf12a836793..d3a586a6f26 100644 --- a/internal/config/parse.go +++ b/internal/config/parse.go @@ -121,6 +121,8 @@ func ParseConfigBytes(data []byte) (*Config, error) { cfg.SanitizeCodexHeaderDefaults() cfg.SanitizeClaudeHeaderDefaults() cfg.SanitizeClaudeKeys() + cfg.SanitizeAntigravity() + cfg.SanitizeDevin() cfg.SanitizeOpenAICompatibility() cfg.OAuthExcludedModels = NormalizeOAuthExcludedModels(cfg.OAuthExcludedModels) cfg.SanitizeOAuthModelAlias() diff --git a/internal/home/client.go b/internal/home/client.go index 8934fbe170c..8c2d49c7908 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -2018,6 +2018,7 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( } return c.endConfigSubscriberLifetimeWithSubscription(errACK, pubsub, "failed ACK") } + heartbeatStarted := time.Now() // A protocol-one ACK means Home already committed this membership. Preserve it if the command probe fails. if len(args) > 1 { c.markMembershipTakeoverEligible() @@ -2065,10 +2066,16 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( }() go func() { defer close(receiveDone) + lastHeartbeat := heartbeatStarted for { _, timeout := c.subscriptionParameters() - event, errReceive := receiveSubscriptionEvent(receiveCtx, pubsub, sampler, timeout) + event, errReceive := receiveSubscriptionEvent(receiveCtx, pubsub, sampler, lastHeartbeat.Add(timeout)) if errReceive == nil { + // Only a persisted heartbeat PONG renews liveness. Config, cluster traffic, + // and RTT echoes cannot extend the deadline while the consumer is busy. + if pong, ok := event.(*redis.Pong); ok && pong != nil && pong.Payload == "" { + lastHeartbeat = time.Now() + } select { case events <- event: continue @@ -2144,16 +2151,17 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( log.Warn("failed to apply config update from home control center, ignoring") } case *redis.Pong: - c.resetReconnectFailures() + if msg != nil && msg.Payload == "" { + c.resetReconnectFailures() + } default: log.Debugf("home subscription returned unsupported message type %T", event) } } } -// receiveSubscriptionEvent consumes telemetry without extending the heartbeat deadline. -func receiveSubscriptionEvent(ctx context.Context, pubsub *redis.PubSub, sampler *clusterRTTSampler, timeout time.Duration) (any, error) { - deadline := time.Now().Add(timeout) +// receiveSubscriptionEvent preserves the absolute heartbeat deadline across ordinary messages. +func receiveSubscriptionEvent(ctx context.Context, pubsub *redis.PubSub, sampler *clusterRTTSampler, deadline time.Time) (any, error) { for { remaining := time.Until(deadline) if remaining <= 0 { @@ -2163,6 +2171,9 @@ func receiveSubscriptionEvent(ctx context.Context, pubsub *redis.PubSub, sampler if errReceive != nil { return nil, errReceive } + if !time.Now().Before(deadline) { + return nil, os.ErrDeadlineExceeded + } if _, ok := event.(*redis.Subscription); ok { return nil, fmt.Errorf("unexpected Home subscription ACK after initialization") } diff --git a/internal/home/client_heartbeat_deadline_test.go b/internal/home/client_heartbeat_deadline_test.go new file mode 100644 index 00000000000..e3c77a5329f --- /dev/null +++ b/internal/home/client_heartbeat_deadline_test.go @@ -0,0 +1,126 @@ +package home + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "strings" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/redis/go-redis/v9" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func TestManagedSubscriberOrdinaryTrafficDoesNotRenewHeartbeat(t *testing.T) { + for _, mode := range []string{"config", "cluster", "blocked config consumer"} { + t.Run(mode, func(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "GET") { + return "$2\r\n{}\r\n" + } + return balanceTestRedisReply(args) + }) + client.SetManagedLifetime(true) + synctest.Test(t, func(t *testing.T) { + defer client.Close() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + releaseConsumer := make(chan struct{}) + var messages atomic.Int32 + sub := redis.NewClient(&redis.Options{ + Addr: "127.0.0.1:8327", Protocol: 2, DisableIdentity: true, MaxRetries: -1, + Dialer: func(context.Context, string, string) (net.Conn, error) { + local, peer := net.Pipe() + go serveRedisCommandTestConn(peer, &redisCommandLog{}, func(args []string) string { + if !strings.EqualFold(args[0], "SUBSCRIBE") { + return balanceTestRedisReply(args) + } + go func() { + ticker := time.NewTicker(10 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + channel, payload := redisChannelConfig, "{}" + if mode == "cluster" { + channel, payload = redisChannelCluster, `{"ok":true,"nodes":[]}` + } + frame := fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$%d\r\n%s\r\n$%d\r\n%s\r\n", len(channel), channel, len(payload), payload) + if _, errWrite := io.WriteString(peer, frame); errWrite != nil { + if !errors.Is(errWrite, io.ErrClosedPipe) && !errors.Is(errWrite, net.ErrClosed) { + t.Errorf("write ordinary subscription message: %v", errWrite) + } + return + } + messages.Add(1) + } + }() + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + }) + return local, nil + }, + }) + defer func() { + if errClose := sub.Close(); errClose != nil { + t.Error(errClose) + } + }() + ready := make(chan struct{}) + done := make(chan error, 1) + go func() { + initial := true + done <- client.RunConfigSubscriberLifetime(ctx, func([]byte) error { + if !initial { + if mode == "blocked config consumer" { + <-releaseConsumer + } + return nil + } + initial = false + if errConfig := client.SetLifecycleConfig(config.CredentialConcurrencyConfig{ + LifecycleConfigRevision: 1, CPAHeartbeatTimeout: time.Minute, + }); errConfig != nil { + return errConfig + } + client.mu.Lock() + previousSub := client.sub + client.sub = sub + client.mu.Unlock() + if errClose := previousSub.Close(); errClose != nil { + return fmt.Errorf("close bootstrap subscription: %w", errClose) + } + return nil + }, func() { close(ready) }) + }() + select { + case <-ready: + case errRun := <-done: + close(releaseConsumer) + t.Fatalf("subscription did not become ready: %v", errRun) + } + // Virtual time keeps this deadline test independent of wall-clock scheduling. + time.Sleep(time.Minute) + synctest.Wait() + heartbeatOK := client.HeartbeatOK() + fenced := client.dispatchFenced.Load() && client.releaseFenced.Load() + close(releaseConsumer) + synctest.Wait() + errRun := <-done + if heartbeatOK || !fenced || !isTimeoutError(errRun) || ctx.Err() != nil { + t.Fatalf("ordinary traffic extended heartbeat liveness: healthy=%t fenced=%t error=%v context=%v", heartbeatOK, fenced, errRun, ctx.Err()) + } + if messages.Load() < 5 { + t.Fatalf("insufficient ordinary traffic before heartbeat expiry: %d messages", messages.Load()) + } + }) + }) + } +} diff --git a/internal/home/client_latency_failure_test.go b/internal/home/client_latency_failure_test.go index bf0175ea933..0eba5188244 100644 --- a/internal/home/client_latency_failure_test.go +++ b/internal/home/client_latency_failure_test.go @@ -92,7 +92,7 @@ func TestClusterLatencyPingFailureBetweenReceives(t *testing.T) { t.Fatalf("failed PING was not invalidated: %+v", samples) } // The new ACK cannot silently keep the old lifetime alive, even if reception wins the race. - event, errReceive := receiveSubscriptionEvent(ctx, pubsub, sampler, time.Second) + event, errReceive := receiveSubscriptionEvent(ctx, pubsub, sampler, time.Now().Add(time.Second)) if errReceive == nil || isTimeoutError(errReceive) || event != nil { t.Fatalf("replacement ACK was accepted: event=%v error=%v", event, errReceive) } diff --git a/internal/home/client_latency_test.go b/internal/home/client_latency_test.go index c2c028622b3..177bf9834a7 100644 --- a/internal/home/client_latency_test.go +++ b/internal/home/client_latency_test.go @@ -21,11 +21,15 @@ func TestSubscriptionTelemetryPreservesHeartbeatDeadline(t *testing.T) { interval time.Duration stale bool heartbeats bool + traffic string }{ {name: "default timeout", timeout: homeSubscriptionReceiveTimeout, interval: 30 * time.Second}, {name: "timeout longer than sampling interval", timeout: 75 * time.Second, interval: 30 * time.Second}, {name: "stale telemetry", timeout: homeSubscriptionReceiveTimeout, interval: time.Second, stale: true}, {name: "normal heartbeats renew deadline", timeout: homeSubscriptionReceiveTimeout, interval: time.Second, heartbeats: true}, + {name: "config traffic cannot renew deadline", timeout: 75 * time.Second, interval: 30 * time.Second, traffic: redisChannelConfig}, + {name: "cluster traffic cannot renew deadline", timeout: 75 * time.Second, interval: 30 * time.Second, traffic: redisChannelCluster}, + {name: "unrelated pong cannot renew deadline", timeout: 75 * time.Second, interval: 30 * time.Second, traffic: "unrelated-pong"}, } { t.Run(test.name, func(t *testing.T) { synctest.Test(t, func(t *testing.T) { @@ -83,13 +87,19 @@ func TestSubscriptionTelemetryPreservesHeartbeatDeadline(t *testing.T) { return } token := "" - if !test.heartbeats || (elapsed != 2*time.Second && elapsed != 4*time.Second) { + if test.traffic == "" && (!test.heartbeats || (elapsed != 2*time.Second && elapsed != 4*time.Second)) { token = sampler.begin(time.Now()) if test.stale { sampler.expire(time.Now(), true) } } - if _, errWrite := io.WriteString(peer, fmt.Sprintf("*2\r\n$4\r\npong\r\n$%d\r\n%s\r\n", len(token), token)); errWrite != nil { + frame := fmt.Sprintf("*2\r\n$4\r\npong\r\n$%d\r\n%s\r\n", len(token), token) + if test.traffic == "unrelated-pong" { + frame = "*2\r\n$4\r\npong\r\n$5\r\nother\r\n" + } else if test.traffic != "" { + frame = fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$%d\r\n%s\r\n$2\r\n{}\r\n", len(test.traffic), test.traffic) + } + if _, errWrite := io.WriteString(peer, frame); errWrite != nil { // The receiver closes the pipe when its heartbeat deadline expires. if !errors.Is(errWrite, io.ErrClosedPipe) && !errors.Is(errWrite, net.ErrClosed) { t.Errorf("write telemetry PONG: %v", errWrite) @@ -100,24 +110,29 @@ func TestSubscriptionTelemetryPreservesHeartbeatDeadline(t *testing.T) { } }() heartbeats := 0 + deadline := started.Add(test.timeout) for { - event, errReceive := receiveSubscriptionEvent(ctx, pubsub, sampler, test.timeout) + event, errReceive := receiveSubscriptionEvent(ctx, pubsub, sampler, deadline) if errReceive != nil { if !isTimeoutError(errReceive) { t.Fatalf("receive error = %v, want heartbeat timeout", errReceive) } break } + if test.traffic != "" { + continue + } pong, ok := event.(*redis.Pong) if !ok || pong.Payload != "" { t.Fatalf("telemetry escaped to the heartbeat consumer: %#v", event) } heartbeats++ + deadline = time.Now().Add(test.timeout) } if elapsed := time.Since(started); elapsed != wantElapsed || heartbeats != wantHeartbeats { t.Fatalf("elapsed=%v heartbeats=%d, want %v and %d", elapsed, heartbeats, wantElapsed, wantHeartbeats) } - if samples := sampler.history.report(time.Now()); len(samples) != 1 || samples[0].Valid == test.stale { + if samples := sampler.history.report(time.Now()); test.traffic == "" && (len(samples) != 1 || samples[0].Valid == test.stale) { t.Fatalf("telemetry was not sampled independently of liveness: %+v", samples) } }) diff --git a/internal/logging/request_capture.go b/internal/logging/request_capture.go index 6a97ea4ae8c..f922639af5f 100644 --- a/internal/logging/request_capture.go +++ b/internal/logging/request_capture.go @@ -1,10 +1,41 @@ package logging -import "context" +import ( + "bytes" + "context" + "fmt" + "net/http" + "sort" + "strings" + "sync" + "time" + "unicode/utf8" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" +) + +const ( + // MaxDeferredAPIRequestBodyBytes bounds total deferred upstream request bodies per HTTP request. + MaxDeferredAPIRequestBodyBytes = 32 << 20 // 32 MiB + // MaxDeferredAPIResponseBodyBytes bounds each upstream attempt's deferred response body. + MaxDeferredAPIResponseBodyBytes = 64 << 10 // 64 KiB + // MaxDeferredAPIResponseHeaderBytes bounds cloned response headers per upstream attempt. + MaxDeferredAPIResponseHeaderBytes = 16 << 10 // 16 KiB + // MaxDeferredAPIResponseErrorBytes bounds each recorded upstream error string. + MaxDeferredAPIResponseErrorBytes = 16 << 10 // 16 KiB + // MaxTotalDeferredAPIResponseBodyBytes bounds total deferred upstream response payload + // (body chunks, headers, and error strings) across retries for a single HTTP request. + MaxTotalDeferredAPIResponseBodyBytes = 512 << 10 // 512 KiB + // MaxDeferredAPIResponseAttempts bounds the number of upstream response attempt slots retained per HTTP request. + MaxDeferredAPIResponseAttempts = 16 + maxDeferredErrorsPerAttempt = 8 + maxDeferredHeaderEntries = 64 +) type deferredAPIRequestCaptureDisabledKey struct{} +type deferredAPICaptureKey struct{} -// WithDeferredAPIRequestCaptureDisabled skips error-only upstream request capture +// WithDeferredAPIRequestCaptureDisabled skips error-only upstream request and response capture // for this context and its descendants. Full request logging is unaffected. // Connection-scoped contexts keep this policy even when asynchronous callbacks // outlive the handler and its pooled Gin context is reused. @@ -12,6 +43,9 @@ func WithDeferredAPIRequestCaptureDisabled(ctx context.Context) context.Context if ctx == nil { ctx = context.Background() } + if capture := DeferredAPICaptureFromContext(ctx); capture != nil { + capture.Close() + } return context.WithValue(ctx, deferredAPIRequestCaptureDisabledKey{}, true) } @@ -23,3 +57,659 @@ func DeferredAPIRequestCaptureDisabled(ctx context.Context) bool { disabled, _ := ctx.Value(deferredAPIRequestCaptureDisabledKey{}).(bool) return disabled } + +// DeferredAPIResponse renders the upstream response attempt sections of an +// error log from details captured while full request logging is disabled. +type DeferredAPIResponse interface { + // HasPayload reports whether any upstream response metadata, body, or error was captured. + HasPayload() bool + // RenderAPIResponse returns the formatted "=== API RESPONSE n ===" attempt + // sections, or nil when no upstream response details were captured. + RenderAPIResponse() []byte +} + +// WithDeferredAPICapture attaches a request-scoped DeferredAPICapture to ctx so +// background stream producers record into the request's own synchronized holder +// without accessing a pooled Gin context. +func WithDeferredAPICapture(ctx context.Context, capture *DeferredAPICapture) context.Context { + if ctx == nil { + ctx = context.Background() + } + if capture == nil { + return ctx + } + return context.WithValue(ctx, deferredAPICaptureKey{}, capture) +} + +// DeferredAPICaptureFromContext returns the request-scoped DeferredAPICapture carried by ctx, if any. +func DeferredAPICaptureFromContext(ctx context.Context) *DeferredAPICapture { + if ctx == nil { + return nil + } + capture, _ := ctx.Value(deferredAPICaptureKey{}).(*DeferredAPICapture) + return capture +} + +// InheritDeferredCapture propagates error-only capture policy and any request-scoped +// DeferredAPICapture holder from src to dst. +func InheritDeferredCapture(dst, src context.Context) context.Context { + if dst == nil { + dst = context.Background() + } + if src == nil { + return dst + } + if DeferredAPIRequestCaptureDisabled(dst) || DeferredAPIRequestCaptureDisabled(src) { + return WithDeferredAPIRequestCaptureDisabled(dst) + } + if DeferredAPICaptureFromContext(dst) == nil { + if capture := DeferredAPICaptureFromContext(src); capture != nil { + dst = WithDeferredAPICapture(dst, capture) + } + } + return dst +} + +type deferredHeaderEntry struct { + key string + value string +} + +type deferredResponseAttempt struct { + index int + capturedAt time.Time + status int + headers []deferredHeaderEntry + headersRecorded bool + headersTruncated bool + body []byte + bodyHasContent bool + prevWasSSEEvent bool + bodyTruncated bool + errors []string + errorsTruncated bool + requestBodyBytes int +} + +func (a *deferredResponseAttempt) markCapturedAt() { + if a.capturedAt.IsZero() { + a.capturedAt = time.Now() + } +} + +func (a *deferredResponseAttempt) hasContent() bool { + return a.status > 0 || a.headersRecorded || a.headersTruncated || a.bodyHasContent || a.bodyTruncated || len(a.errors) > 0 || a.errorsTruncated +} + +// DeferredAPICapture retains bounded upstream request and response details for +// a single HTTP request while full request logging is disabled. All operations +// are synchronized and become no-ops after Close so asynchronous stream workers +// cannot race with Finalize or leak into a reused Gin context. +type DeferredAPICapture struct { + mu sync.Mutex + requests []DeferredAPIRequest + requestBytes int + attempts []*deferredResponseAttempt + omittedAttempts int + responseBytes int + closed bool +} + +// NewDeferredAPICapture constructs a fresh request-scoped DeferredAPICapture. +func NewDeferredAPICapture() *DeferredAPICapture { + return &DeferredAPICapture{} +} + +// Close marks the capture closed and releases buffered upstream payloads. +func (c *DeferredAPICapture) Close() { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + c.requests = nil + c.requestBytes = 0 + c.attempts = nil + c.omittedAttempts = 0 + c.responseBytes = 0 +} + +// IsClosed reports whether Close has been called. +func (c *DeferredAPICapture) IsClosed() bool { + if c == nil { + return true + } + c.mu.Lock() + defer c.mu.Unlock() + return c.closed +} + +// RecordRequest stores a deferred upstream request builder within MaxDeferredAPIRequestBodyBytes +// and opens the paired upstream response attempt slot. +func (c *DeferredAPICapture) RecordRequest(body []byte, build func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) DeferredAPIRequest) (DeferredAPIRequest, int, bool) { + return c.recordRequest(body, false, build) +} + +// RecordOrReplacePendingRequest replaces the trailing request builder when its paired +// response attempt has recorded no response details yet (for example when a websocket +// request is immediately rejected during its HTTP upgrade handshake), or appends a new +// request attempt otherwise. +func (c *DeferredAPICapture) RecordOrReplacePendingRequest(body []byte, build func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) DeferredAPIRequest) (DeferredAPIRequest, int, bool, bool) { + fn, bytesUsed, replaced, ok := c.recordRequestWithReplace(body, true, build) + return fn, bytesUsed, replaced, ok +} + +func (c *DeferredAPICapture) recordRequest(body []byte, replacePending bool, build func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) DeferredAPIRequest) (DeferredAPIRequest, int, bool) { + fn, bytesUsed, _, ok := c.recordRequestWithReplace(body, replacePending, build) + return fn, bytesUsed, ok +} + +func (c *DeferredAPICapture) recordRequestWithReplace(body []byte, replacePending bool, build func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) DeferredAPIRequest) (DeferredAPIRequest, int, bool, bool) { + if c == nil || build == nil { + return nil, 0, false, false + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil, 0, false, false + } + if replacePending && c.omittedAttempts == 0 && len(c.attempts) > 0 && len(c.requests) == len(c.attempts) { + lastAttempt := c.attempts[len(c.attempts)-1] + if lastAttempt != nil && !lastAttempt.hasContent() { + c.requestBytes -= lastAttempt.requestBodyBytes + if c.requestBytes < 0 { + c.requestBytes = 0 + } + capturedAt := time.Now() + remaining := MaxDeferredAPIRequestBodyBytes - c.requestBytes + if remaining < 0 { + remaining = 0 + } + captureLength := len(body) + if captureLength > remaining { + captureLength = remaining + } + capturedBody := bytes.Clone(body[:captureLength]) + bodyEmpty := len(body) == 0 + bodyTruncated := captureLength < len(body) + c.requestBytes += captureLength + lastAttempt.requestBodyBytes = captureLength + fn := build(lastAttempt.index, capturedAt, capturedBody, bodyEmpty, bodyTruncated) + c.requests[len(c.requests)-1] = fn + return fn, c.requestBytes, true, true + } + } + index := len(c.requests) + 1 + if len(c.requests) >= MaxDeferredAPIResponseAttempts { + c.startAttemptLocked(index, 0) + return nil, c.requestBytes, false, false + } + capturedAt := time.Now() + remaining := MaxDeferredAPIRequestBodyBytes - c.requestBytes + if remaining < 0 { + remaining = 0 + } + captureLength := len(body) + if captureLength > remaining { + captureLength = remaining + } + capturedBody := bytes.Clone(body[:captureLength]) + bodyEmpty := len(body) == 0 + bodyTruncated := captureLength < len(body) + c.requestBytes += captureLength + fn := build(index, capturedAt, capturedBody, bodyEmpty, bodyTruncated) + c.requests = append(c.requests, fn) + c.startAttemptLocked(index, captureLength) + return fn, c.requestBytes, false, true +} + +// HasRequests reports whether any deferred upstream requests were captured. +func (c *DeferredAPICapture) HasRequests() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + return !c.closed && len(c.requests) > 0 +} + +// RenderAPIRequests renders all captured upstream request sections. +func (c *DeferredAPICapture) RenderAPIRequests() []byte { + if c == nil { + return nil + } + c.mu.Lock() + requests := append([]DeferredAPIRequest(nil), c.requests...) + omitted := c.omittedAttempts + closed := c.closed + c.mu.Unlock() + if closed || len(requests) == 0 { + return nil + } + var body bytes.Buffer + for _, buildRequest := range requests { + if buildRequest == nil { + continue + } + body.Write(buildRequest()) + } + if omitted > 0 { + body.WriteString(fmt.Sprintf("[API REQUEST ATTEMPTS TRUNCATED: omitted %d additional attempt(s)]\n\n", omitted)) + } + return body.Bytes() +} + +func (c *DeferredAPICapture) startAttemptLocked(index int, requestBodyBytes int) { + if len(c.attempts) >= MaxDeferredAPIResponseAttempts { + c.omittedAttempts++ + return + } + if next := len(c.attempts) + 1; next > index { + index = next + } + c.attempts = append(c.attempts, &deferredResponseAttempt{ + index: index, + requestBodyBytes: requestBodyBytes, + }) +} + +func (c *DeferredAPICapture) latestAttemptLocked() *deferredResponseAttempt { + if c.omittedAttempts > 0 { + return nil + } + if len(c.attempts) == 0 { + c.startAttemptLocked(1, 0) + } + if len(c.attempts) == 0 { + return nil + } + return c.attempts[len(c.attempts)-1] +} + +// RecordResponseMetadata records the HTTP status and bounded cloned headers for the latest attempt. +func (c *DeferredAPICapture) RecordResponseMetadata(status int, headers http.Header) { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return + } + attempt := c.latestAttemptLocked() + if attempt == nil { + return + } + attempt.markCapturedAt() + if status > 0 && attempt.status == 0 { + attempt.status = status + } + if !attempt.headersRecorded { + attempt.headersRecorded = true + if len(headers) > 0 { + budget := min(MaxDeferredAPIResponseHeaderBytes, MaxTotalDeferredAPIResponseBodyBytes-c.responseBytes) + cloned, used, truncated := cloneBoundedHeaders(headers, budget) + attempt.headers = cloned + attempt.headersTruncated = truncated + c.responseBytes += used + } + } +} + +// RecordResponseError appends a bounded transport or stream error to the latest attempt. +func (c *DeferredAPICapture) RecordResponseError(err error) { + if c == nil || err == nil { + return + } + raw := strings.TrimSpace(err.Error()) + if raw == "" { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return + } + attempt := c.latestAttemptLocked() + if attempt == nil { + return + } + attempt.markCapturedAt() + if len(attempt.errors) >= maxDeferredErrorsPerAttempt { + attempt.errorsTruncated = true + return + } + budget := min(MaxDeferredAPIResponseErrorBytes, MaxTotalDeferredAPIResponseBodyBytes-c.responseBytes) + if budget <= 0 { + attempt.errorsTruncated = true + return + } + if len(raw) > budget { + cut := truncateUTF8Boundary(raw, budget) + cloned := strings.Clone(raw[:cut]) + c.responseBytes += cut + attempt.errors = append(attempt.errors, fmt.Sprintf("%s [API RESPONSE ERROR TRUNCATED: captured first %d bytes]", cloned, cut)) + return + } + cloned := strings.Clone(raw) + c.responseBytes += len(cloned) + attempt.errors = append(attempt.errors, cloned) +} + +// AppendResponseChunk appends a bounded response chunk to the latest attempt. +func (c *DeferredAPICapture) AppendResponseChunk(chunk []byte) { + if c == nil { + return + } + data := bytes.TrimSpace(chunk) + if len(data) == 0 { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return + } + attempt := c.latestAttemptLocked() + if attempt == nil { + return + } + attempt.markCapturedAt() + + currentChunkIsSSEEvent := bytes.HasPrefix(data, []byte("event:")) + currentChunkIsSSEData := bytes.HasPrefix(data, []byte("data:")) + + attemptRemaining := MaxDeferredAPIResponseBodyBytes - len(attempt.body) + totalRemaining := MaxTotalDeferredAPIResponseBodyBytes - c.responseBytes + remaining := min(attemptRemaining, totalRemaining) + if remaining <= 0 { + attempt.bodyTruncated = true + return + } + + if attempt.bodyHasContent { + separator := "\n\n" + if attempt.prevWasSSEEvent && currentChunkIsSSEData { + separator = "\n" + } + if len(separator) >= remaining { + attempt.bodyTruncated = true + return + } + attempt.body = append(attempt.body, separator...) + c.responseBytes += len(separator) + remaining -= len(separator) + } + + writeLength := len(data) + if writeLength > remaining { + writeLength = remaining + attempt.bodyTruncated = true + } + if attempt.body == nil { + attempt.body = bytes.Clone(data[:writeLength]) + } else { + attempt.body = append(attempt.body, data[:writeLength]...) + } + c.responseBytes += writeLength + attempt.bodyHasContent = true + attempt.prevWasSSEEvent = currentChunkIsSSEEvent && !attempt.bodyTruncated +} + +// HasPayload reports whether any attempt recorded upstream response details. +func (c *DeferredAPICapture) HasPayload() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return false + } + return c.hasPayloadLocked() +} + +func (c *DeferredAPICapture) hasPayloadLocked() bool { + for _, attempt := range c.attempts { + if attempt != nil && attempt.hasContent() { + return true + } + } + return false +} + +// RenderAPIResponse formats captured upstream response attempts for error logs. +func (c *DeferredAPICapture) RenderAPIResponse() []byte { + if c == nil { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed || !c.hasPayloadLocked() { + return nil + } + var builder strings.Builder + for _, attempt := range c.attempts { + if attempt == nil || !attempt.hasContent() { + continue + } + if builder.Len() > 0 { + builder.WriteByte('\n') + } + capturedAt := attempt.capturedAt + if capturedAt.IsZero() { + capturedAt = time.Now() + } + builder.WriteString(fmt.Sprintf("=== API RESPONSE %d ===\n", attempt.index)) + builder.WriteString(fmt.Sprintf("Timestamp: %s\n\n", capturedAt.Format(time.RFC3339Nano))) + if attempt.status > 0 { + builder.WriteString(fmt.Sprintf("Status: %d\n", attempt.status)) + } + if attempt.headersRecorded || attempt.headersTruncated || attempt.bodyHasContent || attempt.bodyTruncated { + builder.WriteString("Headers:\n") + writeMaskedHeaders(&builder, attempt.headers, attempt.headersTruncated) + builder.WriteByte('\n') + } + if attempt.bodyHasContent || attempt.bodyTruncated { + builder.WriteString("Body:\n") + builder.Write(attempt.body) + if attempt.bodyTruncated { + if len(attempt.body) > 0 && attempt.body[len(attempt.body)-1] != '\n' { + builder.WriteByte('\n') + } + builder.WriteString(fmt.Sprintf("[API RESPONSE BODY TRUNCATED: captured first %d bytes]", len(attempt.body))) + } + if len(attempt.body) == 0 || attempt.body[len(attempt.body)-1] != '\n' || attempt.bodyTruncated { + builder.WriteByte('\n') + } + if len(attempt.errors) > 0 || attempt.errorsTruncated { + builder.WriteByte('\n') + } + } + for i, errText := range attempt.errors { + if i > 0 { + builder.WriteByte('\n') + } + builder.WriteString(fmt.Sprintf("Error: %s\n", errText)) + } + if attempt.errorsTruncated { + if len(attempt.errors) > 0 { + builder.WriteByte('\n') + } + builder.WriteString("[API RESPONSE ERRORS TRUNCATED]\n") + } + } + if c.omittedAttempts > 0 && builder.Len() > 0 { + builder.WriteString(fmt.Sprintf("\n[API RESPONSE ATTEMPTS TRUNCATED: omitted %d additional attempt(s)]\n", c.omittedAttempts)) + } + if builder.Len() == 0 { + return nil + } + return []byte(builder.String()) +} + +// AttemptCount returns the number of recorded response attempt slots. +func (c *DeferredAPICapture) AttemptCount() int { + if c == nil { + return 0 + } + c.mu.Lock() + defer c.mu.Unlock() + return len(c.attempts) +} + +// AttemptCapturedAt returns the captured response timestamp for the given 0-based attempt slot. +func (c *DeferredAPICapture) AttemptCapturedAt(i int) time.Time { + if c == nil { + return time.Time{} + } + c.mu.Lock() + defer c.mu.Unlock() + if i < 0 || i >= len(c.attempts) || c.attempts[i] == nil { + return time.Time{} + } + return c.attempts[i].capturedAt +} + +// AttemptBodyStats returns the captured body length, slice capacity, and truncation flag for the given 0-based attempt slot. +func (c *DeferredAPICapture) AttemptBodyStats(i int) (length int, capacity int, truncated bool) { + if c == nil { + return 0, 0, false + } + c.mu.Lock() + defer c.mu.Unlock() + if i < 0 || i >= len(c.attempts) || c.attempts[i] == nil { + return 0, 0, false + } + attempt := c.attempts[i] + return len(attempt.body), cap(attempt.body), attempt.bodyTruncated +} + +// TotalResponseBodyBytes returns the total captured upstream response bytes across attempts. +func (c *DeferredAPICapture) TotalResponseBodyBytes() int { + if c == nil { + return 0 + } + c.mu.Lock() + defer c.mu.Unlock() + return c.responseBytes +} + +func cloneBoundedHeaders(headers http.Header, budget int) ([]deferredHeaderEntry, int, bool) { + if len(headers) == 0 { + return nil, 0, false + } + if budget <= 0 { + return nil, 0, true + } + truncated := len(headers) > maxDeferredHeaderEntries + limitKeys := min(len(headers), maxDeferredHeaderEntries) + keys := make([]string, 0, limitKeys) + for key := range headers { + pos := sort.SearchStrings(keys, key) + if len(keys) < limitKeys { + keys = append(keys, "") + for i := len(keys) - 1; i > pos; i-- { + keys[i] = keys[i-1] + } + keys[pos] = key + } else if pos < limitKeys { + for i := limitKeys - 1; i > pos; i-- { + keys[i] = keys[i-1] + } + keys[pos] = key + } + } + + entries := make([]deferredHeaderEntry, 0, len(keys)) + used := 0 + stop := false + for _, key := range keys { + if len(entries) >= maxDeferredHeaderEntries || len(key) > budget-used { + truncated = true + break + } + clonedKey := strings.Clone(key) + values := headers[key] + if len(values) == 0 { + entries = append(entries, deferredHeaderEntry{key: clonedKey}) + used += len(clonedKey) + continue + } + for _, value := range values { + if len(entries) >= maxDeferredHeaderEntries { + truncated = true + stop = true + break + } + masked := util.MaskSensitiveHeaderValue(clonedKey, value) + remaining := budget - used - len(clonedKey) + if remaining <= 0 { + truncated = true + stop = true + break + } + if len(masked) > remaining { + cut := truncateUTF8Boundary(masked, remaining) + entries = append(entries, deferredHeaderEntry{ + key: clonedKey, + value: strings.Clone(masked[:cut]), + }) + used += len(clonedKey) + cut + truncated = true + stop = true + break + } + entries = append(entries, deferredHeaderEntry{ + key: clonedKey, + value: strings.Clone(masked), + }) + used += len(clonedKey) + len(masked) + } + if stop { + break + } + } + return entries, used, truncated +} + +func truncateUTF8Boundary(s string, maxBytes int) int { + if maxBytes <= 0 { + return 0 + } + if len(s) <= maxBytes { + return len(s) + } + for maxBytes > 0 && !utf8.RuneStart(s[maxBytes]) { + maxBytes-- + } + if maxBytes == 0 { + return 1 + } + return maxBytes +} + +func writeMaskedHeaders(builder *strings.Builder, headers []deferredHeaderEntry, truncated bool) { + if builder == nil { + return + } + if len(headers) == 0 { + if truncated { + builder.WriteString("[API RESPONSE HEADERS TRUNCATED]\n") + return + } + builder.WriteString("\n") + return + } + for _, entry := range headers { + if entry.value == "" { + builder.WriteString(fmt.Sprintf("%s:\n", entry.key)) + continue + } + builder.WriteString(fmt.Sprintf("%s: %s\n", entry.key, entry.value)) + } + if truncated { + builder.WriteString("[API RESPONSE HEADERS TRUNCATED]\n") + } +} diff --git a/internal/logging/request_logger.go b/internal/logging/request_logger.go index 550a0ae26fc..cdc56790f03 100644 --- a/internal/logging/request_logger.go +++ b/internal/logging/request_logger.go @@ -15,6 +15,7 @@ const ( WebsocketTimelineSourceContextKey = "WEBSOCKET_TIMELINE_SOURCE" APIRequestSourceContextKey = "API_REQUEST_SOURCE" DeferredAPIRequestContextKey = "DEFERRED_API_REQUEST" + DeferredAPIResponseContextKey = "DEFERRED_API_RESPONSE" APIResponseSourceContextKey = "API_RESPONSE_SOURCE" APIResponseCapturedContextKey = "API_RESPONSE_CAPTURED" APIWebsocketTimelineSourceContextKey = "API_WEBSOCKET_TIMELINE_SOURCE" diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 7f82e7dcf39..f02e53f828a 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -174,6 +174,9 @@ type ModelRegistration struct { Providers map[string]int // SuspendedClients tracks temporarily disabled clients keyed by client ID SuspendedClients map[string]string + // SuspendedClientDeadlines records optional expiry deadlines for suspended clients. + // Clients without an entry (or with a zero deadline) stay suspended until resumed. + SuspendedClientDeadlines map[string]time.Time } // ModelRegistryHook provides optional callbacks for external integrations to track model list changes. @@ -216,6 +219,8 @@ type ModelRegistry struct { registrationEpoch atomic.Uint64 // hook is an optional callback sink for model registration changes hook ModelRegistryHook + // nowFunc overrides time.Now when set (used by deterministic tests). + nowFunc func() time.Time } // Global model registry instance @@ -239,6 +244,13 @@ func GetGlobalRegistry() *ModelRegistry { }) return globalRegistry } +func (r *ModelRegistry) now() time.Time { + if r != nil && r.nowFunc != nil { + return r.nowFunc() + } + return time.Now() +} + func (r *ModelRegistry) ensureAvailableModelsCacheLocked() { if r.availableModelsCache == nil { r.availableModelsCache = make(map[string]availableModelsCacheEntry) @@ -553,7 +565,7 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo } r.registrationEpoch.Add(1) - now := time.Now() + now := r.now() if !hadExisting { // Pure addition normally has no prior metadata. Clear orphaned metadata // and refresh its models before publishing the new catalog. @@ -693,6 +705,9 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo if !preserveState || providerChanged || oldCounts[id] == 0 { delete(reg.QuotaExceededClients, clientID) delete(reg.SuspendedClients, clientID) + if reg.SuspendedClientDeadlines != nil { + delete(reg.SuspendedClientDeadlines, clientID) + } } if providerChanged && provider != "" { if _, newlyAdded := addedSet[id]; newlyAdded { @@ -757,6 +772,9 @@ func (r *ModelRegistry) addModelRegistration(modelID, provider string, model *Mo if existing.SuspendedClients == nil { existing.SuspendedClients = make(map[string]string) } + if existing.SuspendedClientDeadlines == nil { + existing.SuspendedClientDeadlines = make(map[string]time.Time) + } if existing.InfoByProvider == nil { existing.InfoByProvider = make(map[string]*ModelInfo) } @@ -774,12 +792,13 @@ func (r *ModelRegistry) addModelRegistration(modelID, provider string, model *Mo } registration := &ModelRegistration{ - Info: cloneModelInfo(model), - InfoByProvider: make(map[string]*ModelInfo), - Count: 1, - LastUpdated: now, - QuotaExceededClients: make(map[string]*time.Time), - SuspendedClients: make(map[string]string), + Info: cloneModelInfo(model), + InfoByProvider: make(map[string]*ModelInfo), + Count: 1, + LastUpdated: now, + QuotaExceededClients: make(map[string]*time.Time), + SuspendedClients: make(map[string]string), + SuspendedClientDeadlines: make(map[string]time.Time), } if provider != "" { registration.Providers = map[string]int{provider: 1} @@ -800,6 +819,9 @@ func (r *ModelRegistry) removeModelRegistration(clientID, modelID, provider stri if !preserveState { delete(registration.QuotaExceededClients, clientID) delete(registration.SuspendedClients, clientID) + if registration.SuspendedClientDeadlines != nil { + delete(registration.SuspendedClientDeadlines, clientID) + } } if registration.Count < 0 { registration.Count = 0 @@ -926,7 +948,7 @@ func (r *ModelRegistry) unregisterClientInternal(clientID string) { return } - now := time.Now() + now := r.now() for _, modelID := range models { if registration, isExists := r.models[modelID]; isExists { registration.Count-- @@ -937,6 +959,9 @@ func (r *ModelRegistry) unregisterClientInternal(clientID string) { if registration.SuspendedClients != nil { delete(registration.SuspendedClients, clientID) } + if registration.SuspendedClientDeadlines != nil { + delete(registration.SuspendedClientDeadlines, clientID) + } if hasProvider && registration.Providers != nil { if count, ok := registration.Providers[provider]; ok { @@ -994,7 +1019,7 @@ func (r *ModelRegistry) SetModelQuotaExceeded(clientID, modelID string) { r.ensureAvailableModelsCacheLocked() if registration, exists := r.models[modelID]; exists { - now := time.Now() + now := r.now() registration.QuotaExceededClients[clientID] = &now r.invalidateAvailableModelsCacheLocked() log.Debugf("Marked model %s as quota exceeded for client %s", modelID, clientID) @@ -1022,6 +1047,9 @@ type ClientModelProjection struct { ModelID string Suspended bool SuspendReason string + // SuspendUntil optionally bounds the suspension in time. Once it passes, + // the client counts as available again without requiring re-projection. + SuspendUntil time.Time QuotaExceeded bool } @@ -1089,7 +1117,7 @@ func (r *ModelRegistry) ApplyClientModelProjections(clientID string, epoch uint6 r.clientGenerations[clientID] = generation r.ensureAvailableModelsCacheLocked() - now := time.Now() + now := r.now() changed := false for _, proj := range projections { modelID := strings.TrimSpace(proj.ModelID) @@ -1109,8 +1137,22 @@ func (r *ModelRegistry) ApplyClientModelProjections(clientID string, epoch uint6 if registration.SuspendedClients == nil { registration.SuspendedClients = make(map[string]string) } - if currentReason, ok := registration.SuspendedClients[clientID]; !ok || currentReason != proj.SuspendReason { + if registration.SuspendedClientDeadlines == nil { + registration.SuspendedClientDeadlines = make(map[string]time.Time) + } + deadlineChanged := false + if currentDeadline, ok := registration.SuspendedClientDeadlines[clientID]; ok { + deadlineChanged = !proj.SuspendUntil.Equal(currentDeadline) + } else { + deadlineChanged = !proj.SuspendUntil.IsZero() + } + if currentReason, ok := registration.SuspendedClients[clientID]; !ok || currentReason != proj.SuspendReason || deadlineChanged { registration.SuspendedClients[clientID] = proj.SuspendReason + if proj.SuspendUntil.IsZero() { + delete(registration.SuspendedClientDeadlines, clientID) + } else { + registration.SuspendedClientDeadlines[clientID] = proj.SuspendUntil + } registration.LastUpdated = now changed = true } @@ -1118,6 +1160,9 @@ func (r *ModelRegistry) ApplyClientModelProjections(clientID string, epoch uint6 if registration.SuspendedClients != nil { if _, ok := registration.SuspendedClients[clientID]; ok { delete(registration.SuspendedClients, clientID) + if registration.SuspendedClientDeadlines != nil { + delete(registration.SuspendedClientDeadlines, clientID) + } registration.LastUpdated = now changed = true } @@ -1348,11 +1393,20 @@ func (r *ModelRegistry) SuspendClientModel(clientID, modelID, reason string) { if registration.SuspendedClients == nil { registration.SuspendedClients = make(map[string]string) } - if _, already := registration.SuspendedClients[clientID]; already { + now := r.now() + _, _, active := clientSuspensionActive(registration, clientID, now) + hadDeadline := false + if registration.SuspendedClientDeadlines != nil { + _, hadDeadline = registration.SuspendedClientDeadlines[clientID] + } + if active && !hadDeadline { return } registration.SuspendedClients[clientID] = reason - registration.LastUpdated = time.Now() + if registration.SuspendedClientDeadlines != nil { + delete(registration.SuspendedClientDeadlines, clientID) + } + registration.LastUpdated = now r.invalidateAvailableModelsCacheLocked() if reason != "" { log.Debugf("Suspended client %s for model %s: %s", clientID, modelID, reason) @@ -1381,7 +1435,10 @@ func (r *ModelRegistry) ResumeClientModel(clientID, modelID string) { return } delete(registration.SuspendedClients, clientID) - registration.LastUpdated = time.Now() + if registration.SuspendedClientDeadlines != nil { + delete(registration.SuspendedClientDeadlines, clientID) + } + registration.LastUpdated = r.now() r.invalidateAvailableModelsCacheLocked() log.Debugf("Resumed client %s for model %s", clientID, modelID) } @@ -1419,15 +1476,16 @@ func (r *ModelRegistry) IsModelSuspendedForClient(clientID, modelID string) bool return false } + now := r.now() r.mutex.RLock() defer r.mutex.RUnlock() registration, exists := r.models[modelID] - if !exists || registration == nil || registration.SuspendedClients == nil { + if !exists || registration == nil { return false } - _, suspended := registration.SuspendedClients[clientID] - return suspended + _, _, active := clientSuspensionActive(registration, clientID, now) + return active } // IsModelQuotaExceededForClient reports whether a model is currently marked quota exceeded for a specific client. @@ -1456,7 +1514,7 @@ func (r *ModelRegistry) IsModelQuotaExceededForClient(clientID, modelID string) // Returns: // - []map[string]any: List of available models in the requested format func (r *ModelRegistry) GetAvailableModels(handlerType string) []map[string]any { - now := time.Now() + now := r.now() r.mutex.RLock() if cache, ok := r.availableModelsCache[handlerType]; ok && (cache.expiresAt.IsZero() || now.Before(cache.expiresAt)) { @@ -1493,6 +1551,34 @@ func isQuotaSuspension(reason string) bool { return strings.EqualFold(reason, "quota") || strings.EqualFold(reason, "credential_quota") } +func clientSuspensionActive(registration *ModelRegistration, clientID string, now time.Time) (string, time.Time, bool) { + if registration == nil || len(registration.SuspendedClients) == 0 { + return "", time.Time{}, false + } + reason, ok := registration.SuspendedClients[clientID] + if !ok { + return "", time.Time{}, false + } + if deadline, hasDeadline := registration.SuspendedClientDeadlines[clientID]; hasDeadline && !deadline.IsZero() { + if !now.Before(deadline) { + return "", time.Time{}, false + } + return reason, deadline, true + } + return reason, time.Time{}, true +} + +func clientQuotaRecoveryAt(registration *ModelRegistration, clientID string, quotaTime *time.Time) (time.Time, bool) { + if registration == nil || quotaTime == nil { + return time.Time{}, false + } + recoveryAt := quotaTime.Add(modelQuotaExceededWindow) + if deadline, hasDeadline := registration.SuspendedClientDeadlines[clientID]; hasDeadline && !deadline.IsZero() && deadline.Before(recoveryAt) { + recoveryAt = deadline + } + return recoveryAt, true +} + func modelRegistrationAvailability(registration *ModelRegistration, now time.Time) (bool, time.Time) { if registration == nil { return false, time.Time{} @@ -1501,11 +1587,11 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim availableClients := registration.Count expiredClients := 0 var expiresAt time.Time - for _, quotaTime := range registration.QuotaExceededClients { - if quotaTime == nil { + for clientID, quotaTime := range registration.QuotaExceededClients { + recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime) + if !ok { continue } - recoveryAt := quotaTime.Add(modelQuotaExceededWindow) if now.Before(recoveryAt) { expiredClients++ if expiresAt.IsZero() || recoveryAt.Before(expiresAt) { @@ -1518,8 +1604,15 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim otherSuspended := 0 quotaAndSuspended := 0 if registration.SuspendedClients != nil { - for clientID, reason := range registration.SuspendedClients { - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + for clientID := range registration.SuspendedClients { + reason, deadline, active := clientSuspensionActive(registration, clientID, now) + if !active { + continue + } + if !deadline.IsZero() && (expiresAt.IsZero() || deadline.Before(expiresAt)) { + expiresAt = deadline + } + if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); ok && now.Before(recoveryAt) { quotaAndSuspended++ } if isQuotaSuspension(reason) { @@ -1543,7 +1636,7 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim // GetAvailableModelInfos returns cloned metadata for all currently available models. func (r *ModelRegistry) GetAvailableModelInfos() []*ModelInfo { - now := time.Now() + now := r.now() r.mutex.RLock() defer r.mutex.RUnlock() @@ -1681,7 +1774,7 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn return nil } - now := time.Now() + now := r.now() result := make([]*ModelInfo, 0, len(providerModels)) for modelID, entry := range providerModels { @@ -1703,20 +1796,24 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { continue } - if quotaTime != nil && now.Sub(*quotaTime) < modelQuotaExceededWindow { + if recoveryAt, okRecovery := clientQuotaRecoveryAt(registration, clientID, quotaTime); okRecovery && now.Before(recoveryAt) { expiredClients++ } } } if registration.SuspendedClients != nil { - for clientID, reason := range registration.SuspendedClients { + for clientID := range registration.SuspendedClients { if clientID == "" { continue } if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { continue } - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + reason, _, active := clientSuspensionActive(registration, clientID, now) + if !active { + continue + } + if recoveryAt, okRecovery := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); okRecovery && now.Before(recoveryAt) { quotaAndSuspended++ } if isQuotaSuspension(reason) { @@ -1759,18 +1856,21 @@ func (r *ModelRegistry) GetModelCount(modelID string) int { defer r.mutex.RUnlock() if registration, exists := r.models[modelID]; exists { - now := time.Now() + now := r.now() // Count clients that have exceeded quota but haven't recovered yet expiredClients := 0 - for _, quotaTime := range registration.QuotaExceededClients { - if quotaTime != nil && now.Sub(*quotaTime) < modelQuotaExceededWindow { + for clientID, quotaTime := range registration.QuotaExceededClients { + if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime); ok && now.Before(recoveryAt) { expiredClients++ } } suspendedClients := 0 for clientID := range registration.SuspendedClients { - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + if _, _, active := clientSuspensionActive(registration, clientID, now); !active { + continue + } + if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); ok && now.Before(recoveryAt) { continue } suspendedClients++ @@ -1991,22 +2091,33 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) } } -// CleanupExpiredQuotas removes expired quota tracking entries +// CleanupExpiredQuotas removes expired quota and bounded suspension entries func (r *ModelRegistry) CleanupExpiredQuotas() { r.mutex.Lock() defer r.mutex.Unlock() - now := time.Now() + now := r.now() invalidated := false for modelID, registration := range r.models { + if registration == nil { + continue + } for clientID, quotaTime := range registration.QuotaExceededClients { - if quotaTime != nil && now.Sub(*quotaTime) >= modelQuotaExceededWindow { + if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime); ok && !now.Before(recoveryAt) { delete(registration.QuotaExceededClients, clientID) invalidated = true log.Debugf("Cleaned up expired quota tracking for model %s, client %s", modelID, clientID) } } + for clientID, deadline := range registration.SuspendedClientDeadlines { + if !deadline.IsZero() && !now.Before(deadline) { + delete(registration.SuspendedClientDeadlines, clientID) + delete(registration.SuspendedClients, clientID) + invalidated = true + log.Debugf("Cleaned up expired suspension for model %s, client %s", modelID, clientID) + } + } } if invalidated { r.invalidateAvailableModelsCacheLocked() diff --git a/internal/registry/model_registry_transient_suspend_deadline_test.go b/internal/registry/model_registry_transient_suspend_deadline_test.go new file mode 100644 index 00000000000..c22398e5a76 --- /dev/null +++ b/internal/registry/model_registry_transient_suspend_deadline_test.go @@ -0,0 +1,194 @@ +package registry + +import ( + "testing" + "time" +) + +func listedIDs(models []map[string]any) map[string]bool { + out := make(map[string]bool, len(models)) + for _, m := range models { + if id, _ := m["id"].(string); id != "" { + out[id] = true + } + } + return out +} + +func infoIDs(infos []*ModelInfo) map[string]bool { + out := make(map[string]bool, len(infos)) + for _, info := range infos { + if info != nil && info.ID != "" { + out[info.ID] = true + } + } + return out +} + +func TestModelRegistry_TransientSuspendDeadline_AutoExpiresCacheAndAllViews(t *testing.T) { + baseNow := time.Date(2026, 11, 11, 1, 5, 0, 0, time.UTC) + mockNow := baseNow + + r := newTestModelRegistry() + r.nowFunc = func() time.Time { return mockNow } + + r.RegisterClient("xai-auth", "xai", []*ModelInfo{ + {ID: "grok-4.7", OwnedBy: "xai", Created: baseNow.Unix()}, + {ID: "grok-4.7-cp", OwnedBy: "xai", Created: baseNow.Unix()}, + {ID: "grok-4.6", OwnedBy: "xai", Created: baseNow.Unix()}, + }) + epoch := r.ClientRegistrationEpoch("xai-auth") + deadline := baseNow.Add(time.Minute) + + if !r.ApplyClientModelProjections("xai-auth", epoch, 1, []ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: deadline}, + {ModelID: "grok-4.7-cp", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: deadline}, + {ModelID: "grok-4.6"}, + }) { + t.Fatal("expected initial projection to apply") + } + + // During cooldown: grok-4.7 and grok-4.7-cp are hidden, grok-4.6 remains listed. + duringOpenAI := listedIDs(r.GetAvailableModels("openai")) + if duringOpenAI["grok-4.7"] || duringOpenAI["grok-4.7-cp"] || !duringOpenAI["grok-4.6"] { + t.Fatalf("unexpected /v1/models set during transient cooldown: %v", duringOpenAI) + } + if cached := r.availableModelsCache["openai"]; !cached.expiresAt.Equal(deadline) { + t.Fatalf("availableModelsCache expiresAt = %v, want %v", cached.expiresAt, deadline) + } + if infos := infoIDs(r.GetAvailableModelInfos()); infos["grok-4.7"] || infos["grok-4.7-cp"] || !infos["grok-4.6"] { + t.Fatalf("unexpected GetAvailableModelInfos during transient cooldown: %v", infos) + } + if provInfos := infoIDs(r.GetAvailableModelsByProvider("xai")); provInfos["grok-4.7"] || provInfos["grok-4.7-cp"] || !provInfos["grok-4.6"] { + t.Fatalf("unexpected GetAvailableModelsByProvider during transient cooldown: %v", provInfos) + } + if got := r.GetModelCount("grok-4.7"); got != 0 { + t.Fatalf("GetModelCount(grok-4.7) during cooldown = %d, want 0", got) + } + if !r.IsModelSuspendedForClient("xai-auth", "grok-4.7") { + t.Fatal("expected IsModelSuspendedForClient to report true during active transient cooldown") + } + + // Still active 1ns before deadline. + mockNow = deadline.Add(-time.Nanosecond) + stillCooling := listedIDs(r.GetAvailableModels("openai")) + if stillCooling["grok-4.7"] || stillCooling["grok-4.7-cp"] { + t.Fatalf("models listed before deadline elapsed: %v", stillCooling) + } + + // Exact deadline instant with NO re-projection, NO request, and NO refresh. + mockNow = deadline + afterOpenAI := listedIDs(r.GetAvailableModels("openai")) + if !afterOpenAI["grok-4.7"] || !afterOpenAI["grok-4.7-cp"] || !afterOpenAI["grok-4.6"] { + t.Fatalf("expected all models listed once transient cooldown expired without re-projection, got %v", afterOpenAI) + } + if infos := infoIDs(r.GetAvailableModelInfos()); !infos["grok-4.7"] || !infos["grok-4.7-cp"] { + t.Fatalf("expected GetAvailableModelInfos to list models after deadline, got %v", infos) + } + if provInfos := infoIDs(r.GetAvailableModelsByProvider("xai")); !provInfos["grok-4.7"] || !provInfos["grok-4.7-cp"] { + t.Fatalf("expected GetAvailableModelsByProvider to list models after deadline, got %v", provInfos) + } + if got := r.GetModelCount("grok-4.7"); got != 1 { + t.Fatalf("GetModelCount(grok-4.7) after cooldown = %d, want 1", got) + } + if got := r.GetModelCount("grok-4.7-cp"); got != 1 { + t.Fatalf("GetModelCount(grok-4.7-cp) after cooldown = %d, want 1", got) + } + if r.IsModelSuspendedForClient("xai-auth", "grok-4.7") { + t.Fatal("expected IsModelSuspendedForClient to report false once transient cooldown expired") + } + + r.CleanupExpiredQuotas() + if _, exists := r.models["grok-4.7"].SuspendedClients["xai-auth"]; exists { + t.Fatal("expected CleanupExpiredQuotas to prune expired transient suspension entry") + } + if _, exists := r.models["grok-4.7"].SuspendedClientDeadlines["xai-auth"]; exists { + t.Fatal("expected CleanupExpiredQuotas to prune expired transient suspension deadline") + } +} + +func TestModelRegistry_OpenEndedSuspensionAndManualOverridePersist(t *testing.T) { + baseNow := time.Date(2026, 11, 11, 2, 0, 0, 0, time.UTC) + mockNow := baseNow + + r := newTestModelRegistry() + r.nowFunc = func() time.Time { return mockNow } + + r.RegisterClient("xai-auth", "xai", []*ModelInfo{{ID: "grok-4.7", OwnedBy: "xai"}}) + epoch := r.ClientRegistrationEpoch("xai-auth") + deadline := baseNow.Add(time.Minute) + + // Start with a bounded transient suspension. + if !r.ApplyClientModelProjections("xai-auth", epoch, 1, []ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "transient 408", SuspendUntil: deadline}, + }) { + t.Fatal("initial transient projection rejected") + } + + // Override with manual open-ended SuspendClientModel while already in SuspendedClients. + r.SuspendClientModel("xai-auth", "grok-4.7", "manual") + + mockNow = baseNow.Add(24 * time.Hour) + if listed := listedIDs(r.GetAvailableModels("openai")); listed["grok-4.7"] { + t.Fatalf("manual SuspendClientModel leaked model after earlier transient deadline: %v", listed) + } + if !r.IsModelSuspendedForClient("xai-auth", "grok-4.7") { + t.Fatal("expected manual open-ended suspension to stay active") + } + + // Resume, then apply an open-ended projection (zero SuspendUntil). + r.ResumeClientModel("xai-auth", "grok-4.7") + if !listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] { + t.Fatal("expected ResumeClientModel to restore listing") + } + + if !r.ApplyClientModelProjections("xai-auth", epoch, 2, []ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "disabled"}, + }) { + t.Fatal("open-ended projection rejected") + } + mockNow = baseNow.Add(48 * time.Hour) + if listed := listedIDs(r.GetAvailableModels("openai")); listed["grok-4.7"] { + t.Fatalf("open-ended projection leaked model: %v", listed) + } +} + +func TestModelRegistry_ShortQuotaDeadlineRecoversWithDisabledBackup(t *testing.T) { + baseNow := time.Date(2026, 11, 11, 3, 0, 0, 0, time.UTC) + mockNow := baseNow + + r := newTestModelRegistry() + r.nowFunc = func() time.Time { return mockNow } + + model := []*ModelInfo{{ID: "grok-4.7", OwnedBy: "xai"}} + r.RegisterClient("c1-quota", "xai", model) + r.RegisterClient("c2-disabled", "xai", model) + + epoch1 := r.ClientRegistrationEpoch("c1-quota") + epoch2 := r.ClientRegistrationEpoch("c2-disabled") + + shortQuotaDeadline := baseNow.Add(time.Minute) + if !r.ApplyClientModelProjections("c1-quota", epoch1, 1, []ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "quota", SuspendUntil: shortQuotaDeadline, QuotaExceeded: true}, + }) { + t.Fatal("quota projection rejected") + } + if !r.ApplyClientModelProjections("c2-disabled", epoch2, 1, []ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "disabled"}, + }) { + t.Fatal("disabled projection rejected") + } + + if listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] { + t.Fatal("expected model hidden while c1 is in 1m quota cooldown and c2 is disabled") + } + + // Advance past the 1m quota deadline (still within the default 5m modelQuotaExceededWindow). + mockNow = shortQuotaDeadline.Add(time.Second) + if !listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] { + t.Fatal("expected model listed once c1's 1m quota deadline passed") + } + if got := r.GetModelCount("grok-4.7"); got != 1 { + t.Fatalf("GetModelCount(grok-4.7) after c1 quota recovery = %d, want 1", got) + } +} diff --git a/internal/registry/models/models.json b/internal/registry/models/models.json index eec7606848e..86272eb6e51 100644 --- a/internal/registry/models/models.json +++ b/internal/registry/models/models.json @@ -1816,6 +1816,39 @@ "web_search": true } }, + { + "id": "gemini-4-carbon", + "object": "model", + "owned_by": "google", + "type": "gemini", + "display_name": "Gemini 4 Carbon", + "name": "gemini-4-carbon", + "description": "Gemini 4 Carbon", + "context_length": 1048576, + "max_completion_tokens": 1048576, + "thinking": { + "min": 128, + "max": 65535, + "dynamic_allowed": true, + "levels": [ + "low", + "medium", + "high" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "audio", + "video" + ], + "supportedOutputModalities": [ + "text" + ], + "native_capabilities": { + "web_search": true + } + }, { "id": "gemini-omni-1.2-flash", "object": "model", @@ -4432,6 +4465,40 @@ "web_search": true } }, + { + "id": "gemini-4-carbon-high", + "object": "model", + "owned_by": "antigravity", + "type": "antigravity", + "supports_web_search": true, + "display_name": "Gemini 4 Carbon", + "name": "gemini-4-carbon-high", + "description": "Gemini 4 Carbon (High)", + "context_length": 1048576, + "max_completion_tokens": 1048576, + "thinking": { + "min": 1, + "max": 65535, + "dynamic_allowed": true, + "levels": [ + "low", + "medium", + "high" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "audio", + "video" + ], + "supportedOutputModalities": [ + "text" + ], + "native_capabilities": { + "web_search": true + } + }, { "id": "gemini-3.8-flash-high", "object": "model", @@ -5013,6 +5080,40 @@ } ], "meta": [ + { + "id": "muse-ember-1.0", + "object": "model", + "created": 1794182400, + "owned_by": "meta", + "type": "meta", + "display_name": "Muse Ember 1.0", + "name": "muse-ember-1.0", + "description": "Meta's smallest and most efficient model designed for complex reasoning execution, tool calls, coding subtasks, routing, extraction, and long document parsing.", + "context_length": 1048576, + "max_completion_tokens": 524288, + "thinking": { + "levels": [ + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "video", + "audio" + ], + "supportedOutputModalities": [ + "text" + ], + "native_capabilities": { + "web_search": true + } + }, { "id": "muse-wonder-1.0", "object": "model", diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index bd0c33ecd63..dd7d0a5c8a1 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -13,6 +13,7 @@ import ( "fmt" "net/http" "strings" + "sync" "time" "github.com/MoeForge/moego/ojson" @@ -50,7 +51,10 @@ const ( // AntigravityExecutor proxies requests to the antigravity upstream. type AntigravityExecutor struct { - cfg *config.Config + cfg *config.Config + matcherMu sync.RWMutex + lastWordsKey string + cachedMatcher *helps.SensitiveWordMatcher } // NewAntigravityExecutor creates a new Antigravity executor instance. @@ -64,19 +68,50 @@ func NewAntigravityExecutor(cfg *config.Config) *AntigravityExecutor { return &AntigravityExecutor{cfg: cfg} } +func (e *AntigravityExecutor) getSensitiveWordMatcher() *helps.SensitiveWordMatcher { + if e == nil || e.cfg == nil { + return nil + } + words := e.cfg.Antigravity.SensitiveWords + replacements := e.cfg.Antigravity.SensitiveWordsReplace + if len(words) == 0 && len(replacements) == 0 { + return nil + } + key := helps.SensitiveWordConfigCacheKey(words, replacements) + e.matcherMu.RLock() + if e.lastWordsKey == key { + m := e.cachedMatcher + e.matcherMu.RUnlock() + return m + } + e.matcherMu.RUnlock() + + e.matcherMu.Lock() + defer e.matcherMu.Unlock() + if e.lastWordsKey == key { + return e.cachedMatcher + } + e.lastWordsKey = key + e.cachedMatcher = helps.BuildAntigravitySensitiveWordMatcher(words, replacements) + return e.cachedMatcher +} + func (e *AntigravityExecutor) obfuscateSensitiveWords(payload []byte) []byte { - if e == nil || e.cfg == nil || len(e.cfg.Antigravity.SensitiveWords) == 0 { + matcher := e.getSensitiveWordMatcher() + if matcher == nil { return payload } - matcher := helps.BuildSensitiveWordMatcher(e.cfg.Antigravity.SensitiveWords) return helps.ObfuscateSensitiveWordsInSystemInstruction(payload, matcher) } func (e *AntigravityExecutor) obfuscateSensitiveWordsDoc(doc *ojson.Doc) bool { - if doc == nil || e == nil || e.cfg == nil || len(e.cfg.Antigravity.SensitiveWords) == 0 { + if doc == nil { + return false + } + matcher := e.getSensitiveWordMatcher() + if matcher == nil { return false } - matcher := helps.BuildSensitiveWordMatcher(e.cfg.Antigravity.SensitiveWords) return helps.ObfuscateSensitiveWordsInSystemInstructionDoc(doc, matcher) } diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go index bb7ab2fcce4..7a55bc55557 100644 --- a/internal/runtime/executor/antigravity_executor_execute.go +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -152,12 +152,12 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) err = errRead return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if httpResp.StatusCode == http.StatusTooManyRequests { decision := decideAntigravity429(bodyBytes) @@ -373,6 +373,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { @@ -386,7 +387,6 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * err = errRead return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if httpResp.StatusCode == http.StatusTooManyRequests { decision := decideAntigravity429(bodyBytes) diff --git a/internal/runtime/executor/antigravity_executor_signature_test.go b/internal/runtime/executor/antigravity_executor_signature_test.go index bca7bc24bfe..0a967feda23 100644 --- a/internal/runtime/executor/antigravity_executor_signature_test.go +++ b/internal/runtime/executor/antigravity_executor_signature_test.go @@ -236,8 +236,8 @@ func TestAntigravitySensitiveWordsObfuscatesSystemInstructionOnly(t *testing.T) payload := []byte(`{"request":{"systemInstruction":{"parts":[{"text":"Use proxy safely"}]},"contents":[{"role":"user","parts":[{"text":"proxy remains unchanged"}]}]}}`) got := executor.obfuscateSensitiveWords(payload) - if systemText := ojson.GetBytes(got, "request.systemInstruction.parts.0.text").String(); systemText != "Use p\u200Broxy safely" { - t.Fatalf("system instruction = %q, want zero-width obfuscation", systemText) + if systemText := ojson.GetBytes(got, "request.systemInstruction.parts.0.text").String(); systemText != "Use p\u034Froxy safely" { + t.Fatalf("system instruction = %q, want CGJ obfuscation", systemText) } if contentText := ojson.GetBytes(got, "request.contents.0.parts.0.text").String(); contentText != "proxy remains unchanged" { t.Fatalf("content text = %q, want unchanged", contentText) @@ -288,12 +288,180 @@ func TestAntigravityStreamObfuscatesSensitiveSystemInstruction(t *testing.T) { body := <-captured got := ojson.GetBytes(body, "request.systemInstruction.parts.0.text").String() - want := "You are H\u200Bermes Agent, an intelligent AI assistant created by N\u200Bous Research." + want := "You are H\u034Fermes Agent, an intelligent AI assistant created by N\u034Fous R\u034Fesearch." if got != want { t.Fatalf("system instruction = %q, want %q; body=%s", got, want, body) } } +func TestAntigravityExecuteObfuscatesCodexIdentityWithCGJ(t *testing.T) { + captured := make(chan []byte, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read request body: %v", errRead) + return + } + captured <- body + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"OK"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1,"totalTokenCount":11}}}`)) + })) + defer server.Close() + + executor := NewAntigravityExecutor(&config.Config{ + Antigravity: config.AntigravityConfig{SensitiveWords: []string{"codex", "GPT-5"}}, + RequestRetry: 1, + }) + auth := &cliproxyauth.Auth{ + Metadata: map[string]any{ + "access_token": "token-123", + "expired": time.Now().Add(24 * time.Hour).Format(time.RFC3339), + "project_id": "project-1", + }, + Attributes: map[string]string{"base_url": server.URL}, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + } + + cases := []struct { + name string + payload string + }{ + { + name: "instructions_field", + payload: `{"model":"gemini-4-carbon-high","instructions":"You are Codex, an agent based on GPT-5.","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Reply only OK."}]}],"max_output_tokens":128,"stream":false}`, + }, + { + name: "initial_developer_message", + payload: `{"model":"gemini-4-carbon-high","input":[{"type":"message","role":"developer","content":[{"type":"input_text","text":"You are Codex, an agent based on GPT-5."}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"Reply only OK."}]}],"max_output_tokens":128,"stream":false}`, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, errExec := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gemini-4-carbon-high", + Payload: []byte(tc.payload), + }, opts) + if errExec != nil { + t.Fatalf("Execute() error = %v", errExec) + } + body := <-captured + got := ojson.GetBytes(body, "request.systemInstruction.parts.0.text").String() + want := "You are C\u034Fodex, an agent based on G\u034FPT-5." + if got != want { + t.Fatalf("systemInstruction = %q, want %q", got, want) + } + if strings.Contains(got, "\u200B") { + t.Fatalf("systemInstruction = %q, must not use U+200B", got) + } + }) + } +} + +func TestAntigravitySensitiveWordsReplaceFromYAML(t *testing.T) { + cfg, errParse := config.ParseConfigBytes([]byte(`antigravity: + sensitive-words: + - "Codex" + - "GPT-5" + sensitive-words-replace: + "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant." +`)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + executor := NewAntigravityExecutor(cfg) + payload := []byte(`{"request":{"systemInstruction":{"parts":[{"text":"You are Codex, an agent based on GPT-5. Avoid mentioning Codex."}]},"contents":[{"role":"user","parts":[{"text":"You are Codex, an agent based on GPT-5."}]}]}}`) + + got := executor.obfuscateSensitiveWords(payload) + systemText := ojson.GetBytes(got, "request.systemInstruction.parts.0.text").String() + wantSystem := "You are a helpful coding assistant. Avoid mentioning C\u034Fodex." + if systemText != wantSystem { + t.Fatalf("systemInstruction = %q, want %q", systemText, wantSystem) + } + contentText := ojson.GetBytes(got, "request.contents.0.parts.0.text").String() + if contentText != "You are Codex, an agent based on GPT-5." { + t.Fatalf("contents = %q, want unchanged", contentText) + } +} + +func TestAntigravitySensitiveWordsReplaceOnlyPreservesToolsAndCountTokens(t *testing.T) { + captured := make(chan []byte, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read request body: %v", errRead) + return + } + captured <- body + w.Header().Set("Content-Type", "application/json") + if strings.Contains(r.URL.Path, "countTokens") { + _, _ = w.Write([]byte(`{"totalTokens":42}`)) + return + } + _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"OK"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1,"totalTokenCount":11}}}`)) + })) + defer server.Close() + + cfg := &config.Config{ + Antigravity: config.AntigravityConfig{ + SensitiveWordsReplace: map[string]string{ + "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant.", + }, + }, + RequestRetry: 1, + } + executor := NewAntigravityExecutor(cfg) + m1 := executor.getSensitiveWordMatcher() + m2 := executor.getSensitiveWordMatcher() + if m1 == nil || m1 != m2 { + t.Fatal("expected cached matcher reuse when replacements are unchanged") + } + + auth := &cliproxyauth.Auth{ + Metadata: map[string]any{ + "access_token": "token-123", + "expired": time.Now().Add(24 * time.Hour).Format(time.RFC3339), + "project_id": "project-1", + }, + Attributes: map[string]string{"base_url": server.URL}, + } + req := cliproxyexecutor.Request{ + Model: "gemini-4-carbon-high", + Payload: []byte(`{"model":"gemini-4-carbon-high","instructions":"You are Codex, an agent based on GPT-5. Keep extra instructions intact.","tools":[{"type":"function","name":"run_codex_check","description":"Codex diagnostic tool","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}],"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"You are Codex, an agent based on GPT-5."}]}],"max_output_tokens":128,"stream":false}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + } + + if _, errExec := executor.Execute(context.Background(), auth, req, opts); errExec != nil { + t.Fatalf("Execute() error = %v", errExec) + } + execBody := <-captured + wantInstruction := "You are a helpful coding assistant. Keep extra instructions intact." + if got := ojson.GetBytes(execBody, "request.systemInstruction.parts.0.text").String(); got != wantInstruction { + t.Fatalf("Execute systemInstruction = %q, want %q", got, wantInstruction) + } + if gotToolDesc := ojson.GetBytes(execBody, "request.tools.0.functionDeclarations.0.description").String(); gotToolDesc != "Codex diagnostic tool" { + t.Fatalf("tool description = %q, want unchanged", gotToolDesc) + } + if gotUser := ojson.GetBytes(execBody, "request.contents.0.parts.0.text").String(); gotUser != "You are Codex, an agent based on GPT-5." { + t.Fatalf("user content = %q, want unchanged", gotUser) + } + + if _, errCount := executor.CountTokens(context.Background(), auth, req, opts); errCount != nil { + t.Fatalf("CountTokens() error = %v", errCount) + } + countBody := <-captured + if got := ojson.GetBytes(countBody, "request.systemInstruction.parts.0.text").String(); got != wantInstruction { + t.Fatalf("CountTokens systemInstruction = %q, want %q", got, wantInstruction) + } +} + func TestAntigravityStreamDoesNotEmitSyntheticTerminalOnReadError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index 31961e4385b..835b3279b37 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -151,6 +151,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { @@ -164,7 +165,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya err = errRead return nil, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if httpResp.StatusCode == http.StatusTooManyRequests { decision := decideAntigravity429(bodyBytes) diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index bde7323e42b..eb8fc8f8d5e 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -140,11 +140,11 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return cliproxyexecutor.Response{}, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { count := ojson.GetBytesNoCopy(bodyBytes, "totalTokens").Int() diff --git a/internal/runtime/executor/antigravity_issue7000_boundary_test.go b/internal/runtime/executor/antigravity_issue7000_boundary_test.go index 538e96625de..f4c9c88974b 100644 --- a/internal/runtime/executor/antigravity_issue7000_boundary_test.go +++ b/internal/runtime/executor/antigravity_issue7000_boundary_test.go @@ -241,7 +241,7 @@ func TestIssue7000AntigravityExecutionAndFinalPayloadRules(t *testing.T) { switch rules { case "none": assertIssue7000PromptSchema(t, request.Get("systemInstruction"), issue7000LiteralSchema) - if !strings.Contains(request.Get("systemInstruction.parts.0.text").String(), "a\u200bnthropic") { + if !strings.Contains(request.Get("systemInstruction.parts.0.text").String(), "a\u034fnthropic") { t.Errorf("original system instruction was lost or not obfuscated: %s", body) } for _, key := range []string{"responseSchema", "responseMimeType"} { diff --git a/internal/runtime/executor/apply_patch_repair_test.go b/internal/runtime/executor/apply_patch_repair_test.go index d1403004a4d..6ec34b9b78b 100644 --- a/internal/runtime/executor/apply_patch_repair_test.go +++ b/internal/runtime/executor/apply_patch_repair_test.go @@ -470,12 +470,20 @@ func TestApplyPatchRepairResponsesSourceTerminalNonStream(t *testing.T) { if len(response.Payload) != 0 { t.Errorf("failed source returned payload: %s", response.Payload) } - assertTask6PatchError(t, errExecute) + wantGatewayMsg := "Invalid apply_patch tool arguments received from upstream." + if provider == "xai" && mode == "empty" { + wantGatewayMsg = xaiEmptyStreamMessage + if status, okStatus := errExecute.(interface{ StatusCode() int }); !okStatus || status.StatusCode() != http.StatusBadGateway || errExecute.Error() != xaiEmptyStreamMessage { + t.Fatalf("expected xai empty stream 502, got %T %v", errExecute, errExecute) + } + } else { + assertTask6PatchError(t, errExecute) + } checkUsage() auth.ID += "-gateway" checkGatewayUsage := task6CaptureFailureUsage(t, auth.ID) result := task6Gateway(t, exec, auth, false) - if result.Code != http.StatusBadGateway || !strings.Contains(result.Body.String(), "Invalid apply_patch tool arguments received from upstream.") { + if result.Code != http.StatusBadGateway || !strings.Contains(result.Body.String(), wantGatewayMsg) { t.Errorf("non-stream gateway did not return clean 502: %d %s", result.Code, result.Body) } checkGatewayUsage() @@ -615,6 +623,10 @@ func TestApplyPatchRepairOrdinaryEmptyHTTPPassthrough(t *testing.T) { if errExecute != nil || len(response.Payload) != 0 { t.Fatalf("ordinary Kimi response changed: %s %v", response.Payload, errExecute) } + } else if provider == "xai" { + if status, okStatus := errExecute.(interface{ StatusCode() int }); !okStatus || status.StatusCode() != http.StatusBadGateway || errExecute.Error() != xaiEmptyStreamMessage { + t.Fatalf("ordinary xAI empty stream changed: %v", errExecute) + } } else if status, okStatus := errExecute.(interface{ StatusCode() int }); !okStatus || status.StatusCode() != http.StatusRequestTimeout { t.Fatalf("ordinary disconnect changed: %v", errExecute) } diff --git a/internal/runtime/executor/claude_executor_cloaking.go b/internal/runtime/executor/claude_executor_cloaking.go index 1ebac79a1cd..469add50dd5 100644 --- a/internal/runtime/executor/claude_executor_cloaking.go +++ b/internal/runtime/executor/claude_executor_cloaking.go @@ -82,7 +82,9 @@ func getCloakConfigFromAuth(auth *cliproxyauth.Auth) (cloakMode string, strictMo strictMode = strings.EqualFold(lookupCloakAttr("cloak_strict_mode"), "true") - if wordsStr := lookupCloakAttr("cloak_sensitive_words"); wordsStr != "" { + if metaWords := claudeauth.ReadMetadataStringSlice(&auth.Metadata, "cloak_sensitive_words"); len(metaWords) > 0 { + sensitiveWords = metaWords + } else if wordsStr := lookupCloakAttr("cloak_sensitive_words"); wordsStr != "" { sensitiveWords = strings.Split(wordsStr, ",") for i := range sensitiveWords { sensitiveWords[i] = strings.TrimSpace(sensitiveWords[i]) @@ -94,6 +96,44 @@ func getCloakConfigFromAuth(auth *cliproxyauth.Auth) (cloakMode string, strictMo return cloakMode, strictMode, sensitiveWords, cacheUserID } +func getCloakReplacementsFromAuth(auth *cliproxyauth.Auth) map[string]string { + if auth == nil { + return nil + } + parseJSONStringMap := func(raw string) map[string]string { + trimmed := strings.TrimSpace(raw) + if trimmed == "" || !strings.HasPrefix(trimmed, "{") || !ojson.Valid(trimmed) { + return nil + } + parsed := ojson.Parse(trimmed) + if !parsed.IsObject() { + return nil + } + out := make(map[string]string) + parsed.ForEach(func(k, v ojson.Result) bool { + if v.Type == ojson.String { + out[k.String()] = v.String() + } + return true + }) + return config.NormalizeSensitiveWordsReplace(out) + } + if auth.Attributes != nil { + if raw := auth.Attributes["cloak_sensitive_words_replace"]; strings.TrimSpace(raw) != "" { + if parsed := parseJSONStringMap(raw); len(parsed) > 0 { + return parsed + } + } + } + if fromMap := claudeauth.ReadMetadataStringMap(&auth.Metadata, "cloak_sensitive_words_replace"); len(fromMap) > 0 { + return config.NormalizeSensitiveWordsReplace(fromMap) + } + if raw := claudeauth.ReadMetadataString(&auth.Metadata, "cloak_sensitive_words_replace"); strings.TrimSpace(raw) != "" { + return parseJSONStringMap(raw) + } + return nil +} + // injectFakeUserID generates and injects a fake user ID into the request metadata. // When useCache is false, a new user ID is generated for every call. func injectFakeUserID(ctx context.Context, payload []byte, apiKey string, useCache bool) ([]byte, error) { @@ -2300,9 +2340,18 @@ type claudeWirePolicy struct { } type claudeCloakSettings struct { - strictMode bool - sensitiveWords []string - cacheUserID bool + strictMode bool + sensitiveWords []string + sensitiveWordsReplace map[string]string + cacheUserID bool + serverSideFallback bool +} + +func (s claudeCloakSettings) buildSensitiveWordMatcher() *helps.SensitiveWordMatcher { + if len(s.sensitiveWords) == 0 && len(s.sensitiveWordsReplace) == 0 { + return nil + } + return helps.BuildSensitiveWordMatcherWithReplacements(s.sensitiveWords, s.sensitiveWordsReplace) } func isClaudeCodeIdentitySystemText(text string) bool { @@ -2415,15 +2464,18 @@ func resolveClaudeWirePolicy(cfg *config.Config, auth *cliproxyauth.Auth, apiKey } cloakCfg := resolveClaudeKeyCloakConfig(cfg, auth) attrMode, attrStrict, attrWords, attrCache := getCloakConfigFromAuth(auth) + attrWordsReplace := getCloakReplacementsFromAuth(auth) cloakMode := "auto" if cfg != nil && cfg.DisableClaudeCloakMode { cloakMode = "never" } settings := claudeCloakSettings{ - strictMode: attrStrict, - sensitiveWords: attrWords, - cacheUserID: attrCache, + strictMode: attrStrict, + sensitiveWords: attrWords, + sensitiveWordsReplace: attrWordsReplace, + cacheUserID: attrCache, + serverSideFallback: cfg != nil && cfg.Claude.ServerSideFallback, } explicitMode := strings.ToLower(strings.TrimSpace(attrMode)) if attrMode != "" { @@ -2440,13 +2492,16 @@ func resolveClaudeWirePolicy(cfg *config.Config, auth *cliproxyauth.Auth, apiKey if len(cloakCfg.SensitiveWords) > 0 { settings.sensitiveWords = cloakCfg.SensitiveWords } + if len(cloakCfg.SensitiveWordsReplace) > 0 { + settings.sensitiveWordsReplace = cloakCfg.SensitiveWordsReplace + } if cloakCfg.CacheUserID != nil { settings.cacheUserID = *cloakCfg.CacheUserID } } fp := resolveClaudeFingerprintPolicy(cfg, auth, apiKey) - cloakConfigured := cloakCfg != nil || attrMode != "" || attrStrict || len(attrWords) > 0 || attrCache + cloakConfigured := cloakCfg != nil || attrMode != "" || attrStrict || len(attrWords) > 0 || len(attrWordsReplace) > 0 || attrCache trimmedMode := strings.ToLower(strings.TrimSpace(cloakMode)) uncloak := !fp.ProfileClaudeCodeCLI && (explicitMode == "never" || explicitMode == "uncloak") explicitNever := trimmedMode == "never" && !uncloak @@ -2602,14 +2657,13 @@ func applyCloakingInternalDoc( turnOriginArgs..., ) - // In native Claude Code 2.1.258, claude-fable-5-1 requests carry: - // "fallbacks": [{"model": "claude-opus-5"}] - model := strings.ToLower(strings.TrimSpace(doc.Get("model").String())) - if isClaudeOpus55Model(model) && !isProbeOrHelper && !doc.Get("fallbacks").Exists() { - _ = doc.SetRawBytes("fallbacks", []byte(`[{"model":"claude-opus-4-8"}]`)) - } - if isClaudeFable51Model(model) && !isProbeOrHelper { - if !doc.Get("fallbacks").Exists() { + // Inject cross-model server-side refusal fallbacks only when explicitly + // enabled via upstream.claude.server-side-fallback. + if settings.serverSideFallback && !isProbeOrHelper && !doc.Get("fallbacks").Exists() { + model := strings.ToLower(strings.TrimSpace(doc.Get("model").String())) + if isClaudeOpus55Model(model) { + _ = doc.SetRawBytes("fallbacks", []byte(`[{"model":"claude-opus-4-8"}]`)) + } else if isClaudeFable51Model(model) { _ = doc.SetRawBytes("fallbacks", []byte(`[{"model":"claude-opus-5"}]`)) } } @@ -2633,10 +2687,11 @@ func applyCloakingInternalDoc( } } - // Apply sensitive word obfuscation - if obfuscateSensitiveWords && len(settings.sensitiveWords) > 0 { - matcher := helps.BuildSensitiveWordMatcher(settings.sensitiveWords) - helps.ObfuscateSensitiveWordsDoc(doc, matcher) + // Apply sensitive word obfuscation and replacements + if obfuscateSensitiveWords { + if matcher := settings.buildSensitiveWordMatcher(); matcher != nil { + helps.ObfuscateSensitiveWordsDoc(doc, matcher) + } } return true, true, nil diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go index ed58582cee2..e9352424684 100644 --- a/internal/runtime/executor/claude_executor_execute.go +++ b/internal/runtime/executor/claude_executor_execute.go @@ -220,6 +220,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, err } } + helps.NormalizeClaudeToolPropertyNamesDoc(doc) sanitizeClaudeMessagesForClaudeUpstreamWithDebugDoc(ctx, doc, baseModel, helps.APIKeyModelIsCompat(req)) if fp.ApplyCLIIdentity { err = applyClaudeCLIIdentityDoc(doc, auth, apiKey, url, claudeSessionID, fp.SynthesizeIdentity) @@ -227,9 +228,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, err } } - if cloaked && len(wireSettings.sensitiveWords) > 0 { - matcher := helps.BuildSensitiveWordMatcher(wireSettings.sensitiveWords) - helps.ObfuscateSensitiveWordsDoc(doc, matcher) + if cloaked { + if matcher := wireSettings.buildSensitiveWordMatcher(); matcher != nil { + helps.ObfuscateSensitiveWordsDoc(doc, matcher) + } } cchBilling := "" if cchSigning { @@ -328,13 +330,15 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, errClassified } b, readErr := io.ReadAll(errBody) + helps.AppendAPIResponseChunk(ctx, e.cfg, b) if readErr != nil { helps.RecordAPIResponseError(ctx, e.cfg, readErr) msg := fmt.Sprintf("failed to read error response body: %v", readErr) helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) + if len(b) == 0 { + b = []byte(msg) + } } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := errBody.Close(); errClose != nil { log.Errorf("response body close error: %v", errClose) @@ -358,11 +362,11 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r } }() data, err := io.ReadAll(decodedBody) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, wrapClaudeFastRequestError(fastRequest, httpResp.StatusCode, err) } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) var streamUsage helps.StreamUsageBuffer if upstreamStream { if errValidate := validateClaudeStreamingResponse(data); errValidate != nil { diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index 59b1221cccc..aa20045d172 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -217,6 +217,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return nil, err } } + helps.NormalizeClaudeToolPropertyNamesDoc(doc) sanitizeClaudeMessagesForClaudeUpstreamWithDebugDoc(ctx, doc, baseModel, helps.APIKeyModelIsCompat(req)) if fp.ApplyCLIIdentity { err = applyClaudeCLIIdentityDoc(doc, auth, apiKey, url, claudeSessionID, fp.SynthesizeIdentity) @@ -224,9 +225,10 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return nil, err } } - if cloaked && len(wireSettings.sensitiveWords) > 0 { - matcher := helps.BuildSensitiveWordMatcher(wireSettings.sensitiveWords) - helps.ObfuscateSensitiveWordsDoc(doc, matcher) + if cloaked { + if matcher := wireSettings.buildSensitiveWordMatcher(); matcher != nil { + helps.ObfuscateSensitiveWordsDoc(doc, matcher) + } } cchBilling := "" if cchSigning { @@ -322,13 +324,15 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return nil, errClassified } b, readErr := io.ReadAll(errBody) + helps.AppendAPIResponseChunk(ctx, e.cfg, b) if readErr != nil { helps.RecordAPIResponseError(ctx, e.cfg, readErr) msg := fmt.Sprintf("failed to read error response body: %v", readErr) helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) + if len(b) == 0 { + b = []byte(msg) + } } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := errBody.Close(); errClose != nil { log.Errorf("response body close error: %v", errClose) diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index f0b91b0cb29..6974c48cc51 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -5889,10 +5889,213 @@ func TestApplyCloaking_PreservesConfiguredStrictModeAndSensitiveWordsWhenModeOmi } } +func TestClaudeCloakSensitiveWordsReplaceFromYAML(t *testing.T) { + cfg, errParse := config.ParseConfigBytes([]byte(`claude-api-key: + - api-key: "key-replace-123" + cloak: + mode: "always" + sensitive-words: + - "proxy" + sensitive-words-replace: + "Claude Code": "Claude Assistant" +`)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "key-replace-123"}} + payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"Use Claude Code through proxy."}]}]}`) + + out, cloaked, errCloaking := applyCloaking( + context.Background(), + cfg, + auth, + payload, + "key-replace-123", + false, + false, + ) + if errCloaking != nil { + t.Fatalf("applyCloaking() error = %v", errCloaking) + } + if !cloaked { + t.Fatal("applyCloaking() cloaked = false, want true") + } + content := ojson.GetBytes(out, "messages.0.content").Array() + got := content[len(content)-1].Get("text").String() + want := "Use Claude Assistant through p\u200Broxy." + if got != want { + t.Fatalf("user message = %q, want %q", got, want) + } +} + +func TestClaudeOAuthMetadataSensitiveWordsReplaceAndDisabledCloaking(t *testing.T) { + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{"api_key": "sk-ant-oat-replace-test"}, + Metadata: map[string]any{ + "cloak_mode": "always", + "cloak_sensitive_words_replace": map[string]any{ + "You are Codex, an agent based on GPT-5.": "You are a coding assistant.", + }, + }, + } + payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"You are Codex, an agent based on GPT-5."}]}]}`) + out, cloaked, errCloak := applyCloaking(context.Background(), &config.Config{}, auth, payload, "sk-ant-oat-replace-test", false, false) + if errCloak != nil { + t.Fatalf("applyCloaking() error = %v", errCloak) + } + if !cloaked { + t.Fatal("expected cloaked = true") + } + content := ojson.GetBytes(out, "messages.0.content").Array() + if got := content[len(content)-1].Get("text").String(); got != "You are a coding assistant." { + t.Fatalf("rewritten message = %q, want %q", got, "You are a coding assistant.") + } + + authNever := &cliproxyauth.Auth{ + Attributes: map[string]string{ + "api_key": "key-never-replace", + "cloak_mode": "never", + "cloak_sensitive_words_replace": `{"Codex":"Assistant"}`, + }, + } + outNever, didCloakNever, errNever := applyCloaking(context.Background(), &config.Config{}, authNever, payload, "key-never-replace", false, false) + if errNever != nil { + t.Fatalf("applyCloaking(never) error = %v", errNever) + } + if didCloakNever { + t.Fatal("expected cloaked = false when cloak_mode=never") + } + if !strings.Contains(string(outNever), "You are Codex, an agent based on GPT-5.") { + t.Fatalf("disabled cloaking should preserve cleartext, got: %s", outNever) + } +} + +func TestHomeV8ClaudeCloakCredentialOptions(t *testing.T) { + // Pure credential_options.cloak WITHOUT duplicate cloak_* attributes or top-level metadata, + // isolating resolveClaudeKeyCloakConfig's Home path directly. + optionsOnlyAuth := &cliproxyauth.Auth{ + Provider: "claude", + Attributes: map[string]string{ + "api_key": "sk-ant-home-key", + "auth_kind": cliproxyauth.AuthKindAPIKey, + }, + Metadata: map[string]any{ + "credential_options": map[string]any{ + "cloak": map[string]any{ + "sensitive-words": []any{"ACME, Inc."}, + "sensitive-words-replace": map[string]any{"You are Codex, an agent based on GPT-5.": "You are a coding assistant."}, + }, + }, + }, + } + payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"You are Codex, an agent based on GPT-5. Contact ACME, Inc. or Beta Inc."}]}]}`) + + // 1. Omitted cloak.mode must respect shared DisableClaudeCloakMode=true. + cfgDisabled := &config.Config{ + Home: config.HomeConfig{Enabled: true}, + DisableClaudeCloakMode: true, + } + outDisabled, didCloakDisabled, errDisabled := applyCloaking(context.Background(), cfgDisabled, optionsOnlyAuth, payload, "sk-ant-home-key", false, false) + if errDisabled != nil { + t.Fatalf("applyCloaking(disabled) error = %v", errDisabled) + } + if didCloakDisabled { + t.Fatalf("expected didCloak=false when DisableClaudeCloakMode=true and cloak.mode is omitted, got body=%s", outDisabled) + } + + // 2. When DisableClaudeCloakMode=false, credential_options.cloak alone enables cloaking + // and preserves comma-containing sensitive-words without splitting on ','. + cfgEnabled := &config.Config{ + Home: config.HomeConfig{Enabled: true}, + } + outEnabled, didCloakEnabled, errEnabled := applyCloaking(context.Background(), cfgEnabled, optionsOnlyAuth, payload, "sk-ant-home-key", false, false) + if errEnabled != nil { + t.Fatalf("applyCloaking(enabled) error = %v", errEnabled) + } + if !didCloakEnabled { + t.Fatal("expected didCloak=true when Home credential_options.cloak is present") + } + content := ojson.GetBytes(outEnabled, "messages.0.content").Array() + gotText := content[len(content)-1].Get("text").String() + wantText := "You are a coding assistant. Contact A\u200BCME, Inc. or Beta Inc." + if gotText != wantText { + t.Fatalf("rewritten text = %q, want %q", gotText, wantText) + } + + // 3. Empty cloak block (cloak: {}) enables CLI cloaking when DisableClaudeCloakMode=false, + // and stays disabled when DisableClaudeCloakMode=true. + emptyCloakAuth := &cliproxyauth.Auth{ + Provider: "claude", + Attributes: map[string]string{"api_key": "sk-ant-empty-cloak", "auth_kind": cliproxyauth.AuthKindAPIKey}, + Metadata: map[string]any{"credential_options": map[string]any{"cloak": map[string]any{}}}, + } + if _, didEmptyEnabled, errEmpty := applyCloaking(context.Background(), cfgEnabled, emptyCloakAuth, payload, "sk-ant-empty-cloak", false, false); errEmpty != nil || !didEmptyEnabled { + t.Fatalf("empty cloak block with DisableClaudeCloakMode=false: didCloak=%v, err=%v", didEmptyEnabled, errEmpty) + } + if _, didEmptyDisabled, errEmptyDis := applyCloaking(context.Background(), cfgDisabled, emptyCloakAuth, payload, "sk-ant-empty-cloak", false, false); errEmptyDis != nil || didEmptyDisabled { + t.Fatalf("empty cloak block with DisableClaudeCloakMode=true: didCloak=%v, err=%v", didEmptyDisabled, errEmptyDis) + } + + // 4. Explicit mode="always" in credential_options.cloak overrides DisableClaudeCloakMode=true. + explicitAlwaysAuth := &cliproxyauth.Auth{ + Provider: "claude", + Attributes: map[string]string{"api_key": "sk-ant-explicit-always", "auth_kind": cliproxyauth.AuthKindAPIKey}, + Metadata: map[string]any{"credential_options": map[string]any{"cloak": map[string]any{"mode": "always"}}}, + } + if _, didExplicit, errExplicit := applyCloaking(context.Background(), cfgDisabled, explicitAlwaysAuth, payload, "sk-ant-explicit-always", false, false); errExplicit != nil || !didExplicit { + t.Fatalf("explicit mode=always failed to override DisableClaudeCloakMode=true: didCloak=%v, err=%v", didExplicit, errExplicit) + } + + // 5. Standalone mode (Home.Enabled == false) ignores Home-only credential_options.cloak, + // while []string cloak_sensitive_words in OAuth metadata works concurrently under claudeDevicePoolMu. + _, didCloakStandalone, errStandalone := applyCloaking(context.Background(), &config.Config{}, explicitAlwaysAuth, payload, "sk-ant-explicit-always", false, false) + if errStandalone != nil || didCloakStandalone { + t.Fatalf("standalone mode consumed Home credential_options.cloak: didCloak=%v, err=%v", didCloakStandalone, errStandalone) + } + oauthSliceAuth := &cliproxyauth.Auth{ + Provider: "claude", + Attributes: map[string]string{"api_key": "sk-ant-oat-slice", "auth_kind": cliproxyauth.AuthKindOAuth}, + Metadata: map[string]any{ + "cloak_mode": "always", + "cloak_sensitive_words": []string{"ACME, Inc."}, + }, + } + outSlice, didSlice, errSlice := applyCloaking(context.Background(), cfgEnabled, oauthSliceAuth, payload, "sk-ant-oat-slice", false, false) + if errSlice != nil || !didSlice { + t.Fatalf("oauth slice applyCloaking: did=%v err=%v", didSlice, errSlice) + } + sliceBlocks := ojson.GetBytes(outSlice, "messages.0.content").Array() + if got := sliceBlocks[len(sliceBlocks)-1].Get("text").String(); !strings.Contains(got, "A\u200BCME, Inc.") || strings.Contains(got, "Beta I\u200Bnc.") { + t.Fatalf("unexpected []string sensitive-words obfuscation: %q", got) + } + + // Concurrent metadata lock test across StoreDeviceIDPool, ReadMetadataHomeCloak, + // ReadMetadataStringSlice, and ReadMetadataStringMap. + done := make(chan struct{}) + go func() { + defer close(done) + for i := 0; i < 50; i++ { + claudeauth.StoreDeviceIDPool(&optionsOnlyAuth.Metadata, []string{"device-1"}) + } + }() + for i := 0; i < 50; i++ { + _, _, words, repls, _, ok := claudeauth.ReadMetadataHomeCloak(&optionsOnlyAuth.Metadata) + if !ok || len(words) != 1 || len(repls) != 1 { + t.Fatalf("ReadMetadataHomeCloak concurrent read failed: ok=%v words=%v repls=%v", ok, words, repls) + } + _ = claudeauth.ReadMetadataStringSlice(&oauthSliceAuth.Metadata, "cloak_sensitive_words") + _ = claudeauth.ReadMetadataStringMap(&optionsOnlyAuth.Metadata, "cloak_sensitive_words_replace") + } + <-done +} + func TestApplyCloaking_Opus55FallbackOnlyForUnconfirmedClients(t *testing.T) { - cfg := &config.Config{ClaudeKey: []config.ClaudeKey{{ - APIKey: "sk-ant-oat-opus55-test", Cloak: &config.CloakConfig{}, - }}} + cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, + ClaudeKey: []config.ClaudeKey{{ + APIKey: "sk-ant-oat-opus55-test", Cloak: &config.CloakConfig{}, + }}, + } auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-ant-oat-opus55-test"}} for _, tt := range []struct { name string @@ -5946,7 +6149,7 @@ func TestClaudeExecutor_CloakedOpus55PairsFallbackAndBetas(t *testing.T) { "api_key": "sk-ant-oat-cloaked-opus55", "base_url": server.URL, }} payload := []byte(`{"model":"claude-opus-5-5","thinking":{"type":"adaptive"},"messages":[{"role":"user","content":"test"}]}`) - executor := NewClaudeExecutor(&config.Config{}) + executor := NewClaudeExecutor(&config.Config{Claude: config.ClaudeConfig{ServerSideFallback: true}}) _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-opus-5-5", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) @@ -6005,6 +6208,7 @@ func TestClaudeOpus55FallbackReconcilesAfterModelOverride(t *testing.T) { func TestApplyCloaking_FableInjectsFallbacksAndDisplayUpdates(t *testing.T) { cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-fable-test", Cloak: &config.CloakConfig{}, @@ -6383,6 +6587,7 @@ func TestClaudeExecutor_PayloadOverrideFableModelLeavesBuiltinsUnchanged(t *test defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-fable-test", Cloak: &config.CloakConfig{}, @@ -6439,6 +6644,7 @@ func TestClaudeExecutor_PayloadOverrideNonFableModelDoesNotInjectFableBuiltins(t defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-fable-test-2", Cloak: &config.CloakConfig{}, @@ -6550,6 +6756,7 @@ func TestClaudeExecutor_PayloadUnrelatedRuleLeavesBuiltinFallbacksUnchanged(t *t defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-fable-test-4", Cloak: &config.CloakConfig{}, @@ -6664,6 +6871,7 @@ func TestClaudeExecutor_PayloadFableProbeOverrideKeepsBuiltinAdditions(t *testin defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-fable-probe-strip-test", Cloak: &config.CloakConfig{}, @@ -6719,6 +6927,7 @@ func TestClaudeExecutor_PayloadProbeToNormalDoesNotReinitializeBuiltins(t *testi defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-declassify-probe-test", Cloak: &config.CloakConfig{}, @@ -7735,6 +7944,7 @@ func TestClaudeExecutor_PayloadFilterFallbacksPreservesRemovalOnFable(t *testing defer server.Close() cfg := &config.Config{ + Claude: config.ClaudeConfig{ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{{ APIKey: "sk-ant-oat-payload-fable-filter-test", Cloak: &config.CloakConfig{}, diff --git a/internal/runtime/executor/claude_executor_tokens.go b/internal/runtime/executor/claude_executor_tokens.go index 0ffa329f5c1..541979658a8 100644 --- a/internal/runtime/executor/claude_executor_tokens.go +++ b/internal/runtime/executor/claude_executor_tokens.go @@ -70,6 +70,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut if rebuildMidSystemMessageEnabled(e.cfg, auth) { rebuildMidSystemMessagesToTopLevelDoc(doc) } + helps.NormalizeClaudeToolPropertyNamesDoc(doc) sanitizeClaudeMessagesForClaudeUpstreamWithDebugDoc(ctx, doc, baseModel, helps.APIKeyModelIsCompat(req)) applyClaudePreservedThinkingBindingDoc(doc) finalize := helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslatedForPayload, req, opts) @@ -212,8 +213,8 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy } } relocateClaudeSystemPromptForCountTokensWithPolicyDoc(doc, settings.strictMode, explicitCacheMode, !policy.OAuth) - if len(settings.sensitiveWords) > 0 { - helps.ObfuscateSensitiveWordsDoc(doc, helps.BuildSensitiveWordMatcher(settings.sensitiveWords)) + if matcher := settings.buildSensitiveWordMatcher(); matcher != nil { + helps.ObfuscateSensitiveWordsDoc(doc, matcher) } } else if policy.Uncloak { uncloakClaudeCodePayloadDoc(doc, confirmedClaudeCode) @@ -233,6 +234,7 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy mcpAliases := resolveClaudeMCPAliasOptions(ctx, opts.Metadata, req.Metadata) prepareClaudeOAuthToolNamesForUpstreamDoc(doc, mcpAliases) } + helps.NormalizeClaudeToolPropertyNamesDoc(doc) sanitizeClaudeMessagesForClaudeUpstreamWithDebugDoc(ctx, doc, baseModel, helps.APIKeyModelIsCompat(req)) applyClaudePreservedThinkingBindingDoc(doc) stripPromptCacheOptionsDoc(doc) @@ -314,13 +316,15 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy return cliproxyexecutor.Response{}, classifyClaudeUpstreamErrorWithCooling(resp.StatusCode, resp.Header, []byte(msg), e.modelLevelCooling()) } b, readErr := io.ReadAll(errBody) + helps.AppendAPIResponseChunk(ctx, e.cfg, b) if readErr != nil { helps.RecordAPIResponseError(ctx, e.cfg, readErr) msg := fmt.Sprintf("failed to read error response body: %v", readErr) helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) + if len(b) == 0 { + b = []byte(msg) + } } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) if errClose := errBody.Close(); errClose != nil { log.Errorf("response body close error: %v", errClose) } @@ -340,11 +344,11 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy } }() data, err := io.ReadAll(decodedBody) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return cliproxyexecutor.Response{}, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) count := ojson.GetBytesNoCopy(data, "input_tokens").Int() out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data) return cliproxyexecutor.Response{Payload: out, Headers: resp.Header.Clone()}, nil diff --git a/internal/runtime/executor/claude_issue7389_test.go b/internal/runtime/executor/claude_issue7389_test.go new file mode 100644 index 00000000000..8342c1c178b --- /dev/null +++ b/internal/runtime/executor/claude_issue7389_test.go @@ -0,0 +1,333 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "regexp" + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" +) + +var claudeIssue7389PropertyKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9_.-]{1,64}$`) + +const claudeIssue7389Schema = `{"type":"object","properties":{"not[labels]":{"type":"string","description":"Labels to exclude"},"state":{"type":"string"}},"required":["not[labels]"]}` + +const claudeIssue7389ProjectSchema = `{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}` + +var claudeIssue7389ResponsesPayload = `{"model":"claude-opus-5-5","input":[` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"list"}]},` + + `{"type":"function_call","call_id":"call_1","name":"list_merge_requests","arguments":"{\"not[labels]\":\"bug\",\"state\":\"opened\"}"},` + + `{"type":"function_call_output","call_id":"call_1","output":"[]"},` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"again"}]}],` + + `"tools":[{"type":"function","name":"list_merge_requests","description":"List merge requests.","parameters":` + claudeIssue7389Schema + `},` + + `{"type":"function","name":"get_project","description":"Get a project.","parameters":` + claudeIssue7389ProjectSchema + `}]}` + +var claudeIssue7389ClaudePayload = `{"model":"claude-opus-5-5","max_tokens":256,"messages":[` + + `{"role":"user","content":"list"},` + + `{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"list_merge_requests","input":{"not[labels]":"bug","state":"opened"}}]},` + + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"[]"},{"type":"text","text":"again"}]}],` + + `"tools":[{"name":"list_merge_requests","description":"List merge requests.","strict":true,"input_schema":` + claudeIssue7389Schema + `},` + + `{"name":"get_project","description":"Get a project.","input_schema":` + claudeIssue7389ProjectSchema + `}]}` + +func TestClaudeIssue7389RejectedPropertyNamesDoNotBreakRequests(t *testing.T) { + sources := []struct { + name string + format sdktranslator.Format + payload string + }{ + {name: "openai responses", format: sdktranslator.FormatOpenAIResponse, payload: claudeIssue7389ResponsesPayload}, + {name: "claude", format: sdktranslator.FormatClaude, payload: claudeIssue7389ClaudePayload}, + } + auths := []struct { + name string + attrs map[string]string + }{ + {name: "api key", attrs: map[string]string{"api_key": "key-7389"}}, + {name: "claude code mcp alias", attrs: map[string]string{"api_key": "key-7389-cli", "fingerprint_profile": "claude-code-cli", "cloak_mode": "always"}}, + } + for _, authCase := range auths { + for _, source := range sources { + for _, stream := range []bool{false, true} { + name := authCase.name + " " + source.name + if stream { + name += " stream" + } + t.Run(name, func(t *testing.T) { + claudeIssue7389RunRoundTrip(t, authCase.attrs, source.format, source.payload, stream) + }) + } + } + } +} + +func claudeIssue7389RunRoundTrip(t *testing.T, attrs map[string]string, format sdktranslator.Format, sourcePayload string, stream bool) { + t.Helper() + var upstreamBody []byte + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + upstreamBody, _ = io.ReadAll(req.Body) + toolName := claudeIssue7389Tool(upstreamBody, "List merge requests.").Get("name").String() + contentType := "application/json" + body := claudeIssue7389JSONResponse(toolName) + if ojson.GetBytes(upstreamBody, "stream").Bool() { + contentType = "text/event-stream" + body = claudeIssue7389SSEResponse(toolName) + } + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{contentType}}, Body: io.NopCloser(strings.NewReader(body)), Request: req}, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: attrs} + payload := sourcePayload + if stream && format == sdktranslator.FormatClaude { + payload = strings.Replace(payload, `"max_tokens":256,`, `"max_tokens":256,"stream":true,`, 1) + } + request := cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: []byte(payload)} + options := cliproxyexecutor.Options{SourceFormat: format, Stream: stream} + + var output []byte + if stream { + result, err := executor.ExecuteStream(ctx, auth, request, options) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + output = append(output, chunk.Payload...) + output = append(output, '\n') + } + } else { + resp, err := executor.Execute(ctx, auth, request, options) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + output = resp.Payload + } + + claudeIssue7389AssertUpstream(t, upstreamBody) + if attrs["fingerprint_profile"] != "" && claudeIssue7389Tool(upstreamBody, "List merge requests.").Get("name").String() == "list_merge_requests" { + t.Fatalf("expected an MCP tool alias upstream: %s", ojson.GetBytes(upstreamBody, "tools").Raw) + } + args := claudeIssue7389ClientArguments(t, format, stream, output) + fields := claudeIssue7389ObjectFields(args) + if fields["not[labels]"] != `"feature"` || fields["state"] != `"closed"` || len(fields) != 2 { + t.Fatalf("client tool arguments = %s, want the original property names; output=%s", args, output) + } +} + +func TestClaudeIssue7389CountTokensDoesNotSendRejectedPropertyNames(t *testing.T) { + var upstreamBody []byte + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + upstreamBody, _ = io.ReadAll(req.Body) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"input_tokens":12}`)), Request: req}, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "key-7389"}} + request := cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: []byte(claudeIssue7389ClaudePayload)} + if _, err := executor.CountTokens(ctx, auth, request, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}); err != nil { + t.Fatalf("CountTokens() error = %v", err) + } + claudeIssue7389AssertUpstream(t, upstreamBody) +} + +// User payload rules are the final semantic barrier: values they configure +// must reach upstream exactly as configured. +func TestClaudeIssue7389PayloadRulesRunAfterSchemaFallback(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { + var upstreamBody []byte + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + upstreamBody, _ = io.ReadAll(req.Body) + body := claudeIssue7389SSEResponse("list_merge_requests") + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(body)), Request: req}, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + models := []config.PayloadModelRule{{Name: "claude-opus-5-5"}} + cfg := &config.Config{Payload: config.PayloadConfig{ + Override: []config.PayloadRule{{Models: models, Params: map[string]any{"tools.0.description": "configured"}}}, + OverrideRaw: []config.PayloadRule{{Models: models, Params: map[string]any{"tools.0.input_schema": `{"type":"object","properties":{"configured[x]":{"type":"string"}}}`}}}, + }} + executor := NewClaudeExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "key-7389"}} + request := cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: []byte(claudeIssue7389ResponsesPayload)} + options := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, Stream: stream} + if stream { + result, err := executor.ExecuteStream(ctx, auth, request, options) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + } + } else if _, err := executor.Execute(ctx, auth, request, options); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if got := ojson.GetBytes(upstreamBody, "tools.0.input_schema").Raw; got != `{"type":"object","properties":{"configured[x]":{"type":"string"}}}` { + t.Fatalf("configured input_schema = %s, want the payload rule value", got) + } + if got := ojson.GetBytes(upstreamBody, "tools.0.description").String(); got != "configured" { + t.Fatalf("configured description = %q, want the payload rule value", got) + } + }) + } +} + +// claudeIssue7389Tool returns the upstream declaration whose description +// starts with prefix, whatever name aliasing the request path applied. +func claudeIssue7389Tool(body []byte, prefix string) ojson.Result { + var found ojson.Result + ojson.GetBytes(body, "tools").ForEach(func(_, tool ojson.Result) bool { + if strings.HasPrefix(tool.Get("description").String(), prefix) { + found = tool + return false + } + return true + }) + return found +} + +// claudeIssue7389AssertPropertyKeys checks every properties map in a schema. +func claudeIssue7389AssertPropertyKeys(t *testing.T, node ojson.Result, propertyMap bool) { + t.Helper() + if node.Type != ojson.JSON { + return + } + object := node.IsObject() + node.ForEach(func(key, value ojson.Result) bool { + if object && propertyMap && !claudeIssue7389PropertyKeyPattern.MatchString(key.String()) { + t.Fatalf("upstream property key %q does not match Claude pattern", key.String()) + } + claudeIssue7389AssertPropertyKeys(t, value, object && !propertyMap && key.String() == "properties") + return true + }) +} + +func claudeIssue7389AssertUpstream(t *testing.T, body []byte) { + t.Helper() + if len(body) == 0 { + t.Fatal("upstream request was not captured") + } + ojson.GetBytes(body, "tools").ForEach(func(_, tool ojson.Result) bool { + claudeIssue7389AssertPropertyKeys(t, tool.Get("input_schema"), false) + return true + }) + + reported := claudeIssue7389Tool(body, "List merge requests.") + if got := reported.Get("input_schema").Raw; got != `{"type":"object"}` { + t.Fatalf("reported tool input_schema = %s, want an open object", got) + } + if reported.Get("strict").Exists() { + t.Fatalf("reported tool keeps strict: %s", reported.Raw) + } + description := reported.Get("description").String() + if !strings.Contains(description, `"not[labels]":{"type":"string","description":"Labels to exclude"}`) || !strings.Contains(description, `"required":["not[labels]"]`) { + t.Fatalf("reported tool description does not carry the original schema: %q", description) + } + + project := claudeIssue7389Tool(body, "Get a project.") + if got := project.Get("description").String(); got != "Get a project." { + t.Fatalf("valid tool description = %q, want unchanged", got) + } + if !project.Get("input_schema.properties.id").Exists() || project.Get("input_schema.required.0").String() != "id" { + t.Fatalf("valid tool input_schema changed: %s", project.Get("input_schema").Raw) + } + + foundHistory := false + ojson.GetBytes(body, "messages").ForEach(func(_, message ojson.Result) bool { + message.Get("content").ForEach(func(_, block ojson.Result) bool { + if block.Get("type").String() != "tool_use" { + return true + } + foundHistory = true + fields := claudeIssue7389ObjectFields([]byte(block.Get("input").Raw)) + if fields["not[labels]"] != `"bug"` || fields["state"] != `"opened"` || len(fields) != 2 { + t.Fatalf("upstream history tool_use input = %s, want the original arguments", block.Get("input").Raw) + } + return true + }) + return true + }) + if !foundHistory { + t.Fatalf("upstream history tool_use not found: %s", body) + } +} + +func claudeIssue7389ObjectFields(raw []byte) map[string]string { + fields := make(map[string]string) + ojson.ParseBytes(raw).ForEach(func(k, v ojson.Result) bool { + fields[strings.Clone(k.String())] = strings.Clone(v.Raw) + return true + }) + return fields +} + +func claudeIssue7389JSONResponse(toolName string) string { + return `{"id":"msg_7389","type":"message","role":"assistant","model":"claude-opus-5-5","content":[{"type":"tool_use","id":"toolu_2","name":` + + string(ojson.AppendJSONString(nil, toolName)) + `,"input":{"not[labels]":"feature","state":"closed"}}],"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":1}}` +} + +func claudeIssue7389SSEResponse(toolName string) string { + var b strings.Builder + writeEvent := func(event, data string) { + b.WriteString("event: " + event + "\ndata: " + data + "\n\n") + } + writeEvent("message_start", `{"type":"message_start","message":{"id":"msg_7389","type":"message","role":"assistant","model":"claude-opus-5-5","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":0}}}`) + writeEvent("content_block_start", `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_2","name":`+string(ojson.AppendJSONString(nil, toolName))+`,"input":{}}}`) + for _, fragment := range []string{"", `{"not[la`, `bels]": "fea`, `ture", "state"`, `: "closed"}`} { + writeEvent("content_block_delta", `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":`+string(ojson.AppendJSONString(nil, fragment))+`}}`) + } + writeEvent("content_block_stop", `{"type":"content_block_stop","index":0}`) + writeEvent("message_delta", `{"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":1}}`) + writeEvent("message_stop", `{"type":"message_stop"}`) + return b.String() +} + +// claudeIssue7389ClientArguments extracts the final tool arguments the client sees. +func claudeIssue7389ClientArguments(t *testing.T, format sdktranslator.Format, stream bool, output []byte) []byte { + t.Helper() + switch { + case format == sdktranslator.FormatClaude && !stream: + return []byte(ojson.GetBytes(output, `content.#(type=="tool_use").input`).Raw) + case format == sdktranslator.FormatClaude && stream: + var args strings.Builder + for _, line := range bytes.Split(output, []byte("\n")) { + line = bytes.TrimSpace(line) + if !bytes.HasPrefix(line, []byte("data:")) { + continue + } + event := ojson.ParseBytes(bytes.TrimSpace(line[len("data:"):])) + if event.Get("delta.type").String() == "input_json_delta" { + args.WriteString(event.Get("delta.partial_json").String()) + } + } + return []byte(args.String()) + case !stream: + return []byte(ojson.GetBytes(output, `output.#(type=="function_call").arguments`).String()) + default: + for _, line := range bytes.Split(output, []byte("\n")) { + line = bytes.TrimSpace(line) + if !bytes.HasPrefix(line, []byte("data:")) { + continue + } + event := ojson.ParseBytes(bytes.TrimSpace(line[len("data:"):])) + if event.Get("type").String() == "response.function_call_arguments.done" { + return []byte(event.Get("arguments").String()) + } + } + } + t.Fatalf("tool arguments not found in client output: %s", output) + return nil +} diff --git a/internal/runtime/executor/claude_issue7440_test.go b/internal/runtime/executor/claude_issue7440_test.go new file mode 100644 index 00000000000..43efc9a55e2 --- /dev/null +++ b/internal/runtime/executor/claude_issue7440_test.go @@ -0,0 +1,400 @@ +package executor + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" +) + +func TestClaudeExecutor_CloakedOAuthOmitsServerSideFallbackByDefault_Issue7440(t *testing.T) { + models := []string{"claude-opus-5-5", "claude-fable-5-1"} + formats := []struct { + name string + format sdktranslator.Format + payload func(model string) []byte + }{ + { + name: "openai-chat", + format: sdktranslator.FormatOpenAI, + payload: func(model string) []byte { + return []byte(fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}]}`, model)) + }, + }, + { + name: "claude-messages", + format: sdktranslator.FormatClaude, + payload: func(model string) []byte { + return []byte(fmt.Sprintf(`{"model":%q,"thinking":{"type":"adaptive"},"messages":[{"role":"user","content":"hello"}]}`, model)) + }, + }, + } + + for _, model := range models { + for _, source := range formats { + for _, stream := range []bool{false, true} { + mode := "execute" + if stream { + mode = "stream" + } + t.Run(model+"/"+source.name+"/"+mode, func(t *testing.T) { + var seenBody []byte + var seenHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenBody, _ = io.ReadAll(r.Body) + seenHeaders = r.Header.Clone() + if ojson.GetBytes(seenBody, "stream").Bool() { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, + "event: message_start\n"+ + `data: {"type":"message_start","message":{"id":"msg_7440","type":"message","role":"assistant","model":"claude-opus-4-8","content":[],"stop_reason":null,"usage":{"input_tokens":5,"output_tokens":0}}}`+"\n\n"+ + "event: content_block_start\n"+ + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`+"\n\n"+ + "event: content_block_delta\n"+ + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`+"\n\n"+ + "event: content_block_stop\n"+ + `data: {"type":"content_block_stop","index":0}`+"\n\n"+ + "event: message_delta\n"+ + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`+"\n\n"+ + "event: message_stop\n"+ + `data: {"type":"message_stop"}`+"\n\n", + ) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_7440","type":"message","model":"claude-opus-4-8","role":"assistant","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":5,"output_tokens":2}}`)) + })) + defer server.Close() + + auth := &cliproxyauth.Auth{ + ID: "oauth-issue-7440-" + model, + Metadata: claudeOAuthTestMetadata(), + Attributes: map[string]string{ + "api_key": "sk-ant-oat-issue-7440", + "base_url": server.URL, + "cloak_mode": "always", + }, + } + executor := NewClaudeExecutor(&config.Config{}) + req := cliproxyexecutor.Request{ + Model: model, + Payload: source.payload(model), + } + opts := cliproxyexecutor.Options{ + SourceFormat: source.format, + Stream: stream, + } + + if stream { + res, errStream := executor.ExecuteStream(context.Background(), auth, req, opts) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v", errStream) + } + var chunks [][]byte + for chunk := range res.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + if len(chunk.Payload) > 0 { + chunks = append(chunks, chunk.Payload) + } + } + if source.format == sdktranslator.FormatOpenAI { + if len(chunks) == 0 { + t.Fatal("expected non-empty translated OpenAI stream chunks") + } + for i, rawChunk := range chunks { + gotModel := ojson.GetBytesNoCopy(rawChunk, "model").String() + if gotModel != "claude-opus-4-8" { + t.Fatalf("chunk[%d] model = %q, want claude-opus-4-8; payload=%s", i, gotModel, rawChunk) + } + } + } + } else { + resp, errExec := executor.Execute(context.Background(), auth, req, opts) + if errExec != nil { + t.Fatalf("Execute() error = %v", errExec) + } + if source.format == sdktranslator.FormatOpenAI { + if gotModel := ojson.GetBytesNoCopy(resp.Payload, "model").String(); gotModel != "claude-opus-4-8" { + t.Fatalf("response model = %q, want claude-opus-4-8; payload=%s", gotModel, resp.Payload) + } + } + } + + if ojson.GetBytesNoCopy(seenBody, "fallbacks").Exists() { + t.Fatalf("default cloaking injected unsolicited fallbacks: %s", ojson.GetBytesNoCopy(seenBody, "fallbacks").Raw) + } + betas := seenHeaders.Get("Anthropic-Beta") + if strings.Contains(betas, claudeServerSideFallbackBeta) { + t.Fatalf("Anthropic-Beta = %q, must not include %s by default", betas, claudeServerSideFallbackBeta) + } + if strings.Contains(betas, claudeFallbackCreditBeta) { + t.Fatalf("Anthropic-Beta = %q, must not include %s by default", betas, claudeFallbackCreditBeta) + } + }) + } + } + } +} + +func TestClaudeExecutor_ServerSideFallbackOptInAndPayloadFilter_Issue7440(t *testing.T) { + optInCfg, errParse := config.ParseConfigBytes([]byte("config-version: 8\nupstream:\n claude:\n server-side-fallback: true\n")) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + + var seenBody []byte + var seenHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenBody, _ = io.ReadAll(r.Body) + seenHeaders = r.Header.Clone() + if ojson.GetBytes(seenBody, "stream").Bool() { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, + "event: message_start\n"+ + `data: {"type":"message_start","message":{"id":"msg_optin","type":"message","role":"assistant","model":"claude-opus-4-8","content":[],"stop_reason":null,"usage":{"input_tokens":4,"output_tokens":0}}}`+"\n\n"+ + "event: content_block_start\n"+ + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`+"\n\n"+ + "event: content_block_delta\n"+ + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}`+"\n\n"+ + "event: content_block_stop\n"+ + `data: {"type":"content_block_stop","index":0}`+"\n\n"+ + "event: message_delta\n"+ + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}`+"\n\n"+ + "event: message_stop\n"+ + `data: {"type":"message_stop"}`+"\n\n", + ) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_optin","type":"message","model":"claude-opus-4-8","role":"assistant","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":4,"output_tokens":1}}`)) + })) + defer server.Close() + + auth := &cliproxyauth.Auth{ + ID: "oauth-optin-7440", + Metadata: claudeOAuthTestMetadata(), + Attributes: map[string]string{ + "api_key": "sk-ant-oat-optin-7440", + "base_url": server.URL, + "cloak_mode": "always", + }, + } + req := cliproxyexecutor.Request{ + Model: "claude-opus-5-5", + Payload: []byte(`{"model":"claude-opus-5-5","messages":[{"role":"user","content":"hello"}]}`), + } + + if _, errExec := NewClaudeExecutor(optInCfg).Execute(context.Background(), auth, req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAI}); errExec != nil { + t.Fatalf("Execute(optIn) error = %v", errExec) + } + if got := ojson.GetBytesNoCopy(seenBody, "fallbacks.0.model").String(); got != "claude-opus-4-8" { + t.Fatalf("opt-in fallbacks.0.model = %q, want claude-opus-4-8; body=%s", got, seenBody) + } + for _, beta := range []string{claudeServerSideFallbackBeta, claudeFallbackCreditBeta} { + if !strings.Contains(seenHeaders.Get("Anthropic-Beta"), beta) { + t.Fatalf("Anthropic-Beta = %q, missing %s when server-side-fallback is enabled", seenHeaders.Get("Anthropic-Beta"), beta) + } + } + + // Payload filter remains the final barrier even when server-side-fallback is enabled. + filteredCfg, errParseFilter := config.ParseConfigBytes([]byte("config-version: 8\nupstream:\n claude:\n server-side-fallback: true\nrequests:\n payload:\n filter:\n - models:\n - name: \"claude-opus-5-5\"\n protocol: \"claude\"\n params:\n - \"fallbacks\"\n")) + if errParseFilter != nil { + t.Fatalf("ParseConfigBytes(filtered) error = %v", errParseFilter) + } + if _, errExec := NewClaudeExecutor(filteredCfg).Execute(context.Background(), auth, req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAI}); errExec != nil { + t.Fatalf("Execute(filtered) error = %v", errExec) + } + if ojson.GetBytesNoCopy(seenBody, "fallbacks").Exists() { + t.Fatalf("payload.filter failed to remove fallbacks: %s", ojson.GetBytesNoCopy(seenBody, "fallbacks").Raw) + } + for _, beta := range []string{claudeServerSideFallbackBeta, claudeFallbackCreditBeta} { + if strings.Contains(seenHeaders.Get("Anthropic-Beta"), beta) { + t.Fatalf("Anthropic-Beta = %q, must omit %s after payload.filter removes fallbacks", seenHeaders.Get("Anthropic-Beta"), beta) + } + } +} + +func TestClaudeExecutor_DefaultNoFallbackPreservesCallerFallbackAndRefusalTransparency_Issue7440(t *testing.T) { + // 1. Explicit caller fallbacks are still honored when server-side-fallback defaults to false. + var seenBody []byte + var seenHeaders http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenBody, _ = io.ReadAll(r.Body) + seenHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_caller_fb","type":"message","model":"claude-sonnet-5","role":"assistant","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":4,"output_tokens":1}}`)) + })) + defer server.Close() + + auth := &cliproxyauth.Auth{ + ID: "oauth-caller-fb-7440", + Metadata: claudeOAuthTestMetadata(), + Attributes: map[string]string{ + "api_key": "sk-ant-oat-caller-fb-7440", + "base_url": server.URL, + "cloak_mode": "always", + }, + } + callerFallbackReq := cliproxyexecutor.Request{ + Model: "claude-opus-5-5", + Payload: []byte(`{"model":"claude-opus-5-5","fallbacks":[{"model":"claude-sonnet-5"}],"messages":[{"role":"user","content":"hello"}]}`), + } + if _, errExec := NewClaudeExecutor(&config.Config{}).Execute(context.Background(), auth, callerFallbackReq, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}); errExec != nil { + t.Fatalf("Execute(caller fallback) error = %v", errExec) + } + if got := ojson.GetBytesNoCopy(seenBody, "fallbacks.0.model").String(); got != "claude-sonnet-5" { + t.Fatalf("caller fallback = %q, want claude-sonnet-5", got) + } + if !strings.Contains(seenHeaders.Get("Anthropic-Beta"), claudeServerSideFallbackBeta) { + t.Fatalf("Anthropic-Beta = %q, want %s for explicit caller fallback", seenHeaders.Get("Anthropic-Beta"), claudeServerSideFallbackBeta) + } + + // 2. When server-side-fallback is off and Opus 5.5 declines with reasoning_extraction refusal, + // both Execute and ExecuteStream surface the refusal transparently on OpenAI Chat Completions. + refusalServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if ojson.GetBytes(body, "stream").Bool() { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, + "event: message_start\n"+ + `data: {"type":"message_start","message":{"id":"msg_refusal_7440","type":"message","role":"assistant","model":"claude-opus-5-5","content":[],"stop_reason":null,"usage":{"input_tokens":6,"output_tokens":0}}}`+"\n\n"+ + "event: message_delta\n"+ + `data: {"type":"message_delta","delta":{"stop_reason":"refusal","stop_sequence":null,"stop_details":{"type":"refusal","category":"reasoning_extraction","explanation":"Request declined under reasoning extraction safeguards."}},"usage":{"output_tokens":0}}`+"\n\n"+ + "event: message_stop\n"+ + `data: {"type":"message_stop"}`+"\n\n", + ) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_refusal_7440","type":"message","model":"claude-opus-5-5","role":"assistant","content":[],"stop_reason":"refusal","stop_sequence":null,"stop_details":{"type":"refusal","category":"reasoning_extraction","explanation":"Request declined under reasoning extraction safeguards."},"usage":{"input_tokens":6,"output_tokens":0}}`)) + })) + defer refusalServer.Close() + + refusalAuth := &cliproxyauth.Auth{ + ID: "oauth-refusal-7440", + Metadata: claudeOAuthTestMetadata(), + Attributes: map[string]string{ + "api_key": "sk-ant-oat-refusal-7440", + "base_url": refusalServer.URL, + "cloak_mode": "always", + }, + } + chatReq := cliproxyexecutor.Request{ + Model: "claude-opus-5-5", + Payload: []byte(`{"model":"claude-opus-5-5","messages":[{"role":"user","content":"hello"}]}`), + } + + _, errRefusalExec := NewClaudeExecutor(&config.Config{}).Execute(context.Background(), refusalAuth, chatReq, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAI}) + if errRefusalExec == nil || !strings.Contains(errRefusalExec.Error(), "reasoning_extraction") { + t.Fatalf("Execute(refusal) err = %v, want reasoning_extraction refusal error", errRefusalExec) + } + + streamRes, errRefusalStream := NewClaudeExecutor(&config.Config{}).ExecuteStream(context.Background(), refusalAuth, chatReq, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAI, Stream: true}) + if errRefusalStream != nil { + t.Fatalf("ExecuteStream(refusal) initial err = %v", errRefusalStream) + } + sawContentFilter := false + var terminalStreamErr error + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + terminalStreamErr = chunk.Err + continue + } + if ojson.GetBytesNoCopy(chunk.Payload, "choices.0.finish_reason").String() == "content_filter" { + sawContentFilter = true + } + if gotModel := ojson.GetBytesNoCopy(chunk.Payload, "model").String(); gotModel != "claude-opus-5-5" { + t.Fatalf("refusal stream chunk model = %q, want claude-opus-5-5", gotModel) + } + } + if !sawContentFilter { + t.Fatal("expected finish_reason=content_filter chunk before refusal termination") + } + if terminalStreamErr == nil || !strings.Contains(terminalStreamErr.Error(), "reasoning_extraction") { + t.Fatalf("refusal stream terminal err = %v, want reasoning_extraction refusal error", terminalStreamErr) + } +} + +func TestClaudeExecutor_OpenAIResponsesSurfacesServingModelAndFallbackWhenOptedIn_Issue7440(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, + "event: message_start\n"+ + `data: {"type":"message_start","message":{"id":"msg_resp_e2e_7440","type":"message","role":"assistant","model":"claude-opus-4-8","content":[],"stop_reason":null,"usage":{"input_tokens":6,"output_tokens":0}}}`+"\n\n"+ + "event: content_block_start\n"+ + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`+"\n\n"+ + "event: content_block_delta\n"+ + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"fallback via responses"}}`+"\n\n"+ + "event: content_block_stop\n"+ + `data: {"type":"content_block_stop","index":0}`+"\n\n"+ + "event: message_delta\n"+ + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":4,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":0},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`+"\n\n"+ + "event: message_stop\n"+ + `data: {"type":"message_stop"}`+"\n\n", + ) + })) + defer server.Close() + + auth := &cliproxyauth.Auth{ + ID: "oauth-responses-fb-7440", + Metadata: claudeOAuthTestMetadata(), + Attributes: map[string]string{ + "api_key": "sk-ant-oat-responses-fb-7440", + "base_url": server.URL, + "cloak_mode": "always", + }, + } + cfg := &config.Config{Claude: config.ClaudeConfig{ServerSideFallback: true}} + exec := NewClaudeExecutor(cfg) + req := cliproxyexecutor.Request{ + Model: "claude-opus-5-5", + Payload: []byte(`{"model":"claude-opus-5-5","input":[{"role":"user","content":"hello"}]}`), + } + + nonStreamResp, errExec := exec.Execute(context.Background(), auth, req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if errExec != nil { + t.Fatalf("Execute(OpenAIResponses) error = %v", errExec) + } + if got := ojson.GetBytesNoCopy(nonStreamResp.Payload, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("Execute(OpenAIResponses) model = %q, want claude-opus-4-8; payload=%s", got, string(nonStreamResp.Payload)) + } + if fb := ojson.GetBytesNoCopy(nonStreamResp.Payload, "fallback"); !fb.Exists() || fb.Get("from.model").String() != "claude-opus-5-5" || fb.Get("to.model").String() != "claude-opus-4-8" { + t.Fatalf("Execute(OpenAIResponses) fallback = %s, want from=5-5 to=4-8", fb.Raw) + } + + streamResp, errStream := exec.ExecuteStream(context.Background(), auth, req, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream(OpenAIResponses) error = %v", errStream) + } + var completedModel string + var sawCompletedFallback bool + for chunk := range streamResp.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + for _, line := range strings.Split(string(chunk.Payload), "\n") { + if !strings.HasPrefix(line, "data: ") { + continue + } + ev := ojson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data: "))) + if ev.Get("type").String() == "response.completed" { + completedModel = ev.Get("response.model").String() + sawCompletedFallback = ev.Get("response.fallback.to.model").String() == "claude-opus-4-8" + } + } + } + if completedModel != "claude-opus-4-8" || !sawCompletedFallback { + t.Fatalf("ExecuteStream(OpenAIResponses) completedModel=%q sawCompletedFallback=%v, want claude-opus-4-8 and true", completedModel, sawCompletedFallback) + } +} diff --git a/internal/runtime/executor/claude_signing.go b/internal/runtime/executor/claude_signing.go index 6d4c7601ac3..4355f7bf24f 100644 --- a/internal/runtime/executor/claude_signing.go +++ b/internal/runtime/executor/claude_signing.go @@ -10,6 +10,7 @@ import ( "github.com/MoeForge/moego/ojson" xxHash64 "github.com/pierrec/xxHash/xxHash64" + claudeauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/claude" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" @@ -650,6 +651,17 @@ func resolveClaudeKeyConfig(cfg *config.Config, auth *cliproxyauth.Auth) *config // resolveClaudeKeyCloakConfig finds the matching ClaudeKey config and returns its CloakConfig. func resolveClaudeKeyCloakConfig(cfg *config.Config, auth *cliproxyauth.Auth) *config.CloakConfig { + if cfg != nil && cfg.Home.Enabled && auth != nil { + if mode, strictMode, words, replacements, cacheUserID, ok := claudeauth.ReadMetadataHomeCloak(&auth.Metadata); ok { + return config.NormalizeCloakConfig(&config.CloakConfig{ + Mode: mode, + StrictMode: strictMode, + SensitiveWords: words, + SensitiveWordsReplace: replacements, + CacheUserID: cacheUserID, + }) + } + } entry := resolveClaudeKeyConfig(cfg, auth) if entry == nil { return nil diff --git a/internal/runtime/executor/claude_user_agent_extensions_test.go b/internal/runtime/executor/claude_user_agent_extensions_test.go index 760c7dcef72..f5e25404862 100644 --- a/internal/runtime/executor/claude_user_agent_extensions_test.go +++ b/internal/runtime/executor/claude_user_agent_extensions_test.go @@ -43,14 +43,15 @@ func testClaudeExecutorUserAgentFields(t *testing.T, stabilize bool) { payload := doc.Bytes() var nativeSystem string for _, testCase := range []struct { - userAgent, stabilizedUserAgent string + userAgent, stabilizedUserAgent, xApp string }{ - {userAgent: "claude-cli/2.1.353 (external, cli)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli)"}, - {userAgent: "claude-cli/2.1.353 (external, cli, workload/cron)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, workload/cron)"}, - {userAgent: "claude-cli/2.1.353 (external, cli, future-feature/value, agent-sdk/0.2.133)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, future-feature/value)"}, - {userAgent: "claude-cli/2.1.353 (external, cli, request-phase/background, agent-sdk/0.2.133, workload/cron)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, request-phase/background, workload/cron)"}, + {userAgent: "claude-cli/2.1.353 (external, cli)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli)", xApp: "cli"}, + {userAgent: "claude-cli/2.1.353 (external, cli)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli)", xApp: "cli-bg"}, + {userAgent: "claude-cli/2.1.353 (external, cli, workload/cron)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, workload/cron)", xApp: "cli"}, + {userAgent: "claude-cli/2.1.353 (external, cli, future-feature/value, agent-sdk/0.2.133)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, future-feature/value)", xApp: "cli"}, + {userAgent: "claude-cli/2.1.353 (external, cli, request-phase/background, agent-sdk/0.2.133, workload/cron)", stabilizedUserAgent: "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, request-phase/background, workload/cron)", xApp: "cli"}, } { - t.Run(testCase.userAgent, func(t *testing.T) { + t.Run(testCase.userAgent+"/"+testCase.xApp, func(t *testing.T) { var upstreamBody []byte var upstreamHeaders http.Header transport := roundTripperFunc(func(request *http.Request) (*http.Response, error) { @@ -72,7 +73,7 @@ func testClaudeExecutorUserAgentFields(t *testing.T, stabilize bool) { ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) headers := http.Header{} headers.Set("User-Agent", testCase.userAgent) - headers.Set("X-App", "cli") + headers.Set("X-App", testCase.xApp) headers.Set("Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20,extended-cache-ttl-2025-04-11") auth := &cliproxyauth.Auth{ID: "ua-fields-" + t.Name(), Metadata: claudeOAuthTestMetadata(), Attributes: map[string]string{"api_key": "sk-ant-oat-ua-fields"}} executor := NewClaudeExecutor(&config.Config{ClaudeHeaderDefaults: config.ClaudeHeaderDefaults{ @@ -129,6 +130,9 @@ func testClaudeExecutorUserAgentFields(t *testing.T, stabilize bool) { if got := upstreamHeaders.Get("User-Agent"); got != wantUserAgent { t.Fatalf("upstream UA = %q, want %q", got, wantUserAgent) } + if got := claudeFingerprintHeaderValue(upstreamHeaders, "X-App"); got != testCase.xApp { + t.Fatalf("upstream X-App = %q, want %q", got, testCase.xApp) + } }) } }) diff --git a/internal/runtime/executor/codex_executor_compact_test.go b/internal/runtime/executor/codex_executor_compact_test.go index abdadae93be..04ff7b09db0 100644 --- a/internal/runtime/executor/codex_executor_compact_test.go +++ b/internal/runtime/executor/codex_executor_compact_test.go @@ -3,6 +3,7 @@ package executor import ( "bytes" "context" + "fmt" "io" "net/http" "net/http/httptest" @@ -13,9 +14,11 @@ import ( "github.com/MoeForge/moego/ojson" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" ) @@ -814,3 +817,1078 @@ func TestCodexExecutorExpandsRecognizedCPACompactionCapsules_Issue7118(t *testin } }) } + +func TestCodexExecutor_IsCompatResponsesCompactionTriggerAndReplay_Issue7406(t *testing.T) { + const marker = "COBALT-7406" + var seenPaths []string + var seenBodies [][]byte + var returnIncomplete bool + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read upstream body: %v", errRead) + http.Error(w, "read error", http.StatusBadRequest) + return + } + seenPaths = append(seenPaths, r.URL.Path) + seenBodies = append(seenBodies, body) + + w.Header().Set("Content-Type", "text/event-stream") + if returnIncomplete { + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_inc\",\"type\":\"message\",\"status\":\"incomplete\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Partial truncated summary.\"}]}}\n\n"+ + "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_inc\",\"object\":\"response\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"},\"output\":[{\"id\":\"msg_inc\",\"type\":\"message\",\"status\":\"incomplete\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Partial truncated summary.\"}]}],\"usage\":{\"input_tokens\":12,\"output_tokens\":6,\"total_tokens\":18}}}\n\n", + ) + return + } + + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"rs_sum\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"Summarizing session context.\"}]}}\n\n"+ + "data: {\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{\"id\":\"msg_sum\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Summary retaining "+marker+".\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_sum\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":14,\"output_tokens\":8,\"total_tokens\":22,\"input_tokens_details\":{\"cached_tokens\":5},\"output_tokens_details\":{\"reasoning_tokens\":3}}}}\n\n", + ) + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{ + { + APIKey: "test-compat-key", + BaseURL: server.URL + "/v1", + Models: []config.CodexModel{ + {Name: "deepseek-v4.1-flash", IsCompat: true}, + {Name: "gpt-5.4", IsCompat: false}, + }, + }, + }, + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "deepseek-v4.1-flash", Protocol: "codex"}}, + Params: map[string]any{ + "max_output_tokens": 2048, + }, + }, + }, + }, + } + + executor := NewCodexExecutor(cfg) + wsExecutor := NewCodexWebsocketsExecutor(cfg) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": server.URL + "/v1", + "api_key": "test-compat-key", + }, + } + + triggerPayload := []byte(`{ + "model": "deepseek-v4.1-flash", + "stream": true, + "input": [ + { + "type": "additional_tools", + "tools": [ + { + "type": "function", + "name": "exec_command", + "description": "Execute a shell command", + "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}} + } + ] + }, + { + "id": "ws_hist_1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "deepseek responses"} + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Remember project marker ` + marker + `."}] + }, + {"type": "compaction_trigger"} + ] + }`) + + assertSummaryRequest := func(t *testing.T, path string, body []byte, label string) { + t.Helper() + if path != "/v1/responses" { + t.Fatalf("%s upstream path = %q, want /v1/responses", label, path) + } + if bytes.Contains(body, []byte("compaction_trigger")) { + t.Fatalf("%s summary request leaked compaction_trigger: %s", label, body) + } + if ojson.GetBytes(body, "tools").Exists() { + t.Fatalf("%s summary request unexpectedly carried tools (including image_generation or web_search): %s", label, body) + } + if gotMax := ojson.GetBytes(body, "max_output_tokens").Int(); gotMax != 2048 { + t.Fatalf("%s max_output_tokens = %d, want 2048 from payload override; body=%s", label, gotMax, body) + } + } + + // 1. Streaming ExecuteStream(compaction_trigger) on is-compat codex endpoint. + streamResult, errStream := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + OriginalRequest: triggerPayload, + }) + if errStream != nil { + t.Fatalf("ExecuteStream(compaction_trigger) error = %v", errStream) + } + var streamBuf bytes.Buffer + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + streamBuf.Write(chunk.Payload) + } + completed := claudeCompactionCompletedResponse(t, streamBuf.String()) + outputItems := completed.Get("response.output").Array() + if len(outputItems) != 1 || outputItems[0].Get("type").String() != "compaction" { + t.Fatalf("expected single compaction item in stream response, got: %s", completed.Raw) + } + capsule := outputItems[0].Get("encrypted_content").String() + if !strings.HasPrefix(capsule, "cpa-ag-compact-v1:") { + t.Fatalf("expected CPA compaction capsule prefix, got %q", capsule) + } + if gotCached := completed.Get("response.usage.input_tokens_details.cached_tokens").Int(); gotCached != 5 { + t.Fatalf("cached_tokens = %d, want 5", gotCached) + } + if gotReasoning := completed.Get("response.usage.output_tokens_details.reasoning_tokens").Int(); gotReasoning != 3 { + t.Fatalf("reasoning_tokens = %d, want 3", gotReasoning) + } + assertSummaryRequest(t, seenPaths[0], seenBodies[0], "ExecuteStream") + + // 2. Non-streaming Execute(compaction_trigger) on is-compat codex endpoint. + nonStreamResp, errExec := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: triggerPayload, + }) + if errExec != nil { + t.Fatalf("Execute(compaction_trigger) error = %v", errExec) + } + if gotObj := ojson.GetBytes(nonStreamResp.Payload, "object").String(); gotObj != "response.compaction" { + t.Fatalf("Execute object = %q, want response.compaction; payload=%s", gotObj, nonStreamResp.Payload) + } + nonStreamItems := ojson.GetBytes(nonStreamResp.Payload, "output").Array() + if len(nonStreamItems) != 1 || nonStreamItems[0].Get("type").String() != "compaction" { + t.Fatalf("expected single compaction item in non-stream response, got: %s", nonStreamResp.Payload) + } + assertSummaryRequest(t, seenPaths[1], seenBodies[1], "Execute") + + // 3. WebSocket executor delegation for compaction_trigger on is-compat codex endpoint. + wsResp, errWSExec := wsExecutor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: triggerPayload, + }) + if errWSExec != nil { + t.Fatalf("wsExecutor.Execute(compaction_trigger) error = %v", errWSExec) + } + if gotObj := ojson.GetBytes(wsResp.Payload, "object").String(); gotObj != "response.compaction" { + t.Fatalf("wsExecutor.Execute object = %q, want response.compaction", gotObj) + } + wsStreamRes, errWSStream := wsExecutor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + OriginalRequest: triggerPayload, + }) + if errWSStream != nil { + t.Fatalf("wsExecutor.ExecuteStream(compaction_trigger) error = %v", errWSStream) + } + var wsStreamBuf bytes.Buffer + for chunk := range wsStreamRes.Chunks { + if chunk.Err != nil { + t.Fatalf("ws stream chunk error = %v", chunk.Err) + } + wsStreamBuf.Write(chunk.Payload) + } + wsCompleted := claudeCompactionCompletedResponse(t, wsStreamBuf.String()) + if wsItems := wsCompleted.Get("response.output").Array(); len(wsItems) != 1 || wsItems[0].Get("type").String() != "compaction" { + t.Fatalf("expected single compaction item in ws stream response, got: %s", wsCompleted.Raw) + } + + // 4. Subsequent replay on is-compat expands CPA capsule and drops foreign compaction items. + replayPayload := []byte(fmt.Sprintf(`{ + "model": "deepseek-v4.1-flash", + "stream": true, + "input": [ + {"id": "cmp_foreign", "type": "compaction", "encrypted_content": "gAAAAABforeign-openai-compaction"}, + {"id": %q, "type": "compaction", "encrypted_content": %q}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the project marker?"}]} + ] + }`, outputItems[0].Get("id").String(), capsule)) + + contStream, errCont := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: replayPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + OriginalRequest: replayPayload, + }) + if errCont != nil { + t.Fatalf("ExecuteStream(replay) error = %v", errCont) + } + for chunk := range contStream.Chunks { + if chunk.Err != nil { + t.Fatalf("replay stream chunk error = %v", chunk.Err) + } + } + replaySeen := seenBodies[len(seenBodies)-1] + if bytes.Contains(replaySeen, []byte("cpa-ag-compact-v1:")) || bytes.Contains(replaySeen, []byte("gAAAAABforeign-openai-compaction")) { + t.Fatalf("expected both CPA capsule and foreign compaction item to be removed from is-compat upstream payload: %s", replaySeen) + } + if !bytes.Contains(replaySeen, []byte(marker)) { + t.Fatalf("expected expanded summary marker %q in replay upstream payload: %s", marker, replaySeen) + } + + // CountTokens on is-compat also drops foreign compaction items and expands CPA capsule cleanly. + if _, errTokens := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: replayPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: replayPayload, + }); errTokens != nil { + t.Fatalf("CountTokens(replay) error = %v", errTokens) + } + + // 5. Chained second compaction_trigger expands prior CPA capsule before summarizing. + secondTriggerPayload := []byte(fmt.Sprintf(`{ + "model": "deepseek-v4.1-flash", + "input": [ + {"id": %q, "type": "compaction", "encrypted_content": %q}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "And remember SILVER-99."}]}, + {"type": "compaction_trigger"} + ] + }`, outputItems[0].Get("id").String(), capsule)) + + secondResp, errSecond := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: secondTriggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: secondTriggerPayload, + }) + if errSecond != nil { + t.Fatalf("Execute(second compaction_trigger) error = %v", errSecond) + } + if gotObj := ojson.GetBytes(secondResp.Payload, "object").String(); gotObj != "response.compaction" { + t.Fatalf("second compaction object = %q, want response.compaction", gotObj) + } + secondSeen := seenBodies[len(seenBodies)-1] + if bytes.Contains(secondSeen, []byte("compaction_trigger")) || bytes.Contains(secondSeen, []byte("cpa-ag-compact-v1:")) { + t.Fatalf("second summary request leaked compaction_trigger or unexpanded capsule: %s", secondSeen) + } + if !bytes.Contains(secondSeen, []byte(marker)) { + t.Fatalf("expected prior summary marker %q in second summary request: %s", marker, secondSeen) + } + + // 6. Incomplete summary must fail closed rather than sealing a truncated capsule. + returnIncomplete = true + if _, errIncExec := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: triggerPayload, + }); errIncExec == nil { + t.Fatal("expected Execute(compaction_trigger) to fail on incomplete summary, got nil") + } + if _, errIncStream := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: triggerPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + OriginalRequest: triggerPayload, + }); errIncStream == nil { + t.Fatal("expected ExecuteStream(compaction_trigger) to fail on incomplete summary, got nil") + } + returnIncomplete = false + + // 7. Native (is-compat: false) Codex model forwards compaction_trigger untouched without local summary interception. + nativePayload := []byte(`{"model":"gpt-5.4","input":[{"type":"message","role":"user","content":"history"},{"type":"compaction_trigger"}]}`) + nativeResp, errNative := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: nativePayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: nativePayload, + }) + if errNative != nil { + t.Fatalf("Execute(native gpt-5.4) error = %v", errNative) + } + nativeSeen := seenBodies[len(seenBodies)-1] + if !bytes.Contains(nativeSeen, []byte("compaction_trigger")) { + t.Fatalf("expected native is-compat=false model to forward compaction_trigger untouched, got: %s", nativeSeen) + } + if gotObj := ojson.GetBytes(nativeResp.Payload, "object").String(); gotObj != "response" { + t.Fatalf("native response object = %q, want response", gotObj) + } +} + +func TestCodexExecutor_IsCompatCompactionValidationAndUsage_Issue7406(t *testing.T) { + var upstreamHits int + var mode string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamHits++ + _, _ = io.ReadAll(r.Body) + switch mode { + case "http500": + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + _, _ = io.WriteString(w, `{"error":{"message":"upstream boom"}}`) + case "incomplete": + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_1\",\"type\":\"message\",\"status\":\"incomplete\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"partial\"}]}}\n\n"+ + "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_inc\",\"object\":\"response\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"},\"output\":[],\"usage\":{\"input_tokens\":20,\"output_tokens\":10,\"total_tokens\":30,\"input_tokens_details\":{\"cached_tokens\":6},\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ) + case "empty_summary": + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"rs_only\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[{\"type\":\"summary_text\",\"text\":\"only thinking\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_empty\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":20,\"output_tokens\":10,\"total_tokens\":30,\"input_tokens_details\":{\"cached_tokens\":6},\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ) + case "empty_incomplete": + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_empty_inc\",\"object\":\"response\",\"status\":\"incomplete\",\"output\":[],\"usage\":{\"input_tokens\":20,\"output_tokens\":10,\"total_tokens\":30,\"input_tokens_details\":{\"cached_tokens\":6},\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ) + case "response_failed": + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_fail\",\"object\":\"response\",\"status\":\"failed\",\"error\":{\"code\":\"server_error\",\"message\":\"upstream failed mid-summary\"},\"output\":[],\"usage\":{\"input_tokens\":20,\"output_tokens\":10,\"total_tokens\":30,\"input_tokens_details\":{\"cached_tokens\":6},\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ) + default: + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_ok\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Summary text ok.\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_ok\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":20,\"output_tokens\":10,\"total_tokens\":30,\"input_tokens_details\":{\"cached_tokens\":6},\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ) + } + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: server.URL, + Models: []config.CodexModel{{Name: "deepseek-v4.1-flash", IsCompat: true}}, + }}, + } + executor := NewCodexExecutor(cfg) + wsExecutor := NewCodexWebsocketsExecutor(cfg) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test-compat-key"}, + } + + // 1. Incremental previous_response_id + compaction_trigger must fail closed without hitting upstream. + incTrigger := []byte(`{"model":"deepseek-v4.1-flash","event_id":"ev_inc","previous_response_id":"resp_real_1","input":[{"type":"compaction_trigger"}]}`) + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: incTrigger} + if _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: incTrigger}, opts); err == nil { + t.Fatal("expected HTTP Execute with previous_response_id + compaction_trigger to fail") + } + wsCtx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + _, errWSInc := wsExecutor.ExecuteStream(wsCtx, auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: incTrigger}, opts) + if !cliproxyexecutor.IsUpstreamWebsocketTurnError(errWSInc) { + t.Fatalf("expected UpstreamWebsocketTurnError on websocket incremental trigger, got %v", errWSInc) + } + + // 2. Empty conversation input (only compaction_trigger + additional_tools) must fail closed. + emptyTrigger := []byte(`{"model":"deepseek-v4.1-flash","input":[{"type":"additional_tools","tools":[]},{"type":"compaction_trigger"}]}`) + if _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: emptyTrigger}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: emptyTrigger}); err == nil { + t.Fatal("expected empty-history compaction_trigger to fail") + } + + // 3. Synthetic resp_ag_compact_* previous_response_id on normal turn must never reach upstream. + synthCont := []byte(`{"model":"deepseek-v4.1-flash","previous_response_id":"resp_ag_compact_12345","input":[{"type":"message","role":"user","content":"hi"}]}`) + if _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: synthCont}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: synthCont}); err == nil { + t.Fatal("expected synthetic previous_response_id on HTTP Execute to fail") + } + if _, err := wsExecutor.ExecuteStream(wsCtx, auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: synthCont}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: synthCont}); !cliproxyexecutor.IsUpstreamWebsocketTurnError(err) { + t.Fatalf("expected UpstreamWebsocketTurnError for synthetic previous_response_id on WS, got %v", err) + } + if upstreamHits != 0 { + t.Fatalf("upstreamHits = %d, want 0 for rejected validation cases", upstreamHits) + } + + // 4. CountTokens semantic consistency: foreign compaction dropped + CPA capsule expanded matches expanded developer message. + capsule, errSeal := helps.SealAntigravityCompaction("Token count summary", "deepseek-v4.1-flash") + if errSeal != nil { + t.Fatalf("SealAntigravityCompaction: %v", errSeal) + } + mixedTokensPayload := []byte(fmt.Sprintf(`{"model":"deepseek-v4.1-flash","input":[{"id":"cmp_foreign","type":"compaction","encrypted_content":"gAAAAABforeign"},{"id":"cmp_ag_compact_1","type":"compaction","encrypted_content":%q},{"type":"message","role":"user","content":[{"type":"input_text","text":"next"}]}]}`, capsule)) + expectedExpandedPayload := []byte(`{"model":"deepseek-v4.1-flash","input":[{"type":"message","role":"developer","content":[{"type":"input_text","text":"Context summary from previous turns:\nToken count summary"}]},{"type":"message","role":"user","content":[{"type":"input_text","text":"next"}]}]}`) + gotCountResp, errGotCount := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: mixedTokensPayload}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if errGotCount != nil { + t.Fatalf("CountTokens(mixed): %v", errGotCount) + } + wantCountResp, errWantCount := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{Model: "deepseek-v4.1-flash", Payload: expectedExpandedPayload}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if errWantCount != nil { + t.Fatalf("CountTokens(expected): %v", errWantCount) + } + if string(gotCountResp.Payload) != string(wantCountResp.Payload) { + t.Fatalf("CountTokens mismatch: got %s, want %s", gotCountResp.Payload, wantCountResp.Payload) + } + + // 5. Usage outcome & token accounting across valid, incomplete, empty_summary, http500, and canceled. + validTrigger := []byte(`{"model":"deepseek-v4.1-flash","input":[{"type":"message","role":"user","content":"history"},{"type":"compaction_trigger"}]}`) + for _, tc := range []struct { + name string + mode string + cancelCtx bool + wantErr bool + wantTokens int64 + }{ + {name: "valid", mode: "ok", wantErr: false, wantTokens: 30}, + {name: "incomplete", mode: "incomplete", wantErr: true, wantTokens: 30}, + {name: "empty_summary", mode: "empty_summary", wantErr: true, wantTokens: 30}, + {name: "empty_incomplete_with_usage", mode: "empty_incomplete", wantErr: true, wantTokens: 30}, + {name: "response_failed_with_usage", mode: "response_failed", wantErr: true, wantTokens: 30}, + {name: "http500", mode: "http500", wantErr: true, wantTokens: 0}, + {name: "canceled", mode: "ok", cancelCtx: true, wantErr: true, wantTokens: 0}, + } { + t.Run("usage_"+tc.name, func(t *testing.T) { + mode = tc.mode + alias := t.Name() + capture := &codexResponseModelUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), codexResponseModelNoopUsagePlugin{}) + }) + callCtx := coreusage.WithRequestedModelAlias(context.Background(), alias) + if tc.cancelCtx { + var cancel context.CancelFunc + callCtx, cancel = context.WithCancel(callCtx) + cancel() + } + _, errExec := executor.Execute(callCtx, auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: validTrigger, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: validTrigger, + }) + if (errExec != nil) != tc.wantErr { + t.Fatalf("Execute err = %v, wantErr = %t", errExec, tc.wantErr) + } + rec := capture.await(t) + if rec.Failed != tc.wantErr { + t.Fatalf("usage record Failed = %t, want %t (fail=%+v)", rec.Failed, tc.wantErr, rec.Fail) + } + if rec.Detail.TotalTokens != tc.wantTokens { + t.Fatalf("usage record TotalTokens = %d, want %d (detail=%+v)", rec.Detail.TotalTokens, tc.wantTokens, rec.Detail) + } + if tc.wantTokens > 0 && (rec.Detail.InputTokens != 20 || rec.Detail.OutputTokens != 10 || rec.Detail.ReasoningTokens != 4 || rec.Detail.CachedTokens != 6) { + t.Fatalf("unexpected token breakdown: %+v", rec.Detail) + } + select { + case extra := <-capture.records: + t.Fatalf("unexpected duplicate usage record emitted: %+v", extra) + default: + } + }) + } +} + +func TestCodexExecutor_IsCompatCompactionMatrix_Issue7406(t *testing.T) { + var lastBody []byte + var lastAuthHeader string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + lastAuthHeader = r.Header.Get("Authorization") + lastBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_m\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Matrix summary.\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_m\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":8,\"output_tokens\":4,\"total_tokens\":12}}}\n\n", + ) + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "outer-key", + BaseURL: server.URL, + Models: []config.CodexModel{{ + Name: "deepseek-v4.1-flash", + IsCompat: true, + Thinking: ®istry.ThinkingSupport{Levels: []string{"low", "medium", "high"}}, + }}, + }}, + Payload: config.PayloadConfig{ + Filter: []config.PayloadFilterRule{{ + Models: []config.PayloadModelRule{{Name: "deepseek-v4.1-flash", Protocol: "codex"}}, + Params: []string{"service_tier"}, + }}, + Override: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "deepseek-v4.1-flash", Protocol: "codex"}}, + Params: map[string]any{"max_output_tokens": 1536}, + }}, + }, + } + executor := NewCodexExecutor(cfg) + siwcAuth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{"base_url": server.URL, "api_key": "outer-key", "use_siwc": "true"}, + Metadata: map[string]any{"siwc": map[string]any{"access_token": "siwc-secret-token"}}, + } + + trigger := []byte(`{"model":"deepseek-v4.1-flash(high)","service_tier":"priority","input":[{"type":"message","role":"user","content":"summarize me"},{"type":"compaction_trigger"}]}`) + resp, err := executor.Execute(context.Background(), siwcAuth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash(high)", + Payload: trigger, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: trigger, + }) + if err != nil { + t.Fatalf("Execute(SIWC + suffix + payload rules) error = %v", err) + } + if ojson.GetBytes(resp.Payload, "object").String() != "response.compaction" { + t.Fatalf("expected response.compaction, got: %s", resp.Payload) + } + if lastAuthHeader != "Bearer siwc-secret-token" { + t.Fatalf("Authorization = %q, want Bearer siwc-secret-token", lastAuthHeader) + } + if gotEffort := ojson.GetBytes(lastBody, "reasoning.effort").String(); gotEffort != "high" { + t.Fatalf("reasoning.effort = %q, want high; body=%s", gotEffort, lastBody) + } + if ojson.GetBytes(lastBody, "service_tier").Exists() { + t.Fatalf("service_tier should be filtered by payload filter rule: %s", lastBody) + } + if gotMax := ojson.GetBytes(lastBody, "max_output_tokens").Int(); gotMax != 1536 { + t.Fatalf("max_output_tokens = %d, want 1536; body=%s", gotMax, lastBody) + } + + // Authoritative ResolvedModelInfo{IsCompat: false} overrides local config IsCompat: true. + overrideResp, errOverride := executor.Execute(context.Background(), siwcAuth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: trigger, + Metadata: map[string]any{ + "cliproxy.resolved_api_key_model_info": ®istry.ModelInfo{ID: "deepseek-v4.1-flash", IsCompat: false}, + }, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: trigger, + }) + if errOverride != nil { + t.Fatalf("Execute(authoritative IsCompat=false) error = %v", errOverride) + } + if !bytes.Contains(lastBody, []byte("compaction_trigger")) { + t.Fatalf("expected authoritative IsCompat=false to forward compaction_trigger untouched: %s", lastBody) + } + if ojson.GetBytes(overrideResp.Payload, "object").String() != "response" { + t.Fatalf("expected normal response when authoritative IsCompat=false, got: %s", overrideResp.Payload) + } +} + +func TestCodexExecutor_IsCompatCompactionDuplexSession_Issue7406(t *testing.T) { + const marker = "DUPLEX-COMPACT-7406" + wsFrames := make(chan []byte, 8) + httpBodies := make(chan []byte, 4) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if websocket.IsWebSocketUpgrade(r) { + conn, errUp := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUp != nil { + t.Errorf("upgrade ws: %v", errUp) + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close upstream websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + // Turn 1: initial WS response.create + _, frame1, errRead1 := conn.ReadMessage() + if errRead1 != nil { + return + } + wsFrames <- frame1 + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_ws_1","output":[]}}`)) + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_ws_1","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ack-1"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`)) + + // Turn 3 (post-compaction root response.create; Turn 2 ran locally via HTTP!) + _, frame2, errRead2 := conn.ReadMessage() + if errRead2 != nil { + return + } + wsFrames <- frame2 + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_ws_3","output":[]}}`)) + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_ws_3","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ack-3"}]}],"usage":{"input_tokens":8,"output_tokens":3,"total_tokens":11}}}`)) + + // Turn 4 (incremental continuation referencing resp_ws_3) + _, frame3, errRead3 := conn.ReadMessage() + if errRead3 != nil { + return + } + wsFrames <- frame3 + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_ws_4","previous_response_id":"resp_ws_3","output":[]}}`)) + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_ws_4","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ack-4"}]}],"usage":{"input_tokens":9,"output_tokens":3,"total_tokens":12}}}`)) + + _, _, _ = conn.ReadMessage() + return + } + + body, _ := io.ReadAll(r.Body) + httpBodies <- body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_sum\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Duplex summary with "+marker+".\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_sum\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":10,\"output_tokens\":5,\"total_tokens\":15}}}\n\n", + ) + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: server.URL, + Models: []config.CodexModel{{Name: "deepseek-v4.1-flash", IsCompat: true}}, + }}, + } + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + + wsExec := NewCodexWebsocketsExecutor(cfg) + const sessID = "duplex-compat-compaction-7406" + t.Cleanup(func() { wsExec.CloseExecutionSession(sessID) }) + + auth := &cliproxyauth.Auth{ + ID: "auth-duplex-compat-7406", + Provider: "codex", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test-compat-key", "websockets": "true"}, + } + + inputCh := make(chan cliproxyexecutor.WebsocketInput, 8) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ctx = cliproxyexecutor.WithWebsocketInput(cliproxyexecutor.WithDownstreamWebsocket(ctx), inputCh) + + turn1 := []byte(`{"type":"response.create","model":"client-alias","instructions":"Keep concise.","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Remember ` + marker + `"}]}]}`) + streamRes, errStream := wsExec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: turn1, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: sessID}, + }) + if errStream != nil { + t.Fatalf("ExecuteStream(turn1) error = %v", errStream) + } + + readNextEvent := func(wantTypes ...string) ojson.Result { + t.Helper() + deadline := time.After(5 * time.Second) + for { + select { + case chunk, ok := <-streamRes.Chunks: + if !ok { + t.Fatalf("duplex stream closed unexpectedly while waiting for %v", wantTypes) + } + if chunk.Err != nil { + t.Fatalf("duplex stream chunk error while waiting for %v: %v", wantTypes, chunk.Err) + } + trimmed := bytes.TrimSpace(chunk.Payload) + for _, line := range bytes.Split(trimmed, []byte("\n")) { + line = bytes.TrimSpace(line) + if bytes.HasPrefix(line, []byte("data:")) { + line = bytes.TrimSpace(line[len("data:"):]) + } + if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) { + continue + } + parsed := ojson.ParseBytes(line) + gotType := parsed.Get("type").String() + wantError := false + for _, want := range wantTypes { + if want == "error" { + wantError = true + } + if gotType == want { + return parsed + } + } + if gotType == "error" && !wantError { + t.Fatalf("unexpected error event while waiting for %v: %s", wantTypes, line) + } + } + case <-deadline: + t.Fatalf("timed out waiting for duplex event %v", wantTypes) + } + } + } + + // Drain Turn 1 completion. + _ = readNextEvent("response.completed") + + // 2a. response.steer carrying compaction_trigger is rejected locally without hitting upstream WS. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.steer","event_id":"ev_steer_bad","previous_response_id":"resp_ws_1","input":[{"type":"compaction_trigger"}]}`)} + steerErrEv := readNextEvent("error") + if steerErrEv.Get("status").Int() != http.StatusBadRequest || steerErrEv.Get("error.event_id").String() != "ev_steer_bad" { + t.Fatalf("unexpected steer error event: %s", steerErrEv.Raw) + } + + // 2b. Incremental response.create with previous_response_id + compaction_trigger returns 409 previous_response_not_found. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","event_id":"ev_inc_compact","previous_response_id":"resp_ws_1","input":[{"type":"compaction_trigger"}]}`)} + incErrEv := readNextEvent("error") + if incErrEv.Get("status").Int() != http.StatusConflict || incErrEv.Get("error.code").String() != "previous_response_not_found" || incErrEv.Get("error.event_id").String() != "ev_inc_compact" { + t.Fatalf("unexpected incremental compaction error event: %s", incErrEv.Raw) + } + + // 2c. Full response.create with updated instructions="INSTRUCTIONS_B" + compaction_trigger executes local HTTP summary. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","model":"client-alias","instructions":"INSTRUCTIONS_B","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Remember ` + marker + `"}]},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ack-1"}]},{"type":"compaction_trigger"}]}`)} + compactCompleted := readNextEvent("response.completed") + synthRespID := compactCompleted.Get("response.id").String() + if !strings.HasPrefix(synthRespID, "resp_ag_compact_") { + t.Fatalf("expected synthetic compaction response ID, got %q (%s)", synthRespID, compactCompleted.Raw) + } + if compactCompleted.Get("response.output.0.type").String() != "compaction" { + t.Fatalf("expected compaction item in duplex compaction output: %s", compactCompleted.Raw) + } + summaryHTTPBody := <-httpBodies + if got := ojson.GetBytes(summaryHTTPBody, "model").String(); got != "deepseek-v4.1-flash" { + t.Fatalf("HTTP summary model = %q, want real model; body=%s", got, summaryHTTPBody) + } + if bytes.Contains(summaryHTTPBody, []byte("compaction_trigger")) || ojson.GetBytes(summaryHTTPBody, "tools").Exists() { + t.Fatalf("duplex HTTP summary body leaked compaction_trigger or tools: %s", summaryHTTPBody) + } + + // Invalid synthetic-parent inputs must be rejected without consuming the parent. + for _, invalid := range []string{`{}`, `42`, `true`, `null`} { + payload := fmt.Sprintf(`{"type":"response.create","event_id":"invalid-input","previous_response_id":%q,"input":%s}`, synthRespID, invalid) + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(payload)} + ev := readNextEvent("error") + if ev.Get("status").Int() != http.StatusBadRequest || ev.Get("error.event_id").String() != "invalid-input" { + t.Fatalf("input %s: expected correlated 400, got %s", invalid, ev.Raw) + } + } + + // 3. response.append without instructions or previous_response_id inherits INSTRUCTIONS_B (not initial Keep concise.) and rewrites to expanded root response.create. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.append","input":"What was the marker?"}`)} + turn3Completed := readNextEvent("response.completed") + if turn3Completed.Get("response.id").String() != "resp_ws_3" { + t.Fatalf("turn 3 response.id = %q, want resp_ws_3", turn3Completed.Get("response.id").String()) + } + firstFrame := <-wsFrames + if got := ojson.GetBytes(firstFrame, "model").String(); got != "deepseek-v4.1-flash" { + t.Fatalf("initial model = %q, want real model; frame=%s", got, firstFrame) + } + wsTurn3 := <-wsFrames + if got := ojson.GetBytes(wsTurn3, "model").String(); got != "deepseek-v4.1-flash" { + t.Fatalf("continuation model = %q, want real model; frame=%s", got, wsTurn3) + } + items := ojson.GetBytes(wsTurn3, "input").Array() + last := items[len(items)-1] + if last.Get("role").String() != "user" || !strings.Contains(last.Raw, "What was the marker?") { + t.Fatalf("synthetic-parent string input lost or not a user message: %s", wsTurn3) + } + if ojson.GetBytes(wsTurn3, "previous_response_id").Exists() { + t.Fatalf("synthetic previous_response_id leaked to upstream WS on turn 3: %s", wsTurn3) + } + if gotInst := ojson.GetBytes(wsTurn3, "instructions").String(); gotInst != "INSTRUCTIONS_B" { + t.Fatalf("turn 3 instructions = %q, want INSTRUCTIONS_B (from compaction turn, not initial); frame=%s", gotInst, wsTurn3) + } + if bytes.Contains(wsTurn3, []byte("cpa-ag-compact-v1:")) || !bytes.Contains(wsTurn3, []byte(marker)) { + t.Fatalf("expected expanded summary marker without raw capsule on turn 3 WS frame: %s", wsTurn3) + } + + // 4. Subsequent incremental continuation referencing resp_ws_3 forwards previous_response_id on upstream WS. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","model":"client-alias","previous_response_id":"resp_ws_3","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Follow up 4"}]}]}`)} + turn4Completed := readNextEvent("response.completed") + if turn4Completed.Get("response.id").String() != "resp_ws_4" { + t.Fatalf("turn 4 response.id = %q, want resp_ws_4", turn4Completed.Get("response.id").String()) + } + wsTurn4 := <-wsFrames + if got := ojson.GetBytes(wsTurn4, "model").String(); got != "deepseek-v4.1-flash" { + t.Fatalf("ordinary continuation model = %q, want real model; frame=%s", got, wsTurn4) + } + if gotPrev := ojson.GetBytes(wsTurn4, "previous_response_id").String(); gotPrev != "resp_ws_3" { + t.Fatalf("turn 4 WS previous_response_id = %q, want resp_ws_3; frame=%s", gotPrev, wsTurn4) + } + + // 5. Unknown synthetic resp_ag_compact_* is rejected locally with 409 previous_response_not_found. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","event_id":"ev_stale_synth","previous_response_id":"resp_ag_compact_unknown","input":[{"type":"message","role":"user","content":"stale"}]}`)} + staleErrEv := readNextEvent("error") + if staleErrEv.Get("status").Int() != http.StatusConflict || staleErrEv.Get("error.code").String() != "previous_response_not_found" { + t.Fatalf("unexpected stale synthetic ID error: %s", staleErrEv.Raw) + } +} + +func TestCodexExecutor_IsCompatCompactionDuplexEOFCancelsHTTP_Issue7406(t *testing.T) { + httpStarted := make(chan struct{}) + httpCanceled := make(chan struct{}) + closeUpstreamWS := make(chan struct{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if websocket.IsWebSocketUpgrade(r) { + conn, errUp := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUp != nil { + return + } + defer func() { _ = conn.Close() }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + _, _, errRead := conn.ReadMessage() + if errRead != nil { + return + } + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_ws_1","output":[]}}`)) + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_ws_1","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}`)) + <-closeUpstreamWS + return + } + + _, _ = io.ReadAll(r.Body) + close(httpStarted) + <-r.Context().Done() + close(httpCanceled) + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: server.URL, + Models: []config.CodexModel{{Name: "deepseek-v4.1-flash", IsCompat: true}}, + }}, + } + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + + wsExec := NewCodexWebsocketsExecutor(cfg) + const sessID = "duplex-compat-eof-7406" + t.Cleanup(func() { wsExec.CloseExecutionSession(sessID) }) + auth := &cliproxyauth.Auth{ + ID: "auth-duplex-eof-7406", + Provider: "codex", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test-compat-key", "websockets": "true"}, + } + + inputCh := make(chan cliproxyexecutor.WebsocketInput, 4) + ctx := cliproxyexecutor.WithWebsocketInput(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), inputCh) + streamRes, errStream := wsExec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: []byte(`{"type":"response.create","model":"deepseek-v4.1-flash","input":[{"type":"message","role":"user","content":"turn1"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: sessID}, + }) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v", errStream) + } + + // Drain Turn 1 response.completed. + for chunk := range streamRes.Chunks { + if ojson.GetBytes(bytes.TrimSpace(chunk.Payload), "type").String() == "response.completed" { + break + } + } + + // Trigger in-flight HTTP compaction, wait at synchronization barrier, then drop upstream WS. + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","event_id":"ev_compact_eof","model":"deepseek-v4.1-flash","input":[{"type":"message","role":"user","content":"turn1"},{"type":"compaction_trigger"}]}`)} + select { + case <-httpStarted: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for duplex HTTP compaction to start") + } + close(closeUpstreamWS) + + select { + case <-httpCanceled: + case <-time.After(5 * time.Second): + t.Fatal("upstream WS EOF did not cancel in-flight duplex HTTP compaction") + } + + var gotTurnErr *cliproxyexecutor.UpstreamWebsocketTurnError + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + if uErr, ok := chunk.Err.(*cliproxyexecutor.UpstreamWebsocketTurnError); ok { + gotTurnErr = uErr + } + } + } + if gotTurnErr == nil { + t.Fatal("expected UpstreamWebsocketTurnError on duplex WS EOF during compaction") + } + if len(gotTurnErr.Abandoned) != 1 || gotTurnErr.Abandoned[0].EventID != "ev_compact_eof" || !gotTurnErr.Abandoned[0].NotSent { + t.Fatalf("expected correlated NotSent abandoned error for ev_compact_eof, got idle=%t abandoned=%+v", gotTurnErr.Idle, gotTurnErr.Abandoned) + } +} + +func TestCodexExecutor_IsCompatCompactionDuplexDeliveryOrder_Issue7406(t *testing.T) { + sendTurn1Complete := make(chan struct{}) + httpStarted := make(chan struct{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if websocket.IsWebSocketUpgrade(r) { + conn, errUp := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUp != nil { + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close upstream websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + _, _, errRead := conn.ReadMessage() + if errRead != nil { + return + } + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_turn1","output":[]}}`)) + <-sendTurn1Complete + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_turn1","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}`)) + _, _, _ = conn.ReadMessage() + return + } + + _, _ = io.ReadAll(r.Body) + close(httpStarted) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_s\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Ordered summary.\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_s\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":4,\"output_tokens\":2,\"total_tokens\":6}}}\n\n", + ) + })) + defer server.Close() + + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: server.URL, + Models: []config.CodexModel{{Name: "deepseek-v4.1-flash", IsCompat: true}}, + }}, + } + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + + wsExec := NewCodexWebsocketsExecutor(cfg) + terminalPending := make(chan struct{}, 1) + compactionReady := make(chan bool, 8) + wsExec.duplexStateObserver = func(event string, ready bool) { + switch event { + case "terminal_pending": + terminalPending <- struct{}{} + case "compaction_wait": + compactionReady <- ready + } + } + const sessID = "duplex-compat-order-7406" + t.Cleanup(func() { wsExec.CloseExecutionSession(sessID) }) + auth := &cliproxyauth.Auth{ + ID: "auth-duplex-order-7406", + Provider: "codex", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test-compat-key", "websockets": "true"}, + } + + inputCh := make(chan cliproxyexecutor.WebsocketInput, 4) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ctx = cliproxyexecutor.WithWebsocketInput(cliproxyexecutor.WithDownstreamWebsocket(ctx), inputCh) + + streamRes, errStream := wsExec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "deepseek-v4.1-flash", + Payload: []byte(`{"type":"response.create","model":"deepseek-v4.1-flash","input":[{"type":"message","role":"user","content":"turn1"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: sessID}, + }) + if errStream != nil { + t.Fatalf("ExecuteStream error = %v", errStream) + } + + // Read Turn 1 response.created so inputReady is open. + firstChunk := <-streamRes.Chunks + if ojson.GetBytes(bytes.TrimSpace(firstChunk.Payload), "type").String() != "response.created" { + t.Fatalf("expected Turn 1 response.created first, got %s", firstChunk.Payload) + } + + // Let the reader finish terminal bookkeeping, but leave its unbuffered send + // blocked. Only then enqueue compaction and observe the writer evaluating its + // fence. This proves both goroutines reached the disputed delivery window. + close(sendTurn1Complete) + select { + case <-terminalPending: + case <-time.After(5 * time.Second): + t.Fatal("reader did not reach terminal delivery") + } + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(`{"type":"response.create","model":"deepseek-v4.1-flash","input":[{"type":"message","role":"user","content":"turn1"},{"type":"compaction_trigger"}]}`)} + select { + case ready := <-compactionReady: + if ready { + t.Fatal("compaction fence opened before prior terminal delivery") + } + case <-time.After(5 * time.Second): + t.Fatal("writer did not evaluate the compaction fence") + } + + secondChunk := <-streamRes.Chunks + if ojson.GetBytes(bytes.TrimSpace(secondChunk.Payload), "type").String() != "response.completed" || + ojson.GetBytes(bytes.TrimSpace(secondChunk.Payload), "response.id").String() != "resp_turn1" { + t.Fatalf("expected Turn 1 response.completed before any Turn 2 synthetic event, got: %s", secondChunk.Payload) + } + + // Now Turn 2 synthetic events must follow in order ending with synthetic response.completed. + var sawSynthCreated, sawSynthCompleted bool + for !sawSynthCompleted { + chunk := <-streamRes.Chunks + for _, line := range bytes.Split(chunk.Payload, []byte("\n")) { + line = bytes.TrimSpace(line) + if bytes.HasPrefix(line, []byte("data:")) { + line = bytes.TrimSpace(line[len("data:"):]) + } + if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) { + continue + } + evType := ojson.GetBytes(line, "type").String() + if evType == "response.created" { + sawSynthCreated = true + } + if evType == "response.completed" { + if !sawSynthCreated { + t.Fatalf("synthetic response.completed arrived before response.created: %s", line) + } + sawSynthCompleted = true + } + } + } +} diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go index dc3f3c5cbb5..fae2f322010 100644 --- a/internal/runtime/executor/codex_executor_execute.go +++ b/internal/runtime/executor/codex_executor_execute.go @@ -17,19 +17,110 @@ import ( log "github.com/sirupsen/logrus" ) +type codexCompactionSummaryContextKey struct{} + +type codexCompactionSummaryState struct { + baseModel string +} + +func withCodexCompactionSummary(ctx context.Context, baseModel string) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, codexCompactionSummaryContextKey{}, &codexCompactionSummaryState{baseModel: baseModel}) +} + +func codexCompactionSummaryStateFromContext(ctx context.Context) *codexCompactionSummaryState { + if ctx == nil { + return nil + } + state, _ := ctx.Value(codexCompactionSummaryContextKey{}).(*codexCompactionSummaryState) + return state +} + +func codexCompactionSummaryFromContext(ctx context.Context) bool { + return codexCompactionSummaryStateFromContext(ctx) != nil +} + +func rejectSyntheticCompactionPreviousResponseID(ctx context.Context, payload []byte) error { + prevID := strings.TrimSpace(ojson.GetBytesNoCopy(payload, "previous_response_id").String()) + if !helps.RecognizedResponsesCompactionResponseID(prevID) { + return nil + } + if cliproxyexecutor.DownstreamWebsocket(ctx) { + return &cliproxyexecutor.UpstreamWebsocketTurnError{ + FullContextRequired: true, + EventID: strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(payload, "event_id").String())), + } + } + return statusErr{ + code: http.StatusBadRequest, + msg: "previous_response_id references a local compaction response; resend full conversation input without previous_response_id", + } +} + +func validateCodexLocalCompactionSource(ctx context.Context, payload []byte) error { + prevID := strings.TrimSpace(ojson.GetBytesNoCopy(payload, "previous_response_id").String()) + reqType := strings.TrimSpace(ojson.GetBytesNoCopy(payload, "type").String()) + if prevID != "" || reqType == "response.append" { + if cliproxyexecutor.DownstreamWebsocket(ctx) { + return &cliproxyexecutor.UpstreamWebsocketTurnError{ + FullContextRequired: true, + EventID: strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(payload, "event_id").String())), + } + } + return statusErr{ + code: http.StatusBadRequest, + msg: "local summary compaction requires full conversation input without previous_response_id", + } + } + if !helps.HasResponsesCompactionConversationInput(payload) { + return statusErr{ + code: http.StatusBadRequest, + msg: "local summary compaction requires non-empty conversation input", + } + } + return nil +} + +func (e *CodexExecutor) expandCodexResponsesCompaction(auth *cliproxyauth.Auth, req *cliproxyexecutor.Request, opts *cliproxyexecutor.Options, baseModel string) (bool, error) { + if req == nil || opts == nil { + return false, nil + } + isCompat := e.resolveCodexModelIsCompat(auth, *req, baseModel) + if isCompat && opts.SourceFormat == sdktranslator.FormatOpenAIResponse && opts.Alt != "responses/compact" { + return true, helps.ExpandResponsesCompactionCapsules(req, opts, e.Identifier()) + } + return isCompat, helps.ExpandRecognizedResponsesCompactionCapsules(req, opts) +} + +func (e *CodexExecutor) shouldCompactCompatResponsesTrigger(isCompat bool, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) bool { + return isCompat && opts.SourceFormat == sdktranslator.FormatOpenAIResponse && opts.Alt != "responses/compact" && helps.HasResponsesCompactionTriggerRequest(req, opts) +} + func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { ctx = helps.EnsureSessionContext(ctx, opts, req.Payload) ctx = helps.WithCodexCommentaryPolicy(ctx, e.cfg) - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { - return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} - } if opts.Alt == "responses/compact" { return e.executeCompact(ctx, auth, req, opts) } if isCodexOpenAIImageRequest(opts) { + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } return e.executeOpenAIImage(ctx, auth, req, opts) } baseModel := thinking.ParseSuffix(req.Model).ModelName + isCompat, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel) + if errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } + if !codexCompactionSummaryFromContext(ctx) && e.shouldCompactCompatResponsesTrigger(isCompat, req, opts) { + return e.executeCodexResponsesCompaction(ctx, auth, req, opts) + } + if errPrev := rejectSyntheticCompactionPreviousResponseID(ctx, req.Payload); errPrev != nil { + return resp, errPrev + } apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) if errSIWC := ensureCodexSIWCCredentialsForRequest(ctx, auth, apiKey, opts.Headers); errSIWC != nil { @@ -50,7 +141,6 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - isCompat := e.resolveCodexModelIsCompat(auth, req, baseModel) targetModel := resolveResponsesSignatureTargetModel(baseModel, req, opts) originalTranslated, body, updatesChanged, err := translateCodexRequestPairWithUpdateIntentContext(ctx, from, to, targetModel, originalPayload, req.Payload, false, isCompat) if err != nil { @@ -76,10 +166,13 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re useSIWC := codexUseSIWCForRequest(ctx, auth, opts.Headers) instructionsResult := ojson.GetBytesNoCopy(req.Payload, "instructions") hasExplicitInstructions := instructionsResult.Exists() && strings.TrimSpace(instructionsResult.String()) != "" - ensureWebSearchToolForHistoryDoc(doc, opts.Headers) + isCompactionSummary := codexCompactionSummaryFromContext(ctx) + if !isCompactionSummary { + ensureWebSearchToolForHistoryDoc(doc, opts.Headers) + } if !useSIWC { normalizeCodexInstructionsDoc(doc, helps.IsNativeCodexRequest(req.Payload, opts)) - if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { + if !isCompactionSummary && (e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff) { checkModel := req.Model if strings.TrimSpace(checkModel) == "" { checkModel = baseModel @@ -153,11 +246,14 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClearReplay != nil { return resp, errClearReplay } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, b, e.modelLevelCooling()) return resp, err @@ -186,9 +282,15 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if streamErr, terminalBody, ok := codexTerminalFailureErrWithCooling(eventData, e.modelLevelCooling()); ok { if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + if detail, hasDetail := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); hasDetail { + reporter.PublishFailureWithDetail(ctx, detail, errClearReplay) + } return resp, errClearReplay } err = streamErr + if detail, hasDetail := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); hasDetail { + reporter.PublishFailureWithDetail(ctx, detail, err) + } return resp, err } @@ -212,6 +314,9 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if helps.IsCodexTerminalEmptyIncomplete(eventData, len(outputItemsByIndex)+len(outputItemsFallback), sawOutputDelta) { err = newCodexEmptyIncompleteStreamError() + if detail, hasDetail := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); hasDetail { + reporter.PublishFailureWithDetail(ctx, detail, err) + } return resp, err } @@ -225,13 +330,43 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 { return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} } - if detail, ok := helps.ParseCodexUsageWithOutputItems(completedData, outputItemsByIndex, outputItemsFallback); ok { - reporter.Publish(ctx, detail) - } - publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, completedData, outputItemsByIndex, outputItemsFallback) + detail, hasDetail := helps.ParseCodexUsageWithOutputItems(completedData, outputItemsByIndex, outputItemsFallback) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } + if summaryState := codexCompactionSummaryStateFromContext(ctx); summaryState != nil { + if eventType == "response.incomplete" { + err = statusErr{code: http.StatusBadGateway, msg: "compaction summary response status is \"incomplete\", want \"completed\""} + if hasDetail { + reporter.PublishFailureWithDetail(ctx, detail, err) + } else { + reporter.PublishFailure(ctx, err) + } + return cliproxyexecutor.Response{}, err + } + compactedPayload, errBuild := helps.BuildResponsesCompactionResponseFromSummary(summaryState.baseModel, out) + if errBuild != nil { + err = statusErr{code: http.StatusBadGateway, msg: errBuild.Error()} + if hasDetail { + reporter.PublishFailureWithDetail(ctx, detail, err) + } else { + reporter.PublishFailure(ctx, err) + } + return cliproxyexecutor.Response{}, err + } + if hasDetail { + reporter.Publish(ctx, detail) + } else { + reporter.EnsurePublished(ctx) + } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, completedData, outputItemsByIndex, outputItemsFallback) + resp = cliproxyexecutor.Response{Payload: compactedPayload, Headers: httpResp.Header.Clone()} + return resp, nil + } + if hasDetail { + reporter.Publish(ctx, detail) + } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, completedData, outputItemsByIndex, outputItemsFallback) resp = cliproxyexecutor.Response{Payload: out, Incomplete: eventType == "response.incomplete", Headers: httpResp.Header.Clone()} return resp, nil } @@ -247,6 +382,51 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re return resp, err } +func (e *CodexExecutor) executeCodexResponsesCompaction(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + var resp cliproxyexecutor.Response + sourcePayload := helps.ResponsesCompactionSourcePayload(req, opts) + if errValidate := validateCodexLocalCompactionSource(ctx, sourcePayload); errValidate != nil { + return resp, errValidate + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + summaryPayload := helps.PrepareAntigravityCompactionSummaryPayload(sourcePayload, baseModel) + summaryReq := req + summaryReq.Payload = summaryPayload + summaryOpts := opts + summaryOpts.Alt = "" + summaryOpts.Stream = false + summaryOpts.OriginalRequest = nil + summaryOpts.SourceFormat = sdktranslator.FormatOpenAIResponse + summaryOpts.ResponseFormat = sdktranslator.FormatOpenAIResponse + + return e.Execute(withCodexCompactionSummary(ctx, baseModel), auth, summaryReq, summaryOpts) +} + +func (e *CodexExecutor) executeCodexResponsesCompactionStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + cliproxyexecutor.MarkLocalWebsocketTurn(ctx) + baseModel := thinking.ParseSuffix(req.Model).ModelName + summaryResp, errSummary := e.executeCodexResponsesCompaction(ctx, auth, req, opts) + if errSummary != nil { + return nil, errSummary + } + chunks, errChunks := helps.BuildResponsesCompactionStreamChunksFromResponse(baseModel, summaryResp.Payload) + if errChunks != nil { + return nil, statusErr{code: http.StatusBadGateway, msg: errChunks.Error()} + } + out := make(chan cliproxyexecutor.StreamChunk, len(chunks)) + for _, chunk := range chunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + close(out) + + headers := summaryResp.Headers.Clone() + if headers == nil { + headers = make(http.Header) + } + headers.Set("Content-Type", "text/event-stream") + return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil +} + func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} @@ -346,18 +526,21 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, b, e.modelLevelCooling()) return resp, err } data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) upstreamData := helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) var param any out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, upstreamData, ¶m) diff --git a/internal/runtime/executor/codex_executor_execute_usage_test.go b/internal/runtime/executor/codex_executor_execute_usage_test.go index 18fd2f1af27..b2910619926 100644 --- a/internal/runtime/executor/codex_executor_execute_usage_test.go +++ b/internal/runtime/executor/codex_executor_execute_usage_test.go @@ -6,10 +6,13 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/MoeForge/moego/ojson" + "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" _ "github.com/router-for-me/CLIProxyAPI/v8/internal/translator" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" @@ -265,3 +268,61 @@ func TestCodexExecutorExecuteCompactRecordsFinalizedPayloadOverrideReasoningEffo t.Fatalf("usage record ReasoningEffort = %q, want %q", record.ReasoningEffort, "high") } } + +func TestCodexExecutorPreservesPartialNon2xxBodyAndReadErrorWhenRequestLogDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, alt := range []string{"", "responses/compact"} { + name := "responses" + if alt != "" { + name = "compact" + } + t.Run(name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hijacker, ok := w.(http.Hijacker) + if !ok { + t.Error("response writer does not support hijacking") + return + } + conn, buf, errHijack := hijacker.Hijack() + if errHijack != nil { + t.Errorf("hijack: %v", errHijack) + return + } + _, _ = buf.WriteString("HTTP/1.1 502 Bad Gateway\r\nContent-Type: application/json\r\nContent-Length: 4096\r\n\r\n{\"error\":{\"message\":\"partial codex 502") + _ = buf.Flush() + _ = conn.Close() + })) + defer server.Close() + + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + executor := NewCodexExecutor(&config.Config{}) + _, errExec := executor.Execute(ctx, codexOAuthTestAuth(server.URL), cliproxyexecutor.Request{ + Model: "gpt-5.5", + Payload: []byte(`{"model":"gpt-5.5","input":[{"type":"message","role":"user","content":"hi"}]}`), + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, Alt: alt}) + if errExec == nil { + t.Fatal("expected Execute to fail on 502") + } + + value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey) + if !exists { + t.Fatal("deferred API response was not recorded") + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil { + t.Fatalf("deferred API response = %#v, want logging.DeferredAPIResponse", value) + } + rendered := string(deferred.RenderAPIResponse()) + for _, want := range []string{ + "Status: 502", + `{"error":{"message":"partial codex 502`, + "Error: unexpected EOF", + } { + if !strings.Contains(rendered, want) { + t.Fatalf("deferred API response missing %q:\n%s", want, rendered) + } + } + }) + } +} diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index 6074124a530..3223e020a16 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -25,13 +25,23 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} } - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { - return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} - } if isCodexOpenAIImageRequest(opts) { + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } return e.executeOpenAIImageStream(ctx, auth, req, opts) } baseModel := thinking.ParseSuffix(req.Model).ModelName + isCompat, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel) + if errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } + if e.shouldCompactCompatResponsesTrigger(isCompat, req, opts) { + return e.executeCodexResponsesCompactionStream(ctx, auth, req, opts) + } + if errPrev := rejectSyntheticCompactionPreviousResponseID(ctx, req.Payload); errPrev != nil { + return nil, errPrev + } apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) if errSIWC := ensureCodexSIWCCredentialsForRequest(ctx, auth, apiKey, opts.Headers); errSIWC != nil { @@ -54,7 +64,6 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - isCompat := e.resolveCodexModelIsCompat(auth, req, baseModel) targetModel := resolveResponsesSignatureTargetModel(baseModel, req, opts) originalTranslated, body, updatesChanged, err := translateCodexRequestPairWithUpdateIntentContext(ctx, from, to, targetModel, originalPayload, req.Payload, true, isCompat) if err != nil { @@ -177,6 +186,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("codex executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if readErr != nil { helps.RecordAPIResponseError(ctx, e.cfg, readErr) return nil, readErr @@ -184,7 +194,6 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, data); errClearReplay != nil { return nil, errClearReplay } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, data, e.modelLevelCooling()) return nil, err diff --git a/internal/runtime/executor/codex_executor_tokens.go b/internal/runtime/executor/codex_executor_tokens.go index 22c73eb72af..e65d9a18f58 100644 --- a/internal/runtime/executor/codex_executor_tokens.go +++ b/internal/runtime/executor/codex_executor_tokens.go @@ -18,10 +18,10 @@ import ( func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { ctx = helps.WithCodexCommentaryPolicy(ctx, e.cfg) - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + baseModel := thinking.ParseSuffix(req.Model).ModelName + if _, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel); errExpand != nil { return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} } - baseModel := thinking.ParseSuffix(req.Model).ModelName from := opts.SourceFormat responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 11e60c3d96b..7f11a7cdaed 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -140,11 +140,11 @@ func (e *CodexExecutor) executeOpenAIImage(ctx context.Context, auth *cliproxyau helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, data, e.modelLevelCooling()) @@ -248,11 +248,11 @@ func (e *CodexExecutor) executeOpenAIImageStream(ctx context.Context, auth *clip if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("codex executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return nil, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, data, e.modelLevelCooling()) return nil, err @@ -399,11 +399,11 @@ func (e *CodexExecutor) executeDirectOpenAIImage(ctx context.Context, auth *clip helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, data, e.modelLevelCooling()) @@ -470,11 +470,11 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("codex executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return nil, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErrWithHeadersAndCooling(httpResp.StatusCode, httpResp.Header, data, e.modelLevelCooling()) return nil, err diff --git a/internal/runtime/executor/codex_websockets_duplex.go b/internal/runtime/executor/codex_websockets_duplex.go index e91333b24f8..29cd1569f25 100644 --- a/internal/runtime/executor/codex_websockets_duplex.go +++ b/internal/runtime/executor/codex_websockets_duplex.go @@ -44,15 +44,25 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( ) *cliproxyexecutor.StreamResult { streamCtx, cancel := context.WithCancel(ctx) // Duplex successors bypass dispatch, including automatic steering responses. + // Reader usage reporting uses ctx so post-cancel failure publishing still has + // request metadata; writer-side local HTTP calls derive from streamCtx so an + // upstream websocket disconnect cancels in-flight HTTP compaction immediately. followupCtx := ctx + writerFollowupCtx := streamCtx if selection := coreusage.HomeSelectionFromContext(ctx); selection != nil { - selection.Reused = true - followupCtx = coreusage.WithHomeSelection(ctx, *selection) + reusedSelection := *selection + reusedSelection.Reused = true + followupCtx = coreusage.WithHomeSelection(ctx, reusedSelection) + writerFollowupCtx = coreusage.WithHomeSelection(streamCtx, reusedSelection) } out := make(chan cliproxyexecutor.StreamChunk) // This first frame was successfully written by ExecuteStream before handoff. log.Infof("codex websockets: request forwarded session=%s auth=%s url=%s event=response.create", sess.sessionID, auth.ID, initial.wsURL) - writeErrors := make(chan error, 1) + type writerFailure struct { + err error + preserveClassification bool + } + writeErrors := make(chan writerFailure, 1) // Explicit creates have their own request settings. Automatic successors // inherit the preceding response settings, just as they do upstream. var metadataMu sync.Mutex @@ -67,9 +77,27 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( acceptedSteers := make(map[string]codexAcceptedSteer) current := initial responseID := "" + latestTurnResponseID := "" responseSettings := make(map[string]*codexWebsocketPrepared) steeringSettings := make(map[string]*codexWebsocketPrepared) var responseOrder []string + firstResponse := true + responseActive := false + terminalDeliveryPending := false + localCompactionDeliveredAny := false + localCompactionItems := make(map[string]string) + var localCompactionOrder []string + storeResponseSettingsLocked := func(id string, snapshot *codexWebsocketPrepared) { + if id == "" || snapshot == nil { + return + } + responseSettings[id] = snapshot + responseOrder = append(responseOrder, id) + if len(responseOrder) > 16 { + delete(responseSettings, responseOrder[0]) + responseOrder = responseOrder[1:] + } + } releaseSteeringSettings := func(parent string) { for _, target := range unacknowledgedSteers { if target.parent == parent { @@ -135,17 +163,32 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } go func() { defer close(writerDone) - fail := func(err error) { + failWithClassification := func(err error, preserveClassification bool) { select { - case writeErrors <- err: + case writeErrors <- writerFailure{err: err, preserveClassification: preserveClassification}: default: } cancel() } - reject := func(message string) bool { + fail := func(err error) { failWithClassification(err, false) } + emitDuplexError := func(status int, errType, code, param, message, eventID string) bool { + errBody := map[string]string{ + "type": errType, + "message": message, + } + if code != "" { + errBody["code"] = code + } + if param != "" { + errBody["param"] = param + } + if eventID != "" { + errBody["event_id"] = eventID + } payload, _ := json.Marshal(map[string]any{ - "type": "error", "status": http.StatusBadRequest, - "error": map[string]string{"type": "invalid_request_error", "message": message}, + "type": "error", + "status": status, + "error": errBody, }) // Reader cleanup joins this writer before closing out. Cancellation must // also release a local error blocked behind a disconnected downstream. @@ -157,6 +200,19 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( return false } } + reject := func(message string) bool { + return emitDuplexError(http.StatusBadRequest, "invalid_request_error", "", "", message, "") + } + rejectPreviousResponseNotFound := func(eventID string) bool { + return emitDuplexError( + http.StatusConflict, + "invalid_request_error", + "previous_response_not_found", + "previous_response_id", + "Previous response is not available on this websocket; resend the full conversation input without previous_response_id", + eventID, + ) + } initialModel := "" if initial != nil { initialModel = strings.TrimSpace(ojson.GetBytesNoCopy(initial.originalPayload, "model").String()) @@ -305,13 +361,15 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( if allowed, keepOpen := authorizePayload(payload, false); !allowed { return keepOpen } + eventID := strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(payload, "event_id").String())) metadataMu.Lock() isAppend := ojson.GetBytesNoCopy(payload, "type").String() == "response.append" prevID := strings.TrimSpace(ojson.GetBytesNoCopy(payload, "previous_response_id").String()) - if isAppend && prevID == "" && responseID != "" { - prevID = responseID + if isAppend && prevID == "" && latestTurnResponseID != "" { + prevID = latestTurnResponseID payload, _ = ojson.SetBytes(payload, "previous_response_id", prevID) } + localCompactionItem, hasLocalCompactionParent := localCompactionItems[prevID] wrongParent := false for _, accepted := range acceptedSteers { if accepted.successorID == "" && prevID != waitingParent { @@ -325,6 +383,9 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } return true } + if prevID != "" && !hasLocalCompactionParent && helps.RecognizedResponsesCompactionResponseID(prevID) { + return rejectPreviousResponseNotFound(eventID) + } model := strings.TrimSpace(ojson.GetBytesNoCopy(payload, "model").String()) if model != "" && model != req.Model && model != boundModel { fail(cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()) @@ -336,11 +397,9 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( if !allowed { return keepOpen } + model = req.Model if model == "" { - model = req.Model - if model == "" { - model = boundModel - } + model = boundModel } createDoc := ojson.NewDoc(payload) _ = createDoc.Set("model", model) @@ -369,11 +428,187 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( _ = createDoc.SetRawBytes("instructions", []byte(inst.Raw)) } } + if hasLocalCompactionParent { + if !createDoc.Get("instructions").Exists() { + metadataMu.Lock() + parentSettings := responseSettings[prevID] + metadataMu.Unlock() + targetSettings := parentSettings + if targetSettings == nil { + targetSettings = initial + } + var inst ojson.Result + if targetSettings != nil { + inst = ojson.GetBytesNoCopy(targetSettings.clientBody, "instructions") + if !inst.Exists() { + inst = ojson.GetBytesNoCopy(targetSettings.originalPayload, "instructions") + } + } + if !inst.Exists() && initial != nil { + inst = ojson.GetBytesNoCopy(initial.clientBody, "instructions") + if !inst.Exists() { + inst = ojson.GetBytesNoCopy(initial.originalPayload, "instructions") + } + } + if inst.Exists() { + _ = createDoc.SetRawBytes("instructions", []byte(inst.Raw)) + } + } + existingInput := createDoc.Get("input") + rebuiltItems := make([]string, 0, 2) + rebuiltItems = append(rebuiltItems, localCompactionItem) + switch { + case existingInput.IsArray(): + for _, item := range existingInput.Array() { + rebuiltItems = append(rebuiltItems, item.Raw) + } + case existingInput.Type == ojson.String: + message := ojson.NewDocString(`{"type":"message","role":"user"}`) + _ = message.Set("content", existingInput.String()) + rebuiltItems = append(rebuiltItems, message.String()) + case existingInput.Exists(): + return emitDuplexError(http.StatusBadRequest, "invalid_request_error", "", "input", "input must be a string or an array", eventID) + } + _ = createDoc.SetRaw("input", "["+strings.Join(rebuiltItems, ",")+"]") + _ = createDoc.Delete("previous_response_id") + _ = createDoc.Set("type", "response.create") + isAppend = false + prevID = "" + } payload = createDoc.Bytes() nextReq, nextOpts := req, opts + // Dispatch already resolved aliases to the bound upstream execution model. + // The client model above is only used for authorization and consistency. + nextReq.Model = req.Model nextReq.Payload = payload nextOpts.OriginalRequest = payload nextOpts.Headers = turnHeaders + turnBaseModel := thinking.ParseSuffix(nextReq.Model).ModelName + if turnIsCompat := e.resolveCodexModelIsCompat(auth, nextReq, turnBaseModel); e.shouldCompactCompatResponsesTrigger(turnIsCompat, nextReq, nextOpts) { + if _, errExpand := e.expandCodexResponsesCompaction(auth, &nextReq, &nextOpts, turnBaseModel); errExpand != nil { + return emitDuplexError(http.StatusBadRequest, "invalid_request_error", "", "", errExpand.Error(), eventID) + } + if prevID != "" || isAppend { + return rejectPreviousResponseNotFound(eventID) + } + if !helps.HasResponsesCompactionConversationInput(nextReq.Payload) { + return emitDuplexError(http.StatusBadRequest, "invalid_request_error", "", "", "local summary compaction requires non-empty conversation input", eventID) + } + if !waitFor(func() bool { + ready := !firstResponse && !responseActive && !terminalDeliveryPending && len(pending) == 0 && len(unacknowledgedSteers) == 0 && !automaticActive + if e.duplexStateObserver != nil { + e.duplexStateObserver("compaction_wait", ready) + } + return ready + }) { + return false + } + if !cliproxyexecutor.WebsocketAuthEnabled(streamCtx, auth.ID) { + fail(fmt.Errorf("websocket credential is no longer enabled")) + return false + } + // Build business settings snapshot from the turn request (before summary + // prompt preparation strips tools/fields) so subsequent response.append + // or synthetic-parent continuations inherit this turn's instructions. + snapshotDoc := ojson.NewDoc([]byte("{}")) + if reasoning := ojson.GetBytesNoCopy(nextReq.Payload, "reasoning"); reasoning.Exists() { + _ = snapshotDoc.SetRawBytes("reasoning", []byte(reasoning.Raw)) + } + if inst := ojson.GetBytesNoCopy(nextReq.Payload, "instructions"); inst.Exists() { + _ = snapshotDoc.SetRawBytes("instructions", []byte(inst.Raw)) + } else if initial != nil { + initInst := ojson.GetBytesNoCopy(initial.clientBody, "instructions") + if !initInst.Exists() { + initInst = ojson.GetBytesNoCopy(initial.originalPayload, "instructions") + } + if initInst.Exists() { + _ = snapshotDoc.SetRawBytes("instructions", []byte(initInst.Raw)) + } + } + if serviceTier := ojson.GetBytesNoCopy(nextReq.Payload, "service_tier"); serviceTier.Exists() { + _ = snapshotDoc.SetRawBytes("service_tier", []byte(serviceTier.Raw)) + } + compactionSettings := &codexWebsocketPrepared{ + to: sdktranslator.FromString("codex"), + clientBody: snapshotDoc.Bytes(), + } + metadataMu.Lock() + localCompactionDeliveredAny = false + metadataMu.Unlock() + compactionResp, errCompact := e.CodexExecutor.executeCodexResponsesCompaction(writerFollowupCtx, auth, nextReq, nextOpts) + if errCompact != nil { + if streamCtx.Err() != nil { + return false + } + var statusHolder interface{ StatusCode() int } + status := http.StatusBadGateway + if errors.As(errCompact, &statusHolder) && statusHolder.StatusCode() > 0 { + status = statusHolder.StatusCode() + } + if status == http.StatusUnauthorized || status == http.StatusForbidden || status == http.StatusTooManyRequests { + failWithClassification(errCompact, true) + return false + } + errType := "server_error" + if status >= 400 && status < 500 { + errType = "invalid_request_error" + } + return emitDuplexError(status, errType, "", "", errCompact.Error(), eventID) + } + chunks, errChunks := helps.BuildResponsesCompactionStreamChunksFromResponse(turnBaseModel, compactionResp.Payload) + if errChunks != nil { + return emitDuplexError(http.StatusBadGateway, "server_error", "", "", errChunks.Error(), eventID) + } + compactedRespID := "" + compactedItemRaw := "" + for i := len(chunks) - 1; i >= 0; i-- { + for _, line := range bytes.Split(chunks[i], []byte("\n")) { + trimmedLine := bytes.TrimSpace(line) + if !bytes.HasPrefix(trimmedLine, dataTag) { + continue + } + eventData := bytes.TrimSpace(trimmedLine[len(dataTag):]) + if ojson.GetBytesNoCopy(eventData, "type").String() == "response.completed" { + compactedRespID = strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(eventData, "response.id").String())) + compactedItemRaw = strings.Clone(ojson.GetBytesNoCopy(eventData, "response.output.0").Raw) + break + } + } + if compactedRespID != "" { + break + } + } + if compactedRespID != "" && compactedItemRaw != "" { + metadataMu.Lock() + latestTurnResponseID = compactedRespID + waitingParent = "" + storeResponseSettingsLocked(compactedRespID, compactionSettings) + localCompactionItems[compactedRespID] = compactedItemRaw + localCompactionOrder = append(localCompactionOrder, compactedRespID) + if len(localCompactionOrder) > 16 { + delete(localCompactionItems, localCompactionOrder[0]) + localCompactionOrder = localCompactionOrder[1:] + } + metadataMu.Unlock() + } + for idx, chunk := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunk}: + if idx == 0 { + metadataMu.Lock() + localCompactionDeliveredAny = true + metadataMu.Unlock() + } + case <-streamCtx.Done(): + return false + } + } + metadataMu.Lock() + localCompactionDeliveredAny = false + metadataMu.Unlock() + consumedPayload = nil + return true + } prepared, errPrepare := e.prepareCodexWebsocketStream(streamCtx, auth, nextReq, nextOpts) if errPrepare != nil { fail(errPrepare) @@ -484,10 +719,19 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( steerOpts := opts steerOpts.OriginalRequest = payload steerOpts.Headers = steerHeaders - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&steerReq, &steerOpts); errExpand != nil { + steerBaseModel := thinking.ParseSuffix(req.Model).ModelName + steerIsCompat, errExpand := e.expandCodexResponsesCompaction(auth, &steerReq, &steerOpts, steerBaseModel) + if errExpand != nil { fail(statusErr{code: http.StatusBadRequest, msg: errExpand.Error()}) return } + if steerIsCompat && helps.HasResponsesCompactionTrigger(steerReq.Payload) { + steerEventID := strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(steerReq.Payload, "event_id").String())) + if !emitDuplexError(http.StatusBadRequest, "invalid_request_error", "", "", "compaction_trigger is not supported in response.steer", steerEventID) { + return + } + continue + } steerDoc := ojson.NewDoc(steerReq.Payload) if codexUseSIWCForRequest(ctx, auth, steerHeaders) { thinking.StripConfigurationUpdatesDoc(steerDoc) @@ -577,8 +821,6 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } } reporter := initialReporter - firstResponse := true - responseActive := false outputItems := make(map[int64][]byte) var outputFallback [][]byte // A business rejection followed by EOF with nothing outstanding stays terminal. @@ -586,14 +828,20 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( for { kind, payload, errRead := readCodexWebsocketMessage(streamCtx, sess, conn, readCh) if errRead != nil { + preserveClassification := false select { - case errWrite := <-writeErrors: - errRead = errWrite + case failure := <-writeErrors: + errRead = failure.err + preserveClassification = failure.preserveClassification default: } if ctx.Err() == nil { var connectionErr error = &codexDuplexConnectionError{cause: errRead} - if lost := codexWebsocketTurnTransportError(errRead); cliproxyexecutor.IsUpstreamWebsocketTurnError(lost) { + if preserveClassification { + // HTTP business failures must retain credential/cooldown policy. + connectionErr = errRead + } + if lost := codexWebsocketTurnTransportError(errRead); !preserveClassification && cliproxyexecutor.IsUpstreamWebsocketTurnError(lost) { // Stop and join the writer first: after that nothing else is sent // upstream and the outcome of every operation is known. cancel() @@ -632,7 +880,10 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( for _, queued := range pendingCreates { abandon(queued, true) } - abandon(consumedPayload, true) + abandon(consumedPayload, !localCompactionDeliveredAny) + if !firstResponse && !responseActive && len(consumedPayload) > 0 && len(turnErr.Abandoned) == 0 && !localCompactionDeliveredAny { + turnErr.NotSent = true + } if !firstResponse && !responseActive && len(turnErr.Abandoned) > 0 { // No running turn owns the loss: the correlated errors report it. turnErr.Idle = true @@ -660,7 +911,9 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } payload = bytes.TrimSpace(payload) eventType := ojson.GetBytesNoCopy(payload, "type").String() + metadataMu.Lock() establishing := firstResponse && eventType == "response.created" + metadataMu.Unlock() if eventType == "response.created" { metadataMu.Lock() parent := strings.Clone(ojson.GetBytesNoCopy(payload, "response.previous_response_id").String()) @@ -694,6 +947,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( current, pending = pending[0], pending[1:] } responseID = createdResponseID + latestTurnResponseID = createdResponseID // Retain response settings, not request history or authorization headers. // In-flight steering pins its parent's settings independently of this window. snapshot := *current @@ -709,22 +963,18 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( _ = snapshotDoc.SetRawBytes("service_tier", []byte(serviceTier.Raw)) } snapshot.clientBody = snapshotDoc.Bytes() - responseSettings[responseID] = &snapshot - responseOrder = append(responseOrder, responseID) - if len(responseOrder) > 16 { - delete(responseSettings, responseOrder[0]) - responseOrder = responseOrder[1:] - } + storeResponseSettingsLocked(responseID, &snapshot) + wasFirstResponse := firstResponse + firstResponse = false + responseActive = true releaseSteeringSettings(parent) metadataMu.Unlock() wakeWriter() - if !firstResponse { + if !wasFirstResponse { reporter = helps.NewExecutorUsageReporter(followupCtx, e, req.Model, auth) reporter.SetTranslatedReasoningEffort(current.clientBody, current.to.String()) reporter.StartResponseTTFT() } - firstResponse = false - responseActive = true outputItems = make(map[int64][]byte) outputFallback = nil } @@ -736,10 +986,10 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( if strings.HasPrefix(eventType, "response.steer.") { id := strings.Clone(ojson.GetBytesNoCopy(payload, "steer.id").String()) parent := strings.Clone(ojson.GetBytesNoCopy(payload, "steer.previous_response_id").String()) + metadataMu.Lock() if parent == "" { parent = responseID } - metadataMu.Lock() consumeSubmission := func() codexSteerSubmission { for i, target := range unacknowledgedSteers { if target.parent == parent || target.parent == "" { @@ -771,7 +1021,10 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } continue } - if !firstResponse && (eventType == "error" || eventType == "response.failed") { + metadataMu.Lock() + isInitialResponse := firstResponse + metadataMu.Unlock() + if !isInitialResponse && (eventType == "error" || eventType == "response.failed") { var credentialErr error if wsErr, ok := parseCodexWebsocketErrorWithCooling(payload, e.modelLevelCooling()); ok { credentialErr = wsErr @@ -791,7 +1044,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } } eventPrepared, eventReporter := current, reporter - if !firstResponse && (eventType == "response.failed" || eventType == "error") { + if !isInitialResponse && (eventType == "response.failed" || eventType == "error") { failedID := ojson.GetBytesNoCopy(payload, "response.id").String() if failedID == "" { failedID = ojson.GetBytesNoCopy(payload, "response_id").String() @@ -806,9 +1059,11 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( eventPrepared, pending = pending[0], pending[1:] eventReporter = helps.NewExecutorUsageReporter(followupCtx, e, req.Model, auth) eventReporter.SetTranslatedReasoningEffort(eventPrepared.clientBody, eventPrepared.to.String()) + terminalDeliveryPending = true } else if !ambiguous { responseActive = false automaticActive = false + terminalDeliveryPending = true } metadataMu.Unlock() wakeWriter() @@ -848,7 +1103,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( terminalBusinessRejection = errors.As(terminalErr, &status) && status.StatusCode() == http.StatusBadRequest helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", terminalErr) eventReporter.PublishFailure(ctx, terminalErr) - if firstResponse { + if isInitialResponse { send(cliproxyexecutor.StreamChunk{Err: terminalErr}) return } @@ -861,11 +1116,11 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( if endedID == "" { endedID = ojson.GetBytesNoCopy(payload, "response_id").String() } + metadataMu.Lock() if endedID == "" && !responseActive { // An id-less failure that ended the running response. endedID = responseID } - metadataMu.Lock() // A steer is complete once the response that consumed it has ended. for id, accepted := range acceptedSteers { if endedID != "" && accepted.successorID == endedID { @@ -876,11 +1131,11 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } if eventType == "response.completed" || eventType == "response.done" || eventType == "response.incomplete" { terminalBusinessRejection = false - responseActive = false metadataMu.Lock() + responseActive = false automaticActive = false + terminalDeliveryPending = true metadataMu.Unlock() - wakeWriter() payload = normalizeCodexWebsocketCompletion(payload) if !current.preserveNativeOutput { payload = patchCodexCompletedOutput(payload, outputItems, outputFallback) @@ -894,10 +1149,19 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( reporter.EnsurePublished(ctx) } publishCodexImageToolUsageWithOutputItems(ctx, reporter, current.clientBody, payload, outputItems, outputFallback) + if e.duplexStateObserver != nil { + e.duplexStateObserver("terminal_pending", false) + } } if !send(cliproxyexecutor.StreamChunk{Payload: helps.EnsureResponsesUsageDetails(payload), Incomplete: eventType == "response.incomplete"}) { return } + if isCodexTerminalEventType(eventType) { + metadataMu.Lock() + terminalDeliveryPending = false + metadataMu.Unlock() + wakeWriter() + } if establishing { // Deliver response.created before any locally generated error so the // downstream handler also observes successful bootstrap first. diff --git a/internal/runtime/executor/codex_websockets_duplex_health_test.go b/internal/runtime/executor/codex_websockets_duplex_health_test.go index fdd09734574..41dffaa6b3c 100644 --- a/internal/runtime/executor/codex_websockets_duplex_health_test.go +++ b/internal/runtime/executor/codex_websockets_duplex_health_test.go @@ -87,3 +87,87 @@ func TestCodexDuplexConnectionTimeoutDoesNotCoolHealthyAccount(t *testing.T) { t.Fatalf("connection timeout cooled the model: %+v", state.LastError) } } + +func TestCodexDuplexSummaryCredentialFailuresThroughManager(t *testing.T) { + for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden, http.StatusTooManyRequests} { + t.Run(http.StatusText(status), func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !websocket.IsWebSocketUpgrade(r) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(`{"error":{"message":"summary credential rejected"}}`)) + return + } + c, errUpgrade := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUpgrade != nil { + t.Error(errUpgrade) + return + } + defer func() { + if errClose := c.Close(); errClose != nil { + t.Errorf("close upstream websocket: %v", errClose) + } + }() + if _, _, errRead := c.ReadMessage(); errRead != nil { + t.Error(errRead) + return + } + for _, payload := range []string{`{"type":"response.created","response":{"id":"r1"}}`, `{"type":"response.completed","response":{"id":"r1","output":[]}}`} { + if errWrite := c.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Error(errWrite) + return + } + } + _, _, _ = c.ReadMessage() + })) + defer upstream.Close() + const id, model = "duplex-summary-health-account", "duplex-summary-health-model" + cfg := &config.Config{CodexKey: []config.CodexKey{{APIKey: "test-key", BaseURL: upstream.URL, Models: []config.CodexModel{{Name: model, IsCompat: true}}}}} + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + exec := NewCodexAutoExecutor(cfg) + manager.RegisterExecutor(exec) + defer exec.CloseExecutionSession("summary-health-session") + if _, errRegister := manager.Register(context.Background(), &auth.Auth{ID: id, Provider: "codex", Status: auth.StatusActive, Attributes: map[string]string{"api_key": "test-key", "base_url": upstream.URL, "websockets": "true"}}); errRegister != nil { + t.Fatal(errRegister) + } + registry.GetGlobalRegistry().RegisterClient(id, "codex", []*registry.ModelInfo{{ID: model, IsCompat: true}}) + defer registry.GetGlobalRegistry().UnregisterClient(id) + ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + defer cancel() + input := make(chan core.WebsocketInput, 1) + ctx = core.WithWebsocketInput(core.WithDownstreamWebsocket(ctx), input) + stream, errStream := manager.ExecuteStream(ctx, []string{"codex"}, core.Request{Model: model, Payload: []byte(`{"model":"duplex-summary-health-model","input":[]}`)}, core.Options{SourceFormat: translator.FormatOpenAIResponse, Metadata: map[string]any{core.ExecutionSessionMetadataKey: "summary-health-session"}}) + if errStream != nil { + t.Fatal(errStream) + } + var terminal error + for chunk := range stream.Chunks { + if ojson.GetBytes(chunk.Payload, "type").String() == "response.completed" { + input <- core.WebsocketInput{Payload: []byte(`{"type":"response.create","input":[{"role":"user","content":"history"},{"type":"compaction_trigger"}]}`)} + } + if chunk.Err != nil { + terminal = chunk.Err + } + } + var statusHolder interface{ StatusCode() int } + if !errors.As(terminal, &statusHolder) || statusHolder.StatusCode() != status { + t.Fatalf("terminal error = %v, want HTTP %d", terminal, status) + } + var scoped interface{ IsRequestScoped() bool } + if errors.As(terminal, &scoped) && scoped.IsRequestScoped() { + t.Fatalf("credential failure incorrectly classified as request-scoped: %v", terminal) + } + current, ok := manager.GetByID(id) + if !ok { + t.Fatal("account disappeared") + } + state := current.ModelStates[model] + if state == nil || state.LastError == nil || state.LastError.HTTPStatus != status || state.LastError.IsRequestScoped() || !state.Unavailable || state.NextRetryAfter.IsZero() { + t.Fatalf("HTTP %d did not retain credential failure state: %+v", status, state) + } + }) + } +} diff --git a/internal/runtime/executor/codex_websockets_duplex_test.go b/internal/runtime/executor/codex_websockets_duplex_test.go index c4d85c2a9b1..1366b1fa8e8 100644 --- a/internal/runtime/executor/codex_websockets_duplex_test.go +++ b/internal/runtime/executor/codex_websockets_duplex_test.go @@ -591,3 +591,67 @@ func TestCodexDuplexQueuedCreateDoesNotBlockSubsequentSteer(t *testing.T) { t.Fatal("server did not finish") } } + +func TestCodexDuplexNativeAliasContinuationUsesRealModel(t *testing.T) { + frames := make(chan []byte, 3) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, errUpgrade := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUpgrade != nil { + t.Error(errUpgrade) + return + } + defer func() { + if errClose := c.Close(); errClose != nil { + t.Errorf("close upstream websocket: %v", errClose) + } + }() + for i := 0; i < 3; i++ { + _, frame, errRead := c.ReadMessage() + if errRead != nil { + return + } + frames <- frame + for _, payload := range []string{`{"type":"response.created","response":{"id":"r1"}}`, `{"type":"response.completed","response":{"id":"r1","output":[]}}`} { + if errWrite := c.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Error(errWrite) + return + } + } + } + _, _, _ = c.ReadMessage() + })) + defer upstream.Close() + cfg := &config.Config{CodexKey: []config.CodexKey{{APIKey: "test-key", BaseURL: upstream.URL, Models: []config.CodexModel{{Name: "real-model", IsCompat: false}}}}} + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + exec := NewCodexWebsocketsExecutor(cfg) + defer exec.CloseExecutionSession("native-alias-session") + ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + defer cancel() + input := make(chan core.WebsocketInput, 1) + ctx = core.WithWebsocketInput(core.WithDownstreamWebsocket(ctx), input) + stream, errStream := exec.ExecuteStream(ctx, &cliproxyauth.Auth{ID: "native-alias-account", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": upstream.URL, "websockets": "true"}}, core.Request{Model: "real-model", Payload: []byte(`{"model":"client-alias","input":[]}`)}, core.Options{SourceFormat: sdktranslator.FromString("codex"), Metadata: map[string]any{core.ExecutionSessionMetadataKey: "native-alias-session"}}) + if errStream != nil { + t.Fatal(errStream) + } + for turn := 0; turn < 3; turn++ { + for { + chunk, ok := <-stream.Chunks + if !ok || chunk.Err != nil { + t.Fatalf("turn %d stream closed or failed: %v", turn, chunk.Err) + } + if ojson.GetBytes(chunk.Payload, "type").String() == "response.completed" { + break + } + } + frame := <-frames + if got := ojson.GetBytes(frame, "model").String(); got != "real-model" { + t.Fatalf("turn %d model = %q, want real-model; frame=%s", turn, got, frame) + } + if turn == 0 { + input <- core.WebsocketInput{Payload: []byte(`{"type":"response.create","model":"client-alias","previous_response_id":"r1","input":"follow up"}`)} + } else if turn == 1 { + input <- core.WebsocketInput{Payload: []byte(`{"type":"response.create","previous_response_id":"r1","input":"follow up again"}`)} + } + } +} diff --git a/internal/runtime/executor/codex_websockets_errors.go b/internal/runtime/executor/codex_websockets_errors.go index 718d56b1523..0e30d980804 100644 --- a/internal/runtime/executor/codex_websockets_errors.go +++ b/internal/runtime/executor/codex_websockets_errors.go @@ -194,15 +194,50 @@ func recordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, resp * } func websocketHandshakeBody(resp *http.Response) []byte { + body, _ := websocketHandshakeBodyWithReadErr(resp) + return body +} + +func websocketHandshakeBodyWithReadErr(resp *http.Response) ([]byte, error) { if resp == nil || resp.Body == nil { - return nil + return nil, nil } - body, _ := io.ReadAll(resp.Body) + body, errRead := io.ReadAll(resp.Body) closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") if len(body) == 0 { - return nil + return nil, errRead } - return body + return body, errRead +} + +func recordWebsocketUpgradeRejection(ctx context.Context, cfg *config.Config, wsReqLog helps.UpstreamRequestLog, resp *http.Response) []byte { + bodyErr, errRead := websocketHandshakeBodyWithReadErr(resp) + if resp != nil { + helps.RecordAPIWebsocketUpgradeRejection(ctx, cfg, websocketUpgradeRequestLog(wsReqLog), resp.StatusCode, resp.Header.Clone(), bodyErr) + if errRead != nil { + helps.RecordAPIResponseError(ctx, cfg, errRead) + } + } + return bodyErr +} + +func recordWebsocketDialFailure(ctx context.Context, cfg *config.Config, wsReqLog helps.UpstreamRequestLog, resp *http.Response, errDial error) []byte { + bodyErr := recordWebsocketUpgradeRejection(ctx, cfg, wsReqLog, resp) + helps.RecordAPIWebsocketError(ctx, cfg, "dial", errDial) + return bodyErr +} + +func recordWebsocketDialRetryFailure(ctx context.Context, cfg *config.Config, wsReqLog helps.UpstreamRequestLog, resp *http.Response, errDial error) []byte { + if resp != nil { + bodyErr := recordWebsocketUpgradeRejection(ctx, cfg, wsReqLog, resp) + helps.RecordAPIWebsocketError(ctx, cfg, "dial_retry", errDial) + return bodyErr + } + if cfg != nil && !cfg.CommercialMode && !cfg.RequestLog { + helps.RecordAPIRequest(ctx, cfg, websocketUpgradeRequestLog(wsReqLog)) + } + helps.RecordAPIWebsocketError(ctx, cfg, "dial_retry", errDial) + return nil } func closeHTTPResponseBody(resp *http.Response, logPrefix string) { diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index bafe8f305f2..e1b44ebbcd0 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -22,14 +22,21 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut ctx = context.Background() } ctx = helps.WithCodexCommentaryPolicy(ctx, e.cfg) - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { - return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} - } if opts.Alt == "responses/compact" { return e.CodexExecutor.executeCompact(ctx, auth, req, opts) } baseModel := thinking.ParseSuffix(req.Model).ModelName + isCompat, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel) + if errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } + if e.shouldCompactCompatResponsesTrigger(isCompat, req, opts) { + return e.CodexExecutor.executeCodexResponsesCompaction(ctx, auth, req, opts) + } + if errPrev := rejectSyntheticCompactionPreviousResponseID(ctx, req.Payload); errPrev != nil { + return resp, errPrev + } apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) if errSIWC := ensureCodexSIWCCredentialsForRequest(ctx, auth, apiKey, opts.Headers); errSIWC != nil { return resp, errSIWC @@ -55,7 +62,6 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - isCompat := e.resolveCodexModelIsCompat(auth, req, baseModel) targetModel := resolveResponsesSignatureTargetModel(baseModel, req, opts) originalTranslated, body, updatesChanged, err := translateCodexRequestPairWithUpdateIntentContext(ctx, from, to, targetModel, originalPayload, req.Payload, false, isCompat) if err != nil { @@ -185,10 +191,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut conn, closer, respHS, errDial = e.ensureUpstreamConnWithSIWC(dialCtx, auth, sess, authID, wsURL, wsHeaders, useSIWC) } if errDial != nil { - bodyErr := websocketHandshakeBody(respHS) - if respHS != nil { - helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) - } + bodyErr := recordWebsocketDialFailure(ctx, e.cfg, wsReqLog, respHS, errDial) if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { if opts.ExecutionLifecycle == nil && !cliproxyexecutor.DownstreamWebsocket(ctx) { return e.CodexExecutor.Execute(ctx, auth, req, opts) @@ -204,7 +207,6 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if respHS != nil && respHS.StatusCode > 0 { return resp, newCodexStatusErrWithHeadersAndCooling(respHS.StatusCode, respHS.Header, bodyErr, e.modelLevelCooling()) } - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) return resp, errDial } if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { @@ -246,8 +248,8 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) if !shouldRetryCodexWebsocketSend(errSend) { - helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) return resp, errSend } @@ -291,8 +293,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut return resp, errSendRetry } } else { - closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) + _ = recordWebsocketDialRetryFailure(ctx, e.cfg, wsReqLog, respHSRetry, errDialRetry) return resp, errDialRetry } } else { diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index a8114f9e865..199bb2d8f7a 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -24,6 +24,9 @@ type CodexWebsocketsExecutor struct { *CodexExecutor store *codexWebsocketSessionStore + + // Set before execution by concurrency tests; never mutated during a session. + duplexStateObserver func(event string, ready bool) } func NewCodexWebsocketsExecutor(cfg *config.Config) *CodexWebsocketsExecutor { diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 1d5702d2098..4d9d28fa730 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -20,6 +20,7 @@ import ( "github.com/gorilla/websocket" internalcache "github.com/router-for-me/CLIProxyAPI/v8/internal/cache" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -3347,3 +3348,208 @@ func TestCodexWebsockets_SendErrorLogsSessionObject(t *testing.T) { t.Fatalf("expected reason=send_error in log output, got: %s", logOutput) } } + +func TestCodexWebsocketReconnectUpgradeRejectionRecordedWhenRequestLogDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + + for _, stream := range []bool{false, true} { + name := "execute" + if stream { + name = "stream" + } + t.Run(name, func(t *testing.T) { + var dials atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if dials.Add(1) == 1 { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + return + } + defer func() { _ = conn.Close() }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "11") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error":{"message":"codex websocket reconnect forbidden"}}`)) + })) + defer server.Close() + + var writes atomic.Int32 + testWebsocketWritePayloadHook = func(conn *websocket.Conn) { + if writes.Add(1) == 1 { + _ = conn.SetWriteDeadline(time.Now().Add(-time.Second)) + } + } + defer func() { testWebsocketWritePayloadHook = nil }() + + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + exec := NewCodexWebsocketsExecutor(&config.Config{ + SDKConfig: config.SDKConfig{ + DisableImageGeneration: config.DisableImageGenerationAll, + }, + }) + auth := &cliproxyauth.Auth{ + ID: "codex-ws-reconnect", + Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}, + } + req := cliproxyexecutor.Request{ + Model: "gpt-5.6-sol", + Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"type":"message","role":"user","content":"reconnect"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "sess-codex-reconnect-" + name, + }, + } + + var errExec error + if stream { + _, errExec = exec.ExecuteStream(ctx, auth, req, opts) + } else { + _, errExec = exec.Execute(ctx, auth, req, opts) + } + if errExec == nil { + t.Fatal("expected reconnect upgrade rejection to fail") + } + + value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey) + if !exists { + t.Fatal("deferred API response was not recorded on codex websocket reconnect rejection") + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil { + t.Fatalf("deferred API response = %#v, want logging.DeferredAPIResponse", value) + } + rendered := string(deferred.RenderAPIResponse()) + for _, want := range []string{ + "=== API RESPONSE 1 ===", + "Status: 101", + "Error: send:", + "=== API RESPONSE 2 ===", + "Status: 403", + "Retry-After: 11", + `{"error":{"message":"codex websocket reconnect forbidden"}}`, + "Error: dial_retry:", + } { + if !strings.Contains(rendered, want) { + t.Fatalf("deferred API response missing %q:\n%s", want, rendered) + } + } + + var transportDials atomic.Int32 + transportServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if transportDials.Add(1) == 1 { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + return + } + defer func() { _ = conn.Close() }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + } + if hijacker, ok := w.(http.Hijacker); ok { + if rawConn, _, errHijack := hijacker.Hijack(); errHijack == nil { + _ = rawConn.Close() + } + } + })) + defer transportServer.Close() + + writes.Store(0) + subGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + subCtx := context.WithValue(context.Background(), "gin", subGinCtx) + subAuth := &cliproxyauth.Auth{ + ID: "codex-ws-transport-" + name, + Attributes: map[string]string{"api_key": "sk-test", "base_url": transportServer.URL}, + } + subOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "sess-codex-transport-" + name, + }, + } + if stream { + _, errExec = exec.ExecuteStream(subCtx, subAuth, req, subOpts) + } else { + _, errExec = exec.Execute(subCtx, subAuth, req, subOpts) + } + if errExec == nil { + t.Fatal("expected transport reconnect failure") + } + subVal, ok := subGinCtx.Get(logging.DeferredAPIResponseContextKey) + if !ok { + t.Fatal("deferred API response missing on transport reconnect failure") + } + subRendered := string(subVal.(logging.DeferredAPIResponse).RenderAPIResponse()) + for _, want := range []string{ + "=== API RESPONSE 1 ===", + "Error: send:", + "=== API RESPONSE 2 ===", + "Error: dial_retry:", + } { + if !strings.Contains(subRendered, want) { + t.Fatalf("transport reconnect response missing %q:\n%s", want, subRendered) + } + } + + initServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hijacker, ok := w.(http.Hijacker) + if !ok { + t.Error("hijacker unsupported") + return + } + rawConn, rawBuf, errHijack := hijacker.Hijack() + if errHijack != nil { + t.Errorf("hijack: %v", errHijack) + return + } + _, _ = rawBuf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: invalid\r\n\r\n") + _ = rawBuf.Flush() + _ = rawConn.Close() + })) + defer initServer.Close() + + initGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + initCtx := context.WithValue(context.Background(), "gin", initGinCtx) + initAuth := &cliproxyauth.Auth{ + ID: "codex-ws-init-" + name, + Attributes: map[string]string{"api_key": "sk-test", "base_url": initServer.URL}, + } + initOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + } + if stream { + _, errExec = exec.ExecuteStream(initCtx, initAuth, req, initOpts) + } else { + _, errExec = exec.Execute(initCtx, initAuth, req, initOpts) + } + if errExec == nil { + t.Fatal("expected initial invalid-101 handshake to fail") + } + initVal, ok := initGinCtx.Get(logging.DeferredAPIResponseContextKey) + if !ok { + t.Fatal("deferred API response missing on initial invalid-101 handshake") + } + initRendered := string(initVal.(logging.DeferredAPIResponse).RenderAPIResponse()) + for _, want := range []string{"Status: 101", "Error: dial:"} { + if !strings.Contains(initRendered, want) { + t.Fatalf("initial handshake response missing %q:\n%s", want, initRendered) + } + } + }) + } +} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index d533e83e9cd..e01dc786f66 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -30,6 +30,31 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } baseModel := thinking.ParseSuffix(req.Model).ModelName + isCompat, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel) + if errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } + if e.shouldCompactCompatResponsesTrigger(isCompat, req, opts) { + executionSessionID := executionSessionIDFromOptions(opts) + var sess *codexWebsocketSession + if executionSessionID != "" { + sess = e.getOrCreateSession(executionSessionID) + if sess != nil { + sess.reqMu.Lock() + defer sess.reqMu.Unlock() + } + } + res, errCompact := e.CodexExecutor.executeCodexResponsesCompactionStream(ctx, auth, req, opts) + if errCompact == nil && sess != nil { + sess.connMu.Lock() + existingConn := sess.conn + sess.connMu.Unlock() + if existingConn != nil { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, existingConn, "local_compaction", nil) + } + } + return res, errCompact + } reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) defer reporter.TrackFailure(ctx, &err) @@ -133,10 +158,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr unlockStreamSession() return nil, codexWebsocketTurnEventID(errDial, originalPayload) } - bodyErr := websocketHandshakeBody(respHS) - if respHS != nil { - helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) - } + bodyErr := recordWebsocketDialFailure(ctx, e.cfg, wsReqLog, respHS, errDial) if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { unlockStreamSession() if opts.ExecutionLifecycle == nil && !cliproxyexecutor.DownstreamWebsocket(ctx) { @@ -154,7 +176,6 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr unlockStreamSession() return nil, newCodexStatusErrWithHeadersAndCooling(respHS.StatusCode, respHS.Header, bodyErr, e.modelLevelCooling()) } - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) unlockStreamSession() return nil, errDial } @@ -220,8 +241,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr // Retry once with a new websocket connection for the same execution session. connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConnWithSIWC(ctx, auth, sess, authID, wsURL, wsHeaders, useSIWC) if errDialRetry != nil || connRetry == nil { - closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) + _ = recordWebsocketDialRetryFailure(ctx, e.cfg, wsReqLog, respHSRetry, errDialRetry) sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errDialRetry @@ -857,10 +877,14 @@ type codexWebsocketPrepared struct { } func (e *CodexWebsocketsExecutor) prepareCodexWebsocketStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*codexWebsocketPrepared, error) { - if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { - return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + if errPrev := rejectSyntheticCompactionPreviousResponseID(ctx, req.Payload); errPrev != nil { + return nil, errPrev } baseModel := thinking.ParseSuffix(req.Model).ModelName + isCompat, errExpand := e.expandCodexResponsesCompaction(auth, &req, &opts, baseModel) + if errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) if errSIWC := ensureCodexSIWCCredentialsForRequest(ctx, auth, apiKey, opts.Headers); errSIWC != nil { return nil, errSIWC @@ -883,7 +907,6 @@ func (e *CodexWebsocketsExecutor) prepareCodexWebsocketStream(ctx context.Contex originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - isCompat := e.resolveCodexModelIsCompat(auth, req, baseModel) targetModel := resolveResponsesSignatureTargetModel(baseModel, req, opts) originalTranslated, body, updatesChanged, err := translateCodexRequestPairWithUpdateIntentContext(ctx, from, to, targetModel, originalPayload, req.Payload, true, isCompat) if err != nil { diff --git a/internal/runtime/executor/devin_executor.go b/internal/runtime/executor/devin_executor.go index 0bc7587dae6..fd3e3dc8bf3 100644 --- a/internal/runtime/executor/devin_executor.go +++ b/internal/runtime/executor/devin_executor.go @@ -280,8 +280,11 @@ func (e *DevinExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - errData, _ := io.ReadAll(io.LimitReader(httpResp.Body, 1<<20)) + errData, errRead := io.ReadAll(io.LimitReader(httpResp.Body, 1<<20)) helps.AppendAPIResponseChunk(ctx, e.cfg, errData) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } return resp, newDevinStatusErrorForAuth(httpResp.StatusCode, httpResp.Header, errData, auth) } @@ -372,9 +375,12 @@ func (e *DevinExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - errData, _ := io.ReadAll(io.LimitReader(httpResp.Body, 1<<20)) + errData, errRead := io.ReadAll(io.LimitReader(httpResp.Body, 1<<20)) _ = httpResp.Body.Close() helps.AppendAPIResponseChunk(ctx, e.cfg, errData) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } return nil, newDevinStatusErrorForAuth(httpResp.StatusCode, httpResp.Header, errData, auth) } @@ -2691,15 +2697,23 @@ func (e *DevinExecutor) getSensitiveWords() []string { return nil } +func (e *DevinExecutor) getSensitiveWordsReplace() map[string]string { + if e != nil && e.cfg != nil && len(e.cfg.Devin.SensitiveWordsReplace) > 0 { + return e.cfg.Devin.SensitiveWordsReplace + } + return nil +} + func (e *DevinExecutor) getSensitiveWordMatcher() *helps.SensitiveWordMatcher { if e == nil { return nil } words := e.getSensitiveWords() - if len(words) == 0 { + replacements := e.getSensitiveWordsReplace() + if len(words) == 0 && len(replacements) == 0 { return nil } - key := strings.Join(words, "\x00") + key := helps.SensitiveWordConfigCacheKey(words, replacements) e.matcherMu.RLock() if e.lastWordsKey == key { m := e.cachedMatcher @@ -2714,7 +2728,7 @@ func (e *DevinExecutor) getSensitiveWordMatcher() *helps.SensitiveWordMatcher { return e.cachedMatcher } e.lastWordsKey = key - e.cachedMatcher = helps.BuildSensitiveWordMatcher(words) + e.cachedMatcher = helps.BuildSensitiveWordMatcherWithReplacements(words, replacements) return e.cachedMatcher } diff --git a/internal/runtime/executor/devin_executor_test.go b/internal/runtime/executor/devin_executor_test.go index bca8e7678e2..77c3bffab7b 100644 --- a/internal/runtime/executor/devin_executor_test.go +++ b/internal/runtime/executor/devin_executor_test.go @@ -142,6 +142,24 @@ func TestDevinExecutor_GetSensitiveWords(t *testing.T) { } } +func TestDevinSensitiveWordsReplaceFromYAML(t *testing.T) { + cfg, errParse := config.ParseConfigBytes([]byte(`devin: + sensitive-words: + - "Codex" + sensitive-words-replace: + "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant." +`)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + e := NewDevinExecutor(cfg) + got := helps.SanitizeDevinSystemPrompt("You are Codex, an agent based on GPT-5.\nFollow project conventions.", e.getSensitiveWordMatcher()) + want := "You are a helpful coding assistant.\nFollow project conventions." + if got != want { + t.Fatalf("SanitizeDevinSystemPrompt() = %q, want %q", got, want) + } +} + func TestParseInteractionsPayload(t *testing.T) { interactionsPayload := []byte(`{ "system_instruction": "You are a helpful coding assistant.", diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 9fdacdb8f0f..8c57616bc39 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -241,8 +241,11 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) @@ -251,11 +254,11 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, err } data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if action != "countTokens" { cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) } @@ -391,8 +394,11 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("gemini executor: close response body error: %v", errClose) @@ -700,11 +706,11 @@ func (e *GeminiExecutor) executeInteractions(ctx context.Context, auth *cliproxy }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = statusErr{code: httpResp.StatusCode, msg: string(data)} @@ -787,11 +793,14 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - data, _ := io.ReadAll(httpResp.Body) + data, errRead := io.ReadAll(httpResp.Body) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("gemini executor: close interactions error response body error: %v", errClose) } helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } return nil, statusErr{code: httpResp.StatusCode, msg: string(data)} } @@ -989,11 +998,11 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone()) data, err := io.ReadAll(resp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return cliproxyexecutor.Response{}, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if resp.StatusCode < 200 || resp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", resp.StatusCode, helps.SummarizeErrorBody(resp.Header.Get("Content-Type"), data)) return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(data)} diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index 43d096350eb..ff87112937d 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -542,8 +542,11 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) @@ -552,11 +555,11 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au return resp, err } data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if action != "countTokens" && !isImagenModel(baseModel) { cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) } @@ -713,8 +716,11 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) @@ -723,11 +729,11 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip return resp, err } data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if action != "countTokens" && !isImagenModel(baseModel) { cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) } @@ -864,8 +870,11 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) @@ -1149,8 +1158,11 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) @@ -1424,17 +1436,20 @@ func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)} } data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return cliproxyexecutor.Response{}, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) count := ojson.GetBytesNoCopy(data, "totalTokens").Int() out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data) return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil @@ -1530,17 +1545,20 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth * }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)} } data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return cliproxyexecutor.Response{}, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) count := ojson.GetBytesNoCopy(data, "totalTokens").Int() out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data) return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil @@ -1835,11 +1853,11 @@ func (e *GeminiVertexExecutor) executeInteractions(ctx context.Context, auth *cl }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = statusErr{code: httpResp.StatusCode, msg: string(data)} @@ -1960,11 +1978,14 @@ func (e *GeminiVertexExecutor) executeInteractionsStream(ctx context.Context, au } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - data, _ := io.ReadAll(httpResp.Body) + data, errRead := io.ReadAll(httpResp.Body) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close interactions error response body error: %v", errClose) } helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } return nil, statusErr{code: httpResp.StatusCode, msg: string(data)} } diff --git a/internal/runtime/executor/helps/antigravity_compaction.go b/internal/runtime/executor/helps/antigravity_compaction.go index eecdf8a6858..48da8e72f0f 100644 --- a/internal/runtime/executor/helps/antigravity_compaction.go +++ b/internal/runtime/executor/helps/antigravity_compaction.go @@ -20,9 +20,10 @@ import ( ) const ( - antigravityCompactionCapsulePrefix = "cpa-ag-compact-v1:" - antigravityCompactionItemIDPrefix = "cmp_ag_compact_" - fixedCompactionKeySecret = "CLIProxyAPI" + antigravityCompactionCapsulePrefix = "cpa-ag-compact-v1:" + antigravityCompactionItemIDPrefix = "cmp_ag_compact_" + antigravityCompactionResponseIDPrefix = "resp_ag_compact_" + fixedCompactionKeySecret = "CLIProxyAPI" ) type antigravityCompactionCapsuleData struct { @@ -36,6 +37,32 @@ func RecognizedAntigravityCompactionCapsule(encryptedContent string) bool { return strings.HasPrefix(strings.TrimSpace(encryptedContent), antigravityCompactionCapsulePrefix) } +// RecognizedResponsesCompactionResponseID reports whether responseID is a locally synthesized CPA compaction response ID. +func RecognizedResponsesCompactionResponseID(responseID string) bool { + return strings.HasPrefix(strings.TrimSpace(responseID), antigravityCompactionResponseIDPrefix) +} + +// HasResponsesCompactionConversationInput reports whether payload carries at least one +// non-control conversation item in input for local summary compaction. +func HasResponsesCompactionConversationInput(payload []byte) bool { + input := ojson.GetBytesNoCopy(payload, "input") + if input.Type == ojson.String { + return strings.TrimSpace(input.String()) != "" + } + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + switch strings.TrimSpace(item.Get("type").String()) { + case "compaction_trigger", "additional_tools", "configuration_update": + continue + default: + return true + } + } + return false +} + func isRecognizedCPACompactionItem(item ojson.Result) bool { if item.Get("type").String() != "compaction" { return false @@ -222,8 +249,8 @@ func PrepareAntigravityCompactionSummaryPayloadDoc(doc *ojson.Doc, modelName str if input.IsArray() { var filtered []string for _, item := range input.Array() { - itemType := item.Get("type").String() - if itemType == "compaction_trigger" || itemType == "additional_tools" { + switch strings.TrimSpace(item.Get("type").String()) { + case "compaction_trigger", "additional_tools", "configuration_update": continue } filtered = append(filtered, item.Raw) @@ -242,6 +269,8 @@ func PrepareAntigravityCompactionSummaryPayloadDoc(doc *ojson.Doc, modelName str // Remove fields not needed or incompatible with summary generation for _, field := range []string{ + "type", + "event_id", "stream", "tools", "tool_choice", @@ -250,6 +279,7 @@ func PrepareAntigravityCompactionSummaryPayloadDoc(doc *ojson.Doc, modelName str "additional_tools", "truncation", "metadata", + "text", } { _ = util.DeleteJSONDocIfExists(doc, field) } diff --git a/internal/runtime/executor/helps/antigravity_compaction_test.go b/internal/runtime/executor/helps/antigravity_compaction_test.go index dc0db241b8b..c30b8ca727f 100644 --- a/internal/runtime/executor/helps/antigravity_compaction_test.go +++ b/internal/runtime/executor/helps/antigravity_compaction_test.go @@ -69,8 +69,17 @@ func TestPrepareAntigravityCompactionSummaryPayload(t *testing.T) { "tools": [{"type": "function", "name": "bash"}], "tool_choice": "auto", "previous_response_id": "resp-old", + "text": { + "format": { + "type": "json_schema", + "name": "strict_bool", + "strict": true, + "schema": {"type": "object", "properties": {"ok": {"type": "boolean"}}, "required": ["ok"], "additionalProperties": false} + } + }, "input": [ {"type": "additional_tools", "tools": [{"type": "function", "name": "exec_command"}]}, + {"type": "configuration_update", "instructions": "Be brief"}, {"type": "message", "role": "user", "content": "hi"}, {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "hello"}]}, {"type": "compaction_trigger"} @@ -94,6 +103,9 @@ func TestPrepareAntigravityCompactionSummaryPayload(t *testing.T) { if ojson.GetBytes(prepared, "previous_response_id").Exists() { t.Fatal("expected previous_response_id to be removed") } + if ojson.GetBytes(prepared, "text").Exists() { + t.Fatal("expected text format constraints to be removed") + } input := ojson.GetBytes(prepared, "input").Array() if len(input) != 3 { @@ -107,6 +119,9 @@ func TestPrepareAntigravityCompactionSummaryPayload(t *testing.T) { if item.Get("type").String() == "additional_tools" { t.Fatal("additional_tools must be removed from summary payload") } + if item.Get("type").String() == "configuration_update" { + t.Fatal("configuration_update must be removed from summary payload") + } } lastItem := input[len(input)-1] diff --git a/internal/runtime/executor/helps/apply_patch_responses.go b/internal/runtime/executor/helps/apply_patch_responses.go index 0f074023c12..2a81c7b2c5b 100644 --- a/internal/runtime/executor/helps/apply_patch_responses.go +++ b/internal/runtime/executor/helps/apply_patch_responses.go @@ -161,6 +161,7 @@ type ApplyPatchResponsesState struct { upstream []byte eventLine []byte active, failed, closed bool + sawPatch bool transportDone bool } @@ -681,10 +682,22 @@ func (s *ApplyPatchResponsesState) Transform(event []byte) ([][]byte, error) { } var out [][]byte for _, e := range events { + if s.isCustomPatchEvent(e) { + s.sawPatch = true + } converted, errTransform := s.Bridge.Transform(e) + if !s.sawPatch && !strings.HasPrefix(ojson.GetBytesNoCopy(e, "type").String(), "response.custom_tool_call_input.") { + for _, c := range converted { + if strings.HasPrefix(ojson.GetBytesNoCopy(c, "type").String(), "response.custom_tool_call_input.") { + s.sawPatch = true + break + } + } + } out = append(out, converted...) if errTransform != nil { s.failed = true + s.sawPatch = true s.byDispatcherKey, s.records, s.upstream = nil, nil, nil return out, errTransform } @@ -696,6 +709,37 @@ func (s *ApplyPatchResponsesState) Transform(event []byte) ([][]byte, error) { return out, nil } +func (s *ApplyPatchResponsesState) isCustomPatchEvent(event []byte) bool { + if s == nil || !s.active { + return false + } + root := util.ParseGJSONBytesNoCopy(event) + if item := root.Get("item"); item.Exists() { + if item.Get("type").String() == "custom_tool_call" { + return false + } + d, ok := s.tools[util.QualifyResponsesNamespaceToolName(item.Get("namespace").String(), item.Get("name").String())] + return ok && applypatch.IsCustomTool(d.Tool) + } + if name := root.Get("name").String(); name != "" { + d, ok := s.tools[util.QualifyResponsesNamespaceToolName(root.Get("namespace").String(), name)] + return ok && applypatch.IsCustomTool(d.Tool) + } + return false +} + +// HasPatchActivity reports whether upstream emitted any custom apply_patch call +// or unfinished patch namespace dispatcher before closing the response. +func (s *ApplyPatchResponsesState) HasPatchActivity() bool { + if s == nil || !s.active { + return false + } + if s.failed || s.sawPatch || s.unfinishedDispatcher() { + return true + } + return s.Bridge != nil && s.Bridge.Finish() != nil +} + // Finish validates source response closure independently of completed tool input. // Common Finish remains argument-only because it also runs inside terminal conversion. func (s *ApplyPatchResponsesState) Finish() error { diff --git a/internal/runtime/executor/helps/claude_client_detection.go b/internal/runtime/executor/helps/claude_client_detection.go index aa1daa6115f..52506d58f10 100644 --- a/internal/runtime/executor/helps/claude_client_detection.go +++ b/internal/runtime/executor/helps/claude_client_detection.go @@ -174,7 +174,7 @@ func DetectClaudeCodeRequest(headers http.Header, payload []byte, countTokens bo entrypoint := userAgent.entrypoint subclient := claudeSubclientForEntrypoint(entrypoint, cfg) detection := ClaudeCodeRequestDetection{ - XAppCLI: headerValue(headers, "X-App") == "cli", + XAppCLI: isNativeClaudeXApp(headerValue(headers, "X-App")), UserAgent: validUserAgent && plausibleClaudeCodeVersion(userAgent.version, cfg), BetasPresent: headerContainsClaudeCodeBeta(headers), Entrypoint: entrypoint, @@ -197,6 +197,15 @@ func DetectClaudeCodeRequest(headers http.Header, payload []byte, countTokens bo return detection } +func isNativeClaudeXApp(value string) bool { + switch value { + case "cli", "cli-bg": + return true + default: + return false + } +} + func claudeCodeHelperBetaProfile(redactThinking bool, trailing ...string) string { betas := []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"} if redactThinking { diff --git a/internal/runtime/executor/helps/claude_client_detection_test.go b/internal/runtime/executor/helps/claude_client_detection_test.go index f4152396c98..1acbaf3d220 100644 --- a/internal/runtime/executor/helps/claude_client_detection_test.go +++ b/internal/runtime/executor/helps/claude_client_detection_test.go @@ -665,3 +665,44 @@ func TestDetectClaudeCodeRequestAcceptsConfiguredNativeEntrypoints(t *testing.T) t.Fatalf("configured entrypoint without strong signals = %#v, want unconfirmed", unconfirmed) } } + +func TestDetectClaudeCodeRequestAcceptsBackgroundSessionXApp(t *testing.T) { + payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID) + + t.Run("messages", func(t *testing.T) { + headers := confirmedClaudeCodeHeaders() + headers.Set("X-App", "cli-bg") + detection := DetectClaudeCodeRequest(headers, payload, false) + if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient || !detection.XAppCLI { + t.Fatalf("detection = %#v, want confirmed native CLI for X-App: cli-bg", detection) + } + }) + + t.Run("count_tokens", func(t *testing.T) { + headers := confirmedClaudeCodeHeaders() + headers.Set("X-App", "cli-bg") + detection := DetectClaudeCodeRequest(headers, []byte(`{"messages":[]}`), true) + if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient || !detection.XAppCLI { + t.Fatalf("detection = %#v, want confirmed count_tokens for X-App: cli-bg", detection) + } + }) + + t.Run("measured_helper", func(t *testing.T) { + headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true)) + headers.Set("X-App", "cli-bg") + detection := DetectClaudeCodeRequest(headers, measuredClaudeCodeMinimalHelperPayload(), false) + if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient || !detection.HelperProfile || !detection.XAppCLI { + t.Fatalf("detection = %#v, want confirmed measured helper for X-App: cli-bg", detection) + } + }) + + t.Run("rejects_unknown_x_app", func(t *testing.T) { + for _, unknownApp := range []string{"cli-other", "web", "sdk"} { + headers := confirmedClaudeCodeHeaders() + headers.Set("X-App", unknownApp) + if detection := DetectClaudeCodeRequest(headers, payload, false); detection.Confirmed || detection.XAppCLI { + t.Fatalf("X-App %q detection = %#v, want unconfirmed", unknownApp, detection) + } + } + }) +} diff --git a/internal/runtime/executor/helps/claude_tool_property_names.go b/internal/runtime/executor/helps/claude_tool_property_names.go new file mode 100644 index 00000000000..adb99d8e2cf --- /dev/null +++ b/internal/runtime/executor/helps/claude_tool_property_names.go @@ -0,0 +1,175 @@ +package helps + +import ( + "fmt" + "strings" + + "github.com/MoeForge/moego/ojson" + log "github.com/sirupsen/logrus" +) + +// Claude rejects a request whose tools declare an input_schema property key +// outside ^[a-zA-Z0-9_.-]{1,64}$, while OpenAI-compatible clients and MCP +// servers commonly use parameter names such as "not[labels]". Because clients +// resend their tools on every turn, one such tool fails the whole conversation. +// +// Claude presents tool schemas to the model as part of its prompt, so such a +// tool is sent with an open object input_schema and its original schema +// embedded verbatim in the description. No parameter is renamed or dropped: +// the model sees every original name and constraint and calls the tool with +// the original names, so neither historical tool_use inputs nor tool calls in +// responses need any rewriting. Claude no longer enforces the embedded schema +// structurally, so strict is removed as structured outputs cannot apply it. +const ( + claudePropertyKeyMaxLength = 64 + claudeToolSchemaMaxInspectDepth = 256 + claudeToolSchemaFallbackPrefix = "The input object must conform to the following JSON Schema. Use its property names exactly as written, including any characters such as brackets:\n" + claudeToolSchemaFallbackSchema = `{"type":"object"}` +) + +// claudePropertyKeyAccepted reports whether Claude accepts name as an +// input_schema property key. +func claudePropertyKeyAccepted(name string) bool { + if len(name) == 0 || len(name) > claudePropertyKeyMaxLength { + return false + } + for i := 0; i < len(name); i++ { + c := name[i] + if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' || c == '.' || c == '-' { + continue + } + return false + } + return true +} + +// claudeSchemaHasRejectedPropertyKey reports whether schema, or any subschema +// reachable through JSON Schema applicator and definition keywords, declares a +// properties key Claude rejects. Keyword names are only interpreted at schema +// positions, so a definition or property named like a keyword is handled by +// its role. Values of other keywords, such as const, default, enum, examples +// and unknown annotations, are data and are not inspected. Schemas nested +// beyond the inspection depth are reported as rejected so they take the +// fallback, which still carries the full schema text. +func claudeSchemaHasRejectedPropertyKey(schema ojson.Result, depth int) bool { + if !schema.IsObject() { + return false + } + if depth > claudeToolSchemaMaxInspectDepth { + return true + } + found := false + subschema := func(value ojson.Result) bool { + found = claudeSchemaHasRejectedPropertyKey(value, depth+1) + return !found + } + subschemas := func(value ojson.Result) { + value.ForEach(func(_, item ojson.Result) bool { return subschema(item) }) + } + schema.ForEach(func(key, value ojson.Result) bool { + switch key.String() { + case "properties": + if !value.IsObject() { + return true + } + value.ForEach(func(name, item ojson.Result) bool { + if !claudePropertyKeyAccepted(name.String()) { + found = true + return false + } + return subschema(item) + }) + case "$defs", "definitions", "patternProperties", "dependentSchemas", "dependencies": + // Maps of named subschemas; dependencies may also hold name arrays, + // which are not objects and are skipped by subschema. + if value.IsObject() { + subschemas(value) + } + case "allOf", "anyOf", "oneOf", "prefixItems": + if value.IsArray() { + subschemas(value) + } + case "items": + if value.IsArray() { + subschemas(value) + } else { + subschema(value) + } + case "additionalProperties", "unevaluatedProperties", "additionalItems", "unevaluatedItems", + "contains", "propertyNames", "not", "if", "then", "else", "contentSchema": + subschema(value) + } + return !found + }) + return found +} + +// claudeToolNeedsSchemaFallback reports whether a Claude client tool +// declaration has an input_schema that Claude would reject. +func claudeToolNeedsSchemaFallback(tool ojson.Result) bool { + if IsClaudeServerToolType(tool.Get("type").String()) { + return false + } + schema := tool.Get("input_schema") + return schema.IsObject() && claudeSchemaHasRejectedPropertyKey(schema, 0) +} + +// NormalizeClaudeToolPropertyNamesDoc applies the schema fallback to every +// client tool declared in tools[] or in an inline tool_addition definition +// whose input_schema has a property key Claude rejects. Tools Claude accepts +// are left unchanged. It must run before user payload rules. +func NormalizeClaudeToolPropertyNamesDoc(doc *ojson.Doc) { + if doc == nil || doc.Len() == 0 { + return + } + var paths []string + doc.Get("tools").ForEach(func(index, tool ojson.Result) bool { + if claudeToolNeedsSchemaFallback(tool) { + paths = append(paths, fmt.Sprintf("tools.%d", index.Int())) + } + return true + }) + doc.Get("messages").ForEach(func(messageIndex, message ojson.Result) bool { + message.Get("content").ForEach(func(partIndex, part ojson.Result) bool { + if part.Get("type").String() != "tool_addition" || part.Get("tool.type").String() != "tool_definition" { + return true + } + if claudeToolNeedsSchemaFallback(part.Get("tool.definition")) { + paths = append(paths, fmt.Sprintf("messages.%d.content.%d.tool.definition", messageIndex.Int(), partIndex.Int())) + } + return true + }) + return true + }) + for _, path := range paths { + applyClaudeToolSchemaFallbackDoc(doc, path) + } +} + +func applyClaudeToolSchemaFallbackDoc(doc *ojson.Doc, path string) { + tool := doc.Get(path) + name := strings.Clone(tool.Get("name").String()) + var description strings.Builder + if existing := tool.Get("description").String(); strings.TrimSpace(existing) != "" { + description.WriteString(existing) + description.WriteString("\n\n") + } + description.WriteString(claudeToolSchemaFallbackPrefix) + description.WriteString(tool.Get("input_schema").Raw) + hasStrict := tool.Get("strict").Exists() + + if errSet := doc.Set(path+".description", description.String()); errSet != nil { + log.Warnf("claude tool %q: failed to embed input_schema in description: %v", name, errSet) + return + } + if errSet := doc.SetRaw(path+".input_schema", claudeToolSchemaFallbackSchema); errSet != nil { + log.Warnf("claude tool %q: failed to replace input_schema: %v", name, errSet) + return + } + if hasStrict { + if errDelete := doc.Delete(path + ".strict"); errDelete != nil { + log.Warnf("claude tool %q: failed to remove strict: %v", name, errDelete) + } + } + log.Debugf("claude tool %q: input_schema has property names Claude rejects; embedded the schema in the description", name) +} diff --git a/internal/runtime/executor/helps/claude_tool_property_names_test.go b/internal/runtime/executor/helps/claude_tool_property_names_test.go new file mode 100644 index 00000000000..0bbaa11f816 --- /dev/null +++ b/internal/runtime/executor/helps/claude_tool_property_names_test.go @@ -0,0 +1,134 @@ +package helps + +import ( + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" +) + +func TestClaudePropertyKeyAccepted(t *testing.T) { + for name, want := range map[string]bool{ + "state": true, + "user.name": true, + "a-b_c": true, + strings.Repeat("x", 64): true, + strings.Repeat("x", 65): false, + "": false, + "not[labels]": false, + "$schema": false, + "a b": false, + "名": false, + } { + if got := claudePropertyKeyAccepted(name); got != want { + t.Errorf("claudePropertyKeyAccepted(%q) = %v, want %v", name, got, want) + } + } +} + +func TestNormalizeClaudeToolPropertyNamesDocEmbedsRejectedSchemas(t *testing.T) { + rejectedSchema := `{"type":"object","properties":{"not[labels]":{"type":"string"},"state":{"type":"string"}},"required":["not[labels]"]}` + nestedSchema := `{"type":"object","$defs":{"f":{"type":"object","properties":{"$schema":{"type":"string"}}}},"properties":{"filter":{"$ref":"#/$defs/f"}}}` + inlineSchema := `{"type":"object","properties":{"items":{"type":"array","items":{"type":"object","properties":{"id[]":{"type":"string"}}}}}}` + propertyNamedDefault := `{"type":"object","properties":{"default":{"type":"object","properties":{"a b":{"type":"string"}}}}}` + instanceOnly := `{"type":"object","properties":{"q":{"type":"object","default":{"properties":{"a b":1}},"examples":[{"properties":{"c d":2}}],"const":{"properties":{"e f":3}},"enum":[{"properties":{"g h":4}}]}}}` + validSchema := `{"type":"object","properties":{"user.name":{"type":"string"}}}` + body := `{"tools":[` + + `{"name":"list_merge_requests","description":"List merge requests.","strict":true,"cache_control":{"type":"ephemeral"},"input_schema":` + rejectedSchema + `,"input_examples":[{"not[labels]":"bug"}]},` + + `{"name":"nested","input_schema":` + nestedSchema + `},` + + `{"name":"property_named_default","description":"Default.","input_schema":` + propertyNamedDefault + `},` + + `{"name":"instance_only","description":"Data.","input_schema":` + instanceOnly + `},` + + `{"name":"valid","description":"Valid.","strict":true,"input_schema":` + validSchema + `},` + + `{"type":"web_search_20250305","name":"web_search","max_uses":1}],` + + `"messages":[` + + `{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"list_merge_requests","input":{"not[labels]":"bug"}}]},` + + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":"[]"}]},` + + `{"role":"system","content":[{"type":"tool_addition","tool":{"type":"tool_definition","definition":{"name":"inline_tool","description":"Inline.","input_schema":` + inlineSchema + `}}}]}]}` + doc := ojson.NewDocString(body) + NormalizeClaudeToolPropertyNamesDoc(doc) + out := doc.Bytes() + + fallbacks := map[string]string{ + "tools.0": "List merge requests.\n\n" + claudeToolSchemaFallbackPrefix + rejectedSchema, + "tools.1": claudeToolSchemaFallbackPrefix + nestedSchema, + "tools.2": "Default.\n\n" + claudeToolSchemaFallbackPrefix + propertyNamedDefault, + "messages.2.content.0.tool.definition": "Inline.\n\n" + claudeToolSchemaFallbackPrefix + inlineSchema, + } + for path, description := range fallbacks { + if got := ojson.GetBytes(out, path+".input_schema").Raw; got != claudeToolSchemaFallbackSchema { + t.Errorf("%s.input_schema = %s, want %s", path, got, claudeToolSchemaFallbackSchema) + } + if got := ojson.GetBytes(out, path+".description").String(); got != description { + t.Errorf("%s.description = %q, want %q", path, got, description) + } + if ojson.GetBytes(out, path+".strict").Exists() { + t.Errorf("%s.strict must be removed", path) + } + } + unchanged := map[string]string{ + "tools.0.cache_control": `{"type":"ephemeral"}`, + "tools.0.input_examples": `[{"not[labels]":"bug"}]`, + "tools.3": `{"name":"instance_only","description":"Data.","input_schema":` + instanceOnly + `}`, + "tools.4": `{"name":"valid","description":"Valid.","strict":true,"input_schema":` + validSchema + `}`, + "tools.5": `{"type":"web_search_20250305","name":"web_search","max_uses":1}`, + "messages.0.content.0.input": `{"not[labels]":"bug"}`, + } + for path, want := range unchanged { + if got := ojson.GetBytes(out, path).Raw; got != want { + t.Errorf("%s = %s, want %s", path, got, want) + } + } +} + +func TestNormalizeClaudeToolPropertyNamesDocLeavesAcceptedRequestsUnchanged(t *testing.T) { + body := `{"tools":[{"name":"t","description":"d","input_schema":{"type":"object","$defs":{"a[b]":{"type":"string"}},"patternProperties":{"^x\\[":{}},"properties":{"file.xml":{"$ref":"#/$defs/a[b]"}}}}],"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"t","input":{"x[1]":2}}]}]}` + doc := ojson.NewDocString(body) + NormalizeClaudeToolPropertyNamesDoc(doc) + if got := doc.String(); got != body { + t.Fatalf("body changed:\n got: %s\nwant: %s", got, body) + } +} + +func TestNormalizeClaudeToolPropertyNamesDocFallsBackForDeepSchemas(t *testing.T) { + depth := claudeToolSchemaMaxInspectDepth + 10 + schema := strings.Repeat(`{"not":`, depth) + `{}` + strings.Repeat(`}`, depth) + doc := ojson.NewDocString(`{"tools":[{"name":"deep","input_schema":` + schema + `}]}`) + NormalizeClaudeToolPropertyNamesDoc(doc) + if got := doc.Get("tools.0.input_schema").Raw; got != claudeToolSchemaFallbackSchema { + t.Fatalf("deep input_schema = %.80s, want the fallback", got) + } + if !strings.HasSuffix(doc.Get("tools.0.description").String(), schema) { + t.Fatal("deep schema must be embedded verbatim in the description") + } +} + +// Keyword names are only interpreted at schema positions; named definitions, +// dependent schemas, patterns and properties are classified by their role. +func TestClaudeSchemaHasRejectedPropertyKeyFollowsSchemaPositions(t *testing.T) { + rejected := `{"type":"object","properties":{"not[labels]":{"type":"string"}}}` + accepted := `{"type":"object","properties":{"value":{"type":"string"}}}` + cases := map[string]struct { + schema string + want bool + }{ + "defs named default": {schema: `{"$defs":{"default":` + rejected + `},"properties":{"f":{"$ref":"#/$defs/default"}}}`, want: true}, + "definitions named enum": {schema: `{"definitions":{"enum":` + rejected + `}}`, want: true}, + "dependent schema named const": {schema: `{"dependentSchemas":{"const":` + rejected + `}}`, want: true}, + "dependencies schema named const": {schema: `{"dependencies":{"const":` + rejected + `,"a":["b"]}}`, want: true}, + "pattern named examples": {schema: `{"patternProperties":{"examples":` + rejected + `}}`, want: true}, + "property named default": {schema: `{"properties":{"default":` + rejected + `}}`, want: true}, + "property named properties": {schema: `{"properties":{"properties":` + rejected + `}}`, want: true}, + "nested applicators": {schema: `{"anyOf":[{"items":[{"not":{"if":{"then":{"contains":` + rejected + `}}}}]}]}`, want: true}, + "tuple items and prefix items": {schema: `{"prefixItems":[true,{"additionalProperties":` + rejected + `}]}`, want: true}, + "defs named properties": {schema: `{"$defs":{"text":{"type":"string"},"properties":{"$ref":"#/$defs/text"}},"properties":{"value":{"$ref":"#/$defs/properties"}},"additionalProperties":false}`, want: false}, + "property named properties valid": {schema: `{"properties":{"properties":` + accepted + `}}`, want: false}, + "defs with rejected definition key": {schema: `{"$defs":{"a[b]":` + accepted + `}}`, want: false}, + "data keywords": {schema: `{"default":{"properties":{"a b":1}},"const":{"properties":{"c d":2}},"enum":[{"properties":{"e f":3}}],"examples":[{"properties":{"g h":4}}]}`, want: false}, + "unknown annotation": {schema: `{"x-ui":{"properties":{"a b":{}}}}`, want: false}, + } + for name, tc := range cases { + if got := claudeSchemaHasRejectedPropertyKey(ojson.Parse(tc.schema), 0); got != tc.want { + t.Errorf("%s: claudeSchemaHasRejectedPropertyKey = %v, want %v", name, got, tc.want) + } + } +} diff --git a/internal/runtime/executor/helps/cloak_obfuscate.go b/internal/runtime/executor/helps/cloak_obfuscate.go index 10c6b5388ae..d70cf14cd3e 100644 --- a/internal/runtime/executor/helps/cloak_obfuscate.go +++ b/internal/runtime/executor/helps/cloak_obfuscate.go @@ -3,7 +3,9 @@ package helps import ( "regexp" "sort" + "strconv" "strings" + "unicode" "unicode/utf8" "github.com/MoeForge/moego/ojson" @@ -12,49 +14,218 @@ import ( // zeroWidthSpace is the Unicode zero-width space character used for obfuscation. const zeroWidthSpace = "\u200B" -// SensitiveWordMatcher holds the compiled regex for matching sensitive words. +// combiningGraphemeJoiner is the Unicode combining grapheme joiner (U+034F) +// used for Antigravity system-instruction obfuscation where U+200B obfuscation +// remains rejected by the upstream endpoint. +const combiningGraphemeJoiner = "\u034F" + +type sensitiveWordRule struct { + target string + lowerTarget string + replacement string + isReplacement bool +} + +// SensitiveWordMatcher holds the compiled regex and replacement state for matching sensitive words. type SensitiveWordMatcher struct { - regex *regexp.Regexp + regex *regexp.Regexp + replaceFn func(string) string + lookupReplacement func(string) (string, bool) } -// BuildSensitiveWordMatcher compiles a regex from the word list. +// SensitiveWordConfigCacheKey builds a deterministic, collision-free cache key +// for a (words, replacements) configuration pair. +func SensitiveWordConfigCacheKey(words []string, replacements map[string]string) string { + var b strings.Builder + b.WriteString("w:") + b.WriteString(strconv.Itoa(len(words))) + b.WriteByte(';') + for _, w := range words { + b.WriteString(strconv.Itoa(len(w))) + b.WriteByte(':') + b.WriteString(w) + b.WriteByte(';') + } + b.WriteString("|r:") + b.WriteString(strconv.Itoa(len(replacements))) + b.WriteByte(';') + if len(replacements) > 0 { + keys := make([]string, 0, len(replacements)) + for k := range replacements { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + v := replacements[k] + b.WriteString(strconv.Itoa(len(k))) + b.WriteByte(':') + b.WriteString(k) + b.WriteByte('=') + b.WriteString(strconv.Itoa(len(v))) + b.WriteByte(':') + b.WriteString(v) + b.WriteByte(';') + } + } + return b.String() +} + +// BuildSensitiveWordMatcher compiles a regex from the word list using U+200B. // Words are sorted by length (longest first) for proper matching. func BuildSensitiveWordMatcher(words []string) *SensitiveWordMatcher { - if len(words) == 0 { + return BuildSensitiveWordMatcherWithReplacements(words, nil) +} + +// BuildSensitiveWordMatcherWithReplacements compiles a matcher combining literal +// replacements and U+200B sensitive-word obfuscation. +func BuildSensitiveWordMatcherWithReplacements(words []string, replacements map[string]string) *SensitiveWordMatcher { + return buildSensitiveWordMatcher(words, replacements, obfuscateWord, zeroWidthSpace) +} + +// BuildAntigravitySensitiveWordMatcher compiles a matcher for Antigravity system +// instructions combining literal replacements and per-subword U+034F obfuscation. +func BuildAntigravitySensitiveWordMatcher(words []string, replacements map[string]string) *SensitiveWordMatcher { + return buildSensitiveWordMatcher(words, replacements, obfuscateAntigravityPhrase, combiningGraphemeJoiner) +} + +func buildSensitiveWordMatcher( + words []string, + replacements map[string]string, + obfuscateFn func(string) string, + marker string, +) *SensitiveWordMatcher { + if len(words) == 0 && len(replacements) == 0 { return nil } - // Filter and normalize words - var validWords []string - for _, w := range words { - w = strings.TrimSpace(w) - if utf8.RuneCountInString(w) >= 2 && !strings.Contains(w, zeroWidthSpace) { - validWords = append(validWords, w) + var rules []sensitiveWordRule + var replacementRules []sensitiveWordRule + var exactReplacements map[string]string + var lowerReplacements map[string]string + seenLower := make(map[string]struct{}, len(words)+len(replacements)) + + if len(replacements) > 0 { + rawKeys := make([]string, 0, len(replacements)) + for k := range replacements { + rawKeys = append(rawKeys, k) } + sort.Strings(rawKeys) + + for _, rawKey := range rawKeys { + target := strings.TrimSpace(rawKey) + if target == "" { + continue + } + if exactReplacements == nil { + exactReplacements = make(map[string]string, len(rawKeys)) + lowerReplacements = make(map[string]string, len(rawKeys)) + } + if _, exists := exactReplacements[target]; !exists || rawKey == target { + exactReplacements[target] = replacements[rawKey] + } + lower := strings.ToLower(target) + if _, exists := lowerReplacements[lower]; !exists { + lowerReplacements[lower] = replacements[rawKey] + } + if _, exists := seenLower[lower]; exists { + continue + } + seenLower[lower] = struct{}{} + rule := sensitiveWordRule{ + target: target, + lowerTarget: lower, + replacement: replacements[rawKey], + isReplacement: true, + } + rules = append(rules, rule) + replacementRules = append(replacementRules, rule) + } + } + + for _, rawWord := range words { + w := strings.TrimSpace(rawWord) + if utf8.RuneCountInString(w) < 2 || strings.Contains(w, marker) { + continue + } + if obfuscateFn(w) == w { + continue + } + lower := strings.ToLower(w) + if _, exists := seenLower[lower]; exists { + continue + } + seenLower[lower] = struct{}{} + rules = append(rules, sensitiveWordRule{ + target: w, + lowerTarget: lower, + }) } - if len(validWords) == 0 { + if len(rules) == 0 { return nil } - // Sort by length (longest first) for proper matching - sort.Slice(validWords, func(i, j int) bool { - return len(validWords[i]) > len(validWords[j]) - }) + sortRulesLongestFirst(rules) + sortRulesLongestFirst(replacementRules) - // Escape and join - escaped := make([]string, len(validWords)) - for i, w := range validWords { - escaped[i] = regexp.QuoteMeta(w) + escaped := make([]string, len(rules)) + for i, r := range rules { + escaped[i] = regexp.QuoteMeta(r.target) } - pattern := "(?i)" + strings.Join(escaped, "|") - re, err := regexp.Compile(pattern) - if err != nil { + pattern := "(?i:" + strings.Join(escaped, "|") + ")" + re, errCompile := regexp.Compile(pattern) + if errCompile != nil { return nil } - return &SensitiveWordMatcher{regex: re} + var lookupReplacement func(string) (string, bool) + replaceFn := obfuscateFn + if len(exactReplacements) > 0 { + lookupReplacement = func(matched string) (string, bool) { + if repl, ok := exactReplacements[matched]; ok { + return repl, true + } + if repl, ok := lowerReplacements[strings.ToLower(matched)]; ok { + return repl, true + } + for _, r := range replacementRules { + if strings.EqualFold(matched, r.target) { + return r.replacement, true + } + } + return "", false + } + replaceFn = func(matched string) string { + if repl, ok := lookupReplacement(matched); ok { + return repl + } + return obfuscateFn(matched) + } + } + + return &SensitiveWordMatcher{ + regex: re, + replaceFn: replaceFn, + lookupReplacement: lookupReplacement, + } +} + +func sortRulesLongestFirst(rules []sensitiveWordRule) { + sort.SliceStable(rules, func(i, j int) bool { + ri := utf8.RuneCountInString(rules[i].target) + rj := utf8.RuneCountInString(rules[j].target) + if ri != rj { + return ri > rj + } + if len(rules[i].target) != len(rules[j].target) { + return len(rules[i].target) > len(rules[j].target) + } + if rules[i].isReplacement != rules[j].isReplacement { + return rules[i].isReplacement + } + return rules[i].lowerTarget < rules[j].lowerTarget + }) } // obfuscateWord inserts a zero-width space after the first grapheme. @@ -63,7 +234,6 @@ func obfuscateWord(word string) string { return word } - // Get first rune r, size := utf8.DecodeRuneInString(word) if r == utf8.RuneError || size >= len(word) { return word @@ -72,15 +242,81 @@ func obfuscateWord(word string) string { return string(r) + zeroWidthSpace + word[size:] } -// obfuscateText replaces all sensitive words in the text. +func isAntigravitySubwordRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '\u034F' || r == '\u200B' +} + +func obfuscateSingleSubwordCGJ(subword string) (string, bool) { + if strings.Contains(subword, combiningGraphemeJoiner) { + return subword, true + } + r, size := utf8.DecodeRuneInString(subword) + if r == utf8.RuneError || size >= len(subword) { + return subword, false + } + return string(r) + combiningGraphemeJoiner + subword[size:], true +} + +// obfuscateAntigravityPhrase inserts U+034F after the first rune of every +// alphanumeric subword (>= 2 runes) within phrase, preserving punctuation and whitespace. +func obfuscateAntigravityPhrase(phrase string) string { + if phrase == "" || strings.Contains(phrase, combiningGraphemeJoiner) { + return phrase + } + + var b strings.Builder + b.Grow(len(phrase) + len(combiningGraphemeJoiner)*2) + subwordStart := -1 + inserted := false + + for i, r := range phrase { + if isAntigravitySubwordRune(r) { + if subwordStart < 0 { + subwordStart = i + } + continue + } + if subwordStart >= 0 { + obf, did := obfuscateSingleSubwordCGJ(phrase[subwordStart:i]) + b.WriteString(obf) + if did { + inserted = true + } + subwordStart = -1 + } + b.WriteRune(r) + } + if subwordStart >= 0 { + obf, did := obfuscateSingleSubwordCGJ(phrase[subwordStart:]) + b.WriteString(obf) + if did { + inserted = true + } + } + + if inserted { + return b.String() + } + + r, size := utf8.DecodeRuneInString(phrase) + if r == utf8.RuneError || size >= len(phrase) { + return phrase + } + return string(r) + combiningGraphemeJoiner + phrase[size:] +} + +// obfuscateText replaces or obfuscates all sensitive words in the text. func (m *SensitiveWordMatcher) obfuscateText(text string) string { if m == nil || m.regex == nil { return text } + if m.replaceFn != nil { + return m.regex.ReplaceAllStringFunc(text, m.replaceFn) + } return m.regex.ReplaceAllStringFunc(text, obfuscateWord) } -// Matches reports whether the text contains any of the configured sensitive words. +// Matches reports whether the text contains any configured sensitive word or replacement target. func (m *SensitiveWordMatcher) Matches(text string) bool { if m == nil || m.regex == nil { return false @@ -88,7 +324,80 @@ func (m *SensitiveWordMatcher) Matches(text string) bool { return m.regex.MatchString(text) } -// ObfuscateText replaces all sensitive words in the text. +// SanitizeDevinText applies configured replacements and Devin line-level +// sensitive-word filtering in a single non-cascading pass: +// - matches handled by sensitive-words-replace emit their literal replacement +// verbatim without re-scanning or triggering line deletion; +// - lines containing an unreplaced sensitive-words match are dropped. +func (m *SensitiveWordMatcher) SanitizeDevinText(text string) string { + if m == nil || m.regex == nil || text == "" { + return text + } + indices := m.regex.FindAllStringIndex(text, -1) + if len(indices) == 0 { + return strings.TrimSpace(text) + } + + var keptLines []string + var currentLine strings.Builder + dropCurrentLine := false + + flushLine := func() { + if !dropCurrentLine { + keptLines = append(keptLines, currentLine.String()) + } + currentLine.Reset() + dropCurrentLine = false + } + + appendLiteral := func(seg string) { + for { + nl := strings.IndexByte(seg, '\n') + if nl < 0 { + currentLine.WriteString(seg) + return + } + currentLine.WriteString(seg[:nl]) + flushLine() + seg = seg[nl+1:] + } + } + + last := 0 + for _, pair := range indices { + start, end := pair[0], pair[1] + if start > last { + appendLiteral(text[last:start]) + } + matched := text[start:end] + if m.lookupReplacement != nil { + if repl, ok := m.lookupReplacement(matched); ok { + currentLine.WriteString(repl) + last = end + continue + } + } + dropCurrentLine = true + for { + nl := strings.IndexByte(matched, '\n') + if nl < 0 { + break + } + flushLine() + dropCurrentLine = true + matched = matched[nl+1:] + } + last = end + } + if last < len(text) { + appendLiteral(text[last:]) + } + flushLine() + + return strings.TrimSpace(strings.Join(keptLines, "\n")) +} + +// ObfuscateText replaces or obfuscates all sensitive words in the text. func (m *SensitiveWordMatcher) ObfuscateText(text string) string { return m.obfuscateText(text) } diff --git a/internal/runtime/executor/helps/cloak_obfuscate_test.go b/internal/runtime/executor/helps/cloak_obfuscate_test.go new file mode 100644 index 00000000000..39cd89b4812 --- /dev/null +++ b/internal/runtime/executor/helps/cloak_obfuscate_test.go @@ -0,0 +1,111 @@ +package helps + +import ( + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" +) + +func TestBuildSensitiveWordMatcherPreservesZeroWidthSpaceDefault(t *testing.T) { + matcher := BuildSensitiveWordMatcher([]string{"proxy", "Claude Code"}) + if matcher == nil { + t.Fatal("expected non-nil matcher") + } + + got := matcher.ObfuscateText("Use proxy with Claude Code.") + want := "Use p\u200Broxy with C\u200Blaude Code." + if got != want { + t.Fatalf("ObfuscateText() = %q, want %q", got, want) + } +} + +func TestBuildAntigravitySensitiveWordMatcherSubwordsAndPunctuation(t *testing.T) { + matcher := BuildAntigravitySensitiveWordMatcher([]string{ + "Codex", + "GPT-5", + "Claude Code", + "", + "claude-code", + "Agent X", + }, nil) + if matcher == nil { + t.Fatal("expected non-nil matcher") + } + + input := "You are Codex, an agent based on GPT-5. Follow in Claude Code / claude-code with Agent X." + got := matcher.ObfuscateText(input) + want := "You are C\u034Fodex, an agent based on G\u034FPT-5. Follow in C\u034Flaude C\u034Fode / c\u034Flaude-c\u034Fode with A\u034Fgent X." + if got != want { + t.Fatalf("ObfuscateText() = %q, want %q", got, want) + } + if strings.Contains(got, "\u200B") { + t.Fatalf("ObfuscateText() = %q, must not contain U+200B", got) + } + if secondPass := matcher.ObfuscateText(got); secondPass != got { + t.Fatalf("second obfuscation pass = %q, want idempotent %q", secondPass, got) + } +} + +func TestSensitiveWordMatcherReplacementsPrecedenceAndEmptyDeletion(t *testing.T) { + matcher := BuildAntigravitySensitiveWordMatcher( + []string{"Claude", "Codex", "GPT-5"}, + map[string]string{ + "You are Codex, an agent based on GPT-5.": "You are a helpful coding assistant.", + "Claude Code": "Claude", + "Codex": "Assistant", + "RemoveMe ": "", + }, + ) + if matcher == nil { + t.Fatal("expected non-nil matcher") + } + + input := "you are codex, an agent based on gpt-5. Built for Claude Code, standalone Claude, Codex, and RemoveMe ." + got := matcher.ObfuscateText(input) + want := "You are a helpful coding assistant. Built for Claude, standalone C\u034Flaude, Assistant, and ." + if got != want { + t.Fatalf("ObfuscateText() = %q, want %q", got, want) + } +} + +func TestObfuscateSensitiveWordsDocReplacementsOnlyAndBillingHeaderSkip(t *testing.T) { + matcher := BuildSensitiveWordMatcherWithReplacements(nil, map[string]string{ + "anthropic": "provider", + "secret": "", + }) + if matcher == nil { + t.Fatal("expected non-nil matcher when only replacements are configured") + } + + doc := ojson.NewDocString(`{ + "system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.321; cc_entrypoint=cli;"}, + {"type": "text", "text": "Keep anthropic rules and drop secret"} + ], + "messages": [ + {"role": "user", "content": "Ask anthropic about secret"} + ] + }`) + if !ObfuscateSensitiveWordsDoc(doc, matcher) { + t.Fatal("expected ObfuscateSensitiveWordsDoc to modify document") + } + if gotBilling := doc.Get("system.0.text").String(); gotBilling != "x-anthropic-billing-header: cc_version=2.1.321; cc_entrypoint=cli;" { + t.Fatalf("billing header = %q, want untouched", gotBilling) + } + if gotSys := doc.Get("system.1.text").String(); gotSys != "Keep provider rules and drop " { + t.Fatalf("system.1.text = %q, want %q", gotSys, "Keep provider rules and drop ") + } + if gotMsg := doc.Get("messages.0.content").String(); gotMsg != "Ask provider about " { + t.Fatalf("messages.0.content = %q, want %q", gotMsg, "Ask provider about ") + } +} + +func TestSensitiveWordConfigCacheKeyCollisionFree(t *testing.T) { + k1 := SensitiveWordConfigCacheKey([]string{"a", "b"}, map[string]string{"x": "1=2"}) + k2 := SensitiveWordConfigCacheKey([]string{"a;1:b"}, map[string]string{"x": "1=2"}) + k3 := SensitiveWordConfigCacheKey([]string{"a", "b"}, map[string]string{"x=2:1": "2"}) + if k1 == k2 || k1 == k3 || k2 == k3 { + t.Fatalf("expected distinct cache keys, got %q, %q, %q", k1, k2, k3) + } +} diff --git a/internal/runtime/executor/helps/devin_wire.go b/internal/runtime/executor/helps/devin_wire.go index ecebd412d22..a85a67e2223 100644 --- a/internal/runtime/executor/helps/devin_wire.go +++ b/internal/runtime/executor/helps/devin_wire.go @@ -617,7 +617,7 @@ func ParseDevinFrame(payload []byte) (DevinFrameResult, error) { } // SanitizeDevinSystemPrompt strips Claude Code attribution and CLI identity headers -// (referencing Antigravity conventions), and applies zero-width obfuscation for configured sensitive words. +// (referencing Antigravity conventions), and applies configured sensitive-word replacements and obfuscation. func SanitizeDevinSystemPrompt(prompt string, matcher *SensitiveWordMatcher) string { if prompt == "" { return "" @@ -648,16 +648,13 @@ func SanitizeDevinSystemPrompt(prompt string, matcher *SensitiveWordMatcher) str if strings.Contains(trimmed, "- Don’t output ANSI escape codes directly — the CLI renderer applies them.") { continue } - if matcher != nil && matcher.Matches(trimmed) { - continue - } kept = append(kept, line) } - res := strings.TrimSpace(strings.Join(kept, "\n")) + res := strings.Join(kept, "\n") if matcher != nil && res != "" { - res = matcher.ObfuscateText(res) + return matcher.SanitizeDevinText(res) } - return res + return strings.TrimSpace(res) } func parseDevinToolCallDelta(data []byte) (DevinToolCallDelta, error) { diff --git a/internal/runtime/executor/helps/devin_wire_test.go b/internal/runtime/executor/helps/devin_wire_test.go index 176f884bc34..fe0c647e03b 100644 --- a/internal/runtime/executor/helps/devin_wire_test.go +++ b/internal/runtime/executor/helps/devin_wire_test.go @@ -159,6 +159,39 @@ func TestSanitizeDevinSystemPrompt_AndSensitiveWords(t *testing.T) { } } +func TestSanitizeDevinSystemPrompt_ReplacementsAndMultiline(t *testing.T) { + matcher := BuildSensitiveWordMatcherWithReplacements( + []string{"Claude", "DROP_LINE_TOKEN"}, + map[string]string{ + "Claude Code": "Claude", + "First Header\nSecond Line": "Merged Header", + }, + ) + rawPrompt := "First Header\r\nSecond Line\r\nUse Claude Code safely.\r\nLine with DROP_LINE_TOKEN should be dropped." + sanitized := SanitizeDevinSystemPrompt(rawPrompt, matcher) + want := "Merged Header\nUse Claude safely." + if sanitized != want { + t.Fatalf("SanitizeDevinSystemPrompt() = %q, want %q", sanitized, want) + } +} + +func TestSanitizeDevinSystemPrompt_MultilineReplacementSinglePassNoDropOrCascade(t *testing.T) { + matcher := BuildSensitiveWordMatcherWithReplacements( + []string{"Claude"}, + map[string]string{ + "First\nSecond": "Claude", + "Multi\nLine": "Alias", + "Alias": "Final", + }, + ) + rawPrompt := "First\r\nSecond\r\nMulti\r\nLine" + sanitized := SanitizeDevinSystemPrompt(rawPrompt, matcher) + want := "Claude\nAlias" + if sanitized != want { + t.Fatalf("SanitizeDevinSystemPrompt() = %q, want %q", sanitized, want) + } +} + func TestBuildDevinGetChatMessageRequest_WithImages(t *testing.T) { prompts := []DevinPrompt{ { diff --git a/internal/runtime/executor/helps/logging_helpers.go b/internal/runtime/executor/helps/logging_helpers.go index 298d4276bdb..07388988b29 100644 --- a/internal/runtime/executor/helps/logging_helpers.go +++ b/internal/runtime/executor/helps/logging_helpers.go @@ -20,13 +20,15 @@ import ( ) const ( - apiAttemptsKey = "API_UPSTREAM_ATTEMPTS" - apiRequestKey = "API_REQUEST" - apiResponseKey = "API_RESPONSE" - apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE" - deferredAPIRequestBytesKey = "DEFERRED_API_REQUEST_BYTES" - creditsUsedKey = "__antigravity_credits_used__" - maxDeferredAPIRequestBodyBytes = 32 << 20 // 32 MiB + apiAttemptsKey = "API_UPSTREAM_ATTEMPTS" + apiRequestKey = "API_REQUEST" + apiResponseKey = "API_RESPONSE" + apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE" + deferredAPIRequestBytesKey = "DEFERRED_API_REQUEST_BYTES" + creditsUsedKey = "__antigravity_credits_used__" + maxDeferredAPIRequestBodyBytes = logging.MaxDeferredAPIRequestBodyBytes + maxDeferredAPIResponseBodyBytes = logging.MaxDeferredAPIResponseBodyBytes + maxTotalDeferredAPIResponseBodyBytes = logging.MaxTotalDeferredAPIResponseBodyBytes ) // UpstreamRequestLog captures the outbound upstream request details for logging. @@ -66,18 +68,26 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ if cfg == nil || cfg.CommercialMode { return } - // Check the immutable policy before accessing the pooled request context. - if !cfg.RequestLog && logging.DeferredAPIRequestCaptureDisabled(ctx) { + // Check the immutable policy and request-scoped capture before accessing the pooled Gin context. + if !cfg.RequestLog { + if logging.DeferredAPIRequestCaptureDisabled(ctx) { + return + } + if capture := logging.DeferredAPICaptureFromContext(ctx); capture != nil { + deferAPIRequestToCapture(capture, info) + return + } + ginCtx := ginContextFrom(ctx) + if ginCtx == nil { + return + } + deferAPIRequest(ginCtx, info) return } ginCtx := ginContextFrom(ctx) if ginCtx == nil { return } - if !cfg.RequestLog { - deferAPIRequest(ginCtx, info) - return - } attempts := getAttempts(ginCtx) index := len(attempts) + 1 @@ -129,47 +139,135 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ } } -func deferAPIRequest(ginCtx *gin.Context, info UpstreamRequestLog) { - if ginCtx == nil { - return - } - var requests []logging.DeferredAPIRequest - if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists { - requests, _ = value.([]logging.DeferredAPIRequest) - } - index := len(requests) + 1 - capturedInfo := info - capturedAt := time.Now() - capturedBytes, _ := ginCtx.Get(deferredAPIRequestBytesKey) - bytesUsed, _ := capturedBytes.(int) - remaining := maxDeferredAPIRequestBodyBytes - bytesUsed - if remaining < 0 { - remaining = 0 - } - captureLength := len(info.Body) - if captureLength > remaining { - captureLength = remaining - } - capturedInfo.Body = bytes.Clone(info.Body[:captureLength]) - bodyEmpty := len(info.Body) == 0 - bodyTruncated := captureLength < len(info.Body) - ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed+captureLength) - requests = append(requests, func() []byte { - builder := newAPIRequestLogBuilder(index, capturedInfo, capturedAt) +func buildDeferredAPIRequestCallback(capturedInfo UpstreamRequestLog, index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) logging.DeferredAPIRequest { + reqInfo := capturedInfo + reqInfo.Body = capturedBody + return func() []byte { + builder := newAPIRequestLogBuilder(index, reqInfo, capturedAt) if bodyEmpty { builder.WriteString("") } else { - builder.Write(capturedInfo.Body) + builder.Write(reqInfo.Body) if bodyTruncated { - builder.WriteString(fmt.Sprintf("\n[API REQUEST BODY TRUNCATED: captured first %d bytes]", captureLength)) + builder.WriteString(fmt.Sprintf("\n[API REQUEST BODY TRUNCATED: captured first %d bytes]", len(reqInfo.Body))) } } builder.WriteString("\n\n") return []byte(builder.String()) + } +} + +func deferAPIRequestToCapture(capture *logging.DeferredAPICapture, info UpstreamRequestLog) (logging.DeferredAPIRequest, int, bool) { + capturedInfo := info + return capture.RecordRequest(info.Body, func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) logging.DeferredAPIRequest { + return buildDeferredAPIRequestCallback(capturedInfo, index, capturedAt, capturedBody, bodyEmpty, bodyTruncated) + }) +} + +func deferOrReplacePendingAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) { + if !deferredResponseCaptureEnabled(ctx, cfg) { + return + } + capturedInfo := info + if capture := logging.DeferredAPICaptureFromContext(ctx); capture != nil { + capture.RecordOrReplacePendingRequest(info.Body, func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) logging.DeferredAPIRequest { + return buildDeferredAPIRequestCallback(capturedInfo, index, capturedAt, capturedBody, bodyEmpty, bodyTruncated) + }) + return + } + ginCtx := ginContextFrom(ctx) + if ginCtx == nil { + return + } + capture := deferredResponseCaptureFrom(ginCtx) + fn, bytesUsed, replaced, ok := capture.RecordOrReplacePendingRequest(info.Body, func(index int, capturedAt time.Time, capturedBody []byte, bodyEmpty, bodyTruncated bool) logging.DeferredAPIRequest { + return buildDeferredAPIRequestCallback(capturedInfo, index, capturedAt, capturedBody, bodyEmpty, bodyTruncated) }) + if !ok { + return + } + var requests []logging.DeferredAPIRequest + if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists { + requests, _ = value.([]logging.DeferredAPIRequest) + } + if replaced && len(requests) > 0 { + requests[len(requests)-1] = fn + } else { + requests = append(requests, fn) + } + ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed) + ginCtx.Set(logging.DeferredAPIRequestContextKey, requests) +} + +func deferAPIRequest(ginCtx *gin.Context, info UpstreamRequestLog) { + if ginCtx == nil { + return + } + capture := deferredResponseCaptureFrom(ginCtx) + fn, bytesUsed, ok := deferAPIRequestToCapture(capture, info) + if !ok { + return + } + var requests []logging.DeferredAPIRequest + if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists { + requests, _ = value.([]logging.DeferredAPIRequest) + } + requests = append(requests, fn) + ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed) ginCtx.Set(logging.DeferredAPIRequestContextKey, requests) } +func deferredResponseCaptureEnabled(ctx context.Context, cfg *config.Config) bool { + if cfg == nil || cfg.CommercialMode || cfg.RequestLog { + return false + } + return !logging.DeferredAPIRequestCaptureDisabled(ctx) +} + +func resolveDeferredAPICapture(ctx context.Context, cfg *config.Config) *logging.DeferredAPICapture { + // Check the immutable policy and request-scoped capture before accessing the pooled Gin context. + if !deferredResponseCaptureEnabled(ctx, cfg) { + return nil + } + if capture := logging.DeferredAPICaptureFromContext(ctx); capture != nil { + return capture + } + ginCtx := ginContextFrom(ctx) + if ginCtx == nil { + return nil + } + return deferredResponseCaptureFrom(ginCtx) +} + +func deferredResponseCaptureFrom(ginCtx *gin.Context) *logging.DeferredAPICapture { + if value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey); exists { + if capture, ok := value.(*logging.DeferredAPICapture); ok && capture != nil { + return capture + } + } + capture := logging.NewDeferredAPICapture() + ginCtx.Set(logging.DeferredAPIResponseContextKey, capture) + return capture +} + +func recordDeferredResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) { + if capture := resolveDeferredAPICapture(ctx, cfg); capture != nil { + capture.RecordResponseMetadata(status, headers) + } +} + +func recordDeferredResponseError(ctx context.Context, cfg *config.Config, err error) { + if capture := resolveDeferredAPICapture(ctx, cfg); capture != nil { + capture.RecordResponseError(err) + } +} + +func appendDeferredResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) { + if capture := resolveDeferredAPICapture(ctx, cfg); capture != nil { + capture.AppendResponseChunk(chunk) + } +} + func newAPIRequestLogBuilder(index int, info UpstreamRequestLog, timestamp time.Time) *strings.Builder { builder := &strings.Builder{} builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index)) @@ -195,6 +293,7 @@ func newAPIRequestLogBuilder(index int, info UpstreamRequestLog, timestamp time. func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) { logging.SetResponseHeaders(ctx, headers) if !requestLogCaptureEnabled(cfg) { + recordDeferredResponseMetadata(ctx, cfg, status, headers) return } ginCtx := ginContextFrom(ctx) @@ -222,7 +321,11 @@ func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status i // RecordAPIResponseError adds an error entry for the latest attempt when no HTTP response is available. func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error) { - if !requestLogCaptureEnabled(cfg) || err == nil { + if err == nil { + return + } + if !requestLogCaptureEnabled(cfg) { + recordDeferredResponseError(ctx, cfg, err) return } ginCtx := ginContextFrom(ctx) @@ -248,6 +351,7 @@ func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error) // AppendAPIResponseChunk appends an upstream response chunk to Gin context for request logging. func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) { if !requestLogCaptureEnabled(cfg) { + appendDeferredResponseChunk(ctx, cfg, chunk) return } data := bytes.TrimSpace(chunk) @@ -291,7 +395,11 @@ func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt // RecordAPIWebsocketRequest stores an upstream websocket request event in Gin context. func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) { - if !requestLogCaptureEnabled(cfg) { + if cfg == nil || cfg.CommercialMode { + return + } + if !cfg.RequestLog { + RecordAPIRequest(ctx, cfg, info) return } ginCtx := ginContextFrom(ctx) @@ -325,6 +433,7 @@ func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info Ups func RecordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, status int, headers http.Header) { logging.SetResponseHeaders(ctx, headers) if !requestLogCaptureEnabled(cfg) { + recordDeferredResponseMetadata(ctx, cfg, status, headers) return } ginCtx := ginContextFrom(ctx) @@ -348,7 +457,17 @@ func RecordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, status // RecordAPIWebsocketUpgradeRejection stores a rejected websocket upgrade as an HTTP attempt. func RecordAPIWebsocketUpgradeRejection(ctx context.Context, cfg *config.Config, info UpstreamRequestLog, status int, headers http.Header, body []byte) { logging.SetResponseHeaders(ctx, headers) - if !requestLogCaptureEnabled(cfg) { + if cfg == nil || cfg.CommercialMode { + return + } + // Check the immutable policy and request-scoped capture before accessing the pooled Gin context. + if !cfg.RequestLog { + if logging.DeferredAPIRequestCaptureDisabled(ctx) { + return + } + deferOrReplacePendingAPIRequest(ctx, cfg, info) + RecordAPIResponseMetadata(ctx, cfg, status, headers) + AppendAPIResponseChunk(ctx, cfg, body) return } ginCtx := ginContextFrom(ctx) @@ -383,6 +502,7 @@ func WebsocketUpgradeRequestURL(rawURL string) string { // AppendAPIWebsocketResponse stores an upstream websocket response frame in Gin context. func AppendAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) { if !requestLogCaptureEnabled(cfg) { + appendDeferredResponseChunk(ctx, cfg, payload) return } data := bytes.TrimSpace(payload) @@ -413,7 +533,15 @@ func AppendCodexAPIWebsocketResponse(ctx context.Context, cfg *config.Config, pa // RecordAPIWebsocketError stores an upstream websocket error event in Gin context. func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) { - if !requestLogCaptureEnabled(cfg) || err == nil { + if err == nil { + return + } + if !requestLogCaptureEnabled(cfg) { + stageErr := err + if trimmed := strings.TrimSpace(stage); trimmed != "" { + stageErr = fmt.Errorf("%s: %w", trimmed, err) + } + recordDeferredResponseError(ctx, cfg, stageErr) return } ginCtx := ginContextFrom(ctx) diff --git a/internal/runtime/executor/helps/logging_helpers_test.go b/internal/runtime/executor/helps/logging_helpers_test.go index 607ecea9cd0..78506199ae9 100644 --- a/internal/runtime/executor/helps/logging_helpers_test.go +++ b/internal/runtime/executor/helps/logging_helpers_test.go @@ -1,10 +1,13 @@ package helps import ( + "bytes" "context" "errors" + "fmt" "net/http" "net/http/httptest" + "runtime" "strings" "sync" "testing" @@ -171,6 +174,281 @@ func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing } } +func TestDeferredAPIResponseCapturedAndAlignedAcrossAttempts(t *testing.T) { + gin.SetMode(gin.TestMode) + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{} + + // Attempt 1: request recorded, then aborted before any upstream response arrives. + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{ + URL: "https://api.example.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"attempt":1}`), + }) + capture := deferredResponseCaptureFrom(ginCtx) + if capture.AttemptCount() != 1 || !capture.AttemptCapturedAt(0).IsZero() { + t.Fatalf("attempt 1 capturedAt = %v, want zero before response arrives", capture.AttemptCapturedAt(0)) + } + if capture.HasPayload() { + t.Fatal("HasPayload = true before any response details were recorded") + } + + // Attempt 2: response metadata, SSE chunks, and read error recorded. + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{ + URL: "https://api.example.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"attempt":2}`), + }) + if !capture.AttemptCapturedAt(1).IsZero() { + t.Fatalf("attempt 2 capturedAt = %v, want zero before response event", capture.AttemptCapturedAt(1)) + } + headers := http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Authorization": []string{"Bearer super-secret-upstream-token"}, + "Set-Cookie": []string{"__cf_bm=super-secret-cf-cookie"}, + } + RecordAPIResponseMetadata(ctx, cfg, http.StatusBadGateway, headers) + headers.Set("Content-Type", "text/mutated") + if capture.AttemptCapturedAt(1).IsZero() { + t.Fatal("attempt 2 capturedAt was not stamped on first response event") + } + AppendAPIResponseChunk(ctx, cfg, []byte("event: response.created")) + AppendAPIResponseChunk(ctx, cfg, []byte(`data: {"type":"response.created"}`)) + RecordAPIResponseError(ctx, cfg, errors.New("stream disconnected")) + + // Attempt 3: transport dial error with no HTTP response metadata. + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{ + URL: "https://api.example.com/v1/responses", + Method: http.MethodPost, + Body: []byte(`{"attempt":3}`), + }) + RecordAPIResponseError(ctx, cfg, errors.New("dial tcp timeout")) + + value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey) + if !exists { + t.Fatal("deferred API response was not stored in Gin context") + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil || !deferred.HasPayload() { + t.Fatalf("deferred API response = %#v, want populated DeferredAPIResponse", value) + } + + rendered := string(deferred.RenderAPIResponse()) + if strings.Contains(rendered, "=== API RESPONSE 1 ===") { + t.Fatalf("attempt 1 had no response content and should not render:\n%s", rendered) + } + for _, want := range []string{ + "=== API RESPONSE 2 ===", + "Status: 502", + "Authorization: Bearer supe...oken", + "Content-Type: text/event-stream", + "event: response.created\ndata: {\"type\":\"response.created\"}", + "Error: stream disconnected\n\n=== API RESPONSE 3 ===", + "Error: dial tcp timeout", + } { + if !strings.Contains(rendered, want) { + t.Fatalf("rendered deferred response missing %q:\n%s", want, rendered) + } + } + for _, secret := range []string{"super-secret-upstream-token", "super-secret-cf-cookie"} { + if strings.Contains(rendered, secret) { + t.Fatalf("rendered deferred response leaked secret %q:\n%s", secret, rendered) + } + } + attempt3Section := rendered[strings.Index(rendered, "=== API RESPONSE 3 ==="):] + if strings.Contains(attempt3Section, "Headers:") { + t.Fatalf("dial error without HTTP response should not fabricate Headers section:\n%s", attempt3Section) + } +} + +func TestDeferredAPIResponseBodyTruncatedAndBoundedAcrossAttempts(t *testing.T) { + gin.SetMode(gin.TestMode) + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{} + + // Pass a 2 MiB slice to verify truncation copies into a bounded buffer rather + // than retaining the multi-megabyte backing array. + largePayload := bytes.Repeat([]byte("a"), 2<<20) + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{Method: http.MethodPost}) + RecordAPIResponseMetadata(ctx, cfg, http.StatusOK, nil) + AppendAPIResponseChunk(ctx, cfg, largePayload) + + capture := deferredResponseCaptureFrom(ginCtx) + attempt1Len, attempt1Cap, attempt1Truncated := capture.AttemptBodyStats(0) + if attempt1Len != maxDeferredAPIResponseBodyBytes { + t.Fatalf("attempt 1 body len = %d, want %d", attempt1Len, maxDeferredAPIResponseBodyBytes) + } + if attempt1Cap > maxDeferredAPIResponseBodyBytes*2 { + t.Fatalf("attempt 1 body cap = %d, retained oversized backing array", attempt1Cap) + } + if !attempt1Truncated { + t.Fatal("attempt 1 bodyTruncated = false, want true") + } + + maxAttempts := maxTotalDeferredAPIResponseBodyBytes/maxDeferredAPIResponseBodyBytes + 2 + for i := 1; i < maxAttempts; i++ { + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{Method: http.MethodPost}) + RecordAPIResponseMetadata(ctx, cfg, http.StatusOK, nil) + AppendAPIResponseChunk(ctx, cfg, largePayload) + } + if capture.TotalResponseBodyBytes() > maxTotalDeferredAPIResponseBodyBytes { + t.Fatalf("total deferred response bytes = %d, want <= %d", capture.TotalResponseBodyBytes(), maxTotalDeferredAPIResponseBodyBytes) + } + lastLen, _, lastTruncated := capture.AttemptBodyStats(capture.AttemptCount() - 1) + if lastLen != 0 || !lastTruncated { + t.Fatalf("exhausted budget attempt len=%d truncated=%t, want len=0 truncated=true", lastLen, lastTruncated) + } + + rendered := string(capture.RenderAPIResponse()) + wantMarker := fmt.Sprintf("[API RESPONSE BODY TRUNCATED: captured first %d bytes]", maxDeferredAPIResponseBodyBytes) + if !strings.Contains(rendered, wantMarker) { + t.Fatalf("rendered response missing truncation marker %q", wantMarker) + } +} + +func TestDeferredAPIResponseBoundsLargeErrorsHeadersAndRetries(t *testing.T) { + gin.SetMode(gin.TestMode) + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{} + + giantErr := errors.New(strings.Repeat("e", 2<<20)) + giantHeaders := http.Header{ + "X-Giant-Header": []string{strings.Repeat("h", 256<<10)}, + } + + for range 24 { + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{Method: http.MethodPost}) + RecordAPIResponseMetadata(ctx, cfg, http.StatusBadGateway, giantHeaders) + RecordAPIResponseError(ctx, cfg, giantErr) + } + + capture := deferredResponseCaptureFrom(ginCtx) + rendered := string(capture.RenderAPIResponse()) + if len(rendered) > maxTotalDeferredAPIResponseBodyBytes+(128<<10) { + t.Fatalf("rendered deferred response len = %d, exceeded bounded budget", len(rendered)) + } + if !strings.Contains(rendered, "[API RESPONSE ERROR TRUNCATED:") { + t.Fatalf("rendered deferred response missing error truncation marker:\n%.512s", rendered) + } + + t.Run("hundred thousand short headers do not allocate unbounded header map", func(t *testing.T) { + manyHeaders := make(http.Header, 100_000) + for i := range 100_000 { + manyHeaders[fmt.Sprintf("X-Short-%06d", i)] = []string{"1"} + } + hdrGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + hdrCtx := context.WithValue(context.Background(), "gin", hdrGinCtx) + RecordAPIRequest(hdrCtx, cfg, UpstreamRequestLog{Method: http.MethodPost}) + + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + RecordAPIResponseMetadata(hdrCtx, cfg, http.StatusBadGateway, manyHeaders) + runtime.ReadMemStats(&after) + + allocated := after.TotalAlloc - before.TotalAlloc + if allocated > 512<<10 { + t.Fatalf("RecordAPIResponseMetadata allocated %d bytes for 100k short headers, want <= %d", allocated, 512<<10) + } + hdrRendered := string(deferredResponseCaptureFrom(hdrGinCtx).RenderAPIResponse()) + if !strings.Contains(hdrRendered, "[API RESPONSE HEADERS TRUNCATED]") { + t.Fatalf("expected header truncation marker:\n%s", hdrRendered) + } + }) + + t.Run("sixty-five short headers retain first sixty-four sorted entries", func(t *testing.T) { + sixtyFive := make(http.Header, 65) + for i := range 65 { + sixtyFive[fmt.Sprintf("X-Cap-%02d", i)] = []string{fmt.Sprintf("val-%02d", i)} + } + capGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + capCtx := context.WithValue(context.Background(), "gin", capGinCtx) + RecordAPIRequest(capCtx, cfg, UpstreamRequestLog{Method: http.MethodPost}) + RecordAPIResponseMetadata(capCtx, cfg, http.StatusBadGateway, sixtyFive) + + capRendered := string(deferredResponseCaptureFrom(capGinCtx).RenderAPIResponse()) + for i := range 64 { + wantLine := fmt.Sprintf("X-Cap-%02d: val-%02d", i, i) + if !strings.Contains(capRendered, wantLine) { + t.Fatalf("rendered headers dropped retained entry %d (%q):\n%s", i, wantLine, capRendered) + } + } + if strings.Contains(capRendered, "X-Cap-64:") { + t.Fatalf("rendered headers unexpectedly retained 65th entry:\n%s", capRendered) + } + if !strings.Contains(capRendered, "[API RESPONSE HEADERS TRUNCATED]") { + t.Fatalf("rendered headers missing truncation marker:\n%s", capRendered) + } + }) +} + +func TestDeferredAPIResponseSkippedWhenFullLogCommercialOrDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("full request logging uses regular capture only", func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}} + + RecordAPIResponseMetadata(ctx, cfg, http.StatusOK, http.Header{"Content-Type": []string{"text/event-stream"}}) + AppendAPIResponseChunk(ctx, cfg, []byte(`data: {"type":"response.created"}`)) + RecordAPIResponseError(ctx, cfg, errors.New("EOF")) + + if _, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey); exists { + t.Fatal("full logging must not populate DeferredAPIResponseContextKey") + } + if _, exists := ginCtx.Get(apiResponseKey); !exists { + t.Fatal("expected regular API_RESPONSE capture") + } + }) + + t.Run("commercial mode retains nothing", func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{CommercialMode: true} + + RecordAPIResponseMetadata(ctx, cfg, http.StatusBadGateway, http.Header{"X-Test": []string{"1"}}) + AppendAPIResponseChunk(ctx, cfg, []byte("upstream error")) + RecordAPIResponseError(ctx, cfg, errors.New("boom")) + RecordAPIWebsocketUpgradeRejection(ctx, cfg, UpstreamRequestLog{URL: "https://example.com"}, http.StatusForbidden, nil, []byte("denied")) + + for _, key := range []string{logging.DeferredAPIRequestContextKey, logging.DeferredAPIResponseContextKey, apiResponseKey} { + if _, exists := ginCtx.Get(key); exists { + t.Fatalf("commercial mode retained %s", key) + } + } + }) + + t.Run("disabled policy survives Gin context reuse", func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + parent := context.WithValue(context.Background(), "gin", ginCtx) + socketCtx, cancel := context.WithCancel(logging.WithDeferredAPIRequestCaptureDisabled(parent)) + cancel() + callbackCtx := context.WithoutCancel(socketCtx) + + RecordAPIRequest(parent, &config.Config{}, UpstreamRequestLog{Body: []byte("new-http-request")}) + RecordAPIResponseMetadata(parent, &config.Config{}, http.StatusBadGateway, nil) + + var callbacks sync.WaitGroup + for range 8 { + callbacks.Go(func() { + RecordAPIResponseMetadata(callbackCtx, &config.Config{}, http.StatusInternalServerError, nil) + AppendAPIResponseChunk(callbackCtx, &config.Config{}, []byte("stale-websocket-chunk")) + RecordAPIResponseError(callbackCtx, &config.Config{}, errors.New("stale-websocket-error")) + RecordAPIWebsocketUpgradeRejection(callbackCtx, &config.Config{}, UpstreamRequestLog{URL: "https://stale.invalid"}, http.StatusForbidden, nil, []byte("stale")) + }) + } + callbacks.Wait() + + rendered := string(deferredResponseCaptureFrom(ginCtx).RenderAPIResponse()) + if strings.Contains(rendered, "stale") { + t.Fatalf("disabled async callbacks polluted reused Gin context:\n%s", rendered) + } + }) +} + func TestAPIResponseAttemptsAreSeparated(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/internal/runtime/executor/helps/model_capabilities.go b/internal/runtime/executor/helps/model_capabilities.go index 24335be931f..4720a0a46f0 100644 --- a/internal/runtime/executor/helps/model_capabilities.go +++ b/internal/runtime/executor/helps/model_capabilities.go @@ -1,6 +1,9 @@ package helps import ( + "fmt" + "strings" + "github.com/MoeForge/moego/ojson" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" @@ -41,10 +44,23 @@ func ApplyRequestThinkingWithModelInfo(body []byte, req cliproxyexecutor.Request source = opts.OriginalRequest } summaryConfig := translatedRequestSummaryConfig(body, req.Payload, originalSource, req.Model, fromFormat, toFormat) + var out []byte + var err error if modelInfo, ok := cliproxyauth.ResolvedModelInfo(req); ok { - return thinking.ApplyThinkingWithModelInfoAndSummary(body, source, req.Model, fromFormat, toFormat, provider, modelInfo, summaryConfig, normalizedUpdatesChanged...) + out, err = thinking.ApplyThinkingWithModelInfoAndSummary(body, source, req.Model, fromFormat, toFormat, provider, modelInfo, summaryConfig, normalizedUpdatesChanged...) + } else { + out, err = thinking.ApplyThinkingWithFallbackModelInfoAndSummary(body, source, req.Model, fromFormat, toFormat, provider, fallbackModelInfo, summaryConfig, normalizedUpdatesChanged...) + } + if err != nil { + return out, err } - return thinking.ApplyThinkingWithFallbackModelInfoAndSummary(body, source, req.Model, fromFormat, toFormat, provider, fallbackModelInfo, summaryConfig, normalizedUpdatesChanged...) + if shouldReconcileClaudeOpenAICompatToolCallReasoning(req, fromFormat, toFormat) && len(out) > 0 { + doc := ojson.NewDoc(out) + if reconcileClaudeOpenAICompatToolCallReasoningDoc(doc) { + out = doc.Bytes() + } + } + return out, nil } // ApplyRequestThinkingWithModelInfoDoc applies request thinking onto doc in place @@ -63,8 +79,66 @@ func ApplyRequestThinkingWithModelInfoDoc(doc *ojson.Doc, req cliproxyexecutor.R source = opts.OriginalRequest } summaryConfig := translatedRequestSummaryConfigDoc(doc, req.Payload, originalSource, req.Model, fromFormat, toFormat) + var err error if modelInfo, ok := cliproxyauth.ResolvedModelInfo(req); ok { - return thinking.ApplyThinkingWithModelInfoAndSummaryDoc(doc, source, req.Model, fromFormat, toFormat, provider, modelInfo, summaryConfig, normalizedUpdatesChanged...) + err = thinking.ApplyThinkingWithModelInfoAndSummaryDoc(doc, source, req.Model, fromFormat, toFormat, provider, modelInfo, summaryConfig, normalizedUpdatesChanged...) + } else { + err = thinking.ApplyThinkingWithFallbackModelInfoAndSummaryDoc(doc, source, req.Model, fromFormat, toFormat, provider, fallbackModelInfo, summaryConfig, normalizedUpdatesChanged...) + } + if err != nil { + return err + } + if shouldReconcileClaudeOpenAICompatToolCallReasoning(req, fromFormat, toFormat) { + reconcileClaudeOpenAICompatToolCallReasoningDoc(doc) + } + return nil +} + +func shouldReconcileClaudeOpenAICompatToolCallReasoning(req cliproxyexecutor.Request, fromFormat, toFormat string) bool { + return APIKeyModelIsCompat(req) && + strings.EqualFold(strings.TrimSpace(fromFormat), "claude") && + strings.EqualFold(strings.TrimSpace(toFormat), "openai") +} + +func reconcileClaudeOpenAICompatToolCallReasoningDoc(doc *ojson.Doc) bool { + if doc == nil || doc.Len() == 0 { + return false + } + messages := doc.Get("messages") + if !messages.Exists() || !messages.IsArray() { + return false + } + reasoningActive := false + if effortRes := doc.Get("reasoning_effort"); effortRes.Exists() { + resolvedEffort := strings.ToLower(strings.TrimSpace(effortRes.String())) + if resolvedEffort != "" && resolvedEffort != string(thinking.LevelNone) && resolvedEffort != "0" && resolvedEffort != "false" { + reasoningActive = true + } + } + + modified := false + for idx, msg := range messages.Array() { + if strings.TrimSpace(msg.Get("role").String()) != "assistant" { + continue + } + toolCalls := msg.Get("tool_calls") + if !toolCalls.Exists() || !toolCalls.IsArray() || len(toolCalls.Array()) == 0 { + continue + } + reasoning := msg.Get("reasoning_content") + if reasoningActive { + if !reasoning.Exists() || (strings.TrimSpace(reasoning.String()) == "" && reasoning.String() != " ") { + if doc.Set(fmt.Sprintf("messages.%d.reasoning_content", idx), " ") == nil { + modified = true + } + } + continue + } + if reasoning.Exists() && strings.TrimSpace(reasoning.String()) == "" { + if doc.Delete(fmt.Sprintf("messages.%d.reasoning_content", idx)) == nil { + modified = true + } + } } - return thinking.ApplyThinkingWithFallbackModelInfoAndSummaryDoc(doc, source, req.Model, fromFormat, toFormat, provider, fallbackModelInfo, summaryConfig, normalizedUpdatesChanged...) + return modified } diff --git a/internal/runtime/executor/helps/model_capabilities_test.go b/internal/runtime/executor/helps/model_capabilities_test.go index ccc39f338b3..01334477e06 100644 --- a/internal/runtime/executor/helps/model_capabilities_test.go +++ b/internal/runtime/executor/helps/model_capabilities_test.go @@ -379,3 +379,125 @@ func TestAPIKeyModelIsCompatUsesAuthoritativeHomeCapabilities(t *testing.T) { } } } + +func TestApplyRequestThinkingReconcilesClaudeOpenAICompatToolCallReasoning(t *testing.T) { + toolCallWithPlaceholder := `{"model":"deepseek-v4-flash","reasoning_effort":"low","messages":[{"role":"assistant","content":"","reasoning_content":" ","tool_calls":[{"id":"call_1","type":"function","function":{"name":"Read","arguments":"{}"}}]}]}` + toolCallWithoutReasoning := `{"model":"deepseek-v4-flash","reasoning_effort":"none","messages":[{"role":"assistant","content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"Read","arguments":"{}"}}]}]}` + toolCallWithRealReasoning := `{"model":"deepseek-v4-flash","reasoning_effort":"low","messages":[{"role":"assistant","content":"","reasoning_content":"actual thought","tool_calls":[{"id":"call_1","type":"function","function":{"name":"Read","arguments":"{}"}}]}]}` + + for _, tc := range []struct { + name string + model string + sourceBody string + translated string + fromFormat string + toFormat string + modelInfo *registry.ModelInfo + wantExists bool + wantReasoning string + }{ + { + name: "unsupported thinking model strips reasoning_effort and removes placeholder", + model: "deepseek-v4-flash", + sourceBody: `{"model":"deepseek-v4-flash","thinking":{"type":"enabled","budget_tokens":1024}}`, + translated: toolCallWithPlaceholder, + fromFormat: "claude", + toFormat: "openai", + modelInfo: ®istry.ModelInfo{ + ID: "deepseek-v4-flash", + IsCompat: true, + Thinking: nil, + }, + wantExists: false, + }, + { + name: "model disallowing zero clamps none to active level and injects placeholder", + model: "deepseek-v4-flash", + sourceBody: `{"model":"deepseek-v4-flash","thinking":{"type":"disabled"}}`, + translated: toolCallWithoutReasoning, + fromFormat: "claude", + toFormat: "openai", + modelInfo: ®istry.ModelInfo{ + ID: "deepseek-v4-flash", + IsCompat: true, + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"low", "medium", "high"}, + ZeroAllowed: false, + }, + }, + wantExists: true, + wantReasoning: " ", + }, + { + name: "real reasoning preserved when suffix disables thinking", + model: "deepseek-v4-flash(none)", + sourceBody: `{"model":"deepseek-v4-flash(none)","thinking":{"type":"enabled","budget_tokens":1024}}`, + translated: toolCallWithRealReasoning, + fromFormat: "claude", + toFormat: "openai", + modelInfo: ®istry.ModelInfo{ + ID: "deepseek-v4-flash", + IsCompat: true, + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"none", "low", "medium", "high"}, + ZeroAllowed: true, + }, + }, + wantExists: true, + wantReasoning: "actual thought", + }, + { + name: "non-compat model is untouched", + model: "deepseek-v4-flash(high)", + sourceBody: `{"model":"deepseek-v4-flash(high)","thinking":{"type":"disabled"}}`, + translated: toolCallWithoutReasoning, + fromFormat: "claude", + toFormat: "openai", + modelInfo: ®istry.ModelInfo{ + ID: "deepseek-v4-flash", + IsCompat: false, + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"low", "medium", "high"}, + }, + }, + wantExists: false, + }, + } { + t.Run(tc.name, func(t *testing.T) { + req := cliproxyexecutor.Request{ + Model: tc.model, + Payload: []byte(tc.sourceBody), + Metadata: map[string]any{ + "cliproxy.resolved_api_key_model_info": tc.modelInfo, + }, + } + opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(tc.sourceBody), + } + + outBytes, errBytes := helps.ApplyRequestThinking([]byte(tc.translated), req, opts, tc.fromFormat, tc.toFormat, "openai-compatibility") + if errBytes != nil { + t.Fatalf("ApplyRequestThinking() error = %v", errBytes) + } + gotBytes := ojson.GetBytes(outBytes, "messages.0.reasoning_content") + if gotBytes.Exists() != tc.wantExists { + t.Fatalf("ApplyRequestThinking() reasoning_content exists = %v, want %v; out=%s", gotBytes.Exists(), tc.wantExists, outBytes) + } + if tc.wantExists && gotBytes.String() != tc.wantReasoning { + t.Fatalf("ApplyRequestThinking() reasoning_content = %q, want %q; out=%s", gotBytes.String(), tc.wantReasoning, outBytes) + } + + doc := ojson.NewDocString(tc.translated) + if errDoc := helps.ApplyRequestThinkingDoc(doc, req, opts, tc.fromFormat, tc.toFormat, "openai-compatibility"); errDoc != nil { + t.Fatalf("ApplyRequestThinkingDoc() error = %v", errDoc) + } + gotDoc := doc.Get("messages.0.reasoning_content") + if gotDoc.Exists() != tc.wantExists { + t.Fatalf("ApplyRequestThinkingDoc() reasoning_content exists = %v, want %v; out=%s", gotDoc.Exists(), tc.wantExists, doc.String()) + } + if tc.wantExists && gotDoc.String() != tc.wantReasoning { + t.Fatalf("ApplyRequestThinkingDoc() reasoning_content = %q, want %q; out=%s", gotDoc.String(), tc.wantReasoning, doc.String()) + } + }) + } +} diff --git a/internal/runtime/executor/helps/response_model.go b/internal/runtime/executor/helps/response_model.go index 20b5881d9c3..bd060a463f8 100644 --- a/internal/runtime/executor/helps/response_model.go +++ b/internal/runtime/executor/helps/response_model.go @@ -46,9 +46,28 @@ func extractClaudeResponseModelEvent(data []byte) (model string, terminal bool) } } return "", false + case "content_block_start": + if ojson.GetBytesNoCopy(data, "content_block.type").String() == "fallback" { + m := ojson.GetBytesNoCopy(data, "content_block.to.model") + if m.Type == ojson.String { + s := strings.TrimSpace(m.String()) + if len(s) <= maxResponseModelLength { + return strings.Clone(s), false + } + } + } + return "", false + case "message_delta": + if s := extractClaudeFallbackIterationModel(ojson.GetBytesNoCopy(data, "usage.iterations")); s != "" { + return s, false + } + return "", false case "message_stop": return "", true case "message": + if s := extractClaudeFallbackIterationModel(ojson.GetBytesNoCopy(data, "usage.iterations")); s != "" { + return s, true + } m := ojson.GetBytesNoCopy(data, "model") if m.Type == ojson.String { s := strings.TrimSpace(m.String()) @@ -76,6 +95,30 @@ func extractClaudeResponseModelEvent(data []byte) (model string, terminal bool) } } +func extractClaudeFallbackIterationModel(iterations ojson.Result) string { + if !iterations.Exists() || !iterations.IsArray() { + return "" + } + var served string + for _, it := range iterations.Array() { + if it.Get("type").String() != "fallback_message" { + continue + } + m := it.Get("model") + if m.Type != ojson.String { + continue + } + s := strings.TrimSpace(m.String()) + if s != "" && len(s) <= maxResponseModelLength { + served = s + } + } + if served == "" { + return "" + } + return strings.Clone(served) +} + // extractGeminiResponseModelEvent extracts the response model from a Google Gemini / Vertex / AIStudio response. func extractGeminiResponseModelEvent(data []byte) (model string, terminal bool) { if !ojson.ValidBytes(data) { diff --git a/internal/runtime/executor/helps/response_model_test.go b/internal/runtime/executor/helps/response_model_test.go index d17184add36..6911eb7bdb8 100644 --- a/internal/runtime/executor/helps/response_model_test.go +++ b/internal/runtime/executor/helps/response_model_test.go @@ -478,3 +478,25 @@ func TestExtractResponseModelClonesFromReusableBuffer(t *testing.T) { }) } } + +func TestExtractClaudeResponseModelMidStreamFallback(t *testing.T) { + lines := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_1","model":"claude-opus-5-5","usage":{"input_tokens":8,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}}}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":0},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":5}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + + reporter := NewUsageReporter(context.Background(), "claude", "claude-opus-5-5", nil) + var buffer StreamUsageBuffer + for _, line := range lines { + reporter.ObserveResponseModel(line) + buffer.ObserveClaudeStream(line) + } + if got := reporter.ResponseModel(); got != "claude-opus-4-8" { + t.Fatalf("reporter.ResponseModel() = %q, want claude-opus-4-8", got) + } + if got := buffer.ResponseModel(); got != "claude-opus-4-8" { + t.Fatalf("buffer.ResponseModel() = %q, want claude-opus-4-8", got) + } +} diff --git a/internal/runtime/executor/helps/thinking_test.go b/internal/runtime/executor/helps/thinking_test.go index 8e1293f188d..cfcec31be42 100644 --- a/internal/runtime/executor/helps/thinking_test.go +++ b/internal/runtime/executor/helps/thinking_test.go @@ -98,6 +98,38 @@ func TestApplyThinkingWithSourcePayloadPreservesOriginalOnlySummary(t *testing.T } } +func TestApplyThinkingWithSourcePayload_AntigravityClaudeThinkingDisplayUpdates(t *testing.T) { + source := []byte(`{"model":"gemini-3.8-flash-high","max_tokens":4096,"thinking":{"type":"adaptive","display":"updates"},"messages":[{"role":"user","content":"Hello, think first."}]}`) + translated := sdktranslator.TranslateRequest( + sdktranslator.FormatClaude, + sdktranslator.FormatAntigravity, + "gemini-3.8-flash-high", + source, + false, + ) + + out, err := helps.ApplyThinkingWithSourcePayload( + translated, + source, + source, + "gemini-3.8-flash-high", + sdktranslator.FormatClaude.String(), + sdktranslator.FormatAntigravity.String(), + "antigravity", + ) + if err != nil { + t.Fatalf("ApplyThinkingWithSourcePayload() error = %v", err) + } + include := ojson.GetBytes(out, "request.generationConfig.thinkingConfig.includeThoughts") + if !include.Exists() || !include.Bool() { + t.Fatalf("includeThoughts not enabled: %s", out) + } + level := ojson.GetBytes(out, "request.generationConfig.thinkingConfig.thinkingLevel").String() + if level != "high" { + t.Fatalf("thinkingLevel = %q, want high: %s", level, out) + } +} + func TestApplyThinkingWithSourcePayload_AntigravityResponsesReasoningSummaryAuto(t *testing.T) { source := []byte(`{"model":"gemini-3.8-flash-high","reasoning":{"effort":"low","summary":"auto"},"input":"hi"}`) translated := sdktranslator.TranslateRequest( diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 1fbbf6963f8..8c346c63a3d 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -19,6 +19,7 @@ import ( "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v8/internal/clienterror" + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" internallogging "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" @@ -42,6 +43,7 @@ type UsageReporter struct { authIndex string authMu sync.RWMutex accessTokenHash string + requestScopedRules []internalconfig.RequestScopedErrorRule authType string apiKey string accessProvider string @@ -154,6 +156,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox reporter.authID = auth.ID reporter.authIndex = auth.EnsureIndex() reporter.accessTokenHash = authAccessTokenSHA256(auth) + reporter.requestScopedRules = cliproxyauth.ResolveRequestScopedErrorRules(auth, cliproxyauth.RequestScopedErrorConfigFromContext(ctx)) } return reporter } @@ -202,11 +205,27 @@ func (r *UsageReporter) UpdateAccessTokenFingerprint(auth *cliproxyauth.Auth) { if r == nil { return } + var updatedRules []internalconfig.RequestScopedErrorRule + if auth != nil { + updatedRules = cliproxyauth.ResolveRequestScopedErrorRules(auth, nil) + } r.authMu.Lock() r.accessTokenHash = authAccessTokenSHA256(auth) + if len(updatedRules) > 0 { + r.requestScopedRules = updatedRules + } r.authMu.Unlock() } +func (r *UsageReporter) requestScopedErrorRules() []internalconfig.RequestScopedErrorRule { + if r == nil { + return nil + } + r.authMu.RLock() + defer r.authMu.RUnlock() + return r.requestScopedRules +} + func (r *UsageReporter) accessTokenFingerprint() string { if r == nil { return "" @@ -555,11 +574,11 @@ func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.De } func (r *UsageReporter) PublishFailure(ctx context.Context, errs ...error) { - r.publishWithOutcome(ctx, usage.Detail{}, true, failFromErrors(errs...)) + r.publishWithOutcome(ctx, usage.Detail{}, true, r.failFromErrors(ctx, errs...)) } func (r *UsageReporter) PublishFailureWithDetail(ctx context.Context, detail usage.Detail, errs ...error) { - r.publishWithOutcome(ctx, detail, true, failFromErrors(errs...)) + r.publishWithOutcome(ctx, detail, true, r.failFromErrors(ctx, errs...)) } func (r *UsageReporter) TrackFailure(ctx context.Context, errPtr *error) { @@ -721,7 +740,15 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f } } +func (r *UsageReporter) failFromErrors(_ context.Context, errs ...error) usage.Failure { + return failFromErrorsWithRules(r.requestScopedErrorRules(), errs...) +} + func failFromErrors(errs ...error) usage.Failure { + return failFromErrorsWithRules(nil, errs...) +} + +func failFromErrorsWithRules(rules []internalconfig.RequestScopedErrorRule, errs ...error) usage.Failure { for _, err := range errs { if err == nil { continue @@ -742,6 +769,12 @@ func failFromErrors(errs ...error) usage.Failure { } if truncatedErr, ok := errors.AsType[cliproxyexecutor.StreamTruncatedError](err); ok && truncatedErr != nil && truncatedErr.IsStreamTruncated() { fail.Code = usage.FailureStreamTruncatedCode + } else if action, okAction := cliproxyauth.MatchRequestScopedErrorRules(rules, err); okAction { + if action == cliproxyauth.RequestScopedActionStop || action == cliproxyauth.RequestScopedActionContinue { + fail.Code = usage.FailureRequestScopedCode + } else if action == cliproxyauth.RequestScopedActionStopAndCooldown || action == cliproxyauth.RequestScopedActionContinueAndCooldown { + fail.Code = usage.FailureForceCooldownCode + } } else if requestErr, ok := errors.AsType[cliproxyexecutor.RequestScopedError](err); ok && requestErr != nil && requestErr.IsRequestScoped() { fail.Code = usage.FailureRequestScopedCode } @@ -1353,7 +1386,7 @@ func (b *StreamUsageBuffer) ObserveClaudeStream(line []byte) { return } payload := jsonPayload(line) - if b.responseModel == "" && len(payload) > 0 { + if len(payload) > 0 { if model, _ := extractClaudeResponseModelEvent(payload); model != "" { b.responseModel = model } diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 386cbe16b4e..f7a871ff137 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -1517,6 +1517,79 @@ func TestFailFromErrorsCarriesRequestScopedCode(t *testing.T) { if plain.Code != "" { t.Fatalf("upstream model rejection gained request-scoped classification: %+v", plain) } + + rules := []config.RequestScopedErrorRule{ + { + Status: http.StatusNotFound, + Match: []string{"model_not_found"}, + Action: cliproxyauth.RequestScopedActionStop, + }, + { + Status: http.StatusBadRequest, + Match: []string{"force_cool_marker"}, + Action: cliproxyauth.RequestScopedActionStopAndCooldown, + }, + } + authWithMetadataRules := &cliproxyauth.Auth{ + ID: "auth-rule-test", + Provider: "codex", + Metadata: map[string]any{"request_scoped_errors": rules}, + } + reporter := NewUsageReporter(context.Background(), "codex", "gpt-5.4", authWithMetadataRules) + var errNotFound error = &cliproxyauth.Error{HTTPStatus: http.StatusNotFound, Message: "model_not_found"} + failMatched := reporter.failFromErrors(context.Background(), errNotFound) + if failMatched.Code != usage.FailureRequestScopedCode || failMatched.StatusCode != http.StatusNotFound { + t.Fatalf("matched request-scoped rule must set fail.Code=%q, got %+v", usage.FailureRequestScopedCode, failMatched) + } + + var errForceCool error = &cliproxyauth.Error{HTTPStatus: http.StatusBadRequest, Message: "force_cool_marker"} + failCooldown := reporter.failFromErrors(context.Background(), errForceCool) + if failCooldown.Code != usage.FailureForceCooldownCode { + t.Fatalf("cooldown action rule must set fail.Code=%q, got %+v", usage.FailureForceCooldownCode, failCooldown) + } + + // A configured cooldown rule must also override an error that intrinsically implements RequestScopedError. + reqScopedForceCool := cliproxyauth.NewRequestScopedError("force_cool_marker", http.StatusBadRequest) + if failOverride := reporter.failFromErrors(context.Background(), reqScopedForceCool); failOverride.Code != usage.FailureForceCooldownCode { + t.Fatalf("cooldown rule must override intrinsic RequestScopedError: got %+v", failOverride) + } + + // Mutating the caller's slice after NewUsageReporter must not affect the reporter's snapshot. + rules[0].Match[0] = "mutated_after_creation" + failIsolated := reporter.failFromErrors(context.Background(), errNotFound) + if failIsolated.Code != usage.FailureRequestScopedCode { + t.Fatalf("reporter rule snapshot was mutated by caller slice modification: %+v", failIsolated) + } + + oauthAuth := &cliproxyauth.Auth{ + ID: "oauth-rule-test", + Provider: "codex", + Attributes: map[string]string{"auth_kind": "oauth"}, + } + cfg := &config.Config{ + OAuthRequestScopedErrors: map[string][]config.RequestScopedErrorRule{ + "codex": { + { + Status: http.StatusNotFound, + Match: []string{"endpoint_unsupported"}, + Action: cliproxyauth.RequestScopedActionContinue, + }, + }, + }, + } + ctxWithCfg := cliproxyauth.WithRequestScopedErrorConfig(context.Background(), cfg) + oauthReporter := NewUsageReporter(ctxWithCfg, "codex", "gpt-5.4", oauthAuth) + // Updating token fingerprint with an OAuth auth that has no per-auth metadata rules + // must preserve the OAuth rules resolved from the context config at reporter creation. + oauthReporter.UpdateAccessTokenFingerprint(&cliproxyauth.Auth{ + ID: "oauth-rule-test", + Provider: "codex", + Attributes: map[string]string{"auth_kind": "oauth", "api_key": "rotated-token"}, + }) + failOAuth := oauthReporter.failFromErrors(ctxWithCfg, &cliproxyauth.Error{HTTPStatus: http.StatusNotFound, Message: "endpoint_unsupported"}) + if failOAuth.Code != usage.FailureRequestScopedCode { + t.Fatalf("OAuth request-scoped rule from context config must set fail.Code=%q after token update, got %+v", usage.FailureRequestScopedCode, failOAuth) + } } func TestParseClaudeUsageExtractsCacheTTLServerToolsIterationsAndCost(t *testing.T) { diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index d0b0cd468b3..80d9159b371 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -210,18 +210,21 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = statusErr{code: httpResp.StatusCode, msg: string(b)} return resp, err } data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) reporter.ObserveResponseModel(data) var param any // Note: TranslateNonStream uses req.Model (original with suffix) to preserve @@ -345,8 +348,11 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("kimi executor: close response body error: %v", errClose) @@ -507,19 +513,22 @@ func (e *KimiExecutor) executeResponses(ctx context.Context, auth *cliproxyauth. helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = statusErr{code: httpResp.StatusCode, msg: string(b)} return resp, err } data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) reporter.ObserveResponseModel(data) originalRequest := req.Payload @@ -633,11 +642,14 @@ func (e *KimiExecutor) executeResponsesStream(ctx context.Context, auth *cliprox helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("kimi executor: close response body error: %v", errClose) } helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = statusErr{code: httpResp.StatusCode, msg: string(b)} return nil, err diff --git a/internal/runtime/executor/meta_executor_execute.go b/internal/runtime/executor/meta_executor_execute.go index ed7e98095be..b08b534358c 100644 --- a/internal/runtime/executor/meta_executor_execute.go +++ b/internal/runtime/executor/meta_executor_execute.go @@ -46,7 +46,14 @@ func (e *MetaExecutor) prepareResponsesRequest(ctx context.Context, auth *clipro isCompat := helps.APIKeyModelIsCompat(req) targetModel := resolveMetaSignatureTargetModel(baseModel, req, opts) originalTranslated := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, targetModel, originalPayload, stream, isCompat) - originalTranslated = helps.SetStringIfDifferent(originalTranslated, "model", baseModel) + origDoc := ojson.NewDoc(originalTranslated) + origChanged := helps.SetStringIfDifferentDoc(origDoc, "model", baseModel) + if syncMetaPromptCacheRetentionDoc(origDoc, originalPayload) { + origChanged = true + } + if origChanged { + originalTranslated = origDoc.Bytes() + } body, err := helps.TranslateRequestReturningError(ctx, opts.Headers, e.cfg, from, to, targetModel, bytes.Clone(req.Payload), stream, isCompat) if err != nil { return nil, err @@ -62,7 +69,7 @@ func (e *MetaExecutor) prepareResponsesRequest(ctx context.Context, auth *clipro helps.SetStringIfDifferentDoc(doc, "model", baseModel) helps.SetBoolIfDifferentDoc(doc, "stream", stream) _ = doc.Delete("generate") - _ = doc.Delete("prompt_cache_retention") + syncMetaPromptCacheRetentionDoc(doc, req.Payload) _ = doc.Delete("safety_identifier") _ = doc.Delete("stream_options") _ = doc.Delete("client_metadata") @@ -95,6 +102,33 @@ func (e *MetaExecutor) prepareResponsesRequest(ctx context.Context, auth *clipro }, nil } +// syncMetaPromptCacheRetentionDoc preserves an explicit non-empty string +// prompt_cache_retention supplied by the caller without injecting a default, +// and removes null, blank, or non-string values before user payload rules run. +func syncMetaPromptCacheRetentionDoc(doc *ojson.Doc, source []byte) bool { + if doc == nil { + return false + } + if retention := extractMetaPromptCacheRetention(source); retention != "" { + return helps.SetStringIfDifferentDoc(doc, "prompt_cache_retention", retention) + } + if !doc.Get("prompt_cache_retention").Exists() { + return false + } + return doc.Delete("prompt_cache_retention") == nil +} + +func extractMetaPromptCacheRetention(source []byte) string { + if len(source) == 0 { + return "" + } + val := ojson.GetBytesNoCopy(source, "prompt_cache_retention") + if val.Type != ojson.String { + return "" + } + return strings.TrimSpace(val.String()) +} + func (e *MetaExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { ctx = helps.EnsureSessionContext(ctx, opts, req.Payload) if opts.Alt == "responses/compact" { @@ -149,11 +183,11 @@ func (e *MetaExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) reporter.ObserveResponseModel(data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) diff --git a/internal/runtime/executor/meta_executor_stream.go b/internal/runtime/executor/meta_executor_stream.go index af94e144ccd..ff439c3160a 100644 --- a/internal/runtime/executor/meta_executor_stream.go +++ b/internal/runtime/executor/meta_executor_stream.go @@ -68,11 +68,11 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("meta executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return nil, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) return nil, wrapMetaUpstreamError(httpResp.StatusCode, data) } diff --git a/internal/runtime/executor/meta_executor_test.go b/internal/runtime/executor/meta_executor_test.go index 55d9f03361d..8b7f3a28801 100644 --- a/internal/runtime/executor/meta_executor_test.go +++ b/internal/runtime/executor/meta_executor_test.go @@ -2247,3 +2247,207 @@ func TestMetaExecutor_ResponsesCompactionTriggerAndReplay_Issue6781(t *testing.T } assertStatusErr(t, errIncStream, http.StatusBadGateway) } + +func TestMetaExecutor_PromptCacheRetentionPassthrough_Issue7424(t *testing.T) { + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read request body: %v", errRead) + http.Error(w, "bad body", http.StatusBadRequest) + return + } + gotBody = bytes.Clone(body) + writeMetaResponsesOK(w, "ok") + })) + defer server.Close() + + auth := &cliproxyauth.Auth{ + Provider: "meta", + Attributes: map[string]string{ + "api_key": "meta-token", + "base_url": server.URL, + }, + } + + runBothModes := func(t *testing.T, exec *MetaExecutor, format, payload string) [][]byte { + t.Helper() + raw := []byte(payload) + var bodies [][]byte + + gotBody = nil + _, errExec := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "muse-ember-1.0", + Payload: raw, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString(format), + OriginalRequest: raw, + }) + if errExec != nil { + t.Fatalf("Execute(%s) error = %v", format, errExec) + } + bodies = append(bodies, bytes.Clone(gotBody)) + + gotBody = nil + streamRes, errStream := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "muse-ember-1.0", + Payload: raw, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString(format), + Stream: true, + OriginalRequest: raw, + }) + if errStream != nil { + t.Fatalf("ExecuteStream(%s) error = %v", format, errStream) + } + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + t.Fatalf("ExecuteStream(%s) chunk error = %v", format, chunk.Err) + } + } + bodies = append(bodies, bytes.Clone(gotBody)) + return bodies + } + + t.Run("omitted or blank values are not sent and never default to 24h", func(t *testing.T) { + exec := NewMetaExecutor(&config.Config{}) + cases := []struct { + name string + format string + payload string + }{ + { + name: "responses omitted", + format: "openai-response", + payload: `{"model":"muse-ember-1.0","input":"hello"}`, + }, + { + name: "responses null", + format: "openai-response", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":null,"input":"hello"}`, + }, + { + name: "responses blank", + format: "openai-response", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":" ","input":"hello"}`, + }, + { + name: "chat completions omitted", + format: "openai", + payload: `{"model":"muse-ember-1.0","messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "codex null", + format: "codex", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":null,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`, + }, + { + name: "codex non-string", + format: "codex", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":24,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + for idx, body := range runBothModes(t, exec, tc.format, tc.payload) { + if ojson.GetBytes(body, "prompt_cache_retention").Exists() { + t.Fatalf("mode %d unexpectedly sent prompt_cache_retention: %s", idx, body) + } + } + }) + } + }) + + t.Run("explicit client prompt_cache_retention passes through", func(t *testing.T) { + exec := NewMetaExecutor(&config.Config{}) + cases := []struct { + name string + format string + payload string + want string + }{ + { + name: "responses 24h", + format: "openai-response", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":"24h","input":"hello"}`, + want: "24h", + }, + { + name: "responses in_memory", + format: "openai-response", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":"in_memory","input":"hello"}`, + want: "in_memory", + }, + { + name: "chat completions 24h", + format: "openai", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":"24h","messages":[{"role":"user","content":"hello"}]}`, + want: "24h", + }, + { + name: "codex 24h", + format: "codex", + payload: `{"model":"muse-ember-1.0","prompt_cache_retention":"24h","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`, + want: "24h", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + for idx, body := range runBothModes(t, exec, tc.format, tc.payload) { + if got := ojson.GetBytes(body, "prompt_cache_retention").String(); got != tc.want { + t.Fatalf("mode %d prompt_cache_retention = %q, want %q; body=%s", idx, got, tc.want, body) + } + } + }) + } + }) + + t.Run("payload rules remain final semantic barrier", func(t *testing.T) { + defaultExec := NewMetaExecutor(&config.Config{ + Payload: config.PayloadConfig{ + Default: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "muse-*", Protocol: "meta"}}, + Params: map[string]any{"prompt_cache_retention": "24h"}, + }}, + }, + }) + for idx, body := range runBothModes(t, defaultExec, "openai-response", `{"model":"muse-ember-1.0","input":"hello"}`) { + if got := ojson.GetBytes(body, "prompt_cache_retention").String(); got != "24h" { + t.Fatalf("default rule on omitted field mode %d = %q, want 24h; body=%s", idx, got, body) + } + } + for idx, body := range runBothModes(t, defaultExec, "openai-response", `{"model":"muse-ember-1.0","prompt_cache_retention":"in_memory","input":"hello"}`) { + if got := ojson.GetBytes(body, "prompt_cache_retention").String(); got != "in_memory" { + t.Fatalf("default rule overwrote explicit client in_memory in mode %d: got %q; body=%s", idx, got, body) + } + } + + overrideExec := NewMetaExecutor(&config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "muse-*", Protocol: "meta"}}, + Params: map[string]any{"prompt_cache_retention": "24h"}, + }}, + }, + }) + for idx, body := range runBothModes(t, overrideExec, "openai-response", `{"model":"muse-ember-1.0","prompt_cache_retention":"in_memory","input":"hello"}`) { + if got := ojson.GetBytes(body, "prompt_cache_retention").String(); got != "24h" { + t.Fatalf("override rule mode %d = %q, want 24h; body=%s", idx, got, body) + } + } + + filterExec := NewMetaExecutor(&config.Config{ + Payload: config.PayloadConfig{ + Filter: []config.PayloadFilterRule{{ + Models: []config.PayloadModelRule{{Name: "muse-*", Protocol: "meta"}}, + Params: []string{"prompt_cache_retention"}, + }}, + }, + }) + for idx, body := range runBothModes(t, filterExec, "openai-response", `{"model":"muse-ember-1.0","prompt_cache_retention":"24h","input":"hello"}`) { + if ojson.GetBytes(body, "prompt_cache_retention").Exists() { + t.Fatalf("filter rule failed to remove prompt_cache_retention in mode %d: %s", idx, body) + } + } + }) +} diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index 17a06945a8a..c4ac7e0cc5d 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -230,18 +230,21 @@ func (e *OpenAICompatExecutor) executeWithCredentials(ctx context.Context, auth }() helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = newOpenAICompatStatusError(httpResp.StatusCode, httpResp.Header, b) return resp, err } body, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, body) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, body) reporter.ObserveResponseModel(body) // Ensure we at least record the request even if upstream doesn't return usage // Translate response back to source format when needed @@ -332,12 +335,12 @@ func (e *OpenAICompatExecutor) executeImages(ctx context.Context, auth *cliproxy helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) body, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, body) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) err = errRead return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, body) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), body)) @@ -470,8 +473,11 @@ func (e *OpenAICompatExecutor) executeStreamWithCredentials(ctx context.Context, } helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) + b, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("openai compat executor: close response body error: %v", errClose) @@ -759,11 +765,11 @@ func (e *OpenAICompatExecutor) executeImagesStream(ctx context.Context, auth *cl if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("openai compat executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, body) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return nil, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, body) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), body)) return nil, statusErr{code: httpResp.StatusCode, msg: string(body)} } diff --git a/internal/runtime/executor/openai_compat_executor_reasoning_test.go b/internal/runtime/executor/openai_compat_executor_reasoning_test.go index d5cda5fc543..05ecc5a0b36 100644 --- a/internal/runtime/executor/openai_compat_executor_reasoning_test.go +++ b/internal/runtime/executor/openai_compat_executor_reasoning_test.go @@ -58,6 +58,142 @@ func TestOpenAICompatExecutorUsesCompatibleClaudeTranslation(t *testing.T) { } } +func TestOpenAICompatExecutorAddsPlaceholderReasoningForToolCallsWithoutThinkingWhenCompat(t *testing.T) { + models := []config.PayloadModelRule{{Name: "deepseek-v4-flash", Protocol: "openai"}} + basePayload := `{"model":"deepseek-v4-flash","thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"user","content":"what time"},{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"get_time","input":{}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1","content":"12:00"}]}]}` + disabledPayload := `{"model":"deepseek-v4-flash","thinking":{"type":"disabled"},"messages":[{"role":"user","content":"what time"},{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"get_time","input":{}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1","content":"12:00"}]}]}` + + for _, tt := range []struct { + name string + reqModel string + reqPayload string + payloadCfg config.PayloadConfig + wantExists bool + wantReasoning string + }{ + { + name: "placeholder injected when thinking enabled", + reqModel: "deepseek-v4-flash", + reqPayload: basePayload, + wantExists: true, + wantReasoning: " ", + }, + { + name: "placeholder omitted when thinking disabled", + reqModel: "deepseek-v4-flash", + reqPayload: disabledPayload, + wantExists: false, + }, + { + name: "model suffix high overrides disabled body and injects placeholder", + reqModel: "deepseek-v4-flash(high)", + reqPayload: disabledPayload, + wantExists: true, + wantReasoning: " ", + }, + { + name: "model suffix none overrides enabled body and removes placeholder", + reqModel: "deepseek-v4-flash(none)", + reqPayload: basePayload, + wantExists: false, + }, + { + name: "final payload override wins", + reqModel: "deepseek-v4-flash", + reqPayload: basePayload, + payloadCfg: config.PayloadConfig{Override: []config.PayloadRule{ + {Models: models, Params: map[string]any{"messages.1.reasoning_content": "custom-override"}}, + }}, + wantExists: true, + wantReasoning: "custom-override", + }, + { + name: "final payload filter is not restored", + reqModel: "deepseek-v4-flash", + reqPayload: basePayload, + payloadCfg: config.PayloadConfig{Filter: []config.PayloadFilterRule{ + {Models: models, Params: []string{"messages.1.reasoning_content"}}, + }}, + wantExists: false, + }, + } { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/stream=%v", tt.name, stream), func(t *testing.T) { + bodies := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read upstream request: %v", errRead) + } + bodies <- body + if stream { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"chatcmpl-test","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`) + })) + t.Cleanup(server.Close) + + executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{Payload: tt.payloadCfg}) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test-key", + }, + } + request := cliproxyexecutor.Request{ + Model: tt.reqModel, + Payload: []byte(tt.reqPayload), + Metadata: map[string]any{ + "cliproxy.resolved_api_key_model_info": ®istry.ModelInfo{ + IsCompat: true, + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"none", "low", "medium", "high"}, + ZeroAllowed: true, + }, + }, + }, + } + options := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + ResponseFormat: sdktranslator.FormatOpenAI, + Stream: stream, + } + + if stream { + result, errStream := executor.ExecuteStream(context.Background(), auth, request, options) + if errStream != nil { + t.Fatalf("ExecuteStream error: %v", errStream) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error: %v", chunk.Err) + } + } + } else if _, errExecute := executor.Execute(context.Background(), auth, request, options); errExecute != nil { + t.Fatalf("Execute error: %v", errExecute) + } + + upstreamBody := <-bodies + assistant := ojson.GetBytes(upstreamBody, "messages.1") + reasoning := assistant.Get("reasoning_content") + if reasoning.Exists() != tt.wantExists { + t.Fatalf("reasoning_content exists = %v, want %v; body=%s", reasoning.Exists(), tt.wantExists, upstreamBody) + } + if tt.wantExists && reasoning.String() != tt.wantReasoning { + t.Fatalf("reasoning_content = %q, want %q; body=%s", reasoning.String(), tt.wantReasoning, upstreamBody) + } + if !assistant.Get("tool_calls").Exists() { + t.Fatalf("tool_calls missing from upstream request: %s", upstreamBody) + } + }) + } + } +} + func TestOpenAICompatExecutorClaudeStandaloneEffortFinalPayload(t *testing.T) { models := []config.PayloadModelRule{{Name: "custom-openai", Protocol: "openai"}} for _, tt := range []struct { diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go index 9f28f59b0ed..a6eb980b127 100644 --- a/internal/runtime/executor/xai_executor_execute.go +++ b/internal/runtime/executor/xai_executor_execute.go @@ -72,33 +72,53 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { data, errRead := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return resp, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) return resp, xaiStatusErr(httpResp.StatusCode, data) } data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools) var upstreamUsage helps.StreamUsageBuffer - for _, line := range bytes.Split(data, []byte("\n")) { - if !bytes.HasPrefix(line, xaiDataTag) { - continue + events := helps.NewCodexSSEScanner(ctx, bytes.NewReader(data), len(data)+1) + sawSSEData := false + var pendingEventName string + var pendingData []byte + hasPendingData := false + + dispatchSSEEvent := func(eventName string, rawEventData []byte) (cliproxyexecutor.Response, bool, error) { + rawEventData = bytes.TrimSpace(rawEventData) + if len(rawEventData) == 0 || bytes.Equal(rawEventData, []byte("[DONE]")) { + return cliproxyexecutor.Response{}, false, nil + } + sawSSEData = true + if !ojson.ValidBytes(rawEventData) { + if errTerminal, ok := xaiTerminalFailureErr(eventName, rawEventData); ok { + upstreamUsage.PublishFailure(ctx, reporter, errTerminal) + return cliproxyexecutor.Response{}, true, errTerminal + } + return cliproxyexecutor.Response{}, false, nil } - eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) + eventData := xaiNormalizeReasoningSummaryData(rawEventData) upstreamUsage.ObserveResponsesStream(eventData) + reporter.ObserveResponseModel(eventData) + if errTerminal, ok := xaiTerminalFailureErr(eventName, eventData); ok { + upstreamUsage.PublishFailure(ctx, reporter, errTerminal) + return cliproxyexecutor.Response{}, true, errTerminal + } prepared.applyPatch.RememberDispatcherEvent(eventData) eventData = namespaceRestorer.restore(eventData) if prepared.webSearchAlias != "" { @@ -106,15 +126,15 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req } eventData = responseFilter.apply(eventData) if len(eventData) == 0 { - continue + return cliproxyexecutor.Response{}, false, nil } - events, errBridge := prepared.applyPatch.Transform(eventData) + transformedEvents, errBridge := prepared.applyPatch.Transform(eventData) if errBridge != nil { errBridge = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} upstreamUsage.PublishFailure(ctx, reporter, errBridge) - return resp, errBridge + return cliproxyexecutor.Response{}, true, errBridge } - for _, eventData := range events { + for _, eventData := range transformedEvents { reporter.ObserveResponseModel(eventData) eventType := ojson.GetBytesNoCopy(eventData, "type").String() switch eventType { @@ -133,7 +153,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 { errTranslation := statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} upstreamUsage.PublishFailure(ctx, reporter, errTranslation) - return cliproxyexecutor.Response{}, errTranslation + return cliproxyexecutor.Response{}, true, errTranslation } if !upstreamUsage.Publish(ctx, reporter) { if detail, ok := helps.ParseCodexUsageWithOutputItems(completedData, outputItemsByIndex, outputItemsFallback); ok { @@ -143,17 +163,92 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req if prepared.responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } - return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil + return cliproxyexecutor.Response{ + Payload: out, + Incomplete: eventType == "response.incomplete", + Headers: httpResp.Header.Clone(), + }, true, nil + } + } + return cliproxyexecutor.Response{}, false, nil + } + + flushPendingEvent := func() (cliproxyexecutor.Response, bool, error) { + if !hasPendingData { + pendingEventName = "" + pendingData = nil + return cliproxyexecutor.Response{}, false, nil + } + eventName := pendingEventName + eventPayload := pendingData + pendingEventName = "" + pendingData = nil + hasPendingData = false + return dispatchSSEEvent(eventName, eventPayload) + } + + for events.Scan() { + line := bytes.TrimRight(events.Bytes(), "\r") + if len(bytes.TrimSpace(line)) == 0 { + if outResp, done, errDispatch := flushPendingEvent(); done { + return outResp, errDispatch + } + continue + } + if bytes.HasPrefix(line, xaiEventTag) { + nextEventName := strings.TrimSpace(string(line[len(xaiEventTag):])) + if hasPendingData && pendingEventName != "" { + if outResp, done, errDispatch := flushPendingEvent(); done { + return outResp, errDispatch + } + } + pendingEventName = nextEventName + continue + } + if !bytes.HasPrefix(line, xaiDataTag) { + continue + } + rawEventData := bytes.TrimSpace(line[len(xaiDataTag):]) + if hasPendingData { + if ojson.ValidBytes(pendingData) || bytes.Equal(pendingData, []byte("[DONE]")) { + if outResp, done, errDispatch := flushPendingEvent(); done { + return outResp, errDispatch + } + } else { + pendingData = append(append(pendingData, '\n'), rawEventData...) + continue } } + pendingData = bytes.Clone(rawEventData) + hasPendingData = true + } + if outResp, done, errDispatch := flushPendingEvent(); done { + return outResp, errDispatch } - if errFinish := prepared.applyPatch.Finish(); errFinish != nil { - errFinish = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} - upstreamUsage.PublishFailure(ctx, reporter, errFinish) - return resp, errFinish + if !sawSSEData { + if trimmed := bytes.TrimSpace(data); len(trimmed) > 0 { + if errTerminal, ok := xaiTerminalFailureErr("", trimmed); ok { + upstreamUsage.PublishFailure(ctx, reporter, errTerminal) + return resp, errTerminal + } + } + } + if prepared.applyPatch.HasPatchActivity() { + if errFinish := prepared.applyPatch.Finish(); errFinish != nil { + errFinish = xaiRequestScopedStreamError{statusErr: statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}} + upstreamUsage.PublishFailure(ctx, reporter, errFinish) + return resp, errFinish + } + } + if !sawSSEData { + errEmpty := newXAIEmptyStreamError() + upstreamUsage.PublishFailure(ctx, reporter, errEmpty) + return resp, errEmpty } - return resp, statusErr{code: http.StatusRequestTimeout, msg: "xai stream error: stream disconnected before response.completed or response.incomplete"} + errIncomplete := newXAIIncompleteStreamError() + upstreamUsage.PublishFailure(ctx, reporter, errIncomplete) + return resp, errIncomplete } func (e *XAIExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { @@ -242,11 +337,11 @@ func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxya helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return nil, nil, nil, nil, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) diff --git a/internal/runtime/executor/xai_executor_media.go b/internal/runtime/executor/xai_executor_media.go index 9ebbdb42bf7..4799fd2c96b 100644 --- a/internal/runtime/executor/xai_executor_media.go +++ b/internal/runtime/executor/xai_executor_media.go @@ -63,11 +63,11 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) @@ -145,11 +145,11 @@ func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) diff --git a/internal/runtime/executor/xai_executor_response.go b/internal/runtime/executor/xai_executor_response.go index a9afe42821b..fdb538ccb9c 100644 --- a/internal/runtime/executor/xai_executor_response.go +++ b/internal/runtime/executor/xai_executor_response.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "sort" + "strconv" "strings" "time" @@ -1313,6 +1314,269 @@ func xaiPatchCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]by // cli-chat-proxy ("Usage resets over a rolling 24-hour window"). const xaiFreeUsageExhaustedCooldown = 24 * time.Hour +const ( + xaiEmptyStreamMessage = "xai stream error: empty upstream stream" + xaiIncompleteStreamMessage = "xai stream error: stream closed before response.completed or response.incomplete" +) + +type xaiRequestScopedStreamError struct { + statusErr +} + +func (xaiRequestScopedStreamError) IsRequestScoped() bool { + return true +} + +func newXAIEmptyStreamError() xaiRequestScopedStreamError { + return xaiRequestScopedStreamError{statusErr: statusErr{ + code: http.StatusBadGateway, + msg: xaiEmptyStreamMessage, + }} +} + +func newXAIIncompleteStreamError() xaiRequestScopedStreamError { + return xaiRequestScopedStreamError{statusErr: statusErr{ + code: http.StatusBadGateway, + msg: xaiIncompleteStreamMessage, + }} +} + +func xaiTerminalFailureErr(eventName string, eventData []byte) (statusErr, bool) { + doc, ok := xaiTerminalFailureDoc(eventName, eventData) + if !ok || doc == nil { + return statusErr{}, false + } + status := xaiTerminalFailureStatus(eventData, doc) + return xaiStatusErr(status, doc.Bytes()), true +} + +func xaiTerminalFailureDoc(eventName string, eventData []byte) (*ojson.Doc, bool) { + trimmed := bytes.TrimSpace(eventData) + if len(trimmed) == 0 { + return nil, false + } + if !ojson.ValidBytes(trimmed) { + if strings.EqualFold(strings.TrimSpace(eventName), "error") { + doc := ojson.NewDocString(`{"error":{}}`) + _ = doc.Set("error.message", string(trimmed)) + return doc, true + } + return nil, false + } + + eventType := strings.TrimSpace(ojson.GetBytesNoCopy(trimmed, "type").String()) + if eventType == "" { + normalizedEventName := strings.TrimSpace(eventName) + if normalizedEventName == "error" || normalizedEventName == "response.failed" { + eventType = normalizedEventName + } else if errField := ojson.GetBytesNoCopy(trimmed, "error"); errField.Exists() && errField.Type != ojson.Null && !ojson.GetBytesNoCopy(trimmed, "output").Exists() { + if errField.Type == ojson.JSON || strings.TrimSpace(errField.String()) != "" { + eventType = "error" + } + } + } + + var doc *ojson.Doc + switch eventType { + case "error": + doc = xaiTerminalErrorDoc(trimmed, "error") + if doc == nil { + doc = xaiTerminalErrorDoc(trimmed, "body.error") + } + if doc == nil { + doc = xaiTerminalTopLevelErrorDoc(trimmed) + } + case "response.failed": + doc = xaiTerminalErrorDoc(trimmed, "response.error") + if doc == nil { + doc = xaiTerminalErrorDoc(trimmed, "error") + } + default: + return nil, false + } + + if doc == nil { + doc = ojson.NewDocString(`{"error":{"message":"upstream stream failed without error details"}}`) + } + if seq := ojson.GetBytesNoCopy(trimmed, "sequence_number"); seq.Exists() && seq.Type == ojson.Number { + _ = doc.Set("sequence_number", seq.Int()) + } + return doc, true +} + +func xaiTerminalErrorDoc(eventData []byte, path string) *ojson.Doc { + errorResult := ojson.GetBytesNoCopy(eventData, path) + if !errorResult.Exists() || errorResult.Type == ojson.Null { + return nil + } + + doc := ojson.NewDocString(`{"error":{}}`) + hasField := false + if errorResult.Type == ojson.JSON { + if errorResult.IsObject() && len(errorResult.Map()) > 0 { + _ = doc.SetRawBytes("error", []byte(errorResult.Raw)) + hasField = true + } + } else if message := strings.TrimSpace(errorResult.String()); message != "" { + _ = doc.Set("error.message", message) + hasField = true + } + + if strings.TrimSpace(doc.Get("error.code").String()) == "" { + for _, codePath := range []string{"response.error.code", "code", "body.error.code"} { + if code := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, codePath).String()); code != "" { + _ = doc.Set("error.code", code) + hasField = true + break + } + } + } + if strings.TrimSpace(doc.Get("error.type").String()) == "" { + for _, typePath := range []string{"response.error.type", "error_type", "body.error.type"} { + if errType := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, typePath).String()); errType != "" { + _ = doc.Set("error.type", errType) + hasField = true + break + } + } + } + if strings.TrimSpace(doc.Get("error.param").String()) == "" { + for _, paramPath := range []string{"response.error.param", "param", "body.error.param"} { + if param := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, paramPath).String()); param != "" { + _ = doc.Set("error.param", param) + hasField = true + break + } + } + } + if strings.TrimSpace(doc.Get("error.message").String()) == "" { + for _, msgPath := range []string{"response.error.message", "message", "body.error.message"} { + if message := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, msgPath).String()); message != "" { + _ = doc.Set("error.message", message) + hasField = true + break + } + } + } + if strings.TrimSpace(doc.Get("error.message").String()) == "" { + if code := strings.TrimSpace(doc.Get("error.code").String()); code != "" { + _ = doc.Set("error.message", code) + hasField = true + } else if errType := strings.TrimSpace(doc.Get("error.type").String()); errType != "" { + _ = doc.Set("error.message", errType) + hasField = true + } + } + if !hasField { + return nil + } + return doc +} + +func xaiTerminalTopLevelErrorDoc(eventData []byte) *ojson.Doc { + message := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, "message").String()) + code := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, "code").String()) + errorType := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, "error_type").String()) + param := strings.TrimSpace(ojson.GetBytesNoCopy(eventData, "param").String()) + if message == "" && code == "" && errorType == "" && param == "" { + return nil + } + + doc := ojson.NewDocString(`{"error":{}}`) + if message != "" { + _ = doc.Set("error.message", message) + } + if code != "" { + _ = doc.Set("error.code", code) + } + if errorType != "" { + _ = doc.Set("error.type", errorType) + } + if param != "" { + _ = doc.Set("error.param", param) + } + if strings.TrimSpace(doc.Get("error.message").String()) == "" { + if code != "" { + _ = doc.Set("error.message", code) + } else if errorType != "" { + _ = doc.Set("error.message", errorType) + } + } + return doc +} + +func xaiTerminalFailureStatus(eventData []byte, doc *ojson.Doc) int { + for _, path := range []string{"status", "status_code", "error.status", "error.status_code", "response.error.status", "response.error.status_code", "body.status", "body.error.status"} { + res := ojson.GetBytesNoCopy(eventData, path) + switch res.Type { + case ojson.Number: + if status := int(res.Int()); status >= 400 && status <= 599 { + return status + } + case ojson.String: + if status, errAtoi := strconv.Atoi(strings.TrimSpace(res.String())); errAtoi == nil && status >= 400 && status <= 599 { + return status + } + } + } + + errorType := strings.ToLower(strings.TrimSpace(doc.Get("error.type").String())) + errorCode := strings.ToLower(strings.TrimSpace(doc.Get("error.code").String())) + errorMsg := strings.ToLower(strings.TrimSpace(doc.Get("error.message").String())) + if status, errAtoi := strconv.Atoi(errorCode); errAtoi == nil && status >= 400 && status <= 599 { + return status + } + + switch { + case errorType == "authentication_error", + errorCode == "invalid_api_key", + errorCode == "unauthorized", + strings.HasPrefix(errorCode, "unauthenticated"), + strings.Contains(errorCode, "bad-credentials"), + strings.Contains(errorMsg, "access token could not be validated"): + return http.StatusUnauthorized + case errorType == "permission_error", + errorCode == "forbidden", + errorCode == "permission_denied", + errorCode == "permission-denied", + strings.HasPrefix(errorCode, "permission_denied:"), + strings.HasPrefix(errorCode, "permission-denied:"): + return http.StatusForbidden + case errorType == "not_found_error", + errorCode == "not_found", + errorCode == "model_not_found": + return http.StatusNotFound + case errorType == "rate_limit_error", + errorCode == "rate_limit_exceeded", + errorCode == "rate_limit", + strings.Contains(errorCode, "free-usage-exhausted"), + strings.Contains(errorCode, "resource_exhausted"), + strings.Contains(errorCode, "resource-exhausted"), + strings.Contains(errorMsg, "free-usage-exhausted"), + strings.Contains(errorMsg, "included free usage"): + return http.StatusTooManyRequests + case errorCode == "cyber_policy", + errorCode == "content_policy_violation", + errorType == "invalid_request_error", + errorType == "invalid_request", + errorType == "bad_request_error", + errorType == "image_generation_user_error", + errorCode == "invalid_value", + errorCode == "invalid_request_error", + errorCode == "invalid_argument", + errorCode == "invalid-argument", + errorCode == "context_length_exceeded", + errorCode == "context_too_large": + return http.StatusBadRequest + case errorType == "service_unavailable_error", + errorCode == "server_is_overloaded", + errorCode == "service_unavailable": + return http.StatusServiceUnavailable + default: + return http.StatusBadGateway + } +} + // xaiStatusErr normalizes upstream xAI error bodies for conductor behavior: // - credential invalidation (403 bad-credentials) is remapped to 401 so the // existing OAuth refresh-once-and-retry path runs instead of payment cooldown @@ -1335,7 +1599,13 @@ func xaiStatusErr(code int, body []byte) statusErr { return err } codeStr := strings.ToLower(ojson.GetBytesNoCopy(body, "code").String()) - msg := strings.ToLower(ojson.GetBytesNoCopy(body, "error").String()) + if codeStr == "" { + codeStr = strings.ToLower(ojson.GetBytesNoCopy(body, "error.code").String()) + } + msg := strings.ToLower(ojson.GetBytesNoCopy(body, "error.message").String()) + if msg == "" { + msg = strings.ToLower(ojson.GetBytesNoCopy(body, "error").String()) + } if msg == "" { msg = strings.ToLower(string(body)) } diff --git a/internal/runtime/executor/xai_executor_speech.go b/internal/runtime/executor/xai_executor_speech.go index 93b678af64d..a185839d8b0 100644 --- a/internal/runtime/executor/xai_executor_speech.go +++ b/internal/runtime/executor/xai_executor_speech.go @@ -119,11 +119,11 @@ func (e *XAIExecutor) executeSpeech(ctx context.Context, auth *cliproxyauth.Auth helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) data, err := io.ReadAll(httpResp.Body) + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go index 5c7c55bf9c8..a517643636e 100644 --- a/internal/runtime/executor/xai_executor_stream.go +++ b/internal/runtime/executor/xai_executor_stream.go @@ -65,11 +65,11 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("xai executor: close response body error: %v", errClose) } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) if errRead != nil { helps.RecordAPIResponseError(ctx, e.cfg, errRead) return nil, errRead } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) return nil, xaiStatusErr(httpResp.StatusCode, data) } diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 581b8bf47ba..8598bbc8a39 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -13,6 +13,7 @@ import ( "net/http/httptest" "runtime" "strings" + "sync/atomic" "testing" "time" @@ -23,8 +24,10 @@ import ( xaiauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/xai" internalcache "github.com/router-for-me/CLIProxyAPI/v8/internal/cache" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" _ "github.com/router-for-me/CLIProxyAPI/v8/internal/translator" + sdkhandlers "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" @@ -8110,3 +8113,551 @@ func TestXAIWebsocketUsagePreservesFilteredSearch(t *testing.T) { t.Fatal("usage not published") } } + +func TestXAIExecutorExecuteReturnsTerminalStreamFailures(t *testing.T) { + t.Parallel() + + want24h := 24 * time.Hour + tests := []struct { + name string + sseBody string + contentType string + wantStatus int + wantMsg string + wantCode string + wantType string + wantRetryAfter *time.Duration + }{ + { + name: "error event with invalid_request_error", + sseBody: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"model\":\"grok-4.7-build\"}}\n\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"code\":\"invalid_value\",\"message\":\"Invalid input.\"},\"sequence_number\":2}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusBadRequest, + wantMsg: "Invalid input.", + wantCode: "invalid_value", + wantType: "invalid_request_error", + }, + { + name: "error event with bad-credentials remaps to 401", + sseBody: "data: {\"type\":\"error\",\"status\":403,\"error\":{\"code\":\"unauthenticated:bad-credentials\",\"message\":\"The OAuth2 access token could not be validated.\"}}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusUnauthorized, + wantMsg: "The OAuth2 access token could not be validated.", + wantCode: "unauthenticated:bad-credentials", + }, + { + name: "error event with free-usage-exhausted maps to 429 with 24h retry-after", + sseBody: "data: {\"type\":\"error\",\"error\":{\"code\":\"subscription:free-usage-exhausted\",\"message\":\"You've used all the included free usage for model grok-4.7 for now.\"}}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusTooManyRequests, + wantMsg: "You've used all the included free usage for model grok-4.7 for now.", + wantCode: "subscription:free-usage-exhausted", + wantRetryAfter: &want24h, + }, + { + name: "event error header with flat xai refusal payload", + sseBody: "event: error\ndata: {\"code\":\"permission-denied\",\"error\":\"I'm sorry, I can't help with that request.\"}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusForbidden, + wantMsg: "I'm sorry, I can't help with that request.", + wantCode: "permission-denied", + }, + { + name: "data line before event error header", + sseBody: "data: {\"code\":\"invalid_argument\",\"message\":\"Invalid input.\"}\nevent: error\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusBadRequest, + wantMsg: "Invalid input.", + wantCode: "invalid_argument", + }, + { + name: "response.failed event with server_error", + sseBody: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"model\":\"grok-4.7-build\"}}\n\ndata: {\"type\":\"response.failed\",\"sequence_number\":3,\"response\":{\"id\":\"resp_1\",\"status\":\"failed\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"Upstream generation failed.\"}}}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusBadGateway, + wantMsg: "Upstream generation failed.", + wantCode: "server_error", + wantType: "server_error", + }, + { + name: "multiline sse data error event", + sseBody: "event: error\ndata: {\"type\":\"error\",\ndata: \"error\":{\"type\":\"invalid_request_error\",\"code\":\"invalid_argument\",\"message\":\"Multiline error.\"}}\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusBadRequest, + wantMsg: "Multiline error.", + wantCode: "invalid_argument", + wantType: "invalid_request_error", + }, + { + name: "plain text data under event error header", + sseBody: "event: error\ndata: upstream gateway error\n\n", + contentType: "text/event-stream", + wantStatus: http.StatusBadGateway, + wantMsg: "upstream gateway error", + }, + { + name: "non-sse json error body on 200", + sseBody: `{"code":"permission-denied","error":"Request blocked."}`, + contentType: "application/json", + wantStatus: http.StatusForbidden, + wantMsg: "Request blocked.", + wantCode: "permission-denied", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", tc.contentType) + _, _ = io.WriteString(w, tc.sseBody) + })) + defer server.Close() + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + payload := []byte(`{"model":"grok-4.7","stream":false,"messages":[{"role":"user","content":"Reply with one word: ready."}]}`) + _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + OriginalRequest: payload, + Stream: false, + }) + if err == nil { + t.Fatal("Execute() error = nil, want terminal stream error") + } + + var statusCoder cliproxyexecutor.StatusError + if !errors.As(err, &statusCoder) || statusCoder.StatusCode() != tc.wantStatus { + t.Fatalf("status = %v, want %d; err=%v", statusCoder, tc.wantStatus, err) + } + var scopedErr cliproxyexecutor.RequestScopedError + if errors.As(err, &scopedErr) && scopedErr.IsRequestScoped() { + t.Fatalf("explicit upstream terminal error must not set IsRequestScoped() = true: %T (%v)", err, err) + } + + clientBody := sdkhandlers.BuildErrorResponseBodyWithError(statusCoder.StatusCode(), err.Error(), err) + if got := ojson.GetBytes(clientBody, "error.message").String(); got != tc.wantMsg { + t.Fatalf("error.message = %q, want %q; body=%s", got, tc.wantMsg, clientBody) + } + if got := ojson.GetBytes(clientBody, "error.code").String(); got != tc.wantCode { + t.Fatalf("error.code = %q, want %q; body=%s", got, tc.wantCode, clientBody) + } + if tc.wantType != "" { + if got := ojson.GetBytes(clientBody, "error.type").String(); got != tc.wantType { + t.Fatalf("error.type = %q, want %q; body=%s", got, tc.wantType, clientBody) + } + } + + type retryAfterProvider interface { + RetryAfter() *time.Duration + } + var raProvider retryAfterProvider + if tc.wantRetryAfter != nil { + if !errors.As(err, &raProvider) || raProvider.RetryAfter() == nil || *raProvider.RetryAfter() != *tc.wantRetryAfter { + t.Fatalf("RetryAfter = %v, want %v; err=%v", raProvider, *tc.wantRetryAfter, err) + } + } + }) + } +} + +func TestXAIExecutorExecuteEmptyStreamReturns502(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + body string + }{ + {name: "zero byte body", body: ""}, + {name: "sse comment only", body: ": keep-alive\n\n"}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, tc.body) + })) + defer server.Close() + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + payload := []byte(`{"model":"grok-4.7","stream":false,"messages":[{"role":"user","content":"hello"}]}`) + _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + OriginalRequest: payload, + Stream: false, + }) + if err == nil { + t.Fatal("Execute() error = nil, want 502 empty stream error") + } + + var statusCoder cliproxyexecutor.StatusError + if !errors.As(err, &statusCoder) || statusCoder.StatusCode() != http.StatusBadGateway { + t.Fatalf("status = %v, want %d; err=%v", statusCoder, http.StatusBadGateway, err) + } + lower := strings.ToLower(err.Error()) + if !strings.Contains(lower, "empty") || !strings.Contains(lower, "stream") { + t.Fatalf("error = %q, want empty stream message", err.Error()) + } + if strings.Contains(lower, "timeout") { + t.Fatalf("error = %q, must not claim timeout", err.Error()) + } + var scopedErr cliproxyexecutor.RequestScopedError + if !errors.As(err, &scopedErr) || !scopedErr.IsRequestScoped() { + t.Fatalf("empty stream error %T (%v) must implement IsRequestScoped() = true", err, err) + } + clientBody := sdkhandlers.BuildErrorResponseBodyWithError(statusCoder.StatusCode(), err.Error(), err) + if got := ojson.GetBytes(clientBody, "error.code").String(); got == "request_timeout" { + t.Fatalf("client body error.code = %q, must not be request_timeout; body=%s", got, clientBody) + } + }) + } +} + +func TestXAIExecutorExecuteIncompleteStreamDoesNotCooldownCredential(t *testing.T) { + var upstreamCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + call := upstreamCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if call == 1 { + _, _ = io.WriteString(w, "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"model\":\"grok-4.7-build\"}}\n\n") + return + } + _, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_2\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.7-build\",\"output\":[{\"id\":\"msg_2\",\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"text\":\"ready\"}]}],\"usage\":{\"input_tokens\":5,\"output_tokens\":1,\"total_tokens\":6}}}\n\n") + })) + defer server.Close() + + cfg := &config.Config{ + RequestRetry: 3, + MaxRetryInterval: 30, + } + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + manager.SetRetryConfig(3, 30*time.Second, 0) + manager.RegisterExecutor(NewXAIExecutor(cfg)) + + authID := "xai-single-cred-issue-7419" + auth := &cliproxyauth.Auth{ + ID: authID, + Provider: "xai", + Status: cliproxyauth.StatusActive, + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("Register() error = %v", errReg) + } + registry.GetGlobalRegistry().RegisterClient(authID, "xai", []*registry.ModelInfo{{ID: "grok-4.7"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + + payload := []byte(`{"model":"grok-4.7","stream":false,"messages":[{"role":"user","content":"hello"}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + OriginalRequest: payload, + Stream: false, + } + _, err := manager.Execute(context.Background(), []string{"xai"}, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, opts) + if err == nil { + t.Fatal("manager.Execute() error = nil, want incomplete stream error") + } + if gotCalls := upstreamCalls.Load(); gotCalls != 1 { + t.Fatalf("upstream calls on request-scoped incomplete stream = %d, want 1", gotCalls) + } + + var statusCoder cliproxyexecutor.StatusError + if !errors.As(err, &statusCoder) || statusCoder.StatusCode() != http.StatusBadGateway { + t.Fatalf("status = %v, want %d; err=%v", statusCoder, http.StatusBadGateway, err) + } + if strings.Contains(strings.ToLower(err.Error()), "timeout") { + t.Fatalf("error = %q, must not claim timeout", err.Error()) + } + var scopedErr cliproxyexecutor.RequestScopedError + if !errors.As(err, &scopedErr) || !scopedErr.IsRequestScoped() { + t.Fatalf("incomplete stream error %T (%v) must implement IsRequestScoped() = true", err, err) + } + clientBody := sdkhandlers.BuildErrorResponseBodyWithError(statusCoder.StatusCode(), err.Error(), err) + if got := ojson.GetBytes(clientBody, "error.code").String(); got == "request_timeout" { + t.Fatalf("client body error.code = %q, must not be request_timeout; body=%s", got, clientBody) + } + + stored, ok := manager.GetByID(authID) + if !ok || stored == nil { + t.Fatalf("GetByID(%q) missing", authID) + } + if state := stored.ModelStates["grok-4.7"]; state != nil && (state.Unavailable || !state.NextRetryAfter.IsZero()) { + t.Fatalf("model state cooled down after incomplete stream: Unavailable=%v NextRetryAfter=%v", state.Unavailable, state.NextRetryAfter) + } + + // Verify the single credential remains immediately selectable without a 60s transient cooldown. + resp, errNext := manager.Execute(context.Background(), []string{"xai"}, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, opts) + if errNext != nil { + t.Fatalf("subsequent manager.Execute() failed on uncooled credential: %v", errNext) + } + if got := ojson.GetBytes(resp.Payload, "choices.0.message.content").String(); got != "ready" { + t.Fatalf("subsequent response content = %q, want %q; payload=%s", got, "ready", resp.Payload) + } + if gotCalls := upstreamCalls.Load(); gotCalls != 2 { + t.Fatalf("total upstream calls = %d, want 2", gotCalls) + } +} + +func TestXAIExecutorExecuteTruncatedCompletedFrameReturnsIncompleteStream(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"status\":\"completed\",\"output\":[{\"id\":\"msg_1\"") + })) + defer server.Close() + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + payload := []byte(`{"model":"grok-4.7","input":"hello"}`) + _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: payload, + Stream: false, + }) + if err == nil { + t.Fatal("Execute() error = nil, want incomplete stream error on truncated response.completed") + } + var scopedErr cliproxyexecutor.RequestScopedError + if !errors.As(err, &scopedErr) || !scopedErr.IsRequestScoped() { + t.Fatalf("truncated stream error %T (%v) must implement IsRequestScoped() = true", err, err) + } + if strings.Contains(err.Error(), "apply_patch") { + t.Fatalf("truncated response.completed misclassified as apply_patch error: %v", err) + } +} + +func TestXAIExecutorExecuteCustomApplyPatchEmptyAndIncompleteStreamsDoNotCooldown(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + firstStream string + wantStatus int + wantMessage string + }{ + { + name: "empty_stream_with_custom_apply_patch", + firstStream: "", + wantStatus: http.StatusBadGateway, + wantMessage: xaiEmptyStreamMessage, + }, + { + name: "response_created_eof_with_custom_apply_patch", + firstStream: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"status\":\"in_progress\"}}\n\n", + wantStatus: http.StatusBadGateway, + wantMessage: xaiIncompleteStreamMessage, + }, + { + name: "partial_apply_patch_call_eof_preserves_patch_error_without_cooldown", + firstStream: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"status\":\"in_progress\"}}\n\n" + + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"a\",\"call_id\":\"c\",\"name\":\"apply_patch\",\"arguments\":\"\"}}\n\n" + + "data: {\"type\":\"response.function_call_arguments.done\",\"output_index\":0,\"item_id\":\"a\",\"call_id\":\"c\",\"arguments\":\"{\\\"input\\\":\\\"valid patch\\\"}\"}\n\n", + wantStatus: http.StatusBadGateway, + wantMessage: "Invalid apply_patch tool arguments received from upstream.", + }, + { + name: "raw_json_error_with_custom_apply_patch", + firstStream: `{"error":{"message":"Upstream overloaded","type":"server_error","code":"server_is_overloaded"}}`, + wantStatus: http.StatusServiceUnavailable, + wantMessage: "Upstream overloaded", + }, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var upstreamCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + call := upstreamCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if call == 1 { + _, _ = io.WriteString(w, tc.firstStream) + return + } + _, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_ok\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.7\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"recovered\"}]}],\"usage\":{\"input_tokens\":3,\"output_tokens\":2,\"total_tokens\":5}}}\n\n") + })) + defer server.Close() + + manager := cliproxyauth.NewManager(nil, &cliproxyauth.RoundRobinSelector{}, nil) + manager.RegisterExecutor(NewXAIExecutor(&config.Config{})) + + authID := "xai-patch-cooldown-" + tc.name + auth := &cliproxyauth.Auth{ + ID: authID, + Provider: "xai", + Status: cliproxyauth.StatusActive, + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("manager.Register() error = %v", errReg) + } + registry.GetGlobalRegistry().RegisterClient(authID, "xai", []*registry.ModelInfo{{ID: "grok-4.7"}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + payload := []byte(`{"model":"grok-4.7","input":"hello","tools":[{"type":"custom","name":"apply_patch"}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: payload, + Stream: false, + } + + _, errFirst := manager.Execute(context.Background(), []string{"xai"}, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, opts) + if errFirst == nil { + t.Fatal("first manager.Execute() error = nil, want failure") + } + var statusErrVal cliproxyexecutor.StatusError + if !errors.As(errFirst, &statusErrVal) || statusErrVal.StatusCode() != tc.wantStatus { + t.Fatalf("first error status = %v (%T), want %d", errFirst, errFirst, tc.wantStatus) + } + if !strings.Contains(errFirst.Error(), tc.wantMessage) { + t.Fatalf("first error = %q, want substring %q", errFirst.Error(), tc.wantMessage) + } + + if tc.name == "raw_json_error_with_custom_apply_patch" { + return + } + + var scopedErr cliproxyexecutor.RequestScopedError + if !errors.As(errFirst, &scopedErr) || !scopedErr.IsRequestScoped() { + t.Fatalf("first error %T (%v) must implement IsRequestScoped() = true", errFirst, errFirst) + } + + resp, errSecond := manager.Execute(context.Background(), []string{"xai"}, cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: payload, + }, opts) + if errSecond != nil { + t.Fatalf("subsequent manager.Execute() failed on uncooled credential: %v", errSecond) + } + if got := ojson.GetBytes(resp.Payload, "output.0.content.0.text").String(); got != "recovered" { + t.Fatalf("subsequent response text = %q, want %q; payload=%s", got, "recovered", resp.Payload) + } + }) + } +} + +func TestXAIExecutorPreservesPartialUpstreamBodyOnReadErrorWhenRequestLogDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + + for _, tc := range []struct { + name string + statusLine string + partial string + }{ + { + name: "200 stream truncated mid-body", + statusLine: "HTTP/1.1 200 OK", + partial: `data: {"type":"response.created","response":{"id":"resp_partial_200"}}`, + }, + { + name: "502 error body truncated mid-body", + statusLine: "HTTP/1.1 502 Bad Gateway", + partial: `{"error":{"message":"partial upstream 502 detail"`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hijacker, ok := w.(http.Hijacker) + if !ok { + t.Error("response writer does not support hijacking") + return + } + conn, buf, errHijack := hijacker.Hijack() + if errHijack != nil { + t.Errorf("hijack: %v", errHijack) + return + } + rawResponse := fmt.Sprintf("%s\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\n\r\n%s", tc.statusLine, tc.partial) + _, _ = buf.WriteString(rawResponse) + _ = buf.Flush() + _ = conn.Close() + })) + defer server.Close() + + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "test-token"}, + } + _, errExec := exec.Execute( + ctx, + auth, + cliproxyexecutor.Request{ + Model: "grok-4.7", + Payload: []byte(`{"model":"grok-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}`), + }, + cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + }, + ) + if errExec == nil { + t.Fatal("expected read error from truncated upstream connection") + } + + value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey) + if !exists { + t.Fatal("deferred API response was not captured") + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil { + t.Fatalf("deferred API response = %#v, want logging.DeferredAPIResponse", value) + } + rendered := string(deferred.RenderAPIResponse()) + if !strings.Contains(rendered, tc.partial) { + t.Fatalf("deferred API response dropped partial body %q on read error:\n%s", tc.partial, rendered) + } + if !strings.Contains(rendered, "Error: unexpected EOF") { + t.Fatalf("deferred API response missing read error:\n%s", rendered) + } + }) + } +} diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index a86beff28ca..09943aaac70 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -580,15 +580,11 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox upstreamHeaders = respHS.Header.Clone() } if errDial != nil { - bodyErr := websocketHandshakeBody(respHS) - if respHS != nil { - helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) - } + bodyErr := recordWebsocketDialFailure(ctx, e.cfg, wsReqLog, respHS, errDial) if respHS != nil && respHS.StatusCode > 0 { unlockStreamSession() return nil, xaiStatusErr(respHS.StatusCode, bodyErr) } - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) unlockStreamSession() return nil, errDial } @@ -631,9 +627,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { - bodyErrRetry := websocketHandshakeBody(respHSRetry) - closeHTTPResponseBody(respHSRetry, "xai websockets executor: close handshake response body error") - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) + bodyErrRetry := recordWebsocketDialRetryFailure(ctx, e.cfg, wsReqLog, respHSRetry, errDialRetry) sess.clearActive(conn, readCh) sess.reqMu.Unlock() if respHSRetry != nil && respHSRetry.StatusCode > 0 { diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 5e3a11950b4..e65295c2a10 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -10,12 +10,15 @@ import ( "net/http/httptest" "strings" "sync" + "sync/atomic" "testing" "time" "github.com/MoeForge/moego/ojson" + "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" _ "github.com/router-for-me/CLIProxyAPI/v8/internal/translator" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -2645,3 +2648,242 @@ func TestApplyXAIWebsocketHeaders_UserAgent(t *testing.T) { } }) } + +func TestXAIWebsocketReconnectUpgradeRejectionRecordedWhenRequestLogDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var dials atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if dials.Add(1) == 1 { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + return + } + defer func() { _ = conn.Close() }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "9") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"xai websocket reconnect rate limited"}}`)) + })) + defer server.Close() + + var writes atomic.Int32 + testWebsocketWritePayloadHook = func(conn *websocket.Conn) { + if writes.Add(1) == 1 { + _ = conn.SetWriteDeadline(time.Now().Add(-time.Second)) + } + } + defer func() { testWebsocketWritePayloadHook = nil }() + + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := context.WithValue(context.Background(), "gin", ginCtx) + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ + ID: "xai-ws-reconnect", + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "xai-secret-key", + "websockets": "true", + }, + } + _, errExec := exec.ExecuteStream( + ctx, + auth, + cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"reconnect"}]}`), + }, + cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "sess-xai-reconnect", + }, + }, + ) + if errExec == nil { + t.Fatal("expected ExecuteStream reconnect rejection to fail") + } + + value, exists := ginCtx.Get(logging.DeferredAPIResponseContextKey) + if !exists { + t.Fatal("deferred API response was not recorded on websocket reconnect rejection") + } + deferred, ok := value.(logging.DeferredAPIResponse) + if !ok || deferred == nil { + t.Fatalf("deferred API response = %#v, want logging.DeferredAPIResponse", value) + } + rendered := string(deferred.RenderAPIResponse()) + for _, want := range []string{ + "=== API RESPONSE 1 ===", + "Status: 101", + "Error: send:", + "=== API RESPONSE 2 ===", + "Status: 429", + "Retry-After: 9", + `{"error":{"message":"xai websocket reconnect rate limited"}}`, + "Error: dial_retry:", + } { + if !strings.Contains(rendered, want) { + t.Fatalf("deferred API response missing %q:\n%s", want, rendered) + } + } + + t.Run("reconnect transport failure without HTTP response", func(t *testing.T) { + var transportDials atomic.Int32 + var transportServer *httptest.Server + transportServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if transportDials.Add(1) == 1 { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + return + } + defer func() { _ = conn.Close() }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + } + hijacker, ok := w.(http.Hijacker) + if ok { + if rawConn, _, errHijack := hijacker.Hijack(); errHijack == nil { + _ = rawConn.Close() + } + } + })) + defer transportServer.Close() + + var transportWrites atomic.Int32 + testWebsocketWritePayloadHook = func(conn *websocket.Conn) { + if transportWrites.Add(1) == 1 { + _ = conn.SetWriteDeadline(time.Now().Add(-time.Second)) + } + } + defer func() { testWebsocketWritePayloadHook = nil }() + + subGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + subCtx := context.WithValue(context.Background(), "gin", subGinCtx) + subExec := NewXAIWebsocketsExecutor(&config.Config{}) + subExec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + _, errSub := subExec.ExecuteStream( + subCtx, + &cliproxyauth.Auth{ + ID: "xai-ws-transport", + Provider: "xai", + Attributes: map[string]string{ + "base_url": transportServer.URL, + "api_key": "xai-secret-key", + "websockets": "true", + }, + }, + cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"reconnect"}]}`), + }, + cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "sess-xai-transport", + }, + }, + ) + if errSub == nil { + t.Fatal("expected transport reconnect failure") + } + subVal, ok := subGinCtx.Get(logging.DeferredAPIResponseContextKey) + if !ok { + t.Fatal("deferred API response missing on transport reconnect failure") + } + subRendered := string(subVal.(logging.DeferredAPIResponse).RenderAPIResponse()) + for _, want := range []string{ + "=== API RESPONSE 1 ===", + "Error: send:", + "=== API RESPONSE 2 ===", + "Error: dial_retry:", + } { + if !strings.Contains(subRendered, want) { + t.Fatalf("transport reconnect response missing %q:\n%s", want, subRendered) + } + } + }) + + t.Run("initial handshake failure with HTTP response records dial error", func(t *testing.T) { + for _, invalid101 := range []bool{false, true} { + caseName := "status_403" + if invalid101 { + caseName = "invalid_101" + } + t.Run(caseName, func(t *testing.T) { + initServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if invalid101 { + hijacker, ok := w.(http.Hijacker) + if !ok { + t.Error("hijacker unsupported") + return + } + rawConn, rawBuf, errHijack := hijacker.Hijack() + if errHijack != nil { + t.Errorf("hijack: %v", errHijack) + return + } + _, _ = rawBuf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: invalid\r\n\r\n") + _ = rawBuf.Flush() + _ = rawConn.Close() + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error":{"message":"initial xai ws forbidden"}}`)) + })) + defer initServer.Close() + + initGinCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + initCtx := context.WithValue(context.Background(), "gin", initGinCtx) + initExec := NewXAIWebsocketsExecutor(&config.Config{}) + initExec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + _, errInit := initExec.ExecuteStream( + initCtx, + &cliproxyauth.Auth{ + ID: "xai-ws-init", + Provider: "xai", + Attributes: map[string]string{ + "base_url": initServer.URL, + "api_key": "xai-secret-key", + "websockets": "true", + }, + }, + cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"init"}]}`), + }, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}, + ) + if errInit == nil { + t.Fatal("expected initial handshake failure") + } + initVal, ok := initGinCtx.Get(logging.DeferredAPIResponseContextKey) + if !ok { + t.Fatal("deferred API response missing on initial handshake failure") + } + initRendered := string(initVal.(logging.DeferredAPIResponse).RenderAPIResponse()) + wantStatus := "Status: 403" + if invalid101 { + wantStatus = "Status: 101" + } + for _, want := range []string{wantStatus, "Error: dial:"} { + if !strings.Contains(initRendered, want) { + t.Fatalf("initial handshake response missing %q:\n%s", want, initRendered) + } + } + }) + } + }) +} diff --git a/internal/thinking/summary.go b/internal/thinking/summary.go index 78d3862afd9..287dc80abd6 100644 --- a/internal/thinking/summary.go +++ b/internal/thinking/summary.go @@ -324,6 +324,9 @@ func applySummaryConfigForProviderDoc(doc *ojson.Doc, docValid bool, format, mod value := "omitted" if enabled { value = "summarized" + if strings.EqualFold(strings.TrimSpace(doc.Get("thinking.display").String()), "updates") { + value = "updates" + } } if doc.Set("thinking.display", value) == nil { modified = true @@ -554,7 +557,7 @@ func claudeSummaryConfigFromGetter(get func(string) ojson.Result, path string) ( return SummaryConfig{}, false } switch strings.ToLower(strings.TrimSpace(value.String())) { - case "summarized": + case "summarized", "updates": return SummaryConfig{Mode: SummaryEnabled, Detail: "auto"}, true case "omitted": return SummaryConfig{Mode: SummaryDisabled}, true diff --git a/internal/thinking/summary_test.go b/internal/thinking/summary_test.go index ba18c9d45ac..3c1deb6bac4 100644 --- a/internal/thinking/summary_test.go +++ b/internal/thinking/summary_test.go @@ -37,13 +37,17 @@ func TestExtractSummaryConfig(t *testing.T) { {name: "responses boolean summary invalid", format: "openai-response", body: `{"reasoning":{"summary":true}}`, wantMode: SummaryUnspecified}, {name: "responses deprecated generate summary", format: "openai-response", body: `{"reasoning":{"generate_summary":"detailed"}}`, wantMode: SummaryEnabled, wantDetail: "detailed"}, {name: "claude summarized", format: "claude", body: `{"thinking":{"type":"adaptive","display":"summarized"}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, + {name: "claude updates", format: "claude", body: `{"thinking":{"type":"adaptive","display":"updates"}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, {name: "claude omitted", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":2048,"display":"omitted"}}`, wantMode: SummaryDisabled}, {name: "claude display without type is invalid", format: "claude", body: `{"thinking":{"display":"summarized"}}`, wantMode: SummaryUnspecified}, + {name: "claude updates without type is invalid", format: "claude", body: `{"thinking":{"display":"updates"}}`, wantMode: SummaryUnspecified}, {name: "claude display with auto type is invalid", format: "claude", body: `{"thinking":{"type":"auto","display":"summarized"}}`, wantMode: SummaryUnspecified}, // ApplySummaryConfig runs before ApplyThinking fills budget_tokens, so an // absent budget must not be read as inactive thinking. {name: "claude enabled display without budget is valid", format: "claude", body: `{"thinking":{"type":"enabled","display":"summarized"}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, + {name: "claude enabled updates without budget is valid", format: "claude", body: `{"thinking":{"type":"enabled","display":"updates"}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, {name: "claude enabled display with zero budget is invalid", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":0,"display":"summarized"}}`, wantMode: SummaryUnspecified}, + {name: "claude enabled updates with zero budget is invalid", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":0,"display":"updates"}}`, wantMode: SummaryUnspecified}, {name: "claude auto compatibility budget summarized", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":-1,"display":"summarized"}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, {name: "claude auto compatibility budget omitted", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":-1,"display":"omitted"}}`, wantMode: SummaryDisabled}, {name: "gemini include true", format: "gemini", body: `{"generationConfig":{"thinkingConfig":{"includeThoughts":true}}}`, wantMode: SummaryEnabled, wantDetail: "auto"}, @@ -107,7 +111,9 @@ func TestApplySummaryConfig(t *testing.T) { {name: "chat disabled updates legacy include_reasoning when present", format: "openai", body: `{"reasoning_effort":"high","include_reasoning":true}`, config: SummaryConfig{Mode: SummaryDisabled}, path: "include_reasoning", want: "false"}, {name: "chat disabled invents no openrouter field", format: "openai", body: `{"reasoning_effort":"high"}`, config: SummaryConfig{Mode: SummaryDisabled}, path: "reasoning", want: ""}, {name: "claude enabled", format: "claude", body: `{"thinking":{"type":"adaptive"}}`, config: SummaryConfig{Mode: SummaryEnabled}, path: "thinking.display", want: "summarized"}, + {name: "claude enabled preserves existing updates", format: "claude", body: `{"thinking":{"type":"adaptive","display":"updates"}}`, config: SummaryConfig{Mode: SummaryEnabled}, path: "thinking.display", want: "updates"}, {name: "claude disabled", format: "claude", body: `{"thinking":{"type":"enabled","budget_tokens":2048}}`, config: SummaryConfig{Mode: SummaryDisabled}, path: "thinking.display", want: "omitted"}, + {name: "claude disabled overrides existing updates", format: "claude", body: `{"thinking":{"type":"adaptive","display":"updates"}}`, config: SummaryConfig{Mode: SummaryDisabled}, path: "thinking.display", want: "omitted"}, {name: "gemini enabled", format: "gemini", config: SummaryConfig{Mode: SummaryEnabled}, path: "generationConfig.thinkingConfig.includeThoughts", want: "true"}, {name: "gemini disabled", format: "gemini", config: SummaryConfig{Mode: SummaryDisabled}, path: "generationConfig.thinkingConfig.includeThoughts", want: "false"}, {name: "antigravity enabled", format: "antigravity", config: SummaryConfig{Mode: SummaryEnabled}, path: "request.generationConfig.thinkingConfig.includeThoughts", want: "true"}, diff --git a/internal/translator/claude/openai/chat-completions/claude_openai_response.go b/internal/translator/claude/openai/chat-completions/claude_openai_response.go index f9a6afb4718..eafdb41dd4e 100644 --- a/internal/translator/claude/openai/chat-completions/claude_openai_response.go +++ b/internal/translator/claude/openai/chat-completions/claude_openai_response.go @@ -25,6 +25,8 @@ var ( type ConvertAnthropicResponseToOpenAIParams struct { CreatedAt int64 ResponseID string + Model string + Fallbacks translatorcommon.ClaudeFallbackChain FinishReason string Usage claudeUsageTokens TrailingUsageSent bool @@ -93,6 +95,13 @@ func (u claudeUsageTokens) OpenAIUsage() (promptTokens, completionTokens, totalT return promptTokens, completionTokens, totalTokens, cachedTokens, cachedCreationTokens } +func effectiveChatResponseModel(servedModel, requestedModel string) string { + if servedModel != "" { + return servedModel + } + return requestedModel +} + // ConvertClaudeResponseToOpenAI converts Claude Code streaming response format to OpenAI Chat Completions format. // This function processes various Claude Code event types and transforms them into OpenAI-compatible JSON responses. // It handles text content, tool calls, reasoning content, and usage metadata, outputting responses that match @@ -126,47 +135,88 @@ func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, original // Base OpenAI streaming response template templateDoc := ojson.NewDoc([]byte(`{"id":"","object":"chat.completion.chunk","created":0,"model":"","choices":[{"index":0,"delta":{},"finish_reason":null}]}`)) - // Set model - if modelName != "" { - _ = templateDoc.Set("model", modelName) + p := (*param).(*ConvertAnthropicResponseToOpenAIParams) + + // Set model (prefer authoritative upstream serving model when known) + if chunkModel := effectiveChatResponseModel(p.Model, modelName); chunkModel != "" { + _ = templateDoc.Set("model", chunkModel) } // Set response ID and creation time - if (*param).(*ConvertAnthropicResponseToOpenAIParams).ResponseID != "" { - _ = templateDoc.Set("id", (*param).(*ConvertAnthropicResponseToOpenAIParams).ResponseID) + if p.ResponseID != "" { + _ = templateDoc.Set("id", p.ResponseID) } - if (*param).(*ConvertAnthropicResponseToOpenAIParams).CreatedAt > 0 { - _ = templateDoc.Set("created", (*param).(*ConvertAnthropicResponseToOpenAIParams).CreatedAt) + if p.CreatedAt > 0 { + _ = templateDoc.Set("created", p.CreatedAt) } switch eventType { case "message_start": // Initialize response with message metadata when a new message begins if message := root.Get("message"); message.Exists() { - (*param).(*ConvertAnthropicResponseToOpenAIParams).ResponseID = strings.Clone(message.Get("id").String()) - (*param).(*ConvertAnthropicResponseToOpenAIParams).CreatedAt = time.Now().Unix() + p.ResponseID = strings.Clone(message.Get("id").String()) + p.CreatedAt = time.Now().Unix() + p.Fallbacks.Reset() + upstreamModel := strings.TrimSpace(message.Get("model").String()) + if upstreamModel != "" { + p.Model = strings.Clone(upstreamModel) + } - _ = templateDoc.Set("id", (*param).(*ConvertAnthropicResponseToOpenAIParams).ResponseID) - _ = templateDoc.Set("model", modelName) - _ = templateDoc.Set("created", (*param).(*ConvertAnthropicResponseToOpenAIParams).CreatedAt) + _ = templateDoc.Set("id", p.ResponseID) + if chunkModel := effectiveChatResponseModel(p.Model, modelName); chunkModel != "" { + _ = templateDoc.Set("model", chunkModel) + } + _ = templateDoc.Set("created", p.CreatedAt) // Set initial role to assistant for the response _ = templateDoc.Set("choices.0.delta.role", "assistant") + if fromModel, toModel, ok := translatorcommon.ClaudePreTokenFallbackModels(originalRequestRawJSON, requestRawJSON, modelName, upstreamModel); ok { + if fbBytes := p.Fallbacks.RecordPreToken(fromModel, toModel, "", nil); len(fbBytes) > 0 { + _ = templateDoc.SetRawBytes("choices.0.delta.fallback", fbBytes) + } + } // Initialize tool calls accumulator for tracking tool call progress - if (*param).(*ConvertAnthropicResponseToOpenAIParams).ToolCallsAccumulator == nil { - (*param).(*ConvertAnthropicResponseToOpenAIParams).ToolCallsAccumulator = make(map[int]*ToolCallAccumulator) + if p.ToolCallsAccumulator == nil { + p.ToolCallsAccumulator = make(map[int]*ToolCallAccumulator) } - (*param).(*ConvertAnthropicResponseToOpenAIParams).NextToolCallIndex = 0 - (*param).(*ConvertAnthropicResponseToOpenAIParams).Usage.Merge(message.Get("usage")) + p.NextToolCallIndex = 0 + p.Usage.Merge(message.Get("usage")) } return [][]byte{templateDoc.Bytes()} case "content_block_start": - // Start of a content block (text, tool use, or reasoning) + // Start of a content block (text, tool use, fallback boundary, or reasoning) if contentBlock := root.Get("content_block"); contentBlock.Exists() { blockType := contentBlock.Get("type").String() + if blockType == "fallback" { + toModel := strings.TrimSpace(contentBlock.Get("to.model").String()) + fromModel := strings.TrimSpace(contentBlock.Get("from.model").String()) + if toModel != "" { + p.Model = strings.Clone(toModel) + _ = templateDoc.Set("model", p.Model) + } + fbBytes := p.Fallbacks.RecordBlock( + contentBlock, + effectiveChatResponseModel(fromModel, modelName), + effectiveChatResponseModel(toModel, p.Model), + "", + nil, + ) + if len(fbBytes) > 0 { + _ = templateDoc.SetRawBytes("choices.0.delta.fallback", fbBytes) + if allBytes := p.Fallbacks.ArrayBytes(); len(allBytes) > 0 { + _ = templateDoc.SetRawBytes("choices.0.delta.fallbacks", allBytes) + } + return [][]byte{templateDoc.Bytes()} + } + if toModel != "" { + return [][]byte{templateDoc.Bytes()} + } + return [][]byte{} + } + if blockType == "tool_use" { // Start of tool call - initialize accumulator to track arguments toolCallID := strings.Clone(contentBlock.Get("id").String()) @@ -266,29 +316,38 @@ func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, original // Handle usage information for token counts if usage := root.Get("usage"); usage.Exists() { - (*param).(*ConvertAnthropicResponseToOpenAIParams).Usage.Merge(usage) - promptTokens, completionTokens, totalTokens, cachedTokens, cachedCreationTokens := (*param).(*ConvertAnthropicResponseToOpenAIParams).Usage.OpenAIUsage() + if iterModel, newHops := p.Fallbacks.RecordIterations(usage.Get("iterations"), modelName, "", nil); iterModel != "" { + p.Model = iterModel + _ = templateDoc.Set("model", p.Model) + if len(newHops) > 0 { + _ = templateDoc.SetRawBytes("choices.0.delta.fallback", p.Fallbacks.Latest()) + if allBytes := p.Fallbacks.ArrayBytes(); len(allBytes) > 0 { + _ = templateDoc.SetRawBytes("choices.0.delta.fallbacks", allBytes) + } + } + } + p.Usage.Merge(usage) + promptTokens, completionTokens, totalTokens, cachedTokens, cachedCreationTokens := p.Usage.OpenAIUsage() _ = templateDoc.Set("usage.prompt_tokens", promptTokens) _ = templateDoc.Set("usage.completion_tokens", completionTokens) _ = templateDoc.Set("usage.total_tokens", totalTokens) _ = templateDoc.Set("usage.prompt_tokens_details.cached_tokens", cachedTokens) _ = templateDoc.Set("usage.prompt_tokens_details.cached_creation_tokens", cachedCreationTokens) _ = templateDoc.Set("usage.prompt_tokens_details.cache_write_tokens", cachedCreationTokens) - (*param).(*ConvertAnthropicResponseToOpenAIParams).Usage.applyReasoningTokensDoc(templateDoc) + p.Usage.applyReasoningTokensDoc(templateDoc) } return [][]byte{templateDoc.Bytes()} case "message_stop": // Final message event - emit standard OpenAI trailing usage chunk with empty choices array if usage was tracked. - p := (*param).(*ConvertAnthropicResponseToOpenAIParams) if p.Usage.HasUsage && !p.TrailingUsageSent { p.TrailingUsageSent = true usageDoc := ojson.NewDoc([]byte(`{"id":"","object":"chat.completion.chunk","created":0,"model":"","choices":[]}`)) if p.ResponseID != "" { _ = usageDoc.Set("id", p.ResponseID) } - if modelName != "" { - _ = usageDoc.Set("model", modelName) + if chunkModel := effectiveChatResponseModel(p.Model, modelName); chunkModel != "" { + _ = usageDoc.Set("model", chunkModel) } if p.CreatedAt > 0 { _ = usageDoc.Set("created", p.CreatedAt) @@ -356,7 +415,7 @@ func mapAnthropicStopReasonToOpenAI(anthropicReason string) string { // // Returns: // - []byte: An OpenAI-compatible JSON response containing all message content and metadata -func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) []byte { +func ConvertClaudeResponseToOpenAINonStream(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) []byte { rawJSON, _ = translatorcommon.ClaudeMessagesJSONToSSE(rawJSON) chunks := make([][]byte, 0) @@ -377,6 +436,8 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina var stopReason string var contentParts []string var reasoningParts []string + var fallbacks translatorcommon.ClaudeFallbackChain + currentContentLen := 0 usageTokens := claudeUsageTokens{} toolCallsAccumulator := make(map[int]*ToolCallAccumulator) @@ -389,7 +450,13 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina // Extract initial message metadata including ID, model, and input token count if message := root.Get("message"); message.Exists() { messageID = message.Get("id").String() - model = message.Get("model").String() + if m := strings.TrimSpace(message.Get("model").String()); m != "" { + model = m + } + if fromModel, toModel, ok := translatorcommon.ClaudePreTokenFallbackModels(originalRequestRawJSON, requestRawJSON, modelName, model); ok { + offset := 0 + fallbacks.RecordPreToken(fromModel, toModel, "content_offset", &offset) + } createdAt = time.Now().Unix() usageTokens.Merge(message.Get("usage")) } @@ -401,6 +468,20 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina if blockType == "thinking" { // Start of thinking/reasoning content - skip for now as it's handled in delta continue + } else if blockType == "fallback" { + toModel := strings.TrimSpace(contentBlock.Get("to.model").String()) + fromModel := strings.TrimSpace(contentBlock.Get("from.model").String()) + if toModel != "" { + model = toModel + } + offset := currentContentLen + fallbacks.RecordBlock( + contentBlock, + effectiveChatResponseModel(fromModel, modelName), + effectiveChatResponseModel(toModel, model), + "content_offset", + &offset, + ) } else if blockType == "tool_use" { // Initialize tool call accumulator for this index index := int(root.Get("index").Int()) @@ -419,7 +500,9 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina case "text_delta": // Accumulate text content if text := delta.Get("text"); text.Exists() { - contentParts = append(contentParts, text.String()) + textStr := text.String() + contentParts = append(contentParts, textStr) + currentContentLen += len(textStr) } case "thinking_delta": // Accumulate reasoning/thinking content @@ -454,6 +537,10 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina } } if usage := root.Get("usage"); usage.Exists() { + offset := 0 + if iterModel, _ := fallbacks.RecordIterations(usage.Get("iterations"), modelName, "content_offset", &offset); iterModel != "" { + model = iterModel + } usageTokens.Merge(usage) } } @@ -473,11 +560,15 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina // Set basic response fields including message ID, creation time, and model _ = outDoc.Set("id", messageID) _ = outDoc.Set("created", createdAt) - _ = outDoc.Set("model", model) + _ = outDoc.Set("model", effectiveChatResponseModel(model, modelName)) // Set message content by combining all text parts messageContent := strings.Join(contentParts, "") _ = outDoc.Set("choices.0.message.content", messageContent) + if latestFB := fallbacks.Latest(); len(latestFB) > 0 { + _ = outDoc.SetRawBytes("choices.0.message.fallback", latestFB) + _ = outDoc.SetRawBytes("choices.0.message.fallbacks", fallbacks.ArrayBytes()) + } // Add reasoning content if available (following OpenAI reasoning format) if len(reasoningParts) > 0 { diff --git a/internal/translator/claude/openai/chat-completions/claude_openai_response_test.go b/internal/translator/claude/openai/chat-completions/claude_openai_response_test.go index 1c1bcd3def7..7ad817d303b 100644 --- a/internal/translator/claude/openai/chat-completions/claude_openai_response_test.go +++ b/internal/translator/claude/openai/chat-completions/claude_openai_response_test.go @@ -619,3 +619,172 @@ func TestConvertClaudeResponseToOpenAI_ReportsZeroReasoningTokensWhenThinkingIsZ } assertReasoningTokens(t, out[0], 0) } + +func TestConvertClaudeResponseToOpenAI_StreamPreservesUpstreamServingModel(t *testing.T) { + ctx := context.Background() + var param any + + events := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_fb_1","model":"claude-opus-4-8","usage":{"input_tokens":10,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"fallback answer"}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":4,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":0},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + + var chunks [][]byte + for _, ev := range events { + chunks = append(chunks, ConvertClaudeResponseToOpenAI(ctx, "claude-opus-5-5", nil, nil, ev, ¶m)...) + } + if len(chunks) == 0 { + t.Fatal("expected stream chunks") + } + sawFallbackDelta := false + for i, chunk := range chunks { + if got := ojson.GetBytesNoCopy(chunk, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("chunk[%d] model = %q, want claude-opus-4-8; payload=%s", i, got, string(chunk)) + } + if fb := ojson.GetBytesNoCopy(chunk, "choices.0.delta.fallback"); fb.Exists() { + sawFallbackDelta = true + if fb.Get("from.model").String() != "claude-opus-5-5" || fb.Get("to.model").String() != "claude-opus-4-8" { + t.Fatalf("chunk[%d] delta.fallback = %s, want from claude-opus-5-5 to claude-opus-4-8", i, fb.Raw) + } + } + } + if !sawFallbackDelta { + t.Fatal("expected pre-token fallback stream to surface choices[0].delta.fallback") + } + + // Also verify when message_start already reports claude-opus-4-8 AND a fallback content_block is emitted. + var paramPreBlock any + preBlockEvents := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_fb_preblock","model":"claude-opus-4-8","usage":{"input_tokens":10,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"reasoning_extraction"}}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"ok"}}`), + []byte(`data: {"type":"content_block_stop","index":1}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}`), + []byte(`data: {"type":"message_stop"}`), + } + sawTriggerChunk := false + for _, ev := range preBlockEvents { + for _, chunk := range ConvertClaudeResponseToOpenAI(ctx, "claude-opus-5-5", nil, nil, ev, ¶mPreBlock) { + if fb := ojson.GetBytesNoCopy(chunk, "choices.0.delta.fallback"); fb.Exists() && fb.Get("trigger.category").String() == "reasoning_extraction" { + sawTriggerChunk = true + } + } + } + if !sawTriggerChunk { + t.Fatal("expected fallback content_block to surface trigger in choices[0].delta.fallback even when message_start already reported target model") + } + + // Pre-token non-stream response via iterations also surfaces choices[0].message.fallback with content_offset=0. + preNonStreamRaw := []byte(`{"id":"msg_fb_1","type":"message","role":"assistant","model":"claude-opus-4-8","content":[{"type":"text","text":"fallback answer"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":4,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":0},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`) + preNonStreamOut := ConvertClaudeResponseToOpenAINonStream(ctx, "claude-opus-5-5", nil, nil, preNonStreamRaw, nil) + preFB := ojson.GetBytesNoCopy(preNonStreamOut, "choices.0.message.fallback") + if !preFB.Exists() || preFB.Get("from.model").String() != "claude-opus-5-5" || preFB.Get("to.model").String() != "claude-opus-4-8" || preFB.Get("content_offset").Int() != 0 { + t.Fatalf("pre-token non-stream choices[0].message.fallback = %s, want from=5-5 to=4-8 content_offset=0", preFB.Raw) + } +} + +func TestConvertClaudeResponseToOpenAI_MidStreamFallbackUpdatesServingModelAndEmitsBoundary(t *testing.T) { + ctx := context.Background() + var param any + + events := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_fb_mid","model":"claude-opus-5-5","usage":{"input_tokens":10,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial "}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"content_block_start","index":1,"content_block":{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}}}`), + []byte(`data: {"type":"content_block_stop","index":1}`), + []byte(`data: {"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"fallback text"}}`), + []byte(`data: {"type":"content_block_stop","index":2}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":6,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":2},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + + var models []string + var rawChunks [][]byte + var content string + for _, ev := range events { + for _, chunk := range ConvertClaudeResponseToOpenAI(ctx, "claude-opus-5-5", nil, nil, ev, ¶m) { + rawChunks = append(rawChunks, chunk) + models = append(models, ojson.GetBytesNoCopy(chunk, "model").String()) + if deltaText := ojson.GetBytesNoCopy(chunk, "choices.0.delta.content"); deltaText.Exists() { + content += deltaText.String() + } + } + } + + // Expected sequence: message_start (5.5), partial delta (5.5), fallback boundary (4.8), fallback delta (4.8), message_delta (4.8), trailing usage (4.8). + wantModels := []string{ + "claude-opus-5-5", + "claude-opus-5-5", + "claude-opus-4-8", + "claude-opus-4-8", + "claude-opus-4-8", + "claude-opus-4-8", + } + if len(models) != len(wantModels) { + t.Fatalf("chunk models = %v (len %d), want %v (len %d)", models, len(models), wantModels, len(wantModels)) + } + for i := range wantModels { + if models[i] != wantModels[i] { + t.Fatalf("chunk[%d] model = %q, want %q (all=%v)", i, models[i], wantModels[i], models) + } + } + boundaryFB := ojson.GetBytesNoCopy(rawChunks[2], "choices.0.delta.fallback") + if !boundaryFB.Exists() || boundaryFB.Get("from.model").String() != "claude-opus-5-5" || boundaryFB.Get("to.model").String() != "claude-opus-4-8" || boundaryFB.Get("trigger.category").String() != "cyber" { + t.Fatalf("boundary chunk[2] delta.fallback = %s, want from=claude-opus-5-5 to=claude-opus-4-8 trigger.category=cyber", boundaryFB.Raw) + } + if content != "partial fallback text" { + t.Fatalf("accumulated content = %q, want %q", content, "partial fallback text") + } + + nonStreamRaw := []byte(`{"id":"msg_fb_mid","type":"message","role":"assistant","model":"claude-opus-5-5","content":[{"type":"text","text":"partial "},{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}},{"type":"text","text":"fallback text"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":6,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":2},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`) + nonStreamOut := ConvertClaudeResponseToOpenAINonStream(ctx, "claude-opus-5-5", nil, nil, nonStreamRaw, nil) + if got := ojson.GetBytesNoCopy(nonStreamOut, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("non-stream model = %q, want claude-opus-4-8; payload=%s", got, string(nonStreamOut)) + } + if gotContent := ojson.GetBytesNoCopy(nonStreamOut, "choices.0.message.content").String(); gotContent != "partial fallback text" { + t.Fatalf("non-stream content = %q, want %q", gotContent, "partial fallback text") + } + msgFB := ojson.GetBytesNoCopy(nonStreamOut, "choices.0.message.fallback") + if !msgFB.Exists() || msgFB.Get("from.model").String() != "claude-opus-5-5" || msgFB.Get("to.model").String() != "claude-opus-4-8" || msgFB.Get("trigger.category").String() != "cyber" || msgFB.Get("content_offset").Int() != int64(len("partial ")) { + t.Fatalf("non-stream choices[0].message.fallback = %s, want from=5-5 to=4-8 trigger=cyber content_offset=8", msgFB.Raw) + } +} + +func TestConvertClaudeResponseToOpenAI_MultiHopFallbackAndStickyDefault(t *testing.T) { + ctx := context.Background() + + // 1. Multi-hop chain: claude-fable-5-1 ("partA ") -> claude-opus-5 ("partB ") -> claude-opus-4-8 ("partC"). + rawJSON := []byte(`{"id":"msg_chat_chain","type":"message","role":"assistant","model":"claude-fable-5-1","content":[{"type":"text","text":"partA "},{"type":"fallback","from":{"model":"claude-fable-5-1"},"to":{"model":"claude-opus-5"},"trigger":{"type":"refusal","category":"reasoning_extraction"}},{"type":"text","text":"partB "},{"type":"fallback","from":{"model":"claude-opus-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}},{"type":"text","text":"partC"}],"stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":9}}`) + out := ConvertClaudeResponseToOpenAINonStream(ctx, "claude-fable-5-1", nil, nil, rawJSON, nil) + if got := ojson.GetBytesNoCopy(out, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("multi-hop Chat model = %q, want claude-opus-4-8", got) + } + hops := ojson.GetBytesNoCopy(out, "choices.0.message.fallbacks").Array() + if len(hops) != 2 { + t.Fatalf("multi-hop Chat choices[0].message.fallbacks len = %d, want 2; payload=%s", len(hops), string(out)) + } + if hops[0].Get("from.model").String() != "claude-fable-5-1" || hops[0].Get("to.model").String() != "claude-opus-5" || hops[0].Get("content_offset").Int() != 6 { + t.Fatalf("Chat hop[0] = %s, want fable-5-1 -> opus-5 at content_offset=6", hops[0].Raw) + } + if hops[1].Get("from.model").String() != "claude-opus-5" || hops[1].Get("to.model").String() != "claude-opus-4-8" || hops[1].Get("content_offset").Int() != 12 { + t.Fatalf("Chat hop[1] = %s, want opus-5 -> opus-4-8 at content_offset=12", hops[1].Raw) + } + + // 2. Sticky routing with "fallbacks": "default" (no fallback block, no usage.iterations). + translatedDefaultReq := []byte(`{"model":"claude-opus-5-5","fallbacks":"default","messages":[{"role":"user","content":"turn 2"}]}`) + var stickyParam any + stickyStart := []byte(`data: {"type":"message_start","message":{"id":"msg_sticky_chat","model":"claude-opus-4-8","usage":{"input_tokens":5,"output_tokens":0}}}`) + stickyChunks := ConvertClaudeResponseToOpenAI(ctx, "claude-opus-5-5", nil, translatedDefaultReq, stickyStart, &stickyParam) + if len(stickyChunks) != 1 || ojson.GetBytesNoCopy(stickyChunks[0], "choices.0.delta.fallback.to.model").String() != "claude-opus-4-8" { + t.Fatalf("sticky default Chat start chunk = %s, want delta.fallback.to.model=claude-opus-4-8", string(stickyChunks[0])) + } +} diff --git a/internal/translator/claude/openai/responses/claude_openai-responses_response.go b/internal/translator/claude/openai/responses/claude_openai-responses_response.go index 19ff756f913..b1b273c78be 100644 --- a/internal/translator/claude/openai/responses/claude_openai-responses_response.go +++ b/internal/translator/claude/openai/responses/claude_openai-responses_response.go @@ -72,6 +72,8 @@ type claudeToResponsesState struct { WebSearchByToolID map[string]*claudeResponsesWebSearchItem WebSearchItems []*claudeResponsesWebSearchItem StopReason string + FallbackModel string + Fallbacks translatorcommon.ClaudeFallbackChain // usage aggregation Usage claudeResponsesUsageTokens } @@ -941,6 +943,8 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin st.ReasoningIndex = -1 st.ReasoningItems = nil st.StopReason = "" + st.FallbackModel = "" + st.Fallbacks.Reset() st.ApplyPatchCalls = make(map[string]*translatorcommon.ApplyPatchCallState) st.FuncItemAdded = make(map[int]bool) st.FuncArgsSent = make(map[int]int) @@ -967,9 +971,19 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin if requestModelName == "" { requestModelName = modelName } + upstreamModel := strings.TrimSpace(msg.Get("model").String()) + if fromModel, toModel, ok := translatorcommon.ClaudePreTokenFallbackModels(originalRequestRawJSON, requestRawJSON, requestModelName, upstreamModel); ok { + st.FallbackModel = toModel + st.Fallbacks.RecordPreToken(fromModel, toModel, "", nil) + requestModelName = toModel + } if requestModelName != "" { _ = createdDoc.Set("response.model", requestModelName) } + if latestFB := st.Fallbacks.Latest(); len(latestFB) > 0 { + _ = createdDoc.SetRawBytes("response.fallback", latestFB) + _ = createdDoc.SetRawBytes("response.fallbacks", st.Fallbacks.ArrayBytes()) + } out = append(out, emitEvent("response.created", createdDoc.Bytes())) // response.in_progress inprogDoc := ojson.NewDocString(`{"type":"response.in_progress","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","output":[]}}`) @@ -979,6 +993,10 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin if requestModelName != "" { _ = inprogDoc.Set("response.model", requestModelName) } + if latestFB := st.Fallbacks.Latest(); len(latestFB) > 0 { + _ = inprogDoc.SetRawBytes("response.fallback", latestFB) + _ = inprogDoc.SetRawBytes("response.fallbacks", st.Fallbacks.ArrayBytes()) + } out = append(out, emitEvent("response.in_progress", inprogDoc.Bytes())) } case "content_block_start": @@ -1018,6 +1036,35 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin } } + if typ == "fallback" { + toModel := strings.TrimSpace(cb.Get("to.model").String()) + fromModel := strings.TrimSpace(cb.Get("from.model").String()) + if toModel != "" { + st.FallbackModel = strings.Clone(toModel) + } + defaultFrom := fromModel + if defaultFrom == "" { + defaultFrom = st.RequestEcho.model + } + if defaultFrom == "" { + defaultFrom = modelName + } + boundaryIdx := st.NextOutputIndex + if fbBytes := st.Fallbacks.RecordBlock(cb, defaultFrom, st.FallbackModel, "output_index", &boundaryIdx); len(fbBytes) > 0 { + inprogDoc := ojson.NewDocString(`{"type":"response.in_progress","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","output":[]}}`) + _ = inprogDoc.Set("sequence_number", nextSeq()) + _ = inprogDoc.Set("response.id", st.ResponseID) + _ = inprogDoc.Set("response.created_at", st.CreatedAt) + if st.FallbackModel != "" { + _ = inprogDoc.Set("response.model", st.FallbackModel) + } + _ = inprogDoc.SetRawBytes("response.fallback", fbBytes) + _ = inprogDoc.SetRawBytes("response.fallbacks", st.Fallbacks.ArrayBytes()) + out = append(out, emitEvent("response.in_progress", inprogDoc.Bytes())) + } + return noSSEOutput(out) + } + if typ == "text" { st.InTextBlock = true outputIndex := st.messageOutputIndex() @@ -1218,10 +1265,28 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin } return noSSEOutput(out) case "message_delta": - st.Usage.Merge(root.Get("usage")) + usage := root.Get("usage") + st.Usage.Merge(usage) if stopReason := root.Get("delta.stop_reason"); stopReason.Exists() { st.StopReason = strings.Clone(stopReason.String()) } + defaultFrom := st.RequestEcho.model + if defaultFrom == "" { + defaultFrom = modelName + } + if iterModel, newHops := st.Fallbacks.RecordIterations(usage.Get("iterations"), defaultFrom, "", nil); iterModel != "" { + st.FallbackModel = iterModel + if len(newHops) > 0 { + inprogDoc := ojson.NewDocString(`{"type":"response.in_progress","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","output":[]}}`) + _ = inprogDoc.Set("sequence_number", nextSeq()) + _ = inprogDoc.Set("response.id", st.ResponseID) + _ = inprogDoc.Set("response.created_at", st.CreatedAt) + _ = inprogDoc.Set("response.model", st.FallbackModel) + _ = inprogDoc.SetRawBytes("response.fallback", st.Fallbacks.Latest()) + _ = inprogDoc.SetRawBytes("response.fallbacks", st.Fallbacks.ArrayBytes()) + return [][]byte{emitEvent("response.in_progress", inprogDoc.Bytes())} + } + } return [][]byte{} case "message_stop": toolStatus := claudeResponsesOutputStatus(st.StopReason) @@ -1258,6 +1323,13 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin } // Inject original request fields into response as per docs/response.completed.json st.RequestEcho.applyDoc(completedDoc, "response.") + if st.FallbackModel != "" { + _ = completedDoc.Set("response.model", st.FallbackModel) + } + if latestFB := st.Fallbacks.Latest(); len(latestFB) > 0 { + _ = completedDoc.SetRawBytes("response.fallback", latestFB) + _ = completedDoc.SetRawBytes("response.fallbacks", st.Fallbacks.ArrayBytes()) + } // Build response.output from aggregated state outputsWrapperDoc := ojson.NewDocString(`{"arr":[]}`) @@ -1391,7 +1463,7 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin } // ConvertClaudeResponseToOpenAIResponsesNonStream aggregates Claude SSE into a single OpenAI Responses JSON. -func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []byte { +func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []byte { rawJSON, nativeModel := translatorcommon.ClaudeMessagesJSONToSSE(rawJSON) // Aggregate Claude SSE lines into a single OpenAI Responses JSON (non-stream) // We follow the same aggregation logic as the streaming variant but produce @@ -1428,10 +1500,12 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ stri // Aggregation state var ( - responseID string - createdAt int64 - stopReason string - usageTokens claudeResponsesUsageTokens + responseID string + createdAt int64 + stopReason string + usageTokens claudeResponsesUsageTokens + fallbackModel string + fallbacks translatorcommon.ClaudeFallbackChain ) type nonStreamOutputItem struct { @@ -1486,6 +1560,14 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ stri responseID = msg.Get("id").String() createdAt = time.Now().Unix() usageTokens.Merge(msg.Get("usage")) + defaultReqModel := st.RequestEcho.model + if defaultReqModel == "" { + defaultReqModel = modelName + } + if fromModel, toModel, ok := translatorcommon.ClaudePreTokenFallbackModels(originalRequestRawJSON, requestRawJSON, defaultReqModel, msg.Get("model").String()); ok { + fallbackModel = toModel + fallbacks.RecordPreToken(fromModel, toModel, "", nil) + } } case "content_block_start": @@ -1499,6 +1581,21 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ stri activeMessageItem = nil } switch typ { + case "fallback": + toModel := strings.TrimSpace(cb.Get("to.model").String()) + fromModel := strings.TrimSpace(cb.Get("from.model").String()) + if toModel != "" { + fallbackModel = strings.Clone(toModel) + } + defaultFrom := fromModel + if defaultFrom == "" { + defaultFrom = st.RequestEcho.model + } + if defaultFrom == "" { + defaultFrom = modelName + } + boundaryIdx := nextOutputIndex + fallbacks.RecordBlock(cb, defaultFrom, fallbackModel, "output_index", &boundaryIdx) case "text": item := activeMessageItem if item == nil { @@ -1648,10 +1745,18 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ stri // Output items are finalized after all deltas have been aggregated. case "message_delta": - usageTokens.Merge(root.Get("usage")) + usage := root.Get("usage") + usageTokens.Merge(usage) if value := root.Get("delta.stop_reason"); value.Exists() { stopReason = value.String() } + defaultFrom := st.RequestEcho.model + if defaultFrom == "" { + defaultFrom = modelName + } + if iterModel, _ := fallbacks.RecordIterations(usage.Get("iterations"), defaultFrom, "", nil); iterModel != "" { + fallbackModel = iterModel + } } } @@ -1675,6 +1780,13 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(ctx context.Context, _ stri if nativeModel != "" { _ = outDoc.Set("model", nativeModel) } + if fallbackModel != "" { + _ = outDoc.Set("model", fallbackModel) + } + if latestFB := fallbacks.Latest(); len(latestFB) > 0 { + _ = outDoc.SetRawBytes("fallback", latestFB) + _ = outDoc.SetRawBytes("fallbacks", fallbacks.ArrayBytes()) + } // Build output array in the order of the original content blocks. outputs := make([][]byte, 0, len(outputItems)) diff --git a/internal/translator/claude/openai/responses/claude_openai-responses_response_test.go b/internal/translator/claude/openai/responses/claude_openai-responses_response_test.go index 948551ab610..d499bd43078 100644 --- a/internal/translator/claude/openai/responses/claude_openai-responses_response_test.go +++ b/internal/translator/claude/openai/responses/claude_openai-responses_response_test.go @@ -2501,3 +2501,221 @@ func BenchmarkConvertClaudeResponseToOpenAIResponses_TextDelta(b *testing.B) { ConvertClaudeResponseToOpenAIResponses(context.Background(), "claude-test", request, nil, delta, ¶m) } } + +func TestConvertClaudeResponseToOpenAIResponses_PreTokenAndMidStreamFallbackTransparency(t *testing.T) { + ctx := context.Background() + + // 1. Pre-token fallback with requestRawJSON fallbacks: created, in_progress, completed, and non-stream all surface claude-opus-4-8 + fallback metadata. + originalReq := []byte(`{"model":"claude-opus-5-5","input":[{"role":"user","content":"hi"}]}`) + translatedReq := []byte(`{"model":"claude-opus-5-5","fallbacks":[{"model":"claude-opus-4-8"}],"messages":[{"role":"user","content":"hi"}]}`) + preEvents := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_resp_pre","model":"claude-opus-4-8","usage":{"input_tokens":8,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"fallback answer"}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":0},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":3}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + var preParam any + var createdData, completedData ojson.Result + for _, ev := range preEvents { + for _, out := range ConvertClaudeResponseToOpenAIResponses(ctx, "claude-opus-5-5", originalReq, translatedReq, ev, &preParam) { + event, data := parseClaudeResponsesSSEEvent(t, out) + switch event { + case "response.created": + createdData = data + case "response.completed": + completedData = data + } + } + } + if got := createdData.Get("response.model").String(); got != "claude-opus-4-8" { + t.Fatalf("pre-token response.created model = %q, want claude-opus-4-8", got) + } + if createdData.Get("response.fallback.from.model").String() != "claude-opus-5-5" || createdData.Get("response.fallback.to.model").String() != "claude-opus-4-8" { + t.Fatalf("pre-token response.created fallback = %s, want from=5-5 to=4-8", createdData.Get("response.fallback").Raw) + } + if got := completedData.Get("response.model").String(); got != "claude-opus-4-8" { + t.Fatalf("pre-token response.completed model = %q, want claude-opus-4-8", got) + } + if completedData.Get("response.fallback.from.model").String() != "claude-opus-5-5" || completedData.Get("response.fallback.to.model").String() != "claude-opus-4-8" { + t.Fatalf("pre-token response.completed fallback = %s, want from=5-5 to=4-8", completedData.Get("response.fallback").Raw) + } + + preNonStreamSSE := bytes.Join(preEvents, []byte("\n\n")) + preNonStreamOut := ConvertClaudeResponseToOpenAIResponsesNonStream(ctx, "claude-opus-5-5", originalReq, translatedReq, preNonStreamSSE, nil) + if got := ojson.GetBytesNoCopy(preNonStreamOut, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("pre-token non-stream Responses model = %q, want claude-opus-4-8; payload=%s", got, string(preNonStreamOut)) + } + if fb := ojson.GetBytesNoCopy(preNonStreamOut, "fallback"); !fb.Exists() || fb.Get("from.model").String() != "claude-opus-5-5" || fb.Get("to.model").String() != "claude-opus-4-8" { + t.Fatalf("pre-token non-stream Responses fallback = %s, want from=5-5 to=4-8", fb.Raw) + } + + // 2. Mid-stream fallback: output[0] closes at fallback boundary, response.in_progress announces fallback + output_index=1, and completed/non-stream preserve serving model + boundary. + midEvents := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_resp_mid","model":"claude-opus-5-5","usage":{"input_tokens":8,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial "}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"content_block_start","index":1,"content_block":{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}}}`), + []byte(`data: {"type":"content_block_stop","index":1}`), + []byte(`data: {"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"fallback text"}}`), + []byte(`data: {"type":"content_block_stop","index":2}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":6,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":2},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + + var midParam any + var boundaryInProgress, midCompleted ojson.Result + for _, ev := range midEvents { + for _, out := range ConvertClaudeResponseToOpenAIResponses(ctx, "claude-opus-5-5", originalReq, translatedReq, ev, &midParam) { + event, data := parseClaudeResponsesSSEEvent(t, out) + if event == "response.in_progress" && data.Get("response.fallback").Exists() { + boundaryInProgress = data + } + if event == "response.completed" { + midCompleted = data + } + } + } + if !boundaryInProgress.Exists() { + t.Fatal("expected mid-stream fallback to emit response.in_progress with response.fallback") + } + if got := boundaryInProgress.Get("response.model").String(); got != "claude-opus-4-8" { + t.Fatalf("boundary response.in_progress model = %q, want claude-opus-4-8", got) + } + if boundaryInProgress.Get("response.fallback.trigger.category").String() != "cyber" || boundaryInProgress.Get("response.fallback.output_index").Int() != 1 { + t.Fatalf("boundary response.in_progress fallback = %s, want trigger=cyber output_index=1", boundaryInProgress.Get("response.fallback").Raw) + } + if got := midCompleted.Get("response.model").String(); got != "claude-opus-4-8" { + t.Fatalf("mid-stream response.completed model = %q, want claude-opus-4-8", got) + } + if midCompleted.Get("response.fallback.output_index").Int() != 1 || len(midCompleted.Get("response.output").Array()) != 2 { + t.Fatalf("mid-stream response.completed = %s, want fallback.output_index=1 and 2 split message items", midCompleted.Get("response").Raw) + } + + midNativeRaw := []byte(`{"id":"msg_resp_mid","type":"message","role":"assistant","model":"claude-opus-5-5","content":[{"type":"text","text":"partial "},{"type":"fallback","from":{"model":"claude-opus-5-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}},{"type":"text","text":"fallback text"}],"stop_reason":"end_turn","usage":{"input_tokens":8,"output_tokens":6,"iterations":[{"type":"message","model":"claude-opus-5-5","output_tokens":2},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":4}]}}`) + midNonStreamOut := ConvertClaudeResponseToOpenAIResponsesNonStream(ctx, "claude-opus-5-5", originalReq, translatedReq, midNativeRaw, nil) + if got := ojson.GetBytesNoCopy(midNonStreamOut, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("mid-stream non-stream Responses model = %q, want claude-opus-4-8; payload=%s", got, string(midNonStreamOut)) + } + midFB := ojson.GetBytesNoCopy(midNonStreamOut, "fallback") + if !midFB.Exists() || midFB.Get("trigger.category").String() != "cyber" || midFB.Get("output_index").Int() != 1 { + t.Fatalf("mid-stream non-stream Responses fallback = %s, want trigger=cyber output_index=1", midFB.Raw) + } +} + +func TestConvertClaudeResponseToOpenAIResponses_DefaultStringFallbackAndStickyRouting(t *testing.T) { + ctx := context.Background() + + // 1. "fallbacks": "default" + sticky routing (upstream serves claude-opus-4-8 directly in message_start with no fallback block and no usage.iterations). + origReq := []byte(`{"model":"claude-opus-5-5","input":[{"role":"user","content":"turn 2"}]}`) + translatedDefaultReq := []byte(`{"model":"claude-opus-5-5","fallbacks":"default","messages":[{"role":"user","content":"turn 2"}]}`) + stickyEvents := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_sticky","model":"claude-opus-4-8","usage":{"input_tokens":12,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"sticky answer"}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`), + []byte(`data: {"type":"message_stop"}`), + } + + var param any + var created, inProgress, completed ojson.Result + for _, ev := range stickyEvents { + for _, out := range ConvertClaudeResponseToOpenAIResponses(ctx, "claude-opus-5-5", origReq, translatedDefaultReq, ev, ¶m) { + event, data := parseClaudeResponsesSSEEvent(t, out) + switch event { + case "response.created": + created = data + case "response.in_progress": + if !inProgress.Exists() { + inProgress = data + } + case "response.completed": + completed = data + } + } + } + if created.Get("response.model").String() != "claude-opus-4-8" || inProgress.Get("response.model").String() != "claude-opus-4-8" || completed.Get("response.model").String() != "claude-opus-4-8" { + t.Fatalf("sticky default fallback models = (%q, %q, %q), want all claude-opus-4-8", + created.Get("response.model").String(), inProgress.Get("response.model").String(), completed.Get("response.model").String()) + } + if created.Get("response.fallback.from.model").String() != "claude-opus-5-5" || created.Get("response.fallback.to.model").String() != "claude-opus-4-8" { + t.Fatalf("sticky created fallback = %s, want from=5-5 to=4-8", created.Get("response.fallback").Raw) + } + + stickyNonStream := ConvertClaudeResponseToOpenAIResponsesNonStream(ctx, "claude-opus-5-5", origReq, translatedDefaultReq, bytes.Join(stickyEvents, []byte("\n\n")), nil) + if got := ojson.GetBytesNoCopy(stickyNonStream, "model").String(); got != "claude-opus-4-8" { + t.Fatalf("sticky non-stream model = %q, want claude-opus-4-8", got) + } + + // 2. Dated snapshot of the primary model (claude-opus-5-5-20260501) must NOT be misclassified as a fallback. + var snapParam any + snapStart := []byte(`data: {"type":"message_start","message":{"id":"msg_snap","model":"claude-opus-5-5-20260501","usage":{"input_tokens":4,"output_tokens":0}}}`) + for _, out := range ConvertClaudeResponseToOpenAIResponses(ctx, "claude-opus-5-5", origReq, translatedDefaultReq, snapStart, &snapParam) { + event, data := parseClaudeResponsesSSEEvent(t, out) + if event == "response.created" && data.Get("response.fallback").Exists() { + t.Fatalf("dated primary snapshot should not emit fallback metadata: %s", data.Get("response.fallback").Raw) + } + } +} + +func TestConvertClaudeResponseToOpenAIResponses_MultiHopFallbackChainPreservesAllBoundaries(t *testing.T) { + ctx := context.Background() + origReq := []byte(`{"model":"claude-fable-5-1","input":[{"role":"user","content":"chain"}]}`) + translatedReq := []byte(`{"model":"claude-fable-5-1","fallbacks":[{"model":"claude-opus-5"},{"model":"claude-opus-4-8"}],"messages":[{"role":"user","content":"chain"}]}`) + + events := [][]byte{ + []byte(`data: {"type":"message_start","message":{"id":"msg_chain","model":"claude-fable-5-1","usage":{"input_tokens":10,"output_tokens":0}}}`), + []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partA "}}`), + []byte(`data: {"type":"content_block_stop","index":0}`), + []byte(`data: {"type":"content_block_start","index":1,"content_block":{"type":"fallback","from":{"model":"claude-fable-5-1"},"to":{"model":"claude-opus-5"},"trigger":{"type":"refusal","category":"reasoning_extraction"}}}`), + []byte(`data: {"type":"content_block_stop","index":1}`), + []byte(`data: {"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"partB "}}`), + []byte(`data: {"type":"content_block_stop","index":2}`), + []byte(`data: {"type":"content_block_start","index":3,"content_block":{"type":"fallback","from":{"model":"claude-opus-5"},"to":{"model":"claude-opus-4-8"},"trigger":{"type":"refusal","category":"cyber"}}}`), + []byte(`data: {"type":"content_block_stop","index":3}`), + []byte(`data: {"type":"content_block_start","index":4,"content_block":{"type":"text","text":""}}`), + []byte(`data: {"type":"content_block_delta","index":4,"delta":{"type":"text_delta","text":"partC"}}`), + []byte(`data: {"type":"content_block_stop","index":4}`), + []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":9,"iterations":[{"type":"message","model":"claude-fable-5-1","output_tokens":3},{"type":"fallback_message","model":"claude-opus-5","output_tokens":3},{"type":"fallback_message","model":"claude-opus-4-8","output_tokens":3}]}}`), + []byte(`data: {"type":"message_stop"}`), + } + + var param any + var completed ojson.Result + for _, ev := range events { + for _, out := range ConvertClaudeResponseToOpenAIResponses(ctx, "claude-fable-5-1", origReq, translatedReq, ev, ¶m) { + event, data := parseClaudeResponsesSSEEvent(t, out) + if event == "response.completed" { + completed = data + } + } + } + if got := completed.Get("response.model").String(); got != "claude-opus-4-8" { + t.Fatalf("multi-hop completed model = %q, want claude-opus-4-8", got) + } + hops := completed.Get("response.fallbacks").Array() + if len(hops) != 2 { + t.Fatalf("multi-hop response.fallbacks len = %d, want 2; raw=%s", len(hops), completed.Get("response.fallbacks").Raw) + } + if hops[0].Get("from.model").String() != "claude-fable-5-1" || hops[0].Get("to.model").String() != "claude-opus-5" || hops[0].Get("output_index").Int() != 1 { + t.Fatalf("hop[0] = %s, want fable-5-1 -> opus-5 at output_index=1", hops[0].Raw) + } + if hops[1].Get("from.model").String() != "claude-opus-5" || hops[1].Get("to.model").String() != "claude-opus-4-8" || hops[1].Get("output_index").Int() != 2 { + t.Fatalf("hop[1] = %s, want opus-5 -> opus-4-8 at output_index=2", hops[1].Raw) + } + if len(completed.Get("response.output").Array()) != 3 { + t.Fatalf("multi-hop response.output len = %d, want 3 message items", len(completed.Get("response.output").Array())) + } + + nonStreamOut := ConvertClaudeResponseToOpenAIResponsesNonStream(ctx, "claude-fable-5-1", origReq, translatedReq, bytes.Join(events, []byte("\n\n")), nil) + nonStreamHops := ojson.GetBytesNoCopy(nonStreamOut, "fallbacks").Array() + if len(nonStreamHops) != 2 || nonStreamHops[0].Get("output_index").Int() != 1 || nonStreamHops[1].Get("output_index").Int() != 2 { + t.Fatalf("non-stream multi-hop fallbacks = %s, want 2 hops with output_index 1 and 2", ojson.GetBytesNoCopy(nonStreamOut, "fallbacks").Raw) + } +} diff --git a/internal/translator/common/claude_native_response.go b/internal/translator/common/claude_native_response.go index b1dd1ae8c77..1cb67dd640f 100644 --- a/internal/translator/common/claude_native_response.go +++ b/internal/translator/common/claude_native_response.go @@ -83,3 +83,233 @@ func ClaudeMessagesJSONToSSE(rawJSON []byte) ([]byte, string) { emit(map[string]any{"type": "message_stop"}) return out.Bytes(), strings.Clone(root.Get("model").String()) } + +// ClaudeModelSameOrSnapshot reports whether primaryModel and upstreamModel +// refer to the same Anthropic model or an 8-digit YYYYMMDD snapshot thereof. +func ClaudeModelSameOrSnapshot(primaryModel, upstreamModel string) bool { + p := strings.ToLower(strings.TrimSpace(primaryModel)) + u := strings.ToLower(strings.TrimSpace(upstreamModel)) + if p == "" || u == "" { + return false + } + if p == u { + return true + } + return isClaudeDatedSnapshotOf(p, u) || isClaudeDatedSnapshotOf(u, p) +} + +func isClaudeDatedSnapshotOf(base, candidate string) bool { + if !strings.HasPrefix(candidate, base+"-") { + return false + } + suffix := candidate[len(base)+1:] + if len(suffix) != 8 { + return false + } + for i := 0; i < len(suffix); i++ { + if suffix[i] < '0' || suffix[i] > '9' { + return false + } + } + return true +} + +func claudeRequestHasFallbacks(raw []byte) bool { + if len(raw) == 0 || !ojson.ValidBytes(raw) { + return false + } + fb := ojson.GetBytesNoCopy(raw, "fallbacks") + if !fb.Exists() { + return false + } + switch fb.Type { + case ojson.String: + return strings.TrimSpace(fb.String()) != "" + case ojson.JSON: + if fb.IsArray() { + return len(fb.Array()) > 0 + } + if fb.IsObject() { + return len(fb.Map()) > 0 + } + } + return false +} + +func claudePayloadStringModel(raw []byte) string { + if len(raw) == 0 || !ojson.ValidBytes(raw) { + return "" + } + m := ojson.GetBytesNoCopy(raw, "model") + if m.Type != ojson.String { + return "" + } + return strings.TrimSpace(m.String()) +} + +func isCanonicalClaudeModelID(model string) bool { + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "claude-") +} + +// ClaudePreTokenFallbackModels detects when an Anthropic message_start model +// already switched to a fallback model (including array fallbacks, +// fallbacks: "default", and sticky-routing turns). +func ClaudePreTokenFallbackModels(originalRequestRawJSON, requestRawJSON []byte, fallbackRequestedModel, upstreamModel string) (fromModel, toModel string, ok bool) { + upstreamModel = strings.TrimSpace(upstreamModel) + if upstreamModel == "" { + return "", "", false + } + primary := claudePayloadStringModel(requestRawJSON) + if primary == "" { + primary = claudePayloadStringModel(originalRequestRawJSON) + } + if primary == "" { + primary = strings.TrimSpace(fallbackRequestedModel) + } + if primary != "" && ClaudeModelSameOrSnapshot(primary, upstreamModel) { + return "", "", false + } + hasFallbacks := claudeRequestHasFallbacks(requestRawJSON) || claudeRequestHasFallbacks(originalRequestRawJSON) + if !hasFallbacks && !(isCanonicalClaudeModelID(primary) && isCanonicalClaudeModelID(upstreamModel)) { + return "", "", false + } + return strings.Clone(primary), strings.Clone(upstreamModel), true +} + +// BuildClaudeFallbackMetadata serializes a structured fallback metadata object +// preserving upstream from/to/trigger objects and an optional boundary index. +func BuildClaudeFallbackMetadata(contentBlock ojson.Result, defaultFromModel, defaultToModel, boundaryKey string, boundaryIndex *int) []byte { + doc := ojson.NewDocString("{}") + hasField := false + if contentBlock.Exists() && contentBlock.Get("from").IsObject() { + _ = doc.SetRawBytes("from", []byte(contentBlock.Get("from").Raw)) + hasField = true + } else if m := strings.TrimSpace(defaultFromModel); m != "" { + _ = doc.Set("from.model", m) + hasField = true + } + if contentBlock.Exists() && contentBlock.Get("to").IsObject() { + _ = doc.SetRawBytes("to", []byte(contentBlock.Get("to").Raw)) + hasField = true + } else if m := strings.TrimSpace(defaultToModel); m != "" { + _ = doc.Set("to.model", m) + hasField = true + } + if contentBlock.Exists() && contentBlock.Get("trigger").IsObject() { + _ = doc.SetRawBytes("trigger", []byte(contentBlock.Get("trigger").Raw)) + hasField = true + } + if !hasField { + return nil + } + if boundaryKey != "" && boundaryIndex != nil { + _ = doc.Set(boundaryKey, *boundaryIndex) + } + return doc.Bytes() +} + +// ClaudeFallbackChain accumulates ordered Anthropic fallback hops across +// pre-token message_start switches, mid-stream fallback content blocks, and +// usage.iterations fallback_message entries without losing earlier boundaries. +type ClaudeFallbackChain struct { + Entries [][]byte + inferredPreTokenSrc string + hasExplicitBlock bool +} + +// Reset clears all accumulated fallback state for a new message. +func (c *ClaudeFallbackChain) Reset() { + c.Entries = nil + c.inferredPreTokenSrc = "" + c.hasExplicitBlock = false +} + +// RecordPreToken records a pre-token or sticky fallback detected at message_start. +func (c *ClaudeFallbackChain) RecordPreToken(fromModel, toModel, boundaryKey string, boundaryIndex *int) []byte { + fb := BuildClaudeFallbackMetadata(ojson.Result{}, fromModel, toModel, boundaryKey, boundaryIndex) + if len(fb) == 0 { + return nil + } + c.Entries = [][]byte{fb} + c.inferredPreTokenSrc = strings.Clone(strings.TrimSpace(fromModel)) + c.hasExplicitBlock = false + return fb +} + +// RecordBlock records an explicit Anthropic content_block of type "fallback". +// If the first explicit block describes the same transition already inferred at +// message_start, it upgrades that entry with trigger/boundary details; otherwise +// it appends a new hop in order. +func (c *ClaudeFallbackChain) RecordBlock(contentBlock ojson.Result, defaultFromModel, defaultToModel, boundaryKey string, boundaryIndex *int) []byte { + fb := BuildClaudeFallbackMetadata(contentBlock, defaultFromModel, defaultToModel, boundaryKey, boundaryIndex) + if len(fb) == 0 { + return nil + } + blockFrom := strings.TrimSpace(contentBlock.Get("from.model").String()) + if blockFrom == "" { + blockFrom = strings.TrimSpace(defaultFromModel) + } + if !c.hasExplicitBlock && len(c.Entries) == 1 && c.inferredPreTokenSrc != "" && ClaudeModelSameOrSnapshot(c.inferredPreTokenSrc, blockFrom) { + c.Entries[0] = fb + } else { + c.Entries = append(c.Entries, fb) + } + c.inferredPreTokenSrc = "" + c.hasExplicitBlock = true + return fb +} + +// RecordIterations reconciles usage.iterations fallback_message entries, +// returning the final fallback serving model and any newly discovered hop +// metadata when explicit fallback content blocks were absent. +func (c *ClaudeFallbackChain) RecordIterations(iterations ojson.Result, defaultFromModel, boundaryKey string, boundaryIndex *int) (finalModel string, newEntries [][]byte) { + if !iterations.Exists() || !iterations.IsArray() { + return "", nil + } + curFrom := strings.TrimSpace(defaultFromModel) + var hops [][]byte + for _, it := range iterations.Array() { + m := strings.TrimSpace(it.Get("model").String()) + if m == "" { + continue + } + switch it.Get("type").String() { + case "message": + curFrom = m + case "fallback_message": + finalModel = strings.Clone(m) + if !c.hasExplicitBlock { + if fb := BuildClaudeFallbackMetadata(ojson.Result{}, curFrom, m, boundaryKey, boundaryIndex); len(fb) > 0 { + hops = append(hops, fb) + } + } + curFrom = m + } + } + if finalModel == "" || c.hasExplicitBlock || len(hops) == 0 { + return finalModel, nil + } + prevLen := len(c.Entries) + if len(hops) > prevLen { + c.Entries = hops + newEntries = hops[prevLen:] + } + c.inferredPreTokenSrc = "" + return finalModel, newEntries +} + +// Latest returns the most recent fallback hop metadata, or nil when none occurred. +func (c *ClaudeFallbackChain) Latest() []byte { + if len(c.Entries) == 0 { + return nil + } + return c.Entries[len(c.Entries)-1] +} + +// ArrayBytes serializes all ordered fallback hops as a JSON array, or nil when empty. +func (c *ClaudeFallbackChain) ArrayBytes() []byte { + if len(c.Entries) == 0 { + return nil + } + return JoinRawArray(c.Entries) +} diff --git a/internal/translator/openai/claude/openai_claude_compat_test.go b/internal/translator/openai/claude/openai_claude_compat_test.go index ed908049f9e..1f3675cfea9 100644 --- a/internal/translator/openai/claude/openai_claude_compat_test.go +++ b/internal/translator/openai/claude/openai_claude_compat_test.go @@ -39,13 +39,97 @@ func TestConvertClaudeRequestToOpenAIWithCompatDoesNotAddReasoningWithoutThinkin result, _ := ConvertClaudeRequestToOpenAIWithCompat("deepseek-v4", payload, false) assistant := ojson.GetBytesNoCopy(result, "messages.0") if assistant.Get("reasoning_content").Exists() { - t.Fatalf("compatible translation added reasoning_content without thinking: %s", result) + t.Fatalf("compatible translation added reasoning_content without request thinking: %s", result) } if !assistant.Get("tool_calls").Exists() { t.Fatalf("tool_calls missing from compatible translation: %s", result) } } +func TestConvertClaudeRequestToOpenAIWithCompatDoesNotAddReasoningWhenThinkingDisabled(t *testing.T) { + for _, tc := range []struct { + name string + payload string + }{ + { + name: "thinking type disabled", + payload: `{"thinking":{"type":"disabled"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "thinking budget zero", + payload: `{"thinking":{"type":"enabled","budget_tokens":0},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "output_config effort none", + payload: `{"output_config":{"effort":"none"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "model suffix none overrides enabled body", + payload: `{"model":"deepseek-v4(none)","thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + result, _ := ConvertClaudeRequestToOpenAIWithCompat("deepseek-v4", []byte(tc.payload), false) + assistant := ojson.GetBytesNoCopy(result, "messages.0") + if assistant.Get("reasoning_content").Exists() { + t.Fatalf("compatible translation added reasoning_content when thinking disabled: %s", result) + } + if !assistant.Get("tool_calls").Exists() { + t.Fatalf("tool_calls missing from compatible translation: %s", result) + } + }) + } +} + +func TestConvertClaudeRequestToOpenAIWithCompatDoesNotAddReasoningWithoutThinkingForTextOnly(t *testing.T) { + payload := []byte(`{"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"assistant","content":[{"type":"text","text":"Plain answer without thinking."}]}]}`) + + result, _ := ConvertClaudeRequestToOpenAIWithCompat("deepseek-v4", payload, false) + assistant := ojson.GetBytesNoCopy(result, "messages.0") + if assistant.Get("reasoning_content").Exists() { + t.Fatalf("compatible translation added reasoning_content to text-only assistant turn: %s", result) + } +} + +func TestConvertClaudeRequestToOpenAIWithCompatAddsPlaceholderReasoningForToolCallsWithoutThinking(t *testing.T) { + for _, tc := range []struct { + name string + payload string + }{ + { + name: "enabled thinking with absent thinking block", + payload: `{"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "enabled thinking with whitespace-only thinking block", + payload: `{"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":" ","signature":""},{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "adaptive thinking with absent thinking block", + payload: `{"thinking":{"type":"adaptive"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "standalone output_config effort with absent thinking block", + payload: `{"output_config":{"effort":"high"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + { + name: "model suffix high overrides disabled body", + payload: `{"model":"deepseek-v4(high)","thinking":{"type":"disabled"},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + result, _ := ConvertClaudeRequestToOpenAIWithCompat("deepseek-v4", []byte(tc.payload), false) + assistant := ojson.GetBytesNoCopy(result, "messages.0") + if got := assistant.Get("reasoning_content").String(); got != " " { + t.Fatalf("reasoning_content = %q, want 1-char placeholder; output: %s", got, result) + } + if !assistant.Get("tool_calls").Exists() { + t.Fatalf("tool_calls missing from compatible translation: %s", result) + } + }) + } +} + func TestConvertClaudeRequestToOpenAIWithCompatPreservesIncompatibleThinking(t *testing.T) { payload := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"reason","signature":"claude#opaque"},{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`) @@ -60,10 +144,13 @@ func TestConvertClaudeRequestToOpenAIWithCompatPreservesIncompatibleThinking(t * } func TestConvertClaudeRequestToOpenAIWithoutCompatDoesNotAddReasoningForToolCalls(t *testing.T) { - payload := []byte(`{"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`) - - result, _ := ConvertClaudeRequestToOpenAI("deepseek-v4", payload, false) - if ojson.GetBytesNoCopy(result, "messages.0.reasoning_content").Exists() { - t.Fatalf("default translation added reasoning_content: %s", result) + for _, payload := range []string{ + `{"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + `{"thinking":{"type":"enabled","budget_tokens":1024},"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}]}`, + } { + result, _ := ConvertClaudeRequestToOpenAI("deepseek-v4", []byte(payload), false) + if ojson.GetBytesNoCopy(result, "messages.0.reasoning_content").Exists() { + t.Fatalf("default translation added reasoning_content: %s", result) + } } } diff --git a/internal/translator/openai/claude/openai_claude_request.go b/internal/translator/openai/claude/openai_claude_request.go index 6a9c93c6cf3..bcc122d4cb8 100644 --- a/internal/translator/openai/claude/openai_claude_request.go +++ b/internal/translator/openai/claude/openai_claude_request.go @@ -114,6 +114,8 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream _ = outDoc.Set("reasoning_effort", effort) } + reasoningActive := isClaudeRequestReasoningActive(outDoc.Get("reasoning_effort"), modelName, root.Get("model").String()) + // Process messages and system. messageCapacity := root.Get("messages.#").Int() if root.Get("system").Exists() { @@ -342,6 +344,11 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream // Add reasoning_content if present if hasReasoning { _ = msgDoc.Set("reasoning_content", reasoningContent) + } else if hasToolCalls && preserveThinkingBlocks && reasoningActive { + // Assistant turn with tool_calls but no replayable thinking text. + // Compat thinking upstreams (DeepSeek, Moonshot/Kimi) require + // reasoning_content to be present on such messages when reasoning is active. + _ = msgDoc.Set("reasoning_content", " ") } // Add tool_calls if present (in same message as content) @@ -545,6 +552,37 @@ func normalizeObjectSchemaProperties(schema any) any { } } +func isClaudeRequestReasoningActive(effortRes ojson.Result, modelName, payloadModel string) bool { + if active, ok := isThinkingSuffixActive(modelName); ok { + return active + } + if active, ok := isThinkingSuffixActive(payloadModel); ok { + return active + } + if !effortRes.Exists() { + return false + } + resolvedEffort := strings.ToLower(strings.TrimSpace(effortRes.String())) + return resolvedEffort != "" && resolvedEffort != string(thinking.LevelNone) && resolvedEffort != "0" && resolvedEffort != "false" +} + +func isThinkingSuffixActive(model string) (bool, bool) { + suffix := thinking.ParseSuffix(strings.TrimSpace(model)) + if !suffix.HasSuffix { + return false, false + } + if mode, ok := thinking.ParseSpecialSuffix(suffix.RawSuffix); ok { + return mode != thinking.ModeNone, true + } + if budget, ok := thinking.ParseNumericSuffix(suffix.RawSuffix); ok { + return budget > 0, true + } + if level, ok := thinking.ParseLevelSuffix(suffix.RawSuffix); ok { + return level != "" && level != thinking.LevelNone, true + } + return false, false +} + func shouldMapClaudeThinkingToGPTReasoning(part ojson.Result, preserveThinkingBlocks ...bool) bool { preserveThinking := len(preserveThinkingBlocks) > 0 && preserveThinkingBlocks[0] if preserveThinking { diff --git a/internal/util/nocopy_invariant_test.go b/internal/util/nocopy_invariant_test.go index 368e68c6e92..222ff0987c3 100644 --- a/internal/util/nocopy_invariant_test.go +++ b/internal/util/nocopy_invariant_test.go @@ -116,6 +116,7 @@ var reviewedInPlaceByteWrites = map[string]reviewedInPlaceByteWrite{ "internal/client/codex/live/tcp_proxy.go": {1, "copies header and payload into a freshly allocated frame"}, "internal/home/client.go": {1, "zeroes a secret buffer after json.Unmarshal has copied every value out"}, "internal/pluginstore/auth.go": {1, "zeroes a locally built credential buffer after base64 encoding copied it out"}, + "sdk/cliproxy/auth/selector_weighted.go": {2, "shifts []weightedPoolMember structs to insert or remove a pool member; no byte of any payload is rewritten"}, } // TestInPlaceByteWritesAreReviewed keeps the set of in-place byte writes small diff --git a/internal/util/provider.go b/internal/util/provider.go index 810a9cb2d0b..e29e5c1e5a1 100644 --- a/internal/util/provider.go +++ b/internal/util/provider.go @@ -215,7 +215,7 @@ func MaskAuthorizationHeader(value string) string { // // Behavior by header key (case-insensitive): // - "Authorization": Preserve the auth type prefix (e.g., "Bearer ") and mask only the credential part. -// - Headers containing "api-key": Mask the entire value using HideAPIKey. +// - Headers containing "api-key", "apikey", "token", "secret", or "cookie": Mask the entire value using HideAPIKey. // - Others: Return the original value unchanged. // // Parameters: @@ -232,7 +232,8 @@ func MaskSensitiveHeaderValue(key, value string) string { case strings.Contains(lowerKey, "api-key"), strings.Contains(lowerKey, "apikey"), strings.Contains(lowerKey, "token"), - strings.Contains(lowerKey, "secret"): + strings.Contains(lowerKey, "secret"), + strings.Contains(lowerKey, "cookie"): return HideAPIKey(value) default: return value diff --git a/internal/util/provider_test.go b/internal/util/provider_test.go index daf45c8db6d..24665d6415a 100644 --- a/internal/util/provider_test.go +++ b/internal/util/provider_test.go @@ -37,3 +37,25 @@ func TestMaskSensitiveQueryMasksValue(t *testing.T) { } }) } + +func TestMaskSensitiveHeaderValueMasksCookiesAndSecrets(t *testing.T) { + for _, tc := range []struct { + key string + value string + wantLeak string + }{ + {key: "Authorization", value: "Bearer super-secret-token-value", wantLeak: "super-secret-token-value"}, + {key: "X-Api-Key", value: "sk-super-secret-api-key", wantLeak: "sk-super-secret-api-key"}, + {key: "Cookie", value: "session=super-secret-cookie", wantLeak: "super-secret-cookie"}, + {key: "Set-Cookie", value: "__cf_bm=super-secret-cf-cookie; Path=/; HttpOnly", wantLeak: "super-secret-cf-cookie"}, + } { + masked := util.MaskSensitiveHeaderValue(tc.key, tc.value) + if strings.Contains(masked, tc.wantLeak) { + t.Fatalf("MaskSensitiveHeaderValue(%q) leaked %q in %q", tc.key, tc.wantLeak, masked) + } + } + + if got := util.MaskSensitiveHeaderValue("X-Request-Id", "req-plain-123"); got != "req-plain-123" { + t.Fatalf("MaskSensitiveHeaderValue(X-Request-Id) = %q, want unchanged", got) + } +} diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index 5ae60d3aeb8..19541a88c2f 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -117,9 +117,15 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if !reflect.DeepEqual(oldCfg.Antigravity.SensitiveWords, newCfg.Antigravity.SensitiveWords) { changes = append(changes, fmt.Sprintf("antigravity.sensitive-words: %d -> %d", len(oldCfg.Antigravity.SensitiveWords), len(newCfg.Antigravity.SensitiveWords))) } + if !equalStringMap(oldCfg.Antigravity.SensitiveWordsReplace, newCfg.Antigravity.SensitiveWordsReplace) { + changes = append(changes, fmt.Sprintf("antigravity.sensitive-words-replace: %d -> %d", len(oldCfg.Antigravity.SensitiveWordsReplace), len(newCfg.Antigravity.SensitiveWordsReplace))) + } if !reflect.DeepEqual(oldCfg.Devin.SensitiveWords, newCfg.Devin.SensitiveWords) { changes = append(changes, fmt.Sprintf("devin.sensitive-words: %d -> %d", len(oldCfg.Devin.SensitiveWords), len(newCfg.Devin.SensitiveWords))) } + if !equalStringMap(oldCfg.Devin.SensitiveWordsReplace, newCfg.Devin.SensitiveWordsReplace) { + changes = append(changes, fmt.Sprintf("devin.sensitive-words-replace: %d -> %d", len(oldCfg.Devin.SensitiveWordsReplace), len(newCfg.Devin.SensitiveWordsReplace))) + } oldEnabled := "" if oldCfg.Antigravity.ConnectionPool.Enabled != nil { oldEnabled = fmt.Sprintf("%t", *oldCfg.Antigravity.ConnectionPool.Enabled) @@ -173,6 +179,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldCfg.Claude.ToolSearch != newCfg.Claude.ToolSearch { changes = append(changes, fmt.Sprintf("claude.tool-search: %t -> %t", oldCfg.Claude.ToolSearch, newCfg.Claude.ToolSearch)) } + if oldCfg.Claude.ServerSideFallback != newCfg.Claude.ServerSideFallback { + changes = append(changes, fmt.Sprintf("claude.server-side-fallback: %t -> %t", oldCfg.Claude.ServerSideFallback, newCfg.Claude.ServerSideFallback)) + } if oldCfg.XAI.InjectXSearch != newCfg.XAI.InjectXSearch { changes = append(changes, fmt.Sprintf("xai.inject-x-search: %t -> %t", oldCfg.XAI.InjectXSearch, newCfg.XAI.InjectXSearch)) } @@ -330,16 +339,25 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { changes = append(changes, fmt.Sprintf("claude[%d].fingerprint-profile: %s -> %s", i, strings.TrimSpace(o.FingerprintProfile), strings.TrimSpace(n.FingerprintProfile))) } changes = appendOptionalIntChange(changes, fmt.Sprintf("claude[%d].request-retry", i), o.RequestRetry, n.RequestRetry) - if o.Cloak != nil && n.Cloak != nil { + switch { + case o.Cloak == nil && n.Cloak != nil: + changes = append(changes, fmt.Sprintf("claude[%d].cloak: added", i)) + case o.Cloak != nil && n.Cloak == nil: + changes = append(changes, fmt.Sprintf("claude[%d].cloak: removed", i)) + case o.Cloak != nil && n.Cloak != nil: if strings.TrimSpace(o.Cloak.Mode) != strings.TrimSpace(n.Cloak.Mode) { changes = append(changes, fmt.Sprintf("claude[%d].cloak.mode: %s -> %s", i, o.Cloak.Mode, n.Cloak.Mode)) } if o.Cloak.StrictMode != n.Cloak.StrictMode { changes = append(changes, fmt.Sprintf("claude[%d].cloak.strict-mode: %t -> %t", i, o.Cloak.StrictMode, n.Cloak.StrictMode)) } - if len(o.Cloak.SensitiveWords) != len(n.Cloak.SensitiveWords) { + if !reflect.DeepEqual(o.Cloak.SensitiveWords, n.Cloak.SensitiveWords) { changes = append(changes, fmt.Sprintf("claude[%d].cloak.sensitive-words: %d -> %d", i, len(o.Cloak.SensitiveWords), len(n.Cloak.SensitiveWords))) } + if !equalStringMap(o.Cloak.SensitiveWordsReplace, n.Cloak.SensitiveWordsReplace) { + changes = append(changes, fmt.Sprintf("claude[%d].cloak.sensitive-words-replace: %d -> %d", i, len(o.Cloak.SensitiveWordsReplace), len(n.Cloak.SensitiveWordsReplace))) + } + changes = appendOptionalBoolChange(changes, fmt.Sprintf("claude[%d].cloak.cache-user-id", i), o.Cloak.CacheUserID, n.Cloak.CacheUserID) } } } @@ -688,7 +706,7 @@ func equalStringMap(a, b map[string]string) bool { return false } for k, v := range a { - if b[k] != v { + if bv, ok := b[k]; !ok || bv != v { return false } } diff --git a/internal/watcher/diff/config_diff_test.go b/internal/watcher/diff/config_diff_test.go index 963a704cc1c..69248fe5f29 100644 --- a/internal/watcher/diff/config_diff_test.go +++ b/internal/watcher/diff/config_diff_test.go @@ -469,7 +469,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { }, }, XAI: config.XAIConfig{InjectXSearch: true}, - Claude: config.ClaudeConfig{ToolSearch: true}, + Claude: config.ClaudeConfig{ToolSearch: true, ServerSideFallback: true}, ClaudeKey: []config.ClaudeKey{ {APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}}, {APIKey: "c2"}, @@ -522,6 +522,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { expectContains(t, details, `antigravity.connection-pool.idle-conn-timeout: "30s" -> "10s"`) expectContains(t, details, "xai.inject-x-search: false -> true") expectContains(t, details, "claude.tool-search: false -> true") + expectContains(t, details, "claude.server-side-fallback: false -> true") expectContains(t, details, "api-keys count: 1 -> 2") expectContains(t, details, "claude-api-key count: 1 -> 2") expectContains(t, details, "codex-api-key count: 1 -> 2") @@ -839,3 +840,70 @@ func TestBuildConfigChangeDetails_KVStore(t *testing.T) { t.Fatalf("whitespace-only path difference produced changes: %v", unchanged) } } + +func TestBuildConfigChangeDetails_SensitiveWordsReplace(t *testing.T) { + oldCfg := &config.Config{ + Antigravity: config.AntigravityConfig{ + SensitiveWordsReplace: map[string]string{"Codex": "Assistant"}, + }, + ClaudeKey: []config.ClaudeKey{{ + APIKey: "k1", + Cloak: &config.CloakConfig{SensitiveWordsReplace: map[string]string{"a": "b"}}, + }}, + } + newCfg := &config.Config{ + Antigravity: config.AntigravityConfig{ + SensitiveWordsReplace: map[string]string{"Codex": "Assistant", "GPT-5": "Model"}, + }, + Devin: config.DevinConfig{ + SensitiveWordsReplace: map[string]string{"Claude Code": "Devin"}, + }, + ClaudeKey: []config.ClaudeKey{{ + APIKey: "k1", + Cloak: &config.CloakConfig{SensitiveWordsReplace: map[string]string{"a": "c"}}, + }}, + } + changes := BuildConfigChangeDetails(oldCfg, newCfg) + expectContains(t, changes, "antigravity.sensitive-words-replace: 1 -> 2") + expectContains(t, changes, "devin.sensitive-words-replace: 0 -> 1") + expectContains(t, changes, "claude[0].cloak.sensitive-words-replace: 1 -> 1") + + emptyA := &config.Config{ClaudeKey: []config.ClaudeKey{{APIKey: "k1", Cloak: &config.CloakConfig{}}}} + emptyB := &config.Config{ + Antigravity: config.AntigravityConfig{SensitiveWordsReplace: map[string]string{}}, + Devin: config.DevinConfig{SensitiveWordsReplace: map[string]string{}}, + ClaudeKey: []config.ClaudeKey{{APIKey: "k1", Cloak: &config.CloakConfig{SensitiveWordsReplace: map[string]string{}}}}, + } + if noDiff := BuildConfigChangeDetails(emptyA, emptyB); len(noDiff) != 0 { + t.Fatalf("expected no diff between nil and empty SensitiveWordsReplace maps, got %v", noDiff) + } + + emptyValOld := &config.Config{ + Antigravity: config.AntigravityConfig{SensitiveWordsReplace: map[string]string{"Codex": ""}}, + Devin: config.DevinConfig{SensitiveWordsReplace: map[string]string{"Codex": ""}}, + ClaudeKey: []config.ClaudeKey{{APIKey: "k1", Cloak: &config.CloakConfig{SensitiveWordsReplace: map[string]string{"Codex": ""}}}}, + } + emptyValNew := &config.Config{ + Antigravity: config.AntigravityConfig{SensitiveWordsReplace: map[string]string{"GPT-5": ""}}, + Devin: config.DevinConfig{SensitiveWordsReplace: map[string]string{"GPT-5": ""}}, + ClaudeKey: []config.ClaudeKey{{APIKey: "k1", Cloak: &config.CloakConfig{SensitiveWordsReplace: map[string]string{"GPT-5": ""}}}}, + } + emptyValChanges := BuildConfigChangeDetails(emptyValOld, emptyValNew) + expectContains(t, emptyValChanges, "antigravity.sensitive-words-replace: 1 -> 1") + expectContains(t, emptyValChanges, "devin.sensitive-words-replace: 1 -> 1") + expectContains(t, emptyValChanges, "claude[0].cloak.sensitive-words-replace: 1 -> 1") + + nilCloakOld := &config.Config{ClaudeKey: []config.ClaudeKey{{APIKey: "k1"}}} + addedCloakNew := &config.Config{ClaudeKey: []config.ClaudeKey{{ + APIKey: "k1", + Cloak: &config.CloakConfig{SensitiveWords: []string{"API"}, SensitiveWordsReplace: map[string]string{"a": "b"}}, + }}} + expectContains(t, BuildConfigChangeDetails(nilCloakOld, addedCloakNew), "claude[0].cloak: added") + expectContains(t, BuildConfigChangeDetails(addedCloakNew, nilCloakOld), "claude[0].cloak: removed") + + sameLenWordsNew := &config.Config{ClaudeKey: []config.ClaudeKey{{ + APIKey: "k1", + Cloak: &config.CloakConfig{SensitiveWords: []string{"proxy"}, SensitiveWordsReplace: map[string]string{"a": "b"}}, + }}} + expectContains(t, BuildConfigChangeDetails(addedCloakNew, sameLenWordsNew), "claude[0].cloak.sensitive-words: 1 -> 1") +} diff --git a/internal/watcher/watcher_test.go b/internal/watcher/watcher_test.go index 22b7b909bc6..af05a88c1db 100644 --- a/internal/watcher/watcher_test.go +++ b/internal/watcher/watcher_test.go @@ -990,7 +990,14 @@ func TestHandleEventConfigChangeSchedulesReload(t *testing.T) { w.handleEvent(fsnotify.Event{Name: configPath, Op: fsnotify.Write}) - time.Sleep(400 * time.Millisecond) + // Wait for the debounced reload instead of sleeping a fixed interval; a + // loaded host can fire the debounce timer late. Waiting for the callback + // also guarantees the reload has stopped reading config.yaml before + // TempDir cleanup runs, which Windows requires. + deadline := time.Now().Add(5 * time.Second) + for atomic.LoadInt32(&reloads) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } if atomic.LoadInt32(&reloads) != 1 { t.Fatalf("expected config change to trigger reload once, got %d", reloads) } diff --git a/sdk/api/handlers/claude/code_handlers_cooldown_test.go b/sdk/api/handlers/claude/code_handlers_cooldown_test.go index 33c8e6e37f1..c29353364b8 100644 --- a/sdk/api/handlers/claude/code_handlers_cooldown_test.go +++ b/sdk/api/handlers/claude/code_handlers_cooldown_test.go @@ -224,3 +224,101 @@ func TestClaudeCooldownMessagesRoute(t *testing.T) { }) } } + +type cooldownClaudeTestExecutor struct { + preFTKeepAliveClaudeExecutor +} + +func (*cooldownClaudeTestExecutor) Identifier() string { return "claude" } + +func TestClaudeCooldownMessagesRoute_PassthroughUnifiedRateLimitHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + chunks := make(chan coreexecutor.StreamChunk) + close(chunks) + executor := &cooldownClaudeTestExecutor{preFTKeepAliveClaudeExecutor: preFTKeepAliveClaudeExecutor{chunks: chunks}} + manager := coreauth.NewManager(nil, coreauth.NewSessionAffinitySelector(&coreauth.RoundRobinSelector{}), nil) + manager.RegisterExecutor(executor) + const model = "issue-7417-cooldown-route-model" + const authID = "issue-7417-cooldown-route-auth" + next := time.Now().Add(7823 * time.Second) + auth := &coreauth.Auth{ + ID: authID, + Provider: executor.Identifier(), + Status: coreauth.StatusActive, + Unavailable: true, + NextRetryAfter: next, + Quota: coreauth.QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: next, + Signals: map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-5h-Status": "rejected", + "Anthropic-Ratelimit-Unified-Reset": strconv.FormatInt(next.Add(-10*time.Second).Unix(), 10), + }, + }, + LastError: &coreauth.Error{Code: "rate_limit_error", Message: "upstream quota exceeded", HTTPStatus: http.StatusTooManyRequests}, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, executor.Identifier(), []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + for _, passthrough := range []bool{false, true} { + for _, streaming := range []bool{false, true} { + name := "passthrough_" + strconv.FormatBool(passthrough) + "_stream_" + strconv.FormatBool(streaming) + t.Run(name, func(t *testing.T) { + handler := NewClaudeCodeAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: passthrough}, manager)) + router := gin.New() + router.POST("/v1/messages", handler.ClaudeMessages) + + requestBody := `{"model":"` + model + `","messages":[{"role":"user","content":"hi"}],"stream":` + strconv.FormatBool(streaming) + `}` + request := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(requestBody)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + + before := time.Now().Unix() + router.ServeHTTP(recorder, request) + after := time.Now().Unix() + + if recorder.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want 429; body=%s", recorder.Code, recorder.Body.String()) + } + seconds, errSeconds := strconv.ParseInt(recorder.Header().Get("Retry-After"), 10, 64) + if errSeconds != nil || seconds <= 0 { + t.Fatalf("Retry-After = %q, want positive seconds", recorder.Header().Get("Retry-After")) + } + + if !passthrough { + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-Status"); got != "" { + t.Fatalf("passthrough=false leaked Anthropic-Ratelimit-Unified-Status = %q", got) + } + return + } + + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-Status = %q, want rejected", got) + } + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "five_hour" { + t.Fatalf("Anthropic-Ratelimit-Unified-Representative-Claim = %q, want five_hour", got) + } + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-Overage-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-Overage-Status = %q, want rejected", got) + } + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-5h-Status = %q, want rejected", got) + } + resetAt, errReset := strconv.ParseInt(recorder.Header().Get("Anthropic-Ratelimit-Unified-Reset"), 10, 64) + if errReset != nil || resetAt < before+seconds || resetAt > after+seconds { + t.Fatalf("Anthropic-Ratelimit-Unified-Reset = %q, want unix epoch in [%d, %d]", recorder.Header().Get("Anthropic-Ratelimit-Unified-Reset"), before+seconds, after+seconds) + } + if got := recorder.Header().Get("Anthropic-Ratelimit-Unified-5h-Reset"); got != strconv.FormatInt(resetAt, 10) { + t.Fatalf("Anthropic-Ratelimit-Unified-5h-Reset = %q, want %d", got, resetAt) + } + }) + } + } +} diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 2f1b89a5072..47c9f9969f7 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -635,6 +635,14 @@ func (h *BaseAPIHandler) GetContextWithCancel(handler interfaces.APIHandler, c * if requestCtx != nil { parentCtx = httptransport.InheritDecorator(parentCtx, requestCtx) + parentCtx = logging.InheritDeferredCapture(parentCtx, requestCtx) + } + if c != nil && logging.DeferredAPICaptureFromContext(parentCtx) == nil && !logging.DeferredAPIRequestCaptureDisabled(parentCtx) { + if value, exists := c.Get(logging.DeferredAPIResponseContextKey); exists { + if capture, ok := value.(*logging.DeferredAPICapture); ok && capture != nil { + parentCtx = logging.WithDeferredAPICapture(parentCtx, capture) + } + } } if requestCtx != nil && logging.GetRequestID(parentCtx) == "" { if requestID := logging.GetRequestID(requestCtx); requestID != "" { diff --git a/sdk/api/handlers/openai/openai_responses_steering_integration_test.go b/sdk/api/handlers/openai/openai_responses_steering_integration_test.go index 359e6af24e7..e144f52ea9e 100644 --- a/sdk/api/handlers/openai/openai_responses_steering_integration_test.go +++ b/sdk/api/handlers/openai/openai_responses_steering_integration_test.go @@ -3,6 +3,8 @@ package openai import ( "bytes" "context" + "fmt" + "io" "net/http" "net/http/httptest" "strings" @@ -197,3 +199,300 @@ func TestResponsesSteeringFullDuplexIntegration(t *testing.T) { }) } } + +func TestResponsesWebsocketIsCompatLocalCompactionIntegration_Issue7406(t *testing.T) { + gin.SetMode(gin.TestMode) + const marker = "WS-HANDLER-COMPACT-7406" + for _, steering := range []bool{false, true} { + t.Run(fmt.Sprintf("steering=%t", steering), func(t *testing.T) { + wsReceived := make(chan []byte, 8) + httpReceived := make(chan []byte, 4) + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if websocket.IsWebSocketUpgrade(r) { + conn, errUp := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUp != nil { + t.Errorf("upgrade ws: %v", errUp) + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + for idx := 1; ; idx++ { + _, frame, errRead := conn.ReadMessage() + if errRead != nil { + return + } + wsReceived <- frame + respID := fmt.Sprintf("resp_ws_%d", idx) + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.created","response":{"id":%q,"output":[]}}`, respID))) + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ws-ok-%d"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`, respID, idx))) + } + } + + body, _ := io.ReadAll(r.Body) + httpReceived <- body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_sum\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Compacted summary keeping "+marker+".\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_sum\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":10,\"output_tokens\":4,\"total_tokens\":14}}}\n\n", + ) + })) + defer upstream.Close() + + modelID := fmt.Sprintf("compat-ws-model-%t", steering) + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: upstream.URL, + Models: []config.CodexModel{{Name: modelID, IsCompat: true}}, + }}, + } + cfg.Codex.ResponseSteering = steering + cfg.CodexResponseSteering = steering + + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + manager.RegisterExecutor(runtimeexecutor.NewCodexAutoExecutor(cfg)) + authID := fmt.Sprintf("compat-ws-auth-%t", steering) + if _, errReg := manager.Register(context.Background(), &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{"api_key": "test-compat-key", "base_url": upstream.URL, "websockets": "true"}, + }); errReg != nil { + t.Fatal(errReg) + } + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: modelID, IsCompat: true}}) + defer registry.GetGlobalRegistry().UnregisterClient(authID) + + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager)) + router := gin.New() + router.GET("/v1/responses", h.ResponsesWebsocket) + downstream := httptest.NewServer(router) + defer downstream.Close() + + conn, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(downstream.URL, "http")+"/v1/responses", nil) + if errDial != nil { + t.Fatal(errDial) + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + readUntil := func(wantType string) ojson.Result { + t.Helper() + for { + _, b, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("downstream read waiting for %s: %v", wantType, errRead) + } + parsed := ojson.ParseBytes(b) + gotType := parsed.Get("type").String() + if gotType == wantType { + return parsed + } + if gotType == "error" && wantType != "error" { + t.Fatalf("unexpected error while waiting for %s: %s", wantType, b) + } + } + } + + // Turn 1: First turn is local compaction_trigger; socket must stay open even when steering=true. + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"PRE_COMPACT_RAW_TEXT"}]},{"type":"compaction_trigger"}]}`, modelID))); errWrite != nil { + t.Fatal(errWrite) + } + compactedEv := readUntil("response.completed") + synthID := compactedEv.Get("response.id").String() + if !strings.HasPrefix(synthID, "resp_ag_compact_") { + t.Fatalf("expected synthetic compaction response ID, got %q (%s)", synthID, compactedEv.Raw) + } + <-httpReceived + + // Turn 2a: Bogus synthetic previous_response_id is rejected locally without hitting upstream. + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"previous_response_id":"resp_ag_compact_wrong","input":[{"type":"message","role":"user","content":"next"}]}`, modelID))); errWrite != nil { + t.Fatal(errWrite) + } + bogusErr := readUntil("error") + if bogusErr.Get("error.code").String() != "previous_response_not_found" { + t.Fatalf("expected previous_response_not_found for wrong synthetic ID, got %s", bogusErr.Raw) + } + + // Turn 2b: Valid synthetic previous_response_id starts a fresh root WS turn with expanded summary and NO pre-compaction raw text or synthetic ID. + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","previous_response_id":%q,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Post-compaction question"}]}]}`, synthID))); errWrite != nil { + t.Fatal(errWrite) + } + turn2Done := readUntil("response.completed") + if turn2Done.Get("response.id").String() != "resp_ws_1" { + t.Fatalf("turn 2 response.id = %q, want resp_ws_1", turn2Done.Get("response.id").String()) + } + wsRoot := <-wsReceived + if ojson.GetBytes(wsRoot, "previous_response_id").Exists() { + t.Fatalf("synthetic previous_response_id leaked to upstream WS: %s", wsRoot) + } + if bytes.Contains(wsRoot, []byte("PRE_COMPACT_RAW_TEXT")) { + t.Fatalf("pre-compaction history leaked alongside summary capsule: %s", wsRoot) + } + if bytes.Contains(wsRoot, []byte("cpa-ag-compact-v1:")) || !bytes.Contains(wsRoot, []byte(marker)) { + t.Fatalf("expected expanded summary marker without raw capsule: %s", wsRoot) + } + + // Turn 3: Normal incremental continuation on upstream WS works seamlessly. + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","previous_response_id":"resp_ws_1","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Incremental follow-up"}]}]}`)); errWrite != nil { + t.Fatal(errWrite) + } + turn3Done := readUntil("response.completed") + if turn3Done.Get("response.id").String() != "resp_ws_2" { + t.Fatalf("turn 3 response.id = %q, want resp_ws_2", turn3Done.Get("response.id").String()) + } + wsCont := <-wsReceived + if gotPrev := ojson.GetBytes(wsCont, "previous_response_id").String(); gotPrev != "resp_ws_1" { + t.Fatalf("turn 3 WS previous_response_id = %q, want resp_ws_1; frame=%s", gotPrev, wsCont) + } + }) + } +} + +func TestResponsesWebsocketExistingWSCompactionRejectsStaleParent_Issue7406(t *testing.T) { + gin.SetMode(gin.TestMode) + wsReceived := make(chan []byte, 8) + var wsConnCount atomic.Int32 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if websocket.IsWebSocketUpgrade(r) { + connNum := wsConnCount.Add(1) + conn, errUp := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUp != nil { + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + for turn := 1; ; turn++ { + _, frame, errRead := conn.ReadMessage() + if errRead != nil { + return + } + wsReceived <- frame + respID := fmt.Sprintf("resp_ws_c%d_t%d", connNum, turn) + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.created","response":{"id":%q,"output":[]}}`, respID))) + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":4,"output_tokens":2,"total_tokens":6}}}`, respID))) + } + } + + _, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"msg_sum\",\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Summary after turn 1.\"}]}}\n\n"+ + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_sum\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":8,\"output_tokens\":4,\"total_tokens\":12}}}\n\n", + ) + })) + defer upstream.Close() + + const modelID = "compat-ws-stale-parent-model" + cfg := &config.Config{ + CodexKey: []config.CodexKey{{ + APIKey: "test-compat-key", + BaseURL: upstream.URL, + Models: []config.CodexModel{{Name: modelID, IsCompat: true}}, + }}, + } + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + manager.RegisterExecutor(runtimeexecutor.NewCodexAutoExecutor(cfg)) + const authID = "compat-ws-stale-parent-auth" + if _, errReg := manager.Register(context.Background(), &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{"api_key": "test-compat-key", "base_url": upstream.URL, "websockets": "true"}, + }); errReg != nil { + t.Fatal(errReg) + } + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: modelID, IsCompat: true}}) + defer registry.GetGlobalRegistry().UnregisterClient(authID) + + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager)) + router := gin.New() + router.GET("/v1/responses", h.ResponsesWebsocket) + downstream := httptest.NewServer(router) + defer downstream.Close() + + conn, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(downstream.URL, "http")+"/v1/responses", nil) + if errDial != nil { + t.Fatal(errDial) + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + readUntil := func(wantType string) ojson.Result { + t.Helper() + for { + _, b, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("downstream read waiting for %s: %v", wantType, errRead) + } + parsed := ojson.ParseBytes(b) + gotType := parsed.Get("type").String() + if gotType == wantType { + return parsed + } + if gotType == "error" && wantType != "error" { + t.Fatalf("unexpected error while waiting for %s: %s", wantType, b) + } + } + } + + // Turn 1: Real upstream WS response -> resp_ws_c1_t1. + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","role":"user","content":"turn 1"}]}`, modelID))) + turn1Done := readUntil("response.completed") + oldWSRespID := turn1Done.Get("response.id").String() + if oldWSRespID != "resp_ws_c1_t1" { + t.Fatalf("turn 1 response.id = %q, want resp_ws_c1_t1", oldWSRespID) + } + <-wsReceived + + // Turn 2: Full local compaction_trigger invalidates c1 and establishes compaction boundary. + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","role":"user","content":"turn 1"},{"type":"compaction_trigger"}]}`, modelID))) + compactDone := readUntil("response.completed") + synthID := compactDone.Get("response.id").String() + if !strings.HasPrefix(synthID, "resp_ag_compact_") { + t.Fatalf("turn 2 response.id = %q, want resp_ag_compact_*", synthID) + } + + // Turn 3a: Stale pre-compaction WS response ID (resp_ws_c1_t1) MUST be rejected with 409 previous_response_not_found, NOT silently grafted onto synthID. + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"previous_response_id":%q,"input":[{"type":"message","role":"user","content":"wrong branch"}]}`, modelID, oldWSRespID))) + staleErr := readUntil("error") + if staleErr.Get("status").Int() != http.StatusConflict || staleErr.Get("error.code").String() != "previous_response_not_found" { + t.Fatalf("expected 409 previous_response_not_found for pre-compaction WS ID %q, got %s", oldWSRespID, staleErr.Raw) + } + + // Turn 3b: Valid continuation referencing synthID dials fresh WS c2 with expanded summary and no previous_response_id. + _ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","previous_response_id":%q,"input":[{"type":"message","role":"user","content":"clean post-compact turn"}]}`, synthID))) + turn3Done := readUntil("response.completed") + if turn3Done.Get("response.id").String() != "resp_ws_c2_t1" { + t.Fatalf("turn 3b response.id = %q, want resp_ws_c2_t1", turn3Done.Get("response.id").String()) + } + wsTurn3 := <-wsReceived + if ojson.GetBytes(wsTurn3, "previous_response_id").Exists() { + t.Fatalf("unexpected previous_response_id on fresh WS root: %s", wsTurn3) + } + if !bytes.Contains(wsTurn3, []byte("Summary after turn 1.")) { + t.Fatalf("expected expanded summary in Turn 3b WS root: %s", wsTurn3) + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 41e6a373af1..125f2045e86 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -133,8 +133,8 @@ func newResponsesWebsocketWriter(conn *websocket.Conn) *responsesWebsocketWriter } // closeForUpstreamError sends a best-effort close frame without waiting behind -// an active downstream data writer. If a data write already owns writeMu, the -// connection is closed immediately so the blocked writer and session can exit. +// an active downstream writer. If a data or Ping write already owns writeMu, +// the connection is closed immediately so the blocked writer and session can exit. func (w *responsesWebsocketWriter) closeForUpstreamError(err error) (bool, error) { if w == nil || w.conn == nil { return false, nil @@ -173,7 +173,15 @@ func (w *responsesWebsocketWriter) writePing() error { if w == nil || w.conn == nil { return errors.New("responses websocket: writer is nil") } - w.writeMu.Lock() + if w.closing.Load() { + return websocket.ErrCloseSent + } + if !w.writeMu.TryLock() { + if w.closing.Load() { + return websocket.ErrCloseSent + } + return nil + } defer w.writeMu.Unlock() if w.closing.Load() { return websocket.ErrCloseSent @@ -181,6 +189,56 @@ func (w *responsesWebsocketWriter) writePing() error { return w.conn.WriteControl(websocket.PingMessage, nil, time.Time{}) } +func (w *responsesWebsocketWriter) keepAlive(ctx context.Context, cancel context.CancelCauseFunc, intervals <-chan time.Duration) { + if errPing := runResponsesWebsocketKeepAlive(ctx, intervals, w.writePing); errPing != nil { + cancel(errPing) + _, _ = w.closeWithoutError() + } +} + +func runResponsesWebsocketKeepAlive(ctx context.Context, intervals <-chan time.Duration, ping func() error) error { + var ticker *time.Ticker + var ticks <-chan time.Time + defer func() { + if ticker != nil { + ticker.Stop() + } + }() + for { + select { + case <-ctx.Done(): + return nil + case interval := <-intervals: + if interval <= 0 { + if ticker != nil { + ticker.Stop() + } + ticks = nil + continue + } + if ticker == nil { + ticker = time.NewTicker(interval) + } else { + ticker.Reset(interval) + } + ticks = ticker.C + case <-ticks: + if errPing := ping(); errPing != nil { + return errPing + } + } + } +} + +func queueResponsesWebsocketKeepAliveInterval(intervals chan time.Duration, interval time.Duration) { + // A single producer replaces a pending interval without waiting for a Ping write. + select { + case <-intervals: + default: + } + intervals <- interval +} + func (w *responsesWebsocketWriter) closeWithPayload(payload []byte) (bool, error) { if w == nil || w.conn == nil { return false, nil @@ -374,6 +432,25 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { defer cancelSocket(nil) c.Request = c.Request.WithContext(socketCtx) upgradeRequest := c.Request.Clone(socketCtx) + // Start Ping at upgrade, before opening an upstream stream can block. + var keepAliveIntervals chan time.Duration + var keepAliveInterval time.Duration + updateKeepAlive := func() { + if h == nil { + return + } + interval := handlers.StreamingKeepAliveInterval(h.CurrentConfig()) + if interval == keepAliveInterval { + return + } + keepAliveInterval = interval + if keepAliveIntervals == nil { + keepAliveIntervals = make(chan time.Duration, 1) + go writer.keepAlive(socketCtx, cancelSocket, keepAliveIntervals) + } + queueResponsesWebsocketKeepAliveInterval(keepAliveIntervals, interval) + } + updateKeepAlive() // The reader owns the client socket so an in-flight interrupt is not stuck // behind the response currently being forwarded. requestLogEnabled := false @@ -472,6 +549,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { lastResponseOutput := []byte("[]") var observedCompaction responsesWebsocketObservedCompactionState lastResponseID := "" + localCompactionBoundary := false // Remains pending until a generating request commits successfully. pendingPrewarmID := "" var lastResponsePendingToolCallIDs []string @@ -568,6 +646,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { if msgType != websocket.TextMessage && msgType != websocket.BinaryMessage { continue } + updateKeepAlive() // log.Infof( // "responses websocket: downstream_in id=%s type=%d event=%s payload=%s", // passthroughSessionID, @@ -776,6 +855,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName) } else if len(lastRequest) == 0 && previousResponseID != "" { errMsg = responsesWebsocketPreviousResponseNotFoundError() + } else if previousResponseID != "" && (localCompactionBoundary || strings.HasPrefix(previousResponseID, "resp_ag_compact_")) && previousResponseID != strings.TrimSpace(lastResponseID) { + errMsg = responsesWebsocketPreviousResponseNotFoundError() } else { requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState( payload, @@ -828,6 +909,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { lastResponseOutput = []byte("[]") observedCompaction.clear() lastResponseID = "" + localCompactionBoundary = false lastResponsePendingToolCallIDs = nil prewarmID, errWrite := writeResponsesWebsocketSyntheticPrewarm(c, writer, requestJSON, wsTimelineLog, passthroughSessionID) if errWrite != nil { @@ -860,6 +942,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { turnStartEpoch := inputEpoch.Load() cliCtx, cliCancel := h.GetContextWithCancel(h, c, executionParent) cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx) + cliCtx = cliproxyexecutor.WithLocalWebsocketTurnTracker(cliCtx) cliCtx = cliproxyexecutor.WithWebsocketInputEpoch(cliCtx, &inputEpoch, msgEpoch) if duplexInput != nil { cliCtx = cliproxyexecutor.WithWebsocketInput(cliCtx, duplexInput) @@ -889,6 +972,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { preserveNativeOutput.Store(false) codexDuplexStream.Store(false) + cliproxyexecutor.ResetLocalWebsocketTurn(cliCtx) authID = strings.TrimSpace(authID) if authID == "" || h == nil || h.AuthManager == nil { return @@ -964,12 +1048,14 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { passthroughSessionID, responsesWebsocketForwardOptions{ preserveCompletionOutput: preserveNativeOutput.Load, - duplexStream: codexDuplexStream.Load, - toolCacheTurn: toolCacheTurn, - suppressError: replayPinnedAuthFailure, - localInterrupt: localInterrupt, + duplexStream: func() bool { + return codexDuplexStream.Load() && !cliproxyexecutor.LocalWebsocketTurnExecuted(cliCtx) + }, + toolCacheTurn: toolCacheTurn, + suppressError: replayPinnedAuthFailure, + localInterrupt: localInterrupt, websocketOwned: func() bool { - return attemptedUpstreamMode == responsesWebsocketUpstreamModeWS + return attemptedUpstreamMode == responsesWebsocketUpstreamModeWS && !cliproxyexecutor.LocalWebsocketTurnExecuted(cliCtx) }, }, ) @@ -1020,6 +1106,15 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { toolCacheTurn.commit() pendingPrewarmID = "" + localTurnExecuted := cliproxyexecutor.LocalWebsocketTurnExecuted(cliCtx) + if localTurnExecuted { + attemptedUpstreamMode = responsesWebsocketUpstreamModeHTTP + if lastAttemptedAuthID != "" { + rememberPinnedAuth(lastAttemptedAuthID, modelName) + } + nextLastRequest = responsesWebsocketLocalCompactionBoundaryRequest(requestJSON) + } + localCompactionBoundary = localTurnExecuted upstreamMode = attemptedUpstreamMode if upstreamMode == responsesWebsocketUpstreamModeWS { upstreamWebsocketAuthID = lastAttemptedAuthID @@ -1093,6 +1188,17 @@ func (s *responsesWebsocketObservedCompactionState) clear() { *s = responsesWebsocketObservedCompactionState{} } +func responsesWebsocketLocalCompactionBoundaryRequest(requestJSON []byte) []byte { + if len(requestJSON) == 0 { + return []byte(`{"input":[]}`) + } + doc := ojson.NewDoc(requestJSON) + _ = doc.Delete("type") + _ = doc.Delete("previous_response_id") + _ = doc.SetRawBytes("input", []byte("[]")) + return doc.Bytes() +} + func responsesWebsocketHTTPReplayRequiredError() error { return cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } diff --git a/sdk/api/handlers/openai/openai_responses_websocket_forward.go b/sdk/api/handlers/openai/openai_responses_websocket_forward.go index 0dfd7702443..8e5a60f34dd 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_forward.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_forward.go @@ -28,7 +28,6 @@ type responsesWebsocketForwardOptions struct { duplexStream func() bool toolCacheTurn *responsesWebsocketToolCacheTurn suppressError func(*interfaces.ErrorMessage) bool - keepAliveInterval *time.Duration localInterrupt *responsesLocalInterrupt // websocketOwned reports whether this turn runs on an upstream // websocket. While true, no interrupt may claim the local HTTP path. @@ -66,20 +65,6 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( downstreamSessionKey = websocketDownstreamSessionKey(c.Request) } - var keepAliveTicker *time.Ticker - var keepAliveC <-chan time.Time - keepAliveInterval := time.Duration(0) - if h != nil { - keepAliveInterval = handlers.StreamingKeepAliveInterval(h.CurrentConfig()) - } - if opts.keepAliveInterval != nil { - keepAliveInterval = *opts.keepAliveInterval - } - if keepAliveInterval > 0 { - keepAliveTicker = time.NewTicker(keepAliveInterval) - defer keepAliveTicker.Stop() - keepAliveC = keepAliveTicker.C - } if opts.localInterrupt != nil { websocketOwned := opts.websocketOwned != nil && opts.websocketOwned() opts.localInterrupt.begin(websocketOwned) @@ -155,13 +140,9 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } return completeResponsesWebsocketLocalInterrupt(writer, wsTimelineLog, cancel, interruptPayload, outputItemsByIndex, outputItemsFallback, pendingToolCallIDs) case <-c.Request.Context().Done(): - cancel(c.Request.Context().Err()) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, c.Request.Context().Err() - case <-keepAliveC: - if errPing := writer.writePing(); errPing != nil { - cancel(errPing) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errPing - } + cause := context.Cause(c.Request.Context()) + cancel(cause) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, cause case errMsg, ok := <-errs: if !ok { errs = nil @@ -212,9 +193,6 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( cancel(nil) return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, nil } - if keepAliveTicker != nil && keepAliveInterval > 0 { - keepAliveTicker.Reset(keepAliveInterval) - } // A queued local interrupt belongs to an in-flight HTTP turn: // consume it before forwarding more upstream payload. Once the diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index c2f60c9a085..4a3b8de3fed 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1,6 +1,7 @@ package openai import ( + "bufio" "bytes" "context" "encoding/json" @@ -8,6 +9,7 @@ import ( "fmt" "io" "maps" + "net" "net/http" "net/http/httptest" "runtime" @@ -16,6 +18,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "unicode/utf8" @@ -7502,195 +7505,111 @@ func TestNormalizeSubsequentRequestAssistantInputTriggersTranscriptReplacement(t } } -func TestForwardResponsesWebsocketEmitsPeriodicPingControlFrames(t *testing.T) { - gin.SetMode(gin.TestMode) - - serverErrCh := make(chan error, 1) - data := make(chan []byte) - errCh := make(chan *interfaces.ErrorMessage) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) - if err != nil { - serverErrCh <- err - return - } - defer func() { _ = conn.Close() }() +type websocketKeepAliveBootstrapExecutor struct { + websocketCaptureExecutor + started chan struct{} + release chan struct{} +} - ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) - ctx.Request = r +func (e *websocketKeepAliveBootstrapExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + close(e.started) + select { + case <-e.release: + return e.websocketCaptureExecutor.ExecuteStream(ctx, auth, req, opts) + case <-ctx.Done(): + return nil, ctx.Err() + } +} - cfg := &sdkconfig.SDKConfig{ - Streaming: sdkconfig.StreamingConfig{ - KeepAliveSeconds: 1, - }, - } - h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(cfg, nil)) +func TestResponsesWebsocketPingsBeforeUpstreamStreamIsReady(t *testing.T) { + gin.SetMode(gin.TestMode) + executor := &websocketKeepAliveBootstrapExecutor{ + started: make(chan struct{}), + release: make(chan struct{}), + } + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ID: "auth-ws-keepalive", Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register auth: %v", err) + } + const model = "ws-keepalive-test-model" + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) - _, _, _, errMsg, errForward := h.forwardResponsesWebsocket( - ctx, - newResponsesWebsocketWriter(conn), - func(...interface{}) {}, - data, - errCh, - newInMemoryWebsocketTimelineLog(), - "session-keepalive-test", - ) - if errMsg != nil { - serverErrCh <- fmt.Errorf("unexpected error message: %v", errMsg.Error) - return - } - if errForward != nil { - serverErrCh <- fmt.Errorf("unexpected forward error: %v", errForward) - return - } - serverErrCh <- nil - })) + cfg := &sdkconfig.SDKConfig{Streaming: sdkconfig.StreamingConfig{KeepAliveSeconds: 1}} + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(cfg, manager)) + router := gin.New() + router.GET("/v1/responses", h.ResponsesWebsocket) + server := httptest.NewServer(router) defer server.Close() - wsURL := "ws" + strings.TrimPrefix(server.URL, "http") - clientConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("dial websocket: %v", err) } - defer func() { _ = clientConn.Close() }() + defer func() { _ = conn.Close() }() - pingReceived := make(chan struct{}, 1) - clientConn.SetPingHandler(func(appData string) error { + pings := make(chan bool, 1) + pong := conn.PingHandler() + conn.SetPingHandler(func(payload string) error { + started := false select { - case pingReceived <- struct{}{}: + case <-executor.started: + started = true default: } - return clientConn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) + select { + case pings <- started: + default: + } + return pong(payload) }) - clientDone := make(chan struct{}) + completed := make(chan error, 1) go func() { - defer close(clientDone) - for { - _, _, errRead := clientConn.ReadMessage() - if errRead != nil { - return - } + _, payload, errRead := conn.ReadMessage() + if errRead == nil && ojson.GetBytesNoCopy(payload, "type").String() != wsEventTypeCompleted { + errRead = fmt.Errorf("expected response.completed, got %s", payload) } + completed <- errRead }() select { - case <-pingReceived: - // Received expected Ping control frame while upstream is waiting. - case <-time.After(1500 * time.Millisecond): - t.Fatal("expected websocket Ping control frame during upstream wait, got none") - } - - // Unblock forwardResponsesWebsocket with terminal completion. - data <- []byte(`{"type":"response.done","response":{"id":"resp-ping-1","output":[]}}`) - close(data) - close(errCh) - - select { - case serverErr := <-serverErrCh: - if serverErr != nil { - t.Fatalf("server error: %v", serverErr) - } - case <-time.After(2 * time.Second): - t.Fatal("server timed out completing forwardResponsesWebsocket") + case <-pings: + case errRead := <-completed: + t.Fatalf("websocket ended before the first turn: %v", errRead) + case <-time.After(5 * time.Second): + t.Fatal("expected Ping after upgrade, before sending a turn") } - - <-clientDone -} - -func TestForwardResponsesWebsocketPingKeepAliveOptionsOverride(t *testing.T) { - gin.SetMode(gin.TestMode) - - serverErrCh := make(chan error, 1) - data := make(chan []byte) - errCh := make(chan *interfaces.ErrorMessage) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) - if err != nil { - serverErrCh <- err - return - } - defer func() { _ = conn.Close() }() - - ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) - ctx.Request = r - - interval := 20 * time.Millisecond - h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)) - - _, _, _, errMsg, errForward := h.forwardResponsesWebsocket( - ctx, - newResponsesWebsocketWriter(conn), - func(...interface{}) {}, - data, - errCh, - newInMemoryWebsocketTimelineLog(), - "session-keepalive-override", - responsesWebsocketForwardOptions{ - keepAliveInterval: &interval, - }, - ) - if errMsg != nil { - serverErrCh <- fmt.Errorf("unexpected error message: %v", errMsg.Error) - return - } - if errForward != nil { - serverErrCh <- fmt.Errorf("unexpected forward error: %v", errForward) - return - } - serverErrCh <- nil - })) - defer server.Close() - - wsURL := "ws" + strings.TrimPrefix(server.URL, "http") - clientConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) - if err != nil { - t.Fatalf("dial websocket: %v", err) + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-keepalive-test-model","input":[]}`)); errWrite != nil { + t.Fatalf("write websocket request: %v", errWrite) } - defer func() { _ = clientConn.Close() }() - - pingReceived := make(chan struct{}, 1) - clientConn.SetPingHandler(func(appData string) error { + deadline := time.After(5 * time.Second) +waitForBootstrapPing: + for { select { - case pingReceived <- struct{}{}: - default: - } - return clientConn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) - }) - - clientDone := make(chan struct{}) - go func() { - defer close(clientDone) - for { - _, _, errRead := clientConn.ReadMessage() - if errRead != nil { - return + case started := <-pings: + if started { + break waitForBootstrapPing } + case errRead := <-completed: + t.Fatalf("websocket ended before the upstream stream was released: %v", errRead) + case <-deadline: + t.Fatal("expected Ping while the upstream stream is blocked") } - }() - - select { - case <-pingReceived: - case <-time.After(500 * time.Millisecond): - t.Fatal("expected websocket Ping control frame via option override, got none") } - data <- []byte(`{"type":"response.done","response":{"id":"resp-override","output":[]}}`) - close(data) - close(errCh) - + close(executor.release) select { - case serverErr := <-serverErrCh: - if serverErr != nil { - t.Fatalf("server error: %v", serverErr) + case errRead := <-completed: + if errRead != nil { + t.Fatalf("read completed response: %v", errRead) } - case <-time.After(2 * time.Second): - t.Fatal("server timed out completing forwardResponsesWebsocket") + case <-time.After(5 * time.Second): + t.Fatal("expected response.completed after releasing the upstream stream") } - - <-clientDone } func TestResponsesWebsocketWriterWritePing(t *testing.T) { @@ -7709,6 +7628,23 @@ func TestResponsesWebsocketWriterWritePing(t *testing.T) { defer func() { _ = conn.Close() }() writer := newResponsesWebsocketWriter(conn) + // A background Ping must skip rather than block behind an active data write. + writer.writeMu.Lock() + pingDone := make(chan error, 1) + go func() { pingDone <- writer.writePing() }() + select { + case errPing := <-pingDone: + writer.writeMu.Unlock() + if errPing != nil { + t.Errorf("writePing() while data write active = %v, want nil", errPing) + } + case <-time.After(time.Second): + writer.writeMu.Unlock() + <-pingDone + t.Error("Ping waited for the data write mutex") + return + } + if errPing := writer.writePing(); errPing != nil { t.Errorf("writePing() error = %v, want nil", errPing) } @@ -7751,81 +7687,275 @@ func TestResponsesWebsocketWriterWritePing(t *testing.T) { } } -func TestForwardResponsesWebsocketPingWriteFailureAbortsSession(t *testing.T) { - gin.SetMode(gin.TestMode) +type blockingWebsocketWriteConn struct { + net.Conn + blockWrites atomic.Bool + writeBlock chan struct{} + closed chan struct{} + closeOnce sync.Once +} - serverErrCh := make(chan error, 1) - cancelledCh := make(chan error, 1) - data := make(chan []byte) - errCh := make(chan *interfaces.ErrorMessage) +func newBlockingWebsocketWriteConn(conn net.Conn) *blockingWebsocketWriteConn { + return &blockingWebsocketWriteConn{ + Conn: conn, + writeBlock: make(chan struct{}, 1), + closed: make(chan struct{}), + } +} - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) - if err != nil { - serverErrCh <- err - return +func (c *blockingWebsocketWriteConn) Write(b []byte) (int, error) { + if c.blockWrites.Load() { + select { + case c.writeBlock <- struct{}{}: + default: } - defer func() { _ = conn.Close() }() + <-c.closed + return 0, net.ErrClosed + } + return c.Conn.Write(b) +} - ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) - ctx.Request = r +func (c *blockingWebsocketWriteConn) Close() error { + c.closeOnce.Do(func() { + close(c.closed) + }) + return c.Conn.Close() +} - interval := 10 * time.Millisecond - h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)) +type hijackableWebsocketResponseWriter struct { + header http.Header + conn net.Conn +} - _, _, _, errMsg, errForward := h.forwardResponsesWebsocket( - ctx, - newResponsesWebsocketWriter(conn), - func(errs ...interface{}) { - if len(errs) > 0 { - if errVal, ok := errs[0].(error); ok { - cancelledCh <- errVal - return - } +func newHijackableWebsocketResponseWriter(conn net.Conn) *hijackableWebsocketResponseWriter { + return &hijackableWebsocketResponseWriter{ + header: make(http.Header), + conn: conn, + } +} + +func (w *hijackableWebsocketResponseWriter) Header() http.Header { return w.header } +func (w *hijackableWebsocketResponseWriter) Write(b []byte) (int, error) { return w.conn.Write(b) } +func (w *hijackableWebsocketResponseWriter) WriteHeader(int) {} +func (w *hijackableWebsocketResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.conn, bufio.NewReadWriter(bufio.NewReader(w.conn), bufio.NewWriter(w.conn)), nil +} + +func TestResponsesWebsocketWriterCloseInterruptsBlockedPing(t *testing.T) { + cases := []struct { + name string + close func(*responsesWebsocketWriter) error + }{ + { + name: "closeForUpstreamError", + close: func(w *responsesWebsocketWriter) error { + matched, errClose := w.closeForUpstreamError(&websocket.CloseError{ + Code: websocket.CloseMessageTooBig, + Text: "message too big", + }) + if !matched && errClose == nil { + return errors.New("message-too-big error did not match") } - cancelledCh <- nil + return errClose }, - data, - errCh, - newInMemoryWebsocketTimelineLog(), - "session-ping-fail", - responsesWebsocketForwardOptions{ - keepAliveInterval: &interval, + }, + { + name: "closeWithPayload", + close: func(w *responsesWebsocketWriter) error { + _, errClose := w.closeWithPayload([]byte(`{"type":"error","status":400}`)) + return errClose }, - ) - if errMsg != nil { - serverErrCh <- fmt.Errorf("unexpected error message: %v", errMsg.Error) + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + clientRaw, serverRaw := net.Pipe() + defer func() { _ = clientRaw.Close() }() + blockingServer := newBlockingWebsocketWriteConn(serverRaw) + defer func() { _ = blockingServer.Close() }() + + serverConnCh := make(chan *websocket.Conn, 1) + serverErrCh := make(chan error, 1) + go func() { + req, errReq := http.NewRequest(http.MethodGet, "http://example.com/v1/responses", nil) + if errReq != nil { + serverErrCh <- errReq + return + } + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Sec-WebSocket-Version", "13") + req.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + rw := newHijackableWebsocketResponseWriter(blockingServer) + conn, errUpgrade := responsesWebsocketUpgrader.Upgrade(rw, req, nil) + if errUpgrade != nil { + serverErrCh <- errUpgrade + return + } + serverConnCh <- conn + }() + + // Drain the HTTP 101 Switching Protocols response on the client side of the pipe. + buf := make([]byte, 512) + if _, errRead := clientRaw.Read(buf); errRead != nil { + t.Fatalf("read upgrade response: %v", errRead) + } + + var serverConn *websocket.Conn + select { + case serverConn = <-serverConnCh: + case errUpgrade := <-serverErrCh: + t.Fatalf("upgrade websocket: %v", errUpgrade) + case <-time.After(time.Second): + t.Fatal("timed out waiting for websocket upgrade") + } + + writer := newResponsesWebsocketWriter(serverConn) + blockingServer.blockWrites.Store(true) + + pingDone := make(chan error, 1) + go func() { + pingDone <- writer.writePing() + }() + + select { + case <-blockingServer.writeBlock: + case <-time.After(time.Second): + t.Fatal("timed out waiting for writePing to enter blocked network write") + } + + closeDone := make(chan error, 1) + go func() { + closeDone <- tc.close(writer) + }() + + select { + case errClose := <-closeDone: + if errClose != nil { + t.Fatalf("close returned error: %v", errClose) + } + case <-time.After(time.Second): + t.Fatal("terminal close blocked behind stuck Ping write") + } + + select { + case errPing := <-pingDone: + if !errors.Is(errPing, net.ErrClosed) { + t.Fatalf("writePing error = %v, want net.ErrClosed", errPing) + } + case <-time.After(time.Second): + t.Fatal("writePing did not exit after connection close") + } + }) + } +} + +func TestResponsesWebsocketKeepAliveUpdatesInterval(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + intervals := make(chan time.Duration, 1) + queueResponsesWebsocketKeepAliveInterval(intervals, time.Second) + queueResponsesWebsocketKeepAliveInterval(intervals, 20*time.Millisecond) + + pings := 0 + done := make(chan error, 1) + go func() { + done <- runResponsesWebsocketKeepAlive(ctx, intervals, func() error { + pings++ + return nil + }) + }() + + synctest.Wait() + time.Sleep(20 * time.Millisecond) + synctest.Wait() + if pings != 1 { + t.Fatalf("pings after replacing pending interval = %d, want 1", pings) + } + + queueResponsesWebsocketKeepAliveInterval(intervals, 10*time.Millisecond) + synctest.Wait() + time.Sleep(10 * time.Millisecond) + synctest.Wait() + if pings != 2 { + t.Fatalf("pings after shortening interval = %d, want 2", pings) + } + + queueResponsesWebsocketKeepAliveInterval(intervals, 0) + synctest.Wait() + time.Sleep(time.Second) + synctest.Wait() + if pings != 2 { + t.Fatalf("pings after disabling keepalive = %d, want 2", pings) + } + + queueResponsesWebsocketKeepAliveInterval(intervals, 10*time.Millisecond) + synctest.Wait() + time.Sleep(10 * time.Millisecond) + synctest.Wait() + if pings != 3 { + t.Fatalf("pings after re-enabling keepalive = %d, want 3", pings) + } + + cancel() + if err := <-done; err != nil { + t.Fatalf("keepalive did not shut down cleanly: %v", err) + } + }) +} + +func TestForwardResponsesWebsocketPreservesPingFailureCause(t *testing.T) { + requestCtx, cancelSocket := context.WithCancelCause(context.Background()) + cancelSocket(net.ErrClosed) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil).WithContext(requestCtx) + + var canceledWith error + _, _, _, _, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( + c, nil, func(args ...interface{}) { + canceledWith, _ = args[0].(error) + }, nil, nil, nil, "ping-write-failure", + ) + if !errors.Is(err, net.ErrClosed) || !errors.Is(canceledWith, net.ErrClosed) { + t.Fatalf("forward error = %v, cancellation cause = %v, want closed connection", err, canceledWith) + } +} + +func TestResponsesWebsocketKeepAliveCancelsOnPingFailure(t *testing.T) { + serverErr := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErr <- err return } - serverErrCh <- errForward + _ = conn.Close() + + ctx, cancel := context.WithCancelCause(r.Context()) + defer cancel(nil) + intervals := make(chan time.Duration, 1) + intervals <- 10 * time.Millisecond + newResponsesWebsocketWriter(conn).keepAlive(ctx, cancel, intervals) + serverErr <- context.Cause(ctx) })) defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") - clientConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("dial websocket: %v", err) } - - // Close client connection immediately so the next server Ping write fails. - _ = clientConn.Close() + defer func() { _ = conn.Close() }() select { - case serverErr := <-serverErrCh: - if serverErr == nil { - t.Fatal("expected error on server ping write failure, got nil") + case errKeepAlive := <-serverErr: + if !errors.Is(errKeepAlive, net.ErrClosed) { + t.Fatalf("keepalive cancellation cause = %v, want closed connection", errKeepAlive) } case <-time.After(2 * time.Second): - t.Fatal("server timed out awaiting ping write abort") - } - - select { - case cancelErr := <-cancelledCh: - if cancelErr == nil { - t.Fatal("expected cancel callback to be invoked with ping error, got nil") - } - case <-time.After(time.Second): - t.Fatal("timed out awaiting cancel callback on ping write failure") + t.Fatal("timed out awaiting cancellation after Ping failed") } } diff --git a/sdk/cliproxy/antigravity_models_registration_test.go b/sdk/cliproxy/antigravity_models_registration_test.go index d3d01b1006e..2149d3a4b30 100644 --- a/sdk/cliproxy/antigravity_models_registration_test.go +++ b/sdk/cliproxy/antigravity_models_registration_test.go @@ -757,14 +757,14 @@ func TestAntigravityBlockedHintPersistenceDoesNotBlockPublicationOrOtherAccounts svc.refreshAntigravityModels(t.Context()) select { case <-authBProbed: - case <-time.After(time.Second): + case <-time.After(5 * time.Second): cancelLifetime() svc.WaitAntigravityProbes() t.Fatal("stalled persistence on authA blocked authB probe") } select { case <-authARefreshed: - case <-time.After(time.Second): + case <-time.After(5 * time.Second): cancelLifetime() svc.WaitAntigravityProbes() t.Fatal("stalled persistence on authA blocked subsequent probe for authA") @@ -778,7 +778,7 @@ func TestAntigravityBlockedHintPersistenceDoesNotBlockPublicationOrOtherAccounts // (which happens after applyAntigravityModelHints published authA's updated catalog). select { case <-store.saveEntered: - case <-time.After(time.Second): + case <-time.After(5 * time.Second): cancelLifetime() svc.WaitAntigravityProbes() t.Fatal("subsequent persistence task did not start after registry publication") @@ -800,7 +800,7 @@ func TestAntigravityBlockedHintPersistenceDoesNotBlockPublicationOrOtherAccounts svc.WaitAntigravityProbes() t.Fatal(errUpdateB) } - case <-time.After(time.Second): + case <-time.After(5 * time.Second): cancelLifetime() svc.WaitAntigravityProbes() t.Fatal("persisting Update on authB blocked while store Save was stalled on authA") @@ -837,7 +837,7 @@ func TestAntigravityBlockedHintPersistenceDoesNotBlockPublicationOrOtherAccounts cancelB() select { case <-canceledUpdateBDone: - case <-time.After(time.Second): + case <-time.After(5 * time.Second): close(store.releaseSave) cancelLifetime() svc.WaitAntigravityProbes() @@ -864,8 +864,9 @@ func TestAntigravityBlockedHintPersistenceDoesNotBlockPublicationOrOtherAccounts }() select { case <-doneWait: - case <-time.After(time.Second): - close(store.releaseSave) + case <-time.After(5 * time.Second): + // releaseSave was already closed above; closing it again would panic and + // hide this failure message. t.Fatal("service context cancellation did not cancel blocked metadata persistence") } } diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 8347f38ba28..4d7411df006 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "maps" "net" "net/http" "regexp" @@ -290,10 +291,17 @@ func (m *Manager) clearDisabledCooldownStates(cfg *internalconfig.Config) bool { if auth == nil { continue } - if !quotaCooldownDisabledForAuthWithConfig(auth, cfg) && !auth.Disabled && auth.Status != StatusDisabled { + isDisabled := auth.Disabled || auth.Status == StatusDisabled + if !isDisabled && !quotaCooldownDisabledForAuthWithConfig(auth, cfg) { continue } - if clearCooldownStateForAuth(auth, now) { + var cleared bool + if isDisabled { + cleared = clearCooldownStateForAuth(auth, now) + } else { + cleared = clearDisabledCooldownStateForAuth(auth, now) + } + if cleared { snapshots = append(snapshots, auth.Clone()) } } @@ -352,6 +360,13 @@ func (m *Manager) RestoreCooldownStates(ctx context.Context) error { } } } + for authID := range snapshotsByID { + if auth := m.auths[authID]; auth != nil && m.cooldownDisabledForAuth(auth) { + if clearDisabledCooldownStateForAuth(auth, now) { + snapshotsByID[authID] = auth.Clone() + } + } + } m.mu.Unlock() if m.scheduler != nil { @@ -369,7 +384,7 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti return false } auth := m.auths[authID] - if auth == nil || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) || hasUnauthorizedAuthFailure(auth) { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled || (m.cooldownDisabledForAuth(auth) && !isForceCooldownRecordAt(record, now)) || hasUnauthorizedAuthFailure(auth) { return false } updatedAt := record.UpdatedAt @@ -382,19 +397,101 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti if quota.Exceeded && quota.NextRecoverAt.IsZero() { quota.NextRecoverAt = record.NextRetryAfter } + quota.cooldownEvaluated = quota.Exceeded + if executorKeyFromAuth(auth) == "claude" || strings.EqualFold(strings.TrimSpace(record.Provider), "claude") { + quota.cooldownSignals = sanitizePersistedClaudeCooldownSignals(record.CooldownSignals) + } + seedInitialForceCooldownDeadlines(auth) + outerForceUntil := recordOuterForceUntil(record) + indepForceUntil := recordIndependentForceUntil(record) if model == "" { + prevNextRetry := auth.NextRetryAfter + prevQuota := auth.Quota.Clone() + prevUpdatedAt := auth.UpdatedAt + prevLastErr := cloneError(auth.LastError) + prevStatusMsg := auth.StatusMessage + hadLiveAuthBlock := auth.Unavailable && (prevNextRetry.After(now) || (prevQuota.Exceeded && prevQuota.NextRecoverAt.After(now))) + liveNewer := hadLiveAuthBlock && prevUpdatedAt.After(updatedAt) + auth.Unavailable = true auth.Status = StatusError - auth.NextRetryAfter = record.NextRetryAfter - applyCooldownFields(&auth.Quota, quota) + if prevNextRetry.After(now) && prevNextRetry.After(record.NextRetryAfter) { + auth.NextRetryAfter = prevNextRetry + } else { + auth.NextRetryAfter = record.NextRetryAfter + } + mergedQuota := quota + prevCredQuota := prevQuota.Exceeded && strings.TrimSpace(prevQuota.Reason) == "credential_quota" && prevQuota.NextRecoverAt.After(now) + recCredQuota := quota.Exceeded && strings.TrimSpace(quota.Reason) == "credential_quota" && quota.NextRecoverAt.After(now) + if prevCredQuota || recCredQuota { + mergedQuota.Exceeded = true + mergedQuota.Reason = "credential_quota" + var credRecover time.Time + if prevCredQuota { + credRecover = prevQuota.NextRecoverAt + } + if recCredQuota && quota.NextRecoverAt.After(credRecover) { + credRecover = quota.NextRecoverAt + } + mergedQuota.NextRecoverAt = credRecover + if prevQuota.BackoffLevel > mergedQuota.BackoffLevel { + mergedQuota.BackoffLevel = prevQuota.BackoffLevel + } + if prevCredQuota && (!recCredQuota || prevQuota.NextRecoverAt.After(quota.NextRecoverAt) || len(mergedQuota.cooldownSignals) == 0) { + mergedQuota.cooldownSignals = prevQuota.cooldownSignals + } + } else if prevQuota.Exceeded && prevQuota.NextRecoverAt.After(now) { + mergedQuota.Exceeded = true + if prevQuota.NextRecoverAt.After(mergedQuota.NextRecoverAt) { + mergedQuota.NextRecoverAt = prevQuota.NextRecoverAt + } + if prevQuota.BackoffLevel > mergedQuota.BackoffLevel { + mergedQuota.BackoffLevel = prevQuota.BackoffLevel + } + if liveNewer && prevQuota.Reason != "" { + mergedQuota.Reason = prevQuota.Reason + } else if mergedQuota.Reason == "" { + mergedQuota.Reason = prevQuota.Reason + } + if prevQuota.NextRecoverAt.After(quota.NextRecoverAt) || len(mergedQuota.cooldownSignals) == 0 { + mergedQuota.cooldownSignals = prevQuota.cooldownSignals + } + } + applyCooldownFields(&auth.Quota, mergedQuota) + auth.Quota = mergeQuotaObservation(auth.Quota, prevQuota) auth.Quota = mergeQuotaObservation(auth.Quota, quota) + if credForce := recordCredQuotaForceUntil(record); credForce.After(now) && credForce.After(auth.credQuotaForceUntil) { + auth.credQuotaForceUntil = credForce + } + authForce := recordAuthLevelForceUntil(record) + if isLegacyAggregatedAuthRecord(record, auth.ModelStates, now) { + authForce = time.Time{} + } + if authForce.After(now) && authForce.After(auth.forceCooldownUntil) { + auth.forceCooldownUntil = authForce + } + auth.forceCooldownSeeded = true auth.Generation++ - auth.UpdatedAt = updatedAt - if reason != "" { - auth.StatusMessage = reason + if liveNewer { + auth.UpdatedAt = prevUpdatedAt + if prevStatusMsg != "" { + auth.StatusMessage = prevStatusMsg + } else if reason != "" { + auth.StatusMessage = reason + } + if prevLastErr != nil { + auth.LastError = prevLastErr + } else { + auth.LastError = cloneError(record.LastError) + } + } else { + auth.UpdatedAt = updatedAt + if reason != "" { + auth.StatusMessage = reason + } + auth.LastError = cloneError(record.LastError) } - auth.LastError = cloneError(record.LastError) return true } @@ -406,6 +503,7 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti // state (reason "quota" whose own error is the explicit daily/weekly // exhaustion message) is credential-derived and must not be captured; // genuine model-scoped 429s carry different messages and stay eligible. + seedModelForceCooldownDeadline(state) liveShadow := state.credentialQuotaShadow.cloneShallow() if liveShadow == nil && !state.credentialQuotaPropagated && state.Quota.Reason != "credential_quota" && isModelStateActiveCooldown(state, now) && !isDevinQuotaExhaustionMessage(devinStateErrorText(state)) { @@ -413,21 +511,29 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti } wasLiveWrapped := state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota" || (auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) + isWrappedRecord := record.IndependentState != nil || record.CredentialQuotaPropagated || strings.TrimSpace(quota.Reason) == "credential_quota" + var recordModelForceUntil time.Time + if !isWrappedRecord && outerForceUntil.After(now) { + recordModelForceUntil = outerForceUntil + } else if isWrappedRecord && outerForceUntil.After(now) && outerForceUntil.After(auth.credQuotaForceUntil) { + auth.credQuotaForceUntil = outerForceUntil + } mergeModelState(state, &ModelState{ - Unavailable: true, - Status: StatusError, - StatusMessage: reason, - NextRetryAfter: record.NextRetryAfter, - Quota: quota, - LastError: cloneError(record.LastError), - UpdatedAt: updatedAt, + Unavailable: true, + Status: StatusError, + StatusMessage: reason, + NextRetryAfter: record.NextRetryAfter, + Quota: quota, + LastError: cloneError(record.LastError), + UpdatedAt: updatedAt, + forceCooldownUntil: recordModelForceUntil, + forceCooldownSeeded: true, }) // The persisted independent layer joins the merge newest-wins; a newer // live independent layer (shadow or freshly captured unwrapped state, // e.g. an own failure recorded during a delayed store load) is never // clobbered by the older record. A wrapped record is any record whose // outer reason is the propagated credential quota. - isWrappedRecord := record.IndependentState != nil || record.CredentialQuotaPropagated || strings.TrimSpace(quota.Reason) == "credential_quota" if !isWrappedRecord && wasLiveWrapped { // An unwrapped independent record landing beneath a live wrap becomes // the model's independent layer (unless it is itself a legacy @@ -435,48 +541,444 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti // presenting the propagated credential quota. if !isDevinQuotaExhaustionMessage(devinRecordErrorText(record)) { recordShadow := &ModelState{ - Unavailable: true, - Status: StatusError, - StatusMessage: reason, - NextRetryAfter: record.NextRetryAfter, - Quota: quota, - LastError: cloneError(record.LastError), - UpdatedAt: updatedAt, + Unavailable: true, + Status: StatusError, + StatusMessage: reason, + NextRetryAfter: record.NextRetryAfter, + Quota: quota, + LastError: cloneError(record.LastError), + UpdatedAt: updatedAt, + forceCooldownUntil: recordModelForceUntil, + forceCooldownSeeded: true, } - preferred := state.credentialQuotaShadow - if preferred == nil || (liveShadow != nil && !liveShadow.UpdatedAt.After(updatedAt)) { - preferred = recordShadow - } else if liveShadow != nil && liveShadow.UpdatedAt.After(updatedAt) { - preferred = liveShadow + preferred := state.credentialQuotaShadow.cloneShallow() + if liveShadow != nil { + preferred = mergeModelState(preferred, liveShadow.cloneShallow()) } + preferred = mergeModelState(preferred, recordShadow) state.credentialQuotaShadow = preferred } state.credentialQuotaPropagated = true applyCooldownFields(&state.Quota, QuotaState{ - Exceeded: true, - Reason: "credential_quota", - NextRecoverAt: auth.Quota.NextRecoverAt, - BackoffLevel: auth.Quota.BackoffLevel, + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: auth.Quota.NextRecoverAt, + BackoffLevel: auth.Quota.BackoffLevel, + cooldownSignals: auth.Quota.cooldownSignals, }) if state.NextRetryAfter.Before(auth.Quota.NextRecoverAt) { state.NextRetryAfter = auth.Quota.NextRecoverAt } + if state.credentialQuotaShadow != nil && state.credentialQuotaShadow.NextRetryAfter.After(state.NextRetryAfter) { + state.NextRetryAfter = state.credentialQuotaShadow.NextRetryAfter + } state.Unavailable = true state.Status = StatusError } else if isWrappedRecord { state.credentialQuotaPropagated = true - if record.IndependentState != nil && (liveShadow == nil || !liveShadow.UpdatedAt.After(record.IndependentState.UpdatedAt)) { - state.credentialQuotaShadow = record.IndependentState.cloneShallow() - } else if liveShadow != nil { - state.credentialQuotaShadow = liveShadow + var restoredShadow *ModelState + if record.IndependentState != nil { + restoredShadow = record.IndependentState.cloneShallow() + restoredShadow.Quota.cooldownEvaluated = restoredShadow.Quota.Exceeded + if indepForceUntil.After(restoredShadow.forceCooldownUntil) { + restoredShadow.forceCooldownUntil = indepForceUntil + } + restoredShadow.forceCooldownSeeded = true + if executorKeyFromAuth(auth) == "claude" || strings.EqualFold(strings.TrimSpace(record.Provider), "claude") { + restoredShadow.Quota.cooldownSignals = sanitizePersistedClaudeCooldownSignals(record.IndependentCooldownSignals) + } + } + state.credentialQuotaShadow = mergeModelState(liveShadow, restoredShadow) + if state.credentialQuotaShadow != nil && state.credentialQuotaShadow.NextRetryAfter.After(state.NextRetryAfter) { + state.NextRetryAfter = state.credentialQuotaShadow.NextRetryAfter } } + auth.forceCooldownSeeded = true auth.Generation++ - auth.UpdatedAt = updatedAt + if updatedAt.After(auth.UpdatedAt) { + auth.UpdatedAt = updatedAt + } updateAggregatedAvailability(auth, now) return true } +func isForceCooldownError(err *Error) bool { + return err != nil && err.Code == ErrorCodeForceCooldown +} + +func recordOuterForceUntil(record CooldownStateRecord) time.Time { + if record.ForceCooldownEvaluated { + return record.ForceCooldownUntil + } + if isForceCooldownError(record.LastError) { + return record.NextRetryAfter + } + return time.Time{} +} + +func recordCredQuotaForceUntil(record CooldownStateRecord) time.Time { + if record.ForceCooldownEvaluated { + return record.CredQuotaForceCooldownUntil + } + if strings.TrimSpace(record.Quota.Reason) == "credential_quota" && isForceCooldownError(record.LastError) { + if !record.Quota.NextRecoverAt.IsZero() { + return record.Quota.NextRecoverAt + } + return record.NextRetryAfter + } + return time.Time{} +} + +func recordAuthLevelForceUntil(record CooldownStateRecord) time.Time { + if record.ForceCooldownEvaluated { + return record.ForceCooldownUntil + } + if strings.TrimSpace(record.Quota.Reason) != "credential_quota" && isForceCooldownError(record.LastError) { + return record.NextRetryAfter + } + return time.Time{} +} + +// isLegacyAggregatedAuthRecord reports whether a pre-ForceCooldownEvaluated +// Model=="" record is merely the updateAggregatedAvailability copy of +// auth.ModelStates rather than an independent model-less cooldown. +func isLegacyAggregatedAuthRecord(record CooldownStateRecord, states map[string]*ModelState, now time.Time) bool { + if record.ForceCooldownEvaluated || len(states) == 0 || strings.TrimSpace(record.Quota.Reason) == "credential_quota" { + return false + } + var earliestRetry time.Time + matchedOrigin := false + hasState := false + for _, st := range states { + if st == nil { + continue + } + hasState = true + if st.Status == StatusDisabled { + continue + } + if !st.Unavailable || st.NextRetryAfter.IsZero() || !st.NextRetryAfter.After(now) { + return false + } + if earliestRetry.IsZero() || st.NextRetryAfter.Before(earliestRetry) { + earliestRetry = st.NextRetryAfter + } + if st.UpdatedAt.Equal(record.UpdatedAt) && cooldownErrorEqual(st.LastError, record.LastError) { + matchedOrigin = true + } + } + return hasState && matchedOrigin && !earliestRetry.IsZero() && record.NextRetryAfter.Equal(earliestRetry) +} + +func recordIndependentForceUntil(record CooldownStateRecord) time.Time { + if record.IndependentState == nil { + return time.Time{} + } + if record.ForceCooldownEvaluated { + return record.IndependentForceCooldownUntil + } + if isForceCooldownError(record.IndependentState.LastError) { + return record.IndependentState.NextRetryAfter + } + return time.Time{} +} + +func isForceCooldownRecordAt(record CooldownStateRecord, now time.Time) bool { + if strings.TrimSpace(record.Model) == "" { + return recordAuthLevelForceUntil(record).After(now) || recordCredQuotaForceUntil(record).After(now) + } + return recordOuterForceUntil(record).After(now) || recordIndependentForceUntil(record).After(now) +} + +// Pure read-only force deadline accessors safe under RLock. +func effectiveCredQuotaForceUntil(auth *Auth) time.Time { + if auth == nil { + return time.Time{} + } + if auth.forceCooldownSeeded { + return auth.credQuotaForceUntil + } + if isForceCooldownError(auth.LastError) && auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" { + return auth.Quota.NextRecoverAt + } + return time.Time{} +} + +func effectiveAuthForceUntil(auth *Auth) time.Time { + if auth == nil { + return time.Time{} + } + if auth.forceCooldownSeeded { + return auth.forceCooldownUntil + } + if len(auth.ModelStates) == 0 && isForceCooldownError(auth.LastError) { + return auth.NextRetryAfter + } + return time.Time{} +} + +func effectiveModelForceUntil(state *ModelState) time.Time { + if state == nil { + return time.Time{} + } + if state.forceCooldownSeeded { + return state.forceCooldownUntil + } + isWrapped := state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota" + if !isWrapped && isForceCooldownError(state.LastError) { + return state.NextRetryAfter + } + return time.Time{} +} + +func seedModelForceCooldownDeadline(state *ModelState) { + if state == nil { + return + } + if state.credentialQuotaShadow != nil { + seedModelForceCooldownDeadline(state.credentialQuotaShadow) + } + if state.forceCooldownSeeded { + return + } + state.forceCooldownSeeded = true + isWrapped := state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota" + if !isWrapped && isForceCooldownError(state.LastError) && !state.NextRetryAfter.IsZero() { + state.forceCooldownUntil = state.NextRetryAfter + } +} + +// seedInitialForceCooldownDeadlines initializes explicit force-cooldown provenance +// once from pre-populated Auth/ModelState struct literals before runtime updates. +func seedInitialForceCooldownDeadlines(auth *Auth) { + if auth == nil { + return + } + for _, state := range auth.ModelStates { + seedModelForceCooldownDeadline(state) + } + if auth.forceCooldownSeeded { + return + } + auth.forceCooldownSeeded = true + if !isForceCooldownError(auth.LastError) { + return + } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && !auth.Quota.NextRecoverAt.IsZero() { + auth.credQuotaForceUntil = auth.Quota.NextRecoverAt + } else if len(auth.ModelStates) == 0 && !auth.NextRetryAfter.IsZero() { + auth.forceCooldownUntil = auth.NextRetryAfter + } +} + +func clampStateToForceDeadline(state *ModelState, now time.Time) (kept bool, mutated bool) { + if state == nil { + return false, false + } + seedModelForceCooldownDeadline(state) + if !state.forceCooldownUntil.After(now) || !state.NextRetryAfter.After(now) { + return false, false + } + if state.NextRetryAfter.After(state.forceCooldownUntil) { + state.NextRetryAfter = state.forceCooldownUntil + mutated = true + } + if state.Quota.Exceeded && state.Quota.NextRecoverAt.After(state.forceCooldownUntil) { + state.Quota.NextRecoverAt = state.forceCooldownUntil + mutated = true + } + return true, mutated +} + +// clearDisabledCooldownStateForAuth clears ordinary cooldown/quota state for an auth +// under disable-cooling while preserving active ErrorCodeForceCooldown deadlines. +func clearDisabledCooldownStateForAuth(auth *Auth, now time.Time) bool { + if auth == nil || hasUnauthorizedAuthFailure(auth) { + return false + } + seedInitialForceCooldownDeadlines(auth) + + keepCredQuotaForce := auth.Quota.Exceeded && + auth.Quota.Reason == "credential_quota" && + auth.credQuotaForceUntil.After(now) && + auth.Quota.NextRecoverAt.After(now) + keepAuthLevelForce := auth.Unavailable && + auth.forceCooldownUntil.After(now) && + auth.NextRetryAfter.After(now) + combinedAuthForceUntil := time.Time{} + if keepCredQuotaForce { + combinedAuthForceUntil = auth.credQuotaForceUntil + } + if keepAuthLevelForce && auth.forceCooldownUntil.After(combinedAuthForceUntil) { + combinedAuthForceUntil = auth.forceCooldownUntil + } + + if len(auth.ModelStates) == 0 { + if keepCredQuotaForce || keepAuthLevelForce { + changed := false + if !keepCredQuotaForce && !auth.credQuotaForceUntil.IsZero() { + auth.credQuotaForceUntil = time.Time{} + changed = true + } + if !keepAuthLevelForce && !auth.forceCooldownUntil.IsZero() { + auth.forceCooldownUntil = time.Time{} + changed = true + } + if keepCredQuotaForce { + if !auth.Quota.NextRecoverAt.Equal(auth.credQuotaForceUntil) { + auth.Quota.NextRecoverAt = auth.credQuotaForceUntil + changed = true + } + } else if auth.Quota.Reason == "credential_quota" { + applyCooldownFields(&auth.Quota, QuotaState{}) + changed = true + } else if auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(combinedAuthForceUntil) { + auth.Quota.NextRecoverAt = combinedAuthForceUntil + changed = true + } + if !auth.NextRetryAfter.Equal(combinedAuthForceUntil) { + auth.NextRetryAfter = combinedAuthForceUntil + changed = true + } + if changed { + auth.Generation++ + auth.UpdatedAt = now + } + return changed + } + return clearCooldownStateForAuth(auth, now) + } + + savedAuthNextRetry := combinedAuthForceUntil + savedAuthQuota := cooldownFieldsOf(auth.Quota) + if !keepCredQuotaForce { + if savedAuthQuota.Reason == "credential_quota" { + savedAuthQuota = QuotaState{} + } else if keepAuthLevelForce && savedAuthQuota.Exceeded && savedAuthQuota.NextRecoverAt.After(auth.forceCooldownUntil) { + savedAuthQuota.NextRecoverAt = auth.forceCooldownUntil + } + } + + changed := false + if keepCredQuotaForce { + if !auth.Quota.NextRecoverAt.Equal(auth.credQuotaForceUntil) { + auth.Quota.NextRecoverAt = auth.credQuotaForceUntil + changed = true + } + if !auth.NextRetryAfter.Equal(combinedAuthForceUntil) { + auth.NextRetryAfter = combinedAuthForceUntil + changed = true + } + } + + for _, state := range auth.ModelStates { + if state == nil { + continue + } + isWrapped := state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota" + if isWrapped { + if keepCredQuotaForce { + var shadowDeadline time.Time + if state.credentialQuotaShadow != nil { + keptShadow, mutatedShadow := clampStateToForceDeadline(state.credentialQuotaShadow, now) + if !keptShadow { + state.credentialQuotaShadow = nil + state.UpdatedAt = now + changed = true + } else { + shadowDeadline = state.credentialQuotaShadow.NextRetryAfter + if mutatedShadow { + state.UpdatedAt = now + changed = true + } + } + } + effectiveOuter := auth.credQuotaForceUntil + if shadowDeadline.After(effectiveOuter) { + effectiveOuter = shadowDeadline + } + if state.Quota.NextRecoverAt.After(auth.credQuotaForceUntil) { + state.Quota.NextRecoverAt = auth.credQuotaForceUntil + state.UpdatedAt = now + changed = true + } + if state.NextRetryAfter.After(effectiveOuter) { + state.NextRetryAfter = effectiveOuter + state.UpdatedAt = now + changed = true + } + continue + } + if shadow := state.credentialQuotaShadow; shadow != nil { + if keptShadow, _ := clampStateToForceDeadline(shadow, now); keptShadow { + preservedSeq := shadow.cooldownRecordedSeq + *state = *shadow.cloneShallow() + state.cooldownRecordedSeq = preservedSeq + state.credentialQuotaPropagated = false + state.credentialQuotaShadow = nil + state.UpdatedAt = now + changed = true + continue + } + } + } else if keptModel, mutatedModel := clampStateToForceDeadline(state, now); keptModel { + if mutatedModel { + state.UpdatedAt = now + changed = true + } + continue + } + if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() || state.credentialQuotaShadow != nil || state.credentialQuotaPropagated || !state.forceCooldownUntil.IsZero() { + state.Unavailable = false + state.NextRetryAfter = time.Time{} + applyCooldownFields(&state.Quota, QuotaState{}) + state.cooldownRecordedSeq = 0 + state.forceCooldownUntil = time.Time{} + state.credentialQuotaShadow = nil + state.credentialQuotaPropagated = false + state.UpdatedAt = now + changed = true + } + } + + if !keepCredQuotaForce { + auth.credQuotaForceUntil = time.Time{} + } + if !keepAuthLevelForce { + auth.forceCooldownUntil = time.Time{} + } + + prevUnavailable := auth.Unavailable + prevNextRetry := auth.NextRetryAfter + prevQuota := cooldownFieldsOf(auth.Quota) + if !keepCredQuotaForce { + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + applyCooldownFields(&auth.Quota, QuotaState{}) + if !keepAuthLevelForce { + auth.cooldownRecordedSeq = 0 + } + updateAggregatedAvailability(auth, now) + } + if keepAuthLevelForce { + auth.Unavailable = true + if savedAuthNextRetry.After(auth.NextRetryAfter) { + auth.NextRetryAfter = savedAuthNextRetry + } + if !keepCredQuotaForce && savedAuthQuota.Exceeded && savedAuthQuota.NextRecoverAt.After(now) && !auth.Quota.NextRecoverAt.After(savedAuthQuota.NextRecoverAt) { + applyCooldownFields(&auth.Quota, savedAuthQuota) + } + } + if auth.Unavailable != prevUnavailable || !auth.NextRetryAfter.Equal(prevNextRetry) || !cooldownQuotaEqual(cooldownFieldsOf(auth.Quota), prevQuota) { + changed = true + } + if changed { + auth.Generation++ + auth.UpdatedAt = now + } + return changed +} + func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { if auth == nil { return false @@ -485,6 +987,9 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { return false } changed := false + auth.forceCooldownUntil = time.Time{} + auth.credQuotaForceUntil = time.Time{} + auth.cooldownRecordedSeq = 0 if auth.Unavailable || !auth.NextRetryAfter.IsZero() || auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.IsZero() { auth.Unavailable = false auth.NextRetryAfter = time.Time{} @@ -496,6 +1001,8 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { if state == nil { continue } + state.forceCooldownUntil = time.Time{} + state.cooldownRecordedSeq = 0 if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() || state.credentialQuotaShadow != nil || state.credentialQuotaPropagated { state.Unavailable = false state.NextRetryAfter = time.Time{} @@ -708,15 +1215,16 @@ func (m *Manager) cooldownStateRecordsSnapshot() []CooldownStateRecord { } func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) []CooldownStateRecord { - if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) { + if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled { return nil } + coolingDisabled := m.cooldownDisabledForAuth(auth) records := make([]CooldownStateRecord, 0, 1+len(auth.ModelStates)) - if record, ok := authCooldownStateRecord(auth, now); ok { + if record, ok := authCooldownStateRecord(auth, now); ok && (!coolingDisabled || isForceCooldownRecordAt(record, now)) { records = append(records, record) } for model, state := range auth.ModelStates { - if record, ok := modelCooldownStateRecord(auth, model, state, now); ok { + if record, ok := modelCooldownStateRecord(auth, model, state, now); ok && (!coolingDisabled || isForceCooldownRecordAt(record, now)) { records = append(records, record) } } @@ -756,6 +1264,15 @@ func cooldownStateRecordEqual(a, b CooldownStateRecord) bool { if !cooldownErrorEqual(a.LastError, b.LastError) { return false } + if a.ForceCooldownEvaluated != b.ForceCooldownEvaluated || + !a.ForceCooldownUntil.Equal(b.ForceCooldownUntil) || + !a.CredQuotaForceCooldownUntil.Equal(b.CredQuotaForceCooldownUntil) || + !a.IndependentForceCooldownUntil.Equal(b.IndependentForceCooldownUntil) { + return false + } + if !maps.Equal(a.CooldownSignals, b.CooldownSignals) || !maps.Equal(a.IndependentCooldownSignals, b.IndependentCooldownSignals) { + return false + } return cooldownIndependentStateEqual(a.IndependentState, b.IndependentState) } @@ -793,16 +1310,34 @@ func authCooldownStateRecord(auth *Auth, now time.Time) (CooldownStateRecord, bo if auth == nil || !auth.Unavailable || auth.NextRetryAfter.IsZero() || !auth.NextRetryAfter.After(now) { return CooldownStateRecord{}, false } + var cooldownSignals map[string]string + if auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) && !auth.NextRetryAfter.After(auth.Quota.NextRecoverAt) { + cooldownSignals = sanitizePersistedClaudeCooldownSignals(auth.Quota.cooldownSignals) + } + var authForceUntil time.Time + if u := effectiveAuthForceUntil(auth); u.After(now) { + authForceUntil = u + } + var credQuotaForceUntil time.Time + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + if u := effectiveCredQuotaForceUntil(auth); u.After(now) { + credQuotaForceUntil = u + } + } return CooldownStateRecord{ - Provider: strings.TrimSpace(auth.Provider), - AuthID: auth.ID, - AuthFile: cooldownAuthFile(auth), - Status: "cooling", - NextRetryAfter: auth.NextRetryAfter, - Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError), - Quota: cooldownFieldsOf(auth.Quota), - LastError: cloneError(auth.LastError), - UpdatedAt: auth.UpdatedAt, + Provider: strings.TrimSpace(auth.Provider), + AuthID: auth.ID, + AuthFile: cooldownAuthFile(auth), + Status: "cooling", + NextRetryAfter: auth.NextRetryAfter, + Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError), + Quota: cooldownFieldsOf(auth.Quota), + LastError: cloneError(auth.LastError), + UpdatedAt: auth.UpdatedAt, + CooldownSignals: cooldownSignals, + ForceCooldownUntil: authForceUntil, + CredQuotaForceCooldownUntil: credQuotaForceUntil, + ForceCooldownEvaluated: true, }, true } @@ -811,20 +1346,45 @@ func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now t if auth == nil || state == nil || model == "" || !state.Unavailable || state.NextRetryAfter.IsZero() || !state.NextRetryAfter.After(now) { return CooldownStateRecord{}, false } + var cooldownSignals map[string]string + liveCredentialWrap := auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) + expiredCredentialWrap := (state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota") && !liveCredentialWrap + if !expiredCredentialWrap && state.Quota.Exceeded && state.Quota.NextRecoverAt.After(now) && + (state.Quota.Reason == "credential_quota" || !state.NextRetryAfter.After(state.Quota.NextRecoverAt)) { + cooldownSignals = sanitizePersistedClaudeCooldownSignals(state.Quota.cooldownSignals) + } + var outerForceUntil time.Time + if liveCredentialWrap { + if u := effectiveCredQuotaForceUntil(auth); u.After(now) { + outerForceUntil = u + } + } else if u := effectiveModelForceUntil(state); u.After(now) { + outerForceUntil = u + } record := CooldownStateRecord{ - Provider: strings.TrimSpace(auth.Provider), - AuthID: auth.ID, - AuthFile: cooldownAuthFile(auth), - Model: model, - Status: "cooling", - NextRetryAfter: state.NextRetryAfter, - Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError), - Quota: cooldownFieldsOf(state.Quota), - LastError: cloneError(state.LastError), - UpdatedAt: state.UpdatedAt, + Provider: strings.TrimSpace(auth.Provider), + AuthID: auth.ID, + AuthFile: cooldownAuthFile(auth), + Model: model, + Status: "cooling", + NextRetryAfter: state.NextRetryAfter, + Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError), + Quota: cooldownFieldsOf(state.Quota), + LastError: cloneError(state.LastError), + UpdatedAt: state.UpdatedAt, + CooldownSignals: cooldownSignals, + ForceCooldownUntil: outerForceUntil, + ForceCooldownEvaluated: true, } if state.credentialQuotaShadow != nil { record.IndependentState = state.credentialQuotaShadow.cloneShallow() + if u := effectiveModelForceUntil(state.credentialQuotaShadow); u.After(now) { + record.IndependentForceCooldownUntil = u + } + if state.credentialQuotaShadow.Quota.Exceeded && state.credentialQuotaShadow.Quota.NextRecoverAt.After(now) && + !state.credentialQuotaShadow.NextRetryAfter.After(state.credentialQuotaShadow.Quota.NextRecoverAt) { + record.IndependentCooldownSignals = sanitizePersistedClaudeCooldownSignals(state.credentialQuotaShadow.Quota.cooldownSignals) + } } record.CredentialQuotaPropagated = state.credentialQuotaPropagated return record, true @@ -994,7 +1554,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { modelState = state resetModelState(state, now) updateAggregatedAvailability(auth, now) - if !hasModelError(auth, now) { + if !hasModelError(auth, now) && !effectiveAuthForceUntil(auth).After(now) { auth.LastError = nil auth.StatusMessage = "" auth.Status = StatusActive @@ -1046,10 +1606,18 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { *state = *fresh } } + seedInitialForceCooldownDeadlines(auth) + seedModelForceCooldownDeadline(state) state.Unavailable = true state.Status = StatusError state.UpdatedAt = now prevModelRetryAfter := state.NextRetryAfter + prevModelForceUntil := state.forceCooldownUntil + prevModelQuota := cooldownFieldsOf(state.Quota) + prevModelQuotaSignals := maps.Clone(state.Quota.cooldownSignals) + var ownModelForceDeadline time.Time + isForceResult := isForceCooldownError(result.Error) + appliedCredQuotaWrap := false if result.Error != nil { state.LastError = cloneError(result.Error) state.StatusMessage = result.Error.Message @@ -1096,6 +1664,9 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { // a still-live genuine quota deadline is left untouched and // keeps blocking through the generic retry branches below. next := transientRateLimitRetryAfter(now, result.RetryAfter, disableCooling) + if isForceResult { + ownModelForceDeadline = next + } if !next.IsZero() && state.Quota.Exceeded && state.Quota.NextRecoverAt.After(next) { next = state.Quota.NextRecoverAt } @@ -1123,8 +1694,12 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { state.NextRetryAfter = next } case 429: + prevStateQuotaActive := state.Quota.Exceeded && state.Quota.NextRecoverAt.After(now) + prevStateRecoverAt := state.Quota.NextRecoverAt + prevStateCooldownSignals := maps.Clone(state.Quota.cooldownSignals) var next time.Time var credentialNext time.Time + keptPriorStateQuota := false backoffLevel := state.Quota.BackoffLevel if result.CredentialScope { backoffLevel = 0 @@ -1133,6 +1708,8 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } if !disableCooling { + reusedOpenWindow := false + var baseBackoffBefore int if result.RetryAfter != nil { cooldown := *result.RetryAfter if cooldown < minQuotaCooldownFloor { @@ -1149,19 +1726,52 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { quotaForFailure.BackoffLevel = 0 } } + reusedOpenWindow = quotaForFailure.NextRecoverAt.After(now) + baseBackoffBefore = quotaForFailure.BackoffLevel next, backoffLevel = quotaCooldownAfterFailure(quotaForFailure, now) } credentialNext = next - if state.Quota.Exceeded && state.Quota.NextRecoverAt.After(next) { - next = state.Quota.NextRecoverAt + if isForceResult { + if result.RetryAfter != nil || !reusedOpenWindow { + ownModelForceDeadline = credentialNext + } else if result.CredentialScope { + if auth.credQuotaForceUntil.After(now) { + ownModelForceDeadline = auth.credQuotaForceUntil + } else { + ownModelForceDeadline, _ = quotaCooldownAfterFailure(QuotaState{BackoffLevel: baseBackoffBefore}, now) + } + } else if prevModelForceUntil.After(now) { + ownModelForceDeadline = prevModelForceUntil + } else { + ownModelForceDeadline, _ = quotaCooldownAfterFailure(QuotaState{BackoffLevel: baseBackoffBefore}, now) + } + } + if prevStateQuotaActive && (prevStateRecoverAt.After(next) || (!result.CredentialScope && result.RetryAfter == nil)) { + next = prevStateRecoverAt + keptPriorStateQuota = true + } else if result.CredentialScope && result.RetryAfter == nil && auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + // quotaCooldownAfterFailure reused the open credential_quota window. + keptPriorStateQuota = prevStateQuotaActive && prevStateRecoverAt.Equal(next) + } + } + var cooldownSignals map[string]string + if !disableCooling { + if keptPriorStateQuota { + cooldownSignals = prevStateCooldownSignals + } else { + cooldownSignals = rejectedClaudeSignalsForCooldown(result.Provider, responseHeaders) + if len(cooldownSignals) == 0 && prevStateQuotaActive && prevStateRecoverAt.Equal(next) { + cooldownSignals = prevStateCooldownSignals + } } } state.NextRetryAfter = next applyCooldownFields(&state.Quota, QuotaState{ - Exceeded: true, - Reason: "quota", - NextRecoverAt: next, - BackoffLevel: backoffLevel, + Exceeded: true, + Reason: "quota", + NextRecoverAt: next, + BackoffLevel: backoffLevel, + cooldownSignals: cooldownSignals, }) // Non-credential own failures re-derive their independent layer // below, so a stale shadow is dropped. A credential-scoped @@ -1169,7 +1779,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { // refreshes the propagated outer layer. if !result.CredentialScope { state.credentialQuotaShadow = nil - } else if !state.credentialQuotaPropagated && savedOuterState.Quota.Reason != "credential_quota" { + } else if !disableCooling && !state.credentialQuotaPropagated && savedOuterState.Quota.Reason != "credential_quota" { // First wrap of the reporting model: preserve any live prior // own block (including an ordinary model-scoped 429, which // quota restoration does not supersede). An outer @@ -1187,12 +1797,25 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } if result.CredentialScope && !disableCooling { + appliedCredQuotaWrap = true + prevAuthCredActive := auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) + prevAuthCredRecoverAt := auth.Quota.NextRecoverAt + prevAuthCredSignals := maps.Clone(auth.Quota.cooldownSignals) + var credCooldownSignals map[string]string + if prevAuthCredActive && (prevAuthCredRecoverAt.After(credentialNext) || result.RetryAfter == nil) { + credCooldownSignals = prevAuthCredSignals + } else { + credCooldownSignals = rejectedClaudeSignalsForCooldown(result.Provider, responseHeaders) + if len(credCooldownSignals) == 0 && prevAuthCredActive && prevAuthCredRecoverAt.Equal(credentialNext) { + credCooldownSignals = prevAuthCredSignals + } + } for _, otherState := range auth.ModelStates { if otherState != nil && otherState != state { // Normalize a naturally expired wrap first (issue #6690) so a // stale marker cannot suppress snapshotting a fresh block. if (otherState.credentialQuotaPropagated || otherState.Quota.Reason == "credential_quota") && - !(auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) { + !prevAuthCredActive { if shadow := otherState.credentialQuotaShadow; shadow != nil && isModelStateActiveCooldown(shadow, now) { preservedSeq := shadow.cooldownRecordedSeq *otherState = *shadow.cloneShallow() @@ -1219,8 +1842,13 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { otherState.Unavailable = true otherState.Status = StatusError otherQuotaNext := credentialNext + otherSignals := credCooldownSignals + if prevAuthCredActive && prevAuthCredRecoverAt.After(otherQuotaNext) { + otherQuotaNext = prevAuthCredRecoverAt + } if otherState.Quota.Exceeded && otherState.Quota.NextRecoverAt.After(otherQuotaNext) { otherQuotaNext = otherState.Quota.NextRecoverAt + otherSignals = maps.Clone(otherState.Quota.cooldownSignals) } otherRetryAfter := otherQuotaNext // Propagation only extends a sibling's still-live @@ -1230,25 +1858,38 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } otherState.NextRetryAfter = otherRetryAfter applyCooldownFields(&otherState.Quota, QuotaState{ - Exceeded: true, - Reason: "credential_quota", - NextRecoverAt: otherQuotaNext, - BackoffLevel: backoffLevel, + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: otherQuotaNext, + BackoffLevel: backoffLevel, + cooldownSignals: otherSignals, }) } } if !wasTerminalUnauthorized { auth.Unavailable = true authNext := credentialNext - if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && - auth.Quota.NextRecoverAt.After(authNext) { - authNext = auth.Quota.NextRecoverAt + if prevAuthCredActive && prevAuthCredRecoverAt.After(authNext) { + authNext = prevAuthCredRecoverAt } auth.Quota.Exceeded = true auth.Quota.Reason = "credential_quota" auth.Quota.NextRecoverAt = authNext auth.Quota.BackoffLevel = backoffLevel + auth.Quota.cooldownSignals = maps.Clone(credCooldownSignals) + auth.Quota.cooldownEvaluated = true auth.NextRetryAfter = authNext + if authNext.After(state.Quota.NextRecoverAt) || (prevAuthCredActive && result.RetryAfter == nil) { + state.Quota.NextRecoverAt = authNext + state.Quota.cooldownSignals = maps.Clone(credCooldownSignals) + if state.NextRetryAfter.Before(authNext) { + state.NextRetryAfter = authNext + } + } + if isForceResult && ownModelForceDeadline.After(auth.credQuotaForceUntil) { + auth.credQuotaForceUntil = ownModelForceDeadline + } + auth.forceCooldownSeeded = true // The reporting model is now part of the wrap too: mark it // so a later credential report from another model cannot // snapshot this credential-derived state as independent. @@ -1268,10 +1909,54 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { state.Unavailable = false state.Quota.Exceeded = false } - if result.Error != nil && result.Error.Code == ErrorCodeForceCooldown && state.NextRetryAfter.IsZero() { + if isForceResult && state.NextRetryAfter.IsZero() { state.NextRetryAfter = now.Add(transientErrorCooldown) state.Unavailable = true } + if !appliedCredQuotaWrap { + if result.CredentialScope && auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.credQuotaForceUntil.After(now) && auth.Quota.NextRecoverAt.After(now) { + lastErr := cloneError(state.LastError) + statusMsg := state.StatusMessage + *state = *savedOuterState + state.LastError = lastErr + state.StatusMessage = statusMsg + state.Unavailable = true + state.Status = StatusError + if state.NextRetryAfter.Before(auth.credQuotaForceUntil) { + state.NextRetryAfter = auth.credQuotaForceUntil + } + state.UpdatedAt = now + } else { + if isForceResult { + if ownModelForceDeadline.IsZero() { + ownModelForceDeadline = state.NextRetryAfter + } + if ownModelForceDeadline.After(prevModelForceUntil) { + state.forceCooldownUntil = ownModelForceDeadline + } else { + state.forceCooldownUntil = prevModelForceUntil + } + } else { + state.forceCooldownUntil = prevModelForceUntil + } + state.forceCooldownSeeded = true + if state.forceCooldownUntil.After(now) { + if state.NextRetryAfter.Before(state.forceCooldownUntil) { + state.NextRetryAfter = state.forceCooldownUntil + } + state.Unavailable = true + state.Status = StatusError + if (!state.Quota.Exceeded || state.Quota.NextRecoverAt.IsZero()) && prevModelQuota.Exceeded && prevModelQuota.NextRecoverAt.After(now) { + restored := prevModelQuota + if disableCooling && restored.NextRecoverAt.After(state.forceCooldownUntil) { + restored.NextRecoverAt = state.forceCooldownUntil + } + restored.cooldownSignals = prevModelQuotaSignals + applyCooldownFields(&state.Quota, restored) + } + } + } + } // A later failure only extends a still-live cooldown; it never // shortens one. A deliberate zero write (disableCooling) still // clears the deadline. @@ -1310,10 +1995,11 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } if !(outer.Quota.Exceeded && outer.Quota.Reason == "credential_quota") { applyCooldownFields(&outer.Quota, QuotaState{ - Exceeded: true, - Reason: "credential_quota", - NextRecoverAt: auth.Quota.NextRecoverAt, - BackoffLevel: auth.Quota.BackoffLevel, + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: auth.Quota.NextRecoverAt, + BackoffLevel: auth.Quota.BackoffLevel, + cooldownSignals: auth.Quota.cooldownSignals, }) } *outerState = *outer @@ -1337,7 +2023,36 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { disableCooling = false } if !wasTerminalUnauthorized { + prevAuthQuotaActive := auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) + prevAuthRecoverAt := auth.Quota.NextRecoverAt + prevAuthCooldownSignals := maps.Clone(auth.Quota.cooldownSignals) + keptPriorAuthQuota := false + if !disableCooling && statusCodeFromResult(result.Error) == 429 && prevAuthQuotaActive { + if result.RetryAfter == nil { + keptPriorAuthQuota = true + } else { + cooldown := *result.RetryAfter + if cooldown < minQuotaCooldownFloor { + cooldown = minQuotaCooldownFloor + } + if prevAuthRecoverAt.After(now.Add(cooldown).Round(0)) { + keptPriorAuthQuota = true + } + } + } applyAuthFailureState(auth, result.Error, result.RetryAfter, now, disableCooling) + if !disableCooling && statusCodeFromResult(result.Error) == 429 && auth.Quota.Exceeded { + auth.Quota.cooldownEvaluated = true + if keptPriorAuthQuota { + auth.Quota.cooldownSignals = prevAuthCooldownSignals + } else if rejectedSignals := rejectedClaudeSignalsForCooldown(result.Provider, responseHeaders); len(rejectedSignals) > 0 { + auth.Quota.cooldownSignals = rejectedSignals + } else if prevAuthQuotaActive && auth.Quota.NextRecoverAt.Equal(prevAuthRecoverAt) { + auth.Quota.cooldownSignals = prevAuthCooldownSignals + } else { + auth.Quota.cooldownSignals = nil + } + } auth.cooldownRecordedSeq = m.currentExecutionSeq() } } @@ -1630,15 +2345,25 @@ func mergeModelState(target, source *ModelState) *ModelState { NextRetryAfter: target.NextRetryAfter, LastError: cloneError(preferred.LastError), Quota: QuotaState{ - Exceeded: target.Quota.Exceeded || source.Quota.Exceeded, - Reason: preferred.Quota.Reason, - NextRecoverAt: target.Quota.NextRecoverAt, - BackoffLevel: target.Quota.BackoffLevel, + Exceeded: target.Quota.Exceeded || source.Quota.Exceeded, + Reason: preferred.Quota.Reason, + NextRecoverAt: target.Quota.NextRecoverAt, + BackoffLevel: target.Quota.BackoffLevel, + cooldownEvaluated: target.Quota.cooldownEvaluated || source.Quota.cooldownEvaluated, }, UpdatedAt: target.UpdatedAt, } merged.Quota = mergeQuotaObservation(merged.Quota, fallback.Quota) merged.Quota = mergeQuotaObservation(merged.Quota, preferred.Quota) + if source.Quota.NextRecoverAt.After(target.Quota.NextRecoverAt) { + merged.Quota.cooldownSignals = maps.Clone(source.Quota.cooldownSignals) + } else if target.Quota.NextRecoverAt.After(source.Quota.NextRecoverAt) { + merged.Quota.cooldownSignals = maps.Clone(target.Quota.cooldownSignals) + } else if len(preferred.Quota.cooldownSignals) > 0 { + merged.Quota.cooldownSignals = maps.Clone(preferred.Quota.cooldownSignals) + } else if len(fallback.Quota.cooldownSignals) > 0 { + merged.Quota.cooldownSignals = maps.Clone(fallback.Quota.cooldownSignals) + } if source.NextRetryAfter.After(merged.NextRetryAfter) { merged.NextRetryAfter = source.NextRetryAfter } @@ -1659,6 +2384,11 @@ func mergeModelState(target, source *ModelState) *ModelState { if source.cooldownRecordedSeq > merged.cooldownRecordedSeq { merged.cooldownRecordedSeq = source.cooldownRecordedSeq } + merged.forceCooldownUntil = target.forceCooldownUntil + if source.forceCooldownUntil.After(merged.forceCooldownUntil) { + merged.forceCooldownUntil = source.forceCooldownUntil + } + merged.forceCooldownSeeded = target.forceCooldownSeeded || source.forceCooldownSeeded if merged.StatusMessage == "" { merged.StatusMessage = fallback.StatusMessage } @@ -1680,12 +2410,12 @@ func mergeModelState(target, source *ModelState) *ModelState { if merged.Quota.Reason == "credential_quota" || merged.Quota.Reason == "quota" || target.credentialQuotaPropagated || source.credentialQuotaPropagated || target.credentialQuotaShadow != nil || source.credentialQuotaShadow != nil { - shadow := fallback.credentialQuotaShadow - if preferred.credentialQuotaShadow != nil { - shadow = preferred.credentialQuotaShadow - } - if shadow != nil { - merged.credentialQuotaShadow = shadow.cloneShallow() + if target.credentialQuotaShadow != nil && source.credentialQuotaShadow != nil { + merged.credentialQuotaShadow = mergeModelState(target.credentialQuotaShadow.cloneShallow(), source.credentialQuotaShadow.cloneShallow()) + } else if preferred.credentialQuotaShadow != nil { + merged.credentialQuotaShadow = preferred.credentialQuotaShadow.cloneShallow() + } else if fallback.credentialQuotaShadow != nil { + merged.credentialQuotaShadow = fallback.credentialQuotaShadow.cloneShallow() } merged.credentialQuotaPropagated = target.credentialQuotaPropagated || source.credentialQuotaPropagated } @@ -1704,6 +2434,7 @@ func resetModelState(state *ModelState, now time.Time) { state.LastError = nil applyCooldownFields(&state.Quota, QuotaState{}) state.cooldownRecordedSeq = 0 + state.forceCooldownUntil = time.Time{} state.credentialQuotaShadow = nil state.credentialQuotaPropagated = false state.UpdatedAt = now @@ -1806,9 +2537,23 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) { } if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { auth.Unavailable = true + if activeAuthForce := effectiveAuthForceUntil(auth); activeAuthForce.After(now) && auth.NextRetryAfter.Before(activeAuthForce) { + auth.NextRetryAfter = activeAuthForce + } return } + activeAuthForce := effectiveAuthForceUntil(auth) + if !activeAuthForce.After(now) && !auth.forceCooldownUntil.IsZero() { + auth.forceCooldownUntil = time.Time{} + } if len(auth.ModelStates) == 0 { + if activeAuthForce.After(now) { + auth.Unavailable = true + if auth.NextRetryAfter.Before(activeAuthForce) { + auth.NextRetryAfter = activeAuthForce + } + return + } clearAggregatedAvailability(auth) return } @@ -1853,6 +2598,13 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) { } } if !hasState { + if activeAuthForce.After(now) { + auth.Unavailable = true + if auth.NextRetryAfter.Before(activeAuthForce) { + auth.NextRetryAfter = activeAuthForce + } + return + } clearAggregatedAvailability(auth) return } @@ -1862,6 +2614,12 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) { } else { auth.NextRetryAfter = time.Time{} } + if activeAuthForce.After(now) { + auth.Unavailable = true + if auth.NextRetryAfter.Before(activeAuthForce) { + auth.NextRetryAfter = activeAuthForce + } + } if quotaExceeded { auth.Quota.Exceeded = true auth.Quota.Reason = "quota" @@ -1873,10 +2631,7 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) { } else if auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) { // Retain active auth-level quota cooldown } else { - auth.Quota.Exceeded = false - auth.Quota.Reason = "" - auth.Quota.NextRecoverAt = time.Time{} - auth.Quota.BackoffLevel = 0 + applyCooldownFields(&auth.Quota, QuotaState{}) } } @@ -1920,13 +2675,12 @@ func clearAuthStateOnSuccess(auth *Auth, now time.Time) { auth.Unavailable = false auth.Status = StatusActive auth.StatusMessage = "" - auth.Quota.Exceeded = false - auth.Quota.Reason = "" - auth.Quota.NextRecoverAt = time.Time{} - auth.Quota.BackoffLevel = 0 + applyCooldownFields(&auth.Quota, QuotaState{}) auth.LastError = nil auth.NextRetryAfter = time.Time{} auth.cooldownRecordedSeq = 0 + auth.forceCooldownUntil = time.Time{} + auth.credQuotaForceUntil = time.Time{} auth.UpdatedAt = now } @@ -3140,7 +3894,17 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati if auth == nil { return } + seedInitialForceCooldownDeadlines(auth) prevAuthRetryAfter := auth.NextRetryAfter + prevAuthForceUntil := auth.forceCooldownUntil + prevAuthQuota := cooldownFieldsOf(auth.Quota) + prevAuthQuotaSignals := maps.Clone(auth.Quota.cooldownSignals) + var prevCredQuotaForceUntil time.Time + if prevAuthQuota.Exceeded && prevAuthQuota.Reason == "credential_quota" && prevAuthQuota.NextRecoverAt.After(now) && auth.credQuotaForceUntil.After(now) { + prevCredQuotaForceUntil = auth.credQuotaForceUntil + } + var ownAuthForceDeadline time.Time + isForceAuth := isForceCooldownError(resultErr) if shouldSkipCredentialCooldown(resultErr) { return } @@ -3164,6 +3928,9 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati // Without a model key the transient capacity rejection degrades to an // auth-level transient cooldown (issue 7017), never quota exhaustion. next := transientRateLimitRetryAfter(now, retryAfter, disableCooling) + if isForceAuth { + ownAuthForceDeadline = next + } if !next.IsZero() && auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(next) { next = auth.Quota.NextRecoverAt } @@ -3233,8 +4000,22 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati cooldown = minQuotaCooldownFloor } next = now.Add(cooldown).Round(0) + if isForceAuth { + ownAuthForceDeadline = next + } } else { + reusedOpenAuthWindow := auth.Quota.NextRecoverAt.After(now) + baseAuthBackoff := auth.Quota.BackoffLevel next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now) + if isForceAuth { + if !reusedOpenAuthWindow { + ownAuthForceDeadline = next + } else if prevAuthForceUntil.After(now) { + ownAuthForceDeadline = prevAuthForceUntil + } else { + ownAuthForceDeadline, _ = quotaCooldownAfterFailure(QuotaState{BackoffLevel: baseAuthBackoff}, now) + } + } } if auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(next) { next = auth.Quota.NextRecoverAt @@ -3254,15 +4035,61 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati auth.Unavailable = !auth.NextRetryAfter.IsZero() } } + if isForceAuth && auth.NextRetryAfter.IsZero() { + auth.NextRetryAfter = now.Add(transientErrorCooldown) + auth.Unavailable = true + } + if isForceAuth { + if ownAuthForceDeadline.IsZero() { + ownAuthForceDeadline = auth.NextRetryAfter + } + if ownAuthForceDeadline.After(prevAuthForceUntil) { + auth.forceCooldownUntil = ownAuthForceDeadline + } else { + auth.forceCooldownUntil = prevAuthForceUntil + } + } else { + auth.forceCooldownUntil = prevAuthForceUntil + } + auth.forceCooldownSeeded = true + if prevCredQuotaForceUntil.After(now) { + if prevCredQuotaForceUntil.After(auth.credQuotaForceUntil) { + auth.credQuotaForceUntil = prevCredQuotaForceUntil + } + if !auth.Quota.Exceeded || auth.Quota.Reason != "credential_quota" || auth.Quota.NextRecoverAt.Before(prevCredQuotaForceUntil) { + restored := prevAuthQuota + if disableCooling && restored.NextRecoverAt.After(prevCredQuotaForceUntil) { + restored.NextRecoverAt = prevCredQuotaForceUntil + } else if !disableCooling && auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(restored.NextRecoverAt) { + restored.NextRecoverAt = auth.Quota.NextRecoverAt + } + restored.cooldownSignals = prevAuthQuotaSignals + applyCooldownFields(&auth.Quota, restored) + } + } + activeAuthForce := auth.forceCooldownUntil + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.credQuotaForceUntil.After(activeAuthForce) { + activeAuthForce = auth.credQuotaForceUntil + } + if activeAuthForce.After(now) { + if auth.NextRetryAfter.Before(activeAuthForce) { + auth.NextRetryAfter = activeAuthForce + } + auth.Unavailable = true + if (!auth.Quota.Exceeded || auth.Quota.NextRecoverAt.IsZero()) && prevAuthQuota.Exceeded && prevAuthQuota.NextRecoverAt.After(now) { + restored := prevAuthQuota + if disableCooling && restored.NextRecoverAt.After(activeAuthForce) { + restored.NextRecoverAt = activeAuthForce + } + restored.cooldownSignals = prevAuthQuotaSignals + applyCooldownFields(&auth.Quota, restored) + } + } // A later failure only extends a still-live credential cooldown; a // deliberate zero write (disableCooling) still clears it. if !auth.NextRetryAfter.IsZero() && prevAuthRetryAfter.After(auth.NextRetryAfter) && prevAuthRetryAfter.After(now) { auth.NextRetryAfter = prevAuthRetryAfter } - if resultErr != nil && resultErr.Code == ErrorCodeForceCooldown && auth.NextRetryAfter.IsZero() { - auth.NextRetryAfter = now.Add(transientErrorCooldown) - auth.Unavailable = true - } } // quotaCooldownAfterFailure returns the recovery deadline and backoff level for diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index b1a8e8c626f..d6441a3956d 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -25,6 +25,14 @@ import ( log "github.com/sirupsen/logrus" ) +func (m *Manager) newUpstreamAttemptContext(ctx context.Context) context.Context { + ctx = logging.WithFreshResponseHeadersHolder(ctx) + if RequestScopedErrorConfigFromContext(ctx) == nil { + ctx = WithRequestScopedErrorConfig(ctx, m.runtimeConfigSnapshot()) + } + return cliproxyexecutor.WithUpstreamAttemptTracker(ctx) +} + func newUpstreamAttemptContext(ctx context.Context) context.Context { ctx = logging.WithFreshResponseHeadersHolder(ctx) return cliproxyexecutor.WithUpstreamAttemptTracker(ctx) @@ -92,8 +100,35 @@ func unwrapUpstreamExecutionAttempt(err error) error { } func unwrapExecutionBoundaryError(err error) error { - err = unwrapRequestStopError(err) - return unwrapUpstreamExecutionAttempt(err) + for err != nil { + switch wrapped := err.(type) { + case requestStopError: + err = wrapped.error + case schedulerRejectError: + err = wrapped.error + case *roundEndedBySchedulerRejectError: + if wrapped == nil || wrapped.cause == nil { + return err + } + err = wrapped.cause + case *upstreamExecutionAttemptError: + if wrapped == nil || wrapped.cause == nil { + return err + } + err = wrapped.cause + default: + return err + } + } + return nil +} + +func roundPickFailureError(lastErr, upstreamErr, errPick error) error { + preferred := preferredExecutionAttemptError(lastErr, upstreamErr) + if isSchedulerRejectError(errPick) { + return markRoundEndedBySchedulerReject(preferred) + } + return preferred } func preferredExecutionAttemptError(fallback, upstream error) error { @@ -122,7 +157,7 @@ func (m *Manager) normalizeSessionAffinityNeverTerminalError(err error, provider if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return err } - if isRequestTerminatedError(err) || isRequestStopError(err) || isRequestInvalidError(err) || isSessionAffinityRetryError(err) { + if isRequestTerminatedError(err) || isRequestStopError(err) || isSchedulerRejectError(err) || isRequestInvalidError(err) || isSessionAffinityRetryError(err) { return err } var authErr *Error @@ -182,6 +217,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye var lastErr error var preferredUpstreamErr error + var terminalRoundRejectedByScheduler bool retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { roundAttempted := make(map[string]struct{}) @@ -193,6 +229,14 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) } + if isSchedulerRejectError(errExec) { + if lastErr == nil && preferredUpstreamErr == nil { + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) + } + terminalRoundRejectedByScheduler = true + break + } + terminalRoundRejectedByScheduler = hasRoundEndedBySchedulerReject(errExec) if hasUpstreamExecutionAttempt(errExec) { preferredUpstreamErr = errExec } @@ -214,7 +258,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) lastErr = unwrapExecutionBoundaryError(lastErr) lastErr = m.normalizeSessionAffinityNeverTerminalError(lastErr, normalized, retryModel, opts) - if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { + if !terminalRoundRejectedByScheduler && hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { return cliproxyexecutor.Response{}, errCredits } else if ok { @@ -256,6 +300,12 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) } + if isSchedulerRejectError(errExec) { + if lastErr == nil && preferredUpstreamErr == nil { + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) + } + break + } if hasUpstreamExecutionAttempt(errExec) { preferredUpstreamErr = errExec } @@ -303,6 +353,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli var lastErr error var preferredUpstreamErr error + var terminalRoundRejectedByScheduler bool homeRetryLimit := -1 retryModel := authSelectionModelFromOptions(opts, req.Model) attempt := 0 @@ -335,6 +386,14 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli if isRequestTerminatedError(errStream) || isRequestStopError(errStream) { return nil, unwrapExecutionBoundaryError(errStream) } + if isSchedulerRejectError(errStream) { + if lastErr == nil && preferredUpstreamErr == nil { + return nil, unwrapExecutionBoundaryError(errStream) + } + terminalRoundRejectedByScheduler = true + break + } + terminalRoundRejectedByScheduler = hasRoundEndedBySchedulerReject(errStream) lastErr = errStream wait, shouldRetry := m.shouldRetryAfterErrorWithAttempted(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry, roundAttempted) if !shouldRetry { @@ -358,7 +417,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli } lastErr = unwrapExecutionBoundaryError(lastErr) lastErr = m.normalizeSessionAffinityNeverTerminalError(lastErr, normalized, retryModel, opts) - if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { + if !terminalRoundRejectedByScheduler && hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { return nil, errCredits } else if ok { @@ -593,7 +652,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) + return cliproxyexecutor.Response{}, roundPickFailureError(lastErr, upstreamErr, errPick) } return cliproxyexecutor.Response{}, errPick } @@ -610,7 +669,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAliasForRequest(execCtx, auth, routeModel, pickOpts.Headers) if len(models) == 0 { @@ -636,7 +695,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req var authErr error didRefreshOnUnauthorized := false for _, upstreamModel := range models { - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -687,7 +746,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized, execOpts.Headers); okRefresh { auth = refreshed didRefreshOnUnauthorized = true - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) executionSeq = m.nextExecutionSeq() startRetry := time.Now() @@ -719,7 +778,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if isCredentialScopedError(errExec) { result.CredentialScope = true } - action, okAction := matchRequestScopedErrorAction(auth, errExec, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, errExec, m.requestScopedErrorConfig(execCtx)) applyRequestScopedActionToResult(action, okAction, &result) if isResponsesCompactAvailabilityNeutralError(execOpts, errExec, result.Error) { m.recordAvailabilityNeutralResult(execCtx, result) @@ -753,7 +812,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req return resp, nil } if authErr != nil { - action, okAction := matchRequestScopedErrorAction(auth, authErr, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, authErr, m.requestScopedErrorConfig(execCtx)) if okAction { if isRequestScopedStop(action, okAction) { return cliproxyexecutor.Response{}, wrapRequestStopError(authErr) @@ -811,7 +870,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) + return cliproxyexecutor.Response{}, roundPickFailureError(lastErr, upstreamErr, errPick) } return cliproxyexecutor.Response{}, errPick } @@ -828,7 +887,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAliasForRequest(execCtx, auth, routeModel, pickOpts.Headers) if len(models) == 0 { @@ -854,7 +913,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, var authErr error didRefreshOnUnauthorized := false for _, upstreamModel := range models { - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -905,7 +964,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized, execOpts.Headers); okRefresh { auth = refreshed didRefreshOnUnauthorized = true - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) executionSeq = m.nextExecutionSeq() startRetry := time.Now() @@ -934,7 +993,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if ra := retryAfterFromError(errExec); ra != nil { result.RetryAfter = ra } - action, okAction := matchRequestScopedErrorAction(auth, errExec, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, errExec, m.requestScopedErrorConfig(execCtx)) applyRequestScopedActionToResult(action, okAction, &result) // Some Anthropic-compatible upstreams and plugin executors do not // implement count_tokens and return a generic endpoint 404 or an @@ -976,7 +1035,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, return resp, nil } if authErr != nil { - action, okAction := matchRequestScopedErrorAction(auth, authErr, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, authErr, m.requestScopedErrorConfig(execCtx)) if okAction { if isRequestScopedStop(action, okAction) { return cliproxyexecutor.Response{}, wrapRequestStopError(authErr) @@ -1069,7 +1128,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if homeMode { return nil, markHomeRetryRoundExhausted(preferredErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errPick)) } - return nil, preferredErr + return nil, roundPickFailureError(lastErr, upstreamErr, errPick) } return nil, errPick } @@ -1161,7 +1220,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string Reused: selection.Retained(), Model: routeModel, CredentialPolicy: credentialPolicyFromContext(ctx), }) } - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.executionModelCandidatesWithAlias(auth, routeModel) // Home owns dispatched credential availability; only local selections need state filtering. if selection == nil { @@ -1282,15 +1341,16 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if hasUpstreamExecutionAttempt(errStream) { upstreamErr = errStream } - action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) - if isRequestScopedStop(action, okAction) || (!okAction && (isRequestInvalidError(errStream) || cliproxyexecutor.IsUpstreamWebsocketTurnError(errStream))) { + action, okAction := matchRequestScopedErrorAction(auth, errStream, m.requestScopedErrorConfig(execCtx)) + stoppedByRule := isRequestStopError(errStream) || isRequestScopedStop(action, okAction) + if stoppedByRule || (!okAction && (isRequestInvalidError(errStream) || cliproxyexecutor.IsUpstreamWebsocketTurnError(errStream))) { // Only redispatch needs a release acknowledgement. Preserve the final // response even when Home cannot acknowledge the released admission. if selection != nil { releaseAttempt() selection.End("request_stopped") } - if isRequestScopedStop(action, okAction) { + if stoppedByRule { return nil, wrapRequestStopError(errStream) } return nil, errStream @@ -2353,7 +2413,7 @@ func (m *Manager) VerifyInference(ctx context.Context, id, model string) (*Auth, Metadata: map[string]any{cliproxyexecutor.CanonicalSessionIDMetadataKey: sessionID}, } - execCtx := newUpstreamAttemptContext(ctx) + execCtx := m.newUpstreamAttemptContext(ctx) execCtx = syncMetadataSessionToContext(execCtx, opts.Metadata) executionSeq := m.nextExecutionSeq() accessToken := authAccessTokenForRequest(execCtx, auth, opts.Headers) diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 26e85115fa5..ddd7a2bd235 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -1438,7 +1438,7 @@ func hasAntigravityProvider(providers []string) bool { } func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool { - if isRequestTerminatedError(lastErr) || isRequestStopError(lastErr) { + if isRequestTerminatedError(lastErr) || isRequestStopError(lastErr) || isSchedulerRejectError(lastErr) || hasRoundEndedBySchedulerReject(lastErr) { return false } status := statusCodeFromError(lastErr) @@ -1497,6 +1497,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy clearSchedulerSelectionID(creditsOpts.Metadata) creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) creditsCtx = coreusage.WithSelectionID(creditsCtx, "") + creditsCtx = m.newUpstreamAttemptContext(creditsCtx) preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) if errPrepare != nil { continue @@ -1508,6 +1509,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy continue } for _, upstreamModel := range models { + creditsCtx = m.newUpstreamAttemptContext(creditsCtx) resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -1563,6 +1565,7 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl clearSchedulerSelectionID(creditsOpts.Metadata) creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) creditsCtx = coreusage.WithSelectionID(creditsCtx, "") + creditsCtx = m.newUpstreamAttemptContext(creditsCtx) preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) if errPrepare != nil { continue diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 8bccc582b0d..11fcf62bd26 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -138,7 +138,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c execCtx = coreusage.WithHomeSelection(execCtx, coreusage.HomeSelection{ Reused: selection.Retained(), Model: routeModel, CredentialPolicy: credentialPolicyFromContext(ctx), }) - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) @@ -177,7 +177,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c continue } for _, upstreamModel := range models { - execCtx = newUpstreamAttemptContext(execCtx) + execCtx = m.newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -278,7 +278,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c if isCredentialScopedError(errExecute) { result.CredentialScope = true } - action, okAction := matchRequestScopedErrorAction(preparedAuth, errExecute, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(preparedAuth, errExecute, m.requestScopedErrorConfig(execCtx)) applyRequestScopedActionToResult(action, okAction, &result) m.reportHomeResult(execCtx, result, preparedAuth) roundTiming.Observe(errExecute) diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index b6bc8a90d6e..1128cdbbd8e 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -112,8 +112,11 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { } auth.UpdatedAt = now cooldownStateChanged := normalizeModelStates(auth) - if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { + seedInitialForceCooldownDeadlines(auth) + if auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged + } else if m.cooldownDisabledForAuth(auth) { + cooldownStateChanged = clearDisabledCooldownStateForAuth(auth, now) || cooldownStateChanged } auth.EnsureIndex() m.mu.Lock() @@ -425,6 +428,9 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.Quota.ObservedAt = observation.ObservedAt auth.Quota.Signals = observation.Signals auth.cooldownRecordedSeq = existing.cooldownRecordedSeq + auth.forceCooldownUntil = existing.forceCooldownUntil + auth.credQuotaForceUntil = existing.credQuotaForceUntil + auth.forceCooldownSeeded = existing.forceCooldownSeeded } if credChanged || mode == updateModeRefresh { applyPreservedRejectedTokens(existing, auth, mode == updateModeRefresh) @@ -452,6 +458,7 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.Unavailable = false auth.NextRetryAfter = time.Time{} auth.cooldownRecordedSeq = 0 + auth.forceCooldownUntil = time.Time{} auth.LastError = nil auth.StatusMessage = "" auth.Status = StatusActive @@ -511,6 +518,25 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.Quota.ObservedAt = observation.ObservedAt auth.Quota.Signals = observation.Signals auth.cooldownRecordedSeq = existing.cooldownRecordedSeq + auth.forceCooldownUntil = existing.forceCooldownUntil + auth.credQuotaForceUntil = existing.credQuotaForceUntil + auth.forceCooldownSeeded = existing.forceCooldownSeeded + if auth.Status == StatusActive { + auth.Status = existing.Status + } + } else if !credChanged && mode == updateModeReplace && existing.forceCooldownUntil.After(now) && !auth.NextRetryAfter.After(now) { + auth.Unavailable = existing.Unavailable + auth.NextRetryAfter = existing.NextRetryAfter + auth.cooldownRecordedSeq = existing.cooldownRecordedSeq + auth.forceCooldownUntil = existing.forceCooldownUntil + auth.credQuotaForceUntil = existing.credQuotaForceUntil + auth.forceCooldownSeeded = existing.forceCooldownSeeded + if auth.LastError == nil && existing.LastError != nil { + auth.LastError = cloneError(existing.LastError) + } + if auth.StatusMessage == "" { + auth.StatusMessage = existing.StatusMessage + } if auth.Status == StatusActive { auth.Status = existing.Status } @@ -535,8 +561,11 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo } auth.UpdatedAt = now cooldownStateChanged = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) || cooldownStateChanged - if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { + seedInitialForceCooldownDeadlines(auth) + if auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged + } else if m.cooldownDisabledForAuth(auth) { + cooldownStateChanged = clearDisabledCooldownStateForAuth(auth, now) || cooldownStateChanged } if existing.Metadata != nil { if persistedHints, hasHints := existing.Metadata["antigravity_capability_hints"]; hasHints { @@ -622,8 +651,17 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo if m.scheduler != nil { schedulerSnapshot = authClone.Clone() } + shouldForgetWeighted := authClone.Disabled || + authClone.Status == StatusDisabled || + (authWeight(existingBeforeSave) > 0 && authWeight(authClone) <= 0) || + executorKeyFromAuth(existingBeforeSave) != executorKeyFromAuth(authClone) + forgetEpoch := authClone.RegistrationEpoch + forgetGen := authClone.Generation m.mu.Unlock() _ = schedulerSnapshot + if shouldForgetWeighted { + m.forgetWeightedAuthUpTo(authClone.ID, forgetEpoch, forgetGen) + } return auth.Clone(), true, cooldownStateChanged, nil } @@ -698,6 +736,7 @@ func (m *Manager) RemoveIf(ctx context.Context, id string, predicate func(*Auth) m.authEpochs[id]++ tombstoneEpoch := m.authEpochs[id] m.mu.Unlock() + m.forgetWeightedAuthUpTo(id, tombstoneEpoch, ^uint64(0)) releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { @@ -733,6 +772,22 @@ func (m *Manager) invalidateSessionAffinity(authID string) { } } +func (m *Manager) forgetWeightedAuthUpTo(authID string, maxEpoch, maxGen uint64) { + if m == nil || authID == "" { + return + } + sel := m.Selector() + if epochForgetter, ok := sel.(interface { + ForgetAuthUpTo(string, uint64, uint64) + }); ok && epochForgetter != nil { + epochForgetter.ForgetAuthUpTo(authID, maxEpoch, maxGen) + return + } + if forgetter, ok := sel.(interface{ ForgetAuth(string) }); ok && forgetter != nil { + forgetter.ForgetAuth(authID) + } +} + // Load resets manager state from the backing store. func (m *Manager) Load(ctx context.Context) error { // Exclude complete Save/publication transactions, including new registrations @@ -772,6 +827,7 @@ func (m *Manager) Load(ctx context.Context) error { } auth.EnsureIndex() _ = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) + seedInitialForceCooldownDeadlines(auth) if len(auth.ModelStates) > 0 { updateAggregatedAvailability(auth, now) } @@ -820,6 +876,9 @@ func (m *Manager) Load(ctx context.Context) error { m.scheduler.RecordRemovalTombstone(rt.id, rt.epoch) } } + for _, rt := range removedTombstones { + m.forgetWeightedAuthUpTo(rt.id, rt.epoch, ^uint64(0)) + } m.structuralEpoch.Add(1) m.syncScheduler() return nil diff --git a/sdk/cliproxy/auth/conductor_models.go b/sdk/cliproxy/auth/conductor_models.go index dd5ab10169c..d0d3f49a75b 100644 --- a/sdk/cliproxy/auth/conductor_models.go +++ b/sdk/cliproxy/auth/conductor_models.go @@ -237,38 +237,98 @@ func (m *Manager) clientModelProjectionForAuth(auth *Auth, routeModel string, no } state := existingModelState(auth, targetKey) - isSuspended := auth.Disabled || auth.Status == StatusDisabled - if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { - isSuspended = true + authDisabled := auth.Disabled || auth.Status == StatusDisabled + stateDisabled := state != nil && state.Status == StatusDisabled + terminalUnauthorized := hasUnauthorizedAuthFailure(auth) + openEnded := authDisabled || stateDisabled || terminalUnauthorized + isSuspended := openEnded + + var credQuotaUntil time.Time + credQuotaSuspended := false + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" { + if auth.Quota.NextRecoverAt.IsZero() { + isSuspended = true + credQuotaSuspended = true + openEnded = true + } else if auth.Quota.NextRecoverAt.After(now) { + isSuspended = true + credQuotaSuspended = true + credQuotaUntil = auth.Quota.NextRecoverAt + } } + isQuotaExceeded := false - var suspendReason string + stateSuspended := false + var stateUntil time.Time if state != nil { - if state.Status == StatusDisabled || state.Unavailable || (!state.NextRetryAfter.IsZero() && state.NextRetryAfter.After(now)) { + if stateDisabled { + stateSuspended = true + } else if blocked, _, retryAfter := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now); blocked { isSuspended = true + stateSuspended = true + if retryAfter.IsZero() { + openEnded = true + } else { + stateUntil = retryAfter + } } if state.Quota.Exceeded && (state.Quota.NextRecoverAt.IsZero() || state.Quota.NextRecoverAt.After(now)) { isQuotaExceeded = true } - if isSuspended { - suspendReason = cooldownReason(state.StatusMessage, state.Quota, state.LastError) - } } - if len(auth.ModelStates) == 0 && auth.Unavailable && auth.NextRetryAfter.After(now) { + + var authFallbackUntil time.Time + if len(auth.ModelStates) == 0 { // With no per-model states, scheduling falls back to the credential-wide // cooldown (isAuthBlockedForModel); the projection must agree with it. // When states exist, unmatched models stay schedulable by design, so the // credential-wide fields must not suspend them here either. - isSuspended = true + if blocked, _, retryAfter := availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now); blocked { + isSuspended = true + if retryAfter.IsZero() { + openEnded = true + } else { + authFallbackUntil = retryAfter + } + } + } + + var suspendReason string + var suspendUntil time.Time + if isSuspended { + switch { + case authDisabled || stateDisabled: + suspendReason = "disabled" + case terminalUnauthorized: + suspendReason = cooldownReason(auth.StatusMessage, QuotaState{}, auth.LastError) + if suspendReason == "" { + suspendReason = "unauthorized" + } + case stateSuspended: + suspendReason = cooldownReason(state.StatusMessage, state.Quota, state.LastError) + suspendUntil = stateUntil + if credQuotaSuspended && (strings.EqualFold(suspendReason, "quota") || strings.EqualFold(suspendReason, "credential_quota")) && credQuotaUntil.After(suspendUntil) { + suspendUntil = credQuotaUntil + } + default: + suspendUntil = authFallbackUntil + if credQuotaUntil.After(suspendUntil) { + suspendUntil = credQuotaUntil + } + } + if suspendReason == "" { + suspendReason = cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError) + } } - if isSuspended && suspendReason == "" { - suspendReason = cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError) + if openEnded || !isSuspended { + suspendUntil = time.Time{} } return registry.ClientModelProjection{ ModelID: targetModel, Suspended: isSuspended, SuspendReason: suspendReason, + SuspendUntil: suspendUntil, QuotaExceeded: isQuotaExceeded, } } diff --git a/sdk/cliproxy/auth/conductor_models_transient_projection_test.go b/sdk/cliproxy/auth/conductor_models_transient_projection_test.go new file mode 100644 index 00000000000..8976df8d6d3 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_models_transient_projection_test.go @@ -0,0 +1,301 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" +) + +func openAIListedModelSet(models []map[string]any) map[string]bool { + out := make(map[string]bool, len(models)) + for _, m := range models { + if id, _ := m["id"].(string); id != "" { + out[id] = true + } + } + return out +} + +func TestManager_ClientModelProjection_TransientCooldownExpiresWithForkedAlias(t *testing.T) { + previousCooldown := transientErrorCooldownSeconds.Load() + transientErrorCooldownSeconds.Store(0) + t.Cleanup(func() { transientErrorCooldownSeconds.Store(previousCooldown) }) + + m := NewManager(nil, nil, nil) + m.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + "xai": {{ + Name: "grok-4.7", + Alias: "grok-4.7-cp", + Fork: true, + ForceMapping: true, + }}, + }) + + auth := &Auth{ + ID: "xai-transient-7421", + Provider: "xai", + Status: StatusActive, + } + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + reg := registry.GetGlobalRegistry() + nowUnix := time.Now().Unix() + reg.RegisterClient(auth.ID, "xai", []*registry.ModelInfo{ + {ID: "grok-4.7", OwnedBy: "xai", Created: nowUnix}, + {ID: "grok-4.7-cp", OwnedBy: "xai", Created: nowUnix}, + {ID: "grok-4.6", OwnedBy: "xai", Created: nowUnix}, + }) + t.Cleanup(func() { reg.UnregisterClient(auth.ID) }) + + m.MarkResult(context.Background(), Result{ + AuthID: auth.ID, + Provider: "xai", + Model: "grok-4.7", + Success: false, + Error: &Error{ + HTTPStatus: http.StatusRequestTimeout, + Message: "xai stream error: stream disconnected", + }, + }) + + snap, ok := m.GetByID(auth.ID) + if !ok || snap == nil { + t.Fatal("expected auth snapshot after MarkResult") + } + state := existingModelState(snap, "grok-4.7") + if state == nil || !state.Unavailable || state.NextRetryAfter.IsZero() { + t.Fatalf("expected active transient model cooldown, got %+v", state) + } + + listedDuring := openAIListedModelSet(reg.GetAvailableModels("openai")) + if listedDuring["grok-4.7"] || listedDuring["grok-4.7-cp"] { + t.Fatalf("expected grok-4.7 and grok-4.7-cp to be hidden during live transient cooldown, got %v", listedDuring) + } + if !listedDuring["grok-4.6"] { + t.Fatalf("expected unaffected sibling grok-4.6 to remain listed, got %v", listedDuring) + } + + during := state.NextRetryAfter.Add(-time.Second) + after := state.NextRetryAfter.Add(time.Second) + + for _, modelID := range []string{"grok-4.7", "grok-4.7-cp"} { + selectionModel := m.selectionModelForAuth(snap, modelID) + blockedDuring, _, _ := isAuthBlockedForModel(snap, selectionModel, during) + if !blockedDuring { + t.Fatalf("expected selector to block %s during transient cooldown", modelID) + } + projDuring := m.clientModelProjectionForAuth(snap, modelID, during) + if !projDuring.Suspended { + t.Fatalf("expected %s to be projected suspended during transient cooldown: %+v", modelID, projDuring) + } + if !projDuring.SuspendUntil.Equal(state.NextRetryAfter) { + t.Fatalf("expected %s SuspendUntil=%v, got %v", modelID, state.NextRetryAfter, projDuring.SuspendUntil) + } + + blockedAfter, _, _ := isAuthBlockedForModel(snap, selectionModel, after) + if blockedAfter { + t.Fatalf("expected selector to unblock %s after transient cooldown expired", modelID) + } + projAfter := m.clientModelProjectionForAuth(snap, modelID, after) + if projAfter.Suspended { + t.Fatalf("expected %s projection Suspended=false once NextRetryAfter passed, got %+v", modelID, projAfter) + } + if !projAfter.SuspendUntil.IsZero() { + t.Fatalf("expected %s SuspendUntil to be zero once expired, got %v", modelID, projAfter.SuspendUntil) + } + } + + // Simulate the recorded transient deadline now lying in the past without a + // new MarkResult or credential refresh: /v1/models must list both models. + epoch := reg.ClientRegistrationEpoch(auth.ID) + staleUntil := time.Now().Add(-time.Second) + if !reg.ApplyClientModelProjections(auth.ID, epoch, snap.Generation, []registry.ClientModelProjection{ + {ModelID: "grok-4.7", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: staleUntil}, + {ModelID: "grok-4.7-cp", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: staleUntil}, + {ModelID: "grok-4.6"}, + }) { + t.Fatal("projection with past deadline rejected") + } + listedAfter := openAIListedModelSet(reg.GetAvailableModels("openai")) + if !listedAfter["grok-4.7"] || !listedAfter["grok-4.7-cp"] || !listedAfter["grok-4.6"] { + t.Fatalf("expected /v1/models to list grok-4.7 and grok-4.7-cp after deadline passed, got %v", listedAfter) + } +} + +func TestManager_ClientModelProjection_OpenEndedAuthGatesStayUnbounded(t *testing.T) { + now := time.Now() + transientUntil := now.Add(time.Minute) + m := NewManager(nil, nil, nil) + + cases := []struct { + name string + auth *Auth + }{ + { + name: "disabled_auth_with_model_transient_deadline", + auth: &Auth{ + ID: "disabled-with-transient", + Provider: "xai", + Disabled: true, + Status: StatusDisabled, + ModelStates: map[string]*ModelState{ + "grok-4.7": { + Unavailable: true, + Status: StatusError, + StatusMessage: "xai stream error: stream disconnected", + NextRetryAfter: transientUntil, + }, + }, + }, + }, + { + name: "disabled_auth_with_active_model_quota", + auth: &Auth{ + ID: "disabled-with-quota", + Provider: "xai", + Disabled: true, + Status: StatusDisabled, + ModelStates: map[string]*ModelState{ + "grok-4.7": { + Unavailable: true, + Status: StatusError, + NextRetryAfter: transientUntil, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: transientUntil, + }, + }, + }, + }, + }, + { + name: "disabled_auth_with_auth_level_deadline", + auth: &Auth{ + ID: "disabled-auth-level", + Provider: "xai", + Disabled: true, + Unavailable: true, + NextRetryAfter: transientUntil, + }, + }, + { + name: "terminal_unauthorized_auth_without_model_state", + auth: &Auth{ + ID: "unauthorized-no-model-state", + Provider: "xai", + Unavailable: true, + Status: StatusError, + LastError: &Error{ + HTTPStatus: http.StatusUnauthorized, + Code: "unauthorized", + Message: "refresh token invalid", + }, + }, + }, + { + name: "terminal_unauthorized_auth_with_model_transient_deadline", + auth: &Auth{ + ID: "unauthorized-with-transient", + Provider: "xai", + Unavailable: true, + Status: StatusError, + LastError: &Error{ + HTTPStatus: http.StatusUnauthorized, + Code: "unauthorized", + Message: "refresh token invalid", + }, + ModelStates: map[string]*ModelState{ + "grok-4.7": { + Unavailable: true, + Status: StatusError, + StatusMessage: "xai stream error: stream disconnected", + NextRetryAfter: transientUntil, + }, + }, + }, + }, + { + name: "disabled_model_state_with_transient_deadline", + auth: &Auth{ + ID: "model-disabled-with-transient", + Provider: "xai", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + "grok-4.7": { + Unavailable: true, + Status: StatusDisabled, + StatusMessage: "model disabled", + NextRetryAfter: transientUntil, + }, + }, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + proj := m.clientModelProjectionForAuth(tc.auth, "grok-4.7", now) + if !proj.Suspended { + t.Fatalf("expected Suspended=true, got %+v", proj) + } + if !proj.SuspendUntil.IsZero() { + t.Fatalf("open-ended gate must not carry finite SuspendUntil, got %v", proj.SuspendUntil) + } + if proj.SuspendReason == "quota" || proj.SuspendReason == "credential_quota" { + t.Fatalf("open-ended disabled/unauthorized gate must not carry quota SuspendReason, got %q", proj.SuspendReason) + } + }) + } +} + +func TestManager_ClientModelProjection_CredentialQuotaIgnoresExpiredModelTransientError(t *testing.T) { + now := time.Now() + recoverAt := now.Add(time.Hour) + m := NewManager(nil, nil, nil) + auth := &Auth{ + ID: "xai-cred-quota-stale-transient", + Provider: "xai", + Status: StatusError, + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: recoverAt, + }, + ModelStates: map[string]*ModelState{ + "grok-4.7": { + Unavailable: true, + Status: StatusError, + StatusMessage: "xai stream error: stream disconnected", + NextRetryAfter: now.Add(-time.Minute), + LastError: &Error{ + HTTPStatus: http.StatusRequestTimeout, + Message: "xai stream error: stream disconnected", + }, + }, + }, + } + + proj := m.clientModelProjectionForAuth(auth, "grok-4.7", now) + if !proj.Suspended { + t.Fatalf("expected model under live credential_quota to be suspended, got %+v", proj) + } + if proj.SuspendReason != "credential_quota" { + t.Fatalf("expected SuspendReason=%q so sole-credential quota stays listed, got %q", "credential_quota", proj.SuspendReason) + } + if !proj.SuspendUntil.Equal(recoverAt) { + t.Fatalf("expected SuspendUntil=%v, got %v", recoverAt, proj.SuspendUntil) + } + + // Sibling model with no ModelState entry also inherits the credential_quota deadline. + siblingProj := m.clientModelProjectionForAuth(auth, "grok-4.6", now) + if !siblingProj.Suspended || siblingProj.SuspendReason != "credential_quota" || !siblingProj.SuspendUntil.Equal(recoverAt) { + t.Fatalf("unexpected sibling projection under credential_quota: %+v", siblingProj) + } +} diff --git a/sdk/cliproxy/auth/conductor_persistence.go b/sdk/cliproxy/auth/conductor_persistence.go index fef0e160024..865a3eaef2f 100644 --- a/sdk/cliproxy/auth/conductor_persistence.go +++ b/sdk/cliproxy/auth/conductor_persistence.go @@ -2,9 +2,11 @@ package auth import ( "context" + "maps" "reflect" "sync" "sync/atomic" + "time" log "github.com/sirupsen/logrus" ) @@ -277,6 +279,11 @@ func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) { if d.IsNil() { d.Set(reflect.MakeMap(d.Type())) } + if field.Name == "ModelStates" && newValue.IsValid() && !newValue.IsNil() { + if ms, ok := newValue.Interface().(*ModelState); ok { + newValue = reflect.ValueOf(ms.Clone()) + } + } d.SetMapIndex(key, newValue) } continue @@ -285,6 +292,39 @@ func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) { d.Set(n) } } + mergeRuntimeCooldownSaveDelta(target, before, after, preserveConcurrent) +} + +func mergeRuntimeCooldownSaveDelta(target, before, after *Auth, preserveConcurrent bool) { + if target == nil || before == nil || after == nil { + return + } + if !after.forceCooldownUntil.Equal(before.forceCooldownUntil) && (!preserveConcurrent || target.forceCooldownUntil.Equal(before.forceCooldownUntil)) { + target.forceCooldownUntil = after.forceCooldownUntil + } + if !after.credQuotaForceUntil.Equal(before.credQuotaForceUntil) && (!preserveConcurrent || target.credQuotaForceUntil.Equal(before.credQuotaForceUntil)) { + target.credQuotaForceUntil = after.credQuotaForceUntil + } + if after.forceCooldownSeeded != before.forceCooldownSeeded && (!preserveConcurrent || target.forceCooldownSeeded == before.forceCooldownSeeded) { + target.forceCooldownSeeded = after.forceCooldownSeeded + } + if after.cooldownRecordedSeq != before.cooldownRecordedSeq && (!preserveConcurrent || target.cooldownRecordedSeq == before.cooldownRecordedSeq) { + target.cooldownRecordedSeq = after.cooldownRecordedSeq + } + if !maps.Equal(after.Quota.cooldownSignals, before.Quota.cooldownSignals) && (!preserveConcurrent || maps.Equal(target.Quota.cooldownSignals, before.Quota.cooldownSignals)) { + target.Quota.cooldownSignals = maps.Clone(after.Quota.cooldownSignals) + } + if after.Quota.cooldownEvaluated != before.Quota.cooldownEvaluated && (!preserveConcurrent || target.Quota.cooldownEvaluated == before.Quota.cooldownEvaluated) { + target.Quota.cooldownEvaluated = after.Quota.cooldownEvaluated + } + if (before.Unavailable || before.Quota.Exceeded) && + after.Unavailable == before.Unavailable && + after.NextRetryAfter.Equal(before.NextRetryAfter) && + !target.Unavailable && target.NextRetryAfter.IsZero() && + (!target.Quota.Exceeded || target.Quota.NextRecoverAt.IsZero()) { + target.forceCooldownUntil = time.Time{} + target.credQuotaForceUntil = time.Time{} + } } func equalSaveValue(a, b reflect.Value) bool { diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index 819dd94785b..6bfc063a65c 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -1014,13 +1014,34 @@ func applyDevinQuotaRestoration(auth, existing *Auth, now time.Time) bool { } resetModelState(state, now) } - auth.Status = StatusActive - auth.StatusMessage = "" - auth.LastError = nil - auth.Unavailable = false - auth.NextRetryAfter = time.Time{} + independentForce := effectiveAuthForceUntil(existing) + auth.credQuotaForceUntil = time.Time{} + auth.forceCooldownSeeded = true + if independentForce.After(now) { + auth.forceCooldownUntil = independentForce + auth.Unavailable = true + auth.NextRetryAfter = independentForce + auth.Status = StatusError + if isForceCooldownError(existing.LastError) { + auth.LastError = cloneError(existing.LastError) + auth.StatusMessage = existing.StatusMessage + } else { + auth.LastError = nil + auth.StatusMessage = ErrorCodeForceCooldown + } + } else { + auth.forceCooldownUntil = time.Time{} + auth.Status = StatusActive + auth.StatusMessage = "" + auth.LastError = nil + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + } if len(auth.ModelStates) > 0 { updateAggregatedAvailability(auth, now) + if auth.Unavailable && auth.Status == StatusActive { + auth.Status = StatusError + } } return true } @@ -1135,7 +1156,7 @@ func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, e return auth, false } // A 401 matched by a configured request-scoped rule is the rule's to handle: no rejection, no refresh. - if _, ruleMatched := matchRequestScopedErrorAction(auth, execErr, m.runtimeConfigSnapshot()); ruleMatched { + if _, ruleMatched := matchRequestScopedErrorAction(auth, execErr, m.requestScopedErrorConfig(ctx)); ruleMatched { return auth, false } failedToken := authAccessTokenForRequest(ctx, auth, clientHeaders...) diff --git a/sdk/cliproxy/auth/conductor_request_scoped_errors.go b/sdk/cliproxy/auth/conductor_request_scoped_errors.go index 41a14ec8def..eb1d1aa5e5e 100644 --- a/sdk/cliproxy/auth/conductor_request_scoped_errors.go +++ b/sdk/cliproxy/auth/conductor_request_scoped_errors.go @@ -1,6 +1,7 @@ package auth import ( + "context" "encoding/json" "errors" "regexp" @@ -11,6 +12,31 @@ import ( cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" ) +type requestScopedErrorConfigContextKey struct{} + +// WithRequestScopedErrorConfig stores the runtime config snapshot on ctx so +// usage reporters can evaluate request-scoped error rules when publishing +// attempt failures. +func WithRequestScopedErrorConfig(ctx context.Context, cfg *internalconfig.Config) context.Context { + if ctx == nil { + ctx = context.Background() + } + if cfg == nil { + return ctx + } + return context.WithValue(ctx, requestScopedErrorConfigContextKey{}, cfg) +} + +// RequestScopedErrorConfigFromContext retrieves the runtime config snapshot +// stored by WithRequestScopedErrorConfig. +func RequestScopedErrorConfigFromContext(ctx context.Context) *internalconfig.Config { + if ctx == nil { + return nil + } + cfg, _ := ctx.Value(requestScopedErrorConfigContextKey{}).(*internalconfig.Config) + return cfg +} + // Request-scoped error actions. const ( RequestScopedActionStop = "stop" @@ -57,6 +83,74 @@ func wrapRequestStopError(err error) error { return requestStopError{error: unwrapRequestStopError(err)} } +type schedulerRejectError struct { + error +} + +func (e schedulerRejectError) Unwrap() error { + return e.error +} + +func (e schedulerRejectError) IsSchedulerReject() bool { + return true +} + +func isSchedulerRejectError(err error) bool { + if err == nil { + return false + } + type rejectChecker interface { + IsSchedulerReject() bool + } + var rc rejectChecker + return errors.As(err, &rc) && rc != nil && rc.IsSchedulerReject() +} + +func unwrapSchedulerRejectError(err error) error { + var rejectErr schedulerRejectError + if errors.As(err, &rejectErr) { + return rejectErr.error + } + return err +} + +func wrapSchedulerRejectError(err error) error { + if err == nil { + return nil + } + return schedulerRejectError{error: unwrapSchedulerRejectError(err)} +} + +type roundEndedBySchedulerRejectError struct { + cause error +} + +func (e *roundEndedBySchedulerRejectError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *roundEndedBySchedulerRejectError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func hasRoundEndedBySchedulerReject(err error) bool { + var marked *roundEndedBySchedulerRejectError + return errors.As(err, &marked) && marked != nil +} + +func markRoundEndedBySchedulerReject(err error) error { + if err == nil || hasRoundEndedBySchedulerReject(err) { + return err + } + return &roundEndedBySchedulerRejectError{cause: err} +} + func (m *Manager) runtimeConfigSnapshot() *internalconfig.Config { if m == nil { return nil @@ -65,6 +159,92 @@ func (m *Manager) runtimeConfigSnapshot() *internalconfig.Config { return cfg } +func (m *Manager) requestScopedErrorConfig(ctx context.Context) *internalconfig.Config { + if cfg := RequestScopedErrorConfigFromContext(ctx); cfg != nil { + return cfg + } + return m.runtimeConfigSnapshot() +} + +func cloneRequestScopedErrorRules(rules []internalconfig.RequestScopedErrorRule) []internalconfig.RequestScopedErrorRule { + if len(rules) == 0 { + return nil + } + out := make([]internalconfig.RequestScopedErrorRule, len(rules)) + for i, r := range rules { + out[i] = internalconfig.RequestScopedErrorRule{ + Status: r.Status, + Action: r.Action, + } + if len(r.Match) > 0 { + out[i].Match = append([]string(nil), r.Match...) + } + if len(r.MatchRegexr) > 0 { + out[i].MatchRegexr = append([]string(nil), r.MatchRegexr...) + } + } + return out +} + +// ResolveRequestScopedErrorRules extracts and deep-copies the effective +// RequestScopedErrorRule list for an auth and config snapshot. +func ResolveRequestScopedErrorRules(auth *Auth, cfg *internalconfig.Config) []internalconfig.RequestScopedErrorRule { + return cloneRequestScopedErrorRules(extractRequestScopedErrorRules(auth, cfg)) +} + +// MatchRequestScopedErrorRules evaluates an error against a pre-resolved rule list. +func MatchRequestScopedErrorRules(rules []internalconfig.RequestScopedErrorRule, err error) (string, bool) { + if err == nil || len(rules) == 0 || cliproxyexecutor.IsUpstreamWebsocketTurnError(err) { + return "", false + } + + statusCode := statusCodeFromError(err) + body := extractErrorBody(err) + + for _, rule := range rules { + if rule.Status <= 0 || rule.Status != statusCode { + continue + } + if len(rule.Match) == 0 && len(rule.MatchRegexr) == 0 { + continue + } + + matched := false + for _, substr := range rule.Match { + if substr != "" && strings.Contains(body, substr) { + matched = true + break + } + } + if !matched { + for _, pattern := range rule.MatchRegexr { + if pattern != "" { + if re, errCompile := regexp.Compile(pattern); errCompile == nil && re.MatchString(body) { + matched = true + break + } + } + } + } + if !matched { + continue + } + + action := strings.ToLower(strings.TrimSpace(rule.Action)) + switch action { + case RequestScopedActionStop, + RequestScopedActionStopAndCooldown, + RequestScopedActionContinue, + RequestScopedActionContinueAndCooldown: + return action, true + default: + continue + } + } + + return "", false +} + // extractRequestScopedErrorRules retrieves the configured RequestScopedErrorRule list for an auth. func extractRequestScopedErrorRules(auth *Auth, cfg *internalconfig.Config) []internalconfig.RequestScopedErrorRule { if auth != nil && auth.Metadata != nil { @@ -181,63 +361,18 @@ func extractErrorBody(err error) string { return err.Error() } -// matchRequestScopedErrorAction evaluates an error against the auth's RequestScopedErrors rules. +// MatchRequestScopedErrorAction evaluates an error against the auth's RequestScopedErrors rules. // If a rule matches, it returns (action, true). // If no rule matches, it returns ("", false). +func MatchRequestScopedErrorAction(auth *Auth, err error, cfg *internalconfig.Config) (string, bool) { + return matchRequestScopedErrorAction(auth, err, cfg) +} + func matchRequestScopedErrorAction(auth *Auth, err error, cfg *internalconfig.Config) (string, bool) { if err == nil || cliproxyexecutor.IsUpstreamWebsocketTurnError(err) { return "", false } - rules := extractRequestScopedErrorRules(auth, cfg) - if len(rules) == 0 { - return "", false - } - - statusCode := statusCodeFromError(err) - body := extractErrorBody(err) - - for _, rule := range rules { - if rule.Status <= 0 || rule.Status != statusCode { - continue - } - if len(rule.Match) == 0 && len(rule.MatchRegexr) == 0 { - continue - } - - matched := false - for _, substr := range rule.Match { - if substr != "" && strings.Contains(body, substr) { - matched = true - break - } - } - if !matched { - for _, pattern := range rule.MatchRegexr { - if pattern != "" { - if re, errCompile := regexp.Compile(pattern); errCompile == nil && re.MatchString(body) { - matched = true - break - } - } - } - } - if !matched { - continue - } - - action := strings.ToLower(strings.TrimSpace(rule.Action)) - switch action { - case RequestScopedActionStop, - RequestScopedActionStopAndCooldown, - RequestScopedActionContinue, - RequestScopedActionContinueAndCooldown: - return action, true - default: - continue - } - } - - return "", false + return MatchRequestScopedErrorRules(extractRequestScopedErrorRules(auth, cfg), err) } func applyRequestScopedActionToResult(action string, okAction bool, result *Result) { diff --git a/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go b/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go index b587b83cb46..0504c9e4836 100644 --- a/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go +++ b/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go @@ -9,6 +9,7 @@ import ( internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" ) @@ -67,6 +68,9 @@ func TestUnwrapExecutionBoundaryErrorRemovesInternalMarkers(t *testing.T) { }{ {name: "attempt inside stop", err: wrapRequestStopError(markUpstreamExecutionAttempt(base))}, {name: "stop inside attempt", err: markUpstreamExecutionAttempt(wrapRequestStopError(base))}, + {name: "scheduler reject", err: wrapSchedulerRejectError(base)}, + {name: "attempt inside round scheduler reject", err: markRoundEndedBySchedulerReject(markUpstreamExecutionAttempt(base))}, + {name: "round scheduler reject inside attempt", err: markUpstreamExecutionAttempt(markRoundEndedBySchedulerReject(base))}, } { t.Run(test.name, func(t *testing.T) { if got := unwrapExecutionBoundaryError(test.err); got != base { @@ -1286,3 +1290,51 @@ func TestRequestScopedErrors_ResponseBodyProvider_MatchesUnderlyingPayload(t *te t.Fatal("expected auth1 to be in cooldown when matching ResponseBody()") } } + +func TestHomeExecutionAttachesRequestScopedErrorConfigToContext(t *testing.T) { + cfg := &internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + OAuthRequestScopedErrors: map[string][]internalconfig.RequestScopedErrorRule{ + "codex": { + { + Status: http.StatusNotFound, + Match: []string{"model_not_found_on_route"}, + Action: RequestScopedActionStop, + }, + }, + }, + } + m := NewManager(nil, nil, nil) + m.SetConfig(cfg) + + var observedAction string + var observedMatched bool + exec := &mockCustomErrorExecutor{ + identifier: "codex", + executeFn: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + err := customStatusError{code: http.StatusNotFound, msg: "model_not_found_on_route"} + observedAction, observedMatched = MatchRequestScopedErrorAction(auth, err, RequestScopedErrorConfigFromContext(ctx)) + return cliproxyexecutor.Response{}, err + }, + } + m.RegisterExecutor(exec) + + executions := executionregistry.New() + dispatcher := &accountedHomeExecutionDispatcher{ + auths: []Auth{{ + ID: "home-oauth-1", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{"auth_kind": "oauth"}, + }}, + } + m.PublishHomeDispatch(dispatcher, executions, 1) + + _, errExec := m.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-5.4"}, cliproxyexecutor.Options{}) + if errExec == nil { + t.Fatal("expected error, got nil") + } + if !observedMatched || observedAction != RequestScopedActionStop { + t.Fatalf("executor context did not match request-scoped rule: matched=%v action=%q", observedMatched, observedAction) + } +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index e379fb32960..0a6baddea1a 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -909,7 +909,10 @@ func (m *Manager) availableAuthsForRouteModelForRequest(ctx context.Context, aut if resetIn < 0 { resetIn = 0 } - return nil, newModelCooldownErrorWithCause(routeModel, providerForError, resetIn, lastCandidateErr) + claudeCooldown, claudeSignals := claudeCooldownMetadataForAuths(ctx, auths, func(candidate *Auth) string { + return m.selectionModelForAuth(candidate, routeModel) + }, now, clientHeaders...) + return nil, newModelCooldownErrorForCandidates(routeModel, providerForError, resetIn, lastCandidateErr, claudeCooldown, claudeSignals) } if unauthorizedCount == len(auths) && len(auths) > 0 { terminalCause := latestUnauthorizedCandidateError(auths) @@ -1024,7 +1027,14 @@ func restoreModelCooldownErrorModel(err error, requestedModel string) error { if !errors.As(err, &cooldownErr) || cooldownErr == nil || cooldownErr.model != "" { return err } - return newModelCooldownErrorWithCause(requestedModel, cooldownErr.provider, cooldownErr.resetIn, cooldownErr.cause) + return newModelCooldownErrorForCandidates( + requestedModel, + cooldownErr.provider, + cooldownErr.resetIn, + cooldownErr.cause, + cooldownErr.claudeCooldown, + cooldownErr.claudeSignals, + ) } func latestUnauthorizedCandidateError(auths []*Auth) error { @@ -1337,7 +1347,7 @@ func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginSc if rejectMessage == "" { rejectMessage = "scheduler rejected candidate selection" } - return nil, true, wrapRequestStopError(&Error{Code: rejectCode, Message: rejectMessage}) + return nil, true, wrapSchedulerRejectError(&Error{Code: rejectCode, Message: rejectMessage}) } if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil { keepSelectionID = true @@ -1779,7 +1789,7 @@ func (m *Manager) shouldRetryAfterErrorWithAttempted(ctx context.Context, opts c if status == http.StatusOK { return 0, false } - if isRequestInvalidError(err) || isRequestStopError(err) { + if isRequestInvalidError(err) || isRequestStopError(err) || isSchedulerRejectError(err) { return 0, false } if m.HomeEnabled() { diff --git a/sdk/cliproxy/auth/conductor_selection_cooldown_test.go b/sdk/cliproxy/auth/conductor_selection_cooldown_test.go index 95217a98704..a63dcda4538 100644 --- a/sdk/cliproxy/auth/conductor_selection_cooldown_test.go +++ b/sdk/cliproxy/auth/conductor_selection_cooldown_test.go @@ -3,10 +3,13 @@ package auth import ( "context" "errors" + "net/http" + "strconv" "strings" "testing" "time" + internallogging "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" ) @@ -813,3 +816,1197 @@ func TestErrorWithCause_ErrorDirectFormat(t *testing.T) { t.Fatalf("Error() = %q, want containing upstream error summary", got) } } + +func TestModelCooldownError_ClaudeUnifiedRateLimitHeaders(t *testing.T) { + t.Parallel() + + t.Run("preserves observed 7d claim under 5h remaining across affinity and scheduler paths", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + name := "scheduler" + if useAffinity { + name = "session_affinity" + } + t.Run(name, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + const model = "issue-7417-claude-7d-short-" + modelID := model + name + authID := "issue-7417-auth-" + name + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + retryDelay := 2 * time.Hour + upstreamReset := time.Now().Add(retryDelay - 15*time.Second).Unix() + markCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(markCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Reset": {strconv.FormatInt(upstreamReset, 10)}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Overage-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(upstreamReset, 10)}, + }) + manager.MarkResult(markCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelID, + Success: false, + CredentialScope: true, + RetryAfter: &retryDelay, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "rate_limit_error", HTTPStatus: http.StatusTooManyRequests}, + }) + + _, errExecute := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("expected cooldown error, got nil") + } + + var cooldownErr *modelCooldownError + if !errors.As(errExecute, &cooldownErr) { + t.Fatalf("Execute() error = %T (%v), want *modelCooldownError", errExecute, errExecute) + } + + before := time.Now().Unix() + headers := cooldownErr.Headers() + after := time.Now().Unix() + retrySeconds, errSeconds := strconv.ParseInt(headers.Get("Retry-After"), 10, 64) + if errSeconds != nil || retrySeconds <= 0 { + t.Fatalf("Retry-After = %q, want positive seconds", headers.Get("Retry-After")) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-Status = %q, want rejected", got) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("Anthropic-Ratelimit-Unified-Representative-Claim = %q, want seven_day", got) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-7d-Status = %q, want rejected", got) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "allowed" { + t.Fatalf("Anthropic-Ratelimit-Unified-5h-Status = %q, want allowed", got) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-Overage-Status"); got != "rejected" { + t.Fatalf("Anthropic-Ratelimit-Unified-Overage-Status = %q, want rejected", got) + } + + unifiedReset, errReset := strconv.ParseInt(headers.Get("Anthropic-Ratelimit-Unified-Reset"), 10, 64) + if errReset != nil { + t.Fatalf("Anthropic-Ratelimit-Unified-Reset = %q, want unix epoch", headers.Get("Anthropic-Ratelimit-Unified-Reset")) + } + if unifiedReset < before+retrySeconds || unifiedReset > after+retrySeconds { + t.Fatalf("Anthropic-Ratelimit-Unified-Reset = %d, want in [%d, %d] aligned with Retry-After %d", unifiedReset, before+retrySeconds, after+retrySeconds, retrySeconds) + } + if got := headers.Get("Anthropic-Ratelimit-Unified-7d-Reset"); got != strconv.FormatInt(unifiedReset, 10) { + t.Fatalf("Anthropic-Ratelimit-Unified-7d-Reset = %q, want %d", got, unifiedReset) + } + + safeHeaders := SafeResponseHeaders(errExecute) + if got := safeHeaders.Get("Anthropic-Ratelimit-Unified-Status"); got != "" { + t.Fatalf("SafeResponseHeaders leaked unified header %q", got) + } + }) + } + }) + + t.Run("infers claim from duration when signals are absent", func(t *testing.T) { + errFiveHour := newModelCooldownError("claude-haiku-4-5-20251001", "claude", 2*time.Hour) + h5 := errFiveHour.Headers() + if got := h5.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "five_hour" { + t.Fatalf("5h fallback claim = %q, want five_hour", got) + } + if got := h5.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "rejected" { + t.Fatalf("5h fallback status = %q, want rejected", got) + } + if got := h5.Get("Anthropic-Ratelimit-Unified-Overage-Status"); got != "rejected" { + t.Fatalf("5h fallback overage status = %q, want rejected", got) + } + + errSevenDay := newModelCooldownError("claude-haiku-4-5-20251001", "claude", 6*time.Hour) + h7 := errSevenDay.Headers() + if got := h7.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("7d fallback claim = %q, want seven_day", got) + } + if got := h7.Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "rejected" { + t.Fatalf("7d fallback status = %q, want rejected", got) + } + }) + + t.Run("preserves allowed shared windows on overage claim and aligns model-family window resets", func(t *testing.T) { + overageErr := newModelCooldownErrorForCandidates("claude-sonnet-4-5-20250929", "claude", 30*time.Second, nil, true, map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-Representative-Claim": "overage", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d-Status": "allowed", + "Anthropic-Ratelimit-Unified-Overage-Status": "rejected", + }) + overageHeaders := overageErr.Headers() + if got := overageHeaders.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "allowed" { + t.Fatalf("overage cooldown overwrote 5h status = %q, want allowed", got) + } + if got := overageHeaders.Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "allowed" { + t.Fatalf("overage cooldown overwrote 7d status = %q, want allowed", got) + } + if got := overageHeaders.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "overage" { + t.Fatalf("overage cooldown claim = %q, want overage", got) + } + + sonnetErr := newModelCooldownErrorForCandidates("claude-sonnet-4-5-20250929", "claude", 2*time.Hour, nil, true, map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-Representative-Claim": "seven_day_sonnet", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d_sonnet-Status": "rejected", + "Anthropic-Ratelimit-Unified-7d_sonnet-Reset": "1700000000", + "Anthropic-Ratelimit-Unified-Overage-Status": "rejected", + }) + sonnetHeaders := sonnetErr.Headers() + if got := sonnetHeaders.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "allowed" { + t.Fatalf("7d_sonnet cooldown overwrote 5h status = %q, want allowed", got) + } + if got := sonnetHeaders.Get("Anthropic-Ratelimit-Unified-7d_sonnet-Reset"); got != sonnetHeaders.Get("Anthropic-Ratelimit-Unified-Reset") { + t.Fatalf("7d_sonnet reset = %q, want %q", got, sonnetHeaders.Get("Anthropic-Ratelimit-Unified-Reset")) + } + + overageNoClaimErr := newModelCooldownErrorForCandidates("claude-sonnet-4-5-20250929", "claude", 30*time.Second, nil, true, map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d-Status": "allowed", + "Anthropic-Ratelimit-Unified-Overage-Status": "rejected", + }) + overageNoClaimHeaders := overageNoClaimErr.Headers() + if got := overageNoClaimHeaders.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "overage" { + t.Fatalf("overage without claim inferred %q, want overage", got) + } + if got := overageNoClaimHeaders.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "allowed" { + t.Fatalf("overage without claim 5h status = %q, want allowed", got) + } + + sonnetNoClaimErr := newModelCooldownErrorForCandidates("claude-sonnet-4-5-20250929", "claude", 2*time.Hour, nil, true, map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d_sonnet-Status": "rejected", + "Anthropic-Ratelimit-Unified-Overage-Status": "rejected", + }) + sonnetNoClaimHeaders := sonnetNoClaimErr.Headers() + if got := sonnetNoClaimHeaders.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day_sonnet" { + t.Fatalf("7d_sonnet without claim inferred %q, want seven_day_sonnet", got) + } + + fiveAllowedOnlyErr := newModelCooldownErrorForCandidates("claude-sonnet-4-5-20250929", "claude", 2*time.Hour, nil, true, map[string]string{ + "Anthropic-Ratelimit-Unified-Status": "rejected", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + }) + fiveAllowedOnlyHeaders := fiveAllowedOnlyErr.Headers() + if got := fiveAllowedOnlyHeaders.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("5h=allowed with 2h reset inferred %q, want seven_day", got) + } + if got := fiveAllowedOnlyHeaders.Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "rejected" { + t.Fatalf("5h=allowed with 2h reset 7d status = %q, want rejected", got) + } + }) + + t.Run("does not emit unified headers for non-Claude or cloudflare cooldowns", func(t *testing.T) { + manager := NewManager(nil, NewSessionAffinitySelector(&RoundRobinSelector{}), nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "codex"}) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + const codexModel = "issue-7417-codex-model" + const codexAuthID = "issue-7417-codex-auth" + const cfModel = "issue-7417-cf-model" + const cfAuthID = "issue-7417-cf-auth" + next := time.Now().Add(30 * time.Minute) + + registry.GetGlobalRegistry().RegisterClient(codexAuthID, "codex", []*registry.ModelInfo{{ID: codexModel}}) + registry.GetGlobalRegistry().RegisterClient(cfAuthID, "claude", []*registry.ModelInfo{{ID: cfModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(codexAuthID) + registry.GetGlobalRegistry().UnregisterClient(cfAuthID) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: codexAuthID, + Provider: "codex", + Status: StatusActive, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next}, + }); errRegister != nil { + t.Fatalf("register codex auth: %v", errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: cfAuthID, + Provider: "claude", + Status: StatusActive, + Unavailable: true, + NextRetryAfter: next, + ModelStates: map[string]*ModelState{ + cfModel: { + Status: StatusError, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{Exceeded: true, Reason: "cloudflare challenge", NextRecoverAt: next}, + }, + }, + }); errRegister != nil { + t.Fatalf("register cloudflare auth: %v", errRegister) + } + + _, errCodex := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: codexModel}, cliproxyexecutor.Options{}) + var codexCooldown *modelCooldownError + if !errors.As(errCodex, &codexCooldown) { + t.Fatalf("codex Execute() error = %T (%v), want *modelCooldownError", errCodex, errCodex) + } + if codexCooldown.provider != "" { + t.Fatalf("codex cooldown provider = %q, want empty mixed provider", codexCooldown.provider) + } + if got := codexCooldown.Headers().Get("Anthropic-Ratelimit-Unified-Status"); got != "" { + t.Fatalf("codex cooldown leaked Anthropic-Ratelimit-Unified-Status = %q", got) + } + + _, errCF := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: cfModel}, cliproxyexecutor.Options{}) + var cfCooldown *modelCooldownError + if !errors.As(errCF, &cfCooldown) { + t.Fatalf("cloudflare Execute() error = %T (%v), want *modelCooldownError", errCF, errCF) + } + if got := cfCooldown.Headers().Get("Anthropic-Ratelimit-Unified-Status"); got != "" { + t.Fatalf("cloudflare challenge cooldown leaked Anthropic-Ratelimit-Unified-Status = %q", got) + } + }) + + t.Run("does not synthesize unified headers for ordinary per-minute rate limits across affinity and scheduler paths", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + allowedWindowModel := "issue-7417-rpm-allowed-" + pathName + allowedWindowAuth := "issue-7417-rpm-allowed-auth-" + pathName + plainRPMModel := "issue-7417-rpm-plain-" + pathName + plainRPMAuth := "issue-7417-rpm-plain-auth-" + pathName + modelQuotaModel := "issue-7417-model-unified-" + pathName + modelQuotaAuth := "issue-7417-model-unified-auth-" + pathName + + registry.GetGlobalRegistry().RegisterClient(allowedWindowAuth, "claude", []*registry.ModelInfo{{ID: allowedWindowModel}}) + registry.GetGlobalRegistry().RegisterClient(plainRPMAuth, "claude", []*registry.ModelInfo{{ID: plainRPMModel}}) + registry.GetGlobalRegistry().RegisterClient(modelQuotaAuth, "claude", []*registry.ModelInfo{{ID: modelQuotaModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(allowedWindowAuth) + registry.GetGlobalRegistry().UnregisterClient(plainRPMAuth) + registry.GetGlobalRegistry().UnregisterClient(modelQuotaAuth) + }) + + for _, id := range []string{allowedWindowAuth, plainRPMAuth, modelQuotaAuth} { + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: id, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register %s: %v", id, errRegister) + } + } + + shortRetry := 30 * time.Second + rpmBody := `{"type":"error","error":{"type":"rate_limit_error","message":"This request would exceed the requests per minute rate limit. Please try again later."}}` + + // 1. Per-minute 429 with allowed 5h window and future unified reset. + allowedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(allowedCtx, http.Header{ + "Retry-After": {"30"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(4*time.Hour).Unix(), 10)}, + }) + manager.MarkResult(allowedCtx, Result{ + AuthID: allowedWindowAuth, + Provider: "claude", + Model: allowedWindowModel, + Success: false, + CredentialScope: false, + RetryAfter: &shortRetry, + Error: &Error{Code: "rate_limit_error", Message: rpmBody, HTTPStatus: http.StatusTooManyRequests}, + }) + + // 2. Per-minute 429 with only Retry-After and no unified rejection headers. + plainCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(plainCtx, http.Header{ + "Retry-After": {"30"}, + }) + manager.MarkResult(plainCtx, Result{ + AuthID: plainRPMAuth, + Provider: "claude", + Model: plainRPMModel, + Success: false, + CredentialScope: false, + RetryAfter: &shortRetry, + Error: &Error{Code: "rate_limit_error", Message: rpmBody, HTTPStatus: http.StatusTooManyRequests}, + }) + + // 3. Model-level cooling (e.g. API-key upstream) that explicitly reports unified 5h rejection. + fiveHourDelay := 2 * time.Hour + rejectedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rejectedCtx, http.Header{ + "Retry-After": {"7200"}, + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(fiveHourDelay).Unix(), 10)}, + }) + manager.MarkResult(rejectedCtx, Result{ + AuthID: modelQuotaAuth, + Provider: "claude", + Model: modelQuotaModel, + Success: false, + CredentialScope: false, + RetryAfter: &fiveHourDelay, + Error: &Error{Code: "rate_limit_error", Message: "quota exhausted", HTTPStatus: http.StatusTooManyRequests}, + }) + + for _, modelID := range []string{allowedWindowModel, plainRPMModel} { + _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + var cdErr *modelCooldownError + if !errors.As(errExec, &cdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", modelID, errExec, errExec) + } + hdr := cdErr.Headers() + if got := hdr.Get("Retry-After"); got == "" { + t.Fatalf("%s Retry-After = empty, want positive seconds", modelID) + } + for key := range hdr { + if strings.HasPrefix(strings.ToLower(key), "anthropic-ratelimit-unified-") { + t.Fatalf("%s ordinary rate-limit cooldown leaked unified header %s=%q", modelID, key, hdr.Get(key)) + } + } + } + + _, errModelQuota := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelQuotaModel}, cliproxyexecutor.Options{}) + var modelQuotaCooldown *modelCooldownError + if !errors.As(errModelQuota, &modelQuotaCooldown) { + t.Fatalf("modelQuota Execute() error = %T (%v), want *modelCooldownError", errModelQuota, errModelQuota) + } + if got := modelQuotaCooldown.Headers().Get("Anthropic-Ratelimit-Unified-Status"); got != "rejected" { + t.Fatalf("modelQuota Anthropic-Ratelimit-Unified-Status = %q, want rejected", got) + } + if got := modelQuotaCooldown.Headers().Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "five_hour" { + t.Fatalf("modelQuota Representative-Claim = %q, want five_hour", got) + } + }) + } + }) + + t.Run("retains unified headers when late success refreshes passive signals during active cooldown", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + credModel := "issue-7417-late-cred-" + pathName + credAuth := "issue-7417-late-cred-auth-" + pathName + restoredModel := "issue-7417-late-restored-" + pathName + restoredAuth := "issue-7417-late-restored-auth-" + pathName + modelScopedModel := "issue-7417-late-model-" + pathName + modelScopedAuth := "issue-7417-late-model-auth-" + pathName + + registry.GetGlobalRegistry().RegisterClient(credAuth, "claude", []*registry.ModelInfo{{ID: credModel}}) + registry.GetGlobalRegistry().RegisterClient(restoredAuth, "claude", []*registry.ModelInfo{{ID: restoredModel}}) + registry.GetGlobalRegistry().RegisterClient(modelScopedAuth, "claude", []*registry.ModelInfo{{ID: modelScopedModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(credAuth) + registry.GetGlobalRegistry().UnregisterClient(restoredAuth) + registry.GetGlobalRegistry().UnregisterClient(modelScopedAuth) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: credAuth, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register credAuth: %v", errRegister) + } + recoverAt := time.Now().Add(2 * time.Hour) + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: restoredAuth, + Provider: "claude", + Status: StatusError, + Unavailable: true, + NextRetryAfter: recoverAt, + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: recoverAt, + }, + }); errRegister != nil { + t.Fatalf("register restoredAuth: %v", errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: modelScopedAuth, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register modelScopedAuth: %v", errRegister) + } + + delay := 2 * time.Hour + rejectedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rejectedCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(delay).Unix(), 10)}, + }) + + // Step 1: Record subscription 429 on credAuth and modelScopedAuth. + staleSeq := manager.nextExecutionSeq() + failureSeq := manager.nextExecutionSeq() + manager.MarkResult(rejectedCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModel, + Success: false, + CredentialScope: true, + RetryAfter: &delay, + ExecutionSeq: failureSeq, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "rate_limit_error", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(rejectedCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: false, + CredentialScope: false, + RetryAfter: &delay, + ExecutionSeq: failureSeq, + Error: &Error{Code: "rate_limit_error", Message: "rate_limit_error", HTTPStatus: http.StatusTooManyRequests}, + }) + + // Step 2: Earlier in-flight requests complete with 200 OK and allowed unified headers. + allowedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(allowedCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"allowed"}, + }) + manager.MarkResult(allowedCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModel, + Success: true, + ExecutionSeq: staleSeq, + }) + manager.MarkResult(allowedCtx, Result{ + AuthID: restoredAuth, + Provider: "claude", + Model: restoredModel, + Success: true, + ExecutionSeq: staleSeq, + }) + manager.MarkResult(allowedCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: true, + ExecutionSeq: staleSeq, + }) + + // Step 3: Subsequent requests during active cooldown still emit unified rate-limit headers. + for _, tc := range []struct { + model string + wantClaim string + }{ + {model: credModel, wantClaim: "seven_day"}, + {model: restoredModel, wantClaim: "five_hour"}, + {model: modelScopedModel, wantClaim: "seven_day"}, + } { + _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: tc.model}, cliproxyexecutor.Options{}) + var cdErr *modelCooldownError + if !errors.As(errExec, &cdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", tc.model, errExec, errExec) + } + hdr := cdErr.Headers() + if got := hdr.Get("Anthropic-Ratelimit-Unified-Status"); got != "rejected" { + t.Fatalf("%s Anthropic-Ratelimit-Unified-Status = %q, want rejected", tc.model, got) + } + if got := hdr.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != tc.wantClaim { + t.Fatalf("%s Representative-Claim = %q, want %s", tc.model, got, tc.wantClaim) + } + } + }) + } + }) + + t.Run("does not emit stale unified headers when 503 follows expired model quota cooldown", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + modelID := "issue-7417-expired-503-" + pathName + authID := "issue-7417-expired-503-auth-" + pathName + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // 1. Model-scoped unified 429 records rejected signals. + quotaDelay := time.Hour + rejectedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rejectedCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"five_hour"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(quotaDelay).Unix(), 10)}, + }) + manager.MarkResult(rejectedCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelID, + Success: false, + CredentialScope: false, + RetryAfter: "aDelay, + Error: &Error{Code: "rate_limit_error", Message: "quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // 2. Simulate natural expiry of the quota deadline without wall-clock sleep. + past := time.Now().Add(-time.Minute) + if _, okModify, errModify := manager.Modify(context.Background(), authID, 0, func(a *Auth) error { + a.NextRetryAfter = past + if state := a.ModelStates[modelID]; state != nil { + state.NextRetryAfter = past + state.Quota.NextRecoverAt = past + } + a.Quota.NextRecoverAt = past + return nil + }); errModify != nil || !okModify { + t.Fatalf("Modify() ok=%v err=%v", okModify, errModify) + } + + // 3. Subsequent request fails with 503 Service Unavailable and Retry-After. + serverDelay := 45 * time.Second + manager.MarkResult(context.Background(), Result{ + AuthID: authID, + Provider: "claude", + Model: modelID, + Success: false, + RetryAfter: &serverDelay, + Error: &Error{Code: "server_error", Message: "overloaded", HTTPStatus: http.StatusServiceUnavailable}, + }) + + // 4. Subsequent selection error must not advertise unified subscription exhaustion. + _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + if errExec == nil { + t.Fatal("expected unavailable/cooldown error after 503, got nil") + } + if carrier, ok := errExec.(interface{ Headers() http.Header }); ok && carrier != nil { + for key := range carrier.Headers() { + if strings.HasPrefix(strings.ToLower(key), "anthropic-ratelimit-unified-") { + t.Fatalf("%s post-expiry 503 leaked unified header %s=%q", pathName, key, carrier.Headers().Get(key)) + } + } + } + }) + } + }) + + t.Run("persists and restores model-scoped unified cooldown headers across restart", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + store := NewFileCooldownStateStore(t.TempDir()) + modelID := "issue-7417-restart-model-" + pathName + authID := "issue-7417-restart-auth-" + pathName + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + writerMgr := NewManager(nil, nil, nil) + writerMgr.SetCooldownStateStore(store) + if _, errRegister := writerMgr.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("writer Register: %v", errRegister) + } + + delay := 2 * time.Hour + rejectedCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rejectedCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(delay).Unix(), 10)}, + }) + writerMgr.MarkResult(rejectedCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelID, + Success: false, + CredentialScope: false, + RetryAfter: &delay, + Error: &Error{Code: "rate_limit_error", Message: "quota", HTTPStatus: http.StatusTooManyRequests}, + }) + writerMgr.PersistCooldownStates(context.Background()) + + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + restoredMgr := NewManager(nil, selector, nil) + restoredMgr.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + restoredMgr.SetCooldownStateStore(store) + if _, errRegister := restoredMgr.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("restored Register: %v", errRegister) + } + if errRestore := restoredMgr.RestoreCooldownStates(context.Background()); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + _, errExec := restoredMgr.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + var cdErr *modelCooldownError + if !errors.As(errExec, &cdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", pathName, errExec, errExec) + } + hdr := cdErr.Headers() + if got := hdr.Get("Anthropic-Ratelimit-Unified-Status"); got != "rejected" { + t.Fatalf("%s restored Anthropic-Ratelimit-Unified-Status = %q, want rejected", pathName, got) + } + if got := hdr.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("%s restored Representative-Claim = %q, want seven_day", pathName, got) + } + if got := hdr.Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "rejected" { + t.Fatalf("%s restored 7d-Status = %q, want rejected", pathName, got) + } + if got := hdr.Get("Anthropic-Ratelimit-Unified-5h-Status"); got != "allowed" { + t.Fatalf("%s restored 5h-Status = %q, want allowed", pathName, got) + } + }) + } + }) + + t.Run("keeps longer cooldown window signals when shorter out-of-order 429 completes later", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + credModelA := "issue-7417-ooo-cred-a-" + pathName + credModelB := "issue-7417-ooo-cred-b-" + pathName + credAuth := "issue-7417-ooo-cred-auth-" + pathName + modelScopedModel := "issue-7417-ooo-model-" + pathName + modelScopedAuth := "issue-7417-ooo-model-auth-" + pathName + rpmFirstModel := "issue-7417-ooo-rpm-" + pathName + rpmFirstAuth := "issue-7417-ooo-rpm-auth-" + pathName + + registry.GetGlobalRegistry().RegisterClient(credAuth, "claude", []*registry.ModelInfo{{ID: credModelA}, {ID: credModelB}}) + registry.GetGlobalRegistry().RegisterClient(modelScopedAuth, "claude", []*registry.ModelInfo{{ID: modelScopedModel}}) + registry.GetGlobalRegistry().RegisterClient(rpmFirstAuth, "claude", []*registry.ModelInfo{{ID: rpmFirstModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(credAuth) + registry.GetGlobalRegistry().UnregisterClient(modelScopedAuth) + registry.GetGlobalRegistry().UnregisterClient(rpmFirstAuth) + }) + + for _, id := range []string{credAuth, modelScopedAuth, rpmFirstAuth} { + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: id, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register %s: %v", id, errRegister) + } + } + + longDelay := 2 * time.Hour + shortDelay := time.Hour + sevenDayCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(sevenDayCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(longDelay).Unix(), 10)}, + }) + fiveHourCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(fiveHourCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"five_hour"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(shortDelay).Unix(), 10)}, + }) + + // Credential-scoped: 2h seven_day on model A, then shorter 1h five_hour on model B. + manager.MarkResult(sevenDayCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModelA, + Success: false, + CredentialScope: true, + RetryAfter: &longDelay, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "7d quota", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModelB, + Success: false, + CredentialScope: true, + RetryAfter: &shortDelay, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "5h quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // Model-scoped: 2h seven_day, then shorter 1h five_hour. + manager.MarkResult(sevenDayCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: false, + CredentialScope: false, + RetryAfter: &longDelay, + Error: &Error{Code: "rate_limit_error", Message: "7d quota", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: false, + CredentialScope: false, + RetryAfter: &shortDelay, + Error: &Error{Code: "rate_limit_error", Message: "5h quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // Longer ordinary 10m RPM 429 followed by shorter 1m unified 429: governing deadline stays ordinary RPM. + rpmLong := 10 * time.Minute + unifiedShort := time.Minute + rpmCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rpmCtx, http.Header{"Retry-After": {"600"}}) + manager.MarkResult(rpmCtx, Result{ + AuthID: rpmFirstAuth, + Provider: "claude", + Model: rpmFirstModel, + Success: false, + CredentialScope: false, + RetryAfter: &rpmLong, + Error: &Error{Code: "rate_limit_error", Message: "rpm", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourCtx, Result{ + AuthID: rpmFirstAuth, + Provider: "claude", + Model: rpmFirstModel, + Success: false, + CredentialScope: false, + RetryAfter: &unifiedShort, + Error: &Error{Code: "rate_limit_error", Message: "5h", HTTPStatus: http.StatusTooManyRequests}, + }) + + for _, m := range []string{credModelA, credModelB, modelScopedModel} { + _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: m}, cliproxyexecutor.Options{}) + var cdErr *modelCooldownError + if !errors.As(errExec, &cdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", m, errExec, errExec) + } + if got := cdErr.Headers().Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("%s Representative-Claim = %q, want seven_day (from governing 2h deadline)", m, got) + } + if got := cdErr.Headers().Get("Anthropic-Ratelimit-Unified-7d-Status"); got != "rejected" { + t.Fatalf("%s 7d-Status = %q, want rejected", m, got) + } + } + + _, errRPM := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: rpmFirstModel}, cliproxyexecutor.Options{}) + var rpmCdErr *modelCooldownError + if !errors.As(errRPM, &rpmCdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", rpmFirstModel, errRPM, errRPM) + } + if got := rpmCdErr.Headers().Get("Anthropic-Ratelimit-Unified-Status"); got != "" { + t.Fatalf("%s leaked unified status %q while 10m ordinary RPM deadline still governs", rpmFirstModel, got) + } + }) + } + }) + + t.Run("reverts to independent model layer after propagated credential quota expires", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + modelA := "issue-7417-shadow-a-" + pathName + modelB := "issue-7417-shadow-b-" + pathName + authID := "issue-7417-shadow-auth-" + pathName + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelA}, {ID: modelB}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // 1. Model A enters a 10m ordinary RPM cooldown (no unified rejection headers). + rpmDelay := 10 * time.Minute + rpmCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(rpmCtx, http.Header{"Retry-After": {"600"}}) + manager.MarkResult(rpmCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelA, + Success: false, + CredentialScope: false, + RetryAfter: &rpmDelay, + Error: &Error{Code: "rate_limit_error", Message: "rpm", HTTPStatus: http.StatusTooManyRequests}, + }) + + // 2. Model B triggers a shorter 1m credential-wide 5h quota cooldown. + credDelay := time.Minute + fiveHourCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(fiveHourCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"five_hour"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(credDelay).Unix(), 10)}, + }) + manager.MarkResult(fiveHourCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelB, + Success: false, + CredentialScope: true, + RetryAfter: &credDelay, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "5h quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // 3. Expire the 1m credential-wide quota deterministically while Model A's 10m RPM block remains live. + past := time.Now().Add(-time.Second) + manager.mu.Lock() + curAuth := manager.auths[authID] + curAuth.Quota.NextRecoverAt = past + curAuth.NextRetryAfter = past + if sb := curAuth.ModelStates[modelB]; sb != nil { + sb.NextRetryAfter = past + sb.Quota.NextRecoverAt = past + } + snapshot := curAuth.Clone() + manager.mu.Unlock() + if manager.scheduler != nil { + manager.scheduler.upsertAuth(snapshot) + } + + // 4. Model A is still cooling down from its independent 10m RPM error and must not emit 5h unified headers. + _, errExecA := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelA}, cliproxyexecutor.Options{}) + var cdErrA *modelCooldownError + if !errors.As(errExecA, &cdErrA) { + t.Fatalf("%s modelA Execute() error = %T (%v), want *modelCooldownError", pathName, errExecA, errExecA) + } + for key := range cdErrA.Headers() { + if strings.HasPrefix(strings.ToLower(key), "anthropic-ratelimit-unified-") { + t.Fatalf("%s modelA leaked expired credential-quota header %s=%q", pathName, key, cdErrA.Headers().Get(key)) + } + } + }) + } + }) + + t.Run("matches longer independent model quota signals and deadline beneath shorter credential quota", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + modelA := "issue-7417-dual-quota-a-" + pathName + modelB := "issue-7417-dual-quota-b-" + pathName + authID := "issue-7417-dual-quota-auth-" + pathName + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelA}, {ID: modelB}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: authID, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + // 1. Model A enters a 2h independent model-scoped seven_day quota cooldown. + twoHours := 2 * time.Hour + sevenDayCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(sevenDayCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(twoHours).Unix(), 10)}, + }) + manager.MarkResult(sevenDayCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelA, + Success: false, + CredentialScope: false, + RetryAfter: &twoHours, + Error: &Error{Code: "rate_limit_error", Message: "7d model quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // 2. Model B triggers a shorter 1h credential-scoped five_hour quota cooldown. + oneHour := time.Hour + fiveHourCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(fiveHourCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"five_hour"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(oneHour).Unix(), 10)}, + }) + manager.MarkResult(fiveHourCtx, Result{ + AuthID: authID, + Provider: "claude", + Model: modelB, + Success: false, + CredentialScope: true, + RetryAfter: &oneHour, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "5h cred quota", HTTPStatus: http.StatusTooManyRequests}, + }) + + // Model A must report seven_day AND ~2h Retry-After / reset (not 1h). + _, errExecA := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelA}, cliproxyexecutor.Options{}) + var cdErrA *modelCooldownError + if !errors.As(errExecA, &cdErrA) { + t.Fatalf("%s modelA Execute() error = %T (%v), want *modelCooldownError", pathName, errExecA, errExecA) + } + hdrA := cdErrA.Headers() + if got := hdrA.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("%s modelA claim = %q, want seven_day", pathName, got) + } + retrySecA, errSecA := strconv.ParseInt(hdrA.Get("Retry-After"), 10, 64) + if errSecA != nil || retrySecA < 7100 { + t.Fatalf("%s modelA Retry-After = %q, want ~7200s (>= 7100s) matching 2h independent deadline", pathName, hdrA.Get("Retry-After")) + } + + // Model B must report five_hour AND ~1h Retry-After. + _, errExecB := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelB}, cliproxyexecutor.Options{}) + var cdErrB *modelCooldownError + if !errors.As(errExecB, &cdErrB) { + t.Fatalf("%s modelB Execute() error = %T (%v), want *modelCooldownError", pathName, errExecB, errExecB) + } + hdrB := cdErrB.Headers() + if got := hdrB.Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "five_hour" { + t.Fatalf("%s modelB claim = %q, want five_hour", pathName, got) + } + retrySecB, errSecB := strconv.ParseInt(hdrB.Get("Retry-After"), 10, 64) + if errSecB != nil || retrySecB > 3660 || retrySecB < 3500 { + t.Fatalf("%s modelB Retry-After = %q, want ~3600s", pathName, hdrB.Get("Retry-After")) + } + }) + } + }) + + t.Run("keeps active cooldown signals when hintless RetryAfter nil 429 reuses existing deadline", func(t *testing.T) { + for _, useAffinity := range []bool{true, false} { + pathName := "scheduler" + if useAffinity { + pathName = "session_affinity" + } + t.Run(pathName, func(t *testing.T) { + var selector Selector + if useAffinity { + selector = NewSessionAffinitySelector(&RoundRobinSelector{}) + } + manager := NewManager(nil, selector, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: "claude"}) + + credModel := "issue-7417-nilhint-cred-" + pathName + credAuth := "issue-7417-nilhint-cred-auth-" + pathName + modelScopedModel := "issue-7417-nilhint-model-" + pathName + modelScopedAuth := "issue-7417-nilhint-model-auth-" + pathName + authOnlyModel := "issue-7417-nilhint-authonly-" + pathName + authOnlyID := "issue-7417-nilhint-authonly-auth-" + pathName + + registry.GetGlobalRegistry().RegisterClient(credAuth, "claude", []*registry.ModelInfo{{ID: credModel}}) + registry.GetGlobalRegistry().RegisterClient(modelScopedAuth, "claude", []*registry.ModelInfo{{ID: modelScopedModel}}) + registry.GetGlobalRegistry().RegisterClient(authOnlyID, "claude", []*registry.ModelInfo{{ID: authOnlyModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(credAuth) + registry.GetGlobalRegistry().UnregisterClient(modelScopedAuth) + registry.GetGlobalRegistry().UnregisterClient(authOnlyID) + }) + + for _, id := range []string{credAuth, modelScopedAuth, authOnlyID} { + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: id, + Provider: "claude", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("register %s: %v", id, errRegister) + } + } + + twoHours := 2 * time.Hour + sevenDayCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(sevenDayCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"seven_day"}, + "Anthropic-Ratelimit-Unified-Reset": {strconv.FormatInt(time.Now().Add(twoHours).Unix(), 10)}, + }) + fiveHourNoHintCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(fiveHourNoHintCtx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-5h-Status": {"rejected"}, + "Anthropic-Ratelimit-Unified-7d-Status": {"allowed"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": {"five_hour"}, + }) + + manager.MarkResult(sevenDayCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModel, + Success: false, + CredentialScope: true, + RetryAfter: &twoHours, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "7d", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourNoHintCtx, Result{ + AuthID: credAuth, + Provider: "claude", + Model: credModel, + Success: false, + CredentialScope: true, + RetryAfter: nil, + Error: &Error{Code: ErrorCodeCredentialQuota, Message: "5h no hint", HTTPStatus: http.StatusTooManyRequests}, + }) + + manager.MarkResult(sevenDayCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: false, + CredentialScope: false, + RetryAfter: &twoHours, + Error: &Error{Code: "rate_limit_error", Message: "7d", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourNoHintCtx, Result{ + AuthID: modelScopedAuth, + Provider: "claude", + Model: modelScopedModel, + Success: false, + CredentialScope: false, + RetryAfter: nil, + Error: &Error{Code: "rate_limit_error", Message: "5h no hint", HTTPStatus: http.StatusTooManyRequests}, + }) + + manager.MarkResult(sevenDayCtx, Result{ + AuthID: authOnlyID, + Provider: "claude", + Success: false, + RetryAfter: &twoHours, + Error: &Error{Code: "rate_limit_error", Message: "7d", HTTPStatus: http.StatusTooManyRequests}, + }) + manager.MarkResult(fiveHourNoHintCtx, Result{ + AuthID: authOnlyID, + Provider: "claude", + Success: false, + RetryAfter: nil, + Error: &Error{Code: "rate_limit_error", Message: "5h no hint", HTTPStatus: http.StatusTooManyRequests}, + }) + + for _, m := range []string{credModel, modelScopedModel, authOnlyModel} { + _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: m}, cliproxyexecutor.Options{}) + var cdErr *modelCooldownError + if !errors.As(errExec, &cdErr) { + t.Fatalf("%s Execute() error = %T (%v), want *modelCooldownError", m, errExec, errExec) + } + if got := cdErr.Headers().Get("Anthropic-Ratelimit-Unified-Representative-Claim"); got != "seven_day" { + t.Fatalf("%s Representative-Claim = %q, want seven_day", m, got) + } + } + }) + } + }) +} diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 47f87de43f9..3c37aced1cc 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -137,7 +137,7 @@ func (m *Manager) recordCanceledUpstreamCooldown(ctx context.Context, auth *Auth if m == nil || auth == nil || retryAfter == nil { return } - action, okAction := matchRequestScopedErrorAction(auth, err, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, err, m.requestScopedErrorConfig(ctx)) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(err), Options: opts, RetryAfter: retryAfter, AccessToken: authAccessTokenForRequest(ctx, auth, opts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.CredentialScope = isCredentialScopedError(err) applyRequestScopedActionToResult(action, okAction, &result) @@ -165,7 +165,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re entry := logEntryWithRequestID(ctx) warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), chunk.Err) rerr := resultErrorFromError(chunk.Err) - action, okAction := matchRequestScopedErrorAction(auth, chunk.Err, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, chunk.Err, m.requestScopedErrorConfig(ctx)) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: opts, AccessToken: accessToken, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(chunk.Err) result.CredentialScope = isCredentialScopedError(chunk.Err) @@ -412,7 +412,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi var upstreamErr error didRefreshOnUnauthorized := false for idx, execModel := range execModels { - ctx = newUpstreamAttemptContext(ctx) + ctx = m.newUpstreamAttemptContext(ctx) resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel @@ -456,7 +456,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi auth = refreshed publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true - ctx = newUpstreamAttemptContext(ctx) + ctx = m.newUpstreamAttemptContext(ctx) ctx = syncMetadataSessionToContext(ctx, execOpts.Metadata) executionSeq = m.nextExecutionSeq() startRetry := time.Now() @@ -489,7 +489,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) if errStream != nil { rerr := resultErrorFromError(errStream) - action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, errStream, m.requestScopedErrorConfig(ctx)) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(errStream) if isCredentialScopedError(errStream) { @@ -538,7 +538,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi auth = refreshed publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true - ctx = newUpstreamAttemptContext(ctx) + ctx = m.newUpstreamAttemptContext(ctx) executionSeq = m.nextExecutionSeq() startRetry := time.Now() retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) @@ -578,7 +578,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } if bootstrapErr != nil { - action, okAction := matchRequestScopedErrorAction(auth, bootstrapErr, m.runtimeConfigSnapshot()) + action, okAction := matchRequestScopedErrorAction(auth, bootstrapErr, m.requestScopedErrorConfig(ctx)) if okAction { rerr := resultErrorFromError(bootstrapErr) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} diff --git a/sdk/cliproxy/auth/cooldown_state.go b/sdk/cliproxy/auth/cooldown_state.go index 28c9c91f048..938d5a2da46 100644 --- a/sdk/cliproxy/auth/cooldown_state.go +++ b/sdk/cliproxy/auth/cooldown_state.go @@ -34,6 +34,24 @@ type CooldownStateRecord struct { // CredentialQuotaPropagated preserves wrap provenance for models that // carry no independent layer (e.g. the reporting model itself). CredentialQuotaPropagated bool `json:"credential_quota_propagated,omitempty"` + // CooldownSignals persists bounded rejected quota headers captured when + // this cooldown was recorded without touching passive Quota.Signals watermarks. + CooldownSignals map[string]string `json:"cooldown_signals,omitempty"` + // IndependentCooldownSignals persists CooldownSignals for IndependentState. + IndependentCooldownSignals map[string]string `json:"independent_cooldown_signals,omitempty"` + // ForceCooldownUntil persists the exact outer ErrorCodeForceCooldown deadline + // (credQuotaForceUntil for wrapped model records, model/auth forceCooldownUntil otherwise), + // preventing longer ordinary extensions or sibling shadows from inflating it on restore. + ForceCooldownUntil time.Time `json:"force_cooldown_until,omitempty"` + // CredQuotaForceCooldownUntil persists auth-level credential_quota force provenance + // independently from model-less auth ForceCooldownUntil when both coexist on an auth. + CredQuotaForceCooldownUntil time.Time `json:"cred_quota_force_cooldown_until,omitempty"` + // IndependentForceCooldownUntil persists the exact ErrorCodeForceCooldown deadline + // for IndependentState beneath a credential_quota wrap. + IndependentForceCooldownUntil time.Time `json:"independent_force_cooldown_until,omitempty"` + // ForceCooldownEvaluated marks records produced by provenance-aware versions so + // a zero ForceCooldownUntil means "explicitly ordinary" rather than legacy unknown. + ForceCooldownEvaluated bool `json:"force_cooldown_evaluated,omitempty"` } // CooldownStateStore persists runtime cooldown state independently from auth tokens. diff --git a/sdk/cliproxy/auth/home_concurrency.go b/sdk/cliproxy/auth/home_concurrency.go index 15d5c291d58..8a3f47a4db8 100644 --- a/sdk/cliproxy/auth/home_concurrency.go +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -310,7 +310,7 @@ func SafeResponseHeaders(err error) http.Header { } var modelCooldown *modelCooldownError if errors.As(err, &modelCooldown) && modelCooldown != nil { - return modelCooldown.Headers() + return modelCooldown.SafeResponseHeaders() } var unavailable *authUnavailableError if errors.As(err, &unavailable) && unavailable != nil { diff --git a/sdk/cliproxy/auth/issue_7431_regression_test.go b/sdk/cliproxy/auth/issue_7431_regression_test.go new file mode 100644 index 00000000000..d3e7c23af4b --- /dev/null +++ b/sdk/cliproxy/auth/issue_7431_regression_test.go @@ -0,0 +1,2575 @@ +package auth + +import ( + "context" + "fmt" + "net/http" + "sync" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestIssue7431_LargeWeightedPoolRotatesAcrossTransientExclusionsAndCooldowns(t *testing.T) { + for _, withAffinity := range []bool{false, true} { + name := "scheduler-fast-path" + if withAffinity { + name = "session-affinity-legacy" + } + t.Run(name, func(t *testing.T) { + const ( + totalCreds = 1080 + model = "issue-7431-large-wrr-model" + compatName = "issue-7431-pool" + totalPicks = 1120 + failingEvery = 40 + ) + providerKey := util.OpenAICompatibleProviderKey(compatName) + disableCooling := true + retryCount := 4 + cfg := &internalconfig.Config{ + RequestRetry: retryCount, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: compatName, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCooling, + Models: []internalconfig.OpenAICompatibilityModel{{Name: model, Alias: model}}, + RequestScopedErrors: []internalconfig.RequestScopedErrorRule{{ + Status: http.StatusTooManyRequests, + Match: []string{"server_busy"}, + Action: RequestScopedActionContinueAndCooldown, + }}, + }}, + } + + weighted := &WeightedRoundRobinSelector{} + var selector Selector = weighted + if withAffinity { + affinity := NewSessionAffinitySelector(weighted) + t.Cleanup(affinity.Stop) + selector = affinity + } + + manager := NewManager(nil, selector, nil) + manager.SetConfig(cfg) + manager.SetRetryConfig(retryCount, 0, 0) + + reg := registry.GetGlobalRegistry() + authIDs := make([]string, totalCreds) + failSet := make(map[string]bool) + for i := 0; i < totalCreds; i++ { + id := fmt.Sprintf("thk-%s-%04d", name, i) + authIDs[i] = id + // First 4 IDs and every 40th ID fail with 429 so transient exclusions churn continuously. + if i < 4 || i%failingEvery == 0 { + failSet[id] = true + } + auth := &Auth{ + ID: id, + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{ + "auth_kind": "apikey", + "api_key": fmt.Sprintf("key-%04d", i), + "compat_name": compatName, + "provider_key": providerKey, + }, + } + reg.RegisterClient(id, providerKey, []*registry.ModelInfo{{ID: model}}) + if _, errReg := manager.Register(WithSkipPersist(t.Context()), auth); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + } + t.Cleanup(func() { + for _, id := range authIDs { + reg.UnregisterClient(id) + } + }) + + selectedCounts := make(map[string]int, totalCreds) + manager.RegisterExecutor(&mockCustomErrorExecutor{ + identifier: providerKey, + executeFn: func(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + selectedCounts[auth.ID]++ + if failSet[auth.ID] { + return cliproxyexecutor.Response{}, customStatusError{ + code: http.StatusTooManyRequests, + msg: "server_busy", + } + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil + }, + }) + + for reqIdx := 0; reqIdx < totalPicks; reqIdx++ { + opts := cliproxyexecutor.Options{} + if withAffinity { + opts.Headers = http.Header{"Session_id": {fmt.Sprintf("req-session-%d", reqIdx)}} + } + if _, errExec := manager.Execute(t.Context(), []string{providerKey}, cliproxyexecutor.Request{Model: model}, opts); errExec != nil { + t.Fatalf("request %d failed: %v", reqIdx, errExec) + } + } + + healthySelected := 0 + for _, id := range authIDs { + if !failSet[id] && selectedCounts[id] > 0 { + healthySelected++ + } + } + wantMinHealthy := totalCreds - len(failSet) - 10 + if healthySelected < wantMinHealthy { + t.Fatalf("large WRR pool collapsed: only %d/%d healthy credentials selected across %d requests", healthySelected, totalCreds-len(failSet), totalPicks) + } + }) + } +} + +func TestIssue7431_SmoothWeightedStateLargePoolRetainsTransientlyExcludedCredits(t *testing.T) { + const members = 2131 + state := &smoothWeightedState{} + full := make(map[string]int64, members) + for i := 0; i < members; i++ { + full[fmt.Sprintf("cred-%04d", i)] = 1 + } + state.prepare(full) + state.current["cred-0000"] = -(members - 1) + + subset := make(map[string]int64, members-1) + for i := 1; i < members; i++ { + subset[fmt.Sprintf("cred-%04d", i)] = 1 + } + state.prepare(subset) + + if got := state.current["cred-0000"]; got != -(members - 1) { + t.Fatalf("transiently excluded credential credit = %d, want %d", got, -(members - 1)) + } +} + +func TestIssue7431_SchedulerRebuildPreservesUnchangedShardCursors(t *testing.T) { + ctx := t.Context() + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.RegisterExecutor(schedulerProviderTestExecutor{provider: "openai"}) + manager.RegisterExecutor(schedulerProviderTestExecutor{provider: "claude"}) + + const model = "issue-7431-rebuild-cursor-model" + reg := registry.GetGlobalRegistry() + var ids []string + for i := 0; i < 5; i++ { + id := fmt.Sprintf("rr-auth-%02d", i) + ids = append(ids, id) + auth := &Auth{ID: id, Provider: "openai", Status: StatusActive} + if _, errReg := manager.Register(WithSkipPersist(ctx), auth); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + reg.RegisterClient(id, "openai", []*registry.ModelInfo{{ID: model}}) + manager.RefreshSchedulerEntry(id) + } + otherAuth := &Auth{ID: "claude-other-01", Provider: "claude", Status: StatusActive} + if _, errReg := manager.Register(WithSkipPersist(ctx), otherAuth); errReg != nil { + t.Fatalf("register other auth: %v", errReg) + } + reg.RegisterClient(otherAuth.ID, "claude", []*registry.ModelInfo{{ID: model}}) + manager.RefreshSchedulerEntry(otherAuth.ID) + t.Cleanup(func() { + for _, id := range ids { + reg.UnregisterClient(id) + } + reg.UnregisterClient(otherAuth.ID) + }) + + for wantIdx := 0; wantIdx < 3; wantIdx++ { + picked, _, errPick := manager.pickNext(ctx, "openai", model, cliproxyexecutor.Options{}, nil) + if errPick != nil || picked == nil { + t.Fatalf("pick %d failed: auth=%v err=%v", wantIdx, picked, errPick) + } + if picked.ID != ids[wantIdx] { + t.Fatalf("pick %d = %q, want %q", wantIdx, picked.ID, ids[wantIdx]) + } + } + + // Simulate scheduler rebuild triggered by unrelated config/provider updates. + manager.scheduler.rebuild(manager.snapshotAuths()) + + pickedAfter, _, errPick := manager.pickNext(ctx, "openai", model, cliproxyexecutor.Options{}, nil) + if errPick != nil || pickedAfter == nil { + t.Fatalf("pick after rebuild failed: auth=%v err=%v", pickedAfter, errPick) + } + if pickedAfter.ID != ids[3] { + t.Fatalf("pick after rebuild reset rotation to %q, want %q", pickedAfter.ID, ids[3]) + } +} + +func TestIssue7431_DisableCoolingPreservesActiveForceCooldownOnReloadAndUpdate(t *testing.T) { + ctx := WithSkipPersist(t.Context()) + const ( + compatName = "issue-7431-force-cool" + model = "issue-7431-force-cool-model" + ) + providerKey := util.OpenAICompatibleProviderKey(compatName) + disableCooling := true + buildConfig := func(extraClaudeKey string) *internalconfig.Config { + cfg := &internalconfig.Config{ + RequestRetry: 2, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: compatName, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCooling, + Models: []internalconfig.OpenAICompatibilityModel{{Name: model, Alias: model}}, + RequestScopedErrors: []internalconfig.RequestScopedErrorRule{{ + Status: http.StatusTooManyRequests, + Match: []string{"rate_limited"}, + Action: RequestScopedActionContinueAndCooldown, + }}, + }}, + } + if extraClaudeKey != "" { + cfg.ClaudeKey = []internalconfig.ClaudeKey{{APIKey: extraClaudeKey}} + } + return cfg + } + + store := &recordingCooldownStateStore{} + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.SetCooldownStateStore(store) + manager.SetConfig(buildConfig("")) + manager.SetRetryConfig(2, 0, 0) + + reg := registry.GetGlobalRegistry() + auth0 := &Auth{ + ID: "force-cool-auth-0", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{ + "auth_kind": "apikey", + "api_key": "k0", + "compat_name": compatName, + "provider_key": providerKey, + }, + } + auth1 := &Auth{ + ID: "force-cool-auth-1", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{ + "auth_kind": "apikey", + "api_key": "k1", + "compat_name": compatName, + "provider_key": providerKey, + }, + } + for _, a := range []*Auth{auth0, auth1} { + reg.RegisterClient(a.ID, providerKey, []*registry.ModelInfo{{ID: model}}) + if _, errReg := manager.Register(ctx, a); errReg != nil { + t.Fatalf("register %s: %v", a.ID, errReg) + } + manager.RefreshSchedulerEntry(a.ID) + } + t.Cleanup(func() { + reg.UnregisterClient(auth0.ID) + reg.UnregisterClient(auth1.ID) + }) + + manager.RegisterExecutor(&mockCustomErrorExecutor{ + identifier: providerKey, + executeFn: func(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if auth.ID == auth0.ID { + return cliproxyexecutor.Response{}, customStatusError{ + code: http.StatusTooManyRequests, + msg: "rate_limited", + } + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil + }, + }) + + if _, errExec := manager.Execute(ctx, []string{providerKey}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}); errExec != nil { + t.Fatalf("initial execute failed: %v", errExec) + } + + now := time.Now() + got0, ok := manager.GetByID(auth0.ID) + if !ok || IsAuthAvailableForModel(got0, model, now) { + t.Fatalf("expected %s to enter force cooldown before reload", auth0.ID) + } + + // 1. Unrelated config reload must NOT erase active ErrorCodeForceCooldown under disable-cooling: true. + manager.SetConfig(buildConfig("sk-ant-unrelated")) + got0AfterReload, _ := manager.GetByID(auth0.ID) + if IsAuthAvailableForModel(got0AfterReload, model, now) { + t.Fatalf("SetConfig erased active ErrorCodeForceCooldown on %s when disable-cooling: true", auth0.ID) + } + pickedAfterReload, _, errPick := manager.pickNext(ctx, providerKey, model, cliproxyexecutor.Options{}, nil) + if errPick != nil || pickedAfterReload == nil || pickedAfterReload.ID != auth1.ID { + t.Fatalf("scheduler selected %v (err=%v) after config reload, want %s while %s is force-cooled", pickedAfterReload, errPick, auth1.ID, auth0.ID) + } + + // 2. Non-credential metadata update on auth0 must NOT erase active ErrorCodeForceCooldown. + updatedAuth0 := got0AfterReload.Clone() + updatedAuth0.Attributes["label_note"] = "reloaded" + if _, errUpd := manager.Update(ctx, updatedAuth0); errUpd != nil { + t.Fatalf("Update(%s) failed: %v", auth0.ID, errUpd) + } + got0AfterUpdate, _ := manager.GetByID(auth0.ID) + if IsAuthAvailableForModel(got0AfterUpdate, model, now) { + t.Fatalf("Update() erased active ErrorCodeForceCooldown on %s when disable-cooling: true", auth0.ID) + } + + // 3. CooldownStateStore persistence & restore must also retain active ErrorCodeForceCooldown under disable-cooling: true. + manager.PersistCooldownStates(ctx) + if len(store.savedRecords()) == 0 { + t.Fatalf("expected active ErrorCodeForceCooldown record to be persisted even with disable-cooling: true") + } + + restoredManager := NewManager(nil, &RoundRobinSelector{}, nil) + restoredManager.SetCooldownStateStore(&recordingCooldownStateStore{load: store.savedRecords()}) + restoredManager.SetConfig(buildConfig("sk-ant-unrelated")) + freshAuth0 := &Auth{ + ID: auth0.ID, + Provider: providerKey, + Status: StatusActive, + Attributes: auth0.Attributes, + } + if _, errReg := restoredManager.Register(ctx, freshAuth0); errReg != nil { + t.Fatalf("register freshAuth0: %v", errReg) + } + if errRestore := restoredManager.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates failed: %v", errRestore) + } + got0Restored, _ := restoredManager.GetByID(auth0.ID) + if IsAuthAvailableForModel(got0Restored, model, now) { + t.Fatalf("RestoreCooldownStates dropped active ErrorCodeForceCooldown under disable-cooling: true") + } +} + +func TestIssue7431_LargeWeightedPoolEdgeCases(t *testing.T) { + ctx := context.Background() + + t.Run("small-subset-does-not-clamp-absent-member-debt", func(t *testing.T) { + sel := &WeightedRoundRobinSelector{} + route := "openai-compatible-edge:model-a" + full := make([]*Auth, 1050) + for i := range full { + w := 10 + if i == 0 { + w = 500 + } + full[i] = &Auth{ + ID: fmt.Sprintf("edge-auth-%04d", i), + Provider: "openai-compatible-edge", + Status: StatusActive, + Attributes: map[string]string{"weight": fmt.Sprintf("%d", w)}, + } + } + picked, res := sel.reserve(route, full, "") + if picked == nil || picked.ID != "edge-auth-0000" { + t.Fatalf("first pick = %v, want edge-auth-0000", picked) + } + res.finish(true) + + // Shrink temporarily to just 2 low-weight members; edge-auth-0000 is absent (e.g. cooled). + subset := []*Auth{full[1], full[2]} + pickedSub, resSub := sel.reserve(route, subset, "") + if pickedSub == nil { + t.Fatal("subset pick returned nil") + } + resSub.finish(true) + + // When edge-auth-0000 returns, its large negative debt (-10490) must still hold it back, + // rather than having been clamped to -200 (subsetTotal=20 * 10). + pickedNext, resNext := sel.reserve(route, full, "") + if pickedNext == nil { + t.Fatal("next full pick returned nil") + } + resNext.finish(true) + if pickedNext.ID == "edge-auth-0000" { + t.Fatalf("edge-auth-0000 was reselected immediately after 2-auth subset; absent debt was prematurely clamped") + } + }) + + t.Run("priority-buckets-isolated-and-forget-auth-prunes-dynamic-pool", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-prio" + const model = "issue-7431-prio-model" + sel := &WeightedRoundRobinSelector{} + manager := NewManager(nil, sel, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: providerKey}) + reg := registry.GetGlobalRegistry() + + for i := 0; i < 1030; i++ { + id := fmt.Sprintf("prio-hi-%04d", i) + auth := &Auth{ + ID: id, + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{ + "provider_key": providerKey, + "priority": "10", + "weight": "5", + }, + } + if _, errReg := manager.Register(ctx, auth); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + reg.RegisterClient(id, providerKey, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + } + lowAuth := &Auth{ + ID: "prio-lo-0000", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{ + "provider_key": providerKey, + "priority": "0", + "weight": "5", + }, + } + if _, errReg := manager.Register(ctx, lowAuth); errReg != nil { + t.Fatalf("register lowAuth: %v", errReg) + } + reg.RegisterClient(lowAuth.ID, providerKey, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(lowAuth.ID) }) + + picked, _, errPick := manager.pickNext(ctx, providerKey, model, cliproxyexecutor.Options{}, nil) + if errPick != nil || picked == nil || picked.ID != "prio-hi-0000" { + t.Fatalf("pick = %v (err=%v), want prio-hi-0000", picked, errPick) + } + + // Deregistering prio-hi-0000 must prune it from largeStates via ForgetAuth. + if !manager.RemoveIf(ctx, "prio-hi-0000", nil) { + t.Fatalf("RemoveIf(prio-hi-0000) returned false") + } + sel.mu.Lock() + for _, state := range sel.largeStates { + if _, found := state.findMember("prio-hi-0000"); found { + sel.mu.Unlock() + t.Fatalf("ForgetAuth failed to prune removed auth prio-hi-0000 from largeStates") + } + if _, foundLow := state.findMember("prio-lo-0000"); foundLow { + sel.mu.Unlock() + t.Fatalf("priority 0 auth leaked into priority 10 largeStates bucket") + } + } + sel.mu.Unlock() + }) +} + +func TestIssue7431_ForceCooldownEdgeCases(t *testing.T) { + ctx := context.Background() + now := time.Now() + + t.Run("mixed-ordinary-and-force-cooldowns-plus-wrapped-shadow-and-reset-quota", func(t *testing.T) { + disableCooling := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{ + { + Name: "openai-compatible-issue-7431-edge", + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCooling, + }, + }, + } + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "force-edge-auth", + Provider: "openai-compatible-issue-7431-edge", + Status: StatusError, + Attributes: map[string]string{ + "provider_key": "openai-compatible-issue-7431-edge", + }, + Unavailable: true, + NextRetryAfter: now.Add(5 * time.Minute), + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: now.Add(5 * time.Minute), + }, + LastError: &Error{Code: credentialQuotaErrorCode, HTTPStatus: 429, Message: "daily limit"}, + ModelStates: map[string]*ModelState{ + "model-force": { + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(5 * time.Minute), + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(5 * time.Minute)}, + credentialQuotaPropagated: true, + credentialQuotaShadow: &ModelState{ + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(2 * time.Minute), + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(2 * time.Minute)}, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 429, Message: "force model cooldown"}, + UpdatedAt: now, + }, + UpdatedAt: now, + }, + "model-ordinary": { + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(5 * time.Minute), + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(5 * time.Minute)}, + credentialQuotaPropagated: true, + credentialQuotaShadow: &ModelState{ + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(2 * time.Minute), + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(2 * time.Minute)}, + LastError: &Error{HTTPStatus: 429, Message: "normal model cooldown"}, + UpdatedAt: now, + }, + UpdatedAt: now, + }, + "model-expired-force": { + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(-time.Second), + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(-time.Second)}, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 429, Message: "expired force"}, + UpdatedAt: now.Add(-time.Minute), + }, + }, + } + if _, errReg := manager.Register(ctx, auth); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + got, _ := manager.GetByID(auth.ID) + if IsAuthAvailableForModel(got, "model-force", now) { + t.Fatal("expected wrapped active ErrorCodeForceCooldown shadow on model-force to survive Register under disable-cooling") + } + if !IsAuthAvailableForModel(got, "model-ordinary", now) { + t.Fatal("expected ordinary non-force cooldown on model-ordinary to be cleared under disable-cooling") + } + if !IsAuthAvailableForModel(got, "model-expired-force", now) { + t.Fatal("expected expired ErrorCodeForceCooldown on model-expired-force to be cleared") + } + + // Explicit admin ResetQuota must still clear ErrorCodeForceCooldown. + if _, _, errReset := manager.ResetQuota(ctx, auth.ID); errReset != nil { + t.Fatalf("ResetQuota: %v", errReset) + } + gotAfterReset, _ := manager.GetByID(auth.ID) + if !IsAuthAvailableForModel(gotAfterReset, "model-force", now) { + t.Fatal("expected explicit ResetQuota to clear ErrorCodeForceCooldown") + } + + // Auth-level (Model == "") ErrorCodeForceCooldown recorded alongside existing ModelStates + // must also survive SetConfig reload under disable-cooling: true. + retry := 2 * time.Minute + manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Success: false, + RetryAfter: &retry, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "force auth cooldown"}, + }) + manager.SetConfig(cfg) + gotAuthForce, _ := manager.GetByID(auth.ID) + if IsAuthAvailableForModel(gotAuthForce, "", now) { + t.Fatal("expected model-less auth-level ErrorCodeForceCooldown to survive SetConfig reload alongside ModelStates") + } + }) + + t.Run("mark-result-interleaving-does-not-corrupt-force-credential-cooldown-provenance", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-provenance" + disableCoolingTrue := true + disableCoolingFalse := false + cfgCoolingOff := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + cfgCoolingOn := &internalconfig.Config{ + DisableCooling: false, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingFalse, + }}, + } + + // Order A: Force credential_quota first, followed by ordinary model error overwriting auth.LastError. + mgrA := NewManager(nil, &RoundRobinSelector{}, nil) + mgrA.SetConfig(cfgCoolingOff) + authA := &Auth{ + ID: "prov-auth-a", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + } + if _, errReg := mgrA.Register(ctx, authA); errReg != nil { + t.Fatalf("Register authA: %v", errReg) + } + fiveMin := 5 * time.Minute + mgrA.MarkResult(ctx, Result{ + AuthID: authA.ID, + Provider: providerKey, + Model: "model-1", + Success: false, + CredentialScope: true, + RetryAfter: &fiveMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force account limit"}, + }) + mgrA.MarkResult(ctx, Result{ + AuthID: authA.ID, + Provider: providerKey, + Model: "model-2", + Success: false, + Error: &Error{Code: "server_error", HTTPStatus: http.StatusInternalServerError, Message: "ordinary 500"}, + }) + mgrA.SetConfig(cfgCoolingOff) + gotA, _ := mgrA.GetByID(authA.ID) + checkTime := time.Now() + if IsAuthAvailableForModel(gotA, "model-1", checkTime) || IsAuthAvailableForModel(gotA, "model-2", checkTime) { + t.Fatalf("ordinary model error overwrote auth.LastError and caused SetConfig to erase active force credential_quota: quota=%+v", gotA.Quota) + } + + // Order B: Ordinary credential_quota first (cooling enabled), followed by model-scoped force cooldown, + // then reload enabling disable-cooling: true. + mgrB := NewManager(nil, &RoundRobinSelector{}, nil) + mgrB.SetConfig(cfgCoolingOn) + authB := &Auth{ + ID: "prov-auth-b", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + } + if _, errReg := mgrB.Register(ctx, authB); errReg != nil { + t.Fatalf("Register authB: %v", errReg) + } + mgrB.MarkResult(ctx, Result{ + AuthID: authB.ID, + Provider: providerKey, + Model: "model-1", + Success: false, + CredentialScope: true, + RetryAfter: &fiveMin, + Error: &Error{Code: credentialQuotaErrorCode, HTTPStatus: http.StatusTooManyRequests, Message: "ordinary account limit"}, + }) + twoMin := 2 * time.Minute + mgrB.MarkResult(ctx, Result{ + AuthID: authB.ID, + Provider: providerKey, + Model: "model-2", + Success: false, + RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force model limit"}, + }) + mgrB.SetConfig(cfgCoolingOff) + gotB, _ := mgrB.GetByID(authB.ID) + checkTimeB := time.Now() + if !IsAuthAvailableForModel(gotB, "model-1", checkTimeB) { + t.Fatalf("model-scoped force error overwrote auth.LastError and falsely preserved ordinary credential_quota on model-1: quota=%+v", gotB.Quota) + } + if IsAuthAvailableForModel(gotB, "model-2", checkTimeB) { + t.Fatal("expected model-2's independent force cooldown to survive after ordinary credential_quota was cleared") + } + + // Persistence + restore after auth.LastError was overwritten in Order A must also preserve force credential_quota. + storeA := &recordingCooldownStateStore{} + mgrA.SetCooldownStateStore(storeA) + mgrA.PersistCooldownStates(ctx) + mgrARestored := NewManager(nil, &RoundRobinSelector{}, nil) + mgrARestored.SetCooldownStateStore(&recordingCooldownStateStore{load: storeA.savedRecords()}) + mgrARestored.SetConfig(cfgCoolingOff) + if _, errReg := mgrARestored.Register(ctx, &Auth{ + ID: authA.ID, + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register restored authA: %v", errReg) + } + if errRestore := mgrARestored.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates authA: %v", errRestore) + } + gotARestored, _ := mgrARestored.GetByID(authA.ID) + if IsAuthAvailableForModel(gotARestored, "model-1", time.Now()) || IsAuthAvailableForModel(gotARestored, "model-2", time.Now()) { + t.Fatalf("RestoreCooldownStates lost force credential_quota after auth.LastError overwrite: quota=%+v", gotARestored.Quota) + } + }) +} + +func TestIssue7431_DynamicPoolBoundsAndReservationLifecycle(t *testing.T) { + ctx := context.Background() + + t.Run("large-states-enforces-max-keys-capacity-bound", func(t *testing.T) { + sel := &WeightedRoundRobinSelector{maxKeys: 4} + pool := make([]*Auth, 1030) + for i := range pool { + pool[i] = &Auth{ + ID: fmt.Sprintf("cap-auth-%04d", i), + Provider: "openai-compatible-cap", + Status: StatusActive, + Attributes: map[string]string{"weight": "1"}, + } + } + _, inFlightRes := sel.reserve("openai-compatible-cap:route-0", pool, "") + if inFlightRes == nil { + t.Fatal("expected reservation on route-0") + } + for r := 1; r <= 6; r++ { + _, res := sel.reserve(fmt.Sprintf("openai-compatible-cap:route-%d", r), pool, "") + res.finish(true) + } + sel.mu.Lock() + gotLargeLen := len(sel.largeStates) + sel.mu.Unlock() + if gotLargeLen > 4 { + t.Fatalf("len(sel.largeStates) = %d, want <= maxKeys (4)", gotLargeLen) + } + inFlightRes.finish(false) + + // All-in-flight eviction must isolate stale reservations from newly created state for the same route. + sel2 := &WeightedRoundRobinSelector{maxKeys: 2} + _, evictResA := sel2.reserve("openai-compatible-cap:a", pool, "") + _, evictResB := sel2.reserve("openai-compatible-cap:b", pool, "") + _, evictResC := sel2.reserve("openai-compatible-cap:c", pool, "") + evictResB.finish(true) + evictResC.finish(true) + // Re-create route "a" and commit one clean pick, then roll back stale evictResA. + _, freshResA := sel2.reserve("openai-compatible-cap:a", pool, "") + freshResA.finish(true) + evictResA.finish(false) + sel2.mu.Lock() + stateA := sel2.largeStates["openai-compatible-cap:a"] + for _, m := range stateA.members { + if m.pending != 0 { + sel2.mu.Unlock() + t.Fatalf("evicted reservation rollback corrupted fresh route state: member %s pending=%d", m.id, m.pending) + } + } + sel2.mu.Unlock() + }) + + t.Run("forget-and-reregister-same-id-does-not-pollute-pending-credit", func(t *testing.T) { + sel := &WeightedRoundRobinSelector{} + const route = "openai-compatible-epoch:model-x" + pool := make([]*Auth, 1030) + for i := range pool { + pool[i] = &Auth{ + ID: fmt.Sprintf("epoch-auth-%04d", i), + Provider: "openai-compatible-epoch", + Status: StatusActive, + Attributes: map[string]string{"weight": "1"}, + } + } + picked1, staleRes := sel.reserve(route, pool, "") + if picked1 == nil || picked1.ID != "epoch-auth-0000" { + t.Fatalf("first pick = %v, want epoch-auth-0000", picked1) + } + + // Permanently remove epoch-auth-0000 while staleRes is still in-flight, + // then re-register a fresh credential with the same ID. + sel.ForgetAuth("epoch-auth-0000") + _, freshRes := sel.reserve(route, pool, "") + freshRes.finish(true) + + // Roll back staleRes; it must NOT write +1029 pending credit onto the newly registered epoch-auth-0000. + staleRes.finish(false) + + sel.mu.Lock() + state := sel.largeStates[route] + var nonZeroPendingID string + var nonZeroPendingVal int64 + if state != nil { + for _, m := range state.members { + if m.pending != 0 { + nonZeroPendingID = m.id + nonZeroPendingVal = m.pending + break + } + } + } + sel.mu.Unlock() + if nonZeroPendingID != "" { + t.Fatalf("member %s has leaked pending=%d after stale reservation rollback across ForgetAuth+re-register, want 0", nonZeroPendingID, nonZeroPendingVal) + } + }) + + t.Run("large-weighted-pool-preserves-progress-across-scheduler-rebuild", func(t *testing.T) { + const ( + providerA = "openai-compatible-wrr-rebuild-a" + providerB = "openai-compatible-wrr-rebuild-b" + model = "issue-7431-wrr-rebuild-model" + ) + sel := &WeightedRoundRobinSelector{} + manager := NewManager(nil, sel, nil) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: providerA}) + manager.RegisterExecutor(&mockCustomErrorExecutor{identifier: providerB}) + reg := registry.GetGlobalRegistry() + + allAuths := make([]*Auth, 0, 1031) + for i := 0; i < 1030; i++ { + id := fmt.Sprintf("wrr-reb-%04d", i) + a := &Auth{ + ID: id, + Provider: providerA, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerA, "weight": "1"}, + } + reg.RegisterClient(id, providerA, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + registered, errReg := manager.Register(WithSkipPersist(ctx), a) + if errReg != nil { + t.Fatalf("Register %s: %v", id, errReg) + } + allAuths = append(allAuths, registered) + } + + for i := 0; i < 3; i++ { + picked, _, errPick := manager.pickNext(ctx, providerA, model, cliproxyexecutor.Options{}, nil) + if errPick != nil || picked == nil || picked.ID != fmt.Sprintf("wrr-reb-%04d", i) { + t.Fatalf("pick %d before rebuild = %v (err=%v), want wrr-reb-%04d", i, picked, errPick, i) + } + } + + unrelated := &Auth{ + ID: "unrelated-b-0000", + Provider: providerB, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerB, "weight": "1"}, + } + reg.RegisterClient(unrelated.ID, providerB, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(unrelated.ID) }) + regB, errReg := manager.Register(WithSkipPersist(ctx), unrelated) + if errReg != nil { + t.Fatalf("Register unrelated: %v", errReg) + } + allAuths = append(allAuths, regB) + manager.scheduler.rebuild(allAuths) + + pickedAfter, _, errPick := manager.pickNext(ctx, providerA, model, cliproxyexecutor.Options{}, nil) + if errPick != nil || pickedAfter == nil || pickedAfter.ID != "wrr-reb-0003" { + t.Fatalf("pick after rebuild = %v (err=%v), want wrr-reb-0003", pickedAfter, errPick) + } + }) +} + +func TestIssue7431_SameLayerForceCooldownAndPersistenceAndEpochRaces(t *testing.T) { + ctx := context.Background() + + t.Run("same-layer-ordinary-error-does-not-clear-active-force-cooldown", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-samelayer" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetConfig(cfg) + + authModel := &Auth{ + ID: "samelayer-model-auth", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + } + authGlobal := &Auth{ + ID: "samelayer-global-auth", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + } + if _, errReg := mgr.Register(WithSkipPersist(ctx), authModel); errReg != nil { + t.Fatalf("register authModel: %v", errReg) + } + if _, errReg := mgr.Register(WithSkipPersist(ctx), authGlobal); errReg != nil { + t.Fatalf("register authGlobal: %v", errReg) + } + + twoMin := 2 * time.Minute + // 1. Same model: force 503 followed by in-flight ordinary 500 under disable-cooling: true. + mgr.MarkResult(ctx, Result{ + AuthID: authModel.ID, + Provider: providerKey, + Model: "model-same", + Success: false, + RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "force 503"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: authModel.ID, + Provider: providerKey, + Model: "model-same", + Success: false, + Error: &Error{Code: "server_error", HTTPStatus: http.StatusInternalServerError, Message: "ordinary 500"}, + }) + now := time.Now() + gotModel, _ := mgr.GetByID(authModel.ID) + if IsAuthAvailableForModel(gotModel, "model-same", now) { + t.Fatalf("ordinary 500 on same model cleared active force cooldown immediately: state=%+v", gotModel.ModelStates["model-same"]) + } + mgr.SetConfig(cfg) + gotModelAfterReload, _ := mgr.GetByID(authModel.ID) + if IsAuthAvailableForModel(gotModelAfterReload, "model-same", now) { + t.Fatalf("SetConfig after ordinary 500 on same model cleared active force cooldown: state=%+v", gotModelAfterReload.ModelStates["model-same"]) + } + + // 2. Model-less auth level (Model == ""): force 503 followed by in-flight ordinary 500 under disable-cooling: true. + mgr.MarkResult(ctx, Result{ + AuthID: authGlobal.ID, + Provider: providerKey, + Success: false, + RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "force 503"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: authGlobal.ID, + Provider: providerKey, + Success: false, + Error: &Error{Code: "server_error", HTTPStatus: http.StatusInternalServerError, Message: "ordinary 500"}, + }) + gotGlobal, _ := mgr.GetByID(authGlobal.ID) + if IsAuthAvailableForModel(gotGlobal, "", now) { + t.Fatalf("ordinary 500 on model-less auth cleared active force cooldown immediately: RetryAfter=%v", gotGlobal.NextRetryAfter) + } + mgr.SetConfig(cfg) + gotGlobalAfterReload, _ := mgr.GetByID(authGlobal.ID) + if IsAuthAvailableForModel(gotGlobalAfterReload, "", now) { + t.Fatalf("SetConfig after ordinary 500 on model-less auth cleared active force cooldown: RetryAfter=%v", gotGlobalAfterReload.NextRetryAfter) + } + }) + + t.Run("persisted-force-deadlines-preserve-exact-unequal-bounds-across-disk-roundtrip", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-disk" + disableCoolingFalse := false + disableCoolingTrue := true + cfgOn := &internalconfig.Config{ + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingFalse, + }}, + } + cfgOff := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + + cdsDir := t.TempDir() + fileStore := NewFileCooldownStateStore(cdsDir) + mgr1 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr1.SetCooldownStateStore(fileStore) + mgr1.SetConfig(cfgOn) + + // Auth 1: 2m force cooldown on model-x, then extended to 10m by an ordinary 429. + a1 := &Auth{ID: "disk-auth-1", Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey}} + // Auth 2: 2m force credential_quota on model-a + 10m force independent shadow on model-b. + a2 := &Auth{ID: "disk-auth-2", Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey}} + if _, errReg := mgr1.Register(WithSkipPersist(ctx), a1); errReg != nil { + t.Fatalf("register a1: %v", errReg) + } + if _, errReg := mgr1.Register(WithSkipPersist(ctx), a2); errReg != nil { + t.Fatalf("register a2: %v", errReg) + } + + twoMin := 2 * time.Minute + tenMin := 10 * time.Minute + mgr1.MarkResult(ctx, Result{ + AuthID: a1.ID, Provider: providerKey, Model: "model-x", RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force 2m"}, + }) + mgr1.MarkResult(ctx, Result{ + AuthID: a1.ID, Provider: providerKey, Model: "model-x", RetryAfter: &tenMin, + Error: &Error{Code: "rate_limited", HTTPStatus: http.StatusTooManyRequests, Message: "ordinary 10m"}, + }) + + mgr1.MarkResult(ctx, Result{ + AuthID: a2.ID, Provider: providerKey, Model: "model-a", CredentialScope: true, RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force cred 2m"}, + }) + mgr1.MarkResult(ctx, Result{ + AuthID: a2.ID, Provider: providerKey, Model: "model-b", RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force model-b 10m"}, + }) + mgr1.PersistCooldownStates(ctx) + + // Restart with disable-cooling: true and restore from real .cds files on disk. + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsDir)) + mgr2.SetConfig(cfgOff) + for _, id := range []string{a1.ID, a2.ID} { + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: id, Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("register %s in mgr2: %v", id, errReg) + } + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + baseNow := time.Now() + gotA1, _ := mgr2.GetByID(a1.ID) + if IsAuthAvailableForModel(gotA1, "model-x", baseNow.Add(time.Minute)) { + t.Fatal("expected disk-auth-1 model-x to remain blocked within 2m force deadline") + } + if !IsAuthAvailableForModel(gotA1, "model-x", baseNow.Add(3*time.Minute)) { + t.Fatalf("disk-auth-1 model-x remained blocked at +3m; 10m ordinary extension inflated restored force deadline: %+v", gotA1.ModelStates["model-x"]) + } + + gotA2, _ := mgr2.GetByID(a2.ID) + if IsAuthAvailableForModel(gotA2, "model-a", baseNow.Add(time.Minute)) { + t.Fatal("expected disk-auth-2 model-a to remain blocked within 2m force credential_quota") + } + if !IsAuthAvailableForModel(gotA2, "model-a", baseNow.Add(3*time.Minute)) { + t.Fatalf("disk-auth-2 model-a remained blocked at +3m; model-b 10m shadow inflated credQuotaForceUntil: quota=%+v", gotA2.Quota) + } + if IsAuthAvailableForModel(gotA2, "model-b", baseNow.Add(5*time.Minute)) { + t.Fatal("expected disk-auth-2 model-b to remain blocked by its 10m independent force shadow at +5m") + } + }) + + t.Run("concurrent-load-persist-and-get-has-no-data-race", func(t *testing.T) { + now := time.Now() + loadStore := &blockingSelectiveAuthStore{ + onList: func() ([]*Auth, error) { + return []*Auth{{ + ID: "load-race-auth", + Provider: "xai", + Status: StatusError, + Unavailable: true, + NextRetryAfter: now.Add(5 * time.Minute), + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "force"}, + ModelStates: map[string]*ModelState{ + "grok-4": { + Unavailable: true, + Status: StatusError, + NextRetryAfter: now.Add(5 * time.Minute), + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "force"}, + UpdatedAt: now, + }, + }, + }}, nil + }, + } + cdsStore := &recordingCooldownStateStore{} + mgr := NewManager(loadStore, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(cdsStore) + if errLoad := mgr.Load(ctx); errLoad != nil { + t.Fatalf("Load: %v", errLoad) + } + start := make(chan struct{}) + var wg sync.WaitGroup + for range 8 { + wg.Add(3) + go func() { + defer wg.Done() + <-start + _ = mgr.Load(ctx) + }() + go func() { + defer wg.Done() + <-start + mgr.PersistCooldownStates(ctx) + }() + go func() { + defer wg.Done() + <-start + _, _ = mgr.GetByID("load-race-auth") + }() + } + close(start) + wg.Wait() + }) + + t.Run("delayed-forget-after-remove-does-not-wipe-reregistered-auth", func(t *testing.T) { + const ( + providerKey = "openai-compatible-issue-7431-forget-race" + model = "issue-7431-forget-race-model" + ) + sel := &WeightedRoundRobinSelector{} + aff := NewSessionAffinitySelector(sel) + t.Cleanup(aff.Stop) + mgr := NewManager(nil, aff, nil) + mgr.RegisterExecutor(&mockCustomErrorExecutor{identifier: providerKey}) + reg := registry.GetGlobalRegistry() + + for i := 0; i < 1030; i++ { + id := fmt.Sprintf("fr-auth-%04d", i) + reg.RegisterClient(id, providerKey, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: id, Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}, + }); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + } + + // Pick once so fr-auth-0000 (epoch 1) enters largeStates. + picked1, _, errPick1 := mgr.pickNext(ctx, providerKey, model, cliproxyexecutor.Options{}, nil) + if errPick1 != nil || picked1 == nil || picked1.ID != "fr-auth-0000" { + t.Fatalf("first pick = %v (err=%v), want fr-auth-0000", picked1, errPick1) + } + + // Remove fr-auth-0000, then immediately re-register fr-auth-0000 (higher RegistrationEpoch) + // and select again so the new instance is active in largeStates with credit. + if !mgr.RemoveIf(ctx, "fr-auth-0000", nil) { + t.Fatal("RemoveIf(fr-auth-0000) returned false") + } + reRegistered, errReReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "fr-auth-0000", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}, + }) + if errReReg != nil { + t.Fatalf("re-register fr-auth-0000: %v", errReReg) + } + // Pick on the pool so re-registered fr-auth-0000 (epoch >= 3) is tracked in largeStates. + if _, _, errPick2 := mgr.pickNext(ctx, providerKey, model, cliproxyexecutor.Options{}, nil); errPick2 != nil { + t.Fatalf("second pick failed: %v", errPick2) + } + + // Simulate a delayed ForgetAuth from the prior RemoveIf (whose tombstoneEpoch < reRegistered.RegistrationEpoch). + staleTombstoneEpoch := reRegistered.RegistrationEpoch - 1 + mgr.forgetWeightedAuthUpTo("fr-auth-0000", staleTombstoneEpoch, ^uint64(0)) + + sel.mu.Lock() + var foundNew bool + for _, state := range sel.largeStates { + if idx, ok := state.findMember("fr-auth-0000"); ok && state.members[idx].current != 0 { + foundNew = true + } + } + sel.mu.Unlock() + if !foundNew { + t.Fatal("delayed forget from earlier RemoveIf wiped re-registered fr-auth-0000 state") + } + }) + + t.Run("ordinary-credential-scope-error-does-not-clear-active-model-force-cooldown", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-cred-ordinary" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + cdsDir := t.TempDir() + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsDir)) + mgr.SetConfig(cfg) + + a := &Auth{ + ID: "cred-ordinary-over-model-force", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + } + if _, errReg := mgr.Register(WithSkipPersist(ctx), a); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + twoMin := 2 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: a.ID, + Provider: providerKey, + Model: "model-1", + RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "force 503"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: a.ID, + Provider: providerKey, + Model: "model-1", + CredentialScope: true, + Error: &Error{Code: credentialQuotaErrorCode, HTTPStatus: http.StatusTooManyRequests, Message: "ordinary cred 429"}, + }) + + now := time.Now() + gotImmediate, _ := mgr.GetByID(a.ID) + if IsAuthAvailableForModel(gotImmediate, "model-1", now) { + t.Fatalf("ordinary CredentialScope 429 erased active model force cooldown immediately: state=%+v", gotImmediate.ModelStates["model-1"]) + } + mgr.SetConfig(cfg) + gotAfterReload, _ := mgr.GetByID(a.ID) + if IsAuthAvailableForModel(gotAfterReload, "model-1", now) { + t.Fatalf("SetConfig after ordinary CredentialScope 429 erased active model force cooldown: state=%+v", gotAfterReload.ModelStates["model-1"]) + } + mgr.PersistCooldownStates(ctx) + mgrRestored := NewManager(nil, &RoundRobinSelector{}, nil) + mgrRestored.SetCooldownStateStore(NewFileCooldownStateStore(cdsDir)) + mgrRestored.SetConfig(cfg) + if _, errReg := mgrRestored.Register(WithSkipPersist(ctx), &Auth{ + ID: a.ID, Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register restored: %v", errReg) + } + if errRestore := mgrRestored.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + gotRestored, _ := mgrRestored.GetByID(a.ID) + if IsAuthAvailableForModel(gotRestored, "model-1", now) { + t.Fatalf("RestoreCooldownStates lost active model force cooldown after ordinary CredentialScope 429: state=%+v", gotRestored.ModelStates["model-1"]) + } + }) + + t.Run("mixed-transport-large-websocket-pool-keeps-credits-when-http-peer-excluded", func(t *testing.T) { + const model = "gpt-5-ws-mixed" + wsCtx := cliproxyexecutor.WithDownstreamWebsocket(ctx) + + // 1. Direct selector / legacy path: excluding the sole HTTP peer in tried must not switch dynamic pool keys. + sel := &WeightedRoundRobinSelector{} + all := make([]*Auth, 0, 1031) + wsOnly := make([]*Auth, 0, 1030) + httpAuth := &Auth{ + ID: "codex-http-0000", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "weight": "1", + "websockets": "false", + }, + } + all = append(all, httpAuth) + for i := 0; i < 1030; i++ { + a := &Auth{ + ID: fmt.Sprintf("codex-ws-%04d", i), + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "weight": "1", + "websockets": "true", + }, + } + all = append(all, a) + wsOnly = append(wsOnly, a) + } + + for i := 0; i < 3; i++ { + picked, errPick := sel.Pick(wsCtx, "codex", model, cliproxyexecutor.Options{}, all) + if errPick != nil || picked == nil || picked.ID != fmt.Sprintf("codex-ws-%04d", i) { + t.Fatalf("selector pick %d with HTTP peer present = %v (err=%v), want codex-ws-%04d", i, picked, errPick, i) + } + } + pickedWithoutHTTP, errPick := sel.Pick(wsCtx, "codex", model, cliproxyexecutor.Options{}, wsOnly) + if errPick != nil || pickedWithoutHTTP == nil || pickedWithoutHTTP.ID != "codex-ws-0003" { + t.Fatalf("selector pick after HTTP peer excluded = %v (err=%v), want codex-ws-0003", pickedWithoutHTTP, errPick) + } + + // 2. Scheduler fast path: cooling the sole HTTP peer must not switch dynamic pool keys. + schedSel := &WeightedRoundRobinSelector{} + mgr := NewManager(nil, schedSel, nil) + mgr.RegisterExecutor(&mockCustomErrorExecutor{identifier: "codex"}) + reg := registry.GetGlobalRegistry() + for _, a := range all { + id := a.ID + reg.RegisterClient(id, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + if _, errReg := mgr.Register(WithSkipPersist(ctx), a.Clone()); errReg != nil { + t.Fatalf("Register %s: %v", id, errReg) + } + } + for i := 0; i < 3; i++ { + picked, _, errPickSched := mgr.pickNext(wsCtx, "codex", model, cliproxyexecutor.Options{}, nil) + if errPickSched != nil || picked == nil || picked.ID != fmt.Sprintf("codex-ws-%04d", i) { + t.Fatalf("scheduler pick %d before HTTP cooldown = %v (err=%v), want codex-ws-%04d", i, picked, errPickSched, i) + } + } + oneMin := time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: httpAuth.ID, + Provider: "codex", + Model: model, + RetryAfter: &oneMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "http down"}, + }) + pickedAfterCool, _, errPickAfter := mgr.pickNext(wsCtx, "codex", model, cliproxyexecutor.Options{}, nil) + if errPickAfter != nil || pickedAfterCool == nil || pickedAfterCool.ID != "codex-ws-0003" { + t.Fatalf("scheduler pick after HTTP peer cooled = %v (err=%v), want codex-ws-0003", pickedAfterCool, errPickAfter) + } + }) + + t.Run("small-exact-pool-isolates-inflight-reservations-across-register-replacement-and-load", func(t *testing.T) { + const ( + providerKey = "openai-compatible-issue-7431-small-epoch" + route = providerKey + ":small-model" + ) + loadItems := []*Auth{ + {ID: "small-0", Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}}, + {ID: "small-1", Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}}, + {ID: "small-2", Provider: providerKey, Status: StatusActive, Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}}, + } + store := &schedulerLoadStore{auths: []*Auth{loadItems[0].Clone(), loadItems[1].Clone(), loadItems[2].Clone()}} + sel := &WeightedRoundRobinSelector{} + mgr := NewManager(store, sel, nil) + if errLoad := mgr.Load(ctx); errLoad != nil { + t.Fatalf("initial Load: %v", errLoad) + } + + currentPool := func() []*Auth { + out := make([]*Auth, 0, 3) + for _, id := range []string{"small-0", "small-1", "small-2"} { + a, _ := mgr.GetByID(id) + out = append(out, a) + } + return out + } + + // 1. In-place Register replacement (without RemoveIf) advances small-0's RegistrationEpoch. + pickedStale, staleRegRes := sel.reserve(route, currentPool(), "") + if pickedStale == nil || pickedStale.ID != "small-0" { + t.Fatalf("initial pick = %v, want small-0", pickedStale) + } + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "small-0", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey, "weight": "1"}, + }); errReg != nil { + t.Fatalf("in-place Register replacement: %v", errReg) + } + // New small-0 (epoch 2) must not inherit staleRegRes's -2 pending debt, and a delayed + // commit of staleRegRes must be ignored rather than charging the new small-0. + pickedFresh, freshRegRes := sel.reserve(route, currentPool(), "") + if pickedFresh == nil || pickedFresh.ID != "small-0" { + t.Fatalf("pick after in-place Register replacement = %v, want fresh small-0 (not suppressed by stale pending)", pickedFresh) + } + freshRegRes.finish(false) + staleRegRes.finish(true) + + sel.mu.Lock() + for _, head := range sel.states { + for st := head; st != nil; st = st.next { + for _, m := range st.members { + if m.pending != 0 || (m.id == "small-0" && m.current != 0) { + sel.mu.Unlock() + t.Fatalf("in-place Register replacement failed to isolate stale reservation: member %s current=%d pending=%d", m.id, m.current, m.pending) + } + } + } + } + sel.mu.Unlock() + + // 2. Manager.Load advances RegistrationEpoch across retained IDs while a reservation is in-flight. + _, staleLoadRes := sel.reserve(route, currentPool(), "") + store.auths = []*Auth{loadItems[0].Clone(), loadItems[1].Clone(), loadItems[2].Clone()} + if errLoad := mgr.Load(ctx); errLoad != nil { + t.Fatalf("second Load: %v", errLoad) + } + pickedAfterLoad, freshLoadRes := sel.reserve(route, currentPool(), "") + if pickedAfterLoad == nil || pickedAfterLoad.ID != "small-0" { + t.Fatalf("pick after Load = %v, want small-0", pickedAfterLoad) + } + freshLoadRes.finish(false) + staleLoadRes.finish(true) + + sel.mu.Lock() + for _, head := range sel.states { + for st := head; st != nil; st = st.next { + for _, m := range st.members { + if m.pending != 0 || m.current != 0 { + sel.mu.Unlock() + t.Fatalf("Load epoch bump failed to isolate stale reservation: member %s current=%d pending=%d", m.id, m.current, m.pending) + } + } + } + } + sel.mu.Unlock() + }) + + t.Run("consecutive-force-429-without-retry-after-preserves-escalated-backoff-across-reload", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-backoff" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetConfig(cfg) + + past := time.Now().Add(-time.Minute) + // Pre-seed BackoffLevel=3 (expired window) so next escalated quota backoff is 1s * 2^3 = 8s (>> 1s base). + authModel := &Auth{ + ID: "backoff-model-auth", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + ModelStates: map[string]*ModelState{ + "m1": { + Status: StatusActive, + Quota: QuotaState{BackoffLevel: 3, NextRecoverAt: past}, + }, + }, + } + authCred := &Auth{ + ID: "backoff-cred-auth", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", BackoffLevel: 3, NextRecoverAt: past}, + } + authGlobal := &Auth{ + ID: "backoff-global-auth", + Provider: providerKey, + Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + Quota: QuotaState{BackoffLevel: 3, NextRecoverAt: past}, + } + for _, a := range []*Auth{authModel, authCred, authGlobal} { + if _, errReg := mgr.Register(WithSkipPersist(ctx), a); errReg != nil { + t.Fatalf("Register %s: %v", a.ID, errReg) + } + } + // Restore BackoffLevel=3 in runtime after initial Register cleared expired QuotaState. + mgr.mu.Lock() + mgr.auths[authModel.ID].ModelStates["m1"].Quota.BackoffLevel = 3 + mgr.auths[authCred.ID].Quota.Exceeded = true + mgr.auths[authCred.ID].Quota.Reason = "credential_quota" + mgr.auths[authCred.ID].Quota.BackoffLevel = 3 + mgr.auths[authGlobal.ID].Quota.BackoffLevel = 3 + mgr.mu.Unlock() + + beforeMark := time.Now() + mgr.MarkResult(ctx, Result{ + AuthID: authModel.ID, Provider: providerKey, Model: "m1", + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force 429 no hint"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: authCred.ID, Provider: providerKey, Model: "m1", CredentialScope: true, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force cred 429 no hint"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: authGlobal.ID, Provider: providerKey, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force global 429 no hint"}, + }) + + mgr.SetConfig(cfg) + probeAt := beforeMark.Add(4 * time.Second) // > 1s base floor, < 8s escalated backoff + gotModel, _ := mgr.GetByID(authModel.ID) + if IsAuthAvailableForModel(gotModel, "m1", probeAt) { + t.Fatalf("SetConfig truncated escalated model force 429 backoff to <=4s (want ~8s): state=%+v", gotModel.ModelStates["m1"]) + } + gotCred, _ := mgr.GetByID(authCred.ID) + if IsAuthAvailableForModel(gotCred, "m1", probeAt) { + t.Fatalf("SetConfig truncated escalated credential_quota force 429 backoff to <=4s (want ~8s): quota=%+v", gotCred.Quota) + } + gotGlobal, _ := mgr.GetByID(authGlobal.ID) + if IsAuthAvailableForModel(gotGlobal, "", probeAt) { + t.Fatalf("SetConfig truncated escalated auth-level force 429 backoff to <=4s (want ~8s): NextRetryAfter=%v", gotGlobal.NextRetryAfter) + } + }) + + t.Run("http-bound-large-pool-keeps-credits-when-last-websocket-peer-excluded", func(t *testing.T) { + const model = "gpt-5-http-bound" + wsCtx := cliproxyexecutor.WithDownstreamWebsocket(ctx) + + // 1. Legacy SessionAffinitySelector path: HTTP-bound session charges 1030 HTTP peers; + // excluding the sole WS peer must keep charging the same "http" dynamic pool. + baseSel := &WeightedRoundRobinSelector{} + affSel := NewSessionAffinitySelector(baseSel) + t.Cleanup(affSel.Stop) + + all := make([]*Auth, 0, 1031) + httpOnly := make([]*Auth, 0, 1030) + wsPeer := &Auth{ + ID: "codex-ws-sole", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "weight": "1", + "websockets": "true", + }, + } + all = append(all, wsPeer) + for i := 0; i < 1030; i++ { + a := &Auth{ + ID: fmt.Sprintf("codex-http-%04d", i), + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "weight": "1", + "websockets": "false", + }, + } + all = append(all, a) + httpOnly = append(httpOnly, a) + } + + stickyOpts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"http-sticky-sess"}}} + // Bind session to codex-http-0000 on an HTTP request first, then send downstream WS requests. + if _, errBind := affSel.Pick(ctx, "codex", model, stickyOpts, httpOnly); errBind != nil { + t.Fatalf("bind HTTP session: %v", errBind) + } + if _, errPickWS1 := affSel.Pick(wsCtx, "codex", model, stickyOpts, all); errPickWS1 != nil { + t.Fatalf("WS pick with WS peer present: %v", errPickWS1) + } + if _, errPickWS2 := affSel.Pick(wsCtx, "codex", model, stickyOpts, httpOnly); errPickWS2 != nil { + t.Fatalf("WS pick after WS peer excluded: %v", errPickWS2) + } + httpBucketKey := dynamicWeightedPoolKey("codex:"+model, 0, "http") + baseSel.mu.Lock() + largeBucketCount := len(baseSel.largeStates) + httpBucket := baseSel.largeStates[httpBucketKey] + var httpBoundDebt int64 + if httpBucket != nil { + if idx, ok := httpBucket.findMember("codex-http-0000"); ok { + httpBoundDebt = httpBucket.members[idx].current + } + } + baseSel.mu.Unlock() + if largeBucketCount != 2 || httpBoundDebt != -2*1029 { + t.Fatalf("expected HTTP-bound WS traffic to accumulate both picks in 'http' bucket (debt=%d, want %d; buckets=%d, want 2)", httpBoundDebt, -2*1029, largeBucketCount) + } + + // 2. Scheduler fast path: pinned HTTP auth on downstream WS keeps using the same "http" dynamic pool after sole WS peer cools. + schedSel := &WeightedRoundRobinSelector{} + mgr := NewManager(nil, schedSel, nil) + mgr.RegisterExecutor(&mockCustomErrorExecutor{identifier: "codex"}) + reg := registry.GetGlobalRegistry() + for _, a := range all { + id := a.ID + reg.RegisterClient(id, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + if _, errReg := mgr.Register(WithSkipPersist(ctx), a.Clone()); errReg != nil { + t.Fatalf("Register %s: %v", id, errReg) + } + } + pinOpts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: "codex-http-0000"}} + if _, _, errP1 := mgr.pickNext(wsCtx, "codex", model, pinOpts, nil); errP1 != nil { + t.Fatalf("scheduler pinned pick before WS cooldown: %v", errP1) + } + oneMin := time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: wsPeer.ID, Provider: "codex", Model: model, RetryAfter: &oneMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "ws down"}, + }) + if _, _, errP2 := mgr.pickNext(wsCtx, "codex", model, pinOpts, nil); errP2 != nil { + t.Fatalf("scheduler pinned pick after WS cooldown: %v", errP2) + } + schedSel.mu.Lock() + schedBucketCount := len(schedSel.largeStates) + schedSel.mu.Unlock() + if schedBucketCount != 1 { + t.Fatalf("scheduler switched dynamic buckets when sole WS peer cooled: got %d buckets, want 1", schedBucketCount) + } + }) + + t.Run("delayed-store-restore-does-not-shorten-newer-live-auth-force-cooldown", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-delayed-restore" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + baseNow := time.Now() + staleRetry := baseNow.Add(2 * time.Minute) + loadEntered := make(chan struct{}) + loadRelease := make(chan struct{}) + slowStore := &issue7431BlockingLoadCooldownStore{ + records: []CooldownStateRecord{{ + Provider: providerKey, + AuthID: "delayed-auth-0", + Status: "cooling", + NextRetryAfter: staleRetry, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "older 2m force"}, + UpdatedAt: baseNow.Add(-time.Minute), + ForceCooldownUntil: staleRetry, + ForceCooldownEvaluated: true, + }}, + entered: loadEntered, + release: loadRelease, + } + + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(slowStore) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "delayed-auth-0", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + restoreDone := make(chan error, 1) + go func() { + restoreDone <- mgr.RestoreCooldownStates(ctx) + }() + <-loadEntered + + // Record a newer 10-minute auth-level force cooldown while store.Load is in flight. + tenMin := 10 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: "delayed-auth-0", Provider: providerKey, RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "newer 10m force"}, + }) + close(loadRelease) + if errRestore := <-restoreDone; errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + got, _ := mgr.GetByID("delayed-auth-0") + if IsAuthAvailableForModel(got, "", baseNow.Add(5*time.Minute)) { + t.Fatalf("delayed RestoreCooldownStates shortened newer 10m auth-level force cooldown to %v", got.NextRetryAfter) + } + }) + + t.Run("dynamic-pool-prunes-churned-members-and-shrinks-backing-array-despite-inflight-reservation", func(t *testing.T) { + sel := &WeightedRoundRobinSelector{} + const route = "openai-compatible-churn:model-c" + makeBatch := func(batch int) []*Auth { + out := make([]*Auth, 1030) + for i := range out { + out[i] = &Auth{ + ID: fmt.Sprintf("churn-b%02d-%04d", batch, i), + Provider: "openai-compatible-churn", + Status: StatusActive, + Attributes: map[string]string{"weight": "1"}, + } + } + return out + } + + _, pinnedInFlight := sel.reserve(route, makeBatch(0), "") + if pinnedInFlight == nil { + t.Fatal("expected in-flight reservation on batch 0") + } + for b := 1; b <= 12; b++ { + _, res := sel.reserve(route, makeBatch(b), "") + res.finish(true) + } + + sel.mu.Lock() + state := sel.largeStates[route] + gotLen := len(state.members) + gotCap := cap(state.members) + sel.mu.Unlock() + if gotLen > 3200 || gotCap > 4096 { + t.Fatalf("dynamic pool members grew unbounded while 1 reservation remained in-flight: len=%d cap=%d", gotLen, gotCap) + } + + pinnedInFlight.finish(false) + sel.mu.Lock() + for _, m := range state.members { + if m.pending != 0 { + sel.mu.Unlock() + t.Fatalf("member %s has non-zero pending=%d after pinned reservation rollback", m.id, m.pending) + } + } + sel.mu.Unlock() + }) + + t.Run("ordinary-modelless-429-does-not-wipe-credential-scoped-force-cooldown", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-modelless-429" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + cdsPath := t.TempDir() + "/modelless-429.cds" + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "ml429-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + baseNow := time.Now() + fiveMin := 5 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: "ml429-auth", Provider: providerKey, Model: "m1", CredentialScope: true, RetryAfter: &fiveMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "force cred 429"}, + }) + // Ordinary model-less 429 under disable-cooling must not clobber credential_quota force state. + mgr.MarkResult(ctx, Result{ + AuthID: "ml429-auth", Provider: providerKey, + Error: &Error{Code: "rate_limited", HTTPStatus: http.StatusTooManyRequests, Message: "ordinary global 429"}, + }) + + gotAfterMark, _ := mgr.GetByID("ml429-auth") + if IsAuthAvailableForModel(gotAfterMark, "m1", baseNow.Add(2*time.Minute)) || IsAuthAvailableForModel(gotAfterMark, "", baseNow.Add(2*time.Minute)) { + t.Fatalf("ordinary model-less 429 cleared active credential-scoped force cooldown: auth=%+v", gotAfterMark) + } + + mgr.SetConfig(cfg) + gotAfterReload, _ := mgr.GetByID("ml429-auth") + if IsAuthAvailableForModel(gotAfterReload, "m1", baseNow.Add(2*time.Minute)) || IsAuthAvailableForModel(gotAfterReload, "", baseNow.Add(2*time.Minute)) { + t.Fatalf("SetConfig after ordinary model-less 429 cleared credential-scoped force cooldown: auth=%+v", gotAfterReload) + } + + mgr.PersistCooldownStates(ctx) + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr2.SetConfig(cfg) + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: "ml429-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register mgr2: %v", errReg) + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + gotRestored, _ := mgr2.GetByID("ml429-auth") + if IsAuthAvailableForModel(gotRestored, "m1", baseNow.Add(2*time.Minute)) || IsAuthAvailableForModel(gotRestored, "", baseNow.Add(2*time.Minute)) { + t.Fatalf("restored auth lost credential-scoped force cooldown after ordinary model-less 429: auth=%+v", gotRestored) + } + }) + + t.Run("coexisting-credential-and-auth-force-cooldowns-preserve-longer-deadline-across-reload-and-disk", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-dual-force" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + cdsPath := t.TempDir() + "/dual-force.cds" + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr.SetConfig(cfg) + for _, id := range []string{"dual-short-cred-long-auth", "dual-long-cred-short-auth"} { + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: id, Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register %s: %v", id, errReg) + } + } + + baseNow := time.Now() + twoMin := 2 * time.Minute + tenMin := 10 * time.Minute + + // Order A: 2m credential-scoped force 429, then 10m model-less force 503. + mgr.MarkResult(ctx, Result{ + AuthID: "dual-short-cred-long-auth", Provider: providerKey, Model: "m1", CredentialScope: true, RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "2m cred force"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: "dual-short-cred-long-auth", Provider: providerKey, RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "10m auth force"}, + }) + + // Order B: 10m credential-scoped force 429, then 2m model-less force 429. + mgr.MarkResult(ctx, Result{ + AuthID: "dual-long-cred-short-auth", Provider: providerKey, Model: "m1", CredentialScope: true, RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "10m cred force"}, + }) + mgr.MarkResult(ctx, Result{ + AuthID: "dual-long-cred-short-auth", Provider: providerKey, RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "2m auth force"}, + }) + + mgr.SetConfig(cfg) + for _, id := range []string{"dual-short-cred-long-auth", "dual-long-cred-short-auth"} { + got, _ := mgr.GetByID(id) + if IsAuthAvailableForModel(got, "", baseNow.Add(5*time.Minute)) { + t.Fatalf("%s truncated longer 10m force deadline on SetConfig: NextRetryAfter=%v", id, got.NextRetryAfter) + } + } + if gotLongCred, _ := mgr.GetByID("dual-long-cred-short-auth"); IsAuthAvailableForModel(gotLongCred, "m1", baseNow.Add(5*time.Minute)) { + t.Fatalf("dual-long-cred-short-auth lost 10m model wrap on SetConfig: %+v", gotLongCred) + } + + mgr.PersistCooldownStates(ctx) + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr2.SetConfig(cfg) + for _, id := range []string{"dual-short-cred-long-auth", "dual-long-cred-short-auth"} { + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: id, Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register mgr2 %s: %v", id, errReg) + } + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + mgr2.SetConfig(cfg) + for _, id := range []string{"dual-short-cred-long-auth", "dual-long-cred-short-auth"} { + got, _ := mgr2.GetByID(id) + if IsAuthAvailableForModel(got, "", baseNow.Add(5*time.Minute)) { + t.Fatalf("%s lost longer 10m force deadline across disk persistence + SetConfig: NextRetryAfter=%v", id, got.NextRetryAfter) + } + } + if gotLongCred2, _ := mgr2.GetByID("dual-long-cred-short-auth"); IsAuthAvailableForModel(gotLongCred2, "m1", baseNow.Add(5*time.Minute)) { + t.Fatalf("dual-long-cred-short-auth lost 10m model wrap across disk persistence + SetConfig: %+v", gotLongCred2) + } + }) + + t.Run("concurrent-restore-during-update-save-preserves-force-provenance-across-reload", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-update-save" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + saveEntered := make(chan struct{}) + saveRelease := make(chan struct{}) + var saveOnce sync.Once + authStore := &blockingSelectiveAuthStore{ + onSave: func(a *Auth) { + if a != nil && a.ID == "save-race-auth" { + saveOnce.Do(func() { + close(saveEntered) + <-saveRelease + }) + } + }, + } + baseNow := time.Now() + tenMinUntil := baseNow.Add(10 * time.Minute) + cdsStore := &recordingCooldownStateStore{ + load: []CooldownStateRecord{{ + Provider: providerKey, + AuthID: "save-race-auth", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "restored 10m force"}, + UpdatedAt: baseNow, + ForceCooldownUntil: tenMinUntil, + ForceCooldownEvaluated: true, + }}, + } + + mgr := NewManager(authStore, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(cdsStore) + mgr.SetConfig(cfg) + initial, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "save-race-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + Metadata: map[string]any{"rev": "1"}, + }) + if errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + upd := initial.Clone() + upd.Metadata["rev"] = "2" + updateDone := make(chan error, 1) + go func() { + _, errUpd := mgr.Update(ctx, upd) + updateDone <- errUpd + }() + <-saveEntered + + if errRestore := mgr.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates during Update save: %v", errRestore) + } + close(saveRelease) + if errUpd := <-updateDone; errUpd != nil { + t.Fatalf("Update: %v", errUpd) + } + + // Reload config: restored 10m force provenance must have survived mergeAuthSaveDelta. + mgr.SetConfig(cfg) + got, _ := mgr.GetByID("save-race-auth") + if IsAuthAvailableForModel(got, "", baseNow.Add(5*time.Minute)) { + t.Fatalf("RestoreCooldownStates force provenance lost during concurrent Update Store.Save + SetConfig: %+v", got) + } + }) + + t.Run("modelless-force-cooldown-survives-ordinary-model-error-reload-and-restore", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-agg-force" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + cdsPath := t.TempDir() + "/agg-force.cds" + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "agg-force-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + baseNow := time.Now() + tenMin := 10 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: "agg-force-auth", Provider: providerKey, RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "10m auth force"}, + }) + gotBeforeModelErr, _ := mgr.GetByID("agg-force-auth") + wantDeadline := gotBeforeModelErr.NextRetryAfter + + // Subsequent ordinary model-level 500 under disable-cooling must not clear auth-level force cooldown. + mgr.MarkResult(ctx, Result{ + AuthID: "agg-force-auth", Provider: providerKey, Model: "m1", + Error: &Error{Code: "upstream_500", HTTPStatus: http.StatusInternalServerError, Message: "ordinary 500"}, + }) + gotAfterModelErr, _ := mgr.GetByID("agg-force-auth") + if IsAuthAvailableForModel(gotAfterModelErr, "", baseNow.Add(5*time.Minute)) || + IsAuthAvailableForModel(gotAfterModelErr, "m1", baseNow.Add(5*time.Minute)) || + !gotAfterModelErr.NextRetryAfter.Equal(wantDeadline) { + t.Fatalf("ordinary model 500 wiped/altered 10m auth-level force cooldown: got=%v want=%v unavailable=%v", gotAfterModelErr.NextRetryAfter, wantDeadline, gotAfterModelErr.Unavailable) + } + + mgr.SetConfig(cfg) + gotAfterReload, _ := mgr.GetByID("agg-force-auth") + if IsAuthAvailableForModel(gotAfterReload, "", baseNow.Add(5*time.Minute)) || + IsAuthAvailableForModel(gotAfterReload, "m1", baseNow.Add(5*time.Minute)) || + !gotAfterReload.NextRetryAfter.Equal(wantDeadline) { + t.Fatalf("SetConfig after ordinary model 500 lost 10m auth-level force cooldown: got=%v want=%v", gotAfterReload.NextRetryAfter, wantDeadline) + } + + mgr.PersistCooldownStates(ctx) + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr2.SetConfig(cfg) + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: "agg-force-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register mgr2: %v", errReg) + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + gotRestored, _ := mgr2.GetByID("agg-force-auth") + if IsAuthAvailableForModel(gotRestored, "", baseNow.Add(5*time.Minute)) || + IsAuthAvailableForModel(gotRestored, "m1", baseNow.Add(5*time.Minute)) || + !gotRestored.NextRetryAfter.Equal(wantDeadline) { + t.Fatalf("restored auth lost 10m auth-level force cooldown: got=%v want=%v", gotRestored.NextRetryAfter, wantDeadline) + } + }) + + t.Run("forget-auth-shrinks-large-dynamic-pool-backing-array-with-inflight-reservation", func(t *testing.T) { + sel := &WeightedRoundRobinSelector{} + const route = "openai-compatible-shrink:model-s" + hugePool := make([]*Auth, 5000) + for i := range hugePool { + hugePool[i] = &Auth{ + ID: fmt.Sprintf("shrink-auth-%04d", i), + Provider: "openai-compatible-shrink", + Status: StatusActive, + Attributes: map[string]string{"weight": "1"}, + } + } + + // Hold an in-flight reservation on the retained 1030 prefix while seeding 5000 members. + _, seedRes := sel.reserve(route, hugePool, "") + seedRes.finish(true) + _, inFlightRes := sel.reserve(route, hugePool[:1030], "") + if inFlightRes == nil { + t.Fatal("expected in-flight reservation on retained 1030 subset") + } + + // Permanently remove 3970 credentials so 1030 (> 1024) remain. + for i := 1030; i < len(hugePool); i++ { + sel.ForgetAuth(hugePool[i].ID) + } + + sel.mu.Lock() + state := sel.largeStates[route] + gotLen := len(state.members) + gotCap := cap(state.members) + sel.mu.Unlock() + if gotLen != 1030 || gotCap > 2060 { + t.Fatalf("ForgetAuth left oversized backing array after shrinking 5000->1030: len=%d cap=%d (want len=1030, cap<=2060)", gotLen, gotCap) + } + + inFlightRes.finish(false) + sel.mu.Lock() + for _, m := range state.members { + if m.pending != 0 || m.inFlight != 0 { + sel.mu.Unlock() + t.Fatalf("member %s has leaked pending=%d inFlight=%d after rollback", m.id, m.pending, m.inFlight) + } + } + sel.mu.Unlock() + }) + + t.Run("ordinary-credential-quota-plus-longer-auth-force-preserves-auth-force-across-disk-restore", func(t *testing.T) { + cdsPath := t.TempDir() + "/ord-cred-long-force.cds" + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "ord-cred-force-auth", Provider: "xai", Status: StatusActive, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + baseNow := time.Now() + twoMin := 2 * time.Minute + tenMin := 10 * time.Minute + // 1. Ordinary 2m credential-scoped 429 on m1 (CredQuotaForceCooldownUntil == zero). + mgr.MarkResult(ctx, Result{ + AuthID: "ord-cred-force-auth", Provider: "xai", Model: "m1", CredentialScope: true, RetryAfter: &twoMin, + Error: &Error{Code: "rate_limited", HTTPStatus: http.StatusTooManyRequests, Message: "ordinary 2m cred quota"}, + }) + // 2. Independent 10m model-less force 503 (ForceCooldownUntil == +10m). + mgr.MarkResult(ctx, Result{ + AuthID: "ord-cred-force-auth", Provider: "xai", RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "10m auth force"}, + }) + mgr.PersistCooldownStates(ctx) + + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: "ord-cred-force-auth", Provider: "xai", Status: StatusActive, + }); errReg != nil { + t.Fatalf("Register mgr2: %v", errReg) + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + // At +3m: 2m ordinary credential_quota has expired, but 10m independent auth force must still block m1 and "". + got, _ := mgr2.GetByID("ord-cred-force-auth") + if IsAuthAvailableForModel(got, "m1", baseNow.Add(3*time.Minute)) || IsAuthAvailableForModel(got, "", baseNow.Add(3*time.Minute)) { + t.Fatalf("restored auth lost 10m independent auth force after 2m ordinary credential_quota expired: %+v", got) + } + }) + + t.Run("devin-quota-restoration-preserves-independent-modelless-force-cooldown", func(t *testing.T) { + now := time.Now() + existing := &Auth{ + ID: "devin-dual-auth", + Provider: "devin", + Status: StatusError, + Unavailable: true, + UpdatedAt: now, + NextRetryAfter: now.Add(10 * time.Minute), + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "10m force 503"}, + StatusMessage: "transient upstream error", + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: now.Add(2 * time.Minute), + ObservedAt: now.Add(-time.Minute), + }, + forceCooldownUntil: now.Add(10 * time.Minute), + credQuotaForceUntil: now.Add(2 * time.Minute), + forceCooldownSeeded: true, + ModelStates: map[string]*ModelState{ + "devin-m1": { + Status: StatusError, + Unavailable: true, + NextRetryAfter: now.Add(2 * time.Minute), + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(2 * time.Minute)}, + credentialQuotaPropagated: true, + }, + }, + } + refreshed := existing.Clone() + refreshed.Quota.ObservedAt = now.Add(time.Second) + refreshed.Quota.Signals = map[string]string{ + "daily_quota_remaining_percent": "80", + "weekly_quota_remaining_percent": "90", + } + + evalNow := now.Add(2 * time.Second) + if !applyDevinQuotaRestoration(refreshed, existing, evalNow) { + t.Fatal("expected applyDevinQuotaRestoration to return true") + } + if IsAuthAvailableForModel(refreshed, "devin-m1", now.Add(5*time.Minute)) || + IsAuthAvailableForModel(refreshed, "", now.Add(5*time.Minute)) { + t.Fatalf("applyDevinQuotaRestoration wiped unrelated 10m independent auth force cooldown: %+v", refreshed) + } + }) + + t.Run("long-ordinary-credential-cooldown-plus-short-auth-force-clears-ordinary-tail-on-disable-cooling-reload", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-long-ord-short-force" + disableCoolingFalse := false + disableCoolingTrue := true + cfgCoolOn := &internalconfig.Config{ + DisableCooling: false, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingFalse, + }}, + } + cfgCoolOff := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + + cdsPath := t.TempDir() + "/long-ord-short-force.cds" + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr.SetConfig(cfgCoolOn) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "reverse-dual-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + baseNow := time.Now() + twentyMin := 20 * time.Minute + twoMin := 2 * time.Minute + + // 1. Ordinary 20m credential-scoped 429 while cooling is enabled. + mgr.MarkResult(ctx, Result{ + AuthID: "reverse-dual-auth", Provider: providerKey, Model: "m1", CredentialScope: true, RetryAfter: &twentyMin, + Error: &Error{Code: "rate_limited", HTTPStatus: http.StatusTooManyRequests, Message: "ordinary 20m cred 429"}, + }) + // 2. Independent 2m model-less force 503. + mgr.MarkResult(ctx, Result{ + AuthID: "reverse-dual-auth", Provider: providerKey, RetryAfter: &twoMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusServiceUnavailable, Message: "2m auth force 503"}, + }) + mgr.PersistCooldownStates(ctx) + + // 3. Reload with disable-cooling: true: 2m force must stay active at +1m, 20m ordinary tail must clear by +3m. + mgr.SetConfig(cfgCoolOff) + gotReload, _ := mgr.GetByID("reverse-dual-auth") + if IsAuthAvailableForModel(gotReload, "m1", baseNow.Add(time.Minute)) || IsAuthAvailableForModel(gotReload, "", baseNow.Add(time.Minute)) { + t.Fatalf("2m auth force cooldown prematurely cleared at +1m: %+v", gotReload) + } + if !IsAuthAvailableForModel(gotReload, "m1", baseNow.Add(3*time.Minute)) || !IsAuthAvailableForModel(gotReload, "", baseNow.Add(3*time.Minute)) { + t.Fatalf("20m ordinary credential cooldown survived disable-cooling reload past 2m force window (+3m): NextRetryAfter=%v Quota=%+v", gotReload.NextRetryAfter, gotReload.Quota) + } + + // 4. Also verify restoring pre-reload .cds directly into a disable-cooling: true manager clamps to 2m. + mgr2 := NewManager(nil, &RoundRobinSelector{}, nil) + mgr2.SetCooldownStateStore(NewFileCooldownStateStore(cdsPath)) + mgr2.SetConfig(cfgCoolOff) + if _, errReg := mgr2.Register(WithSkipPersist(ctx), &Auth{ + ID: "reverse-dual-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register mgr2: %v", errReg) + } + if errRestore := mgr2.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + gotRestored, _ := mgr2.GetByID("reverse-dual-auth") + if IsAuthAvailableForModel(gotRestored, "m1", baseNow.Add(time.Minute)) || IsAuthAvailableForModel(gotRestored, "", baseNow.Add(time.Minute)) { + t.Fatalf("restored 2m auth force cooldown prematurely cleared at +1m: %+v", gotRestored) + } + if !IsAuthAvailableForModel(gotRestored, "m1", baseNow.Add(3*time.Minute)) || !IsAuthAvailableForModel(gotRestored, "", baseNow.Add(3*time.Minute)) { + t.Fatalf("restored 20m ordinary credential cooldown survived disable-cooling restore past 2m force window (+3m): NextRetryAfter=%v Quota=%+v", gotRestored.NextRetryAfter, gotRestored.Quota) + } + }) + + t.Run("delayed-store-restore-preserves-longer-persisted-independent-shadow-force-deadline", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-delayed-shadow" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + baseNow := time.Now() + twoMinUntil := baseNow.Add(2 * time.Minute) + tenMinUntil := baseNow.Add(10 * time.Minute) + + loadEntered := make(chan struct{}) + loadRelease := make(chan struct{}) + slowStore := &issue7431BlockingLoadCooldownStore{ + records: []CooldownStateRecord{ + { + Provider: providerKey, + AuthID: "delayed-shadow-auth", + Status: "cooling", + NextRetryAfter: twoMinUntil, + Reason: "credential_quota", + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: twoMinUntil}, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 429, Message: "2m cred force"}, + UpdatedAt: baseNow.Add(-time.Minute), + CredQuotaForceCooldownUntil: twoMinUntil, + ForceCooldownEvaluated: true, + }, + { + Provider: providerKey, + AuthID: "delayed-shadow-auth", + Model: "m1", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: "credential_quota", + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: twoMinUntil}, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 429, Message: "2m cred force"}, + UpdatedAt: baseNow.Add(-time.Minute), + CredentialQuotaPropagated: true, + IndependentState: &ModelState{ + Status: StatusError, + Unavailable: true, + NextRetryAfter: tenMinUntil, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "older 10m model force"}, + UpdatedAt: baseNow.Add(-2 * time.Minute), + }, + ForceCooldownUntil: twoMinUntil, + IndependentForceCooldownUntil: tenMinUntil, + ForceCooldownEvaluated: true, + }, + }, + entered: loadEntered, + release: loadRelease, + } + + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(slowStore) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "delayed-shadow-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + restoreDone := make(chan error, 1) + go func() { + restoreDone <- mgr.RestoreCooldownStates(ctx) + }() + <-loadEntered + + // Newer 3m model force arrives while store.Load is blocked. + threeMin := 3 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: "delayed-shadow-auth", Provider: providerKey, Model: "m1", RetryAfter: &threeMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "newer 3m model force"}, + }) + close(loadRelease) + if errRestore := <-restoreDone; errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + mgr.SetConfig(cfg) + got, _ := mgr.GetByID("delayed-shadow-auth") + if IsAuthAvailableForModel(got, "m1", baseNow.Add(5*time.Minute)) { + t.Fatalf("delayed RestoreCooldownStates truncated 10m persisted independent shadow force deadline to %v", got.ModelStates["m1"].NextRetryAfter) + } + + recStore := &recordingCooldownStateStore{} + mgr.SetCooldownStateStore(recStore) + mgr.PersistCooldownStates(ctx) + foundTenMin := false + for _, rec := range recStore.savedRecords() { + if rec.Model == "m1" && rec.NextRetryAfter.Equal(tenMinUntil) && rec.IndependentForceCooldownUntil.Equal(tenMinUntil) { + foundTenMin = true + } + } + if !foundTenMin { + t.Fatalf("re-persisted records lost 10m independent force deadline: %+v", recStore.savedRecords()) + } + }) + + t.Run("delayed-restore-with-stale-model-and-auth-records-preserves-newer-credential-force-cooldown", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-stale-model-auth" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + baseNow := time.Now() + twoMinUntil := baseNow.Add(2 * time.Minute) + threeMinUntil := baseNow.Add(3 * time.Minute) + + loadEntered := make(chan struct{}) + loadRelease := make(chan struct{}) + slowStore := &issue7431BlockingLoadCooldownStore{ + records: []CooldownStateRecord{ + { + Provider: providerKey, + AuthID: "stale-combo-auth", + Model: "m-stale", + Status: "cooling", + NextRetryAfter: twoMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "older 2m model force"}, + UpdatedAt: baseNow.Add(-2 * time.Minute), + ForceCooldownUntil: twoMinUntil, + ForceCooldownEvaluated: true, + }, + { + Provider: providerKey, + AuthID: "stale-combo-auth", + Status: "cooling", + NextRetryAfter: threeMinUntil, + Reason: "quota", + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: threeMinUntil}, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 429, Message: "older 3m auth force 429"}, + UpdatedAt: baseNow.Add(-time.Minute), + ForceCooldownUntil: threeMinUntil, + ForceCooldownEvaluated: true, + }, + }, + entered: loadEntered, + release: loadRelease, + } + + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(slowStore) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "stale-combo-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + + restoreDone := make(chan error, 1) + go func() { + restoreDone <- mgr.RestoreCooldownStates(ctx) + }() + <-loadEntered + + // Newer 10m credential-scoped force 429 arrives while store.Load is blocked. + tenMin := 10 * time.Minute + mgr.MarkResult(ctx, Result{ + AuthID: "stale-combo-auth", Provider: providerKey, Model: "m-live", CredentialScope: true, RetryAfter: &tenMin, + Error: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: http.StatusTooManyRequests, Message: "newer 10m cred force 429"}, + }) + close(loadRelease) + if errRestore := <-restoreDone; errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + mgr.SetConfig(cfg) + got, _ := mgr.GetByID("stale-combo-auth") + if IsAuthAvailableForModel(got, "m-live", baseNow.Add(5*time.Minute)) || + IsAuthAvailableForModel(got, "m-other", baseNow.Add(5*time.Minute)) || + IsAuthAvailableForModel(got, "", baseNow.Add(5*time.Minute)) { + t.Fatalf("delayed RestoreCooldownStates shortened newer 10m credential-scoped force cooldown: NextRetryAfter=%v Quota=%+v", got.NextRetryAfter, got.Quota) + } + }) + + t.Run("cooldown-state-record-equal-detects-cred-quota-force-cooldown-until-change", func(t *testing.T) { + now := time.Now().Round(0) + recA := CooldownStateRecord{ + Provider: "xai", + AuthID: "eq-auth", + Status: "cooling", + NextRetryAfter: now.Add(10 * time.Minute), + Reason: "credential_quota", + Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(10 * time.Minute)}, + UpdatedAt: now, + ForceCooldownEvaluated: true, + CredQuotaForceCooldownUntil: time.Time{}, + } + recB := recA + recB.CredQuotaForceCooldownUntil = now.Add(10 * time.Minute) + if cooldownStateRecordEqual(recA, recB) { + t.Fatal("cooldownStateRecordEqual ignored CredQuotaForceCooldownUntil difference") + } + }) + + t.Run("legacy-dual-model-and-aggregated-auth-force-records-do-not-block-healthy-sibling-model", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-legacy-agg" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + baseNow := time.Now() + tenMinUntil := baseNow.Add(10 * time.Minute) + forceErr := &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "model-a force 503"} + + // Legacy store snapshot (ForceCooldownEvaluated == false): + // 1) legacy-dual-auth: model-a failed with force 503 AND old code persisted an aggregated Model=="" record. + // 2) legacy-global-auth: genuine model-less force 503 (only Model=="" record). + legacyStore := &recordingCooldownStateStore{ + load: []CooldownStateRecord{ + { + Provider: providerKey, + AuthID: "legacy-dual-auth", + Model: "model-a", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: cloneError(forceErr), + UpdatedAt: baseNow, + }, + { + Provider: providerKey, + AuthID: "legacy-dual-auth", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: cloneError(forceErr), + UpdatedAt: baseNow, + }, + { + Provider: providerKey, + AuthID: "legacy-global-auth", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "global force 503"}, + UpdatedAt: baseNow, + }, + }, + } + + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(legacyStore) + mgr.SetConfig(cfg) + for _, id := range []string{"legacy-dual-auth", "legacy-global-auth"} { + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: id, Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register %s: %v", id, errReg) + } + } + if errRestore := mgr.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + + gotDual, _ := mgr.GetByID("legacy-dual-auth") + if IsAuthAvailableForModel(gotDual, "model-a", baseNow.Add(2*time.Minute)) { + t.Fatal("expected legacy model-a force cooldown to stay blocked") + } + if !IsAuthAvailableForModel(gotDual, "model-b", baseNow.Add(2*time.Minute)) { + t.Fatalf("legacy aggregated auth record falsely promoted model-a force cooldown to block sibling model-b: %+v", gotDual) + } + // Subsequent success on model-a must clear the credential completely. + mgr.MarkResult(ctx, Result{AuthID: "legacy-dual-auth", Provider: providerKey, Model: "model-a", Success: true}) + gotDualAfterSuccess, _ := mgr.GetByID("legacy-dual-auth") + if !IsAuthAvailableForModel(gotDualAfterSuccess, "model-a", baseNow.Add(2*time.Minute)) || + !IsAuthAvailableForModel(gotDualAfterSuccess, "", baseNow.Add(2*time.Minute)) { + t.Fatalf("model-a success failed to clear legacy-dual-auth: %+v", gotDualAfterSuccess) + } + + // Genuine legacy model-less force record must still block all models. + gotGlobal, _ := mgr.GetByID("legacy-global-auth") + if IsAuthAvailableForModel(gotGlobal, "model-b", baseNow.Add(2*time.Minute)) || + IsAuthAvailableForModel(gotGlobal, "", baseNow.Add(2*time.Minute)) { + t.Fatalf("expected genuine legacy model-less force cooldown to remain blocked: %+v", gotGlobal) + } + }) + + t.Run("legacy-longer-model-force-plus-shorter-independent-auth-force-preserves-auth-block-for-sibling-models", func(t *testing.T) { + const providerKey = "openai-compatible-issue-7431-legacy-coexist" + disableCoolingTrue := true + cfg := &internalconfig.Config{ + DisableCooling: true, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: providerKey, + BaseURL: "https://example.invalid/v1", + DisableCooling: &disableCoolingTrue, + }}, + } + baseNow := time.Now() + thirtyMinUntil := baseNow.Add(30 * time.Minute) + tenMinUntil := baseNow.Add(10 * time.Minute) + + legacyStore := &recordingCooldownStateStore{ + load: []CooldownStateRecord{ + { + Provider: providerKey, + AuthID: "legacy-coexist-auth", + Model: "model-a", + Status: "cooling", + NextRetryAfter: thirtyMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "30m model-a force"}, + UpdatedAt: baseNow.Add(-2 * time.Minute), + }, + { + Provider: providerKey, + AuthID: "legacy-coexist-auth", + Status: "cooling", + NextRetryAfter: tenMinUntil, + Reason: ErrorCodeForceCooldown, + LastError: &Error{Code: ErrorCodeForceCooldown, HTTPStatus: 503, Message: "10m auth force"}, + UpdatedAt: baseNow.Add(-time.Minute), + }, + }, + } + + mgr := NewManager(nil, &RoundRobinSelector{}, nil) + mgr.SetCooldownStateStore(legacyStore) + mgr.SetConfig(cfg) + if _, errReg := mgr.Register(WithSkipPersist(ctx), &Auth{ + ID: "legacy-coexist-auth", Provider: providerKey, Status: StatusActive, + Attributes: map[string]string{"provider_key": providerKey}, + }); errReg != nil { + t.Fatalf("Register: %v", errReg) + } + if errRestore := mgr.RestoreCooldownStates(ctx); errRestore != nil { + t.Fatalf("RestoreCooldownStates: %v", errRestore) + } + mgr.SetConfig(cfg) + + got, _ := mgr.GetByID("legacy-coexist-auth") + if IsAuthAvailableForModel(got, "model-b", baseNow.Add(5*time.Minute)) { + t.Fatalf("expected 10m legacy independent auth force to block healthy model-b at +5m: %+v", got) + } + if !IsAuthAvailableForModel(got, "model-b", baseNow.Add(15*time.Minute)) { + t.Fatalf("expected model-b to recover once 10m auth force expired at +15m: %+v", got) + } + if IsAuthAvailableForModel(got, "model-a", baseNow.Add(15*time.Minute)) { + t.Fatalf("expected model-a 30m force cooldown to stay blocked at +15m: %+v", got.ModelStates["model-a"]) + } + }) +} + +type issue7431BlockingLoadCooldownStore struct { + records []CooldownStateRecord + entered chan struct{} + release chan struct{} + once sync.Once +} + +func (s *issue7431BlockingLoadCooldownStore) Load(context.Context) ([]CooldownStateRecord, error) { + s.once.Do(func() { + close(s.entered) + <-s.release + }) + return cloneCooldownStateRecords(s.records), nil +} + +func (s *issue7431BlockingLoadCooldownStore) Save(context.Context, []CooldownStateRecord) error { + return nil +} diff --git a/sdk/cliproxy/auth/quota_signals.go b/sdk/cliproxy/auth/quota_signals.go index fea368e5319..071158c6226 100644 --- a/sdk/cliproxy/auth/quota_signals.go +++ b/sdk/cliproxy/auth/quota_signals.go @@ -1,6 +1,7 @@ package auth import ( + "maps" "net/http" "sort" "strings" @@ -86,6 +87,32 @@ func applyCooldownFields(dst *QuotaState, cooldown QuotaState) { dst.Reason = cooldown.Reason dst.NextRecoverAt = cooldown.NextRecoverAt dst.BackoffLevel = cooldown.BackoffLevel + dst.cooldownSignals = maps.Clone(cooldown.cooldownSignals) + dst.cooldownEvaluated = cooldown.Exceeded || cooldown.cooldownEvaluated +} + +func rejectedClaudeSignalsForCooldown(provider string, headers http.Header) map[string]string { + if !strings.EqualFold(strings.TrimSpace(provider), "claude") { + return nil + } + signals := collectQuotaSignals("claude", headers) + if !hasRejectedClaudeUnifiedSignals(signals) { + return nil + } + return signals +} + +func sanitizePersistedClaudeCooldownSignals(raw map[string]string) map[string]string { + if len(raw) == 0 { + return nil + } + headers := make(http.Header, len(raw)) + for key, val := range raw { + if strings.HasPrefix(strings.ToLower(strings.TrimSpace(key)), "anthropic-ratelimit-unified-") { + headers.Set(key, val) + } + } + return rejectedClaudeSignalsForCooldown("claude", headers) } // collectQuotaSignals builds the bounded snapshot for a single response. diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index b62e332d143..9908e7a6aa1 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -116,6 +116,7 @@ type readyBucketCursorState struct { func snapshotReadyViewCursors(view readyView) readyViewCursorState { state := readyViewCursorState{lastPicked: view.lastPicked} + state.weightedState.peak = view.weightedState.peak if len(view.weightedState.current) > 0 { state.weightedState.current = make(map[string]int64, len(view.weightedState.current)) for authID, current := range view.weightedState.current { @@ -136,6 +137,7 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { return } view.lastPicked = state.lastPicked + view.weightedState.peak = max(state.weightedState.peak, len(view.flat)) weights := scheduledWeightVector(view.flat) if len(state.weightedState.current) == 0 || weightsConfigChanged(state.weightedState.weights, weights) { return @@ -154,7 +156,7 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { // newAuthScheduler constructs an empty scheduler configured for the supplied selector strategy. func newAuthScheduler(selector Selector) *authScheduler { - weighted, _ := selector.(*WeightedRoundRobinSelector) + weighted := UnwrapWeightedRoundRobinSelector(selector) return &authScheduler{ strategy: selectorStrategy(selector), weightedSelector: weighted, @@ -188,7 +190,7 @@ func (s *authScheduler) setSelector(selector Selector) { s.mu.Lock() defer s.mu.Unlock() s.strategy = selectorStrategy(selector) - s.weightedSelector, _ = selector.(*WeightedRoundRobinSelector) + s.weightedSelector = UnwrapWeightedRoundRobinSelector(selector) clear(s.mixedCursors) clear(s.mixedWeightedStates) } @@ -281,10 +283,11 @@ func (s *authScheduler) rebuild(auths []*Auth) { s.mu.Lock() defer s.mu.Unlock() - // Capture existing auth metadata so that in-flight MarkResult updates (which advanced - // Generation and updated cooldown/quota states) are preserved during rebuild. + // Capture existing auth metadata and model-shard cursors so that in-flight MarkResult + // updates and rotation progress for unchanged shards survive rebuilds. existingMetas := make(map[string]*scheduledAuthMeta) - for _, pState := range s.providers { + existingShardCursors := make(map[string]map[string]map[int]readyBucketCursorState, len(s.providers)) + for providerKey, pState := range s.providers { if pState == nil { continue } @@ -293,19 +296,78 @@ func (s *authScheduler) rebuild(auths []*Auth) { existingMetas[id] = meta } } + if len(pState.modelShards) > 0 { + modelCursors := make(map[string]map[int]readyBucketCursorState, len(pState.modelShards)) + for modelKey, shard := range pState.modelShards { + if shard == nil || len(shard.readyByPriority) == 0 { + continue + } + bucketCursors := make(map[int]readyBucketCursorState, len(shard.readyByPriority)) + for priority, bucket := range shard.readyByPriority { + if bucket == nil { + continue + } + bucketCursors[priority] = readyBucketCursorState{ + all: snapshotReadyViewCursors(bucket.all), + ws: snapshotReadyViewCursors(bucket.ws), + } + } + if len(bucketCursors) > 0 { + modelCursors[modelKey] = bucketCursors + } + } + if len(modelCursors) > 0 { + existingShardCursors[providerKey] = modelCursors + } + } } + existingMixedCursors := s.mixedCursors s.providers = make(map[string]*providerScheduler) s.authProviders = make(map[string]string) if s.authGenerations == nil { s.authGenerations = make(map[string]scheduledGenerationMeta) } - s.mixedCursors = make(map[string]int) + s.mixedCursors = make(map[string]int, len(existingMixedCursors)) s.mixedWeightedStates = make(map[string]*smoothWeightedState) now := time.Now() for _, auth := range auths { s.upsertAuthRebuildLocked(auth, existingMetas, now) } + for providerKey, modelCursors := range existingShardCursors { + pState := s.providers[providerKey] + if pState == nil { + continue + } + for modelKey, bucketCursors := range modelCursors { + shard := pState.ensureModelLocked(modelKey, now) + if shard == nil { + continue + } + for priority, cursorState := range bucketCursors { + if bucket := shard.readyByPriority[priority]; bucket != nil { + restoreReadyViewCursors(&bucket.all, cursorState.all) + restoreReadyViewCursors(&bucket.ws, cursorState.ws) + } + } + } + } + for cursorKey, cursor := range existingMixedCursors { + providersPart := cursorKey + if colon := strings.IndexByte(cursorKey, ':'); colon >= 0 { + providersPart = cursorKey[:colon] + } + allPresent := providersPart != "" + for _, providerKey := range strings.Split(providersPart, ",") { + if s.providers[providerKey] == nil { + allPresent = false + break + } + } + if allPresent { + s.mixedCursors[cursorKey] = cursor + } + } } // upsertAuth incrementally synchronizes one auth into the scheduler (lifecycle/update). @@ -440,11 +502,18 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo bucket := shard.readyByPriority[priority] entries := bucket.all.flat nonWebsocketOnly := false - if cliproxyexecutor.DownstreamWebsocket(ctx) && providerPrefersWebsocketTransport(providerKey) && bucket.ws.pickFirst(accountingPredicate) != nil { - if pinnedAuthID != "" && !shard.entries[pinnedAuthID].meta.websocketEnabled { - nonWebsocketOnly = true + transportClass := "" + if cliproxyexecutor.DownstreamWebsocket(ctx) && providerPrefersWebsocketTransport(providerKey) { + if bucket.ws.pickFirst(accountingPredicate) != nil { + if pinnedAuthID != "" && !shard.entries[pinnedAuthID].meta.websocketEnabled { + nonWebsocketOnly = true + transportClass = "http" + } else { + entries = bucket.ws.flat + transportClass = "ws" + } } else { - entries = bucket.ws.flat + transportClass = "http" } } namespace := providerKey @@ -452,7 +521,7 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo namespace = "mixed" } route := namespace + ":" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) - return s.pickWeightedPoolLocked(route, entries, accountingPredicate, pinnedAuthID, nonWebsocketOnly) + return s.pickWeightedPoolLocked(route, priority, transportClass, entries, accountingPredicate, pinnedAuthID, nonWebsocketOnly) } if picked := shard.pickReadyLocked(preferWebsocket, strategy, predicate); picked != nil { return picked, nil @@ -604,7 +673,7 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s return nil, "", ctx.Err() } route := "mixed:" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) - picked, errPick := s.pickWeightedPoolLocked(route, entries, predicate, pinnedAuthID, false) + picked, errPick := s.pickWeightedPoolLocked(route, bestPriority, "", entries, predicate, pinnedAuthID, false) if errPick != nil { return nil, "", errPick } @@ -682,6 +751,11 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st var latestAuthID string var latestAuthErr error + allClaude := true + var earliestClaudeTime time.Time + var earliestClaudeID string + var earliestClaudeSignals map[string]string + for _, providerKey := range providers { providerState := s.providers[providerKey] if providerState == nil { @@ -698,6 +772,18 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st if !localEarliest.IsZero() && (earliest.IsZero() || localEarliest.Before(earliest)) { earliest = localEarliest } + if localTotal > 0 { + localClaude, localClaudeEarliest, localClaudeID, localSignals := shard.claudeCooldownSummaryLocked(providerKey, model, now, predicate) + if !localClaude { + allClaude = false + } else if allClaude && !localClaudeEarliest.IsZero() { + if earliestClaudeTime.IsZero() || localClaudeEarliest.Before(earliestClaudeTime) || (localClaudeEarliest.Equal(earliestClaudeTime) && localClaudeID < earliestClaudeID) { + earliestClaudeTime = localClaudeEarliest + earliestClaudeID = localClaudeID + earliestClaudeSignals = localSignals + } + } + } candModelErr, candModelTime, candModelAuthID, candAuthErr, candAuthTime, candAuthAuthID := shard.candidateErrorsLocked(model, predicate) if candModelErr != nil { if latestModelErr == nil || candModelTime.After(latestModelTime) || (candModelTime.Equal(latestModelTime) && candModelAuthID > latestModelAuthID) { @@ -730,7 +816,11 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st if resetIn < 0 { resetIn = 0 } - return newModelCooldownErrorWithCause(model, "", resetIn, lastCandidateErr) + var claudeSignals map[string]string + if allClaude { + claudeSignals = earliestClaudeSignals + } + return newModelCooldownErrorForCandidates(model, "", resetIn, lastCandidateErr, allClaude, claudeSignals) } if unauthorizedCount == total { terminalCause := latestAuthErr @@ -943,6 +1033,9 @@ func isCredentialBlocked(auth *Auth, supportedModelCount int, now time.Time) boo if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { return true } + if effectiveAuthForceUntil(auth).After(now) && auth.NextRetryAfter.After(now) { + return true + } if len(auth.ModelStates) == 0 { return auth.Unavailable || auth.Quota.Exceeded || (!auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now)) } @@ -1475,7 +1568,8 @@ func (m *modelScheduler) unavailableErrorLocked(provider, model string, predicat if resetIn < 0 { resetIn = 0 } - return newModelCooldownErrorWithCause(model, providerForError, resetIn, lastCandidateErr) + claudeCooldown, _, _, claudeSignals := m.claudeCooldownSummaryLocked(provider, model, now, predicate) + return newModelCooldownErrorForCandidates(model, providerForError, resetIn, lastCandidateErr, claudeCooldown, claudeSignals) } if unauthorizedCount == total { terminalCause := authErr @@ -1591,6 +1685,57 @@ func (m *modelScheduler) availabilitySummaryLocked(predicate func(*scheduledAuth return total, cooldownCount, unauthorizedCount, earliest } +func (m *modelScheduler) claudeCooldownSummaryLocked(fallbackProvider, model string, now time.Time, predicate func(*scheduledAuth) bool) (bool, time.Time, string, map[string]string) { + if m == nil { + return false, time.Time{}, "", nil + } + matched := 0 + var earliest time.Time + var earliestID string + var earliestAuth *Auth + var earliestSignals map[string]string + for _, entry := range m.entries { + if predicate != nil && !predicate(entry) { + continue + } + matched++ + if entry == nil || entry.auth == nil { + return false, time.Time{}, "", nil + } + entryProvider := "" + if entry.meta != nil && entry.meta.providerKey != "" { + entryProvider = entry.meta.providerKey + } else if resolved := executorKeyFromAuth(entry.auth); resolved != "" { + entryProvider = resolved + } else { + entryProvider = strings.ToLower(strings.TrimSpace(fallbackProvider)) + } + if entryProvider != "claude" { + return false, time.Time{}, "", nil + } + checkModel := model + if existingModelState(entry.auth, checkModel) == nil && m.modelKey != "" { + checkModel = m.modelKey + } + okClaude, signals := claudeUnifiedCooldownForAuth(entry.auth, checkModel, now) + if !okClaude { + return false, time.Time{}, "", nil + } + if entry.state == scheduledStateCooldown && !entry.nextRetryAt.IsZero() { + if earliestAuth == nil || entry.nextRetryAt.Before(earliest) || (entry.nextRetryAt.Equal(earliest) && entry.auth.ID < earliestID) { + earliest = entry.nextRetryAt + earliestID = entry.auth.ID + earliestAuth = entry.auth + earliestSignals = signals + } + } + } + if matched == 0 || earliestAuth == nil { + return false, time.Time{}, "", nil + } + return true, earliest, earliestID, earliestSignals +} + // rebuildIndexesLocked reconstructs ready and blocked views from the current entry map. func (m *modelScheduler) rebuildIndexesLocked() { cursorStates := make(map[int]readyBucketCursorState, len(m.readyByPriority)) diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 59e31aeda4c..7ae060c229e 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -3098,6 +3098,515 @@ func TestManagerPluginSchedulerSelectionID(t *testing.T) { }) } +func TestManagerPluginSchedulerRejectAfterFailedAttempt_Issue7434(t *testing.T) { + t.Run("transient_error_reruns_same_credential_in_next_retry_round", func(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + prevTransient := transientErrorCooldownSeconds.Load() + quotaCooldownDisabled.Store(false) + SetTransientErrorCooldownSeconds(-1) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + transientErrorCooldownSeconds.Store(prevTransient) + }) + + for _, mode := range []string{"execute", "count", "stream"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(nil, &FillFirstSelector{}, nil) + manager.SetRetryConfig(1, 0, 0) + + authAID := "issue7434-transient-a-" + mode + authBID := "issue7434-transient-b-" + mode + modelID := "claude-sonnet-4-5" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authAID, "claude", []*registry.ModelInfo{{ID: modelID}}) + reg.RegisterClient(authBID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authAID) + reg.UnregisterClient(authBID) + }) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authAID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authAID, errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authBID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authBID, errRegister) + } + + var attemptedAuths []string + manager.executors["claude"] = &callbackSchedulerTestExecutor{ + schedulerTestExecutor: schedulerTestExecutor{provider: "claude"}, + execute: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + attemptedAuths = append(attemptedAuths, auth.ID) + if len(attemptedAuths) == 1 { + return cliproxyexecutor.Response{}, &Error{Code: "upstream_error", Message: "internal server error", HTTPStatus: http.StatusInternalServerError, Retryable: true} + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil + }, + countTokens: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + attemptedAuths = append(attemptedAuths, auth.ID) + if len(attemptedAuths) == 1 { + return cliproxyexecutor.Response{}, &Error{Code: "upstream_error", Message: "internal server error", HTTPStatus: http.StatusInternalServerError, Retryable: true} + } + return cliproxyexecutor.Response{Payload: []byte(`{"input_tokens":12}`)}, nil + }, + executeStream: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + attemptedAuths = append(attemptedAuths, auth.ID) + if len(attemptedAuths) == 1 { + errCh := make(chan cliproxyexecutor.StreamChunk, 1) + errCh <- cliproxyexecutor.StreamChunk{Err: &Error{Code: "upstream_error", Message: "bootstrap stream 500", HTTPStatus: http.StatusInternalServerError, Retryable: true}} + close(errCh) + return &cliproxyexecutor.StreamResult{Chunks: errCh}, nil + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"message_stop\"}\n\n")} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + }, + } + + scheduler := &fakePluginScheduler{ + pick: func(_ context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { + if prevSelected, _ := req.Options.Metadata[cliproxyexecutor.SelectedAuthMetadataKey].(string); prevSelected != "" { + return pluginapi.SchedulerPickResponse{ + Handled: true, + Reject: true, + RejectReason: "session pinned to " + prevSelected, + }, true, nil + } + for _, c := range req.Candidates { + if c.ID == authAID { + return pluginapi.SchedulerPickResponse{Handled: true, AuthID: authAID}, true, nil + } + } + return pluginapi.SchedulerPickResponse{Handled: true, Reject: true, RejectReason: "pinned credential missing"}, true, nil + }, + } + manager.SetPluginScheduler(scheduler) + + switch mode { + case "execute": + resp, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + if errExec != nil { + t.Fatalf("Execute() error = %v, want nil", errExec) + } + if string(resp.Payload) != `{"ok":true}` { + t.Fatalf("Execute() payload = %s, want {\"ok\":true}", string(resp.Payload)) + } + case "count": + resp, errCount := manager.ExecuteCount(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + if errCount != nil { + t.Fatalf("ExecuteCount() error = %v, want nil", errCount) + } + if string(resp.Payload) != `{"input_tokens":12}` { + t.Fatalf("ExecuteCount() payload = %s, want {\"input_tokens\":12}", string(resp.Payload)) + } + case "stream": + streamRes, errStream := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v, want nil", errStream) + } + for range streamRes.Chunks { + } + } + + if len(attemptedAuths) != 2 || attemptedAuths[0] != authAID || attemptedAuths[1] != authAID { + t.Fatalf("attemptedAuths = %v, want [%s %s]", attemptedAuths, authAID, authAID) + } + if len(scheduler.requests) != 3 { + t.Fatalf("len(scheduler.requests) = %d, want 3 (round0 pick, round0 rejected re-pick, round1 pick)", len(scheduler.requests)) + } + }) + } + }) + + t.Run("rate_limit_429_withholds_cooled_credential_in_next_round_so_plugin_switches", func(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + }) + + for _, mode := range []string{"execute", "stream"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(nil, &FillFirstSelector{}, nil) + manager.SetRetryConfig(1, 0, 0) + + authAID := "issue7434-429-a-" + mode + authBID := "issue7434-429-b-" + mode + modelID := "claude-sonnet-4-5" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authAID, "claude", []*registry.ModelInfo{{ID: modelID}}) + reg.RegisterClient(authBID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authAID) + reg.UnregisterClient(authBID) + }) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authAID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authAID, errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authBID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authBID, errRegister) + } + + var attemptedAuths []string + manager.executors["claude"] = &callbackSchedulerTestExecutor{ + schedulerTestExecutor: schedulerTestExecutor{provider: "claude"}, + execute: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + attemptedAuths = append(attemptedAuths, auth.ID) + if auth.ID == authAID { + return cliproxyexecutor.Response{}, &Error{Code: "rate_limited", Message: "rate limit exceeded", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil + }, + executeStream: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + attemptedAuths = append(attemptedAuths, auth.ID) + if auth.ID == authAID { + return nil, &Error{Code: "rate_limited", Message: "rate limit exceeded", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"message_stop\"}\n\n")} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + }, + } + + scheduler := &fakePluginScheduler{ + pick: func(_ context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { + if prevSelected, _ := req.Options.Metadata[cliproxyexecutor.SelectedAuthMetadataKey].(string); prevSelected != "" { + return pluginapi.SchedulerPickResponse{ + Handled: true, + Reject: true, + RejectReason: "do not switch mid-round", + }, true, nil + } + for _, c := range req.Candidates { + if c.ID == authAID { + return pluginapi.SchedulerPickResponse{Handled: true, AuthID: authAID}, true, nil + } + } + return pluginapi.SchedulerPickResponse{Handled: true, AuthID: req.Candidates[0].ID}, true, nil + }, + } + manager.SetPluginScheduler(scheduler) + + if mode == "execute" { + if _, errExec := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}); errExec != nil { + t.Fatalf("Execute() error = %v, want nil", errExec) + } + } else { + streamRes, errStream := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v, want nil", errStream) + } + for range streamRes.Chunks { + } + } + + if len(attemptedAuths) != 2 || attemptedAuths[0] != authAID || attemptedAuths[1] != authBID { + t.Fatalf("attemptedAuths = %v, want [%s %s]", attemptedAuths, authAID, authBID) + } + if len(scheduler.requests) != 3 { + t.Fatalf("len(scheduler.requests) = %d, want 3", len(scheduler.requests)) + } + }) + } + }) + + t.Run("preserves_upstream_429_when_retries_exhausted_or_sticky_plugin_rejects_next_round", func(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + }) + + for _, retryRounds := range []int{0, 1} { + name := "request_retry_0" + wantPicks := 2 + if retryRounds == 1 { + name = "request_retry_1_sticky_reject_in_round_1" + wantPicks = 3 + } + t.Run(name, func(t *testing.T) { + for _, mode := range []string{"execute", "count", "stream"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(nil, &FillFirstSelector{}, nil) + manager.SetRetryConfig(retryRounds, 1, 0) + + authAID := "issue7434-preserve-a-" + name + "-" + mode + authBID := "issue7434-preserve-b-" + name + "-" + mode + modelID := "claude-sonnet-4-5" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authAID, "claude", []*registry.ModelInfo{{ID: modelID}}) + reg.RegisterClient(authBID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authAID) + reg.UnregisterClient(authBID) + }) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authAID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authAID, errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authBID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authBID, errRegister) + } + + manager.executors["claude"] = &callbackSchedulerTestExecutor{ + schedulerTestExecutor: schedulerTestExecutor{provider: "claude"}, + execute: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + return cliproxyexecutor.Response{}, &Error{Code: "rate_limited", Message: "upstream 429", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + }, + countTokens: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + return cliproxyexecutor.Response{}, &Error{Code: "rate_limited", Message: "upstream 429", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + }, + executeStream: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + return nil, &Error{Code: "rate_limited", Message: "upstream 429", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + }, + } + + scheduler := &fakePluginScheduler{ + pick: func(_ context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { + for _, c := range req.Candidates { + if c.ID == authAID { + return pluginapi.SchedulerPickResponse{Handled: true, AuthID: authAID}, true, nil + } + } + return pluginapi.SchedulerPickResponse{ + Handled: true, + Reject: true, + RejectReason: "bound credential unavailable", + }, true, nil + }, + } + manager.SetPluginScheduler(scheduler) + + var gotErr error + switch mode { + case "execute": + _, gotErr = manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + case "count": + _, gotErr = manager.ExecuteCount(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + case "stream": + _, gotErr = manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{Stream: true}) + } + if gotErr == nil { + t.Fatal("error = nil, want upstream 429 error") + } + var authErr *Error + if !errors.As(gotErr, &authErr) { + t.Fatalf("error = %T (%v), want *auth.Error", gotErr, gotErr) + } + if authErr.Code != "rate_limited" || authErr.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("error = %s (status %d, msg %q), want rate_limited (status 429)", authErr.Code, authErr.HTTPStatus, authErr.Message) + } + if len(scheduler.requests) != wantPicks { + t.Fatalf("len(scheduler.requests) = %d, want %d", len(scheduler.requests), wantPicks) + } + }) + } + }) + } + }) + + t.Run("first_pick_reject_stops_without_retry_and_unwraps_boundary_error", func(t *testing.T) { + for _, mode := range []string{"select", "execute", "count", "stream"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(nil, &FillFirstSelector{}, nil) + manager.SetRetryConfig(2, 0, 0) + + authID := "issue7434-first-reject-" + mode + modelID := "claude-sonnet-4-5" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authID, Provider: "claude"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authID, errRegister) + } + + execCalled := false + manager.executors["claude"] = &callbackSchedulerTestExecutor{ + schedulerTestExecutor: schedulerTestExecutor{provider: "claude"}, + execute: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + execCalled = true + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil + }, + countTokens: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + execCalled = true + return cliproxyexecutor.Response{Payload: []byte(`{"input_tokens":1}`)}, nil + }, + executeStream: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + execCalled = true + ch := make(chan cliproxyexecutor.StreamChunk, 1) + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + }, + } + + scheduler := &fakePluginScheduler{ + resp: pluginapi.SchedulerPickResponse{ + Handled: true, + Reject: true, + RejectCode: "policy_denied", + RejectReason: "connection reset by peer policy guard", + }, + handled: true, + } + manager.SetPluginScheduler(scheduler) + + var gotErr error + switch mode { + case "select": + _, gotErr = manager.SelectAuth(context.Background(), "claude", modelID, cliproxyexecutor.Options{}) + case "execute": + _, gotErr = manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + case "count": + _, gotErr = manager.ExecuteCount(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + case "stream": + _, gotErr = manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{Stream: true}) + } + if gotErr == nil { + t.Fatal("error = nil, want rejection error") + } + authErr, okType := gotErr.(*Error) + if !okType || authErr == nil { + t.Fatalf("error concrete type = %T (%v), want unwrapped *auth.Error", gotErr, gotErr) + } + if authErr.Code != "policy_denied" || authErr.Message != "connection reset by peer policy guard" { + t.Fatalf("error = %s: %s, want policy_denied: connection reset by peer policy guard", authErr.Code, authErr.Message) + } + if execCalled { + t.Fatal("executor was called after first-pick reject") + } + if len(scheduler.requests) != 1 { + t.Fatalf("len(scheduler.requests) = %d, want 1", len(scheduler.requests)) + } + }) + } + }) + + t.Run("antigravity_credits_fallback_not_triggered_after_attempt_failure_and_plugin_reject", func(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + }) + + for _, retryRounds := range []int{0, 1} { + roundCase := "request_retry_0" + if retryRounds == 1 { + roundCase = "request_retry_1_sticky_reject_in_round_1" + } + t.Run(roundCase, func(t *testing.T) { + for _, mode := range []string{"execute", "stream"} { + t.Run(mode, func(t *testing.T) { + manager := NewManager(nil, &FillFirstSelector{}, nil) + manager.SetConfig(&internalconfig.Config{ + QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true}, + }) + manager.SetRetryConfig(retryRounds, 0, 0) + + authAID := "issue7434-ag-a-" + roundCase + "-" + mode + authBID := "issue7434-ag-b-" + roundCase + "-" + mode + modelID := "claude-sonnet-4-5" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authAID, "antigravity", []*registry.ModelInfo{{ID: modelID}}) + reg.RegisterClient(authBID, "antigravity", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authAID) + reg.UnregisterClient(authBID) + antigravityCreditsHintByAuth.Delete(authAID) + antigravityCreditsHintByAuth.Delete(authBID) + }) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authAID, Provider: "antigravity"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authAID, errRegister) + } + if _, errRegister := manager.Register(context.Background(), &Auth{ID: authBID, Provider: "antigravity"}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", authBID, errRegister) + } + SetAntigravityCreditsHint(authBID, AntigravityCreditsHint{ + Known: true, + Available: true, + CreditAmount: 100, + UpdatedAt: time.Now(), + }) + + var attemptedAuths []string + creditsRequested := false + manager.executors["antigravity"] = &callbackSchedulerTestExecutor{ + schedulerTestExecutor: schedulerTestExecutor{provider: "antigravity"}, + execute: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + attemptedAuths = append(attemptedAuths, auth.ID) + if AntigravityCreditsRequested(ctx) { + creditsRequested = true + return cliproxyexecutor.Response{Payload: []byte(`{"credits":true}`)}, nil + } + return cliproxyexecutor.Response{}, &Error{Code: "rate_limited", Message: "ag 429", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + }, + executeStream: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + attemptedAuths = append(attemptedAuths, auth.ID) + if AntigravityCreditsRequested(ctx) { + creditsRequested = true + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte("data: ok\n\n")} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + } + return nil, &Error{Code: "rate_limited", Message: "ag 429", HTTPStatus: http.StatusTooManyRequests, Retryable: true} + }, + } + + scheduler := &fakePluginScheduler{ + pick: func(_ context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { + for _, c := range req.Candidates { + if c.ID == authAID { + return pluginapi.SchedulerPickResponse{Handled: true, AuthID: authAID}, true, nil + } + } + return pluginapi.SchedulerPickResponse{ + Handled: true, + Reject: true, + RejectReason: "do not bypass bound credential via credits fallback", + }, true, nil + }, + } + manager.SetPluginScheduler(scheduler) + + var gotErr error + if mode == "execute" { + _, gotErr = manager.Execute(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{}) + } else { + _, gotErr = manager.ExecuteStream(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: modelID}, cliproxyexecutor.Options{Stream: true}) + } + if creditsRequested { + t.Fatalf("Antigravity credits fallback ran despite scheduler Reject (attemptedAuths=%v)", attemptedAuths) + } + if len(attemptedAuths) != 1 || attemptedAuths[0] != authAID { + t.Fatalf("attemptedAuths = %v, want [%s]", attemptedAuths, authAID) + } + if gotErr == nil { + t.Fatal("error = nil, want upstream 429 error") + } + authErr, okType := gotErr.(*Error) + if !okType || authErr == nil { + t.Fatalf("error concrete type = %T (%v), want unwrapped *auth.Error", gotErr, gotErr) + } + if authErr.Code != "rate_limited" || authErr.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("error = %s (status %d), want rate_limited (status 429)", authErr.Code, authErr.HTTPStatus) + } + }) + } + }) + } + }) +} + type callbackSchedulerTestExecutor struct { schedulerTestExecutor execute func(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) diff --git a/sdk/cliproxy/auth/scheduler_weighted.go b/sdk/cliproxy/auth/scheduler_weighted.go index 4b37cc5402c..3b72a7e99e1 100644 --- a/sdk/cliproxy/auth/scheduler_weighted.go +++ b/sdk/cliproxy/auth/scheduler_weighted.go @@ -4,7 +4,7 @@ package auth // snapshots without cloning the whole pool. Legacy and fast-path selections use // the same selector-owned credits, including pending reservations. // The scheduler mutex must be held while reading entries and selecting an auth. -func (s *authScheduler) pickWeightedPoolLocked(route string, entries []*scheduledAuth, predicate func(*scheduledAuth) bool, pinnedAuthID string, nonWebsocketOnly bool) (*Auth, error) { +func (s *authScheduler) pickWeightedPoolLocked(route string, priority int, transportClass string, entries []*scheduledAuth, predicate func(*scheduledAuth) bool, pinnedAuthID string, nonWebsocketOnly bool) (*Auth, error) { if s.weightedSelector == nil { return nil, &Error{Code: "auth_unavailable", Message: "weighted selector is not configured"} } @@ -25,7 +25,7 @@ func (s *authScheduler) pickWeightedPoolLocked(route string, entries []*schedule ids = append(ids, entry.auth.ID) weights = append(weights, entry.meta.weight) } - picked, reservation := s.weightedSelector.reservePool(route, auths, weights, ids, pinnedAuthID) + picked, reservation := s.weightedSelector.reservePoolScoped(route, priority, transportClass, auths, weights, ids, pinnedAuthID) if picked == nil { return nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} } diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 24c76560e03..dda37508db3 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "hash/fnv" + "maps" "math" "net/http" "regexp" @@ -43,11 +44,31 @@ type WeightedRoundRobinSelector struct { states map[[32]byte]*weightedPoolState stateCount int maxKeys int + // largeStates holds one shared dynamic accumulator per (route, priority, transport) + // bucket once a pool exceeds maxExactWeightedPoolMembers, preserving rotation credits + // across transient retry exclusions and cooldowns. + largeStates map[string]*weightedPoolState + largeOrder []string +} + +// UnwrapWeightedRoundRobinSelector extracts the underlying *WeightedRoundRobinSelector +// from a bare selector or a SessionAffinitySelector wrapper. +func UnwrapWeightedRoundRobinSelector(selector Selector) *WeightedRoundRobinSelector { + if weighted, ok := selector.(*WeightedRoundRobinSelector); ok { + return weighted + } + if affinity, ok := selector.(*SessionAffinitySelector); ok && affinity != nil { + if weighted, okFallback := affinity.fallback.(*WeightedRoundRobinSelector); okFallback { + return weighted + } + } + return nil } type smoothWeightedState struct { current map[string]int64 weights map[string]int64 + peak int } type weightedSelectorStateModelKey struct{} @@ -83,10 +104,12 @@ const ( ) type modelCooldownError struct { - model string - resetIn time.Duration - provider string - cause error + model string + resetIn time.Duration + provider string + cause error + claudeCooldown bool + claudeSignals map[string]string } // NewModelCooldownError creates an error representing model-level cooldown. @@ -99,17 +122,27 @@ func newModelCooldownError(model, provider string, resetIn time.Duration) *model } func newModelCooldownErrorWithCause(model, provider string, resetIn time.Duration, cause error) *modelCooldownError { + return newModelCooldownErrorForCandidates(model, provider, resetIn, cause, isClaudeCooldownProvider(provider), nil) +} + +func newModelCooldownErrorForCandidates(model, provider string, resetIn time.Duration, cause error, claudeCooldown bool, claudeSignals map[string]string) *modelCooldownError { if resetIn < 0 { resetIn = 0 } return &modelCooldownError{ - model: model, - provider: provider, - resetIn: resetIn, - cause: cause, + model: model, + provider: provider, + resetIn: resetIn, + cause: cause, + claudeCooldown: claudeCooldown, + claudeSignals: maps.Clone(claudeSignals), } } +func isClaudeCooldownProvider(provider string) bool { + return strings.EqualFold(strings.TrimSpace(provider), "claude") +} + func (e *modelCooldownError) IsModelCooldown() bool { return true } @@ -352,17 +385,419 @@ func (e *modelCooldownError) StatusCode() int { return http.StatusTooManyRequests } -func (e *modelCooldownError) Headers() http.Header { +func (e *modelCooldownError) resetSeconds() int64 { + if e == nil || e.resetIn <= 0 { + return 0 + } + seconds := int64(e.resetIn / time.Second) + if e.resetIn%time.Second > 0 { + seconds++ + } + return seconds +} + +// SafeResponseHeaders exposes only the locally trusted cooldown headers. +func (e *modelCooldownError) SafeResponseHeaders() http.Header { headers := make(http.Header) headers.Set("Content-Type", "application/json") - resetSeconds := int(math.Ceil(e.resetIn.Seconds())) - if resetSeconds < 0 { - resetSeconds = 0 + headers.Set("Retry-After", strconv.FormatInt(e.resetSeconds(), 10)) + return headers +} + +func (e *modelCooldownError) Headers() http.Header { + headers := e.SafeResponseHeaders() + if e != nil && e.claudeCooldown && e.resetIn > 0 { + e.appendClaudeUnifiedHeaders(headers) } - headers.Set("Retry-After", strconv.Itoa(resetSeconds)) return headers } +func (e *modelCooldownError) appendClaudeUnifiedHeaders(headers http.Header) { + resetSeconds := e.resetSeconds() + if resetSeconds <= 0 { + return + } + resetAt := time.Now().Unix() + resetSeconds + resetStr := strconv.FormatInt(resetAt, 10) + + if hasRejectedClaudeUnifiedSignals(e.claudeSignals) { + for key, value := range e.claudeSignals { + if strings.HasPrefix(strings.ToLower(strings.TrimSpace(key)), "anthropic-ratelimit-unified-") { + headers.Set(key, value) + } + } + headers.Set("Anthropic-Ratelimit-Unified-Status", "rejected") + headers.Set("Anthropic-Ratelimit-Unified-Reset", resetStr) + + rejWindows := rejectedClaudeUnifiedWindows(headers) + claim := strings.ToLower(strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Representative-Claim"))) + if len(rejWindows) == 0 && claim == "" && isClaudeOverageOnlyHeaders(headers) { + claim = "overage" + headers.Set("Anthropic-Ratelimit-Unified-Representative-Claim", claim) + } + if len(rejWindows) == 0 && claim != "overage" { + fallbackWindow := fallbackClaudeWindowForClaim(headers, claim, e.resetIn) + if fallbackWindow != "" { + statusHeader := claudeUnifiedWindowHeader(fallbackWindow, "Status") + if strings.TrimSpace(headers.Get(statusHeader)) == "" { + headers.Set(statusHeader, "rejected") + rejWindows = append(rejWindows, fallbackWindow) + } + } + } + + if len(rejWindows) == 1 { + headers.Set(claudeUnifiedWindowHeader(rejWindows[0], "Reset"), resetStr) + } else if len(rejWindows) > 1 { + latestHeader := "" + var latestEpoch int64 + for _, window := range rejWindows { + resetHeader := claudeUnifiedWindowHeader(window, "Reset") + raw := strings.TrimSpace(headers.Get(resetHeader)) + if raw == "" { + headers.Set(resetHeader, resetStr) + continue + } + if parsed, errParse := strconv.ParseInt(raw, 10, 64); errParse == nil && (latestHeader == "" || parsed >= latestEpoch) { + latestHeader = resetHeader + latestEpoch = parsed + } + } + if latestHeader != "" { + headers.Set(latestHeader, resetStr) + } + } + + if claim == "" { + if inferred := inferClaudeRepresentativeClaim(headers, rejWindows, e.resetIn); inferred != "" { + headers.Set("Anthropic-Ratelimit-Unified-Representative-Claim", inferred) + } + } + if strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Overage-Status")) == "" { + headers.Set("Anthropic-Ratelimit-Unified-Overage-Status", "rejected") + } + return + } + + claim := "five_hour" + if e.resetIn > 5*time.Hour { + claim = "seven_day" + } + headers.Set("Anthropic-Ratelimit-Unified-Status", "rejected") + headers.Set("Anthropic-Ratelimit-Unified-Reset", resetStr) + headers.Set("Anthropic-Ratelimit-Unified-Representative-Claim", claim) + headers.Set("Anthropic-Ratelimit-Unified-Overage-Status", "rejected") + if claim == "five_hour" { + headers.Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected") + headers.Set("Anthropic-Ratelimit-Unified-5h-Reset", resetStr) + } else { + headers.Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") + headers.Set("Anthropic-Ratelimit-Unified-7d-Reset", resetStr) + } +} + +func claudeUnifiedWindowName(headerKey string) (string, bool) { + lower := strings.ToLower(strings.TrimSpace(headerKey)) + rest, hasPrefix := strings.CutPrefix(lower, "anthropic-ratelimit-unified-") + if !hasPrefix { + return "", false + } + window, hasSuffix := strings.CutSuffix(rest, "-status") + if !hasSuffix || window == "" || window == "overage" { + return "", false + } + return window, true +} + +func claudeUnifiedWindowHeader(window, field string) string { + return http.CanonicalHeaderKey("Anthropic-Ratelimit-Unified-" + window + "-" + field) +} + +func rejectedClaudeUnifiedWindows(headers http.Header) []string { + var windows []string + for key := range headers { + window, ok := claudeUnifiedWindowName(key) + if ok && strings.EqualFold(strings.TrimSpace(headers.Get(key)), "rejected") { + windows = append(windows, window) + } + } + sort.Strings(windows) + return windows +} + +func isAllowedClaudeWindowStatus(status string) bool { + lower := strings.ToLower(strings.TrimSpace(status)) + return lower == "allowed" || lower == "allowed_warning" +} + +func isClaudeUtilizationBelowCap(raw string) bool { + raw = strings.TrimSpace(raw) + if raw == "" { + return false + } + u, errParse := strconv.ParseFloat(raw, 64) + if errParse != nil || math.IsNaN(u) || math.IsInf(u, 0) { + return false + } + return u >= 0 && u < 1.0 +} + +func isClaudeOverageOnlyHeaders(headers http.Header) bool { + overageRejected := strings.EqualFold(strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Overage-Status")), "rejected") || + strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Overage-Disabled-Reason")) != "" + if !overageRejected { + return false + } + status5h := strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-5h-Status")) + status7d := strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-7d-Status")) + allowed5h := isAllowedClaudeWindowStatus(status5h) + allowed7d := isAllowedClaudeWindowStatus(status7d) + if allowed5h && allowed7d { + return true + } + if allowed7d && status5h == "" { + util5h := strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-5h-Utilization")) + if util5h == "" || isClaudeUtilizationBelowCap(util5h) { + return true + } + } + if allowed5h && status7d == "" && isClaudeUtilizationBelowCap(headers.Get("Anthropic-Ratelimit-Unified-7d-Utilization")) { + return true + } + return false +} + +func fallbackClaudeWindowForClaim(headers http.Header, claim string, resetIn time.Duration) string { + switch { + case claim == "seven_day_overage_included": + return "7d_oi" + case strings.HasPrefix(claim, "seven_day_"): + if suffix := strings.TrimPrefix(claim, "seven_day_"); suffix != "" { + return "7d_" + suffix + } + case claim == "seven_day": + return "7d" + case claim == "five_hour": + return "5h" + case claim == "": + allowed5h := isAllowedClaudeWindowStatus(headers.Get("Anthropic-Ratelimit-Unified-5h-Status")) + allowed7d := isAllowedClaudeWindowStatus(headers.Get("Anthropic-Ratelimit-Unified-7d-Status")) + switch { + case allowed5h && allowed7d: + return "" + case allowed5h: + return "7d" + case allowed7d: + return "5h" + case resetIn > 5*time.Hour: + return "7d" + default: + return "5h" + } + } + return "" +} + +func inferClaudeRepresentativeClaim(headers http.Header, windows []string, resetIn time.Duration) string { + has5h := false + has7d := false + has7dOI := false + modelFamily7d := "" + otherWindow := "" + for _, window := range windows { + switch { + case window == "7d_oi": + has7dOI = true + case strings.HasPrefix(window, "7d_"): + if modelFamily7d == "" { + modelFamily7d = "seven_day_" + strings.TrimPrefix(window, "7d_") + } + case window == "7d": + has7d = true + case window == "5h": + has5h = true + default: + if otherWindow == "" { + otherWindow = window + } + } + } + switch { + case has7dOI: + return "seven_day_overage_included" + case modelFamily7d != "": + return modelFamily7d + case has7d: + return "seven_day" + case has5h: + return "five_hour" + case otherWindow != "": + return otherWindow + } + if strings.EqualFold(strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Overage-Status")), "rejected") || + strings.TrimSpace(headers.Get("Anthropic-Ratelimit-Unified-Overage-Disabled-Reason")) != "" { + return "overage" + } + allowed5h := isAllowedClaudeWindowStatus(headers.Get("Anthropic-Ratelimit-Unified-5h-Status")) + allowed7d := isAllowedClaudeWindowStatus(headers.Get("Anthropic-Ratelimit-Unified-7d-Status")) + switch { + case allowed5h && allowed7d: + return "" + case allowed5h: + return "seven_day" + case allowed7d: + return "five_hour" + case resetIn > 5*time.Hour: + return "seven_day" + default: + return "five_hour" + } +} + +func hasRejectedClaudeUnifiedSignals(signals map[string]string) bool { + if len(signals) == 0 { + return false + } + if strings.EqualFold(strings.TrimSpace(signals["Anthropic-Ratelimit-Unified-Status"]), "rejected") { + return true + } + for key, val := range signals { + if _, ok := claudeUnifiedWindowName(key); ok && strings.EqualFold(strings.TrimSpace(val), "rejected") { + return true + } + } + return false +} + +func claudeSignalsFromEvaluatedQuota(q QuotaState) (bool, map[string]string) { + if hasRejectedClaudeUnifiedSignals(q.cooldownSignals) { + return true, maps.Clone(q.cooldownSignals) + } + if q.cooldownEvaluated { + return false, nil + } + if hasRejectedClaudeUnifiedSignals(q.Signals) { + return true, maps.Clone(q.Signals) + } + return false, nil +} + +func claudeUnifiedCooldownForAuth(auth *Auth, model string, now time.Time) (bool, map[string]string) { + if auth == nil || executorKeyFromAuth(auth) != "claude" { + return false, nil + } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + state := existingModelState(auth, model) + if state != nil && state.NextRetryAfter.After(auth.Quota.NextRecoverAt) { + if shadow := state.credentialQuotaShadow; shadow != nil { + if !shadow.Quota.Exceeded || !shadow.Quota.NextRecoverAt.After(now) || shadow.NextRetryAfter.After(shadow.Quota.NextRecoverAt) { + return false, nil + } + if strings.EqualFold(strings.TrimSpace(shadow.Quota.Reason), "cloudflare challenge") { + return false, nil + } + return claudeSignalsFromEvaluatedQuota(shadow.Quota) + } + return false, nil + } + if ok, signals := claudeSignalsFromEvaluatedQuota(auth.Quota); ok { + return true, signals + } + if state != nil && !state.Quota.NextRecoverAt.Before(auth.Quota.NextRecoverAt) { + if ok, signals := claudeSignalsFromEvaluatedQuota(state.Quota); ok { + return true, signals + } + } + return true, nil + } + if len(auth.ModelStates) > 0 { + if model == "" { + if !auth.Unavailable { + return false, nil + } + keys := make([]string, 0, len(auth.ModelStates)) + for key := range auth.ModelStates { + keys = append(keys, key) + } + sort.Strings(keys) + var earliestRecover time.Time + var earliestSignals map[string]string + found := false + for _, key := range keys { + st := auth.ModelStates[key] + if st != nil && (st.credentialQuotaPropagated || st.Quota.Reason == "credential_quota") { + st = st.credentialQuotaShadow + } + if st == nil || !st.Quota.Exceeded || !st.Quota.NextRecoverAt.After(now) || st.NextRetryAfter.After(st.Quota.NextRecoverAt) { + return false, nil + } + if strings.EqualFold(strings.TrimSpace(st.Quota.Reason), "cloudflare challenge") { + return false, nil + } + okSt, sig := claudeSignalsFromEvaluatedQuota(st.Quota) + if !okSt { + return false, nil + } + if !found || st.Quota.NextRecoverAt.Before(earliestRecover) { + found = true + earliestRecover = st.Quota.NextRecoverAt + earliestSignals = sig + } + } + return found, earliestSignals + } + state := existingModelState(auth, model) + if state != nil && (state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota") { + state = state.credentialQuotaShadow + } + if state == nil || !state.Quota.Exceeded || !state.Quota.NextRecoverAt.After(now) || state.NextRetryAfter.After(state.Quota.NextRecoverAt) { + return false, nil + } + if strings.EqualFold(strings.TrimSpace(state.Quota.Reason), "cloudflare challenge") { + return false, nil + } + return claudeSignalsFromEvaluatedQuota(state.Quota) + } + if !auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.After(now) || auth.NextRetryAfter.After(auth.Quota.NextRecoverAt) { + return false, nil + } + if strings.EqualFold(strings.TrimSpace(auth.Quota.Reason), "cloudflare challenge") { + return false, nil + } + return claudeSignalsFromEvaluatedQuota(auth.Quota) +} + +func claudeCooldownMetadataForAuths(ctx context.Context, auths []*Auth, modelForAuth func(*Auth) string, now time.Time, clientHeaders ...http.Header) (bool, map[string]string) { + if len(auths) == 0 { + return false, nil + } + var earliestAuth *Auth + var earliestSignals map[string]string + var earliestNext time.Time + for _, candidate := range auths { + checkModel := "" + if modelForAuth != nil { + checkModel = modelForAuth(candidate) + } + okClaude, signals := claudeUnifiedCooldownForAuth(candidate, checkModel, now) + if !okClaude { + return false, nil + } + blocked, reason, next := isAuthBlockedForModelForRequest(ctx, candidate, checkModel, now, clientHeaders...) + if blocked && reason == blockReasonCooldown && next.After(now) { + if earliestAuth == nil || next.Before(earliestNext) || (next.Equal(earliestNext) && candidate.ID < earliestAuth.ID) { + earliestAuth = candidate + earliestSignals = signals + earliestNext = next + } + } + } + if earliestAuth == nil { + return false, nil + } + return true, earliestSignals +} + func authPriority(auth *Auth) int { if auth == nil || auth.Attributes == nil { return 0 @@ -558,7 +993,10 @@ func getAvailableAuthsWithPriorityModeForRequest(ctx context.Context, auths []*A if resetIn < 0 { resetIn = 0 } - return nil, newModelCooldownError(model, providerForError, resetIn) + claudeCooldown, claudeSignals := claudeCooldownMetadataForAuths(ctx, auths, func(*Auth) string { + return model + }, now, clientHeaders...) + return nil, newModelCooldownErrorForCandidates(model, providerForError, resetIn, nil, claudeCooldown, claudeSignals) } return nil, newAuthUnavailableError(earliest, now) } @@ -701,9 +1139,10 @@ func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model s return picked, nil } -// maxSmoothWeightedStateEntries bounds a single accumulator map so credentials that are -// removed permanently cannot leak entries. Real pools stay far below this bound, so the -// transient subsets produced by retry exclusions and cooldowns are never pruned. +// maxSmoothWeightedStateEntries is the floor bound on a single accumulator map so +// permanently removed credentials cannot leak entries on small pools. Larger pools scale +// their prune bound with the historical peak size (2*peak) so transient retry exclusions +// and cooldowns never wipe accumulated credits. const maxSmoothWeightedStateEntries = 1024 // prepare syncs the configured weights into the state without discarding accumulated @@ -712,6 +1151,9 @@ const maxSmoothWeightedStateEntries = 1024 // affinity), because discarding credits there would collapse selection onto the first // candidate in slice order. func (s *smoothWeightedState) prepare(weights map[string]int64) { + if len(weights) > s.peak { + s.peak = len(weights) + } if s.current == nil || weightsConfigChanged(s.weights, weights) { s.current = make(map[string]int64, len(weights)) } @@ -724,10 +1166,18 @@ func (s *smoothWeightedState) prepare(weights map[string]int64) { s.pruneStale(weights) } +func (s *smoothWeightedState) pruneLimit() int { + if scaled := 2 * s.peak; scaled > maxSmoothWeightedStateEntries { + return scaled + } + return maxSmoothWeightedStateEntries +} + // pruneStale drops entries for credentials outside the current candidate set, but only -// once a map exceeds the safety bound, so ordinary transient exclusions keep their credits. +// once a map exceeds the scaled safety bound, so transient exclusions keep their credits. func (s *smoothWeightedState) pruneStale(weights map[string]int64) { - if len(s.current) <= maxSmoothWeightedStateEntries && len(s.weights) <= maxSmoothWeightedStateEntries { + limit := s.pruneLimit() + if len(s.current) <= limit && len(s.weights) <= limit { return } for authID := range s.current { @@ -818,8 +1268,46 @@ func isAuthBlockedForModelForRequest(ctx context.Context, auth *Auth, model stri if exp, ok := auth.AccessTokenExpirationTimeForRequest(ctx, clientHeaders...); ok && !exp.IsZero() && !exp.After(now) { return true, blockReasonOther, time.Time{} } + if authForce := effectiveAuthForceUntil(auth); authForce.After(now) && auth.NextRetryAfter.After(now) { + next := authForce + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(next) { + next = auth.Quota.NextRecoverAt + } + if model != "" && len(auth.ModelStates) > 0 { + modelKey := canonicalModelLookupKey(model) + for stateModel, state := range auth.ModelStates { + if state == nil || canonicalModelLookupKey(stateModel) != modelKey { + continue + } + if state.Status == StatusDisabled { + return true, blockReasonDisabled, time.Time{} + } + stateBlocked, _, stateNext := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now) + if stateBlocked && stateNext.After(next) { + next = stateNext + } + } + } + return true, blockReasonCooldown, next + } if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { - return true, blockReasonCooldown, auth.Quota.NextRecoverAt + credNext := auth.Quota.NextRecoverAt + if model != "" && len(auth.ModelStates) > 0 { + modelKey := canonicalModelLookupKey(model) + for stateModel, state := range auth.ModelStates { + if state == nil || canonicalModelLookupKey(stateModel) != modelKey { + continue + } + if state.Status == StatusDisabled { + return true, blockReasonDisabled, time.Time{} + } + stateBlocked, _, next := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now) + if stateBlocked && next.After(credNext) { + credNext = next + } + } + } + return true, blockReasonCooldown, credNext } unavailable := codexRefreshUnavailableForRequest(ctx, auth, now, clientHeaders...) if model != "" { @@ -1633,6 +2121,30 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) { } } +// ForgetAuth forwards permanent credential removal to the fallback selector so +// weighted round-robin accounting can prune deregistered credentials without +// disturbing credits of transiently excluded credentials. +func (s *SessionAffinitySelector) ForgetAuth(authID string) { + s.ForgetAuthUpTo(authID, ^uint64(0), ^uint64(0)) +} + +// ForgetAuthUpTo forwards epoch-gated permanent credential removal to the fallback selector +// so delayed removals never evict a newly re-registered credential with the same ID. +func (s *SessionAffinitySelector) ForgetAuthUpTo(authID string, maxEpoch, maxGen uint64) { + if s == nil || authID == "" { + return + } + if epochForgetter, ok := s.fallback.(interface { + ForgetAuthUpTo(string, uint64, uint64) + }); ok && epochForgetter != nil { + epochForgetter.ForgetAuthUpTo(authID, maxEpoch, maxGen) + return + } + if forgetter, ok := s.fallback.(interface{ ForgetAuth(string) }); ok && forgetter != nil { + forgetter.ForgetAuth(authID) + } +} + // LookupAffinity observes the current session affinity binding without side effects. // It never selects a credential, creates a binding, rebinds, or refreshes TTL. // Optional authFilters allow callers (such as Manager) to exclude credentials that do not match the requested provider. diff --git a/sdk/cliproxy/auth/selector_weighted.go b/sdk/cliproxy/auth/selector_weighted.go index a984fa79d84..c048c3b0d2c 100644 --- a/sdk/cliproxy/auth/selector_weighted.go +++ b/sdk/cliproxy/auth/selector_weighted.go @@ -6,6 +6,7 @@ import ( "encoding/binary" "math" "sort" + "strconv" "time" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -22,13 +23,17 @@ type weightedTransportProviderKey struct{} // weightedReservation separates speculative credit from committed, clamped credit. // Releasing one request must never undo accounting committed by another request. type weightedReservation struct { - selector *WeightedRoundRobinSelector - state *weightedPoolState - generation uint64 - member int - total int64 - bound bool - finished bool + selector *WeightedRoundRobinSelector + state *weightedPoolState + generation uint64 + member int + selectedID string + selectedEpoch uint64 + candidateIDs []string + candidateEpochs []uint64 + total int64 + bound bool + finished bool } // pickWithAffinityReservation normalizes weighted candidates and defers accounting @@ -80,6 +85,7 @@ func (s *WeightedRoundRobinSelector) reserveForRequest(ctx context.Context, prov } var available []*Auth authID := "" + transportClass := "" if bound == nil { var errAvailable error available, errAvailable = getSelectorAvailableAuths(ctx, positiveWeightAuths(auths), provider, model, time.Now()) @@ -87,6 +93,13 @@ func (s *WeightedRoundRobinSelector) reserveForRequest(ctx context.Context, prov return nil, nil, errAvailable } available = preferCodexWebsocketAuths(ctx, transportProvider, available) + if cliproxyexecutor.DownstreamWebsocket(ctx) && transportProvider == "codex" && len(available) > 0 { + if authWebsocketsEnabled(available[0]) { + transportClass = "ws" + } else { + transportClass = "http" + } + } } else { authID = bound.ID // A binding may bypass priority or transport preference, but its traffic @@ -107,12 +120,21 @@ func (s *WeightedRoundRobinSelector) reserveForRequest(ctx context.Context, prov } } available = nonWebsocket + transportClass = "http" } else { available = preferred + if cliproxyexecutor.DownstreamWebsocket(ctx) && transportProvider == "codex" && len(available) > 0 { + if authWebsocketsEnabled(available[0]) { + transportClass = "ws" + } else { + transportClass = "http" + } + } } } + priority := inferredWeightedPoolPriority(available) key := provider + ":" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) - picked, pending := s.reserve(key, available, authID) + picked, pending := s.reserveScoped(key, priority, transportClass, available, authID) if picked == nil { return nil, nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} } @@ -122,20 +144,92 @@ func (s *WeightedRoundRobinSelector) reserveForRequest(ctx context.Context, prov // weightedPoolMember stores each ID once, with parallel accounting values in a // single allocation. Members are sorted by ID; reservations retain only an index. type weightedPoolMember struct { - id string - weight int64 - current int64 - pending int64 + id string + weight int64 + current int64 + pending int64 + epoch uint64 + regEpoch uint64 + regGen uint64 + inFlight uint32 +} + +type weightedAuthVersion struct { + regEpoch uint64 + regGen uint64 +} + +func extractWeightedAuthVersions(available []*Auth, ids []string) []weightedAuthVersion { + versions := make([]weightedAuthVersion, len(ids)) + if len(available) == len(ids) { + sortedMatch := true + for i, auth := range available { + if auth == nil || auth.ID != ids[i] { + sortedMatch = false + break + } + versions[i] = weightedAuthVersion{regEpoch: auth.RegistrationEpoch, regGen: auth.Generation} + } + if sortedMatch { + return versions + } + clear(versions) + } + for _, auth := range available { + if auth == nil { + continue + } + idx := sort.SearchStrings(ids, auth.ID) + if idx < len(ids) && ids[idx] == auth.ID { + v := weightedAuthVersion{regEpoch: auth.RegistrationEpoch, regGen: auth.Generation} + if v.regEpoch > versions[idx].regEpoch || (v.regEpoch == versions[idx].regEpoch && v.regGen > versions[idx].regGen) { + versions[idx] = v + } + } + } + return versions +} + +func memberOlderOrEqual(member weightedPoolMember, maxEpoch, maxGen uint64) bool { + if member.regEpoch != maxEpoch { + return member.regEpoch < maxEpoch + } + return member.regGen <= maxGen } // weightedPoolState is separate from the scheduler's transient-subset accumulator: // each selector pool has fixed membership and shares its weights across reservations. type weightedPoolState struct { - route string - members []weightedPoolMember - inFlight uint64 - generation uint64 - next *weightedPoolState // Exact-identity collision chain for the digest key. + route string + members []weightedPoolMember + inFlight uint64 + generation uint64 + next *weightedPoolState // Exact-identity collision chain for the digest key. + dynamic bool // Shared bucket-scoped accumulator for pools > maxExactWeightedPoolMembers. + peak int // Largest candidate count observed on a dynamic pool. + nextMemberEpoch uint64 // Monotonic member identity counter isolating removed/re-added IDs. + lastActiveIDs []string // Most recent active candidate IDs for post-drain compaction. +} + +// maxExactWeightedPoolMembers is the candidate-set size above which exact-membership +// pools switch to a shared bucket-scoped dynamic accumulator so transient retry +// exclusions and cooldowns cannot collapse smooth WRR onto the lowest Auth.IDs. +const maxExactWeightedPoolMembers = maxSmoothWeightedStateEntries + +func inferredWeightedPoolPriority(available []*Auth) int { + for _, auth := range available { + if auth != nil { + return authPriority(auth) + } + } + return 0 +} + +func dynamicWeightedPoolKey(route string, priority int, transportClass string) string { + if priority == 0 && transportClass == "" { + return route + } + return route + "\x00p=" + strconv.Itoa(priority) + "\x00t=" + transportClass } func sortedWeightedIDs(weights map[string]int64) []string { @@ -182,17 +276,56 @@ func (s *weightedPoolState) matches(route string, ids []string) bool { return true } -func (s *weightedPoolState) prepare(weights []int64) { +func (s *weightedPoolState) prepareWithVersions(weights []int64, versions []weightedAuthVersion) { + weightChanged := false + epochReplaced := false for i, member := range s.members { if member.weight != weights[i] { - s.generation++ - s.inFlight = 0 - for i := range s.members { - s.members[i].weight = weights[i] + weightChanged = true + break + } + if member.regEpoch > 0 && versions[i].regEpoch > member.regEpoch { + epochReplaced = true + } + } + if weightChanged { + s.generation++ + s.inFlight = 0 + for i := range s.members { + s.members[i].weight = weights[i] + s.members[i].current = 0 + s.members[i].pending = 0 + s.members[i].regEpoch = versions[i].regEpoch + s.members[i].regGen = versions[i].regGen + } + return + } + if epochReplaced { + s.generation++ + s.inFlight = 0 + var totalWeight int64 + for i := range s.members { + s.members[i].pending = 0 + if s.members[i].regEpoch > 0 && versions[i].regEpoch > s.members[i].regEpoch { s.members[i].current = 0 - s.members[i].pending = 0 } - return + if versions[i].regEpoch > s.members[i].regEpoch || + (versions[i].regEpoch == s.members[i].regEpoch && versions[i].regGen > s.members[i].regGen) { + s.members[i].regEpoch = versions[i].regEpoch + s.members[i].regGen = versions[i].regGen + } + totalWeight = saturatingAddInt64(totalWeight, s.members[i].weight) + } + if totalWeight > 0 { + s.boundCredits(totalWeight) + } + return + } + for i := range s.members { + if versions[i].regEpoch > s.members[i].regEpoch || + (versions[i].regEpoch == s.members[i].regEpoch && versions[i].regGen > s.members[i].regGen) { + s.members[i].regEpoch = versions[i].regEpoch + s.members[i].regGen = versions[i].regGen } } } @@ -200,22 +333,35 @@ func (s *weightedPoolState) prepare(weights []int64) { // reserve selects and reserves credit atomically. A non-empty authID preserves a // sticky binding; otherwise committed and pending credits both guide selection. func (s *WeightedRoundRobinSelector) reserve(route string, available []*Auth, authID string) (*Auth, *weightedReservation) { + return s.reserveScoped(route, inferredWeightedPoolPriority(available), "", available, authID) +} + +func (s *WeightedRoundRobinSelector) reserveScoped(route string, priority int, transportClass string, available []*Auth, authID string) (*Auth, *weightedReservation) { weightVector := authWeightVector(available) ids := sortedWeightedIDs(weightVector) weights := make([]int64, len(ids)) for i, id := range ids { weights[i] = weightVector[id] } - return s.reservePool(route, available, weights, ids, authID) + return s.reservePoolScoped(route, priority, transportClass, available, weights, ids, authID) } // reservePool consumes a normalized pool. Scheduler ready views already supply // ID-sorted candidates and cached weights, avoiding sorting and cloning auths. // ids must be sorted, with one positive weight per ID in the parallel slice. func (s *WeightedRoundRobinSelector) reservePool(route string, available []*Auth, weights []int64, ids []string, authID string) (*Auth, *weightedReservation) { - key := weightedPoolDigest(route, ids) + return s.reservePoolScoped(route, inferredWeightedPoolPriority(available), "", available, weights, ids, authID) +} + +func (s *WeightedRoundRobinSelector) reservePoolScoped(route string, priority int, transportClass string, available []*Auth, weights []int64, ids []string, authID string) (*Auth, *weightedReservation) { + dKey := dynamicWeightedPoolKey(route, priority, transportClass) + versions := extractWeightedAuthVersions(available, ids) s.mu.Lock() defer s.mu.Unlock() + if len(ids) > maxExactWeightedPoolMembers || s.largeStates[dKey] != nil { + return s.reserveDynamicPoolLocked(dKey, route, available, weights, ids, versions, authID) + } + key := weightedPoolDigest(route, ids) if s.states == nil { s.states = make(map[[32]byte]*weightedPoolState) } @@ -224,22 +370,28 @@ func (s *WeightedRoundRobinSelector) reservePool(route string, available []*Auth state = state.next } if state == nil { - limit := s.maxKeys - if limit <= 0 { - limit = 4096 - } - if s.stateCount >= limit { - s.states = make(map[[32]byte]*weightedPoolState) - s.stateCount = 0 + limit := s.poolKeyLimit() + if s.stateCount+len(s.largeStates) >= limit { + if s.stateCount > 0 { + s.states = make(map[[32]byte]*weightedPoolState) + s.stateCount = 0 + } else { + s.evictOneLargePoolLocked() + } } state = &weightedPoolState{route: route, members: make([]weightedPoolMember, len(ids)), next: s.states[key]} for i, id := range ids { - state.members[i] = weightedPoolMember{id: id, weight: weights[i]} + state.members[i] = weightedPoolMember{ + id: id, + weight: weights[i], + regEpoch: versions[i].regEpoch, + regGen: versions[i].regGen, + } } s.states[key] = state s.stateCount++ } else { - state.prepare(weights) + state.prepareWithVersions(weights, versions) } var picked *Auth var pickedCurrent, total int64 @@ -279,6 +431,366 @@ func (s *WeightedRoundRobinSelector) reservePool(route string, available []*Auth } } +func (s *WeightedRoundRobinSelector) poolKeyLimit() int { + if s.maxKeys > 0 { + return s.maxKeys + } + return 4096 +} + +func (s *WeightedRoundRobinSelector) removeLargeOrderKeyLocked(target string) { + for i, key := range s.largeOrder { + if key == target { + s.largeOrder = append(s.largeOrder[:i], s.largeOrder[i+1:]...) + return + } + } +} + +func (s *WeightedRoundRobinSelector) evictOneLargePoolLocked() { + if len(s.largeOrder) == 0 { + for k := range s.largeStates { + delete(s.largeStates, k) + return + } + return + } + for i, key := range s.largeOrder { + st := s.largeStates[key] + if st == nil || st.inFlight == 0 { + delete(s.largeStates, key) + s.largeOrder = append(s.largeOrder[:i], s.largeOrder[i+1:]...) + return + } + } + oldest := s.largeOrder[0] + s.largeOrder = s.largeOrder[1:] + delete(s.largeStates, oldest) +} + +func (s *WeightedRoundRobinSelector) reserveDynamicPoolLocked(dKey, route string, available []*Auth, weights []int64, ids []string, versions []weightedAuthVersion, authID string) (*Auth, *weightedReservation) { + if s.largeStates == nil { + s.largeStates = make(map[string]*weightedPoolState) + } + state := s.largeStates[dKey] + if state == nil { + limit := s.poolKeyLimit() + for len(s.largeStates)+s.stateCount >= limit { + if s.stateCount > 0 { + s.states = make(map[[32]byte]*weightedPoolState) + s.stateCount = 0 + if len(s.largeStates) < limit { + break + } + } + if len(s.largeStates) == 0 { + break + } + s.evictOneLargePoolLocked() + } + state = &weightedPoolState{route: route, dynamic: true} + s.largeStates[dKey] = state + s.largeOrder = append(s.largeOrder, dKey) + } + state.syncDynamicMembers(ids, weights, versions) + + var picked *Auth + var pickedEpoch uint64 + var pickedCurrent, total int64 + for _, auth := range available { + if auth == nil { + continue + } + idIndex := sort.SearchStrings(ids, auth.ID) + if idIndex == len(ids) || ids[idIndex] != auth.ID || weights[idIndex] <= 0 { + continue + } + memberIndex, okMember := state.findMember(auth.ID) + if !okMember { + continue + } + weight := weights[idIndex] + total = saturatingAddInt64(total, weight) + if authID != "" && auth.ID != authID { + continue + } + member := &state.members[memberIndex] + current := saturatingAddInt64(member.current, member.pending) + current = saturatingAddInt64(current, weight) + if picked == nil || current > pickedCurrent { + picked, pickedEpoch, pickedCurrent = auth, member.epoch, current + } + } + if picked == nil { + return nil, nil + } + candidateIDs := append([]string(nil), ids...) + state.lastActiveIDs = candidateIDs + candidateEpochs := make([]uint64, len(ids)) + for i, id := range ids { + if memberIndex, okMember := state.findMember(id); okMember { + candidateEpochs[i] = state.members[memberIndex].epoch + } + } + fullTotalWeight := max(total, state.totalMemberWeight()) + if authID != "" && !state.dynamicBoundChargeFits(candidateIDs, candidateEpochs, picked.ID, pickedEpoch, total, fullTotalWeight) { + return picked, nil + } + state.adjustDynamicPending(candidateIDs, candidateEpochs, picked.ID, pickedEpoch, total, false) + state.inFlight++ + return picked, &weightedReservation{ + selector: s, + state: state, + generation: state.generation, + selectedID: picked.ID, + selectedEpoch: pickedEpoch, + candidateIDs: candidateIDs, + candidateEpochs: candidateEpochs, + total: total, + bound: authID != "", + } +} + +func (s *weightedPoolState) findMember(id string) (int, bool) { + index := sort.Search(len(s.members), func(i int) bool { return s.members[i].id >= id }) + if index < len(s.members) && s.members[index].id == id { + return index, true + } + return index, false +} + +func (s *weightedPoolState) findMemberWithEpoch(id string, epoch uint64) (int, bool) { + index, ok := s.findMember(id) + if !ok || s.members[index].epoch != epoch { + return index, false + } + return index, true +} + +func (s *weightedPoolState) totalMemberWeight() int64 { + var total int64 + for _, member := range s.members { + total = saturatingAddInt64(total, member.weight) + } + return total +} + +func (s *weightedPoolState) syncDynamicMembers(ids []string, weights []int64, versions []weightedAuthVersion) { + if len(ids) > s.peak { + s.peak = len(ids) + } + changed := false + for i, id := range ids { + if index, ok := s.findMember(id); ok && s.members[index].weight != weights[i] { + changed = true + break + } + } + if changed { + s.generation++ + s.inFlight = 0 + for i := range s.members { + s.members[i].current = 0 + s.members[i].pending = 0 + s.members[i].inFlight = 0 + } + } + for i, id := range ids { + ver := versions[i] + index, ok := s.findMember(id) + if ok { + if s.members[index].regEpoch > 0 && ver.regEpoch > s.members[index].regEpoch { + s.nextMemberEpoch++ + s.members[index] = weightedPoolMember{ + id: id, + weight: weights[i], + epoch: s.nextMemberEpoch, + regEpoch: ver.regEpoch, + regGen: ver.regGen, + } + continue + } + s.members[index].weight = weights[i] + if ver.regEpoch > s.members[index].regEpoch || + (ver.regEpoch == s.members[index].regEpoch && ver.regGen > s.members[index].regGen) { + s.members[index].regEpoch = ver.regEpoch + s.members[index].regGen = ver.regGen + } + continue + } + s.nextMemberEpoch++ + s.members = append(s.members, weightedPoolMember{}) + copy(s.members[index+1:], s.members[index:]) + s.members[index] = weightedPoolMember{ + id: id, + weight: weights[i], + epoch: s.nextMemberEpoch, + regEpoch: ver.regEpoch, + regGen: ver.regGen, + } + } + s.pruneDynamicMembers(ids) +} + +func (s *weightedPoolState) pruneDynamicMembers(presentIDs []string) { + pruneBound := max(maxSmoothWeightedStateEntries, 2*s.peak) + if len(s.members) <= pruneBound { + return + } + present := make(map[string]struct{}, len(presentIDs)) + for _, id := range presentIDs { + present[id] = struct{}{} + } + keepCount := 0 + for _, member := range s.members { + if member.inFlight > 0 { + keepCount++ + continue + } + if _, ok := present[member.id]; ok { + keepCount++ + } + } + if keepCount == len(s.members) { + return + } + kept := make([]weightedPoolMember, 0, keepCount) + for _, member := range s.members { + if member.inFlight > 0 { + kept = append(kept, member) + continue + } + if _, ok := present[member.id]; ok { + kept = append(kept, member) + } + } + s.members = kept +} + +func (s *weightedPoolState) adjustDynamicPending(ids []string, epochs []uint64, selectedID string, selectedEpoch uint64, subsetTotal int64, release bool) { + for i, id := range ids { + index, ok := s.findMemberWithEpoch(id, epochs[i]) + if !ok { + continue + } + member := &s.members[index] + delta := member.weight + if id == selectedID && epochs[i] == selectedEpoch { + delta -= subsetTotal + } + if release { + delta = -delta + if member.inFlight > 0 { + member.inFlight-- + } + } else { + member.inFlight++ + } + member.pending = saturatingAddInt64(member.pending, delta) + } +} + +func (s *weightedPoolState) dynamicBoundChargeFits(ids []string, epochs []uint64, selectedID string, selectedEpoch uint64, subsetTotal, fullTotalWeight int64) bool { + limit := int64(math.MaxInt64) + if fullTotalWeight <= math.MaxInt64/maxSmoothWeightedCreditRounds { + limit = fullTotalWeight * maxSmoothWeightedCreditRounds + } + for i, id := range ids { + index, ok := s.findMemberWithEpoch(id, epochs[i]) + if !ok { + continue + } + member := &s.members[index] + delta := member.weight + if id == selectedID && epochs[i] == selectedEpoch { + delta -= subsetTotal + } + current := saturatingAddInt64(member.current, member.pending) + next := saturatingAddInt64(current, delta) + if next < -limit || next > limit { + return false + } + } + return true +} + +func (s *weightedPoolState) chargeDynamic(ids []string, epochs []uint64, selectedID string, selectedEpoch uint64, subsetTotal int64, bound bool) { + fullTotalWeight := max(subsetTotal, s.totalMemberWeight()) + if !bound || s.dynamicBoundChargeFits(ids, epochs, selectedID, selectedEpoch, subsetTotal, fullTotalWeight) { + for i, id := range ids { + if index, ok := s.findMemberWithEpoch(id, epochs[i]); ok { + member := &s.members[index] + member.current = saturatingAddInt64(member.current, member.weight) + } + } + if index, ok := s.findMemberWithEpoch(selectedID, selectedEpoch); ok { + s.members[index].current = saturatingAddInt64(s.members[index].current, -subsetTotal) + } + } + if s.inFlight == 0 { + s.boundCredits(fullTotalWeight) + } +} + +// ForgetAuth removes permanently deregistered credentials from weighted pool state. +// Transient exclusions (retry tried sets and cooldowns) must never call ForgetAuth. +func (s *WeightedRoundRobinSelector) ForgetAuth(authID string) { + s.ForgetAuthUpTo(authID, ^uint64(0), ^uint64(0)) +} + +// ForgetAuthUpTo removes weighted pool entries for authID whose observed +// (RegistrationEpoch, Generation) is no newer than (maxEpoch, maxGen), so a delayed +// removal or disablement cannot evict a newly re-registered or re-enabled credential. +func (s *WeightedRoundRobinSelector) ForgetAuthUpTo(authID string, maxEpoch, maxGen uint64) { + if s == nil || authID == "" { + return + } + s.mu.Lock() + defer s.mu.Unlock() + for dKey, state := range s.largeStates { + if index, ok := state.findMember(authID); ok && memberOlderOrEqual(state.members[index], maxEpoch, maxGen) { + last := len(state.members) - 1 + copy(state.members[index:], state.members[index+1:]) + state.members[last] = weightedPoolMember{} + state.members = state.members[:last] + if state.peak > len(state.members) { + state.peak = len(state.members) + } + if len(state.members) > 0 && cap(state.members) > 2*max(maxSmoothWeightedStateEntries, len(state.members)) { + compacted := make([]weightedPoolMember, len(state.members)) + copy(compacted, state.members) + state.members = compacted + } + state.lastActiveIDs = nil + } + if len(state.members) == 0 || (state.inFlight == 0 && state.peak <= maxExactWeightedPoolMembers && len(state.members) <= maxExactWeightedPoolMembers) { + delete(s.largeStates, dKey) + s.removeLargeOrderKeyLocked(dKey) + } + } + for key, head := range s.states { + var previous *weightedPoolState + for state := head; state != nil; { + next := state.next + if index, ok := state.findMember(authID); ok && memberOlderOrEqual(state.members[index], maxEpoch, maxGen) { + if previous == nil { + if next == nil { + delete(s.states, key) + } else { + s.states[key] = next + } + } else { + previous.next = next + } + s.stateCount-- + } else { + previous = state + } + state = next + } + } +} + // finish confirms or releases a reservation exactly once under the selector lock. func (r *weightedReservation) finish(committed bool) { if r == nil { @@ -293,6 +805,19 @@ func (r *weightedReservation) finish(committed bool) { if r.state.generation != r.generation { return } + if r.state.dynamic { + r.state.adjustDynamicPending(r.candidateIDs, r.candidateEpochs, r.selectedID, r.selectedEpoch, r.total, true) + r.state.inFlight-- + if committed { + r.state.chargeDynamic(r.candidateIDs, r.candidateEpochs, r.selectedID, r.selectedEpoch, r.total, r.bound) + } else if r.state.inFlight == 0 { + r.state.boundCredits(max(r.total, r.state.totalMemberWeight())) + } + if r.state.inFlight == 0 { + r.state.pruneDynamicMembers(r.state.lastActiveIDs) + } + return + } r.state.adjustPending(r.member, r.total, true) r.state.inFlight-- if committed { diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 7cf8d91aa7c..69c41e6dd86 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -151,7 +151,10 @@ type Auth struct { // stamped at or below it were dispatched before the cooldown existed and // must not clear it. Runtime only: zero after restart so restored // cooldowns can still be cleared by any success. - cooldownRecordedSeq uint64 `json:"-"` + cooldownRecordedSeq uint64 `json:"-"` + forceCooldownUntil time.Time `json:"-"` + credQuotaForceUntil time.Time `json:"-"` + forceCooldownSeeded bool `json:"-"` recentRequests recentRequestRing `json:"-"` indexAssigned bool `json:"-"` @@ -355,6 +358,10 @@ type QuotaState struct { // Cooldown transitions must use applyCooldownFields so they cannot replace // this snapshot. Signals map[string]string `json:"signals,omitempty"` + // cooldownSignals retains rejected quota headers captured when the active + // cooldown was recorded so out-of-order completions cannot overwrite them. + cooldownSignals map[string]string + cooldownEvaluated bool } // Clone returns an independent copy of the quota state. @@ -366,6 +373,12 @@ func (q QuotaState) Clone() QuotaState { copyQuota.Signals[key] = value } } + if len(q.cooldownSignals) > 0 { + copyQuota.cooldownSignals = make(map[string]string, len(q.cooldownSignals)) + for key, value := range q.cooldownSignals { + copyQuota.cooldownSignals[key] = value + } + } return copyQuota } @@ -390,7 +403,9 @@ type ModelState struct { // stamped at or below it were dispatched before the cooldown existed and // must not clear it. Runtime only: zero after restart so restored // cooldowns can still be cleared by any success. - cooldownRecordedSeq uint64 `json:"-"` + cooldownRecordedSeq uint64 `json:"-"` + forceCooldownUntil time.Time `json:"-"` + forceCooldownSeeded bool `json:"-"` // credentialQuotaShadow holds this model's independent (own-failure) // cooldown layer while the outer fields present the auth-level // credential quota propagated over it, so quota restoration can drop diff --git a/sdk/cliproxy/executor/context.go b/sdk/cliproxy/executor/context.go index 30cbebf544f..eca0abf837b 100644 --- a/sdk/cliproxy/executor/context.go +++ b/sdk/cliproxy/executor/context.go @@ -8,11 +8,16 @@ import ( type downstreamWebsocketContextKey struct{} type requireUpstreamWebsocketContextKey struct{} type upstreamAttemptTrackerContextKey struct{} +type localWebsocketTurnTrackerContextKey struct{} type upstreamAttemptTracker struct { attempted atomic.Bool } +type localWebsocketTurnTracker struct { + executed atomic.Bool +} + // WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection. func WithDownstreamWebsocket(ctx context.Context) context.Context { if ctx == nil { @@ -77,3 +82,48 @@ func UpstreamAttempted(ctx context.Context) bool { tracker, ok := ctx.Value(upstreamAttemptTrackerContextKey{}).(*upstreamAttemptTracker) return ok && tracker != nil && tracker.attempted.Load() } + +// WithLocalWebsocketTurnTracker installs a tracker that records when a downstream +// websocket turn was fulfilled locally over HTTP rather than establishing or +// owning an upstream websocket response. +func WithLocalWebsocketTurnTracker(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, localWebsocketTurnTrackerContextKey{}, &localWebsocketTurnTracker{}) +} + +// MarkLocalWebsocketTurn records that the current downstream websocket turn was +// answered locally without creating an upstream websocket response. +func MarkLocalWebsocketTurn(ctx context.Context) { + if ctx == nil { + return + } + tracker, ok := ctx.Value(localWebsocketTurnTrackerContextKey{}).(*localWebsocketTurnTracker) + if !ok || tracker == nil { + return + } + tracker.executed.Store(true) +} + +// ResetLocalWebsocketTurn clears the local websocket turn marker before a new auth attempt. +func ResetLocalWebsocketTurn(ctx context.Context) { + if ctx == nil { + return + } + tracker, ok := ctx.Value(localWebsocketTurnTrackerContextKey{}).(*localWebsocketTurnTracker) + if !ok || tracker == nil { + return + } + tracker.executed.Store(false) +} + +// LocalWebsocketTurnExecuted reports whether the current downstream websocket turn +// was fulfilled locally rather than on an upstream websocket. +func LocalWebsocketTurnExecuted(ctx context.Context) bool { + if ctx == nil { + return false + } + tracker, ok := ctx.Value(localWebsocketTurnTrackerContextKey{}).(*localWebsocketTurnTracker) + return ok && tracker != nil && tracker.executed.Load() +} diff --git a/sdk/cliproxy/service_auth_sync_test.go b/sdk/cliproxy/service_auth_sync_test.go index 9ba47e70a15..2eeb3258a45 100644 --- a/sdk/cliproxy/service_auth_sync_test.go +++ b/sdk/cliproxy/service_auth_sync_test.go @@ -1044,21 +1044,32 @@ func TestHandleAuthUpdates_SameRevisionWaitDoesNotWaitForOtherAuthInBatch(t *tes service := &Service{cfg: &config.Config{}, coreManager: manager} - bStarted := make(chan struct{}) - bBlock := make(chan struct{}) + // Registration workers run in parallel, so either auth's task can be the + // second to reach the hook. Block that one and treat the other as the + // unrelated auth whose same-revision wait must not stall. + secondStarted := make(chan struct{}) + secondBlock := make(chan struct{}) var started atomic.Int32 modelRegistrationTaskHook = func() { if started.Add(1) == 2 { - close(bStarted) - <-bBlock + close(secondStarted) + <-secondBlock } } + finishedBatch := make(chan struct{}) t.Cleanup(func() { modelRegistrationTaskHook = nil select { - case <-bBlock: + case <-secondBlock: default: - close(bBlock) + close(secondBlock) + } + // Let the released registration finish before the registry cleanup + // above runs; otherwise it re-registers models into the global + // registry after the client was unregistered and leaks into later tests. + select { + case <-finishedBatch: + case <-time.After(5 * time.Second): } }) @@ -1067,30 +1078,46 @@ func TestHandleAuthUpdates_SameRevisionWaitDoesNotWaitForOtherAuthInBatch(t *tes updateB := watcher.AuthUpdate{Action: watcher.AuthUpdateActionModify, ID: authBID, Auth: authB} updateB.SetRevision(1) - finishedBatch := make(chan struct{}) go func() { defer close(finishedBatch) service.handleAuthUpdates(context.Background(), []watcher.AuthUpdate{updateA, updateB}) }() select { - case <-bStarted: + case <-secondStarted: case <-time.After(2 * time.Second): t.Fatal("second auth registration in batch did not start") } - doneA := make(chan struct{}) + // The unblocked task registers its auth's models; the blocked one has not + // run yet. Use that to tell which auth finished. + var finishedUpdate watcher.AuthUpdate + deadline := time.Now().Add(2 * time.Second) + for finishedUpdate.ID == "" { + switch { + case len(reg.GetModelsForClient(authAID)) > 0: + finishedUpdate = updateA + case len(reg.GetModelsForClient(authBID)) > 0: + finishedUpdate = updateB + case time.Now().After(deadline): + t.Fatal("unblocked auth registration in batch did not finish") + default: + time.Sleep(5 * time.Millisecond) + } + } + + doneFinished := make(chan struct{}) go func() { - service.handleAuthUpdate(context.Background(), updateA) - close(doneA) + service.handleAuthUpdate(context.Background(), finishedUpdate) + close(doneFinished) }() select { - case <-doneA: + case <-doneFinished: case <-time.After(2 * time.Second): - t.Fatal("auth A hook wait blocked on unrelated auth B registration") + t.Fatalf("auth %s hook wait blocked on unrelated auth registration in the same batch", finishedUpdate.ID) } - close(bBlock) + close(secondBlock) select { case <-finishedBatch: case <-time.After(2 * time.Second): diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index e0e2d8d07d7..672548dbfe4 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -104,10 +104,18 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { } func newRoutingSelector(state routingRuntimeState) coreauth.Selector { + return newRoutingSelectorWithExisting(state, nil) +} + +func newRoutingSelectorWithExisting(state routingRuntimeState, existing coreauth.Selector) coreauth.Selector { var selector coreauth.Selector switch state.strategy { case "weighted-round-robin": - selector = &coreauth.WeightedRoundRobinSelector{} + if existingWeighted := coreauth.UnwrapWeightedRoundRobinSelector(existing); existingWeighted != nil { + selector = existingWeighted + } else { + selector = &coreauth.WeightedRoundRobinSelector{} + } case "fill-first": selector = &coreauth.FillFirstSelector{} default: @@ -270,7 +278,7 @@ func (s *Service) applyManagerConfig(ctx context.Context, commit configCommit) b } routingState := normalizedRoutingRuntimeState(commit.cfg) if s.appliedRoutingState == nil || *s.appliedRoutingState != routingState { - s.coreManager.SetSelector(newRoutingSelector(routingState)) + s.coreManager.SetSelector(newRoutingSelectorWithExisting(routingState, s.coreManager.Selector())) s.appliedRoutingState = &routingState } s.applyRetryConfig(commit.cfg) diff --git a/sdk/cliproxy/service_issue_7431_test.go b/sdk/cliproxy/service_issue_7431_test.go new file mode 100644 index 00000000000..9b396819f24 --- /dev/null +++ b/sdk/cliproxy/service_issue_7431_test.go @@ -0,0 +1,146 @@ +package cliproxy + +import ( + "context" + "fmt" + "net/http" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +type noopServiceIssue7431Executor struct { + provider string +} + +func (e noopServiceIssue7431Executor) Identifier() string { return e.provider } + +func (e noopServiceIssue7431Executor) Execute(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil +} + +func (e noopServiceIssue7431Executor) ExecuteStream(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} + +func (e noopServiceIssue7431Executor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e noopServiceIssue7431Executor) CountTokens(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e noopServiceIssue7431Executor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestIssue7431_ServiceApplyManagerConfigPreservesWeightedSelectorAcrossAffinityReload(t *testing.T) { + ctx := coreauth.WithSkipPersist(t.Context()) + const ( + provider = "issue-7431-svc-wrr" + model = "issue-7431-svc-model" + ) + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(noopServiceIssue7431Executor{provider: provider}) + service := &Service{coreManager: manager} + + initialCfg := &internalconfig.Config{ + Routing: internalconfig.RoutingConfig{ + Strategy: "weighted-round-robin", + SessionAffinity: false, + }, + } + if !service.applyManagerConfig(ctx, configCommit{cfg: initialCfg, sequence: 1}) { + t.Fatal("initial applyManagerConfig failed") + } + + reg := registry.GetGlobalRegistry() + var auths []*coreauth.Auth + for i := 0; i < 4; i++ { + id := fmt.Sprintf("svc-wrr-auth-%02d", i) + a := &coreauth.Auth{ID: id, Provider: provider, Status: coreauth.StatusActive} + auths = append(auths, a) + reg.RegisterClient(id, provider, []*registry.ModelInfo{{ID: model}}) + if _, errReg := manager.Register(ctx, a); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + manager.RefreshSchedulerEntry(id) + } + t.Cleanup(func() { + for _, a := range auths { + reg.UnregisterClient(a.ID) + } + }) + + var lastPickedID string + manager.RegisterExecutor(&funcIssue7431Executor{ + provider: provider, + onExec: func(a *coreauth.Auth) { + lastPickedID = a.ID + }, + }) + + var pickedBefore []string + for i := 0; i < 2; i++ { + if _, errExec := manager.Execute(ctx, []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}); errExec != nil { + t.Fatalf("execute %d failed: %v", i, errExec) + } + pickedBefore = append(pickedBefore, lastPickedID) + } + if pickedBefore[0] != auths[0].ID || pickedBefore[1] != auths[1].ID { + t.Fatalf("picks before reload = %v, want [%s %s]", pickedBefore, auths[0].ID, auths[1].ID) + } + + // Reload routing config changing session-affinity while keeping strategy: weighted-round-robin. + reloadedCfg := &internalconfig.Config{ + Routing: internalconfig.RoutingConfig{ + Strategy: "weighted-round-robin", + SessionAffinity: true, + SessionAffinityTTL: "30m", + }, + } + if !service.applyManagerConfig(ctx, configCommit{cfg: reloadedCfg, sequence: 2}) { + t.Fatal("reload applyManagerConfig failed") + } + + if _, errExec := manager.Execute(ctx, []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}); errExec != nil { + t.Fatalf("execute after reload failed: %v", errExec) + } + if lastPickedID != auths[2].ID { + t.Fatalf("pick after affinity config reload reset WRR to %q, want %q", lastPickedID, auths[2].ID) + } +} + +type funcIssue7431Executor struct { + provider string + onExec func(*coreauth.Auth) +} + +func (e funcIssue7431Executor) Identifier() string { return e.provider } + +func (e funcIssue7431Executor) Execute(_ context.Context, auth *coreauth.Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if e.onExec != nil { + e.onExec(auth) + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil +} + +func (e funcIssue7431Executor) ExecuteStream(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} + +func (e funcIssue7431Executor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e funcIssue7431Executor) CountTokens(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e funcIssue7431Executor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, nil +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 4ddaebc04d7..dd1f84ec455 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -90,6 +90,10 @@ const FailureStreamTruncatedCode = "stream_truncated" // the selected credential, including local capability validation failures. const FailureRequestScopedCode = "request_scoped" +// FailureForceCooldownCode marks failures matched by a request-scoped error rule +// configured with a cooldown action (stop-and-cooldown or continue-and-cooldown). +const FailureForceCooldownCode = "force_cooldown" + // Failure holds HTTP failure metadata for an upstream request attempt. type Failure struct { StatusCode int