diff --git a/README.md b/README.md index 9735daa45f3..eba2e893ddf 100644 --- a/README.md +++ b/README.md @@ -134,7 +134,7 @@ You can access the following providers locally and with multiple CLI accounts th Devin - Devin models (SWE-2, GPT-6 Astra, Claude Fable 5.1, etc.). Connect a Devin account with --devin-login to route the models available to that account. + Devin models (SWE-2, SWE-2 Flash, GPT-6 Astra, Claude Fable 5.1, etc.). Connect a Devin account with --devin-login to route the models available to that account. @@ -406,7 +406,7 @@ A Cloudflare-based LLM gateway and proxy server: CLIProxyAPI deployed to Cloudfl ### [Hermes CLIProxyAPI Provider](https://github.com/duu261/hermes-cliproxyapi-provider) -Model-provider plugin for [Hermes Agent](https://github.com/NousResearch/hermes-agent) that registers one Hermes provider per CLIProxyAPI provider, each on the wire its executor accepts (Anthropic Messages, Responses or Chat Completions), with context windows and reasoning levels read from CLIProxyAPI. Optional read-only Management API pass for exact provider grouping. Listed in the [Hermes plugin catalog](https://hermes-agent.nousresearch.com/docs/plugins/cliproxyapi-provider): `hermes plugins install cliproxyapi-provider`. +Model-provider plugin for [Hermes Agent](https://github.com/NousResearch/hermes-agent) that registers one Hermes provider per CLIProxyAPI provider, each on the wire its executor accepts (Anthropic Messages, Responses or Chat Completions), with context windows and reasoning levels read from CLIProxyAPI. Optional read-only Management API pass for exact provider grouping. A companion plugin adds a `/cpa` chat command for each pool's account health and, on request, each account's live quota. Listed in the [Hermes plugin catalog](https://hermes-agent.nousresearch.com/docs/plugins): `hermes plugins install cliproxyapi-provider` and `hermes plugins install cpa-status`. ### [Ebb](https://github.com/lucasnsaraujo/ebb) diff --git a/README_CN.md b/README_CN.md index 2132c49d765..62a385aa0f5 100644 --- a/README_CN.md +++ b/README_CN.md @@ -40,7 +40,7 @@ CLIProxyAPI 是一个为 CLI 提供 OpenAI/Gemini/Claude/Codex/Grok 兼容 API Devin - Devin 提供的模型（SWE-2、GPT-6 Astra、Claude Fable 5.1 等）。使用 --devin-login 连接 Devin 账户，即可路由该账户可用的模型。 + Devin 提供的模型（SWE-2、SWE-2 Flash、GPT-6 Astra、Claude Fable 5.1 等）。使用 --devin-login 连接 Devin 账户，即可路由该账户可用的模型。 @@ -301,7 +301,7 @@ Claude Code 状态栏，按当前 CPA 实例展示 Codex / Grok / Antigravity / ### [Hermes CLIProxyAPI Provider](https://github.com/duu261/hermes-cliproxyapi-provider) -[Hermes Agent](https://github.com/NousResearch/hermes-agent) 的模型提供方插件：为每个 CLIProxyAPI 提供方注册一个 Hermes 提供方，各自使用其执行器接受的协议（Anthropic Messages、Responses 或 Chat Completions），上下文窗口与推理等级读取自 CLIProxyAPI。可选只读 Management API 查询，用于精确的提供方分组。已收录于 [Hermes 插件目录](https://hermes-agent.nousresearch.com/docs/plugins/cliproxyapi-provider)：`hermes plugins install cliproxyapi-provider`。 +[Hermes Agent](https://github.com/NousResearch/hermes-agent) 的模型提供方插件：为每个 CLIProxyAPI 提供方注册一个 Hermes 提供方，各自使用其执行器接受的协议（Anthropic Messages、Responses 或 Chat Completions），上下文窗口与推理等级读取自 CLIProxyAPI。可选只读 Management API 查询，用于精确的提供方分组。配套插件提供 `/cpa` 聊天命令，可查看每个号池的账号状态，并可按需查看每个账号的实时额度。已收录于 [Hermes 插件目录](https://hermes-agent.nousresearch.com/docs/plugins)：`hermes plugins install cliproxyapi-provider` 与 `hermes plugins install cpa-status`。 ### [Ebb](https://github.com/lucasnsaraujo/ebb) diff --git a/README_JA.md b/README_JA.md index 8eca9f71b95..49034a5dab0 100644 --- a/README_JA.md +++ b/README_JA.md @@ -40,7 +40,7 @@ CLIProxyAPI は、CLI向けのOpenAI/Gemini/Claude/Codex/Grok互換APIインタ Devin - Devinのモデル（SWE-2、GPT-6 Astra、Claude Fable 5.1など）。--devin-loginでDevinアカウントに接続し、そのアカウントで利用可能なモデルにリクエストを送れます。 + Devinのモデル（SWE-2、SWE-2 Flash、GPT-6 Astra、Claude Fable 5.1など）。--devin-loginでDevinアカウントに接続し、そのアカウントで利用可能なモデルにリクエストを送れます。 @@ -300,7 +300,7 @@ Cloudflare ベースの LLM ゲートウェイ兼プロキシサーバー。CLIP ### [Hermes CLIProxyAPI Provider](https://github.com/duu261/hermes-cliproxyapi-provider) -[Hermes Agent](https://github.com/NousResearch/hermes-agent) 向けのモデルプロバイダープラグイン。CLIProxyAPI のプロバイダーごとに Hermes プロバイダーを 1 つ登録し、それぞれの executor が受け付けるプロトコル（Anthropic Messages、Responses、Chat Completions）を使います。コンテキストウィンドウと推論レベルは CLIProxyAPI から取得します。正確なプロバイダーグループ分けのための読み取り専用 Management API 照会はオプションです。[Hermes プラグインカタログ](https://hermes-agent.nousresearch.com/docs/plugins/cliproxyapi-provider)に掲載されています：`hermes plugins install cliproxyapi-provider`。 +[Hermes Agent](https://github.com/NousResearch/hermes-agent) 向けのモデルプロバイダープラグイン。CLIProxyAPI のプロバイダーごとに Hermes プロバイダーを 1 つ登録し、それぞれの executor が受け付けるプロトコル（Anthropic Messages、Responses、Chat Completions）を使います。コンテキストウィンドウと推論レベルは CLIProxyAPI から取得します。正確なプロバイダーグループ分けのための読み取り専用 Management API 照会はオプションです。コンパニオンプラグインの `/cpa` チャットコマンドで、各プールのアカウント状態と、必要に応じて各アカウントのリアルタイムのクォータを確認できます。[Hermes プラグインカタログ](https://hermes-agent.nousresearch.com/docs/plugins)に掲載されています：`hermes plugins install cliproxyapi-provider` および `hermes plugins install cpa-status`。 ### [Ebb](https://github.com/lucasnsaraujo/ebb) diff --git a/cmd/fetch_codex_models/main.go b/cmd/fetch_codex_models/main.go index ba9408f5d17..df5cb4771d4 100644 --- a/cmd/fetch_codex_models/main.go +++ b/cmd/fetch_codex_models/main.go @@ -197,7 +197,12 @@ func ensureAccessToken(ctx context.Context, store *sdkauth.FileTokenStore, auth return "", false, fmt.Errorf("missing access_token and refresh_token") } + journal, errJournal := store.RefreshJournalFor(auth) + if errJournal != nil { + return "", false, errJournal + } svc := codexauth.NewCodexAuthWithProxyURL(nil, auth.ProxyURL) + svc.SetRefreshJournal(journal) tokenData, errRefresh := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) if errRefresh != nil { return "", false, errRefresh diff --git a/cmd/server/main.go b/cmd/server/main.go index 960d5e53009..047fea26a1c 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -110,6 +110,7 @@ func main() { var claudeLogin bool var noBrowser bool var oauthCallbackPort int + var loginProxyURL string var antigravityLogin bool var kimiLogin bool var kimiAILogin bool @@ -140,6 +141,7 @@ func main() { flag.BoolVar(&claudeLogin, "claude-login", false, "Login to Claude using OAuth") flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth") flag.IntVar(&oauthCallbackPort, "oauth-callback-port", 0, "Override OAuth callback port (defaults to provider-specific port)") + flag.StringVar(&loginProxyURL, "login-proxy", "", "Proxy URL used for the OAuth login flow and saved as the new credential's proxy_url (http, https, socks5, socks5h)") flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth") flag.BoolVar(&kimiLogin, "kimi-login", false, "Login to Kimi (.com) using OAuth") flag.BoolVar(&kimiAILogin, "kimi-ai-login", false, "Login to Kimi.ai using OAuth") @@ -656,6 +658,7 @@ func main() { options := &cmd.LoginOptions{ NoBrowser: noBrowser, CallbackPort: oauthCallbackPort, + ProxyURL: loginProxyURL, } commandMode := vertexImport != "" || siwcBind != "" || antigravityLogin || codexLogin || codexDeviceLogin || claudeLogin || kimiLogin || kimiAILogin || xaiLogin || devinLogin || metaLogin diff --git a/config.example.yaml b/config.example.yaml index 90311aeb898..2acab753c82 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -170,6 +170,19 @@ routing: strategy: "round-robin" # round-robin (default), weighted-round-robin, fill-first # weighted-round-robin uses each credential's integer weight (default 1, maximum 1,000,000). # Non-positive weights exclude the credential while this strategy is active. + # Sticky requests (including inherited and LCP bindings) also consume weighted credits + # without moving existing bindings. Compensation is scoped to each eligible credential set; + # priority tiers, transport preferences, and retry/cooldown subsets do not share credits. + # Pinned requests use the eligible pool before pinning, within the pinned credential + # priority tier and transport class, even when session-affinity is disabled. + # Returning to a pool resumes its retained credits. + # Settled credits are capped at +/-32 times the eligible pool total; ratios are approximate. + # In-flight offsets are preserved until the pool settles, regardless of commit order. + # Sticky requests keep their bindings but stop adding compensation when the full + # credit delta cannot fit the cap; new sessions keep earning credit to avoid starvation. + # Credits are process-local and reset on restart or selector replacement. + # In Home mode, compatible Home versions charge retained WS turns from usage after completion; + # 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, @@ -198,6 +211,8 @@ routing: # - "quota": transient errors (408, 5xx, empty stream, buffer limit) fail over for that # request only while keeping the session bound to its original credential; quota and auth # errors (429, 401/402/403, 404/model support, invalid_grant) rebind as normal. + # - "never": never move or fail over an established session to another credential; + # keep the binding and return the error (or retry the same credential) instead. # Ignored when routing.session-affinity is false. session-affinity-rebind-on: "any" @@ -512,9 +527,9 @@ upstream: # the same exact software baseline rather than learning newer client versions (unless the # credential resolves to cloak_mode: never, which preserves the confirmed caller's headers). # header-defaults: - # user-agent: "claude-cli/2.1.280 (external, cli)" - # package-version: "0.112.1" - # runtime-version: "v26.3.0" + # user-agent: "claude-cli/2.1.321 (external, cli)" + # package-version: "0.128.0" + # runtime-version: "v26.12.1" # os: "MacOS" # arch: "arm64" # timeout: "600" diff --git a/docs/codex-refresh-recovery.md b/docs/codex-refresh-recovery.md new file mode 100644 index 00000000000..bd84b35196e --- /dev/null +++ b/docs/codex-refresh-recovery.md @@ -0,0 +1,34 @@ +# Codex refresh recovery + +A lost refresh response does not establish whether the provider consumed the refresh token. CLIProxyAPI records a claim **before** sending a Codex or local SIWC refresh request. Transport failures, unreadable or invalid successful responses, and HTTP responses without a qualified no-effect contract stop replay. A 429 or 5xx status alone is not such a contract. Redirects cannot replay refresh POSTs. + +## Recovery + +- An unresolved outcome is reported as `refresh_outcome_unknown`, with its failure stage. An unexpired access token remains usable; the refresh grant requires recovery. +- Scheduled and forced refreshes cannot resubmit a claimed grant. There is no lease expiry or time-based reset. +- A successful owner saves a receipt before returning its new credentials. Late callers reuse unexpired receipts and follow subsequent recorded rotations rather than resubmitting old tokens. An expired receipt requires recovery, including when a successful response returned the same refresh token; it is never reported as successful renewal. +- Adopt credentials from the successful refresh owner, or sign in again to obtain a **new** grant. Re-importing the same token, deleting an auth record, or restarting the server is not recovery. +- If a process stops between recording its claim and saving its result, the outcome remains uncertain. A failed journal write also fails closed; do not delete the claim to retry. + +## Persistence and deployment + +The journal must accompany credentials across backups, migrations, and restores: + +| Credential storage | Authoritative journal | +| --- | --- | +| Files | `.codex-refresh/` under the auth directory | +| PostgreSQL | `_refresh_journal` in the configured schema | +| Git | Immutable `refs/tags/cliproxy-refresh/` references in the same remote | +| S3-compatible storage | `codex-refresh/` under the configured object prefix | + +Receipts contain credentials. Protect them with the same access controls as auth files. Do not publish the Git repository, bucket, database, or journal directory. Claims and receipts are retained rather than expiring automatically. Revoke the associated grants before administratively removing their journal history. Restoring older credentials without their complete journal history is unsafe. + +On Unix, file claims synchronize their canonical directory ancestry before an exchange; those directories must permit synchronization. Windows preserves claims across process restarts but does not support directory fsync, so this is not a power-loss durability guarantee. + +PostgreSQL initializes the journal before refresh workers start, using a transaction-scoped advisory lock to coordinate concurrent initialization. It requires permission to create the journal table and acquire that lock. Git requires permission to create the journal tags; all journal pushes are non-force. S3 backends must implement atomic `If-None-Match: *` writes. A capability check rejects backends that ignore that condition. Storage failure never falls back to refreshing against a local mirror. + +All processes sharing credentials must use the same authoritative journal and this safety protocol. Older deployments or external clients refreshing the same grant cannot be coordinated by a journal they do not participate in. + +Direct SDK users can attach a `codex.RefreshJournal` using `WithRefreshJournal` or `SetRefreshJournal`. Without a configured auth directory or explicit journal, protection is process-local only; a caller persisting credentials across restarts must also provide a durable journal. + +CLIProxyAPIHome owns its separate outer-token refresh implementation. These changes cover refreshes performed in this repository (including local SIWC refreshes in Home mode); the Home service requires its own corresponding update. diff --git a/docs/management-api-v8.md b/docs/management-api-v8.md index f91304442d9..3dbc15daf73 100644 --- a/docs/management-api-v8.md +++ b/docs/management-api-v8.md @@ -64,6 +64,48 @@ The Home-owned revision fields `credentials.concurrency.observation-barrier-revision`, and `plugins.auth-revision` cannot be changed through these endpoints. +### Weighted routing with session affinity + +Set `routing.strategy` to `weighted-round-robin` to use per-credential integer +`weight` values (default `1`, maximum `1000000`; non-positive values are excluded). +With `routing.session-affinity` enabled, sticky requests, parent inheritance, +Guardian selections, and LCP matches also consume weighted credits without moving +existing bindings. New or reselected sessions favor under-served credentials +within the same eligible credential set. For example, with weights `3:1`, four +sticky requests to the second credential can be compensated by twelve new +requests to the first when both remain eligible. + +Different priority tiers, transport preferences, and retry/cooldown candidate +sets keep separate credits; returning to a set resumes its retained credits. +Pinned requests account against the eligible pool before pinning, within the +pinned credential's priority tier and transport class, even when session affinity +is disabled. Pinned and ordinary requests share the same weighted credits. A +lower-priority binding never shifts the request ratio among higher-priority credentials. + +Settled credits are capped at ±32 times the eligible pool total, so request +ratios are approximate rather than guaranteed. While reservations are in flight, +committed balances can temporarily exceed the cap; their pending offsets remain +intact until the pool settles, including when the last reservation rolls back. Sticky requests still use their binding at +the cap, but stop adding compensation when their full credit delta cannot fit. +New sessions continue repaying debt and earning credit, so excessive sticky +traffic does not starve smaller-weight peers. + +Credits are local to the CPA process and reset on restart, selector replacement, +or bounded pool-cache eviction. Concurrent selections see pending credit; local +selection/binding success confirms it, while failure releases it. In Home mode, +scheduling and credit accounting remain owned by Home. With matching CPA and Home +versions, retained WebSocket responses (including Codex duplex successors) report +`home_selection.reused` through usage so Home can charge later turns without +redispatching or acquiring another concurrency slot. The scheduling model is +carried separately from the upstream model. First-turn primary usage does not +charge again. Split additional-model billing records (such as image usage) carry +`home_selection.additional` and never add scheduling charges, even when they have +independent execution IDs. Each reused execution is charged once through primary +usage; redelivery of that execution ID is deduplicated. +Compensation arrives after completion and uses Home's current eligible +pool and API-key permissions, so it remains approximate. Old peers ignore the +optional metadata; missing usage cannot compensate a reused turn. + ### Codex multi-agent configuration migration `client.codex.optimize-multi-agent-v2` is the sole runtime setting. Loading older @@ -114,7 +156,7 @@ retain the corresponding business operation's fields. | `/oauth/callback` | GET, POST | Submit an OAuth callback. | | `/plugins` | GET | List installed plugins. | | `/plugins/` | DELETE | Delete a plugin. | -| `/plugins/store` | GET | List the plugin store. | +| `/plugins/store` | GET | List the plugin store (supports optional `?refresh=true` to bypass cached release checks). | | `/plugins/store//install` | POST | Install or update a plugin. | | `/plugins//quota` | GET, POST, DELETE | Read, fetch, or reset plugin quota. | diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index 8305345fe94..b8d8d75d498 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -792,6 +792,11 @@ func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth, quota authFileQu entry["request_retry"] = requestRetry } if quota.v8 { + // OAuth login entry the credential was signed in through; multi-entry plugins store it + // next to the provider, and the deprecated v0 route keeps its frozen response shape. + if oauthProvider := authOAuthProvider(auth); oauthProvider != "" { + entry["oauth_provider"] = oauthProvider + } nowUTC := time.Now().UTC() if exp, okExp := auth.AccessTokenExpirationTime(); okExp && !exp.IsZero() && exp.Unix() > 0 { entry["access_token_expires_at"] = exp.UTC().Format(time.RFC3339) @@ -980,6 +985,27 @@ func authProjectID(auth *coreauth.Auth) string { return "" } +func authOAuthProvider(auth *coreauth.Auth) string { + if auth == nil { + return "" + } + for _, key := range []string{"oauth_provider", "portal"} { + if auth.Attributes != nil { + if v := strings.ToLower(strings.TrimSpace(auth.Attributes[key])); v != "" { + return v + } + } + if auth.Metadata != nil { + if raw, ok := auth.Metadata[key].(string); ok { + if v := strings.ToLower(strings.TrimSpace(raw)); v != "" { + return v + } + } + } + } + return "" +} + func extractCodexIDTokenClaims(auth *coreauth.Auth) gin.H { if auth == nil || auth.Metadata == nil { return nil diff --git a/internal/api/handlers/management/auth_files_oauth_provider_test.go b/internal/api/handlers/management/auth_files_oauth_provider_test.go new file mode 100644 index 00000000000..c6dea51994a --- /dev/null +++ b/internal/api/handlers/management/auth_files_oauth_provider_test.go @@ -0,0 +1,457 @@ +package management + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" + + "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/pluginhost" + "github.com/router-for-me/CLIProxyAPI/v8/internal/watcher/synthesizer" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" +) + +func TestAuthFilesList_ReportsOAuthProvider(t *testing.T) { + cases := []struct { + id string + attributes map[string]string + metadata map[string]any + want string + }{ + {id: "builtin"}, + {id: "attributes-normalized", attributes: map[string]string{"oauth_provider": " Vendor-CN "}, want: "vendor-cn"}, + {id: "metadata-normalized", metadata: map[string]any{"oauth_provider": " Vendor-CN "}, want: "vendor-cn"}, + {id: "attributes-portal-normalized", attributes: map[string]string{"portal": " Vendor-CN "}, want: "vendor-cn"}, + {id: "metadata-portal-normalized", metadata: map[string]any{"portal": " Vendor-CN "}, want: "vendor-cn"}, + { + id: "attributes-oauth-over-metadata-oauth", + attributes: map[string]string{"oauth_provider": " Vendor ", "portal": "ignored-attributes"}, + metadata: map[string]any{"oauth_provider": "vendor-cn", "portal": "ignored-metadata"}, + want: "vendor", + }, + { + id: "metadata-oauth-over-attributes-portal", + attributes: map[string]string{"portal": "ignored-attributes"}, + metadata: map[string]any{"oauth_provider": " Vendor-CN ", "portal": "ignored-metadata"}, + want: "vendor-cn", + }, + { + id: "attributes-portal-over-metadata-portal", + attributes: map[string]string{"portal": " Vendor-CN "}, + metadata: map[string]any{"portal": "ignored-metadata"}, + want: "vendor-cn", + }, + { + id: "blank-attributes-oauth", + attributes: map[string]string{"oauth_provider": " ", "portal": "ignored"}, + metadata: map[string]any{"oauth_provider": " Vendor-CN "}, + want: "vendor-cn", + }, + { + id: "blank-metadata-oauth", + attributes: map[string]string{"oauth_provider": " ", "portal": " Vendor-CN "}, + metadata: map[string]any{"oauth_provider": " ", "portal": "ignored"}, + want: "vendor-cn", + }, + { + id: "blank-attributes-portal", + attributes: map[string]string{"oauth_provider": " ", "portal": " "}, + metadata: map[string]any{"oauth_provider": " ", "portal": " Vendor-CN "}, + want: "vendor-cn", + }, + {id: "numeric-oauth-falls-back", attributes: map[string]string{"portal": "vendor-cn"}, metadata: map[string]any{"oauth_provider": 42}, want: "vendor-cn"}, + {id: "boolean-oauth-falls-back", metadata: map[string]any{"oauth_provider": true, "portal": "vendor-cn"}, want: "vendor-cn"}, + {id: "numeric-portal", metadata: map[string]any{"portal": 42}}, + {id: "boolean-portal", metadata: map[string]any{"portal": true}}, + {id: "non-string-hints", metadata: map[string]any{"oauth_provider": 42, "portal": true}}, + {id: "blank-hints", attributes: map[string]string{"oauth_provider": " ", "portal": " "}, metadata: map[string]any{"oauth_provider": " ", "portal": " "}}, + } + manager := coreauth.NewManager(nil, nil, nil) + for _, tc := range cases { + attributes := tc.attributes + if attributes == nil { + attributes = make(map[string]string) + } + // Runtime-only entries do not require files on disk. + attributes["runtime_only"] = "true" + if _, errRegister := manager.Register(context.Background(), &coreauth.Auth{ + ID: tc.id, FileName: tc.id + ".json", Provider: "vendor", Attributes: attributes, Metadata: tc.metadata, + }); errRegister != nil { + t.Fatalf("register auth %s: %v", tc.id, errRegister) + } + } + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager) + router := oauthProviderTestRouter(h) + files := listOAuthProviderTestFiles(t, router, "/v8/management/credentials") + if len(files) != len(cases) { + t.Fatalf("v8 files count = %d, want %d", len(files), len(cases)) + } + for _, tc := range cases { + t.Run(tc.id, func(t *testing.T) { + assertOAuthProviderTestFile(t, files, tc.id, "vendor", tc.want) + }) + } + legacyFiles := listOAuthProviderTestFiles(t, router, "/v0/management/auth-files") + if len(legacyFiles) != len(cases) { + t.Fatalf("v0 files count = %d, want %d", len(legacyFiles), len(cases)) + } + for id, file := range legacyFiles { + if value := file.Get("oauth_provider"); value.Exists() { + t.Errorf("%s: v0 must not expose oauth_provider, got %s", id, value.Raw) + } + } +} + +func TestAuthFilesOAuthProvider_LoginLifecycle(t *testing.T) { + for _, tc := range []struct { + portal string + provider string + identifiers []string + want string + }{ + {portal: "vendor", provider: "vendor", identifiers: []string{"vendor", "vendor-cn"}, want: "vendor"}, + {portal: "vendor-cn", provider: "vendor", identifiers: []string{"vendor", "vendor-cn"}, want: "vendor-cn"}, + {portal: "single-vendor", provider: "single-vendor", identifiers: []string{"single-vendor"}}, + } { + for _, multipleAuths := range []bool{false, true} { + responseKind := "auth" + if multipleAuths { + responseKind = "auths" + } + t.Run(tc.portal+"/"+responseKind, func(t *testing.T) { + authDir := t.TempDir() + state := "state-oauth-provider-" + tc.portal + "-" + responseKind + t.Cleanup(func() { CancelOAuthSession(state) }) + provider := &testAuthProvider{ + identifier: tc.identifiers[0], identifiers: tc.identifiers, + startLogin: func(_ context.Context, req pluginapi.AuthLoginStartRequest) (pluginapi.AuthLoginStartResponse, error) { + if req.Provider != tc.portal { + t.Fatalf("start provider = %q, want %q", req.Provider, tc.portal) + } + return pluginapi.AuthLoginStartResponse{ + Provider: req.Provider, State: state, URL: "https://login.example/" + req.Provider, ExpiresAt: time.Now().Add(time.Hour), + }, nil + }, + pollLogin: func(_ context.Context, req pluginapi.AuthLoginPollRequest) (pluginapi.AuthLoginPollResponse, error) { + if req.Provider != tc.portal { + t.Fatalf("poll provider = %q, want %q", req.Provider, tc.portal) + } + data := pluginapi.AuthData{Provider: tc.provider, ID: "account.json", FileName: "account.json"} + resp := pluginapi.AuthLoginPollResponse{Status: pluginapi.AuthLoginStatusSuccess, Auth: data} + if multipleAuths { + resp.Auths = []pluginapi.AuthData{data} + } + return resp, nil + }, + } + host := pluginhost.New() + host.RegisterPluginForTest("oauth-provider", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{AuthProvider: provider}}) + if got := host.HasMultipleAuthProviders(tc.portal); got != (len(tc.identifiers) > 1) { + t.Fatalf("HasMultipleAuthProviders(%q) = %t", tc.portal, got) + } + + // Exercise management conversion independently of the host's stamping. + rawResp, errPoll := provider.PollLogin(context.Background(), pluginapi.AuthLoginPollRequest{Provider: tc.portal}) + if errPoll != nil { + t.Fatalf("plugin poll: %v", errPoll) + } + records := pluginLoginPollAuthsWithProvider(host, tc.portal, rawResp) + if len(records) != 1 { + t.Fatalf("management poll records = %d, want 1", len(records)) + } + assertOAuthProviderTestAuth(t, records[0], tc.want) + + resp, handled, errHostPoll := host.PollLogin(context.Background(), tc.portal, state) + if errHostPoll != nil || !handled { + t.Fatalf("host poll handled=%t error=%v", handled, errHostPoll) + } + assertOAuthProviderTestAuth(t, host.AuthDataToCoreAuth(resp.Auth, "", ""), tc.want) + for _, data := range resp.Auths { + assertOAuthProviderTestAuth(t, host.AuthDataToCoreAuth(data, "", ""), tc.want) + } + + manager := coreauth.NewManager(nil, nil, nil) + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir, Port: 8317}, manager) + h.SetPluginHost(host) + router := oauthProviderTestRouter(h) + start := oauthProviderTestGET(t, router, "/v8/management/oauth/auth-url?provider="+tc.portal) + if got := start.Get("state").String(); got != state { + t.Fatalf("login state = %q, want %q: %s", got, state, start.Raw) + } + poll := oauthProviderTestGET(t, router, "/v8/management/get-auth-status?state="+state) + if got := poll.Get("status").String(); got != "ok" { + t.Fatalf("poll status = %q: %s", got, poll.Raw) + } + auth, ok := manager.GetByID("account.json") + if !ok { + t.Fatal("signed-in credential was not registered") + } + if got := authOAuthProvider(auth); got != tc.want { + t.Fatalf("registered oauth_provider = %q, want %q", got, tc.want) + } + if tc.want == "" { + assertOAuthProviderTestAuth(t, auth, "") + } + files := listOAuthProviderTestFiles(t, router, "/v8/management/credentials") + assertOAuthProviderTestFile(t, files, "account.json", tc.provider, tc.want) + saved, errRead := os.ReadFile(filepath.Join(authDir, "account.json")) + if errRead != nil { + t.Fatalf("read saved auth: %v", errRead) + } + value := ojson.GetBytes(saved, "oauth_provider") + if value.String() != tc.want || (tc.want == "" && value.Exists()) { + t.Fatalf("saved oauth_provider = %s, want %q", value.Raw, tc.want) + } + }) + } + } +} + +func TestPluginLoginPollAuths_PreservesStorageOAuthProviderThroughHost(t *testing.T) { + for _, responseKind := range []string{"auth", "auths"} { + for _, key := range []string{"oauth_provider", "portal"} { + t.Run(responseKind+"/"+key, func(t *testing.T) { + host := pluginhost.New() + host.RegisterPluginForTest("vendor-plugin", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: &testAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + pollLogin: func(_ context.Context, req pluginapi.AuthLoginPollRequest) (pluginapi.AuthLoginPollResponse, error) { + if req.Provider != "vendor" { + t.Fatalf("poll provider = %q, want vendor", req.Provider) + } + doc := ojson.NewDocString(`{"type":"vendor"}`) + doc.Set(key, " Vendor-CN ") + data := pluginapi.AuthData{Provider: "vendor", StorageJSON: doc.Bytes()} + resp := pluginapi.AuthLoginPollResponse{Status: pluginapi.AuthLoginStatusSuccess} + if responseKind == "auths" { + resp.Auths = []pluginapi.AuthData{data} + } else { + resp.Auth = data + } + return resp, nil + }, + }, + }}) + resp, handled, errPoll := host.PollLogin(context.Background(), "vendor", "state") + if errPoll != nil || !handled { + t.Fatalf("PollLogin() handled=%t error=%v", handled, errPoll) + } + records := pluginLoginPollAuthsWithProvider(host, "vendor", resp) + if len(records) != 1 { + t.Fatalf("poll records = %d, want 1", len(records)) + } + assertOAuthProviderTestAuth(t, records[0], "vendor-cn") + storage := records[0].Storage.(interface{ RawJSON() []byte }) + if got := ojson.GetBytes(storage.RawJSON(), "oauth_provider").String(); got != "vendor-cn" { + t.Fatalf("stored oauth_provider = %q, want vendor-cn", got) + } + }) + } + } +} + +func TestPluginLoginPollAuths_ExplicitOAuthProviderHints(t *testing.T) { + for _, tc := range []struct { + name string + data pluginapi.AuthData + }{ + {name: "attribute-oauth", data: pluginapi.AuthData{Attributes: map[string]string{"oauth_provider": " Vendor "}}}, + {name: "metadata-oauth", data: pluginapi.AuthData{Metadata: map[string]any{"oauth_provider": " Vendor "}}}, + {name: "attribute-portal", data: pluginapi.AuthData{Attributes: map[string]string{"portal": " Vendor "}}}, + {name: "metadata-portal", data: pluginapi.AuthData{Metadata: map[string]any{"portal": " Vendor "}}}, + {name: "storage-oauth", data: pluginapi.AuthData{StorageJSON: []byte(`{"type":"vendor","oauth_provider":" Vendor "}`)}}, + {name: "storage-portal-fallback", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":42,"portal":" Vendor "}`)}}, + { + name: "attribute-oauth-over-all", + data: pluginapi.AuthData{ + Attributes: map[string]string{"oauth_provider": " Vendor ", "portal": "vendor-cn"}, + Metadata: map[string]any{"oauth_provider": "vendor-cn", "portal": "vendor-cn"}, + StorageJSON: []byte(`{"oauth_provider":"vendor-cn","portal":"vendor-cn"}`), + }, + }, + { + name: "metadata-oauth-over-attribute-portal", + data: pluginapi.AuthData{ + Attributes: map[string]string{"portal": "vendor-cn"}, + Metadata: map[string]any{"oauth_provider": " Vendor ", "portal": "vendor-cn"}, + StorageJSON: []byte(`{"oauth_provider":"vendor-cn","portal":"vendor-cn"}`), + }, + }, + { + name: "attribute-portal-over-metadata-portal", + data: pluginapi.AuthData{ + Attributes: map[string]string{"portal": " Vendor "}, + Metadata: map[string]any{"portal": "vendor-cn"}, + StorageJSON: []byte(`{"oauth_provider":"vendor-cn","portal":"vendor-cn"}`), + }, + }, + { + name: "metadata-portal-over-storage-oauth", + data: pluginapi.AuthData{ + Metadata: map[string]any{"portal": " Vendor "}, + StorageJSON: []byte(`{"oauth_provider":"vendor-cn","portal":"vendor-cn"}`), + }, + }, + {name: "storage-oauth-over-storage-portal", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":" Vendor ","portal":"vendor-cn"}`)}}, + } { + t.Run(tc.name, func(t *testing.T) { + tc.data.Provider = "vendor" + records := pluginLoginPollAuthsWithProvider(pluginhost.New(), "vendor-cn", pluginapi.AuthLoginPollResponse{ + Status: pluginapi.AuthLoginStatusSuccess, Auth: tc.data, + }) + if len(records) != 1 { + t.Fatalf("poll records = %d, want 1", len(records)) + } + assertOAuthProviderTestAuth(t, records[0], "vendor") + }) + } +} + +func TestAuthFilesOAuthProvider_ColdStartLifecycle(t *testing.T) { + for _, parserKind := range []string{"plugin", "file-fallback", "unhandled-plugin", "core-plugin-parser"} { + t.Run(parserKind, func(t *testing.T) { + authDir := t.TempDir() + path := filepath.Join(authDir, "cold.json") + raw := []byte(`{"type":"vendor","oauth_provider":"vendor"}`) + if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil { + t.Fatalf("write cold-start auth: %v", errWrite) + } + host := pluginhost.New() + ctx := &synthesizer.SynthesisContext{AuthDir: authDir, Now: time.Now(), StrictFileParsing: true} + switch parserKind { + case "plugin": + host.RegisterPluginForTest("vendor-plugin", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: &oauthProviderParseTestProvider{testAuthProvider: testAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + }}, + }}) + ctx.PluginAuthParser = host + case "unhandled-plugin": + ctx.PluginAuthParser = host + case "core-plugin-parser": + // The synthesizer must also support parsers that omit file attribution. + ctx.PluginAuthParser = oauthProviderFileParser(func(_ context.Context, req pluginapi.AuthParseRequest) (*coreauth.Auth, bool, error) { + auth := host.AuthDataToCoreAuth(pluginapi.AuthData{Provider: "vendor", StorageJSON: []byte(`{"type":"vendor"}`)}, req.Path, req.FileName) + auth.Metadata = nil + return auth, true, nil + }) + } + auths, errSynthesize := synthesizer.SynthesizeAuthFile(ctx, path, raw) + if errSynthesize != nil || len(auths) != 1 { + t.Fatalf("SynthesizeAuthFile() count=%d error=%v", len(auths), errSynthesize) + } + auth := auths[0] + if got := authOAuthProvider(auth); got != "vendor" { + t.Fatalf("cold-start oauth_provider = %q, want vendor", got) + } + if parserKind == "plugin" || parserKind == "core-plugin-parser" { + assertOAuthProviderTestAuth(t, auth, "vendor") + storage, ok := auth.Storage.(interface{ RawJSON() []byte }) + if !ok || ojson.GetBytes(storage.RawJSON(), "oauth_provider").String() != "vendor" { + t.Fatal("plugin storage did not retain oauth_provider=vendor") + } + } + manager := coreauth.NewManager(nil, nil, nil) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register synthesized auth: %v", errRegister) + } + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + files := listOAuthProviderTestFiles(t, oauthProviderTestRouter(h), "/v8/management/credentials") + assertOAuthProviderTestFile(t, files, auth.ID, "vendor", "vendor") + }) + } +} + +type oauthProviderParseTestProvider struct { + testAuthProvider +} + +func (p *oauthProviderParseTestProvider) ParseAuth(_ context.Context, req pluginapi.AuthParseRequest) (pluginapi.AuthParseResponse, error) { + return pluginapi.AuthParseResponse{Handled: true, Auth: pluginapi.AuthData{Provider: req.Provider, StorageJSON: req.RawJSON}}, nil +} + +type oauthProviderFileParser func(context.Context, pluginapi.AuthParseRequest) (*coreauth.Auth, bool, error) + +func (f oauthProviderFileParser) ParseAuth(ctx context.Context, req pluginapi.AuthParseRequest) (*coreauth.Auth, bool, error) { + return f(ctx, req) +} + +func oauthProviderTestRouter(h *Handler) *gin.Engine { + router := gin.New() + v8 := router.Group("/v8/management", func(c *gin.Context) { + c.Set(ConfigV8ContextKey, true) + c.Next() + }) + v8.GET("/credentials", h.ListAuthFiles) + v8.GET("/oauth/auth-url", h.StartOAuthV8) + v8.GET("/get-auth-status", h.GetAuthStatus) + router.GET("/v0/management/auth-files", h.ListAuthFiles) + return router +} + +func oauthProviderTestGET(t *testing.T, router http.Handler, path string) ojson.Result { + t.Helper() + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil)) + if rec.Code != http.StatusOK || !ojson.ValidBytes(rec.Body.Bytes()) { + t.Fatalf("GET %s status=%d body=%s", path, rec.Code, rec.Body.String()) + } + return ojson.Parse(rec.Body.String()) +} + +func listOAuthProviderTestFiles(t *testing.T, router http.Handler, path string) map[string]ojson.Result { + t.Helper() + resp := oauthProviderTestGET(t, router, path) + files := resp.Get("files") + if !files.IsArray() || len(files.Array()) == 0 { + t.Fatalf("GET %s returned no files: %s", path, resp.Raw) + } + out := make(map[string]ojson.Result) + for _, file := range files.Array() { + out[file.Get("id").String()] = file + } + return out +} + +func assertOAuthProviderTestFile(t *testing.T, files map[string]ojson.Result, id, provider, want string) { + t.Helper() + file, ok := files[id] + if !ok { + t.Fatalf("missing credential %s", id) + } + value := file.Get("oauth_provider") + if value.String() != want || (want == "" && value.Exists()) { + t.Errorf("%s: oauth_provider = %s, want %q", id, value.Raw, want) + } + if file.Get("type").String() != provider || file.Get("provider").String() != provider { + t.Errorf("%s: credential provider changed: %s", id, file.Raw) + } +} + +func assertOAuthProviderTestAuth(t *testing.T, auth *coreauth.Auth, want string) { + t.Helper() + if auth == nil { + t.Fatal("missing auth") + } + if want == "" { + for _, key := range []string{"oauth_provider", "portal"} { + if value, exists := auth.Metadata[key]; exists { + t.Errorf("single-entry metadata must omit %s, got %#v", key, value) + } + if value, exists := auth.Attributes[key]; exists { + t.Errorf("single-entry attributes must omit %s, got %q", key, value) + } + } + return + } + if auth.Metadata["oauth_provider"] != want || auth.Attributes["oauth_provider"] != want { + t.Errorf("oauth_provider metadata=%#v attributes=%q, want %q", auth.Metadata["oauth_provider"], auth.Attributes["oauth_provider"], want) + } +} diff --git a/internal/api/handlers/management/auth_files_provider_oauth.go b/internal/api/handlers/management/auth_files_provider_oauth.go index 24070794597..0d057981660 100644 --- a/internal/api/handlers/management/auth_files_provider_oauth.go +++ b/internal/api/handlers/management/auth_files_provider_oauth.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/MoeForge/moego/ojson" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/antigravity" "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/claude" @@ -110,7 +111,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) { } // Generate random state parameter - state, err := misc.GenerateRandomState() + state, err := claude.GenerateState() if err != nil { log.Errorf("Failed to generate state parameter: %v", err) c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) @@ -1568,10 +1569,12 @@ func pluginLoginPollAuthsWithProvider(host *pluginhost.Host, provider string, re authData.Provider = trimmedSessionProvider } effectivePortal := trimmedSessionProvider + hasExplicitHint := false for _, key := range []string{"oauth_provider", "portal"} { if authData.Attributes != nil { if v := strings.ToLower(strings.TrimSpace(authData.Attributes[key])); v != "" { effectivePortal = v + hasExplicitHint = true break } } @@ -1579,24 +1582,34 @@ func pluginLoginPollAuthsWithProvider(host *pluginhost.Host, provider string, re if raw, ok := authData.Metadata[key].(string); ok { if v := strings.ToLower(strings.TrimSpace(raw)); v != "" { effectivePortal = v + hasExplicitHint = true + break + } + } + } + } + if !hasExplicitHint && ojson.ValidBytes(authData.StorageJSON) { + for _, key := range []string{"oauth_provider", "portal"} { + value := ojson.GetBytes(authData.StorageJSON, key) + if value.Type == ojson.String { + if hint := strings.ToLower(strings.TrimSpace(value.String())); hint != "" { + effectivePortal = hint + hasExplicitHint = true break } } } } - if effectivePortal != "" && !strings.EqualFold(strings.TrimSpace(authData.Provider), effectivePortal) { + isMultiEntry := host.HasMultipleAuthProviders(effectivePortal) || host.HasMultipleAuthProviders(authData.Provider) + if effectivePortal != "" && (!strings.EqualFold(strings.TrimSpace(authData.Provider), effectivePortal) || hasExplicitHint || isMultiEntry) { if authData.Metadata == nil { authData.Metadata = make(map[string]any) } - if _, exists := authData.Metadata["oauth_provider"]; !exists { - authData.Metadata["oauth_provider"] = effectivePortal - } + authData.Metadata["oauth_provider"] = effectivePortal if authData.Attributes == nil { authData.Attributes = make(map[string]string) } - if strings.TrimSpace(authData.Attributes["oauth_provider"]) == "" { - authData.Attributes["oauth_provider"] = effectivePortal - } + authData.Attributes["oauth_provider"] = effectivePortal } record := host.AuthDataToCoreAuth(authData, "", "") if record == nil { diff --git a/internal/api/handlers/management/oauth_claude_authorize_test.go b/internal/api/handlers/management/oauth_claude_authorize_test.go new file mode 100644 index 00000000000..8f4bd2c61d1 --- /dev/null +++ b/internal/api/handlers/management/oauth_claude_authorize_test.go @@ -0,0 +1,142 @@ +package management + +import ( + "bytes" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/MoeForge/moego/base64" + "github.com/MoeForge/moego/ojson" + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/claude" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +type claudeAuthorizationProbe struct { + *claude.ClaudeAuth + authURL string + state string +} + +func (p *claudeAuthorizationProbe) GenerateAuthURL(state string, pkceCodes *claude.PKCECodes) (string, string, error) { + authURL, returnedState, errGenerate := p.ClaudeAuth.GenerateAuthURL(state, pkceCodes) + if errGenerate != nil { + return "", "", errGenerate + } + p.authURL, p.state = authURL, returnedState + // Stop before starting a background callback waiter that could outlive this test. + return "", "", errors.New("authorization URL inspected") +} + +func TestClaudeOAuthV8CallbackPreservesURLSafeState(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, form := range []string{"json", "redirect_url", "query"} { + t.Run(form, func(t *testing.T) { + authDir := t.TempDir() + handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil) + router := gin.New() + router.POST("/v8/management/oauth/callback", handler.PostOAuthCallbackV8) + router.GET("/v8/management/oauth/callback", handler.GetOAuthCallbackV8) + state := strings.Repeat("Ab9-_", 8) + "AbA" + RegisterOAuthSession(state, "anthropic") + defer CancelOAuthSession(state) + + callbackRequest := func(callbackState string) *http.Request { + values := map[string]string{"provider": "claude", "code": "callback-code", "state": callbackState} + query := url.Values{} + for key, value := range values { + query.Set(key, value) + } + if form == "query" { + return httptest.NewRequest(http.MethodGet, "/v8/management/oauth/callback?"+query.Encode(), nil) + } + if form == "redirect_url" { + values = map[string]string{"redirect_url": "http://localhost:54545/callback?" + query.Encode()} + } + doc := ojson.NewDocString(`{}`) + for key, value := range values { + if errSet := doc.Set(key, value); errSet != nil { + t.Fatalf("set callback field %s: %v", key, errSet) + } + } + request := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/callback", bytes.NewReader(doc.Bytes())) + request.Header.Set("Content-Type", "application/json") + return request + } + + wrongState := httptest.NewRecorder() + router.ServeHTTP(wrongState, callbackRequest(strings.ToLower(state))) + if wrongState.Code != http.StatusNotFound { + t.Fatalf("wrong-state callback status = %d, want 404: %s", wrongState.Code, wrongState.Body.String()) + } + entries, errReadDir := os.ReadDir(authDir) + if errReadDir != nil || len(entries) != 0 { + t.Fatalf("wrong-state callback wrote a file: entries = %v, error = %v", entries, errReadDir) + } + + response := httptest.NewRecorder() + router.ServeHTTP(response, callbackRequest(state)) + if response.Code != http.StatusOK { + t.Fatalf("callback status = %d, want 200: %s", response.Code, response.Body.String()) + } + payload, errRead := os.ReadFile(filepath.Join(authDir, ".oauth-anthropic-"+state+".oauth")) + if errRead != nil { + t.Fatalf("read callback file: %v", errRead) + } + if got := ojson.GetBytesNoCopy(payload, "state").String(); got != state { + t.Fatalf("callback file state = %q, want %q", got, state) + } + if got := ojson.GetBytesNoCopy(payload, "code").String(); got != "callback-code" { + t.Fatalf("callback file code = %q, want callback-code", got) + } + }) + } +} + +func TestStartOAuthV8ClaudeUsesNativeAuthorizationParameters(t *testing.T) { + gin.SetMode(gin.TestMode) + probe := &claudeAuthorizationProbe{ClaudeAuth: &claude.ClaudeAuth{}} + originalFactory := newClaudeOAuthService + newClaudeOAuthService = func(*config.Config, string) claudeOAuthService { return probe } + defer func() { newClaudeOAuthService = originalFactory }() + + handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, nil) + router := gin.New() + router.POST("/v8/management/oauth/auth-url", handler.StartOAuthV8) + + request := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"claude"}`)) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + if response.Code != http.StatusInternalServerError { + t.Fatalf("authorization probe status = %d, want 500: %s", response.Code, response.Body.String()) + } + state := probe.state + if len(state) != 43 { + t.Errorf("state length = %d, want 43", len(state)) + } + decoded, errDecode := base64.RawURLEncoding.DecodeString(state) + if errDecode != nil || len(decoded) != 32 { + t.Errorf("state must encode 32 random bytes as unpadded base64url: decoded length = %d, error = %v", len(decoded), errDecode) + } + if errValidate := ValidateOAuthState(state); errValidate != nil { + t.Errorf("generated state cannot be used by the callback handler: %v", errValidate) + } + parsed, errParse := url.Parse(probe.authURL) + if errParse != nil { + t.Fatalf("parse authorization URL: %v", errParse) + } + if got := parsed.Query().Get("state"); got != state { + t.Error("authorization URL did not preserve the generated state") + } + const wantScope = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload" + if got := parsed.Query().Get("scope"); got != wantScope { + t.Errorf("authorization scope = %q, want %q", got, wantScope) + } +} diff --git a/internal/api/handlers/management/plugin_store.go b/internal/api/handlers/management/plugin_store.go index ec2470f80f9..c9295f0c51f 100644 --- a/internal/api/handlers/management/plugin_store.go +++ b/internal/api/handlers/management/plugin_store.go @@ -158,7 +158,7 @@ func (h *Handler) ListPluginStore(c *gin.Context) { } } client := h.newPluginStoreClient(proxyURL, "", storeAuth) - latestVersions := h.latestPluginVersions(c.Request.Context(), client, latestInput) + latestVersions := h.latestPluginVersions(c.Request.Context(), client, latestInput, pluginStoreForceRefresh(c)) entries := make([]pluginStoreListEntry, 0, len(plugins)) for index, item := range plugins { @@ -220,6 +220,14 @@ func (h *Handler) ListPluginStore(c *gin.Context) { }) } +func pluginStoreForceRefresh(c *gin.Context) bool { + if c == nil { + return false + } + refresh := strings.TrimSpace(c.Query("refresh")) + return strings.EqualFold(refresh, "true") || refresh == "1" +} + func (h *Handler) InstallPluginFromStore(c *gin.Context) { h.installPluginFromStore(c, runtime.GOOS, runtime.GOARCH) } diff --git a/internal/api/handlers/management/plugin_store_release.go b/internal/api/handlers/management/plugin_store_release.go index 1e64dcf05ae..62e05465231 100644 --- a/internal/api/handlers/management/plugin_store_release.go +++ b/internal/api/handlers/management/plugin_store_release.go @@ -13,11 +13,13 @@ import ( const ( pluginReleaseCacheTTL = time.Hour pluginReleaseFailureCacheTTL = 30 * time.Second + pluginReleaseRefreshFloor = 10 * time.Second pluginReleaseConcurrency = 2 ) type pluginReleaseCacheEntry struct { version string + checkedAt time.Time nextCheckAt time.Time } @@ -64,7 +66,7 @@ func (cache *pluginReleaseCache) cached(client pluginstore.Client, plugin plugin } // latestPluginVersions preserves catalog order, including skipped placeholders. -func (h *Handler) latestPluginVersions(ctx context.Context, client pluginstore.Client, plugins []pluginstore.Plugin) []string { +func (h *Handler) latestPluginVersions(ctx context.Context, client pluginstore.Client, plugins []pluginstore.Plugin, force bool) []string { versions := make([]string, len(plugins)) var wg sync.WaitGroup for index, plugin := range plugins { @@ -74,14 +76,14 @@ func (h *Handler) latestPluginVersions(ctx context.Context, client pluginstore.C wg.Add(1) go func(index int) { defer wg.Done() - versions[index] = h.latestPluginVersion(ctx, client, plugins[index]) + versions[index] = h.latestPluginVersion(ctx, client, plugins[index], force) }(index) } wg.Wait() return versions } -func (h *Handler) latestPluginVersion(ctx context.Context, client pluginstore.Client, plugin pluginstore.Plugin) string { +func (h *Handler) latestPluginVersion(ctx context.Context, client pluginstore.Client, plugin pluginstore.Plugin, force bool) string { if pluginstore.PluginInstallType(plugin) != pluginstore.InstallTypeGitHubRelease { return "" } @@ -95,7 +97,12 @@ func (h *Handler) latestPluginVersion(ctx context.Context, client pluginstore.Cl for { cache.mu.Lock() entry = cache.entries[key] - if cache.now().Before(entry.nextCheckAt) || ctx.Err() != nil { + now := cache.now() + fresh := now.Before(entry.nextCheckAt) + if force { + fresh = !entry.checkedAt.IsZero() && now.Before(entry.checkedAt.Add(pluginReleaseRefreshFloor)) + } + if fresh || ctx.Err() != nil { cache.mu.Unlock() return entry.version } @@ -137,18 +144,21 @@ func (h *Handler) latestPluginVersion(ctx context.Context, client pluginstore.Cl version, errRelease = pluginstore.ReleaseVersion(release) } if ctx.Err() == nil { - ttl := pluginReleaseFailureCacheTTL - if errRelease != nil { - log.WithError(errRelease).WithField("plugin_id", plugin.ID).Warn("pluginstore: failed to fetch latest release") - } else { - entry.version = version - ttl = pluginReleaseCacheTTL - } // Start the TTL when the lookup finishes, not when it was queued. - entry.nextCheckAt = cache.now().Add(ttl) + now := cache.now() + entry.checkedAt = now var rateLimit *pluginstore.RateLimitError - if errors.As(errRelease, &rateLimit) { + if errRelease == nil { + entry.version = version + entry.nextCheckAt = now.Add(pluginReleaseCacheTTL) + } else if errors.As(errRelease, &rateLimit) { + log.WithError(errRelease).WithField("plugin_id", plugin.ID).Warn("pluginstore: failed to fetch latest release") entry.nextCheckAt = rateLimit.RetryAt + } else { + log.WithError(errRelease).WithField("plugin_id", plugin.ID).Warn("pluginstore: failed to fetch latest release") + if next := now.Add(pluginReleaseFailureCacheTTL); next.After(entry.nextCheckAt) { + entry.nextCheckAt = next + } } } } diff --git a/internal/api/handlers/management/plugin_store_release_test.go b/internal/api/handlers/management/plugin_store_release_test.go index 82bf26250e8..d0ea54ea840 100644 --- a/internal/api/handlers/management/plugin_store_release_test.go +++ b/internal/api/handlers/management/plugin_store_release_test.go @@ -134,7 +134,7 @@ func TestPluginReleaseCacheRetainsSuccessAndUsesCompletionTime(t *testing.T) { })} check := func(want string, wantCalls int) { t.Helper() - if got := h.latestPluginVersion(context.Background(), client, plugin); got != want || calls != wantCalls { + if got := h.latestPluginVersion(context.Background(), client, plugin, false); got != want || calls != wantCalls { t.Fatalf("version=%q calls=%d, want %q/%d", got, calls, want, wantCalls) } } @@ -176,11 +176,11 @@ func TestPluginReleaseCacheCoalescesConcurrentLookups(t *testing.T) { h := &Handler{} plugin := pluginstore.Plugin{Repository: "https://github.com/test/plugin"} results := make(chan string, 21) - go func() { results <- h.latestPluginVersion(context.Background(), client, plugin) }() + go func() { results <- h.latestPluginVersion(context.Background(), client, plugin, false) }() <-started for range 20 { ctx := &pluginReleaseWaitContext{Context: context.Background(), waiting: make(chan struct{})} - go func() { results <- h.latestPluginVersion(ctx, client, plugin) }() + go func() { results <- h.latestPluginVersion(ctx, client, plugin, false) }() <-ctx.waiting } close(release) @@ -211,11 +211,11 @@ func TestPluginReleaseCacheCanceledLeaderHandsOffToFollower(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() leaderResult := make(chan string, 1) - go func() { leaderResult <- h.latestPluginVersion(ctx, client, plugin) }() + go func() { leaderResult <- h.latestPluginVersion(ctx, client, plugin, false) }() <-started waitCtx := &pluginReleaseWaitContext{Context: context.Background(), waiting: make(chan struct{})} followerResult := make(chan string, 1) - go func() { followerResult <- h.latestPluginVersion(waitCtx, client, plugin) }() + go func() { followerResult <- h.latestPluginVersion(waitCtx, client, plugin, false) }() <-waitCtx.waiting cancel() if got := <-leaderResult; got != "" { @@ -237,13 +237,13 @@ func TestPluginReleaseCacheCanceledFollowerDoesNotCancelLeader(t *testing.T) { h := &Handler{} plugin := pluginstore.Plugin{Repository: "https://github.com/test/plugin"} leaderResult := make(chan string, 1) - go func() { leaderResult <- h.latestPluginVersion(context.Background(), client, plugin) }() + go func() { leaderResult <- h.latestPluginVersion(context.Background(), client, plugin, false) }() <-started ctx, cancel := context.WithCancel(context.Background()) defer cancel() waitCtx := &pluginReleaseWaitContext{Context: ctx, waiting: make(chan struct{})} followerResult := make(chan string, 1) - go func() { followerResult <- h.latestPluginVersion(waitCtx, client, plugin) }() + go func() { followerResult <- h.latestPluginVersion(waitCtx, client, plugin, false) }() <-waitCtx.waiting cancel() if got := <-followerResult; got != "" { @@ -269,7 +269,7 @@ func TestPluginReleaseCacheSharesConcurrencyAndCancelsQueuedLookup(t *testing.T) results := make(chan string, pluginReleaseConcurrency) for index := range pluginReleaseConcurrency { go func() { - results <- h.latestPluginVersion(context.Background(), client, pluginstore.Plugin{Repository: fmt.Sprintf("https://github.com/test/plugin-%d", index)}) + results <- h.latestPluginVersion(context.Background(), client, pluginstore.Plugin{Repository: fmt.Sprintf("https://github.com/test/plugin-%d", index)}, false) }() } for range pluginReleaseConcurrency { @@ -280,7 +280,7 @@ func TestPluginReleaseCacheSharesConcurrencyAndCancelsQueuedLookup(t *testing.T) waitCtx := &pluginReleaseWaitContext{Context: ctx, waiting: make(chan struct{})} queuedResult := make(chan string, 1) plugin := pluginstore.Plugin{Repository: "https://github.com/test/queued"} - go func() { queuedResult <- h.latestPluginVersion(waitCtx, client, plugin) }() + go func() { queuedResult <- h.latestPluginVersion(waitCtx, client, plugin, false) }() <-waitCtx.waiting cancel() if got := <-queuedResult; got != "" || calls.Load() != pluginReleaseConcurrency { @@ -290,7 +290,7 @@ func TestPluginReleaseCacheSharesConcurrencyAndCancelsQueuedLookup(t *testing.T) for range pluginReleaseConcurrency { <-results } - if got := h.latestPluginVersion(context.Background(), client, plugin); got != "1.0.0" || calls.Load() != pluginReleaseConcurrency+1 { + if got := h.latestPluginVersion(context.Background(), client, plugin, false); got != "1.0.0" || calls.Load() != pluginReleaseConcurrency+1 { t.Fatalf("canceled lookup was negatively cached: version=%q calls=%d", got, calls.Load()) } } @@ -310,7 +310,7 @@ func TestPluginReleaseCacheSeparatesCredentialsAndEgress(t *testing.T) { client.ResolvedAuth = []pluginstore.ResolvedAuthConfig{{Match: "https://api.github.com/", Type: pluginstore.AuthTypeGitHubToken, Token: pluginstore.Secret(token)}} client.ResolvedAuthExpiresAt = time.Now().Add(time.Hour) } - if got := h.latestPluginVersion(context.Background(), client, plugin); got != "1.0.0" { + if got := h.latestPluginVersion(context.Background(), client, plugin, false); got != "1.0.0" { t.Fatalf("version = %q", got) } } @@ -318,7 +318,7 @@ func TestPluginReleaseCacheSeparatesCredentialsAndEgress(t *testing.T) { t.Fatalf("calls = %d, want 3", calls.Load()) } client.NetworkScope = "other-proxy" - h.latestPluginVersion(context.Background(), client, plugin) + h.latestPluginVersion(context.Background(), client, plugin, false) if calls.Load() != 4 { t.Fatalf("calls after proxy change = %d, want 4", calls.Load()) } @@ -352,7 +352,7 @@ func TestPluginReleaseChecksStopQueuedRequestsDuringCooldown(t *testing.T) { h.pluginReleases.entries[pluginReleaseKey(client, plugins[index])] = pluginReleaseCacheEntry{version: "1.0.0"} } result := make(chan []string, 1) - go func() { result <- h.latestPluginVersions(context.Background(), client, plugins) }() + go func() { result <- h.latestPluginVersions(context.Background(), client, plugins, false) }() for range pluginReleaseConcurrency { <-started } @@ -463,3 +463,132 @@ func TestListPluginStoreUsesCachedUninstalledVersionWithoutRefresh(t *testing.T) t.Fatalf("plugins=%#v requests=%v", body.Plugins, httpClient.counts) } } + +func listPluginStoreWithQueryForTest(t *testing.T, h *Handler, query string) pluginStoreListResponse { + t.Helper() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v8/management/plugins/store"+query, nil) + h.ListPluginStore(c) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d; body=%s", rec.Code, rec.Body.String()) + } + var body pluginStoreListResponse + if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil { + t.Fatal(errDecode) + } + return body +} + +func TestListPluginStoreRefreshBypassesCacheWithFloor(t *testing.T) { + t.Parallel() + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + releaseTag := "v0.1.0" + var releaseCalls int + h := &Handler{ + cfg: &config.Config{Plugins: config.PluginsConfig{ + Dir: writeManagementPluginFile(t, "sample-provider-v0.1.0"), + }}, + pluginReleases: pluginReleaseCache{nowFunc: func() time.Time { return now }}, + pluginStoreHTTPClient: pluginReleaseDoerFunc(func(req *http.Request) (*http.Response, error) { + if req.URL.Host != "api.github.com" { + return fakePluginStoreHTTPClient{pluginstore.DefaultRegistryURL: registryJSON(t)}.Do(req) + } + releaseCalls++ + return pluginReleaseResponse(releaseTag), nil + }), + } + + body := listPluginStoreWithQueryForTest(t, h, "") + if body.Plugins[0].Version != "0.1.0" || body.Plugins[0].UpdateAvailable || releaseCalls != 1 { + t.Fatalf("initial entry=%#v releaseCalls=%d", body.Plugins[0], releaseCalls) + } + + releaseTag = "v0.2.0" + now = now.Add(time.Minute) + body = listPluginStoreWithQueryForTest(t, h, "") + if body.Plugins[0].Version != "0.1.0" || body.Plugins[0].UpdateAvailable || releaseCalls != 1 { + t.Fatalf("cached entry=%#v releaseCalls=%d, want cached 0.1.0 without network call", body.Plugins[0], releaseCalls) + } + + body = listPluginStoreWithQueryForTest(t, h, "?refresh=true") + if body.Plugins[0].Version != "0.2.0" || !body.Plugins[0].UpdateAvailable || releaseCalls != 2 { + t.Fatalf("refreshed entry=%#v releaseCalls=%d, want 0.2.0 with update_available=true and 2 calls", body.Plugins[0], releaseCalls) + } + + releaseTag = "v0.3.0" + now = now.Add(5 * time.Second) + body = listPluginStoreWithQueryForTest(t, h, "?refresh=true") + if body.Plugins[0].Version != "0.2.0" || releaseCalls != 2 { + t.Fatalf("floor-throttled entry=%#v releaseCalls=%d, want 0.2.0 and 2 calls within 10s floor", body.Plugins[0], releaseCalls) + } + + now = now.Add(5 * time.Second) + body = listPluginStoreWithQueryForTest(t, h, "?refresh=1") + if body.Plugins[0].Version != "0.3.0" || !body.Plugins[0].UpdateAvailable || releaseCalls != 3 { + t.Fatalf("refresh=1 entry=%#v releaseCalls=%d, want 0.3.0 and 3 calls after 10s floor", body.Plugins[0], releaseCalls) + } +} + +func TestListPluginStoreRefreshPreservesRateLimiterAndSkipsUninstalled(t *testing.T) { + t.Parallel() + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + reset := time.Now().Add(time.Hour).Truncate(time.Second) + registry := pluginstore.Registry{ + SchemaVersion: pluginstore.SchemaVersion, + Plugins: []pluginstore.Plugin{ + { + ID: "installed-plugin", Name: "Installed", Description: "Installed plugin", Author: "owner", + Version: "0.1.0", Repository: "https://github.com/owner/installed-plugin", + }, + { + ID: "uninstalled-plugin", Name: "Uninstalled", Description: "Uninstalled plugin", Author: "owner", + Version: "0.1.0", Repository: "https://github.com/owner/uninstalled-plugin", + }, + }, + } + registryBytes, errMarshal := json.Marshal(registry) + if errMarshal != nil { + t.Fatal(errMarshal) + } + var releaseCalls int + rateLimited := false + h := &Handler{ + cfg: &config.Config{Plugins: config.PluginsConfig{ + Dir: writeManagementPluginFile(t, "installed-plugin-v0.1.0"), + }}, + pluginReleases: pluginReleaseCache{nowFunc: func() time.Time { return now }}, + pluginStoreRateLimiter: &pluginstore.GitHubRateLimiter{}, + pluginStoreHTTPClient: pluginReleaseDoerFunc(func(req *http.Request) (*http.Response, error) { + if req.URL.Host != "api.github.com" { + return fakePluginStoreHTTPClient{pluginstore.DefaultRegistryURL: registryBytes}.Do(req) + } + releaseCalls++ + if !rateLimited { + return pluginReleaseResponse("v0.2.0"), nil + } + headers := make(http.Header) + headers.Set("X-RateLimit-Remaining", "0") + headers.Set("X-RateLimit-Reset", strconv.FormatInt(reset.Unix(), 10)) + return &http.Response{StatusCode: http.StatusForbidden, Header: headers, Body: io.NopCloser(strings.NewReader("limited"))}, nil + }), + } + + body := listPluginStoreWithQueryForTest(t, h, "?refresh=true") + if body.Plugins[0].Version != "0.2.0" || body.Plugins[1].Version != "0.1.0" || releaseCalls != 1 { + t.Fatalf("first refresh plugins=%#v releaseCalls=%d, want installed=0.2.0 uninstalled=0.1.0 calls=1", body.Plugins, releaseCalls) + } + + rateLimited = true + now = now.Add(time.Minute) + body = listPluginStoreWithQueryForTest(t, h, "?refresh=true") + if body.Plugins[0].Version != "0.2.0" || releaseCalls != 2 { + t.Fatalf("rate-limited refresh plugins=%#v releaseCalls=%d, want cached 0.2.0 and 2 calls", body.Plugins, releaseCalls) + } + + now = now.Add(time.Minute) + body = listPluginStoreWithQueryForTest(t, h, "?refresh=true") + if body.Plugins[0].Version != "0.2.0" || releaseCalls != 2 { + t.Fatalf("cooldown refresh plugins=%#v releaseCalls=%d, want rate limiter to block network call during cooldown", body.Plugins, releaseCalls) + } +} diff --git a/internal/api/handlers/management/plugin_store_test.go b/internal/api/handlers/management/plugin_store_test.go index 4669b73935e..18b7a613ca0 100644 --- a/internal/api/handlers/management/plugin_store_test.go +++ b/internal/api/handlers/management/plugin_store_test.go @@ -16,6 +16,7 @@ import ( "strings" "sync" "testing" + "time" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" @@ -402,6 +403,74 @@ func TestListPluginStoreShowsLatestReleaseVersionAndCaches(t *testing.T) { } } +func TestListPluginStoreRefreshQueryBypassesCache(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + httpClient := &countingPluginStoreHTTPClient{responses: fakePluginStoreHTTPClient{ + "https://registry.example/registry.json": registryJSON(t), + "https://api.github.com/repos/author-name/cliproxy-sample-provider-plugin/releases/latest": []byte(`{ + "tag_name": "v0.1.0", + "assets": [] + }`), + }} + h := &Handler{ + cfg: &config.Config{ + Plugins: config.PluginsConfig{ + Enabled: true, + Dir: writeManagementPluginFile(t, "sample-provider-v0.1.0"), + }, + }, + configFilePath: writeTestConfigFile(t), + pluginStoreRegistryURL: "https://registry.example/registry.json", + pluginStoreHTTPClient: httpClient, + pluginReleases: pluginReleaseCache{nowFunc: func() time.Time { return now }}, + } + + listWithQuery := func(query string) pluginStoreListResponse { + t.Helper() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v8/management/plugins/store"+query, nil) + h.ListPluginStore(c) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + var body pluginStoreListResponse + if errDecode := json.Unmarshal(rec.Body.Bytes(), &body); errDecode != nil { + t.Fatalf("Unmarshal() error = %v; body=%s", errDecode, rec.Body.String()) + } + return body + } + + body := listWithQuery("") + if body.Plugins[0].Version != "0.1.0" || body.Plugins[0].UpdateAvailable { + t.Fatalf("initial plugin = %#v, want 0.1.0 without update", body.Plugins[0]) + } + + httpClient.mu.Lock() + httpClient.responses["https://api.github.com/repos/author-name/cliproxy-sample-provider-plugin/releases/latest"] = []byte(`{ + "tag_name": "v0.2.0", + "assets": [] + }`) + httpClient.mu.Unlock() + + now = now.Add(pluginReleaseRefreshFloor) + body = listWithQuery("?refresh=false") + if body.Plugins[0].Version != "0.1.0" || body.Plugins[0].UpdateAvailable { + t.Fatalf("refresh=false plugin = %#v, want cached 0.1.0", body.Plugins[0]) + } + + body = listWithQuery("?refresh=true") + if body.Plugins[0].Version != "0.2.0" || !body.Plugins[0].UpdateAvailable { + t.Fatalf("refresh=true plugin = %#v, want refreshed 0.2.0 with update_available=true", body.Plugins[0]) + } + releaseCalls := httpClient.count("https://api.github.com/repos/author-name/cliproxy-sample-provider-plugin/releases/latest") + if releaseCalls != 2 { + t.Fatalf("latest release fetched %d times, want 2", releaseCalls) + } +} + func TestListPluginStoreFallsBackToRegistryVersion(t *testing.T) { t.Parallel() diff --git a/internal/api/middleware/request_logging.go b/internal/api/middleware/request_logging.go index 5ffb153b6fb..33c01abd8e8 100644 --- a/internal/api/middleware/request_logging.go +++ b/internal/api/middleware/request_logging.go @@ -51,7 +51,8 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc { } loggerEnabled := logger.IsEnabled() - captureBody := shouldCaptureRequestBody(loggerEnabled, c.Request) + hasMatchedRoute := c.FullPath() != "" + captureBody := hasMatchedRoute && shouldCaptureRequestBody(loggerEnabled, c.Request) // Capture request metadata before authentication without reading the request body. requestInfo, err := captureRequestInfo(c, false) @@ -67,7 +68,7 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc { wrapper.logOnErrorOnly = true } c.Writer = wrapper - attachRequestLogSources(c, logger, loggerEnabled) + attachRequestLogSources(c, logger, loggerEnabled && hasMatchedRoute) deferredCapture := attachDeferredRequestBodyCapture(c.Request, logger, requestInfo, false, false) if captureBody { c.Set(authenticatedRequestBodyCapturerContextKey, func() { diff --git a/internal/api/middleware/request_logging_test.go b/internal/api/middleware/request_logging_test.go index b1bf861ad39..a0d20aeba76 100644 --- a/internal/api/middleware/request_logging_test.go +++ b/internal/api/middleware/request_logging_test.go @@ -995,3 +995,183 @@ func TestResponseWriterLogOnErrorOnlyZeroBufferUntilActionableError(t *testing.T t.Fatalf("wrapper.body.String() = %q, want %q", got, errFrame) } } + +func TestRequestLoggingMiddleware_Unrouted404DoesNotEvictErrorDumpRing(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("unrouted 404 does not create error log or evict existing error dumps when request-log is false", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 5) + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/messages", func(c *gin.Context) { + c.JSON(http.StatusBadGateway, gin.H{"error": "upstream failure"}) + }) + router.NoRoute(func(c *gin.Context) { + c.AbortWithStatus(http.StatusNotFound) + }) + + realFailReq := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4-5"}`)) + realFailReq.Header.Set("Content-Type", "application/json") + realFailResp := httptest.NewRecorder() + router.ServeHTTP(realFailResp, realFailReq) + if realFailResp.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", realFailResp.Code, http.StatusBadGateway) + } + + for i := 0; i < 10; i++ { + probeReq := httptest.NewRequest(http.MethodPost, "/v1/files", strings.NewReader(`{"purpose":"user_data"}`)) + probeReq.Header.Set("Content-Type", "application/json") + probeResp := httptest.NewRecorder() + router.ServeHTTP(probeResp, probeReq) + if probeResp.Code != http.StatusNotFound { + t.Fatalf("probe %d status = %d, want %d", i, probeResp.Code, http.StatusNotFound) + } + } + + entries, errRead := os.ReadDir(logsDir) + if errRead != nil { + t.Fatalf("read logs dir: %v", errRead) + } + var names []string + for _, entry := range entries { + names = append(names, entry.Name()) + } + if len(names) != 1 || !strings.HasPrefix(names[0], "error-v1-messages-") { + t.Fatalf("expected only the original error-v1-messages-*.log dump to survive, got: %v", names) + } + }) + + t.Run("unrouted 404 does not create error log when request-log is true", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(true, logsDir, "", 5) + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/messages", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + router.NoRoute(func(c *gin.Context) { + c.AbortWithStatus(http.StatusNotFound) + }) + + probeReq := httptest.NewRequest(http.MethodPost, "/v1/files", strings.NewReader(`{"purpose":"user_data"}`)) + probeReq.Header.Set("Content-Type", "application/json") + probeResp := httptest.NewRecorder() + router.ServeHTTP(probeResp, probeReq) + if probeResp.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d", probeResp.Code, http.StatusNotFound) + } + + entries, errRead := os.ReadDir(logsDir) + if errRead != nil && !os.IsNotExist(errRead) { + t.Fatalf("read logs dir: %v", errRead) + } + if len(entries) != 0 { + var names []string + for _, entry := range entries { + names = append(names, entry.Name()) + } + t.Fatalf("expected 0 files for unrouted 404 when request-log=true, got: %v", names) + } + }) + + t.Run("routed 404 still creates error log when request-log is false", func(t *testing.T) { + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(false, logsDir, "", 5) + + router := gin.New() + router.Use(RequestLoggingMiddleware(logger)) + router.POST("/v1/messages", func(c *gin.Context) { + c.JSON(http.StatusNotFound, gin.H{"error": "model not found"}) + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"missing-model"}`)) + req.Header.Set("Content-Type", "application/json") + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) + if resp.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d", resp.Code, http.StatusNotFound) + } + + entries, errRead := os.ReadDir(logsDir) + if errRead != nil { + t.Fatalf("read logs dir: %v", errRead) + } + if len(entries) != 1 || !strings.HasPrefix(entries[0].Name(), "error-v1-messages-") { + t.Fatalf("expected 1 error-v1-messages-*.log for routed 404, got: %+v", entries) + } + }) + + t.Run("unrouted streaming 404 does not initialize stream writer or leak temp files when request-log is true", func(t *testing.T) { + logsDir := t.TempDir() + fileLogger := logging.NewFileRequestLogger(true, logsDir, "", 5) + spy := &spyRequestLogger{underlying: fileLogger} + + router := gin.New() + router.Use(RequestLoggingMiddleware(spy)) + router.NoRoute(func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Writer.WriteHeader(http.StatusNotFound) + _, _ = c.Writer.Write([]byte("event: error\ndata: not found\n\n")) + }) + + probeReq := httptest.NewRequest(http.MethodPost, "/v1/files", strings.NewReader(`{"stream":true}`)) + probeReq.Header.Set("Content-Type", "application/json") + probeResp := httptest.NewRecorder() + router.ServeHTTP(probeResp, probeReq) + if probeResp.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d", probeResp.Code, http.StatusNotFound) + } + if spy.lastStreamingRequestID != "" || spy.lastStreamWriter != nil { + t.Fatalf("expected LogStreamingRequest not to be called for unrouted 404, got requestID=%q writer=%v", spy.lastStreamingRequestID, spy.lastStreamWriter) + } + + entries, errRead := os.ReadDir(logsDir) + if errRead != nil && !os.IsNotExist(errRead) { + t.Fatalf("read logs dir: %v", errRead) + } + if len(entries) != 0 { + var names []string + for _, entry := range entries { + names = append(names, entry.Name()) + } + t.Fatalf("expected 0 files/temp dirs for unrouted streaming 404, got: %v", names) + } + }) + + t.Run("unrouted streaming 200 to uncommitted 404 transition discards stream writer without creating error log", func(t *testing.T) { + logsDir := t.TempDir() + fileLogger := logging.NewFileRequestLogger(true, logsDir, "", 5) + + router := gin.New() + router.Use(RequestLoggingMiddleware(fileLogger)) + router.NoRoute(func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Writer.WriteHeader(http.StatusOK) + c.Writer.WriteHeader(http.StatusNotFound) + _, _ = c.Writer.Write([]byte("data: not found\n\n")) + }) + + probeReq := httptest.NewRequest(http.MethodPost, "/v1/files", strings.NewReader(`{"stream":true}`)) + probeReq.Header.Set("Content-Type", "application/json") + probeResp := httptest.NewRecorder() + router.ServeHTTP(probeResp, probeReq) + if probeResp.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d", probeResp.Code, http.StatusNotFound) + } + + entries, errRead := os.ReadDir(logsDir) + if errRead != nil && !os.IsNotExist(errRead) { + t.Fatalf("read logs dir: %v", errRead) + } + if len(entries) != 0 { + var names []string + for _, entry := range entries { + names = append(names, entry.Name()) + } + t.Fatalf("expected 0 files/temp dirs after discarded streaming 404, got: %v", names) + } + }) +} diff --git a/internal/api/middleware/response_writer.go b/internal/api/middleware/response_writer.go index 72c977df121..57323930e1e 100644 --- a/internal/api/middleware/response_writer.go +++ b/internal/api/middleware/response_writer.go @@ -165,18 +165,22 @@ func (w *ResponseWriterWrapper) closeConnectionIfRequestBodyUnconsumed(statusCod } func (w *ResponseWriterWrapper) shouldBufferResponseBody() bool { + status := w.effectiveStatusCode() + apiErrors := extractAPIResponseErrors(w.ginContext) + if status < http.StatusBadRequest { + if apiErrStatus := actionableAPIResponseErrorStatus(apiErrors); apiErrStatus >= http.StatusBadRequest { + status = apiErrStatus + } + } + if !w.panicked && isUnroutedNotFound(w.ginContext, status, apiErrors) { + return false + } if w.logger != nil && w.logger.IsEnabled() { return true } if !w.logOnErrorOnly { return false } - status := w.effectiveStatusCode() - if status < http.StatusBadRequest { - if apiErrStatus := actionableAPIResponseErrorStatus(extractAPIResponseErrors(w.ginContext)); apiErrStatus >= http.StatusBadRequest { - status = apiErrStatus - } - } return status >= http.StatusBadRequest && status != clienterror.StatusClientClosedRequest } @@ -232,25 +236,29 @@ func (w *ResponseWriterWrapper) WriteHeader(statusCode int) { // If streaming, initialize streaming log writer once and refresh uncommitted status updates. if !w.finalized && w.isStreaming && w.logger != nil && w.logger.IsEnabled() { if w.streamWriter == nil { - streamWriter, err := w.logger.LogStreamingRequest( - w.requestInfo.URL, - w.requestInfo.Method, - w.requestInfo.Headers, - w.currentCapturedRequestBody(), - w.requestInfo.RequestID, - ) - if err == nil { - chunkChan := make(chan []byte, 100) // Buffered channel for async writes - doneChan := make(chan struct{}) - w.streamWriter = streamWriter - w.chunkChannel = chunkChan - w.streamDone = doneChan - - // Start async chunk processor with captured locals so Finalize clearing fields cannot race. - go w.processStreamingChunks(doneChan, chunkChan, streamWriter) - - // Write status immediately - _ = streamWriter.WriteStatus(statusCode, w.headers) + if !w.panicked && isUnroutedNotFound(w.ginContext, statusCode, extractAPIResponseErrors(w.ginContext)) { + w.isStreaming = false + } else { + streamWriter, err := w.logger.LogStreamingRequest( + w.requestInfo.URL, + w.requestInfo.Method, + w.requestInfo.Headers, + w.currentCapturedRequestBody(), + w.requestInfo.RequestID, + ) + if err == nil { + chunkChan := make(chan []byte, 100) // Buffered channel for async writes + doneChan := make(chan struct{}) + w.streamWriter = streamWriter + w.chunkChannel = chunkChan + w.streamDone = doneChan + + // Start async chunk processor with captured locals so Finalize clearing fields cannot race. + go w.processStreamingChunks(doneChan, chunkChan, streamWriter) + + // Write status immediately + _ = streamWriter.WriteStatus(statusCode, w.headers) + } } } else { _ = w.streamWriter.WriteStatus(statusCode, w.headers) @@ -397,9 +405,10 @@ func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error { apiRequestSource := w.extractAPIRequestSource(c) apiResponseSource := w.extractAPIResponseSource(c) apiWebsocketTimelineSource := w.extractAPIWebsocketTimelineSource(c) - if !w.logger.IsEnabled() && !forceLog { + if (!w.logger.IsEnabled() && !forceLog) || (!w.panicked && isUnroutedNotFound(c, finalStatusCode, slicesAPIResponseError)) { + errDiscard := w.discardStreamingLogWriter() cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource, apiWebsocketTimelineSource) - return nil + return errDiscard } if w.isStreaming && w.streamWriter != nil { @@ -823,6 +832,26 @@ func mergeFileBodySource(payload []byte, source *logging.FileBodySource) ([]byte return buf.Bytes(), nil } +func (w *ResponseWriterWrapper) discardStreamingLogWriter() error { + if w == nil || w.streamWriter == nil { + return nil + } + if w.chunkChannel != nil { + close(w.chunkChannel) + w.chunkChannel = nil + } + if w.streamDone != nil { + <-w.streamDone + w.streamDone = nil + } + streamWriter := w.streamWriter + w.streamWriter = nil + if discarder, ok := streamWriter.(interface{ Discard() error }); ok { + return discarder.Discard() + } + return streamWriter.Close() +} + func cleanupFileBodySources(sources ...*logging.FileBodySource) { for _, source := range sources { if source == nil { @@ -886,6 +915,50 @@ func isContextCanceled(c *gin.Context) bool { return ctx != nil && errors.Is(ctx.Err(), context.Canceled) } +func isUnroutedNotFound(c *gin.Context, statusCode int, apiErrors []*interfaces.ErrorMessage) bool { + if statusCode != http.StatusNotFound { + return false + } + if hasActionableAPIResponseErrors(apiErrors) { + return false + } + if c == nil || c.FullPath() != "" { + return false + } + if hasRecordedUpstreamActivity(c) { + return false + } + return true +} + +func hasRecordedUpstreamActivity(c *gin.Context) bool { + if c == nil { + return false + } + for _, key := range []string{"API_REQUEST", "API_RESPONSE", "API_WEBSOCKET_TIMELINE"} { + if value, exists := c.Get(key); exists { + if data, ok := value.([]byte); ok && len(data) > 0 { + return true + } + } + } + if value, exists := c.Get(logging.DeferredAPIRequestContextKey); exists { + if requests, ok := value.([]logging.DeferredAPIRequest); ok && len(requests) > 0 { + return true + } + } + for _, key := range []string{ + logging.APIRequestSourceContextKey, + logging.APIResponseSourceContextKey, + logging.APIWebsocketTimelineSourceContextKey, + } { + if source := extractFileBodySource(c, key); source != nil && source.HasPayload() { + return true + } + } + return false +} + func hasActionableError(c *gin.Context, statusCode int, apiErrors []*interfaces.ErrorMessage) bool { if hasActionableAPIResponseErrors(apiErrors) { return true @@ -893,6 +966,9 @@ func hasActionableError(c *gin.Context, statusCode int, apiErrors []*interfaces. if statusCode == clienterror.StatusClientClosedRequest { return false } + if isUnroutedNotFound(c, statusCode, apiErrors) { + return false + } if isContextCanceled(c) && statusCode < http.StatusBadRequest { return false } diff --git a/internal/api/middleware/response_writer_test.go b/internal/api/middleware/response_writer_test.go index 70814fc2536..5565a69ee84 100644 --- a/internal/api/middleware/response_writer_test.go +++ b/internal/api/middleware/response_writer_test.go @@ -284,6 +284,17 @@ func TestHasActionableError(t *testing.T) { apiErrors: []*interfaces.ErrorMessage{{StatusCode: 0, Error: fmt.Errorf("read: %w", context.Canceled)}}, want: false, }, + { + name: "404 unrouted without api errors", + statusCode: http.StatusNotFound, + want: false, + }, + { + name: "404 with actionable upstream api error", + statusCode: http.StatusNotFound, + apiErrors: []*interfaces.ErrorMessage{{StatusCode: http.StatusNotFound, Error: errors.New("upstream 404")}}, + want: true, + }, } for _, tc := range tests { @@ -384,3 +395,119 @@ func TestFinalizeIncludes500InForceLog(t *testing.T) { t.Fatalf("expected 1 logged call for 500 status, got: %v", logger.loggedCalls) } } + +func TestFinalizeExcludesUnrouted404FromForceLog(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/files", nil) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + statusCode: http.StatusNotFound, + requestInfo: &RequestInfo{ + URL: "/v1/files", + Method: "POST", + RequestID: "req-404", + Timestamp: time.Now(), + }, + } + + if err := wrapper.Finalize(c); err != nil { + t.Fatalf("Finalize error: %v", err) + } + if len(logger.loggedCalls) != 0 { + t.Fatalf("expected 0 logged calls for unrouted 404, got %d: %v", len(logger.loggedCalls), logger.loggedCalls) + } +} + +func TestFinalizeLogsUnrouted404WhenUpstreamEvidenceOrPanicExists(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("deferred api request recorded", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/files", nil) + c.Set(logging.DeferredAPIRequestContextKey, []logging.DeferredAPIRequest{ + func() []byte { return []byte("=== API REQUEST 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 err := wrapper.Finalize(c); err != nil { + t.Fatalf("Finalize error: %v", err) + } + if len(logger.loggedCalls) != 1 || logger.loggedCalls[0] != http.StatusNotFound { + t.Fatalf("expected 1 logged call when deferred API request exists, got: %v", logger.loggedCalls) + } + }) + + t.Run("file body source with payload recorded", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/files", nil) + + source, errSource := logging.NewFileBodySourceInDir(t.TempDir(), "api-request") + if errSource != nil { + t.Fatalf("NewFileBodySourceInDir: %v", errSource) + } + defer func() { _ = source.Cleanup() }() + part, errPart := source.CreatePart("req") + if errPart != nil { + t.Fatalf("CreatePart: %v", errPart) + } + if _, errWrite := part.WriteString("upstream request payload"); errWrite != nil { + t.Fatalf("WriteString: %v", errWrite) + } + if errClose := part.Close(); errClose != nil { + t.Fatalf("Close: %v", errClose) + } + c.Set(logging.APIRequestSourceContextKey, source) + + 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 err := wrapper.Finalize(c); err != nil { + t.Fatalf("Finalize error: %v", err) + } + if len(logger.loggedCalls) != 1 || logger.loggedCalls[0] != http.StatusNotFound { + t.Fatalf("expected 1 logged call when file body source payload exists, got: %v", logger.loggedCalls) + } + }) + + t.Run("panicked handler on 404 still logs", func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/files", nil) + + logger := &recordingRequestLogger{enabled: false} + wrapper := &ResponseWriterWrapper{ + ResponseWriter: c.Writer, + logger: logger, + logOnErrorOnly: true, + panicked: true, + statusCode: http.StatusNotFound, + requestInfo: &RequestInfo{URL: "/v1/files", Method: http.MethodPost, Timestamp: time.Now()}, + } + if err := wrapper.Finalize(c); err != nil { + t.Fatalf("Finalize error: %v", err) + } + if len(logger.loggedCalls) != 1 || logger.loggedCalls[0] != http.StatusNotFound { + t.Fatalf("expected 1 logged call when panicked=true, got: %v", logger.loggedCalls) + } + }) +} diff --git a/internal/auth/claude/anthropic_auth.go b/internal/auth/claude/anthropic_auth.go index b3b2618fd38..c6dd901e0f2 100644 --- a/internal/auth/claude/anthropic_auth.go +++ b/internal/auth/claude/anthropic_auth.go @@ -34,6 +34,10 @@ const ( RedirectURI = "http://localhost:54545/callback" ClaudeOAuthScope = "user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload" + // Native Claude Code includes the Console scope when authorizing a new login. + // Keep the established token refresh scope unchanged. + claudeOAuthAuthorizeScope = "org:create_api_key " + ClaudeOAuthScope + claudeRefreshMinBackoff = 5 * time.Second claudeRefreshMaxBackoff = 5 * time.Minute claudeRefreshTimeout = 30 * time.Second @@ -330,7 +334,7 @@ func (o *ClaudeAuth) GenerateAuthURL(state string, pkceCodes *PKCECodes) (string "client_id": {ClientID}, "response_type": {"code"}, "redirect_uri": {RedirectURI}, - "scope": {ClaudeOAuthScope}, + "scope": {claudeOAuthAuthorizeScope}, "code_challenge": {pkceCodes.CodeChallenge}, "code_challenge_method": {"S256"}, "state": {state}, diff --git a/internal/auth/claude/anthropic_auth_test.go b/internal/auth/claude/anthropic_auth_test.go index aaa8e2ddd77..90def0d818a 100644 --- a/internal/auth/claude/anthropic_auth_test.go +++ b/internal/auth/claude/anthropic_auth_test.go @@ -129,6 +129,7 @@ func TestExchangeCodeForTokensUsesNative220ControlPlaneShape(t *testing.T) { var order []string headers := make(map[string]http.Header) var tokenBody []byte + state := strings.Repeat("Ab9-_", 8) + "AbA" auth := &ClaudeAuth{ httpClient: &http.Client{ @@ -168,7 +169,7 @@ func TestExchangeCodeForTokensUsesNative220ControlPlaneShape(t *testing.T) { }, } - bundle, errExchange := auth.ExchangeCodeForTokens(t.Context(), "auth-code", "state-value", &PKCECodes{CodeVerifier: "verifier"}) + bundle, errExchange := auth.ExchangeCodeForTokens(t.Context(), "auth-code", state, &PKCECodes{CodeVerifier: "verifier"}) if errExchange != nil { t.Fatalf("ExchangeCodeForTokens() error = %v", errExchange) } @@ -184,7 +185,7 @@ func TestExchangeCodeForTokensUsesNative220ControlPlaneShape(t *testing.T) { } // Key order mirrors the captured native exchange body. - wantBody := `{"grant_type":"authorization_code","code":"auth-code","redirect_uri":"` + RedirectURI + `","client_id":"` + ClientID + `","code_verifier":"verifier","state":"state-value"}` + wantBody := `{"grant_type":"authorization_code","code":"auth-code","redirect_uri":"` + RedirectURI + `","client_id":"` + ClientID + `","code_verifier":"verifier","state":"` + state + `"}` if got := string(tokenBody); got != wantBody { t.Fatalf("exchange body = %q, want %q", got, wantBody) } @@ -499,7 +500,7 @@ func TestRefreshTokensUsesNative220ControlPlaneShape(t *testing.T) { if errRead != nil { t.Fatal(errRead) } - wantBody := `{"client_id":"` + ClientID + `","grant_type":"refresh_token","refresh_token":"` + refreshToken + `","scope":"` + ClaudeOAuthScope + `"}` + wantBody := `{"client_id":"` + ClientID + `","grant_type":"refresh_token","refresh_token":"` + refreshToken + `","scope":"user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload"}` if got := string(body); got != wantBody { t.Fatalf("refresh body = %q, want %q", got, wantBody) } diff --git a/internal/auth/claude/oauth_authorize_test.go b/internal/auth/claude/oauth_authorize_test.go new file mode 100644 index 00000000000..7f213d40b80 --- /dev/null +++ b/internal/auth/claude/oauth_authorize_test.go @@ -0,0 +1,50 @@ +package claude + +import ( + "net/url" + "strings" + "testing" + + "github.com/MoeForge/moego/base64" +) + +func TestGenerateStateMatchesNativeClaudeCode(t *testing.T) { + seen := make(map[string]struct{}) + for range 4 { + state, errGenerate := GenerateState() + if errGenerate != nil { + t.Fatalf("GenerateState() error = %v", errGenerate) + } + if len(state) != 43 { + t.Fatalf("state length = %d, want 43", len(state)) + } + decoded, errDecode := base64.RawURLEncoding.DecodeString(state) + if errDecode != nil || len(decoded) != 32 { + t.Fatalf("state must encode 32 random bytes as unpadded base64url: decoded length = %d, error = %v", len(decoded), errDecode) + } + if _, duplicate := seen[state]; duplicate { + t.Fatal("GenerateState() reused a state") + } + seen[state] = struct{}{} + } +} + +func TestGenerateAuthURLUsesNativeClaudeCodeScopes(t *testing.T) { + state := strings.Repeat("Ab9-_", 8) + "AbA" + auth := &ClaudeAuth{} + authURL, returnedState, errGenerate := auth.GenerateAuthURL(state, &PKCECodes{CodeChallenge: "challenge"}) + if errGenerate != nil { + t.Fatalf("GenerateAuthURL() error = %v", errGenerate) + } + parsed, errParse := url.Parse(authURL) + if errParse != nil { + t.Fatalf("parse authorization URL: %v", errParse) + } + const wantScope = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload" + if got := parsed.Query().Get("scope"); got != wantScope { + t.Errorf("authorization scope = %q, want %q", got, wantScope) + } + if returnedState != state || parsed.Query().Get("state") != state { + t.Error("authorization URL did not preserve the supplied state") + } +} diff --git a/internal/auth/claude/pkce.go b/internal/auth/claude/pkce.go index e8b11f9a9c9..5c71aa22a91 100644 --- a/internal/auth/claude/pkce.go +++ b/internal/auth/claude/pkce.go @@ -11,6 +11,15 @@ import ( "github.com/MoeForge/moego/base64" ) +// GenerateState returns a 43-character OAuth state matching native Claude Code. +func GenerateState() (string, error) { + stateBytes := make([]byte, 32) + if _, err := rand.Read(stateBytes); err != nil { + return "", fmt.Errorf("failed to generate Claude OAuth state: %w", err) + } + return base64.RawURLEncoding.EncodeToString(stateBytes), nil +} + // GeneratePKCECodes generates a PKCE code verifier and challenge pair // following RFC 7636 specifications for OAuth 2.0 PKCE extension. // This provides additional security for the OAuth flow by ensuring that diff --git a/internal/auth/codex/openai.go b/internal/auth/codex/openai.go index a942fd0d3fa..92cd6cbb783 100644 --- a/internal/auth/codex/openai.go +++ b/internal/auth/codex/openai.go @@ -31,16 +31,17 @@ type CodexTokenData struct { // CodexSIWCTokenData holds the OAuth token information obtained from OpenAI's Sign in with ChatGPT (SIWC) flow. type CodexSIWCTokenData struct { - AccessToken string `json:"access_token"` - RefreshToken string `json:"refresh_token"` - IDToken string `json:"id_token,omitempty"` - ClientID string `json:"client_id"` - Scope string `json:"scope,omitempty"` - Scopes []string `json:"scopes,omitempty"` - Expire string `json:"expired"` - ExpiresAt string `json:"expires_at,omitempty"` - LastRefresh string `json:"last_refresh"` - DeviceID string `json:"device_id,omitempty"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token,omitempty"` + ClientID string `json:"client_id"` + Scope string `json:"scope,omitempty"` + Scopes []string `json:"scopes,omitempty"` + Expire string `json:"expired"` + ExpiresAt string `json:"expires_at,omitempty"` + EarliestRefreshAt any `json:"earliest_refresh_at,omitempty"` + LastRefresh string `json:"last_refresh"` + DeviceID string `json:"device_id,omitempty"` } // ToMap converts CodexSIWCTokenData into a map suitable for credential metadata persistence. @@ -58,6 +59,9 @@ func (d *CodexSIWCTokenData) ToMap() map[string]any { if d.ExpiresAt != "" { out["expires_at"] = d.ExpiresAt } + if d.EarliestRefreshAt != nil { + out["earliest_refresh_at"] = d.EarliestRefreshAt + } if d.IDToken != "" { out["id_token"] = d.IDToken } diff --git a/internal/auth/codex/openai_auth.go b/internal/auth/codex/openai_auth.go index d5246acec50..c1d8737de30 100644 --- a/internal/auth/codex/openai_auth.go +++ b/internal/auth/codex/openai_auth.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/MoeForge/moego/ojson" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" log "github.com/sirupsen/logrus" @@ -35,6 +36,7 @@ const ( type CodexAuth struct { httpClient *http.Client siwcAgentNameHint string + refreshJournal RefreshJournal } var codexRefreshGroup singleflight.Group @@ -59,9 +61,14 @@ func NewCodexAuthWithProxyURL(cfg *config.Config, proxyURL string) *CodexAuth { agentNameHint = cfg.Codex.EffectiveSignInWithChatGPTIdentifier() } sdkCfg.ProxyURL = effectiveProxyURL + var journal RefreshJournal + if cfg != nil && strings.TrimSpace(cfg.AuthDir) != "" { + journal = &FileRefreshJournal{Dir: cfg.AuthDir} + } return &CodexAuth{ httpClient: util.SetProxy(&sdkCfg, &http.Client{}), siwcAgentNameHint: agentNameHint, + refreshJournal: journal, } } @@ -205,7 +212,9 @@ func (o *CodexAuth) RefreshTokens(ctx context.Context, refreshToken string) (*Co result, err, _ := codexRefreshGroup.Do(refreshToken, func() (interface{}, error) { refreshCtx, cancelRefresh := context.WithTimeout(context.WithoutCancel(ctx), codexRefreshTimeout) defer cancelRefresh() - return o.refreshTokensSingleFlight(refreshCtx, refreshToken) + return journalRefresh(refreshCtx, o.journal(ctx), ClientID, refreshToken, func() (*CodexTokenData, error) { + return o.refreshTokensSingleFlight(refreshCtx, refreshToken) + }, decodeRefreshReceipt) }) if err != nil { return nil, err @@ -227,15 +236,15 @@ func (o *CodexAuth) refreshTokensSingleFlight(ctx context.Context, refreshToken req, errReq := http.NewRequestWithContext(ctx, "POST", TokenURL, strings.NewReader(data.Encode())) if errReq != nil { - return nil, fmt.Errorf("failed to create refresh request: %w", errReq) + return nil, &RefreshOutcomeUnknownError{Stage: "request-create", Cause: errReq} } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - resp, errDo := o.httpClient.Do(req) + resp, errDo := o.doRefreshRequest(req) if errDo != nil { - return nil, fmt.Errorf("token refresh request failed: %w", errDo) + return nil, &RefreshOutcomeUnknownError{Stage: "transport", Cause: errDo} } defer func() { if errClose := resp.Body.Close(); errClose != nil { @@ -245,11 +254,15 @@ func (o *CodexAuth) refreshTokensSingleFlight(ctx context.Context, refreshToken body, errRead := io.ReadAll(resp.Body) if errRead != nil { - return nil, fmt.Errorf("failed to read refresh response: %w", errRead) + return nil, &RefreshOutcomeUnknownError{Stage: "response-read", Cause: errRead} } if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("token refresh failed with status %d: %s", resp.StatusCode, string(body)) + errResponse := fmt.Errorf("token refresh failed with status %d: %s", resp.StatusCode, string(body)) + if isNonRetryableRefreshErr(errResponse) { + return nil, errResponse + } + return nil, &RefreshOutcomeUnknownError{Stage: "response-status", Cause: errResponse} } var tokenResp struct { @@ -260,8 +273,16 @@ func (o *CodexAuth) refreshTokensSingleFlight(ctx context.Context, refreshToken ExpiresIn int `json:"expires_in"` } - if errUnmarshal := json.Unmarshal(body, &tokenResp); errUnmarshal != nil { - return nil, fmt.Errorf("failed to parse refresh response: %w", errUnmarshal) + if !ojson.ValidBytes(body) { + return nil, &RefreshOutcomeUnknownError{Stage: "response-decode"} + } + response := ojson.ParseBytes(body) + tokenResp.AccessToken = response.Get("access_token").String() + tokenResp.RefreshToken = response.Get("refresh_token").String() + tokenResp.IDToken = response.Get("id_token").String() + tokenResp.ExpiresIn = int(response.Get("expires_in").Int()) + if response.Get("access_token").Type != ojson.String || strings.TrimSpace(tokenResp.AccessToken) == "" || response.Get("expires_in").Type != ojson.Number || tokenResp.ExpiresIn <= 0 || response.Get("refresh_token").Type != ojson.String || strings.TrimSpace(tokenResp.RefreshToken) == "" { + return nil, &RefreshOutcomeUnknownError{Stage: "response-decode"} } // Extract account ID from ID token @@ -347,6 +368,9 @@ func isNonRetryableRefreshErr(err error) bool { if err == nil { return false } + if IsRefreshOutcomeUnknown(err) { + return true + } raw := strings.ToLower(err.Error()) for _, token := range []string{ "refresh_token_reused", diff --git a/internal/auth/codex/openai_auth_test.go b/internal/auth/codex/openai_auth_test.go index aeaeb5e1c70..9689cbef793 100644 --- a/internal/auth/codex/openai_auth_test.go +++ b/internal/auth/codex/openai_auth_test.go @@ -65,6 +65,9 @@ func TestRefreshTokens_UsesIndependentTimeout(t *testing.T) { func resetCodexRefreshGroupForTest() { codexRefreshGroup = singleflight.Group{} + processRefreshJournal.mu.Lock() + processRefreshJournal.receipts = nil + processRefreshJournal.mu.Unlock() } func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) { @@ -340,10 +343,10 @@ func TestRefreshTokensWithRetry_TerminalGrantErrorsOnlyAttemptOnce(t *testing.T) wantCalls: 1, }, { - name: "transient_503_retries", + name: "unqualified_503_does_not_retry", statusCode: http.StatusServiceUnavailable, body: `{"error":"service_unavailable"}`, - wantCalls: 2, + wantCalls: 1, }, } diff --git a/internal/auth/codex/refresh_journal.go b/internal/auth/codex/refresh_journal.go new file mode 100644 index 00000000000..04485c1272a --- /dev/null +++ b/internal/auth/codex/refresh_journal.go @@ -0,0 +1,244 @@ +package codex + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/MoeForge/moego/ojson" +) + +// RefreshJournal records consumption before sending a refresh grant. Claims never +// expire: an absent response is not evidence that the provider did not rotate it. +// Implementations must atomically create claims in the authoritative store, not +// a disposable local mirror. Receipts contain credentials and require the same +// protection as the credential store. Complete is called only by the claimant. +const PendingRefreshReceipt = `{"stage":"pending"}` + +type RefreshJournal interface { + RefreshClaim(context.Context, string) (bool, error) + RefreshRead(context.Context, string) ([]byte, error) + RefreshComplete(context.Context, string, []byte) error +} + +// RefreshOutcomeUnknownError preserves the stage without exposing transport URLs, +// response bodies, or credentials in diagnostic messages. +type RefreshOutcomeUnknownError struct { + Stage string + Cause error +} + +func (e *RefreshOutcomeUnknownError) Error() string { + return "refresh_outcome_unknown: " + e.Stage + "; reconcile the refresh result or sign in again" +} + +func (e *RefreshOutcomeUnknownError) Unwrap() error { return e.Cause } + +// IsRefreshOutcomeUnknown also recognizes wrapped partition errors. +func IsRefreshOutcomeUnknown(err error) bool { + var outcome *RefreshOutcomeUnknownError + return errors.As(err, &outcome) +} + +type refreshJournalContextKey struct{} + +// WithRefreshJournal binds refresh attempts to their authoritative credential store. +func WithRefreshJournal(ctx context.Context, journal RefreshJournal) context.Context { + return context.WithValue(ctx, refreshJournalContextKey{}, journal) +} + +// SetRefreshJournal configures the journal for callers outside the auth manager. +func (o *CodexAuth) SetRefreshJournal(journal RefreshJournal) { o.refreshJournal = journal } + +func refreshGrantKey(clientID, token string) string { + return fmt.Sprintf("%x", sha256.Sum256([]byte(clientID+"\x00"+token))) +} + +func (o *CodexAuth) journal(ctx context.Context) RefreshJournal { + if journal, ok := ctx.Value(refreshJournalContextKey{}).(RefreshJournal); ok && journal != nil { + return journal + } + if o.refreshJournal != nil { + return o.refreshJournal + } + return &processRefreshJournal +} + +// journalRefresh returns a saved successful receipt to late callers. An incomplete +// claim is deliberately not reclaimed, even after a crash or a storage timeout. +func journalRefresh[T any](ctx context.Context, journal RefreshJournal, clientID, grant string, run func() (T, error), decode func(ojson.Result) (T, error)) (T, error) { + var zero T + key := refreshGrantKey(clientID, grant) + claimed, errClaim := journal.RefreshClaim(ctx, key) + if errClaim != nil { + return zero, &RefreshOutcomeUnknownError{Stage: "claim", Cause: errClaim} + } + if !claimed { + var value T + var expires string + seen := make(map[string]bool) + for len(seen) < 1024 { + if seen[key] { + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-cycle"} + } + seen[key] = true + receipt, errRead := journal.RefreshRead(ctx, key) + if errRead != nil { + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-read", Cause: errRead} + } + if len(receipt) == 0 && len(seen) > 1 { + return usableRefreshReceipt(value, expires) + } + if !ojson.ValidBytes(receipt) { + return zero, &RefreshOutcomeUnknownError{Stage: "pending"} + } + doc := ojson.ParseBytes(receipt) + if stage := doc.Get("stage").String(); stage != "" { + return zero, &RefreshOutcomeUnknownError{Stage: stage} + } + if code := doc.Get("rejected").String(); code != "" { + return zero, fmt.Errorf("token refresh rejected: %s", code) + } + var errDecode error + value, errDecode = decode(doc.Get("value")) + if errDecode != nil { + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-decode", Cause: errDecode} + } + expires = doc.Get("value.expired").String() + nextGrant := doc.Get("value.refresh_token").String() + if nextGrant == "" || refreshGrantKey(clientID, nextGrant) == key { + return usableRefreshReceipt(value, expires) + } + key = refreshGrantKey(clientID, nextGrant) + } + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-chain-limit"} + } + + value, errRun := run() + doc := ojson.NewDoc(nil) + var errSet error + if errRun == nil { + errSet = doc.Set("value", value) + } else { + var outcome *RefreshOutcomeUnknownError + if errors.As(errRun, &outcome) { + errSet = doc.Set("stage", outcome.Stage) + } else { + errSet = doc.Set("rejected", refreshRejectionCode(errRun)) + } + } + if errSet != nil { + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-encode", Cause: errSet} + } + // Use an independent acquisition context for the receipt, just as for the + // exchange itself. Caller cancellation must not discard a successful rotation. + if errComplete := journal.RefreshComplete(ctx, key, doc.Bytes()); errComplete != nil { + return zero, errors.Join(errRun, &RefreshOutcomeUnknownError{Stage: "receipt-write", Cause: errComplete}) + } + return value, errRun +} + +func usableRefreshReceipt[T any](value T, expires string) (T, error) { + expiry, errParse := time.Parse(time.RFC3339, expires) + if errParse != nil || !expiry.After(time.Now()) { + var zero T + return zero, &RefreshOutcomeUnknownError{Stage: "receipt-expired", Cause: errParse} + } + return value, nil +} + +func decodeRefreshReceipt(value ojson.Result) (*CodexTokenData, error) { + if !value.IsObject() || value.Get("access_token").String() == "" || value.Get("expired").String() == "" { + return nil, fmt.Errorf("invalid refresh receipt") + } + return &CodexTokenData{ + AccessToken: value.Get("access_token").String(), + RefreshToken: value.Get("refresh_token").String(), + IDToken: value.Get("id_token").String(), + AccountID: value.Get("account_id").String(), + Email: value.Get("email").String(), + Expire: value.Get("expired").String(), + PlanType: value.Get("plan_type").String(), + }, nil +} + +func decodeSIWCRefreshReceipt(value ojson.Result) (*CodexSIWCTokenData, error) { + if !value.IsObject() || value.Get("access_token").String() == "" || value.Get("refresh_token").String() == "" || value.Get("client_id").String() == "" { + return nil, fmt.Errorf("invalid SIWC refresh receipt") + } + scope := value.Get("scope").String() + return &CodexSIWCTokenData{ + AccessToken: value.Get("access_token").String(), + RefreshToken: value.Get("refresh_token").String(), + IDToken: value.Get("id_token").String(), + ClientID: value.Get("client_id").String(), + Scope: scope, + Scopes: strings.Fields(scope), + Expire: value.Get("expired").String(), + ExpiresAt: value.Get("expires_at").String(), + EarliestRefreshAt: value.Get("earliest_refresh_at").Value(), + LastRefresh: value.Get("last_refresh").String(), + DeviceID: value.Get("device_id").String(), + }, nil +} + +func refreshRejectionCode(err error) string { + raw := strings.ToLower(err.Error()) + for _, code := range []string{"refresh_token_reused", "refresh_token_expired", "refresh_token_invalidated", "token_revoked", "invalid_grant", "invalid_client"} { + if strings.Contains(raw, code) { + return code + } + } + return "401 unauthorized" +} + +// Direct SDK callers without a configured store still cannot replay an unresolved +// grant in this process. Server and utility callers attach a durable journal. +type memoryRefreshJournal struct { + mu sync.Mutex + receipts map[string][]byte +} + +var processRefreshJournal memoryRefreshJournal + +func (j *memoryRefreshJournal) RefreshClaim(_ context.Context, key string) (bool, error) { + j.mu.Lock() + defer j.mu.Unlock() + if j.receipts == nil { + j.receipts = make(map[string][]byte) + } + if _, exists := j.receipts[key]; exists { + return false, nil + } + j.receipts[key] = nil + return true, nil +} + +func (j *memoryRefreshJournal) RefreshRead(_ context.Context, key string) ([]byte, error) { + j.mu.Lock() + defer j.mu.Unlock() + receipt, exists := j.receipts[key] + if exists && receipt == nil { + return []byte(PendingRefreshReceipt), nil + } + return append([]byte(nil), receipt...), nil +} + +func (j *memoryRefreshJournal) RefreshComplete(_ context.Context, key string, receipt []byte) error { + j.mu.Lock() + defer j.mu.Unlock() + j.receipts[key] = append([]byte(nil), receipt...) + return nil +} + +// doRefreshRequest forbids 307/308 redirects from silently replaying a POST body. +func (o *CodexAuth) doRefreshRequest(req *http.Request) (*http.Response, error) { + client := *o.httpClient + client.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + return client.Do(req) +} diff --git a/internal/auth/codex/refresh_journal_file.go b/internal/auth/codex/refresh_journal_file.go new file mode 100644 index 00000000000..e4d1c1195f8 --- /dev/null +++ b/internal/auth/codex/refresh_journal_file.go @@ -0,0 +1,141 @@ +package codex + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + + log "github.com/sirupsen/logrus" +) + +// FileRefreshJournal stores immutable claims outside credential JSON files so +// watcher reloads and credential replacement cannot erase unresolved outcomes. +type FileRefreshJournal struct{ Dir string } + +func (j *FileRefreshJournal) path(key string) (string, error) { + if j == nil || j.Dir == "" || len(key) != 64 { + return "", fmt.Errorf("invalid refresh journal directory or key") + } + for _, char := range key { + if !(char >= '0' && char <= '9' || char >= 'a' && char <= 'f') { + return "", fmt.Errorf("invalid refresh journal key") + } + } + return filepath.Join(j.Dir, ".codex-refresh", key), nil +} + +func (j *FileRefreshJournal) RefreshClaim(ctx context.Context, key string) (bool, error) { + if errContext := ctx.Err(); errContext != nil { + return false, errContext + } + path, errPath := j.path(key) + if errPath != nil { + return false, errPath + } + if errMkdir := os.MkdirAll(filepath.Dir(path), 0o700); errMkdir != nil { + return false, errMkdir + } + if errMkdir := os.Mkdir(path, 0o700); errMkdir != nil { + if errors.Is(errMkdir, os.ErrExist) { + return false, nil + } + return false, errMkdir + } + if errSync := syncRefreshClaimDirectories(path); errSync != nil { + return false, errSync + } + return true, nil +} + +func (j *FileRefreshJournal) RefreshRead(ctx context.Context, key string) ([]byte, error) { + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + path, errPath := j.path(key) + if errPath != nil { + return nil, errPath + } + data, errRead := os.ReadFile(filepath.Join(path, "receipt")) + if errors.Is(errRead, os.ErrNotExist) { + if _, errStat := os.Stat(path); errStat == nil { + return []byte(PendingRefreshReceipt), nil + } else if !errors.Is(errStat, os.ErrNotExist) { + return nil, errStat + } + return nil, nil + } + return data, errRead +} + +func (j *FileRefreshJournal) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + path, errPath := j.path(key) + if errPath != nil { + return errPath + } + file, errCreate := os.CreateTemp(path, ".receipt-") + if errCreate != nil { + return errCreate + } + defer func() { + if errRemove := os.Remove(file.Name()); errRemove != nil && !errors.Is(errRemove, os.ErrNotExist) { + log.Warnf("remove refresh receipt temporary file: %v", errRemove) + } + }() + _, errWrite := file.Write(receipt) + if errWrite == nil { + errWrite = file.Sync() + } + errClose := file.Close() + if errWrite != nil || errClose != nil { + return errors.Join(errWrite, errClose) + } + if errRename := os.Rename(file.Name(), filepath.Join(path, "receipt")); errRename != nil { + return errRename + } + return syncRefreshDirectory(path) +} + +// Sync every directory entry that can have been created by MkdirAll before +// permitting an exchange. Canonicalize symlinks so newly-created target +// ancestors, rather than just the spelling of the configured path, are durable. +func syncRefreshClaimDirectories(path string) error { + canonical, errCanonical := filepath.EvalSymlinks(path) + if errCanonical != nil { + return errCanonical + } + dir, errAbs := filepath.Abs(canonical) + if errAbs != nil { + return errAbs + } + for { + if errSync := syncRefreshDirectory(dir); errSync != nil { + return errSync + } + parent := filepath.Dir(dir) + if parent == dir { + return nil + } + dir = parent + } +} + +func syncRefreshDirectory(path string) error { + // Windows does not support fsync on directory handles. Atomic directory + // creation still survives a process restart without a volatile sidecar. + if runtime.GOOS == "windows" { + return nil + } + dir, errOpen := os.Open(path) + if errOpen != nil { + return errOpen + } + errSync := dir.Sync() + errClose := dir.Close() + return errors.Join(errSync, errClose) +} diff --git a/internal/auth/codex/refresh_journal_file_test.go b/internal/auth/codex/refresh_journal_file_test.go new file mode 100644 index 00000000000..49a71e9bfb1 --- /dev/null +++ b/internal/auth/codex/refresh_journal_file_test.go @@ -0,0 +1,49 @@ +package codex + +import ( + "context" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +func TestFileRefreshClaimRequiresParentSync(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("directory fsync is unavailable on Windows") + } + dir := filepath.Join(t.TempDir(), "auths") + if errMkdir := os.Mkdir(dir, 0o300); errMkdir != nil { + t.Fatal(errMkdir) + } + t.Cleanup(func() { + if errChmod := os.Chmod(dir, 0o700); errChmod != nil { + t.Error(errChmod) + } + }) + if handle, errOpen := os.Open(dir); errOpen == nil { + if errClose := handle.Close(); errClose != nil { + t.Fatal(errClose) + } + t.Skip("test requires directory read permissions to be enforced") + } + journal := &FileRefreshJournal{Dir: dir} + claimed, errClaim := journal.RefreshClaim(context.Background(), strings.Repeat("e", 64)) + if claimed || errClaim == nil { + t.Fatalf("claim allowed OAuth without synchronizing parent directory: claimed=%v, error=%v", claimed, errClaim) + } +} + +func TestFileRefreshClaimCreatesDurableAncestorDirectories(t *testing.T) { + journal := &FileRefreshJournal{Dir: filepath.Join(t.TempDir(), "new", "nested", "auths")} + key := strings.Repeat("f", 64) + claimed, errClaim := journal.RefreshClaim(context.Background(), key) + if !claimed || errClaim != nil { + t.Fatalf("create journal hierarchy: %v, %v", claimed, errClaim) + } + claimed, errClaim = journal.RefreshClaim(context.Background(), key) + if claimed || errClaim != nil { + t.Fatalf("claim was recreated: %v, %v", claimed, errClaim) + } +} diff --git a/internal/auth/codex/refresh_outcome_test.go b/internal/auth/codex/refresh_outcome_test.go new file mode 100644 index 00000000000..4af721f16ae --- /dev/null +++ b/internal/auth/codex/refresh_outcome_test.go @@ -0,0 +1,161 @@ +package codex + +import ( + "context" + "errors" + "io" + "net/http" + "os" + "os/exec" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func TestRefreshOutcomeSurvivesLaterCallAndNewClient(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + cfg := &config.Config{AuthDir: t.TempDir()} + calls := 0 + transport := roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + if errParse := req.ParseForm(); errParse != nil { + t.Fatal(errParse) + } + if req.Form.Get("refresh_token") == "fresh-authorization" { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":"fresh-access","refresh_token":"rotated-fresh","expires_in":3600}`)), Request: req}, nil + } + return nil, errors.New("response lost after possible acceptance") + }) + for attempt := 0; attempt < 2; attempt++ { + auth := NewCodexAuth(cfg) + auth.httpClient.Transport = transport + _, errRefresh := auth.RefreshTokensWithRetry(context.Background(), "unresolved-grant", 1) + if errRefresh == nil { + t.Fatal("expected unresolved refresh error") + } + } + if calls != 1 { + t.Fatalf("unresolved grant replayed by later caller: calls = %d, want 1", calls) + } + auth := NewCodexAuth(cfg) + auth.httpClient.Transport = transport + tokens, errRefresh := auth.RefreshTokens(context.Background(), "fresh-authorization") + if errRefresh != nil || tokens == nil || tokens.AccessToken != "fresh-access" { + t.Fatalf("new authorization must recover: tokens = %v, error = %v", tokens, errRefresh) + } + if calls != 2 { + t.Fatalf("new grant calls = %d, want 2", calls) + } +} + +func TestRefreshJournalSurvivesProcessRestart(t *testing.T) { + if dir := os.Getenv("CLIPROXY_TEST_REFRESH_JOURNAL"); dir != "" { + auth := NewCodexAuth(&config.Config{AuthDir: dir}) + auth.httpClient.Transport = roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("unresolved grant replayed after restart") + return nil, nil + }) + _, errRefresh := auth.RefreshTokens(context.Background(), "restart-grant") + var outcome *RefreshOutcomeUnknownError + if !errors.As(errRefresh, &outcome) || outcome.Stage != "transport" { + t.Fatalf("lost stored outcome: %v", errRefresh) + } + return + } + dir := t.TempDir() + auth := NewCodexAuth(&config.Config{AuthDir: dir}) + auth.httpClient.Transport = roundTripFunc(func(*http.Request) (*http.Response, error) { return nil, io.ErrUnexpectedEOF }) + if _, errRefresh := auth.RefreshTokens(context.Background(), "restart-grant"); !IsRefreshOutcomeUnknown(errRefresh) { + t.Fatal(errRefresh) + } + binary, errExecutable := os.Executable() + if errExecutable != nil { + t.Fatal(errExecutable) + } + cmd := exec.Command(binary, "-test.run=^TestRefreshJournalSurvivesProcessRestart$") + cmd.Env = append(os.Environ(), "CLIPROXY_TEST_REFRESH_JOURNAL="+dir) + if output, errRun := cmd.CombinedOutput(); errRun != nil { + t.Fatalf("restart: %v\n%s", errRun, output) + } +} + +func TestRefreshReceiptFollowsLaterRotation(t *testing.T) { + auth := NewCodexAuth(&config.Config{AuthDir: t.TempDir()}) + calls := 0 + auth.httpClient.Transport = roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + if errParse := req.ParseForm(); errParse != nil { + t.Fatal(errParse) + } + body := `{"access_token":"access-second","refresh_token":"third","expires_in":3600}` + if req.Form.Get("refresh_token") == "first" { + body = `{"access_token":"access-first","refresh_token":"second","expires_in":3600}` + } + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: req}, nil + }) + for _, token := range []string{"first", "second"} { + if _, errRefresh := auth.RefreshTokens(context.Background(), token); errRefresh != nil { + t.Fatal(errRefresh) + } + } + result, errRefresh := auth.RefreshTokens(context.Background(), "first") + if errRefresh != nil || result == nil || result.RefreshToken != "third" || calls != 2 { + t.Fatalf("late caller received stale rotation: %+v, %v, calls=%d", result, errRefresh, calls) + } +} + +func TestSIWCUnknownOutcomeMustNotReplay(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + calls := 0 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: &qualificationLostBody{}, Request: req}, nil + })}} + for i := 0; i < 2; i++ { + _, errRefresh := auth.RefreshSIWCTokensWithRetry(context.Background(), "unresolved-siwc", "issued-client", "device", 3) + if errRefresh == nil { + t.Fatal("expected unknown SIWC outcome") + } + } + if calls != 1 { + t.Fatalf("SIWC unresolved grant calls = %d, want 1", calls) + } +} + +func TestRefreshRedirectCannotReplayGrant(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + calls := 0 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + if calls == 1 { + return &http.Response{StatusCode: 307, Header: http.Header{"Location": []string{"https://auth.openai.com/redirected"}}, Body: io.NopCloser(strings.NewReader("")), Request: req}, nil + } + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":"new","refresh_token":"new-refresh","expires_in":3600}`)), Request: req}, nil + })}} + _, errRefresh := auth.RefreshTokensWithRetry(context.Background(), "redirect-grant", 3) + if !IsRefreshOutcomeUnknown(errRefresh) || calls != 1 { + t.Fatalf("redirect replayed grant: calls=%d, error=%v", calls, errRefresh) + } +} + +func TestRefreshHTTPFailureDoesNotProveReplaySafe(t *testing.T) { + for _, status := range []int{http.StatusTooManyRequests, http.StatusServiceUnavailable} { + t.Run(http.StatusText(status), func(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + calls := 0 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"error":"unavailable"}`)), Request: req}, nil + })}} + _, errRefresh := auth.RefreshTokensWithRetry(context.Background(), "unknown-http-outcome", 2) + if errRefresh == nil || calls != 1 { + t.Fatalf("status %d cannot authorize replay: calls = %d, error = %v", status, calls, errRefresh) + } + }) + } +} diff --git a/internal/auth/codex/refresh_patch.go b/internal/auth/codex/refresh_patch.go index 2e831dbfddf..ee0631f7d32 100644 --- a/internal/auth/codex/refresh_patch.go +++ b/internal/auth/codex/refresh_patch.go @@ -4,7 +4,7 @@ import "reflect" var codexSIWCCredentialKeys = []string{ "access_token", "refresh_token", "id_token", "client_id", "device_id", - "expired", "expires_at", "last_refresh", "scope", "scopes", + "expired", "expires_at", "earliest_refresh_at", "last_refresh", "scope", "scopes", } // CredentialRefreshPatch is an immutable, partition-specific refresh result. @@ -119,9 +119,11 @@ func applyCredentialFields(disk, before, after map[string]any) { disk[key] = value } } - if _, hadExpiresAt := before["expires_at"]; hadExpiresAt { - if _, hasExpiresAt := after["expires_at"]; !hasExpiresAt { - delete(disk, "expires_at") + for _, key := range []string{"expires_at", "earliest_refresh_at"} { + if _, had := before[key]; had { + if _, has := after[key]; !has { + delete(disk, key) + } } } } diff --git a/internal/auth/codex/refresh_receipt_expiry_test.go b/internal/auth/codex/refresh_receipt_expiry_test.go new file mode 100644 index 00000000000..a9dcfafa2cb --- /dev/null +++ b/internal/auth/codex/refresh_receipt_expiry_test.go @@ -0,0 +1,52 @@ +package codex + +import ( + "context" + "io" + "net/http" + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" +) + +func TestExpiredSameGrantReceiptRequiresRecovery(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + journal := &memoryRefreshJournal{} + calls := 0 + auth := &CodexAuth{refreshJournal: journal, httpClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Request: req, Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"same-grant","expires_in":3600}`))}, nil + })}} + ctx := context.Background() + if _, errRefresh := auth.RefreshTokens(ctx, "same-grant"); errRefresh != nil { + t.Fatal(errRefresh) + } + value, errRefresh := auth.RefreshTokens(ctx, "same-grant") + if errRefresh != nil || value == nil || value.AccessToken != "new-access" || calls != 1 { + t.Fatalf("unexpired receipt was not reused: value=%v, error=%v, calls=%d", value, errRefresh, calls) + } + key := refreshGrantKey(ClientID, "same-grant") + receipt, errRead := journal.RefreshRead(ctx, key) + if errRead != nil { + t.Fatal(errRead) + } + // Advance receipt age deterministically without a wall-clock sleep. + doc := ojson.NewDoc(receipt) + if errSet := doc.Set("value.expired", "2000-01-01T00:00:00Z"); errSet != nil { + t.Fatal(errSet) + } + if errComplete := journal.RefreshComplete(ctx, key, doc.Bytes()); errComplete != nil { + t.Fatal(errComplete) + } + for attempt := 0; attempt < 2; attempt++ { + value, errRefresh := auth.RefreshTokens(ctx, "same-grant") + if !IsRefreshOutcomeUnknown(errRefresh) || value != nil { + t.Fatalf("expired receipt returned success: value=%v, error=%v", value, errRefresh) + } + if calls != 1 { + t.Fatalf("replayed consumed grant: upstream calls=%d", calls) + } + } +} diff --git a/internal/auth/codex/relay_response_loss_qualification_test.go b/internal/auth/codex/relay_response_loss_qualification_test.go new file mode 100644 index 00000000000..906cfee1307 --- /dev/null +++ b/internal/auth/codex/relay_response_loss_qualification_test.go @@ -0,0 +1,45 @@ +package codex + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" +) + +type qualificationLostBody struct{} + +func (*qualificationLostBody) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } +func (*qualificationLostBody) Close() error { return nil } + +func TestRelayQualificationAmbiguousRefreshMustNotReplay(t *testing.T) { + for _, stage := range []string{"accepted-response-read", "accepted-json-decode", "ambiguous-transport"} { + t.Run(stage, func(t *testing.T) { + resetCodexRefreshGroupForTest() + defer resetCodexRefreshGroupForTest() + var calls atomic.Int32 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls.Add(1) + if stage == "ambiguous-transport" { + return nil, errors.New("synthetic reset after request acceptance unknown") + } + var body io.ReadCloser = &qualificationLostBody{} + if stage == "accepted-json-decode" { + body = io.NopCloser(strings.NewReader(`{"access_token":`)) + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: body, Request: req}, nil + })}} + _, err := auth.RefreshTokensWithRetry(context.Background(), "synthetic-one-use-refresh-token", 3) + if err == nil { + t.Fatal("expected an ambiguous refresh error") + } + t.Logf("OBSERVED_REFRESH_TRANSPORT_CALLS=%d", calls.Load()) + if calls.Load() != 1 { + t.Fatalf("one-use refresh replay: got %d transport calls, want 1", calls.Load()) + } + }) + } +} diff --git a/internal/auth/codex/siwc_auth.go b/internal/auth/codex/siwc_auth.go index f8587090603..42c4cc3efcd 100644 --- a/internal/auth/codex/siwc_auth.go +++ b/internal/auth/codex/siwc_auth.go @@ -3,18 +3,20 @@ package codex import ( "context" "crypto/rand" - "encoding/json" "errors" "fmt" - "github.com/MoeForge/moego/base64" "io" + "math" "net/http" "net/url" "regexp" "slices" + "strconv" "strings" "time" + "github.com/MoeForge/moego/base64" + "github.com/MoeForge/moego/ojson" log "github.com/sirupsen/logrus" "golang.org/x/sync/singleflight" ) @@ -95,15 +97,34 @@ func (o *CodexAuth) GenerateSIWCAuthURL(state, nonce, deviceID, redirectURI stri return fmt.Sprintf("%s?%s", SIWCAuthorizeURL, params.Encode()), nil } -type siwcTokenResponse struct { - AccessToken string `json:"access_token"` - RefreshToken string `json:"refresh_token"` - IDToken string `json:"id_token"` - Scope string `json:"scope"` - ExpiresIn int `json:"expires_in"` +// GenerateSIWCReauthURL reuses the issued registration when replacing an unusable grant. +func (o *CodexAuth) GenerateSIWCReauthURL(clientID, state, nonce, deviceID, redirectURI string, pkceCodes *PKCECodes) (string, error) { + clientID = strings.TrimSpace(clientID) + if clientID == "" || clientID == SIWCDynamicClientID { + return "", fmt.Errorf("SIWC reauthorization requires the saved issued client ID") + } + rawURL, errGenerate := o.GenerateSIWCAuthURL(state, nonce, deviceID, redirectURI, pkceCodes) + if errGenerate != nil { + return "", errGenerate + } + parsed, errParse := url.Parse(rawURL) + if errParse != nil { + return "", fmt.Errorf("failed to parse SIWC authorization URL: %w", errParse) + } + query := parsed.Query() + query.Set("client_id", clientID) + query.Del("agent_name_hint") + parsed.RawQuery = query.Encode() + return parsed.String(), nil } -func (o *CodexAuth) requestSIWCToken(ctx context.Context, form url.Values, clientID, deviceID string, requireIDToken bool) (*CodexSIWCTokenData, error) { +func (o *CodexAuth) requestSIWCToken(ctx context.Context, form url.Values, clientID, deviceID string, requireIDToken bool) (result *CodexSIWCTokenData, errResult error) { + stage := "request-create" + defer func() { + if form.Get("grant_type") == "refresh_token" && errResult != nil && !isNonRetryableRefreshErr(errResult) { + errResult = &RefreshOutcomeUnknownError{Stage: stage, Cause: errResult} + } + }() req, errReq := http.NewRequestWithContext(ctx, http.MethodPost, siwcTokenEndpointURL, strings.NewReader(form.Encode())) if errReq != nil { return nil, fmt.Errorf("failed to create SIWC token request: %w", errReq) @@ -111,7 +132,12 @@ func (o *CodexAuth) requestSIWCToken(ctx context.Context, form url.Values, clien req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - resp, errDo := o.httpClient.Do(req) + stage = "transport" + request := o.httpClient.Do + if form.Get("grant_type") == "refresh_token" { + request = o.doRefreshRequest + } + resp, errDo := request(req) if errDo != nil { return nil, fmt.Errorf("SIWC token request failed: %w", errDo) } @@ -121,57 +147,100 @@ func (o *CodexAuth) requestSIWCToken(ctx context.Context, form url.Values, clien } }() + stage = "response-read" body, errRead := io.ReadAll(resp.Body) if errRead != nil { return nil, fmt.Errorf("failed to read SIWC token response: %w", errRead) } if resp.StatusCode != http.StatusOK { + stage = "response-status" + if requestID := strings.TrimSpace(resp.Header.Get("x-request-id")); requestID != "" { + return nil, fmt.Errorf("OpenAI OAuth token request failed (%d, request_id=%q): %s", resp.StatusCode, requestID, strings.TrimSpace(string(body))) + } return nil, fmt.Errorf("OpenAI OAuth token request failed (%d): %s", resp.StatusCode, strings.TrimSpace(string(body))) } - var tokenResp siwcTokenResponse - if errUnmarshal := json.Unmarshal(body, &tokenResp); errUnmarshal != nil { - return nil, fmt.Errorf("failed to parse SIWC token response: %w", errUnmarshal) + stage = "response-decode" + tokenResp := ojson.ParseBytes(body) + if !ojson.ValidBytes(body) || !tokenResp.IsObject() { + return nil, fmt.Errorf("failed to parse SIWC token response: expected a JSON object") } - if strings.TrimSpace(tokenResp.AccessToken) == "" { + accessToken := tokenResp.Get("access_token") + if accessToken.Type != ojson.String || strings.TrimSpace(accessToken.Str) == "" { return nil, fmt.Errorf("OpenAI OAuth token response has invalid access_token") } - if strings.TrimSpace(tokenResp.RefreshToken) == "" { + refreshToken := tokenResp.Get("refresh_token") + if refreshToken.Type != ojson.String || strings.TrimSpace(refreshToken.Str) == "" { return nil, fmt.Errorf("OpenAI OAuth token response has invalid refresh_token") } - if requireIDToken && strings.TrimSpace(tokenResp.IDToken) == "" { + idToken := tokenResp.Get("id_token") + if idToken.Type != ojson.Null && idToken.Type != ojson.String { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid id_token") + } + if requireIDToken && strings.TrimSpace(idToken.Str) == "" { return nil, fmt.Errorf("OpenAI OAuth token response did not contain an ID token") } - if strings.TrimSpace(tokenResp.Scope) == "" { + scope := tokenResp.Get("scope") + // An omitted, unchanged scope on refresh must not discard a rotated refresh token. + if (requireIDToken || scope.Exists()) && (scope.Type != ojson.String || strings.TrimSpace(scope.Str) == "") { return nil, fmt.Errorf("OpenAI OAuth token response has invalid scope") } - if tokenResp.ExpiresIn <= 0 { + expiresInResult := tokenResp.Get("expires_in") + maxExpiresInSeconds := int64(^uint64(0)>>1) / int64(time.Second) + var expiresIn int64 + switch expiresInResult.Type { + case ojson.Number: + f, errParse := strconv.ParseFloat(expiresInResult.Raw, 64) + if errParse != nil || f <= 0 || math.IsInf(f, 0) || math.IsNaN(f) || f != math.Trunc(f) || f > float64(maxExpiresInSeconds) { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") + } + expiresIn = int64(f) + case ojson.String: + raw := strings.TrimSpace(expiresInResult.Str) + if raw == "" { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") + } + i, errParse := strconv.ParseInt(raw, 10, 64) + if errParse == nil { + if i <= 0 || i > maxExpiresInSeconds { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") + } + expiresIn = i + break + } + f, errParse := strconv.ParseFloat(raw, 64) + if errParse != nil || f <= 0 || math.IsInf(f, 0) || math.IsNaN(f) || f != math.Trunc(f) || f > float64(maxExpiresInSeconds) { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") + } + expiresIn = int64(f) + default: return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") } - scopes := strings.Fields(tokenResp.Scope) - if !slices.Contains(scopes, SIWCDirectTokenScope) { + scopes := strings.Fields(scope.Str) + if scope.Exists() && !slices.Contains(scopes, SIWCDirectTokenScope) { return nil, fmt.Errorf("OpenAI OAuth grant did not include %s", SIWCDirectTokenScope) } now := time.Now() - rawTTL := time.Duration(tokenResp.ExpiresIn) * time.Second + rawTTL := time.Duration(expiresIn) * time.Second ttl := rawTTL - SIWCExpiryMargin if ttl <= 0 { ttl = rawTTL } return &CodexSIWCTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - IDToken: tokenResp.IDToken, - ClientID: strings.TrimSpace(clientID), - Scope: strings.TrimSpace(tokenResp.Scope), - Scopes: scopes, - Expire: now.Add(ttl).Format(time.RFC3339), - ExpiresAt: now.Add(rawTTL).Format(time.RFC3339), - LastRefresh: now.Format(time.RFC3339), - DeviceID: strings.TrimSpace(deviceID), + AccessToken: accessToken.Str, + RefreshToken: refreshToken.Str, + IDToken: idToken.Str, + ClientID: strings.TrimSpace(clientID), + Scope: strings.TrimSpace(scope.Str), + Scopes: scopes, + Expire: now.Add(ttl).Format(time.RFC3339), + ExpiresAt: now.Add(rawTTL).Format(time.RFC3339), + EarliestRefreshAt: tokenResp.Get("earliest_refresh_at").Value(), + LastRefresh: now.Format(time.RFC3339), + DeviceID: strings.TrimSpace(deviceID), }, nil } @@ -257,13 +326,18 @@ func (o *CodexAuth) RefreshSIWCTokens(ctx context.Context, refreshToken, clientI result, errSF, _ := siwcRefreshGroup.Do(sfKey, func() (any, error) { refreshCtx, cancelRefresh := context.WithTimeout(context.WithoutCancel(ctx), codexRefreshTimeout) defer cancelRefresh() + if errReadiness := siwcRefreshReadiness(refreshCtx); errReadiness != nil { + return nil, errReadiness + } form := url.Values{ "grant_type": {"refresh_token"}, "client_id": {clientID}, "refresh_token": {refreshToken}, "resource": {SIWCResource}, } - return o.requestSIWCToken(refreshCtx, form, clientID, deviceID, false) + return journalRefresh(refreshCtx, o.journal(ctx), clientID, refreshToken, func() (*CodexSIWCTokenData, error) { + return o.requestSIWCToken(refreshCtx, form, clientID, deviceID, false) + }, decodeSIWCRefreshReceipt) }) if errSF != nil { return nil, WrapSIWCRefreshError(errSF) @@ -293,6 +367,9 @@ func (o *CodexAuth) RefreshSIWCTokensWithRetry(ctx context.Context, refreshToken if errRefresh == nil { return tokenData, nil } + if errors.Is(errRefresh, context.Canceled) || errors.Is(errRefresh, ErrSIWCRefreshDeferred) { + return nil, errRefresh + } if isNonRetryableRefreshErr(errRefresh) { log.Warnf("SIWC token refresh attempt %d failed with non-retryable error: %v", attempt+1, errRefresh) return nil, errRefresh diff --git a/internal/auth/codex/siwc_power_darwin.go b/internal/auth/codex/siwc_power_darwin.go new file mode 100644 index 00000000000..99b747f1b8c --- /dev/null +++ b/internal/auth/codex/siwc_power_darwin.go @@ -0,0 +1,47 @@ +package codex + +import ( + "context" + "fmt" + "os/exec" + "slices" + "strings" + "time" +) + +// ReadSIWCSystemState reads the macOS power capabilities used by SIWC renewal. +var ReadSIWCSystemState = func(ctx context.Context) ([]byte, error) { + return exec.CommandContext(ctx, "/usr/bin/pmset", "-g", "systemstate").Output() +} + +func siwcRefreshReadiness(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + // shortcut: checks entry state; use native sleep coordination if in-flight sleep still breaks renewal. + state, errState := ReadSIWCSystemState(ctx) + if errState != nil { + return fmt.Errorf("SIWC renewal deferred: macOS power state unavailable: %w", errState) + } + foundCapabilities := false + var fields []string + for _, line := range strings.Split(string(state), "\n") { + capabilities, ok := strings.CutPrefix(strings.TrimSpace(line), "Current System Capabilities are:") + if !ok { + continue + } + foundCapabilities = true + fields = strings.Fields(capabilities) + if slices.Contains(fields, "Graphics") && slices.Contains(fields, "Network") { + return nil + } + } + if !foundCapabilities { + raw := strings.TrimSpace(string(state)) + raw = strings.Join(strings.Fields(strings.ReplaceAll(raw, "\n", " ")), " ") + if len(raw) > 200 { + raw = raw[:200] + "..." + } + return fmt.Errorf("SIWC renewal deferred: unexpected pmset output: %q: %w", raw, ErrSIWCRefreshDeferred) + } + return fmt.Errorf("SIWC renewal deferred until interactive macOS wake: missing Graphics/Network (%v): %w", fields, ErrSIWCRefreshDeferred) +} diff --git a/internal/auth/codex/siwc_power_darwin_test.go b/internal/auth/codex/siwc_power_darwin_test.go new file mode 100644 index 00000000000..c255f8dfae6 --- /dev/null +++ b/internal/auth/codex/siwc_power_darwin_test.go @@ -0,0 +1,77 @@ +package codex + +import ( + "context" + "errors" + "io" + "net/http" + "os" + "strings" + "testing" +) + +func TestMain(m *testing.M) { + ReadSIWCSystemState = func(context.Context) ([]byte, error) { + return []byte("Current System Capabilities are: CPU Graphics Audio Network\nCurrent Power State: 4\n"), nil + } + os.Exit(m.Run()) +} + +func TestSIWCRefreshDefersDuringMaintenanceWakeThenRenews(t *testing.T) { + for _, state := range []string{ + "Current System Capabilities are: CPU Network\nCurrent Power State: 2\n", + "Current System Capabilities are: CPU Graphics Audio\n", + "Graphics Network", + } { + t.Run(strings.TrimSpace(state), func(t *testing.T) { + previous := ReadSIWCSystemState + t.Cleanup(func() { ReadSIWCSystemState = previous }) + currentState := state + checks := 0 + ReadSIWCSystemState = func(context.Context) ([]byte, error) { + checks++ + return []byte(currentState), nil + } + calls := 0 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"new-refresh","expires_in":3600,"scope":"chatgpt.tokens.use.direct"}`))}, nil + })}} + auth.SetRefreshJournal(&memoryRefreshJournal{}) + token, errRefresh := auth.RefreshSIWCTokensWithRetry(context.Background(), "sleep-refresh", "oaiapp_sleep", "", 3) + if token != nil || errRefresh == nil || calls != 0 || checks != 1 { + t.Fatalf("maintenance wake: token returned = %t, error = %v, requests = %d, checks = %d; want deferral without token requests or retries", token != nil, errRefresh, calls, checks) + } + if errors.Is(errRefresh, context.Canceled) { + t.Fatalf("maintenance wake: refresh returned context.Canceled; want a dedicated deferral error: %v", errRefresh) + } + if strings.TrimSpace(state) == "Graphics Network" && !strings.Contains(errRefresh.Error(), "unexpected pmset output") { + t.Fatalf("maintenance wake: error = %v; want parse diagnostics for unexpected pmset output", errRefresh) + } + var partition *SIWCRefreshError + if !errors.As(errRefresh, &partition) { + t.Fatalf("deferral lost the SIWC error partition: %v", errRefresh) + } + currentState = "Current System Capabilities are: CPU Graphics Audio Network\nCurrent Power State: 4\n" + token, errRefresh = auth.RefreshSIWCTokensWithRetry(context.Background(), "sleep-refresh", "oaiapp_sleep", "", 3) + if errRefresh != nil || token == nil || token.RefreshToken != "new-refresh" || calls != 1 { + t.Fatalf("interactive wake: token returned = %t, error = %v, requests = %d; want successful renewal", token != nil, errRefresh, calls) + } + }) + } +} + +func TestSIWCRefreshDoesNotConsumeGrantWhenPowerStateUnavailable(t *testing.T) { + previous := ReadSIWCSystemState + t.Cleanup(func() { ReadSIWCSystemState = previous }) + errProbe := errors.New("power-state probe unavailable") + ReadSIWCSystemState = func(context.Context) ([]byte, error) { return nil, errProbe } + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("token request sent without a known power state") + return nil, nil + })}} + _, errRefresh := auth.RefreshSIWCTokens(context.Background(), "refresh", "oaiapp_probe", "") + if !errors.Is(errRefresh, errProbe) { + t.Fatalf("power-state error was lost: %v", errRefresh) + } +} diff --git a/internal/auth/codex/siwc_power_other.go b/internal/auth/codex/siwc_power_other.go new file mode 100644 index 00000000000..c036209e4e1 --- /dev/null +++ b/internal/auth/codex/siwc_power_other.go @@ -0,0 +1,7 @@ +//go:build !darwin + +package codex + +import "context" + +func siwcRefreshReadiness(context.Context) error { return nil } diff --git a/internal/auth/codex/siwc_refresh_deferral.go b/internal/auth/codex/siwc_refresh_deferral.go new file mode 100644 index 00000000000..e364ac1315a --- /dev/null +++ b/internal/auth/codex/siwc_refresh_deferral.go @@ -0,0 +1,6 @@ +package codex + +import "errors" + +// ErrSIWCRefreshDeferred indicates SIWC renewal should be deferred and retried later. +var ErrSIWCRefreshDeferred = errors.New("siwc refresh deferred") diff --git a/internal/auth/codex/siwc_renewal_test.go b/internal/auth/codex/siwc_renewal_test.go new file mode 100644 index 00000000000..6002a4234a4 --- /dev/null +++ b/internal/auth/codex/siwc_renewal_test.go @@ -0,0 +1,140 @@ +package codex + +import ( + "context" + "io" + "net/http" + "reflect" + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" +) + +func TestSIWCTokenResponseValidatesScope(t *testing.T) { + for _, tc := range []struct { + name, scope string + exchange bool + wantError bool + }{ + {name: "refresh omits unchanged scope"}, + {name: "refresh explicitly clears scope", scope: `""`, wantError: true}, + {name: "refresh returns null scope", scope: `null`, wantError: true}, + {name: "refresh removes direct permission", scope: `"openid profile"`, wantError: true}, + {name: "refresh returns malformed scope", scope: `123`, wantError: true}, + {name: "refresh grants direct permission", scope: `"openid chatgpt.tokens.use.direct"`}, + {name: "exchange requires explicit scope", exchange: true, wantError: true}, + {name: "exchange grants direct permission", exchange: true, scope: `"openid chatgpt.tokens.use.direct"`}, + } { + t.Run(tc.name, func(t *testing.T) { + doc := ojson.NewDocString(`{"access_token":"access","refresh_token":"refresh","id_token":"id","expires_in":3600}`) + if tc.scope != "" { + if errSet := doc.SetRaw("scope", tc.scope); errSet != nil { + t.Fatal(errSet) + } + } + body := doc.String() + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + })}} + auth.SetRefreshJournal(&memoryRefreshJournal{}) + var err error + if tc.exchange { + _, err = auth.ExchangeSIWCCodeForTokens(context.Background(), "code", "oaiapp_scope", "", "", &PKCECodes{CodeVerifier: "verifier"}) + } else { + _, err = auth.RefreshSIWCTokens(context.Background(), "refresh", "oaiapp_scope", "") + } + if (err != nil) != tc.wantError { + t.Fatalf("error = %v, wantError = %v", err, tc.wantError) + } + }) + } +} + +func TestSIWCTokenResponseValidatesExpiresIn(t *testing.T) { + for _, tc := range []struct { + name string + expiresInRaw string + wantError bool + }{ + {name: "accept float json number", expiresInRaw: "3600.0"}, + {name: "accept scientific notation", expiresInRaw: "3.6e3"}, + {name: "reject fractional value", expiresInRaw: "3600.5", wantError: true}, + {name: "reject oversized seconds", expiresInRaw: "1e20", wantError: true}, + {name: "reject infinity string", expiresInRaw: `"Inf"`, wantError: true}, + } { + t.Run(tc.name, func(t *testing.T) { + doc := ojson.NewDocString(`{"access_token":"access","refresh_token":"refresh","expires_in":3600}`) + if errSet := doc.SetRaw("expires_in", tc.expiresInRaw); errSet != nil { + t.Fatal(errSet) + } + body := doc.String() + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + })}} + auth.SetRefreshJournal(&memoryRefreshJournal{}) + token, errRefresh := auth.RefreshSIWCTokens(context.Background(), "refresh", "oaiapp_expires_in", "") + if tc.wantError { + if errRefresh == nil || token != nil { + t.Fatalf("error = %v, token returned = %t; want error", errRefresh, token != nil) + } + return + } + if errRefresh != nil || token == nil { + t.Fatalf("error = %v, token returned = %t; want success", errRefresh, token != nil) + } + }) + } +} + +func TestSIWCRefreshReceiptPreservesServerWindow(t *testing.T) { + for _, window := range []string{`"2026-11-07T13:07:09Z"`, `1794056829`, `null`} { + t.Run(window, func(t *testing.T) { + doc := ojson.NewDocString(`{"access_token":"access","refresh_token":"rotated","expires_in":3600}`) + if errSet := doc.SetRaw("earliest_refresh_at", window); errSet != nil { + t.Fatal(errSet) + } + body := doc.String() + calls := 0 + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + })}} + auth.SetRefreshJournal(&memoryRefreshJournal{}) + for attempt := 0; attempt < 2; attempt++ { + token, errRefresh := auth.RefreshSIWCTokens(context.Background(), "receipt-grant", "oaiapp_receipt", "") + if errRefresh != nil { + t.Fatal(errRefresh) + } + if !reflect.DeepEqual(token.EarliestRefreshAt, ojson.Parse(window).Value()) { + t.Fatalf("attempt %d: earliest_refresh_at = %v, want %s", attempt, token.EarliestRefreshAt, window) + } + } + if calls != 1 { + t.Fatalf("receipt was not reused: requests = %d, want 1", calls) + } + }) + } +} + +func TestSIWCRefreshErrorRetainsRequestID(t *testing.T) { + auth := &CodexAuth{httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"X-Request-Id": []string{"req-siwc-test"}}, + Body: io.NopCloser(strings.NewReader(`{"error":"invalid_grant"}`)), + }, nil + })}} + _, errRefresh := auth.RefreshSIWCTokens(context.Background(), "refresh", "oaiapp_request_id", "") + if errRefresh == nil || !strings.Contains(errRefresh.Error(), "req-siwc-test") || !strings.Contains(errRefresh.Error(), "invalid_grant") { + t.Fatalf("upstream correlation details missing: %v", errRefresh) + } +} + +func TestGenerateSIWCReauthURLRejectsUnissuedClient(t *testing.T) { + for _, clientID := range []string{"", " ", SIWCDynamicClientID} { + if _, errGenerate := (&CodexAuth{}).GenerateSIWCReauthURL(clientID, "state", "nonce", "e61bbe28-07ef-466d-8e5d-a344f94ab305", "", &PKCECodes{CodeChallenge: "challenge"}); errGenerate == nil { + t.Errorf("accepted an unissued client ID %q", clientID) + } + } +} diff --git a/internal/cache/antigravity_reasoning_replay_cache.go b/internal/cache/antigravity_reasoning_replay_cache.go index a53017f400f..d1aaad73e9d 100644 --- a/internal/cache/antigravity_reasoning_replay_cache.go +++ b/internal/cache/antigravity_reasoning_replay_cache.go @@ -4,13 +4,13 @@ import ( "bytes" "context" "crypto/rand" - "encoding/json" "fmt" "sort" "strings" "sync" "time" + "github.com/MoeForge/moego/base64" "github.com/MoeForge/moego/ojson" homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" @@ -165,9 +165,28 @@ func GetAntigravityReasoningReplayItemsRequired(ctx context.Context, modelName, return items, found, errGet } +// PeekAntigravityReasoningReplayItemsRequired retrieves existing replay items +// without creating an absence reservation tombstone on a cache miss. Use this +// for read-only paths such as CountTokens that never commit or delete a turn. +func PeekAntigravityReasoningReplayItemsRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, bool, error) { + items, _, found, errGet := getAntigravityReasoningReplayItemsWithSnapshot(ctx, modelName, sessionKey, false) + return items, found, errGet +} + +// PeekAntigravityReasoningReplayItemsWithSnapshotRequired retrieves existing +// replay items and their snapshot without creating an absence reservation +// tombstone on a cache miss. +func PeekAntigravityReasoningReplayItemsWithSnapshotRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, AntigravityReasoningReplaySnapshot, bool, error) { + return getAntigravityReasoningReplayItemsWithSnapshot(ctx, modelName, sessionKey, false) +} + // GetAntigravityReasoningReplayItemsWithSnapshotRequired retrieves replay items // and the exact cache state that guarded this request. func GetAntigravityReasoningReplayItemsWithSnapshotRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, AntigravityReasoningReplaySnapshot, bool, error) { + return getAntigravityReasoningReplayItemsWithSnapshot(ctx, modelName, sessionKey, true) +} + +func getAntigravityReasoningReplayItemsWithSnapshot(ctx context.Context, modelName, sessionKey string, reserveOnMiss bool) ([][]byte, AntigravityReasoningReplaySnapshot, bool, error) { key := antigravityReasoningReplayCacheKey(modelName, sessionKey) if key == "" { return nil, AntigravityReasoningReplaySnapshot{}, false, nil @@ -190,6 +209,9 @@ func GetAntigravityReasoningReplayItemsWithSnapshotRequired(ctx context.Context, found = true break } + if !reserveOnMiss { + return nil, AntigravityReasoningReplaySnapshot{loaded: true}, false, nil + } reservation := newAntigravityReasoningReplayTombstone() swapped, errReserve := client.KVCompareAndSwap(ctx, kvKey, nil, false, reservation, AntigravityReasoningReplayCacheTTL) if errReserve != nil { @@ -233,11 +255,17 @@ func GetAntigravityReasoningReplayItemsWithSnapshotRequired(ctx context.Context, defer antigravityReasoningReplayMu.Unlock() entry, ok := antigravityReasoningReplayEntries[key] if !ok { + if !reserveOnMiss { + return nil, AntigravityReasoningReplaySnapshot{loaded: true, evictionEpoch: antigravityReasoningReplayEvictionEpoch}, false, nil + } return nil, reserveAntigravityReasoningReplayAbsentLocked(key, now), false, nil } if now.Sub(entry.Timestamp) > AntigravityReasoningReplayCacheTTL { antigravityReasoningReplayEvictionEpoch++ delete(antigravityReasoningReplayEntries, key) + if !reserveOnMiss { + return nil, AntigravityReasoningReplaySnapshot{loaded: true, evictionEpoch: antigravityReasoningReplayEvictionEpoch}, false, nil + } return nil, reserveAntigravityReasoningReplayAbsentLocked(key, now), false, nil } entry.Timestamp = now @@ -425,6 +453,48 @@ func newAntigravityReasoningReplayGeneration() string { return fmt.Sprintf("%x", nonce[:]) } +func marshalReasoningReplayRawJSONArray(items [][]byte) ([]byte, error) { + doc := ojson.NewDocString("[]") + for i, item := range items { + trimmed := bytes.TrimSpace(item) + if len(trimmed) == 0 || !ojson.ValidBytes(trimmed) { + return nil, fmt.Errorf("invalid replay item json at index %d", i) + } + if errSet := doc.SetRawBytes("-1", trimmed); errSet != nil { + return nil, errSet + } + } + return doc.Bytes(), nil +} + +func unmarshalReasoningReplayRawJSONArray(raw []byte) ([][]byte, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || !ojson.ValidBytes(trimmed) { + return nil, fmt.Errorf("invalid reasoning replay KV payload (please upgrade CLIProxyAPI to a matching version)") + } + root := util.ParseGJSONBytesNoCopy(trimmed) + if !root.IsArray() { + return nil, fmt.Errorf("unsupported reasoning replay KV format (please upgrade CLIProxyAPI to a matching version)") + } + elems := root.Array() + items := make([][]byte, 0, len(elems)) + for _, elem := range elems { + switch { + case elem.IsObject(): + items = append(items, []byte(elem.Raw)) + case elem.Type == ojson.String: + decoded, errDecode := base64.StdEncoding.DecodeString(elem.String()) + if errDecode != nil { + return nil, fmt.Errorf("invalid legacy base64 reasoning replay item (please upgrade CLIProxyAPI): %w", errDecode) + } + items = append(items, decoded) + default: + return nil, fmt.Errorf("unsupported reasoning replay KV element type (please upgrade CLIProxyAPI to a matching version)") + } + } + return items, nil +} + func marshalAntigravityReasoningReplayHomeValue(items [][]byte, branch string) ([]byte, error) { if branch == "" { branch = newAntigravityReasoningReplayGeneration() @@ -436,11 +506,14 @@ func marshalAntigravityReasoningReplayHomeValue(items [][]byte, branch string) ( stored := make([][]byte, 0, len(items)+1) stored = append(stored, doc.Bytes()) stored = append(stored, items...) - return json.Marshal(stored) + return marshalReasoningReplayRawJSONArray(stored) } func decodeAntigravityReasoningReplayHomeValue(raw []byte) (items [][]byte, deleted bool, generation, branch string, ok bool) { - if errUnmarshal := json.Unmarshal(raw, &items); errUnmarshal != nil { + var errUnmarshal error + items, errUnmarshal = unmarshalReasoningReplayRawJSONArray(raw) + if errUnmarshal != nil { + log.Warnf("antigravity reasoning replay KV decode failed: %v", errUnmarshal) return nil, false, "", "", false } if len(items) == 0 || strings.TrimSpace(ojson.GetBytesNoCopy(items[0], "type").String()) != antigravityReasoningReplayGenerationItemType { @@ -470,7 +543,7 @@ func newAntigravityReasoningReplayTombstone() []byte { _ = doc.Set("type", antigravityReasoningReplayGenerationItemType) _ = doc.Set("generation", newAntigravityReasoningReplayGeneration()) _ = doc.Set("branch", newAntigravityReasoningReplayGeneration()) - raw, _ := json.Marshal([][]byte{doc.Bytes()}) + raw, _ := marshalReasoningReplayRawJSONArray([][]byte{doc.Bytes()}) return raw } diff --git a/internal/cache/antigravity_reasoning_replay_cache_test.go b/internal/cache/antigravity_reasoning_replay_cache_test.go index c0215c7f919..9cef703e984 100644 --- a/internal/cache/antigravity_reasoning_replay_cache_test.go +++ b/internal/cache/antigravity_reasoning_replay_cache_test.go @@ -3,12 +3,12 @@ package cache import ( "bytes" "context" - "encoding/json" "fmt" "sync" "testing" "time" + "github.com/MoeForge/moego/base64" "github.com/MoeForge/moego/ojson" homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" ) @@ -494,21 +494,29 @@ func TestAntigravityReasoningReplayLocalAbsenceReservationsStayWithinEntryBound( } } -func TestAntigravityReasoningReplayHomeWritesRemainLegacyArrayReadable(t *testing.T) { +func TestAntigravityReasoningReplayHomeWritesUnwrappedJSONArrayAndReadsLegacyBase64(t *testing.T) { client := newFakeAntigravityReasoningReplayKVClient() useFakeAntigravityReasoningReplayKVClient(t, client, true) - const model, session = "gemini-3.6-flash-high", "home-legacy-readable" + const model, session = "gemini-3.6-flash-high", "home-unwrapped-json-array" item := antigravityReplayTestItem("legacy-readable-signature-123456") if !CacheAntigravityReasoningReplayItems(model, session, [][]byte{item}) { t.Fatal("Home write failed") } raw := client.values[antigravityReasoningReplayKVKey(model, session)] - var legacyItems [][]byte - if errUnmarshal := json.Unmarshal(raw, &legacyItems); errUnmarshal != nil { - t.Fatalf("new Home value is not readable as legacy [][]byte: %v", errUnmarshal) + root := ojson.ParseBytes(raw) + if !root.IsArray() || len(root.Array()) != 2 || !root.Get("0").IsObject() || !root.Get("1").IsObject() { + t.Fatalf("new Home value must be an unwrapped JSON object array without base64 double encoding: %s", raw) } - if len(legacyItems) != 2 || ojson.GetBytes(legacyItems[0], "type").String() != antigravityReasoningReplayGenerationItemType || !bytes.Contains(legacyItems[1], []byte("legacy-readable-signature")) { - t.Fatalf("legacy-readable Home array malformed: %q", legacyItems) + if root.Get("0.type").String() != antigravityReasoningReplayGenerationItemType || root.Get("1.thoughtSignature").String() != "legacy-readable-signature-123456" { + t.Fatalf("unwrapped Home array malformed: %s", raw) + } + + // Verify backward-compatible decoding of legacy base64-encoded [][]byte arrays. + legacyRaw := []byte(`["` + base64.StdEncoding.EncodeToString(item) + `"]`) + client.values[antigravityReasoningReplayKVKey(model, session)] = legacyRaw + gotLegacy, _, foundLegacy, errLegacy := GetAntigravityReasoningReplayItemsWithSnapshotRequired(context.Background(), model, session) + if errLegacy != nil || !foundLegacy || len(gotLegacy) != 1 || ojson.GetBytes(gotLegacy[0], "thoughtSignature").String() != "legacy-readable-signature-123456" { + t.Fatalf("legacy base64 Home array read = %q, found=%v err=%v", gotLegacy, foundLegacy, errLegacy) } } @@ -518,7 +526,7 @@ func TestAntigravityReasoningReplayHomeReadNormalizesAndRejectsMixedInvalidChain const model, session = "gemini-3.6-flash-high", "home-validation" key := antigravityReasoningReplayKVKey(model, session) valid := []byte(`{"type":"function_call_part","name":"run","args":{"b":2,"a":1},"targetOccurrence":1,"thoughtSignature":"valid-home-signature-123456"}`) - raw, errMarshal := json.Marshal([][]byte{valid}) + raw, errMarshal := marshalReasoningReplayRawJSONArray([][]byte{valid}) if errMarshal != nil { t.Fatal(errMarshal) } @@ -534,7 +542,7 @@ func TestAntigravityReasoningReplayHomeReadNormalizesAndRejectsMixedInvalidChain t.Fatalf("valid Home read expire count = %d, want 1", client.expireCount) } - invalidRaw, errInvalidMarshal := json.Marshal([][]byte{valid, []byte(`{"type":"unknown"}`)}) + invalidRaw, errInvalidMarshal := marshalReasoningReplayRawJSONArray([][]byte{valid, []byte(`{"type":"unknown"}`)}) if errInvalidMarshal != nil { t.Fatal(errInvalidMarshal) } diff --git a/internal/cache/codex_reasoning_replay_cache.go b/internal/cache/codex_reasoning_replay_cache.go index f9898789816..8ff8d640f95 100644 --- a/internal/cache/codex_reasoning_replay_cache.go +++ b/internal/cache/codex_reasoning_replay_cache.go @@ -2,7 +2,6 @@ package cache import ( "context" - "encoding/json" "sort" "strings" "sync" @@ -88,7 +87,7 @@ func CacheCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, se log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errClient) return false } - raw, errMarshal := json.Marshal(normalized) + raw, errMarshal := marshalReasoningReplayRawJSONArray(normalized) if errMarshal != nil { log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errMarshal) return false @@ -146,13 +145,15 @@ func AppendCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, s } var existing [][]byte if found { - if errUnmarshal := json.Unmarshal(existingRaw, &existing); errUnmarshal != nil { - log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errUnmarshal) + var errUnmarshal error + existing, errUnmarshal = unmarshalReasoningReplayRawJSONArray(existingRaw) + if errUnmarshal != nil { + log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:* (please upgrade CLIProxyAPI): %v", errUnmarshal) return false } } combined := appendCodexReasoningReplayTurn(existing, normalized) - raw, errMarshal := json.Marshal(combined) + raw, errMarshal := marshalReasoningReplayRawJSONArray(combined) if errMarshal != nil { log.Errorf("home kv best-effort codex reasoning replay append failed prefix=cpa:codex:*: %v", errMarshal) return false @@ -267,8 +268,9 @@ func GetCodexReasoningReplayItemsRequired(ctx context.Context, modelName, sessio if errGet != nil || !found { return nil, false, errGet } - var homeItems [][]byte - if errUnmarshal := json.Unmarshal(raw, &homeItems); errUnmarshal != nil { + homeItems, errUnmarshal := unmarshalReasoningReplayRawJSONArray(raw) + if errUnmarshal != nil { + log.Warnf("codex reasoning replay KV decode failed (please upgrade CLIProxyAPI): %v", errUnmarshal) return nil, false, errUnmarshal } if _, errExpire := client.KVExpire(ctx, codexReasoningReplayKVKey(modelName, sessionKey), CodexReasoningReplayCacheTTL); errExpire != nil { diff --git a/internal/cache/codex_reasoning_replay_cache_test.go b/internal/cache/codex_reasoning_replay_cache_test.go index 441a2d928d4..3a01a89a07e 100644 --- a/internal/cache/codex_reasoning_replay_cache_test.go +++ b/internal/cache/codex_reasoning_replay_cache_test.go @@ -3,14 +3,13 @@ package cache import ( "bytes" "context" - "encoding/json" "errors" "fmt" - "github.com/MoeForge/moego/base64" "sync" "testing" "time" + "github.com/MoeForge/moego/base64" "github.com/MoeForge/moego/ojson" homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" ) @@ -130,7 +129,7 @@ func validCodexReasoningReplayItemForTest(seed byte) []byte { func mustCodexReasoningReplayJSON(t *testing.T, items [][]byte) []byte { t.Helper() - raw, errMarshal := json.Marshal(items) + raw, errMarshal := marshalReasoningReplayRawJSONArray(items) if errMarshal != nil { t.Fatalf("marshal replay items: %v", errMarshal) } diff --git a/internal/cache/xai_reasoning_replay_cache.go b/internal/cache/xai_reasoning_replay_cache.go index b786aeba711..b33614cc2b9 100644 --- a/internal/cache/xai_reasoning_replay_cache.go +++ b/internal/cache/xai_reasoning_replay_cache.go @@ -2,7 +2,6 @@ package cache import ( "context" - "encoding/json" "sort" "strings" "sync" @@ -101,7 +100,7 @@ func StoreXAIReasoningReplayItems(ctx context.Context, modelName, sessionKey str log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errClient) return XAIReasoningReplayStoreBackendError } - raw, errMarshal := json.Marshal(normalized) + raw, errMarshal := marshalReasoningReplayRawJSONArray(normalized) if errMarshal != nil { log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errMarshal) return XAIReasoningReplayStoreBackendError @@ -164,8 +163,9 @@ func GetXAIReasoningReplayItemsRequired(ctx context.Context, modelName, sessionK if errGet != nil || !found { return nil, false, errGet } - var homeItems [][]byte - if errUnmarshal := json.Unmarshal(raw, &homeItems); errUnmarshal != nil { + homeItems, errUnmarshal := unmarshalReasoningReplayRawJSONArray(raw) + if errUnmarshal != nil { + log.Warnf("xai reasoning replay KV decode failed (please upgrade CLIProxyAPI): %v", errUnmarshal) return nil, false, errUnmarshal } if _, errExpire := client.KVExpire(ctx, xaiReasoningReplayKVKey(modelName, sessionKey), XAIReasoningReplayCacheTTL); errExpire != nil { diff --git a/internal/cache/xai_reasoning_replay_cache_test.go b/internal/cache/xai_reasoning_replay_cache_test.go index e3faf53fd42..cdeefe1388c 100644 --- a/internal/cache/xai_reasoning_replay_cache_test.go +++ b/internal/cache/xai_reasoning_replay_cache_test.go @@ -3,12 +3,11 @@ package cache import ( "context" "crypto/sha256" - "encoding/json" "errors" - "github.com/MoeForge/moego/base64" "testing" "time" + "github.com/MoeForge/moego/base64" "github.com/MoeForge/moego/ojson" homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" ) @@ -90,7 +89,7 @@ func useFakeXAIReasoningReplayKVClient(t *testing.T, client *fakeXAIReasoningRep func mustXAIReasoningReplayJSON(t *testing.T, items [][]byte) []byte { t.Helper() - raw, err := json.Marshal(items) + raw, err := marshalReasoningReplayRawJSONArray(items) if err != nil { t.Fatalf("marshal replay items: %v", err) } diff --git a/internal/cmd/anthropic_login.go b/internal/cmd/anthropic_login.go index e2d240b1fbe..f01aa344451 100644 --- a/internal/cmd/anthropic_login.go +++ b/internal/cmd/anthropic_login.go @@ -34,6 +34,7 @@ func DoClaudeLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/cmd/antigravity_login.go b/internal/cmd/antigravity_login.go index e5f06262e13..7dac2c47528 100644 --- a/internal/cmd/antigravity_login.go +++ b/internal/cmd/antigravity_login.go @@ -24,6 +24,7 @@ func DoAntigravityLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/cmd/devin_login.go b/internal/cmd/devin_login.go index b12e8ffc002..284000b3ae7 100644 --- a/internal/cmd/devin_login.go +++ b/internal/cmd/devin_login.go @@ -24,6 +24,7 @@ func DoDevinLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/cmd/kimi_login.go b/internal/cmd/kimi_login.go index 870986130fa..c69382dc823 100644 --- a/internal/cmd/kimi_login.go +++ b/internal/cmd/kimi_login.go @@ -39,6 +39,7 @@ func doKimiLoginWithProvider(cfg *config.Config, options *LoginOptions, provider NoBrowser: options.NoBrowser, Metadata: map[string]string{}, Prompt: options.Prompt, + ProxyURL: options.ProxyURL, } record, savedPath, err := manager.Login(context.Background(), provider, cfg, authOpts) diff --git a/internal/cmd/meta_login.go b/internal/cmd/meta_login.go index 365b6c01b54..37641c25d90 100644 --- a/internal/cmd/meta_login.go +++ b/internal/cmd/meta_login.go @@ -24,6 +24,7 @@ func DoMetaLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/cmd/openai_device_login.go b/internal/cmd/openai_device_login.go index 57415669cbe..040c82f1cb3 100644 --- a/internal/cmd/openai_device_login.go +++ b/internal/cmd/openai_device_login.go @@ -34,6 +34,7 @@ func DoCodexDeviceLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{ codexLoginModeMetadataKey: codexLoginModeDevice, }, diff --git a/internal/cmd/openai_login.go b/internal/cmd/openai_login.go index a9574e10c4e..f0b1b2e52f6 100644 --- a/internal/cmd/openai_login.go +++ b/internal/cmd/openai_login.go @@ -24,6 +24,8 @@ type LoginOptions struct { // Prompt allows the caller to provide interactive input when needed. Prompt func(prompt string) (string, error) + + ProxyURL string } // DoCodexLogin triggers the Codex OAuth flow through the shared authentication manager. @@ -48,6 +50,7 @@ func DoCodexLogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/cmd/siwc_bind.go b/internal/cmd/siwc_bind.go index f186fb0ee0c..fe476b1a053 100644 --- a/internal/cmd/siwc_bind.go +++ b/internal/cmd/siwc_bind.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/MoeForge/moego/ojson" "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/browser" @@ -157,6 +158,17 @@ func DoSIWCBind(cfg *config.Config, target string, options *LoginOptions) { log.Errorf("siwc-bind: %v", errDevice) return } + credentialBytes, errRead := os.ReadFile(authFilePath) + if errRead != nil { + log.Errorf("siwc-bind: failed to read saved registration: %v", errRead) + return + } + registration := ojson.GetBytes(credentialBytes, "siwc.client_id") + if !ojson.ValidBytes(credentialBytes) || (registration.Exists() && registration.Type != ojson.String) { + log.Error("siwc-bind: invalid saved registration") + return + } + savedClientID := strings.TrimSpace(registration.Str) callbackPort := codex.SIWCCallbackPort if options.CallbackPort > 0 { @@ -200,6 +212,9 @@ func DoSIWCBind(cfg *config.Config, target string, options *LoginOptions) { authSvc := codex.NewCodexAuth(cfg) authURL, errURL := authSvc.GenerateSIWCAuthURL(state, nonce, deviceID, redirectURI, pkceCodes) + if savedClientID != "" { + authURL, errURL = authSvc.GenerateSIWCReauthURL(savedClientID, state, nonce, deviceID, redirectURI, pkceCodes) + } if errURL != nil { log.Errorf("siwc-bind: failed to generate authorization URL: %v", errURL) return @@ -305,12 +320,20 @@ waitForCallback: fmt.Println("Sign in with ChatGPT authentication failed: state mismatch") return } - if strings.TrimSpace(result.ClientID) == "" { + issuedClientID := strings.TrimSpace(result.ClientID) + if savedClientID != "" { + if issuedClientID != "" && issuedClientID != savedClientID { + fmt.Println("Sign in with ChatGPT authentication failed: callback client ID does not match the saved registration") + return + } + issuedClientID = savedClientID + } + if issuedClientID == "" || issuedClientID == codex.SIWCDynamicClientID { fmt.Println("Sign in with ChatGPT authentication failed: callback did not contain an issued client ID") return } - siwcData, errExchange := authSvc.ExchangeSIWCCodeForTokens(context.Background(), result.Code, result.ClientID, redirectURI, deviceID, pkceCodes) + siwcData, errExchange := authSvc.ExchangeSIWCCodeForTokens(context.Background(), result.Code, issuedClientID, redirectURI, deviceID, pkceCodes) if errExchange != nil { fmt.Printf("Sign in with ChatGPT token exchange failed: %v\n", errExchange) return diff --git a/internal/cmd/siwc_reauth_test.go b/internal/cmd/siwc_reauth_test.go new file mode 100644 index 00000000000..5f32f348b11 --- /dev/null +++ b/internal/cmd/siwc_reauth_test.go @@ -0,0 +1,143 @@ +package cmd + +import ( + "bufio" + "context" + "io" + "net" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +type siwcBindRoundTripper func(*http.Request) (*http.Response, error) + +func (f siwcBindRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestDoSIWCBindReusesSavedRegistration(t *testing.T) { + for _, tc := range []struct { + name, returned string + wantSuccess bool + }{ + {"callback omits unchanged client", "", true}, + {"callback confirms saved client", "oaiapp_saved", true}, + {"callback changes client", "oaiapp_other", false}, + } { + t.Run(tc.name, func(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "codex.json") + const initial = `{"type":"codex","access_token":"outer","refresh_token":"outer-refresh","siwc":{"access_token":"old","refresh_token":"old-refresh","client_id":"oaiapp_saved","device_id":"e61bbe28-07ef-466d-8e5d-a344f94ab305"}}` + if errWrite := os.WriteFile(path, []byte(initial), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatal(errListen) + } + port := listener.Addr().(*net.TCPAddr).Port + if errClose := listener.Close(); errClose != nil { + t.Fatal(errClose) + } + previousTransport := http.DefaultTransport + t.Cleanup(func() { http.DefaultTransport = previousTransport }) + exchanges := 0 + http.DefaultTransport = siwcBindRoundTripper(func(req *http.Request) (*http.Response, error) { + if req.URL.String() != codex.SIWCTokenURL { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("127.0.0.1"))}, nil + } + exchanges++ + if errParse := req.ParseForm(); errParse != nil { + return nil, errParse + } + if req.Form.Get("client_id") != "oaiapp_saved" { + t.Errorf("token exchange did not use the saved registration") + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":"new","refresh_token":"new-refresh","id_token":"id","expires_in":3600,"scope":"openid chatgpt.tokens.use.direct"}`))}, nil + }) + reader, writer, errPipe := os.Pipe() + if errPipe != nil { + t.Fatal(errPipe) + } + previousStdout := os.Stdout + os.Stdout = writer + t.Cleanup(func() { + os.Stdout = previousStdout + if errClose := reader.Close(); errClose != nil { + t.Error(errClose) + } + }) + done := make(chan struct{}) + go func() { + defer close(done) + defer func() { + if errClose := writer.Close(); errClose != nil { + t.Error(errClose) + } + }() + DoSIWCBind(&config.Config{AuthDir: dir}, path, &LoginOptions{NoBrowser: true, CallbackPort: port}) + }() + scanner := bufio.NewScanner(reader) + foundURL := false + for scanner.Scan() { + if !strings.HasPrefix(scanner.Text(), codex.SIWCAuthorizeURL+"?") { + continue + } + foundURL = true + authURL, errParse := url.Parse(scanner.Text()) + if errParse != nil { + t.Fatal(errParse) + } + query := authURL.Query() + if query.Get("client_id") != "oaiapp_saved" || query.Has("agent_name_hint") || query.Get("ext_agent_host_id") != "urn:uuid:e61bbe28-07ef-466d-8e5d-a344f94ab305" || query.Get("code_challenge") == "" { + t.Error("authorization URL did not preserve the issued registration, host identity, and PKCE") + } + callback := "http://127.0.0.1:" + strconv.Itoa(port) + codex.SIWCCallbackPath + values := url.Values{"code": {"code"}, "state": {query.Get("state")}, "client_id": {tc.returned}} + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, callback+"?"+values.Encode(), nil) + if errRequest != nil { + cancel() + t.Fatal(errRequest) + } + callbackClient := &http.Client{Transport: previousTransport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + response, errCallback := callbackClient.Do(req) + if errCallback == nil { + if errClose := response.Body.Close(); errClose != nil { + t.Error(errClose) + } + } + cancel() + if errCallback != nil { + t.Fatal(errCallback) + } + } + if errScan := scanner.Err(); errScan != nil { + t.Fatal(errScan) + } + <-done + if !foundURL { + t.Fatal("authorization URL was not printed") + } + raw, errRead := os.ReadFile(path) + if errRead != nil { + t.Fatal(errRead) + } + if tc.wantSuccess { + if exchanges != 1 || ojson.GetBytes(raw, "siwc.refresh_token").String() != "new-refresh" || ojson.GetBytes(raw, "siwc.client_id").String() != "oaiapp_saved" { + t.Fatal("reauthorization did not persist the replacement under the saved client") + } + } else if exchanges != 0 || string(raw) != initial { + t.Fatal("mismatched callback exchanged tokens or changed the credential") + } + }) + } +} diff --git a/internal/cmd/xai_login.go b/internal/cmd/xai_login.go index ee4e9e8a545..0cba7b04941 100644 --- a/internal/cmd/xai_login.go +++ b/internal/cmd/xai_login.go @@ -24,6 +24,7 @@ func DoXAILogin(cfg *config.Config, options *LoginOptions) { authOpts := &sdkAuth.LoginOptions{ NoBrowser: options.NoBrowser, CallbackPort: options.CallbackPort, + ProxyURL: options.ProxyURL, Metadata: map[string]string{}, Prompt: promptFn, } diff --git a/internal/config/config_types.go b/internal/config/config_types.go index 7c77b7cdabf..c186d21ab18 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -559,6 +559,8 @@ type RoutingConfig struct { // - "quota": transient errors (408, 5xx, empty stream, buffer limit) fail over for that // request only while keeping the session bound to its original credential; quota and // auth errors (429, 401/402/403, 404/model support, invalid_grant) rebind as normal. + // - "never": never move or fail over an established session to another credential; + // keep the binding and return the error (or retry the same credential) instead. // Ignored when SessionAffinity is false. SessionAffinityRebindOn string `yaml:"session-affinity-rebind-on,omitempty" json:"session-affinity-rebind-on,omitempty"` } diff --git a/internal/logging/request_logger_home.go b/internal/logging/request_logger_home.go index 7aa6c0207ce..b283c44c5bc 100644 --- a/internal/logging/request_logger_home.go +++ b/internal/logging/request_logger_home.go @@ -257,3 +257,15 @@ func (w *homeStreamingLogWriter) Close() error { } return client.RPushRequestLog(context.Background(), raw) } + +func (w *homeStreamingLogWriter) Discard() error { + if w == nil { + return nil + } + if w.chunkChan != nil { + close(w.chunkChan) + <-w.doneChan + w.chunkChan = nil + } + return nil +} diff --git a/internal/logging/request_logger_streaming.go b/internal/logging/request_logger_streaming.go index 7c3505d4d60..15c67163848 100644 --- a/internal/logging/request_logger_streaming.go +++ b/internal/logging/request_logger_streaming.go @@ -258,6 +258,15 @@ func (w *FileStreamingLogWriter) Close() error { return writeErr } +// Discard stops the streaming writer and removes temporary files without creating a final log file. +func (w *FileStreamingLogWriter) Discard() error { + if w == nil { + return nil + } + w.logFilename = "" + return w.Close() +} + // asyncWriter runs in a goroutine to buffer chunks from the channel. // It continuously reads chunks from the channel and appends them to a temp file for later assembly. func (w *FileStreamingLogWriter) asyncWriter() { diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index e0419bb37f6..48771f958d6 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -1142,6 +1142,7 @@ func (a *executorAdapter) Refresh(ctx context.Context, auth *coreauth.Auth) (ref return nil, errRefresh } data := pluginResp.Auth + explicitPortalHint := authDataPortalHint(&data) if strings.TrimSpace(data.Provider) == "" { data.Provider = authProvider(auth) } @@ -1176,7 +1177,7 @@ func (a *executorAdapter) Refresh(ctx context.Context, auth *coreauth.Auth) (ref } data.Attributes = attributes } - preserveAuthOAuthProviderMetadata(&data, auth) + preserveAuthOAuthProviderMetadata(&data, auth, explicitPortalHint) preserveFileAuthPriority(&data, auth) preserveAuthDisabledState(&data, auth) if len(data.StorageJSON) == 0 { diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 90cbd557bfd..9a80469391b 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2793,6 +2793,130 @@ func TestUsageAdapterPropagatesRequestIDAndTraceID(t *testing.T) { } } +func TestUsageAdapterPropagatesModalityTokenDetails(t *testing.T) { + var gotRecord pluginapi.UsageRecord + plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { + gotRecord = record + }) + host := newHostWithRecords(capabilityRecord{ + id: "usage-modality-details", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + UsagePlugin: plugin, + }}, + }) + adapter := &usageAdapter{ + host: host, + pluginID: "usage-modality-details", + plugin: plugin, + } + + textIn := int64(30) + imageIn := int64(50) + audioIn := int64(0) + videoIn := int64(18) + fileIn := int64(12) + cachedText := int64(25) + cachedImage := int64(15) + cachedAudio := int64(0) + textOut := int64(0) + imageOut := int64(1056) + audioOut := int64(42) + + adapter.HandleUsage(context.Background(), coreusage.Record{ + Provider: "openai", + Model: "gpt-image-2.5-flare", + Detail: coreusage.Detail{ + InputTokens: 110, + OutputTokens: 1098, + ReasoningTokens: 10, + CachedTokens: 40, + CacheReadTokens: 40, + CacheCreationTokens: 5, + TotalTokens: 1208, + TextInputTokens: &textIn, + ImageInputTokens: &imageIn, + AudioInputTokens: &audioIn, + VideoInputTokens: &videoIn, + FileInputTokens: &fileIn, + CachedTextTokens: &cachedText, + CachedImageTokens: &cachedImage, + CachedAudioTokens: &cachedAudio, + TextOutputTokens: &textOut, + ImageOutputTokens: &imageOut, + AudioOutputTokens: &audioOut, + }, + }) + + // Mutate caller-owned source values after HandleUsage to verify independent pointer copies. + textIn = 999 + imageIn = 999 + audioIn = 999 + videoIn = 999 + fileIn = 999 + cachedText = 999 + cachedImage = 999 + cachedAudio = 999 + textOut = 999 + imageOut = 999 + audioOut = 999 + + got := gotRecord.Detail + if got.InputTokens != 110 || got.OutputTokens != 1098 || got.ReasoningTokens != 10 || + got.CachedTokens != 40 || got.CacheReadTokens != 40 || got.CacheCreationTokens != 5 || got.TotalTokens != 1208 { + t.Fatalf("aggregate Detail = %+v", got) + } + assertInt64Ptr := func(name string, ptr *int64, src *int64, want int64) { + t.Helper() + if ptr == nil { + t.Fatalf("%s = nil, want %d", name, want) + } + if ptr == src { + t.Fatalf("%s shares pointer %p with source record", name, ptr) + } + if *ptr != want { + t.Fatalf("%s = %d, want %d", name, *ptr, want) + } + } + assertInt64Ptr("TextInputTokens", got.TextInputTokens, &textIn, 30) + assertInt64Ptr("ImageInputTokens", got.ImageInputTokens, &imageIn, 50) + assertInt64Ptr("AudioInputTokens", got.AudioInputTokens, &audioIn, 0) + assertInt64Ptr("VideoInputTokens", got.VideoInputTokens, &videoIn, 18) + assertInt64Ptr("FileInputTokens", got.FileInputTokens, &fileIn, 12) + assertInt64Ptr("CachedTextTokens", got.CachedTextTokens, &cachedText, 25) + assertInt64Ptr("CachedImageTokens", got.CachedImageTokens, &cachedImage, 15) + assertInt64Ptr("CachedAudioTokens", got.CachedAudioTokens, &cachedAudio, 0) + assertInt64Ptr("TextOutputTokens", got.TextOutputTokens, &textOut, 0) + assertInt64Ptr("ImageOutputTokens", got.ImageOutputTokens, &imageOut, 1056) + assertInt64Ptr("AudioOutputTokens", got.AudioOutputTokens, &audioOut, 42) + + // Verify absent counters remain nil and distinct from explicit zero. + adapter.HandleUsage(context.Background(), coreusage.Record{ + Provider: "openai", + Model: "gpt-image-2.5-flare", + Detail: coreusage.Detail{ + InputTokens: 80, + TextInputTokens: coreusage.Int64Ptr(30), + ImageInputTokens: coreusage.Int64Ptr(50), + TextOutputTokens: coreusage.Int64Ptr(0), + }, + }) + got = gotRecord.Detail + if got.TextInputTokens == nil || *got.TextInputTokens != 30 { + t.Fatalf("TextInputTokens = %v, want 30", got.TextInputTokens) + } + if got.ImageInputTokens == nil || *got.ImageInputTokens != 50 { + t.Fatalf("ImageInputTokens = %v, want 50", got.ImageInputTokens) + } + if got.TextOutputTokens == nil || *got.TextOutputTokens != 0 { + t.Fatalf("TextOutputTokens = %v, want explicit 0", got.TextOutputTokens) + } + if got.AudioInputTokens != nil || got.VideoInputTokens != nil || got.FileInputTokens != nil || + got.CachedTextTokens != nil || got.CachedImageTokens != nil || got.CachedAudioTokens != nil || + got.ImageOutputTokens != nil || got.AudioOutputTokens != nil { + t.Fatalf("expected absent modality counters to remain nil, got %+v", got) + } +} + func TestUsageAdapterPreservesExplicitGenerateFalse(t *testing.T) { var gotGenerate bool plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { diff --git a/internal/pluginhost/adapters_usage_translation.go b/internal/pluginhost/adapters_usage_translation.go index 2dcab50abb0..e064fd82cbb 100644 --- a/internal/pluginhost/adapters_usage_translation.go +++ b/internal/pluginhost/adapters_usage_translation.go @@ -640,6 +640,17 @@ func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) CacheReadTokens: record.Detail.CacheReadTokens, CacheCreationTokens: record.Detail.CacheCreationTokens, TotalTokens: record.Detail.TotalTokens, + TextInputTokens: coreusage.CloneInt64Ptr(record.Detail.TextInputTokens), + ImageInputTokens: coreusage.CloneInt64Ptr(record.Detail.ImageInputTokens), + AudioInputTokens: coreusage.CloneInt64Ptr(record.Detail.AudioInputTokens), + VideoInputTokens: coreusage.CloneInt64Ptr(record.Detail.VideoInputTokens), + FileInputTokens: coreusage.CloneInt64Ptr(record.Detail.FileInputTokens), + CachedTextTokens: coreusage.CloneInt64Ptr(record.Detail.CachedTextTokens), + CachedImageTokens: coreusage.CloneInt64Ptr(record.Detail.CachedImageTokens), + CachedAudioTokens: coreusage.CloneInt64Ptr(record.Detail.CachedAudioTokens), + TextOutputTokens: coreusage.CloneInt64Ptr(record.Detail.TextOutputTokens), + ImageOutputTokens: coreusage.CloneInt64Ptr(record.Detail.ImageOutputTokens), + AudioOutputTokens: coreusage.CloneInt64Ptr(record.Detail.AudioOutputTokens), }, ResponseHeaders: cloneHeader(record.ResponseHeaders), }) diff --git a/internal/pluginhost/auth_provider.go b/internal/pluginhost/auth_provider.go index 116e51b81fa..49049a18ea9 100644 --- a/internal/pluginhost/auth_provider.go +++ b/internal/pluginhost/auth_provider.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "github.com/MoeForge/moego/ojson" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" @@ -209,6 +210,20 @@ func (h *Host) HasAuthProvider(provider string) bool { return h.authProviderRecord(provider) != nil } +// HasMultipleAuthProviders reports whether the active owning plugin declares multiple OAuth entries. +func (h *Host) HasMultipleAuthProviders(provider string) bool { + record := h.authProviderRecord(provider) + return record != nil && h.recordHasMultipleAuthProviders(*record) +} + +func (h *Host) recordHasMultipleAuthProviders(record capabilityRecord) bool { + if h == nil || !h.recordCurrent(record) { + return false + } + identifiers, okIdentifiers := h.callRecordAuthProviderIdentifiers(record) + return okIdentifiers && len(identifiers) > 1 +} + // SameAuthProvider reports whether providerA and providerB belong to the same active plugin auth provider. func (h *Host) SameAuthProvider(providerA, providerB string) bool { providerA = normalizeProviderID(providerA) @@ -500,16 +515,12 @@ func (h *Host) ParseAuth(ctx context.Context, req pluginapi.AuthParseRequest) (* } func extractRawJSONOAuthProviderHint(raw []byte) string { - if len(bytes.TrimSpace(raw)) == 0 { - return "" - } - var payload map[string]any - if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil || payload == nil { + if !ojson.ValidBytes(raw) { return "" } for _, key := range []string{"oauth_provider", "portal"} { - if value, ok := payload[key].(string); ok { - if normalized := normalizeProviderID(value); normalized != "" { + if value := ojson.GetBytes(raw, key); value.Type == ojson.String { + if normalized := normalizeProviderID(value.String()); normalized != "" { return normalized } } @@ -618,6 +629,7 @@ func (h *Host) callSingleProviderParseAuths(ctx context.Context, record capabili return nil, false, nil } datas := pluginAuthParseResponseAuths(resp) + multiEntry := h.recordHasMultipleAuthProviders(record) auths = make([]*coreauth.Auth, 0, len(datas)) for _, data := range datas { if strings.TrimSpace(data.Provider) == "" { @@ -629,8 +641,11 @@ func (h *Host) callSingleProviderParseAuths(ctx context.Context, record capabili if normalizeProviderID(data.Provider) == "" { return nil, true, fmt.Errorf("auth provider %s returned auth without provider", record.id) } - effectivePortal := firstNonEmpty(authDataPortalHint(&data), portalIdentifier) - stampAuthDataOAuthProvider(&data, effectivePortal) + explicitHint := firstNonEmpty(authDataPortalHint(&data), extractRawJSONOAuthProviderHint(req.RawJSON)) + effectivePortal := firstNonEmpty(explicitHint, portalIdentifier) + if explicitHint != "" || multiEntry || normalizeProviderID(data.Provider) != effectivePortal { + stampAuthDataOAuthProvider(&data, effectivePortal) + } parsed := h.AuthDataToCoreAuth(data, req.Path, req.FileName) if parsed == nil { return nil, true, fmt.Errorf("auth provider %s returned invalid auth data", record.id) @@ -825,12 +840,21 @@ func (h *Host) callPollLogin(ctx context.Context, record capabilityRecord, provi return pluginapi.AuthLoginPollResponse{}, true, errPoll } if resp.Status == pluginapi.AuthLoginStatusSuccess { + multiEntry := h.recordHasMultipleAuthProviders(record) for index := range resp.Auths { - stampAuthDataOAuthProvider(&resp.Auths[index], matchedPortal) - h.recordCredentialProvider(record.id, resp.Auths[index].Provider, matchedPortal) + explicitHint := authDataPortalHint(&resp.Auths[index]) + effectivePortal := firstNonEmpty(explicitHint, matchedPortal) + if multiEntry || explicitHint != "" || normalizeProviderID(resp.Auths[index].Provider) != effectivePortal { + stampAuthDataOAuthProvider(&resp.Auths[index], effectivePortal) + } + h.recordCredentialProvider(record.id, resp.Auths[index].Provider, effectivePortal) + } + explicitHint := authDataPortalHint(&resp.Auth) + effectivePortal := firstNonEmpty(explicitHint, matchedPortal) + if multiEntry || explicitHint != "" || normalizeProviderID(resp.Auth.Provider) != effectivePortal { + stampAuthDataOAuthProvider(&resp.Auth, effectivePortal) } - stampAuthDataOAuthProvider(&resp.Auth, matchedPortal) - h.recordCredentialProvider(record.id, resp.Auth.Provider, matchedPortal) + h.recordCredentialProvider(record.id, resp.Auth.Provider, effectivePortal) } return resp, true, nil } @@ -853,7 +877,7 @@ func authDataPortalHint(data *pluginapi.AuthData) string { } } } - return "" + return extractRawJSONOAuthProviderHint(data.StorageJSON) } func stampAuthDataOAuthProvider(data *pluginapi.AuthData, portalIdentifier string) { @@ -864,22 +888,17 @@ func stampAuthDataOAuthProvider(data *pluginapi.AuthData, portalIdentifier strin if portalIdentifier == "" { return } - credProvider := normalizeProviderID(data.Provider) - if credProvider == "" || credProvider == portalIdentifier { + if normalizeProviderID(data.Provider) == "" { return } if data.Metadata == nil { data.Metadata = make(map[string]any) } - if _, exists := data.Metadata["oauth_provider"]; !exists { - data.Metadata["oauth_provider"] = portalIdentifier - } + data.Metadata["oauth_provider"] = portalIdentifier if data.Attributes == nil { data.Attributes = make(map[string]string) } - if strings.TrimSpace(data.Attributes["oauth_provider"]) == "" { - data.Attributes["oauth_provider"] = portalIdentifier - } + data.Attributes["oauth_provider"] = portalIdentifier } func authPortalHints(auth *coreauth.Auth) []string { @@ -972,6 +991,7 @@ func (h *Host) RefreshAuth(ctx context.Context, auth *coreauth.Auth) (refreshed return nil, true, errRefresh } data := pluginResp.Auth + explicitPortalHint := authDataPortalHint(&data) if strings.TrimSpace(data.Provider) == "" { data.Provider = authProvider(auth) } @@ -1006,7 +1026,7 @@ func (h *Host) RefreshAuth(ctx context.Context, auth *coreauth.Auth) (refreshed } data.Attributes = attributes } - preserveAuthOAuthProviderMetadata(&data, auth) + preserveAuthOAuthProviderMetadata(&data, auth, explicitPortalHint) preserveFileAuthPriority(&data, auth) preserveAuthDisabledState(&data, auth) if len(data.StorageJSON) == 0 { @@ -1031,10 +1051,19 @@ func (h *Host) RefreshAuth(ctx context.Context, auth *coreauth.Auth) (refreshed return next, true, nil } -func preserveAuthOAuthProviderMetadata(data *pluginapi.AuthData, auth *coreauth.Auth) { - if data == nil || auth == nil || len(auth.Metadata) == 0 { +// explicitPortalHint must be resolved before merging existing auth fields. +func preserveAuthOAuthProviderMetadata(data *pluginapi.AuthData, auth *coreauth.Auth, explicitPortalHint string) { + if data == nil || auth == nil { return } + hint := explicitPortalHint + if hint == "" { + hint = authDataPortalHint(&pluginapi.AuthData{ + Attributes: auth.Attributes, + Metadata: auth.Metadata, + StorageJSON: storageJSONFromAuth(auth), + }) + } for _, key := range []string{"oauth_provider", "portal"} { raw, ok := auth.Metadata[key].(string) if !ok || strings.TrimSpace(raw) == "" { @@ -1047,6 +1076,17 @@ func preserveAuthOAuthProviderMetadata(data *pluginapi.AuthData, auth *coreauth. data.Metadata[key] = raw } } + if hint == "" { + return + } + if data.Metadata == nil { + data.Metadata = make(map[string]any) + } + data.Metadata["oauth_provider"] = hint + if data.Attributes == nil { + data.Attributes = make(map[string]string) + } + data.Attributes["oauth_provider"] = hint } func preserveAuthDisabledState(data *pluginapi.AuthData, auth *coreauth.Auth) { @@ -1248,6 +1288,10 @@ func pluginAuthDataToCoreAuth(data pluginapi.AuthData, path, fileName string, au if attributes == nil { attributes = make(map[string]string) } + if hint := authDataPortalHint(&data); hint != "" { + metadata["oauth_provider"] = hint + attributes["oauth_provider"] = hint + } path = strings.TrimSpace(path) if path != "" { if attributes[coreauth.AttributePath] == "" { diff --git a/internal/pluginhost/auth_provider_oauth_provider_test.go b/internal/pluginhost/auth_provider_oauth_provider_test.go new file mode 100644 index 00000000000..aa7401cffa9 --- /dev/null +++ b/internal/pluginhost/auth_provider_oauth_provider_test.go @@ -0,0 +1,381 @@ +package pluginhost + +import ( + "bytes" + "context" + "os" + "path/filepath" + "testing" + + "github.com/MoeForge/moego/ojson" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" +) + +func TestHasMultipleAuthProviders(t *testing.T) { + host := newHostWithRecords( + capabilityRecord{id: "multi", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, credentialProviders: []string{"canonical-vendor"}}, + }}}, + capabilityRecord{id: "separate", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProviders: []pluginapi.AuthProvider{fakeAuthProvider{identifier: "portal-one"}, fakeAuthProvider{identifier: "portal-two"}}, + }}}, + capabilityRecord{id: "single", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{identifier: "single-vendor", identifiers: []string{"single-vendor", " SINGLE-VENDOR ", ""}}, + }}}, + ) + for _, tc := range []struct { + provider string + want bool + }{ + {provider: " Vendor ", want: true}, + {provider: "vendor-cn", want: true}, + {provider: "canonical-vendor", want: true}, + {provider: "portal-one", want: true}, + {provider: "portal-two", want: true}, + {provider: "single-vendor"}, + {provider: "missing"}, + {}, + } { + if got := host.HasMultipleAuthProviders(tc.provider); got != tc.want { + t.Errorf("HasMultipleAuthProviders(%q) = %t, want %t", tc.provider, got, tc.want) + } + } + var nilHost *Host + if nilHost.HasMultipleAuthProviders("vendor") { + t.Fatal("nil host reports multiple auth providers") + } + record := host.authProviderRecord("vendor") + setHostSnapshotForTest(host, false) + if host.HasMultipleAuthProviders("vendor") || host.recordHasMultipleAuthProviders(*record) { + t.Fatal("inactive plugin reports multiple auth providers") + } +} + +func TestParseAuthsOAuthProviderAttribution(t *testing.T) { + for _, tc := range []struct { + name string + request string + identifiers []string + data pluginapi.AuthData + raw []byte + want string + }{ + {name: "multi-primary", request: "vendor", identifiers: []string{"vendor", "vendor-cn"}, want: "vendor"}, + {name: "multi-secondary", request: "vendor-cn", identifiers: []string{"vendor", "vendor-cn"}, want: "vendor-cn"}, + {name: "single-without-hint", request: "vendor", identifiers: []string{"vendor"}}, + {name: "attribute-hint", request: "vendor", identifiers: []string{"vendor"}, data: pluginapi.AuthData{Attributes: map[string]string{"oauth_provider": " Vendor "}}, want: "vendor"}, + {name: "metadata-hint", request: "vendor", identifiers: []string{"vendor"}, data: pluginapi.AuthData{Metadata: map[string]any{"portal": " Vendor "}}, want: "vendor"}, + {name: "storage-hint", request: "vendor", identifiers: []string{"vendor"}, data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":" Vendor "}`)}, want: "vendor"}, + {name: "raw-hint", request: "vendor", identifiers: []string{"vendor"}, raw: []byte(`{"oauth_provider":" Vendor "}`), want: "vendor"}, + { + name: "data-over-storage-over-request", request: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + data: pluginapi.AuthData{Attributes: map[string]string{"portal": " Vendor "}, StorageJSON: []byte(`{"oauth_provider":"vendor-cn"}`)}, + raw: []byte(`{"oauth_provider":"vendor-cn"}`), want: "vendor", + }, + { + name: "storage-over-request", request: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":true,"portal":" Vendor-CN "}`)}, + raw: []byte(`{"oauth_provider":"vendor"}`), want: "vendor-cn", + }, + } { + t.Run(tc.name, func(t *testing.T) { + host := newHostWithRecords(capabilityRecord{id: "vendor-plugin", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{ + identifier: "vendor", identifiers: tc.identifiers, + parseAuth: func(context.Context, pluginapi.AuthParseRequest) (pluginapi.AuthParseResponse, error) { + data := tc.data + data.Provider = "vendor" + return pluginapi.AuthParseResponse{Handled: true, Auth: data}, nil + }, + }, + }}}) + auths, handled, errParse := host.ParseAuths(context.Background(), pluginapi.AuthParseRequest{Provider: tc.request, RawJSON: tc.raw}) + if errParse != nil || !handled || len(auths) != 1 { + t.Fatalf("ParseAuths() handled=%t count=%d error=%v", handled, len(auths), errParse) + } + auth := auths[0] + if tc.want == "" { + if _, exists := auth.Metadata["oauth_provider"]; exists { + t.Errorf("single-entry parser added metadata: %#v", auth.Metadata) + } + if _, exists := auth.Attributes["oauth_provider"]; exists { + t.Errorf("single-entry parser added attributes: %#v", auth.Attributes) + } + } else if auth.Metadata["oauth_provider"] != tc.want || auth.Attributes["oauth_provider"] != tc.want { + t.Errorf("parsed attribution metadata=%#v attributes=%#v, want %q", auth.Metadata, auth.Attributes, tc.want) + } + }) + } +} + +func TestPluginAuthDataToCoreAuthStorageOAuthProvider(t *testing.T) { + for _, tc := range []struct { + name string + data pluginapi.AuthData + want string + }{ + {name: "raw-oauth", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":" Vendor "}`)}, want: "vendor"}, + {name: "raw-portal", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":" ","portal":" Vendor-CN "}`)}, want: "vendor-cn"}, + {name: "invalid-hints-fall-back", data: pluginapi.AuthData{Metadata: map[string]any{"oauth_provider": 42}, Attributes: map[string]string{"portal": " "}, StorageJSON: []byte(`{"oauth_provider":"vendor"}`)}, want: "vendor"}, + {name: "metadata-wins", data: pluginapi.AuthData{Metadata: map[string]any{"portal": "vendor-cn"}, StorageJSON: []byte(`{"oauth_provider":"vendor"}`)}, want: "vendor-cn"}, + {name: "attributes-win", data: pluginapi.AuthData{Attributes: map[string]string{"oauth_provider": "vendor-cn"}, StorageJSON: []byte(`{"oauth_provider":"vendor"}`)}, want: "vendor-cn"}, + {name: "invalid-json", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":"vendor"`)}}, + {name: "non-string", data: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":42,"portal":true}`)}}, + } { + t.Run(tc.name, func(t *testing.T) { + tc.data.Provider = "vendor" + auth := pluginAuthDataToCoreAuth(tc.data, "", "", "") + hint := authDataPortalHint(&pluginapi.AuthData{Metadata: auth.Metadata, Attributes: auth.Attributes}) + if hint != tc.want { + t.Fatalf("converted hint = %q, want %q", hint, tc.want) + } + if tc.want != "" { + if auth.Metadata["oauth_provider"] != tc.want || auth.Attributes["oauth_provider"] != tc.want { + t.Fatalf("hint not copied to both maps: metadata=%#v attributes=%#v", auth.Metadata, auth.Attributes) + } + } + }) + } +} + +func TestPreserveAuthOAuthProviderFromAttributes(t *testing.T) { + for _, key := range []string{"oauth_provider", "portal"} { + for _, metadata := range []map[string]any{nil, {"access_token": "old"}, {"oauth_provider": " ", "portal": true}} { + t.Run(key, func(t *testing.T) { + auth := &coreauth.Auth{Provider: "vendor", Metadata: metadata, Attributes: map[string]string{key: " Vendor "}} + data := pluginapi.AuthData{Provider: "vendor", Metadata: map[string]any{"access_token": "new"}} + preserveAuthOAuthProviderMetadata(&data, auth, authDataPortalHint(&data)) + if data.Metadata["oauth_provider"] != "vendor" || data.Attributes["oauth_provider"] != "vendor" { + t.Fatalf("preserved metadata=%#v attributes=%#v, want vendor", data.Metadata, data.Attributes) + } + if data.Metadata["access_token"] != "new" { + t.Fatal("preservation replaced refreshed token") + } + }) + } + } +} + +func TestPollLoginRecordsEffectiveOAuthProvider(t *testing.T) { + for _, responseKind := range []string{"auth", "auths"} { + t.Run(responseKind, func(t *testing.T) { + host := newHostWithRecords(capabilityRecord{id: "vendor-plugin", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + pollLogin: func(context.Context, pluginapi.AuthLoginPollRequest) (pluginapi.AuthLoginPollResponse, error) { + data := pluginapi.AuthData{Provider: "vendor", StorageJSON: []byte(`{"oauth_provider":"vendor-cn"}`)} + resp := pluginapi.AuthLoginPollResponse{Status: pluginapi.AuthLoginStatusSuccess} + if responseKind == "auths" { + resp.Auths = []pluginapi.AuthData{data} + } else { + resp.Auth = data + } + return resp, nil + }, + }, + }}}) + _, handled, errPoll := host.PollLogin(context.Background(), "vendor", "state") + if errPoll != nil || !handled { + t.Fatalf("PollLogin() handled=%t error=%v", handled, errPoll) + } + if got, ok := host.recordedCredentialProviderPortal("vendor-plugin", "vendor"); !ok || got != "vendor-cn" { + t.Fatalf("recorded portal = %q, present=%t, want vendor-cn", got, ok) + } + }) + } +} + +func TestExecutorAdapterRefreshOAuthProvider(t *testing.T) { + for _, tc := range []struct { + name string + metadata map[string]any + want string + }{ + {name: "existing-effective-hint", metadata: map[string]any{"access_token": "new"}, want: "vendor"}, + {name: "plugin-explicit-hint", metadata: map[string]any{"oauth_provider": "vendor-cn", "access_token": "new"}, want: "vendor-cn"}, + } { + t.Run(tc.name, func(t *testing.T) { + record := normalizeTestCapabilityRecord(capabilityRecord{id: "vendor-plugin", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + refreshAuth: func(context.Context, pluginapi.AuthRefreshRequest) (pluginapi.AuthRefreshResponse, error) { + return pluginapi.AuthRefreshResponse{Auth: pluginapi.AuthData{Metadata: tc.metadata}}, nil + }, + }, + }}}) + host := newHostWithRecords(record) + adapter := newExecutorAdapterForRecordForTest(host, record, &fakeExecutor{}, nil, nil) + refreshed, errRefresh := adapter.Refresh(context.Background(), &coreauth.Auth{ + Provider: "vendor", Attributes: map[string]string{"oauth_provider": "vendor"}, Metadata: map[string]any{"portal": "vendor-cn"}, + }) + if errRefresh != nil || refreshed == nil { + t.Fatalf("Refresh() auth=%#v error=%v", refreshed, errRefresh) + } + if refreshed.Metadata["oauth_provider"] != tc.want || refreshed.Attributes["oauth_provider"] != tc.want { + t.Fatalf("refreshed attribution metadata=%#v attributes=%#v, want %q", refreshed.Metadata, refreshed.Attributes, tc.want) + } + storage := refreshed.Storage.(interface{ RawJSON() []byte }) + if got := ojson.GetBytes(storage.RawJSON(), "oauth_provider").String(); got != tc.want { + t.Fatalf("refreshed storage hint = %q, want %q", got, tc.want) + } + }) + } +} + +func TestRefreshAuthOAuthProviderSurvivesColdStart(t *testing.T) { + for _, tc := range []struct { + name string + attributes map[string]string + metadata map[string]any + rawJSON []byte + response pluginapi.AuthData + want string + }{ + { + name: "attribute-oauth-over-metadata-portal", + attributes: map[string]string{"oauth_provider": "vendor"}, + metadata: map[string]any{"portal": "vendor-cn"}, + want: "vendor", + }, + { + name: "metadata-oauth-over-attribute-portal", + attributes: map[string]string{"portal": "vendor-cn"}, + metadata: map[string]any{"oauth_provider": "vendor"}, + want: "vendor", + }, + { + name: "attribute-portal-over-metadata-portal", + attributes: map[string]string{"portal": "vendor"}, + metadata: map[string]any{"portal": "vendor-cn"}, + want: "vendor", + }, + { + name: "metadata-portal-over-storage-oauth", + metadata: map[string]any{"portal": "vendor"}, + rawJSON: []byte(`{"oauth_provider":"vendor-cn"}`), + want: "vendor", + }, + { + name: "storage-only", + rawJSON: []byte(`{"oauth_provider":"vendor-cn"}`), + want: "vendor-cn", + }, + { + name: "plugin-attribute-oauth-over-existing", + attributes: map[string]string{"oauth_provider": "vendor"}, + metadata: map[string]any{"oauth_provider": "vendor"}, + response: pluginapi.AuthData{Attributes: map[string]string{"oauth_provider": "vendor-cn"}}, + want: "vendor-cn", + }, + { + name: "plugin-metadata-oauth-over-existing", + attributes: map[string]string{"oauth_provider": "vendor"}, + response: pluginapi.AuthData{Metadata: map[string]any{"oauth_provider": "vendor-cn"}}, + want: "vendor-cn", + }, + { + name: "plugin-attribute-portal-over-existing", + attributes: map[string]string{"oauth_provider": "vendor"}, + response: pluginapi.AuthData{Attributes: map[string]string{"portal": "vendor-cn"}}, + want: "vendor-cn", + }, + { + name: "plugin-metadata-portal-over-existing", + attributes: map[string]string{"oauth_provider": "vendor"}, + response: pluginapi.AuthData{Metadata: map[string]any{"portal": "vendor-cn"}}, + want: "vendor-cn", + }, + { + name: "plugin-storage-oauth-over-existing", + attributes: map[string]string{"oauth_provider": "vendor"}, + metadata: map[string]any{"portal": "vendor"}, + response: pluginapi.AuthData{StorageJSON: []byte(`{"oauth_provider":"vendor-cn"}`)}, + want: "vendor-cn", + }, + } { + t.Run(tc.name, func(t *testing.T) { + newHost := func() *Host { + return newHostWithRecords(capabilityRecord{id: "vendor-plugin", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + refreshAuth: func(context.Context, pluginapi.AuthRefreshRequest) (pluginapi.AuthRefreshResponse, error) { + data := tc.response + data.Metadata = cloneAnyMap(data.Metadata) + if data.Metadata == nil { + data.Metadata = make(map[string]any) + } + data.Metadata["access_token"] = "new" + return pluginapi.AuthRefreshResponse{Auth: data}, nil + }, + parseAuth: func(_ context.Context, req pluginapi.AuthParseRequest) (pluginapi.AuthParseResponse, error) { + return pluginapi.AuthParseResponse{Handled: true, Auth: pluginapi.AuthData{Provider: "vendor", StorageJSON: req.RawJSON}}, nil + }, + }, + }}}) + } + auth := &coreauth.Auth{Provider: "vendor", Attributes: tc.attributes, Metadata: tc.metadata} + if len(tc.rawJSON) > 0 { + auth.Storage = &pluginTokenStorage{provider: "vendor", rawJSON: tc.rawJSON} + } + refreshed, handled, errRefresh := newHost().RefreshAuth(context.Background(), auth) + if errRefresh != nil || !handled || refreshed == nil { + t.Fatalf("RefreshAuth() handled=%t auth=%#v error=%v", handled, refreshed, errRefresh) + } + if refreshed.Metadata["oauth_provider"] != tc.want || refreshed.Attributes["oauth_provider"] != tc.want { + t.Fatalf("refreshed attribution metadata=%#v attributes=%#v, want %q", refreshed.Metadata, refreshed.Attributes, tc.want) + } + if refreshed.Metadata["access_token"] != "new" { + t.Fatal("refresh lost the updated token") + } + raw := refreshed.Storage.(interface{ RawJSON() []byte }).RawJSON() + if got := ojson.GetBytes(raw, "oauth_provider").String(); got != tc.want { + t.Fatalf("refreshed storage hint = %q, want %q", got, tc.want) + } + path := filepath.Join(t.TempDir(), "account.json") + if errSaveTokenToFile := refreshed.Storage.SaveTokenToFile(path); errSaveTokenToFile != nil { + t.Fatalf("save refreshed auth: %v", errSaveTokenToFile) + } + saved, errRead := os.ReadFile(path) + if errRead != nil { + t.Fatalf("read saved auth: %v", errRead) + } + if !bytes.Equal(saved, raw) { + t.Fatalf("saved payload = %s, want RawJSON() %s", saved, raw) + } + reloaded, handled, errParse := newHost().ParseAuth(context.Background(), pluginapi.AuthParseRequest{ + Provider: "vendor", Path: path, FileName: "account.json", RawJSON: saved, + }) + if errParse != nil || !handled || reloaded == nil { + t.Fatalf("cold-start ParseAuth() handled=%t auth=%#v error=%v", handled, reloaded, errParse) + } + if reloaded.Metadata["oauth_provider"] != tc.want || reloaded.Attributes["oauth_provider"] != tc.want { + t.Fatalf("cold-start attribution metadata=%#v attributes=%#v, want %q", reloaded.Metadata, reloaded.Attributes, tc.want) + } + }) + } +} + +func TestRefreshAuthPreservesAttributeOnlyOAuthProvider(t *testing.T) { + host := newHostWithRecords(capabilityRecord{id: "vendor-plugin", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + AuthProvider: fakeAuthProvider{ + identifier: "vendor", identifiers: []string{"vendor", "vendor-cn"}, + refreshAuth: func(context.Context, pluginapi.AuthRefreshRequest) (pluginapi.AuthRefreshResponse, error) { + return pluginapi.AuthRefreshResponse{Auth: pluginapi.AuthData{Provider: "vendor", Metadata: map[string]any{"access_token": "new"}}}, nil + }, + }, + }}}) + auth, handled, errRefresh := host.RefreshAuth(context.Background(), &coreauth.Auth{ + Provider: "vendor", Attributes: map[string]string{"portal": " Vendor "}, + }) + if errRefresh != nil || !handled || auth == nil { + t.Fatalf("RefreshAuth() handled=%t auth=%#v error=%v", handled, auth, errRefresh) + } + if auth.Metadata["oauth_provider"] != "vendor" || auth.Attributes["oauth_provider"] != "vendor" { + t.Fatalf("refreshed attribution metadata=%#v attributes=%#v", auth.Metadata, auth.Attributes) + } + storage := auth.Storage.(interface{ RawJSON() []byte }) + if got := ojson.GetBytes(storage.RawJSON(), "oauth_provider").String(); got != "vendor" { + t.Fatalf("refreshed storage hint = %q, want vendor", got) + } +} diff --git a/internal/pluginhost/loader_unix.go b/internal/pluginhost/loader_unix.go index d31f7b77804..1489bc84240 100644 --- a/internal/pluginhost/loader_unix.go +++ b/internal/pluginhost/loader_unix.go @@ -144,8 +144,15 @@ var ( hostCallbackEntries sync.Map probedPluginShadows sync.Map pluginProbeExecutable = os.Executable + unixDlopen = defaultUnixDlopen ) +func defaultUnixDlopen(path string) unsafe.Pointer { + cPath := C.CString(path) + defer C.free(unsafe.Pointer(cPath)) + return C.cliproxy_dlopen(cPath) +} + func init() { if probePath := os.Getenv(pluginDlopenProbeEnv); probePath != "" { runPluginDlopenProbe(probePath) @@ -170,21 +177,13 @@ func defaultPluginLoader() pluginLoader { } func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, error) { - loadPath, errShadow := shadowCopyPlugin(file) + loadPath, releaseGuard, errShadow := acquireShadowPluginLoad(file) if errShadow != nil { return nil, errShadow } - if errProbe := probePluginLibrary(file.Path, loadPath); errProbe != nil { - removeShadowPlugin(loadPath) - return nil, errProbe - } - cPath := C.CString(loadPath) - defer C.free(unsafe.Pointer(cPath)) - - handle := C.cliproxy_dlopen(cPath) - if handle == nil { - removeShadowPlugin(loadPath) - return nil, fmt.Errorf("dlopen %s: %s", file.Path, dlerrorString()) + handle, loadPath, errOpen := openUnixShadowPluginWithFallback(file, loadPath, releaseGuard) + if errOpen != nil { + return nil, errOpen } // From here on the image must never be unloaded: dlopen may already // have started the plugin runtime, and a Go c-shared runtime cannot be @@ -324,6 +323,50 @@ func (c *dynamicLibraryClient) close() { } } +func openUnixShadowPluginWithFallback(file pluginFile, loadPath string, releaseGuard func()) (unsafe.Pointer, string, error) { + errProbe := probePluginLibrary(file.Path, loadPath) + if errProbe != nil { + releaseGuard() + if _, cached := probedPluginShadows.Load(filepath.Base(loadPath)); cached { + return nil, "", errProbe + } + if fallbackHandle, fallbackPath, okFallback := tryOpenUnixFallbackShadowPlugin(file, loadPath); okFallback { + return fallbackHandle, fallbackPath, nil + } + return nil, "", errProbe + } + handle := unixDlopen(loadPath) + firstErr := dlerrorString() + releaseGuard() + if handle != nil { + return handle, loadPath, nil + } + probedPluginShadows.Delete(filepath.Base(loadPath)) + + if fallbackHandle, fallbackPath, okFallback := tryOpenUnixFallbackShadowPlugin(file, loadPath); okFallback { + return fallbackHandle, fallbackPath, nil + } + return nil, "", fmt.Errorf("dlopen %s: %s", file.Path, firstErr) +} + +func tryOpenUnixFallbackShadowPlugin(file pluginFile, failedLoadPath string) (unsafe.Pointer, string, bool) { + fallbackPath, releaseFallbackGuard, errFallback := acquireShadowPluginFallbackLoad(file, failedLoadPath) + if errFallback != nil { + return nil, "", false + } + if errFallbackProbe := probePluginLibrary(file.Path, fallbackPath); errFallbackProbe != nil { + releaseFallbackGuard() + return nil, "", false + } + fallbackHandle := unixDlopen(fallbackPath) + releaseFallbackGuard() + if fallbackHandle != nil { + return fallbackHandle, fallbackPath, true + } + probedPluginShadows.Delete(filepath.Base(fallbackPath)) + return nil, "", false +} + func dlerrorString() string { errText := C.cliproxy_dlerror() if errText == nil { @@ -418,7 +461,11 @@ func probePluginLibrary(sourcePath, loadPath string) error { execPath = os.Args[0] } - statusFile, errStatus := os.CreateTemp(filepath.Dir(loadPath), ".cliproxy-probe-*.status") + statusDir, errStatusDir := shadowPluginDir() + if errStatusDir != nil { + statusDir = os.TempDir() + } + statusFile, errStatus := os.CreateTemp(statusDir, ".cliproxy-probe-*.status") if errStatus != nil { return fmt.Errorf("dlopen %s: create preflight probe status file: %w", sourcePath, errStatus) } diff --git a/internal/pluginhost/loader_unix_test.go b/internal/pluginhost/loader_unix_test.go index e6246587caa..160636caae8 100644 --- a/internal/pluginhost/loader_unix_test.go +++ b/internal/pluginhost/loader_unix_test.go @@ -14,6 +14,7 @@ import ( "strings" "syscall" "testing" + "unsafe" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" ) @@ -167,6 +168,88 @@ func testShutdownCallCount(t *testing.T, client pluginClient) int { return count } +func TestUnixPluginLoaderFallsBackToSharedCacheWhenLocalShadowDirNoexec(t *testing.T) { + dir := t.TempDir() + pluginPath := filepath.Join(dir, "noexec-fallback"+testPluginLibraryExtension()) + buildReloadTestPlugin(t, 1, pluginPath) + + localShadowDir := filepath.Clean(filepath.Join(dir, shadowPluginLocalDirName)) + origDlopen := unixDlopen + t.Cleanup(func() { + unixDlopen = origDlopen + }) + unixDlopen = func(path string) unsafe.Pointer { + if filepath.Clean(filepath.Dir(path)) == localShadowDir { + return nil + } + return origDlopen(path) + } + + host := New() + file := pluginFile{ID: "noexec-fallback", Path: pluginPath, Version: "1"} + client, errOpen := dynamicLibraryLoader{}.Open(file, host) + if errOpen != nil { + t.Fatalf("Open() with simulated noexec local .shadow error = %v, want fallback success", errOpen) + } + defer client.Shutdown() + dlc, okType := client.(*dynamicLibraryClient) + if !okType { + t.Fatalf("Open() client type = %T, want *dynamicLibraryClient", client) + } + t.Cleanup(func() { + removeShadowPluginDir(filepath.Dir(dlc.tempPath)) + }) + wantFallbackDir, errFallbackDir := shadowPluginFallbackDir(pluginPath) + if errFallbackDir != nil { + t.Fatalf("shadowPluginFallbackDir() error = %v", errFallbackDir) + } + if filepath.Clean(filepath.Dir(dlc.tempPath)) != filepath.Clean(wantFallbackDir) { + t.Fatalf("loaded shadow dir = %q, want fallback dir %q", filepath.Dir(dlc.tempPath), wantFallbackDir) + } + resp, errCall := client.Call(context.Background(), "test.marker", nil) + if errCall != nil { + t.Fatalf("Call() fallback plugin error = %v", errCall) + } + if string(resp) != "v1" { + t.Fatalf("Call() fallback plugin response = %q, want v1", resp) + } +} + +func TestUnixPluginLoaderLoadsFromReadOnlyPopulatedLocalShadowDir(t *testing.T) { + dir := t.TempDir() + pluginPath := filepath.Join(dir, "readonly-shadow"+testPluginLibraryExtension()) + buildReloadTestPlugin(t, 1, pluginPath) + file := pluginFile{ID: "readonly-shadow", Path: pluginPath, Version: "1"} + + shadowPath, errShadow := shadowCopyPlugin(file) + if errShadow != nil { + t.Fatalf("shadowCopyPlugin() pre-populate error = %v", errShadow) + } + localShadowDir := filepath.Dir(shadowPath) + probedPluginShadows.Delete(filepath.Base(shadowPath)) + + if errChmod := os.Chmod(localShadowDir, 0o500); errChmod != nil { + t.Fatalf("Chmod(localShadowDir) error = %v", errChmod) + } + t.Cleanup(func() { + _ = os.Chmod(localShadowDir, 0o700) + }) + + host := New() + client, errOpen := dynamicLibraryLoader{}.Open(file, host) + if errOpen != nil { + t.Fatalf("Open() with read-only populated .shadow error = %v, want success", errOpen) + } + defer client.Shutdown() + resp, errCall := client.Call(context.Background(), "test.marker", nil) + if errCall != nil { + t.Fatalf("Call() plugin error = %v", errCall) + } + if string(resp) != "v1" { + t.Fatalf("Call() plugin response = %q, want v1", resp) + } +} + func TestUnixPluginOpenFailureKeepsImageUsableOnRetry(t *testing.T) { dir := t.TempDir() pluginPath := filepath.Join(dir, "failing"+testPluginLibraryExtension()) diff --git a/internal/pluginhost/loader_unsupported.go b/internal/pluginhost/loader_unsupported.go index 303d106c57b..442b37039fb 100644 --- a/internal/pluginhost/loader_unsupported.go +++ b/internal/pluginhost/loader_unsupported.go @@ -13,3 +13,5 @@ func (unsupportedLoader) Open(file pluginFile, host *Host) (pluginClient, error) func defaultPluginLoader() pluginLoader { return unsupportedLoader{} } + +func cleanupStaleShadowPlugins(string, []pluginFile) {} diff --git a/internal/pluginhost/loader_windows.go b/internal/pluginhost/loader_windows.go index bfffeed7183..cb6edb29956 100644 --- a/internal/pluginhost/loader_windows.go +++ b/internal/pluginhost/loader_windows.go @@ -71,14 +71,25 @@ func defaultPluginLoader() pluginLoader { } func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, error) { - loadPath, errShadow := shadowCopyPlugin(file) + loadPath, releaseGuard, errShadow := acquireShadowPluginLoad(file) if errShadow != nil { return nil, errShadow } dll, errLoad := syscall.LoadDLL(loadPath) + releaseGuard() if errLoad != nil { - removeShadowPlugin(loadPath) - return nil, errLoad + if fallbackPath, releaseFallbackGuard, errFallback := acquireShadowPluginFallbackLoad(file, loadPath); errFallback == nil { + fallbackDLL, errFallbackLoad := syscall.LoadDLL(fallbackPath) + releaseFallbackGuard() + if errFallbackLoad == nil { + dll = fallbackDLL + loadPath = fallbackPath + errLoad = nil + } + } + if errLoad != nil { + return nil, errLoad + } } proc, errProc := dll.FindProc("cliproxy_plugin_init") if errProc != nil { diff --git a/internal/pluginhost/platform.go b/internal/pluginhost/platform.go index b3bb636e1fc..f4c928e6782 100644 --- a/internal/pluginhost/platform.go +++ b/internal/pluginhost/platform.go @@ -270,21 +270,26 @@ func cleanupUnselectedPluginFiles(root string, loaded []pluginFile) error { paths[filepath.Clean(file.Path)] = struct{}{} } var errs []error + keptCandidates := make([]pluginFile, 0, len(candidates)) for _, candidate := range candidates { paths := loadedByID[candidate.ID] if len(paths) == 0 { + keptCandidates = append(keptCandidates, candidate) continue } if _, selected := paths[filepath.Clean(candidate.Path)]; selected { + keptCandidates = append(keptCandidates, candidate) continue } if errRemove := os.Remove(candidate.Path); errRemove != nil && !errors.Is(errRemove, os.ErrNotExist) { errs = append(errs, errRemove) + keptCandidates = append(keptCandidates, candidate) log.WithError(errRemove).Warnf("pluginhost: failed to remove old plugin file %s", candidate.Path) continue } log.WithFields(pluginLogFields(candidate.ID, "", candidate.Version, candidate.Path)).Info("pluginhost: old plugin file removed") } + cleanupStaleShadowPlugins(root, append(keptCandidates, loaded...)) return errors.Join(errs...) } diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 332674bfdf5..ed96ba2e300 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -53,6 +53,7 @@ type rpcError struct { message string statusCode int retryAfterSeconds int + credentialScoped bool } func (e rpcError) Error() string { @@ -67,6 +68,22 @@ func (e rpcError) StatusCode() int { return e.statusCode } +func (e rpcError) IsCredentialScoped() bool { + return e.credentialScoped || strings.EqualFold(strings.TrimSpace(e.Code), pluginabi.ErrorCodeCredentialQuota) +} + +func (e rpcError) IsRequestScoped() bool { + return !e.IsCredentialScoped() && strings.EqualFold(strings.TrimSpace(e.Code), pluginabi.ErrorCodeRequestScoped) +} + +func (e rpcError) IsTransientRateLimit() bool { + return !e.IsCredentialScoped() && strings.EqualFold(strings.TrimSpace(e.Code), pluginabi.ErrorCodeTransientRateLimit) +} + +func (e rpcError) IsStreamTruncated() bool { + return !e.IsCredentialScoped() && strings.EqualFold(strings.TrimSpace(e.Code), pluginabi.ErrorCodeStreamTruncated) +} + func (e rpcError) RetryAfter() *time.Duration { if e.retryAfterSeconds <= 0 { return nil @@ -440,6 +457,7 @@ func decodeEnvelopeResult[T any](envelope pluginabi.Envelope) (T, error) { message: message, statusCode: envelope.Error.HTTPStatus, retryAfterSeconds: envelope.Error.RetryAfterSeconds, + credentialScoped: envelope.Error.CredentialScoped, } } return zero, fmt.Errorf("plugin call failed") diff --git a/internal/pluginhost/rpc_client_error_test.go b/internal/pluginhost/rpc_client_error_test.go index 8baefe6f9ff..929bcb139df 100644 --- a/internal/pluginhost/rpc_client_error_test.go +++ b/internal/pluginhost/rpc_client_error_test.go @@ -9,6 +9,8 @@ import ( "testing" "time" + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/redisqueue" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -492,3 +494,335 @@ func TestPluginCountTokensNotImplementedDoesNotTriggerCooldown(t *testing.T) { t.Fatal("Execute returned empty payload") } } + +type resultCapturingHook struct { + results []coreauth.Result +} + +func (h *resultCapturingHook) OnAuthRegistered(context.Context, *coreauth.Auth) {} +func (h *resultCapturingHook) OnAuthUpdated(context.Context, *coreauth.Auth) {} +func (h *resultCapturingHook) OnResult(_ context.Context, result coreauth.Result) { + h.results = append(h.results, result) +} + +func TestPluginExecutor429CredentialScopedCoolsEntireCredentialAndSiblingModels(t *testing.T) { + coreauth.SetQuotaCooldownDisabled(false) + + prevQueueEnabled := redisqueue.Enabled() + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(true) + t.Cleanup(func() { + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(prevQueueEnabled) + }) + + for _, tc := range []struct { + name string + envelopeJSON string + wantCode string + }{ + { + name: "credential_scoped flag", + envelopeJSON: `{"ok":false,"error":{"code":"rate_limited","message":"account limited, resets at 2026-01-01T00:00:00Z","http_status":429,"retry_after_seconds":120,"credential_scoped":true}}`, + wantCode: "rate_limited", + }, + { + name: "credential_quota error code", + envelopeJSON: `{"ok":false,"error":{"code":"credential_quota","message":"account limited, resets at 2026-01-01T00:00:00Z","http_status":429,"retry_after_seconds":120}}`, + wantCode: "credential_quota", + }, + { + name: "credential_scoped overrides transient_rate_limit code", + envelopeJSON: `{"ok":false,"error":{"code":"transient_rate_limit","message":"account limited, resets at 2026-01-01T00:00:00Z","http_status":429,"retry_after_seconds":120,"credential_scoped":true}}`, + wantCode: "credential_quota", + }, + { + name: "credential_scoped overrides request_scoped code", + envelopeJSON: `{"ok":false,"error":{"code":"request_scoped","message":"account limited, resets at 2026-01-01T00:00:00Z","http_status":429,"retry_after_seconds":120,"credential_scoped":true}}`, + wantCode: "credential_quota", + }, + } { + t.Run(tc.name, func(t *testing.T) { + subscriber, unsubscribe := redisqueue.SubscribeErrors() + defer unsubscribe() + + _, errDecode := decodeRPCEnvelope[pluginapi.ExecutorResponse]([]byte(tc.envelopeJSON)) + if errDecode == nil { + t.Fatal("decodeRPCEnvelope returned nil error") + } + var scoped interface{ IsCredentialScoped() bool } + if !errors.As(errDecode, &scoped) || !scoped.IsCredentialScoped() { + t.Fatalf("decoded plugin error %T does not report IsCredentialScoped() = true", errDecode) + } + + var ( + provider = "plugin-provider-cred-scope-" + tc.name + modelA = "plugin-model-a-" + tc.name + modelB = "plugin-model-b-" + tc.name + modelC = "plugin-model-c-" + tc.name + authID = "auth-cred-scope-" + tc.name + execCalled int + ) + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-cred-scope-" + tc.name}) + host := newHostWithRecords(executorRecord) + exec := &fakeExecutor{ + identifier: provider, + execute: func(context.Context, pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + execCalled++ + return pluginapi.ExecutorResponse{}, errDecode + }, + } + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = provider + + hook := &resultCapturingHook{} + authMgr := coreauth.NewManager(nil, nil, hook) + authMgr.SetRetryConfig(0, 0, 0) + authMgr.RegisterExecutor(adapter) + auth := &coreauth.Auth{ + ID: authID, + Provider: provider, + Status: coreauth.StatusActive, + FileName: authID + ".json", + Attributes: map[string]string{"type": "oauth"}, + ModelStates: map[string]*coreauth.ModelState{ + modelB: {Status: coreauth.StatusActive}, + }, + } + if _, errReg := authMgr.Register(context.Background(), auth); errReg != nil { + t.Fatalf("Register auth: %v", errReg) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{ + {ID: modelA}, + {ID: modelB}, + {ID: modelC}, + }) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + opts := coreexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + + before := time.Now() + _, errExecA := authMgr.Execute(context.Background(), []string{provider}, coreexecutor.Request{ + Model: modelA, + Payload: []byte(`{"model":"` + modelA + `"}`), + }, opts) + if errExecA == nil { + t.Fatal("Execute(modelA) returned nil error, want 429") + } + if execCalled != 1 { + t.Fatalf("execCalled after modelA = %d, want 1", execCalled) + } + if len(hook.results) != 1 || !hook.results[0].CredentialScope { + t.Fatalf("recorded result = %+v, want CredentialScope=true", hook.results) + } + + var eventPayload []byte + select { + case eventPayload = <-subscriber: + case <-time.After(time.Second): + t.Fatal("timeout waiting for published error event") + } + statusCode := int(ojson.GetBytes(eventPayload, "status_code").Int()) + eventCode := ojson.GetBytes(eventPayload, "code").String() + if statusCode != http.StatusTooManyRequests || eventCode != tc.wantCode { + t.Fatalf("error event status=%d code=%q, want 429 %q", statusCode, eventCode, tc.wantCode) + } + unavailable := ojson.GetBytes(eventPayload, "auth_status.unavailable").Bool() + quotaExceeded := ojson.GetBytes(eventPayload, "auth_status.quota.exceeded").Bool() + quotaReason := ojson.GetBytes(eventPayload, "auth_status.quota.reason").String() + if !unavailable || !quotaExceeded || quotaReason != "credential_quota" { + t.Fatalf("error event auth_status = %s, want unavailable with quota.reason=credential_quota", string(eventPayload)) + } + + updated, okAuth := authMgr.GetByID(authID) + if !okAuth || updated == nil { + t.Fatalf("GetByID(%q) missing", authID) + } + if !updated.Unavailable || !updated.Quota.Exceeded || updated.Quota.Reason != "credential_quota" { + t.Fatalf("updated auth quota = %+v (unavailable=%v), want unavailable with reason=credential_quota", updated.Quota, updated.Unavailable) + } + if !updated.Quota.NextRecoverAt.After(before.Add(110 * time.Second)) { + t.Fatalf("updated auth NextRecoverAt = %v, want > %v", updated.Quota.NextRecoverAt, before.Add(110*time.Second)) + } + stateB := updated.ModelStates[modelB] + if stateB == nil || !stateB.Unavailable || !stateB.Quota.Exceeded || stateB.Quota.Reason != "credential_quota" { + t.Fatalf("sibling modelB state = %+v, want unavailable with reason=credential_quota", stateB) + } + + // Sending requests for sibling modelB (pre-existing ModelState) and modelC (unvisited model) + // must NOT select the cooled credential again. + for _, siblingModel := range []string{modelB, modelC} { + _, errSibling := authMgr.Execute(context.Background(), []string{provider}, coreexecutor.Request{ + Model: siblingModel, + Payload: []byte(`{"model":"` + siblingModel + `"}`), + }, opts) + if errSibling == nil { + t.Fatalf("Execute(%s) returned nil error, want cooldown rejection", siblingModel) + } + if execCalled != 1 { + t.Fatalf("execCalled after %s = %d, want 1 (credential must not be selected again for sibling model)", siblingModel, execCalled) + } + } + }) + } +} + +func TestStreamCallbacksPreserveCredentialScopedPluginErrors(t *testing.T) { + detail := `{"code":"rate_limited","message":"account limited","http_status":429,"retry_after_seconds":60,"credential_scoped":true}` + for _, tc := range []struct { + name string + send func(host *Host, streamID string) error + }{ + {"emit", func(host *Host, streamID string) error { + _, err := host.callHostStreamEmit(context.Background(), []byte(`{"stream_id":"`+streamID+`","error_detail":`+detail+`}`)) + return err + }}, + {"close", func(host *Host, streamID string) error { + _, err := host.callHostStreamClose([]byte(`{"stream_id":"` + streamID + `","error_detail":` + detail + `}`)) + return err + }}, + } { + t.Run(tc.name, func(t *testing.T) { + host := New() + streamID, chunks, cleanup := host.streams.open(context.Background()) + t.Cleanup(cleanup) + if err := tc.send(host, streamID); err != nil { + t.Fatalf("stream %s callback error = %v", tc.name, err) + } + chunk := <-chunks + if chunk.Err == nil { + t.Fatal("stream error was dropped") + } + var scoped interface{ IsCredentialScoped() bool } + if !errors.As(chunk.Err, &scoped) || !scoped.IsCredentialScoped() { + t.Fatalf("stream error %T does not report IsCredentialScoped() = true", chunk.Err) + } + }) + } +} + +func TestPluginExecutorStream429CredentialScopedCoolsEntireCredential(t *testing.T) { + coreauth.SetQuotaCooldownDisabled(false) + + for _, tc := range []struct { + name string + detail string + }{ + { + name: "rate_limited with credential_scoped", + detail: `{"code":"rate_limited","message":"account limited","http_status":429,"retry_after_seconds":90,"credential_scoped":true}`, + }, + { + name: "transient_rate_limit overridden by credential_scoped", + detail: `{"code":"transient_rate_limit","message":"account limited","http_status":429,"retry_after_seconds":90,"credential_scoped":true}`, + }, + { + name: "request_scoped overridden by credential_scoped", + detail: `{"code":"request_scoped","message":"account limited","http_status":429,"retry_after_seconds":90,"credential_scoped":true}`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + var ( + provider = "plugin-provider-stream-cred-scope-" + tc.name + modelA = "plugin-stream-model-a-" + tc.name + modelB = "plugin-stream-model-b-" + tc.name + authID = "auth-stream-cred-scope-" + tc.name + ) + executorRecord := normalizeTestCapabilityRecord(capabilityRecord{id: "executor-stream-cred-scope-" + tc.name}) + host := newHostWithRecords(executorRecord) + + var streamCalls int + exec := &fakeExecutor{ + identifier: provider, + executeStream: func(ctx context.Context, _ pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + streamCalls++ + streamID, chunks, cleanup := host.streams.open(ctx) + if _, errClose := host.callHostStreamClose([]byte(`{"stream_id":"` + streamID + `","error_detail":` + tc.detail + `}`)); errClose != nil { + cleanup() + return pluginapi.ExecutorStreamResponse{}, errClose + } + return pluginapi.ExecutorStreamResponse{ + Chunks: cleanupWhenStreamDone(ctx, chunks, cleanup), + }, nil + }, + } + adapter := newExecutorAdapterForRecordForTest(host, executorRecord, exec, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + []sdktranslator.Format{sdktranslator.FormatOpenAI}, + ) + adapter.provider = provider + + authMgr := coreauth.NewManager(nil, nil, nil) + authMgr.SetRetryConfig(0, 0, 0) + authMgr.RegisterExecutor(adapter) + auth := &coreauth.Auth{ + ID: authID, + Provider: provider, + Status: coreauth.StatusActive, + FileName: authID + ".json", + Attributes: map[string]string{"type": "oauth"}, + } + if _, errReg := authMgr.Register(context.Background(), auth); errReg != nil { + t.Fatalf("Register auth: %v", errReg) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{ + {ID: modelA}, + {ID: modelB}, + }) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + opts := coreexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + + resStreamA, errStreamA := authMgr.ExecuteStream(context.Background(), []string{provider}, coreexecutor.Request{ + Model: modelA, + Payload: []byte(`{"model":"` + modelA + `","stream":true}`), + }, opts) + if errStreamA == nil { + if resStreamA == nil || resStreamA.Chunks == nil { + t.Fatal("ExecuteStream(modelA) returned nil error and nil chunks") + } + chunk := <-resStreamA.Chunks + errStreamA = chunk.Err + } + if errStreamA == nil { + t.Fatal("ExecuteStream(modelA) returned nil stream error, want 429") + } + if streamCalls != 1 { + t.Fatalf("streamCalls after modelA = %d, want 1", streamCalls) + } + + updated, okAuth := authMgr.GetByID(authID) + if !okAuth || updated == nil { + t.Fatalf("GetByID(%q) missing", authID) + } + if !updated.Unavailable || !updated.Quota.Exceeded || updated.Quota.Reason != "credential_quota" { + t.Fatalf("updated auth quota = %+v (unavailable=%v), want unavailable with reason=credential_quota", updated.Quota, updated.Unavailable) + } + + _, errStreamB := authMgr.ExecuteStream(context.Background(), []string{provider}, coreexecutor.Request{ + Model: modelB, + Payload: []byte(`{"model":"` + modelB + `","stream":true}`), + }, opts) + if errStreamB == nil { + t.Fatal("ExecuteStream(modelB) returned nil error, want cooldown rejection") + } + if streamCalls != 1 { + t.Fatalf("streamCalls after modelB = %d, want 1 (credential must not be selected again for sibling model)", streamCalls) + } + }) + } +} diff --git a/internal/pluginhost/rpc_schema_test.go b/internal/pluginhost/rpc_schema_test.go index bf5a1e83b8e..7d5b7e4b47b 100644 --- a/internal/pluginhost/rpc_schema_test.go +++ b/internal/pluginhost/rpc_schema_test.go @@ -11,6 +11,7 @@ import ( "testing" sdkaccess "github.com/router-for-me/CLIProxyAPI/v8/sdk/access" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) @@ -906,6 +907,105 @@ func TestRPCSchedulerAndUsageSelectionIDCompatibility(t *testing.T) { } } +func TestUsageRecordRPCModalityTokenCounters(t *testing.T) { + var deliveredRaw []byte + var deliveredRecord pluginapi.UsageRecord + rpcAdapter := &rpcPluginAdapter{ + id: "rpc-usage-modality", + client: &lifecycleTestClient{call: func(_ context.Context, method string, request []byte) ([]byte, error) { + if method != pluginabi.MethodUsageHandle { + return nil, io.ErrUnexpectedEOF + } + deliveredRaw = bytes.Clone(request) + if errUnmarshal := json.Unmarshal(request, &deliveredRecord); errUnmarshal != nil { + return nil, errUnmarshal + } + return marshalRPCResult(rpcEmptyResponse{}) + }}, + } + + host := newHostWithRecords(capabilityRecord{ + id: "rpc-usage-modality", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + UsagePlugin: rpcAdapter, + }}, + }) + rpcAdapter.host = host + adapter := &usageAdapter{ + host: host, + pluginID: "rpc-usage-modality", + plugin: rpcAdapter, + } + + adapter.HandleUsage(context.Background(), coreusage.Record{ + Provider: "openai", + Model: "gpt-image-2.5-flare", + Detail: coreusage.Detail{ + InputTokens: 80, + OutputTokens: 1056, + CachedTokens: 40, + CacheReadTokens: 40, + TotalTokens: 1136, + TextInputTokens: coreusage.Int64Ptr(30), + ImageInputTokens: coreusage.Int64Ptr(50), + CachedTextTokens: coreusage.Int64Ptr(25), + CachedImageTokens: coreusage.Int64Ptr(15), + TextOutputTokens: coreusage.Int64Ptr(0), + ImageOutputTokens: coreusage.Int64Ptr(1056), + }, + }) + + got := deliveredRecord.Detail + if got.InputTokens != 80 || got.OutputTokens != 1056 || got.CachedTokens != 40 || got.CacheReadTokens != 40 || got.TotalTokens != 1136 { + t.Fatalf("delivered aggregate Detail = %+v", got) + } + if got.TextInputTokens == nil || *got.TextInputTokens != 30 { + t.Fatalf("delivered TextInputTokens = %v, want 30", got.TextInputTokens) + } + if got.ImageInputTokens == nil || *got.ImageInputTokens != 50 { + t.Fatalf("delivered ImageInputTokens = %v, want 50", got.ImageInputTokens) + } + if got.CachedTextTokens == nil || *got.CachedTextTokens != 25 { + t.Fatalf("delivered CachedTextTokens = %v, want 25", got.CachedTextTokens) + } + if got.CachedImageTokens == nil || *got.CachedImageTokens != 15 { + t.Fatalf("delivered CachedImageTokens = %v, want 15", got.CachedImageTokens) + } + if got.TextOutputTokens == nil || *got.TextOutputTokens != 0 { + t.Fatalf("delivered TextOutputTokens = %v, want explicit 0", got.TextOutputTokens) + } + if got.ImageOutputTokens == nil || *got.ImageOutputTokens != 1056 { + t.Fatalf("delivered ImageOutputTokens = %v, want 1056", got.ImageOutputTokens) + } + if got.AudioInputTokens != nil || got.VideoInputTokens != nil || got.FileInputTokens != nil || + got.CachedAudioTokens != nil || got.AudioOutputTokens != nil { + t.Fatalf("expected absent modality counters to remain nil over RPC, got %+v", got) + } + + // Legacy plugins compiled against the 7-field UsageDetail must still unmarshal cleanly. + type legacyUsageDetail struct { + InputTokens int64 + OutputTokens int64 + ReasoningTokens int64 + CachedTokens int64 + CacheReadTokens int64 + CacheCreationTokens int64 + TotalTokens int64 + } + type legacyUsageRecordWithDetail struct { + Provider string + Model string + Detail legacyUsageDetail + } + var legacyRec legacyUsageRecordWithDetail + if errLegacy := json.Unmarshal(deliveredRaw, &legacyRec); errLegacy != nil { + t.Fatalf("Unmarshal(legacyUsageRecordWithDetail) error = %v", errLegacy) + } + if legacyRec.Detail.InputTokens != 80 || legacyRec.Detail.OutputTokens != 1056 || legacyRec.Detail.TotalTokens != 1136 { + t.Fatalf("legacyUsageRecordWithDetail.Detail = %+v", legacyRec.Detail) + } +} + type recordingRawPatchClient struct { schemaVersion uint32 methods []string diff --git a/internal/pluginhost/shadow.go b/internal/pluginhost/shadow.go index 3884b11e25e..2438e3b7ced 100644 --- a/internal/pluginhost/shadow.go +++ b/internal/pluginhost/shadow.go @@ -10,9 +10,12 @@ import ( "io" "os" "path/filepath" + "runtime" "strconv" "strings" "sync" + + log "github.com/sirupsen/logrus" ) // Plugin dynamic libraries are loaded from content-addressed shadow copies. @@ -25,21 +28,82 @@ const ( shadowPluginPrefix = "cliproxy-plugin-" shadowPluginTempPrefix = ".cliproxy-plugin-" shadowPluginProcessDirPrefix = "pid-" + shadowPluginCacheDirName = "cache" + shadowPluginLocalDirName = ".shadow" shadowPluginDigestLength = 32 ) -var shadowPluginCleanupOnce sync.Once +type cachedPluginDigest struct { + size int64 + modTimeNano int64 + digest string +} + +var ( + shadowPluginCleanupOnce sync.Once + cleanedShadowTempDirs sync.Map + pluginDigestCache sync.Map + pruneShadowBeforeLockHook func(shadowFileName string) +) func shadowCopyPlugin(file pluginFile) (string, error) { - dir, errDir := shadowPluginDir() + path, release, errShadow := shadowCopyPluginInternal(file, false) + if release != nil { + release() + } + return path, errShadow +} + +func acquireShadowPluginLoad(file pluginFile) (string, func(), error) { + return shadowCopyPluginInternal(file, true) +} + +func acquireShadowPluginFallbackLoad(file pluginFile, failedLoadPath string) (string, func(), error) { + source := filepath.Clean(file.Path) + dir, errDir := shadowPluginFallbackDir(source) if errDir != nil { - return "", errDir + return "", nil, errDir + } + if failedLoadPath != "" && filepath.Clean(filepath.Dir(failedLoadPath)) == filepath.Clean(dir) { + return "", nil, fmt.Errorf("shadow plugin already in fallback directory") } + size, digest, errHash := cachedOrHashPluginFile(source) + if errHash != nil { + return "", nil, errHash + } + cleanupShadowTempFilesOnce(dir) + return shadowCopyPluginToDirWithGuard(file, source, dir, size, digest, true) +} + +func shadowCopyPluginInternal(file pluginFile, holdGuard bool) (string, func(), error) { shadowPluginCleanupOnce.Do(func() { - removeStaleShadowPlugins(dir) - removeStaleShadowPluginDirs(filepath.Dir(dir), os.Getpid()) + removeStaleShadowPluginDirs(shadowPluginBaseDir(), os.Getpid()) }) - return shadowCopyPluginToDir(file, dir) + source := filepath.Clean(file.Path) + size, digest, errHash := hashPluginFile(source) + if errHash != nil { + return "", nil, errHash + } + if localDir, errLocalDir := shadowPluginLocalDir(source); errLocalDir == nil { + cleanupShadowTempFilesOnce(localDir) + if target, release, errCopy := shadowCopyPluginToDirWithGuard(file, source, localDir, size, digest, holdGuard); errCopy == nil { + return target, release, nil + } + } + dir, errDir := shadowPluginFallbackDir(source) + if errDir != nil { + return "", nil, errDir + } + cleanupShadowTempFilesOnce(dir) + return shadowCopyPluginToDirWithGuard(file, source, dir, size, digest, holdGuard) +} + +func cleanupShadowTempFilesOnce(dir string) { + cleaned := filepath.Clean(dir) + if _, loaded := cleanedShadowTempDirs.LoadOrStore(cleaned, struct{}{}); loaded { + return + } + removeStaleShadowTempFiles(cleaned) } // removeStaleShadowPluginDirs removes shadow directories left behind by @@ -72,7 +136,70 @@ func removeStaleShadowPluginDirs(baseDir string, currentPID int) { func shadowCopyPluginToDir(file pluginFile, dir string) (string, error) { source := filepath.Clean(file.Path) - tmp, errTemp := os.CreateTemp(dir, shadowPluginTempPrefix+file.ID+"-*"+filepath.Ext(source)) + size, digest, errHash := hashPluginFile(source) + if errHash != nil { + return "", errHash + } + return shadowCopyPluginToDirWithDigest(file, dir, size, digest) +} + +func shadowCopyPluginToDirWithGuard(file pluginFile, sourcePath string, dir string, size int64, digest string, holdGuard bool) (string, func(), error) { + releaseGuard := func() {} + if holdGuard { + sourceDir := filepath.Dir(filepath.Clean(sourcePath)) + unlockCoord, errLock := acquireShadowCoordLock(sourceDir, dir, false, false) + if errLock != nil { + return "", nil, errLock + } + releaseGuard = unlockCoord + } + target, errCopy := shadowCopyPluginToDirWithDigest(file, dir, size, digest) + if errCopy != nil { + releaseGuard() + return "", nil, errCopy + } + return target, releaseGuard, nil +} + +func acquireShadowCoordLock(sourceDir string, fallbackDir string, exclusive bool, nonBlocking bool) (func(), error) { + baseDir := fallbackDir + if strings.TrimSpace(sourceDir) != "" { + baseDir = filepath.Join(shadowPluginBaseDir(), shadowPluginCacheDirName, shadowDirectoryDigest(sourceDir)) + } + if errMkdir := os.MkdirAll(baseDir, 0o700); errMkdir != nil { + return nil, errMkdir + } + lockPath := filepath.Join(baseDir, ".cliproxy-coord.lock") + lockFile, errOpen := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0o600) + if errOpen != nil { + return nil, errOpen + } + if errLock := lockShadowCoordFile(lockFile, exclusive, nonBlocking); errLock != nil { + if errClose := lockFile.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: coord lock file close failed") + } + return nil, errLock + } + var once sync.Once + return func() { + once.Do(func() { + unlockShadowCoordFile(lockFile) + if errClose := lockFile.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: coord lock file close failed") + } + }) + }, nil +} + +func shadowCopyPluginToDirWithDigest(file pluginFile, dir string, size int64, digest string) (string, error) { + source := filepath.Clean(file.Path) + ext := filepath.Ext(source) + target := shadowPluginPath(dir, file.ID, digest, ext) + if shadowPluginMatches(target, size, digest) { + return target, nil + } + + tmp, errTemp := os.CreateTemp(dir, shadowPluginTempPattern(os.Getpid(), file.ID, ext)) if errTemp != nil { return "", errTemp } @@ -86,56 +213,185 @@ func shadowCopyPluginToDir(file pluginFile, dir string) (string, error) { in, errOpen := os.Open(source) if errOpen != nil { - _ = tmp.Close() + if errClose := tmp.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: failed to close temp shadow plugin file") + } return "", errOpen } defer func() { - _ = in.Close() + if errClose := in.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: failed to close plugin source file") + } }() hasher := sha256.New() - size, errCopy := io.Copy(io.MultiWriter(tmp, hasher), in) + copiedSize, errCopy := io.Copy(io.MultiWriter(tmp, hasher), in) if errCopy != nil { - _ = tmp.Close() + if errClose := tmp.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: failed to close temp shadow plugin file") + } return "", errCopy } if errClose := tmp.Close(); errClose != nil { return "", errClose } - digest := hex.EncodeToString(hasher.Sum(nil)) - target := shadowPluginPath(dir, file.ID, digest, filepath.Ext(source)) + copiedDigest := hex.EncodeToString(hasher.Sum(nil)) + target = shadowPluginPath(dir, file.ID, copiedDigest, ext) + if errPublish := publishShadowTempFile(tmpName, target, copiedSize, copiedDigest); errPublish != nil { + return "", errPublish + } + removeTemp = false + return target, nil +} + +func publishShadowTempFile(tmpName string, target string, size int64, digest string) error { + if errPublish := publishShadowPlugin(tmpName, target); errPublish == nil { + return nil + } if shadowPluginMatches(target, size, digest) { - return target, nil + removeShadowPlugin(tmpName) + return nil + } + if errRename := os.Rename(tmpName, target); errRename == nil { + return nil + } + if shadowPluginMatches(target, size, digest) { + removeShadowPlugin(tmpName) + return nil } if errRemove := os.Remove(target); errRemove != nil && !errors.Is(errRemove, os.ErrNotExist) { if shadowPluginMatches(target, size, digest) { - return target, nil + removeShadowPlugin(tmpName) + return nil } removeShadowPlugin(target) - return "", fmt.Errorf("remove stale shadow plugin: %w", errRemove) + return fmt.Errorf("remove stale shadow plugin: %w", errRemove) } if errRename := os.Rename(tmpName, target); errRename != nil { if shadowPluginMatches(target, size, digest) { - return target, nil + removeShadowPlugin(tmpName) + return nil } - return "", fmt.Errorf("move shadow plugin: %w", errRename) + return fmt.Errorf("move shadow plugin: %w", errRename) } - removeTemp = false - return target, nil + return nil +} + +func hashPluginFile(path string) (int64, string, error) { + cleaned := filepath.Clean(path) + file, errOpen := os.Open(cleaned) + if errOpen != nil { + return 0, "", errOpen + } + defer func() { + if errClose := file.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: failed to close plugin file after hashing") + } + }() + hasher := sha256.New() + size, errCopy := io.Copy(hasher, file) + if errCopy != nil { + return 0, "", errCopy + } + digest := hex.EncodeToString(hasher.Sum(nil)) + if info, errStat := file.Stat(); errStat == nil && info.Mode().IsRegular() { + pluginDigestCache.Store(cleaned, cachedPluginDigest{ + size: size, + modTimeNano: info.ModTime().UnixNano(), + digest: digest, + }) + } + return size, digest, nil +} + +func cachedOrHashPluginFile(path string) (int64, string, error) { + cleaned := filepath.Clean(path) + if info, errStat := os.Stat(cleaned); errStat == nil && info.Mode().IsRegular() { + if raw, okCache := pluginDigestCache.Load(cleaned); okCache { + if cached, okType := raw.(cachedPluginDigest); okType && + cached.size == info.Size() && + cached.modTimeNano == info.ModTime().UnixNano() && + cached.digest != "" { + return cached.size, cached.digest, nil + } + } + } + return hashPluginFile(cleaned) +} + +func shadowPluginBaseDir() string { + return filepath.Join(os.TempDir(), "cliproxy-pluginhost") +} + +func shadowPluginLocalDir(sourcePath string) (string, error) { + dir := filepath.Join(filepath.Dir(filepath.Clean(sourcePath)), shadowPluginLocalDirName) + if errMkdir := os.MkdirAll(dir, 0o700); errMkdir != nil { + return "", errMkdir + } + return dir, nil } func shadowPluginDir() (string, error) { - dir := filepath.Join(os.TempDir(), "cliproxy-pluginhost", shadowPluginProcessDirName(os.Getpid())) + dir := filepath.Join(shadowPluginBaseDir(), shadowPluginCacheDirName) + if errMkdir := os.MkdirAll(dir, 0o700); errMkdir != nil { + return "", errMkdir + } + return dir, nil +} + +func shadowPluginFallbackDir(sourcePath string) (string, error) { + baseDir, errDir := shadowPluginDir() + if errDir != nil { + return "", errDir + } + sourceDir := filepath.Dir(filepath.Clean(sourcePath)) + dir := filepath.Join(baseDir, shadowDirectoryDigest(sourceDir)) if errMkdir := os.MkdirAll(dir, 0o700); errMkdir != nil { return "", errMkdir } return dir, nil } +func shadowDirectoryDigest(dir string) string { + cleaned := filepath.Clean(dir) + if absDir, errAbs := filepath.Abs(cleaned); errAbs == nil { + cleaned = filepath.Clean(absDir) + } + sum := sha256.Sum256([]byte(cleaned)) + return hex.EncodeToString(sum[:])[:shadowPluginDigestLength] +} + func shadowPluginProcessDirName(pid int) string { return fmt.Sprintf("%s%d", shadowPluginProcessDirPrefix, pid) } -func removeStaleShadowPlugins(dir string) { +func shadowPluginTempPattern(pid int, id string, extension string) string { + return fmt.Sprintf("%s%s%d-%s-*%s", shadowPluginTempPrefix, shadowPluginProcessDirPrefix, pid, id, extension) +} + +func shadowTempFileOwnedByLiveProcess(name string) bool { + if !strings.HasPrefix(name, shadowPluginTempPrefix) { + return false + } + rest := strings.TrimPrefix(name, shadowPluginTempPrefix) + if !strings.HasPrefix(rest, shadowPluginProcessDirPrefix) { + return false + } + rest = strings.TrimPrefix(rest, shadowPluginProcessDirPrefix) + sep := strings.IndexByte(rest, '-') + if sep <= 0 { + return false + } + pid, errParse := strconv.Atoi(rest[:sep]) + if errParse != nil || pid <= 0 { + return false + } + if pid == os.Getpid() { + return true + } + return shadowPluginProcessAlive(pid) +} + +func removeStaleShadowTempFiles(dir string) { entries, errRead := os.ReadDir(dir) if errRead != nil { return @@ -145,12 +401,178 @@ func removeStaleShadowPlugins(dir string) { continue } name := entry.Name() - if strings.HasPrefix(name, shadowPluginPrefix) || strings.HasPrefix(name, shadowPluginTempPrefix) { + if strings.HasPrefix(name, shadowPluginTempPrefix) && !shadowTempFileOwnedByLiveProcess(name) { removeShadowPlugin(filepath.Join(dir, name)) } } } +func shadowDirHasLiveProcessActivity(entries []os.DirEntry) bool { + for _, entry := range entries { + if entry == nil || entry.IsDir() { + continue + } + name := entry.Name() + if strings.HasPrefix(name, shadowPluginTempPrefix) && shadowTempFileOwnedByLiveProcess(name) { + return true + } + } + return false +} + +func removeStaleShadowPlugins(dir string, keepNames ...string) { + pruneShadowDirForSourceDir(dir, "", keepNames) +} + +func shadowSourceDirHasLiveProcessActivity(dir string, sourceDir string) bool { + if entries, errRead := os.ReadDir(dir); errRead == nil && shadowDirHasLiveProcessActivity(entries) { + return true + } + if sourceDir == "" { + return false + } + fallbackDir := filepath.Join(shadowPluginBaseDir(), shadowPluginCacheDirName, shadowDirectoryDigest(sourceDir)) + if filepath.Clean(fallbackDir) == filepath.Clean(dir) { + return false + } + fallbackEntries, errRead := os.ReadDir(fallbackDir) + if errRead != nil { + return false + } + return shadowDirHasLiveProcessActivity(fallbackEntries) +} + +func pruneShadowDirForSourceDir(dir string, sourceDir string, keepNames []string) { + entries, errRead := os.ReadDir(dir) + if errRead != nil { + return + } + keep := make(map[string]struct{}, len(keepNames)) + for _, name := range keepNames { + if trimmed := strings.TrimSpace(name); trimmed != "" { + keep[filepath.Base(trimmed)] = struct{}{} + } + } + for _, entry := range entries { + if entry == nil || entry.IsDir() { + continue + } + name := entry.Name() + if strings.HasPrefix(name, shadowPluginTempPrefix) { + if !shadowTempFileOwnedByLiveProcess(name) { + removeShadowPlugin(filepath.Join(dir, name)) + } + continue + } + if strings.HasPrefix(name, shadowPluginPrefix) { + if _, okKeep := keep[name]; okKeep { + continue + } + if hook := pruneShadowBeforeLockHook; hook != nil { + hook(name) + } + removeStaleShadowPluginCandidate(dir, sourceDir, name) + } + } +} + +func removeStaleShadowPluginCandidate(dir string, sourceDir string, name string) { + unlockCoord, errLock := acquireShadowCoordLock(sourceDir, dir, true, true) + if errLock != nil { + return + } + defer unlockCoord() + if shadowSourceDirHasLiveProcessActivity(dir, sourceDir) { + return + } + if sourceDir != "" && shadowMatchesCurrentSourceDir(sourceDir, name) { + return + } + removeShadowPlugin(filepath.Join(dir, name)) +} + +func shadowMatchesCurrentSourceDir(sourceDir string, shadowFileName string) bool { + entries, errRead := os.ReadDir(sourceDir) + if errRead != nil { + return false + } + ext := filepath.Ext(shadowFileName) + for _, entry := range entries { + if entry == nil || !entry.Type().IsRegular() { + continue + } + filePath := filepath.Join(sourceDir, entry.Name()) + file, okFile := pluginFileFromPath(filePath, ext) + if !okFile { + continue + } + if !strings.HasPrefix(shadowFileName, shadowPluginPrefix+file.ID+"-") { + continue + } + _, digest, errHash := cachedOrHashPluginFile(filePath) + if errHash != nil { + continue + } + if filepath.Base(shadowPluginPath("", file.ID, digest, ext)) == shadowFileName { + return true + } + } + return false +} + +type shadowDirTarget struct { + shadowDir string + sourceDir string +} + +func cleanupStaleShadowPlugins(root string, activeFiles []pluginFile) { + keepNames := make([]string, 0, len(activeFiles)) + seenDirs := make(map[string]struct{}) + targets := make([]shadowDirTarget, 0, len(activeFiles)*2+4) + addShadowDirsForSourceDir := func(sourceDir string) { + cleanedSourceDir := filepath.Clean(sourceDir) + for _, dir := range []string{ + filepath.Join(cleanedSourceDir, shadowPluginLocalDirName), + filepath.Join(shadowPluginBaseDir(), shadowPluginCacheDirName, shadowDirectoryDigest(cleanedSourceDir)), + } { + cleaned := filepath.Clean(dir) + if _, exists := seenDirs[cleaned]; exists { + continue + } + seenDirs[cleaned] = struct{}{} + targets = append(targets, shadowDirTarget{ + shadowDir: cleaned, + sourceDir: cleanedSourceDir, + }) + } + } + for _, candidateDir := range candidateDirs(root, runtime.GOOS, runtime.GOARCH) { + addShadowDirsForSourceDir(candidateDir) + } + seenSources := make(map[string]struct{}, len(activeFiles)) + for _, file := range activeFiles { + id := strings.TrimSpace(file.ID) + path := strings.TrimSpace(file.Path) + if id == "" || path == "" { + continue + } + source := filepath.Clean(path) + if _, seen := seenSources[source]; seen { + continue + } + seenSources[source] = struct{}{} + addShadowDirsForSourceDir(filepath.Dir(source)) + _, digest, errHash := cachedOrHashPluginFile(source) + if errHash != nil { + continue + } + keepNames = append(keepNames, filepath.Base(shadowPluginPath("", id, digest, filepath.Ext(source)))) + } + for _, target := range targets { + pruneShadowDirForSourceDir(target.shadowDir, target.sourceDir, keepNames) + } +} + func shadowPluginPath(dir string, id string, digest string, extension string) string { if len(digest) > shadowPluginDigestLength { digest = digest[:shadowPluginDigestLength] @@ -171,7 +593,9 @@ func shadowPluginMatches(path string, size int64, digest string) bool { return false } defer func() { - _ = file.Close() + if errClose := file.Close(); errClose != nil { + log.WithError(errClose).Debug("pluginhost: failed to close shadow plugin file") + } }() hasher := sha256.New() if _, errCopy := io.Copy(hasher, file); errCopy != nil { diff --git a/internal/pluginhost/shadow_test.go b/internal/pluginhost/shadow_test.go index 42301947169..0408fdd488d 100644 --- a/internal/pluginhost/shadow_test.go +++ b/internal/pluginhost/shadow_test.go @@ -12,6 +12,7 @@ import ( "runtime" "strings" "testing" + "time" ) func testPluginLibraryExtension() string { @@ -25,17 +26,509 @@ func testPluginLibraryExtension() string { } } -func TestShadowPluginDirIsProcessScoped(t *testing.T) { +func TestShadowPluginDirUsesSharedCache(t *testing.T) { dir, errDir := shadowPluginDir() if errDir != nil { t.Fatalf("shadowPluginDir() error = %v", errDir) } - want := filepath.Join(os.TempDir(), "cliproxy-pluginhost", fmt.Sprintf("pid-%d", os.Getpid())) + want := filepath.Join(os.TempDir(), "cliproxy-pluginhost", "cache") if dir != want { t.Fatalf("shadowPluginDir() = %q, want %q", dir, want) } } +func TestShadowCopyPluginReusesAcrossProcessRestartAndUsesLocalShadowDir(t *testing.T) { + pluginDir := t.TempDir() + source := filepath.Join(pluginDir, "alpha-v1.0.0"+testPluginLibraryExtension()) + content := []byte("plugin-v1") + if errWrite := os.WriteFile(source, content, 0o644); errWrite != nil { + t.Fatalf("WriteFile() error = %v", errWrite) + } + file := pluginFile{ID: "alpha", Path: source, Version: "1.0.0"} + + first, errFirst := shadowCopyPlugin(file) + if errFirst != nil { + t.Fatalf("shadowCopyPlugin() first error = %v", errFirst) + } + wantShadowDir := filepath.Join(pluginDir, ".shadow") + if filepath.Dir(first) != wantShadowDir { + t.Fatalf("shadowCopyPlugin() dir = %q, want %q", filepath.Dir(first), wantShadowDir) + } + + past := time.Date(2025, 1, 2, 3, 4, 5, 0, time.UTC) + if errCh := os.Chtimes(first, past, past); errCh != nil { + t.Fatalf("Chtimes(first) error = %v", errCh) + } + if errCh := os.Chtimes(wantShadowDir, past, past); errCh != nil { + t.Fatalf("Chtimes(shadowDir) error = %v", errCh) + } + + childPath := runChildShadowCopyPlugin(t, file) + if childPath != first { + t.Fatalf("child process shadowCopyPlugin() = %q, want reused %q", childPath, first) + } + info, errStat := os.Stat(first) + if errStat != nil { + t.Fatalf("Stat(%s) after child restart error = %v", first, errStat) + } + if !info.ModTime().Equal(past) { + t.Fatalf("shadow file ModTime = %v, want unchanged %v", info.ModTime(), past) + } + dirInfo, errDirStat := os.Stat(wantShadowDir) + if errDirStat != nil { + t.Fatalf("Stat(%s) after child restart error = %v", wantShadowDir, errDirStat) + } + if !dirInfo.ModTime().Equal(past) { + t.Fatalf("shadow dir ModTime = %v, want unchanged %v (no temp file should be created on cache hit)", dirInfo.ModTime(), past) + } +} + +func TestShadowCopyPluginToDirSkipsTempWriteWhenShadowMatches(t *testing.T) { + dir := t.TempDir() + source := filepath.Join(t.TempDir(), "alpha"+testPluginLibraryExtension()) + if errWrite := os.WriteFile(source, []byte("plugin-v1"), 0o644); errWrite != nil { + t.Fatalf("WriteFile() error = %v", errWrite) + } + file := pluginFile{ID: "alpha", Path: source} + + first, errFirst := shadowCopyPluginToDir(file, dir) + if errFirst != nil { + t.Fatalf("shadowCopyPluginToDir() first error = %v", errFirst) + } + past := time.Date(2025, 1, 2, 3, 4, 5, 0, time.UTC) + if errCh := os.Chtimes(first, past, past); errCh != nil { + t.Fatalf("Chtimes(first) error = %v", errCh) + } + if errCh := os.Chtimes(dir, past, past); errCh != nil { + t.Fatalf("Chtimes(dir) error = %v", errCh) + } + + second, errSecond := shadowCopyPluginToDir(file, dir) + if errSecond != nil { + t.Fatalf("shadowCopyPluginToDir() second error = %v", errSecond) + } + if second != first { + t.Fatalf("second shadow path = %q, want %q", second, first) + } + dirInfo, errDirStat := os.Stat(dir) + if errDirStat != nil { + t.Fatalf("Stat(dir) error = %v", errDirStat) + } + if !dirInfo.ModTime().Equal(past) { + t.Fatalf("shadow dir ModTime = %v, want unchanged %v (temp file was written before checking shadowPluginMatches)", dirInfo.ModTime(), past) + } +} + +func TestShadowCopyPluginFallsBackToSharedCacheWhenLocalShadowDirUnavailable(t *testing.T) { + pluginDir := t.TempDir() + // Create a regular file named ".shadow" so creating the local .shadow directory fails. + blockerPath := filepath.Join(pluginDir, shadowPluginLocalDirName) + if errWrite := os.WriteFile(blockerPath, []byte("blocker"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(blocker) error = %v", errWrite) + } + source := filepath.Join(pluginDir, "fallback-v1.0.0"+testPluginLibraryExtension()) + if errWrite := os.WriteFile(source, []byte("plugin-fallback"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(source) error = %v", errWrite) + } + file := pluginFile{ID: "fallback", Path: source, Version: "1.0.0"} + + got, errShadow := shadowCopyPlugin(file) + if errShadow != nil { + t.Fatalf("shadowCopyPlugin() error = %v", errShadow) + } + t.Cleanup(func() { + removeShadowPlugin(got) + }) + wantDir, errDir := shadowPluginFallbackDir(source) + if errDir != nil { + t.Fatalf("shadowPluginFallbackDir() error = %v", errDir) + } + if filepath.Dir(got) != wantDir { + t.Fatalf("shadowCopyPlugin() fallback dir = %q, want %q", filepath.Dir(got), wantDir) + } +} + +func TestShadowCopyPluginFallbackCacheIsolatesDistinctPluginDirs(t *testing.T) { + ext := testPluginLibraryExtension() + rootA := t.TempDir() + platformDirA := filepath.Join(rootA, runtime.GOOS, runtime.GOARCH) + if errMkdir := os.MkdirAll(platformDirA, 0o755); errMkdir != nil { + t.Fatalf("MkdirAll(platformDirA) error = %v", errMkdir) + } + if errWrite := os.WriteFile(filepath.Join(platformDirA, shadowPluginLocalDirName), []byte("blocker"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(blockerA) error = %v", errWrite) + } + sourceA := filepath.Join(platformDirA, "alpha-v1.0.0"+ext) + if errWrite := os.WriteFile(sourceA, []byte("alpha-install-A"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(sourceA) error = %v", errWrite) + } + fileA := pluginFile{ID: "alpha", Path: sourceA, Version: "1.0.0"} + + rootB := t.TempDir() + platformDirB := filepath.Join(rootB, runtime.GOOS, runtime.GOARCH) + if errMkdir := os.MkdirAll(platformDirB, 0o755); errMkdir != nil { + t.Fatalf("MkdirAll(platformDirB) error = %v", errMkdir) + } + if errWrite := os.WriteFile(filepath.Join(platformDirB, shadowPluginLocalDirName), []byte("blocker"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(blockerB) error = %v", errWrite) + } + sourceB := filepath.Join(platformDirB, "alpha-v1.0.0"+ext) + if errWrite := os.WriteFile(sourceB, []byte("alpha-install-B"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(sourceB) error = %v", errWrite) + } + fileB := pluginFile{ID: "alpha", Path: sourceB, Version: "1.0.0"} + + shadowA, errShadowA := shadowCopyPlugin(fileA) + if errShadowA != nil { + t.Fatalf("shadowCopyPlugin(fileA) error = %v", errShadowA) + } + t.Cleanup(func() { removeShadowPluginDir(filepath.Dir(shadowA)) }) + + shadowB, errShadowB := shadowCopyPlugin(fileB) + if errShadowB != nil { + t.Fatalf("shadowCopyPlugin(fileB) error = %v", errShadowB) + } + t.Cleanup(func() { removeShadowPluginDir(filepath.Dir(shadowB)) }) + + if filepath.Dir(shadowA) == filepath.Dir(shadowB) { + t.Fatalf("fallback cache dirs collided across installations: %q", filepath.Dir(shadowA)) + } + + if errCleanup := cleanupUnselectedPluginFiles(rootA, []pluginFile{fileA}); errCleanup != nil { + t.Fatalf("cleanupUnselectedPluginFiles(rootA) error = %v", errCleanup) + } + if _, errStat := os.Stat(shadowB); errStat != nil { + t.Fatalf("Stat(shadowB %s) after rootA cleanup error = %v, want kept", shadowB, errStat) + } +} + +func TestShadowDirectoryDigestResolvesRelativePathsByWorkingDirectory(t *testing.T) { + workA := t.TempDir() + workB := t.TempDir() + relDir := filepath.Join("plugins", runtime.GOOS, runtime.GOARCH) + + // t.Chdir restores the original directory in a cleanup registered after the + // TempDir cleanups, so it runs first and Windows can remove the directories. + t.Chdir(workA) + digestA := shadowDirectoryDigest(relDir) + + t.Chdir(workB) + digestB := shadowDirectoryDigest(relDir) + + if digestA == digestB { + t.Fatalf("shadowDirectoryDigest(%q) collided across working directories: %s", relDir, digestA) + } +} + +func TestShadowCopyPluginConcurrentPublishNeverUnlinksExistingTarget(t *testing.T) { + dir := t.TempDir() + source := filepath.Join(t.TempDir(), "alpha"+testPluginLibraryExtension()) + content := []byte("concurrent-plugin-content") + if errWrite := os.WriteFile(source, content, 0o644); errWrite != nil { + t.Fatalf("WriteFile() error = %v", errWrite) + } + file := pluginFile{ID: "alpha", Path: source} + + first, errFirst := shadowCopyPluginToDir(file, dir) + if errFirst != nil { + t.Fatalf("shadowCopyPluginToDir() initial error = %v", errFirst) + } + infoBefore, errStatBefore := os.Stat(first) + if errStatBefore != nil { + t.Fatalf("Stat(first) error = %v", errStatBefore) + } + size, digest, errHash := hashPluginFile(source) + if errHash != nil { + t.Fatalf("hashPluginFile() error = %v", errHash) + } + + // Simulate a racer that already wrote its temp file before winner published target, + // and now attempts publishShadowTempFile against the already-published target. + racerTmp := filepath.Join(dir, shadowPluginTempPrefix+"racer"+testPluginLibraryExtension()) + if errWrite := os.WriteFile(racerTmp, content, 0o644); errWrite != nil { + t.Fatalf("WriteFile(racerTmp) error = %v", errWrite) + } + if errPublish := publishShadowTempFile(racerTmp, first, size, digest); errPublish != nil { + t.Fatalf("publishShadowTempFile() error = %v", errPublish) + } + if _, errStatTmp := os.Stat(racerTmp); !os.IsNotExist(errStatTmp) { + t.Fatalf("Stat(racerTmp) error = %v, want not exist", errStatTmp) + } + infoAfter, errStatAfter := os.Stat(first) + if errStatAfter != nil { + t.Fatalf("Stat(first) after race error = %v", errStatAfter) + } + if !os.SameFile(infoBefore, infoAfter) { + t.Fatal("existing target shadow file was unlinked/replaced by racer") + } +} + +func TestCleanupUnselectedPluginFilesDoesNotConfusePrefixPluginIDs(t *testing.T) { + root := t.TempDir() + platformDir := filepath.Join(root, runtime.GOOS, runtime.GOARCH) + if errMkdir := os.MkdirAll(platformDir, 0o755); errMkdir != nil { + t.Fatalf("MkdirAll() error = %v", errMkdir) + } + ext := testPluginLibraryExtension() + alphaOldPath := filepath.Join(platformDir, "alpha-v1.0.0"+ext) + alphaNewPath := filepath.Join(platformDir, "alpha-v2.0.0"+ext) + alphaBetaPath := filepath.Join(platformDir, "alpha-beta-v1.0.0"+ext) + for path, content := range map[string]string{ + alphaOldPath: "alpha-v1", + alphaNewPath: "alpha-v2", + alphaBetaPath: "alpha-beta-v1", + } { + if errWrite := os.WriteFile(path, []byte(content), 0o644); errWrite != nil { + t.Fatalf("WriteFile(%s) error = %v", path, errWrite) + } + } + alphaOldShadow, errOld := shadowCopyPlugin(pluginFile{ID: "alpha", Path: alphaOldPath, Version: "1.0.0"}) + if errOld != nil { + t.Fatalf("shadowCopyPlugin(alpha old) error = %v", errOld) + } + alphaNewFile := pluginFile{ID: "alpha", Path: alphaNewPath, Version: "2.0.0"} + alphaNewShadow, errNew := shadowCopyPlugin(alphaNewFile) + if errNew != nil { + t.Fatalf("shadowCopyPlugin(alpha new) error = %v", errNew) + } + alphaBetaShadow, errBeta := shadowCopyPlugin(pluginFile{ID: "alpha-beta", Path: alphaBetaPath, Version: "1.0.0"}) + if errBeta != nil { + t.Fatalf("shadowCopyPlugin(alpha-beta) error = %v", errBeta) + } + removedPath := filepath.Join(platformDir, "removed-v1.0.0"+ext) + if errWrite := os.WriteFile(removedPath, []byte("removed-v1"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(removedPath) error = %v", errWrite) + } + removedShadow, errRemoved := shadowCopyPlugin(pluginFile{ID: "removed", Path: removedPath, Version: "1.0.0"}) + if errRemoved != nil { + t.Fatalf("shadowCopyPlugin(removed) error = %v", errRemoved) + } + if errRemove := os.Remove(removedPath); errRemove != nil { + t.Fatalf("Remove(removedPath) error = %v", errRemove) + } + + if errCleanup := cleanupUnselectedPluginFiles(root, []pluginFile{alphaNewFile}); errCleanup != nil { + t.Fatalf("cleanupUnselectedPluginFiles() error = %v", errCleanup) + } + if _, errStat := os.Stat(alphaOldShadow); !os.IsNotExist(errStat) { + t.Fatalf("Stat(alphaOldShadow %s) error = %v, want not exist", alphaOldShadow, errStat) + } + if _, errStat := os.Stat(removedShadow); !os.IsNotExist(errStat) { + t.Fatalf("Stat(removedShadow %s) error = %v, want not exist", removedShadow, errStat) + } + if _, errStat := os.Stat(alphaNewShadow); errStat != nil { + t.Fatalf("Stat(alphaNewShadow %s) error = %v, want kept", alphaNewShadow, errStat) + } + if _, errStat := os.Stat(alphaBetaShadow); errStat != nil { + t.Fatalf("Stat(alphaBetaShadow %s) error = %v, want kept for sibling plugin", alphaBetaShadow, errStat) + } +} + +func TestPruneShadowDirPreservesConcurrentlyInstalledPluginAndLiveGuard(t *testing.T) { + root := t.TempDir() + platformDir := filepath.Join(root, runtime.GOOS, runtime.GOARCH) + if errMkdir := os.MkdirAll(platformDir, 0o755); errMkdir != nil { + t.Fatalf("MkdirAll() error = %v", errMkdir) + } + ext := testPluginLibraryExtension() + v1Path := filepath.Join(platformDir, "alpha-v1.0.0"+ext) + if errWrite := os.WriteFile(v1Path, []byte("alpha-v1"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(v1Path) error = %v", errWrite) + } + v1File := pluginFile{ID: "alpha", Path: v1Path, Version: "1.0.0"} + v1Shadow, errV1 := shadowCopyPlugin(v1File) + if errV1 != nil { + t.Fatalf("shadowCopyPlugin(v1) error = %v", errV1) + } + + // Simulate another process installing v2 and publishing its shadow file after keepNames was computed for v1. + v2Path := filepath.Join(platformDir, "alpha-v2.0.0"+ext) + if errWrite := os.WriteFile(v2Path, []byte("alpha-v2"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(v2Path) error = %v", errWrite) + } + v2File := pluginFile{ID: "alpha", Path: v2Path, Version: "2.0.0"} + v2Shadow, errV2 := shadowCopyPlugin(v2File) + if errV2 != nil { + t.Fatalf("shadowCopyPlugin(v2) error = %v", errV2) + } + + shadowDir := filepath.Dir(v1Shadow) + pruneShadowDirForSourceDir(shadowDir, platformDir, []string{filepath.Base(v1Shadow)}) + if _, errStat := os.Stat(v2Shadow); errStat != nil { + t.Fatalf("Stat(v2Shadow %s) error = %v, want preserved because v2 exists in sourceDir", v2Shadow, errStat) + } + + // Write v2 content to a separate temp source before removing v2Path from platformDir so we can + // also test the interleaving where v2 was removed from platformDir while another process holds a guard. + detachedV2Path := filepath.Join(t.TempDir(), "alpha-v2.0.0"+ext) + if errWrite := os.WriteFile(detachedV2Path, []byte("alpha-v2"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(detachedV2Path) error = %v", errWrite) + } + v2File = pluginFile{ID: "alpha", Path: detachedV2Path, Version: "2.0.0"} + if errRemove := os.Remove(v2Path); errRemove != nil { + t.Fatalf("Remove(v2Path) error = %v", errRemove) + } + _, releaseGuard, errAcquire := acquireShadowPluginLoad(v1File) + if errAcquire != nil { + t.Fatalf("acquireShadowPluginLoad() error = %v", errAcquire) + } + pruneShadowDirForSourceDir(shadowDir, platformDir, []string{filepath.Base(v1Shadow)}) + if _, errStat := os.Stat(v2Shadow); errStat != nil { + releaseGuard() + t.Fatalf("Stat(v2Shadow %s) while load guard active error = %v, want preserved", v2Shadow, errStat) + } + releaseGuard() + + pruneShadowDirForSourceDir(shadowDir, platformDir, []string{filepath.Base(v1Shadow)}) + if _, errStat := os.Stat(v2Shadow); !os.IsNotExist(errStat) { + t.Fatalf("Stat(v2Shadow %s) after guard released error = %v, want not exist", v2Shadow, errStat) + } + + // Verify deterministic interleaving: pruneShadowDirForSourceDir has already read directory entries, + // and before it acquires the coord lock to delete a candidate shadow file, another caller acquires a load guard. + v2ShadowRe, errV2Re := shadowCopyPluginToDir(v2File, shadowDir) + if errV2Re != nil { + t.Fatalf("shadowCopyPluginToDir(v2) error = %v", errV2Re) + } + var releaseInterleavedGuard func() + origHook := pruneShadowBeforeLockHook + t.Cleanup(func() { + pruneShadowBeforeLockHook = origHook + if releaseInterleavedGuard != nil { + releaseInterleavedGuard() + } + }) + pruneShadowBeforeLockHook = func(shadowFileName string) { + if shadowFileName == filepath.Base(v2ShadowRe) && releaseInterleavedGuard == nil { + _, release, errLoad := acquireShadowPluginLoad(v1File) + if errLoad != nil { + t.Errorf("interleaved acquireShadowPluginLoad() error = %v", errLoad) + return + } + releaseInterleavedGuard = release + } + } + pruneShadowDirForSourceDir(shadowDir, platformDir, []string{filepath.Base(v1Shadow)}) + pruneShadowBeforeLockHook = origHook + if _, errStat := os.Stat(v2ShadowRe); errStat != nil { + t.Fatalf("Stat(v2ShadowRe %s) after interleaved guard acquisition error = %v, want preserved", v2ShadowRe, errStat) + } + // Even when keepNames is empty (e.g., all source files were removed or failed hashing), + // an active load guard must still prevent pruning shadow files. + pruneShadowDirForSourceDir(shadowDir, platformDir, nil) + if _, errStat := os.Stat(v2ShadowRe); errStat != nil { + t.Fatalf("Stat(v2ShadowRe %s) with empty keepNames while guard active error = %v, want preserved", v2ShadowRe, errStat) + } + if releaseInterleavedGuard != nil { + releaseInterleavedGuard() + releaseInterleavedGuard = nil + } + + // Verify that if a process crashes/exits abnormally while holding the OS coordination lock, + // the kernel automatically releases the lock so subsequent loads and pruning recover immediately. + cmd := exec.Command(os.Args[0], "-test.run=^TestShadowHelperHoldLockAndExit$") + cmd.Env = append( + os.Environ(), + "GO_WANT_SHADOW_HOLD_LOCK_EXIT_HELPER=1", + "CLIPROXY_SHADOW_HELPER_SOURCE_DIR="+platformDir, + "CLIPROXY_SHADOW_HELPER_SHADOW_DIR="+shadowDir, + ) + if output, errRun := cmd.CombinedOutput(); errRun == nil { + t.Fatalf("helper unexpectedly exited with status 0: %s", string(output)) + } + _, releaseAfterCrash, errAfterCrash := acquireShadowPluginLoad(v1File) + if errAfterCrash != nil { + t.Fatalf("acquireShadowPluginLoad() after crashing lock holder error = %v", errAfterCrash) + } + releaseAfterCrash() + pruneShadowDirForSourceDir(shadowDir, platformDir, []string{filepath.Base(v1Shadow)}) + if _, errStat := os.Stat(v2ShadowRe); !os.IsNotExist(errStat) { + t.Fatalf("Stat(v2ShadowRe %s) after crashing lock holder and prune error = %v, want not exist", v2ShadowRe, errStat) + } +} + +func TestRemoveStaleShadowTempFilesPreservesLiveProcessTempFile(t *testing.T) { + dir := t.TempDir() + ext := testPluginLibraryExtension() + currentPID := os.Getpid() + deadPID := startExitedProcessPID(t) + if deadPID == currentPID { + t.Skip("exited helper process reused the test pid") + } + liveOther, liveOtherPID := startWaitingProcess(t) + if liveOtherPID == currentPID || liveOtherPID == deadPID { + t.Skip("live helper process pid collided with another test pid") + } + defer stopWaitingProcess(t, liveOther) + + selfTemp := filepath.Join(dir, fmt.Sprintf("%s%s%d-alpha-self%s", shadowPluginTempPrefix, shadowPluginProcessDirPrefix, currentPID, ext)) + liveTemp := filepath.Join(dir, fmt.Sprintf("%s%s%d-alpha-live%s", shadowPluginTempPrefix, shadowPluginProcessDirPrefix, liveOtherPID, ext)) + deadTemp := filepath.Join(dir, fmt.Sprintf("%s%s%d-alpha-dead%s", shadowPluginTempPrefix, shadowPluginProcessDirPrefix, deadPID, ext)) + legacyTemp := filepath.Join(dir, shadowPluginTempPrefix+"alpha-legacy"+ext) + for _, path := range []string{selfTemp, liveTemp, deadTemp, legacyTemp} { + if errWrite := os.WriteFile(path, []byte("temp"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(%s) error = %v", path, errWrite) + } + } + + removeStaleShadowTempFiles(dir) + + for _, deadPath := range []string{deadTemp, legacyTemp} { + if _, errStat := os.Stat(deadPath); !os.IsNotExist(errStat) { + t.Fatalf("Stat(%s) error = %v, want not exist", deadPath, errStat) + } + } + for _, livePath := range []string{selfTemp, liveTemp} { + if _, errStat := os.Stat(livePath); errStat != nil { + t.Fatalf("Stat(%s) error = %v, want kept for live process", livePath, errStat) + } + } +} + +func TestCleanupUnselectedPluginFilesPrunesStaleShadowPlugins(t *testing.T) { + root := t.TempDir() + platformDir := filepath.Join(root, runtime.GOOS, runtime.GOARCH) + if errMkdir := os.MkdirAll(platformDir, 0o755); errMkdir != nil { + t.Fatalf("MkdirAll() error = %v", errMkdir) + } + ext := testPluginLibraryExtension() + oldPath := filepath.Join(platformDir, "alpha-v1.0.0"+ext) + newPath := filepath.Join(platformDir, "alpha-v2.0.0"+ext) + if errWrite := os.WriteFile(oldPath, []byte("alpha-v1"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(oldPath) error = %v", errWrite) + } + if errWrite := os.WriteFile(newPath, []byte("alpha-v2"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(newPath) error = %v", errWrite) + } + oldShadow, errOldShadow := shadowCopyPlugin(pluginFile{ID: "alpha", Path: oldPath, Version: "1.0.0"}) + if errOldShadow != nil { + t.Fatalf("shadowCopyPlugin(old) error = %v", errOldShadow) + } + newFile := pluginFile{ID: "alpha", Path: newPath, Version: "2.0.0"} + newShadow, errNewShadow := shadowCopyPlugin(newFile) + if errNewShadow != nil { + t.Fatalf("shadowCopyPlugin(new) error = %v", errNewShadow) + } + tempShadow := filepath.Join(filepath.Dir(newShadow), shadowPluginTempPrefix+"alpha-stale"+ext) + if errWrite := os.WriteFile(tempShadow, []byte("partial"), 0o644); errWrite != nil { + t.Fatalf("WriteFile(tempShadow) error = %v", errWrite) + } + + if errCleanup := cleanupUnselectedPluginFiles(root, []pluginFile{newFile}); errCleanup != nil { + t.Fatalf("cleanupUnselectedPluginFiles() error = %v", errCleanup) + } + + if _, errStat := os.Stat(oldShadow); !os.IsNotExist(errStat) { + t.Fatalf("Stat(oldShadow %s) error = %v, want not exist", oldShadow, errStat) + } + if _, errStat := os.Stat(tempShadow); !os.IsNotExist(errStat) { + t.Fatalf("Stat(tempShadow %s) error = %v, want not exist", tempShadow, errStat) + } + if _, errStat := os.Stat(newShadow); errStat != nil { + t.Fatalf("Stat(newShadow %s) error = %v, want kept", newShadow, errStat) + } +} + func TestShadowCopyPluginReusesContentAddressedShadow(t *testing.T) { dir := t.TempDir() source := filepath.Join(t.TempDir(), "alpha"+testPluginLibraryExtension()) @@ -218,6 +711,44 @@ func TestShadowHelperProcessExits(t *testing.T) { } } +func runChildShadowCopyPlugin(t *testing.T, file pluginFile) string { + t.Helper() + cmd := exec.Command(os.Args[0], "-test.run=^TestShadowHelperCopyPlugin$") + cmd.Env = append( + os.Environ(), + "CLIPROXY_SHADOW_COPY_HELPER=1", + "CLIPROXY_SHADOW_COPY_ID="+file.ID, + "CLIPROXY_SHADOW_COPY_PATH="+file.Path, + "CLIPROXY_SHADOW_COPY_VERSION="+file.Version, + ) + out, errRun := cmd.CombinedOutput() + if errRun != nil { + t.Fatalf("child shadowCopyPlugin helper error = %v\n%s", errRun, out) + } + for _, line := range strings.Split(string(out), "\n") { + if strings.HasPrefix(line, "SHADOW_PATH=") { + return strings.TrimSpace(strings.TrimPrefix(line, "SHADOW_PATH=")) + } + } + t.Fatalf("child shadowCopyPlugin output missing SHADOW_PATH: %s", out) + return "" +} + +func TestShadowHelperCopyPlugin(t *testing.T) { + if os.Getenv("CLIPROXY_SHADOW_COPY_HELPER") != "1" { + t.Skip("shadow copy helper process entry point") + } + got, errShadow := shadowCopyPlugin(pluginFile{ + ID: os.Getenv("CLIPROXY_SHADOW_COPY_ID"), + Path: os.Getenv("CLIPROXY_SHADOW_COPY_PATH"), + Version: os.Getenv("CLIPROXY_SHADOW_COPY_VERSION"), + }) + if errShadow != nil { + t.Fatalf("shadowCopyPlugin() error = %v", errShadow) + } + fmt.Printf("SHADOW_PATH=%s\n", got) +} + func startWaitingProcess(t *testing.T) (*exec.Cmd, int) { t.Helper() cmd := exec.Command(os.Args[0], "-test.run=^TestShadowHelperProcessWaits$") @@ -248,6 +779,22 @@ func stopWaitingProcess(t *testing.T, cmd *exec.Cmd) { _ = cmd.Wait() } +func TestShadowHelperHoldLockAndExit(t *testing.T) { + if os.Getenv("GO_WANT_SHADOW_HOLD_LOCK_EXIT_HELPER") != "1" { + t.Skip("crash lock helper entry point") + } + sourceDir := os.Getenv("CLIPROXY_SHADOW_HELPER_SOURCE_DIR") + shadowDir := os.Getenv("CLIPROXY_SHADOW_HELPER_SHADOW_DIR") + if sourceDir == "" || shadowDir == "" { + os.Exit(2) + } + // Acquire an exclusive coordination lock and exit abnormally without unlocking. + if _, errLock := acquireShadowCoordLock(sourceDir, shadowDir, true, false); errLock != nil { + os.Exit(3) + } + os.Exit(1) +} + func TestShadowHelperProcessWaits(t *testing.T) { if os.Getenv("CLIPROXY_SHADOW_HELPER_WAIT") != "1" { t.Skip("waiting helper process entry point") diff --git a/internal/pluginhost/shadow_unix.go b/internal/pluginhost/shadow_unix.go index 39e626e7464..c793a696547 100644 --- a/internal/pluginhost/shadow_unix.go +++ b/internal/pluginhost/shadow_unix.go @@ -8,6 +8,34 @@ import ( "syscall" ) +func lockShadowCoordFile(f *os.File, exclusive bool, nonBlocking bool) error { + how := syscall.LOCK_SH + if exclusive { + how = syscall.LOCK_EX + } + if nonBlocking { + how |= syscall.LOCK_NB + } + for { + errFlock := syscall.Flock(int(f.Fd()), how) + if !errors.Is(errFlock, syscall.EINTR) { + return errFlock + } + } +} + +func unlockShadowCoordFile(f *os.File) { + _ = syscall.Flock(int(f.Fd()), syscall.LOCK_UN) +} + +func publishShadowPlugin(tmpPath, targetPath string) error { + if errLink := os.Link(tmpPath, targetPath); errLink != nil { + return errLink + } + _ = os.Remove(tmpPath) + return nil +} + func removeShadowPlugin(path string) { if path == "" { return diff --git a/internal/pluginhost/shadow_windows.go b/internal/pluginhost/shadow_windows.go index 5efe995f874..80958cfe3cb 100644 --- a/internal/pluginhost/shadow_windows.go +++ b/internal/pluginhost/shadow_windows.go @@ -11,6 +11,35 @@ import ( "golang.org/x/sys/windows" ) +func lockShadowCoordFile(f *os.File, exclusive bool, nonBlocking bool) error { + var flags uint32 + if exclusive { + flags |= windows.LOCKFILE_EXCLUSIVE_LOCK + } + if nonBlocking { + flags |= windows.LOCKFILE_FAIL_IMMEDIATELY + } + var overlapped windows.Overlapped + return windows.LockFileEx(windows.Handle(f.Fd()), flags, 0, 1, 0, &overlapped) +} + +func unlockShadowCoordFile(f *os.File) { + var overlapped windows.Overlapped + _ = windows.UnlockFileEx(windows.Handle(f.Fd()), 0, 1, 0, &overlapped) +} + +func publishShadowPlugin(tmpPath, targetPath string) error { + fromPtr, errFrom := windows.UTF16PtrFromString(tmpPath) + if errFrom != nil { + return errFrom + } + toPtr, errTo := windows.UTF16PtrFromString(targetPath) + if errTo != nil { + return errTo + } + return windows.MoveFileEx(fromPtr, toPtr, 0) +} + func removeShadowPlugin(path string) { if path == "" { return diff --git a/internal/pluginhost/stream_bridge.go b/internal/pluginhost/stream_bridge.go index 6cd5b1639b1..943cc628633 100644 --- a/internal/pluginhost/stream_bridge.go +++ b/internal/pluginhost/stream_bridge.go @@ -79,6 +79,7 @@ func rpcStreamError(message string, detail *pluginabi.Error) error { message: text, statusCode: detail.HTTPStatus, retryAfterSeconds: detail.RetryAfterSeconds, + credentialScoped: detail.CredentialScoped, } } if message != "" { diff --git a/internal/redisqueue/home_usage_test.go b/internal/redisqueue/home_usage_test.go new file mode 100644 index 00000000000..5821fee3cd9 --- /dev/null +++ b/internal/redisqueue/home_usage_test.go @@ -0,0 +1,41 @@ +package redisqueue + +import ( + "testing" + + "github.com/MoeForge/moego/ojson" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" +) + +func TestUsageQueueCarriesHomeSelection(t *testing.T) { + withEnabledQueue(t, func() { + for _, selection := range []*coreusage.HomeSelection{ + nil, + {Model: "route-model"}, + {Additional: true, Model: "route-model"}, + {Reused: true, Additional: true, Model: "route-model"}, + {Reused: true, Model: "route-model", CredentialPolicy: "codex-alpha-search-v1"}, + } { + (&usageQueuePlugin{}).HandleUsage(t.Context(), coreusage.Record{ + RequestID: "response-unique-id", Provider: "codex", Model: "upstream-model", HomeSelection: selection, + }) + items := PopOldest(1) + if len(items) != 1 { + t.Fatalf("queued=%d", len(items)) + } + got := ojson.GetBytes(items[0], "home_selection") + if selection == nil { + if got.Exists() { + t.Fatal("legacy usage acquired a Home marker") + } + continue + } + if got.Get("reused").Bool() != selection.Reused || got.Get("additional").Bool() != selection.Additional || got.Get("model").String() != selection.Model || got.Get("credential_policy").String() != selection.CredentialPolicy { + t.Fatalf("lost selection metadata: %s", got.Raw) + } + if ojson.GetBytes(items[0], "execution_id").String() != "response-unique-id" { + t.Fatal("usage lost the per-response replay identity") + } + } + }) +} diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index a479a4778b2..acbdf4f8a85 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -177,6 +177,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec AccessProvider: accessProvider, RequestID: requestID, ExecutionID: executionID, + HomeSelection: record.HomeSelection, TraceID: traceID, SessionID: sessionID, ParentSessionID: parentSessionID, @@ -222,6 +223,7 @@ type queuedUsageDetail struct { AccessProvider string `json:"access_provider,omitempty"` RequestID string `json:"request_id"` ExecutionID string `json:"execution_id,omitempty"` + HomeSelection *coreusage.HomeSelection `json:"home_selection,omitempty"` TraceID string `json:"trace_id,omitempty"` SessionID string `json:"session_id,omitempty"` ParentSessionID string `json:"parent_session_id,omitempty"` diff --git a/internal/registry/models/models.json b/internal/registry/models/models.json index d8e576221be..eec7606848e 100644 --- a/internal/registry/models/models.json +++ b/internal/registry/models/models.json @@ -278,6 +278,35 @@ "text" ] }, + { + "id": "claude-fable-5-2", + "object": "model", + "created": 1794250800, + "owned_by": "anthropic", + "type": "claude", + "display_name": "Claude Fable 5.2", + "description": "Claude Fable 5.2 extends Claude Fable 5 at the same input and output prices, with cache reads at a quarter of the cost, and brings stronger long-running agentic coding, multistep research, and document, spreadsheet, and slide work.", + "context_length": 1000000, + "max_completion_tokens": 128000, + "thinking": { + "zero_allowed": false, + "dynamic_allowed": true, + "levels": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "supportedInputModalities": [ + "text", + "image" + ], + "supportedOutputModalities": [ + "text" + ] + }, { "id": "claude-opus-5-5", "native_capabilities": { @@ -1035,6 +1064,36 @@ "native_capabilities": { "web_search": true } + }, + { + "id": "gemini-omni-1.2-flash", + "object": "model", + "owned_by": "google", + "type": "gemini", + "display_name": "Gemini Omni 1.2 Flash", + "name": "models/gemini-omni-1.2-flash", + "version": "1.2", + "description": "Gemini Omni 1.2 Flash", + "context_length": 131072, + "max_completion_tokens": 65536, + "thinking": { + "min": 128, + "max": 32768, + "dynamic_allowed": true, + "levels": [ + "low", + "high" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "video" + ], + "supportedOutputModalities": [ + "text", + "video" + ] } ], "vertex": [ @@ -1756,6 +1815,35 @@ "native_capabilities": { "web_search": true } + }, + { + "id": "gemini-omni-1.2-flash", + "object": "model", + "owned_by": "google", + "type": "gemini", + "display_name": "Gemini Omni 1.2 Flash", + "name": "gemini-omni-1.2-flash", + "description": "Gemini Omni 1.2 Flash", + "context_length": 131072, + "max_completion_tokens": 260640, + "thinking": { + "min": 128, + "max": 32768, + "dynamic_allowed": true, + "levels": [ + "low", + "high" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "video" + ], + "supportedOutputModalities": [ + "text", + "video" + ] } ], "gemini-cli": [ @@ -2590,6 +2678,36 @@ "native_capabilities": { "web_search": true } + }, + { + "id": "gemini-omni-1.2-flash", + "object": "model", + "owned_by": "google", + "type": "gemini", + "display_name": "Gemini Omni 1.2 Flash", + "name": "models/gemini-omni-1.2-flash", + "version": "1.2", + "description": "Gemini Omni 1.2 Flash", + "context_length": 131072, + "max_completion_tokens": 65536, + "thinking": { + "min": 128, + "max": 32768, + "dynamic_allowed": true, + "levels": [ + "low", + "high" + ] + }, + "supportedInputModalities": [ + "text", + "image", + "video" + ], + "supportedOutputModalities": [ + "text", + "video" + ] } ], "codex-free": [ @@ -5160,4 +5278,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/runtime/executor/aistudio_executor.go b/internal/runtime/executor/aistudio_executor.go index a827a4e911d..7acfd20b6aa 100644 --- a/internal/runtime/executor/aistudio_executor.go +++ b/internal/runtime/executor/aistudio_executor.go @@ -186,8 +186,14 @@ func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, helps.AppendAPIResponseChunk(ctx, e.cfg, wsResp.Body) } if wsResp.Status < 200 || wsResp.Status >= 300 { + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, body.replayScope, wsResp.Status, wsResp.Body); errClear != nil { + logAntigravityReasoningReplayDegraded(body.replayScope, "invalidate", errClear) + } return resp, statusErr{code: wsResp.Status, msg: string(wsResp.Body)} } + if body.action != "countTokens" { + cacheAntigravityReasoningReplayFromResponse(ctx, body.replayScope, body.replayRequestPayload, wsResp.Body) + } reporter.ObserveResponseModel(wsResp.Body) responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) var param any @@ -284,20 +290,23 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth reporter.StartResponseTTFT() metadataLogged = true } - var body bytes.Buffer + var errBody bytes.Buffer if len(firstEvent.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, firstEvent.Payload) - body.Write(firstEvent.Payload) + errBody.Write(firstEvent.Payload) } if firstEvent.Type == wsrelay.MessageTypeStreamEnd { - return nil, statusErr{code: firstEvent.Status, msg: body.String()} + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, body.replayScope, firstEvent.Status, errBody.Bytes()); errClear != nil { + logAntigravityReasoningReplayDegraded(body.replayScope, "invalidate", errClear) + } + return nil, statusErr{code: firstEvent.Status, msg: errBody.String()} } for event := range wsStream { if event.Err != nil { helps.RecordAPIResponseError(ctx, e.cfg, event.Err) - if body.Len() == 0 { - body.WriteString(event.Err.Error()) + if errBody.Len() == 0 { + errBody.WriteString(event.Err.Error()) } break } @@ -309,14 +318,18 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth if len(event.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) - body.Write(event.Payload) + errBody.Write(event.Payload) } if event.Type == wsrelay.MessageTypeStreamEnd { break } } - return nil, statusErr{code: firstEvent.Status, msg: body.String()} + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, body.replayScope, firstEvent.Status, errBody.Bytes()); errClear != nil { + logAntigravityReasoningReplayDegraded(body.replayScope, "invalidate", errClear) + } + return nil, statusErr{code: firstEvent.Status, msg: errBody.String()} } + replayAccumulator := newAntigravityReasoningReplayAccumulator(body.replayScope, body.replayRequestPayload) out := make(chan cliproxyexecutor.StreamChunk, 1) go func(first wsrelay.StreamEvent) { defer close(out) @@ -336,12 +349,22 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth _ = helps.StopTranslatedClaudeStream(ctx, &streamErr, reporter, &streamUsage, out) }() var param any + var replayCommitted bool + commitReplay := func() { + if !replayCommitted && replayAccumulator != nil && replayAccumulator.terminal && streamErr == nil && helps.ApplyPatchTranslationError(param) == nil { + replayAccumulator.Commit(ctx) + replayCommitted = true + } + } helps.InitializeApplyPatchStream(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, ¶m) metadataLogged := false var relayFrames helps.RelaySSEBuffer processRelayFrames := func(payloads [][]byte) bool { for _, payload := range payloads { for _, frame := range splitRelayFrames(payload) { + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(frame) + } reporter.ObserveResponseModel(frame) streamUsage.ObserveGeminiStream(frame) filtered := helps.FilterSSEUsageMetadata(frame) @@ -350,6 +373,9 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && antigravityStreamChunksHaveTerminalCommitEvent(lines) { + commitReplay() + } for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: @@ -382,6 +408,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + commitReplay() for _, line := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(line)}: @@ -435,6 +462,9 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) } if event.Status < http.StatusOK || event.Status >= http.StatusMultipleChoices { + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, body.replayScope, event.Status, event.Payload); errClear != nil { + logAntigravityReasoningReplayDegraded(body.replayScope, "invalidate", errClear) + } errResponse := statusErr{code: event.Status, msg: string(event.Payload)} helps.RecordAPIResponseError(ctx, e.cfg, errResponse) reporter.PublishFailure(ctx, errResponse) @@ -444,6 +474,12 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth } return false } + if replayAccumulator != nil { + replayAccumulator.observeResponsePayload(event.Payload) + if !replayAccumulator.terminal && (ojson.GetBytesNoCopy(event.Payload, "response.candidates.0").Exists() || ojson.GetBytesNoCopy(event.Payload, "candidates.0").Exists()) { + replayAccumulator.terminal = true + } + } reporter.ObserveResponseModel(event.Payload) streamUsage.Observe(helps.ParseGeminiUsage(event.Payload), true) lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, event.Payload, ¶m, claudeInputTokens) @@ -451,6 +487,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + commitReplay() for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: @@ -561,9 +598,11 @@ func (e *AIStudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) } type translatedPayload struct { - payload []byte - action string - toFormat sdktranslator.Format + payload []byte + action string + toFormat sdktranslator.Format + replayScope antigravityReasoningReplayScope + replayRequestPayload []byte } func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) ([]byte, translatedPayload, error) { @@ -608,6 +647,19 @@ func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexe } normalizeAIStudioThinkingLevelDoc(doc) internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if antigravityUsesReasoningReplayCache(baseModel) { + var changedReplay bool + var errReplay error + replayScope, replayRequestPayload, changedReplay, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return nil, translatedPayload{}, errReplay + } + if changedReplay { + internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") + } + } if action == "countTokens" { _ = doc.Delete("generationConfig") _ = doc.Delete("tools") @@ -615,7 +667,13 @@ func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexe } helps.ApplyPayloadConfigWithRequestDoc(e.cfg, baseModel, to.String(), from.String(), "", doc, originalTranslated, requestedModel, requestPath, opts.Headers) payload = doc.Bytes() - return payload, translatedPayload{payload: payload, action: action, toFormat: to}, nil + return payload, translatedPayload{ + payload: payload, + action: action, + toFormat: to, + replayScope: replayScope, + replayRequestPayload: replayRequestPayload, + }, nil } // normalizeAIStudioThinkingLevel normalizes thinking levels to Google's canonical uppercase enum values. diff --git a/internal/runtime/executor/antigravity_executor_request.go b/internal/runtime/executor/antigravity_executor_request.go index 2a7df38a20f..0516d09fc27 100644 --- a/internal/runtime/executor/antigravity_executor_request.go +++ b/internal/runtime/executor/antigravity_executor_request.go @@ -782,14 +782,22 @@ func generateSessionID() string { } func generateStableSessionID(payload []byte) string { - return generateStableSessionIDFromContents(util.GetGJSONBytesNoCopy(payload, "request.contents")) + contents := util.GetGJSONBytesNoCopy(payload, "request.contents") + if !contents.IsArray() { + contents = util.GetGJSONBytesNoCopy(payload, "contents") + } + return generateStableSessionIDFromContents(contents) } func generateStableSessionIDDoc(doc *ojson.Doc) string { if doc == nil { return generateSessionID() } - return generateStableSessionIDFromContents(doc.Get("request.contents")) + contents := doc.Get("request.contents") + if !contents.IsArray() { + contents = doc.Get("contents") + } + return generateStableSessionIDFromContents(contents) } func generateStableSessionIDFromContents(contents ojson.Result) string { diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index 9be7ccf93ce..31961e4385b 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -358,14 +358,8 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya } } } - if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && responseFormat == sdktranslator.FormatOpenAIResponse { - for _, chunk := range chunks { - for _, line := range bytes.Split(chunk, []byte("\n")) { - if ojson.GetBytesNoCopy(helps.JSONPayload(line), "type").String() == "response.completed" { - commitReplay() - } - } - } + if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && antigravityStreamChunksHaveTerminalCommitEvent(chunks) { + commitReplay() } for i := range chunks { select { diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index a076b45537a..bde7323e42b 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -66,7 +66,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut e.obfuscateSensitiveWordsDoc(doc) helps.PrepareAntigravityClaudeStructuredOutputDoc(baseModel, doc) sanitizeAntigravityGeminiRequestSignaturesDoc(baseModel, doc) - if _, _, errReplay := prepareAntigravityGeminiReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc); errReplay != nil { + if _, errReplay := prepareAntigravityGeminiCountTokensReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc); errReplay != nil { return cliproxyexecutor.Response{}, errReplay } ensureAntigravityGeminiLeadingUserContentDoc(baseModel, doc) diff --git a/internal/runtime/executor/antigravity_reasoning_replay.go b/internal/runtime/executor/antigravity_reasoning_replay.go index 24f5c55da88..de83bb24415 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay.go +++ b/internal/runtime/executor/antigravity_reasoning_replay.go @@ -19,9 +19,11 @@ import ( homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" internalsignature "github.com/router-for-me/CLIProxyAPI/v8/internal/signature" + translatorcommon "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/common" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/session" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" log "github.com/sirupsen/logrus" ) @@ -39,14 +41,22 @@ func antigravityReplayLogKey(value string) string { // antigravityCountClaudeToolProvenanceIDs reports how many reserved // Claude-facing provenance IDs are still present in a Gemini-shaped payload. func antigravityCountClaudeToolProvenanceIDs(payload []byte) int { - return antigravityCountClaudeToolProvenanceIDsResult(util.GetGJSONBytesNoCopy(payload, "request.contents")) + contents := util.GetGJSONBytesNoCopy(payload, "request.contents") + if !contents.IsArray() { + contents = util.GetGJSONBytesNoCopy(payload, "contents") + } + return antigravityCountClaudeToolProvenanceIDsResult(contents) } func antigravityCountClaudeToolProvenanceIDsDoc(doc *ojson.Doc) int { if doc == nil { return 0 } - return antigravityCountClaudeToolProvenanceIDsResult(doc.Get("request.contents")) + contents := doc.Get("request.contents") + if !contents.IsArray() { + contents = doc.Get("contents") + } + return antigravityCountClaudeToolProvenanceIDsResult(contents) } func antigravityCountClaudeToolProvenanceIDsResult(contents ojson.Result) int { @@ -88,6 +98,27 @@ func (s antigravityReasoningReplayScope) valid() bool { return strings.TrimSpace(s.modelName) != "" && strings.TrimSpace(s.sessionKey) != "" } +func antigravityFormatCarriesNativeThoughtSignatures(format sdktranslator.Format) bool { + switch strings.ToLower(strings.TrimSpace(format.String())) { + case sdktranslator.FormatGemini.String(), + sdktranslator.FormatInteractions.String(): + return true + default: + return false + } +} + +func antigravityRequestIsSymmetricForReasoningReplay(opts cliproxyexecutor.Options) bool { + source := strings.TrimSpace(opts.SourceFormat.String()) + if source == "" { + return false + } + if !antigravityFormatCarriesNativeThoughtSignatures(opts.SourceFormat) { + return false + } + return antigravityFormatCarriesNativeThoughtSignatures(cliproxyexecutor.ResponseFormatOrSource(opts)) +} + func antigravityReasoningReplayScopeFromPayload(modelName string, payload []byte) antigravityReasoningReplayScope { sessionID := antigravityReplaySessionIDFromPayload(payload) if sessionID == "" { @@ -127,6 +158,9 @@ func antigravityReasoningReplayScopeFromDoc(modelName string, doc *ojson.Doc) an } func antigravityReasoningReplayScopeFromRequest(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, payload []byte) antigravityReasoningReplayScope { + if antigravityRequestIsSymmetricForReasoningReplay(opts) && !antigravityPayloadHasClaudeToolProvenanceID(payload) { + return antigravityReasoningReplayScope{} + } // Prefer an explicit downstream session over a provider sessionId synthesized // from request text. This keeps identical prompts in separate client sessions // from sharing an opaque Gemini reasoning chain. @@ -157,6 +191,9 @@ func antigravityReasoningReplayScopeFromRequest(ctx context.Context, modelName s } func antigravityReasoningReplayScopeFromRequestDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) antigravityReasoningReplayScope { + if antigravityRequestIsSymmetricForReasoningReplay(opts) && !antigravityDocHasClaudeToolProvenanceID(doc) { + return antigravityReasoningReplayScope{} + } scope := antigravityReasoningReplayScope{modelName: modelName} if sessionKey := antigravityReasoningReplayClientSessionKey(ctx, req, opts); sessionKey != "" { scope.sessionKey = sessionKey @@ -351,12 +388,117 @@ func prepareAntigravityGeminiReasoningReplayPayload(ctx context.Context, modelNa return doc.Bytes(), scope, nil } +// antigravityStreamChunksHaveTerminalCommitEvent reports whether translated SSE +// chunks carry a client-visible terminal completion/truncation event before +// which accumulated reasoning replay items must be committed to KV. +func antigravityStreamChunksHaveTerminalCommitEvent(chunks [][]byte) bool { + for _, chunk := range chunks { + for _, line := range bytes.Split(chunk, []byte("\n")) { + payload := helps.JSONPayload(line) + if len(payload) == 0 { + continue + } + switch ojson.GetBytesNoCopy(payload, "type").String() { + case "response.completed", "response.incomplete", "response.done", "message_stop": + return true + } + if ojson.GetBytesNoCopy(payload, "event_type").String() == "interaction.completed" { + return true + } + for _, path := range []string{ + "choices.0.finish_reason", + "response.choices.0.finish_reason", + "candidates.0.finishReason", + "response.candidates.0.finishReason", + } { + if finishReason := ojson.GetBytesNoCopy(payload, path); finishReason.Exists() && strings.TrimSpace(finishReason.String()) != "" { + return true + } + } + } + } + return false +} + +func isGeminiCountTokensAction(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) bool { + if req.Metadata != nil { + if action, _ := req.Metadata["action"].(string); strings.EqualFold(strings.TrimSpace(action), "countTokens") { + return true + } + } + if opts.Metadata != nil { + if action, _ := opts.Metadata["action"].(string); strings.EqualFold(strings.TrimSpace(action), "countTokens") { + return true + } + } + return false +} + +func prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) (antigravityReasoningReplayScope, []byte, bool, error) { + return prepareGeminiTopLevelReasoningReplayPayloadDocWithMode(ctx, modelName, req, opts, doc, !isGeminiCountTokensAction(req, opts)) +} + +func prepareGeminiTopLevelCountTokensReasoningReplayPayloadDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) (bool, error) { + _, _, changed, err := prepareGeminiTopLevelReasoningReplayPayloadDocWithMode(ctx, modelName, req, opts, doc, false) + return changed, err +} + +func prepareGeminiTopLevelReasoningReplayPayloadDocWithMode(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc, reserveOnMiss bool) (antigravityReasoningReplayScope, []byte, bool, error) { + if doc == nil || !antigravityUsesReasoningReplayCache(modelName) { + return antigravityReasoningReplayScope{}, nil, false, nil + } + if doc.Get("request.contents").Exists() { + scope, changed, err := prepareAntigravityGeminiReasoningReplayPayloadDocWithOptions(ctx, modelName, req, opts, doc, true, reserveOnMiss) + var replayRequestPayload []byte + if reserveOnMiss && scope.valid() { + replayRequestPayload, _ = antigravityReplayRequestEnvelopeFromDoc(doc) + } + return scope, replayRequestPayload, changed, err + } + if !doc.Get("contents").Exists() { + return antigravityReasoningReplayScope{}, nil, false, nil + } + wrappedBytes, ok := antigravityReplayRequestEnvelopeFromDoc(doc) + if !ok { + return antigravityReasoningReplayScope{}, nil, false, nil + } + wrappedDoc := ojson.NewDoc(wrappedBytes) + for _, key := range []string{"sessionId", "session_id"} { + if val := doc.Get(key); val.Exists() { + _ = wrappedDoc.SetRawBytes(key, []byte(val.Raw)) + } + } + scope, changed, err := prepareAntigravityGeminiReasoningReplayPayloadDocWithOptions(ctx, modelName, req, opts, wrappedDoc, false, reserveOnMiss) + if err != nil { + return scope, nil, false, err + } + if changed { + if updatedContents := wrappedDoc.Get("request.contents"); updatedContents.Exists() { + _ = doc.SetRawBytes("contents", []byte(updatedContents.Raw)) + } + } + var replayRequestPayload []byte + if reserveOnMiss && scope.valid() { + replayRequestPayload, _ = antigravityReplayRequestEnvelopeFromDoc(doc) + } + return scope, replayRequestPayload, changed, nil +} + func prepareAntigravityGeminiReasoningReplayPayloadDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) (antigravityReasoningReplayScope, bool, error) { + return prepareAntigravityGeminiReasoningReplayPayloadDocWithOptions(ctx, modelName, req, opts, doc, true, !isGeminiCountTokensAction(req, opts)) +} + +func prepareAntigravityGeminiCountTokensReasoningReplayPayloadDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) (bool, error) { + _, changed, err := prepareAntigravityGeminiReasoningReplayPayloadDocWithOptions(ctx, modelName, req, opts, doc, true, false) + return changed, err +} + +func prepareAntigravityGeminiReasoningReplayPayloadDocWithOptions(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc, normalizeFunctionResponseRoles, reserveOnMiss bool) (antigravityReasoningReplayScope, bool, error) { if doc == nil || !antigravityUsesReasoningReplayCache(modelName) { return antigravityReasoningReplayScope{}, false, nil } originalContentsRaw := doc.Get("request.contents").Raw - scope, replayApplied, errReplay := applyAntigravityReasoningReplayCacheDoc(ctx, modelName, req, opts, doc) + scope, replayApplied, errReplay := applyAntigravityReasoningReplayCacheDocWithMode(ctx, modelName, req, opts, doc, reserveOnMiss) changed := replayApplied if errReplay != nil { logAntigravityReasoningReplayDegraded(scope, "read", errReplay) @@ -366,7 +508,7 @@ func prepareAntigravityGeminiReasoningReplayPayloadDoc(ctx context.Context, mode changed = false replayApplied = false } - if normalizeAntigravityGeminiFunctionResponseRolesDoc(doc) { + if normalizeFunctionResponseRoles && normalizeAntigravityGeminiFunctionResponseRolesDoc(doc) { changed = true } if antigravityDocHasClaudeToolProvenanceID(doc) { @@ -381,7 +523,7 @@ func prepareAntigravityGeminiReasoningReplayPayloadDoc(ctx context.Context, mode changed = true } } - if normalizeAntigravityGeminiFunctionResponseRolesDoc(doc) { + if normalizeFunctionResponseRoles && normalizeAntigravityGeminiFunctionResponseRolesDoc(doc) { changed = true } } @@ -400,7 +542,7 @@ func prepareAntigravityGeminiReasoningReplayPayloadDoc(ctx context.Context, mode origWrapped = []byte(`{"request":{"contents":` + originalContentsRaw + `}}`) } originalPairingValid := internalsignature.ValidateGeminiFunctionCallPairing(origWrapped) == nil - if replayApplied && originalPairingValid && scope.valid() { + if replayApplied && originalPairingValid && scope.valid() && reserveOnMiss { if _, errDelete := internalcache.DeleteAntigravityReasoningReplayItemsIfUnchanged(ctx, scope.modelName, scope.sessionKey, scope.cacheSnapshot); errDelete != nil { logAntigravityReasoningReplayDegraded(scope, "invalidate", errDelete) } @@ -443,6 +585,10 @@ func applyAntigravityReasoningReplayCache(ctx context.Context, modelName string, } func applyAntigravityReasoningReplayCacheDoc(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc) (antigravityReasoningReplayScope, bool, error) { + return applyAntigravityReasoningReplayCacheDocWithMode(ctx, modelName, req, opts, doc, true) +} + +func applyAntigravityReasoningReplayCacheDocWithMode(ctx context.Context, modelName string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, doc *ojson.Doc, reserveOnMiss bool) (antigravityReasoningReplayScope, bool, error) { if doc == nil { return antigravityReasoningReplayScope{}, false, nil } @@ -450,7 +596,17 @@ func applyAntigravityReasoningReplayCacheDoc(ctx context.Context, modelName stri if !scope.valid() { return scope, false, nil } - items, snapshot, ok, err := internalcache.GetAntigravityReasoningReplayItemsWithSnapshotRequired(ctx, scope.modelName, scope.sessionKey) + var ( + items [][]byte + snapshot internalcache.AntigravityReasoningReplaySnapshot + ok bool + err error + ) + if reserveOnMiss { + items, snapshot, ok, err = internalcache.GetAntigravityReasoningReplayItemsWithSnapshotRequired(ctx, scope.modelName, scope.sessionKey) + } else { + items, snapshot, ok, err = internalcache.PeekAntigravityReasoningReplayItemsWithSnapshotRequired(ctx, scope.modelName, scope.sessionKey) + } scope.cacheSnapshot = snapshot reservedBefore := antigravityCountClaudeToolProvenanceIDsDoc(doc) if err != nil || !ok || len(items) == 0 { @@ -490,9 +646,14 @@ func antigravityReplayRequestEnvelopeFromDoc(doc *ojson.Doc) ([]byte, bool) { if requestNode := doc.Get("request"); requestNode.IsObject() { return []byte(`{"request":` + requestNode.Raw + `}`), true } + prefix := "request." contents := doc.Get("request.contents") if !contents.Exists() { - return nil, false + prefix = "" + contents = doc.Get("contents") + if !contents.Exists() { + return nil, false + } } var buf bytes.Buffer buf.WriteString(`{"request":{`) @@ -501,10 +662,10 @@ func antigravityReplayRequestEnvelopeFromDoc(doc *ojson.Doc) ([]byte, bool) { key string path string }{ - {key: "systemInstruction", path: "request.systemInstruction"}, - {key: "tools", path: "request.tools"}, - {key: "toolConfig", path: "request.toolConfig"}, - {key: "contents", path: "request.contents"}, + {key: "systemInstruction", path: prefix + "systemInstruction"}, + {key: "tools", path: prefix + "tools"}, + {key: "toolConfig", path: prefix + "toolConfig"}, + {key: "contents", path: prefix + "contents"}, } { val := doc.Get(field.path) if !val.Exists() { @@ -1972,14 +2133,22 @@ func antigravityFunctionCallMatchesReplayItem(functionCall, itemResult ojson.Res } func antigravityPayloadHasClaudeToolProvenanceID(payload []byte) bool { - return antigravityContentsHasClaudeToolProvenanceID(util.GetGJSONBytesNoCopy(payload, "request.contents")) + contents := util.GetGJSONBytesNoCopy(payload, "request.contents") + if !contents.IsArray() { + contents = util.GetGJSONBytesNoCopy(payload, "contents") + } + return antigravityContentsHasClaudeToolProvenanceID(contents) } func antigravityDocHasClaudeToolProvenanceID(doc *ojson.Doc) bool { if doc == nil { return false } - return antigravityContentsHasClaudeToolProvenanceID(doc.Get("request.contents")) + contents := doc.Get("request.contents") + if !contents.IsArray() { + contents = doc.Get("contents") + } + return antigravityContentsHasClaudeToolProvenanceID(contents) } func antigravityContentsHasClaudeToolProvenanceID(contents ojson.Result) bool { @@ -2528,6 +2697,12 @@ type antigravityPendingThoughtSignature struct { targetKind string } +type antigravityStreamingReplayToolCall struct { + call *translatorcommon.VertexStreamingToolCall + partIndex int + signature string +} + type antigravityReasoningReplayAccumulator struct { scope antigravityReasoningReplayScope responseContextHash string @@ -2544,6 +2719,8 @@ type antigravityReasoningReplayAccumulator struct { thoughtPartIndex int lastResponseKind string pendingSignatures []antigravityPendingThoughtSignature + activeStreamingFC *antigravityStreamingReplayToolCall + hasIncompleteFC bool itemBytes int overflow bool terminal bool @@ -2674,9 +2851,20 @@ func (a *antigravityReasoningReplayAccumulator) appendItem(item []byte) { a.itemBytes += len(item) } -func (a *antigravityReasoningReplayAccumulator) attachDetachedSignatureToLastFunctionCall(signature string) { +func (a *antigravityReasoningReplayAccumulator) attachDetachedSignatureToLastFunctionCall(signature string) bool { if a == nil || signature == "" { - return + return false + } + if a.activeStreamingFC != nil { + if a.activeStreamingFC.signature != "" { + return false + } + if a.itemBytes+len(signature) > internalcache.AntigravityReasoningReplayCacheMaxBytesPerEntry { + a.overflow = true + return false + } + a.activeStreamingFC.signature = signature + return true } for itemIndex := len(a.items) - 1; itemIndex >= 0; itemIndex-- { item := util.ParseGJSONBytesNoCopy(a.items[itemIndex]) @@ -2684,21 +2872,22 @@ func (a *antigravityReasoningReplayAccumulator) attachDetachedSignatureToLastFun continue } if strings.TrimSpace(item.Get("thoughtSignature").String()) != "" { - return + return false } updated, errSet := ojson.SetBytes(a.items[itemIndex], "thoughtSignature", signature) if errSet != nil { - return + return false } delta := len(updated) - len(a.items[itemIndex]) if a.itemBytes+delta > internalcache.AntigravityReasoningReplayCacheMaxBytesPerEntry { a.overflow = true - return + return false } a.items[itemIndex] = updated a.itemBytes += delta - return + return true } + return false } func (a *antigravityReasoningReplayAccumulator) ObserveSSELine(line []byte) { @@ -2759,56 +2948,17 @@ func (a *antigravityReasoningReplayAccumulator) observeResponsePayload(payload [ signature = "" } if fc := part.Get("functionCall"); fc.Exists() { - if a.lastResponseKind == "text" || a.lastResponseKind == "thought" { - a.flushPendingThoughtSignaturesForKind(a.lastResponseKind) - } - if signature != "" { - remainingPending := a.pendingSignatures[:0] - for _, pending := range a.pendingSignatures { - if pending.targetKind != "" { - remainingPending = append(remainingPending, pending) - } - } - a.pendingSignatures = remainingPending - } - if signature == "" { - for pendingIndex := len(a.pendingSignatures) - 1; pendingIndex >= 0; pendingIndex-- { - if a.pendingSignatures[pendingIndex].targetKind == "" { - signature = a.pendingSignatures[pendingIndex].signature - a.pendingSignatures = append(a.pendingSignatures[:pendingIndex], a.pendingSignatures[pendingIndex+1:]...) - break - } - } - } - if signature == "" { - signature = a.claimPrecedingTurnThoughtSignature() - } - keys := antigravityReplayToolCallKeysFromPart(fc) - for _, key := range keys { - dedupeKey := key + "\x00" + signature - if signature == "" { - dedupeKey = fmt.Sprintf("%s\x00part:%d", key, pi) - } - if a.seenFC[dedupeKey] { - return true - } - a.seenFC[dedupeKey] = true - } - occurrenceKey := antigravityFunctionCallKey(fc.Get("name").String(), fc.Get("args").Raw, "") - occurrence := a.functionCallOccurrences[occurrenceKey] - if occurrenceKey != "" { - a.functionCallOccurrences[occurrenceKey] = occurrence + 1 - } - item := buildAntigravityFunctionCallPartItem(a.contentIndex, pi, occurrence, fc, signature) - if len(item) > 0 { - a.appendItem(antigravitySetReplayItemContextHashValue(item, a.responseContextHash)) - if signature != "" { - a.seenSignatures[signature] = true - } - } - a.lastResponseKind = "function_call" + a.observeFunctionCallPart(pi, fc, signature) return true } + if a.activeStreamingFC != nil { + text := part.Get("text") + hasSemanticText := text.Exists() && text.String() != "" + if hasSemanticText { + a.hasIncompleteFC = true + a.activeStreamingFC = nil + } + } targetKind := "" if part.Get("thought").Bool() { @@ -2818,8 +2968,7 @@ func (a *antigravityReasoningReplayAccumulator) observeResponsePayload(payload [ hasSemanticText := text.Exists() && text.String() != "" signatureOnly := signature != "" && !hasSemanticText if signatureOnly && a.lastResponseKind == "function_call" { - if !a.seenSignatures[signature] { - a.attachDetachedSignatureToLastFunctionCall(signature) + if !a.seenSignatures[signature] && a.attachDetachedSignatureToLastFunctionCall(signature) { a.seenSignatures[signature] = true } return true @@ -3007,15 +3156,163 @@ func (a *antigravityReasoningReplayAccumulator) appendPendingThoughtSignatures() a.pendingSignatures = nil } +func (a *antigravityReasoningReplayAccumulator) resolveFunctionCallSignature(signature string) string { + if a.lastResponseKind == "text" || a.lastResponseKind == "thought" { + a.flushPendingThoughtSignaturesForKind(a.lastResponseKind) + } + if signature != "" { + remainingPending := a.pendingSignatures[:0] + for _, pending := range a.pendingSignatures { + if pending.targetKind != "" { + remainingPending = append(remainingPending, pending) + } + } + a.pendingSignatures = remainingPending + } + if signature == "" { + for pendingIndex := len(a.pendingSignatures) - 1; pendingIndex >= 0; pendingIndex-- { + if a.pendingSignatures[pendingIndex].targetKind == "" { + signature = a.pendingSignatures[pendingIndex].signature + a.pendingSignatures = append(a.pendingSignatures[:pendingIndex], a.pendingSignatures[pendingIndex+1:]...) + break + } + } + } + if signature == "" { + signature = a.claimPrecedingTurnThoughtSignature() + } + return signature +} + +func (a *antigravityReasoningReplayAccumulator) recordCompletedFunctionCall(pi int, fc ojson.Result, signature string) { + keys := antigravityReplayToolCallKeysFromPart(fc) + for _, key := range keys { + dedupeKey := key + "\x00" + signature + if signature == "" { + dedupeKey = fmt.Sprintf("%s\x00part:%d", key, pi) + } + if a.seenFC[dedupeKey] { + return + } + a.seenFC[dedupeKey] = true + } + occurrenceKey := antigravityFunctionCallKey(fc.Get("name").String(), fc.Get("args").Raw, "") + occurrence := a.functionCallOccurrences[occurrenceKey] + if occurrenceKey != "" { + a.functionCallOccurrences[occurrenceKey] = occurrence + 1 + } + item := buildAntigravityFunctionCallPartItem(a.contentIndex, pi, occurrence, fc, signature) + if len(item) > 0 { + a.appendItem(antigravitySetReplayItemContextHashValue(item, a.responseContextHash)) + if signature != "" { + a.seenSignatures[signature] = true + } + } + a.lastResponseKind = "function_call" +} + +func (a *antigravityReasoningReplayAccumulator) flushCompletedStreamingFunctionCall(active *antigravityStreamingReplayToolCall) { + if a == nil || active == nil || active.call == nil || active.call.Assembler == nil { + return + } + if active.call.Assembler.Err() != nil { + a.hasIncompleteFC = true + return + } + argsRaw := strings.TrimSpace(active.call.Assembler.FullArguments()) + if argsRaw == "" { + argsRaw = "{}" + } + if !ojson.Valid(argsRaw) { + a.hasIncompleteFC = true + return + } + fcDoc := ojson.NewDocString(`{"name":""}`) + _ = fcDoc.Set("name", active.call.RawName) + if active.call.UpstreamID != "" { + _ = fcDoc.Set("id", active.call.UpstreamID) + } + _ = fcDoc.SetRaw("args", argsRaw) + a.recordCompletedFunctionCall(active.partIndex, ojson.ParseBytes(fcDoc.Bytes()), active.signature) +} + +func (a *antigravityReasoningReplayAccumulator) observeFunctionCallPart(pi int, fc ojson.Result, signature string) { + if a == nil { + return + } + if active := a.activeStreamingFC; active != nil && active.call != nil && active.call.CanContinueWith(fc) { + if signature != "" && active.signature == "" { + active.signature = signature + } + if upstreamID := strings.TrimSpace(fc.Get("id").String()); upstreamID != "" && active.call.UpstreamID == "" { + active.call.UpstreamID = strings.Clone(upstreamID) + } + _, done := active.call.Assembler.PushFunctionCall(fc) + if active.call.Assembler.Err() != nil { + a.hasIncompleteFC = true + a.activeStreamingFC = nil + return + } + if done { + a.activeStreamingFC = nil + a.flushCompletedStreamingFunctionCall(active) + } + a.lastResponseKind = "function_call" + return + } + if translatorcommon.IsVertexStreamingFunctionCall(fc) { + if a.activeStreamingFC != nil { + a.hasIncompleteFC = true + a.activeStreamingFC = nil + } + rawName := strings.Clone(strings.TrimSpace(fc.Get("name").String())) + if rawName == "" { + a.hasIncompleteFC = true + return + } + upstreamID := strings.Clone(strings.TrimSpace(fc.Get("id").String())) + resolvedSig := a.resolveFunctionCallSignature(signature) + call := translatorcommon.NewVertexStreamingToolCall(rawName, rawName, upstreamID, upstreamID, 0) + _, done := call.Assembler.PushFunctionCall(fc) + if call.Assembler.Err() != nil { + a.hasIncompleteFC = true + return + } + active := &antigravityStreamingReplayToolCall{ + call: call, + partIndex: pi, + signature: resolvedSig, + } + if done { + a.flushCompletedStreamingFunctionCall(active) + } else { + a.activeStreamingFC = active + } + a.lastResponseKind = "function_call" + return + } + if a.activeStreamingFC != nil { + a.hasIncompleteFC = true + a.activeStreamingFC = nil + } + rawName := strings.TrimSpace(fc.Get("name").String()) + if rawName == "" && !fc.Get("args").Exists() { + return + } + resolvedSig := a.resolveFunctionCallSignature(signature) + a.recordCompletedFunctionCall(pi, fc, resolvedSig) +} + func (a *antigravityReasoningReplayAccumulator) Commit(ctx context.Context) { if a == nil || !a.scope.valid() { return } log.Debugf("antigravity replay: accumulator commit terminal=%t overflow=%t items=%d (session=%s)", a.terminal, a.overflow, len(a.items), antigravityReplayLogKey(a.scope.sessionKey)) - if !a.terminal { - // No terminal finishReason means the stream never completed, so this turn - // contributes nothing to the ledger and its tool IDs become unresolvable. + if !a.terminal || a.activeStreamingFC != nil || a.hasIncompleteFC { + // No terminal finishReason (or an unclosed/corrupted streaming tool call) + // means the stream never completed cleanly, so this turn contributes + // nothing to the ledger and its tool IDs become unresolvable. return } if a.overflow { diff --git a/internal/runtime/executor/antigravity_reasoning_replay_test.go b/internal/runtime/executor/antigravity_reasoning_replay_test.go index fc3b7545870..16036897074 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay_test.go +++ b/internal/runtime/executor/antigravity_reasoning_replay_test.go @@ -12,6 +12,7 @@ import ( internalcache "github.com/router-for-me/CLIProxyAPI/v8/internal/cache" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" internalsignature "github.com/router-for-me/CLIProxyAPI/v8/internal/signature" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -1955,6 +1956,155 @@ func TestPrepareAntigravityGeminiReasoningReplayStillRejectsBrokenPairing(t *tes } } +func TestAntigravityReasoningReplaySkipsSymmetricGeminiAndInteractionsButKeepsAsymmetricAlias(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache) + + const model = "gemini-3-flash" + payload := []byte(`{"sessionId":"sess-symmetry","request":{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Beijing"}]}]}}`) + responseSSE := []byte(`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":""},{"thoughtSignature":"EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX","functionCall":{"name":"get_weather","args":{"city":"Beijing"},"id":"call_1"}}]},"finishReason":"STOP"}]}}`) + + for _, format := range []sdktranslator.Format{sdktranslator.FormatGemini, sdktranslator.FormatInteractions} { + t.Run("symmetric_"+format.String(), func(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + req := cliproxyexecutor.Request{Model: model, Payload: payload} + opts := cliproxyexecutor.Options{ + SourceFormat: format, + ResponseFormat: format, + } + _, scope, errPrepare := prepareAntigravityGeminiReasoningReplayPayload(context.Background(), model, req, opts, payload) + if errPrepare != nil { + t.Fatalf("prepare error: %v", errPrepare) + } + if scope.valid() { + t.Fatalf("symmetric format %q must not allocate reasoning replay KV scope, got %+v", format, scope) + } + if acc := newAntigravityReasoningReplayAccumulator(scope, payload); acc != nil { + t.Fatalf("symmetric format %q must not create replay accumulator", format) + } + }) + } + + t.Run("asymmetric_claude_haiku_alias_to_gemini_3_flash", func(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + // Downstream client sends Claude /v1/messages with model="claude-haiku-4-5-20251001", + // which model-alias resolves to baseModel="gemini-3-flash" inside AntigravityExecutor. + req := cliproxyexecutor.Request{Model: model, Payload: payload} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + ResponseFormat: sdktranslator.FormatClaude, + Metadata: map[string]any{ + cliproxyexecutor.RequestedModelMetadataKey: "claude-haiku-4-5-20251001", + }, + } + _, scope, errPrepare := prepareAntigravityGeminiReasoningReplayPayload(context.Background(), model, req, opts, payload) + if errPrepare != nil { + t.Fatalf("prepare error: %v", errPrepare) + } + if !scope.valid() { + t.Fatal("asymmetric Claude -> Gemini alias request must allocate reasoning replay KV scope") + } + acc := newAntigravityReasoningReplayAccumulator(scope, payload) + if acc == nil { + t.Fatal("expected non-nil accumulator for asymmetric Claude -> Gemini request") + } + acc.ObserveSSELine(responseSSE) + acc.Commit(context.Background()) + + items, ok := internalcache.GetAntigravityReasoningReplayItems(scope.modelName, scope.sessionKey) + if !ok || len(items) == 0 { + t.Fatalf("expected asymmetric replay items to be persisted in KV, ok=%v len=%d", ok, len(items)) + } + }) + + t.Run("top_level_gemini_doc_replays_asymmetric_and_skips_symmetric", func(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + topLevelTurn1 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Beijing"}]}]}`) + topLevelTurn2 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Beijing"}]},{"role":"model","parts":[{"thoughtSignature":"skip_thought_signature_validator","functionCall":{"name":"get_weather","args":{"city":"Beijing"},"id":"call_1"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_1","response":{"result":"sunny"}}}]}]}`) + topLevelResponse := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":""},{"thoughtSignature":"EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX","functionCall":{"name":"get_weather","args":{"city":"Beijing"},"id":"call_1"}}]},"finishReason":"STOP"}]}`) + + symDoc := ojson.NewDoc(topLevelTurn1) + symScope, _, _, errSym := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: topLevelTurn1}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini, ResponseFormat: sdktranslator.FormatGemini}, + symDoc, + ) + if errSym != nil { + t.Fatalf("symmetric top-level prepare error: %v", errSym) + } + if symScope.valid() { + t.Fatalf("symmetric top-level Gemini must not allocate replay scope, got %+v", symScope) + } + + headers := make(http.Header) + headers.Set("Session-Id", "sess-non-ag-gemini") + asymOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Headers: headers, + } + doc1 := ojson.NewDoc(topLevelTurn1) + scope1, replayReq1, _, err1 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: topLevelTurn1}, + asymOpts, + doc1, + ) + if err1 != nil || !scope1.valid() { + t.Fatalf("asymmetric top-level prepare turn 1: scope=%+v err=%v", scope1, err1) + } + cacheAntigravityReasoningReplayFromResponse(context.Background(), scope1, replayReq1, topLevelResponse) + + doc2 := ojson.NewDoc(topLevelTurn2) + scope2, replayReq2, changed2, err2 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: topLevelTurn2}, + asymOpts, + doc2, + ) + if err2 != nil || !changed2 { + t.Fatalf("asymmetric top-level prepare turn 2: changed=%v err=%v", changed2, err2) + } + if gotSig := doc2.Get("contents.1.parts.0.thoughtSignature").String(); gotSig != "EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX" { + t.Fatalf("replayed signature = %q, want native thoughtSignature", gotSig) + } + if gotRole := doc2.Get("contents.2.role").String(); gotRole != "user" { + t.Fatalf("non-Antigravity functionResponse role = %q, want user", gotRole) + } + + // Turn 2 simulates Vertex stripping wire-only functionCall.id/functionResponse.id + // after prepareGeminiTopLevelReasoningReplayPayloadDoc; the captured replayRequestPayload + // from prepareGeminiTopLevelReasoningReplayPayloadDoc must still keep canonical call_id + // context so Turn 3 matches both call_1 and call_2 signatures. + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc2, asymOpts.SourceFormat.String()) + topLevelResponse2 := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thought":true,"text":""},{"thoughtSignature":"EsMTCsATARFNMg/XNVix5lDpkKaHR7Xg","functionCall":{"name":"get_weather","args":{"city":"Shanghai"},"id":"call_2"}}]},"finishReason":"STOP"}]}`) + cacheAntigravityReasoningReplayFromResponse(context.Background(), scope2, replayReq2, topLevelResponse2) + + topLevelTurn3 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Beijing"}]},{"role":"model","parts":[{"thoughtSignature":"skip_thought_signature_validator","functionCall":{"name":"get_weather","args":{"city":"Beijing"},"id":"call_1"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_1","response":{"result":"sunny"}}}]},{"role":"model","parts":[{"thoughtSignature":"skip_thought_signature_validator","functionCall":{"name":"get_weather","args":{"city":"Shanghai"},"id":"call_2"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_2","response":{"result":"cloudy"}}}]}]}`) + doc3 := ojson.NewDoc(topLevelTurn3) + _, _, changed3, err3 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: topLevelTurn3}, + asymOpts, + doc3, + ) + if err3 != nil || !changed3 { + t.Fatalf("asymmetric top-level prepare turn 3: changed=%v err=%v", changed3, err3) + } + if gotSig1 := doc3.Get("contents.1.parts.0.thoughtSignature").String(); gotSig1 != "EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX" { + t.Fatalf("turn 3 call_1 signature = %q, want preserved signature", gotSig1) + } + if gotSig2 := doc3.Get("contents.3.parts.0.thoughtSignature").String(); gotSig2 != "EsMTCsATARFNMg/XNVix5lDpkKaHR7Xg" { + t.Fatalf("turn 3 call_2 signature = %q, want second signature", gotSig2) + } + }) +} + func TestPrepareAntigravityGeminiReasoningReplayDegradesWhenReplayBreaksPairing(t *testing.T) { internalcache.ClearAntigravityReasoningReplayCache() t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache) @@ -1982,3 +2132,174 @@ func TestPrepareAntigravityGeminiReasoningReplayDegradesWhenReplayBreaksPairing( t.Fatalf("out = %s, want original payload %s", out, payload) } } + +func TestAntigravityStreamChunksHaveTerminalCommitEvent(t *testing.T) { + cases := []struct { + name string + chunks [][]byte + want bool + }{ + { + name: "openai responses completed", + chunks: [][]byte{[]byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n")}, + want: true, + }, + { + name: "openai responses incomplete", + chunks: [][]byte{[]byte("event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n")}, + want: true, + }, + { + name: "openai chat finish_reason stop", + chunks: [][]byte{[]byte("data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")}, + want: true, + }, + { + name: "openai chat null finish_reason", + chunks: [][]byte{[]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}\n\n")}, + want: false, + }, + { + name: "claude message_stop", + chunks: [][]byte{[]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")}, + want: true, + }, + { + name: "bare done frame only", + chunks: [][]byte{[]byte("data: [DONE]\n\n")}, + want: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := antigravityStreamChunksHaveTerminalCommitEvent(tc.chunks); got != tc.want { + t.Fatalf("antigravityStreamChunksHaveTerminalCommitEvent() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestAntigravityReasoningReplayAccumulatorAssemblesVertexStreamingPartialArgs(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache) + + const model = "gemini-3.8-flash" + asymOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Headers: http.Header{"X-Session-Id": []string{"sess-vertex-partial-args"}}, + } + + turn1 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Weather in New York?"}]}]}`) + doc1 := ojson.NewDoc(turn1) + scope1, replayReq1, _, err1 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: turn1}, + asymOpts, + doc1, + ) + if err1 != nil || !scope1.valid() { + t.Fatalf("prepare turn 1: scope=%+v err=%v", scope1, err1) + } + + acc := newAntigravityReasoningReplayAccumulator(scope1, replayReq1) + acc.ObserveSSELine([]byte(`data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call_ny","name":"get_weather","partialArgs":[{"jsonPath":"$.city","stringValue":"New ","willContinue":true}],"willContinue":true}}]}}]}`)) + acc.ObserveSSELine([]byte(`data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"text":"","thoughtSignature":"EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX"}]}}]}`)) + acc.ObserveSSELine([]byte(`data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"partialArgs":[{"jsonPath":"$.city","stringValue":"York"}],"willContinue":false}}]},"finishReason":"STOP"}]}`)) + acc.Commit(context.Background()) + + turn2 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Weather in New York?"}]},{"role":"model","parts":[{"thoughtSignature":"skip_thought_signature_validator","functionCall":{"name":"get_weather","args":{"city":"New York"},"id":"call_ny"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_ny","response":{"result":"sunny"}}}]}]}`) + doc2 := ojson.NewDoc(turn2) + _, _, changed2, err2 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: turn2}, + asymOpts, + doc2, + ) + if err2 != nil || !changed2 { + t.Fatalf("prepare turn 2 after partialArgs stream: changed=%v err=%v", changed2, err2) + } + if gotSig := doc2.Get("contents.1.parts.0.thoughtSignature").String(); gotSig != "EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX" { + t.Fatalf("replayed signature = %q, want native signature from partialArgs stream", gotSig) + } +} + +func TestNonVertexSanitizesReplayedVertexBareTinkSignature(t *testing.T) { + // Verify that a Vertex bare-Tink thoughtSignature replayed from shared KV is + // downgraded to skip_thought_signature_validator when post-replay non-Vertex + // sanitization runs (as in GeminiExecutor and AIStudioExecutor). + const vertexBareTinkSig = "AY89a18f0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + internalcache.ClearAntigravityReasoningReplayCache() + const model = "gemini-3.8-flash" + asymOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Headers: http.Header{"X-Session-Id": []string{"sess-vertex-to-gemini"}}, + } + turn1 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather"}]}]}`) + doc1 := ojson.NewDoc(turn1) + scope1, replayReq1, _, err1 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: turn1}, + asymOpts, + doc1, + ) + if err1 != nil || !scope1.valid() { + t.Fatalf("prepare turn 1: scope=%+v err=%v", scope1, err1) + } + vertexResp := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"` + vertexBareTinkSig + `","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_tink"}}]},"finishReason":"STOP"}]}`) + cacheAntigravityReasoningReplayFromResponse(context.Background(), scope1, replayReq1, vertexResp) + + turn2 := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather"}]},{"role":"model","parts":[{"thoughtSignature":"skip_thought_signature_validator","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_tink"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_tink","response":{"result":"clear"}}}]}]}`) + nonVertexDoc := ojson.NewDoc(turn2) + _, _, changed2, err2 := prepareGeminiTopLevelReasoningReplayPayloadDoc( + context.Background(), + model, + cliproxyexecutor.Request{Model: model, Payload: turn2}, + asymOpts, + nonVertexDoc, + ) + if err2 != nil || !changed2 { + t.Fatalf("prepare turn 2: changed=%v err=%v", changed2, err2) + } + internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(nonVertexDoc, "contents") + if got := nonVertexDoc.Get("contents.1.parts.0.thoughtSignature").String(); got != "skip_thought_signature_validator" { + t.Fatalf("non-Vertex post-replay sanitized signature = %q, want skip_thought_signature_validator", got) + } +} + +func TestCountTokensReasoningReplayPeeksWithoutReservingTombstone(t *testing.T) { + internalcache.ClearAntigravityReasoningReplayCache() + t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache) + + const model = "gemini-3.8-flash" + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Headers: http.Header{"X-Session-Id": []string{"sess-count-tokens-no-tombstone"}}, + } + payload := []byte(`{"contents":[{"role":"user","parts":[{"text":"Count tokens only"}]}]}`) + doc := ojson.NewDoc(payload) + req := cliproxyexecutor.Request{ + Model: model, + Payload: payload, + Metadata: map[string]any{"action": "countTokens"}, + } + scope, replayReq, changed, err := prepareGeminiTopLevelReasoningReplayPayloadDoc(context.Background(), model, req, opts, doc) + if err != nil || changed || len(replayReq) != 0 { + t.Fatalf("countTokens prepare = scope:%+v len(replayReq):%d changed:%v err:%v", scope, len(replayReq), changed, err) + } + // A subsequent real generation turn must still be able to cache best-effort + // without being blocked by an unclosed countTokens reservation. + item := []byte(`{"type":"thought_signature","contentIndex":0,"partIndex":0,"targetOccurrence":0,"thoughtSignature":"EmcKZQFpFH0THU+YNRY5uVwt/0bo2hFo1RmXsMhlX9FCg2n6GssBe5sMV37KB6OrgeKKEP9fxN261QX"}`) + if !internalcache.CacheAntigravityReasoningReplayItems(model, scope.sessionKey, [][]byte{item}) { + t.Fatal("failed to cache replay item after countTokens peek") + } + got, ok := internalcache.GetAntigravityReasoningReplayItems(model, scope.sessionKey) + if !ok || len(got) != 1 { + t.Fatalf("GetAntigravityReasoningReplayItems after countTokens peek = %v, ok=%v", got, ok) + } +} diff --git a/internal/runtime/executor/claude_cloaked_cache_repro_test.go b/internal/runtime/executor/claude_cloaked_cache_repro_test.go index d5bd0502e6c..6bdd3cad8af 100644 --- a/internal/runtime/executor/claude_cloaked_cache_repro_test.go +++ b/internal/runtime/executor/claude_cloaked_cache_repro_test.go @@ -726,3 +726,232 @@ func TestClaudeCloakedColdStartToolContinuationUsesDeterministicPromptID(t *test t.Fatalf("cold start tool continuation prompt IDs differ between failure and retry: %q != %q", p1, p2) } } + +func TestClaudeCloakedTopLevelCacheControl_Issue7366(t *testing.T) { + newCloakedOAuthEnv := func(t *testing.T) (context.Context, *cliproxyauth.Auth, *ClaudeExecutor, *[][]byte) { + t.Helper() + var capturedBodies [][]byte + call := 0 + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + body, errRead := io.ReadAll(req.Body) + if errRead != nil { + t.Fatal(errRead) + } + capturedBodies = append(capturedBodies, body) + call++ + response := fmt.Sprintf(`{"id":"msg_%d","type":"message","model":"claude-opus-5-5","role":"assistant","content":[{"type":"text","text":"OK"}],"usage":{"input_tokens":10,"output_tokens":2}}`, call) + header := make(http.Header) + header.Set("Content-Type", "application/json") + header.Set("request-id", fmt.Sprintf("req_issue7366_%d", call)) + return &http.Response{ + StatusCode: http.StatusOK, + Header: header, + Body: io.NopCloser(strings.NewReader(response)), + Request: req, + }, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + auth := &cliproxyauth.Auth{ + ID: "test-issue7366-oauth", + Attributes: map[string]string{"api_key": "sk-ant-oat-test-issue7366-key"}, + Metadata: map[string]any{ + "account_uuid": "11111111-2222-4333-8444-555555555555", + claudeauth.ClaudeDeviceIDsMetadataKey: []string{ + "0000000000000000000000000000000000000000000000000000000000000001", + }, + }, + } + return ctx, auth, NewClaudeExecutor(&config.Config{}), &capturedBodies + } + + t.Run("Error1_SingleTurnBareTopLevelCacheControlNoTTLMismatch", func(t *testing.T) { + ctx, auth, exec, captured := newCloakedOAuthEnv(t) + payload := []byte(`{"model":"claude-opus-5-5","max_tokens":64,"cache_control":{"type":"ephemeral"},"system":[{"type":"text","text":"You are a helpful assistant.","cache_control":{"type":"ephemeral"}}],"messages":[{"role":"user","content":[{"type":"text","text":"What time is it?"}]}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Headers: http.Header{"Anthropic-Version": []string{"2023-06-01"}}, + } + if _, err := exec.Execute(ctx, auth, cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: payload}, opts); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if len(*captured) != 1 { + t.Fatalf("captured %d bodies, want 1", len(*captured)) + } + seenBody := (*captured)[0] + if ojson.GetBytes(seenBody, "cache_control").Exists() { + t.Fatalf("expected top-level cache_control to be materialized and removed on cloaked request, got body: %s", seenBody) + } + if got := ojson.GetBytes(seenBody, "system.1.cache_control.ttl").String(); got != "1h" { + t.Fatalf("system.1.cache_control.ttl = %q, want 1h: %s", got, seenBody) + } + if got := ojson.GetBytes(seenBody, "messages.0.content.1.cache_control.ttl").String(); got != "1h" { + t.Fatalf("messages.0.content.1.cache_control.ttl = %q, want 1h: %s", got, seenBody) + } + if got := ojson.GetBytes(seenBody, "messages.1.content.0.cache_control.ttl").String(); got != "1h" { + t.Fatalf("messages.1.content.0.cache_control.ttl = %q, want 1h: %s", got, seenBody) + } + }) + + t.Run("Error2_ToolResultTurnWith1hTopLevelCacheControlStaysWithin4Breakpoints", func(t *testing.T) { + ctx, auth, exec, captured := newCloakedOAuthEnv(t) + payload := []byte(`{"model":"claude-opus-5-5","max_tokens":64,"cache_control":{"type":"ephemeral","ttl":"1h"},"system":[{"type":"text","text":"You are a helpful assistant.","cache_control":{"type":"ephemeral","ttl":"1h"}}],"tools":[{"name":"get_time","description":"Returns the current time.","input_schema":{"type":"object","properties":{}}}],"messages":[{"role":"user","content":[{"type":"text","text":"What time is it?"}]},{"role":"assistant","content":[{"type":"tool_use","id":"toolu_01","name":"get_time","input":{}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01","content":"12:00"}]}]}`) + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Headers: http.Header{"Anthropic-Version": []string{"2023-06-01"}}, + } + if _, err := exec.Execute(ctx, auth, cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: payload}, opts); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if len(*captured) != 1 { + t.Fatalf("captured %d bodies, want 1", len(*captured)) + } + seenBody := (*captured)[0] + if ojson.GetBytes(seenBody, "cache_control").Exists() { + t.Fatalf("expected top-level cache_control to be materialized onto last cacheable block and removed, got body: %s", seenBody) + } + if total := countCacheControls(seenBody); total != 4 { + t.Fatalf("total cache_control count = %d, want 4: %s", total, seenBody) + } + for _, path := range []string{ + "system.1.cache_control.ttl", + "messages.0.content.1.cache_control.ttl", + "messages.1.content.0.cache_control.ttl", + "messages.3.content.0.cache_control.ttl", + } { + if got := ojson.GetBytes(seenBody, path).String(); got != "1h" { + t.Fatalf("%s = %q, want 1h: %s", path, got, seenBody) + } + } + }) + + t.Run("StreamPath_MaterializesTopLevelCacheControlAndStaysWithin4Breakpoints", func(t *testing.T) { + var seenBody []byte + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + var errRead error + seenBody, errRead = io.ReadAll(req.Body) + if errRead != nil { + t.Fatal(errRead) + } + header := make(http.Header) + header.Set("Content-Type", "text/event-stream") + sse := "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_s1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-opus-5-5\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" + return &http.Response{ + StatusCode: http.StatusOK, + Header: header, + Body: io.NopCloser(strings.NewReader(sse)), + Request: req, + }, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + auth := &cliproxyauth.Auth{ + ID: "test-issue7366-stream-oauth", + Attributes: map[string]string{"api_key": "sk-ant-oat-test-issue7366-stream-key"}, + Metadata: map[string]any{ + "account_uuid": "11111111-2222-4333-8444-555555555555", + claudeauth.ClaudeDeviceIDsMetadataKey: []string{ + "0000000000000000000000000000000000000000000000000000000000000001", + }, + }, + } + exec := NewClaudeExecutor(&config.Config{}) + payload := []byte(`{"model":"claude-opus-5-5","max_tokens":64,"stream":true,"cache_control":{"type":"ephemeral","ttl":"1h"},"system":[{"type":"text","text":"You are a helpful assistant.","cache_control":{"type":"ephemeral","ttl":"1h"}}],"tools":[{"name":"get_time","description":"Returns the current time.","input_schema":{"type":"object","properties":{}}}],"messages":[{"role":"user","content":[{"type":"text","text":"What time is it?"}]},{"role":"assistant","content":[{"type":"tool_use","id":"toolu_01","name":"get_time","input":{}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01","content":"12:00"}]}]}`) + res, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: payload}, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Headers: http.Header{"Anthropic-Version": []string{"2023-06-01"}}, + }) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + if res != nil && res.Chunks != nil { + for range res.Chunks { + } + } + if ojson.GetBytes(seenBody, "cache_control").Exists() { + t.Fatalf("expected top-level cache_control to be removed in stream path, got: %s", seenBody) + } + if total := countCacheControls(seenBody); total != 4 { + t.Fatalf("total cache_control count = %d, want 4: %s", total, seenBody) + } + }) + + t.Run("DirectAPIKeyPassthroughPreservesTopLevelCacheControlWithoutSyntheticBlocks", func(t *testing.T) { + var seenBody []byte + transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + seenBody, _ = io.ReadAll(req.Body) + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"msg_direct","type":"message","model":"claude-opus-5-5","role":"assistant","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`)), + Request: req, + }, nil + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{ + "api_key": "sk-ant-api03-direct-top-cc-test", + "base_url": "https://api.anthropic.com", + }, + } + payload := []byte(`{"model":"claude-opus-5-5","max_tokens":32,"cache_control":{"type":"ephemeral"},"system":[{"type":"text","text":"caller system"}],"messages":[{"role":"user","content":"hello"}]}`) + if _, err := NewClaudeExecutor(&config.Config{}).Execute(ctx, auth, cliproxyexecutor.Request{Model: "claude-opus-5-5", Payload: payload}, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + Headers: http.Header{"User-Agent": {"custom-sdk/1.0"}}, + }); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if got := ojson.GetBytes(seenBody, "cache_control.type").String(); got != "ephemeral" { + t.Fatalf("direct API key top-level cache_control.type = %q, want ephemeral: %s", got, seenBody) + } + if ojson.GetBytes(seenBody, "system.0.cache_control").Exists() || ojson.GetBytes(seenBody, "messages.0.content.0.cache_control").Exists() { + t.Fatalf("direct API key request with top-level cache_control must not gain synthetic block breakpoints: %s", seenBody) + } + }) + + t.Run("HelpersRespectAndMaterializeTopLevelCacheControl", func(t *testing.T) { + withScope := []byte(`{"cache_control":{"type":"ephemeral","scope":"global"},"system":"Sys","messages":[{"role":"user","content":"Hello"}]}`) + ensured := ensureCacheControl(withScope) + if ojson.GetBytes(ensured, "cache_control").Exists() { + t.Fatalf("ensureCacheControl should remove materialized top-level cache_control: %s", ensured) + } + if got := ojson.GetBytes(ensured, "messages.0.content.0.cache_control.scope").String(); got != "global" { + t.Fatalf("messages.0.content.0.cache_control.scope = %q, want global: %s", got, ensured) + } + + // When top-level cache_control remains (e.g. direct call to enforceCacheControlLimit / upgrade / normalize / strip), + // helpers still account for it. + fiveMarkers := []byte(`{ + "cache_control":{"type":"ephemeral"}, + "system":[ + {"type":"text","text":"s0","cache_control":{"type":"ephemeral"}}, + {"type":"text","text":"s1","cache_control":{"type":"ephemeral"}} + ], + "messages":[{ + "role":"user", + "content":[ + {"type":"text","text":"m0","cache_control":{"type":"ephemeral"}}, + {"type":"text","text":"m1","cache_control":{"type":"ephemeral"}} + ] + }] + }`) + limited := enforceCacheControlLimit(fiveMarkers, 4) + if total := countCacheControls(limited); total != 4 { + t.Fatalf("countCacheControls(limited) = %d, want 4: %s", total, limited) + } + upgraded := upgradeClaudeCacheControlTTL(limited, claudeCacheControlTTL1h) + if got := ojson.GetBytes(upgraded, "cache_control.ttl").String(); got != "1h" { + t.Fatalf("top-level cache_control.ttl after upgrade = %q, want 1h: %s", got, upgraded) + } + stripped := stripClaudeCacheControlTTL(upgraded) + if ojson.GetBytes(stripped, "cache_control.ttl").Exists() { + t.Fatalf("top-level cache_control.ttl should be stripped: %s", stripped) + } + mixedOrder := []byte(`{ + "cache_control":{"type":"ephemeral","ttl":"1h"}, + "messages":[{"role":"user","content":[{"type":"text","text":"m0","cache_control":{"type":"ephemeral"}}]}] + }`) + normalized := normalizeCacheControlTTL(mixedOrder) + if ojson.GetBytes(normalized, "cache_control.ttl").Exists() { + t.Fatalf("normalizeCacheControlTTL must downgrade top-level 1h TTL after a 5m block: %s", normalized) + } + }) +} diff --git a/internal/runtime/executor/claude_cloaked_date_pin_test.go b/internal/runtime/executor/claude_cloaked_date_pin_test.go index c4955f005d5..86608119f52 100644 --- a/internal/runtime/executor/claude_cloaked_date_pin_test.go +++ b/internal/runtime/executor/claude_cloaked_date_pin_test.go @@ -345,17 +345,19 @@ func TestClaudeCloakedDateReminderProbeOverrideDoesNotReinitializeSessionDate(t } func TestClaudeCloakedDateReminderPersistsAcrossRestartWithLocalBBoltKV(t *testing.T) { - // Create the temp dir before registering the shutdown cleanup. Cleanups run - // LIFO, so the store must close before TempDir removal, or Windows fails to - // delete the still-open kv.db. - dbPath := filepath.Join(t.TempDir(), "kv.db") - homekv.ClearCurrent() - _ = kvstore.ShutdownLocal() + if errShutdown := kvstore.ShutdownLocal(); errShutdown != nil { + t.Fatalf("kvstore.ShutdownLocal() setup error = %v", errShutdown) + } helps.ResetClaudeDiagnosticsForTest() + + // Cleanups run in LIFO order: close bbolt before TempDir removes its files. + dbPath := filepath.Join(t.TempDir(), "kv.db") t.Cleanup(func() { homekv.ClearCurrent() - _ = kvstore.ShutdownLocal() + if errShutdown := kvstore.ShutdownLocal(); errShutdown != nil { + t.Errorf("kvstore.ShutdownLocal() cleanup error = %v", errShutdown) + } helps.ResetClaudeDiagnosticsForTest() }) diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index c20855ab686..a90251d9025 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -447,6 +447,25 @@ func (e *ClaudeExecutor) applyNativeToolSearchFallback(ctx context.Context, auth return ctx, req, opts } +// pinResponsesSystemItemPlacement decides, before translation, whether the +// Responses translator keeps mid-session system items in place as role=system +// turns, which keeps the prompt-cache prefix of earlier turns stable. The turn +// needs the cloaked Claude Code wire, which sends the mid-conversation system +// beta, and a model that accepts it. Strict cloaking drops every caller system +// instruction and rebuild_mid_system_message folds the turns into the system +// field, so both keep the top-level path, as does an uncloaked request, which +// forwards the caller's betas unchanged. +func (e *ClaudeExecutor) pinResponsesSystemItemPlacement(ctx context.Context, auth *cliproxyauth.Auth, from sdktranslator.Format, upstreamModel string, policy claudeWirePolicy, settings claudeCloakSettings) context.Context { + if from != sdktranslator.FormatOpenAIResponse { + return ctx + } + inPlace := policy.Cloak && + !settings.strictMode && + !util.IsClaudeLegacySystemReminderModel(upstreamModel) && + !rebuildMidSystemMessageEnabled(e.cfg, auth) + return responsesclaude.WithTopLevelSystemItems(ctx, !inPlace) +} + func (e *ClaudeExecutor) modelLevelCooling() bool { return e != nil && e.cfg != nil && e.cfg.Claude.ModelLevelCooling } diff --git a/internal/runtime/executor/claude_executor_beta_passthrough_test.go b/internal/runtime/executor/claude_executor_beta_passthrough_test.go index 5c37b17c1de..7c351ee5160 100644 --- a/internal/runtime/executor/claude_executor_beta_passthrough_test.go +++ b/internal/runtime/executor/claude_executor_beta_passthrough_test.go @@ -21,7 +21,7 @@ func fixtureAuth() *cliproxyauth.Auth { // A caller the cloak does not recognize (a Claude Code newer than the pinned // profile) keeps betas the proxy does not manage: dropping them fails whole // turns whose features need the missing authorization (#5738). per-turn-control -// is now assembled for fable-5-1, so this uses a beta that is still unmanaged. +// and message-threads are now assembled, so this uses a beta that is still unmanaged. func TestApplyClaudeHeaders_ForwardsUnmanagedCallerBetas(t *testing.T) { t.Parallel() @@ -31,16 +31,16 @@ func TestApplyClaudeHeaders_ForwardsUnmanagedCallerBetas(t *testing.T) { t.Fatalf("NewRequest() error = %v", errReq) } incoming := http.Header{} - incoming.Set("Anthropic-Beta", "message-threads-2026-08-12") + incoming.Set("Anthropic-Beta", "future-feature-2030-01-01") if errHeaders := applyClaudeHeaders(req, auth, auth.Attributes[cliproxyauth.AttributeAPIKey], false, nil, []byte(`{"model":"claude-fable-5-1"}`), &config.Config{}, incoming, false); errHeaders != nil { t.Fatalf("applyClaudeHeaders() error = %v", errHeaders) } betas := req.Header.Get("Anthropic-Beta") - if !strings.Contains(betas, "message-threads-2026-08-12") { + if !strings.Contains(betas, "future-feature-2030-01-01") { t.Fatalf("Anthropic-Beta = %q, want unmanaged caller beta forwarded", betas) } - if !strings.Contains(betas, "mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24") { - t.Fatalf("Anthropic-Beta = %q, want fable-5-1 per-turn-control between mid-conversation-system and tool-changes", betas) + if !strings.Contains(betas, "prompt-caching-scope-2026-01-05,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24") { + t.Fatalf("Anthropic-Beta = %q, want fable-5-1 per-turn-control right after prompt-caching-scope", betas) } } diff --git a/internal/runtime/executor/claude_executor_cli_2321_test.go b/internal/runtime/executor/claude_executor_cli_2321_test.go new file mode 100644 index 00000000000..84c1bd5cadc --- /dev/null +++ b/internal/runtime/executor/claude_executor_cli_2321_test.go @@ -0,0 +1,109 @@ +package executor + +import ( + "strings" + "testing" +) + +// claudeCode2321MainBetas is the Anthropic-Beta value native Claude Code 2.1.321 +// sent for an interactive cli main-thread turn on Fable 5.2, Fable 5.1, Opus 5.5 +// and Sonnet 5.5 (captured against api.anthropic.com on 2026-11-09 with OAuth, +// in an auto-mode session, hence afk-mode). +const claudeCode2321MainBetas = "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05,per-turn-control-2026-07-01," + + "mid-conversation-tool-changes-2026-07-01,inline-tools-2026-09-15,advisor-tool-2026-03-01," + + "thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,dangerous-tool-use-2026-09-03,thinking-binding-controls-2026-08-01," + + "thinking-display-updates-2026-08-18,afk-mode-2026-01-31,extended-cache-ttl-2025-04-11," + + "cache-diagnosis-2026-04-07,message-threads-2026-08-12" + +// claudeCode2321MainBody keeps only the fields of the captured main-thread body +// that gate betas: a role=system turn with an inline tool_addition, a deferred +// tool, the advisor tool, safeguards, thread, diagnostics, and a 1h cache TTL. +func claudeCode2321MainBody(model, display string) []byte { + return []byte(`{"model":"` + model + `","max_tokens":128000,` + + `"messages":[{"role":"user","content":[{"type":"text","text":"say ok"}]},` + + `{"role":"system","content":[{"type":"text","text":"tools changed"},` + + `{"type":"tool_addition","tool":{"type":"tool_definition","definition":{"name":"mcp__docs__guide","input_schema":{"type":"object"}}}}]}],` + + `"system":[{"type":"text","text":"You are Claude Code, Anthropic's official CLI for Claude."},` + + `{"type":"text","text":"prompt","cache_control":{"type":"ephemeral","ttl":"1h","scope":"global"}}],` + + `"tools":[{"name":"Read","input_schema":{"type":"object"}},` + + `{"name":"DeferredToolPlaceholder","input_schema":{"type":"object"},"defer_loading":true},` + + `{"type":"advisor_20260301","name":"advisor","model":"claude-opus-5-5","defer_loading":true}],` + + `"thinking":{"type":"adaptive","display":"` + display + `"},` + + `"context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},` + + `"safeguards":{"type":"dangerous_tool_use"},"output_config":{"effort":"medium"},` + + `"thread":{"type":"create"},"diagnostics":{"previous_message_id":null},` + + `"metadata":{"user_id":"user"},"stream":true}`) +} + +func TestClaudeCodeCLIBetas_MatchesClaudeCode2321Captures(t *testing.T) { + // Haiku 5.5 moves claude-code-20250219 behind prompt-caching-scope. + haiku55 := strings.Replace(claudeCode2321MainBetas, + "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05,", + "oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05,claude-code-20250219,", 1) + // The sdk-cli entrypoint sends thinking.display=omitted. + sdkCLI := withoutClaudeBeta(claudeCode2321MainBetas, claudeThinkingDisplayUpdatesBeta) + + tests := []struct { + name string + body []byte + native string + }{ + {name: "fable-5-2 cli", body: claudeCode2321MainBody("claude-fable-5-2", "updates"), native: claudeCode2321MainBetas}, + {name: "fable-5-1 cli", body: claudeCode2321MainBody("claude-fable-5-1", "updates"), native: claudeCode2321MainBetas}, + {name: "opus-5-5 cli", body: claudeCode2321MainBody("claude-opus-5-5", "updates"), native: claudeCode2321MainBetas}, + {name: "sonnet-5-5 cli", body: claudeCode2321MainBody("claude-sonnet-5-5", "updates"), native: claudeCode2321MainBetas}, + {name: "haiku-5-5 cli", body: claudeCode2321MainBody("claude-haiku-5-5", "updates"), native: haiku55}, + {name: "fable-5-2 sdk-cli", body: claudeCode2321MainBody("claude-fable-5-2", "omitted"), native: sdkCLI}, + { + name: "haiku-4-5 quota probe", + body: []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"quota"}],"metadata":{"user_id":"user"}}`), + native: "oauth-2025-04-20,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,context-management-2025-06-27,prompt-caching-scope-2026-01-05", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + requested := claudeRequestedBetas(tt.native, nil) + if got := claudeCodeCLIBetas(tt.body, requested, true); got != tt.native { + t.Fatalf("with native betas requested:\n got: %s\nwant: %s", got, tt.native) + } + // A caller that sends no betas gets the same list minus afk-mode, + // which Claude Code only sends from auto-mode sessions. + want := withoutClaudeBeta(tt.native, claudeAFKModeBeta) + if got := claudeCodeCLIBetas(tt.body, nil, true); got != want { + t.Fatalf("with no betas requested:\n got: %s\nwant: %s", got, want) + } + }) + } +} + +// The 2.1.321 capture of the prompt "say ok" carried cc_version=2.1.321.20d. +func TestComputeFingerprintMatchesClaudeCode2321Capture(t *testing.T) { + if got := computeFingerprint("say ok", "2.1.321"); got != "20d" { + t.Fatalf("computeFingerprint() = %q, want 20d", got) + } +} + +func TestClaudeIncludePerTurnTiming_2321CatalogModels(t *testing.T) { + for _, tc := range []struct { + model string + want bool + }{ + {model: "claude-fable-5-2", want: true}, + {model: "claude-fable-5-2[1m]", want: true}, + {model: "claude-fable-5-1", want: true}, + {model: "claude-opus-5-5", want: true}, + {model: "claude-sonnet-5-5", want: true}, + {model: "claude-mythos-5-2", want: true}, + {model: "claude-haiku-5-5", want: false}, + {model: "claude-fable-5", want: false}, + } { + t.Run(tc.model, func(t *testing.T) { + body := []byte(`{"model":"` + tc.model + `","output_config":{"timing":{"type":"enabled"}}}`) + if got := claudeIncludePerTurnTiming(body, nil); got != tc.want { + t.Fatalf("claudeIncludePerTurnTiming(%s) = %v, want %v", tc.model, got, tc.want) + } + }) + } +} diff --git a/internal/runtime/executor/claude_executor_cloaking.go b/internal/runtime/executor/claude_executor_cloaking.go index 9bf62699e31..1ebac79a1cd 100644 --- a/internal/runtime/executor/claude_executor_cloaking.go +++ b/internal/runtime/executor/claude_executor_cloaking.go @@ -696,12 +696,12 @@ Report what actually happened, not what you intended. When you say something is func checkSystemInstructionsWithMode(payload []byte, strictMode bool, keepCallerSystemTopLevel ...bool) []byte { keepAtTopLevel := len(keepCallerSystemTopLevel) > 0 && keepCallerSystemTopLevel[0] - return checkSystemInstructionsWithSigningMode(payload, strictMode, false, "2.1.280", "cli", "", keepAtTopLevel) + return checkSystemInstructionsWithSigningMode(payload, strictMode, false, "2.1.321", "cli", "", keepAtTopLevel) } func checkSystemInstructionsWithModeDoc(doc *ojson.Doc, strictMode bool, keepCallerSystemTopLevel ...bool) bool { keepAtTopLevel := len(keepCallerSystemTopLevel) > 0 && keepCallerSystemTopLevel[0] - return checkSystemInstructionsWithSigningModeDoc(doc, strictMode, false, "2.1.280", "cli", "", keepAtTopLevel) + return checkSystemInstructionsWithSigningModeDoc(doc, strictMode, false, "2.1.321", "cli", "", keepAtTopLevel) } // checkSystemInstructionsWithSigningMode keeps the top-level system in Claude @@ -945,51 +945,15 @@ func relocateClaudeSystemPromptForCountTokensWithPolicyDoc(doc *ojson.Doc, stric return true } -// claudeLegacySystemReminderModels lists the official Anthropic model IDs and -// aliases that reject a mid-conversation role=system message. Entries mirror the -// "claude" provider in internal/registry/models/models.json plus Anthropic's own -// bare and "-latest" aliases. Other providers' synthetic IDs do not belong here. -var claudeLegacySystemReminderModels = map[string]struct{}{ - "claude-3-5-haiku-20241022": {}, - "claude-3-5-haiku-latest": {}, - "claude-3-7-sonnet-20250219": {}, - "claude-3-7-sonnet-latest": {}, - "claude-haiku-4-5": {}, - "claude-haiku-4-5-20251001": {}, - "claude-opus-4": {}, - "claude-opus-4-20250514": {}, - "claude-opus-4-1": {}, - "claude-opus-4-1-20250805": {}, - "claude-opus-4-5": {}, - "claude-opus-4-5-20251101": {}, - "claude-opus-4-6": {}, - "claude-opus-4-7": {}, - "claude-sonnet-4": {}, - "claude-sonnet-4-20250514": {}, - "claude-sonnet-4-5": {}, - "claude-sonnet-4-5-20250929": {}, - "claude-sonnet-4-6": {}, -} - func claudeUsesLegacySystemReminder(payload []byte) bool { - model := strings.ToLower(strings.TrimSpace(ojson.GetBytesNoCopy(payload, "model").String())) - if slash := strings.LastIndexByte(model, '/'); slash >= 0 { - model = model[slash+1:] - } - _, legacy := claudeLegacySystemReminderModels[model] - return legacy + return util.IsClaudeLegacySystemReminderModel(ojson.GetBytesNoCopy(payload, "model").String()) } func claudeUsesLegacySystemReminderDoc(doc *ojson.Doc) bool { if doc == nil { return false } - model := strings.ToLower(strings.TrimSpace(doc.Get("model").String())) - if slash := strings.LastIndexByte(model, '/'); slash >= 0 { - model = model[slash+1:] - } - _, legacy := claudeLegacySystemReminderModels[model] - return legacy + return util.IsClaudeLegacySystemReminderModel(doc.Get("model").String()) } // claudeCallerSystemBlockError reports a caller system block that Claude cannot @@ -2733,7 +2697,7 @@ func ensureCacheControlDoc(doc *ojson.Doc) bool { if doc == nil { return false } - changed := false + changed, materializedInMessages := materializeTopLevelCacheControlDoc(doc) if !claudeDocHasCacheableSystem(doc) { if injectToolsCacheControlDoc(doc) { changed = true @@ -2742,12 +2706,114 @@ func ensureCacheControlDoc(doc *ojson.Doc) bool { if injectSystemCacheControlDoc(doc) { changed = true } - if injectMessagesCacheControlDoc(doc) { + if !materializedInMessages && injectMessagesCacheControlDoc(doc) { changed = true } return changed } +func isBareEphemeralClaudeCacheControl(cc ojson.Result) bool { + if !isValidClaudeCacheControl(cc) { + return false + } + return len(cc.Map()) == 1 +} + +func buildTextBlockWithRawCacheControl(text, rawCC string) string { + if rawCC != "" { + return `{"type":"text","text":` + marshalJSONStringWithoutHTMLEscape(text) + `,"cache_control":` + rawCC + `}` + } + return buildTextBlock(text, nil) +} + +func applyTopLevelCacheControlToBlockDoc(doc *ojson.Doc, blockPath string, topCC ojson.Result) bool { + ccPath := blockPath + ".cache_control" + existingCC := doc.Get(ccPath) + if !isValidClaudeCacheControl(existingCC) || (isBareEphemeralClaudeCacheControl(existingCC) && !isBareEphemeralClaudeCacheControl(topCC)) { + if errSet := doc.SetRawBytes(ccPath, []byte(topCC.Raw)); errSet != nil { + return false + } + } + return doc.Delete("cache_control") == nil +} + +// materializeTopLevelCacheControlDoc moves a top-level cache_control (Anthropic +// automatic prompt caching) onto the last cacheable block in the request and +// removes the top-level field so it participates in the 4-breakpoint limit and +// TTL normalization/upgrades without duplicating the final breakpoint. +func materializeTopLevelCacheControlDoc(doc *ojson.Doc) (bool, bool) { + if doc == nil || doc.Len() == 0 { + return false, false + } + topCC := doc.Get("cache_control") + if !isValidClaudeCacheControl(topCC) { + return false, false + } + + var scanner claudeStreamDocScanner + scanner.scan(doc, false, false) + if scanner.messagesIsArray && scanner.lastMessageIndex >= 0 { + if scanner.lastEligibleIndex >= 0 && scanner.lastMsgRole == "system" && scanner.lastMsgContentIsNonEmptyString { + contentPath := fmt.Sprintf("messages.%d.content", scanner.lastMessageIndex) + finalContent := doc.Get(contentPath) + if finalContent.Type == ojson.String && strings.TrimSpace(finalContent.String()) != "" { + newContent := "[" + buildTextBlockWithRawCacheControl(finalContent.String(), topCC.Raw) + "]" + if doc.SetRawBytes(contentPath, []byte(newContent)) == nil && doc.Delete("cache_control") == nil { + return true, true + } + } + } + if scanner.lastEligibleIndex >= 0 { + if scanner.lastEligibleContentIsArray && scanner.lastEligibleTargetBlockIndex >= 0 { + blockPath := fmt.Sprintf("messages.%d.content.%d", scanner.lastEligibleIndex, scanner.lastEligibleTargetBlockIndex) + if applyTopLevelCacheControlToBlockDoc(doc, blockPath, topCC) { + return true, true + } + } + if scanner.lastEligibleContentIsNonEmptyStr { + contentPath := fmt.Sprintf("messages.%d.content", scanner.lastEligibleIndex) + contentRes := doc.Get(contentPath) + if contentRes.Type == ojson.String && strings.TrimSpace(contentRes.String()) != "" { + newContent := "[" + buildTextBlockWithRawCacheControl(contentRes.String(), topCC.Raw) + "]" + if doc.SetRawBytes(contentPath, []byte(newContent)) == nil && doc.Delete("cache_control") == nil { + return true, true + } + } + } + } + } + + system := doc.Get("system") + if system.IsArray() { + if lastSysIdx := lastCacheableClaudeSystemIndex(system); lastSysIdx >= 0 { + return applyTopLevelCacheControlToBlockDoc(doc, fmt.Sprintf("system.%d", lastSysIdx), topCC), false + } + } else if system.Type == ojson.String && strings.TrimSpace(system.String()) != "" { + newSystem := "[" + buildTextBlockWithRawCacheControl(system.String(), topCC.Raw) + "]" + if doc.SetRawBytes("system", []byte(newSystem)) == nil { + return doc.Delete("cache_control") == nil, false + } + } + + tools := doc.Get("tools") + if tools.IsArray() { + lastEligibleToolIndex := -1 + idx := 0 + tools.ForEach(func(_, tool ojson.Result) bool { + if !tool.Get("defer_loading").Bool() { + lastEligibleToolIndex = idx + } + idx++ + return true + }) + if lastEligibleToolIndex >= 0 { + return applyTopLevelCacheControlToBlockDoc(doc, fmt.Sprintf("tools.%d", lastEligibleToolIndex), topCC), false + } + } + + return false, false +} + // claudePayloadHasCacheableSystem reports whether the payload has a system prompt // that injectSystemCacheControl can actually host a breakpoint on. An absent key, an // empty array, an array of only empty/whitespace text blocks, and an empty string @@ -2828,6 +2894,20 @@ func upgradeClaudeCacheControlTTL(payload []byte, ttl string) []byte { return doc.Bytes() } +func claudeCacheControlPath(path string) string { + if path == "" { + return "cache_control" + } + return path + ".cache_control" +} + +func claudeCacheControlTTLPath(path string) string { + if path == "" { + return "cache_control.ttl" + } + return path + ".cache_control.ttl" +} + func upgradeClaudeCacheControlTTLDoc(doc *ojson.Doc, ttl string) bool { if ttl == "" || doc == nil || doc.Len() == 0 { return false @@ -2849,7 +2929,7 @@ func upgradeClaudeCacheControlTTLDoc(doc *ojson.Doc, ttl string) bool { upgraded += `,"scope":` + scope.Raw } upgraded += "}" - if errSet := doc.SetRawBytes(path+".cache_control", []byte(upgraded)); errSet == nil { + if errSet := doc.SetRawBytes(claudeCacheControlPath(path), []byte(upgraded)); errSet == nil { changed = true } return true @@ -2882,7 +2962,7 @@ func stripClaudeCacheControlTTLDoc(doc *ojson.Doc) bool { if !cacheControl.IsObject() || !cacheControl.Get("ttl").Exists() { return true } - if errDel := doc.Delete(path + ".cache_control.ttl"); errDel == nil { + if errDel := doc.Delete(claudeCacheControlTTLPath(path)); errDel == nil { changed = true } return true @@ -2897,6 +2977,7 @@ const ( claudeSecTools claudeSecSystem claudeSecMessages + claudeSecTopCacheControl ) const ( @@ -2972,6 +3053,7 @@ type claudeStreamDocScanner struct { sysCCBuf [8]claudeScannedCC msgCC []claudeScannedCC msgCCBuf [16]claudeScannedCC + topCC []byte messagesIsArray bool messagesLen int @@ -3010,6 +3092,7 @@ func (s *claudeStreamDocScanner) scan(doc *ojson.Doc, collectCC, stopOnFirstUser s.toolCC = s.toolCCBuf[:0] s.sysCC = s.sysCCBuf[:0] s.msgCC = s.msgCCBuf[:0] + s.topCC = nil _, _ = doc.WriteTo(s) } @@ -3077,6 +3160,8 @@ func (s *claudeStreamDocScanner) finishCapture(p []byte, endIdx int) { s.sysCC = append(s.sysCC, claudeScannedCC{idx1: s.elemIdx, raw: stored}) case claudeSecMessages: s.msgCC = append(s.msgCC, claudeScannedCC{idx1: s.elemIdx, idx2: s.blockIdx, raw: stored}) + case claudeSecTopCacheControl: + s.topCC = stored } } @@ -3204,6 +3289,8 @@ func (s *claudeStreamDocScanner) finishKey() { s.topSection = claudeSecSystem case "messages": s.topSection = claudeSecMessages + case "cache_control": + s.topSection = claudeSecTopCacheControl default: s.topSection = claudeSecNone } @@ -3248,6 +3335,10 @@ func (s *claudeStreamDocScanner) finishKey() { func (s *claudeStreamDocScanner) startValue(c byte, idx int) { switch s.depth { + case 1: + if s.topSection == claudeSecTopCacheControl && s.collectCC && s.isObj[1] { + s.startCapture(1, idx) + } case 2: if s.topSection != claudeSecNone && !s.isObj[2] { if s.topSection == claudeSecMessages { @@ -3654,6 +3745,14 @@ func forEachClaudeCacheControlPathDoc(doc *ojson.Doc, visit func(path string, ca return } } + if len(scanner.topCC) > 0 { + ccRes := util.ParseGJSONBytesNoCopy(scanner.topCC) + if ccRes.Exists() { + if !visit("", ccRes) { + return + } + } + } } // forEachClaudeCacheControlBlock walks every block that can carry cache_control @@ -3764,11 +3863,15 @@ func shouldEnsureCacheControlDoc(doc *ojson.Doc, cloaked, confirmedClaudeCode bo } func countCacheControls(payload []byte) int { - return countCacheControlsResult( + count := countCacheControlsResult( ojson.GetBytesNoCopy(payload, "system"), ojson.GetBytesNoCopy(payload, "tools"), ojson.GetBytesNoCopy(payload, "messages"), ) + if isValidClaudeCacheControl(ojson.GetBytesNoCopy(payload, "cache_control")) { + count++ + } + return count } func countCacheControlsDoc(doc *ojson.Doc) int { @@ -3865,7 +3968,7 @@ func normalizeCacheControlTTLDoc(doc *ojson.Doc) bool { if !seen5m { return true } - if errDel := doc.Delete(path + ".cache_control.ttl"); errDel == nil { + if errDel := doc.Delete(claudeCacheControlTTLPath(path)); errDel == nil { modified = true } return true @@ -3911,13 +4014,15 @@ func enforceCacheControlLimitDoc(doc *ojson.Doc, maxBlocks int) bool { maxBlocks-- } - var toolPaths, systemPaths, messagePaths []string + var toolPaths, systemPaths, messagePaths, topPaths []string forEachClaudeCacheControlPathDoc(doc, func(path string, cc ojson.Result) bool { if !isValidClaudeCacheControl(cc) { return true } - ccPath := path + ".cache_control" + ccPath := claudeCacheControlPath(path) switch { + case path == "": + topPaths = append(topPaths, ccPath) case strings.HasPrefix(path, "tools."): toolPaths = append(toolPaths, ccPath) case strings.HasPrefix(path, "system."): @@ -3928,7 +4033,7 @@ func enforceCacheControlLimitDoc(doc *ojson.Doc, maxBlocks int) bool { return true }) - total := len(toolPaths) + len(systemPaths) + len(messagePaths) + total := len(toolPaths) + len(systemPaths) + len(messagePaths) + len(topPaths) if total <= maxBlocks { return false } @@ -3972,6 +4077,10 @@ func enforceCacheControlLimitDoc(doc *ojson.Doc, maxBlocks int) bool { if len(toolPaths) > 0 { deleteList(toolPaths[len(toolPaths)-1:]) } + if excess <= 0 { + return changed + } + deleteList(topPaths) return changed } diff --git a/internal/runtime/executor/claude_executor_diagnostics_test.go b/internal/runtime/executor/claude_executor_diagnostics_test.go index 5f3e7cdbfd0..b8c96baa4ad 100644 --- a/internal/runtime/executor/claude_executor_diagnostics_test.go +++ b/internal/runtime/executor/claude_executor_diagnostics_test.go @@ -193,7 +193,7 @@ func TestClaudeExecutorContinuityAdvancesRequestIDAndPromptIDInBillingHeader(t * // Verify Turn 1: h1 := capturedBillingHeaders[0] - if !strings.Contains(h1, "cc_version=2.1.280.") || !strings.Contains(h1, "cc_entrypoint=cli;") || !strings.Contains(h1, "cch=") { + if !strings.Contains(h1, "cc_version=2.1.321.") || !strings.Contains(h1, "cc_entrypoint=cli;") || !strings.Contains(h1, "cch=") { t.Fatalf("h1 invalid: %s", h1) } if strings.Contains(h1, "cc_prev_req=") { diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go index 51c978f66bd..ed58582cee2 100644 --- a/internal/runtime/executor/claude_executor_execute.go +++ b/internal/runtime/executor/claude_executor_execute.go @@ -68,6 +68,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r incomingHeaders, claudeCodeDetection := detectIncomingClaudeCodeRequest(ctx, opts.Headers, originalPayload, false, e.cfg) wirePolicy, wireSettings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, claudeCodeDetection.Confirmed) confirmedClaudeCode := isEffectiveConfirmedClaudeCode(claudeCodeDetection, wirePolicy) + ctx = e.pinResponsesSystemItemPlacement(ctx, auth, from, upstreamModel, wirePolicy, wireSettings) claudeSessionID := "" if fp.ProfileClaudeCodeCLI { claudeSessionID = helps.ClaudeAgentSessionUUIDForRequest(incomingHeaders, originalPayload, req.Payload, confirmedClaudeCode, opts.Metadata, req.Metadata) diff --git a/internal/runtime/executor/claude_executor_native_helper_test.go b/internal/runtime/executor/claude_executor_native_helper_test.go index da7893617c1..e453efb4519 100644 --- a/internal/runtime/executor/claude_executor_native_helper_test.go +++ b/internal/runtime/executor/claude_executor_native_helper_test.go @@ -30,7 +30,7 @@ func claudeNativeHelperHeaders(betas, compression string) http.Header { "Accept": {"application/json"}, "Accept-Encoding": {compression}, "Content-Type": {"application/json"}, - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {betas}, "Anthropic-Version": {"2023-06-01"}, @@ -39,8 +39,8 @@ func claudeNativeHelperHeaders(betas, compression string) http.Header { "X-Client-Request-Id": {"66666666-7777-4888-8999-aaaaaaaaaaaa"}, "X-Stainless-Lang": {"js"}, "X-Stainless-Runtime": {"node"}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-OS": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, "X-Stainless-Retry-Count": {"0"}, diff --git a/internal/runtime/executor/claude_executor_request.go b/internal/runtime/executor/claude_executor_request.go index 425010e677b..82bd6c81755 100644 --- a/internal/runtime/executor/claude_executor_request.go +++ b/internal/runtime/executor/claude_executor_request.go @@ -54,6 +54,8 @@ const ( claudeCacheDiagnosisBeta = "cache-diagnosis-2026-04-07" claudeRedactThinkingBeta = "redact-thinking-2026-02-12" claudeAFKModeBeta = "afk-mode-2026-01-31" + claudeThinkingTokenCountBeta = "thinking-token-count-2026-05-13" + claudeMessageThreadsBeta = "message-threads-2026-08-12" ) // claudeCodeCLIConstantBetas are the betas Claude Code sends on every @@ -66,7 +68,6 @@ const ( var claudeCodeCLIConstantBetas = []string{ "interleaved-thinking-2025-05-14", claudeRedactThinkingBeta, - "thinking-token-count-2026-05-13", "context-management-2025-06-27", "prompt-caching-scope-2026-01-05", } @@ -112,6 +113,8 @@ var claudeManagedBetaSet = func() map[string]bool { claudeCacheDiagnosisBeta, claudeRedactThinkingBeta, claudeAFKModeBeta, + claudeThinkingTokenCountBeta, + claudeMessageThreadsBeta, } managed = append(managed, claudeCodeCLIConstantBetas...) managed = append(managed, claudeCodeTrailingBetas...) @@ -127,52 +130,61 @@ func isManagedClaudeBeta(beta string) bool { } // claudeCodeCLIBetas assembles the Anthropic-Beta baseline the way Claude Code -// 2.1.280 does: the list is per-request, not a fixed string. requested holds the +// 2.1.321 does: the list is per-request, not a fixed string. requested holds the // betas the caller asked for, which decide the capability flags below. // // Verified against api.anthropic.com with native 2.1.258 captures on interactive, // non-interactive, subagent, and multi-model paths (Sonnet, Opus, Fable, Haiku). -// Claude Code 2.1.280 (measured 2026-09-23, binary 80abbfe) inserts -// mid-conversation-tool-changes immediately after mid-conversation-system on the -// same non-legacy models. The betas named in issue #6054 are feature-gated in -// that binary, so they are emitted only for the model capability or body shape +// Claude Code 2.1.321 (measured 2026-11-09 on cli and sdk-cli entrypoints for +// Fable 5.2, Fable 5.1, Opus 5.5, Sonnet 5.5, Haiku 5.5 and the Haiku 4.5 quota +// probe) no longer sends mid-conversation-system even with role=system turns, +// moves thinking-token-count behind advisor-tool, sends effort ahead of +// dangerous-tool-use, and adds message-threads for bodies carrying thread. +// Feature-gated betas are emitted only for the model capability or body shape // that actually sends them, in the same relative order: // -// 1 claude-code-20250219 +// 1 claude-code-20250219 after prompt-caching-scope on Haiku 5.5; never on quota probes // 2 oauth-2025-04-20 OAuth credentials only // 3 context-1m-2025-08-07 [1m] model variants only // 4 interleaved-thinking-2025-05-14 // 5 redact-thinking-2026-02-12 cli entrypoint, no thinking.display -// 6 thinking-token-count-2026-05-13 -// 7 context-management-2025-06-27 -// 8 prompt-caching-scope-2026-01-05 -// 9 mid-conversation-system-2026-04-07 models accepting a role=system turn -// 10 per-turn-control-2026-07-01 opus-5-5 and fable-5-1, or requested -// 11 timing-2026-09-09 per-turn timing body, or requested -// 12 mid-conversation-tool-changes-2026-07-01 same models as mid-conversation-system -// 13 inline-tools-2026-09-15 inline tool_addition blocks, or requested -// 14 advisor-tool-2026-03-01 requests declaring advisor tools or requesting advisor beta +// 6 context-management-2025-06-27 +// 7 prompt-caching-scope-2026-01-05 +// 8 mid-conversation-system-2026-04-07 requested only, on models accepting a role=system turn +// 9 per-turn-control-2026-07-01 per-turn-control models, or requested +// 10 timing-2026-09-09 per-turn timing body, or requested +// 11 mid-conversation-tool-changes-2026-07-01 models accepting a role=system turn, except Sonnet 5 +// 12 inline-tools-2026-09-15 inline tool_addition blocks, or requested +// 13 advisor-tool-2026-03-01 requests declaring advisor tools or requesting advisor beta +// 14 thinking-token-count-2026-05-13 every request except quota probes // 15 advanced-tool-use-2025-11-20 requests using tool search or another advanced tool-use feature -// 16 mid-conversation-system-clear-at-2026-08-21 messages with clear_at, or requested -// 17 dangerous-tool-use-2026-09-03 safeguards body, or requested -// 18 effort-2025-11-24 effort-supporting models with active thinking +// 16 mid-conversation-system-clear-at-2026-08-21 progress-display models, messages with clear_at, or requested +// 17 effort-2025-11-24 effort-supporting models with active thinking +// 18 dangerous-tool-use-2026-09-03 safeguards body, or requested // 19 server-side-fallback-2026-06-01 requests with fallbacks or requested // 20 fallback-credit-2026-06-01 requests with fallback tokens, fallbacks, or requested // 21 structured-outputs-2025-12-15 structured output requests -// 22 thinking-binding-controls-2026-08-01 thinking.block_binding, or requested +// 22 thinking-binding-controls-2026-08-01 thinking.block_binding, adaptive progress-display models, or requested // 23 thinking-display-updates-2026-08-18 requests with thinking.display=updates -// 24 thinking-resumption-2026-07-17 requested only; the 2.1.280 flag defaults off +// 24 thinking-resumption-2026-07-17 requested only; the flag defaults off // 25 fast-mode-2026-02-01 speed:fast requests only // 26 afk-mode-2026-01-31 auto-mode sessions, forwarded when the caller sends it // 27 extended-cache-ttl-2025-04-11 OAuth credentials (omitted on subagent & probe) // 28 prompt-caching-evict-2026-05-12 evict_on_complete, or requested // 29 cache-diagnosis-2026-04-07 requests with diagnostics only +// 30 message-threads-2026-08-12 requests with a thread object, or requested // // An empty body keeps the optimistic role=system default, matching the cloaking // policy for unknown and future model IDs. func claudeCodeCLIBetas(body []byte, requested map[string]bool, oauthToken bool) string { - betas := make([]string, 0, len(claudeCodeCLIConstantBetas)+len(claudeCodeTrailingBetas)+10) - betas = append(betas, claudeCodeBeta) + betas := make([]string, 0, len(claudeCodeCLIConstantBetas)+len(claudeCodeTrailingBetas)+12) + isProbeOrHelper := helps.IsClaudeProbeOrHelperRequest(body) + isQuotaProbe := helps.IsClaudeQuotaProbeRequest(body) + model := ojson.GetBytesNoCopy(body, "model").String() + claudeCodeAfterScope := isClaudeHaiku55Model(model) + if !isQuotaProbe && !claudeCodeAfterScope { + betas = append(betas, claudeCodeBeta) + } if oauthToken { betas = append(betas, claudeOAuthBeta) } @@ -186,15 +198,20 @@ func claudeCodeCLIBetas(body []byte, requested map[string]bool, oauthToken bool) } betas = append(betas, beta) } + if !isQuotaProbe && claudeCodeAfterScope { + betas = append(betas, claudeCodeBeta) + } if !claudeUsesLegacySystemReminder(body) { - betas = append(betas, claudeMidConvSystemBeta) + if requested[claudeMidConvSystemBeta] { + betas = append(betas, claudeMidConvSystemBeta) + } if claudeIncludePerTurnControl(body, requested) { betas = append(betas, claudePerTurnControlBeta) } if claudeIncludePerTurnTiming(body, requested) { betas = append(betas, claudePerTurnTimingBeta) } - if !isClaudeSonnet5Model(ojson.GetBytesNoCopy(body, "model").String()) { + if !isClaudeSonnet5Model(model) { betas = append(betas, claudeMidConvToolChangesBeta) } if claudeIncludeInlineTools(body, requested) { @@ -213,19 +230,21 @@ func claudeCodeCLIBetas(body []byte, requested map[string]bool, oauthToken bool) if requested[claudeAdvisorToolBeta] || claudeBodyHasAdvisorTool(body) { betas = append(betas, claudeAdvisorToolBeta) } + if !isQuotaProbe || requested[claudeThinkingTokenCountBeta] { + betas = append(betas, claudeThinkingTokenCountBeta) + } if requested[claudeAdvancedToolUseBeta] || claudeBodyUsesAdvancedToolUse(body) { betas = append(betas, claudeAdvancedToolUseBeta) } if !claudeUsesLegacySystemReminder(body) && claudeIncludeMidConvClearAt(body, requested) { betas = append(betas, claudeMidConvSystemClearAtBeta) } - if requested[claudeDangerousToolUseBeta] || ojson.GetBytesNoCopy(body, "safeguards").Exists() { - betas = append(betas, claudeDangerousToolUseBeta) - } if claudeRequestSupportsEffort(body, requested) { betas = append(betas, claudeEffortBeta) } - isProbeOrHelper := helps.IsClaudeProbeOrHelperRequest(body) + if requested[claudeDangerousToolUseBeta] || ojson.GetBytesNoCopy(body, "safeguards").Exists() { + betas = append(betas, claudeDangerousToolUseBeta) + } if !isProbeOrHelper && (requested[claudeServerSideFallbackBeta] || ojson.GetBytesNoCopy(body, "fallbacks").Exists()) { betas = append(betas, claudeServerSideFallbackBeta) } @@ -274,6 +293,9 @@ func claudeCodeCLIBetas(body []byte, requested map[string]bool, oauthToken bool) if diagnostics := ojson.GetBytesNoCopy(body, "diagnostics"); diagnostics.IsObject() { betas = append(betas, claudeCacheDiagnosisBeta) } + if requested[claudeMessageThreadsBeta] || ojson.GetBytesNoCopy(body, "thread").IsObject() { + betas = append(betas, claudeMessageThreadsBeta) + } return strings.Join(betas, ",") } @@ -289,7 +311,7 @@ func claudeCanonicalModel(model string) string { return model } -// claudeModelHasPerTurnEffort reports models whose 2.1.280 catalog capability +// claudeModelHasPerTurnEffort reports models whose 2.1.321 catalog capability // per_turn_effort puts per-turn-control-2026-07-01 on every first-party request. func isClaudeOpus55Model(model string) bool { model = claudeCanonicalModel(model) @@ -337,7 +359,7 @@ func isClaudeFable51OnlyModel(model string) bool { // prefix fails binding validation unless the request opts into dropping them. // See https://platform.claude.com/docs/en/build-with-claude/preserved-thinking func isClaudePreservedThinkingModel(model string) bool { - return isClaudeOpus55Model(model) || isClaudeSonnet55Model(model) || isClaudeHaiku55Model(model) || isClaudeFable51OnlyModel(model) + return isClaudeOpus55Model(model) || isClaudeSonnet55Model(model) || isClaudeHaiku55Model(model) || isClaudeFable51OnlyModel(model) || isClaudeFable52Model(model) } // claudeThinkingTypeAcceptsBlockBinding reports whether the body's thinking @@ -432,10 +454,22 @@ func applyClaudePreservedThinkingBindingDoc(doc *ojson.Doc) bool { return doc.Set("thinking.block_binding.prefix_mismatch_behavior", "drop_block") == nil } -// claudeModelUsesProgressDisplay reports the 2.1.280 interactive CLI models that -// send thinking.display=updates unless the caller already chose a display mode. +// claudeModelUsesProgressDisplay reports the interactive CLI models that send +// thinking.display=updates unless the caller already chose a display mode. +// Claude Code 2.1.321 added Fable 5.2 and Haiku 5.5. func claudeModelUsesProgressDisplay(model string) bool { - return isClaudeOpus55Model(model) || isClaudeFable51Model(model) || isClaudeSonnet5Model(model) || isClaudeSonnet55Model(model) + return isClaudeOpus55Model(model) || isClaudeFable51Model(model) || isClaudeFable52Model(model) || + isClaudeSonnet5Model(model) || isClaudeSonnet55Model(model) || isClaudeHaiku55Model(model) +} + +func isClaudeFable52Model(model string) bool { + model = claudeCanonicalModel(model) + for _, prefix := range []string{"claude-fable-5-2", "claude-fable-5.2"} { + if model == prefix || strings.HasPrefix(model, prefix+"-") || strings.HasPrefix(model, prefix+"[") { + return true + } + } + return false } // applyClaudeCloakThinkingDisplay fills the latest CLI display only when the @@ -468,7 +502,11 @@ func applyClaudeCloakThinkingDisplayDoc(doc *ojson.Doc, callerOwned bool) bool { } func claudeModelHasPerTurnEffort(model string) bool { - return isClaudeOpus55Model(model) || strings.HasPrefix(claudeCanonicalModel(model), "claude-fable-5-1") + return isClaudeOpus55Model(model) || + strings.HasPrefix(claudeCanonicalModel(model), "claude-fable-5-1") || + isClaudeFable52Model(model) || + isClaudeSonnet55Model(model) || + isClaudeHaiku55Model(model) } // claudeModelHasPerTurnTiming reports models whose catalog lists per_turn_timing. @@ -476,9 +514,17 @@ func claudeModelHasPerTurnEffort(model string) bool { // is set, so the beta follows the body or an explicit caller request. func claudeModelHasPerTurnTiming(model string) bool { model = claudeCanonicalModel(model) - return claudeModelHasPerTurnEffort(model) || strings.HasPrefix(model, "claude-mythos-5-1") + if isClaudeHaiku55Model(model) { + return false + } + return claudeModelHasPerTurnEffort(model) || + strings.HasPrefix(model, "claude-mythos-5-1") || + strings.HasPrefix(model, "claude-mythos-5-2") || + strings.HasPrefix(model, "claude-mythos-5.2") } +// claudeIncludePerTurnControl also covers the models Claude Code 2.1.321 sends +// per-turn-control-2026-07-01 for on every main-thread request. func claudeIncludePerTurnControl(body []byte, requested map[string]bool) bool { if requested[claudePerTurnControlBeta] { return true @@ -546,7 +592,8 @@ func claudeRequestSupportsEffort(body []byte, requested map[string]bool) bool { return false } model := strings.ToLower(strings.TrimSpace(ojson.GetBytesNoCopy(body, "model").String())) - if isClaudeHaikuModel(model) { + // Haiku 5.5 is the first Haiku that Claude Code 2.1.321 sends effort for. + if isClaudeHaikuModel(model) && !isClaudeHaiku55Model(model) { return false } thinkingType := strings.ToLower(strings.TrimSpace(ojson.GetBytesNoCopy(body, "thinking.type").String())) @@ -1234,7 +1281,7 @@ func applyClaudeHeadersWithNativeProfile( baseBetas = withClaudeAdvisorToolBeta(baseBetas) } // Measured Haiku helper requests already carry their own cache-ttl beta - // choice (omitting extended-cache-ttl on 2.1.220/2.1.280 probes/title + // choice (omitting extended-cache-ttl on 2.1.220/2.1.280/2.1.321 probes/title // helpers), while helpers arriving via ANTHROPIC_BASE_URL still need // oauth-2025-04-20 when CPA authenticates upstream with OAuth. // Native Claude Code subagents and probes also omit extended-cache-ttl. @@ -1332,7 +1379,7 @@ func applyClaudeHeadersWithNativeProfile( baseBetas = withoutClaudeBeta(baseBetas, claudeExtendedCacheTTLBeta) } } - // Enforce strict native Claude Code 2.1.280 model & turn beta gating: + // Enforce strict native Claude Code 2.1.321 model & turn beta gating: if !claudeRequestSupportsEffort(body, nil) { baseBetas = withoutClaudeBeta(baseBetas, claudeEffortBeta) } @@ -1456,7 +1503,7 @@ func applyClaudeHeadersWithNativeProfile( identityHeader("Anthropic-Version", "2023-06-01") identityHeader("Anthropic-Dangerous-Direct-Browser-Access", "true") identityHeader("X-App", "cli") - // Values below match Claude Code 2.1.280 / @anthropic-ai/sdk 0.112.1. + // Values below match Claude Code 2.1.321 / @anthropic-ai/sdk 0.128.0. identityHeader("X-Stainless-Retry-Count", "0") identityHeader("X-Stainless-Runtime", "node") identityHeader("X-Stainless-Lang", "js") @@ -1574,7 +1621,7 @@ func applyClaudeHeadersWithNativeProfile( if errDeviceProfile != nil { return errDeviceProfile } - helps.ApplyClaudeDeviceProfileHeaders(r, deviceProfile) + helps.ApplyClaudeDeviceProfileHeaders(r, deviceProfile, helps.HeaderValueCaseInsensitive(incomingHeaders, "User-Agent")) } else { helps.ApplyClaudeDefaultDeviceProfileHeaders(r, cfg) } diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index 171e94ef1ae..59b1221cccc 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -73,6 +73,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A incomingHeaders, claudeCodeDetection := detectIncomingClaudeCodeRequest(ctx, opts.Headers, originalPayload, false, e.cfg) wirePolicy, wireSettings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, claudeCodeDetection.Confirmed) confirmedClaudeCode := isEffectiveConfirmedClaudeCode(claudeCodeDetection, wirePolicy) + ctx = e.pinResponsesSystemItemPlacement(ctx, auth, from, upstreamModel, wirePolicy, wireSettings) claudeSessionID := "" if fp.ProfileClaudeCodeCLI { claudeSessionID = helps.ClaudeAgentSessionUUIDForRequest(incomingHeaders, originalPayload, req.Payload, confirmedClaudeCode, opts.Metadata, req.Metadata) diff --git a/internal/runtime/executor/claude_executor_stream_terminal_test.go b/internal/runtime/executor/claude_executor_stream_terminal_test.go index b68f5a1d8b5..8c45659ecc6 100644 --- a/internal/runtime/executor/claude_executor_stream_terminal_test.go +++ b/internal/runtime/executor/claude_executor_stream_terminal_test.go @@ -2,9 +2,14 @@ package executor import ( "context" + "io" + "net" "net/http" "net/http/httptest" + "net/http/httptrace" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -224,3 +229,161 @@ func TestClaudeExecutor_ExecuteStream_Passthrough_ClientDisconnectAfterTerminalE t.Fatal("timed out waiting for usage record") } } + +func TestClaudeExecutor_ExecuteStream_DoesNotDrainOrReuseUpstreamConnectionAfterMessageStop(t *testing.T) { + const streamData = "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_123","type":"message","role":"assistant","content":[],"model":"claude-opus-5","stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}` + "\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":1}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + + var ( + newConns atomic.Int32 + release = make(chan struct{}, 4) + closed = make(chan struct{}, 4) + done = make(chan struct{}) + mu sync.Mutex + reused []bool + ) + defer close(done) + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(streamData)) + w.(http.Flusher).Flush() + // Hold the chunked terminator back until the client either attempts a + // post-message_stop drain read or closes the connection. + select { + case <-release: + case <-r.Context().Done(): + case <-done: + } + })) + server.Config.ConnState = func(_ net.Conn, state http.ConnState) { + switch state { + case http.StateNew: + newConns.Add(1) + case http.StateClosed, http.StateHijacked: + select { + case closed <- struct{}{}: + default: + } + } + } + server.Start() + defer server.Close() + + baseTransport := &http.Transport{} + t.Cleanup(baseTransport.CloseIdleConnections) + transport := postMessageStopReadReleaseRoundTripper{ + base: baseTransport, + total: len(streamData), + release: release, + } + trace := &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + mu.Lock() + reused = append(reused, info.Reused) + mu.Unlock() + }, + PutIdleConn: func(err error) { + if err == nil { + select { + case closed <- struct{}{}: + default: + } + } + }, + } + + exec := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{ + "api_key": "key-no-drain", + "base_url": server.URL, + }, + } + payload := []byte(`{"model":"claude-opus-5","messages":[{"role":"user","content":"hi"}],"stream":true}`) + + for i := 0; i < 2; i++ { + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(transport)) + ctx = httptrace.WithClientTrace(ctx, trace) + result, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-opus-5", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatClaude, + ResponseFormat: sdktranslator.FormatClaude, + }) + if err != nil { + t.Fatalf("round %d ExecuteStream() error = %v", i+1, err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("round %d chunk error = %v", i+1, chunk.Err) + } + } + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatalf("round %d connection did not settle", i+1) + } + } + + mu.Lock() + gotReused := append([]bool(nil), reused...) + mu.Unlock() + if got := newConns.Load(); got != 2 { + t.Fatalf("stream upstream connections = %d, want 2 (reused = %v)", got, gotReused) + } + if len(gotReused) != 2 || gotReused[0] || gotReused[1] { + t.Fatalf("stream GotConn reused = %v, want [false false]", gotReused) + } +} + +type postMessageStopReadReleaseRoundTripper struct { + base http.RoundTripper + total int + release chan<- struct{} +} + +func (rt postMessageStopReadReleaseRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := rt.base.RoundTrip(req) + if err != nil { + return nil, err + } + resp.Body = &postMessageStopReadReleaseBody{ + ReadCloser: resp.Body, + total: rt.total, + release: rt.release, + } + return resp, nil +} + +type postMessageStopReadReleaseBody struct { + io.ReadCloser + total int + read int + release chan<- struct{} + once sync.Once +} + +func (b *postMessageStopReadReleaseBody) Read(p []byte) (int, error) { + if b.read >= b.total { + b.once.Do(func() { + select { + case b.release <- struct{}{}: + default: + } + }) + } + n, err := b.ReadCloser.Read(p) + b.read += n + return n, err +} diff --git a/internal/runtime/executor/claude_executor_subagent_ttl_regression_test.go b/internal/runtime/executor/claude_executor_subagent_ttl_regression_test.go index 0d995a078f8..50b25dcb6de 100644 --- a/internal/runtime/executor/claude_executor_subagent_ttl_regression_test.go +++ b/internal/runtime/executor/claude_executor_subagent_ttl_regression_test.go @@ -55,7 +55,7 @@ func TestClaudeExecutor_SubagentPreserves1hTTLAndExtendedCacheBetaWhenRequested( }] }`) incomingHeaders := http.Header{} - incomingHeaders.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + incomingHeaders.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") incomingHeaders.Set("X-Claude-Code-Agent-Id", "agent-sub-123") incomingHeaders.Set("Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20,extended-cache-ttl-2025-04-11") @@ -132,7 +132,7 @@ func TestClaudeExecutor_SubagentPreserves1hTTLAndExtendedCacheBetaWhenRequested_ }] }`) incomingHeaders := http.Header{} - incomingHeaders.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + incomingHeaders.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") incomingHeaders.Set("X-Claude-Code-Agent-Id", "agent-sub-stream-456") incomingHeaders.Set("Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20,extended-cache-ttl-2025-04-11") @@ -202,7 +202,7 @@ func TestClaudeExecutor_SubagentWithout1hKeepsDefault5m(t *testing.T) { }] }`) incomingHeaders := http.Header{} - incomingHeaders.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + incomingHeaders.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") incomingHeaders.Set("X-Claude-Code-Agent-Id", "agent-sub-default-789") incomingHeaders.Set("Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20") @@ -268,7 +268,7 @@ func TestClaudeExecutor_APIKeySubagentPreserves1hTTLAndInjectsBeta(t *testing.T) }] }`) incomingHeaders := http.Header{} - incomingHeaders.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + incomingHeaders.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") incomingHeaders.Set("X-Claude-Code-Agent-Id", "agent-sub-apikey-101") incomingHeaders.Set("Anthropic-Beta", "claude-code-20250219") // Note: NO extended-cache-ttl beta in incoming header @@ -334,7 +334,7 @@ func TestClaudeExecutor_APIKeySubagentPreserves1hTTLAndInjectsBeta_Stream(t *tes }] }`) incomingHeaders := http.Header{} - incomingHeaders.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + incomingHeaders.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") incomingHeaders.Set("X-Claude-Code-Agent-Id", "agent-sub-apikey-stream-202") incomingHeaders.Set("Anthropic-Beta", "claude-code-20250219") // Note: NO extended-cache-ttl beta in incoming header diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index c9d39fed00e..f0b91b0cb29 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -816,7 +816,7 @@ func TestClaudeExecutor_NonClaudeRequestUsesClaudeCode220CLIFingerprint(t *testi t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, "claude-cli/2.1.280 (external, cli)", "0.112.1", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, "claude-cli/2.1.321 (external, cli)", "0.128.0", "v26.12.1", "MacOS", "arm64") if got := seenHeaders.Get("X-App"); got != "cli" { t.Fatalf("X-App = %q, want cli", got) } @@ -828,8 +828,8 @@ func TestClaudeExecutor_NonClaudeRequestUsesClaudeCode220CLIFingerprint(t *testi if len(system) != 2 { t.Fatalf("system block count = %d, want 2: %s", len(system), seenBody) } - if got := system[0].Get("text").String(); got != "x-anthropic-billing-header: cc_version=2.1.280.d7b; cc_entrypoint=cli;" { - t.Fatalf("billing header = %q, want 2.1.280 CLI fingerprint", got) + if got := system[0].Get("text").String(); got != "x-anthropic-billing-header: cc_version=2.1.321.cf7; cc_entrypoint=cli;" { + t.Fatalf("billing header = %q, want 2.1.321 CLI fingerprint", got) } if got := system[1].Get("text").String(); got != claudeCodeCLIIdentity { t.Fatalf("system[1].text = %q, want official CLI identity", got) @@ -877,12 +877,12 @@ func TestClaudeExecutor_ConfirmedNewerPatchClaudeCodeRequestPreservesInteractive const userID = `{"device_id":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","account_uuid":"","session_id":"11111111-2222-4333-8444-555555555555"}` payload := []byte(`{"model":"claude-opus-4-6","system":[{"type":"text","text":"interactive-system","cache_control":{"type":"ephemeral"}}],"messages":[{"role":"user","content":"x"}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `}}`) incoming := http.Header{ - "User-Agent": {"claude-cli/2.1.281 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.322 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,thinking-token-count-2026-05-13,context-management-2025-06-27,prompt-caching-scope-2026-01-05,effort-2025-11-24"}, "X-Claude-Code-Session-Id": {sessionID}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -904,7 +904,7 @@ func TestClaudeExecutor_ConfirmedNewerPatchClaudeCodeRequestPreservesInteractive t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, "claude-cli/2.1.281 (external, cli)", "0.112.1", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, "claude-cli/2.1.322 (external, cli)", "0.128.0", "v26.12.1", "MacOS", "arm64") if got := ojson.GetBytes(seenBody, "system.0.text").String(); got != "interactive-system" { t.Fatalf("system.0.text = %q, want confirmed client system preserved", got) } @@ -947,7 +947,7 @@ func TestClaudeExecutor_ConfirmedClaudeCodeWithoutCacheControlPreservesContent(t const userID = `{"device_id":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","account_uuid":"","session_id":"11111111-2222-4333-8444-555555555555"}` payload := []byte(`{"model":"claude-opus-4-6","messages":[{"role":"user","content":"x"}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `}}`) incoming := http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}, "X-Claude-Code-Session-Id": {sessionID}, @@ -1006,7 +1006,7 @@ func TestClaudeExecutor_ConfirmedCLIForwardsNativeOpus55Shape(t *testing.T) { const betas = "claude-code-20250219,oauth-2025-04-20,mid-conversation-system-clear-at-2026-08-21,thinking-binding-controls-2026-08-01,extended-cache-ttl-2025-04-11" payload := []byte(`{"model":"claude-opus-5-5","system":[{"type":"text","text":` + fmt.Sprintf("%q", billing) + `},{"type":"text","text":"native identity"},{"type":"text","text":"native model block"}],"messages":[{"role":"user","content":"hello"}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `},"fallbacks":[{"model":"claude-opus-4-8"}],"thinking":{"type":"adaptive","display":"updates"}}`) incoming := http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {betas}, "X-Claude-Code-Session-Id": {"11111111-2222-4333-8444-555555555555"}, @@ -1062,7 +1062,7 @@ func TestClaudeExecutor_ConfirmedVSCodeAgentSDKRequestPreservesIdentity(t *testi const sessionID = "22222222-3333-4444-8555-666666666666" const userID = `{"device_id":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","account_uuid":"","session_id":"22222222-3333-4444-8555-666666666666"}` - const vscodeUA = "claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)" + const vscodeUA = "claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)" const billingHeader = "x-anthropic-billing-header: cc_version=2.1.258.9cb; cc_entrypoint=claude-vscode;" payload := []byte(`{"model":"claude-opus-4-6","system":[{"type":"text","text":` + fmt.Sprintf("%q", billingHeader) + `},{"type":"text","text":"You are a Claude agent, built on Anthropic's Claude Agent SDK.","cache_control":{"type":"ephemeral","ttl":"1h"}},{"type":"text","text":"vscode-agent-system"}],"messages":[{"role":"user","content":"x"}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `}}`) incoming := http.Header{ @@ -1071,8 +1071,8 @@ func TestClaudeExecutor_ConfirmedVSCodeAgentSDKRequestPreservesIdentity(t *testi "Anthropic-Beta": {"claude-code-20250219,interleaved-thinking-2025-05-14"}, "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -1095,7 +1095,7 @@ func TestClaudeExecutor_ConfirmedVSCodeAgentSDKRequestPreservesIdentity(t *testi t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, vscodeUA, "0.112.1", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, vscodeUA, "0.128.0", "v26.12.1", "MacOS", "arm64") if got := seenHeaders.Get("Anthropic-Dangerous-Direct-Browser-Access"); got != "true" { t.Fatalf("Anthropic-Dangerous-Direct-Browser-Access = %q, want preserved true", got) } @@ -1147,7 +1147,7 @@ func TestClaudeExecutor_CopiedVSCodeAgentSDKHeadersWithoutMetadataAreCloaked(t * SourceFormat: sdktranslator.FormatClaude, OriginalRequest: payload, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)"}, + "User-Agent": {"claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}, }, @@ -1156,7 +1156,7 @@ func TestClaudeExecutor_CopiedVSCodeAgentSDKHeadersWithoutMetadataAreCloaked(t * t.Fatalf("Execute() error = %v", errExecute) } - if got := seenHeaders.Get("User-Agent"); got != "claude-cli/2.1.280 (external, cli)" { + if got := seenHeaders.Get("User-Agent"); got != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("User-Agent = %q, want CLI cloak", got) } if got := ojson.GetBytes(seenBody, "system.#").Int(); got != 3 { @@ -1201,7 +1201,7 @@ func TestClaudeExecutor_AgentSDKEntrypointWithStrongSignalsUsesCLICloak(t *testi SourceFormat: sdktranslator.FormatClaude, OriginalRequest: payload, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, sdk-ts, agent-sdk/0.3.220)"}, + "User-Agent": {"claude-cli/2.1.321 (external, sdk-ts, agent-sdk/0.3.220)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}, }, @@ -1210,7 +1210,7 @@ func TestClaudeExecutor_AgentSDKEntrypointWithStrongSignalsUsesCLICloak(t *testi t.Fatalf("Execute() error = %v", errExecute) } - if got := seenHeaders.Get("User-Agent"); got != "claude-cli/2.1.280 (external, cli)" { + if got := seenHeaders.Get("User-Agent"); got != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("User-Agent = %q, want CLI cloak", got) } if got := ojson.GetBytes(seenBody, "system.0.text").String(); !strings.Contains(got, "cc_entrypoint=cli;") { @@ -1256,11 +1256,11 @@ func TestClaudeExecutor_ConfirmedVSCodeOAuthPreservesToolNames(t *testing.T) { SourceFormat: sdktranslator.FormatClaude, OriginalRequest: payload, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)"}, + "User-Agent": {"claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, }, }) if errExecute != nil { @@ -1300,7 +1300,7 @@ func TestClaudeExecutor_ConfirmedSDKTSAgentSDKRequestPassesThroughWithRealProfil const sessionID = "33333333-4444-4555-8666-777777777777" const userID = `{"account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","device_id":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","session_id":"33333333-4444-4555-8666-777777777777"}` - const sdkTSUA = "claude-cli/2.1.304 (external, sdk-ts, agent-sdk/0.3.276)" + const sdkTSUA = "claude-cli/2.1.345 (external, sdk-ts, agent-sdk/0.3.276)" const billingHeader = "x-anthropic-billing-header: cc_version=2.1.304.179; cc_entrypoint=sdk-ts;" payload := []byte(`{"model":"claude-sonnet-5-5","max_tokens":1024,"system":[{"type":"text","text":` + fmt.Sprintf("%q", billingHeader) + `},{"type":"text","text":"You are Claude Code, Anthropic's official CLI for Claude, running within the Claude Agent SDK.","cache_control":{"type":"ephemeral","ttl":"1h"}},{"type":"text","text":"t3-code-system-prompt","cache_control":{"type":"ephemeral","ttl":"1h"}}],"tools":[{"name":"Agent","description":"subagent","input_schema":{"type":"object"}},{"name":"Bash","description":"shell","input_schema":{"type":"object"}}],"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `},"thinking":{"type":"adaptive","display":"summarized"},"output_config":{"effort":"low"}}`) incoming := http.Header{ @@ -1310,7 +1310,7 @@ func TestClaudeExecutor_ConfirmedSDKTSAgentSDKRequestPassesThroughWithRealProfil "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, "X-Stainless-Package-Version": {"0.128.0"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -1340,7 +1340,7 @@ func TestClaudeExecutor_ConfirmedSDKTSAgentSDKRequestPassesThroughWithRealProfil t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, sdkTSUA, "0.128.0", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, sdkTSUA, "0.128.0", "v26.12.1", "MacOS", "arm64") if got := ojson.GetBytes(seenBody, "system.#").Int(); got != 3 { t.Fatalf("system block count = %d, want 3: %s", got, seenBody) } @@ -1381,7 +1381,7 @@ func TestClaudeExecutor_ConfiguredNativeEntrypointPassesThroughWithRealProfile(t const sessionID = "33333333-4444-4555-8666-777777777777" const userID = `{"account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","device_id":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","session_id":"33333333-4444-4555-8666-777777777777"}` - const t3CodeUA = "claude-cli/2.1.316 (external, t3-code, agent-sdk/0.3.276)" + const t3CodeUA = "claude-cli/2.1.357 (external, t3-code, agent-sdk/0.3.276)" const billingHeader = "x-anthropic-billing-header: cc_version=2.1.316.179; cc_entrypoint=t3-code;" payload := []byte(`{"model":"claude-sonnet-5-5","max_tokens":1024,"system":[{"type":"text","text":` + fmt.Sprintf("%q", billingHeader) + `},{"type":"text","text":"You are a Claude agent, built on Anthropic's Claude Agent SDK.","cache_control":{"type":"ephemeral","ttl":"1h"}},{"type":"text","text":"t3-code-system-prompt","cache_control":{"type":"ephemeral","ttl":"1h"}}],"tools":[{"name":"Agent","description":"subagent","input_schema":{"type":"object"}},{"name":"Bash","description":"shell","input_schema":{"type":"object"}}],"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `},"thinking":{"type":"adaptive","display":"summarized"},"output_config":{"effort":"low"}}`) incoming := http.Header{ @@ -1391,7 +1391,7 @@ func TestClaudeExecutor_ConfiguredNativeEntrypointPassesThroughWithRealProfile(t "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, "X-Stainless-Package-Version": {"0.128.0"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -1425,7 +1425,7 @@ func TestClaudeExecutor_ConfiguredNativeEntrypointPassesThroughWithRealProfile(t t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, t3CodeUA, "0.128.0", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, t3CodeUA, "0.128.0", "v26.12.1", "MacOS", "arm64") if got := ojson.GetBytes(seenBody, "system.#").Int(); got != 3 { t.Fatalf("system block count = %d, want 3: %s", got, seenBody) } @@ -1466,7 +1466,7 @@ func TestClaudeExecutor_CloakModeNeverPreservesConfirmedShapeClaudeCodeCallerHea const sessionID = "44444444-5555-4666-8777-888888888888" const userID = `{"account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","device_id":"dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd","session_id":"44444444-5555-4666-8777-888888888888"}` - const callerUA = "claude-cli/2.1.304 (external, sdk-py, agent-sdk/0.3.276)" + const callerUA = "claude-cli/2.1.345 (external, sdk-py, agent-sdk/0.3.276)" payload := []byte(`{"model":"claude-haiku-5-5","max_tokens":1024,"system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.304.c80; cc_entrypoint=sdk-py"},{"type":"text","text":"You are a Claude agent, built on Anthropic's Claude Agent SDK."},{"type":"text","text":"You are naming a coding session."}],"tools":[],"messages":[{"role":"user","content":[{"type":"text","text":"name session"}]}],"metadata":{"user_id":` + fmt.Sprintf("%q", userID) + `}}`) incoming := http.Header{ "User-Agent": {callerUA}, @@ -1475,7 +1475,7 @@ func TestClaudeExecutor_CloakModeNeverPreservesConfirmedShapeClaudeCodeCallerHea "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, "X-Stainless-Package-Version": {"0.128.0"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -1506,7 +1506,7 @@ func TestClaudeExecutor_CloakModeNeverPreservesConfirmedShapeClaudeCodeCallerHea t.Fatalf("Execute() error = %v", errExecute) } - assertClaudeFingerprint(t, seenHeaders, callerUA, "0.128.0", "v26.3.0", "MacOS", "arm64") + assertClaudeFingerprint(t, seenHeaders, callerUA, "0.128.0", "v26.12.1", "MacOS", "arm64") if bytes.Contains(seenBody, []byte(`"cache_control"`)) { t.Fatalf("cloak_mode: never confirmed-shape helper unexpectedly injected cache_control: %s", seenBody) } @@ -2791,7 +2791,7 @@ func TestClaudeExecutor_CountTokensUpstreamConfirmedVSCodePreservesCustomTool(t }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatClaude, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)"}, + "User-Agent": {"claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}, }, @@ -2987,7 +2987,7 @@ func TestClaudeExecutor_CountTokensConfirmedNativePreservesMeasuredOAuthBody(t * _, errCount := executor.countTokensUpstream(ctx, auth, cliproxyexecutor.Request{Model: "claude-opus-5", Payload: payload}, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatClaude, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {incomingBetas}, }, @@ -5732,7 +5732,7 @@ func TestClaudeExecutor_RebuildMidSystemMessageDisabledByDefault(t *testing.T) { }} payload := []byte(`{"system":[{"type":"text","text":"Top rule","cache_control":{"type":"ephemeral"}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},{"role":"system","content":"Mid rule"},{"role":"user","content":[{"type":"text","text":"continue"}]}],"metadata":{"user_id":"{\"device_id\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"account_uuid\":\"\",\"session_id\":\"11111111-2222-4333-8444-555555555555\"}"}}`) ctx := contextWithGinHeaders(map[string]string{ - "User-Agent": "claude-cli/2.1.280 (external, cli)", + "User-Agent": "claude-cli/2.1.321 (external, cli)", "X-App": "cli", "Anthropic-Beta": "claude-code-20250219", }) @@ -5778,7 +5778,7 @@ func TestClaudeExecutor_RebuildMidSystemMessageOptInMovesSystemMessages(t *testi }} payload := []byte(`{"system":"Top rule","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},{"role":"system","content":"Mid string rule"},{"role":"assistant","content":[{"type":"text","text":"ok"}]},{"role":"system","content":[{"type":"text","text":"Mid array rule","cache_control":{"type":"ephemeral"}}]},{"role":"user","content":[{"type":"text","text":"continue"}]}],"metadata":{"user_id":"{\"device_id\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"account_uuid\":\"\",\"session_id\":\"11111111-2222-4333-8444-555555555555\"}"}}`) ctx := contextWithGinHeaders(map[string]string{ - "User-Agent": "claude-cli/2.1.280 (external, cli)", + "User-Agent": "claude-cli/2.1.321 (external, cli)", "X-App": "cli", "Anthropic-Beta": "claude-code-20250219", }) @@ -6360,7 +6360,7 @@ func TestClaudeExecutor_SubagentAndProbeOmit1hCacheTTLAndBeta(t *testing.T) { }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatClaude, Headers: http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-Claude-Code-Agent-Id": {"subagent-confirmed-456"}, "Anthropic-Beta": {"claude-code-20250219,oauth-2025-04-20,effort-2025-11-24"}, }, @@ -8603,7 +8603,7 @@ func TestClaudeExecutor_ExecuteOAuthCustomToolMCPAliasRoundTrip(t *testing.T) { if _, ok := claudeBillingCCHDigitsOffset(upstreamBody); !ok { t.Fatalf("Claude OAuth custom BaseURL body is missing CCH: %s", upstreamBody) } - if got := upstreamHeaders.Get("User-Agent"); got != "claude-cli/2.1.280 (external, cli)" { + if got := upstreamHeaders.Get("User-Agent"); got != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("Messages User-Agent = %q, want CLI identity", got) } wantBetas := claudeCodeCLIBetas(payload, nil, true) @@ -8683,7 +8683,7 @@ func TestClaudeExecutor_ExecuteStreamOAuthCustomToolMCPAliasRoundTrip(t *testing if _, ok := claudeBillingCCHDigitsOffset(upstreamBody); !ok { t.Fatalf("streaming Claude OAuth custom BaseURL body is missing CCH: %s", upstreamBody) } - if got := upstreamHeaders.Get("User-Agent"); got != "claude-cli/2.1.280 (external, cli)" { + if got := upstreamHeaders.Get("User-Agent"); got != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("streaming User-Agent = %q, want CLI identity", got) } wantBetas := claudeCodeCLIBetas(payload, nil, true) @@ -8825,7 +8825,7 @@ func TestInsertClaudeMidConversationSystemMessages_IsIdempotent(t *testing.T) { // baseline to Claude Code 2.1.220 behavior captured against api.anthropic.com. // The OAuth profile was reverified on 2026-08-03 with two distinct accounts. func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { - const constants = "claude-code-20250219,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,thinking-token-count-2026-05-13,context-management-2025-06-27,prompt-caching-scope-2026-01-05" + const constants = "claude-code-20250219,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,context-management-2025-06-27,prompt-caching-scope-2026-01-05" tests := []struct { name string @@ -8837,7 +8837,7 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { { name: "legacy model without tools omits both conditional betas", body: `{"model":"claude-opus-4-6"}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "context 1m sits right after claude-code, not at the end", @@ -8845,8 +8845,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { requested: map[string]bool{claudeContext1MBeta: true}, want: "claude-code-20250219,context-1m-2025-08-07," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,effort-2025-11-24", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "opus-5 1m variant reproduces the full observed order", @@ -8858,8 +8858,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { }, want: "claude-code-20250219,context-1m-2025-08-07," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + "advanced-tool-use-2025-11-20,effort-2025-11-24," + "server-side-fallback-2026-06-01,fallback-credit-2026-06-01", }, @@ -8867,48 +8867,52 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { name: "structured outputs trails effort", body: `{"model":"claude-opus-4-6"}`, requested: map[string]bool{claudeStructuredOutputsBeta: true}, - want: constants + ",effort-2025-11-24,structured-outputs-2025-12-15", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24,structured-outputs-2025-12-15", }, { name: "unknown caller beta is not smuggled into the baseline", body: `{"model":"claude-opus-4-6"}`, requested: map[string]bool{"totally-made-up-2030-01-01": true}, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "claude-sonnet-5 accepts role=system", body: `{"model":"claude-sonnet-5"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", }, { name: "claude-sonnet-5-5 carries tool changes beta unlike sonnet-5", body: `{"model":"claude-sonnet-5-5"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", + want: constants + ",per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", }, { name: "claude-opus-4-8 accepts role=system", body: `{"model":"claude-opus-4-8"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "claude-fable-5 accepts role=system", body: `{"model":"claude-fable-5"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "opus-5-5 carries capability betas without optional body fields", body: `{"model":"claude-opus-5-5","thinking":{"type":"adaptive"}}`, - want: constants + ",mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,thinking-binding-controls-2026-08-01", + want: constants + ",per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,thinking-binding-controls-2026-08-01", }, { name: "opus-5-5 without thinking still carries clear-at capability", body: `{"model":"claude-opus-5-5"}`, - want: constants + ",mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", + want: constants + ",per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", }, { name: "fable-5-1 carries per-turn-control but not timing unless the body asks", body: `{"model":"claude-fable-5-1"}`, - want: constants + ",mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", + want: constants + ",per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24", }, { name: "opus-5-5 timing and the other 2.1.280 gated betas keep wire order", @@ -8916,15 +8920,15 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { requested: map[string]bool{ claudeThinkingResumptionBeta: true, }, - want: constants + ",mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,timing-2026-09-09," + - "mid-conversation-tool-changes-2026-07-01,inline-tools-2026-09-15," + - "mid-conversation-system-clear-at-2026-08-21,dangerous-tool-use-2026-09-03,effort-2025-11-24," + - "thinking-binding-controls-2026-08-01,thinking-resumption-2026-07-17,prompt-caching-evict-2026-05-12", + want: constants + ",per-turn-control-2026-07-01,timing-2026-09-09,mid-conversation-tool-changes-2026-07-01," + + "inline-tools-2026-09-15,thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,dangerous-tool-use-2026-09-03,thinking-binding-controls-2026-08-01," + + "thinking-resumption-2026-07-17,prompt-caching-evict-2026-05-12", }, { name: "claude-opus-4-7 stays on the reminder path", body: `{"model":"claude-opus-4-7"}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "oauth uses tool search and the current cache TTL trailer", @@ -8932,8 +8936,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,advanced-tool-use-2025-11-20," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20," + "effort-2025-11-24,extended-cache-ttl-2025-04-11", }, { @@ -8945,110 +8949,111 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { claudeServerSideFallbackBeta: true, claudeFallbackCreditBeta: true, }, - want: "claude-code-20250219,oauth-2025-04-20,context-1m-2025-08-07," + - "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01," + - "advanced-tool-use-2025-11-20,effort-2025-11-24," + - "server-side-fallback-2026-06-01,fallback-credit-2026-06-01," + - "extended-cache-ttl-2025-04-11", + want: "claude-code-20250219,oauth-2025-04-20," + + "context-1m-2025-08-07,interleaved-thinking-2025-05-14," + + "redact-thinking-2026-02-12,context-management-2025-06-27," + + "prompt-caching-scope-2026-01-05,mid-conversation-tool-changes-2026-07-01," + + "thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20," + + "effort-2025-11-24,server-side-fallback-2026-06-01," + + "fallback-credit-2026-06-01,extended-cache-ttl-2025-04-11", }, { name: "api key path sends neither oauth beta", body: `{"model":"claude-opus-4-6"}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "claude-haiku-4-5-20251001 stays on the reminder path and omits effort", body: `{"model":"claude-haiku-4-5-20251001"}`, - want: constants, + want: constants + ",thinking-token-count-2026-05-13", }, { name: "legacy model with inline tools no longer adds advanced tool use", body: `{"model":"claude-sonnet-4-6","tools":[{"name":"Read"}]}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "deferred tool adds advanced tool use", body: `{"model":"claude-sonnet-4-6","tools":[{"name":"Read","defer_loading":true}]}`, - want: constants + ",advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "tool search server tool adds advanced tool use", body: `{"model":"claude-sonnet-4-6","tools":[{"type":"tool_search_tool_regex_20251119","name":"tool_search_tool_regex"},{"name":"Read"}]}`, - want: constants + ",advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "tool use examples add advanced tool use", body: `{"model":"claude-sonnet-4-6","tools":[{"name":"Read","input_examples":[{"path":"a.go"}]}]}`, - want: constants + ",advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "programmatic tool calling adds advanced tool use", body: `{"model":"claude-sonnet-4-6","tools":[{"name":"Read","allowed_callers":["code_execution_20250825"]}]}`, - want: constants + ",advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "requested advanced tool use is honored for inline tools", body: `{"model":"claude-sonnet-4-6","tools":[{"name":"Read"}]}`, requested: map[string]bool{claudeAdvancedToolUseBeta: true}, - want: constants + ",advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "role=system model without tools adds mid conversation system only", body: `{"model":"claude-opus-5"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "role=system model with tool search adds both in wire order", body: `{"model":"claude-opus-5","tools":[{"name":"Read","defer_loading":true}]}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "empty tools array does not add advanced tool use", body: `{"model":"claude-opus-4-6","tools":[]}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "unknown future model keeps the optimistic role=system default", body: `{"model":"claude-future-9"}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "thinking display summarized drops redact-thinking", body: `{"model":"claude-opus-5","thinking":{"type":"adaptive","display":"summarized"}}`, want: "claude-code-20250219,interleaved-thinking-2025-05-14," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + "effort-2025-11-24", }, { name: "thinking display omitted drops redact-thinking as well", body: `{"model":"claude-opus-4-6","thinking":{"type":"enabled","budget_tokens":2048,"display":"omitted"}}`, want: "claude-code-20250219,interleaved-thinking-2025-05-14," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,effort-2025-11-24", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "thinking without display keeps redact-thinking", body: `{"model":"claude-opus-4-6","thinking":{"type":"adaptive"}}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "blank display value keeps redact-thinking", body: `{"model":"claude-opus-4-6","thinking":{"type":"adaptive","display":" "}}`, - want: constants + ",effort-2025-11-24", + want: constants + ",thinking-token-count-2026-05-13,effort-2025-11-24", }, { name: "advisor tool beta requested placed before advanced-tool-use", body: `{"model":"claude-opus-5","tools":[{"name":"Read","defer_loading":true}]}`, requested: map[string]bool{"advisor-tool-2026-03-01": true}, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,advisor-tool-2026-03-01,advanced-tool-use-2025-11-20,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,advisor-tool-2026-03-01,thinking-token-count-2026-05-13," + + "advanced-tool-use-2025-11-20,effort-2025-11-24", }, { name: "body with advisor server tool automatically adds advisor-tool beta", body: `{"model":"claude-opus-5","tools":[{"type":"advisor_20260301","name":"advisor"}]}`, - want: constants + ",mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01,advisor-tool-2026-03-01,effort-2025-11-24", + want: constants + ",mid-conversation-tool-changes-2026-07-01,advisor-tool-2026-03-01,thinking-token-count-2026-05-13,effort-2025-11-24", }, { // Captured 2026-09-02 from Claude Code 2.1.258 (cli entrypoint, OAuth, @@ -9063,10 +9068,12 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { }, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + - "advisor-tool-2026-03-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,thinking-binding-controls-2026-08-01," + - "afk-mode-2026-01-31,extended-cache-ttl-2025-04-11", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + + "advisor-tool-2026-03-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24," + + "thinking-binding-controls-2026-08-01,afk-mode-2026-01-31," + + "extended-cache-ttl-2025-04-11", }, { name: "afk-mode sits between fast-mode and extended-cache-ttl", @@ -9075,8 +9082,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { requested: map[string]bool{claudeAFKModeBeta: true}, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + "effort-2025-11-24,fast-mode-2026-02-01," + "afk-mode-2026-01-31,extended-cache-ttl-2025-04-11", }, @@ -9086,22 +9093,25 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-tool-changes-2026-07-01," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + "effort-2025-11-24,extended-cache-ttl-2025-04-11", }, { name: "thinking display updates emits thinking-display-updates beta and drops redact-thinking", body: `{"model":"claude-fable-5-1","thinking":{"type":"adaptive","display":"updates"}}`, want: "claude-code-20250219,interleaved-thinking-2025-05-14," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + - "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,thinking-binding-controls-2026-08-01,thinking-display-updates-2026-08-18", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,thinking-binding-controls-2026-08-01," + + "thinking-display-updates-2026-08-18", }, { name: "body with fallbacks automatically adds server-side-fallback beta", body: `{"model":"claude-fable-5-1","fallbacks":[{"model":"claude-opus-5"}]}`, - want: constants + ",mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,server-side-fallback-2026-06-01", + want: constants + ",per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01,thinking-token-count-2026-05-13," + + "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,server-side-fallback-2026-06-01", }, { name: "subagent request omits extended-cache-ttl beta", @@ -9109,18 +9119,17 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + "effort-2025-11-24", }, { name: "probe request max_tokens=1 omits effort and extended-cache-ttl betas", body: `{"model":"claude-sonnet-5","max_tokens":1}`, oauth: true, - want: "claude-code-20250219,oauth-2025-04-20," + - "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21", + want: "oauth-2025-04-20,interleaved-thinking-2025-05-14," + + "redact-thinking-2026-02-12,context-management-2025-06-27," + + "prompt-caching-scope-2026-01-05,mid-conversation-system-clear-at-2026-08-21", }, { name: "haiku model omits effort beta even if requested", @@ -9129,9 +9138,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05," + - "extended-cache-ttl-2025-04-11", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,extended-cache-ttl-2025-04-11", }, { name: "disabled thinking omits effort beta even if requested", @@ -9140,8 +9148,8 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21," + + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + "extended-cache-ttl-2025-04-11", }, { @@ -9150,9 +9158,10 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21," + - "effort-2025-11-24,fallback-credit-2026-06-01,extended-cache-ttl-2025-04-11", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,fallback-credit-2026-06-01," + + "extended-cache-ttl-2025-04-11", }, { name: "oauth body with fallbacks automatically adds fallback-credit beta", @@ -9160,10 +9169,11 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + - "mid-conversation-system-clear-at-2026-08-21,effort-2025-11-24,server-side-fallback-2026-06-01,fallback-credit-2026-06-01," + - "extended-cache-ttl-2025-04-11", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "per-turn-control-2026-07-01,mid-conversation-tool-changes-2026-07-01," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,server-side-fallback-2026-06-01," + + "fallback-credit-2026-06-01,extended-cache-ttl-2025-04-11", }, { name: "requested fallback-credit beta is honored", @@ -9172,9 +9182,10 @@ func TestClaudeCodeCLIBetas_MatchesObservedClientMatrix(t *testing.T) { oauth: true, want: "claude-code-20250219,oauth-2025-04-20," + "interleaved-thinking-2025-05-14,redact-thinking-2026-02-12," + - "thinking-token-count-2026-05-13,context-management-2025-06-27," + - "prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07,mid-conversation-system-clear-at-2026-08-21," + - "effort-2025-11-24,fallback-credit-2026-06-01,extended-cache-ttl-2025-04-11", + "context-management-2025-06-27,prompt-caching-scope-2026-01-05," + + "thinking-token-count-2026-05-13,mid-conversation-system-clear-at-2026-08-21," + + "effort-2025-11-24,fallback-credit-2026-06-01," + + "extended-cache-ttl-2025-04-11", }, } @@ -10585,7 +10596,7 @@ func runClaudeExecutorCapture(t *testing.T, cfg *config.Config, cloakMode string } func TestClaudeExecutor_PreservedThinkingReplayInjectsDropBlock(t *testing.T) { - for _, model := range []string{"claude-haiku-5-5", "claude-opus-5-5", "claude-sonnet-5-5", "claude-fable-5-1"} { + for _, model := range []string{"claude-haiku-5-5", "claude-opus-5-5", "claude-sonnet-5-5", "claude-fable-5-1", "claude-fable-5-2", "claude-fable-5-2[1m]"} { t.Run(model, func(t *testing.T) { seenBody, seenHeaders := runClaudeExecutorCapture(t, nil, "", false, model, preservedThinkingReplayPayload(model)) if got := ojson.GetBytes(seenBody, "thinking.block_binding.prefix_mismatch_behavior").String(); got != "drop_block" { @@ -10642,7 +10653,7 @@ func TestClaudeExecutor_PreservedThinkingRequiresSignedBlock(t *testing.T) { } func TestClaudeExecutor_PreservedThinkingOmittedThinkingSynthesizesAdaptiveBinding(t *testing.T) { - for _, model := range []string{"claude-haiku-5-5", "claude-opus-5-5", "claude-sonnet-5-5", "claude-fable-5-1"} { + for _, model := range []string{"claude-haiku-5-5", "claude-opus-5-5", "claude-sonnet-5-5", "claude-fable-5-1", "claude-fable-5-2"} { for _, tc := range []struct { name string block string diff --git a/internal/runtime/executor/claude_executor_tokens.go b/internal/runtime/executor/claude_executor_tokens.go index 5fd9fd6b969..0ffa329f5c1 100644 --- a/internal/runtime/executor/claude_executor_tokens.go +++ b/internal/runtime/executor/claude_executor_tokens.go @@ -11,6 +11,7 @@ import ( "github.com/MoeForge/moego/ojson" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" + responsesclaude "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/claude/openai/responses" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -36,6 +37,11 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut ctx, req, opts = e.applyNativeToolSearchFallback(ctx, auth, from, upstreamModel, req, opts) responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) to := sdktranslator.FromString("claude") + // Local estimation sends nothing upstream, so in-place system turns gain no + // cache benefit; keep the top-level shape its request validation accepts. + if from == sdktranslator.FormatOpenAIResponse { + ctx = responsesclaude.WithTopLevelSystemItems(ctx, true) + } originalPayload := req.Payload if len(opts.OriginalRequest) > 0 { @@ -160,8 +166,9 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy originalPayload = opts.OriginalRequest } incomingHeaders, claudeCodeDetection := detectIncomingClaudeCodeRequest(ctx, opts.Headers, originalPayload, true, e.cfg) - initialPolicy, _ := resolveClaudeWirePolicy(e.cfg, auth, apiKey, claudeCodeDetection.Confirmed) + initialPolicy, initialSettings := resolveClaudeWirePolicy(e.cfg, auth, apiKey, claudeCodeDetection.Confirmed) confirmedClaudeCode := isEffectiveConfirmedClaudeCode(claudeCodeDetection, initialPolicy) + ctx = e.pinResponsesSystemItemPlacement(ctx, auth, from, upstreamModel, initialPolicy, initialSettings) claudeSessionID := "" if fp.ProfileClaudeCodeCLI { claudeSessionID = helps.ClaudeAgentSessionUUIDForRequest(incomingHeaders, originalPayload, req.Payload, confirmedClaudeCode, opts.Metadata, req.Metadata) diff --git a/internal/runtime/executor/claude_fingerprint_policy_test.go b/internal/runtime/executor/claude_fingerprint_policy_test.go index a187eb43d18..b0a45477951 100644 --- a/internal/runtime/executor/claude_fingerprint_policy_test.go +++ b/internal/runtime/executor/claude_fingerprint_policy_test.go @@ -1249,14 +1249,14 @@ func TestClaudeExecutor_OfficialAPIKeyCloakNeverStripsClaudeCodeIdentity(t *test const userID = `{"device_id":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","account_uuid":"11111111-2222-4333-8444-555555555555","session_id":"11111111-2222-4333-8444-555555555555"}` const billing = "x-anthropic-billing-header: cc_version=2.1.280.d7b; cc_entrypoint=cli; cch=00000; cc_turn_origin=human;" incoming := http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,thinking-token-count-2026-05-13,context-management-2025-06-27,prompt-caching-scope-2026-01-05,effort-2025-11-24,cache-diagnosis-2026-04-07"}, "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, "X-Claude-Code-Request-Class": {"main"}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } @@ -1461,13 +1461,13 @@ func TestClaudeExecutor_GlobalDisableClaudeCloakModePreservesConfirmedCallerIden const userID = `{"device_id":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","account_uuid":"11111111-2222-4333-8444-555555555555","session_id":"11111111-2222-4333-8444-555555555555"}` const billing = "x-anthropic-billing-header: cc_version=2.1.280.d7b; cc_entrypoint=cli; cch=00000; cc_turn_origin=human;" incoming := http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219,interleaved-thinking-2025-05-14"}, "Anthropic-Dangerous-Direct-Browser-Access": {"true"}, "X-Claude-Code-Session-Id": {sessionID}, - "X-Stainless-Package-Version": {"0.112.1"}, - "X-Stainless-Runtime-Version": {"v26.3.0"}, + "X-Stainless-Package-Version": {"0.128.0"}, + "X-Stainless-Runtime-Version": {"v26.12.1"}, "X-Stainless-Os": {"MacOS"}, "X-Stainless-Arch": {"arm64"}, } diff --git a/internal/runtime/executor/claude_responses_system_placement_test.go b/internal/runtime/executor/claude_responses_system_placement_test.go new file mode 100644 index 00000000000..d61ca2a1e94 --- /dev/null +++ b/internal/runtime/executor/claude_responses_system_placement_test.go @@ -0,0 +1,243 @@ +package executor + +import ( + "context" + "errors" + "net/http" + "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" +) + +// responsesMidSessionPayload is a Responses request whose developer item +// arrives after the conversation started, the shape Codex sends on later turns. +func responsesMidSessionPayload(model string) []byte { + return []byte(`{"model":"` + model + `","instructions":"Top rule","input":[` + + `{"role":"user","content":[{"type":"input_text","text":"hi"}]},` + + `{"role":"assistant","content":[{"type":"output_text","text":"ok"}]},` + + `{"role":"developer","content":[{"type":"input_text","text":"Mid rule"}]},` + + `{"role":"user","content":[{"type":"input_text","text":"go"}]}]}`) +} + +func responsesPlacementAuth(attributes map[string]string) *cliproxyauth.Auth { + auth := midSystemAuth() + for key, value := range attributes { + auth.Attributes[key] = value + } + return auth +} + +type responsesPlacementCall func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, auth *cliproxyauth.Auth, model string, payload []byte) error + +var responsesPlacementCalls = map[string]responsesPlacementCall{ + "execute": func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, auth *cliproxyauth.Auth, model string, payload []byte) error { + t.Helper() + _, err := ex.Execute(ctx, auth, cliproxyexecutor.Request{Model: model, Payload: payload}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + return err + }, + "stream": func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, auth *cliproxyauth.Auth, model string, payload []byte) error { + t.Helper() + result, err := ex.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: model, Payload: payload}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + return err + } + for chunk := range result.Chunks { + if chunk.Err != nil { + return chunk.Err + } + } + return nil + }, + "count_tokens": func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, auth *cliproxyauth.Auth, model string, payload []byte) error { + t.Helper() + _, err := ex.CountTokens(ctx, auth, cliproxyexecutor.Request{Model: model, Payload: payload}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + return err + }, +} + +func claudeSystemTurnTexts(body []byte) []string { + var texts []string + ojson.GetBytes(body, "messages").ForEach(func(_, message ojson.Result) bool { + if message.Get("role").String() == "system" { + texts = append(texts, claudeMessageContentText(message.Get("content"))) + } + return true + }) + return texts +} + +// With the cloaked wire and a model that accepts mid-conversation system turns, +// the developer item reaches the upstream right after the user message it +// preceded, so the messages of earlier turns keep their cached prefix. +func TestResponsesMidSessionSystemItemStaysInPlace(t *testing.T) { + const model = "claude-sonnet-5" + for name, call := range responsesPlacementCalls { + t.Run(name, func(t *testing.T) { + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + if err := call(t, ex, upstream.context(t, nil), midSystemAuth(), model, responsesMidSessionPayload(model)); err != nil { + t.Fatalf("call error = %v", err) + } + if !upstream.called { + t.Fatal("expected the request to reach the upstream") + } + messages := ojson.GetBytes(upstream.body, "messages").Array() + last := len(messages) - 1 + if last < 1 || messages[last].Get("role").String() != "system" || messages[last-1].Get("role").String() != "user" { + t.Fatalf("want the developer turn right after the last user message; body=%s", upstream.body) + } + if got := claudeMessageContentText(messages[last].Get("content")); got != "Mid rule" { + t.Fatalf("in-place system turn = %q, want %q; body=%s", got, "Mid rule", upstream.body) + } + }) + } +} + +// A model that rejects mid-conversation system turns keeps the top-level path, +// so the executor sees the same body as before in-place placement: both +// instructions in source order, placed together by the existing cloaking. +func TestResponsesMidSessionSystemItemKeepsTopLevelPathForLegacyModel(t *testing.T) { + const model = "claude-sonnet-4-5" + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + if err := responsesPlacementCalls["execute"](t, ex, upstream.context(t, nil), midSystemAuth(), model, responsesMidSessionPayload(model)); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if texts := claudeSystemTurnTexts(upstream.body); len(texts) != 0 { + t.Fatalf("legacy body carries system turns %q; body=%s", texts, upstream.body) + } + body := string(upstream.body) + top, mid := strings.Index(body, "Top rule"), strings.Index(body, "Mid rule") + if top < 0 || mid < 0 || mid < top { + t.Fatalf("want both instructions in source order; body=%s", upstream.body) + } +} + +// Strict cloaking drops every caller system instruction. A translated +// mid-session developer item must not survive it as an in-place turn. +func TestResponsesMidSessionSystemItemDroppedByStrictCloaking(t *testing.T) { + const model = "claude-sonnet-5" + for name, call := range responsesPlacementCalls { + t.Run(name, func(t *testing.T) { + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + auth := responsesPlacementAuth(map[string]string{"cloak_strict_mode": "true"}) + if err := call(t, ex, upstream.context(t, nil), auth, model, responsesMidSessionPayload(model)); err != nil { + t.Fatalf("call error = %v", err) + } + if !upstream.called { + t.Fatal("expected the request to reach the upstream") + } + body := string(upstream.body) + if strings.Contains(body, "Top rule") || strings.Contains(body, "Mid rule") { + t.Fatalf("strict cloaking must drop every caller instruction; body=%s", upstream.body) + } + }) + } +} + +// An uncloaked request forwards the caller's betas, so it never gains the +// mid-conversation system beta and keeps both instructions top-level. +func TestResponsesMidSessionSystemItemUncloakedStaysTopLevel(t *testing.T) { + const model = "claude-sonnet-5" + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + auth := responsesPlacementAuth(map[string]string{"cloak_mode": "never"}) + if err := responsesPlacementCalls["execute"](t, ex, upstream.context(t, nil), auth, model, responsesMidSessionPayload(model)); err != nil { + t.Fatalf("Execute() error = %v", err) + } + if texts := claudeSystemTurnTexts(upstream.body); len(texts) != 0 { + t.Fatalf("uncloaked body carries system turns %q; body=%s", texts, upstream.body) + } + system := ojson.GetBytes(upstream.body, "system.#.text").Raw + if system != `["Top rule","Mid rule"]` { + t.Fatalf("system = %s, want both instructions in source order; body=%s", system, upstream.body) + } +} + +// An unsupported system part keeps failing explicitly on the top-level path, +// even when a text-only developer item would otherwise stay in place. With the +// rebuild opt-in, an in-place turn would be folded through a text-only filter +// that drops the unsupported part, so the opt-in must keep the top-level path. +func TestResponsesUnsupportedSystemPartStillRejected(t *testing.T) { + for _, test := range []struct { + name string + model string + attributes map[string]string + }{ + {name: "legacy model", model: "claude-sonnet-4-5"}, + {name: "supported model", model: "claude-sonnet-5"}, + {name: "rebuild opt-in", model: "claude-sonnet-5", attributes: map[string]string{"rebuild_mid_system_message": "true"}}, + } { + t.Run(test.name, func(t *testing.T) { + model := test.model + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + payload := []byte(`{"model":"` + model + `","input":[` + + `{"role":"user","content":[{"type":"input_text","text":"hi"}]},` + + `{"role":"developer","content":[{"type":"input_image","image_url":"https://example.com/a.png"}]},` + + `{"role":"developer","content":[{"type":"input_text","text":"Mid rule"}]},` + + `{"role":"user","content":[{"type":"input_text","text":"go"}]}]}`) + err := responsesPlacementCalls["execute"](t, ex, upstream.context(t, nil), responsesPlacementAuth(test.attributes), model, payload) + if err == nil { + t.Fatalf("Execute() error = nil, want the unsupported part rejected; body=%s", upstream.body) + } + if upstream.called { + t.Fatalf("upstream must not be called; body=%s", upstream.body) + } + var blockErr claudeCallerSystemBlockError + if !errors.As(err, &blockErr) || blockErr.StatusCode() != http.StatusBadRequest { + t.Fatalf("error = %v, want the unsupported system block rejected with a 400", err) + } + }) + } +} + +// Local estimation sends nothing upstream, so it keeps every system item at the +// top level, where its request validation accepts them, even for a credential +// whose Messages requests would carry an in-place turn. +func TestResponsesMidSessionSystemItemLocalCountTokensKeepsTopLevel(t *testing.T) { + const model = "claude-sonnet-5" + for _, attributes := range []map[string]string{ + {"base_url": "https://example.invalid"}, + {"base_url": "https://example.invalid", "cloak_strict_mode": "true"}, + } { + upstream := &midSystemUpstream{} + ex := NewClaudeExecutor(midSystemConfig()) + auth := responsesPlacementAuth(attributes) + resp, err := ex.CountTokens(upstream.context(t, nil), auth, cliproxyexecutor.Request{Model: model, Payload: responsesMidSessionPayload(model)}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("CountTokens(%v) error = %v, want a local estimate", attributes, err) + } + if upstream.called { + t.Fatalf("CountTokens(%v) must estimate locally for a custom base URL", attributes) + } + if len(resp.Payload) == 0 { + t.Fatalf("CountTokens(%v) returned an empty payload", attributes) + } + } +} + +// A payload rule that switches the model after translation is outside the +// placement decision. The existing validation answers it explicitly instead of +// sending a turn the model rejects. +func TestResponsesMidSessionSystemItemLateLegacyOverrideRejected(t *testing.T) { + upstream := &midSystemUpstream{} + cfg := midSystemConfig() + cfg.Payload.Override = []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "*"}}, + Params: map[string]any{"model": "claude-haiku-4-5-20251001"}, + }} + ex := NewClaudeExecutor(cfg) + err := responsesPlacementCalls["execute"](t, ex, upstream.context(t, nil), midSystemAuth(), "claude-sonnet-5", responsesMidSessionPayload("claude-sonnet-5")) + assertMidSystemRejected(t, err, upstream) +} diff --git a/internal/runtime/executor/claude_thinking_replay_test.go b/internal/runtime/executor/claude_thinking_replay_test.go index 9968b8e22a4..546a1a3a0cb 100644 --- a/internal/runtime/executor/claude_thinking_replay_test.go +++ b/internal/runtime/executor/claude_thinking_replay_test.go @@ -33,17 +33,17 @@ func claudeReplayTestAuth(baseURL string) *cliproxyauth.Auth { func claudeReplayTestRequest(payload []byte, sessionID string, isCompat bool, source sdktranslator.Format) (cliproxyexecutor.Request, cliproxyexecutor.Options) { return cliproxyexecutor.Request{ - Model: "claude-synthetic-4772", - Payload: payload, - Metadata: map[string]any{ - claudeReplayResolvedModelInfoKey: ®istry.ModelInfo{IsCompat: isCompat}, - }, - }, cliproxyexecutor.Options{ - SourceFormat: source, - Metadata: map[string]any{ - cliproxyexecutor.ExecutionSessionMetadataKey: sessionID, - }, - } + Model: "claude-synthetic-4772", + Payload: payload, + Metadata: map[string]any{ + claudeReplayResolvedModelInfoKey: ®istry.ModelInfo{IsCompat: isCompat}, + }, + }, cliproxyexecutor.Options{ + SourceFormat: source, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: sessionID, + }, + } } func TestClaudeThinkingReplayEnabledRequiresCompatClaudeAPIKey(t *testing.T) { diff --git a/internal/runtime/executor/claude_user_agent_extensions_test.go b/internal/runtime/executor/claude_user_agent_extensions_test.go new file mode 100644 index 00000000000..760c7dcef72 --- /dev/null +++ b/internal/runtime/executor/claude_user_agent_extensions_test.go @@ -0,0 +1,136 @@ +package executor + +import ( + "context" + "io" + "net/http" + "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 TestClaudeExecutorUserAgentFieldsPreserveNativePayload(t *testing.T) { + for _, testCase := range []struct { + name string + stabilize bool + }{ + {name: "default"}, + {name: "stabilized", stabilize: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + testClaudeExecutorUserAgentFields(t, testCase.stabilize) + }) + } +} + +func testClaudeExecutorUserAgentFields(t *testing.T, stabilize bool) { + t.Helper() + for _, path := range []string{"messages", "stream", "count_tokens"} { + t.Run(path, func(t *testing.T) { + doc := ojson.NewDoc([]byte(`{"model":"claude-sonnet-4-6","system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.312; cc_entrypoint=cli; cch=00000;"},{"type":"text","text":"Stable caller instructions","cache_control":{"type":"ephemeral","ttl":"1h"}}],"messages":[{"role":"user","content":[{"type":"text","text":"Check the task","cache_control":{"type":"ephemeral","ttl":"1h"}}]}],"tools":[{"name":"Read","description":"Read a file","input_schema":{"type":"object","properties":{}}}]}`)) + if path != "count_tokens" { + if errSet := doc.Set("metadata.user_id", claudeNativeHelperUserID); errSet != nil { + t.Fatal(errSet) + } + } else if errDelete := doc.Delete("system"); errDelete != nil { + t.Fatal(errDelete) + } + payload := doc.Bytes() + var nativeSystem string + for _, testCase := range []struct { + userAgent, stabilizedUserAgent 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)"}, + } { + t.Run(testCase.userAgent, func(t *testing.T) { + var upstreamBody []byte + var upstreamHeaders http.Header + transport := roundTripperFunc(func(request *http.Request) (*http.Response, error) { + body, errRead := io.ReadAll(request.Body) + if errRead != nil { + return nil, errRead + } + upstreamBody, upstreamHeaders = body, request.Header.Clone() + contentType := "application/json" + responseBody := `{"id":"msg_ua_fields","type":"message","role":"assistant","model":"claude-sonnet-4-6","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}` + if path == "count_tokens" { + responseBody = `{"input_tokens":1}` + } else if path == "stream" { + contentType = "text/event-stream" + responseBody = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_ua_fields\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-6\",\"content\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" + } + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {contentType}}, Body: io.NopCloser(strings.NewReader(responseBody)), Request: request}, nil + }) + 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("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{ + UserAgent: "claude-cli/2.1.353 (external, cli)", StabilizeDeviceProfile: &stabilize, + }}) + request := cliproxyexecutor.Request{Model: "claude-sonnet-4-6", Payload: payload} + options := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, OriginalRequest: payload, Headers: headers} + switch path { + case "messages": + if _, errExecute := executor.Execute(ctx, auth, request, options); errExecute != nil { + t.Fatal(errExecute) + } + case "stream": + result, errStream := executor.ExecuteStream(ctx, auth, request, options) + if errStream != nil { + t.Fatal(errStream) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + case "count_tokens": + if _, errCount := executor.CountTokens(ctx, auth, request, options); errCount != nil { + t.Fatal(errCount) + } + } + if upstreamBody == nil { + t.Fatal("upstream was not called") + } + for _, field := range []string{"messages", "tools"} { + if got, want := ojson.GetBytes(upstreamBody, field).Raw, ojson.GetBytes(payload, field).Raw; got != want { + t.Fatalf("native %s changed\n got: %s\nwant: %s", field, got, want) + } + } + // Signing may update billing, but both UAs must produce the same prefix. + system := ojson.GetBytes(upstreamBody, "system") + if path == "count_tokens" { + if system.Exists() { + t.Fatalf("native count_tokens gained a system prompt: %s", system.Raw) + } + } else if len(system.Array()) != 2 || system.Get("1").Raw != ojson.GetBytes(payload, "system.1").Raw { + t.Fatalf("native system blocks were rebuilt: %s", system.Raw) + } + if nativeSystem == "" { + nativeSystem = system.Raw + } else if system.Raw != nativeSystem { + t.Fatalf("extended UA changed the native prefix\n got: %s\nwant: %s", system.Raw, nativeSystem) + } + wantUserAgent := testCase.userAgent + if stabilize { + wantUserAgent = testCase.stabilizedUserAgent + } + if got := upstreamHeaders.Get("User-Agent"); got != wantUserAgent { + t.Fatalf("upstream UA = %q, want %q", got, wantUserAgent) + } + }) + } + }) + } +} diff --git a/internal/runtime/executor/codex_executor_auth.go b/internal/runtime/executor/codex_executor_auth.go index 6d630e4cb3f..fdaccd3d078 100644 --- a/internal/runtime/executor/codex_executor_auth.go +++ b/internal/runtime/executor/codex_executor_auth.go @@ -14,6 +14,7 @@ import ( codexauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + sdkauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/auth" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" log "github.com/sirupsen/logrus" @@ -140,7 +141,11 @@ func resolveCodexAuthDiskPath(cfg *config.Config, auth *cliproxyauth.Auth) (stri return "", fmt.Errorf("codex executor: stat credential file %s: %w", candidate, errStat) } if !info.IsDir() { - return candidate, nil + canonical, errCanonical := filepath.EvalSymlinks(candidate) + if errCanonical != nil { + return "", fmt.Errorf("codex executor: resolve credential path: %w", errCanonical) + } + return filepath.Abs(canonical) } } return "", nil @@ -292,6 +297,10 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au needSIWC = hasSIWC } } + // Explicit refresh requests must also respect the server's renewal window. + if earliest, okEarliest := auth.SIWCEarliestRefreshTime(); okEarliest && now.Before(earliest) { + needSIWC = false + } suppressedRejectedOuter := false if !diskRotatedOuter && auth.RejectedRefreshToken != "" && outerRefreshToken == strings.TrimSpace(auth.RejectedRefreshToken) { if needOuter { @@ -329,6 +338,15 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au } svc := codexauth.NewCodexAuthWithProxyURL(cfg, auth.ProxyURL) + if diskPath != "" && (cfg == nil || cfg.AuthDir == "") { + svc.SetRefreshJournal(&codexauth.FileRefreshJournal{Dir: filepath.Dir(diskPath)}) + } + store := sdkauth.GetTokenStore() + if _, isFileStore := store.(*sdkauth.FileTokenStore); !isFileStore { + if journal, okJournal := store.(codexauth.RefreshJournal); okJournal { + svc.SetRefreshJournal(journal) + } + } outerRefreshed := false var outerErr error @@ -433,6 +451,7 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au mergedSIWC = make(map[string]any) } delete(mergedSIWC, "expires_at") + delete(mergedSIWC, "earliest_refresh_at") for k, v := range siwcTD.ToMap() { mergedSIWC[k] = v } @@ -497,6 +516,9 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au } auth.Metadata = mergedDisk } + if siwcRefreshed { + log.Debug("SIWC token renewal completed") + } return auth, nil } diff --git a/internal/runtime/executor/codex_executor_auth_test.go b/internal/runtime/executor/codex_executor_auth_test.go index 4046cf937d9..aebb4d25d21 100644 --- a/internal/runtime/executor/codex_executor_auth_test.go +++ b/internal/runtime/executor/codex_executor_auth_test.go @@ -1,6 +1,7 @@ package executor import ( + "bytes" "context" "crypto/tls" "encoding/json" @@ -12,6 +13,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -58,27 +60,34 @@ func startCodexMockOAuthServers(t *testing.T, idToken string) (proxyURL string, mockProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodConnect { - destConn, errDial := net.Dial("tcp", mockAuthServer.Listener.Addr().String()) - if errDial != nil { - http.Error(w, errDial.Error(), http.StatusServiceUnavailable) - return - } - w.WriteHeader(http.StatusOK) hijacker, ok := w.(http.Hijacker) if !ok { - _ = destConn.Close() http.Error(w, "Hijacking not supported", http.StatusInternalServerError) return } - clientConn, _, errHijack := hijacker.Hijack() + clientConn, clientBuf, errHijack := hijacker.Hijack() if errHijack != nil { + return + } + destConn, errDial := net.Dial("tcp", mockAuthServer.Listener.Addr().String()) + if errDial != nil { + _, _ = clientConn.Write([]byte("HTTP/1.1 503 Service Unavailable\r\n\r\n")) + _ = clientConn.Close() + return + } + if _, errWrite := clientConn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n")); errWrite != nil { _ = destConn.Close() + _ = clientConn.Close() return } + var clientReader io.Reader = clientConn + if clientBuf != nil && clientBuf.Reader != nil { + clientReader = clientBuf.Reader + } go func() { defer func() { _ = destConn.Close() }() defer func() { _ = clientConn.Close() }() - _, _ = io.Copy(destConn, clientConn) + _, _ = io.Copy(destConn, clientReader) }() go func() { defer func() { _ = destConn.Close() }() @@ -110,7 +119,7 @@ func TestCodexExecutorRefresh_MissingPlanTypeDefaultsToFree(t *testing.T) { proxyURL, teardown := startCodexMockOAuthServers(t, idTokenWithoutPlan) defer teardown() - cfg := &config.Config{} + cfg := &config.Config{AuthDir: t.TempDir()} executor := NewCodexExecutor(cfg) storage := &codexauth.CodexTokenStorage{ @@ -186,7 +195,7 @@ func TestCodexExecutorRefresh_ExtractsPlanTypeWhenPresent(t *testing.T) { proxyURL, teardown := startCodexMockOAuthServers(t, idTokenWithPlan) defer teardown() - cfg := &config.Config{} + cfg := &config.Config{AuthDir: t.TempDir()} executor := NewCodexExecutor(cfg) storage := &codexauth.CodexTokenStorage{ @@ -262,7 +271,7 @@ func TestCodexExecutorRefresh_EmptyAttributesSnapshotIsolation(t *testing.T) { proxyURL, teardown := startCodexMockOAuthServers(t, idTokenWithoutPlan) defer teardown() - cfg := &config.Config{} + cfg := &config.Config{AuthDir: t.TempDir()} executor := NewCodexExecutor(cfg) emptyAttrs := map[string]string{} @@ -279,7 +288,7 @@ func TestCodexExecutorRefresh_EmptyAttributesSnapshotIsolation(t *testing.T) { ProxyURL: proxyURL, Attributes: emptyAttrs, Metadata: map[string]any{ - "refresh_token": "valid-refresh-token", + "refresh_token": "valid-refresh-token-empty-attrs", }, } @@ -294,6 +303,9 @@ func TestCodexExecutorRefresh_EmptyAttributesSnapshotIsolation(t *testing.T) { if got := refreshed.Attributes["plan_type"]; got != "free" { t.Errorf("refreshed.Attributes[plan_type] = %q, want free", got) } + if got := refreshed.Metadata["account_id"]; got != "acc-codex-empty-attrs" { + t.Errorf("refreshed.Metadata[account_id] = %v, want acc-codex-empty-attrs", got) + } // Verify original empty map was not modified if _, exists := emptyAttrs["plan_type"]; exists { t.Errorf("original empty Attributes map was modified, plan_type = %q", emptyAttrs["plan_type"]) @@ -340,26 +352,34 @@ func startDualCodexAndSIWCMockOAuthServer(t *testing.T, idToken string, failSIWC mockProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodConnect { - destConn, errDial := net.Dial("tcp", mockAuthServer.Listener.Addr().String()) - if errDial != nil { - http.Error(w, errDial.Error(), http.StatusServiceUnavailable) - return - } - w.WriteHeader(http.StatusOK) hijacker, ok := w.(http.Hijacker) if !ok { - _ = destConn.Close() + http.Error(w, "Hijacking not supported", http.StatusInternalServerError) return } - clientConn, _, errHijack := hijacker.Hijack() + clientConn, clientBuf, errHijack := hijacker.Hijack() if errHijack != nil { + return + } + destConn, errDial := net.Dial("tcp", mockAuthServer.Listener.Addr().String()) + if errDial != nil { + _, _ = clientConn.Write([]byte("HTTP/1.1 503 Service Unavailable\r\n\r\n")) + _ = clientConn.Close() + return + } + if _, errWrite := clientConn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n")); errWrite != nil { _ = destConn.Close() + _ = clientConn.Close() return } + var clientReader io.Reader = clientConn + if clientBuf != nil && clientBuf.Reader != nil { + clientReader = clientBuf.Reader + } go func() { defer func() { _ = destConn.Close() }() defer func() { _ = clientConn.Close() }() - _, _ = io.Copy(destConn, clientConn) + _, _ = io.Copy(destConn, clientReader) }() go func() { defer func() { _ = destConn.Close() }() @@ -389,7 +409,7 @@ func TestCodexExecutorRefreshSelectiveAndDual(t *testing.T) { proxyURL, outerCount, siwcCount, teardown := startDualCodexAndSIWCMockOAuthServer(t, idToken, &failSIWC) defer teardown() - executor := NewCodexExecutor(&config.Config{}) + executor := NewCodexExecutor(&config.Config{AuthDir: t.TempDir()}) now := time.Now() // 1. Only SIWC needs refresh (outer expires in 5 days, SIWC expires in 2 minutes) @@ -1115,7 +1135,7 @@ func TestCodexExecutorRefreshHomeModeRefreshesSIWCWhenOnlySIWCDue(t *testing.T) defer teardown() now := time.Now().UTC() - cfg := &config.Config{} + cfg := &config.Config{AuthDir: t.TempDir()} cfg.Home.Enabled = true executor := NewCodexExecutor(cfg) @@ -1698,3 +1718,120 @@ func TestCodexExecutorForceRefreshRejectedRefreshTokenWithoutDiskRotation_Issue7 t.Fatal("expected forced Refresh on unchanged rejected refresh_token without disk rotation to return an error, got nil") } } + +type bufferedConnectClientHelloConn struct { + net.Conn + mu sync.Mutex + connectBuf []byte + realRespRead bool + synthResp bytes.Reader +} + +func (c *bufferedConnectClientHelloConn) Write(p []byte) (int, error) { + c.mu.Lock() + if !c.realRespRead && (len(c.connectBuf) == 0 && strings.HasPrefix(string(p), "CONNECT ") || len(c.connectBuf) > 0 && !bytes.HasSuffix(c.connectBuf, []byte("\r\n\r\n"))) { + c.connectBuf = append(c.connectBuf, p...) + if bytes.HasSuffix(c.connectBuf, []byte("\r\n\r\n")) { + c.synthResp.Reset([]byte("HTTP/1.1 200 Connection established\r\n\r\n")) + } + c.mu.Unlock() + return len(p), nil + } + if len(c.connectBuf) > 0 && !c.realRespRead && len(p) > 0 { + c.realRespRead = true + combined := append(append([]byte(nil), c.connectBuf...), p[0]) + c.mu.Unlock() + if _, errWriteFirst := c.Conn.Write(combined); errWriteFirst != nil { + return 0, errWriteFirst + } + var window [4]byte + for { + copy(window[:3], window[1:4]) + if _, errRead := io.ReadFull(c.Conn, window[3:4]); errRead != nil { + return 1, errRead + } + if string(window[:]) == "\r\n\r\n" { + break + } + } + if len(p) == 1 { + return 1, nil + } + n, errWriteRest := c.Conn.Write(p[1:]) + return n + 1, errWriteRest + } + c.mu.Unlock() + return c.Conn.Write(p) +} + +func (c *bufferedConnectClientHelloConn) Read(p []byte) (int, error) { + c.mu.Lock() + if c.synthResp.Len() > 0 { + n, errRead := c.synthResp.Read(p) + c.mu.Unlock() + return n, errRead + } + c.mu.Unlock() + return c.Conn.Read(p) +} + +func TestCodexMockOAuthProxyPreservesBufferedClientHello(t *testing.T) { + idToken := makeTestCodexRefreshJWT("plus", "acc-buffered-hello") + installBufferedDial := func(t *testing.T) { + t.Helper() + baseTransport, ok := http.DefaultTransport.(*http.Transport) + if !ok || baseTransport == nil { + t.Fatal("expected *http.Transport as DefaultTransport") + } + baseTransport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + rawConn, errDial := (&net.Dialer{}).DialContext(ctx, network, addr) + if errDial != nil { + return nil, errDial + } + return &bufferedConnectClientHelloConn{Conn: rawConn}, nil + } + } + + t.Run("single", func(t *testing.T) { + proxyURL, teardown := startCodexMockOAuthServers(t, idToken) + defer teardown() + installBufferedDial(t) + + executor := NewCodexExecutor(&config.Config{AuthDir: t.TempDir()}) + auth := &cliproxyauth.Auth{ + ID: "codex-buffered-hello-single", + Provider: "codex", + ProxyURL: proxyURL, + Metadata: map[string]any{ + "type": "codex", + "access_token": "old-access", + "refresh_token": "valid-refresh-buffered-hello-single", + }, + } + if _, errRefresh := executor.Refresh(context.Background(), auth); errRefresh != nil { + t.Fatalf("executor.Refresh error = %v", errRefresh) + } + }) + + t.Run("dual", func(t *testing.T) { + failSIWC := false + proxyURL, _, _, teardown := startDualCodexAndSIWCMockOAuthServer(t, idToken, &failSIWC) + defer teardown() + installBufferedDial(t) + + executor := NewCodexExecutor(&config.Config{AuthDir: t.TempDir()}) + auth := &cliproxyauth.Auth{ + ID: "codex-buffered-hello-dual", + Provider: "codex", + ProxyURL: proxyURL, + Metadata: map[string]any{ + "type": "codex", + "access_token": "old-access", + "refresh_token": "valid-refresh-buffered-hello-dual", + }, + } + if _, errRefresh := executor.Refresh(context.Background(), auth); errRefresh != nil { + t.Fatalf("executor.Refresh error = %v", errRefresh) + } + }) +} diff --git a/internal/runtime/executor/codex_executor_claude_output_budget_test.go b/internal/runtime/executor/codex_executor_claude_output_budget_test.go index 871d5393830..6afbf4f4208 100644 --- a/internal/runtime/executor/codex_executor_claude_output_budget_test.go +++ b/internal/runtime/executor/codex_executor_claude_output_budget_test.go @@ -89,7 +89,9 @@ func TestCodexClaudeOutputBudgetBoundariesAndPayloadRules(t *testing.T) { {name: "fallback field", field: `,"max_output_tokens":1`, want: "16"}, {name: "zero remains invalid", field: `,"max_tokens":0`, want: "0"}, {name: "negative remains invalid", field: `,"max_tokens":-1`, want: "-1"}, - {name: "fractional input is not normalized", field: `,"max_tokens":1.5`, want: "1"}, + // The Claude translator truncates 1.5 to the integer 1 before the executor + // sees it, so the translated payload carries a sub-minimum integer budget. + {name: "fractional input truncated by translator", field: `,"max_tokens":1.5`, want: "16"}, {name: "non-numeric input", field: `,"max_tokens":"1"`}, { name: "override matches normalized value", field: `,"max_tokens":1`, want: "8", @@ -99,7 +101,7 @@ func TestCodexClaudeOutputBudgetBoundariesAndPayloadRules(t *testing.T) { name: "final filter wins", field: `,"max_tokens":1`, payload: config.PayloadConfig{Filter: []config.PayloadFilterRule{{Models: modelRule, Params: []string{"max_output_tokens"}}}}, }, - {name: "native responses unchanged", field: `,"max_output_tokens":1`, from: sdktranslator.FormatOpenAIResponse, want: "1"}, + {name: "native responses normalized", field: `,"max_output_tokens":1`, from: sdktranslator.FormatOpenAIResponse, want: "16"}, } for _, tc := range cases { for _, useWebsocket := range []bool{false, true} { diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go index 31b4c909ce8..dc3f3c5cbb5 100644 --- a/internal/runtime/executor/codex_executor_execute.go +++ b/internal/runtime/executor/codex_executor_execute.go @@ -58,14 +58,14 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re } originalTranslated = helps.SetStringIfDifferent(originalTranslated, "model", baseModel) doc := ojson.NewDoc(body) - normalizeCodexClaudeOutputBudgetDoc(doc, from, req.Payload) + normalizeCodexOutputBudgetDoc(doc) helps.SetStringIfDifferentDoc(doc, "model", baseModel) if err = helps.ApplyRequestThinkingDoc(doc, req, opts, from.String(), to.String(), e.Identifier(), updatesChanged); err != nil { return resp, err } - ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslated, req, opts)) + ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslated, req, opts, auth)) helps.SetStringIfDifferentDoc(doc, "model", baseModel) helps.SetBoolIfDifferentDoc(doc, "stream", true) _ = doc.Delete("previous_response_id") @@ -80,7 +80,11 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if !useSIWC { normalizeCodexInstructionsDoc(doc, helps.IsNativeCodexRequest(req.Payload, opts)) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - ensureImageGenerationToolDoc(doc, baseModel, auth, opts.Headers) + checkModel := req.Model + if strings.TrimSpace(checkModel) == "" { + checkModel = baseModel + } + ensureImageGenerationToolDoc(doc, checkModel, auth, opts.Headers) if !codexAuthUsesAPIKey(auth) { normalizeCodexImageGenerationToolsDoc(doc) } diff --git a/internal/runtime/executor/codex_executor_imagegen_test.go b/internal/runtime/executor/codex_executor_imagegen_test.go index 405a2d158db..a43243ee9bd 100644 --- a/internal/runtime/executor/codex_executor_imagegen_test.go +++ b/internal/runtime/executor/codex_executor_imagegen_test.go @@ -11,6 +11,8 @@ 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" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" @@ -152,8 +154,14 @@ func TestEnsureImageGenerationTool_ResponsesLiteFalseMetadataStillInjectsTool(t liteHeaders := make(http.Header) liteHeaders.Set(codexResponsesLiteHeader, "false") liteResult := ensureImageGenerationTool(liteBody, liteModel, nil, liteHeaders) - if ojson.GetBytes(liteResult, "tools").Exists() { - t.Fatalf("for %s: expected no auto-injected tools even when responses-lite is explicitly false, got %s", liteModel, liteResult) + if got := ojson.GetBytes(liteResult, "tools.0.type").String(); got != "image_generation" { + t.Fatalf("for %s: tools.0.type = %q, want image_generation when responses-lite is false and reasoning is absent; body=%s", liteModel, got, liteResult) + } + + reasoningBody := []byte(`{"model":"` + liteModel + `","reasoning":{"effort":"medium"},"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"false"},"input":"hello"}`) + reasoningResult := ensureImageGenerationTool(reasoningBody, liteModel, nil, liteHeaders) + if ojson.GetBytes(reasoningResult, "tools").Exists() { + t.Fatalf("for %s: expected no auto-injected tools when reasoning.effort is present, got %s", liteModel, reasoningResult) } } } @@ -300,7 +308,7 @@ func TestEnsureImageGenerationTool_FreeCodexAuthDoesNotInjectTool(t *testing.T) } func TestEnsureImageGenerationTool_ResponsesLiteModelsDoNotInjectToolByDefault(t *testing.T) { - for _, model := range []string{ + liteModels := []string{ "gpt-6-astra", "gpt-6-sol", "gpt-6.1-sol", @@ -312,21 +320,56 @@ func TestEnsureImageGenerationTool_ResponsesLiteModelsDoNotInjectToolByDefault(t "gpt-5.6-luna", "gpt-reserve", "codex-auto-review", + } + + for _, model := range liteModels { + t.Run("standard_without_reasoning/"+model, func(t *testing.T) { + body := []byte(`{"model":"` + model + `","input":"hello"}`) + result := ensureImageGenerationTool(body, model, nil, nil) + if got := ojson.GetBytes(result, "tools.0.type").String(); got != "image_generation" { + t.Fatalf("for %s without reasoning: tools.0.type = %q, want image_generation; body=%s", model, got, result) + } + + withFunctionTool := []byte(`{"model":"` + model + `","input":"Generate a photorealistic image of a mountain lake at sunset.","tools":[{"type":"function","name":"get_weather","description":"Get current weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}`) + withFunctionResult := ensureImageGenerationTool(withFunctionTool, model, nil, nil) + if got := ojson.GetBytes(withFunctionResult, "tools.1.type").String(); got != "image_generation" { + t.Fatalf("for %s with function tool and no reasoning: tools.1.type = %q, want image_generation; body=%s", model, got, withFunctionResult) + } + }) + + t.Run("with_reasoning_effort/"+model, func(t *testing.T) { + body := []byte(`{"model":"` + model + `","reasoning":{"effort":"medium"},"input":"hello"}`) + result := ensureImageGenerationTool(body, model, nil, nil) + if ojson.GetBytes(result, "tools").Exists() { + t.Fatalf("expected no auto-injected tools for %s with reasoning.effort, got %s", model, result) + } + }) + + t.Run("with_configuration_update_reasoning/"+model, func(t *testing.T) { + body := []byte(`{"model":"` + model + `","input":[{"type":"configuration_update","reasoning":{"effort":"high"}},{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) + result := ensureImageGenerationTool(body, model, nil, nil) + if ojson.GetBytes(result, "tools").Exists() { + t.Fatalf("expected no auto-injected tools for %s with configuration_update reasoning, got %s", model, result) + } + }) + } + + for _, suffixedModel := range []string{ "codex/gpt-6-astra(medium)", "codex/gpt-5.6-sol(high)", } { - t.Run(model, func(t *testing.T) { - body := []byte(`{"model":"` + model + `","input":"hello"}`) - result := ensureImageGenerationTool(body, model, nil, nil) + t.Run("with_thinking_suffix/"+suffixedModel, func(t *testing.T) { + body := []byte(`{"model":"` + suffixedModel + `","input":"hello"}`) + result := ensureImageGenerationTool(body, suffixedModel, nil, nil) if ojson.GetBytes(result, "tools").Exists() { - t.Fatalf("expected no auto-injected tools for %s, got %s", model, result) + t.Fatalf("expected no auto-injected tools for %s, got %s", suffixedModel, result) } }) } for _, nonLiteModel := range []string{"gpt-5.5", "gpt-5.4"} { t.Run("non-lite/"+nonLiteModel, func(t *testing.T) { - body := []byte(`{"model":"` + nonLiteModel + `","input":"hello"}`) + body := []byte(`{"model":"` + nonLiteModel + `","reasoning":{"effort":"medium"},"input":"hello"}`) result := ensureImageGenerationTool(body, nonLiteModel, nil, nil) if got := ojson.GetBytes(result, "tools.0.type").String(); got != "image_generation" { t.Fatalf("for %s: tools.0.type = %q, want image_generation; body=%s", nonLiteModel, got, result) @@ -335,6 +378,473 @@ func TestEnsureImageGenerationTool_ResponsesLiteModelsDoNotInjectToolByDefault(t } } +func TestCodexExecutorStandardResponsesInjectsImageGenerationWithoutReasoningAndSkipsWithReasoning(t *testing.T) { + for _, transport := range []string{"http", "websocket"} { + for _, stream := range []bool{false, true} { + streamLabel := "false" + if stream { + streamLabel = "true" + } + for _, tc := range []struct { + name string + model string + sourceFmt sdktranslator.Format + payload string + reqMetadata map[string]any + reqHeaders http.Header + authAttrs map[string]string + cfg config.Config + overrideHeaders map[string]string + wantTools int + wantImgIdx string + wantToolType string + wantEffort string + }{ + { + name: "standard_responses_with_function_tool_no_reasoning", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"Generate a photorealistic image of a mountain lake at sunset.","tools":[{"type":"function","name":"get_weather","description":"Get current weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}`, + wantTools: 2, + wantImgIdx: "tools.1.type", + }, + { + name: "standard_responses_bare_no_reasoning", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"Generate a photorealistic image of a mountain lake at sunset."}`, + wantTools: 1, + wantImgIdx: "tools.0.type", + }, + { + name: "standard_responses_with_reasoning_effort", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","reasoning":{"effort":"medium"},"input":"Generate a photorealistic image of a mountain lake at sunset."}`, + wantTools: 0, + }, + { + name: "standard_responses_with_model_thinking_suffix", + model: "gpt-6-astra(medium)", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra(medium)","input":"Generate a photorealistic image of a mountain lake at sunset."}`, + wantTools: 0, + }, + { + name: "openai_chat_completions_defaults_reasoning_effort", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAI, + payload: `{"model":"gpt-6-astra","messages":[{"role":"user","content":"Generate a photorealistic image of a mountain lake at sunset."}]}`, + wantTools: 0, + }, + { + name: "payload_default_sets_reasoning_effort_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Default: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{"reasoning.effort": "medium"}, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "payload_override_sets_reasoning_effort_with_function_tool_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello","tools":[{"type":"function","name":"get_weather","description":"Get current weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{"reasoning.effort": "high"}, + }, + }, + }, + }, + wantTools: 1, + wantImgIdx: "tools.0.type", + wantToolType: "function", + wantEffort: "high", + }, + { + name: "payload_conditional_exist_tools_sets_reasoning_effort_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra", Exist: []string{"tools"}}}, + Params: map[string]any{"reasoning.effort": "medium"}, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "payload_conditional_exist_prompt_cache_key_sets_reasoning_effort_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + reqMetadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "sess-prompt-cache-7384"}, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra", Exist: []string{"prompt_cache_key"}}}, + Params: map[string]any{"reasoning.effort": "medium"}, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "payload_override_switches_to_lite_model_with_reasoning_skips_auto_injection", + model: "gpt-5.4", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-5.4","reasoning":{"effort":"medium"},"input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-5.4"}}, + Params: map[string]any{"model": "gpt-6-astra"}, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "model_override_header_enables_responses_lite_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + overrideHeaders: map[string]string{codexResponsesLiteHeader: "true"}, + wantTools: 0, + }, + { + name: "payload_override_sets_reasoning_effort_and_tool_choice_auto_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{ + "reasoning.effort": "medium", + "tool_choice": "auto", + }, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "client_lite_header_overridden_false_by_auth_with_payload_rule_keeps_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + reqHeaders: http.Header{codexResponsesLiteHeader: []string{"true"}}, + authAttrs: map[string]string{"header:" + codexResponsesLiteHeader: "false"}, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{"service_tier": "default"}, + }, + }, + }, + }, + wantTools: 1, + wantImgIdx: "tools.0.type", + }, + { + name: "payload_default_image_tool_choice_overridden_to_auto_with_reasoning_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Default: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{"tool_choice": map[string]any{"type": "image_generation"}}, + }, + }, + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{ + "tool_choice": "auto", + "reasoning.effort": "medium", + }, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "auth_lite_header_true_overridden_false_by_initial_model_then_switched_to_model_without_override_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + authAttrs: map[string]string{"header:" + codexResponsesLiteHeader: "true"}, + overrideHeaders: map[string]string{codexResponsesLiteHeader: "false"}, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{"model": "gpt-5.5"}, + }, + }, + }, + }, + wantTools: 0, + }, + { + name: "payload_override_function_tool_then_filter_deletes_it_with_reasoning_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello","tools":[{"type":"function","name":"get_weather","description":"Get current weather","parameters":{"type":"object"}}]}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{ + "tools.0.description": "Updated weather", + "reasoning.effort": "medium", + }, + }, + }, + Filter: []config.PayloadFilterRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: []string{"tools.0"}, + }, + }, + }, + }, + wantTools: 0, + wantEffort: "medium", + }, + { + name: "payload_override_image_tool_then_filter_deletes_earlier_tool_with_reasoning_keeps_configured_image_tool", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello","tools":[{"type":"function","name":"get_weather","description":"Get current weather","parameters":{"type":"object"}}]}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{ + "tools.1.output_format": "webp", + "reasoning.effort": "medium", + }, + }, + }, + Filter: []config.PayloadFilterRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: []string{"tools.0"}, + }, + }, + }, + }, + wantTools: 1, + wantImgIdx: "tools.0.type", + wantEffort: "medium", + }, + { + name: "payload_override_raw_escaped_configuration_update_reasoning_skips_auto_injection", + model: "gpt-6-astra", + sourceFmt: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-6-astra","input":"hello"}`, + cfg: config.Config{ + Payload: config.PayloadConfig{ + OverrideRaw: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6-astra"}}, + Params: map[string]any{ + "input.-1": `{"type":"\u0063onfiguration_update","reasoning":{"effort":"medium"}}`, + }, + }, + }, + }, + }, + wantTools: 0, + }, + } { + t.Run(tc.name+"/"+transport+"/stream="+streamLabel, func(t *testing.T) { + capturedBody := make(chan []byte, 1) + completedEvent := []byte(`{"type":"response.completed","response":{"id":"resp_7384","object":"response","status":"completed","model":"gpt-6-astra","output":[],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}`) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if transport == "websocket" { + upgrader := websocket.Upgrader{} + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { + _ = conn.Close() + }() + _, msg, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read websocket message: %v", errRead) + return + } + capturedBody <- msg + if errWrite := conn.WriteMessage(websocket.TextMessage, completedEvent); errWrite != nil { + t.Errorf("write websocket message: %v", errWrite) + } + return + } + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read http body: %v", errRead) + return + } + capturedBody <- body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write(append(append([]byte("data: "), completedEvent...), '\n', '\n')) + })) + defer server.Close() + + cfg := tc.cfg + var exec cliproxyauth.ProviderExecutor = NewCodexExecutor(&cfg) + if transport == "websocket" { + exec = NewCodexWebsocketsExecutor(&cfg) + } + if len(tc.overrideHeaders) > 0 { + reg := registry.GetGlobalRegistry() + clientID := t.Name() + "/override-header" + reg.RegisterClient(clientID, "codex", []*registry.ModelInfo{{ + ID: tc.model, + Type: "codex", + Config: ®istry.ModelConfig{OverrideHeader: tc.overrideHeaders}, + }}) + t.Cleanup(func() { reg.UnregisterClient(clientID) }) + } + attrs := map[string]string{ + "base_url": server.URL, + "plan_type": "pro", + } + for k, v := range tc.authAttrs { + attrs[k] = v + } + auth := &cliproxyauth.Auth{ + ID: t.Name(), + Provider: "codex", + Attributes: attrs, + Metadata: map[string]any{ + "access_token": "test-token", + "account_id": "acct-1", + }, + } + req := cliproxyexecutor.Request{ + Model: tc.model, + Payload: []byte(tc.payload), + Metadata: tc.reqMetadata, + } + finalized := 0 + var headers http.Header + if tc.reqHeaders != nil { + headers = make(http.Header, len(tc.reqHeaders)) + for k, vals := range tc.reqHeaders { + for _, v := range vals { + headers.Add(k, v) + } + } + } + opts := cliproxyexecutor.Options{ + SourceFormat: tc.sourceFmt, + Stream: stream, + Headers: headers, + Metadata: map[string]any{ + cliproxyexecutor.FinalizedResponseFormatCallbackMetadataKey: func(string) { finalized++ }, + }, + } + evalsBefore := helps.PayloadConfigEvaluationCount() + var errExec error + if stream { + var res *cliproxyexecutor.StreamResult + res, errExec = exec.ExecuteStream(context.Background(), auth, req, opts) + if res != nil { + for chunk := range res.Chunks { + if chunk.Err != nil { + errExec = chunk.Err + } + } + } + } else { + _, errExec = exec.Execute(context.Background(), auth, req, opts) + } + + if gotEvals := helps.PayloadConfigEvaluationCount() - evalsBefore; gotEvals != 1 { + t.Fatalf("payload rule engine evaluated %d times, want 1", gotEvals) + } + if finalized != 1 { + t.Fatalf("payload finalizer evaluated %d times, want 1", finalized) + } + if errExec != nil { + t.Fatalf("execute error = %v", errExec) + } + gotBody := <-capturedBody + if tc.wantEffort != "" { + if gotEffort := ojson.GetBytes(gotBody, "reasoning.effort").String(); gotEffort != tc.wantEffort { + t.Fatalf("reasoning.effort = %q, want %q; body=%s", gotEffort, tc.wantEffort, gotBody) + } + } + tools := ojson.GetBytes(gotBody, "tools") + if tc.wantTools == 0 { + if tools.Exists() { + t.Fatalf("expected no tools, got %s (body=%s)", tools.Raw, gotBody) + } + return + } + if !tools.IsArray() || len(tools.Array()) != tc.wantTools { + t.Fatalf("tools count = %d, want %d; body=%s", len(tools.Array()), tc.wantTools, gotBody) + } + wantType := tc.wantToolType + if wantType == "" { + wantType = "image_generation" + } + if gotType := ojson.GetBytes(gotBody, tc.wantImgIdx).String(); gotType != wantType { + t.Fatalf("%s = %q, want %s; body=%s", tc.wantImgIdx, gotType, wantType, gotBody) + } + }) + } + } + } +} + func TestCodexExecutorHostedImageGenerationSIWCAndAPIKeyPreserveUpstreamModelAndParams(t *testing.T) { for _, credKind := range []string{"api_key", "siwc"} { for _, transport := range []string{"http", "websocket"} { diff --git a/internal/runtime/executor/codex_executor_output_budget_issue7304_test.go b/internal/runtime/executor/codex_executor_output_budget_issue7304_test.go new file mode 100644 index 00000000000..855d0dcdecc --- /dev/null +++ b/internal/runtime/executor/codex_executor_output_budget_issue7304_test.go @@ -0,0 +1,237 @@ +package executor + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "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/translator" + 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" +) + +// codexMinimumBudgetUpstream is a simplified lower-bound-only model of the ChatGPT +// Codex backend for this regression test: it records every body, rejects any numeric +// max_output_tokens below 16 (integer or fractional, so the floor is the only rule +// exercised), and accepts absent or non-numeric values. It does not model the real +// backend's full input validation. +func codexMinimumBudgetUpstream(t *testing.T) (*httptest.Server, <-chan []byte) { + t.Helper() + requests := make(chan []byte, 4) + completed := `{"type":"response.completed","response":{"id":"resp_1","object":"response","status":"completed","model":"gpt-5.6-sol","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hi"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + accepts := func(body []byte) bool { + value := ojson.GetBytes(body, "max_output_tokens") + return !value.Exists() || value.Type != ojson.Number || value.Int() >= 16 + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Upgrade") == "websocket" { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade: %v", errUpgrade) + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + _, body, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read websocket request: %v", errRead) + return + } + requests <- body + if !accepts(body) { + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","error":{"type":"invalid_request_error","message":"Invalid 'max_output_tokens': integer below minimum value. Expected a value >= 16."}}`)) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(completed)); errWrite != nil { + t.Errorf("write websocket response: %v", errWrite) + } + return + } + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read HTTP request: %v", errRead) + return + } + requests <- body + if !accepts(body) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"error":{"message":"Invalid 'max_output_tokens': integer below minimum value. Expected a value >= 16.","type":"invalid_request_error","param":"max_output_tokens"}}`) + return + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: "+completed+"\n\n") + })) + t.Cleanup(server.Close) + return server, requests +} + +// Issue 7304: operator-triggered inference verification must succeed against the +// real Codex minimum even though it uses a minimal synthetic prompt budget. +func TestCodexVerifyInferenceMinimumOutputBudget_Issue7304(t *testing.T) { + server, requests := codexMinimumBudgetUpstream(t) + + executor := NewCodexExecutor(&config.Config{}) + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + + ctx := cliproxyauth.WithSkipPersist(context.Background()) + auth := &cliproxyauth.Auth{ + ID: "codex-verify-7304.json", + FileName: "codex-verify-7304.json", + Provider: "codex", + Status: cliproxyauth.StatusActive, + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{ + "type": "codex", + "email": "verify-7304@example.com", + "access_token": "at-verify-7304", + "refresh_token": "rt-verify-7304", + "expired": time.Now().UTC().Add(24 * time.Hour).Format(time.RFC3339), + }, + } + if _, errReg := manager.Register(ctx, auth); errReg != nil { + t.Fatalf("Register auth: %v", errReg) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, "codex", []*registry.ModelInfo{{ID: "gpt-5.6-sol"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + + verified, errVerify := manager.VerifyInference(ctx, auth.ID, "gpt-5.6-sol") + if errVerify != nil { + t.Fatalf("VerifyInference must succeed against the Codex >=16 minimum, got error: %v", errVerify) + } + if verified == nil || !verified.InferenceVerified { + t.Fatalf("expected InferenceVerified=true after verification, got %+v", verified) + } + body := <-requests + if got := ojson.GetBytes(body, "max_output_tokens").Raw; got != "16" { + t.Fatalf("upstream max_output_tokens = %q, want 16; body=%s", got, body) + } +} + +// Issue 7304: small positive integer budgets from every Codex source format must be +// normalized to the upstream minimum before user payload rules; invalid budgets +// (0, negative, fractional, non-numeric) stay untouched and no budget is injected. +func TestCodexSmallOutputBudgetsAllSourceFormats_Issue7304(t *testing.T) { + modelRule := []config.PayloadModelRule{{Name: "gpt-5.6-sol", Protocol: "codex"}} + matchedRule := []config.PayloadModelRule{{ + Name: "gpt-5.6-sol", Protocol: "codex", Match: []map[string]any{{"max_output_tokens": 16}}, + }} + + cases := []struct { + name string + from sdktranslator.Format + payload string + want string + rules config.PayloadConfig + }{ + {name: "responses max_output_tokens 1", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":1,"input":"Hi"}`, want: "16"}, + {name: "responses max_output_tokens 15", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":15,"input":"Hi"}`, want: "16"}, + {name: "responses legacy max_completion_tokens 1", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_completion_tokens":1,"input":"Hi"}`, want: "16"}, + {name: "responses legacy max_tokens 1", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_tokens":1,"input":"Hi"}`, want: "16"}, + {name: "responses 16 unchanged", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":16,"input":"Hi"}`, want: "16"}, + {name: "responses 1024 unchanged", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":1024,"input":"Hi"}`, want: "1024"}, + {name: "responses zero untouched", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":0,"input":"Hi"}`, want: "0"}, + {name: "responses negative untouched", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":-5,"input":"Hi"}`, want: "-5"}, + {name: "responses fractional untouched", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":1.5,"input":"Hi"}`, want: "1.5"}, + {name: "responses integral float normalized", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":15.0,"input":"Hi"}`, want: "16"}, + {name: "responses scientific integral normalized", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":1e1,"input":"Hi"}`, want: "16"}, + {name: "responses string untouched", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","max_output_tokens":"1","input":"Hi"}`, want: `"1"`}, + {name: "responses absent", from: sdktranslator.FormatOpenAIResponse, payload: `{"model":"gpt-5.6-sol","input":"Hi"}`}, + {name: "chat max_tokens 1", from: sdktranslator.FormatOpenAI, payload: `{"model":"gpt-5.6-sol","max_tokens":1,"messages":[{"role":"user","content":"Hi"}]}`, want: "16"}, + {name: "chat max_completion_tokens 1", from: sdktranslator.FormatOpenAI, payload: `{"model":"gpt-5.6-sol","max_completion_tokens":1,"messages":[{"role":"user","content":"Hi"}]}`, want: "16"}, + {name: "chat max_output_tokens 1", from: sdktranslator.FormatOpenAI, payload: `{"model":"gpt-5.6-sol","max_output_tokens":1,"messages":[{"role":"user","content":"Hi"}]}`, want: "16"}, + {name: "claude max_tokens 1", from: sdktranslator.FormatClaude, payload: `{"model":"gpt-5.6-sol","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"Hi"}]}]}`, want: "16"}, + {name: "gemini maxOutputTokens 1", from: sdktranslator.FormatGemini, payload: `{"model":"gpt-5.6-sol","contents":[{"role":"user","parts":[{"text":"Hi"}]}],"generationConfig":{"maxOutputTokens":1}}`, want: "16"}, + {name: "gemini top-level max_tokens 1", from: sdktranslator.FormatGemini, payload: `{"model":"gpt-5.6-sol","contents":[{"role":"user","parts":[{"text":"Hi"}]}],"max_tokens":1}`, want: "16"}, + {name: "interactions max_output_tokens 1", from: sdktranslator.FormatInteractions, payload: `{"model":"gpt-5.6-sol","input":"Hi","max_output_tokens":1}`, want: "16"}, + {name: "interactions generation_config 1", from: sdktranslator.FormatInteractions, payload: `{"model":"gpt-5.6-sol","input":"Hi","generation_config":{"max_output_tokens":1}}`, want: "16"}, + { + name: "responses payload override still final", from: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-5.6-sol","max_output_tokens":1,"input":"Hi"}`, want: "8", + rules: config.PayloadConfig{Override: []config.PayloadRule{{Models: matchedRule, Params: map[string]any{"max_output_tokens": 8}}}}, + }, + { + name: "responses payload filter still final", from: sdktranslator.FormatOpenAIResponse, + payload: `{"model":"gpt-5.6-sol","max_output_tokens":1,"input":"Hi"}`, + rules: config.PayloadConfig{Filter: []config.PayloadFilterRule{{Models: modelRule, Params: []string{"max_output_tokens"}}}}, + }, + } + + for _, tc := range cases { + for _, useWebsocket := range []bool{false, true} { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/websocket=%t/stream=%t", tc.name, useWebsocket, stream), func(t *testing.T) { + server, requests := codexMinimumBudgetUpstream(t) + cfg := &config.Config{Payload: tc.rules} + var executor cliproxyauth.ProviderExecutor = NewCodexExecutor(cfg) + if useWebsocket { + executor = NewCodexWebsocketsExecutor(cfg) + } + req := cliproxyexecutor.Request{Model: "gpt-5.6-sol", Payload: []byte(tc.payload)} + opts := cliproxyexecutor.Options{SourceFormat: tc.from, Stream: stream} + auth := codexOAuthTestAuth(server.URL) + auth.ID = "codex-output-budget-7304-test" + + // wantAccept mirrors the simplified upstream model above: any + // numeric budget below 16 is rejected (the simulator does not + // distinguish integer from fractional), so accepts cover normalized, + // valid, absent, and non-numeric values. Accepted requests must + // succeed; rejected ones must surface a budget error. + wantAccept := true + if tc.want != "" { + if parsed, errParse := strconv.ParseFloat(tc.want, 64); errParse == nil && parsed < 16 { + wantAccept = false + } + } + var execErr error + if stream { + result, errStream := executor.ExecuteStream(context.Background(), auth, req, opts) + execErr = errStream + if errStream == nil { + for chunk := range result.Chunks { + if chunk.Err != nil && execErr == nil { + execErr = chunk.Err + } + } + } + } else { + _, execErr = executor.Execute(context.Background(), auth, req, opts) + } + if wantAccept && execErr != nil { + t.Fatalf("execution must succeed for an upstream-acceptable budget, got error: %v", execErr) + } + if !wantAccept && execErr == nil { + t.Fatal("execution must fail when the untouched budget stays below the upstream minimum") + } + if !wantAccept && execErr != nil && !strings.Contains(execErr.Error(), "max_output_tokens") { + t.Fatalf("rejected execution error must propagate the upstream budget error, got: %v", execErr) + } + select { + case body := <-requests: + if got := ojson.GetBytes(body, "max_output_tokens").Raw; got != tc.want { + t.Fatalf("upstream max_output_tokens = %q, want %q; body=%s", got, tc.want, body) + } + case <-time.After(10 * time.Second): + t.Fatal("upstream received no request within 10s") + } + }) + } + } + } +} diff --git a/internal/runtime/executor/codex_executor_request.go b/internal/runtime/executor/codex_executor_request.go index a87a36555be..1acd7d79495 100644 --- a/internal/runtime/executor/codex_executor_request.go +++ b/internal/runtime/executor/codex_executor_request.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "io" + "math" "net/http" neturl "net/url" "runtime" @@ -67,26 +68,29 @@ func translateCodexRequestPairWithUpdateIntentContext(ctx context.Context, from, return originalTranslated, body, changed, err } -// normalizeCodexClaudeOutputBudgetDoc bridges Claude's smaller valid budgets to Codex's minimum. -func normalizeCodexClaudeOutputBudgetDoc(doc *ojson.Doc, from sdktranslator.Format, payload []byte) { - if doc == nil || from != sdktranslator.FormatClaude { +// codexMinOutputTokens is the smallest integer max_output_tokens value the ChatGPT +// Codex backend accepts; smaller integers are rejected with invalid_request_error. +const codexMinOutputTokens = 16 + +// normalizeCodexOutputBudgetDoc raises small positive integer output budgets to the +// Codex minimum. It judges the translated payload (max_output_tokens) rather than the +// client payload so every source format and budget field (max_tokens, +// max_completion_tokens, generation_config, ...) is covered uniformly. Budgets that +// are missing, non-numeric, non-integral, or <= 0 are left untouched for the upstream +// to reject. This runs before user payload rules so overrides and filters stay final. +func normalizeCodexOutputBudgetDoc(doc *ojson.Doc) { + if doc == nil { return } - requested := ojson.GetBytesNoCopy(payload, "max_tokens") - if !requested.Exists() { - requested = ojson.GetBytesNoCopy(payload, "max_output_tokens") - } - if requested.Type != ojson.Number { + translated := doc.Get("max_output_tokens") + if translated.Type != ojson.Number { return } - requestedTokens, errParse := strconv.ParseInt(requested.Raw, 10, 64) - if errParse != nil || requestedTokens < 1 || requestedTokens >= 16 { + tokens, errFloat := strconv.ParseFloat(translated.Raw, 64) + if errFloat != nil || tokens < 1 || tokens >= codexMinOutputTokens || tokens != math.Trunc(tokens) { return } - translated := doc.Get("max_output_tokens") - if translated.Type == ojson.Number && translated.Int() > 0 && translated.Int() < 16 { - _ = doc.Set("max_output_tokens", 16) - } + _ = doc.Set("max_output_tokens", codexMinOutputTokens) } // PrepareRequest injects Codex credentials into the outgoing HTTP request. @@ -804,29 +808,61 @@ func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth return doc.Bytes() } -func ensureImageGenerationToolDoc(doc *ojson.Doc, baseModel string, auth *cliproxyauth.Auth, headers http.Header) bool { - if doc == nil || isCodexResponsesLiteRequestDoc(doc, headers) || isCodexResponsesLiteModel(baseModel) { +func codexRequestHasReasoningForHostedImageDoc(doc *ojson.Doc, model string) bool { + return helps.CodexRequestHasReasoningForHostedImageDoc(doc, model) +} + +func isCodexResponsesLiteRequestForAuthDoc(doc *ojson.Doc, modelName string, auth *cliproxyauth.Auth, headers http.Header) bool { + return helps.CodexResponsesLiteRequestForAuthDoc(doc, modelName, auth, headers) +} + +func codexToolsAlreadyContainImageGenDoc(doc *ojson.Doc) bool { + if doc == nil { return false } - if strings.HasSuffix(baseModel, "spark") { + tools := doc.Get("tools") + if !tools.Exists() || !tools.IsArray() { return false } - if isCodexFreePlanAuth(auth) { - return false + for _, t := range tools.Array() { + if t.Get("type").String() == "image_generation" || isImageGenerationFunctionTool(t) { + return true + } } + return false +} +func appendImageGenerationToolDoc(doc *ojson.Doc) bool { + if doc == nil { + return false + } tools := doc.Get("tools") if !tools.Exists() || !tools.IsArray() { return doc.SetRawBytes("tools", imageGenToolArrayJSON) == nil } - for _, t := range tools.Array() { - if t.Get("type").String() == "image_generation" || isImageGenerationFunctionTool(t) { - return false - } - } return doc.SetRawBytes("tools.-1", imageGenToolJSON) == nil } +func ensureImageGenerationToolDoc(doc *ojson.Doc, baseModel string, auth *cliproxyauth.Auth, headers http.Header) bool { + if doc == nil || isCodexFreePlanAuth(auth) || codexToolsAlreadyContainImageGenDoc(doc) { + return false + } + modelName := strings.TrimSpace(thinking.ParseSuffix(baseModel).ModelName) + if modelName == "" { + modelName = strings.TrimSpace(thinking.ParseSuffix(doc.Get("model").String()).ModelName) + } + if isCodexResponsesLiteRequestForAuthDoc(doc, modelName, auth, headers) { + return false + } + if isCodexResponsesLiteModel(modelName) && codexRequestHasReasoningForHostedImageDoc(doc, baseModel) { + return false + } + if strings.HasSuffix(strings.ToLower(modelName), "spark") { + return false + } + return appendImageGenerationToolDoc(doc) +} + func resolveCodexGPTImage2BaseModel(cfg *config.Config) string { if cfg == nil { return codexOpenAIImagesMainModel diff --git a/internal/runtime/executor/codex_executor_siwc_darwin_test.go b/internal/runtime/executor/codex_executor_siwc_darwin_test.go new file mode 100644 index 00000000000..cca10dabd1d --- /dev/null +++ b/internal/runtime/executor/codex_executor_siwc_darwin_test.go @@ -0,0 +1,16 @@ +package executor + +import ( + "context" + "os" + "testing" + + codexauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestMain(m *testing.M) { + codexauth.ReadSIWCSystemState = func(context.Context) ([]byte, error) { + return []byte("Current System Capabilities are: CPU Graphics Audio Network\nCurrent Power State: 4\n"), nil + } + os.Exit(m.Run()) +} diff --git a/internal/runtime/executor/codex_executor_siwc_renewal_test.go b/internal/runtime/executor/codex_executor_siwc_renewal_test.go new file mode 100644 index 00000000000..114d3122b63 --- /dev/null +++ b/internal/runtime/executor/codex_executor_siwc_renewal_test.go @@ -0,0 +1,115 @@ +package executor + +import ( + "context" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/MoeForge/moego/ojson" + codexauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestCodexExecutorSIWCRetainsRotatedTokensWhenScopeIsOmitted(t *testing.T) { + previousTransport := http.DefaultTransport + t.Cleanup(func() { http.DefaultTransport = previousTransport }) + cycle := 0 + http.DefaultTransport = roundTripperFunc(func(req *http.Request) (*http.Response, error) { + if errParse := req.ParseForm(); errParse != nil { + t.Fatal(errParse) + } + if got, want := req.Form.Get("refresh_token"), fmt.Sprintf("refresh-%d", cycle); got != want { + t.Fatalf("submitted refresh token = %q, want saved replacement %q", got, want) + } + cycle++ + body := fmt.Sprintf(`{"access_token":"access-%d","refresh_token":"refresh-%d","expires_in":3600}`, cycle, cycle) + if cycle == 1 { + doc := ojson.NewDocString(body) + if errSet := doc.Set("earliest_refresh_at", 946684800); errSet != nil { + t.Fatal(errSet) + } + body = doc.String() + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + }) + dir := t.TempDir() + path := filepath.Join(dir, "codex.json") + const grantedScope = "openid chatgpt.tokens.use.direct" + initial := map[string]any{ + "type": "codex", "use_siwc": true, + "siwc": map[string]any{ + "access_token": "access-0", "refresh_token": "refresh-0", "client_id": "oaiapp_test", + "scope": grantedScope, "scopes": strings.Fields(grantedScope), + "expires_at": time.Now().Add(-time.Hour).Format(time.RFC3339), + "earliest_refresh_at": "1999-01-01T00:00:00Z", + }, + } + if errWrite := codexauth.WriteJSONFileAtomic(path, initial); errWrite != nil { + t.Fatal(errWrite) + } + executor := NewCodexExecutor(&config.Config{AuthDir: dir}) + metadata := initial + for wantCycle := 1; wantCycle <= 2; wantCycle++ { + auth := &cliproxyauth.Auth{ID: "codex.json", FileName: "codex.json", Provider: "codex", Metadata: metadata} + if _, errRefresh := executor.Refresh(cliproxyauth.WithForceRefresh(context.Background()), auth); errRefresh != nil { + t.Fatalf("renewal %d: %v", wantCycle, errRefresh) + } + raw, errRead := os.ReadFile(path) + if errRead != nil { + t.Fatal(errRead) + } + saved := ojson.ParseBytes(raw) + if got, want := saved.Get("siwc.refresh_token").String(), fmt.Sprintf("refresh-%d", wantCycle); got != want { + t.Fatalf("persisted refresh token = %q, want %q", got, want) + } + if saved.Get("siwc.scope").String() != grantedScope || saved.Get("siwc.scopes.1").String() != codexauth.SIWCDirectTokenScope { + t.Fatal("previously granted permissions were lost") + } + window := saved.Get("siwc.earliest_refresh_at") + if wantCycle == 1 && (window.Type != ojson.Number || window.Int() != 946684800) { + t.Fatalf("new server renewal window was not persisted: %s", window.Raw) + } + if wantCycle == 2 && window.Exists() { + t.Fatalf("previous token's renewal window survived rotation: %s", window.Raw) + } + metadata = saved.Value().(map[string]any) + } +} + +func TestCodexExecutorSIWCExplicitRefreshHonorsServerWindow(t *testing.T) { + previousTransport := http.DefaultTransport + t.Cleanup(func() { http.DefaultTransport = previousTransport }) + calls := 0 + http.DefaultTransport = roundTripperFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusBadRequest, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"error":"invalid_grant"}`))}, nil + }) + for _, force := range []bool{false, true} { + t.Run(fmt.Sprintf("force=%t", force), func(t *testing.T) { + ctx := context.Background() + if force { + ctx = cliproxyauth.WithForceRefresh(ctx) + } + auth := &cliproxyauth.Auth{Provider: "codex", Metadata: map[string]any{ + "siwc": map[string]any{ + "access_token": "access", "refresh_token": "refresh", "client_id": "oaiapp_window", + "expires_at": time.Now().Add(-time.Hour).Format(time.RFC3339), + "earliest_refresh_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + }} + if _, errRefresh := NewCodexExecutor(&config.Config{}).Refresh(ctx, auth); errRefresh != nil { + t.Fatalf("refresh before the server window: %v", errRefresh) + } + if calls != 0 { + t.Fatalf("sent %d token requests before the server window", calls) + } + }) + } +} diff --git a/internal/runtime/executor/codex_executor_sse_events_test.go b/internal/runtime/executor/codex_executor_sse_events_test.go index 200eef69b58..7555e5b4b47 100644 --- a/internal/runtime/executor/codex_executor_sse_events_test.go +++ b/internal/runtime/executor/codex_executor_sse_events_test.go @@ -576,7 +576,7 @@ func TestCodexVerifyInferenceTranslatesPayloadAndRequiresCompletedEvent_Issue578 } cfg.Codex.StreamBootstrapBuffering = false - // 2. Real completed stream succeeds and the translated body preserves VerifyInference's max_output_tokens=1 while setting stream=true. + // 2. Real completed stream succeeds and the translated body carries VerifyInference's max_output_tokens=16 while setting stream=true. sseResponse = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"r2\",\"model\":\"gpt-5.4\"}}\n\n" + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"pong\"}\n\n" + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r2\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"pong\"}]}]}}\n\n" @@ -594,8 +594,8 @@ func TestCodexVerifyInferenceTranslatesPayloadAndRequiresCompletedEvent_Issue578 if upstreamJSON["stream"] != true || upstreamJSON["store"] != false { t.Fatalf("expected translated Codex body with stream=true and store=false, got: %s", string(capturedBody)) } - if maxOutputTokens := upstreamJSON["max_output_tokens"]; maxOutputTokens != float64(1) { - t.Fatalf("translated Codex body must preserve max_output_tokens=1, got: %s", string(capturedBody)) + if maxOutputTokens := upstreamJSON["max_output_tokens"]; maxOutputTokens != float64(16) { + t.Fatalf("translated Codex body must carry max_output_tokens=16, got: %s", string(capturedBody)) } } diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index 4a38a5934b9..6074124a530 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -62,14 +62,14 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } originalTranslated = helps.SetStringIfDifferent(originalTranslated, "model", baseModel) doc := ojson.NewDoc(body) - normalizeCodexClaudeOutputBudgetDoc(doc, from, req.Payload) + normalizeCodexOutputBudgetDoc(doc) helps.SetStringIfDifferentDoc(doc, "model", baseModel) if err = helps.ApplyRequestThinkingDoc(doc, req, opts, from.String(), to.String(), e.Identifier(), updatesChanged); err != nil { return nil, err } - ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslated, req, opts)) + ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslated, req, opts, auth)) _ = doc.Delete("previous_response_id") _ = doc.Delete("generate") _ = doc.Delete("prompt_cache_retention") @@ -92,7 +92,11 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if !useSIWC { normalizeCodexInstructionsDoc(doc, preserveNativeOutput) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - ensureImageGenerationToolDoc(doc, baseModel, auth, opts.Headers) + checkModel := req.Model + if strings.TrimSpace(checkModel) == "" { + checkModel = baseModel + } + ensureImageGenerationToolDoc(doc, checkModel, auth, opts.Headers) if !codexAuthUsesAPIKey(auth) { normalizeCodexImageGenerationToolsDoc(doc) } diff --git a/internal/runtime/executor/codex_refresh_explicit_path_test.go b/internal/runtime/executor/codex_refresh_explicit_path_test.go new file mode 100644 index 00000000000..438d78cf02d --- /dev/null +++ b/internal/runtime/executor/codex_refresh_explicit_path_test.go @@ -0,0 +1,43 @@ +package executor + +import ( + "context" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + sdkauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/auth" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestCodexRefreshWithExplicitFilePathStore(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "credential.json") + if errWrite := os.WriteFile(path, []byte(`{"type":"codex","access_token":"valid","refresh_token":"explicit-file-grant","expired":"2099-01-01T00:00:00Z"}`), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + calls := 0 + original := http.DefaultTransport + http.DefaultTransport = roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"new-refresh","expires_in":3600}`)), Request: req}, nil + }) + t.Cleanup(func() { http.DefaultTransport = original }) + store := sdkauth.NewFileTokenStore() + manager := cliproxyauth.NewManager(store, nil, nil) + manager.RegisterExecutor(NewCodexExecutor(&config.Config{AuthDir: dir})) + auth := &cliproxyauth.Auth{ID: path, FileName: path, Provider: "codex", Status: cliproxyauth.StatusActive, Attributes: map[string]string{cliproxyauth.AttributePath: path}, Metadata: map[string]any{"access_token": "valid", "refresh_token": "explicit-file-grant", "expired": "2099-01-01T00:00:00Z"}} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); errRefresh != nil { + t.Fatalf("valid explicit-path file store cannot refresh: %v, upstream calls=%d", errRefresh, calls) + } + if calls != 1 { + t.Fatalf("upstream calls=%d, want 1", calls) + } +} diff --git a/internal/runtime/executor/codex_refresh_outcome_test.go b/internal/runtime/executor/codex_refresh_outcome_test.go new file mode 100644 index 00000000000..66d2b3ecc35 --- /dev/null +++ b/internal/runtime/executor/codex_refresh_outcome_test.go @@ -0,0 +1,52 @@ +package executor + +import ( + "context" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + codexauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestCodexUnknownRefreshSurvivesForcedRefreshAndFileAlias(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "credential.json") + if errWrite := os.WriteFile(path, []byte(`{"type":"codex","access_token":"valid","refresh_token":"alias-grant","expired":"2099-01-01T00:00:00Z"}`), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + alias := filepath.Join(t.TempDir(), "alias.json") + if errSymlink := os.Symlink(path, alias); errSymlink != nil { + t.Skipf("symlink unavailable: %v", errSymlink) + } + calls := 0 + original := http.DefaultTransport + http.DefaultTransport = roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"access_token":`)), Request: req}, nil + }) + t.Cleanup(func() { http.DefaultTransport = original }) + for _, filename := range []string{path, alias} { + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(NewCodexExecutor(&config.Config{})) + auth := &cliproxyauth.Auth{ID: filename, FileName: filename, Provider: "codex", Status: cliproxyauth.StatusActive, Attributes: map[string]string{cliproxyauth.AttributePath: filename}, Metadata: map[string]any{"access_token": "valid", "refresh_token": "alias-grant", "expired": "2099-01-01T00:00:00Z"}} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); !codexauth.IsRefreshOutcomeUnknown(errRefresh) { + t.Fatalf("expected unknown outcome, got %v", errRefresh) + } + updated, _ := manager.GetByID(auth.ID) + if !updated.RequiresSignIn || updated.Unavailable { + t.Fatal("must require refresh recovery but preserve valid access") + } + } + if calls != 1 { + t.Fatalf("force refresh or file alias replayed unresolved grant: calls=%d, want 1", calls) + } +} diff --git a/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go index 965f5b0ad6f..6ab82d5fcf2 100644 --- a/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go +++ b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go @@ -46,12 +46,12 @@ func codexTestAuth(baseURL string) *cliproxyauth.Auth { func codexTestRequest() (cliproxyexecutor.Request, cliproxyexecutor.Options) { return cliproxyexecutor.Request{ - Model: "gpt-5.6-terra", - Payload: []byte(`{"model":"gpt-5.6-terra","input":"hello"}`), - }, cliproxyexecutor.Options{ - SourceFormat: sdktranslator.FromString("openai-response"), - Stream: true, - } + Model: "gpt-5.6-terra", + Payload: []byte(`{"model":"gpt-5.6-terra","input":"hello"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + } } // codexSSEServer streams the supplied event payloads as an HTTP 200 SSE response. @@ -92,11 +92,11 @@ func codexWebsocketServer(t *testing.T, frames ...string) *httptest.Server { func codexWebsocketRequest() (cliproxyexecutor.Request, cliproxyexecutor.Options) { return cliproxyexecutor.Request{ - Model: "gpt-5.6-terra", - Payload: []byte(`{"model":"gpt-5.6-terra","input":[{"type":"message","role":"user","content":"hello"}]}`), - }, cliproxyexecutor.Options{ - SourceFormat: sdktranslator.FromString("openai-response"), - } + Model: "gpt-5.6-terra", + Payload: []byte(`{"model":"gpt-5.6-terra","input":[{"type":"message","role":"user","content":"hello"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + } } // drainChunks collects every payload and the first error from a stream result. diff --git a/internal/runtime/executor/codex_websockets_duplex.go b/internal/runtime/executor/codex_websockets_duplex.go index 41aa876a158..e91333b24f8 100644 --- a/internal/runtime/executor/codex_websockets_duplex.go +++ b/internal/runtime/executor/codex_websockets_duplex.go @@ -17,6 +17,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" 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" log "github.com/sirupsen/logrus" ) @@ -42,6 +43,12 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( initialReporter *helps.UsageReporter, headers http.Header, unlock func(), ) *cliproxyexecutor.StreamResult { streamCtx, cancel := context.WithCancel(ctx) + // Duplex successors bypass dispatch, including automatic steering responses. + followupCtx := ctx + if selection := coreusage.HomeSelectionFromContext(ctx); selection != nil { + selection.Reused = true + followupCtx = coreusage.WithHomeSelection(ctx, *selection) + } 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) @@ -712,7 +719,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( metadataMu.Unlock() wakeWriter() if !firstResponse { - reporter = helps.NewExecutorUsageReporter(ctx, e, req.Model, auth) + reporter = helps.NewExecutorUsageReporter(followupCtx, e, req.Model, auth) reporter.SetTranslatedReasoningEffort(current.clientBody, current.to.String()) reporter.StartResponseTTFT() } @@ -797,7 +804,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( ambiguous := failedID == "" && ((len(pending) > 0 && responseActive) || len(unacknowledgedSteers) > 0) if len(pending) > 0 && !currentFailure && !ambiguous { eventPrepared, pending = pending[0], pending[1:] - eventReporter = helps.NewExecutorUsageReporter(ctx, e, req.Model, auth) + eventReporter = helps.NewExecutorUsageReporter(followupCtx, e, req.Model, auth) eventReporter.SetTranslatedReasoningEffort(eventPrepared.clientBody, eventPrepared.to.String()) } else if !ambiguous { responseActive = false diff --git a/internal/runtime/executor/codex_websockets_duplex_test.go b/internal/runtime/executor/codex_websockets_duplex_test.go index 4790bad0306..c4d85c2a9b1 100644 --- a/internal/runtime/executor/codex_websockets_duplex_test.go +++ b/internal/runtime/executor/codex_websockets_duplex_test.go @@ -15,6 +15,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" core "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" ) @@ -95,6 +96,11 @@ func TestCodexDuplexSteeringLifecycle(t *testing.T) { defer cancel() input := make(chan core.WebsocketInput, 8) ctx = core.WithWebsocketInput(core.WithDownstreamWebsocket(ctx), input) + ctx = coreusage.WithRequestedModelAlias(ctx, t.Name()) + ctx = coreusage.WithHomeSelection(ctx, coreusage.HomeSelection{Model: "home-route"}) + capture := &codexResponseModelUsageCapture{alias: t.Name(), records: make(chan coreusage.Record, 8)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { coreusage.RegisterNamedPlugin(t.Name(), codexResponseModelNoopUsagePlugin{}) }) exec := NewCodexWebsocketsExecutor(&config.Config{Codex: config.CodexConfig{ResponseSteering: true}}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "A", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": server.URL, "websockets": "true"}} @@ -161,6 +167,17 @@ func TestCodexDuplexSteeringLifecycle(t *testing.T) { if connections.Load() != 1 { t.Fatalf("connections=%d", connections.Load()) } + ids := make(map[string]bool) + for turn := range 3 { + record := capture.await(t) + if record.HomeSelection == nil || record.HomeSelection.Reused != (turn > 0) || record.HomeSelection.Model != "home-route" { + t.Fatalf("turn %d accounting=%+v", turn, record.HomeSelection) + } + if ids[record.RequestID] { + t.Fatal("duplex responses share an execution ID") + } + ids[record.RequestID] = true + } for _, key := range []string{"accepted", "boundary", "pending", "successor", "failed"} { if !seen[key] { t.Fatalf("missing %s", key) diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index c9092c92400..bafe8f305f2 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -63,7 +63,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } originalTranslated = helps.SetStringIfDifferent(originalTranslated, "model", baseModel) doc := ojson.NewDoc(body) - normalizeCodexClaudeOutputBudgetDoc(doc, from, req.Payload) + normalizeCodexOutputBudgetDoc(doc) helps.SetStringIfDifferentDoc(doc, "model", baseModel) if err = helps.ApplyRequestThinkingDoc(doc, req, opts, from.String(), to.String(), e.Identifier(), updatesChanged); err != nil { @@ -72,7 +72,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut instructionsResult := ojson.GetBytesNoCopy(req.Payload, "instructions") hasExplicitInstructions := instructionsResult.Exists() && strings.TrimSpace(instructionsResult.String()) != "" - ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, "codex-websockets", baseModel, to.String(), "", originalTranslated, req, opts)) + ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, "codex-websockets", baseModel, to.String(), "", originalTranslated, req, opts, auth)) helps.SetStringIfDifferentDoc(doc, "model", baseModel) helps.SetBoolIfDifferentDoc(doc, "stream", true) _ = doc.Delete("prompt_cache_retention") @@ -81,7 +81,11 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if !useSIWC { normalizeCodexInstructionsDoc(doc, nativeRequest) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - ensureImageGenerationToolDoc(doc, baseModel, auth, opts.Headers) + checkModel := req.Model + if strings.TrimSpace(checkModel) == "" { + checkModel = baseModel + } + ensureImageGenerationToolDoc(doc, checkModel, auth, opts.Headers) if !codexAuthUsesAPIKey(auth) { normalizeCodexImageGenerationToolsDoc(doc) } diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 0c580548227..1d5702d2098 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -9,6 +9,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" "sync/atomic" "testing" "time" @@ -2641,14 +2642,18 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} serverPongCh := make(chan string, 1) inWriteHook := make(chan struct{}) + var inWriteHookOnce sync.Once pongDeliveredDuringWrite := make(chan struct{}) + var pongDeliveredOnce sync.Once + serverDone := make(chan struct{}) + defer close(serverDone) testWebsocketWritePayloadHook = func(conn *websocket.Conn) { - close(inWriteHook) + inWriteHookOnce.Do(func() { close(inWriteHook) }) // Wait until server confirms pong was received before allowing write to finish. select { case <-pongDeliveredDuringWrite: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Error("timed out waiting for pong delivery while payload write was held in hook") } } @@ -2669,25 +2674,31 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) { // Start server reader loop so server processes control frames. readErrCh := make(chan error, 1) + payloadReceived := make(chan struct{}) go func() { + signaled := false for { if _, _, errRead := conn.ReadMessage(); errRead != nil { readErrCh <- errRead return } + if !signaled { + signaled = true + close(payloadReceived) + } } }() // Wait until client has entered writeMessage and is actively holding writeMu. select { case <-inWriteHook: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Errorf("timed out waiting for client write hook") return } // Upstream sends Ping WHILE client payload write is in progress holding writeMu. - _ = conn.WriteControl(websocket.PingMessage, []byte("session-ping"), time.Now().Add(time.Second)) + _ = conn.WriteControl(websocket.PingMessage, []byte("session-ping"), time.Now().Add(5*time.Second)) // Server asserts Pong arrives while client write is still blocked in the hook. select { @@ -2695,15 +2706,25 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_WithSession(t *testing.T) { if got != "session-ping" { t.Errorf("unexpected pong payload: got %q, want session-ping", got) } - close(pongDeliveredDuringWrite) - case <-time.After(2 * time.Second): + pongDeliveredOnce.Do(func() { close(pongDeliveredDuringWrite) }) + case <-time.After(5 * time.Second): t.Errorf("pong was not received while payload write was in progress") return } + // The held client write is released together with the pong. Wait until the + // request frame arrives before writing the terminal event and closing. + select { + case <-payloadReceived: + case <-time.After(5 * time.Second): + t.Errorf("timed out waiting for client payload after pong") + return + } + // Now send terminal response. respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`) _ = conn.WriteMessage(websocket.TextMessage, respPayload) + <-serverDone })) defer server.Close() @@ -2741,13 +2762,17 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} serverPongCh := make(chan string, 1) inWriteHook := make(chan struct{}) + var inWriteHookOnce sync.Once pongDeliveredDuringWrite := make(chan struct{}) + var pongDeliveredOnce sync.Once + serverDone := make(chan struct{}) + defer close(serverDone) testWebsocketWritePayloadHook = func(conn *websocket.Conn) { - close(inWriteHook) + inWriteHookOnce.Do(func() { close(inWriteHook) }) select { case <-pongDeliveredDuringWrite: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Error("timed out waiting for pong delivery while payload write was held in hook") } } @@ -2766,24 +2791,30 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) { return nil }) + payloadReceived := make(chan struct{}) go func() { + signaled := false for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } + if !signaled { + signaled = true + close(payloadReceived) + } } }() // Wait until client has entered writeMessage on sessionless path. select { case <-inWriteHook: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Errorf("timed out waiting for client write hook") return } // Upstream sends Ping WHILE client payload write is in progress. - _ = conn.WriteControl(websocket.PingMessage, []byte("sessionless-ping"), time.Now().Add(time.Second)) + _ = conn.WriteControl(websocket.PingMessage, []byte("sessionless-ping"), time.Now().Add(5*time.Second)) // Server asserts Pong arrives while client write is still in progress. select { @@ -2791,14 +2822,22 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_Sessionless(t *testing.T) { if got != "sessionless-ping" { t.Errorf("unexpected pong payload: got %q, want sessionless-ping", got) } - close(pongDeliveredDuringWrite) - case <-time.After(2 * time.Second): + pongDeliveredOnce.Do(func() { close(pongDeliveredDuringWrite) }) + case <-time.After(5 * time.Second): t.Errorf("pong was not received while payload write was in progress on sessionless connection") return } + select { + case <-payloadReceived: + case <-time.After(5 * time.Second): + t.Errorf("timed out waiting for client payload after pong on sessionless connection") + return + } + respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`) _ = conn.WriteMessage(websocket.TextMessage, respPayload) + <-serverDone })) defer server.Close() @@ -2833,13 +2872,17 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_NonstreamSessionless(t *testi upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} serverPongCh := make(chan string, 1) inWriteHook := make(chan struct{}) + var inWriteHookOnce sync.Once pongDeliveredDuringWrite := make(chan struct{}) + var pongDeliveredOnce sync.Once + serverDone := make(chan struct{}) + defer close(serverDone) testWebsocketWritePayloadHook = func(conn *websocket.Conn) { - close(inWriteHook) + inWriteHookOnce.Do(func() { close(inWriteHook) }) select { case <-pongDeliveredDuringWrite: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Error("timed out waiting for pong delivery while nonstream payload write was held in hook") } } @@ -2858,24 +2901,30 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_NonstreamSessionless(t *testi return nil }) + payloadReceived := make(chan struct{}) go func() { + signaled := false for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } + if !signaled { + signaled = true + close(payloadReceived) + } } }() // Wait until client has entered writeMessage on nonstream path. select { case <-inWriteHook: - case <-time.After(2 * time.Second): + case <-time.After(5 * time.Second): t.Errorf("timed out waiting for client write hook") return } // Upstream sends Ping WHILE client payload write is in progress. - _ = conn.WriteControl(websocket.PingMessage, []byte("nonstream-sessionless-ping"), time.Now().Add(time.Second)) + _ = conn.WriteControl(websocket.PingMessage, []byte("nonstream-sessionless-ping"), time.Now().Add(5*time.Second)) // Server asserts Pong arrives while client write is still in progress. select { @@ -2883,14 +2932,22 @@ func TestCodexWebsockets_KeepalivePingDuringUpload_NonstreamSessionless(t *testi if got != "nonstream-sessionless-ping" { t.Errorf("unexpected pong payload: got %q, want nonstream-sessionless-ping", got) } - close(pongDeliveredDuringWrite) - case <-time.After(2 * time.Second): + pongDeliveredOnce.Do(func() { close(pongDeliveredDuringWrite) }) + case <-time.After(5 * time.Second): t.Errorf("pong was not received while payload write was in progress on nonstream sessionless connection") return } + select { + case <-payloadReceived: + case <-time.After(5 * time.Second): + t.Errorf("timed out waiting for client payload after pong on nonstream sessionless connection") + return + } + respPayload := []byte(`{"type":"response.completed","response":{"id":"resp-1","status":"completed","output":[]}}`) _ = conn.WriteMessage(websocket.TextMessage, respPayload) + <-serverDone })) defer server.Close() diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index abffde7380b..d533e83e9cd 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -891,7 +891,7 @@ func (e *CodexWebsocketsExecutor) prepareCodexWebsocketStream(ctx context.Contex } originalTranslated = helps.SetStringIfDifferent(originalTranslated, "model", baseModel) doc := ojson.NewDoc(body) - normalizeCodexClaudeOutputBudgetDoc(doc, from, req.Payload) + normalizeCodexOutputBudgetDoc(doc) helps.SetStringIfDifferentDoc(doc, "model", baseModel) if err = helps.ApplyRequestThinkingDoc(doc, req, opts, from.String(), to.String(), e.Identifier(), updatesChanged); err != nil { @@ -900,13 +900,17 @@ func (e *CodexWebsocketsExecutor) prepareCodexWebsocketStream(ctx context.Contex instructionsResult := ojson.GetBytesNoCopy(req.Payload, "instructions") hasExplicitInstructions := instructionsResult.Exists() && strings.TrimSpace(instructionsResult.String()) != "" - ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, "codex-websockets", baseModel, to.String(), "", originalTranslated, req, opts)) + ctx = helps.WithPayloadFinalizer(ctx, helps.NewPayloadFinalizer(e.cfg, "codex-websockets", baseModel, to.String(), "", originalTranslated, req, opts, auth)) helps.SetStringIfDifferentDoc(doc, "model", baseModel) ensureWebSearchToolForHistoryDoc(doc, opts.Headers) if !useSIWC { normalizeCodexInstructionsDoc(doc, preserveNativeOutput) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - ensureImageGenerationToolDoc(doc, baseModel, auth, opts.Headers) + checkModel := req.Model + if strings.TrimSpace(checkModel) == "" { + checkModel = baseModel + } + ensureImageGenerationToolDoc(doc, checkModel, auth, opts.Headers) if !codexAuthUsesAPIKey(auth) { normalizeCodexImageGenerationToolsDoc(doc) } diff --git a/internal/runtime/executor/devin_executor.go b/internal/runtime/executor/devin_executor.go index 88209c5458d..0bc7587dae6 100644 --- a/internal/runtime/executor/devin_executor.go +++ b/internal/runtime/executor/devin_executor.go @@ -1231,7 +1231,7 @@ func (e *DevinExecutor) streamDevinFramesForAuth( _ = completedDoc.Set("interaction.finish_reason", completionFinishReason) } if finalUsage != nil { - totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens + totalInput := devinTotalInputTokens(finalUsage) totalOutput := finalUsage.CompletionTokens totalTokens := totalInput + totalOutput @@ -1662,7 +1662,7 @@ func consumeDevinFramesToInteractionsForAuth(body io.Reader, model, chatModelUID } if finalUsage != nil { - totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens + totalInput := devinTotalInputTokens(finalUsage) totalOutput := finalUsage.CompletionTokens totalTokens := totalInput + totalOutput @@ -2875,6 +2875,17 @@ func resolveDevinSessionAndCascadeIDs(ctx context.Context, sessionID, cascadeID return sessionID, cascadeID } +// devinTotalInputTokens returns the full prompt size. Devin reports the prompt in +// three disjoint parts: uncached (field 2), cache write (field 4) and cache read +// (field 5). A cache miss on the GPT route is almost all cache write, so leaving +// it out reported a few tokens for a full prompt. +func devinTotalInputTokens(u *helps.DevinUsage) int64 { + if u == nil { + return 0 + } + return u.PromptTokens + u.CachedTokens + u.CacheWriteTokens +} + func normalizeDevinUUID(raw string) string { raw = strings.TrimSpace(raw) if raw == "" { diff --git a/internal/runtime/executor/devin_executor_test.go b/internal/runtime/executor/devin_executor_test.go index 2364a790ea0..bca8e7678e2 100644 --- a/internal/runtime/executor/devin_executor_test.go +++ b/internal/runtime/executor/devin_executor_test.go @@ -21,6 +21,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/util" 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" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" log "github.com/sirupsen/logrus" logtest "github.com/sirupsen/logrus/hooks/test" @@ -3196,8 +3197,9 @@ func TestRegressionIssue5910_UsageStatsCacheWriteTokensInResponses(t *testing.T) } root := ojson.ParseBytes(interactionsJSON) - if root.Get("usage.total_input_tokens").Int() != 53 { - t.Errorf("usage.total_input_tokens = %d, want 53", root.Get("usage.total_input_tokens").Int()) + // total input = uncached 3 + cache read 50 + cache write 14361 + if root.Get("usage.total_input_tokens").Int() != 14414 { + t.Errorf("usage.total_input_tokens = %d, want 14414", root.Get("usage.total_input_tokens").Int()) } if root.Get("usage.cache_write_tokens").Int() != 14361 { t.Errorf("usage.cache_write_tokens = %d, want 14361", root.Get("usage.cache_write_tokens").Int()) @@ -3206,8 +3208,11 @@ func TestRegressionIssue5910_UsageStatsCacheWriteTokensInResponses(t *testing.T) if detail.CacheCreationTokens != 14361 { t.Errorf("ParseInteractionsUsage CacheCreationTokens = %d, want 14361", detail.CacheCreationTokens) } - if detail.InputTokens != 53 { - t.Errorf("ParseInteractionsUsage InputTokens = %d, want 53", detail.InputTokens) + if detail.InputTokens != 14414 { + t.Errorf("ParseInteractionsUsage InputTokens = %d, want 14414", detail.InputTokens) + } + if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete || detail.TokenBreakdown.Input.UncachedTokens != 3 { + t.Errorf("token breakdown = %+v, want complete with 3 uncached", detail.TokenBreakdown) } // 2. Streaming test @@ -3265,8 +3270,11 @@ func TestRegressionIssue5910_UsageStatsCacheWriteTokensInResponses(t *testing.T) if sDetail.CacheCreationTokens != 14361 { t.Errorf("ParseInteractionsStreamUsage CacheCreationTokens = %d, want 14361", sDetail.CacheCreationTokens) } - if sDetail.InputTokens != 53 { - t.Errorf("ParseInteractionsStreamUsage InputTokens = %d, want 53", sDetail.InputTokens) + if sDetail.InputTokens != 14414 { + t.Errorf("ParseInteractionsStreamUsage InputTokens = %d, want 14414", sDetail.InputTokens) + } + if got := cRoot.Get("interaction.usage.total_input_tokens").Int(); got != 14414 { + t.Errorf("interaction.usage.total_input_tokens = %d, want 14414", got) } } diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index f832fb54e79..9fdacdb8f0f 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -176,6 +176,19 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r if action != "countTokens" { helps.EnsureGeminiTrailingUserContentDoc(doc, "contents") } + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if antigravityUsesReasoningReplayCache(baseModel) { + var changedReplay bool + var errReplay error + replayScope, replayRequestPayload, changedReplay, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return resp, errReplay + } + if changedReplay { + internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") + } + } baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, baseModel, action) if opts.Alt != "" && action != "countTokens" { @@ -231,6 +244,9 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r b, _ := io.ReadAll(httpResp.Body) 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 errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } err = statusErr{code: httpResp.StatusCode, msg: string(b)} return resp, err } @@ -240,6 +256,9 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, err } helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if action != "countTokens" { + cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) + } reporter.ObserveResponseModel(data) if blockMsg, blocked := geminiPromptBlockErrorMessage(data); blocked { err = statusErr{code: http.StatusBadGateway, msg: blockMsg} @@ -309,6 +328,19 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A capGeminiMaxOutputTokensDoc(doc, baseModel) internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") helps.EnsureGeminiBoundaryUserContentDoc(doc, "contents") + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if antigravityUsesReasoningReplayCache(baseModel) { + var changedReplay bool + var errReplay error + replayScope, replayRequestPayload, changedReplay, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return nil, errReplay + } + if changedReplay { + internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") + } + } baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, baseModel, "streamGenerateContent") @@ -365,9 +397,13 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("gemini executor: close response body error: %v", errClose) } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } err = statusErr{code: httpResp.StatusCode, msg: string(b)} return nil, err } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, replayRequestPayload) out := make(chan cliproxyexecutor.StreamChunk, 1) go func() { defer close(out) @@ -396,6 +432,13 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A }() var param any var sawFinishReason bool + var replayCommitted bool + commitReplay := func() { + if !replayCommitted && replayAccumulator != nil && replayAccumulator.terminal && streamErr == nil && helps.ApplyPatchTranslationError(param) == nil { + replayAccumulator.Commit(ctx) + replayCommitted = true + } + } var pendingFilterPayload []byte var pendingFilterMsg string helps.InitializeApplyPatchStream(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, ¶m) @@ -405,6 +448,9 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && antigravityStreamChunksHaveTerminalCommitEvent(lines) { + commitReplay() + } for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -420,6 +466,9 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } reporter.ObserveResponseModel(line) rawPayload := helps.JSONPayload(line) if len(rawPayload) > 0 && len(pendingFilterPayload) > 0 { @@ -517,6 +566,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + commitReplay() for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -879,6 +929,15 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut helps.SetStringIfDifferentDoc(doc, "model", baseModel) internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") helps.EnsureGeminiLeadingUserContentDoc(doc, "contents") + if antigravityUsesReasoningReplayCache(baseModel) { + changedReplay, errReplay := prepareGeminiTopLevelCountTokensReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return cliproxyexecutor.Response{}, errReplay + } + if changedReplay { + internalsignature.SanitizeNonVertexGeminiRequestThoughtSignaturesDoc(doc, "contents") + } + } baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, baseModel, "countTokens") diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index bcc774cfcf5..43d096350eb 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -456,7 +456,6 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au fixGeminiImageAspectRatioDoc(baseModel, doc) helps.SetStringIfDifferentDoc(doc, "model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") } @@ -471,6 +470,18 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au helps.DropGeminiTrailingThoughtOnlyContentDoc(doc, "contents") helps.EnsureGeminiTrailingUserContentDoc(doc, "contents") } + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + replayScope, replayRequestPayload, _, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return resp, errReplay + } + } + if !isImagenModel(baseModel) { + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, opts.SourceFormat.String()) + } baseURL := vertexBaseURL(location) url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, baseModel, action) if opts.Alt != "" && action != "countTokens" { @@ -534,6 +545,9 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au b, _ := io.ReadAll(httpResp.Body) 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 errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } err = statusErr{code: httpResp.StatusCode, msg: string(b)} return resp, err } @@ -543,6 +557,9 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au return resp, errRead } helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if action != "countTokens" && !isImagenModel(baseModel) { + cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) + } reporter.ObserveResponseModel(data) // For Imagen models, convert response to Gemini format before translation @@ -612,7 +629,6 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) helps.SetStringIfDifferentDoc(doc, "model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") action := getVertexAction(baseModel, false) @@ -626,6 +642,16 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip helps.DropGeminiTrailingThoughtOnlyContentDoc(doc, "contents") helps.EnsureGeminiTrailingUserContentDoc(doc, "contents") } + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + replayScope, replayRequestPayload, _, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return resp, errReplay + } + } + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) // For API key auth, use simpler URL format without project/location. // The helper honors base URLs that already encode "/publishers/..." so @@ -690,6 +716,9 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip b, _ := io.ReadAll(httpResp.Body) 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 errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } err = statusErr{code: httpResp.StatusCode, msg: string(b)} return resp, err } @@ -699,6 +728,9 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip return resp, errRead } helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if action != "countTokens" && !isImagenModel(baseModel) { + cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, replayRequestPayload, data) + } reporter.ObserveResponseModel(data) if blockMsg, blocked := geminiPromptBlockErrorMessage(data); blocked { err = statusErr{code: http.StatusBadGateway, msg: blockMsg} @@ -757,12 +789,21 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) helps.SetStringIfDifferentDoc(doc, "model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") action := getVertexAction(baseModel, true) helps.DropGeminiTrailingThoughtOnlyContentDoc(doc, "contents") helps.EnsureGeminiBoundaryUserContentDoc(doc, "contents") + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + replayScope, replayRequestPayload, _, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return nil, errReplay + } + } + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) baseURL := vertexBaseURL(location) url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, baseModel, action) // Imagen models don't support streaming, skip SSE params @@ -829,9 +870,13 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } return nil, statusErr{code: httpResp.StatusCode, msg: string(b)} } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, replayRequestPayload) out := make(chan cliproxyexecutor.StreamChunk, 1) go func() { defer close(out) @@ -853,9 +898,16 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte }() var terminalDelivered bool var sawFinishReason bool + var replayCommitted bool var pendingFilterLine []byte var pendingFilterMsg string var param any + commitReplay := func() { + if !replayCommitted && replayAccumulator != nil && replayAccumulator.terminal && streamErr == nil && helps.ApplyPatchTranslationError(param) == nil { + replayAccumulator.Commit(ctx) + replayCommitted = true + } + } helps.InitializeApplyPatchStream(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, ¶m) emitTranslatedLine := func(line []byte) bool { lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, bytes.Clone(line), ¶m, claudeInputTokens) @@ -863,6 +915,9 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && antigravityStreamChunksHaveTerminalCommitEvent(lines) { + commitReplay() + } for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -878,6 +933,9 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } reporter.ObserveResponseModel(line) rawPayload := helps.JSONPayload(line) if len(rawPayload) > 0 && len(pendingFilterLine) > 0 { @@ -970,6 +1028,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + commitReplay() for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -1018,12 +1077,21 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) helps.SetStringIfDifferentDoc(doc, "model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") action := getVertexAction(baseModel, true) helps.DropGeminiTrailingThoughtOnlyContentDoc(doc, "contents") helps.EnsureGeminiBoundaryUserContentDoc(doc, "contents") + replayScope := antigravityReasoningReplayScope{} + var replayRequestPayload []byte + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + replayScope, replayRequestPayload, _, errReplay = prepareGeminiTopLevelReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc) + if errReplay != nil { + return nil, errReplay + } + } + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) // For API key auth, use simpler URL format without project/location url := vertexAPIKeyModelsURL(baseURL, baseModel, action) // Imagen models don't support streaming, skip SSE params @@ -1087,9 +1155,13 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("vertex executor: close response body error: %v", errClose) } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClear != nil { + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } return nil, statusErr{code: httpResp.StatusCode, msg: string(b)} } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, replayRequestPayload) out := make(chan cliproxyexecutor.StreamChunk, 1) go func() { defer close(out) @@ -1111,9 +1183,16 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth }() var terminalDelivered bool var sawFinishReason bool + var replayCommitted bool var pendingFilterLine []byte var pendingFilterMsg string var param any + commitReplay := func() { + if !replayCommitted && replayAccumulator != nil && replayAccumulator.terminal && streamErr == nil && helps.ApplyPatchTranslationError(param) == nil { + replayAccumulator.Commit(ctx) + replayCommitted = true + } + } helps.InitializeApplyPatchStream(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, ¶m) emitTranslatedLine := func(line []byte) bool { lines := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, bytes.Clone(line), ¶m, claudeInputTokens) @@ -1121,6 +1200,9 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + if replayAccumulator != nil && replayAccumulator.terminal && !replayCommitted && antigravityStreamChunksHaveTerminalCommitEvent(lines) { + commitReplay() + } for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -1136,6 +1218,9 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } reporter.ObserveResponseModel(line) rawPayload := helps.JSONPayload(line) if len(rawPayload) > 0 && len(pendingFilterLine) > 0 { @@ -1228,6 +1313,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { return statusErr{code: http.StatusBadGateway, msg: msg} }) + commitReplay() for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: lines[i]}: @@ -1269,13 +1355,18 @@ func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context fixGeminiImageAspectRatioDoc(baseModel, doc) _ = doc.Set("model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) respCtx := context.WithValue(ctx, "alt", opts.Alt) _ = doc.Delete("tools") _ = doc.Delete("generationConfig") _ = doc.Delete("safetySettings") internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") helps.EnsureGeminiLeadingUserContentDoc(doc, "contents") + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + if _, errReplay := prepareGeminiTopLevelCountTokensReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc); errReplay != nil { + return cliproxyexecutor.Response{}, errReplay + } + } + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) baseURL := vertexBaseURL(location) url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, baseModel, "countTokens") @@ -1373,13 +1464,18 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth * fixGeminiImageAspectRatioDoc(baseModel, doc) _ = doc.Set("model", baseModel) - helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) respCtx := context.WithValue(ctx, "alt", opts.Alt) _ = doc.Delete("tools") _ = doc.Delete("generationConfig") _ = doc.Delete("safetySettings") internalsignature.SanitizeGeminiRequestThoughtSignaturesDoc(doc, "contents") helps.EnsureGeminiLeadingUserContentDoc(doc, "contents") + if !isImagenModel(baseModel) && antigravityUsesReasoningReplayCache(baseModel) { + if _, errReplay := prepareGeminiTopLevelCountTokensReasoningReplayPayloadDoc(ctx, baseModel, req, opts, doc); errReplay != nil { + return cliproxyexecutor.Response{}, errReplay + } + } + helps.StripVertexOpenAIResponsesToolCallIDsDoc(doc, from.String()) // For API key auth, use simpler URL format without project/location url := vertexAPIKeyModelsURL(baseURL, baseModel, "countTokens") diff --git a/internal/runtime/executor/helps/claude_client_detection.go b/internal/runtime/executor/helps/claude_client_detection.go index 14475c560f8..aa1daa6115f 100644 --- a/internal/runtime/executor/helps/claude_client_detection.go +++ b/internal/runtime/executor/helps/claude_client_detection.go @@ -26,11 +26,7 @@ const ( claudeDefaultStainlessTimeout = "600" ) -var ( - claudeCodeUserAgentPattern = regexp.MustCompile(`(?i)^claude-cli/`) - claudeCodeUserAgentDetailsPattern = regexp.MustCompile(`(?i)^claude-cli/\S+\s+\(external,\s*([^,)]+)(?:,\s*agent-sdk/([^,)]+))?`) - claudeCodeNativeUserAgentPattern = regexp.MustCompile(`(?i)^claude-cli/[0-9]+\.[0-9]+\.[0-9]+\s+\(external,\s*[^,)]+(?:,\s*agent-sdk/[0-9]+\.[0-9]+\.[0-9]+)?\)$`) -) +var claudeCodeUserAgentPattern = regexp.MustCompile(`(?i)^claude-cli/`) var claudeCodeSubclientByEntrypoint = map[string]string{ "cli": "claude-code-cli", @@ -174,16 +170,16 @@ func DetectClaudeCodeRequest(headers http.Header, payload []byte, countTokens bo if len(configs) > 0 { cfg = configs[0] } - userAgent := headerValue(headers, "User-Agent") - entrypoint, agentSDKVersion := parseClaudeCodeUserAgentDetails(userAgent) + userAgent, validUserAgent := parseClaudeCodeUserAgent(headerValue(headers, "User-Agent")) + entrypoint := userAgent.entrypoint subclient := claudeSubclientForEntrypoint(entrypoint, cfg) detection := ClaudeCodeRequestDetection{ XAppCLI: headerValue(headers, "X-App") == "cli", - UserAgent: plausibleClaudeCodeUserAgent(userAgent, cfg), + UserAgent: validUserAgent && plausibleClaudeCodeVersion(userAgent.version, cfg), BetasPresent: headerContainsClaudeCodeBeta(headers), Entrypoint: entrypoint, Subclient: subclient, - AgentSDKVersion: agentSDKVersion, + AgentSDKVersion: userAgent.agentSDKVersion, } metadataUserID := ojson.GetBytesNoCopy(payload, "metadata.user_id") @@ -690,26 +686,19 @@ func IsClaudeCodeUserAgent(userAgent string) bool { } func plausibleClaudeCodeUserAgent(userAgent string, cfg *config.Config) bool { - userAgent = strings.TrimSpace(userAgent) - if !claudeCodeUserAgentPattern.MatchString(userAgent) || !claudeCodeNativeUserAgentPattern.MatchString(userAgent) { - return false - } - candidate, okCandidate := parseClaudeCLIVersion(userAgent) + parsed, ok := parseClaudeCodeUserAgent(userAgent) + return ok && plausibleClaudeCodeVersion(parsed.version, cfg) +} + +func plausibleClaudeCodeVersion(candidate claudeCLIVersion, cfg *config.Config) bool { baseline, okBaseline := parseClaudeCLIVersion(defaultClaudeDeviceProfile(cfg).UserAgent) // Patch releases (>= baseline.patch) within the release line preserve native passthrough. - return okCandidate && okBaseline && plausibleClaudeCLIVersion(candidate, baseline) + return okBaseline && plausibleClaudeCLIVersion(candidate, baseline) } func parseClaudeCodeUserAgentDetails(userAgent string) (entrypoint, agentSDKVersion string) { - matches := claudeCodeUserAgentDetailsPattern.FindStringSubmatch(strings.TrimSpace(userAgent)) - if len(matches) < 2 { - return "", "" - } - entrypoint = strings.ToLower(strings.TrimSpace(matches[1])) - if len(matches) >= 3 { - agentSDKVersion = strings.TrimSpace(matches[2]) - } - return entrypoint, agentSDKVersion + parsed, _ := parseClaudeCodeUserAgent(userAgent) + return parsed.entrypoint, parsed.agentSDKVersion } func headerValue(headers http.Header, name string) string { diff --git a/internal/runtime/executor/helps/claude_client_detection_test.go b/internal/runtime/executor/helps/claude_client_detection_test.go index 7e0c91566a1..f4152396c98 100644 --- a/internal/runtime/executor/helps/claude_client_detection_test.go +++ b/internal/runtime/executor/helps/claude_client_detection_test.go @@ -18,7 +18,7 @@ func claudeCodeDetectionPayload(userID string) []byte { func confirmedClaudeCodeHeaders() http.Header { return http.Header{ - "User-Agent": {"claude-cli/2.1.280 (external, cli)"}, + "User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219,interleaved-thinking-2025-05-14"}, } @@ -84,7 +84,7 @@ func TestDetectClaudeCodeRequestRequiresAllFourMessageSignals(t *testing.T) { func TestDetectClaudeCodeRequestAcceptsNewerPatchInMeasuredReleaseLine(t *testing.T) { headers := confirmedClaudeCodeHeaders() - headers.Set("User-Agent", "claude-cli/2.1.281 (external, cli)") + headers.Set("User-Agent", "claude-cli/2.1.322 (external, cli)") payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID) detection := DetectClaudeCodeRequest(headers, payload, false) @@ -118,9 +118,9 @@ func TestDetectClaudeCodeRequestRejectsEachMissingMessageSignal(t *testing.T) { headers http.Header body []byte }{ - {name: "x-app", headers: http.Header{"User-Agent": {"claude-cli/2.1.280 (external, cli)"}, "Anthropic-Beta": {"claude-code-20250219"}}, body: payload}, + {name: "x-app", headers: http.Header{"User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "Anthropic-Beta": {"claude-code-20250219"}}, body: payload}, {name: "user-agent", headers: http.Header{"User-Agent": {"curl/8.7.1"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, body: payload}, - {name: "betas", headers: http.Header{"User-Agent": {"claude-cli/2.1.280 (external, cli)"}, "X-App": {"cli"}}, body: payload}, + {name: "betas", headers: http.Header{"User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}}, body: payload}, {name: "metadata", headers: confirmedClaudeCodeHeaders(), body: []byte(`{"messages":[]}`)}, } { t.Run(test.name, func(t *testing.T) { @@ -141,16 +141,16 @@ func TestDetectClaudeCodeRequestClassifiesEntrypoints(t *testing.T) { agentSDKVersion string native bool }{ - {name: "cli", userAgent: "claude-cli/2.1.280 (external, cli)", entrypoint: "cli", subclient: "claude-code-cli", native: true}, - {name: "vscode-agent-sdk", userAgent: "claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)", entrypoint: "claude-vscode", subclient: "claude-code-vscode", agentSDKVersion: "0.3.220", native: true}, - {name: "sdk-cli", userAgent: "claude-cli/2.1.280 (external, sdk-cli)", entrypoint: "sdk-cli", subclient: "claude-code-cli-sdk", native: true}, - {name: "sdk-ts", userAgent: "claude-cli/2.1.280 (external, sdk-ts, agent-sdk/0.3.220)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", agentSDKVersion: "0.3.220", native: true}, - {name: "sdk-py", userAgent: "claude-cli/2.1.280 (external, sdk-py, agent-sdk/0.1.0)", entrypoint: "sdk-py", subclient: "claude-code-sdk-py", agentSDKVersion: "0.1.0"}, - {name: "desktop", userAgent: "claude-cli/2.1.280 (external, claude-desktop)", entrypoint: "claude-desktop", subclient: "claude-desktop"}, - {name: "desktop-third-party-inference", userAgent: "claude-cli/2.1.280 (external, claude-desktop-3p)", entrypoint: "claude-desktop-3p", subclient: "claude-desktop-3p"}, - {name: "remote", userAgent: "claude-cli/2.1.280 (external, remote)", entrypoint: "remote", subclient: "claude-remote"}, - {name: "github-action", userAgent: "claude-cli/2.1.280 (external, claude-code-github-action)", entrypoint: "claude-code-github-action", subclient: "claude-code-gh-action"}, - {name: "unknown", userAgent: "claude-cli/2.1.280 (external, copied-client)", entrypoint: "copied-client"}, + {name: "cli", userAgent: "claude-cli/2.1.321 (external, cli)", entrypoint: "cli", subclient: "claude-code-cli", native: true}, + {name: "vscode-agent-sdk", userAgent: "claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)", entrypoint: "claude-vscode", subclient: "claude-code-vscode", agentSDKVersion: "0.3.220", native: true}, + {name: "sdk-cli", userAgent: "claude-cli/2.1.321 (external, sdk-cli)", entrypoint: "sdk-cli", subclient: "claude-code-cli-sdk", native: true}, + {name: "sdk-ts", userAgent: "claude-cli/2.1.321 (external, sdk-ts, agent-sdk/0.3.220)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", agentSDKVersion: "0.3.220", native: true}, + {name: "sdk-py", userAgent: "claude-cli/2.1.321 (external, sdk-py, agent-sdk/0.1.0)", entrypoint: "sdk-py", subclient: "claude-code-sdk-py", agentSDKVersion: "0.1.0"}, + {name: "desktop", userAgent: "claude-cli/2.1.321 (external, claude-desktop)", entrypoint: "claude-desktop", subclient: "claude-desktop"}, + {name: "desktop-third-party-inference", userAgent: "claude-cli/2.1.321 (external, claude-desktop-3p)", entrypoint: "claude-desktop-3p", subclient: "claude-desktop-3p"}, + {name: "remote", userAgent: "claude-cli/2.1.321 (external, remote)", entrypoint: "remote", subclient: "claude-remote"}, + {name: "github-action", userAgent: "claude-cli/2.1.321 (external, claude-code-github-action)", entrypoint: "claude-code-github-action", subclient: "claude-code-gh-action"}, + {name: "unknown", userAgent: "claude-cli/2.1.321 (external, copied-client)", entrypoint: "copied-client"}, } { t.Run(test.name, func(t *testing.T) { headers := confirmedClaudeCodeHeaders() @@ -171,7 +171,7 @@ func TestDetectClaudeCodeRequestClassifiesEntrypoints(t *testing.T) { func TestDetectClaudeCodeCountTokensAllowsMissingMetadata(t *testing.T) { headers := confirmedClaudeCodeHeaders() - headers.Set("User-Agent", "claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)") + headers.Set("User-Agent", "claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)") detection := DetectClaudeCodeRequest(headers, []byte(`{"messages":[]}`), true) if !detection.Confirmed { t.Fatalf("detection = %#v, want confirmed", detection) @@ -412,7 +412,7 @@ func TestDetectClaudeCodeRequestRejectsMalformedNativeSignals(t *testing.T) { {name: "older patch user agent", headers: http.Header{"User-Agent": {"claude-cli/2.1.257 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID}, {name: "unmeasured next-minor user agent", headers: http.Header{"User-Agent": {"claude-cli/2.2.0 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID}, {name: "implausible future user agent", headers: http.Header{"User-Agent": {"claude-cli/999.0.0 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID}, - {name: "unrelated beta", headers: http.Header{"User-Agent": {"claude-cli/2.1.280 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"anything"}}, userID: validClaudeCodeMetadataUserID}, + {name: "unrelated beta", headers: http.Header{"User-Agent": {"claude-cli/2.1.321 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"anything"}}, userID: validClaudeCodeMetadataUserID}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { @@ -607,11 +607,11 @@ func TestDetectClaudeCodeRequestRecognizesAgentSDKHaiku55Helpers(t *testing.T) { payload []byte cfg *config.Config }{ - {name: "sdk-ts title helper", userAgent: "claude-cli/2.1.304 (external, sdk-ts, agent-sdk/0.3.276)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", betas: titleBetas, payload: titlePayload}, - {name: "sdk-ts web_search helper", userAgent: "claude-cli/2.1.304 (external, sdk-ts, agent-sdk/0.3.276)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", betas: webSearchBetas, payload: webSearchPayload}, + {name: "sdk-ts title helper", userAgent: "claude-cli/2.1.345 (external, sdk-ts, agent-sdk/0.3.276)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", betas: titleBetas, payload: titlePayload}, + {name: "sdk-ts web_search helper", userAgent: "claude-cli/2.1.345 (external, sdk-ts, agent-sdk/0.3.276)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", betas: webSearchBetas, payload: webSearchPayload}, { name: "configured t3-code title helper", - userAgent: "claude-cli/2.1.316 (external, t3-code, agent-sdk/0.3.276)", + userAgent: "claude-cli/2.1.357 (external, t3-code, agent-sdk/0.3.276)", entrypoint: "t3-code", subclient: "claude-code-t3-code", betas: titleBetas, @@ -639,7 +639,7 @@ func TestDetectClaudeCodeRequestRecognizesAgentSDKHaiku55Helpers(t *testing.T) { func TestDetectClaudeCodeRequestAcceptsConfiguredNativeEntrypoints(t *testing.T) { payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID) headers := confirmedClaudeCodeHeaders() - headers.Set("User-Agent", "claude-cli/2.1.316 (external, t3-code, agent-sdk/0.3.276)") + headers.Set("User-Agent", "claude-cli/2.1.357 (external, t3-code, agent-sdk/0.3.276)") if detection := DetectClaudeCodeRequest(headers, payload, false); detection.Confirmed || detection.NativeClient { t.Fatalf("default detection = %#v, want unconfirmed without native-entrypoints config", detection) @@ -659,7 +659,7 @@ func TestDetectClaudeCodeRequestAcceptsConfiguredNativeEntrypoints(t *testing.T) } missingBetaHeaders := confirmedClaudeCodeHeaders() - missingBetaHeaders.Set("User-Agent", "claude-cli/2.1.316 (external, t3-code, agent-sdk/0.3.276)") + missingBetaHeaders.Set("User-Agent", "claude-cli/2.1.357 (external, t3-code, agent-sdk/0.3.276)") missingBetaHeaders.Del("Anthropic-Beta") if unconfirmed := DetectClaudeCodeRequest(missingBetaHeaders, payload, false, cfg); unconfirmed.Confirmed || unconfirmed.StrongSignals { t.Fatalf("configured entrypoint without strong signals = %#v, want unconfirmed", unconfirmed) diff --git a/internal/runtime/executor/helps/claude_code_user_agent.go b/internal/runtime/executor/helps/claude_code_user_agent.go new file mode 100644 index 00000000000..2d88eb082c5 --- /dev/null +++ b/internal/runtime/executor/helps/claude_code_user_agent.go @@ -0,0 +1,92 @@ +package helps + +import "strings" + +type claudeCodeUserAgent struct { + version claudeCLIVersion + entrypoint string + agentSDKVersion string + profileUserAgent string + requestFields []string +} + +// parseClaudeCodeUserAgent separates client identity from opaque request fields. +// Parsing a valid UA does not establish native-client eligibility. +func parseClaudeCodeUserAgent(userAgent string) (claudeCodeUserAgent, bool) { + for _, char := range userAgent { + if (char < ' ' && char != '\t') || char == '\x7f' { + return claudeCodeUserAgent{}, false + } + } + userAgent = strings.TrimSpace(userAgent) + prefix, comment, ok := strings.Cut(userAgent, "(") + if !ok || !strings.HasSuffix(comment, ")") || prefix == strings.TrimSpace(prefix) { + return claudeCodeUserAgent{}, false + } + product := strings.TrimSpace(prefix) + versionText, okVersion := strings.CutPrefix(product, "claude-cli/") + if !okVersion || !claudePackageVersionPattern.MatchString(versionText) { + return claudeCodeUserAgent{}, false + } + version, okVersion := parseClaudeCLIVersion(product) + if !okVersion { + return claudeCodeUserAgent{}, false + } + comment = strings.TrimSuffix(comment, ")") + if strings.ContainsAny(comment, "()") { + return claudeCodeUserAgent{}, false + } + fields := strings.Split(comment, ",") + if len(fields) < 2 || !strings.EqualFold(fields[0], "external") { + return claudeCodeUserAgent{}, false + } + entrypoint := strings.ToLower(strings.TrimSpace(fields[1])) + if entrypoint == "" { + return claudeCodeUserAgent{}, false + } + parsed := claudeCodeUserAgent{ + version: version, entrypoint: entrypoint, profileUserAgent: userAgent, + } + stableFields := []string{fields[0], fields[1]} + seen := make(map[string]struct{}) + for _, field := range fields[2:] { + name, value, okField := strings.Cut(strings.TrimSpace(field), "/") + if !okField || !validClaudeCodeUserAgentFieldName(name) || strings.TrimSpace(value) == "" { + return claudeCodeUserAgent{}, false + } + name = strings.ToLower(name) + if _, duplicate := seen[name]; duplicate { + return claudeCodeUserAgent{}, false + } + seen[name] = struct{}{} + if name == "agent-sdk" { + if !claudePackageVersionPattern.MatchString(value) { + return claudeCodeUserAgent{}, false + } + parsed.agentSDKVersion = value + stableFields = append(stableFields, field) + } else { + parsed.requestFields = append(parsed.requestFields, field) + } + } + if len(parsed.requestFields) > 0 { + // Request fields must never become a credential's stored device identity. + parsed.profileUserAgent = prefix + "(" + strings.Join(stableFields, ",") + ")" + } + return parsed, true +} + +func validClaudeCodeUserAgentFieldName(name string) bool { + if name == "" { + return false + } + // Extension names use HTTP token syntax; values remain opaque. + for _, char := range name { + if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') || + strings.ContainsRune("!#$%&'*+-.^_`|~", char) { + continue + } + return false + } + return true +} diff --git a/internal/runtime/executor/helps/claude_code_user_agent_test.go b/internal/runtime/executor/helps/claude_code_user_agent_test.go new file mode 100644 index 00000000000..a531f036b89 --- /dev/null +++ b/internal/runtime/executor/helps/claude_code_user_agent_test.go @@ -0,0 +1,226 @@ +package helps + +import ( + "context" + "net/http/httptest" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestDetectClaudeCodeRequestAcceptsUserAgentFields(t *testing.T) { + for _, testCase := range []struct { + name, userAgent, sdkVersion string + }{ + {name: "cron", userAgent: "claude-cli/2.1.353 (external, cli, workload/cron)"}, + {name: "field before SDK", userAgent: "claude-cli/2.1.353 (external, cli, workload/cron, agent-sdk/0.2.133)", sdkVersion: "0.2.133"}, + {name: "field after SDK", userAgent: "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, workload/cron)", sdkVersion: "0.2.133"}, + {name: "unknown fields", userAgent: "claude-cli/2.1.353 (external, cli, future-feature/value, workload/background)"}, + } { + for _, countTokens := range []bool{false, true} { + path := "messages" + if countTokens { + path = "count_tokens" + } + t.Run(testCase.name+"/"+path, func(t *testing.T) { + headers := confirmedClaudeCodeHeaders() + headers.Set("User-Agent", testCase.userAgent) + payload := []byte(`{}`) + if !countTokens { + payload = claudeCodeDetectionPayload(validClaudeCodeMetadataUserID) + } + detection := DetectClaudeCodeRequest(headers, payload, countTokens) + if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient { + t.Fatalf("detection = %+v, want confirmed native CLI", detection) + } + if detection.Entrypoint != "cli" || detection.AgentSDKVersion != testCase.sdkVersion { + t.Fatalf("identity = %q/%q, want cli/%q", detection.Entrypoint, detection.AgentSDKVersion, testCase.sdkVersion) + } + }) + } + } +} + +func TestParseClaudeCodeUserAgent(t *testing.T) { + for _, testCase := range []struct { + name, userAgent, entrypoint, sdkVersion, profile string + }{ + {name: "plain", userAgent: "claude-cli/2.1.353 (external, cli)", entrypoint: "cli"}, + {name: "SDK", userAgent: "claude-cli/2.1.353 (external, sdk-ts, agent-sdk/0.2.133)", entrypoint: "sdk-ts", sdkVersion: "0.2.133"}, + {name: "cron", userAgent: "claude-cli/2.1.353 (external, cli, workload/cron)", entrypoint: "cli", profile: "claude-cli/2.1.353 (external, cli)"}, + {name: "SDK after metadata", userAgent: "claude-cli/2.1.353 (external, sdk-ts, workload/cron, agent-sdk/0.2.133, future-feature/value)", entrypoint: "sdk-ts", sdkVersion: "0.2.133", profile: "claude-cli/2.1.353 (external, sdk-ts, agent-sdk/0.2.133)"}, + {name: "token field name", userAgent: "claude-cli/2.1.353 (external, cli, future_field.v2/value/path)", entrypoint: "cli", profile: "claude-cli/2.1.353 (external, cli)"}, + } { + t.Run(testCase.name, func(t *testing.T) { + parsed, ok := parseClaudeCodeUserAgent(testCase.userAgent) + if !ok || parsed.version != (claudeCLIVersion{major: 2, minor: 1, patch: 353}) || parsed.entrypoint != testCase.entrypoint || parsed.agentSDKVersion != testCase.sdkVersion { + t.Fatalf("parsed = %+v, valid = %v", parsed, ok) + } + profile := testCase.profile + if profile == "" { + profile = testCase.userAgent + } + if parsed.profileUserAgent != profile { + t.Fatalf("profile UA = %q, want %q", parsed.profileUserAgent, profile) + } + }) + } +} + +func TestParseClaudeCodeUserAgentRejectsMalformedFields(t *testing.T) { + for _, userAgent := range []string{ + "claude-cli/2.1.353(external, cli)", + "claude-cli/2.1.353-extra (external, cli)", + "claude-cli/999999999999999999999999.1.0 (external, cli)", + "claude-cli/2.1.353 (external, )", + "claude-cli/2.1.353 (external, cli, foo)", + "claude-cli/2.1.353 (external, cli, /cron)", + "claude-cli/2.1.353 (external, cli, workload/)", + "claude-cli/2.1.353 (external, cli, bad name/value)", + "claude-cli/2.1.353 (external, cli, workload/(cron))", + "claude-cli/2.1.353 (external, cli, workload/cron, WORKLOAD/background)", + "claude-cli/2.1.353 (external, cli, agent-sdk/dev)", + "claude-cli/2.1.353 (external, cli, agent-sdk/0.2.133, agent-sdk/0.2.134)", + "claude-cli/2.1.353 (external, cli, workload/cron) trailing", + "claude-cli/2.1.353 (external, cli, workload/cr\non)", + "claude-cli/2.1.353 (external, cli)\r\n", + } { + t.Run(userAgent, func(t *testing.T) { + if parsed, ok := parseClaudeCodeUserAgent(userAgent); ok { + t.Fatalf("accepted malformed UA: %+v", parsed) + } + }) + } +} + +func TestDetectClaudeCodeRequestUserAgentFieldsKeepNativeGuards(t *testing.T) { + for _, testCase := range []struct { + name, userAgent, missingHeader string + missingMetadata, allowEntrypoint, wantConfirmed bool + }{ + {name: "old patch", userAgent: "claude-cli/2.1.320 (external, cli, workload/cron)"}, + {name: "unknown release line", userAgent: "claude-cli/2.2.0 (external, cli, workload/cron)"}, + {name: "unknown entrypoint", userAgent: "claude-cli/2.1.353 (external, other-client, workload/cron)"}, + {name: "missing app", missingHeader: "X-App"}, + {name: "missing beta", missingHeader: "Anthropic-Beta"}, + {name: "missing metadata", missingMetadata: true}, + {name: "configured entrypoint", userAgent: "claude-cli/2.1.353 (external, other-client, workload/cron)", allowEntrypoint: true, wantConfirmed: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + headers := confirmedClaudeCodeHeaders() + userAgent := testCase.userAgent + if userAgent == "" { + userAgent = "claude-cli/2.1.353 (external, cli, workload/cron)" + } + headers.Set("User-Agent", userAgent) + headers.Del(testCase.missingHeader) + payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID) + if testCase.missingMetadata { + payload = []byte(`{}`) + } + cfg := &config.Config{} + if testCase.allowEntrypoint { + cfg.ClaudeHeaderDefaults.NativeEntrypoints = []string{"other-client"} + } + detection := DetectClaudeCodeRequest(headers, payload, false, cfg) + if detection.Confirmed != testCase.wantConfirmed { + t.Fatalf("detection = %+v, want confirmed = %v", detection, testCase.wantConfirmed) + } + }) + } +} + +func TestClaudeDeviceProfileUserAgentFieldsStayRequestScoped(t *testing.T) { + for _, homeMode := range []bool{false, true} { + name := "local" + if homeMode { + name = "home" + } + t.Run(name, func(t *testing.T) { + ResetClaudeDeviceProfileCache() + t.Cleanup(ResetClaudeDeviceProfileCache) + client := newFakeClaudeDeviceProfileKVClient() + useFakeClaudeDeviceProfileKVClient(t, client, homeMode, nil) + auth := &cliproxyauth.Auth{ID: "request-fields-" + name} + baseUA := "claude-cli/2.1.321 (external, cli, agent-sdk/0.2.133)" + cronUA := "claude-cli/2.1.321 (external, cli, workload/cron, agent-sdk/0.2.133, future-feature/value)" + for _, testCase := range []struct{ userAgent, want string }{ + {userAgent: cronUA, want: "claude-cli/2.1.321 (external, cli, agent-sdk/0.2.133, workload/cron, future-feature/value)"}, + {userAgent: baseUA, want: baseUA}, + {userAgent: "claude-cli/2.1.321 (external, cli, future-feature/next)", want: "claude-cli/2.1.321 (external, cli, agent-sdk/0.2.133, future-feature/next)"}, + {want: baseUA}, + } { + headers := claudeDeviceHeaders(testCase.userAgent) + profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "dummy-key", headers, nil) + if errProfile != nil { + t.Fatal(errProfile) + } + if profile.UserAgent != baseUA { + t.Fatalf("profile UA = %q after %q, want %q", profile.UserAgent, testCase.userAgent, baseUA) + } + request := httptest.NewRequest("POST", "https://example.invalid/messages", nil) + ApplyClaudeDeviceProfileHeaders(request, profile, testCase.userAgent) + if got := request.Header.Get("User-Agent"); got != testCase.want { + t.Fatalf("request UA = %q, want only current fields in %q", got, testCase.want) + } + } + if homeMode && client.setCount != 1 { + t.Fatalf("stored %d profiles, want one shared stable identity", client.setCount) + } + }) + } +} + +func TestApplyClaudeDeviceProfileHeadersKeepsCurrentUserAgentFields(t *testing.T) { + profile := defaultClaudeDeviceProfile(nil) + profile.UserAgent = "claude-cli/2.1.321 (external, cli, agent-sdk/0.3.220)" + for _, testCase := range []struct { + name, userAgent, want string + }{ + {name: "current fields", userAgent: "claude-cli/2.1.353 (external, cli, workload/cron, agent-sdk/0.2.133, future-feature/value)", want: "claude-cli/2.1.321 (external, cli, agent-sdk/0.3.220, workload/cron, future-feature/value)"}, + {name: "next request fields", userAgent: "claude-cli/2.1.353 (external, cli, future-feature/next)", want: "claude-cli/2.1.321 (external, cli, agent-sdk/0.3.220, future-feature/next)"}, + {name: "ordinary request", userAgent: "claude-cli/2.1.353 (external, cli)", want: profile.UserAgent}, + {name: "malformed request", userAgent: "claude-cli/2.1.353 (external, cli, foo)", want: profile.UserAgent}, + {name: "no request", want: profile.UserAgent}, + } { + t.Run(testCase.name, func(t *testing.T) { + request := httptest.NewRequest("POST", "https://example.invalid/messages", nil) + ApplyClaudeDeviceProfileHeaders(request, profile, testCase.userAgent) + if got := request.Header.Get("User-Agent"); got != testCase.want { + t.Fatalf("upstream UA = %q, want %q", got, testCase.want) + } + for name, want := range map[string]string{ + "X-Stainless-Package-Version": profile.PackageVersion, + "X-Stainless-Runtime-Version": profile.RuntimeVersion, + "X-Stainless-Os": profile.OS, + "X-Stainless-Arch": profile.Arch, + } { + if got := request.Header.Get(name); got != want { + t.Fatalf("%s = %q, want stabilized %q", name, got, want) + } + } + }) + } +} + +func FuzzParseClaudeCodeUserAgent(f *testing.F) { + for _, seed := range []string{ + "claude-cli/2.1.353 (external, cli)", + "claude-cli/2.1.353 (external, cli, workload/cron, agent-sdk/0.2.133)", + "claude-cli/2.1.353 (external, cli, agent-sdk/dev)", + "", + } { + f.Add(seed) + } + f.Fuzz(func(t *testing.T, userAgent string) { + parsed, ok := parseClaudeCodeUserAgent(userAgent) + if !ok { + return + } + stable, okStable := parseClaudeCodeUserAgent(parsed.profileUserAgent) + if !okStable || stable.version != parsed.version || stable.entrypoint != parsed.entrypoint || stable.agentSDKVersion != parsed.agentSDKVersion || stable.profileUserAgent != parsed.profileUserAgent || len(stable.requestFields) != 0 { + t.Fatalf("unstable parsed identity: %+v -> %+v (valid = %v)", parsed, stable, okStable) + } + }) +} diff --git a/internal/runtime/executor/helps/claude_device_profile.go b/internal/runtime/executor/helps/claude_device_profile.go index b203836f121..e77968e913b 100644 --- a/internal/runtime/executor/helps/claude_device_profile.go +++ b/internal/runtime/executor/helps/claude_device_profile.go @@ -20,9 +20,9 @@ import ( ) const ( - defaultClaudeFingerprintUserAgent = "claude-cli/2.1.280 (external, cli)" - defaultClaudeFingerprintPackageVersion = "0.112.1" - defaultClaudeFingerprintRuntimeVersion = "v26.3.0" + defaultClaudeFingerprintUserAgent = "claude-cli/2.1.321 (external, cli)" + defaultClaudeFingerprintPackageVersion = "0.128.0" + defaultClaudeFingerprintRuntimeVersion = "v26.12.1" defaultClaudeFingerprintOS = "MacOS" defaultClaudeFingerprintArch = "arm64" claudeDeviceProfileTTL = 7 * 24 * time.Hour @@ -304,9 +304,8 @@ func extractClaudeDeviceProfile(headers http.Header, cfg *config.Config) (Claude return ClaudeDeviceProfile{}, false } - userAgent := strings.TrimSpace(headers.Get("User-Agent")) - version, ok := parseClaudeCLIVersion(userAgent) - if !ok || !claudeCodeNativeUserAgentPattern.MatchString(userAgent) { + userAgent, ok := parseClaudeCodeUserAgent(headers.Get("User-Agent")) + if !ok { return ClaudeDeviceProfile{}, false } @@ -320,12 +319,12 @@ func extractClaudeDeviceProfile(headers http.Header, cfg *config.Config) (Claude runtimeVersion = baseline.RuntimeVersion } profile := ClaudeDeviceProfile{ - UserAgent: userAgent, + UserAgent: userAgent.profileUserAgent, PackageVersion: packageVersion, RuntimeVersion: runtimeVersion, OS: firstNonEmptyHeader(headers, "X-Stainless-Os", baseline.OS), Arch: firstNonEmptyHeader(headers, "X-Stainless-Arch", baseline.Arch), - version: version, + version: userAgent.version, hasVersion: true, } return profile, true @@ -618,10 +617,23 @@ func (value claudeDeviceProfileKVValue) ToProfile() ClaudeDeviceProfile { return profile } -func ApplyClaudeDeviceProfileHeaders(r *http.Request, profile ClaudeDeviceProfile) { +// ApplyClaudeDeviceProfileHeaders applies a stable software identity. An optional +// UA from a confirmed native caller supplies only this request's extension fields. +func ApplyClaudeDeviceProfileHeaders(r *http.Request, profile ClaudeDeviceProfile, requestUserAgent ...string) { if r == nil { return } + userAgent := profile.UserAgent + if len(requestUserAgent) > 0 { + stable, validStable := parseClaudeCodeUserAgent(profile.UserAgent) + current, validCurrent := parseClaudeCodeUserAgent(requestUserAgent[0]) + if validStable && validCurrent { + userAgent = stable.profileUserAgent + if len(current.requestFields) > 0 { + userAgent = strings.TrimSuffix(userAgent, ")") + "," + strings.Join(current.requestFields, ",") + ")" + } + } + } for _, headerName := range []string{ "User-Agent", "X-Stainless-Package-Version", @@ -631,21 +643,21 @@ func ApplyClaudeDeviceProfileHeaders(r *http.Request, profile ClaudeDeviceProfil } { r.Header.Del(headerName) } - r.Header.Set("User-Agent", profile.UserAgent) + r.Header.Set("User-Agent", userAgent) r.Header.Set("X-Stainless-Package-Version", profile.PackageVersion) r.Header.Set("X-Stainless-Runtime-Version", profile.RuntimeVersion) r.Header.Set("X-Stainless-Os", profile.OS) r.Header.Set("X-Stainless-Arch", profile.Arch) } -// DefaultClaudeVersion returns the version string (e.g. "2.1.280") from the +// DefaultClaudeVersion returns the version string (e.g. "2.1.321") from the // current baseline device profile. It extracts the version from the User-Agent. func DefaultClaudeVersion(cfg *config.Config) string { profile := defaultClaudeDeviceProfile(cfg) if version, ok := parseClaudeCLIVersion(profile.UserAgent); ok { return strconv.Itoa(version.major) + "." + strconv.Itoa(version.minor) + "." + strconv.Itoa(version.patch) } - return "2.1.280" + return "2.1.321" } func ApplyClaudeDefaultDeviceProfileHeaders(r *http.Request, cfg *config.Config) { diff --git a/internal/runtime/executor/helps/claude_device_profile_test.go b/internal/runtime/executor/helps/claude_device_profile_test.go index d1b2777e717..1ab910173d4 100644 --- a/internal/runtime/executor/helps/claude_device_profile_test.go +++ b/internal/runtime/executor/helps/claude_device_profile_test.go @@ -130,7 +130,7 @@ func TestResolveClaudeDeviceProfileLocalUsesBaselineForInvalidSignals(t *testing func TestResolveClaudeDeviceProfileLocalKeepsExactMeasuredSoftwareBaseline(t *testing.T) { ResetClaudeDeviceProfileCache() auth := &cliproxyauth.Auth{ID: "auth-newer-patch-signals"} - headers := claudeDeviceHeaders("claude-cli/2.1.281 (external, cli)") + headers := claudeDeviceHeaders("claude-cli/2.1.322 (external, cli)") profile := resolveClaudeDeviceProfileLocal(auth, "api-key", headers, nil) baseline := defaultClaudeDeviceProfile(nil) @@ -280,7 +280,7 @@ func TestResolveClaudeDeviceProfileRequiredHomeSeparatesVSCodeAgentSDKFromCLI(t if errCLI != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() CLI error = %v", errCLI) } - vscodeUA := "claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)" + vscodeUA := "claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)" vscodeProfile, errVSCode := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders(vscodeUA), nil) if errVSCode != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() VSCode error = %v", errVSCode) @@ -362,19 +362,19 @@ func TestResolveClaudeDeviceProfilePreservesConfirmedClientAtBaselineVersion(t * client := newFakeClaudeDeviceProfileKVClient() useFakeClaudeDeviceProfileKVClient(t, client, false, nil) auth := &cliproxyauth.Auth{ID: "auth-baseline-entrypoint"} - headers := claudeDeviceHeaders("claude-cli/2.1.280 (external, cli)") - headers.Set("X-Stainless-Package-Version", "0.112.1") - headers.Set("X-Stainless-Runtime-Version", "v26.3.0") + headers := claudeDeviceHeaders("claude-cli/2.1.321 (external, cli)") + headers.Set("X-Stainless-Package-Version", "0.128.0") + headers.Set("X-Stainless-Runtime-Version", "v26.12.1") profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", headers, nil) if errProfile != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile) } - if profile.UserAgent != "claude-cli/2.1.280 (external, cli)" { + if profile.UserAgent != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("UserAgent = %q, want confirmed cli entrypoint preserved", profile.UserAgent) } - if profile.PackageVersion != "0.112.1" || profile.RuntimeVersion != "v26.3.0" { - t.Fatalf("software profile = %s/%s, want 0.112.1/v26.3.0", profile.PackageVersion, profile.RuntimeVersion) + if profile.PackageVersion != "0.128.0" || profile.RuntimeVersion != "v26.12.1" { + t.Fatalf("software profile = %s/%s, want 0.128.0/v26.12.1", profile.PackageVersion, profile.RuntimeVersion) } } @@ -384,27 +384,27 @@ func TestResolveClaudeDeviceProfileSeparatesVSCodeAgentSDKFromCLI(t *testing.T) useFakeClaudeDeviceProfileKVClient(t, client, false, nil) auth := &cliproxyauth.Auth{ID: "auth-subclient-isolation"} - cliHeaders := claudeDeviceHeaders("claude-cli/2.1.280 (external, cli)") - cliHeaders.Set("X-Stainless-Package-Version", "0.112.1") - cliHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0") + cliHeaders := claudeDeviceHeaders("claude-cli/2.1.321 (external, cli)") + cliHeaders.Set("X-Stainless-Package-Version", "0.128.0") + cliHeaders.Set("X-Stainless-Runtime-Version", "v26.12.1") cliProfile, errCLI := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", cliHeaders, nil) if errCLI != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() CLI error = %v", errCLI) } - vscodeUA := "claude-cli/2.1.280 (external, claude-vscode, agent-sdk/0.3.220)" + vscodeUA := "claude-cli/2.1.321 (external, claude-vscode, agent-sdk/0.3.220)" vscodeHeaders := claudeDeviceHeaders(vscodeUA) - vscodeHeaders.Set("X-Stainless-Package-Version", "0.112.1") - vscodeHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0") + vscodeHeaders.Set("X-Stainless-Package-Version", "0.128.0") + vscodeHeaders.Set("X-Stainless-Runtime-Version", "v26.12.1") vscodeProfile, errVSCode := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", vscodeHeaders, nil) if errVSCode != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() VSCode error = %v", errVSCode) } - sdkTSUA := "claude-cli/2.1.280 (external, sdk-ts, agent-sdk/0.3.220)" + sdkTSUA := "claude-cli/2.1.321 (external, sdk-ts, agent-sdk/0.3.220)" sdkTSHeaders := claudeDeviceHeaders(sdkTSUA) - sdkTSHeaders.Set("X-Stainless-Package-Version", "0.112.1") - sdkTSHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0") + sdkTSHeaders.Set("X-Stainless-Package-Version", "0.128.0") + sdkTSHeaders.Set("X-Stainless-Runtime-Version", "v26.12.1") sdkTSProfile, errSDKTS := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", sdkTSHeaders, nil) if errSDKTS != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() sdk-ts error = %v", errSDKTS) @@ -415,25 +415,25 @@ func TestResolveClaudeDeviceProfileSeparatesVSCodeAgentSDKFromCLI(t *testing.T) NativeEntrypoints: []string{"t3-code", "custom-ide"}, }, } - t3CodeUA := "claude-cli/2.1.280 (external, t3-code, agent-sdk/0.3.276)" + t3CodeUA := "claude-cli/2.1.321 (external, t3-code, agent-sdk/0.3.276)" t3CodeHeaders := claudeDeviceHeaders(t3CodeUA) - t3CodeHeaders.Set("X-Stainless-Package-Version", "0.112.1") - t3CodeHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0") + t3CodeHeaders.Set("X-Stainless-Package-Version", "0.128.0") + t3CodeHeaders.Set("X-Stainless-Runtime-Version", "v26.12.1") t3CodeProfile, errT3Code := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", t3CodeHeaders, cfg) if errT3Code != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() t3-code error = %v", errT3Code) } - customIDEUA := "claude-cli/2.1.280 (external, custom-ide, agent-sdk/0.3.276)" + customIDEUA := "claude-cli/2.1.321 (external, custom-ide, agent-sdk/0.3.276)" customIDEHeaders := claudeDeviceHeaders(customIDEUA) - customIDEHeaders.Set("X-Stainless-Package-Version", "0.112.1") - customIDEHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0") + customIDEHeaders.Set("X-Stainless-Package-Version", "0.128.0") + customIDEHeaders.Set("X-Stainless-Runtime-Version", "v26.12.1") customIDEProfile, errCustomIDE := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", customIDEHeaders, cfg) if errCustomIDE != nil { t.Fatalf("ResolveClaudeDeviceProfileRequired() custom-ide error = %v", errCustomIDE) } - if cliProfile.UserAgent != "claude-cli/2.1.280 (external, cli)" { + if cliProfile.UserAgent != "claude-cli/2.1.321 (external, cli)" { t.Fatalf("CLI UserAgent = %q, want CLI profile", cliProfile.UserAgent) } if vscodeProfile.UserAgent != vscodeUA { diff --git a/internal/runtime/executor/helps/claude_diagnostics.go b/internal/runtime/executor/helps/claude_diagnostics.go index 25ca3718894..194c12c46a4 100644 --- a/internal/runtime/executor/helps/claude_diagnostics.go +++ b/internal/runtime/executor/helps/claude_diagnostics.go @@ -796,6 +796,12 @@ func IsClaudeProbeOrHelperRequest(body []byte) bool { return isClaudeProbeRequest(body) || isClaudeTitleHelperRequest(body) } +// IsClaudeQuotaProbeRequest reports the max_tokens=1 quota probe alone, +// excluding the session-title helper that IsClaudeProbeOrHelperRequest also covers. +func IsClaudeQuotaProbeRequest(body []byte) bool { + return isClaudeProbeRequest(body) +} + // IsClaudeSubagentRequest reports whether the incoming request originates from // or represents a Claude Code subagent. func IsClaudeSubagentRequest(headers http.Header, body []byte) bool { @@ -833,6 +839,9 @@ func ClaudePayloadHas1hTTL(payload []byte) bool { if len(payload) == 0 || !jsonMayContainASCII(payload, `"1h"`) || !ojson.ValidBytes(payload) { return false } + if topCC := ojson.GetBytesNoCopy(payload, "cache_control"); topCC.IsObject() && topCC.Get("ttl").String() == "1h" { + return true + } has1h := false checkBlock := func(item ojson.Result) bool { cc := item.Get("cache_control") diff --git a/internal/runtime/executor/helps/claude_json_prefilter_fuzz_test.go b/internal/runtime/executor/helps/claude_json_prefilter_fuzz_test.go index bb206afa3c5..2cc2328e976 100644 --- a/internal/runtime/executor/helps/claude_json_prefilter_fuzz_test.go +++ b/internal/runtime/executor/helps/claude_json_prefilter_fuzz_test.go @@ -58,6 +58,9 @@ func referenceClaudePayloadHas1hTTL(payload []byte) bool { if len(payload) == 0 || !ojson.ValidBytes(payload) { return false } + if topCC := ojson.GetBytes(payload, "cache_control"); topCC.IsObject() && topCC.Get("ttl").String() == "1h" { + return true + } has1h := false checkBlock := func(item ojson.Result) bool { cc := item.Get("cache_control") diff --git a/internal/runtime/executor/helps/codex_hosted_images.go b/internal/runtime/executor/helps/codex_hosted_images.go index 48318e7d2dd..1e9f924276f 100644 --- a/internal/runtime/executor/helps/codex_hosted_images.go +++ b/internal/runtime/executor/helps/codex_hosted_images.go @@ -1,6 +1,7 @@ package helps import ( + "context" "net/http" "strings" @@ -8,9 +9,12 @@ import ( "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/thinking" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" ) +const codexResponsesLiteHeaderName = "X-OpenAI-Internal-Codex-Responses-Lite" + // CodexModelUsesResponsesLite identifies models whose native Codex route uses Responses Lite. func CodexModelUsesResponsesLite(model string) bool { slug := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(model).ModelName)) @@ -22,7 +26,7 @@ func CodexModelUsesResponsesLite(model string) bool { // CodexResponsesLiteRequestDoc reads the route hints without changing the payload. func CodexResponsesLiteRequestDoc(doc *ojson.Doc, headers http.Header) bool { - if strings.EqualFold(strings.TrimSpace(headers.Get("X-OpenAI-Internal-Codex-Responses-Lite")), "true") { + if strings.EqualFold(strings.TrimSpace(headers.Get(codexResponsesLiteHeaderName)), "true") { return true } if doc == nil { @@ -32,15 +36,82 @@ func CodexResponsesLiteRequestDoc(doc *ojson.Doc, headers http.Header) bool { return value.Type == ojson.True || (value.Type == ojson.String && strings.EqualFold(strings.TrimSpace(value.String()), "true")) } +// CodexResponsesLiteRequestForAuthDoc reports whether the request uses the +// Codex Responses-Lite protocol after resolving client_metadata, client +// headers, credential header: attributes, and models.json override_header. +func CodexResponsesLiteRequestForAuthDoc(doc *ojson.Doc, modelName string, auth *cliproxyauth.Auth, headers http.Header) bool { + if CodexResponsesLiteRequestDoc(doc, nil) { + return true + } + var effective string + if headers != nil { + effective = strings.TrimSpace(headers.Get(codexResponsesLiteHeaderName)) + } + if auth != nil && len(auth.Attributes) > 0 { + resolved := (&http.Request{Header: http.Header{}}).WithContext(context.Background()) + util.ApplyCustomHeadersFromAttrs(resolved, auth.Attributes, headers) + if opLite := strings.TrimSpace(resolved.Header.Get(codexResponsesLiteHeaderName)); opLite != "" { + effective = opLite + } + } + if overrides := registry.ModelOverrideHeaders(modelName); len(overrides) > 0 { + for key, value := range overrides { + if strings.EqualFold(strings.TrimSpace(key), codexResponsesLiteHeaderName) { + effective = strings.TrimSpace(value) + } + } + } + return strings.EqualFold(effective, "true") +} + +// CodexRequestHasReasoningForHostedImageDoc reports whether the request carries +// a reasoning configuration (top-level reasoning.effort, configuration_update +// input item, or model thinking suffix). Upstream Codex /responses binds hosted +// image_generation to gpt-image-2-codex, which rejects complexity whenever +// reasoning effort is present on Responses-Lite models (#7270), whereas +// standard Responses requests without reasoning effort succeed (#7384). +func CodexRequestHasReasoningForHostedImageDoc(doc *ojson.Doc, model string) bool { + if thinking.HasThinkingSuffix(model) { + return true + } + if doc == nil { + return false + } + if thinking.HasThinkingSuffix(doc.Get("model").String()) { + return true + } + if reasoning := doc.Get("reasoning"); reasoning.Exists() && reasoning.Type != ojson.Null { + if reasoning.IsObject() { + if effort := reasoning.Get("effort"); effort.Exists() && effort.Type != ojson.Null && strings.TrimSpace(effort.String()) != "" { + return true + } + } else if strings.TrimSpace(reasoning.String()) != "" { + return true + } + } + if input := doc.Get("input"); input.IsArray() { + for _, item := range input.Array() { + if strings.TrimSpace(item.Get("type").String()) != "configuration_update" { + continue + } + if effort := item.Get("reasoning.effort"); effort.Exists() && effort.Type != ojson.Null && strings.TrimSpace(effort.String()) != "" { + return true + } + } + } + return false +} + // ValidateCodexHostedImageGenerationDoc checks the final, configured Responses payload // after all outbound headers, including credential and model overrides, are resolved. -// Native Codex Responses Lite cannot serve forced hosted images. API-key, SIWC, +// Native Codex Responses Lite cannot serve forced hosted images, while API-key, SIWC, // explicit passthrough, and standalone Images API requests retain their own contracts. func ValidateCodexHostedImageGenerationDoc(doc *ojson.Doc, upstreamHeaders http.Header, cfg *config.Config, nativeCodex bool) error { if !nativeCodex || doc == nil || (cfg != nil && cfg.DisableImageGeneration == config.DisableImageGenerationPassthrough) { return nil } - if !CodexModelUsesResponsesLite(doc.Get("model").String()) && !CodexResponsesLiteRequestDoc(doc, upstreamHeaders) { + model := doc.Get("model").String() + if !CodexModelUsesResponsesLite(model) && !CodexResponsesLiteRequestDoc(doc, upstreamHeaders) { return nil } if !codexRequestForcesImageGenerationDoc(doc) { diff --git a/internal/runtime/executor/helps/home_usage_test.go b/internal/runtime/executor/helps/home_usage_test.go new file mode 100644 index 00000000000..acf5ca9b189 --- /dev/null +++ b/internal/runtime/executor/helps/home_usage_test.go @@ -0,0 +1,50 @@ +package helps + +import ( + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" +) + +func TestUsageReporterHomeAdditionalModelAccounting(t *testing.T) { + for _, mode := range []string{"local", "initial", "reused"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + if mode != "local" { + ctx = usage.WithHomeSelection(ctx, usage.HomeSelection{ + Reused: mode == "reused", Model: "route-model", CredentialPolicy: "codex-alpha-search-v1", + }) + } + reporter := NewUsageReporter(ctx, "codex", "text-model", nil) + primary := reporter.buildRecord(usage.Detail{TotalTokens: 12}, false) + image, ok := reporter.buildAdditionalModelRecord("image-model", usage.Detail{TotalTokens: 5}) + if !ok { + t.Fatal("missing image usage") + } + if image.RequestID == "" || image.RequestID == primary.RequestID { + t.Fatal("primary and image usage must retain independent execution IDs") + } + if image.Model != "image-model" || image.Detail.TotalTokens != 5 || primary.Detail.TotalTokens != 12 { + t.Fatal("scheduling metadata changed billable usage") + } + if mode == "local" { + if primary.HomeSelection != nil || image.HomeSelection != nil { + t.Fatal("local execution acquired Home accounting metadata") + } + return + } + if image.HomeSelection == nil || !image.HomeSelection.Additional || image.HomeSelection.Reused != (mode == "reused") { + t.Fatalf("image accounting metadata=%+v", image.HomeSelection) + } + if image.HomeSelection.Model != "route-model" || image.HomeSelection.CredentialPolicy != "codex-alpha-search-v1" { + t.Fatalf("image lost dispatch scope: %+v", image.HomeSelection) + } + if primary.HomeSelection == nil || primary.HomeSelection.Additional || primary.HomeSelection.Reused != (mode == "reused") { + t.Fatalf("image changed primary accounting: %+v", primary.HomeSelection) + } + if usage.HomeSelectionFromContext(ctx).Additional || reporter.buildRecord(usage.Detail{}, false).HomeSelection.Additional { + t.Fatal("image mutated context or reporter metadata") + } + }) + } +} diff --git a/internal/runtime/executor/helps/payload_finalizer.go b/internal/runtime/executor/helps/payload_finalizer.go index a18cd6fdce5..b1d45f20869 100644 --- a/internal/runtime/executor/helps/payload_finalizer.go +++ b/internal/runtime/executor/helps/payload_finalizer.go @@ -9,6 +9,7 @@ import ( "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" ) @@ -47,11 +48,15 @@ func (f PayloadFinalizer) FinalizeDoc(doc *ojson.Doc) { } // NewPayloadFinalizer snapshots matching context before built-in request mutations. -func NewPayloadFinalizer(cfg *config.Config, executor, model, protocol, root string, original []byte, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) PayloadFinalizer { +func NewPayloadFinalizer(cfg *config.Config, executor, model, protocol, root string, original []byte, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, auths ...*cliproxyauth.Auth) PayloadFinalizer { requestedModel := PayloadRequestedModel(opts, req.Model) requestPath := PayloadRequestPath(opts) headers := opts.Headers.Clone() source := append([]byte(nil), original...) + var auth *cliproxyauth.Auth + if len(auths) > 0 { + auth = auths[0] + } var formatCallback func(string) if opts.Metadata != nil { formatCallback, _ = opts.Metadata[cliproxyexecutor.FinalizedResponseFormatCallbackMetadataKey].(func(string)) @@ -61,7 +66,7 @@ func NewPayloadFinalizer(cfg *config.Config, executor, model, protocol, root str key := &body[:1][0] if v, ok := payloadDocProbes.LoadAndDelete(key); ok { if doc, okDoc := v.(*ojson.Doc); okDoc && doc != nil { - ApplyPayloadConfigWithRequestForExecutorDoc(cfg, executor, model, protocol, opts.SourceFormat.String(), root, doc, source, requestedModel, requestPath, headers) + applyPayloadConfigWithTrackedPathsForExecutorAndAuthDoc(cfg, executor, model, protocol, opts.SourceFormat.String(), root, doc, source, requestedModel, requestPath, headers, auth) if formatCallback != nil { formatCallback(strings.Clone(strings.TrimSpace(doc.Get(buildPayloadPath(root, "response_format")).String()))) } @@ -69,7 +74,12 @@ func NewPayloadFinalizer(cfg *config.Config, executor, model, protocol, root str return nil } } - out := ApplyPayloadConfigWithRequestForExecutor(cfg, executor, model, protocol, opts.SourceFormat.String(), root, body, source, requestedModel, requestPath, headers) + doc := ojson.NewDoc(body) + _, mutated := applyPayloadConfigWithTrackedPathsForExecutorAndAuthDoc(cfg, executor, model, protocol, opts.SourceFormat.String(), root, doc, source, requestedModel, requestPath, headers, auth) + out := body + if mutated { + out = doc.Bytes() + } if formatCallback != nil { formatCallback(strings.Clone(strings.TrimSpace(ojson.GetBytesNoCopy(out, buildPayloadPath(root, "response_format")).String()))) } diff --git a/internal/runtime/executor/helps/payload_helpers.go b/internal/runtime/executor/helps/payload_helpers.go index b53bafba4e8..f6e4150a3ce 100644 --- a/internal/runtime/executor/helps/payload_helpers.go +++ b/internal/runtime/executor/helps/payload_helpers.go @@ -6,14 +6,24 @@ import ( "reflect" "strconv" "strings" + "sync/atomic" "github.com/MoeForge/moego/ojson" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" ) +var payloadConfigEvaluations atomic.Int64 + +// PayloadConfigEvaluationCount returns the cumulative number of payload rule +// engine evaluations across the process. +func PayloadConfigEvaluationCount() int64 { + return payloadConfigEvaluations.Load() +} + // ApplyPayloadConfigWithRoot behaves like applyPayloadConfig but treats all parameter // paths as relative to the provided root path and restricts matches to the given // protocol when supplied. Defaults are checked @@ -91,6 +101,11 @@ func ApplyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target } func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, targetExecutor, model, protocol, fromProtocol, root string, doc *ojson.Doc, original []byte, requestedModel, requestPath string, headers http.Header, trackedPaths ...string) (map[string]bool, bool) { + return applyPayloadConfigWithTrackedPathsForExecutorAndAuthDoc(cfg, targetExecutor, model, protocol, fromProtocol, root, doc, original, requestedModel, requestPath, headers, nil, trackedPaths...) +} + +func applyPayloadConfigWithTrackedPathsForExecutorAndAuthDoc(cfg *config.Config, targetExecutor, model, protocol, fromProtocol, root string, doc *ojson.Doc, original []byte, requestedModel, requestPath string, headers http.Header, auth *cliproxyauth.Auth, trackedPaths ...string) (map[string]bool, bool) { + payloadConfigEvaluations.Add(1) touched := make(map[string]bool) if doc == nil || doc.Len() == 0 { return touched, false @@ -133,6 +148,59 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target model = strings.TrimSpace(model) requestedModel = strings.TrimSpace(requestedModel) if model != "" || requestedModel != "" { + toolsPath := buildPayloadPath(root, "tools") + toolChoicePath := buildPayloadPath(root, "tool_choice") + autoInjectedToolIdx := -1 + if isCodexTargetExecutor(targetExecutor) && + cfg.DisableImageGeneration == config.DisableImageGenerationOff && + len(original) > 0 && + !rawToolsContainImageGeneration(original, toolsPath) { + autoInjectedToolIdx = findLastImageGenerationToolIndexDoc(doc, toolsPath) + } + origTools := ojson.GetBytesNoCopy(original, toolsPath) + originalHadEmptyTools := len(original) > 0 && origTools.Exists() && origTools.IsArray() && len(origTools.Array()) == 0 + userWroteToolChoice := false + noteRuleWrite := func(resolvedPath string) { + markTouched(resolvedPath) + if toolChoicePath != "" && payloadRuleTargetsPath(resolvedPath, toolChoicePath) { + userWroteToolChoice = true + return + } + if autoInjectedToolIdx < 0 || toolsPath == "" || !payloadRuleTargetsPath(resolvedPath, toolsPath) { + return + } + if resolvedPath == toolsPath { + autoInjectedToolIdx = -1 + return + } + rest := strings.TrimPrefix(resolvedPath, toolsPath+".") + if idx := strings.IndexByte(rest, '.'); idx >= 0 { + rest = rest[:idx] + } + if elemIdx, errIdx := strconv.Atoi(rest); errIdx == nil && elemIdx == autoInjectedToolIdx { + autoInjectedToolIdx = -1 + } + } + noteRuleDelete := func(resolvedPath string) { + if autoInjectedToolIdx < 0 || toolsPath == "" || !payloadRuleTargetsPath(resolvedPath, toolsPath) { + return + } + if resolvedPath == toolsPath { + autoInjectedToolIdx = -1 + return + } + rest := strings.TrimPrefix(resolvedPath, toolsPath+".") + if strings.Contains(rest, ".") { + return + } + if delIdx, errIdx := strconv.Atoi(rest); errIdx == nil { + if delIdx == autoInjectedToolIdx { + autoInjectedToolIdx = -1 + } else if delIdx >= 0 && delIdx < autoInjectedToolIdx { + autoInjectedToolIdx-- + } + } + } candidates := payloadModelCandidates(model, requestedModel) appliedDefaults := make(map[string]struct{}) var createdByDefault map[string]struct{} @@ -223,7 +291,7 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target } mutated = true appliedDefaults[resolvedPath] = struct{}{} - markTouched(resolvedPath) + noteRuleWrite(resolvedPath) } } } @@ -255,7 +323,7 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target } mutated = true appliedDefaults[resolvedPath] = struct{}{} - markTouched(resolvedPath) + noteRuleWrite(resolvedPath) } } } @@ -276,7 +344,7 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target mutated = true } if applied { - markTouched(resolvedPath) + noteRuleWrite(resolvedPath) } } } @@ -302,7 +370,7 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target mutated = true } if applied { - markTouched(resolvedPath) + noteRuleWrite(resolvedPath) } } } @@ -327,16 +395,126 @@ func applyPayloadConfigWithTrackedPathsForExecutorDoc(cfg *config.Config, target } if existed { mutated = true + noteRuleDelete(resolvedPath) } markTouched(resolvedPath) } } } + if autoInjectedToolIdx >= 0 && !(userWroteToolChoice && docToolChoiceSelectsImageGeneration(doc, toolChoicePath)) { + if codexConfiguredPayloadIncompatibleWithAutoImageGenDoc(doc, root, model, auth, headers) { + if retractAutoInjectedImageGenerationToolAtIndexDoc(doc, toolsPath, autoInjectedToolIdx, originalHadEmptyTools) { + mutated = true + } + } + } } } return touched, mutated } +func findLastImageGenerationToolIndexDoc(doc *ojson.Doc, toolsPath string) int { + if doc == nil { + return -1 + } + tools := doc.Get(toolsPath) + if !tools.Exists() || !tools.IsArray() { + return -1 + } + items := tools.Array() + for idx := len(items) - 1; idx >= 0; idx-- { + if strings.TrimSpace(items[idx].Get("type").String()) == "image_generation" { + return idx + } + } + return -1 +} + +func rawToolsContainImageGeneration(raw []byte, toolsPath string) bool { + if len(raw) == 0 { + return false + } + tools := ojson.GetBytesNoCopy(raw, toolsPath) + if !tools.Exists() || !tools.IsArray() { + return false + } + for _, tool := range tools.Array() { + if strings.TrimSpace(tool.Get("type").String()) == "image_generation" { + return true + } + } + return false +} + +func docToolChoiceSelectsImageGeneration(doc *ojson.Doc, toolChoicePath string) bool { + if doc == nil || toolChoicePath == "" { + return false + } + choice := doc.Get(toolChoicePath) + if !choice.Exists() { + return false + } + if choice.Type == ojson.String { + return strings.EqualFold(strings.TrimSpace(choice.String()), "image_generation") + } + if !choice.IsObject() { + return false + } + switch strings.ToLower(strings.TrimSpace(choice.Get("type").String())) { + case "image_generation": + return true + case "tool": + return strings.EqualFold(strings.TrimSpace(choice.Get("name").String()), "image_generation") + case "allowed_tools": + allowed := choice.Get("tools") + if !allowed.IsArray() { + return false + } + for _, tool := range allowed.Array() { + if strings.EqualFold(strings.TrimSpace(tool.Get("type").String()), "image_generation") { + return true + } + } + } + return false +} + +func codexConfiguredPayloadIncompatibleWithAutoImageGenDoc(doc *ojson.Doc, root, fallbackModel string, auth *cliproxyauth.Auth, headers http.Header) bool { + if doc == nil { + return false + } + configuredModel := strings.TrimSpace(doc.Get(buildPayloadPath(root, "model")).String()) + if configuredModel == "" { + configuredModel = strings.TrimSpace(fallbackModel) + } + modelName := strings.TrimSpace(thinking.ParseSuffix(configuredModel).ModelName) + if CodexResponsesLiteRequestForAuthDoc(doc, modelName, auth, headers) { + return true + } + if strings.HasSuffix(strings.ToLower(modelName), "spark") { + return true + } + return CodexModelUsesResponsesLite(modelName) && CodexRequestHasReasoningForHostedImageDoc(doc, configuredModel) +} + +func retractAutoInjectedImageGenerationToolAtIndexDoc(doc *ojson.Doc, toolsPath string, targetIdx int, originalHadEmptyTools bool) bool { + tools := doc.Get(toolsPath) + if !tools.Exists() || !tools.IsArray() { + return false + } + toolItems := tools.Array() + if targetIdx < 0 || targetIdx >= len(toolItems) { + return false + } + if strings.TrimSpace(toolItems[targetIdx].Get("type").String()) != "image_generation" { + return false + } + if len(toolItems) == 1 && !originalHadEmptyTools { + return doc.Delete(toolsPath) == nil + } + return doc.Delete(toolsPath+"."+strconv.Itoa(targetIdx)) == nil +} + // ApplyPayloadConfigWithRequestTracked applies payload config and reports whether // an applied rule targeted trackedPath or one of its descendants. func ApplyPayloadConfigWithRequestTracked(cfg *config.Config, model, protocol, fromProtocol, root string, payload, original []byte, requestedModel string, requestPath string, headers http.Header, trackedPath string) ([]byte, bool) { diff --git a/internal/runtime/executor/helps/proxy_helpers.go b/internal/runtime/executor/helps/proxy_helpers.go index b0464c633eb..cd725e88a62 100644 --- a/internal/runtime/executor/helps/proxy_helpers.go +++ b/internal/runtime/executor/helps/proxy_helpers.go @@ -460,8 +460,7 @@ func authMetadataString(metadata map[string]any, key string) string { // must bypass the shared cache (zero/negative pool bounds or a fenced config // generation). It uses the kind's own builder so an unpooled uTLS transport keeps // its TLS fingerprint. DisableKeepAlives is always set: no cache owns this transport, -// so any connection left in its idle list would survive until an idle timer fires (or -// forever on a transport built with IdleConnTimeout=0, such as the Claude Code one). +// so any connection left in its idle list would survive until an idle timer fires. func unpooledTransport(build func() (*http.Transport, error)) *http.Transport { tr, errBuild := build() if errBuild != nil { diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index fbff997aae2..1fbbf6963f8 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -32,6 +32,7 @@ type UsageReporter struct { requestID string traceID string selectionID string + homeSelection *usage.HomeSelection provider string baseURL string executorType string @@ -131,6 +132,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox requestID: uuid.NewString(), traceID: traceID, selectionID: usage.SelectionIDFromContext(ctx), + homeSelection: usage.HomeSelectionFromContext(ctx), provider: provider, baseURL: baseURL, model: model, @@ -542,6 +544,13 @@ func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.De } rec := r.buildRecordForModel(model, detail, false, usage.Failure{}) rec.RequestID = uuid.NewString() + if rec.HomeSelection != nil { + // Split billing details do not represent another scheduled request. + // Mark them separately without changing the primary usage metadata. + selection := *rec.HomeSelection + selection.Additional = true + rec.HomeSelection = &selection + } return rec, true } @@ -680,6 +689,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f RequestID: r.requestID, TraceID: r.traceID, SelectionID: r.selectionID, + HomeSelection: r.homeSelection, Provider: r.provider, BaseURL: r.baseURL, ExecutorType: r.executorType, diff --git a/internal/runtime/executor/helps/utls_client.go b/internal/runtime/executor/helps/utls_client.go index cdf1a63706a..42cf15b6e0b 100644 --- a/internal/runtime/executor/helps/utls_client.go +++ b/internal/runtime/executor/helps/utls_client.go @@ -374,8 +374,15 @@ func newClaudeCodeRoundTripperForDialerAndRoots(dialer proxy.Dialer, roots *x509 // so resumption never crosses proxy boundaries. sessionCache := tls.NewLRUClientSessionCache(claudeCodeSessionCacheCapacity) + // Default to short-lived connections (MaxIdleConnsPerHost < 0) so OAuth and + // unpooled API keys never park idle connections that intermediate NATs or + // transparent proxies may silently drop, while keeping DisableKeepAlives false + // so net/http does not overwrite the native Connection: keep-alive wire header. + // Credentials that explicitly opt into connection-pool override MaxIdleConnsPerHost + // and IdleConnTimeout in pooledTransport. transport := &http.Transport{ - ForceAttemptHTTP2: false, + ForceAttemptHTTP2: false, + MaxIdleConnsPerHost: -1, DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { var ( conn net.Conn diff --git a/internal/runtime/executor/helps/utls_proxy_pool_test.go b/internal/runtime/executor/helps/utls_proxy_pool_test.go index e5f3ec3fe7a..6eda4370492 100644 --- a/internal/runtime/executor/helps/utls_proxy_pool_test.go +++ b/internal/runtime/executor/helps/utls_proxy_pool_test.go @@ -2,11 +2,15 @@ package helps import ( "context" + gotls "crypto/tls" + "crypto/x509" "io" "net" "net/http" "net/http/httptest" + "net/http/httptrace" "strconv" + "sync" "sync/atomic" "testing" @@ -96,15 +100,147 @@ func TestUtlsAnthropicTransportPoolsPerCredential(t *testing.T) { t.Fatal("max-idle-conns-per-host: 0 must disable keep-alives on the Anthropic transport") } - // Credentials that did not opt in keep the historical shared transport. - plain := func(id string) http.RoundTripper { + // Credentials that did not opt in keep the shared Claude Code transport with + // idle connection pooling disabled, while opted-in credentials enable pooling. + plain := func(id string) *http.Transport { t.Helper() client := NewUtlsHTTPClient(context.Background(), cfg, &cliproxyauth.Auth{ID: id}, 0) - return client.Transport.(*fallbackRoundTripper).anthropicRoundTripper() + tr, ok := client.Transport.(*fallbackRoundTripper).anthropicRoundTripper().(*http.Transport) + if !ok { + t.Fatalf("unpooled anthropic transport = %T, want *http.Transport", client.Transport.(*fallbackRoundTripper).anthropicRoundTripper()) + } + return tr } if plain("key-a") != plain("key-b") { t.Fatal("unpooled credentials must keep sharing the cached Claude Code transport") } + if plain("key-a").MaxIdleConnsPerHost >= 0 { + t.Fatalf("unpooled Claude Code transport MaxIdleConnsPerHost = %d, want < 0 to disable idle connection reuse", plain("key-a").MaxIdleConnsPerHost) + } +} + +// TestUtlsClaudeCodeDefaultDisablesConnectionReusePreservingExplicitPool verifies +// that by default (including OAuth and unpooled API keys) the Claude Code uTLS +// transport closes connections after each request instead of leaving them in an +// idle pool, while an API key that explicitly configures connection-pool still +// reuses its pooled connection. +func TestUtlsClaudeCodeDefaultDisablesConnectionReusePreservingExplicitPool(t *testing.T) { + ResetProxyTransports() + t.Cleanup(ResetProxyTransports) + + certificate := newResumptionTestCertificate(t) + roots := x509.NewCertPool() + roots.AddCert(certificate.Leaf) + + var ( + newConns atomic.Int32 + headerMu sync.Mutex + connHeader []string + ) + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + headerMu.Lock() + connHeader = append(connHeader, r.Header.Get("Connection")) + headerMu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + server.TLS = &gotls.Config{ + Certificates: []gotls.Certificate{certificate}, + MinVersion: gotls.VersionTLS13, + NextProtos: []string{"http/1.1"}, + } + server.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + newConns.Add(1) + } + } + server.StartTLS() + defer server.Close() + + serverAddr := server.Listener.Addr().String() + dialer := contextDialerFunc(func(ctx context.Context, network, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, network, serverAddr) + }) + + doTwoRequests := func(t *testing.T, tr *http.Transport) []bool { + t.Helper() + t.Cleanup(tr.CloseIdleConnections) + var ( + mu sync.Mutex + reused []bool + ) + trace := &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + mu.Lock() + reused = append(reused, info.Reused) + mu.Unlock() + }, + } + client := &http.Client{Transport: tr} + for i := 0; i < 2; i++ { + req, errReq := http.NewRequestWithContext( + httptrace.WithClientTrace(context.Background(), trace), + http.MethodPost, + "https://api.anthropic.com/v1/messages?beta=true", + nil, + ) + if errReq != nil { + t.Fatalf("NewRequest: %v", errReq) + } + req.Header.Set("Connection", "keep-alive") + resp, errDo := client.Do(req) + if errDo != nil { + t.Fatalf("round %d Do: %v", i+1, errDo) + } + if _, errCopy := io.Copy(io.Discard, resp.Body); errCopy != nil { + t.Fatalf("round %d read body: %v", i+1, errCopy) + } + if errClose := resp.Body.Close(); errClose != nil { + t.Fatalf("round %d close body: %v", i+1, errClose) + } + } + mu.Lock() + defer mu.Unlock() + return append([]bool(nil), reused...) + } + + // 1. Default / OAuth transport: must NOT reuse connections across requests, + // while keeping the wire Connection: keep-alive header intact. + defaultTransport := newClaudeCodeRoundTripperForDialerAndRoots(dialer, roots) + reusedDefault := doTwoRequests(t, defaultTransport) + if got := newConns.Load(); got != 2 { + t.Fatalf("default transport connections = %d, want 2 (reused = %v)", got, reusedDefault) + } + if len(reusedDefault) != 2 || reusedDefault[0] || reusedDefault[1] { + t.Fatalf("default transport GotConn reused = %v, want [false false]", reusedDefault) + } + headerMu.Lock() + gotHeaders := append([]string(nil), connHeader...) + headerMu.Unlock() + for i, h := range gotHeaders { + if h != "keep-alive" { + t.Fatalf("request %d Connection header = %q, want \"keep-alive\"", i+1, h) + } + } + + // 2. Explicit connection-pool opt-in on claude-api-key: MUST still reuse + // connections across requests. + newConns.Store(0) + cfg := testConfig() + pooled, ok := pooledTransport(proxyPoolKindClaudeCode, "direct", cfg, poolOptInAuth("claude-key-pooled"), func() (*http.Transport, error) { + return newClaudeCodeRoundTripperForDialerAndRoots(dialer, roots), nil + }) + if !ok || pooled == nil { + t.Fatal("expected pooled Claude Code transport for opted-in credential") + } + reusedPooled := doTwoRequests(t, pooled) + if got := newConns.Load(); got != 1 { + t.Fatalf("pooled transport connections = %d, want 1 (reused = %v)", got, reusedPooled) + } + if len(reusedPooled) != 2 || reusedPooled[0] || !reusedPooled[1] { + t.Fatalf("pooled transport GotConn reused = %v, want [false true]", reusedPooled) + } } // TestUtlsLazyPoolsDoNotEvictUsedPool pins that a request only populates the pool diff --git a/internal/runtime/executor/openai_compat_executor_reasoning_test.go b/internal/runtime/executor/openai_compat_executor_reasoning_test.go index da3fa8c86ed..d5cda5fc543 100644 --- a/internal/runtime/executor/openai_compat_executor_reasoning_test.go +++ b/internal/runtime/executor/openai_compat_executor_reasoning_test.go @@ -2,6 +2,7 @@ package executor import ( "context" + "fmt" "io" "net/http" "net/http/httptest" @@ -56,3 +57,80 @@ func TestOpenAICompatExecutorUsesCompatibleClaudeTranslation(t *testing.T) { 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 { + name string + payload config.PayloadConfig + wantEffort string + }{ + {name: "standalone effort preserved", wantEffort: "high"}, + { + name: "final override wins", + payload: config.PayloadConfig{Override: []config.PayloadRule{ + {Models: models, Params: map[string]any{"reasoning_effort": "low"}}, + }}, + wantEffort: "low", + }, + { + name: "final filter is not restored", + payload: config.PayloadConfig{Filter: []config.PayloadFilterRule{ + {Models: models, Params: []string{"reasoning_effort"}}, + }}, + }, + } { + 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-1\",\"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-1","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.payload}) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test-key", + }, + } + req := cliproxyexecutor.Request{ + Model: "custom-openai", + Payload: []byte(`{"model":"custom-openai","max_tokens":16,"output_config":{"effort":"high"},"messages":[{"role":"user","content":"hi"}]}`), + } + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, Stream: stream} + if stream { + result, errStream := executor.ExecuteStream(context.Background(), auth, req, opts) + if errStream != nil { + t.Fatalf("ExecuteStream: %v", errStream) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk: %v", chunk.Err) + } + } + } else if _, errExecute := executor.Execute(context.Background(), auth, req, opts); errExecute != nil { + t.Fatalf("Execute: %v", errExecute) + } + body := <-bodies + effort := ojson.GetBytes(body, "reasoning_effort") + if effort.String() != tt.wantEffort || (tt.wantEffort == "" && effort.Exists()) { + t.Fatalf("upstream reasoning_effort = %s, want %q; body=%s", effort.Raw, tt.wantEffort, body) + } + }) + } + } +} diff --git a/internal/runtime/executor/payload_barrier_test.go b/internal/runtime/executor/payload_barrier_test.go index ac3b9ad95a8..c47d7af4cd6 100644 --- a/internal/runtime/executor/payload_barrier_test.go +++ b/internal/runtime/executor/payload_barrier_test.go @@ -21,7 +21,7 @@ import ( func TestPayloadBarrierCodexImageFilter(t *testing.T) { for _, transport := range []string{"http", "websocket"} { for _, stream := range []bool{false, true} { - for _, model := range []string{"gpt-5.4", "gpt-5.5"} { + for _, model := range []string{"gpt-5.4", "gpt-5.5", "gpt-6-astra", "gpt-6-sol"} { t.Run(transport+"/"+model+map[bool]string{false: "/execute", true: "/stream"}[stream], func(t *testing.T) { captured := make(chan []byte, 1) completed := []byte(`{"type":"response.completed","response":{"id":"resp_barrier","status":"completed","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) @@ -59,7 +59,7 @@ func TestPayloadBarrierCodexImageFilter(t *testing.T) { _, _ = w.Write(append(append([]byte("data: "), completed...), '\n', '\n')) })) defer server.Close() - cfg := &config.Config{Payload: config.PayloadConfig{Filter: []config.PayloadFilterRule{{Models: []config.PayloadModelRule{{Name: "gpt-5.4"}}, Params: []string{`tools.#(type=="image_generation")#`, "instructions", "prompt_cache_key"}}}}} + cfg := &config.Config{Payload: config.PayloadConfig{Filter: []config.PayloadFilterRule{{Models: []config.PayloadModelRule{{Name: "gpt-5.4"}, {Name: "gpt-6-astra"}}, Params: []string{`tools.#(type=="image_generation")#`, "instructions", "prompt_cache_key"}}}}} var executor cliproxyauth.ProviderExecutor = NewCodexExecutor(cfg) if transport == "websocket" { executor = NewCodexWebsocketsExecutor(cfg) @@ -81,11 +81,12 @@ func TestPayloadBarrierCodexImageFilter(t *testing.T) { t.Fatal(errExecute) } body := <-captured + filteredModel := model == "gpt-5.4" || model == "gpt-6-astra" image := ojson.GetBytes(body, `tools.#(type=="image_generation")`).Exists() - if image != (model == "gpt-5.5") { + if image != !filteredModel { t.Fatalf("image tool presence = %v: %s", image, body) } - if model == "gpt-5.4" && (ojson.GetBytes(body, "instructions").Exists() || ojson.GetBytes(body, "prompt_cache_key").Exists()) { + if filteredModel && (ojson.GetBytes(body, "instructions").Exists() || ojson.GetBytes(body, "prompt_cache_key").Exists()) { t.Fatalf("late injection undid filter: %s", body) } }) diff --git a/internal/runtime/executor/testdata/tool_search_disabled/dev_executor_golden.json b/internal/runtime/executor/testdata/tool_search_disabled/dev_executor_golden.json index aabd8c4b887..2080bdff494 100644 --- a/internal/runtime/executor/testdata/tool_search_disabled/dev_executor_golden.json +++ b/internal/runtime/executor/testdata/tool_search_disabled/dev_executor_golden.json @@ -3,8 +3,8 @@ "gateway/execute": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"type\":\"tool_search\",\"execution\":\"client\",\"name\":\"ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"},\"cache_control\":{\"type\":\"ephemeral\"}}]}", "gateway/stream": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"type\":\"tool_search\",\"execution\":\"client\",\"name\":\"ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"},\"cache_control\":{\"type\":\"ephemeral\"}}]}", "oauth/count_tokens": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":\"create a calendar event\"}],\"stream\":true,\"tools\":[{\"execution\":\"client\",\"name\":\"mcp__select_sight__disorder_ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"mcp__select_sight__push_exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"mcp__select_sight__elephant_lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}]}", - "oauth/execute": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"\u003csystem-reminder\u003e\\nAs you answer the user's questions, you can use the following context:\\n# currentDate\\nToday's date is 2026-10-11.\\n\\n IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.\\n\u003c/system-reminder\u003e\\n\\n\"},{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"execution\":\"client\",\"name\":\"mcp__select_sight__disorder_ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"mcp__select_sight__push_exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"mcp__select_sight__elephant_lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}],\"system\":[{\"type\":\"text\",\"text\":\"x-anthropic-billing-header: cc_version=2.1.280.283; cc_entrypoint=cli; cch=0; cc_prompt_id=ebb599d3-6b98-4fcc-bc9c-99d7dcbd62b9; cc_turn_origin=human;\"},{\"type\":\"text\",\"text\":\"You are Claude Code, Anthropic's official CLI for Claude.\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}],\"diagnostics\":{\"previous_message_id\":null}}", - "oauth/stream": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"\u003csystem-reminder\u003e\\nAs you answer the user's questions, you can use the following context:\\n# currentDate\\nToday's date is 2026-10-11.\\n\\n IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.\\n\u003c/system-reminder\u003e\\n\\n\"},{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"execution\":\"client\",\"name\":\"mcp__select_sight__disorder_ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"mcp__select_sight__push_exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"mcp__select_sight__elephant_lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}],\"system\":[{\"type\":\"text\",\"text\":\"x-anthropic-billing-header: cc_version=2.1.280.283; cc_entrypoint=cli; cch=0; cc_prompt_id=ebb599d3-6b98-4fcc-bc9c-99d7dcbd62b9; cc_turn_origin=human;\"},{\"type\":\"text\",\"text\":\"You are Claude Code, Anthropic's official CLI for Claude.\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}],\"diagnostics\":{\"previous_message_id\":null}}", + "oauth/execute": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"\u003csystem-reminder\u003e\\nAs you answer the user's questions, you can use the following context:\\n# currentDate\\nToday's date is 2026-10-11.\\n\\n IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.\\n\u003c/system-reminder\u003e\\n\\n\"},{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"execution\":\"client\",\"name\":\"mcp__select_sight__disorder_ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"mcp__select_sight__push_exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"mcp__select_sight__elephant_lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}],\"system\":[{\"type\":\"text\",\"text\":\"x-anthropic-billing-header: cc_version=2.1.321.882; cc_entrypoint=cli; cch=0; cc_prompt_id=ebb599d3-6b98-4fcc-bc9c-99d7dcbd62b9; cc_turn_origin=human;\"},{\"type\":\"text\",\"text\":\"You are Claude Code, Anthropic's official CLI for Claude.\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}],\"diagnostics\":{\"previous_message_id\":null}}", + "oauth/stream": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"\u003csystem-reminder\u003e\\nAs you answer the user's questions, you can use the following context:\\n# currentDate\\nToday's date is 2026-10-11.\\n\\n IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.\\n\u003c/system-reminder\u003e\\n\\n\"},{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"execution\":\"client\",\"name\":\"mcp__select_sight__disorder_ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"mcp__select_sight__push_exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"mcp__select_sight__elephant_lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}],\"system\":[{\"type\":\"text\",\"text\":\"x-anthropic-billing-header: cc_version=2.1.321.882; cc_entrypoint=cli; cch=0; cc_prompt_id=ebb599d3-6b98-4fcc-bc9c-99d7dcbd62b9; cc_turn_origin=human;\"},{\"type\":\"text\",\"text\":\"You are Claude Code, Anthropic's official CLI for Claude.\",\"cache_control\":{\"type\":\"ephemeral\",\"ttl\":\"1h\"}}],\"diagnostics\":{\"previous_message_id\":null}}", "official/count_tokens": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":\"create a calendar event\"}],\"stream\":true,\"tools\":[{\"type\":\"tool_search\",\"execution\":\"client\",\"name\":\"ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"}}]}", "official/execute": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"type\":\"tool_search\",\"execution\":\"client\",\"name\":\"ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"},\"cache_control\":{\"type\":\"ephemeral\"}}]}", "official/stream": "{\"model\":\"claude-sonnet-4-5\",\"max_tokens\":32000,\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"create a calendar event\",\"cache_control\":{\"type\":\"ephemeral\"}}]}],\"metadata\":{},\"stream\":true,\"tools\":[{\"type\":\"tool_search\",\"execution\":\"client\",\"name\":\"ClientSearch\",\"description\":\"Named search.\"},{\"name\":\"exec_command\",\"description\":\"Run a command.\",\"input_schema\":{\"properties\":{\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"],\"type\":\"object\"}},{\"name\":\"lookup_order\",\"description\":\"Look up an order.\",\"input_schema\":{\"properties\":{\"id\":{\"type\":\"string\"}},\"required\":[\"id\"],\"type\":\"object\"},\"cache_control\":{\"type\":\"ephemeral\"}}]}" diff --git a/internal/store/git_refresh_alias_test.go b/internal/store/git_refresh_alias_test.go new file mode 100644 index 00000000000..d319b49397f --- /dev/null +++ b/internal/store/git_refresh_alias_test.go @@ -0,0 +1,19 @@ +package store + +import ( + "path/filepath" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestGitRefreshJournalCoordinatesRepositoryAliases(t *testing.T) { + root := t.TempDir() + setupGitRemoteRepository(t, root, "main", testBranchSpec{name: "main", contents: "initial\n"}) + remote := filepath.Join(root, "seed") + journals := []codex.RefreshJournal{ + NewGitTokenStore(remote, "", "", "main"), + NewGitTokenStore(filepath.Join(remote, ".git"), "", "", "main"), + } + exerciseSharedRefreshJournal(t, journals) +} diff --git a/internal/store/git_refresh_journal.go b/internal/store/git_refresh_journal.go new file mode 100644 index 00000000000..c0b9292977d --- /dev/null +++ b/internal/store/git_refresh_journal.go @@ -0,0 +1,183 @@ +package store + +import ( + "context" + "crypto/rand" + "errors" + "fmt" + "path/filepath" + "strings" + "time" + + "github.com/go-git/go-git/v6" + "github.com/go-git/go-git/v6/config" + "github.com/go-git/go-git/v6/plumbing" + "github.com/go-git/go-git/v6/plumbing/object" + "github.com/go-git/go-git/v6/plumbing/transport" + "github.com/go-git/go-git/v6/storage/filesystem" + "github.com/go-git/go-git/v6/storage/memory" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + log "github.com/sirupsen/logrus" +) + +// Refresh journals use immutable tags, not credential-branch writes. Ordinary +// pulls, mirror recovery, and auth saves therefore cannot erase a claim. Each +// claimant creates a different root commit; a non-force push can publish only one. +func (s *GitTokenStore) RefreshClaim(ctx context.Context, key string) (claimed bool, errResult error) { + errResult = s.withRefreshRepository(ctx, func(repo *git.Repository) error { + ref := refreshGitRef("claim", key) + if hash, errLookup := s.refreshRemoteRef(ctx, repo, ref); errLookup != nil { + return errLookup + } else if !hash.IsZero() { + return nil + } + nonce := make([]byte, 32) + if _, errRand := rand.Read(nonce); errRand != nil { + return errRand + } + errPublish := s.publishRefreshRecord(ctx, repo, ref, fmt.Sprintf("refresh claim %x", nonce)) + if errPublish == nil { + claimed = true + return nil + } + // A conflicting or uncertain push never grants permission to send OAuth. + if hash, errLookup := s.refreshRemoteRef(ctx, repo, ref); errLookup == nil && !hash.IsZero() { + return nil + } + return errPublish + }) + return claimed, errResult +} + +func (s *GitTokenStore) RefreshRead(ctx context.Context, key string) (receipt []byte, errResult error) { + errResult = s.withRefreshRepository(ctx, func(repo *git.Repository) error { + ref := refreshGitRef("receipt", key) + hash, errLookup := s.refreshRemoteRef(ctx, repo, ref) + if errLookup != nil { + return errLookup + } + if hash.IsZero() { + claim, errClaim := s.refreshRemoteRef(ctx, repo, refreshGitRef("claim", key)) + if errClaim != nil { + return errClaim + } + if !claim.IsZero() { + receipt = []byte(codex.PendingRefreshReceipt) + } + return nil + } + if errFetch := repo.FetchContext(ctx, &git.FetchOptions{RemoteName: "origin", ClientOptions: s.gitClientOptions(), RefSpecs: []config.RefSpec{config.RefSpec(ref.String() + ":" + ref.String())}, Tags: git.NoTags}); errFetch != nil && !errors.Is(errFetch, git.NoErrAlreadyUpToDate) { + return errFetch + } + commit, errCommit := repo.CommitObject(hash) + if errCommit != nil { + return errCommit + } + receipt = []byte(commit.Message) + return nil + }) + return receipt, errResult +} + +func (s *GitTokenStore) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + return s.withRefreshRepository(ctx, func(repo *git.Repository) error { + return s.publishRefreshRecord(ctx, repo, refreshGitRef("receipt", key), string(receipt)) + }) +} + +func refreshGitRef(kind, key string) plumbing.ReferenceName { + return plumbing.ReferenceName("refs/tags/cliproxy-refresh/" + kind + "/" + key) +} + +func (s *GitTokenStore) withRefreshRepository(ctx context.Context, run func(*git.Repository) error) (errResult error) { + if s == nil || strings.TrimSpace(s.remote) == "" { + return fmt.Errorf("git refresh journal is unavailable") + } + // Refresh can hold a credential file lock while Save holds s.mu waiting + // for that same file. Journal I/O must neither acquire s.mu nor touch the + // credential mirror. Its immutable remote records need only a temporary + // object database, not a checkout of the credential branch. + repo, errInit := git.Init(memory.NewStorage(), nil) + if errInit != nil { + return errInit + } + defer func() { errResult = errors.Join(errResult, repo.Close()) }() + if _, errRemote := repo.CreateRemote(&config.RemoteConfig{Name: "origin", URLs: []string{s.remote}}); errRemote != nil { + return errRemote + } + // go-git's in-process file receive-pack does not serialize competing + // reference updates. Lock on the shared remote itself for that transport. + endpoint, errParse := transport.ParseURL(s.remote) + if errParse != nil { + return errParse + } + if endpoint.Scheme == "file" { + // Use the transport's own resolution of bare repositories, .git + // directories and gitfiles so aliases share the same physical lock. + remoteStorage, errLoad := transport.DefaultLoader.Load(endpoint) + if errLoad != nil { + return errLoad + } + remoteFilesystem, ok := remoteStorage.(*filesystem.Storage) + if !ok { + return fmt.Errorf("local git refresh journal requires filesystem storage") + } + defer func() { errResult = errors.Join(errResult, remoteFilesystem.Close()) }() + canonical, errCanonical := filepath.EvalSymlinks(remoteFilesystem.Filesystem().Root()) + if errCanonical != nil { + return errCanonical + } + return codex.WithFileLock(filepath.Join(canonical, ".cliproxy-refresh"), func() error { return run(repo) }) + } + return run(repo) +} + +func (s *GitTokenStore) refreshRemoteRef(ctx context.Context, repo *git.Repository, name plumbing.ReferenceName) (plumbing.Hash, error) { + remote, errRemote := repo.Remote("origin") + if errRemote != nil { + return plumbing.ZeroHash, errRemote + } + refs, errList := remote.ListContext(ctx, &git.ListOptions{ClientOptions: s.gitClientOptions()}) + if errList != nil { + return plumbing.ZeroHash, errList + } + for _, ref := range refs { + if ref.Name() == name { + return ref.Hash(), nil + } + } + return plumbing.ZeroHash, nil +} + +func (s *GitTokenStore) publishRefreshRecord(ctx context.Context, repo *git.Repository, ref plumbing.ReferenceName, message string) error { + tree := &object.Tree{} + treeObject := repo.Storer.NewEncodedObject() + if errEncode := tree.Encode(treeObject); errEncode != nil { + return errEncode + } + treeHash, errTree := repo.Storer.SetEncodedObject(treeObject) + if errTree != nil { + return errTree + } + signature := object.Signature{Name: "CLIProxyAPI", Email: "refresh@cliproxyapi.local", When: time.Now()} + commit := &object.Commit{Author: signature, Committer: signature, Message: message, TreeHash: treeHash} + encoded := repo.Storer.NewEncodedObject() + if errEncode := commit.Encode(encoded); errEncode != nil { + return errEncode + } + hash, errStore := repo.Storer.SetEncodedObject(encoded) + if errStore != nil { + return errStore + } + // A unique local source avoids updating a tag fetched from another claimant. + localRef := plumbing.ReferenceName("refs/cliproxy-refresh-pending/" + hash.String()) + if errReference := repo.Storer.SetReference(plumbing.NewHashReference(localRef, hash)); errReference != nil { + return errReference + } + defer func() { + if errRemove := repo.Storer.RemoveReference(localRef); errRemove != nil { + log.Warnf("remove temporary refresh journal reference: %v", errRemove) + } + }() + return repo.PushContext(ctx, &git.PushOptions{RemoteName: "origin", ClientOptions: s.gitClientOptions(), RefSpecs: []config.RefSpec{config.RefSpec(localRef.String() + ":" + ref.String())}}) +} diff --git a/internal/store/git_refresh_lock_test.go b/internal/store/git_refresh_lock_test.go new file mode 100644 index 00000000000..ddb23a488ae --- /dev/null +++ b/internal/store/git_refresh_lock_test.go @@ -0,0 +1,66 @@ +package store + +import ( + "context" + "io" + "net/http" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/go-git/go-git/v6" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +type journalRoundTripper func(*http.Request) (*http.Response, error) + +func (r journalRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return r(req) } + +func TestGitRefreshDoesNotDeadlockWithConcurrentSave(t *testing.T) { + root := t.TempDir() + remote := setupGitRemoteRepository(t, root, "main", testBranchSpec{name: "main", contents: "initial\n"}) + store := NewGitTokenStore(remote, "", "", "main") + store.SetBaseDir(filepath.Join(root, "mirror", "auths")) + if errEnsure := store.EnsureRepository(); errEnsure != nil { + t.Fatal(errEnsure) + } + auth := &coreauth.Auth{ID: "journal.json", FileName: "journal.json", Provider: "codex", Metadata: map[string]any{"type": "codex", "access_token": "old-access", "refresh_token": "journal-lock-grant", "expired": "2020-01-01T00:00:00Z"}} + if _, errSave := store.Save(context.Background(), auth); errSave != nil { + t.Fatal(errSave) + } + transportEntered := make(chan struct{}) + allowResponse := make(chan struct{}) + originalTransport := http.DefaultTransport + http.DefaultTransport = journalRoundTripper(func(req *http.Request) (*http.Response, error) { + close(transportEntered) + <-allowResponse + return &http.Response{StatusCode: 200, Header: make(http.Header), Request: req, Body: io.NopCloser(strings.NewReader(`{"access_token":"new-access","refresh_token":"new-grant","expires_in":3600}`))}, nil + }) + t.Cleanup(func() { http.DefaultTransport = originalTransport }) + refreshDone := make(chan error, 1) + go func() { + _, errRefresh := executor.NewCodexExecutor(&config.Config{AuthDir: store.AuthDir()}).Refresh(codex.WithRefreshJournal(context.Background(), store), auth.Clone()) + refreshDone <- errRefresh + }() + <-transportEntered + saveHoldingStoreMutex := make(chan struct{}) + var once sync.Once + store.pullFn = func(context.Context, *git.Worktree, *git.PullOptions) error { + once.Do(func() { close(saveHoldingStoreMutex) }) + return git.NoErrAlreadyUpToDate + } + saveDone := make(chan error, 1) + go func() { _, errSave := store.Save(context.Background(), auth.Clone()); saveDone <- errSave }() + <-saveHoldingStoreMutex + close(allowResponse) + if errRefresh := <-refreshDone; errRefresh != nil { + t.Fatal(errRefresh) + } + if errSave := <-saveDone; errSave != nil { + t.Fatal(errSave) + } +} diff --git a/internal/store/object_refresh_journal.go b/internal/store/object_refresh_journal.go new file mode 100644 index 00000000000..02ed032c4ee --- /dev/null +++ b/internal/store/object_refresh_journal.go @@ -0,0 +1,138 @@ +package store + +import ( + "bytes" + "context" + "crypto/rand" + "fmt" + "io" + "net/http" + "path" + "time" + + "github.com/minio/minio-go/v7" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + log "github.com/sirupsen/logrus" +) + +// RefreshClaim requires S3 conditional creation, not a last-writer-wins upload. +// Backends that ignore If-None-Match fail the capability probe before any OAuth +// request. A claim is never deleted or retried after an uncertain PUT response. +func (s *ObjectTokenStore) RefreshClaim(ctx context.Context, key string) (bool, error) { + if s == nil || s.client == nil { + return false, fmt.Errorf("object refresh journal is unavailable") + } + if errProbe := s.ensureRefreshConditionalCreate(ctx); errProbe != nil { + return false, errProbe + } + return s.createRefreshObject(ctx, s.refreshObjectKey(key, "claim"), []byte("pending")) +} + +func (s *ObjectTokenStore) RefreshRead(ctx context.Context, key string) ([]byte, error) { + object, errGet := s.client.GetObject(ctx, s.cfg.Bucket, s.refreshObjectKey(key, "receipt"), minio.GetObjectOptions{}) + if errGet != nil { + return nil, errGet + } + data, errRead := io.ReadAll(object) + errClose := object.Close() + if errRead != nil { + response := minio.ToErrorResponse(errRead) + if response.StatusCode == http.StatusNotFound { + _, errStat := s.client.StatObject(ctx, s.cfg.Bucket, s.refreshObjectKey(key, "claim"), minio.StatObjectOptions{}) + if errStat == nil { + return []byte(codex.PendingRefreshReceipt), nil + } + if minio.ToErrorResponse(errStat).StatusCode == http.StatusNotFound { + return nil, nil + } + return nil, errStat + } + return nil, errRead + } + return data, errClose +} + +func (s *ObjectTokenStore) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + created, errCreate := s.createRefreshObject(ctx, s.refreshObjectKey(key, "receipt"), receipt) + if errCreate != nil { + return errCreate + } + if !created { + return fmt.Errorf("refresh receipt already exists") + } + return nil +} + +func (s *ObjectTokenStore) refreshObjectKey(key, kind string) string { + return path.Join(s.cfg.Prefix, "codex-refresh", key, kind) +} + +func (s *ObjectTokenStore) createRefreshObject(ctx context.Context, key string, data []byte) (bool, error) { + headers := http.Header{"If-None-Match": []string{"*"}} + // The bundled minio PutObjectOptions quotes "*" as an ETag; use signed + // headers to send the actual HTTP wildcard, not the literal ETag "*". + url, errPresign := s.client.PresignHeader(ctx, http.MethodPut, s.cfg.Bucket, key, time.Minute, nil, headers) + if errPresign != nil { + return false, fmt.Errorf("sign conditional refresh journal write: %w", errPresign) + } + req, errRequest := http.NewRequestWithContext(ctx, http.MethodPut, url.String(), bytes.NewReader(data)) + if errRequest != nil { + return false, fmt.Errorf("create conditional refresh journal request") + } + req.Header = headers + client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, errDo := client.Do(req) + if errDo != nil { + return false, fmt.Errorf("conditional refresh journal write outcome unknown") + } + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Warnf("close refresh journal response: %v", errClose) + } + }() + if resp.StatusCode == http.StatusPreconditionFailed { + return false, nil + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return false, fmt.Errorf("conditional refresh journal write returned HTTP %d", resp.StatusCode) + } + return true, nil +} + +func (s *ObjectTokenStore) ensureRefreshConditionalCreate(ctx context.Context) error { + s.refreshProbeMu.Lock() + defer s.refreshProbeMu.Unlock() + if s.refreshProbeReady { + return nil + } + // Cache only positive evidence. A temporary storage outage must not + // permanently reject new authorizations after the backend recovers. + if errProbe := s.probeRefreshConditionalCreate(ctx); errProbe != nil { + return errProbe + } + s.refreshProbeReady = true + return nil +} + +func (s *ObjectTokenStore) probeRefreshConditionalCreate(ctx context.Context) error { + nonce := make([]byte, 32) + if _, errRand := rand.Read(nonce); errRand != nil { + return errRand + } + key := s.refreshObjectKey(fmt.Sprintf("probe-%x", nonce), "claim") + first, errFirst := s.createRefreshObject(ctx, key, []byte("probe")) + if errFirst != nil || !first { + return fmt.Errorf("refresh journal conditional-create probe failed: %v", errFirst) + } + second, errSecond := s.createRefreshObject(ctx, key, []byte("probe")) + if errSecond != nil { + return errSecond + } + if second { + return fmt.Errorf("object backend does not enforce conditional creation; refresh refused") + } + if errRemove := s.client.RemoveObject(ctx, s.cfg.Bucket, key, minio.RemoveObjectOptions{}); errRemove != nil { + log.Warnf("remove refresh journal capability probe: %v", errRemove) + } + return nil +} diff --git a/internal/store/object_refresh_journal_test.go b/internal/store/object_refresh_journal_test.go new file mode 100644 index 00000000000..88715e63f47 --- /dev/null +++ b/internal/store/object_refresh_journal_test.go @@ -0,0 +1,134 @@ +package store + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestObjectRefreshJournalConditionalCreation(t *testing.T) { + for _, ignoreCondition := range []bool{false, true} { + t.Run(strconv.FormatBool(ignoreCondition), func(t *testing.T) { + var mu sync.Mutex + objects := map[string][]byte{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + switch r.Method { + case http.MethodPut: + if r.Header.Get("If-None-Match") != "*" { + t.Error("claim did not send the HTTP wildcard condition") + } + if _, exists := objects[r.URL.Path]; exists && !ignoreCondition { + w.WriteHeader(http.StatusPreconditionFailed) + return + } + data, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Error(errRead) + w.WriteHeader(500) + return + } + objects[r.URL.Path] = data + w.Header().Set("ETag", `"synthetic-etag"`) + w.WriteHeader(http.StatusOK) + case http.MethodGet, http.MethodHead: + data, exists := objects[r.URL.Path] + if !exists { + w.WriteHeader(404) + return + } + w.Header().Set("ETag", `"synthetic-etag"`) + w.Header().Set("Content-Length", strconv.Itoa(len(data))) + w.Header().Set("Last-Modified", "Thu, 01 Jan 2026 00:00:00 GMT") + _, _ = w.Write(data) + case http.MethodDelete: + delete(objects, r.URL.Path) + w.WriteHeader(http.StatusNoContent) + default: + t.Errorf("unexpected method %s", r.Method) + w.WriteHeader(500) + } + })) + defer server.Close() + var journals []codex.RefreshJournal + for i := 0; i < 2; i++ { + store, errNew := NewObjectTokenStore(ObjectStoreConfig{Endpoint: strings.TrimPrefix(server.URL, "http://"), Bucket: "test-bucket", AccessKey: "synthetic", SecretKey: "synthetic", Region: "us-east-1", PathStyle: true, LocalRoot: t.TempDir()}) + if errNew != nil { + t.Fatal(errNew) + } + journals = append(journals, store) + } + if ignoreCondition { + claimed, errClaim := journals[0].RefreshClaim(context.Background(), strings.Repeat("b", 64)) + if errClaim == nil || claimed { + t.Fatalf("unsupported condition must fail closed: %v, %v", claimed, errClaim) + } + return + } + exerciseSharedRefreshJournal(t, journals) + }) + } +} + +func exerciseSharedRefreshJournal(t *testing.T, journals []codex.RefreshJournal) { + t.Helper() + ctx := context.Background() + key := strings.Repeat("b", 64) + var mu sync.Mutex + winners := 0 + var owner codex.RefreshJournal + var wg sync.WaitGroup + for _, journal := range journals { + wg.Add(1) + go func() { + defer wg.Done() + claimed, errClaim := journal.RefreshClaim(ctx, key) + if errClaim != nil { + t.Errorf("claim: %v", errClaim) + return + } + if claimed { + mu.Lock() + winners++ + owner = journal + mu.Unlock() + } + }() + } + wg.Wait() + if winners != 1 { + t.Fatalf("claim winners = %d, want 1", winners) + } + for _, journal := range journals { + claimed, errClaim := journal.RefreshClaim(ctx, key) + if errClaim != nil || claimed { + t.Fatalf("unresolved claim was reclaimed: %v, %v", claimed, errClaim) + } + pending, errRead := journal.RefreshRead(ctx, key) + if errRead != nil || string(pending) != codex.PendingRefreshReceipt { + t.Fatalf("pending receipt = %s, error %v", pending, errRead) + } + } + receipt := []byte(`{"stage":"response-read"}`) + if errComplete := owner.RefreshComplete(ctx, key, receipt); errComplete != nil { + t.Fatal(errComplete) + } + for _, journal := range journals { + data, errRead := journal.RefreshRead(ctx, key) + if errRead != nil || string(data) != string(receipt) { + t.Fatalf("shared receipt = %s, error %v", data, errRead) + } + claimed, errClaim := journal.RefreshClaim(ctx, key) + if claimed || errClaim != nil { + t.Fatalf("completed unknown outcome was reclaimed: %v, %v", claimed, errClaim) + } + } +} diff --git a/internal/store/object_refresh_probe_test.go b/internal/store/object_refresh_probe_test.go new file mode 100644 index 00000000000..337f34d9028 --- /dev/null +++ b/internal/store/object_refresh_probe_test.go @@ -0,0 +1,39 @@ +package store + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +func TestRefreshProbeRecoversFromTransientOutage(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := calls.Add(1) + if n == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + if n == 3 { + w.WriteHeader(http.StatusPreconditionFailed) + return + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + store, errNew := NewObjectTokenStore(ObjectStoreConfig{Endpoint: strings.TrimPrefix(server.URL, "http://"), Bucket: "test-bucket", AccessKey: "synthetic", SecretKey: "synthetic", Region: "us-east-1", PathStyle: true, LocalRoot: t.TempDir()}) + if errNew != nil { + t.Fatal(errNew) + } + ctx := context.Background() + if claimed, errClaim := store.RefreshClaim(ctx, strings.Repeat("a", 64)); errClaim == nil || claimed { + t.Fatal("expected first transient failure to fail closed") + } + claimed, errClaim := store.RefreshClaim(ctx, strings.Repeat("b", 64)) + if errClaim != nil || !claimed { + t.Fatalf("healthy backend cannot recover even with new authorization: claimed=%v error=%v HTTP calls=%d", claimed, errClaim, calls.Load()) + } +} diff --git a/internal/store/objectstore.go b/internal/store/objectstore.go index 39b8b089825..22db97cb6b2 100644 --- a/internal/store/objectstore.go +++ b/internal/store/objectstore.go @@ -44,13 +44,15 @@ type ObjectStoreConfig struct { // ObjectTokenStore persists configuration and authentication metadata using an S3-compatible object storage backend. // Files are mirrored to a local workspace so existing file-based flows continue to operate. type ObjectTokenStore struct { - client *minio.Client - cfg ObjectStoreConfig - spoolRoot string - configPath string - authDir string - mu sync.Mutex - putObjectFn func(ctx context.Context, key string, data []byte, contentType string) error + client *minio.Client + cfg ObjectStoreConfig + spoolRoot string + configPath string + authDir string + mu sync.Mutex + putObjectFn func(ctx context.Context, key string, data []byte, contentType string) error + refreshProbeMu sync.Mutex + refreshProbeReady bool } // NewObjectTokenStore initializes an object storage backed token store. diff --git a/internal/store/postgres_refresh_integration_test.go b/internal/store/postgres_refresh_integration_test.go new file mode 100644 index 00000000000..f0c41b0e498 --- /dev/null +++ b/internal/store/postgres_refresh_integration_test.go @@ -0,0 +1,103 @@ +package store + +import ( + "context" + "database/sql" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "sync" + "testing" +) + +func TestPostgresConcurrentInitialRefreshClaims(t *testing.T) { + if os.Getenv("CLIPROXY_TEST_POSTGRES") != "1" || runtime.GOOS == "windows" || os.Geteuid() == 0 { + t.Skip("requires opt-in CLIPROXY_TEST_POSTGRES=1 and a non-root Unix PostgreSQL installation") + } + postgres, errLookup := exec.LookPath("postgres") + if errLookup != nil { + t.Skip("postgres is not installed") + } + bin := filepath.Dir(postgres) + root, errTemp := os.MkdirTemp("", "cpa-refresh-pg-") + if errTemp != nil { + t.Fatal(errTemp) + } + t.Cleanup(func() { + if errRemove := os.RemoveAll(root); errRemove != nil { + t.Error(errRemove) + } + }) + data := filepath.Join(root, "data") + cmdInit := exec.Command(filepath.Join(bin, "initdb"), "-D", data, "-U", "refresh_test", "-A", "trust", "--no-locale") + cmdInit.Env = append(os.Environ(), "LC_ALL=C", "LANG=C") + if output, errRun := cmdInit.CombinedOutput(); errRun != nil { + t.Fatalf("initdb: %v: %s", errRun, output) + } + cmdStart := exec.Command(filepath.Join(bin, "pg_ctl"), "-D", data, "-l", filepath.Join(root, "server.log"), "-o", "-h '' -k "+root, "-w", "start") + cmdStart.Env = append(os.Environ(), "LC_ALL=C", "LANG=C") + if output, errRun := cmdStart.CombinedOutput(); errRun != nil { + serverLog, _ := os.ReadFile(filepath.Join(root, "server.log")) + t.Fatalf("pg_ctl start: %v: %s\n%s", errRun, output, serverLog) + } + t.Cleanup(func() { + cmdStop := exec.Command(filepath.Join(bin, "pg_ctl"), "-D", data, "-m", "immediate", "-w", "stop") + cmdStop.Env = append(os.Environ(), "LC_ALL=C", "LANG=C") + if output, errRun := cmdStop.CombinedOutput(); errRun != nil { + t.Errorf("pg_ctl stop: %v: %s", errRun, output) + } + }) + db, errOpen := sql.Open("pgx", "host="+root+" user=refresh_test dbname=postgres sslmode=disable") + if errOpen != nil { + t.Fatal(errOpen) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Error(errClose) + } + }) + const parallel = 16 + db.SetMaxOpenConns(parallel) + db.SetMaxIdleConns(parallel) + ctx := context.Background() + var warm []*sql.Conn + for index := 0; index < parallel; index++ { + conn, errConn := db.Conn(ctx) + if errConn != nil { + t.Fatal(errConn) + } + warm = append(warm, conn) + } + for _, conn := range warm { + if errClose := conn.Close(); errClose != nil { + t.Fatal(errClose) + } + } + for round := 0; round < 3; round++ { + start := make(chan struct{}) + results := make(chan error, parallel) + var wg sync.WaitGroup + for index := 0; index < parallel; index++ { + wg.Add(1) + go func() { + store := &PostgresStore{db: db, cfg: PostgresStoreConfig{AuthTable: fmt.Sprintf("auths_%d", round)}} + defer wg.Done() + <-start + claimed, errClaim := store.RefreshClaim(ctx, fmt.Sprintf("%064x", index+1)) + if errClaim != nil { + results <- errClaim + } else if !claimed { + results <- fmt.Errorf("distinct grant was not claimed") + } + }() + } + close(start) + wg.Wait() + close(results) + for errClaim := range results { + t.Errorf("round %d concurrent initial claim failed: %v", round, errClaim) + } + } +} diff --git a/internal/store/postgres_refresh_journal.go b/internal/store/postgres_refresh_journal.go new file mode 100644 index 00000000000..52d2b8f9f36 --- /dev/null +++ b/internal/store/postgres_refresh_journal.go @@ -0,0 +1,95 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "fmt" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + log "github.com/sirupsen/logrus" +) + +func (s *PostgresStore) refreshTable() string { + name := s.cfg.AuthTable + if name == "" { + name = defaultAuthTable + } + return s.fullTableName(name + "_refresh_journal") +} + +// RefreshClaim uses a unique database key, independent of disposable auth mirrors. +func (s *PostgresStore) RefreshClaim(ctx context.Context, key string) (bool, error) { + if s == nil || s.db == nil { + return false, fmt.Errorf("postgres refresh journal is unavailable") + } + if errSchema := s.ensureRefreshSchema(ctx); errSchema != nil { + return false, errSchema + } + table := s.refreshTable() + result, errInsert := s.db.ExecContext(ctx, "INSERT INTO "+table+" (id) VALUES ($1) ON CONFLICT (id) DO NOTHING", key) + if errInsert != nil { + return false, errInsert + } + count, errCount := result.RowsAffected() + return count == 1, errCount +} + +// Bootstrap prepares the table before refresh workers start. The lazy path also +// supports embedded callers. A database-wide transaction lock serializes DDL +// across instances, including differently-spelled schema/search_path aliases. +// Only successful initialization is cached, so temporary failures can recover. +func (s *PostgresStore) ensureRefreshSchema(ctx context.Context) error { + s.refreshSchemaMu.Lock() + defer s.refreshSchemaMu.Unlock() + if s.refreshSchemaReady { + return nil + } + tx, errBegin := s.db.BeginTx(ctx, nil) + if errBegin != nil { + return fmt.Errorf("postgres refresh schema: begin: %w", errBegin) + } + defer func() { + if errRollback := tx.Rollback(); errRollback != nil && !errors.Is(errRollback, sql.ErrTxDone) { + log.Warnf("postgres refresh schema: rollback: %v", errRollback) + } + }() + if _, errLock := tx.ExecContext(ctx, "SELECT pg_advisory_xact_lock(hashtext('cliproxyapi:codex-refresh-schema')::bigint)"); errLock != nil { + return fmt.Errorf("postgres refresh schema: lock: %w", errLock) + } + if _, errCreate := tx.ExecContext(ctx, "CREATE TABLE IF NOT EXISTS "+s.refreshTable()+" (id TEXT PRIMARY KEY, receipt BYTEA)"); errCreate != nil { + return fmt.Errorf("postgres refresh schema: create: %w", errCreate) + } + if errCommit := tx.Commit(); errCommit != nil { + return fmt.Errorf("postgres refresh schema: commit: %w", errCommit) + } + s.refreshSchemaReady = true + return nil +} + +func (s *PostgresStore) RefreshRead(ctx context.Context, key string) ([]byte, error) { + var receipt []byte + errRead := s.db.QueryRowContext(ctx, "SELECT receipt FROM "+s.refreshTable()+" WHERE id = $1", key).Scan(&receipt) + if errors.Is(errRead, sql.ErrNoRows) { + return nil, nil + } + if errRead == nil && receipt == nil { + return []byte(codex.PendingRefreshReceipt), nil + } + return receipt, errRead +} + +func (s *PostgresStore) RefreshComplete(ctx context.Context, key string, receipt []byte) error { + result, errUpdate := s.db.ExecContext(ctx, "UPDATE "+s.refreshTable()+" SET receipt = $2 WHERE id = $1 AND receipt IS NULL", key, receipt) + if errUpdate != nil { + return errUpdate + } + count, errCount := result.RowsAffected() + if errCount != nil { + return errCount + } + if count != 1 { + return fmt.Errorf("refresh receipt already completed or claim missing") + } + return nil +} diff --git a/internal/store/postgres_refresh_journal_test.go b/internal/store/postgres_refresh_journal_test.go new file mode 100644 index 00000000000..c9d387a6058 --- /dev/null +++ b/internal/store/postgres_refresh_journal_test.go @@ -0,0 +1,165 @@ +package store + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + "io" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +var refreshSQLDriverID atomic.Int64 + +type refreshSQLDriver struct { + mu sync.Mutex + schemaMu sync.Mutex + failDDL bool + ddlCalls int + rows map[string][]byte +} + +type refreshSQLConn struct { + database *refreshSQLDriver + schemaLocked bool +} +type refreshSQLTx struct{ conn *refreshSQLConn } + +func (tx *refreshSQLTx) Commit() error { + if tx.conn.schemaLocked { + tx.conn.schemaLocked = false + tx.conn.database.schemaMu.Unlock() + } + return nil +} +func (tx *refreshSQLTx) Rollback() error { return tx.Commit() } + +type refreshSQLRows struct { + receipt []byte + found bool +} + +func (d *refreshSQLDriver) Open(string) (driver.Conn, error) { + return &refreshSQLConn{database: d}, nil +} +func (c *refreshSQLConn) Close() error { return nil } +func (c *refreshSQLConn) Prepare(string) (driver.Stmt, error) { + return nil, fmt.Errorf("unexpected prepare") +} +func (c *refreshSQLConn) Begin() (driver.Tx, error) { return &refreshSQLTx{conn: c}, nil } +func (c *refreshSQLConn) ExecContext(_ context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + if strings.HasPrefix(query, "SELECT pg_advisory_xact_lock(") { + c.database.schemaMu.Lock() + c.schemaLocked = true + return driver.RowsAffected(0), nil + } + c.database.mu.Lock() + defer c.database.mu.Unlock() + if strings.HasPrefix(query, "CREATE TABLE IF NOT EXISTS ") { + if !c.schemaLocked { + return nil, fmt.Errorf("DDL must hold a database transaction lock") + } + c.database.ddlCalls++ + if c.database.failDDL { + c.database.failDDL = false + return nil, fmt.Errorf("temporary schema failure") + } + return driver.RowsAffected(0), nil + } + key := args[0].Value.(string) + if strings.HasPrefix(query, "INSERT INTO ") && strings.Contains(query, "ON CONFLICT (id) DO NOTHING") { + if _, exists := c.database.rows[key]; exists { + return driver.RowsAffected(0), nil + } + c.database.rows[key] = nil + return driver.RowsAffected(1), nil + } + if strings.HasPrefix(query, "UPDATE ") && strings.Contains(query, "AND receipt IS NULL") { + if data, exists := c.database.rows[key]; !exists || data != nil { + return driver.RowsAffected(0), nil + } + c.database.rows[key] = append([]byte(nil), args[1].Value.([]byte)...) + return driver.RowsAffected(1), nil + } + return nil, fmt.Errorf("unexpected journal SQL: %s", query) +} +func (c *refreshSQLConn) QueryContext(_ context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + if !strings.HasPrefix(query, "SELECT receipt FROM ") { + return nil, fmt.Errorf("unexpected journal query: %s", query) + } + c.database.mu.Lock() + defer c.database.mu.Unlock() + data, exists := c.database.rows[args[0].Value.(string)] + return &refreshSQLRows{receipt: append([]byte(nil), data...), found: exists}, nil +} +func (r *refreshSQLRows) Columns() []string { return []string{"receipt"} } +func (r *refreshSQLRows) Close() error { return nil } +func (r *refreshSQLRows) Next(values []driver.Value) error { + if !r.found { + return io.EOF + } + r.found = false + values[0] = r.receipt + return nil +} + +func TestPostgresRefreshJournalSurvivesMirrorReplacement(t *testing.T) { + database := &refreshSQLDriver{rows: make(map[string][]byte)} + name := fmt.Sprintf("refresh_journal_%d", refreshSQLDriverID.Add(1)) + sql.Register(name, database) + var journals []codex.RefreshJournal + for i := 0; i < 2; i++ { + db, errOpen := sql.Open(name, "") + if errOpen != nil { + t.Fatal(errOpen) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Error(errClose) + } + }) + journals = append(journals, &PostgresStore{db: db, cfg: PostgresStoreConfig{Schema: "isolated", AuthTable: "auths"}, authDir: t.TempDir()}) + } + exerciseSharedRefreshJournal(t, journals) + // A new client with a newly-created mirror still sees the database claim. + prior := journals[0].(*PostgresStore) + rebuilt := &PostgresStore{db: prior.db, cfg: prior.cfg, authDir: t.TempDir()} + claimed, errClaim := rebuilt.RefreshClaim(context.Background(), strings.Repeat("b", 64)) + if claimed || errClaim != nil { + t.Fatalf("mirror rebuild reclaimed grant: %v, %v", claimed, errClaim) + } +} + +func TestPostgresRefreshSchemaRetriesFailureAndCachesSuccess(t *testing.T) { + database := &refreshSQLDriver{rows: make(map[string][]byte), failDDL: true} + name := fmt.Sprintf("refresh_journal_%d", refreshSQLDriverID.Add(1)) + sql.Register(name, database) + db, errOpen := sql.Open(name, "") + if errOpen != nil { + t.Fatal(errOpen) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Error(errClose) + } + }) + store := &PostgresStore{db: db, cfg: PostgresStoreConfig{AuthTable: "auths"}} + key := strings.Repeat("a", 64) + if claimed, errClaim := store.RefreshClaim(context.Background(), key); claimed || errClaim == nil { + t.Fatalf("failed initialization allowed claim: %v, %v", claimed, errClaim) + } + if claimed, errClaim := store.RefreshClaim(context.Background(), key); !claimed || errClaim != nil { + t.Fatalf("recovered initialization did not allow claim: %v, %v", claimed, errClaim) + } + if claimed, errClaim := store.RefreshClaim(context.Background(), strings.Repeat("b", 64)); !claimed || errClaim != nil { + t.Fatalf("subsequent claim failed: %v, %v", claimed, errClaim) + } + if database.ddlCalls != 2 { + t.Fatalf("DDL calls=%d, want one failure and one success", database.ddlCalls) + } +} diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index 69bc318bdaa..1db51c28e7c 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -40,14 +40,16 @@ type PostgresStoreConfig struct { // PostgresStore persists configuration and authentication metadata using PostgreSQL as backend // while mirroring data to a local workspace so existing file-based workflows continue to operate. type PostgresStore struct { - db *sql.DB - cfg PostgresStoreConfig - spoolRoot string - configPath string - authDir string - cooldownStore *postgresCooldownStateStore - mu sync.Mutex - persistAuthFn func(ctx context.Context, relID string, data []byte) error + db *sql.DB + cfg PostgresStoreConfig + spoolRoot string + configPath string + authDir string + cooldownStore *postgresCooldownStateStore + mu sync.Mutex + persistAuthFn func(ctx context.Context, relID string, data []byte) error + refreshSchemaMu sync.Mutex + refreshSchemaReady bool } // NewPostgresStore establishes a connection to PostgreSQL and prepares the local workspace. @@ -163,7 +165,7 @@ func (s *PostgresStore) EnsureSchema(ctx context.Context) error { `, cooldownTable)); err != nil { return fmt.Errorf("postgres store: create cooldown table: %w", err) } - return nil + return s.ensureRefreshSchema(ctx) } // Bootstrap synchronizes configuration and auth records between PostgreSQL and the local workspace. diff --git a/internal/store/refresh_journal_test.go b/internal/store/refresh_journal_test.go new file mode 100644 index 00000000000..62727286da7 --- /dev/null +++ b/internal/store/refresh_journal_test.go @@ -0,0 +1,91 @@ +package store + +import ( + "context" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestGitRefreshJournalCoordinatesIndependentMirrors(t *testing.T) { + root := filepath.Join(t.TempDir(), "owner@local") + if errMkdir := os.Mkdir(root, 0o700); errMkdir != nil { + t.Fatal(errMkdir) + } + remote := setupGitRemoteRepository(t, root, "main", testBranchSpec{name: "main", contents: "initial\n"}) + var journals []codex.RefreshJournal + for _, name := range []string{"first", "second"} { + store := NewGitTokenStore(remote, "", "", "main") + store.SetBaseDir(filepath.Join(root, name, "auths")) + if errEnsure := store.EnsureRepository(); errEnsure != nil { + t.Fatal(errEnsure) + } + journal, ok := any(store).(codex.RefreshJournal) + if !ok { + t.Fatal("remote store lacks an authoritative refresh journal") + } + journals = append(journals, journal) + } + ctx := context.Background() + key := strings.Repeat("a", 64) + var winners atomic.Int32 + var owner codex.RefreshJournal + var mu sync.Mutex + var wg sync.WaitGroup + for _, journal := range journals { + wg.Add(1) + go func() { + defer wg.Done() + claimed, errClaim := journal.RefreshClaim(ctx, key) + if errClaim != nil { + t.Errorf("claim: %v", errClaim) + return + } + if claimed { + winners.Add(1) + mu.Lock() + owner = journal + mu.Unlock() + } + }() + } + wg.Wait() + if winners.Load() != 1 { + t.Fatalf("claim winners = %d, want 1", winners.Load()) + } + for _, journal := range journals { + claimed, errClaim := journal.RefreshClaim(ctx, key) + if errClaim != nil || claimed { + t.Fatalf("unresolved claim was reclaimed: %v, %v", claimed, errClaim) + } + } + receipt := []byte(`{"value":{"access_token":"synthetic-new"}}`) + if errComplete := owner.RefreshComplete(ctx, key, receipt); errComplete != nil { + t.Fatal(errComplete) + } + for _, journal := range journals { + data, errRead := journal.RefreshRead(ctx, key) + if errRead != nil || string(data) != string(receipt) { + t.Fatalf("shared receipt = %s, error %v", data, errRead) + } + } +} + +func TestRemoteStoresProvideAuthoritativeRefreshJournal(t *testing.T) { + for name, store := range map[string]any{ + "postgres": &PostgresStore{}, + "git": &GitTokenStore{}, + "object": &ObjectTokenStore{}, + } { + t.Run(name, func(t *testing.T) { + if _, ok := store.(codex.RefreshJournal); !ok { + t.Fatal("remote credential store cannot preserve refresh claims across mirror rebuilds") + } + }) + } +} diff --git a/internal/thinking/claude_standalone_effort_test.go b/internal/thinking/claude_standalone_effort_test.go index 16ac37e498a..2231dd52f98 100644 --- a/internal/thinking/claude_standalone_effort_test.go +++ b/internal/thinking/claude_standalone_effort_test.go @@ -7,7 +7,9 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" _ "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking/provider/codex" + _ "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking/provider/openai" codexclaude "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/codex/claude" + openaiclaude "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/openai/claude" ) func TestApplyThinking_ClaudeStandaloneEffort(t *testing.T) { @@ -135,3 +137,59 @@ func TestClaudeToCodexStandaloneEffort(t *testing.T) { }) } } + +func TestClaudeToOpenAIStandaloneEffort(t *testing.T) { + tests := []struct { + name string + fields string + suffix string + wantBoundEffort string + wantUnbound string + }{ + {name: "standalone high", fields: `,"output_config":{"effort":"high"}`, wantBoundEffort: "high", wantUnbound: "high"}, + {name: "standalone low", fields: `,"output_config":{"effort":"low"}`, wantBoundEffort: "low", wantUnbound: "low"}, + {name: "normalized high", fields: `,"output_config":{"effort":" HIGH "}`, wantBoundEffort: "high", wantUnbound: "high"}, + {name: "adaptive high", fields: `,"thinking":{"type":"adaptive"},"output_config":{"effort":"high"}`, wantBoundEffort: "high", wantUnbound: "high"}, + {name: "enabled budget high", fields: `,"thinking":{"type":"enabled","budget_tokens":16000}`, wantBoundEffort: "high", wantUnbound: "high"}, + {name: "budget precedes effort", fields: `,"thinking":{"type":"enabled","budget_tokens":8192},"output_config":{"effort":"high"}`, wantBoundEffort: "medium", wantUnbound: "medium"}, + {name: "disabled precedes effort", fields: `,"thinking":{"type":"disabled"},"output_config":{"effort":"high"}`, wantBoundEffort: "low", wantUnbound: "none"}, + {name: "suffix precedes standalone effort", fields: `,"output_config":{"effort":"high"}`, suffix: "(low)", wantBoundEffort: "low", wantUnbound: "low"}, + {name: "suffix precedes disabled", fields: `,"thinking":{"type":"disabled"},"output_config":{"effort":"low"}`, suffix: "(high)", wantBoundEffort: "high", wantUnbound: "high"}, + {name: "missing effort omits field", wantBoundEffort: "", wantUnbound: ""}, + {name: "blank effort omits field", fields: `,"output_config":{"effort":" "}`, wantBoundEffort: "", wantUnbound: ""}, + {name: "non-string effort omits field", fields: `,"output_config":{"effort":123}`, wantBoundEffort: "", wantUnbound: ""}, + } + info := ®istry.ModelInfo{ + ID: "custom-openai", + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"low", "medium", "high", "xhigh"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + source := []byte(`{"model":"custom-openai","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`[:len(`{"model":"custom-openai","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)-1] + tt.fields + `}`) + translated, errTranslate := openaiclaude.ConvertClaudeRequestToOpenAI("custom-openai", source, false) + if errTranslate != nil { + t.Fatalf("ConvertClaudeRequestToOpenAI returned error: %v", errTranslate) + } + for _, bound := range []bool{false, true} { + var out []byte + var errApply error + want := tt.wantUnbound + if bound { + want = tt.wantBoundEffort + out, errApply = thinking.ApplyThinkingWithModelInfo(translated, source, "custom-openai"+tt.suffix, "claude", "openai", "openai", info) + } else { + out, errApply = thinking.ApplyThinkingWithSourceAndSummary(translated, source, "custom-openai"+tt.suffix, "claude", "openai", "openai", thinking.SummaryConfig{}) + } + if errApply != nil { + t.Fatalf("bound=%v: ApplyThinking returned error: %v", bound, errApply) + } + effort := ojson.GetBytes(out, "reasoning_effort") + if effort.String() != want || (want == "" && effort.Exists()) { + t.Fatalf("bound=%v: reasoning_effort = %s, want %q; body=%s", bound, effort.Raw, want, out) + } + } + }) + } +} diff --git a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go index bd4ff2d6f2c..a0f5559d4ff 100644 --- a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go +++ b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go @@ -142,6 +142,10 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq partResults := partsResult.Array() hasToolCallsArray := false imageCount := 0 + var textContent strings.Builder + var reasoningContent strings.Builder + hasTextContent := false + hasReasoningContent := false for i := 0; i < len(partResults); i++ { partResult := partResults[i] partTextResult := partResult.Get("text") @@ -156,7 +160,8 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq } hasThoughtSignature := thoughtSignatureResult.Exists() && thoughtSignatureResult.String() != "" - hasContentPayload := partTextResult.Exists() || functionCallResult.Exists() || inlineDataResult.Exists() + hasTextPayload := partTextResult.Exists() && partTextResult.String() != "" + hasContentPayload := hasTextPayload || functionCallResult.Exists() || inlineDataResult.Exists() // Ignore encrypted thoughtSignature but keep any actual content in the same part. if hasThoughtSignature && !hasContentPayload { @@ -164,13 +169,15 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq } if partTextResult.Exists() { - textContent := partTextResult.String() + text := partTextResult.String() // Handle text content, distinguishing between regular content and reasoning/thoughts. if partResult.Get("thought").Bool() { - _ = templateDoc.Set("choices.0.delta.reasoning_content", textContent) + hasReasoningContent = true + reasoningContent.WriteString(text) } else { - _ = templateDoc.Set("choices.0.delta.content", textContent) + hasTextContent = true + textContent.WriteString(text) } _ = templateDoc.Set("choices.0.delta.role", "assistant") } else if functionCallResult.Exists() { @@ -222,6 +229,12 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq _ = templateDoc.SetRawBytes("choices.0.delta.images.-1", imageDoc.Bytes()) } } + if hasTextContent { + _ = templateDoc.Set("choices.0.delta.content", textContent.String()) + } + if hasReasoningContent { + _ = templateDoc.Set("choices.0.delta.reasoning_content", reasoningContent.String()) + } } // Determine finish_reason only on the final chunk, which upstream marks with diff --git a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go index 9faa86dc309..855b569df0a 100644 --- a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go +++ b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go @@ -434,3 +434,70 @@ func TestConvertAntigravityResponseToOpenAI_OmitToolCallsWhenAbsentInStreamingDe } } } + +func TestConvertAntigravityResponseToOpenAI_MultiPartThoughtAndTrailingSignature(t *testing.T) { + ctx := context.Background() + + t.Run("thought followed by empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"response":{"candidates":[{"content":{"parts":[{"thought":true,"text":"Analyzing the request..."},{"thought":true,"text":"","thoughtSignature":"EjQK..."}]}}]}}`) + out := ConvertAntigravityResponseToOpenAI(ctx, "gemini-3.1-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Analyzing the request..." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Analyzing the request...", out[0]) + } + if content := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content"); content.Type != ojson.Null { + t.Fatalf("delta.content = %s (type %v), want null; payload=%s", content.Raw, content.Type, out[0]) + } + if role := ojson.GetBytesNoCopy(out[0], "choices.0.delta.role").String(); role != "assistant" { + t.Fatalf("delta.role = %q, want assistant; payload=%s", role, out[0]) + } + }) + + t.Run("thought followed by non-thought empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"response":{"candidates":[{"content":{"parts":[{"thought":true,"text":"Analyzing the request..."},{"text":"","thoughtSignature":"EjQK..."}]}}]}}`) + out := ConvertAntigravityResponseToOpenAI(ctx, "gemini-3.1-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Analyzing the request..." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Analyzing the request...", out[0]) + } + if content := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content"); content.Type != ojson.Null { + t.Fatalf("delta.content = %s (type %v), want null; payload=%s", content.Raw, content.Type, out[0]) + } + }) + + t.Run("visible text followed by empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Final answer"},{"text":"","thoughtSignature":"EjQK..."}]}}]}}`) + out := ConvertAntigravityResponseToOpenAI(ctx, "gemini-3.1-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content").String(); got != "Final answer" { + t.Fatalf("delta.content = %q, want %q; payload=%s", got, "Final answer", out[0]) + } + if reasoning := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content"); reasoning.Type != ojson.Null { + t.Fatalf("delta.reasoning_content = %s (type %v), want null; payload=%s", reasoning.Raw, reasoning.Type, out[0]) + } + }) + + t.Run("multiple thought parts and multiple text parts in same chunk", func(t *testing.T) { + var param any + chunk := []byte(`{"response":{"candidates":[{"content":{"parts":[{"thought":true,"text":"Step 1. "},{"thought":true,"text":"Step 2."},{"text":"Hello "},{"text":"world"}]}}]}}`) + out := ConvertAntigravityResponseToOpenAI(ctx, "gemini-3.1-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Step 1. Step 2." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Step 1. Step 2.", out[0]) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content").String(); got != "Hello world" { + t.Fatalf("delta.content = %q, want %q; payload=%s", got, "Hello world", out[0]) + } + }) +} diff --git a/internal/translator/claude/openai/responses/claude_openai-responses_request.go b/internal/translator/claude/openai/responses/claude_openai-responses_request.go index 22c94b73292..0a3898c12b2 100644 --- a/internal/translator/claude/openai/responses/claude_openai-responses_request.go +++ b/internal/translator/claude/openai/responses/claude_openai-responses_request.go @@ -27,7 +27,9 @@ const ( // into a Claude Messages API request using only gjson/sjson for JSON handling. // It supports: // - instructions, input[].role==system and input[].role==developer -> separate -// top-level system blocks, in source order +// top-level system blocks, in source order; text-only system and developer +// items after the first conversation item stay in place as role=system +// turns unless WithTopLevelSystemItems pins the top-level path // - input[].type==message with input_text/output_text -> user/assistant messages // - function_call/custom_tool_call -> assistant tool_use // - function_call_output/custom_tool_call_output -> user tool_result @@ -145,83 +147,60 @@ func convertOpenAIResponsesRequestToClaude(ctx context.Context, modelName string _ = outDoc.Set("speed", "fast") } - // System-level inputs become canonical top-level Claude system blocks in - // source order: instructions first, then every input item whose role is - // system or developer. Each source block stays a separate Claude block and - // keeps operator authority; the Claude executor decides the final placement - // (mid-conversation role=system messages, or system reminders on legacy - // models), so this layer must not merge, trim or downgrade them to user text. + // Instructions and leading system-level items become separate top-level + // Claude blocks in source order. Eligible later items keep their position in + // the conversation as role=system turns, because clients such as Codex append + // them on later turns and moving them ahead of the history would invalidate + // Claude's prompt-cache prefix. The executor pins every item to the top + // level before translation when model compatibility or wire policy requires + // it, and still owns the final placement of the leading blocks. Both paths + // keep source text, block boundaries and cache_control, and never downgrade + // system-level items to user text. messageCapacity := root.Get("input.#").Int() if messageCapacity == 0 && root.Get("input").Type == ojson.String { messageCapacity = 1 } messageBlocks := common.NewRawArrayItems(messageCapacity) systemBlocks := make([][]byte, 0, 4) - appendSystemText := func(text string, cacheSource ojson.Result) { + appendSystemText := func(text string) { if strings.TrimSpace(text) == "" { return } - block := []byte(`{"type":"text","text":""}`) - block, _ = ojson.SetBytes(block, "text", text) - if cacheSource.Exists() { - block = common.AttachCacheControl(block, cacheSource) - } - systemBlocks = append(systemBlocks, block) + blockDoc := ojson.NewDoc([]byte(`{"type":"text","text":""}`)) + _ = blockDoc.Set("text", text) + systemBlocks = append(systemBlocks, blockDoc.Bytes()) } if instr := root.Get("instructions"); instr.Type == ojson.String { - appendSystemText(instr.String(), ojson.Result{}) + appendSystemText(instr.String()) + } + type responsesSystemInputEntry struct { + blocks [][]byte + eligible bool + emitted bool } + topLevelSystemItems := topLevelSystemItemsFor(ctx, modelName) + var systemInputEntries []responsesSystemInputEntry + systemInputEntryByIndex := map[int]int{} if input := root.Get("input"); input.IsArray() { + conversationStarted := false + itemIndex := -1 input.ForEach(func(_, item ojson.Result) bool { + itemIndex++ if !isResponsesSystemLevelRole(item.Get("role").String()) { + conversationStarted = true return true } - startIdx := len(systemBlocks) - content := item.Get("content") - if content.Type == ojson.String { - appendSystemText(content.String(), ojson.Result{}) - } else if content.IsArray() { - content.ForEach(func(_, part ojson.Result) bool { - switch part.Get("type").String() { - case "input_text", "output_text", "text": - txt := part.Get("text").String() - if strings.TrimSpace(txt) == "" { - if part.Get("cache_control").Exists() && len(systemBlocks) > startIdx { - prevIdx := len(systemBlocks) - 1 - if !ojson.GetBytesNoCopy(systemBlocks[prevIdx], "cache_control").Exists() { - systemBlocks[prevIdx] = common.AttachCacheControl(systemBlocks[prevIdx], part) - } - } - return true - } - appendSystemText(txt, part) - default: - if block := responsesSystemUnsupportedBlock(part); len(block) > 0 { - systemBlocks = append(systemBlocks, block) - } - } - return true - }) - } - // Item-level cache_control applies to the last block this item produced. - if item.Get("cache_control").Exists() && len(systemBlocks) > startIdx { - lastIdx := len(systemBlocks) - 1 - if !ojson.GetBytesNoCopy(systemBlocks[lastIdx], "cache_control").Exists() { - systemBlocks[lastIdx] = common.AttachCacheControl(systemBlocks[lastIdx], item) - } - } + blocks := responsesSystemItemToClaudeBlocks(item) + eligible := !topLevelSystemItems && conversationStarted && isResponsesMessageItem(item) && claudeBlocksAreAllText(blocks) + systemInputEntryByIndex[itemIndex] = len(systemInputEntries) + systemInputEntries = append(systemInputEntries, responsesSystemInputEntry{ + blocks: blocks, + eligible: eligible, + }) return true }) } - formatResult := root.Get("text.format") - if !formatResult.Exists() { - formatResult = root.Get("response_format") - } - if formatInstruction := common.BuildClaudeStructuredOutputInstruction(formatResult); formatInstruction != "" { - appendSystemText(formatInstruction, ojson.Result{}) - } - names := buildClaudeToolNames(root, toolSearch) winners := responsesToolWinners(root, toolSearch) // Native tool search declares discovered tools as deferred and replays @@ -233,8 +212,23 @@ func convertOpenAIResponsesRequestToClaude(ctx context.Context, modelName string var pendingParts [][]byte var pendingToolUseParts [][]byte var turnHasEmptyWebSearchResult bool + // Anthropic accepts a role=system turn only after a user message and before + // an assistant message or the end of messages. Mid-session system items are + // therefore queued until a user message is emitted, which may be the user + // message still being accumulated or a later one, and placed right after it. + var pendingSystemEntryIndices []int appendMessage := func(msg []byte) { messageBlocks = append(messageBlocks, msg) + if len(pendingSystemEntryIndices) == 0 || ojson.GetBytesNoCopy(msg, "role").String() != "user" { + return + } + for _, entryIdx := range pendingSystemEntryIndices { + systemDoc := ojson.NewDoc([]byte(`{"role":"system","content":[]}`)) + _ = systemDoc.SetRawBytes("content", common.JoinRawArray(systemInputEntries[entryIdx].blocks)) + messageBlocks = append(messageBlocks, systemDoc.Bytes()) + systemInputEntries[entryIdx].emitted = true + } + pendingSystemEntryIndices = nil } flushPendingMessage := func() { if pendingRole == "" { @@ -352,8 +346,11 @@ func convertOpenAIResponsesRequestToClaude(ctx context.Context, modelName string webSearchHistory := scanResponsesWebSearchHistory(inputItems) for itemIndex, item := range inputItems { - // System-level items already became top-level system blocks. + // System-level items are top-level blocks unless queued as in-place turns. if isResponsesSystemLevelRole(item.Get("role").String()) { + if entryIdx, ok := systemInputEntryByIndex[itemIndex]; ok && systemInputEntries[entryIdx].eligible { + pendingSystemEntryIndices = append(pendingSystemEntryIndices, entryIdx) + } continue } typ := item.Get("type").String() @@ -655,6 +652,23 @@ func convertOpenAIResponsesRequestToClaude(ctx context.Context, modelName string } } flushPendingMessage() + // Every system-level input item that was not emitted as an in-place turn + // (leading items, ineligible items, and queued items with no following user + // message) becomes a top-level system block in its original source order. + for _, entry := range systemInputEntries { + if !entry.emitted { + systemBlocks = append(systemBlocks, entry.blocks...) + } + } + pendingSystemEntryIndices = nil + // The structured-output instruction stays the last system block. + formatResult := root.Get("text.format") + if !formatResult.Exists() { + formatResult = root.Get("response_format") + } + if formatInstruction := common.BuildClaudeStructuredOutputInstruction(formatResult); formatInstruction != "" { + appendSystemText(formatInstruction) + } if len(unmappedItemTypes) > 0 { log.Warnf("responses->claude: dropped input items of unmapped types %v (model=%s)", unmappedItemTypes, modelName) } @@ -841,6 +855,131 @@ func isResponsesSystemLevelRole(role string) bool { } } +type topLevelSystemItemsContextKey struct{} + +// WithTopLevelSystemItems pins, for one request, whether every system-level +// input item becomes a top-level system block. The executor sets it before +// translation for upstreams where an in-place role=system turn cannot be used: +// models that reject the turn, strict cloaking (which drops every caller system +// instruction) and requests whose mid-conversation system turns are rebuilt +// into the system field. +func WithTopLevelSystemItems(ctx context.Context, topLevel bool) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, topLevelSystemItemsContextKey{}, topLevel) +} + +// topLevelSystemItemsFor returns the value pinned in ctx or, when ctx carries +// none, whether modelName rejects mid-conversation role=system turns. +func topLevelSystemItemsFor(ctx context.Context, modelName string) bool { + if ctx != nil { + if topLevel, ok := ctx.Value(topLevelSystemItemsContextKey{}).(bool); ok { + return topLevel + } + } + return util.IsClaudeLegacySystemReminderModel(modelName) +} + +// responsesSystemItemToClaudeBlocks converts one system-level input item into +// Claude system blocks: one text block per non-blank text part, a blank part's +// cache_control moved to the preceding block, the item's cache_control on its +// last block, and a typed marker for each unsupported part so the executor can +// report it instead of dropping it. +func responsesSystemItemToClaudeBlocks(item ojson.Result) [][]byte { + var blockDocs []*ojson.Doc + attachCacheControlToLast := func(source ojson.Result) { + if !source.Get("cache_control").Exists() || len(blockDocs) == 0 { + return + } + lastDoc := blockDocs[len(blockDocs)-1] + if !lastDoc.Get("cache_control").Exists() { + attachClaudeCacheControlDoc(lastDoc, source) + } + } + appendText := func(text string, cacheSource ojson.Result) { + if strings.TrimSpace(text) == "" { + attachCacheControlToLast(cacheSource) + return + } + blockDoc := ojson.NewDoc([]byte(`{"type":"text","text":""}`)) + _ = blockDoc.Set("text", text) + if cacheSource.Exists() { + attachClaudeCacheControlDoc(blockDoc, cacheSource) + } + blockDocs = append(blockDocs, blockDoc) + } + content := item.Get("content") + if content.Type == ojson.String { + appendText(content.String(), ojson.Result{}) + } else if content.IsArray() { + content.ForEach(func(_, part ojson.Result) bool { + switch part.Get("type").String() { + case "input_text", "output_text", "text": + appendText(part.Get("text").String(), part) + default: + if blockDoc := responsesSystemUnsupportedBlockDoc(part); blockDoc != nil { + blockDocs = append(blockDocs, blockDoc) + } + } + return true + }) + } + attachCacheControlToLast(item) + if len(blockDocs) == 0 { + return nil + } + blocks := make([][]byte, len(blockDocs)) + for i, blockDoc := range blockDocs { + blocks[i] = blockDoc.Bytes() + } + return blocks +} + +func attachClaudeCacheControlDoc(doc *ojson.Doc, src ojson.Result) { + if doc == nil { + return + } + cc := src.Get("cache_control") + if !cc.Exists() || cc.Type == ojson.Null || !cc.IsObject() { + return + } + typ := cc.Get("type") + if typ.Type != ojson.String || typ.String() != "ephemeral" { + return + } + switch doc.Get("type").String() { + case "thinking", "redacted_thinking": + return + case "text": + if strings.TrimSpace(doc.Get("text").String()) == "" { + return + } + } + _ = doc.SetRaw("cache_control", cc.Raw) +} + +// isResponsesMessageItem reports whether an input item is a message, which is +// the only item shape that can become a role=system turn. +func isResponsesMessageItem(item ojson.Result) bool { + typ := item.Get("type").String() + return typ == "" || typ == "message" +} + +// claudeBlocksAreAllText reports whether blocks is non-empty and holds only +// text blocks, the content a role=system turn can carry besides tool changes. +func claudeBlocksAreAllText(blocks [][]byte) bool { + if len(blocks) == 0 { + return false + } + for _, block := range blocks { + if ojson.GetBytesNoCopy(block, "type").String() != "text" { + return false + } + } + return true +} + // dropUnsupportedClaudeAssistantPrefill removes a trailing assistant message for // Claude model families that reject assistant message prefill (e.g., Fable, // Opus 5, Sonnet 4.6, Haiku 5.5), except when the message ends on a server-tool @@ -977,14 +1116,14 @@ func claudeModelRejectsAssistantPrefill(modelName string) bool { // rejected request, and the marker lets the Claude executor fail the request with // the offending type named. The original payload is not copied because the // request can never succeed. -func responsesSystemUnsupportedBlock(part ojson.Result) []byte { +func responsesSystemUnsupportedBlockDoc(part ojson.Result) *ojson.Doc { partType := strings.TrimSpace(part.Get("type").String()) if partType == "" { return nil } - block := []byte(`{"type":""}`) - block, _ = ojson.SetBytes(block, "type", partType) - return block + blockDoc := ojson.NewDoc([]byte(`{"type":""}`)) + _ = blockDoc.Set("type", partType) + return blockDoc } // convertResponsesReasoningToClaudeThinking rebuilds one Claude thinking block diff --git a/internal/translator/claude/openai/responses/claude_openai-responses_system_placement_test.go b/internal/translator/claude/openai/responses/claude_openai-responses_system_placement_test.go new file mode 100644 index 00000000000..a645936c976 --- /dev/null +++ b/internal/translator/claude/openai/responses/claude_openai-responses_system_placement_test.go @@ -0,0 +1,271 @@ +package responses + +import ( + "context" + "strings" + "testing" + + "github.com/MoeForge/moego/ojson" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" +) + +const placementFormatJSON = `"text":{"format":{"type":"json_schema","name":"facts","schema":{"type":"object","properties":{"facts":{"type":"string"}}}}}` + +func convertPlacement(t *testing.T, ctx context.Context, model, body string) ojson.Result { + t.Helper() + out, errConvert := ConvertOpenAIResponsesRequestToClaudeWithCompatContext(ctx, model, []byte(body), false) + if errConvert != nil { + t.Fatal(errConvert) + } + return util.ParseGJSONBytesNoCopy(out) +} + +func placementMessages(root ojson.Result) [][]byte { + var messages [][]byte + root.Get("messages").ForEach(func(_, message ojson.Result) bool { + messages = append(messages, []byte(message.Raw)) + return true + }) + return messages +} + +// Codex appends system-level items to the end of input on later turns. The +// messages translated for one turn must stay a prefix of the next turn's +// messages and the top-level system blocks must not change, or Claude's prompt +// cache misses on everything after the first changed block. +func TestResponsesMidSessionSystemItemsKeepPromptPrefixBetweenTurns(t *testing.T) { + turnTwo := `[ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "base instructions"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "first request"}]}, + {"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "ok"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "first answer"}]}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "2026-11-06"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "second request"}]} + ]` + turnThree := strings.TrimSuffix(strings.TrimSpace(turnTwo), "]") + `, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "second answer"}]}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "changed"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "third request"}]} + ]` + + two := convertPlacement(t, context.Background(), "claude-opus-5", `{"model":"claude-opus-5","input":`+turnTwo+`}`) + three := convertPlacement(t, context.Background(), "claude-opus-5", `{"model":"claude-opus-5","input":`+turnThree+`}`) + + if two.Get("system").Raw != three.Get("system").Raw { + t.Fatalf("top-level system changed between turns:\n%s\n%s", two.Get("system").Raw, three.Get("system").Raw) + } + twoMessages := two.Get("messages").Array() + threeMessages := three.Get("messages").Array() + if len(threeMessages) <= len(twoMessages) { + t.Fatalf("turn three has %d messages, want more than turn two's %d", len(threeMessages), len(twoMessages)) + } + for idx, message := range twoMessages { + if message.Raw != threeMessages[idx].Raw { + t.Fatalf("messages[%d] changed between turns:\n%s\n%s", idx, message.Raw, threeMessages[idx].Raw) + } + } + if got := three.Get("messages.#.role").Raw; got != `["user","assistant","user","assistant","user","system","assistant","user","system"]` { + t.Fatalf("roles = %s", got) + } + if got := three.Get("messages.5.content.0.text").String(); got != "2026-11-06" { + t.Fatalf("messages[5] text = %q", got) + } + if got := three.Get("messages.8.content.0.text").String(); got != "changed" { + t.Fatalf("messages[8] text = %q", got) + } + if problems := claudeMessageInvariantProblems(placementMessages(three)); len(problems) > 0 { + t.Fatalf("invariant problems: %v", problems) + } +} + +// A queued turn attaches to the next user message emitted, which may be the +// user message still being accumulated when the item arrives. +func TestResponsesMidSessionSystemItemFollowsPendingUserMessage(t *testing.T) { + root := convertPlacement(t, context.Background(), "claude-opus-5", `{ + "model": "claude-opus-5", + "instructions": "I1", + "input": [ + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": "S1"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "U1"}]}, + {"type": "message", "role": "developer", "content": "D1"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "A1"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "U2"}]} + ] + }`) + if got := root.Get("system.#.text").Raw; got != `["I1","S1"]` { + t.Fatalf("system = %s, want only the leading items", got) + } + if got := root.Get("messages.#.role").Raw; got != `["user","system","assistant","user"]` { + t.Fatalf("roles = %s, messages: %s", got, root.Get("messages").Raw) + } + if got := root.Get("messages.1.content.0.text").String(); got != "D1" { + t.Fatalf("in-place turn text = %q, want D1", got) + } +} + +// A system-level item between a tool call and its output must not separate +// tool_use from tool_result. It follows the tool_result message. +func TestResponsesMidSessionSystemItemWaitsForToolResult(t *testing.T) { + root := convertPlacement(t, context.Background(), "claude-opus-5", `{ + "model": "claude-opus-5", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "run it"}]}, + {"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": "{}"}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "prefix approved"}]}, + {"type": "function_call_output", "call_id": "call_1", "output": "done"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "finished"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "thanks"}]} + ] + }`) + if got := root.Get("messages.#.role").Raw; got != `["user","assistant","user","system","assistant","user"]` { + t.Fatalf("roles = %s, messages: %s", got, root.Get("messages").Raw) + } + if got := root.Get("messages.2.content.0.type").String(); got != "tool_result" { + t.Fatalf("messages[2] first block = %q, want tool_result", got) + } + if root.Get("system").Exists() { + t.Fatalf("system must stay empty: %s", root.Get("system").Raw) + } + if problems := claudeMessageInvariantProblems(placementMessages(root)); len(problems) > 0 { + t.Fatalf("invariant problems: %v", problems) + } +} + +// Items with a part a role=system turn cannot carry, items that are not +// messages, and items no user message follows keep the top-level path. The +// structured-output instruction stays the last system block. +func TestResponsesMidSessionSystemItemTopLevelFallbacks(t *testing.T) { + root := convertPlacement(t, context.Background(), "claude-opus-5", `{ + "model": "claude-opus-5", + `+placementFormatJSON+`, + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "custom_note", "role": "developer", "content": "typed"}, + {"type": "message", "role": "developer", "content": [ + {"type": "input_text", "text": "see image"}, + {"type": "input_image", "image_url": "https://example.com/a.png"} + ]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "next"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "trailing"}]}, + {"type": "custom_note", "role": "developer", "content": "after_trailing"} + ] + }`) + if strings.Contains(root.Get("messages").Raw, `"role":"system"`) { + t.Fatalf("no item may become a role=system turn: %s", root.Get("messages").Raw) + } + types := root.Get("system.#.type").Raw + if types != `["text","text","input_image","text","text","text"]` { + t.Fatalf("system block types = %s, system: %s", types, root.Get("system").Raw) + } + if got := root.Get("system.0.text").String(); got != "typed" { + t.Fatalf("system[0] = %q, want the non-message item", got) + } + if got := root.Get("system.3.text").String(); got != "trailing" { + t.Fatalf("system[3] = %q, want the trailing item before after_trailing", got) + } + if got := root.Get("system.4.text").String(); got != "after_trailing" { + t.Fatalf("system[4] = %q, want after_trailing in source order", got) + } + if got := root.Get("system.5.text").String(); !strings.Contains(got, "facts") { + t.Fatalf("structured-output instruction must stay last, got %q", got) + } +} + +// A request carrying only a structured-output format keeps its fallback user +// message now that the instruction is appended after the input loop. +func TestResponsesStructuredOutputOnlyKeepsFallbackUserMessage(t *testing.T) { + root := convertPlacement(t, context.Background(), "claude-opus-5", `{"model":"claude-opus-5",`+placementFormatJSON+`}`) + if got := len(root.Get("system").Array()); got != 1 { + t.Fatalf("system blocks = %d, want 1: %s", got, root.Get("system").Raw) + } + if got := root.Get("messages.#.role").Raw; got != `["user"]` { + t.Fatalf("roles = %s, want the fallback user message", got) + } +} + +// In-place turns keep cache_control exactly as top-level blocks do: block-level +// markers stay, a blank part's marker moves to the preceding block, and the +// item's marker lands on the last block unless that block already has one. +func TestResponsesMidSessionSystemItemKeepsCacheControl(t *testing.T) { + root := convertPlacement(t, context.Background(), "claude-opus-5", `{ + "model": "claude-opus-5", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "message", "role": "developer", "cache_control": {"type": "ephemeral", "ttl": "1h"}, "content": [ + {"type": "input_text", "text": "first", "cache_control": {"type": "ephemeral"}}, + {"type": "input_text", "text": "second"}, + {"type": "input_text", "text": " ", "cache_control": {"type": "ephemeral", "ttl": "5m"}} + ]}, + {"type": "message", "role": "developer", "cache_control": {"type": "ephemeral"}, "content": [ + {"type": "input_text", "text": "third"} + ]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "go"}]} + ] + }`) + if got := root.Get("messages.#.role").Raw; got != `["user","system","system"]` { + t.Fatalf("roles = %s, messages: %s", got, root.Get("messages").Raw) + } + first := root.Get("messages.1.content") + if got := first.Get("0.cache_control.type").String(); got != "ephemeral" { + t.Fatalf("block-level marker lost: %s", first.Raw) + } + if got := first.Get("1.cache_control.ttl").String(); got != "5m" { + t.Fatalf("blank part marker must move to the preceding block: %s", first.Raw) + } + if got := len(first.Array()); got != 2 { + t.Fatalf("blank part must not become a block: %s", first.Raw) + } + if got := root.Get("messages.2.content.0.cache_control.type").String(); got != "ephemeral" { + t.Fatalf("item-level marker lost: %s", root.Get("messages.2").Raw) + } +} + +// Without a pinned decision, the requested model decides; a pinned decision +// overrides it in both directions, and the top-level path keeps source order. +func TestResponsesSystemItemPlacementFollowsPinnedDecision(t *testing.T) { + body := `{ + "model": "m", + "instructions": "I1", + ` + placementFormatJSON + `, + "input": [ + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": "S1"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "U1"}]}, + {"type": "message", "role": "developer", "content": "D1"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "A1"}]}, + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": "S2"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "U2"}]} + ] + }` + for _, test := range []struct { + name string + ctx context.Context + model string + wantTopLevel bool + }{ + {name: "legacy model by default", ctx: context.Background(), model: "claude-sonnet-4-5", wantTopLevel: true}, + {name: "supported model by default", ctx: context.Background(), model: "claude-opus-5", wantTopLevel: false}, + {name: "pinned top-level", ctx: WithTopLevelSystemItems(context.Background(), true), model: "claude-opus-5", wantTopLevel: true}, + {name: "pinned in place", ctx: WithTopLevelSystemItems(context.Background(), false), model: "claude-sonnet-4-5", wantTopLevel: false}, + } { + t.Run(test.name, func(t *testing.T) { + root := convertPlacement(t, test.ctx, test.model, body) + texts := root.Get("system.#.text").Array() + if len(texts) == 0 || !strings.Contains(texts[len(texts)-1].String(), "facts") { + t.Fatalf("structured-output instruction must stay last: %s", root.Get("system").Raw) + } + gotSystem := root.Get("system.#.text").Raw + gotRoles := root.Get("messages.#.role").Raw + if test.wantTopLevel { + if !strings.HasPrefix(gotSystem, `["I1","S1","D1","S2",`) || strings.Contains(gotRoles, "system") { + t.Fatalf("want every item top-level in source order; system=%s roles=%s", gotSystem, gotRoles) + } + return + } + if !strings.HasPrefix(gotSystem, `["I1","S1",`) || gotRoles != `["user","system","assistant","user","system"]` { + t.Fatalf("want D1 and S2 in place; system=%s roles=%s", gotSystem, gotRoles) + } + }) + } +} diff --git a/internal/translator/codex/claude/codex_claude_request.go b/internal/translator/codex/claude/codex_claude_request.go index 00fee4af765..8bfed46f657 100644 --- a/internal/translator/codex/claude/codex_claude_request.go +++ b/internal/translator/codex/claude/codex_claude_request.go @@ -84,6 +84,19 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, var declaredToolsByName map[string]ojson.Result declaredToolsIndexed := false hasAdditionalTools := false + emittedToolNames := make(map[string]struct{}) + if toolsResult.IsArray() { + toolArr := toolsResult.Array() + for i := 0; i < len(toolArr); i++ { + tool := toolArr[i] + if isClaudeWebSearchToolType(tool.Get("type").String()) || tool.Get("defer_loading").Bool() { + continue + } + if name := strings.TrimSpace(tool.Get("name").String()); name != "" { + emittedToolNames[name] = struct{}{} + } + } + } // Process system messages and convert them to input content format. systemsResult := rootResult.Get("system") @@ -144,7 +157,7 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, inputItems = append(inputItems, message) } } - if additionalTools, ok := convertClaudeSystemAdditionalTools(systemContent, toolsResult, &declaredToolsByName, &declaredToolsIndexed, toolNameMap); ok { + if additionalTools, ok := convertClaudeSystemAdditionalTools(systemContent, toolsResult, &declaredToolsByName, &declaredToolsIndexed, toolNameMap, emittedToolNames); ok { hasAdditionalTools = true if len(pendingToolUseIDs) > 0 { pendingSystemReminders = append(pendingSystemReminders, additionalTools) @@ -165,6 +178,43 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, // Empty text is forwarded but does not count: it can hide an emptied attachment. turnSendable := 0 bufferedSendable := 0 + var turnDiscoveredTools [][]byte + + appendDiscoveredToolByReference := func(refName string) { + refName = strings.TrimSpace(refName) + if refName == "" { + return + } + ensureClaudeDeclaredToolsIndexed(toolsResult, &declaredToolsByName, &declaredToolsIndexed) + if declaredToolsByName == nil { + return + } + def, exists := declaredToolsByName[refName] + if !exists || !def.IsObject() || isClaudeWebSearchToolType(def.Get("type").String()) { + return + } + origName := strings.TrimSpace(def.Get("name").String()) + if origName == "" { + return + } + if _, alreadyEmitted := emittedToolNames[origName]; alreadyEmitted { + return + } + emittedToolNames[origName] = struct{}{} + turnDiscoveredTools = append(turnDiscoveredTools, convertClaudeFunctionToolToCodex(def, toolNameMap)) + } + + flushTurnDiscoveredTools := func() { + if len(turnDiscoveredTools) == 0 { + return + } + hasAdditionalTools = true + turnSendable++ + itemDoc := ojson.NewDocString(`{"type":"additional_tools","role":"developer","tools":[]}`) + _ = itemDoc.SetRawBytes("tools", translatorcommon.JoinRawArray(turnDiscoveredTools)) + inputItems = append(inputItems, itemDoc.Bytes()) + turnDiscoveredTools = turnDiscoveredTools[:0] + } flushMessage := func() { if len(contentItems) > 0 { @@ -278,14 +328,17 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, switch contentType { case "text": + flushTurnDiscoveredTools() if len(pendingSystemReminders) > 0 { inputItems = append(inputItems, pendingSystemReminders...) pendingSystemReminders = nil } appendTextContent(messageContentResult.Get("text").String()) case "thinking": + flushTurnDiscoveredTools() appendReasoningContent(messageContentResult) case "image": + flushTurnDiscoveredTools() if len(pendingSystemReminders) > 0 { inputItems = append(inputItems, pendingSystemReminders...) pendingSystemReminders = nil @@ -296,6 +349,7 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, drops.Drop(contentType) } case "document", "container_upload": + flushTurnDiscoveredTools() if len(pendingSystemReminders) > 0 { inputItems = append(inputItems, pendingSystemReminders...) pendingSystemReminders = nil @@ -306,6 +360,7 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, drops.Drop(contentType) } case "tool_use": + flushTurnDiscoveredTools() pendingCommentaryMirror = "" flushMessage() if id := messageContentResult.Get("id").String(); id != "" { @@ -324,6 +379,36 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, } _ = functionCallMessageDoc.Set("arguments", messageContentResult.Get("input").Raw) inputItems = append(inputItems, functionCallMessageDoc.Bytes()) + case "tool_reference": + if len(pendingSystemReminders) > 0 { + inputItems = append(inputItems, pendingSystemReminders...) + pendingSystemReminders = nil + } + flushMessage() + refName := messageContentResult.Get("tool_name").String() + if refName == "" { + refName = messageContentResult.Get("name").String() + } + appendDiscoveredToolByReference(refName) + case "tool_search_tool_result": + if len(pendingSystemReminders) > 0 { + inputItems = append(inputItems, pendingSystemReminders...) + pendingSystemReminders = nil + } + flushMessage() + toolRefs := messageContentResult.Get("content.tool_references") + if toolRefs.IsArray() { + for _, refPart := range toolRefs.Array() { + if refPart.Get("type").String() != "tool_reference" { + continue + } + refName := refPart.Get("tool_name").String() + if refName == "" { + refName = refPart.Get("name").String() + } + appendDiscoveredToolByReference(refName) + } + } case "tool_result": flushMessage() functionCallOutputMessageDoc := ojson.NewDocString(`{"type":"function_call_output"}`) @@ -333,10 +418,14 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, if contentResult.IsArray() { contentResults := contentResult.Array() toolResultContentItems := make([][]byte, 0, len(contentResults)) + hasRichToolResultItems := false + rewroteToolRefName := false + var rewrittenToolRefs [][]byte for k := 0; k < len(contentResults); k++ { - toolResultContentType := contentResults[k].Get("type").String() + part := contentResults[k] + toolResultContentType := part.Get("type").String() if toolResultContentType == "image" { - sourceResult := contentResults[k].Get("source") + sourceResult := part.Get("source") if sourceResult.Exists() { data := sourceResult.Get("data").String() if data == "" { @@ -355,16 +444,56 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, toolResultContent := []byte(`{"type":"input_image","image_url":""}`) toolResultContent, _ = ojson.SetBytes(toolResultContent, "image_url", dataURL) toolResultContentItems = append(toolResultContentItems, toolResultContent) + hasRichToolResultItems = true } } + if rewroteToolRefName { + rewrittenToolRefs = append(rewrittenToolRefs, []byte(part.Raw)) + } } else if toolResultContentType == "text" { toolResultContent := []byte(`{"type":"input_text","text":""}`) - toolResultContent, _ = ojson.SetBytes(toolResultContent, "text", contentResults[k].Get("text").String()) + toolResultContent, _ = ojson.SetBytes(toolResultContent, "text", part.Get("text").String()) toolResultContentItems = append(toolResultContentItems, toolResultContent) + hasRichToolResultItems = true + if rewroteToolRefName { + rewrittenToolRefs = append(rewrittenToolRefs, []byte(part.Raw)) + } + } else if toolResultContentType == "tool_reference" { + refField := "tool_name" + refName := strings.TrimSpace(part.Get("tool_name").String()) + if refName == "" { + refField = "name" + refName = strings.TrimSpace(part.Get("name").String()) + } + appendDiscoveredToolByReference(refName) + refBytes := []byte(part.Raw) + if short, ok := toolNameMap[refName]; ok && short != refName { + refDoc := ojson.NewDoc(refBytes) + if refDoc.Set(refField, short) == nil { + refBytes = refDoc.Bytes() + if !rewroteToolRefName { + rewroteToolRefName = true + rewrittenToolRefs = make([][]byte, 0, len(contentResults)) + for prev := 0; prev < k; prev++ { + rewrittenToolRefs = append(rewrittenToolRefs, []byte(contentResults[prev].Raw)) + } + } + } + } + if rewroteToolRefName { + rewrittenToolRefs = append(rewrittenToolRefs, refBytes) + } + toolResultContentDoc := ojson.NewDocString(`{"type":"input_text","text":""}`) + _ = toolResultContentDoc.Set("text", string(refBytes)) + toolResultContentItems = append(toolResultContentItems, toolResultContentDoc.Bytes()) + } else if rewroteToolRefName { + rewrittenToolRefs = append(rewrittenToolRefs, []byte(part.Raw)) } } - if len(toolResultContentItems) > 0 { + if hasRichToolResultItems { _ = functionCallOutputMessageDoc.SetRawBytes("output", translatorcommon.JoinRawArray(toolResultContentItems)) + } else if rewroteToolRefName && len(rewrittenToolRefs) == len(contentResults) { + _ = functionCallOutputMessageDoc.Set("output", string(translatorcommon.JoinRawArray(rewrittenToolRefs))) } else { _ = functionCallOutputMessageDoc.Set("output", messageContentResult.Get("content").String()) } @@ -377,6 +506,7 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, } } flushMessage() + flushTurnDiscoveredTools() if len(pendingSystemReminders) > 0 { inputItems = append(inputItems, pendingSystemReminders...) pendingSystemReminders = nil @@ -407,6 +537,10 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, _ = templateDoc.SetRawBytes("tool_choice", convertClaudeToolChoiceToCodex(rootResult.Get("tool_choice"), toolNameMap, webSearchToolNames)) } if toolsResult.IsArray() { + forcedToolName := "" + if tc := rootResult.Get("tool_choice"); tc.IsObject() && tc.Get("type").String() == "tool" { + forcedToolName = strings.TrimSpace(tc.Get("name").String()) + } toolResults := toolsResult.Array() toolItems = make([][]byte, 0, len(toolResults)) for i := 0; i < len(toolResults); i++ { @@ -417,6 +551,12 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, toolItems = append(toolItems, convertClaudeWebSearchToolToCodex(toolResult)) continue } + if toolResult.Get("defer_loading").Bool() { + name := strings.TrimSpace(toolResult.Get("name").String()) + if _, alreadyEmitted := emittedToolNames[name]; alreadyEmitted || forcedToolName == "" || name != forcedToolName { + continue + } + } toolItems = append(toolItems, convertClaudeFunctionToolToCodex(toolResult, toolNameMap)) } } @@ -739,7 +879,7 @@ func ensureClaudeDeclaredToolsIndexed(toolsResult ojson.Result, declaredToolsByN *declaredToolsByName = make(map[string]ojson.Result, len(toolArr)) } for i := 0; i < len(toolArr); i++ { - if name := toolArr[i].Get("name").String(); name != "" { + if name := strings.TrimSpace(toolArr[i].Get("name").String()); name != "" { if _, exists := (*declaredToolsByName)[name]; !exists { (*declaredToolsByName)[name] = toolArr[i] } @@ -749,7 +889,7 @@ func ensureClaudeDeclaredToolsIndexed(toolsResult ojson.Result, declaredToolsByN // convertClaudeSystemAdditionalTools converts mid-conversation tool_addition blocks // inside a role=system message into a Responses additional_tools input item. -func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, declaredToolsByName *map[string]ojson.Result, declaredToolsIndexed *bool, toolNameMap map[string]string) ([]byte, bool) { +func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, declaredToolsByName *map[string]ojson.Result, declaredToolsIndexed *bool, toolNameMap map[string]string, emittedToolNames map[string]struct{}) ([]byte, bool) { if !content.IsArray() { return nil, false } @@ -763,6 +903,7 @@ func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, decla } toolSpec := part.Get("tool") var def ojson.Result + isReference := false switch toolSpec.Get("type").String() { case "tool_definition": def = toolSpec.Get("definition") @@ -776,6 +917,7 @@ func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, decla } } case "tool_reference": + isReference = true ensureClaudeDeclaredToolsIndexed(toolsResult, declaredToolsByName, declaredToolsIndexed) refName := toolSpec.Get("name").String() if refName == "" { @@ -803,6 +945,14 @@ func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, decla if origName == "" { continue } + if isReference && emittedToolNames != nil { + if _, alreadyEmitted := emittedToolNames[origName]; alreadyEmitted { + continue + } + } + if emittedToolNames != nil { + emittedToolNames[origName] = struct{}{} + } converted := convertClaudeFunctionToolToCodex(def, toolNameMap) if indexByName == nil { indexByName = make(map[string]int) @@ -817,60 +967,132 @@ func convertClaudeSystemAdditionalTools(content, toolsResult ojson.Result, decla if len(toolItems) == 0 { return nil, false } - item := []byte(`{"type":"additional_tools","role":"developer","tools":[]}`) - item, _ = ojson.SetRawBytes(item, "tools", translatorcommon.JoinRawArray(toolItems)) - return item, true + itemDoc := ojson.NewDocString(`{"type":"additional_tools","role":"developer","tools":[]}`) + _ = itemDoc.SetRawBytes("tools", translatorcommon.JoinRawArray(toolItems)) + return itemDoc.Bytes(), true } -func collectClaudeDeclaredToolNames(rootResult ojson.Result, hasToolAddition bool) []string { +func collectClaudeDeclaredToolNames(rootResult ojson.Result, hasDynamicTools bool) []string { var names []string + var deferredNames []string + var deferredSet map[string]struct{} tools := rootResult.Get("tools") if tools.IsArray() { arr := tools.Array() names = make([]string, 0, len(arr)) for i := 0; i < len(arr); i++ { n := arr[i].Get("name").String() - if n != "" { - names = append(names, n) + if n == "" { + continue + } + if hasDynamicTools && arr[i].Get("defer_loading").Bool() { + deferredNames = append(deferredNames, n) + if deferredSet == nil { + deferredSet = make(map[string]struct{}) + } + deferredSet[n] = struct{}{} + continue } + names = append(names, n) } } - if !hasToolAddition { + if !hasDynamicTools { return names } - messages := rootResult.Get("messages") - if !messages.IsArray() { - return names - } - for _, msg := range messages.Array() { - if msg.Get("role").String() != "system" { - continue + appendDeferredRef := func(refName string) { + refName = strings.TrimSpace(refName) + if refName == "" || deferredSet == nil { + return } - content := msg.Get("content") - if !content.IsArray() { - continue + if _, ok := deferredSet[refName]; ok { + names = append(names, refName) } - for _, part := range content.Array() { - if part.Get("type").String() != "tool_addition" { + } + messages := rootResult.Get("messages") + if messages.IsArray() { + for _, msg := range messages.Array() { + content := msg.Get("content") + if !content.IsArray() { continue } - tool := part.Get("tool") - if tool.Get("type").String() == "tool_definition" || tool.Get("definition").IsObject() { - def := tool.Get("definition") - if isClaudeWebSearchToolType(def.Get("type").String()) { - continue + if msg.Get("role").String() == "system" { + for _, part := range content.Array() { + if part.Get("type").String() != "tool_addition" { + continue + } + tool := part.Get("tool") + switch tool.Get("type").String() { + case "tool_definition": + def := tool.Get("definition") + if isClaudeWebSearchToolType(def.Get("type").String()) { + continue + } + if n := def.Get("name").String(); n != "" { + names = append(names, n) + } + case "tool_reference": + refName := tool.Get("name").String() + if refName == "" { + refName = tool.Get("tool_name").String() + } + appendDeferredRef(refName) + default: + if def := tool.Get("definition"); def.IsObject() && !isClaudeWebSearchToolType(def.Get("type").String()) { + if n := def.Get("name").String(); n != "" { + names = append(names, n) + } + } + } } - if n := def.Get("name").String(); n != "" { - names = append(names, n) + continue + } + for _, part := range content.Array() { + switch part.Get("type").String() { + case "tool_reference": + refName := part.Get("tool_name").String() + if refName == "" { + refName = part.Get("name").String() + } + appendDeferredRef(refName) + case "tool_search_tool_result": + toolRefs := part.Get("content.tool_references") + if toolRefs.IsArray() { + for _, refPart := range toolRefs.Array() { + if refPart.Get("type").String() != "tool_reference" { + continue + } + refName := refPart.Get("tool_name").String() + if refName == "" { + refName = refPart.Get("name").String() + } + appendDeferredRef(refName) + } + } + case "tool_result": + nested := part.Get("content") + if !nested.IsArray() { + continue + } + for _, nestedPart := range nested.Array() { + if nestedPart.Get("type").String() != "tool_reference" { + continue + } + refName := nestedPart.Get("tool_name").String() + if refName == "" { + refName = nestedPart.Get("name").String() + } + appendDeferredRef(refName) + } } } } } + names = append(names, deferredNames...) return names } -func buildReverseMapFromClaudeOriginalToShortFromResult(rootResult ojson.Result, hasToolAddition bool) map[string]string { - names := collectClaudeDeclaredToolNames(rootResult, hasToolAddition) +func buildReverseMapFromClaudeOriginalToShortFromResult(rootResult ojson.Result, hasDynamicTools bool) map[string]string { + names := collectClaudeDeclaredToolNames(rootResult, hasDynamicTools) if len(names) == 0 { return map[string]string{} } @@ -879,6 +1101,8 @@ func buildReverseMapFromClaudeOriginalToShortFromResult(rootResult ojson.Result, func claudePayloadMayContainToolAddition(rawJSON []byte) bool { return bytes.Contains(rawJSON, []byte("tool_addition")) || + bytes.Contains(rawJSON, []byte("defer_loading")) || + bytes.Contains(rawJSON, []byte("tool_reference")) || bytes.Contains(rawJSON, []byte(`\u`)) || bytes.Contains(rawJSON, []byte(`\U`)) } diff --git a/internal/translator/codex/claude/codex_claude_request_test.go b/internal/translator/codex/claude/codex_claude_request_test.go index 451256f2e84..e4ca4f2a7c8 100644 --- a/internal/translator/codex/claude/codex_claude_request_test.go +++ b/internal/translator/codex/claude/codex_claude_request_test.go @@ -1529,13 +1529,13 @@ func TestConvertClaudeRequestToCodex_MidConversationToolAdditions(t *testing.T) t.Fatalf("ConvertClaudeRequestToCodex returned error: %v", err) } - // Top-level tools[] must remain untouched (2 tools: Read and deferred_lookup) to preserve the prompt cache prefix. + // Top-level tools[] must include only non-deferred tools (Read) to preserve the prompt cache prefix. topTools := ojson.GetBytesNoCopy(translated, "tools").Array() - if len(topTools) != 2 { - t.Fatalf("expected 2 top-level tools, got %d: %s", len(topTools), ojson.GetBytesNoCopy(translated, "tools").Raw) + if len(topTools) != 1 { + t.Fatalf("expected 1 top-level tool, got %d: %s", len(topTools), ojson.GetBytesNoCopy(translated, "tools").Raw) } - if topTools[0].Get("name").String() != "Read" || topTools[1].Get("name").String() != "deferred_lookup" { - t.Fatalf("unexpected top-level tool names: %q, %q", topTools[0].Get("name").String(), topTools[1].Get("name").String()) + if topTools[0].Get("name").String() != "Read" { + t.Fatalf("unexpected top-level tool name: %q", topTools[0].Get("name").String()) } // Expected input items: @@ -1593,3 +1593,501 @@ func TestConvertClaudeRequestToCodex_MidConversationToolAdditions(t *testing.T) t.Fatalf("expected shortened MCP tool name <= 64 chars, got %q (len=%d)", shortLongName, len(shortLongName)) } } + +func TestConvertClaudeRequestToCodex_ToolSearchDeferredLoadingPreservesPromptCachePrefix(t *testing.T) { + beforeDiscoveryJSON := `{ + "model": "gpt-6-astra", + "system": "You are Claude Code.", + "tools": [ + { + "name": "ToolSearch", + "description": "Search deferred tools", + "input_schema": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "number"} + }, + "required": ["query"] + } + }, + { + "name": "Workflow", + "description": "Run workflow", + "input_schema": {"type": "object", "properties": {"name": {"type": "string"}}} + }, + { + "name": "DeferredToolPlaceholder", + "description": "Placeholder for deferred tools", + "defer_loading": true, + "input_schema": {"type": "object", "properties": {}} + }, + { + "name": "Write", + "description": "Write file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}} + } + ], + "messages": [ + {"role": "user", "content": "Load WebSearch and WebFetch"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_search_1", + "name": "ToolSearch", + "input": {"query": "select:WebSearch,WebFetch", "max_results": 2} + } + ] + } + ] + }` + + afterDiscoveryJSON := `{ + "model": "gpt-6-astra", + "system": "You are Claude Code.", + "tools": [ + { + "name": "ToolSearch", + "description": "Search deferred tools", + "input_schema": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "number"} + }, + "required": ["query"] + } + }, + { + "name": "WebFetch", + "description": "Fetch a URL", + "defer_loading": true, + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"url": {"type": "string"}}, + "required": ["url"] + } + }, + { + "name": "WebSearch", + "description": "Search the web", + "defer_loading": true, + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"] + } + }, + { + "name": "Workflow", + "description": "Run workflow", + "input_schema": {"type": "object", "properties": {"name": {"type": "string"}}} + }, + { + "name": "DeferredToolPlaceholder", + "description": "Placeholder for deferred tools", + "defer_loading": true, + "input_schema": {"type": "object", "properties": {}} + }, + { + "name": "Write", + "description": "Write file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}} + } + ], + "messages": [ + {"role": "user", "content": "Load WebSearch and WebFetch"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_search_1", + "name": "ToolSearch", + "input": {"query": "select:WebSearch,WebFetch", "max_results": 2} + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_search_1", + "content": [ + {"type": "tool_reference", "tool_name": "WebSearch"}, + {"type": "tool_reference", "tool_name": "WebFetch"}, + {"type": "tool_reference", "tool_name": "WebSearch"} + ] + } + ] + } + ] + }` + + beforeOut, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(beforeDiscoveryJSON), false) + if err != nil { + t.Fatalf("beforeDiscovery ConvertClaudeRequestToCodex error: %v", err) + } + afterOut, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(afterDiscoveryJSON), false) + if err != nil { + t.Fatalf("afterDiscovery ConvertClaudeRequestToCodex error: %v", err) + } + + beforeToolsRaw := ojson.GetBytesNoCopy(beforeOut, "tools").Raw + afterToolsRaw := ojson.GetBytesNoCopy(afterOut, "tools").Raw + if beforeToolsRaw != afterToolsRaw { + t.Fatalf("top-level tools[] changed after ToolSearch discovery:\nbefore: %s\nafter: %s", beforeToolsRaw, afterToolsRaw) + } + + topTools := ojson.GetBytesNoCopy(afterOut, "tools").Array() + if len(topTools) != 3 { + t.Fatalf("expected 3 non-deferred top-level tools (ToolSearch, Workflow, Write), got %d: %s", len(topTools), afterToolsRaw) + } + wantTopNames := []string{"ToolSearch", "Workflow", "Write"} + for i, wantName := range wantTopNames { + if got := topTools[i].Get("name").String(); got != wantName { + t.Fatalf("topTools[%d].name = %q, want %q", i, got, wantName) + } + } + + beforeInputs := ojson.GetBytesNoCopy(beforeOut, "input").Array() + afterInputs := ojson.GetBytesNoCopy(afterOut, "input").Array() + if len(beforeInputs) != 3 { + t.Fatalf("expected 3 input items before discovery, got %d: %s", len(beforeInputs), ojson.GetBytesNoCopy(beforeOut, "input").Raw) + } + if len(afterInputs) != 5 { + t.Fatalf("expected 5 input items after discovery (developer, user, function_call, function_call_output, additional_tools), got %d: %s", len(afterInputs), ojson.GetBytesNoCopy(afterOut, "input").Raw) + } + for i := 0; i < len(beforeInputs); i++ { + if beforeInputs[i].Raw != afterInputs[i].Raw { + t.Fatalf("input[%d] prefix changed after discovery:\nbefore: %s\nafter: %s", i, beforeInputs[i].Raw, afterInputs[i].Raw) + } + } + + if got := afterInputs[3].Get("type").String(); got != "function_call_output" { + t.Fatalf("afterInputs[3].type = %q, want function_call_output", got) + } + addTools := afterInputs[4] + if got := addTools.Get("type").String(); got != "additional_tools" { + t.Fatalf("afterInputs[4].type = %q, want additional_tools: %s", got, addTools.Raw) + } + if got := addTools.Get("role").String(); got != "developer" { + t.Fatalf("afterInputs[4].role = %q, want developer: %s", got, addTools.Raw) + } + discovered := addTools.Get("tools").Array() + if len(discovered) != 2 { + t.Fatalf("expected 2 deduplicated discovered tools in additional_tools, got %d: %s", len(discovered), addTools.Raw) + } + if got := discovered[0].Get("name").String(); got != "WebSearch" { + t.Fatalf("discovered[0].name = %q, want WebSearch (preserving reference order)", got) + } + if got := discovered[1].Get("name").String(); got != "WebFetch" { + t.Fatalf("discovered[1].name = %q, want WebFetch (preserving reference order)", got) + } + for i, tool := range discovered { + if tool.Get("defer_loading").Exists() { + t.Fatalf("discovered[%d] still has defer_loading: %s", i, tool.Raw) + } + if tool.Get("input_schema").Exists() || tool.Get("parameters.$schema").Exists() { + t.Fatalf("discovered[%d] still has unnormalized schema fields: %s", i, tool.Raw) + } + if tool.Get("strict").Type != ojson.False { + t.Fatalf("discovered[%d] expected strict=false: %s", i, tool.Raw) + } + } +} + +func TestConvertClaudeRequestToCodex_RepeatedToolSearchAndLongMCPToolMapping(t *testing.T) { + // Two >64-char MCP tools that collide under shortenNameIfNeeded ("mcp__shared-long-server_deliver_final_subagent_handback_report"): + // eagerLongTool is non-deferred; deferredLongTool is deferred and inserted earlier in tools[] after discovery. + eagerLongTool := "mcp__shared-long-server-11111111-2222-3333-4444-555555555555_deliver_final_subagent_handback_report" + deferredLongTool := "mcp__shared-long-server-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee_deliver_final_subagent_handback_report" + + beforeJSON := fmt.Sprintf(`{ + "model": "gpt-6-astra", + "tools": [ + {"name": "ToolSearch", "input_schema": {"type": "object"}}, + {"name": %q, "description": "Eager MCP tool", "input_schema": {"type": "object"}} + ], + "messages": [ + {"role": "user", "content": "Search tools"} + ] + }`, eagerLongTool) + + afterJSON := fmt.Sprintf(`{ + "model": "gpt-6-astra", + "tools": [ + {"name": "ToolSearch", "input_schema": {"type": "object"}}, + {"name": %q, "description": "Deferred MCP tool inserted before eager tool", "defer_loading": true, "input_schema": {"type": "object"}}, + {"name": %q, "description": "Eager MCP tool", "input_schema": {"type": "object"}} + ], + "messages": [ + {"role": "user", "content": "Search tools"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call_ts_1", "name": "ToolSearch", "input": {"query": "action"}} + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_ts_1", + "content": [ + {"type": "tool_reference", "tool_name": %q} + ] + } + ] + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call_ts_2", "name": "ToolSearch", "input": {"query": "action again"}}, + {"type": "tool_use", "id": "call_mcp_1", "name": %q, "input": {}} + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_ts_2", + "content": [ + {"type": "tool_reference", "tool_name": %q} + ] + }, + { + "type": "tool_result", + "tool_use_id": "call_mcp_1", + "content": "done" + } + ] + } + ] + }`, deferredLongTool, eagerLongTool, deferredLongTool, deferredLongTool, deferredLongTool) + + beforeOut, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(beforeJSON), false) + if err != nil { + t.Fatalf("before ConvertClaudeRequestToCodex error: %v", err) + } + afterOut, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(afterJSON), false) + if err != nil { + t.Fatalf("after ConvertClaudeRequestToCodex error: %v", err) + } + + // Top-level tools[] must stay byte-identical even when a colliding deferred tool is inserted earlier in tools[]. + if gotBefore, gotAfter := ojson.GetBytesNoCopy(beforeOut, "tools").Raw, ojson.GetBytesNoCopy(afterOut, "tools").Raw; gotBefore != gotAfter { + t.Fatalf("top-level tools[] changed when colliding deferred tool was inserted:\nbefore: %s\nafter: %s", gotBefore, gotAfter) + } + + inputs := ojson.GetBytesNoCopy(afterOut, "input").Array() + // Count additional_tools items: only the first ToolSearch result should emit additional_tools; + // the second ToolSearch result references the same already-loaded tool and must not emit a duplicate additional_tools item. + var addToolsItems []ojson.Result + for _, item := range inputs { + if item.Get("type").String() == "additional_tools" { + addToolsItems = append(addToolsItems, item) + } + } + if len(addToolsItems) != 1 { + t.Fatalf("expected exactly 1 additional_tools item across repeated references, got %d: %s", len(addToolsItems), ojson.GetBytesNoCopy(afterOut, "input").Raw) + } + + deferredShortName := addToolsItems[0].Get("tools.0.name").String() + eagerShortName := ojson.GetBytesNoCopy(afterOut, "tools.1.name").String() + if deferredShortName == "" || deferredShortName == deferredLongTool || len(deferredShortName) > 64 { + t.Fatalf("expected shortened deferred tool name <= 64 chars, got %q", deferredShortName) + } + if deferredShortName == eagerShortName { + t.Fatalf("expected colliding deferred and eager tools to have distinct short names, both got %q", deferredShortName) + } + + if !strings.HasSuffix(deferredShortName, "_1") { + t.Fatalf("expected deferred colliding tool to receive _1 suffix so eager tool keeps base name, got eager=%q deferred=%q", eagerShortName, deferredShortName) + } + + // Verify the assistant tool_use for deferredLongTool and the response reverse map use the same shortened name. + if gotCallName := ojson.GetBytesNoCopy(afterOut, `input.#(call_id=="call_mcp_1").name`).String(); gotCallName != deferredShortName { + t.Fatalf("function_call name = %q, want %q", gotCallName, deferredShortName) + } + revMap := buildReverseMapFromClaudeOriginalShortToOriginal([]byte(afterJSON)) + if gotOrig := revMap[deferredShortName]; gotOrig != deferredLongTool { + t.Fatalf("reverse map[%q] = %q, want %q", deferredShortName, gotOrig, deferredLongTool) + } + if gotOrig := revMap[eagerShortName]; gotOrig != eagerLongTool { + t.Fatalf("reverse map[%q] = %q, want %q", eagerShortName, gotOrig, eagerLongTool) + } + // Verify function_call_output.output also rewrote tool_reference.tool_name to the shortened name. + var ts1Output string + for _, item := range inputs { + if item.Get("type").String() == "function_call_output" && item.Get("call_id").String() == "call_ts_1" { + ts1Output = item.Get("output").String() + break + } + } + if !strings.Contains(ts1Output, deferredShortName) { + t.Fatalf("expected function_call_output.output to use shortened tool_name %q, got %s", deferredShortName, ts1Output) + } +} + +func TestConvertClaudeRequestToCodex_ToolResultMixedContentAndUnknownToolReference(t *testing.T) { + inputJSON := `{ + "model": "gpt-6-astra", + "tools": [ + {"name": "ToolSearch", "input_schema": {"type": "object"}}, + {"name": "WebSearch", "defer_loading": true, "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}} + ], + "messages": [ + {"role": "user", "content": "Find tools"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call_1", "name": "ToolSearch", "input": {"query": "web"}} + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + {"type": "text", "text": "Matched tools:"}, + {"type": "tool_reference", "tool_name": "WebSearch"}, + {"type": "tool_reference", "tool_name": "NonExistentTool"} + ] + } + ] + } + ] + }` + + out, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(inputJSON), false) + if err != nil { + t.Fatalf("ConvertClaudeRequestToCodex error: %v", err) + } + + outputParts := ojson.GetBytesNoCopy(out, `input.#(type=="function_call_output").output`).Array() + if len(outputParts) != 3 { + t.Fatalf("expected 3 output parts in mixed tool_result, got %d: %s", len(outputParts), ojson.GetBytesNoCopy(out, `input.#(call_id=="call_1").output`).Raw) + } + if outputParts[0].Get("text").String() != "Matched tools:" { + t.Fatalf("unexpected outputParts[0]: %s", outputParts[0].Raw) + } + if !strings.Contains(outputParts[1].Get("text").String(), "WebSearch") { + t.Fatalf("expected WebSearch tool_reference in outputParts[1], got: %s", outputParts[1].Raw) + } + + addTools := ojson.GetBytesNoCopy(out, `input.#(type=="additional_tools").tools`).Array() + if len(addTools) != 1 || addTools[0].Get("name").String() != "WebSearch" { + t.Fatalf("expected only known WebSearch in additional_tools, got: %s", ojson.GetBytesNoCopy(out, `input.#(type=="additional_tools")`).Raw) + } +} + +func TestConvertClaudeRequestToCodex_StandaloneToolReferenceAndSearchResultOrdering(t *testing.T) { + longDeferredTool := "mcp__shared-long-server-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee_deliver_final_subagent_handback_report" + inputJSON := fmt.Sprintf(`{ + "model": "gpt-6-astra", + "tools": [ + {"name": "ToolSearch", "input_schema": {"type": "object"}}, + {"name": "WebSearch", "defer_loading": true, "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}}, + {"name": "WebFetch", "defer_loading": true, "input_schema": {"type": "object", "properties": {"url": {"type": "string"}}}}, + {"name": %q, "defer_loading": true, "input_schema": {"type": "object"}} + ], + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "before standalone ref"}, + {"type": "tool_reference", "tool_name": "WebSearch"}, + {"type": "text", "text": "between ref and search result"}, + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [ + {"type": "tool_reference", "tool_name": "WebFetch"} + ] + } + }, + {"type": "text", "text": "after search result"} + ] + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call_img_ref", "name": "ToolSearch", "input": {"query": "long"}} + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_img_ref", + "content": [ + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "aGVsbG8="} + }, + {"type": "tool_reference", "tool_name": %q} + ] + } + ] + } + ] + }`, longDeferredTool, longDeferredTool) + + out, err := ConvertClaudeRequestToCodex("gpt-6-astra", []byte(inputJSON), false) + if err != nil { + t.Fatalf("ConvertClaudeRequestToCodex error: %v", err) + } + + inputs := ojson.GetBytesNoCopy(out, "input").Array() + // Expected input sequence: + // 0: message ("before standalone ref") + // 1: additional_tools (WebSearch) + // 2: message ("between ref and search result") + // 3: additional_tools (WebFetch) + // 4: message ("after search result") + // 5: function_call (call_img_ref) + // 6: function_call_output (call_img_ref with input_image + shortened tool_reference input_text) + // 7: additional_tools (shortened longDeferredTool) + if len(inputs) != 8 { + t.Fatalf("expected 8 input items, got %d: %s", len(inputs), ojson.GetBytesNoCopy(out, "input").Raw) + } + if inputs[0].Get("type").String() != "message" || inputs[0].Get("content.0.text").String() != "before standalone ref" { + t.Fatalf("inputs[0] = %s, want message 'before standalone ref'", inputs[0].Raw) + } + if inputs[1].Get("type").String() != "additional_tools" || inputs[1].Get("tools.0.name").String() != "WebSearch" { + t.Fatalf("inputs[1] = %s, want additional_tools(WebSearch)", inputs[1].Raw) + } + if inputs[2].Get("type").String() != "message" || inputs[2].Get("content.0.text").String() != "between ref and search result" { + t.Fatalf("inputs[2] = %s, want message 'between ref and search result'", inputs[2].Raw) + } + if inputs[3].Get("type").String() != "additional_tools" || inputs[3].Get("tools.0.name").String() != "WebFetch" { + t.Fatalf("inputs[3] = %s, want additional_tools(WebFetch)", inputs[3].Raw) + } + if inputs[4].Get("type").String() != "message" || inputs[4].Get("content.0.text").String() != "after search result" { + t.Fatalf("inputs[4] = %s, want message 'after search result'", inputs[4].Raw) + } + if inputs[6].Get("type").String() != "function_call_output" { + t.Fatalf("inputs[6] = %s, want function_call_output", inputs[6].Raw) + } + outArr := inputs[6].Get("output").Array() + if len(outArr) != 2 || outArr[0].Get("type").String() != "input_image" || outArr[1].Get("type").String() != "input_text" { + t.Fatalf("expected [input_image, input_text] in mixed tool_result output, got: %s", inputs[6].Get("output").Raw) + } + shortLong := inputs[7].Get("tools.0.name").String() + if shortLong == "" || shortLong == longDeferredTool || !strings.Contains(outArr[1].Get("text").String(), shortLong) { + t.Fatalf("expected shortened tool name %q in mixed tool_result input_text, got: %s", shortLong, outArr[1].Raw) + } +} diff --git a/internal/translator/codex/openai/chat-completions/codex_openai_request.go b/internal/translator/codex/openai/chat-completions/codex_openai_request.go index 389a65fce69..d46c9438904 100644 --- a/internal/translator/codex/openai/chat-completions/codex_openai_request.go +++ b/internal/translator/codex/openai/chat-completions/codex_openai_request.go @@ -85,7 +85,7 @@ func convertOpenAIRequestToCodex(modelName string, inputRawJSON []byte, stream b case "function": functionToolNames[tool.Get("function.name").String()] = struct{}{} case "custom": - customToolNames[tool.Get("name").String()] = struct{}{} + customToolNames[customToolDefinition(tool).Get("name").String()] = struct{}{} } } // A normalized function envelope cannot disambiguate declarations that share a name. @@ -424,15 +424,21 @@ func convertOpenAIRequestToCodex(modelName string, inputRawJSON []byte, stream b t := arr[i] toolType := t.Get("type").String() if toolType == "custom" { - item := []byte(t.Raw) - name := t.Get("name").String() + definition := customToolDefinition(t) + itemDoc := ojson.NewDocString(definition.Raw) + _ = itemDoc.Set("type", "custom") + if grammar := definition.Get("format.grammar"); definition.Get("format.type").String() == "grammar" && grammar.IsObject() { + _ = itemDoc.SetRaw("format", grammar.Raw) + _ = itemDoc.Set("format.type", "grammar") + } + name := definition.Get("name").String() if short, ok := originalToolNameMap[name]; ok { name = short } else { name = shortenNameIfNeeded(name) } - item, _ = ojson.SetBytes(item, "name", name) - toolItems = append(toolItems, item) + _ = itemDoc.Set("name", name) + toolItems = append(toolItems, itemDoc.Bytes()) continue } @@ -487,7 +493,7 @@ func convertOpenAIRequestToCodex(modelName string, inputRawJSON []byte, stream b case tc.IsObject(): tcType := tc.Get("type").String() if tcType == "function" || tcType == "custom" { - name := tc.Get("name").String() + name := customToolDefinition(tc).Get("name").String() if tcType == "function" { name = tc.Get("function.name").String() if _, custom := customToolNames[name]; custom { @@ -682,6 +688,14 @@ func shortenNameIfNeeded(name string) string { return sanitized[:limit] } +// customToolDefinition accepts both Chat Completions and Responses declarations. +func customToolDefinition(tool ojson.Result) ojson.Result { + if custom := tool.Get("custom"); custom.IsObject() { + return custom + } + return tool +} + // collectRequestToolNames extracts unique tool names across tools declarations, // tool_choice, and historical assistant tool_calls in a deterministic order. func collectRequestToolNames(rawJSON []byte) []string { @@ -706,7 +720,7 @@ func collectRequestToolNames(rawJSON []byte) []string { case "function": addName(tool.Get("function.name").String()) case "custom": - addName(tool.Get("name").String()) + addName(customToolDefinition(tool).Get("name").String()) } } } @@ -722,7 +736,7 @@ func collectRequestToolNames(rawJSON []byte) []string { } addName(fnName) case "custom": - addName(tc.Get("name").String()) + addName(customToolDefinition(tc).Get("name").String()) } } diff --git a/internal/translator/codex/openai/chat-completions/codex_openai_request_test.go b/internal/translator/codex/openai/chat-completions/codex_openai_request_test.go index 96b41fea73f..09093c18b8a 100644 --- a/internal/translator/codex/openai/chat-completions/codex_openai_request_test.go +++ b/internal/translator/codex/openai/chat-completions/codex_openai_request_test.go @@ -1,6 +1,7 @@ package chat_completions import ( + "reflect" "regexp" "strings" "testing" @@ -848,6 +849,86 @@ func TestToolNameShortening(t *testing.T) { } } +func TestNestedCustomToolDeclaration(t *testing.T) { + for _, tt := range []struct { + name string + format string + want string + }{ + {"grammar", `,"format":{"type":"grammar","grammar":{"syntax":"lark","definition":"start: patch"}}`, `{"type":"grammar","syntax":"lark","definition":"start: patch"}`}, + {"text", `,"format":{"type":"text"}`, `{"type":"text"}`}, + {"default", "", ""}, + } { + t.Run(tt.name, func(t *testing.T) { + input := []byte(`{ + "messages":[ + {"role":"assistant","tool_calls":[{"id":"call_patch","type":"function","function":{"name":"apply_patch","arguments":"{\"input\":\"patch\"}"}}]}, + {"role":"tool","tool_call_id":"call_patch","content":"done"} + ], + "tools":[ + {"type":"function","function":{"name":"bash"}}, + {"type":"function","function":{"name":"read_bash"}}, + {"type":"function","function":{"name":"stop_bash"}}, + {"type":"function","function":{"name":"list_bash"}}, + {"type":"custom","custom":{"name":"apply_patch","description":"Edit files"` + tt.format + `}} + ], + "tool_choice":{"type":"custom","custom":{"name":"apply_patch"}} + }`) + out, errConvert := ConvertOpenAIRequestToCodex("codex-test", input, true) + if errConvert != nil { + t.Fatal(errConvert) + } + tool := ojson.GetBytesNoCopy(out, "tools.4") + if tool.Get("custom").Exists() { + t.Fatalf("unknown parameter tools[4].custom would be sent upstream: %s", tool.Raw) + } + if tool.Get("type").String() != "custom" || tool.Get("name").String() != "apply_patch" || tool.Get("description").String() != "Edit files" { + t.Fatalf("unexpected custom declaration: %s", tool.Raw) + } + if got, want := tool.Get("format").Value(), ojson.Parse(tt.want).Value(); !reflect.DeepEqual(got, want) { + t.Fatalf("format = %v, want %v", got, want) + } + if got := ojson.GetBytesNoCopy(out, "tool_choice.name").String(); got != "apply_patch" { + t.Fatalf("tool_choice.name = %q, want apply_patch", got) + } + if got := ojson.GetBytesNoCopy(out, "input.0.type").String(); got != "custom_tool_call" { + t.Fatalf("history call type = %q, want custom_tool_call", got) + } + if got := ojson.GetBytesNoCopy(out, "input.0.input").String(); got != "patch" { + t.Fatalf("history input = %q, want patch", got) + } + if got := ojson.GetBytesNoCopy(out, "input.1.type").String(); got != "custom_tool_call_output" { + t.Fatalf("history output type = %q, want custom_tool_call_output", got) + } + }) + } +} + +func TestNestedCustomToolNameShortening(t *testing.T) { + name := strings.Repeat("custom_tool_", 7) + input := []byte(`{"messages":[{"role":"assistant","tool_calls":[{"id":"call_long","type":"function","function":{"name":"` + name + `","arguments":"raw"}}]},{"role":"tool","tool_call_id":"call_long","content":"ok"}],"tools":[{"type":"custom","custom":{"name":"` + name + `"}}],"tool_choice":{"type":"custom","custom":{"name":"` + name + `"}}}`) + out, errConvert := ConvertOpenAIRequestToCodex("codex-test", input, true) + if errConvert != nil { + t.Fatal(errConvert) + } + shortName := ojson.GetBytesNoCopy(out, "tools.0.name").String() + if shortName == "" || shortName == name || len(shortName) > 64 { + t.Fatalf("unexpected shortened custom name: %q", shortName) + } + if got := ojson.GetBytesNoCopy(out, "tool_choice.name").String(); got != shortName { + t.Fatalf("tool_choice.name = %q, want %q", got, shortName) + } + if got := ojson.GetBytesNoCopy(out, "input.0.type").String(); got != "custom_tool_call" { + t.Fatalf("input.0.type = %q, want custom_tool_call", got) + } + if got := ojson.GetBytesNoCopy(out, "input.0.name").String(); got != shortName { + t.Fatalf("input.0.name = %q, want %q", got, shortName) + } + if got := buildReverseMapFromOriginalOpenAI(input)[shortName]; got != name { + t.Fatalf("reverse mapping = %q, want %q", got, name) + } +} + func TestCustomToolNameShortening(t *testing.T) { longName := "a_very_long_custom_tool_name_that_exceeds_sixty_four_characters_limit_test" if len(longName) <= 64 { diff --git a/internal/translator/codex/openai/chat-completions/codex_openai_response.go b/internal/translator/codex/openai/chat-completions/codex_openai_response.go index 5938d1fc6c8..f97e4a5bd96 100644 --- a/internal/translator/codex/openai/chat-completions/codex_openai_response.go +++ b/internal/translator/codex/openai/chat-completions/codex_openai_response.go @@ -58,13 +58,38 @@ func (m *codexChatToolMetadata) resolve(original []byte) { } } m.customPatchWinners = map[string]bool{} - for name, winner := range util.CollectResponsesToolWinners(util.ParseGJSONBytesNoCopy(original)) { + for name, winner := range collectChatToolWinners(original) { if applypatch.IsCustomTool(winner.Tool) { m.customPatchWinners[strings.Clone(name)] = true } } } +// collectChatToolWinners normalizes nested Chat custom declarations for Responses metadata. +func collectChatToolWinners(original []byte) map[string]util.ResponsesToolDescriptor { + root := util.ParseGJSONBytesNoCopy(original) + var doc *ojson.Doc + for i, tool := range root.Get("tools").Array() { + if tool.Get("type").String() != "custom" { + continue + } + definition := customToolDefinition(tool) + if definition.Raw == tool.Raw { + continue + } + if doc == nil { + doc = ojson.NewDoc(original) + } + path := "tools." + strconv.Itoa(i) + _ = doc.SetRaw(path, definition.Raw) + _ = doc.Set(path+".type", "custom") + } + if doc != nil { + return util.CollectResponsesToolWinnersDoc(doc) + } + return util.CollectResponsesToolWinners(root) +} + // restoreToolName maps a shortened upstream tool name back to the // client-declared original, matching buildReverseMapFromOriginalOpenAI. func (m *codexChatToolMetadata) restoreToolName(original []byte, name string) string { @@ -867,7 +892,7 @@ func isOriginalCustomPatch(original []byte, item ojson.Result) bool { return false } } - winner, ok := util.CollectResponsesToolWinners(util.ParseGJSONBytesNoCopy(original))[name] + winner, ok := collectChatToolWinners(original)[name] return ok && applypatch.IsCustomTool(winner.Tool) } diff --git a/internal/translator/codex/openai/chat-completions/codex_openai_response_test.go b/internal/translator/codex/openai/chat-completions/codex_openai_response_test.go index 47166e356d7..5f3616cad07 100644 --- a/internal/translator/codex/openai/chat-completions/codex_openai_response_test.go +++ b/internal/translator/codex/openai/chat-completions/codex_openai_response_test.go @@ -990,9 +990,19 @@ func TestConvertCodexResponseToOpenAI_RestoresNormalizedToolNames(t *testing.T) // Only the original winning custom patch declaration allows the JSON wrapper. func TestApplyPatchCustomChatCompletionsWrapper(t *testing.T) { - for _, toolType := range []string{"custom", "function"} { - t.Run(toolType, func(t *testing.T) { - request := []byte(`{"tools":[{"type":"` + toolType + `","name":"apply_patch"}]}`) + for _, tc := range []struct { + name string + tools string + wrap bool + }{ + {"flat custom", `[{"type":"custom","name":"apply_patch"}]`, true}, + {"nested custom", `[{"type":"custom","custom":{"name":"apply_patch"}}]`, true}, + {"function", `[{"type":"function","name":"apply_patch"}]`, false}, + {"custom before function", `[{"type":"custom","custom":{"name":"apply_patch"}},{"type":"function","function":{"name":"apply_patch"}}]`, false}, + {"function before custom", `[{"type":"function","function":{"name":"apply_patch"}},{"type":"custom","custom":{"name":"apply_patch"}}]`, false}, + } { + t.Run(tc.name, func(t *testing.T) { + request := []byte(`{"tools":` + tc.tools + `}`) var param any var arguments strings.Builder send := func(event string) { @@ -1001,38 +1011,49 @@ func TestApplyPatchCustomChatCompletionsWrapper(t *testing.T) { } } send(`{"type":"response.output_item.added","output_index":0,"item":{"type":"custom_tool_call","id":"a","call_id":"c","name":"apply_patch","input":""}}`) - for _, delta := range []string{"*** Begin Patch\n", "+中文😀 \"\n", "*** End Patch\n"} { + for _, delta := range []string{"*** Begin Patch\n", "+中文😀 \"\\path\\\n", "*** End Patch\n"} { send(`{"type":"response.custom_tool_call_input.delta","item_id":"a","output_index":0,"delta":` + string(mustJSONMarshal(t, delta)) + `}`) } - patch := "*** Begin Patch\n+中文😀 \"\n*** End Patch\n" + patch := "*** Begin Patch\n+中文😀 \"\\path\\\n*** End Patch\n" send(`{"type":"response.custom_tool_call_input.done","item_id":"a","output_index":0,"input":` + string(mustJSONMarshal(t, patch)) + `}`) send(`{"type":"response.output_item.done","output_index":0,"item":{"type":"custom_tool_call","id":"a","call_id":"c","name":"apply_patch","input":` + string(mustJSONMarshal(t, patch)) + `}}`) want := patch - if toolType == "custom" { + if tc.wrap { want = `{"input":` + string(mustJSONMarshal(t, patch)) + `}` } if arguments.String() != want { - t.Fatalf("arguments=%q want=%q", arguments.String(), want) + t.Errorf("arguments=%q want=%q", arguments.String(), want) } var nonStream any out := ConvertCodexResponseToOpenAINonStream(t.Context(), "model", request, request, []byte(`{"type":"response.completed","response":{"output":[{"type":"custom_tool_call","name":"apply_patch","call_id":"c","input":`+string(mustJSONMarshal(t, patch))+`}]}}`), &nonStream) if got := ojson.GetBytesNoCopy(out, "choices.0.message.tool_calls.0.function.arguments").String(); got != want { - t.Fatalf("nonstream arguments=%q want=%q", got, want) + t.Errorf("nonstream arguments=%q want=%q", got, want) + } + if got := isOriginalCustomPatch(request, ojson.Parse(`{"type":"custom_tool_call","name":"apply_patch"}`)); got != tc.wrap { + t.Errorf("isOriginalCustomPatch=%v want=%v", got, tc.wrap) } }) } } func TestApplyPatchCustomChatCompletionsDoneFallback(t *testing.T) { - request := []byte(`{"tools":[{"type":"custom","name":"apply_patch"}]}`) - for _, added := range []bool{false, true} { + for _, tc := range []struct { + tools string + added bool + }{ + {`[{"type":"custom","name":"apply_patch"}]`, false}, + {`[{"type":"custom","name":"apply_patch"}]`, true}, + {`[{"type":"custom","custom":{"name":"apply_patch"}}]`, false}, + {`[{"type":"custom","custom":{"name":"apply_patch"}}]`, true}, + } { + request := []byte(`{"tools":` + tc.tools + `}`) var param any - if added { + if tc.added { _ = ConvertCodexResponseToOpenAI(t.Context(), "m", request, request, []byte(`data: {"type":"response.output_item.added","output_index":0,"item":{"type":"custom_tool_call","id":"a","call_id":"c","name":"apply_patch","input":""}}`), ¶m) } out := ConvertCodexResponseToOpenAI(t.Context(), "m", request, request, []byte(`data: {"type":"response.output_item.done","output_index":0,"item":{"type":"custom_tool_call","id":"a","call_id":"c","name":"apply_patch","input":"p"}}`), ¶m) if len(out) != 1 || ojson.GetBytesNoCopy(out[0], "choices.0.delta.tool_calls.0.function.arguments").String() != `{"input":"p"}` { - t.Fatalf("fallback: %s", out) + t.Errorf("tools=%s added=%v fallback: %s", tc.tools, tc.added, out) } } } @@ -1040,14 +1061,43 @@ func TestApplyPatchCustomChatCompletionsDoneFallback(t *testing.T) { // This round trip requires the request converter to unwrap only a winning custom // patch's normalized function envelope, while explicit custom inputs remain raw. func TestApplyPatchChatCompletionNativeHistoryRoundTrip(t *testing.T) { - original := []byte(`{"messages":[{"role":"user","content":"patch"}],"tools":[{"type":"custom","name":"apply_patch"}]}`) - response := []byte(`{"type":"response.completed","response":{"output":[{"type":"custom_tool_call","call_id":"c","name":"apply_patch","input":"p"}]}}`) - out := ConvertCodexResponseToOpenAINonStream(t.Context(), "m", original, original, response, nil) - message := ojson.GetBytesNoCopy(out, "choices.0.message") - followup := []byte(`{"messages":[` + message.Raw + `,{"role":"tool","tool_call_id":"c","content":"ok"}],"tools":[{"type":"custom","name":"apply_patch"}]}`) - request, _ := ConvertOpenAIRequestToCodex("m", followup, true) - if got := ojson.GetBytesNoCopy(request, "input.0.input").String(); got != "p" { - t.Fatalf("normalized function history must restore raw patch before native Codex: got %q, request=%s", got, request) + patch := "*** Begin Patch\n+中文😀 \"quoted\" \\path\\\n*** End Patch" + wrappedPatch := `{"input":` + string(mustJSONMarshal(t, patch)) + `}` + for _, tools := range []string{ + `[{"type":"custom","name":"apply_patch"}]`, + `[{"type":"custom","custom":{"name":"apply_patch","description":"Edit files","format":{"type":"grammar","grammar":{"syntax":"lark","definition":"start: patch"}}}}]`, + } { + original := []byte(`{"messages":[{"role":"user","content":"patch"}],"tools":` + tools + `}`) + initialReq, errInitial := ConvertOpenAIRequestToCodex("m", original, true) + if errInitial != nil { + t.Fatal(errInitial) + } + if ojson.GetBytesNoCopy(initialReq, "tools.0.custom").Exists() { + t.Fatalf("tools=%s leaked nested custom declaration: %s", tools, initialReq) + } + response := []byte(`{"type":"response.completed","response":{"output":[{"type":"custom_tool_call","call_id":"c","name":"apply_patch","input":` + string(mustJSONMarshal(t, patch)) + `}]}}`) + out := ConvertCodexResponseToOpenAINonStream(t.Context(), "m", original, initialReq, response, nil) + message := ojson.GetBytesNoCopy(out, "choices.0.message") + if got := message.Get("tool_calls.0.function.arguments").String(); got != wrappedPatch { + t.Fatalf("tools=%s arguments=%q, want %q", tools, got, wrappedPatch) + } + followup := []byte(`{"messages":[{"role":"user","content":"patch"},` + message.Raw + `,{"role":"tool","tool_call_id":"c","content":"ok"},{"role":"assistant","content":"Patch applied."}],"tools":` + tools + `}`) + request, errConvert := ConvertOpenAIRequestToCodex("m", followup, true) + if errConvert != nil { + t.Fatal(errConvert) + } + if got := ojson.GetBytesNoCopy(request, "input.1.type").String(); got != "custom_tool_call" { + t.Fatalf("tools=%s history call type=%q, want custom_tool_call", tools, got) + } + if got := ojson.GetBytesNoCopy(request, "input.1.input").String(); got != patch { + t.Fatalf("normalized function history must restore raw patch before native Codex: got %q, want %q, request=%s", got, patch, request) + } + if got := ojson.GetBytesNoCopy(request, "input.2.type").String(); got != "custom_tool_call_output" { + t.Fatalf("tools=%s history output type=%q, want custom_tool_call_output", tools, got) + } + if got := ojson.GetBytesNoCopy(request, "input.3.content.0.text").String(); got != "Patch applied." { + t.Fatalf("tools=%s assistant continuation=%q, want Patch applied.", tools, got) + } } } diff --git a/internal/translator/codex/openai/chat-completions/tool_cache_perf_test.go b/internal/translator/codex/openai/chat-completions/tool_cache_perf_test.go index 96a4e78c01f..60e13ca75e6 100644 --- a/internal/translator/codex/openai/chat-completions/tool_cache_perf_test.go +++ b/internal/translator/codex/openai/chat-completions/tool_cache_perf_test.go @@ -126,6 +126,7 @@ func TestCodexChatToolMetadataEquivalence(t *testing.T) { tcRequest(3), []byte(`{"messages":[],"tools":{"type":"function","function":{"name":"apply_patch"}}}`), []byte(`{"messages":[],"tools":[{"type":"custom","name":"apply_patch","format":{"type":"grammar","syntax":"lark","definition":"start: patch"}}]}`), + []byte(`{"messages":[],"tools":[{"type":"custom","custom":{"name":"apply_patch","format":{"type":"grammar","grammar":{"syntax":"lark","definition":"start: patch"}}}}]}`), []byte(`{"messages":[],"tools":[{"type":"function","function":{"name":"apply_patch"}},{"type":"custom","name":"apply_patch","format":{"type":"grammar","syntax":"lark","definition":"start: patch"}}]}`), []byte(`{"messages":[],"input":[{"type":"additional_tools","tools":[{"type":"custom","name":"apply_patch","format":{"type":"grammar","syntax":"lark","definition":"start: patch"}}]}]}`), []byte(`{"messages":[],"tools":[{"type":"namespace","name":"editor","tools":[{"type":"custom","name":"apply_patch","format":{"type":"grammar","syntax":"lark","definition":"start: patch"}}]}]}`), diff --git a/internal/translator/codex/openai/responses/codex_openai-responses_request.go b/internal/translator/codex/openai/responses/codex_openai-responses_request.go index aee02e15a7f..a1ad3fa7324 100644 --- a/internal/translator/codex/openai/responses/codex_openai-responses_request.go +++ b/internal/translator/codex/openai/responses/codex_openai-responses_request.go @@ -74,6 +74,7 @@ func ConvertOpenAIResponsesRequestToCodex(modelName string, inputRawJSON []byte, inputResult = convertSystemRoleToDeveloperDoc(inputResult, ensureDoc) normalizeCodexBuiltinToolsDoc(rawJSON, ensureDoc) normalizeEmptyFunctionCallArgumentsDoc(inputResult, ensureDoc) + translatorcommon.SanitizeOrphanResponsesToolOutputsInputDoc(inputResult, ensureDoc) if doc != nil { return doc.Bytes(), nil diff --git a/internal/translator/codex/openai/responses/codex_openai-responses_request_test.go b/internal/translator/codex/openai/responses/codex_openai-responses_request_test.go index bede0ae168c..2769fa913d8 100644 --- a/internal/translator/codex/openai/responses/codex_openai-responses_request_test.go +++ b/internal/translator/codex/openai/responses/codex_openai-responses_request_test.go @@ -1098,3 +1098,86 @@ func TestConvertOpenAIResponsesRequestToCodexParallelToolCalls(t *testing.T) { } }) } + +func TestConvertOpenAIResponsesRequestToCodexSanitizesOrphanToolOutputs(t *testing.T) { + raw := []byte(`{ + "model":"gpt-5.4", + "input":[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}, + {"type":"function_call_output","id":"fco_1","name":"codex_apps_open_page","namespace":"agent","output":"{\"page_id\":null}"}, + {"type":"custom_tool_call_output","id":"fco_2","name":"custom_ctx","output":[{"type":"input_text","text":"part one"},{"type":"input_text","text":"part two"},{"type":"input_image","image_url":"https://example.com/page.png"}]}, + {"type":"function_call","call_id":"call_ok","name":"shell","arguments":""}, + {"type":"function_call_output","call_id":"call_ok","output":"ok"}, + {"type":"function_call_output","call_id":" ","output":"blank id"}, + {"type":"function_call_output","id":"fco_3","name":"create_thread","namespace":"codex_app","output":"task"}, + {"type":"function_call","call_id":"call_pair","name":"shell","arguments":"{}"}, + {"type":"function_call_output","output":"paired without call_id"} + ] + }`) + + out, errConvert := ConvertOpenAIResponsesRequestToCodex("gpt-5.4", raw, false) + if errConvert != nil { + t.Fatalf("ConvertOpenAIResponsesRequestToCodex returned error: %v", errConvert) + } + input := ojson.GetBytesNoCopy(out, "input").Array() + if len(input) != 9 { + t.Fatalf("input len = %d, want 9", len(input)) + } + + if got := input[1].Get("type").String(); got != "message" { + t.Fatalf("input[1].type = %q, want %q", got, "message") + } + if got := input[1].Get("role").String(); got != "user" { + t.Fatalf("input[1].role = %q, want %q", got, "user") + } + want1 := "Tool output from codex_apps_open_page:\n{\"page_id\":null}" + if got := input[1].Get("content.0.text").String(); got != want1 { + t.Fatalf("input[1].content.0.text = %q, want %q", got, want1) + } + + if got := input[2].Get("type").String(); got != "message" { + t.Fatalf("input[2].type = %q, want %q", got, "message") + } + want2 := "Tool output from custom_ctx:\npart one\npart two" + if got := input[2].Get("content.0.text").String(); got != want2 { + t.Fatalf("input[2].content.0.text = %q, want %q", got, want2) + } + if got := input[2].Get("content.1.type").String(); got != "input_image" { + t.Fatalf("input[2].content.1.type = %q, want %q", got, "input_image") + } + if got := input[2].Get("content.1.image_url").String(); got != "https://example.com/page.png" { + t.Fatalf("input[2].content.1.image_url = %q, want %q", got, "https://example.com/page.png") + } + + if got := input[3].Get("arguments").String(); got != "{}" { + t.Fatalf("input[3].arguments = %q, want %q", got, "{}") + } + if got := input[4].Get("type").String(); got != "function_call_output" { + t.Fatalf("input[4].type = %q, want %q", got, "function_call_output") + } + if got := input[4].Get("call_id").String(); got != "call_ok" { + t.Fatalf("input[4].call_id = %q, want %q", got, "call_ok") + } + + if got := input[5].Get("type").String(); got != "message" { + t.Fatalf("input[5].type = %q, want %q", got, "message") + } + want5 := "Tool output from tool:\nblank id" + if got := input[5].Get("content.0.text").String(); got != want5 { + t.Fatalf("input[5].content.0.text = %q, want %q", got, want5) + } + + if got := input[6].Get("type").String(); got != "function_call_output" { + t.Fatalf("input[6].type = %q, want %q", got, "function_call_output") + } + if got := input[6].Get("namespace").String(); got != "codex_app" { + t.Fatalf("input[6].namespace = %q, want %q", got, "codex_app") + } + + if got := input[8].Get("type").String(); got != "function_call_output" { + t.Fatalf("input[8].type = %q, want %q", got, "function_call_output") + } + if got := input[8].Get("call_id").String(); got != "call_pair" { + t.Fatalf("input[8].call_id = %q, want %q", got, "call_pair") + } +} diff --git a/internal/translator/common/responses.go b/internal/translator/common/responses.go index 51e5e4cd2a5..434d06b41b8 100644 --- a/internal/translator/common/responses.go +++ b/internal/translator/common/responses.go @@ -1,6 +1,7 @@ package common import ( + "fmt" "strings" "github.com/MoeForge/moego/ojson" @@ -131,7 +132,7 @@ func NormalizeResponsesToolCallOutputs(items []ojson.Result) []ojson.Result { expectedName := pendingCallNames[pendingID] if expectedName != "" { for outIdx, out := range outputs { - if !used[outIdx] && ExtractResponsesCallID(out) == "" { + if !used[outIdx] && ExtractResponsesCallID(out) == "" && !isCodexDelegationToolOutput(out) { outName := strings.TrimSpace(out.Get("name").String()) if outName != "" && outName == expectedName { used[outIdx] = true @@ -150,7 +151,7 @@ func NormalizeResponsesToolCallOutputs(items []ojson.Result) []ojson.Result { continue } for outIdx, out := range outputs { - if !used[outIdx] && ExtractResponsesCallID(out) == "" { + if !used[outIdx] && ExtractResponsesCallID(out) == "" && !isCodexDelegationToolOutput(out) { outName := strings.TrimSpace(out.Get("name").String()) expectedName := pendingCallNames[pendingID] if outName == "" || expectedName == "" || outName == expectedName { @@ -187,3 +188,189 @@ func NormalizeResponsesToolCallOutputs(items []ojson.Result) []ojson.Result { return normalized } + +// SanitizeOrphanResponsesToolOutputsDoc rewrites standalone function_call_output +// and custom_tool_call_output items that carry no call_id into user message +// items. Codex Desktop Companion Page context is injected as function_call_output +// items without call_id, which strict Responses upstreams reject with HTTP 400 +// ("responses function_call_output call_id is required"). +func SanitizeOrphanResponsesToolOutputsDoc(doc *ojson.Doc) bool { + if doc == nil { + return false + } + return SanitizeOrphanResponsesToolOutputsInputDoc(doc.Get("input"), func() *ojson.Doc { + return doc + }) +} + +// SanitizeOrphanResponsesToolOutputsInputDoc pairs tool outputs missing call_id +// with preceding unmatched tool calls when possible, and rewrites remaining +// standalone function_call_output and custom_tool_call_output items without +// call_id into user message items in-place on the document returned by ensureDoc. +func SanitizeOrphanResponsesToolOutputsInputDoc(inputResult ojson.Result, ensureDoc func() *ojson.Doc) bool { + if !inputResult.IsArray() || ensureDoc == nil { + return false + } + inputItems := inputResult.Array() + if len(inputItems) == 0 { + return false + } + + normalizedItems := NormalizeResponsesToolCallOutputs(inputItems) + + changed := false + pendingMissingCallIDOutputType := "" + for i := 0; i < len(inputItems); i++ { + item := normalizedItems[i] + if !item.IsObject() { + pendingMissingCallIDOutputType = "" + continue + } + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType == "function_call" || itemType == "custom_tool_call" { + if ExtractResponsesCallID(item) == "" { + pendingMissingCallIDOutputType = itemType + "_output" + } else { + pendingMissingCallIDOutputType = "" + } + continue + } + if itemType != "function_call_output" && itemType != "custom_tool_call_output" { + pendingMissingCallIDOutputType = "" + continue + } + if isCodexDelegationToolOutput(item) { + pendingMissingCallIDOutputType = "" + continue + } + if resolvedCallID := responsesToolOutputExplicitCallID(item); resolvedCallID != "" { + pendingMissingCallIDOutputType = "" + origCallID := strings.TrimSpace(inputItems[i].Get("call_id").String()) + if origCallID != resolvedCallID { + itemPath := fmt.Sprintf("input.%d.call_id", i) + if errSet := ensureDoc().Set(itemPath, resolvedCallID); errSet == nil { + changed = true + } + } + continue + } + if pendingMissingCallIDOutputType == itemType && + strings.TrimSpace(item.Get("name").String()) == "" && + strings.TrimSpace(item.Get("id").String()) == "" { + pendingMissingCallIDOutputType = "" + continue + } + pendingMissingCallIDOutputType = "" + + itemPath := fmt.Sprintf("input.%d", i) + if setOrphanResponsesToolOutputMessageDoc(ensureDoc(), itemPath, item) { + changed = true + } + } + return changed +} + +func responsesToolOutputExplicitCallID(item ojson.Result) string { + for _, key := range []string{"call_id", "tool_call_id", "callId"} { + val := item.Get(key) + if val.Type == ojson.String { + if s := strings.TrimSpace(val.String()); s != "" { + return s + } + } + } + return "" +} + +func isCodexDelegationToolOutput(item ojson.Result) bool { + if strings.TrimSpace(item.Get("type").String()) != "function_call_output" { + return false + } + if strings.TrimSpace(item.Get("namespace").String()) != "codex_app" { + return false + } + switch strings.TrimSpace(item.Get("name").String()) { + case "create_thread", "send_message_to_thread": + return true + default: + return false + } +} + +func setOrphanResponsesToolOutputMessageDoc(targetDoc *ojson.Doc, itemPath string, item ojson.Result) bool { + if targetDoc == nil { + return false + } + fullText, extraParts := buildOrphanResponsesToolOutputContent(item) + if errSet := targetDoc.SetRawBytes(itemPath, []byte(`{"type":"message","role":"user","content":[{"type":"input_text","text":""}]}`)); errSet != nil { + return false + } + if errSet := targetDoc.Set(itemPath+".content.0.text", fullText); errSet != nil { + return false + } + for idx, partRaw := range extraParts { + partPath := fmt.Sprintf("%s.content.%d", itemPath, idx+1) + if errSet := targetDoc.SetRawBytes(partPath, partRaw); errSet != nil { + return false + } + } + return true +} + +func buildOrphanResponsesToolOutputContent(item ojson.Result) (string, [][]byte) { + name := strings.TrimSpace(item.Get("name").String()) + if name == "" { + name = "tool" + } + outputText, extraParts := extractOrphanResponsesToolOutputParts(item.Get("output")) + if outputText == "" { + outputText = "(empty)" + } + return "Tool output from " + name + ":\n" + outputText, extraParts +} + +func extractOrphanResponsesToolOutputParts(output ojson.Result) (string, [][]byte) { + if !output.Exists() || output.Type == ojson.Null { + return "", nil + } + if output.Type == ojson.String { + return strings.TrimSpace(output.String()), nil + } + if output.IsArray() { + items := output.Array() + parts := make([]string, 0, len(items)) + var extraParts [][]byte + for i := 0; i < len(items); i++ { + part := items[i] + if part.Type == ojson.String { + if s := strings.TrimSpace(part.String()); s != "" { + parts = append(parts, s) + } + continue + } + if !part.IsObject() { + continue + } + matchedText := false + for _, key := range []string{"text", "input_text", "output_text"} { + val := part.Get(key) + if val.Type == ojson.String { + matchedText = true + if s := strings.TrimSpace(val.String()); s != "" { + parts = append(parts, s) + } + break + } + } + if !matchedText { + if raw := strings.TrimSpace(part.Raw); raw != "" { + extraParts = append(extraParts, []byte(raw)) + } + } + } + if len(parts) > 0 || len(extraParts) > 0 { + return strings.Join(parts, "\n"), extraParts + } + } + return strings.TrimSpace(output.Raw), nil +} diff --git a/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go b/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go index 088399b67db..24063d6bdb6 100644 --- a/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go +++ b/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go @@ -215,6 +215,10 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR if partsResult.IsArray() { partResults := partsResult.Array() + var textContent strings.Builder + var reasoningContent strings.Builder + hasTextContent := false + hasReasoningContent := false for i := 0; i < len(partResults); i++ { partResult := partResults[i] partTextResult := partResult.Get("text") @@ -236,10 +240,14 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR } hasThoughtSignature := thoughtSignatureResult.Exists() && thoughtSignatureResult.String() != "" - hasContentPayload := partTextResult.Exists() || functionCallResult.Exists() || inlineDataResult.Exists() + hasTextPayload := partTextResult.Exists() && partTextResult.String() != "" + hasContentPayload := hasTextPayload || functionCallResult.Exists() || inlineDataResult.Exists() // Skip pure thoughtSignature parts but keep any actual payload in the same part. if hasThoughtSignature && !hasContentPayload { + if p.ActiveToolCall[candidateIndex] != nil { + sawNoopFunctionCall = true + } continue } @@ -256,9 +264,11 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR setAssistantRole() // Handle text content, distinguishing between regular content and reasoning/thoughts. if partResult.Get("thought").Bool() { - _ = templateDoc.Set("choices.0.delta.reasoning_content", text) + hasReasoningContent = true + reasoningContent.WriteString(text) } else { - _ = templateDoc.Set("choices.0.delta.content", text) + hasTextContent = true + textContent.WriteString(text) } } else if functionCallResult.Exists() { // Handle function call content. @@ -386,6 +396,12 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR _ = templateDoc.SetRawBytes("choices.0.delta.images.-1", imageDoc.Bytes()) } } + if hasTextContent { + _ = templateDoc.Set("choices.0.delta.content", textContent.String()) + } + if hasReasoningContent { + _ = templateDoc.Set("choices.0.delta.reasoning_content", reasoningContent.String()) + } } upstreamFinishReason := p.UpstreamFinishReason[candidateIndex] @@ -523,6 +539,10 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina if !inlineDataResult.Exists() { inlineDataResult = partResult.Get("inline_data") } + thoughtSignatureResult := partResult.Get("thoughtSignature") + if !thoughtSignatureResult.Exists() { + thoughtSignatureResult = partResult.Get("thought_signature") + } // Speech-to-text models (gemini-3.5-transcribe) deliver the // transcript in an audioTranscription part instead of text. @@ -531,6 +551,15 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina partTextResult = audioTranscriptionResult.Get("text") } + hasThoughtSignature := thoughtSignatureResult.Exists() && thoughtSignatureResult.String() != "" + hasTextPayload := partTextResult.Exists() && partTextResult.String() != "" + hasContentPayload := hasTextPayload || functionCallResult.Exists() || inlineDataResult.Exists() + + // Skip pure thoughtSignature parts but keep any actual payload in the same part. + if hasThoughtSignature && !hasContentPayload { + continue + } + if partTextResult.Exists() { // Append text content, distinguishing between regular content and reasoning. if partResult.Get("thought").Bool() { diff --git a/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go b/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go index 9bb578da756..211f72df397 100644 --- a/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go +++ b/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go @@ -573,6 +573,84 @@ func TestConvertGeminiResponseToOpenAI_InputBufferReuse(t *testing.T) { } } +func TestConvertGeminiResponseToOpenAI_MultiPartThoughtAndTrailingSignature(t *testing.T) { + ctx := context.Background() + + t.Run("thought followed by empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"candidates":[{"content":{"parts":[{"thought":true,"text":"Analyzing the request..."},{"thought":true,"text":"","thoughtSignature":"EjQK..."}]}}]}`) + out := ConvertGeminiResponseToOpenAI(ctx, "gemini-2.5-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Analyzing the request..." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Analyzing the request...", out[0]) + } + if content := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content"); content.Type != ojson.Null { + t.Fatalf("delta.content = %s (type %v), want null; payload=%s", content.Raw, content.Type, out[0]) + } + if role := ojson.GetBytesNoCopy(out[0], "choices.0.delta.role").String(); role != "assistant" { + t.Fatalf("delta.role = %q, want assistant; payload=%s", role, out[0]) + } + }) + + t.Run("thought followed by non-thought empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"candidates":[{"content":{"parts":[{"thought":true,"text":"Analyzing the request..."},{"text":"","thoughtSignature":"EjQK..."}]}}]}`) + out := ConvertGeminiResponseToOpenAI(ctx, "gemini-2.5-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Analyzing the request..." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Analyzing the request...", out[0]) + } + if content := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content"); content.Type != ojson.Null { + t.Fatalf("delta.content = %s (type %v), want null; payload=%s", content.Raw, content.Type, out[0]) + } + }) + + t.Run("visible text followed by empty-text thoughtSignature part", func(t *testing.T) { + var param any + chunk := []byte(`{"candidates":[{"content":{"parts":[{"text":"Final answer"},{"text":"","thoughtSignature":"EjQK..."}]}}]}`) + out := ConvertGeminiResponseToOpenAI(ctx, "gemini-2.5-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content").String(); got != "Final answer" { + t.Fatalf("delta.content = %q, want %q; payload=%s", got, "Final answer", out[0]) + } + if reasoning := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content"); reasoning.Type != ojson.Null { + t.Fatalf("delta.reasoning_content = %s (type %v), want null; payload=%s", reasoning.Raw, reasoning.Type, out[0]) + } + }) + + t.Run("multiple thought parts and multiple text parts in same chunk", func(t *testing.T) { + var param any + chunk := []byte(`{"candidates":[{"content":{"parts":[{"thought":true,"text":"Step 1. "},{"thought":true,"text":"Step 2."},{"text":"Hello "},{"text":"world"}]}}]}`) + out := ConvertGeminiResponseToOpenAI(ctx, "gemini-2.5-pro", nil, nil, chunk, ¶m) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.reasoning_content").String(); got != "Step 1. Step 2." { + t.Fatalf("delta.reasoning_content = %q, want %q; payload=%s", got, "Step 1. Step 2.", out[0]) + } + if got := ojson.GetBytesNoCopy(out[0], "choices.0.delta.content").String(); got != "Hello world" { + t.Fatalf("delta.content = %q, want %q; payload=%s", got, "Hello world", out[0]) + } + }) + + t.Run("non-stream pure thoughtSignature carrier with empty text does not emit empty string", func(t *testing.T) { + raw := []byte(`{"candidates":[{"content":{"parts":[{"thought":true,"text":"Thinking..."},{"text":"","thoughtSignature":"EjQK..."},{"functionCall":{"name":"get_time","args":{"tz":"UTC"}}}]},"finishReason":"STOP"}]}`) + out := ConvertGeminiResponseToOpenAINonStream(ctx, "gemini-2.5-pro", nil, nil, raw, nil) + if got := ojson.GetBytesNoCopy(out, "choices.0.message.reasoning_content").String(); got != "Thinking..." { + t.Fatalf("message.reasoning_content = %q, want %q; payload=%s", got, "Thinking...", out) + } + if content := ojson.GetBytesNoCopy(out, "choices.0.message.content"); content.Type != ojson.Null { + t.Fatalf("message.content = %s (type %v), want null; payload=%s", content.Raw, content.Type, out) + } + }) +} + func BenchmarkConvertGeminiResponseToOpenAI_TextDelta(b *testing.B) { request := []byte(`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"` + strings.Repeat("x", 1024*1024) + `"}]}`) delta := []byte(`{"candidates":[{"content":{"parts":[{"text":"token "}]}}],"modelVersion":"gemini-2.5-flash","responseId":"resp_bench"}`) diff --git a/internal/translator/openai/claude/openai_claude_request.go b/internal/translator/openai/claude/openai_claude_request.go index 31db99aa49c..6a9c93c6cf3 100644 --- a/internal/translator/openai/claude/openai_claude_request.go +++ b/internal/translator/openai/claude/openai_claude_request.go @@ -72,44 +72,46 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream // Stream _ = outDoc.Set("stream", stream) - // Thinking: Convert Claude thinking.budget_tokens to OpenAI reasoning_effort + // Thinking: Convert Claude thinking.budget_tokens and output_config.effort to OpenAI reasoning_effort + effort := "" + if v := root.Get("output_config.effort"); v.Exists() && v.Type == ojson.String { + effort = strings.ToLower(strings.TrimSpace(v.String())) + } if thinkingConfig := root.Get("thinking"); thinkingConfig.Exists() && thinkingConfig.IsObject() { if thinkingType := thinkingConfig.Get("type"); thinkingType.Exists() { switch thinkingType.String() { case "enabled": if budgetTokens := thinkingConfig.Get("budget_tokens"); budgetTokens.Exists() { budget := int(budgetTokens.Int()) - if effort, ok := thinking.ConvertBudgetToLevel(budget); ok && effort != "" { - _ = outDoc.Set("reasoning_effort", effort) + if mappedEffort, ok := thinking.ConvertBudgetToLevel(budget); ok && mappedEffort != "" { + _ = outDoc.Set("reasoning_effort", mappedEffort) } - } else if v := root.Get("output_config.effort"); v.Exists() && v.Type == ojson.String && strings.TrimSpace(v.String()) != "" { + } else if effort != "" { // Some Claude-compatible clients pair manual thinking with output_config.effort. // Preserve that explicit level when there is no legacy token budget to map. - _ = outDoc.Set("reasoning_effort", strings.ToLower(strings.TrimSpace(v.String()))) + _ = outDoc.Set("reasoning_effort", effort) } else { // No budget_tokens or explicit effort specified; preserve the enabled-thinking default. - if effort, ok := thinking.ConvertBudgetToLevel(-1); ok && effort != "" { - _ = outDoc.Set("reasoning_effort", effort) + if mappedEffort, ok := thinking.ConvertBudgetToLevel(-1); ok && mappedEffort != "" { + _ = outDoc.Set("reasoning_effort", mappedEffort) } } case "adaptive", "auto": // Adaptive thinking can carry an explicit effort in output_config.effort (Claude 4.6). // Pass through directly; ApplyThinking handles clamping to target model's levels. - effort := "" - if v := root.Get("output_config.effort"); v.Exists() && v.Type == ojson.String { - effort = strings.ToLower(strings.TrimSpace(v.String())) - } if effort != "" { _ = outDoc.Set("reasoning_effort", effort) } else { _ = outDoc.Set("reasoning_effort", string(thinking.LevelXHigh)) } case "disabled": - if effort, ok := thinking.ConvertBudgetToLevel(0); ok && effort != "" { - _ = outDoc.Set("reasoning_effort", effort) + if mappedEffort, ok := thinking.ConvertBudgetToLevel(0); ok && mappedEffort != "" { + _ = outDoc.Set("reasoning_effort", mappedEffort) } } } + } else if !thinkingConfig.Exists() && effort != "" { + _ = outDoc.Set("reasoning_effort", effort) } // Process messages and system. diff --git a/internal/translator/openai/claude/openai_claude_request_test.go b/internal/translator/openai/claude/openai_claude_request_test.go index 3112f6ada59..3ffd694893d 100644 --- a/internal/translator/openai/claude/openai_claude_request_test.go +++ b/internal/translator/openai/claude/openai_claude_request_test.go @@ -1679,6 +1679,60 @@ func TestConvertClaudeRequestToOpenAI_EnabledThinkingEffort(t *testing.T) { } } +func TestConvertClaudeRequestToOpenAI_StandaloneOutputConfigEffort(t *testing.T) { + tests := []struct { + name string + inputJSON string + wantEffort string + }{ + { + name: "standalone high effort without thinking", + inputJSON: `{"model":"gpt-5.4","max_tokens":16,"output_config":{"effort":"high"},"messages":[{"role":"user","content":"hi"}]}`, + wantEffort: "high", + }, + { + name: "standalone effort is trimmed and lowercased", + inputJSON: `{"output_config":{"effort":" HIGH "}}`, + wantEffort: "high", + }, + { + name: "empty effort without thinking omits reasoning_effort", + inputJSON: `{"output_config":{"effort":""}}`, + wantEffort: "", + }, + { + name: "whitespace-only effort without thinking omits reasoning_effort", + inputJSON: `{"output_config":{"effort":" "}}`, + wantEffort: "", + }, + { + name: "non-string effort without thinking omits reasoning_effort", + inputJSON: `{"output_config":{"effort":123}}`, + wantEffort: "", + }, + { + name: "empty thinking object does not use standalone effort", + inputJSON: `{"thinking":{},"output_config":{"effort":"high"}}`, + wantEffort: "", + }, + { + name: "unknown thinking type does not use standalone effort", + inputJSON: `{"thinking":{"type":"unknown"},"output_config":{"effort":"high"}}`, + wantEffort: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + out, _ := ConvertClaudeRequestToOpenAI("test-model", []byte(tt.inputJSON), false) + effort := ojson.GetBytesNoCopy(out, "reasoning_effort") + if effort.String() != tt.wantEffort || (tt.wantEffort == "" && effort.Exists()) { + t.Fatalf("reasoning_effort = %s, want %q; body=%s", effort.Raw, tt.wantEffort, out) + } + }) + } +} + func TestConvertClaudeRequestToOpenAI_NormalizesBooleanSubschemas(t *testing.T) { inputJSON := []byte(`{ "model": "claude-3-opus", diff --git a/internal/translator/openai/claude/openai_claude_response.go b/internal/translator/openai/claude/openai_claude_response.go index 14e18b7a73d..6f9871c4d1b 100644 --- a/internal/translator/openai/claude/openai_claude_response.go +++ b/internal/translator/openai/claude/openai_claude_response.go @@ -440,7 +440,7 @@ func convertOpenAINonStreamingToAnthropic(rawJSON []byte) [][]byte { var contentBlocks [][]byte for _, reasoningText := range collectOpenAIObjectReasoningTexts(choice.Get("message")) { - if reasoningText == "" { + if strings.TrimSpace(reasoningText) == "" { continue } blockDoc := ojson.NewDocString(`{"type":"thinking","thinking":""}`) @@ -563,10 +563,11 @@ func emitStreamingThinkingDelta(param *ConvertOpenAIResponseToAnthropicParams, r if reasoningText == "" { return } + // Anthropic rejects whitespace-only thinking blocks, so whitespace may only extend an open one. if param.OpenToolCallIndex != -1 { if n := len(param.InterleavedContentChunks); n > 0 && param.InterleavedContentChunks[n-1].Type == "thinking" { param.InterleavedContentChunks[n-1].Text += reasoningText - } else { + } else if strings.TrimSpace(reasoningText) != "" { param.InterleavedContentChunks = append(param.InterleavedContentChunks, InterleavedContentChunk{ Type: "thinking", Text: strings.Clone(reasoningText), @@ -574,6 +575,9 @@ func emitStreamingThinkingDelta(param *ConvertOpenAIResponseToAnthropicParams, r } return } + if !param.ThinkingContentBlockStarted && strings.TrimSpace(reasoningText) == "" { + return + } stopTextContentBlock(param, results) if !param.ThinkingContentBlockStarted { if param.ThinkingContentBlockIndex == -1 { @@ -695,7 +699,7 @@ func appendOpenAINonStreamingMessageContent(contentResult ojson.Result, toolName }) } case "reasoning", "thinking": - if thinkingText := extractOpenAIContentPartThinkingText(item); thinkingText != "" { + if thinkingText := extractOpenAIContentPartThinkingText(item); thinkingBuilder.Len() > 0 || strings.TrimSpace(thinkingText) != "" { flushText() thinkingBuilder.WriteString(thinkingText) } @@ -968,7 +972,7 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina if message := choice.Get("message"); message.Exists() { for _, reasoningText := range collectOpenAIObjectReasoningTexts(message) { - if reasoningText == "" { + if strings.TrimSpace(reasoningText) == "" { continue } blockDoc := ojson.NewDocString(`{"type":"thinking","thinking":""}`) diff --git a/internal/translator/openai/claude/openai_claude_response_test.go b/internal/translator/openai/claude/openai_claude_response_test.go index 457667d16cb..fdb781a8ef7 100644 --- a/internal/translator/openai/claude/openai_claude_response_test.go +++ b/internal/translator/openai/claude/openai_claude_response_test.go @@ -1798,3 +1798,73 @@ func TestNonStream_ContentPartsArrayThinkingAndBlockOrder(t *testing.T) { t.Fatalf("ConvertOpenAIResponseToClaudeNonStream block order = %s, want [thinking, text]", string(outTopLevel)) } } + +func TestStreaming_WhitespaceReasoningNeverOpensThinkingBlock(t *testing.T) { + toolCall := `{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"Read","arguments":"{}"}}]}` + tests := []struct { + name string + deltas []string + want []string + }{ + {"whitespace only then text", []string{`{"reasoning_content":" "}`, `{"content":"Hi"}`}, []string{"text:Hi"}}, + {"whitespace inside a thought", []string{`{"reasoning_content":"a"}`, `{"reasoning_content":"\n\n"}`, `{"reasoning_content":"b"}`}, []string{"thinking:a\n\nb"}}, + {"whitespace after text closed the thought", []string{`{"reasoning_content":"a"}`, `{"content":"x"}`, `{"reasoning_content":" "}`, `{"content":"y"}`}, []string{"thinking:a", "text:xy"}}, + {"whitespace only while a tool call is open", []string{toolCall, `{"reasoning_content":" "}`}, []string{"tool_use:"}}, + {"whitespace inside a thought while a tool call is open", []string{toolCall, `{"reasoning_content":"a"}`, `{"reasoning_content":" "}`, `{"reasoning_content":"b"}`}, []string{"tool_use:", "thinking:a b"}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var chunks []string + for _, delta := range tt.deltas { + chunks = append(chunks, `{"id":"c1","model":"m","choices":[{"index":0,"delta":`+delta+`,"finish_reason":null}]}`) + } + chunks = append(chunks, `{"id":"c1","model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`) + events := runStream(t, streamReq, chunks...) + assertSequentialContentBlocks(t, events) + + var got []string + for _, e := range events { + switch e.Type { + case "content_block_start": + got = append(got, ojson.Get(e.Payload, "content_block.type").String()+":") + case "content_block_delta": + got[len(got)-1] += ojson.Get(e.Payload, "delta.thinking").String() + ojson.Get(e.Payload, "delta.text").String() + } + } + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", tt.want) { + t.Fatalf("blocks = %q, want %q", got, tt.want) + } + }) + } +} + +func TestNonStream_WhitespaceReasoningNeverOpensThinkingBlock(t *testing.T) { + tests := []struct { + name string + message string + want []string + }{ + {"whitespace only reasoning_content", `{"role":"assistant","content":"Hi","reasoning_content":" "}`, []string{"text:Hi"}}, + {"whitespace only thinking part", `{"role":"assistant","content":[{"type":"thinking","thinking":" "},{"type":"text","text":"Hi"}]}`, []string{"text:Hi"}}, + {"whitespace inside thinking parts", `{"role":"assistant","content":[{"type":"thinking","thinking":"a"},{"type":"thinking","thinking":" "},{"type":"thinking","thinking":"b"},{"type":"text","text":"Hi"}]}`, []string{"thinking:a b", "text:Hi"}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp := []byte(`{"id":"c1","model":"m","choices":[{"index":0,"message":` + tt.message + `,"finish_reason":"stop"}]}`) + var paramAny any + outputs := [][]byte{ + ConvertOpenAIResponseToClaudeNonStream(context.Background(), "", nil, nil, resp, nil), + ConvertOpenAIResponseToClaude(context.Background(), "", []byte(`{"stream":false}`), nil, append([]byte("data: "), resp...), ¶mAny)[0], + } + for i, out := range outputs { + var got []string + for _, block := range ojson.GetBytesNoCopy(out, "content").Array() { + got = append(got, block.Get("type").String()+":"+block.Get("thinking").String()+block.Get("text").String()) + } + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", tt.want) { + t.Fatalf("outputs[%d] blocks = %q, want %q (out=%s)", i, got, tt.want, out) + } + } + }) + } +} diff --git a/internal/util/claude_model.go b/internal/util/claude_model.go index 656cc5a3ce8..99e9601d616 100644 --- a/internal/util/claude_model.go +++ b/internal/util/claude_model.go @@ -13,3 +13,42 @@ func IsClaudeThinkingModel(model string) bool { lower := strings.ToLower(model) return strings.Contains(lower, "claude") && strings.Contains(lower, "thinking") } + +// claudeLegacySystemReminderModels lists the official Anthropic model IDs and +// aliases that reject a mid-conversation role=system message. Entries mirror the +// "claude" provider in internal/registry/models/models.json plus Anthropic's own +// bare and "-latest" aliases. Other providers' synthetic IDs do not belong here. +var claudeLegacySystemReminderModels = map[string]struct{}{ + "claude-3-5-haiku-20241022": {}, + "claude-3-5-haiku-latest": {}, + "claude-3-7-sonnet-20250219": {}, + "claude-3-7-sonnet-latest": {}, + "claude-haiku-4-5": {}, + "claude-haiku-4-5-20251001": {}, + "claude-opus-4": {}, + "claude-opus-4-20250514": {}, + "claude-opus-4-1": {}, + "claude-opus-4-1-20250805": {}, + "claude-opus-4-5": {}, + "claude-opus-4-5-20251101": {}, + "claude-opus-4-6": {}, + "claude-opus-4-7": {}, + "claude-sonnet-4": {}, + "claude-sonnet-4-20250514": {}, + "claude-sonnet-4-5": {}, + "claude-sonnet-4-5-20250929": {}, + "claude-sonnet-4-6": {}, +} + +// IsClaudeLegacySystemReminderModel reports whether model is a known Claude +// model that rejects a mid-conversation role=system message. Matching ignores +// case, surrounding whitespace and any provider prefix. Unknown and future IDs +// are not listed, so callers treat them as accepting the turn. +func IsClaudeLegacySystemReminderModel(model string) bool { + normalized := strings.ToLower(strings.TrimSpace(model)) + if slash := strings.LastIndexByte(normalized, '/'); slash >= 0 { + normalized = normalized[slash+1:] + } + _, legacy := claudeLegacySystemReminderModels[normalized] + return legacy +} diff --git a/internal/util/claude_model_test.go b/internal/util/claude_model_test.go index 7bcd96e4f65..39a8de8f30b 100644 --- a/internal/util/claude_model_test.go +++ b/internal/util/claude_model_test.go @@ -62,3 +62,28 @@ func TestIsClaudeModel(t *testing.T) { }) } } + +func TestIsClaudeLegacySystemReminderModel(t *testing.T) { + tests := []struct { + name string + model string + expected bool + }{ + {"listed model", "claude-sonnet-4-5", true}, + {"listed snapshot", "claude-haiku-4-5-20251001", true}, + {"mixed case and whitespace", " Claude-Opus-4-6 ", true}, + {"provider prefix", "anthropic/claude-sonnet-4-6", true}, + {"supported model", "claude-opus-5-5", false}, + {"unknown future model", "claude-sonnet-9", false}, + {"unlisted suffix", "claude-sonnet-4-5-thinking", false}, + {"empty string", "", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsClaudeLegacySystemReminderModel(tt.model); got != tt.expected { + t.Errorf("IsClaudeLegacySystemReminderModel(%q) = %v, expected %v", tt.model, got, tt.expected) + } + }) + } +} diff --git a/internal/util/nocopy_invariant_test.go b/internal/util/nocopy_invariant_test.go index e4c64f09ce9..368e68c6e92 100644 --- a/internal/util/nocopy_invariant_test.go +++ b/internal/util/nocopy_invariant_test.go @@ -104,6 +104,7 @@ type reviewedInPlaceByteWrite struct { } var reviewedInPlaceByteWrites = map[string]reviewedInPlaceByteWrite{ + "internal/kvstore/bolt_store.go": {1, "copies into a freshly allocated envelope before publication; no no-copy JSON result can alias the destination"}, "internal/translator/common/apply_patch_responses.go": {1, "shifts [][]byte item references in a private output slice; no byte of any JSON payload is rewritten"}, "internal/translator/common/vertex_partial_args.go": {1, "shifts []string object keys to insert a key; no byte of any payload is rewritten"}, "internal/runtime/executor/claude_signing.go": {2, "writes CCH digits into bytes.Clone(body); the caller's body is never touched"}, @@ -113,7 +114,6 @@ var reviewedInPlaceByteWrites = map[string]reviewedInPlaceByteWrite{ "internal/runtime/executor/helps/claude_credential_identity.go": {1, "copies key bytes into a local stack-allocated quote array for json.Unmarshal"}, "internal/runtime/executor/helps/claude_mcp_alias.go": {1, "copies an HMAC sum into a local fixed-size digest array"}, "internal/client/codex/live/tcp_proxy.go": {1, "copies header and payload into a freshly allocated frame"}, - "internal/kvstore/bolt_store.go": {1, "copies a stored value into a freshly allocated envelope buffer; the caller's value is never touched"}, "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"}, } diff --git a/internal/watcher/auth_file_snapshot.go b/internal/watcher/auth_file_snapshot.go index 4162e6060db..10c56971dae 100644 --- a/internal/watcher/auth_file_snapshot.go +++ b/internal/watcher/auth_file_snapshot.go @@ -18,7 +18,7 @@ func (w *Watcher) SnapshotAuthFile(path string) ([]*coreauth.Auth, error) { } w.clientsMutex.RLock() cfg := w.config - authDir := w.authDir + authDir := w.currentAuthDirLocked() parser := w.pluginAuthParser w.clientsMutex.RUnlock() if cfg == nil || authDir == "" { diff --git a/internal/watcher/clients.go b/internal/watcher/clients.go index 98bf690dfde..e8ddbccebe8 100644 --- a/internal/watcher/clients.go +++ b/internal/watcher/clients.go @@ -94,7 +94,7 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string parser := w.pluginAuthParser w.clientsMutex.RUnlock() - if resolvedAuthDir, errResolveAuthDir := util.ResolveAuthDir(cfg.AuthDir); errResolveAuthDir != nil { + if resolvedAuthDir, errResolveAuthDir := w.resolveAuthDirForConfig(cfg); errResolveAuthDir != nil { log.Errorf("failed to resolve auth directory for hash cache: %v", errResolveAuthDir) } else if resolvedAuthDir != "" { entries, errReadDir := os.ReadDir(resolvedAuthDir) @@ -254,7 +254,7 @@ func (w *Watcher) rereadDiscardedFileSnapshot(auth *coreauth.Auth) { w.clientsMutex.Lock() existing := w.currentAuths[auth.ID] - if existing == nil || !existing.SameFileSnapshot(auth) { + if existing == nil || !existing.SameFileSnapshot(auth) || !w.isPathInCurrentAuthDirLocked(path) { w.clientsMutex.Unlock() return } @@ -312,6 +312,9 @@ func (w *Watcher) invalidateAuthFileHashLocked(path string) string { } func (w *Watcher) addOrUpdateClientLocked(path string) { + if !w.isPathInCurrentAuthDir(path) { + return + } w.observeAuthFile(path) data, errRead := os.ReadFile(path) if errRead != nil { @@ -341,12 +344,13 @@ func (w *Watcher) addOrUpdateClientLocked(path string) { w.clientsMutex.Unlock() return } + if !w.isPathInCurrentAuthDirLocked(normalized) { + w.clientsMutex.Unlock() + return + } cfg := w.config - authDir := w.authDir + authDir := w.currentAuthDirLocked() parser := w.pluginAuthParser - if w.fileAuthsByPath == nil { - w.fileAuthsByPath = make(map[string]map[string]*coreauth.Auth) - } if prev, ok := w.lastAuthHashes[normalized]; ok && prev == curHash { log.Debugf("auth file unchanged (hash match), skipping reload: %s", filepath.Base(path)) w.clientsMutex.Unlock() @@ -360,23 +364,6 @@ func (w *Watcher) addOrUpdateClientLocked(path string) { oldAuth = cached.Clone() } } - - // Update caches - if w.lastAuthHashes == nil { - w.lastAuthHashes = make(map[string]string) - } - w.lastAuthHashes[normalized] = curHash - if cacheAuthContents { - if w.lastAuthContents == nil { - w.lastAuthContents = make(map[string]*coreauth.Auth) - } - w.lastAuthContents[normalized] = &newAuth - } - - oldByID := make(map[string]*coreauth.Auth, len(w.fileAuthsByPath[normalized])) - for id, a := range w.fileAuthsByPath[normalized] { - oldByID[id] = a - } w.clientsMutex.Unlock() // Compute and log field changes @@ -403,6 +390,27 @@ func (w *Watcher) addOrUpdateClientLocked(path string) { } newByID := authSliceToMap(generated) w.clientsMutex.Lock() + if w.config == nil || !w.isPathInCurrentAuthDirLocked(normalized) || (authDir != "" && !w.sameNormalizedAuthPath(authDir, w.currentAuthDirLocked())) { + w.clientsMutex.Unlock() + return + } + if w.fileAuthsByPath == nil { + w.fileAuthsByPath = make(map[string]map[string]*coreauth.Auth) + } + if w.lastAuthHashes == nil { + w.lastAuthHashes = make(map[string]string) + } + w.lastAuthHashes[normalized] = curHash + if cacheAuthContents { + if w.lastAuthContents == nil { + w.lastAuthContents = make(map[string]*coreauth.Auth) + } + w.lastAuthContents[normalized] = &newAuth + } + oldByID := make(map[string]*coreauth.Auth, len(w.fileAuthsByPath[normalized])) + for id, a := range w.fileAuthsByPath[normalized] { + oldByID[id] = a + } if len(newByID) > 0 { w.fileAuthsByPath[normalized] = authIDSet(newByID) } else { @@ -426,9 +434,16 @@ func (w *Watcher) removeClient(path string) { } func (w *Watcher) removeClientLocked(path string) { + if !w.isPathInCurrentAuthDir(path) { + return + } w.observeAuthFile(path) normalized := w.normalizeAuthPath(path) w.clientsMutex.Lock() + if !w.isPathInCurrentAuthDirLocked(normalized) { + w.clientsMutex.Unlock() + return + } oldByID := make(map[string]*coreauth.Auth, len(w.fileAuthsByPath[normalized])) for id, a := range w.fileAuthsByPath[normalized] { oldByID[id] = a @@ -504,11 +519,22 @@ func authIDSet(auths map[string]*coreauth.Auth) map[string]*coreauth.Auth { return set } +func (w *Watcher) resolveAuthDirForConfig(cfg *config.Config) (string, error) { + raw := "" + if cfg != nil { + raw = strings.TrimSpace(cfg.AuthDir) + } + if raw == "" { + raw = w.currentAuthDir() + } + return util.ResolveAuthDir(raw) +} + func (w *Watcher) loadFileClients(cfg *config.Config) int { authFileCount := 0 successfulAuthCount := 0 - authDir, errResolveAuthDir := util.ResolveAuthDir(cfg.AuthDir) + authDir, errResolveAuthDir := w.resolveAuthDirForConfig(cfg) if errResolveAuthDir != nil { log.Errorf("failed to resolve auth directory: %v", errResolveAuthDir) return 0 diff --git a/internal/watcher/config_reload.go b/internal/watcher/config_reload.go index b9d5aabbfe0..5738909d67b 100644 --- a/internal/watcher/config_reload.go +++ b/internal/watcher/config_reload.go @@ -7,6 +7,7 @@ import ( "encoding/hex" "os" "reflect" + "strings" "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" @@ -102,22 +103,32 @@ func (w *Watcher) reloadConfig() bool { return false } + resolvedNewAuthDir := "" if w.mirroredAuthDir != "" { newConfig.AuthDir = w.mirroredAuthDir + resolvedNewAuthDir = w.mirroredAuthDir } else { if resolvedAuthDir, errResolveAuthDir := util.ResolveAuthDir(newConfig.AuthDir); errResolveAuthDir != nil { log.Errorf("failed to resolve auth directory from config: %v", errResolveAuthDir) } else { newConfig.AuthDir = resolvedAuthDir + resolvedNewAuthDir = resolvedAuthDir } } + w.authRescanMu.Lock() w.clientsMutex.Lock() var oldConfig *config.Config _ = yaml.Unmarshal(w.oldConfigYaml, &oldConfig) + prevAuthDir := w.authDir w.oldConfigYaml, _ = yaml.Marshal(newConfig) w.config = newConfig + if strings.TrimSpace(resolvedNewAuthDir) != "" { + w.authDir = resolvedNewAuthDir + } + curAuthDir := w.authDir w.clientsMutex.Unlock() + w.authRescanMu.Unlock() var affectedOAuthProviders []string if oldConfig != nil { @@ -141,7 +152,10 @@ func (w *Watcher) reloadConfig() bool { } } - authDirChanged := oldConfig == nil || oldConfig.AuthDir != newConfig.AuthDir + authDirChanged := oldConfig == nil || oldConfig.AuthDir != newConfig.AuthDir || prevAuthDir != curAuthDir + if authDirChanged { + w.syncAuthDirWatch() + } retryConfigChanged := oldConfig != nil && (oldConfig.RequestRetry != newConfig.RequestRetry || oldConfig.MaxRetryInterval != newConfig.MaxRetryInterval || oldConfig.MaxRetryCredentials != newConfig.MaxRetryCredentials) forceAuthRefresh := oldConfig != nil && (oldConfig.ForceModelPrefix != newConfig.ForceModelPrefix || !reflect.DeepEqual(oldConfig.OAuthModelAlias, newConfig.OAuthModelAlias) || !reflect.DeepEqual(oldConfig.OAuthSettings, newConfig.OAuthSettings) || retryConfigChanged) diff --git a/internal/watcher/dispatcher.go b/internal/watcher/dispatcher.go index 1e165b452a0..026a917aa04 100644 --- a/internal/watcher/dispatcher.go +++ b/internal/watcher/dispatcher.go @@ -147,17 +147,35 @@ func (w *Watcher) dispatchPersistedAuthUpdateWithRevision(update *AuthUpdate) (b } func (w *Watcher) refreshAuthState(force bool) { - w.clientsMutex.Lock() - w.activeAuthScans++ - cfg := w.config - authDir := w.authDir - parser := w.pluginAuthParser - previous := maps.Clone(w.authRevisions) - previousFiles := maps.Clone(w.fileObservations) - w.clientsMutex.Unlock() - auths, incompleteFiles := snapshotCoreAuthsFunc(cfg, authDir, parser) - w.clientsMutex.Lock() - w.activeAuthScans-- + var ( + cfg *config.Config + authDir string + parser synthesizer.PluginAuthParser + previous map[string]uint64 + previousFiles map[string]uint64 + auths []*coreauth.Auth + incompleteFiles map[string]struct{} + ) + for { + w.clientsMutex.Lock() + w.activeAuthScans++ + cfg = w.config + authDir = w.currentAuthDirLocked() + parser = w.pluginAuthParser + previous = maps.Clone(w.authRevisions) + previousFiles = maps.Clone(w.fileObservations) + w.clientsMutex.Unlock() + + auths, incompleteFiles = snapshotCoreAuthsFunc(cfg, authDir, parser) + + w.clientsMutex.Lock() + w.activeAuthScans-- + currentAuthDir := w.currentAuthDirLocked() + if authDir == currentAuthDir || w.sameNormalizedAuthPath(authDir, currentAuthDir) { + break + } + w.clientsMutex.Unlock() + } var incompletePaths map[string]struct{} if len(incompleteFiles) > 0 { incompletePaths = make(map[string]struct{}, len(incompleteFiles)*2) @@ -267,6 +285,16 @@ func (w *Watcher) refreshAuthState(force bool) { source := auth.Attributes[coreauth.AttributeSource] return source != "" && !strings.HasPrefix(source, "config:") } + belongsToCurrentAuthDir := func(auth *coreauth.Auth) bool { + if !isFileBackedAuth(auth) || authDir == "" { + return true + } + rawPath := fileBackedAuthPath(auth) + if rawPath == "" || filepath.Dir(w.normalizeAuthPath(rawPath)) == "." { + return true + } + return w.isPathInAuthDir(rawPath, authDir) + } for index, auth := range auths { if auth != nil && (changedDuringScan(auth) || incompleteDuringScan(auth)) { auths[index] = nil @@ -282,7 +310,7 @@ func (w *Watcher) refreshAuthState(force bool) { } } for id, auth := range providerReloadAuths { - if auth == nil || previous[id] != w.authRevisions[id] { + if auth == nil || previous[id] != w.authRevisions[id] || !belongsToCurrentAuthDir(auth) { continue } if _, exists := w.currentAuths[id]; exists { @@ -306,7 +334,7 @@ func (w *Watcher) refreshAuthState(force bool) { var rereadPaths []string rereadSeen := make(map[string]struct{}) for _, auth := range w.currentAuths { - if auth == nil { + if auth == nil || !belongsToCurrentAuthDir(auth) { continue } if !changedDuringScan(auth) && !incompleteDuringScan(auth) { @@ -346,9 +374,9 @@ func (w *Watcher) refreshAuthState(force bool) { } } -func (w *Watcher) indexFileBackedAuthLocked(auth *coreauth.Auth) { - if auth == nil || auth.ID == "" || len(auth.Attributes) == 0 { - return +func fileBackedAuthPath(auth *coreauth.Auth) string { + if auth == nil || len(auth.Attributes) == 0 { + return "" } path := auth.Attributes[coreauth.AttributePath] if path == "" { @@ -357,7 +385,14 @@ func (w *Watcher) indexFileBackedAuthLocked(auth *coreauth.Auth) { path = source } } - normalized := w.normalizeAuthPath(path) + return path +} + +func (w *Watcher) indexFileBackedAuthLocked(auth *coreauth.Auth) { + if auth == nil || auth.ID == "" { + return + } + normalized := w.normalizeAuthPath(fileBackedAuthPath(auth)) if normalized == "" { return } @@ -414,7 +449,12 @@ func (w *Watcher) prepareAuthUpdatesLocked(auths []*coreauth.Auth, force bool) [ } for path, pathAuths := range w.fileAuthsByPath { for id := range pathAuths { - if _, ok := newState[id]; !ok { + auth, ok := newState[id] + if !ok { + delete(pathAuths, id) + continue + } + if authPath := fileBackedAuthPath(auth); authPath != "" && !w.sameNormalizedAuthPath(path, authPath) { delete(pathAuths, id) } } diff --git a/internal/watcher/events.go b/internal/watcher/events.go index f6a55d04fc5..5b025275dd4 100644 --- a/internal/watcher/events.go +++ b/internal/watcher/events.go @@ -33,11 +33,17 @@ func (w *Watcher) start(ctx context.Context) error { } log.Debugf("watching config file: %s", w.configPath) - if errAddAuthDir := w.watcher.Add(w.authDir); errAddAuthDir != nil { - log.Errorf("failed to watch auth directory %s: %v", w.authDir, errAddAuthDir) + w.watchMu.Lock() + authDir := w.currentAuthDir() + if errAddAuthDir := w.watcher.Add(authDir); errAddAuthDir != nil { + w.watchMu.Unlock() + log.Errorf("failed to watch auth directory %s: %v", authDir, errAddAuthDir) return errAddAuthDir } - log.Debugf("watching auth directory: %s", w.authDir) + w.started = true + w.watchedAuthDir = authDir + w.watchMu.Unlock() + log.Debugf("watching auth directory: %s", authDir) go w.processEvents(ctx) @@ -45,6 +51,42 @@ func (w *Watcher) start(ctx context.Context) error { return nil } +func (w *Watcher) syncAuthDirWatch() { + if w == nil || w.watcher == nil || w.stopped.Load() { + return + } + w.watchMu.Lock() + defer w.watchMu.Unlock() + if !w.started || w.stopped.Load() { + return + } + target := strings.TrimSpace(w.currentAuthDir()) + oldWatched := strings.TrimSpace(w.watchedAuthDir) + if oldWatched == target { + return + } + if oldWatched != "" && !w.sameNormalizedAuthPath(oldWatched, w.configPath) && !w.sameNormalizedAuthPath(oldWatched, target) { + if errRemove := w.watcher.Remove(oldWatched); errRemove != nil { + log.Debugf("failed to remove previous auth directory watch %s: %v", oldWatched, errRemove) + } + w.watchedAuthDir = "" + } + if target == "" { + w.watchedAuthDir = "" + return + } + if w.watchedAuthDir != "" && w.sameNormalizedAuthPath(w.watchedAuthDir, target) { + w.watchedAuthDir = target + return + } + if errAdd := w.watcher.Add(target); errAdd != nil { + log.Errorf("failed to watch auth directory %s: %v", target, errAdd) + return + } + w.watchedAuthDir = target + log.Debugf("watching auth directory: %s", target) +} + func (w *Watcher) processEvents(ctx context.Context) { for { select { @@ -69,7 +111,7 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { configOps := fsnotify.Write | fsnotify.Create | fsnotify.Rename normalizedName := w.normalizeAuthPath(event.Name) normalizedConfigPath := w.normalizeAuthPath(w.configPath) - normalizedAuthDir := w.normalizeAuthPath(w.authDir) + normalizedAuthDir := w.normalizeAuthPath(w.currentAuthDir()) isConfigEvent := normalizedName == normalizedConfigPath && event.Op&configOps != 0 authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename isAuthJSON := filepath.Dir(normalizedName) == normalizedAuthDir && strings.HasSuffix(normalizedName, ".json") && event.Op&authOps != 0 @@ -91,6 +133,9 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { // Handle auth directory changes incrementally (.json only) w.authRescanMu.Lock() defer w.authRescanMu.Unlock() + if !w.isPathInAuthDir(normalizedName, w.currentAuthDir()) { + return + } w.observeAuthFile(event.Name) if event.Op&(fsnotify.Remove|fsnotify.Rename) != 0 { @@ -101,6 +146,9 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { // Atomic replace on some platforms may surface as Rename (or Remove) before the new file is ready. // Wait briefly; if the path exists again, treat as an update instead of removal. time.Sleep(replaceCheckDelay) + if !w.isPathInAuthDir(normalizedName, w.currentAuthDir()) { + return + } statSucceeded := false if _, statErr := os.Stat(event.Name); statErr == nil { statSucceeded = true @@ -110,6 +158,9 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { // before treating a known auth file as permanently removed. for attempt := 0; attempt < 3; attempt++ { time.Sleep(25 * time.Millisecond) + if !w.isPathInAuthDir(normalizedName, w.currentAuthDir()) { + return + } if _, retryErr := os.Stat(event.Name); retryErr == nil { statSucceeded = true break @@ -152,6 +203,9 @@ func (w *Watcher) observeAuthFile(path string) { } w.clientsMutex.Lock() defer w.clientsMutex.Unlock() + if !w.isPathInCurrentAuthDirLocked(normalized) { + return + } if w.fileObservations == nil { w.fileObservations = make(map[string]uint64) } diff --git a/internal/watcher/synthesizer/file.go b/internal/watcher/synthesizer/file.go index 1a40f508b70..a95b8b1d4f3 100644 --- a/internal/watcher/synthesizer/file.go +++ b/internal/watcher/synthesizer/file.go @@ -126,6 +126,15 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) ([ if len(auths) == 0 { return nil, nil } + filePortalHint := "" + for _, key := range []string{"oauth_provider", "portal"} { + if raw, ok := metadata[key].(string); ok { + if hint := strings.ToLower(strings.TrimSpace(raw)); hint != "" { + filePortalHint = hint + break + } + } + } perAccountModelAliases := extractOAuthModelAliasesFromMetadata(metadata) disabled, _ := metadata["disabled"].(bool) for index, auth := range auths { @@ -133,6 +142,29 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) ([ continue } coreauth.NormalizeCredentialMetadata(auth.Metadata) + if filePortalHint != "" { + hasPortalHint := false + for _, key := range []string{"oauth_provider", "portal"} { + raw, _ := auth.Metadata[key].(string) + if strings.TrimSpace(auth.Attributes[key]) != "" || strings.TrimSpace(raw) != "" { + hasPortalHint = true + break + } + } + if !hasPortalHint { + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Metadata["oauth_provider"] = filePortalHint + auth.Attributes["oauth_provider"] = filePortalHint + if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(auth.Metadata) + } + } + } if len(auths) > 1 { coreauth.MarkPluginVirtualAuth(auth, fullPath, index) } diff --git a/internal/watcher/synthesizer/file_oauth_provider_test.go b/internal/watcher/synthesizer/file_oauth_provider_test.go new file mode 100644 index 00000000000..8d5b95b2ae2 --- /dev/null +++ b/internal/watcher/synthesizer/file_oauth_provider_test.go @@ -0,0 +1,88 @@ +package synthesizer + +import ( + "context" + "path/filepath" + "testing" + "time" + + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" +) + +func TestSynthesizeAuthFilePluginOAuthProviderHint(t *testing.T) { + for _, tc := range []struct { + name string + raw string + attributes map[string]string + metadata map[string]any + want string + inherit bool + }{ + {name: "primary", raw: `{"type":"vendor","oauth_provider":" Vendor "}`, want: "vendor", inherit: true}, + {name: "oauth-before-portal", raw: `{"type":"vendor","oauth_provider":" Vendor ","portal":"vendor-cn"}`, want: "vendor", inherit: true}, + {name: "blank-oauth-fallback", raw: `{"type":"vendor","oauth_provider":" ","portal":" Vendor-CN "}`, want: "vendor-cn", inherit: true}, + {name: "non-string-oauth-fallback", raw: `{"type":"vendor","oauth_provider":42,"portal":" Vendor "}`, want: "vendor", inherit: true}, + {name: "invalid-parsed-hints", raw: `{"type":"vendor","oauth_provider":"vendor"}`, attributes: map[string]string{"oauth_provider": " ", "portal": " "}, metadata: map[string]any{"oauth_provider": true, "portal": 42}, want: "vendor", inherit: true}, + {name: "attributes-oauth-preserved", raw: `{"type":"vendor","oauth_provider":"vendor"}`, attributes: map[string]string{"oauth_provider": "vendor-cn"}, want: "vendor-cn"}, + {name: "metadata-oauth-preserved", raw: `{"type":"vendor","oauth_provider":"vendor"}`, metadata: map[string]any{"oauth_provider": "vendor-cn"}, want: "vendor-cn"}, + {name: "attributes-portal-preserved", raw: `{"type":"vendor","oauth_provider":"vendor"}`, attributes: map[string]string{"portal": "vendor-cn"}, want: "vendor-cn"}, + {name: "metadata-portal-preserved", raw: `{"type":"vendor","oauth_provider":"vendor"}`, metadata: map[string]any{"portal": "vendor-cn"}, want: "vendor-cn"}, + {name: "no-hint", raw: `{"type":"vendor"}`}, + {name: "invalid-file-hints", raw: `{"type":"vendor","oauth_provider":true,"portal":42}`}, + } { + t.Run(tc.name, func(t *testing.T) { + storage := &oauthProviderMetadataStorage{} + parser := multiAuthParserFunc(func(context.Context, pluginapi.AuthParseRequest) ([]*coreauth.Auth, bool, error) { + return []*coreauth.Auth{ + {Provider: "vendor", Metadata: tc.metadata, Attributes: tc.attributes, Storage: storage}, + {Provider: "vendor"}, + }, true, nil + }) + authDir := t.TempDir() + auths, errSynthesize := SynthesizeAuthFile(&SynthesisContext{ + AuthDir: authDir, PluginAuthParser: parser, Now: time.Now(), StrictFileParsing: true, + }, filepath.Join(authDir, "vendor.json"), []byte(tc.raw)) + if errSynthesize != nil || len(auths) != 2 { + t.Fatalf("SynthesizeAuthFile() count=%d error=%v", len(auths), errSynthesize) + } + auth := auths[0] + hint := "" + for _, key := range []string{"oauth_provider", "portal"} { + if hint = auth.Attributes[key]; hint != "" { + break + } + if hint, _ = auth.Metadata[key].(string); hint != "" { + break + } + } + if hint != tc.want { + t.Fatalf("synthesized hint = %q, want %q", hint, tc.want) + } + if tc.inherit { + for index, parsed := range auths { + if parsed.Metadata["oauth_provider"] != tc.want || parsed.Attributes["oauth_provider"] != tc.want { + t.Errorf("auth %d did not inherit hint: metadata=%#v attributes=%#v", index, parsed.Metadata, parsed.Attributes) + } + } + if storage.metadata["oauth_provider"] != tc.want { + t.Fatalf("storage metadata=%#v, want oauth_provider=%q", storage.metadata, tc.want) + } + } else if storage.metadata != nil { + t.Fatalf("existing hint unexpectedly changed storage metadata: %#v", storage.metadata) + } + }) + } +} + +type oauthProviderMetadataStorage struct { + metadata map[string]any +} + +func (s *oauthProviderMetadataStorage) SetMetadata(metadata map[string]any) { + s.metadata = metadata +} + +func (s *oauthProviderMetadataStorage) SaveTokenToFile(string) error { + return nil +} diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 60a5f68f242..bca9c0a6f95 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -4,6 +4,7 @@ package watcher import ( "context" + "path/filepath" "strings" "sync" "sync/atomic" @@ -11,6 +12,7 @@ import ( "github.com/fsnotify/fsnotify" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" "github.com/router-for-me/CLIProxyAPI/v8/internal/watcher/synthesizer" "gopkg.in/yaml.v3" @@ -33,6 +35,9 @@ type authDirProvider interface { type Watcher struct { configPath string authDir string + watchMu sync.Mutex + started bool + watchedAuthDir string config *config.Config clientsMutex sync.RWMutex authRescanMu sync.Mutex @@ -169,9 +174,76 @@ func (w *Watcher) Stop() error { func (w *Watcher) SetConfig(cfg *config.Config) { cloned := cfg.CloneForRuntime() w.clientsMutex.Lock() - defer w.clientsMutex.Unlock() + if cloned != nil { + if w.mirroredAuthDir != "" { + cloned.AuthDir = w.mirroredAuthDir + w.authDir = w.mirroredAuthDir + } else if trimmedAuthDir := strings.TrimSpace(cloned.AuthDir); trimmedAuthDir != "" { + if resolvedAuthDir, errResolveAuthDir := util.ResolveAuthDir(trimmedAuthDir); errResolveAuthDir == nil && resolvedAuthDir != "" { + cloned.AuthDir = resolvedAuthDir + w.authDir = resolvedAuthDir + } else if w.authDir == "" { + w.authDir = trimmedAuthDir + } + } + } w.config = cloned w.oldConfigYaml, _ = yaml.Marshal(cloned) + w.clientsMutex.Unlock() + w.syncAuthDirWatch() +} + +func (w *Watcher) currentAuthDirLocked() string { + if w == nil { + return "" + } + if authDir := strings.TrimSpace(w.authDir); authDir != "" { + return authDir + } + if w.config != nil { + if cfgAuthDir := strings.TrimSpace(w.config.AuthDir); cfgAuthDir != "" { + if resolved, errResolve := util.ResolveAuthDir(cfgAuthDir); errResolve == nil && resolved != "" { + return resolved + } + return cfgAuthDir + } + } + return "" +} + +func (w *Watcher) currentAuthDir() string { + if w == nil { + return "" + } + w.clientsMutex.RLock() + defer w.clientsMutex.RUnlock() + return w.currentAuthDirLocked() +} + +func (w *Watcher) isPathInAuthDir(path, authDir string) bool { + normalizedPath := w.normalizeAuthPath(path) + normalizedDir := w.normalizeAuthPath(authDir) + if normalizedPath == "" || normalizedDir == "" { + return false + } + return w.sameNormalizedAuthPath(filepath.Dir(normalizedPath), normalizedDir) +} + +func (w *Watcher) isPathInCurrentAuthDirLocked(path string) bool { + authDir := w.currentAuthDirLocked() + if authDir == "" { + return true + } + return w.isPathInAuthDir(path, authDir) +} + +func (w *Watcher) isPathInCurrentAuthDir(path string) bool { + if w == nil { + return false + } + w.clientsMutex.RLock() + defer w.clientsMutex.RUnlock() + return w.isPathInCurrentAuthDirLocked(path) } // SetPluginAuthParser updates the plugin auth parser used for file auth synthesis. @@ -215,7 +287,7 @@ func (w *Watcher) RereadDiscardedFileSnapshot(auth *coreauth.Auth) { func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth { w.clientsMutex.RLock() cfg := w.config - authDir := w.authDir + authDir := w.currentAuthDirLocked() parser := w.pluginAuthParser w.clientsMutex.RUnlock() return snapshotCoreAuths(cfg, authDir, parser) diff --git a/internal/watcher/watcher_test.go b/internal/watcher/watcher_test.go index 9df65e6b7e2..22b7b909bc6 100644 --- a/internal/watcher/watcher_test.go +++ b/internal/watcher/watcher_test.go @@ -20,6 +20,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/watcher/synthesizer" sdkAuth "github.com/router-for-me/CLIProxyAPI/v8/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" "gopkg.in/yaml.v3" ) @@ -1885,3 +1886,361 @@ func TestScheduleProcessEventsStopsOnContextDone(t *testing.T) { func hexString(data []byte) string { return strings.ToLower(fmt.Sprintf("%x", data)) } + +func TestReloadConfigUpdatesAuthDirWatchAndSnapshots(t *testing.T) { + tmpDir := t.TempDir() + dirA := filepath.Join(tmpDir, "auth-a") + dirB := filepath.Join(tmpDir, "auth-b") + if errMkdirA := os.MkdirAll(dirA, 0o755); errMkdirA != nil { + t.Fatalf("failed to create auth dir A: %v", errMkdirA) + } + if errMkdirB := os.MkdirAll(dirB, 0o755); errMkdirB != nil { + t.Fatalf("failed to create auth dir B: %v", errMkdirB) + } + + credAPath := filepath.Join(dirA, "cred-a.json") + if errWriteA := os.WriteFile(credAPath, []byte(`{"type":"codex","email":"a@example.com"}`), 0o600); errWriteA != nil { + t.Fatalf("failed to write cred-a.json: %v", errWriteA) + } + sharedAPath := filepath.Join(dirA, "shared.json") + if errWriteSharedA := os.WriteFile(sharedAPath, []byte(`{"type":"codex","email":"shared-a@example.com"}`), 0o600); errWriteSharedA != nil { + t.Fatalf("failed to write shared.json in dirA: %v", errWriteSharedA) + } + credBPath := filepath.Join(dirB, "cred-b.json") + if errWriteB := os.WriteFile(credBPath, []byte(`{"type":"codex","email":"b@example.com"}`), 0o600); errWriteB != nil { + t.Fatalf("failed to write cred-b.json: %v", errWriteB) + } + sharedBPath := filepath.Join(dirB, "shared.json") + if errWriteSharedB := os.WriteFile(sharedBPath, []byte(`{"type":"codex","email":"shared-b@example.com"}`), 0o600); errWriteSharedB != nil { + t.Fatalf("failed to write shared.json in dirB: %v", errWriteSharedB) + } + + configPath := filepath.Join(tmpDir, "config.yaml") + writeConfig := func(authDir string) { + t.Helper() + cfg := &config.Config{ + AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), + } + data, errMarshal := yaml.Marshal(cfg) + if errMarshal != nil { + t.Fatalf("failed to marshal config: %v", errMarshal) + } + if errWriteCfg := os.WriteFile(configPath, data, 0o600); errWriteCfg != nil { + t.Fatalf("failed to write config: %v", errWriteCfg) + } + } + writeConfig(dirA) + + w, errNewWatcher := NewWatcher(configPath, dirA, nil) + if errNewWatcher != nil { + t.Fatalf("NewWatcher failed: %v", errNewWatcher) + } + defer func() { + if errStop := w.Stop(); errStop != nil { + t.Errorf("failed to stop watcher: %v", errStop) + } + }() + + queue := make(chan AuthUpdate, 16) + w.SetAuthUpdateQueue(queue) + w.SetConfig(&config.Config{ + AuthDir: dirA, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + if errStart := w.Start(ctx); errStart != nil { + t.Fatalf("Start failed: %v", errStart) + } + + initialSeen := make(map[string]bool) + for len(initialSeen) < 2 { + select { + case initial := <-queue: + if initial.Action != AuthUpdateActionAdd { + t.Fatalf("unexpected initial auth update: %+v", initial) + } + initialSeen[initial.ID] = true + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for initial updates, got %v", initialSeen) + } + } + if !initialSeen["cred-a.json"] || !initialSeen["shared.json"] { + t.Fatalf("expected cred-a.json and shared.json in initial updates, got %v", initialSeen) + } + + // Switch auth-dir from dirA to dirB at runtime. + writeConfig(dirB) + w.ReloadConfigIfChanged() + + seenAddB := false + seenModifyShared := false + seenDeleteA := false + for !seenAddB || !seenModifyShared || !seenDeleteA { + select { + case update := <-queue: + if update.ID == "cred-b.json" && update.Action == AuthUpdateActionAdd { + seenAddB = true + } + if update.ID == "shared.json" && update.Action == AuthUpdateActionModify && update.Auth != nil && update.Auth.Label == "shared-b@example.com" { + seenModifyShared = true + } + if update.ID == "cred-a.json" && update.Action == AuthUpdateActionDelete { + seenDeleteA = true + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for auth-dir switch updates (seenAddB=%v, seenModifyShared=%v, seenDeleteA=%v)", seenAddB, seenModifyShared, seenDeleteA) + } + } + + snapshot := w.SnapshotCoreAuths() + if len(snapshot) != 2 { + t.Fatalf("expected SnapshotCoreAuths to return 2 entries from dirB, got %+v", snapshot) + } + w.clientsMutex.RLock() + _, staleSharedIndexed := w.fileAuthsByPath[w.normalizeAuthPath(sharedAPath)] + _, newSharedIndexed := w.fileAuthsByPath[w.normalizeAuthPath(sharedBPath)] + w.clientsMutex.RUnlock() + if staleSharedIndexed || !newSharedIndexed { + t.Fatalf("expected fileAuthsByPath to drop dirA shared.json (%v) and keep dirB shared.json (%v)", staleSharedIndexed, newSharedIndexed) + } + + fileSnapshot, errSnapshotFile := w.SnapshotAuthFile(credBPath) + if errSnapshotFile != nil { + t.Fatalf("SnapshotAuthFile(credBPath) failed: %v", errSnapshotFile) + } + if len(fileSnapshot) != 1 || fileSnapshot[0].ID != "cred-b.json" { + t.Fatalf("expected SnapshotAuthFile(credBPath) ID=cred-b.json, got %+v", fileSnapshot) + } + + // Events for the old directory must be ignored, while events in the new directory must be processed. + credA2Path := filepath.Join(dirA, "cred-a2.json") + if errWriteA2 := os.WriteFile(credA2Path, []byte(`{"type":"codex","email":"a2@example.com"}`), 0o600); errWriteA2 != nil { + t.Fatalf("failed to write cred-a2.json: %v", errWriteA2) + } + w.handleEvent(fsnotify.Event{Name: credA2Path, Op: fsnotify.Create}) + + credB2Path := filepath.Join(dirB, "cred-b2.json") + if errWriteB2 := os.WriteFile(credB2Path, []byte(`{"type":"codex","email":"b2@example.com"}`), 0o600); errWriteB2 != nil { + t.Fatalf("failed to write cred-b2.json: %v", errWriteB2) + } + + select { + case update := <-queue: + if update.Action != AuthUpdateActionAdd || update.ID != "cred-b2.json" { + t.Fatalf("expected add update for cred-b2.json, got %+v", update) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for fsnotify update from new auth-dir dirB") + } +} + +func TestSetConfigUpdatesAuthDir(t *testing.T) { + tmpDir := t.TempDir() + dirA := filepath.Join(tmpDir, "auth-a") + dirB := filepath.Join(tmpDir, "auth-b") + if errMkdirA := os.MkdirAll(dirA, 0o755); errMkdirA != nil { + t.Fatalf("failed to create auth dir A: %v", errMkdirA) + } + if errMkdirB := os.MkdirAll(dirB, 0o755); errMkdirB != nil { + t.Fatalf("failed to create auth dir B: %v", errMkdirB) + } + credBPath := filepath.Join(dirB, "cred-b.json") + if errWriteB := os.WriteFile(credBPath, []byte(`{"type":"codex","email":"b@example.com"}`), 0o600); errWriteB != nil { + t.Fatalf("failed to write cred-b.json: %v", errWriteB) + } + + w, errNewWatcher := NewWatcher(filepath.Join(tmpDir, "config.yaml"), dirA, nil) + if errNewWatcher != nil { + t.Fatalf("NewWatcher failed: %v", errNewWatcher) + } + defer func() { + if errStop := w.Stop(); errStop != nil { + t.Errorf("failed to stop watcher: %v", errStop) + } + }() + + w.SetConfig(&config.Config{AuthDir: dirB}) + + snapshot := w.SnapshotCoreAuths() + if len(snapshot) != 1 || snapshot[0].ID != "cred-b.json" { + t.Fatalf("expected SnapshotCoreAuths to return cred-b.json from dirB, got %+v", snapshot) + } + fileSnapshot, errSnapshotFile := w.SnapshotAuthFile(credBPath) + if errSnapshotFile != nil { + t.Fatalf("SnapshotAuthFile failed: %v", errSnapshotFile) + } + if len(fileSnapshot) != 1 || fileSnapshot[0].ID != "cred-b.json" { + t.Fatalf("expected SnapshotAuthFile ID=cred-b.json, got %+v", fileSnapshot) + } +} + +type blockingAuthDirSwitchParser struct { + blockPath string + entered chan struct{} + release chan struct{} + once sync.Once +} + +func (p *blockingAuthDirSwitchParser) ParseAuth(_ context.Context, req pluginapi.AuthParseRequest) (*coreauth.Auth, bool, error) { + if filepath.Clean(req.Path) == filepath.Clean(p.blockPath) { + p.once.Do(func() { + close(p.entered) + }) + <-p.release + } + return nil, false, nil +} + +func TestAddOrUpdateClientDiscardsStaleDirEventAcrossAuthDirSwitch(t *testing.T) { + tmpDir := t.TempDir() + dirA := filepath.Join(tmpDir, "auth-a") + dirB := filepath.Join(tmpDir, "auth-b") + if errMkdirA := os.MkdirAll(dirA, 0o755); errMkdirA != nil { + t.Fatalf("failed to create auth dir A: %v", errMkdirA) + } + if errMkdirB := os.MkdirAll(dirB, 0o755); errMkdirB != nil { + t.Fatalf("failed to create auth dir B: %v", errMkdirB) + } + + lateAPath := filepath.Join(dirA, "late-a.json") + if errWriteLateA := os.WriteFile(lateAPath, []byte(`{"type":"codex","email":"late-a@example.com"}`), 0o600); errWriteLateA != nil { + t.Fatalf("failed to write late-a.json: %v", errWriteLateA) + } + credBPath := filepath.Join(dirB, "cred-b.json") + if errWriteB := os.WriteFile(credBPath, []byte(`{"type":"codex","email":"b@example.com"}`), 0o600); errWriteB != nil { + t.Fatalf("failed to write cred-b.json: %v", errWriteB) + } + + parser := &blockingAuthDirSwitchParser{ + blockPath: lateAPath, + entered: make(chan struct{}), + release: make(chan struct{}), + } + + queue := make(chan AuthUpdate, 16) + w := &Watcher{ + authDir: dirA, + lastAuthHashes: make(map[string]string), + fileAuthsByPath: make(map[string]map[string]*coreauth.Auth), + pluginAuthParser: parser, + } + w.SetConfig(&config.Config{AuthDir: dirA}) + w.SetAuthUpdateQueue(queue) + defer w.stopDispatch() + + done := make(chan struct{}) + go func() { + w.addOrUpdateClientLocked(lateAPath) + close(done) + }() + + select { + case <-parser.entered: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for addOrUpdateClientLocked to reach synthesis") + } + + // Switch auth-dir to dirB while the dirA event is paused mid-synthesis. + w.SetConfig(&config.Config{AuthDir: dirB}) + w.refreshAuthState(false) + + select { + case update := <-queue: + if update.Action != AuthUpdateActionAdd || update.ID != "cred-b.json" { + t.Fatalf("expected add update for cred-b.json, got %+v", update) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for cred-b.json update") + } + + close(parser.release) + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for paused addOrUpdateClientLocked to finish") + } + + w.clientsMutex.RLock() + currentCount := len(w.currentAuths) + _, hasCredB := w.currentAuths["cred-b.json"] + _, hasLateA := w.lastAuthHashes[w.normalizeAuthPath(lateAPath)] + w.clientsMutex.RUnlock() + if currentCount != 1 || !hasCredB || hasLateA { + t.Fatalf("stale dirA event polluted state after switch: currentCount=%d, hasCredB=%v, hasLateA=%v", currentCount, hasCredB, hasLateA) + } +} + +func TestConcurrentSetConfigKeepsWatchedAuthDirOnLatestConfig(t *testing.T) { + tmpDir := t.TempDir() + dirA := filepath.Join(tmpDir, "auth-a") + dirB := filepath.Join(tmpDir, "auth-b") + dirC := filepath.Join(tmpDir, "auth-c") + for _, dir := range []string{dirA, dirB, dirC} { + if errMkdir := os.MkdirAll(dir, 0o755); errMkdir != nil { + t.Fatalf("failed to create auth dir %s: %v", dir, errMkdir) + } + } + configPath := filepath.Join(tmpDir, "config.yaml") + if errWriteCfg := os.WriteFile(configPath, []byte("auth_dir: "+dirA+"\n"), 0o600); errWriteCfg != nil { + t.Fatalf("failed to write config: %v", errWriteCfg) + } + + w, errNewWatcher := NewWatcher(configPath, dirA, nil) + if errNewWatcher != nil { + t.Fatalf("NewWatcher failed: %v", errNewWatcher) + } + defer func() { + if errStop := w.Stop(); errStop != nil { + t.Errorf("failed to stop watcher: %v", errStop) + } + }() + + queue := make(chan AuthUpdate, 16) + w.SetAuthUpdateQueue(queue) + w.SetConfig(&config.Config{AuthDir: dirA}) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if errStart := w.Start(ctx); errStart != nil { + t.Fatalf("Start failed: %v", errStart) + } + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(2) + go func() { + defer wg.Done() + w.SetConfig(&config.Config{AuthDir: dirB}) + }() + go func() { + defer wg.Done() + w.SetConfig(&config.Config{AuthDir: dirC}) + }() + } + wg.Wait() + w.SetConfig(&config.Config{AuthDir: dirC}) + + w.watchMu.Lock() + watched := w.watchedAuthDir + w.watchMu.Unlock() + if !w.sameNormalizedAuthPath(watched, dirC) || !w.sameNormalizedAuthPath(w.currentAuthDir(), dirC) { + t.Fatalf("expected watchedAuthDir=%q and currentAuthDir=%q to match dirC=%q", watched, w.currentAuthDir(), dirC) + } + + credCPath := filepath.Join(dirC, "cred-c.json") + if errWriteC := os.WriteFile(credCPath, []byte(`{"type":"codex","email":"c@example.com"}`), 0o600); errWriteC != nil { + t.Fatalf("failed to write cred-c.json: %v", errWriteC) + } + select { + case update := <-queue: + if update.Action != AuthUpdateActionAdd || update.ID != "cred-c.json" { + t.Fatalf("expected add update for cred-c.json, got %+v", update) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for fsnotify update from dirC") + } +} diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index d365430d7db..2f1b89a5072 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -158,9 +158,19 @@ func BuildErrorResponseBodyWithError(status int, errText string, err error) []by return payload } -var streamingKeepAliveTickerFactory = func(interval time.Duration) (<-chan time.Time, func()) { +var streamingKeepAliveTickerFactory func(interval time.Duration) (<-chan time.Time, func()) +var streamingKeepAliveTickerWithResetFactory func(interval time.Duration) (<-chan time.Time, func(), func()) + +func defaultStreamingKeepAliveTickerWithReset(interval time.Duration) (<-chan time.Time, func(), func()) { ticker := time.NewTicker(interval) - return ticker.C, ticker.Stop + reset := func() { + ticker.Reset(interval) + select { + case <-ticker.C: + default: + } + } + return ticker.C, ticker.Stop, reset } // StreamingKeepAliveInterval returns the streaming keep-alive interval for this server (SSE heartbeats and WebSocket Ping frames). @@ -182,13 +192,27 @@ func StreamingSendKeepAliveBeforeFT(cfg *config.SDKConfig) bool { return cfg != nil && cfg.Streaming.SendKeepAliveBeforeFT && StreamingKeepAliveInterval(cfg) > 0 } +// NewStreamingKeepAliveTickerWithReset returns a ticker channel, a stop function, and a reset function +// for streaming keep-alives. When interval <= 0, it returns a nil channel and no-op functions. +func NewStreamingKeepAliveTickerWithReset(interval time.Duration) (<-chan time.Time, func(), func()) { + if interval <= 0 { + return nil, func() {}, func() {} + } + if streamingKeepAliveTickerWithResetFactory != nil { + return streamingKeepAliveTickerWithResetFactory(interval) + } + if streamingKeepAliveTickerFactory != nil { + keepAliveC, stop := streamingKeepAliveTickerFactory(interval) + return keepAliveC, stop, func() {} + } + return defaultStreamingKeepAliveTickerWithReset(interval) +} + // NewStreamingKeepAliveTicker returns a ticker channel and a stop function for streaming keep-alives. // When interval <= 0, it returns a nil channel and a no-op stop function. func NewStreamingKeepAliveTicker(interval time.Duration) (<-chan time.Time, func()) { - if interval <= 0 { - return nil, func() {} - } - return streamingKeepAliveTickerFactory(interval) + keepAliveC, stop, _ := NewStreamingKeepAliveTickerWithReset(interval) + return keepAliveC, stop } // NewPreFirstTokenKeepAliveTicker returns a ticker channel and a stop function for pre-first-token SSE keep-alives. @@ -212,6 +236,18 @@ func SetStreamingKeepAliveTickerFactoryForTest(factory func(interval time.Durati } } +// SetStreamingKeepAliveTickerWithResetFactoryForTest overrides the keep-alive ticker-with-reset factory +// for deterministic tests and returns a restore function. +func SetStreamingKeepAliveTickerWithResetFactoryForTest(factory func(interval time.Duration) (<-chan time.Time, func(), func())) func() { + prev := streamingKeepAliveTickerWithResetFactory + if factory != nil { + streamingKeepAliveTickerWithResetFactory = factory + } + return func() { + streamingKeepAliveTickerWithResetFactory = prev + } +} + // NonStreamingKeepAliveInterval returns the keep-alive interval for non-streaming responses. // Returning 0 disables keep-alives (default when unset). func NonStreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration { diff --git a/sdk/api/handlers/openai/openai_images_handlers.go b/sdk/api/handlers/openai/openai_images_handlers.go index d70cc841404..7e2dde146ad 100644 --- a/sdk/api/handlers/openai/openai_images_handlers.go +++ b/sdk/api/handlers/openai/openai_images_handlers.go @@ -79,11 +79,46 @@ func setImagesSSEHeaders(c *gin.Context) { } func (h *OpenAIAPIHandler) newImagesStreamKeepAliveTicker() (func(), <-chan time.Time) { + stop, keepAliveC, _ := h.newImagesStreamKeepAliveTickerWithReset() + return stop, keepAliveC +} + +func (h *OpenAIAPIHandler) newImagesStreamKeepAliveTickerWithReset() (func(), <-chan time.Time, func()) { if h == nil || h.BaseAPIHandler == nil { - return func() {}, nil + return func() {}, nil, func() {} } - keepAliveC, stop := handlers.NewStreamingKeepAliveTicker(handlers.StreamingKeepAliveInterval(h.CurrentConfig())) - return stop, keepAliveC + keepAliveC, stop, reset := handlers.NewStreamingKeepAliveTickerWithReset(handlers.StreamingKeepAliveInterval(h.CurrentConfig())) + return stop, keepAliveC, reset +} + +func advanceImagesSSENewlineCount(prev int, chunk []byte) int { + count := prev + for _, b := range chunk { + switch b { + case '\r': + continue + case '\n': + if count < 2 { + count++ + } + default: + count = 0 + } + } + return count +} + +func ensureImagesSSEEventBoundary(w io.Writer, trailingNewlines int) int { + if w == nil { + return 2 + } + switch trailingNewlines { + case 0: + _, _ = w.Write([]byte("\n\n")) + case 1: + _, _ = w.Write([]byte("\n")) + } + return 2 } func (h *OpenAIAPIHandler) newImagesPreFirstTokenKeepAliveTicker() (func(), <-chan time.Time) { @@ -1417,10 +1452,11 @@ func (h *OpenAIAPIHandler) streamRoutedImages(c *gin.Context, imageReq []byte, i stopKeepAlive() setImagesSSEHeaders(c) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) + trailingNewlines := advanceImagesSSENewlineCount(2, chunk) _, _ = c.Writer.Write(chunk) flusher.Flush() streamStarted = true - h.forwardRawImageStream(cliCtx, c, func(err error) { cliCancel(err) }, dataChan, errChan) + h.forwardRawImageStreamWithState(cliCtx, c, func(err error) { cliCancel(err) }, dataChan, errChan, trailingNewlines) return case <-keepAliveC: setImagesSSEHeaders(c) @@ -1432,9 +1468,42 @@ func (h *OpenAIAPIHandler) streamRoutedImages(c *gin.Context, imageReq []byte, i } func (h *OpenAIAPIHandler) forwardRawImageStream(ctx context.Context, c *gin.Context, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) { - stopKeepAlive, keepAliveC := h.newImagesStreamKeepAliveTicker() + h.forwardRawImageStreamWithState(ctx, c, cancel, data, errs, 2) +} + +func (h *OpenAIAPIHandler) forwardRawImageStreamWithState(ctx context.Context, c *gin.Context, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage, trailingNewlines int) { + stopKeepAlive, keepAliveC, resetKeepAlive := h.newImagesStreamKeepAliveTickerWithReset() defer stopKeepAlive() + writeErrorEvent := func(errMsg *interfaces.ErrorMessage) { + trailingNewlines = ensureImagesSSEEventBoundary(c.Writer, trailingNewlines) + writeImagesStreamErrorEvent(c, errMsg) + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + } + + handleChunk := func(chunk []byte, ok bool) bool { + if !ok { + if errMsg, hasPendingError := handlers.PendingStreamError(errs); hasPendingError { + writeErrorEvent(errMsg) + cancel(errMsg.Error) + return true + } + cancel(nil) + return true + } + if len(chunk) > 0 { + resetKeepAlive() + trailingNewlines = advanceImagesSSENewlineCount(trailingNewlines, chunk) + _, _ = c.Writer.Write(chunk) + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + } + return false + } + for { select { case <-c.Request.Context().Done(): @@ -1445,32 +1514,19 @@ func (h *OpenAIAPIHandler) forwardRawImageStream(ctx context.Context, c *gin.Con return case errMsg, ok := <-errs: if ok && errMsg != nil { - writeImagesStreamErrorEvent(c, errMsg) - if flusher, ok := c.Writer.(http.Flusher); ok { - flusher.Flush() - } + writeErrorEvent(errMsg) cancel(errMsg.Error) return } errs = nil case chunk, ok := <-data: - if !ok { - if errMsg, hasPendingError := handlers.PendingStreamError(errs); hasPendingError { - writeImagesStreamErrorEvent(c, errMsg) - if flusher, ok := c.Writer.(http.Flusher); ok { - flusher.Flush() - } - cancel(errMsg.Error) - return - } - cancel(nil) + if handleChunk(chunk, ok) { return } - _, _ = c.Writer.Write(chunk) - if flusher, ok := c.Writer.(http.Flusher); ok { - flusher.Flush() - } case <-keepAliveC: + if trailingNewlines < 2 { + continue + } if flusher, ok := c.Writer.(http.Flusher); ok { writeImagesStreamKeepAlive(c, flusher) } @@ -1557,17 +1613,29 @@ func (h *OpenAIAPIHandler) streamOpenAICompatImages(c *gin.Context, compatReq [] stopKeepAlive() setImagesSSEHeaders(c) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) + trailingNewlines := advanceImagesSSENewlineCount(2, chunk) _, _ = c.Writer.Write(chunk) flusher.Flush() streamStarted = true h.ForwardStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan, handlers.StreamForwardOptions{ NormalizeTerminalError: sanitizeResponsesStreamErrorMessage, WriteChunk: func(next []byte) { + if len(next) == 0 { + return + } + trailingNewlines = advanceImagesSSENewlineCount(trailingNewlines, next) _, _ = c.Writer.Write(next) }, WriteTerminalError: func(errMsg *interfaces.ErrorMessage) { + trailingNewlines = ensureImagesSSEEventBoundary(c.Writer, trailingNewlines) writeImagesStreamErrorEvent(c, errMsg) }, + WriteKeepAlive: func() { + if trailingNewlines < 2 { + return + } + _, _ = c.Writer.Write([]byte(": keep-alive\n\n")) + }, }) return case <-keepAliveC: @@ -2049,7 +2117,7 @@ func (h *OpenAIAPIHandler) forwardImagesStream(ctx context.Context, c *gin.Conte if responseFormat == "" { responseFormat = "b64_json" } - stopKeepAlive, keepAliveC := h.newImagesStreamKeepAliveTicker() + stopKeepAlive, keepAliveC, resetKeepAlive := h.newImagesStreamKeepAliveTickerWithReset() defer stopKeepAlive() emitError := func(errMsg *interfaces.ErrorMessage) *interfaces.ErrorMessage { @@ -2089,6 +2157,7 @@ func (h *OpenAIAPIHandler) forwardImagesStream(ctx context.Context, c *gin.Conte } else { _ = dataDoc.Set("b64_json", b64) } + resetKeepAlive() writeEvent(eventName, dataDoc.Bytes()) case "response.completed": results, _, usageRaw, _, err := extractImagesFromResponsesCompleted(payload) @@ -2111,6 +2180,7 @@ func (h *OpenAIAPIHandler) forwardImagesStream(ctx context.Context, c *gin.Conte if len(usageRaw) > 0 && json.Valid(usageRaw) { _ = dataDoc.SetRawBytes("usage", usageRaw) } + resetKeepAlive() writeEvent(eventName, dataDoc.Bytes()) } return true, nil @@ -2138,6 +2208,29 @@ func (h *OpenAIAPIHandler) forwardImagesStream(ctx context.Context, c *gin.Conte } } + handleDataChunk := func(chunk []byte, ok bool) bool { + if !ok { + for _, frame := range acc.Flush() { + if handleFrame(frame) { + return true + } + } + if errMsg, hasPendingError := handlers.PendingStreamError(errs); hasPendingError { + errMsg = emitError(errMsg) + cancel(errMsg.Error) + return true + } + cancel(nil) + return true + } + for _, frame := range acc.AddChunk(chunk) { + if handleFrame(frame) { + return true + } + } + return false + } + for { select { case <-c.Request.Context().Done(): @@ -2151,25 +2244,9 @@ func (h *OpenAIAPIHandler) forwardImagesStream(ctx context.Context, c *gin.Conte } errs = nil case chunk, ok := <-data: - if !ok { - for _, frame := range acc.Flush() { - if handleFrame(frame) { - return - } - } - if errMsg, hasPendingError := handlers.PendingStreamError(errs); hasPendingError { - errMsg = emitError(errMsg) - cancel(errMsg.Error) - return - } - cancel(nil) + if handleDataChunk(chunk, ok) { return } - for _, frame := range acc.AddChunk(chunk) { - if handleFrame(frame) { - return - } - } case <-keepAliveC: writeImagesStreamKeepAlive(c, flusher) } diff --git a/sdk/api/handlers/openai/openai_images_handlers_test.go b/sdk/api/handlers/openai/openai_images_handlers_test.go index 2dd9859c909..e4dcbbd57fe 100644 --- a/sdk/api/handlers/openai/openai_images_handlers_test.go +++ b/sdk/api/handlers/openai/openai_images_handlers_test.go @@ -13,6 +13,7 @@ import ( "strings" "sync" "testing" + "time" "github.com/MoeForge/moego/ojson" "github.com/gin-gonic/gin" @@ -1299,3 +1300,260 @@ func TestXAIImagesPayloadDefaultAndOverrideRules(t *testing.T) { } }) } + +type rawImageStreamTestExecutor struct { + provider string + chunks <-chan coreexecutor.StreamChunk +} + +func (e *rawImageStreamTestExecutor) Identifier() string { return e.provider } + +func (e *rawImageStreamTestExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} +} + +func (e *rawImageStreamTestExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + return &coreexecutor.StreamResult{Chunks: e.chunks}, nil +} + +func (e *rawImageStreamTestExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e *rawImageStreamTestExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} +} + +func (e *rawImageStreamTestExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented"} +} + +func TestImagesStreamKeepAliveDoesNotSplitCompletedEventMidLine_Issue7330(t *testing.T) { + gin.SetMode(gin.TestMode) + + const ( + keepAliveFrame = "event: keepalive\ndata: {\"type\":\"keepalive\"}\n\n" + completedPart1 = "event: image_generation.completed\ndata: {\"type\":\"image_generation.completed\",\"b64_json\":\"iVBORw0KGgoAAAA" + completedPart2 = "BBBB\",\"size\":\"1536x1024\"}\n\n" + wantCompleted = "event: image_generation.completed\ndata: {\"type\":\"image_generation.completed\",\"b64_json\":\"iVBORw0KGgoAAAABBBB\",\"size\":\"1536x1024\"}\n\n" + ) + + t.Run("forwardRawImageStream suppresses keep-alive mid-line and writes at event boundary", func(t *testing.T) { + tickChan := make(chan time.Time) + restore := handlers.SetStreamingKeepAliveTickerFactoryForTest(func(time.Duration) (<-chan time.Time, func()) { + return tickChan, func() {} + }) + defer restore() + + h := NewOpenAIAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{ + Streaming: sdkconfig.StreamingConfig{KeepAliveSeconds: 20}, + }, nil)) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, imagesGenerationsPath, nil) + + data := make(chan []byte) + errs := make(chan *interfaces.ErrorMessage) + close(errs) + + go func() { + data <- []byte(keepAliveFrame) + data <- []byte(completedPart1) + // Tick while in the middle of the image_generation.completed data: line. + tickChan <- time.Unix(1, 0) + data <- []byte(completedPart2) + // Tick after the SSE event has completed with \n\n. + tickChan <- time.Unix(2, 0) + close(data) + }() + + var canceled error + h.forwardRawImageStream(context.Background(), c, func(err error) { canceled = err }, data, errs) + if canceled != nil { + t.Fatalf("unexpected cancel error: %v", canceled) + } + + wantBody := keepAliveFrame + wantCompleted + ": keep-alive\n\n" + if got := recorder.Body.String(); got != wantBody { + t.Fatalf("stream body = %q, want %q", got, wantBody) + } + }) + + for _, tc := range []struct { + name string + provider string + model string + modelType string + }{ + { + name: "routed codex image stream", + provider: "codex", + model: defaultImagesToolModel, + modelType: "codex", + }, + { + name: "openai-compatibility image stream", + provider: "openai-compatibility", + model: "compat-image-7330", + modelType: registry.OpenAIImageModelType, + }, + } { + t.Run(tc.name, func(t *testing.T) { + tickChan := make(chan time.Time) + resetCh := make(chan struct{}, 8) + restore := handlers.SetStreamingKeepAliveTickerWithResetFactoryForTest(func(time.Duration) (<-chan time.Time, func(), func()) { + return tickChan, func() {}, func() { + resetCh <- struct{}{} + } + }) + defer restore() + + chunks := make(chan coreexecutor.StreamChunk) + exec := &rawImageStreamTestExecutor{provider: tc.provider, chunks: chunks} + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + manager.RegisterExecutor(exec) + auth := &coreauth.Auth{ + ID: "auth-7330-" + tc.provider, + Provider: tc.provider, + Status: coreauth.StatusActive, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("manager.Register: %v", errRegister) + } + modelRegistry := registry.GetGlobalRegistry() + modelRegistry.RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{ + {ID: tc.model, Object: "model", OwnedBy: tc.provider, Type: tc.modelType}, + }) + manager.RefreshSchedulerEntry(auth.ID) + t.Cleanup(func() { + modelRegistry.UnregisterClient(auth.ID) + }) + + go func() { + // First chunk itself is an incomplete SSE event (first 32KB read of image_generation.completed). + chunks <- coreexecutor.StreamChunk{Payload: []byte(completedPart1)} + // Keep-alive tick fires before the second chunk of the same line arrives. + tickChan <- time.Unix(1, 0) + chunks <- coreexecutor.StreamChunk{Payload: []byte(completedPart2)} + // Wait until the handler receives completedPart2 and resets the keep-alive timer. + <-resetCh + // Keep-alive tick fires at the completed event boundary. + tickChan <- time.Unix(2, 0) + close(chunks) + }() + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{ + Streaming: sdkconfig.StreamingConfig{KeepAliveSeconds: 20}, + }, manager) + handler := NewOpenAIAPIHandler(base) + body := strings.NewReader("{\"model\":\"" + tc.model + "\",\"prompt\":\"a red fox\",\"stream\":true}") + resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations) + if resp.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusOK, resp.Body.String()) + } + + wantBody := wantCompleted + ": keep-alive\n\n" + if got := resp.Body.String(); got != wantBody { + t.Fatalf("stream body = %q, want %q", got, wantBody) + } + }) + } + + t.Run("resets keep-alive on each chunk and ensures SSE event boundary before error", func(t *testing.T) { + tickChan := make(chan time.Time) + data := make(chan []byte) + errs := make(chan *interfaces.ErrorMessage, 1) + resetCount := 0 + restore := handlers.SetStreamingKeepAliveTickerWithResetFactoryForTest(func(time.Duration) (<-chan time.Time, func(), func()) { + return tickChan, func() {}, func() { + resetCount++ + } + }) + defer restore() + + h := NewOpenAIAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{ + Streaming: sdkconfig.StreamingConfig{KeepAliveSeconds: 20}, + }, nil)) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, imagesGenerationsPath, nil) + + go func() { + data <- []byte(completedPart1) + tickChan <- time.Unix(1, 0) + errs <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errors.New("upstream interrupted mid-frame")} + close(errs) + close(data) + }() + + var canceled error + h.forwardRawImageStream(context.Background(), c, func(err error) { canceled = err }, data, errs) + if canceled == nil || !strings.Contains(canceled.Error(), "upstream interrupted mid-frame") { + t.Fatalf("unexpected cancel error: %v", canceled) + } + if resetCount != 1 { + t.Fatalf("resetCount = %d, want 1", resetCount) + } + got := recorder.Body.String() + if strings.Contains(got, ": keep-alive") { + t.Fatalf("expected mid-frame keep-alive tick to be suppressed, got %q", got) + } + if !strings.Contains(got, completedPart1+"\n\nevent: error\ndata: ") { + t.Fatalf("expected error event to be preceded by SSE event boundary \\n\\n, got %q", got) + } + }) + + t.Run("forwardImagesStream still writes keep-alive when upstream chunks produce no downstream output", func(t *testing.T) { + tickChan := make(chan time.Time) + resetCount := 0 + restore := handlers.SetStreamingKeepAliveTickerWithResetFactoryForTest(func(time.Duration) (<-chan time.Time, func(), func()) { + return tickChan, func() {}, func() { + resetCount++ + } + }) + defer restore() + + h := NewOpenAIAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{ + Streaming: sdkconfig.StreamingConfig{KeepAliveSeconds: 20}, + }, nil)) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, imagesGenerationsPath, nil) + flusher, ok := c.Writer.(http.Flusher) + if !ok { + t.Fatal("expected gin writer to implement http.Flusher") + } + + data := make(chan []byte) + errs := make(chan *interfaces.ErrorMessage) + close(errs) + + firstChunk := []byte("event: response.created\ndata: {\"type\":\"response.created\"}\n\n") + go func() { + // Ignored non-image event and buffered incomplete frame produce no downstream output. + data <- []byte("event: response.in_progress\ndata: {\"type\":\"response.in_progress\"}\n\n") + data <- []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"image_generation_call\",") + // Keep-alive tick must still be written while downstream has received no output. + tickChan <- time.Unix(1, 0) + data <- []byte("\"result\":\"aW1hZ2U=\"}]}}\n\n") + close(data) + }() + + var canceled error + writeEvent := func(eventName string, dataJSON []byte) { + _, _ = c.Writer.Write([]byte("event: " + eventName + "\ndata: " + string(dataJSON) + "\n\n")) + flusher.Flush() + } + h.forwardImagesStream(context.Background(), c, flusher, func(err error) { canceled = err }, data, errs, firstChunk, "b64_json", "image_generation", writeEvent) + if canceled != nil { + t.Fatalf("unexpected cancel error: %v", canceled) + } + if resetCount != 1 { + t.Fatalf("resetCount = %d, want 1 (only when downstream image event is written)", resetCount) + } + got := recorder.Body.String() + if !strings.HasPrefix(got, ": keep-alive\n\nevent: image_generation.completed\ndata: ") { + t.Fatalf("expected keep-alive before completed event when earlier chunks produced no downstream output, got %q", got) + } + }) +} diff --git a/sdk/api/handlers/openai/openai_responses_handlers.go b/sdk/api/handlers/openai/openai_responses_handlers.go index 6eb68e89164..0dd4aea02a2 100644 --- a/sdk/api/handlers/openai/openai_responses_handlers.go +++ b/sdk/api/handlers/openai/openai_responses_handlers.go @@ -25,6 +25,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/httpwire" "github.com/router-for-me/CLIProxyAPI/v8/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + translatorcommon "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/common" sdkaccess "github.com/router-for-me/CLIProxyAPI/v8/sdk/access" "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" @@ -935,7 +936,8 @@ func (h *OpenAIResponsesAPIHandler) prepareCodexRequestPayload(c *gin.Context, p doc := ojson.NewDoc(payload) toolsChanged := h.prepareCodexMultiAgentV2ToolsDoc(c, doc) orphanChanged := h.prepareCodexOrphanDelegationDoc(c, doc) - if len(payload) == 0 || (!toolsChanged && !orphanChanged) { + toolOutputChanged := translatorcommon.SanitizeOrphanResponsesToolOutputsDoc(doc) + if len(payload) == 0 || (!toolsChanged && !orphanChanged && !toolOutputChanged) { return payload } return doc.Bytes() @@ -986,6 +988,7 @@ func (h *OpenAIResponsesAPIHandler) Compact(c *gin.Context) { doc := ojson.NewDoc(rawJSON) orphanChanged := h.prepareCodexOrphanDelegationDoc(c, doc) + toolOutputChanged := translatorcommon.SanitizeOrphanResponsesToolOutputsDoc(doc) streamResult := doc.Get("stream") if streamResult.Type == ojson.True { @@ -1002,7 +1005,7 @@ func (h *OpenAIResponsesAPIHandler) Compact(c *gin.Context) { _ = doc.Delete("stream") streamDeleted = true } - if orphanChanged || streamDeleted { + if orphanChanged || toolOutputChanged || streamDeleted { rawJSON = doc.Bytes() } diff --git a/sdk/api/handlers/openai/openai_responses_interrupt_test.go b/sdk/api/handlers/openai/openai_responses_interrupt_test.go index 9d9cbf4c4bf..9298ee32a01 100644 --- a/sdk/api/handlers/openai/openai_responses_interrupt_test.go +++ b/sdk/api/handlers/openai/openai_responses_interrupt_test.go @@ -270,8 +270,10 @@ func (e *homeInterruptExecutor) InterruptExecutionSession(ctx context.Context, s } type blockingHTTPInterruptExecutor struct { - calls atomic.Int32 - canceled chan struct{} + calls atomic.Int32 + canceled atomic.Bool + canceledCh chan struct{} + cancelOnce sync.Once } func (*blockingHTTPInterruptExecutor) Identifier() string { return "codex" } @@ -283,6 +285,16 @@ func (*blockingHTTPInterruptExecutor) Execute(context.Context, *coreauth.Auth, c func (e *blockingHTTPInterruptExecutor) ExecuteStream(ctx context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { call := e.calls.Add(1) chunks := make(chan coreexecutor.StreamChunk, 2) + if call == 1 { + context.AfterFunc(ctx, func() { + e.canceled.Store(true) + e.cancelOnce.Do(func() { + if e.canceledCh != nil { + close(e.canceledCh) + } + }) + }) + } go func() { defer close(chunks) if call > 1 { @@ -292,7 +304,12 @@ func (e *blockingHTTPInterruptExecutor) ExecuteStream(ctx context.Context, _ *co } chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"type":"response.created","response":{"id":"r-http"}}`)} <-ctx.Done() - close(e.canceled) + e.canceled.Store(true) + e.cancelOnce.Do(func() { + if e.canceledCh != nil { + close(e.canceledCh) + } + }) }() return &coreexecutor.StreamResult{Chunks: chunks}, nil } @@ -313,11 +330,45 @@ func (*blockingHTTPInterruptExecutor) InterruptExecutionSession(context.Context, return coreexecutor.ErrNoActiveUpstreamWebsocket } +type cancelOrderTimelineAppender struct { + onAppend func(eventType string) +} + +func (a cancelOrderTimelineAppender) Append(eventType string, _ []byte, _ time.Time) { + if a.onAppend != nil { + a.onAppend(eventType) + } +} + +func TestCompleteResponsesWebsocketLocalInterruptCancelsBeforeWritingIncomplete(t *testing.T) { + canceled := false + cancelCalledBeforeWrite := false + timeline := cancelOrderTimelineAppender{ + onAppend: func(eventType string) { + if eventType == "response" { + cancelCalledBeforeWrite = canceled + } + }, + } + _, _, _, _, _ = completeResponsesWebsocketLocalInterrupt( + &responsesWebsocketWriter{}, + timeline, + func(...any) { canceled = true }, + []byte(`{"type":"response.interrupt","response_id":"r-http"}`), + nil, + nil, + nil, + ) + if !cancelCalledBeforeWrite { + t.Fatal("expected completeResponsesWebsocketLocalInterrupt to cancel upstream HTTP context before writing response.incomplete downstream") + } +} + // TestResponsesInterruptStopsHTTPUpstream verifies that an in-flight interrupt // cancels an HTTP upstream turn and lets the same socket continue. func TestResponsesInterruptStopsHTTPUpstream(t *testing.T) { gin.SetMode(gin.TestMode) - executor := &blockingHTTPInterruptExecutor{canceled: make(chan struct{})} + executor := &blockingHTTPInterruptExecutor{canceledCh: make(chan struct{})} manager := coreauth.NewManager(nil, nil, nil) manager.SetConfig(&config.Config{}) manager.RegisterExecutor(executor) @@ -373,10 +424,13 @@ func TestResponsesInterruptStopsHTTPUpstream(t *testing.T) { t.Fatalf("interrupt reason = %q, payload %s", got, interrupted) } select { - case <-executor.canceled: + case <-executor.canceledCh: case <-time.After(5 * time.Second): t.Fatal("http upstream was not canceled") } + if !executor.canceled.Load() { + t.Fatal("http upstream was not canceled") + } if errSend := client.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"previous_response_id":"r-http","input":[]}`, model))); errSend != nil { t.Fatal(errSend) } diff --git a/sdk/api/handlers/openai/openai_responses_multi_agent_test.go b/sdk/api/handlers/openai/openai_responses_multi_agent_test.go index d1de0e41ddc..fbdeb5a14b1 100644 --- a/sdk/api/handlers/openai/openai_responses_multi_agent_test.go +++ b/sdk/api/handlers/openai/openai_responses_multi_agent_test.go @@ -330,3 +330,105 @@ func TestResponsesOrphanCodexDelegationCompatibility(t *testing.T) { t.Fatalf("input.0.type = %q, want function_call_output without header; captured=%s", itemType, capturedNoHeader) } } + +func TestResponsesSanitizesOrphanToolOutputsThroughHTTPAndCompact(t *testing.T) { + gin.SetMode(gin.TestMode) + + const modelID = "test-responses-orphan-tool-output-model" + authID := "test-responses-orphan-tool-output-auth" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: modelID}}) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + coreManager := coreauth.NewManager(nil, nil, nil) + executor := &responsesMultiAgentCaptureExecutor{} + coreManager.RegisterExecutor(executor) + if _, errRegister := coreManager.Register(context.Background(), &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + baseHandler := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreManager) + responsesHandler := NewOpenAIResponsesAPIHandler(baseHandler) + + router := gin.New() + router.POST("/v1/responses", responsesHandler.Responses) + router.POST("/v1/responses/compact", responsesHandler.Compact) + + payload := fmt.Sprintf(`{ + "model": %q, + "stream": false, + "input": [ + {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}, + {"type":"function_call_output","id":"fco_1","name":"codex_apps_open_page","namespace":"agent","output":"{\"page_id\":null}"}, + {"type":"custom_tool_call_output","id":"fco_2","name":"custom_ctx","output":[{"type":"input_text","text":"part one"},{"type":"input_text","text":"part two"},{"type":"input_image","image_url":"https://example.com/page.png"}]}, + {"type":"function_call","call_id":"call_ok","name":"shell","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_ok","output":"ok"}, + {"type":"function_call_output","call_id":" ","output":"blank id"}, + {"type":"function_call_output","id":"fco_3","name":"create_thread","namespace":"codex_app","output":"task"}, + {"type":"function_call","call_id":"call_pair","name":"shell","arguments":"{}"}, + {"type":"function_call_output","output":"paired without call_id"} + ] + }`, modelID) + + for idx, path := range []string{"/v1/responses", "/v1/responses/compact"} { + req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(payload)) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("%s status = %d, want 200; body=%s", path, rec.Code, rec.Body.String()) + } + payloads := executor.Payloads() + if len(payloads) != idx+1 { + t.Fatalf("%s captured payload count = %d, want %d", path, len(payloads), idx+1) + } + input := ojson.GetBytesNoCopy(payloads[idx], "input").Array() + if len(input) != 9 { + t.Fatalf("%s input len = %d, want 9", path, len(input)) + } + if got := input[1].Get("type").String(); got != "message" { + t.Fatalf("%s input[1].type = %q, want message", path, got) + } + want1 := "Tool output from codex_apps_open_page:\n{\"page_id\":null}" + if got := input[1].Get("content.0.text").String(); got != want1 { + t.Fatalf("%s input[1].content.0.text = %q, want %q", path, got, want1) + } + if got := input[2].Get("type").String(); got != "message" { + t.Fatalf("%s input[2].type = %q, want message", path, got) + } + want2 := "Tool output from custom_ctx:\npart one\npart two" + if got := input[2].Get("content.0.text").String(); got != want2 { + t.Fatalf("%s input[2].content.0.text = %q, want %q", path, got, want2) + } + if got := input[2].Get("content.1.type").String(); got != "input_image" { + t.Fatalf("%s input[2].content.1.type = %q, want input_image", path, got) + } + if got := input[2].Get("content.1.image_url").String(); got != "https://example.com/page.png" { + t.Fatalf("%s input[2].content.1.image_url = %q, want https://example.com/page.png", path, got) + } + if got := input[4].Get("type").String(); got != "function_call_output" { + t.Fatalf("%s input[4].type = %q, want function_call_output", path, got) + } + if got := input[5].Get("type").String(); got != "message" { + t.Fatalf("%s input[5].type = %q, want message", path, got) + } + want5 := "Tool output from tool:\nblank id" + if got := input[5].Get("content.0.text").String(); got != want5 { + t.Fatalf("%s input[5].content.0.text = %q, want %q", path, got, want5) + } + if got := input[6].Get("type").String(); got != "function_call_output" { + t.Fatalf("%s input[6].type = %q, want function_call_output", path, got) + } + if got := input[8].Get("type").String(); got != "function_call_output" { + t.Fatalf("%s input[8].type = %q, want function_call_output", path, got) + } + if got := input[8].Get("call_id").String(); got != "call_pair" { + t.Fatalf("%s input[8].call_id = %q, want call_pair", path, got) + } + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_forward.go b/sdk/api/handlers/openai/openai_responses_websocket_forward.go index b5485d41e6f..0dfd7702443 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_forward.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_forward.go @@ -369,11 +369,10 @@ func completeResponsesWebsocketLocalInterrupt( cancel(errSet) return output, responseID, sortedStringSet(pendingToolCallIDs), nil, errSet } + cancel(context.Canceled) if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, incompleteDoc.Bytes(), time.Now()); errWrite != nil { - cancel(errWrite) return output, responseID, sortedStringSet(pendingToolCallIDs), nil, errWrite } - cancel(context.Canceled) return output, responseID, sortedStringSet(pendingToolCallIDs), nil, nil } diff --git a/sdk/api/handlers/openai/openai_responses_websocket_requests.go b/sdk/api/handlers/openai/openai_responses_websocket_requests.go index f958de90e78..ed967bf1de2 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_requests.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_requests.go @@ -107,15 +107,21 @@ func normalizeResponseSubsequentRequest(rawJSON []byte, lastRequest []byte, last return normalized, bytes.Clone(normalized), nil } + // A request without previous_response_id that leaves the last response's tool + // calls unanswered cannot extend that response: appending it to the stale turn + // would send those function_call items without outputs. Treat it as a + // replacement transcript in every upstream mode (e.g. a client-side compaction + // summary that carries no compaction markers). + prev := strings.TrimSpace(ojson.GetBytesNoCopy(rawJSON, "previous_response_id").String()) + if prev == "" && !inputSatisfiesPendingToolCalls(nextInput, lastResponsePendingToolCallIDs) { + normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) + return normalized, bytes.Clone(normalized), nil + } + // Websocket v2 mode uses response.create with previous_response_id + incremental input. // Do not expand it into a full input transcript; upstream expects the incremental payload. if allowIncrementalInputWithPreviousResponseID { - prev := strings.TrimSpace(ojson.GetBytesNoCopy(rawJSON, "previous_response_id").String()) if prev == "" { - if !inputSatisfiesPendingToolCalls(nextInput, lastResponsePendingToolCallIDs) { - normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) - return normalized, bytes.Clone(normalized), nil - } prev = strings.TrimSpace(lastResponseID) } if prev != "" { diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index ce0cfaca0d6..c2f60c9a085 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1587,6 +1587,52 @@ func TestNormalizeResponsesWebsocketRequestSkipsPreviousResponseIDWhenPendingOut } } +func TestNormalizeResponsesWebsocketRequestReplacesTranscriptWhenPendingOutputIsMissingWithoutIncrementalUpstream(t *testing.T) { + lastRequest := []byte(`{"model":"test-model","stream":true,"instructions":"be helpful","input":[{"type":"message","id":"msg-1"}]}`) + lastResponseOutput := []byte(`[ + {"type":"function_call","id":"fc-1","call_id":"call-1"} + ]`) + raw := []byte(`{"type":"response.create","input":[{"type":"message","role":"user","id":"summary-1","content":"compacted summary"}]}`) + + normalized, next, errMsg := normalizeResponsesWebsocketRequestWithIncrementalState(raw, lastRequest, lastResponseOutput, "resp-1", []string{"call-1"}, false, false) + if errMsg != nil { + t.Fatalf("unexpected error: %v", errMsg.Error) + } + if ojson.GetBytesNoCopy(normalized, "previous_response_id").Exists() { + t.Fatalf("previous_response_id must not be set on a transcript replacement: %s", normalized) + } + input := ojson.GetBytesNoCopy(normalized, "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "summary-1" { + t.Fatalf("replacement must not merge the stale turn with its unanswered function_call: %s", normalized) + } + if got := ojson.GetBytesNoCopy(normalized, "model").String(); got != "test-model" { + t.Fatalf("model = %q, want test-model", got) + } + if got := ojson.GetBytesNoCopy(normalized, "instructions").String(); got != "be helpful" { + t.Fatalf("instructions = %q, want be helpful", got) + } + if !bytes.Equal(next, normalized) { + t.Fatalf("next request snapshot should match normalized request") + } +} + +func TestNormalizeResponsesWebsocketRequestMergesWhenPendingOutputIsPresentWithoutIncrementalUpstream(t *testing.T) { + lastRequest := []byte(`{"model":"test-model","stream":true,"instructions":"be helpful","input":[{"type":"message","id":"msg-1"}]}`) + lastResponseOutput := []byte(`[ + {"type":"function_call","id":"fc-1","call_id":"call-1"} + ]`) + raw := []byte(`{"type":"response.create","input":[{"type":"function_call_output","call_id":"call-1","id":"tool-out-1"}]}`) + + normalized, _, errMsg := normalizeResponsesWebsocketRequestWithIncrementalState(raw, lastRequest, lastResponseOutput, "resp-1", []string{"call-1"}, false, false) + if errMsg != nil { + t.Fatalf("unexpected error: %v", errMsg.Error) + } + input := ojson.GetBytesNoCopy(normalized, "input").Array() + if len(input) != 3 || input[0].Get("id").String() != "msg-1" || input[1].Get("id").String() != "fc-1" || input[2].Get("id").String() != "tool-out-1" { + t.Fatalf("answered pending call must keep the merged transcript: %s", normalized) + } +} + func TestNormalizeResponsesWebsocketRequestReplacesCodexLocalCompactionTranscript(t *testing.T) { lastRequest := []byte(`{"model":"gpt-5.6-sol","stream":true,"instructions":"be helpful","input":[ {"type":"message","role":"user","id":"old-user","content":[{"type":"input_text","text":"old prompt"}]}, diff --git a/sdk/api/handlers/stream_forwarder.go b/sdk/api/handlers/stream_forwarder.go index fb5828fc168..23263b5370d 100644 --- a/sdk/api/handlers/stream_forwarder.go +++ b/sdk/api/handlers/stream_forwarder.go @@ -152,68 +152,79 @@ func (h *BaseAPIHandler) ForwardStream(c *gin.Context, flusher http.Flusher, can if opts.KeepAliveInterval != nil { keepAliveInterval = *opts.KeepAliveInterval } - keepAliveC, stopKeepAlive := NewStreamingKeepAliveTicker(keepAliveInterval) + keepAliveC, stopKeepAlive, resetKeepAlive := NewStreamingKeepAliveTickerWithReset(keepAliveInterval) defer stopKeepAlive() var terminalErr *interfaces.ErrorMessage - for { - select { - case <-c.Request.Context().Done(): - cancel(c.Request.Context().Err()) - return - case chunk, ok := <-data: - if !ok { - // Prefer surfacing a terminal error if one is pending. - if terminalErr == nil { - if errMsg, ok := PendingStreamError(errs); ok { - terminalErr = errMsg - if opts.NormalizeTerminalError != nil { - terminalErr = opts.NormalizeTerminalError(terminalErr) - } - } - } - if terminalErr == nil && opts.CloseError != nil { - terminalErr = opts.CloseError() - } - if terminalErr != nil { - if opts.WriteTerminalError != nil { - opts.WriteTerminalError(terminalErr) - } - if !flush() { - return + handleChunk := func(chunk []byte, ok bool) bool { + if !ok { + // Prefer surfacing a terminal error if one is pending. + if terminalErr == nil { + if errMsg, ok := PendingStreamError(errs); ok { + terminalErr = errMsg + if opts.NormalizeTerminalError != nil { + terminalErr = opts.NormalizeTerminalError(terminalErr) } - cancel(terminalErr.Error) - return } - if opts.WriteDone != nil { - opts.WriteDone() + } + if terminalErr == nil && opts.CloseError != nil { + terminalErr = opts.CloseError() + } + if terminalErr != nil { + if opts.WriteTerminalError != nil { + opts.WriteTerminalError(terminalErr) } if !flush() { - return + return true } - cancel(nil) - return + cancel(terminalErr.Error) + return true + } + if opts.WriteDone != nil { + opts.WriteDone() } - writeChunk(chunk) if !flush() { - return + return true } - if opts.ChunkError != nil { - chunkErr := opts.ChunkError() + cancel(nil) + return true + } + beforeSize := c.Writer.Size() + writeChunk(chunk) + if c.Writer.Size() > beforeSize { + resetKeepAlive() + } + if !flush() { + return true + } + if opts.ChunkError != nil { + chunkErr := opts.ChunkError() + if chunkErr != nil { + if opts.NormalizeTerminalError != nil { + chunkErr = opts.NormalizeTerminalError(chunkErr) + } if chunkErr != nil { - if opts.NormalizeTerminalError != nil { - chunkErr = opts.NormalizeTerminalError(chunkErr) - } - if chunkErr != nil { - cancel(chunkErr.Error) - } else { - cancel(nil) - } - return + cancel(chunkErr.Error) + } else { + cancel(nil) } + return true } - if opts.ChunkDone != nil && opts.ChunkDone() { - cancel(nil) + } + if opts.ChunkDone != nil && opts.ChunkDone() { + cancel(nil) + return true + } + return false + } + + for { + select { + case <-c.Request.Context().Done(): + cancel(c.Request.Context().Err()) + return + case chunk, ok := <-data: + if handleChunk(chunk, ok) { return } case errMsg, ok := <-errs: diff --git a/sdk/api/handlers/stream_forwarder_test.go b/sdk/api/handlers/stream_forwarder_test.go index f77475481fe..904e7905a0c 100644 --- a/sdk/api/handlers/stream_forwarder_test.go +++ b/sdk/api/handlers/stream_forwarder_test.go @@ -167,3 +167,53 @@ func TestSSEJSONValidationStateSplitCRLFVariants(t *testing.T) { }) } } + +func TestForwardStreamResetsKeepAliveOnEachChunk(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + + tickChan := make(chan time.Time) + data := make(chan []byte, 3) + errs := make(chan *interfaces.ErrorMessage) + close(errs) + + resetCount := 0 + restore := SetStreamingKeepAliveTickerWithResetFactoryForTest(func(time.Duration) (<-chan time.Time, func(), func()) { + return tickChan, func() {}, func() { + resetCount++ + } + }) + defer restore() + + data <- []byte("buffered-without-output") + data <- []byte("data: {\"part\":1}\n\n") + data <- []byte("data: {\"part\":2}\n\n") + close(data) + + interval := 20 * time.Second + h := &BaseAPIHandler{} + var canceled error + h.ForwardStream(c, recorder, func(err error) { + canceled = err + }, data, errs, StreamForwardOptions{ + KeepAliveInterval: &interval, + WriteChunk: func(chunk []byte) { + if string(chunk) == "buffered-without-output" { + return + } + _, _ = c.Writer.Write(chunk) + }, + }) + + if canceled != nil { + t.Fatalf("unexpected cancel error: %v", canceled) + } + if resetCount != 2 { + t.Fatalf("resetCount = %d, want 2", resetCount) + } + if got := recorder.Body.String(); got != "data: {\"part\":1}\n\ndata: {\"part\":2}\n\n" { + t.Fatalf("body = %q, want two data chunks", got) + } +} diff --git a/sdk/auth/claude.go b/sdk/auth/claude.go index 346a65bc057..6204b1624fa 100644 --- a/sdk/auth/claude.go +++ b/sdk/auth/claude.go @@ -56,7 +56,7 @@ func (a *ClaudeAuthenticator) Login(ctx context.Context, cfg *config.Config, opt return nil, fmt.Errorf("claude pkce generation failed: %w", err) } - state, err := misc.GenerateRandomState() + state, err := claude.GenerateState() if err != nil { return nil, fmt.Errorf("claude state generation failed: %w", err) } diff --git a/sdk/auth/interfaces.go b/sdk/auth/interfaces.go index 4388de8019f..646c9d57bac 100644 --- a/sdk/auth/interfaces.go +++ b/sdk/auth/interfaces.go @@ -19,6 +19,7 @@ type LoginOptions struct { CallbackPort int Metadata map[string]string Prompt func(prompt string) (string, error) + ProxyURL string } // Authenticator manages login and optional refresh flows for a provider. diff --git a/sdk/auth/manager.go b/sdk/auth/manager.go index 4cbefd98e85..62e9f57cc49 100644 --- a/sdk/auth/manager.go +++ b/sdk/auth/manager.go @@ -12,6 +12,7 @@ import ( xaiauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/xai" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/proxyutil" ) // Manager aggregates authenticators and coordinates persistence via a token store. @@ -56,6 +57,17 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config return nil, "", fmt.Errorf("cliproxy auth: authenticator %s not registered", provider) } + loginProxyURL := "" + if opts != nil { + loginProxyURL = strings.TrimSpace(opts.ProxyURL) + } + if loginProxyURL != "" { + if !proxyutil.ValidRequestProxy(loginProxyURL) { + return nil, "", fmt.Errorf("cliproxy auth: invalid login proxy URL %q", proxyutil.Redact(loginProxyURL)) + } + cfg = configWithLoginProxy(cfg, loginProxyURL) + } + record, err := auth.Login(ctx, cfg, opts) if err != nil { return nil, "", err @@ -127,6 +139,7 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config if adoptXAICredential { adoptLegacyXAIFileName(record, legacyXAICredential) } + applyLoginProxyToRecord(record, loginProxyURL) savedPath, err := m.store.Save(coreauth.WithAuthCreationIntent(ctx), record) if err != nil { @@ -159,6 +172,33 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config return record, savedPath, nil } +func configWithLoginProxy(cfg *config.Config, proxyURL string) *config.Config { + if cfg == nil { + return &config.Config{SDKConfig: config.SDKConfig{ProxyURL: proxyURL}} + } + cloned := cfg.CloneForRuntime() + if cloned == nil { + copied := *cfg + cloned = &copied + } + cloned.ProxyURL = proxyURL + return cloned +} + +func applyLoginProxyToRecord(record *coreauth.Auth, proxyURL string) { + if record == nil || proxyURL == "" { + return + } + record.ProxyURL = proxyURL + if record.Metadata == nil { + record.Metadata = make(map[string]any) + } + record.Metadata["proxy_url"] = proxyURL + if setter, ok := record.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(record.Metadata) + } +} + func metadataStringValue(metadata map[string]any, key string) string { value, _ := metadata[key].(string) return strings.TrimSpace(value) diff --git a/sdk/auth/manager_login_proxy_test.go b/sdk/auth/manager_login_proxy_test.go new file mode 100644 index 00000000000..6c031f9981b --- /dev/null +++ b/sdk/auth/manager_login_proxy_test.go @@ -0,0 +1,127 @@ +package auth + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +type proxyCapturingAuthenticator struct { + provider string + record *coreauth.Auth + seenCfg *config.Config +} + +func (a *proxyCapturingAuthenticator) Provider() string { + return a.provider +} + +func (a *proxyCapturingAuthenticator) Login(_ context.Context, cfg *config.Config, _ *LoginOptions) (*coreauth.Auth, error) { + a.seenCfg = cfg + return a.record, nil +} + +func (a *proxyCapturingAuthenticator) RefreshLead() *time.Duration { + return nil +} + +func TestManagerLoginWithProxyOverridesFlowAndPersists(t *testing.T) { + authDir := t.TempDir() + fileName := "demo-proxy.json" + proxyURL := "socks5://user:pass@127.0.0.1:1080" + + authenticator := &proxyCapturingAuthenticator{ + provider: "demo", + record: &coreauth.Auth{ + ID: fileName, + FileName: fileName, + Provider: "demo", + Metadata: map[string]any{"type": "demo", "email": "user@example.com"}, + }, + } + store := NewFileTokenStore() + store.SetBaseDir(authDir) + mgr := NewManager(store, authenticator) + cfg := &config.Config{AuthDir: authDir, SDKConfig: config.SDKConfig{ProxyURL: "http://global.example:8080"}} + + record, _, err := mgr.Login(context.Background(), "demo", cfg, &LoginOptions{ProxyURL: proxyURL}) + if err != nil { + t.Fatalf("Login error: %v", err) + } + if authenticator.seenCfg == nil || authenticator.seenCfg.ProxyURL != proxyURL { + t.Fatalf("authenticator cfg proxy = %v, want %q", authenticator.seenCfg, proxyURL) + } + if cfg.ProxyURL != "http://global.example:8080" { + t.Fatalf("caller cfg mutated: %q", cfg.ProxyURL) + } + if record.ProxyURL != proxyURL { + t.Fatalf("record.ProxyURL = %q, want %q", record.ProxyURL, proxyURL) + } + + raw, errRead := os.ReadFile(filepath.Join(authDir, fileName)) + if errRead != nil { + t.Fatalf("read saved auth: %v", errRead) + } + var saved map[string]any + if errUnmarshal := json.Unmarshal(raw, &saved); errUnmarshal != nil { + t.Fatalf("unmarshal saved auth: %v", errUnmarshal) + } + if saved["proxy_url"] != proxyURL { + t.Fatalf("saved proxy_url = %v, want %q", saved["proxy_url"], proxyURL) + } +} + +func TestManagerLoginWithoutProxyLeavesRecordUntouched(t *testing.T) { + authDir := t.TempDir() + fileName := "demo-noproxy.json" + authenticator := &proxyCapturingAuthenticator{ + provider: "demo", + record: &coreauth.Auth{ + ID: fileName, + FileName: fileName, + Provider: "demo", + Metadata: map[string]any{"type": "demo"}, + }, + } + store := NewFileTokenStore() + store.SetBaseDir(authDir) + mgr := NewManager(store, authenticator) + cfg := &config.Config{AuthDir: authDir} + + if _, _, err := mgr.Login(context.Background(), "demo", cfg, &LoginOptions{}); err != nil { + t.Fatalf("Login error: %v", err) + } + if authenticator.seenCfg != cfg { + t.Fatalf("authenticator received a different cfg without a login proxy") + } + raw, errRead := os.ReadFile(filepath.Join(authDir, fileName)) + if errRead != nil { + t.Fatalf("read saved auth: %v", errRead) + } + var saved map[string]any + if errUnmarshal := json.Unmarshal(raw, &saved); errUnmarshal != nil { + t.Fatalf("unmarshal saved auth: %v", errUnmarshal) + } + if _, ok := saved["proxy_url"]; ok { + t.Fatalf("saved auth unexpectedly contains proxy_url: %v", saved["proxy_url"]) + } +} + +func TestManagerLoginRejectsInvalidProxy(t *testing.T) { + authenticator := &proxyCapturingAuthenticator{provider: "demo", record: &coreauth.Auth{ID: "x.json", Provider: "demo"}} + mgr := NewManager(NewFileTokenStore(), authenticator) + for _, bad := range []string{"ftp://127.0.0.1:21", "direct", "not a url", "http://127.0.0.1:99999"} { + if _, _, err := mgr.Login(context.Background(), "demo", &config.Config{AuthDir: t.TempDir()}, &LoginOptions{ProxyURL: bad}); err == nil { + t.Fatalf("proxy %q: expected error", bad) + } + } + if authenticator.seenCfg != nil { + t.Fatalf("authenticator ran despite invalid proxy") + } +} diff --git a/sdk/auth/refresh_journal.go b/sdk/auth/refresh_journal.go new file mode 100644 index 00000000000..2e29d3cc339 --- /dev/null +++ b/sdk/auth/refresh_journal.go @@ -0,0 +1,38 @@ +package auth + +import ( + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +// RefreshJournalFor snapshots the authoritative directory for this refresh. A +// FileTokenStore may use explicit credential paths without configuring baseDir. +func (s *FileTokenStore) RefreshJournalFor(auth *coreauth.Auth) (codex.RefreshJournal, error) { + if s == nil { + return nil, fmt.Errorf("auth filestore: nil refresh journal store") + } + dir := s.baseDirSnapshot() + if dir == "" { + path, errPath := s.resolveAuthPath(auth) + if errPath != nil { + return nil, errPath + } + canonical, errCanonical := filepath.EvalSymlinks(path) + if errCanonical == nil { + path = canonical + } else if !errors.Is(errCanonical, os.ErrNotExist) { + return nil, errCanonical + } + dir = filepath.Dir(path) + } + absolute, errAbs := filepath.Abs(dir) + if errAbs != nil { + return nil, errAbs + } + return &codex.FileRefreshJournal{Dir: absolute}, nil +} diff --git a/sdk/cliproxy/auth/codex_refresh_selection.go b/sdk/cliproxy/auth/codex_refresh_selection.go new file mode 100644 index 00000000000..b9fdf4281c1 --- /dev/null +++ b/sdk/cliproxy/auth/codex_refresh_selection.go @@ -0,0 +1,52 @@ +package auth + +import ( + "context" + "net/http" + "strings" + "time" +) + +// refreshAccessError keeps an expired access token blocked without turning an +// uncertain refresh exchange into evidence of an upstream access-token rejection. +// wasTerminalUnauthorized describes the state before applying this refresh result. +func refreshAccessError(auth *Auth, message string, accessTokenRejected, wasTerminalUnauthorized bool) *Error { + if auth != nil && strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") && + !accessTokenRejected && !wasTerminalUnauthorized && + (auth.RejectedAccessToken == "" || auth.RejectedAccessToken != authAccessToken(auth)) && + auth.RefreshError != nil && auth.RefreshError.Code == "refresh_outcome_unknown" { + outcome := cloneError(auth.RefreshError) + if outcome.HTTPStatus == http.StatusUnauthorized { + // A diagnostic 401 must not turn this recovery state into access rejection. + outcome.HTTPStatus = http.StatusServiceUnavailable + } + return outcome + } + return &Error{Code: "unauthorized", Message: message, HTTPStatus: http.StatusUnauthorized} +} + +// codexRefreshMissingAccessForRequest keeps model-level availability from +// admitting a missing token while its grant still requires recovery. +func codexRefreshMissingAccessForRequest(ctx context.Context, auth *Auth, headers ...http.Header) bool { + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") || !auth.RequiresSignIn || + auth.RefreshError == nil || auth.RefreshError.Code != "refresh_outcome_unknown" { + return false + } + return authAccessTokenForRequest(ctx, auth, headers...) == "" +} + +// codexRefreshUnavailableForRequest lets a request use a healthy outer token +// while SIWC requires recovery. Ordinary cooldowns and model failures still apply. +func codexRefreshUnavailableForRequest(ctx context.Context, auth *Auth, now time.Time, headers ...http.Header) bool { + if !auth.Unavailable || !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") || + auth.LastError == nil || auth.LastError.Code != "refresh_outcome_unknown" || !auth.RequiresSignIn || + !IsCodexUseSIWC(auth) || IsCodexUseSIWCForRequest(ctx, auth, headers...) || auth.NextRetryAfter.After(now) || + auth.RejectedSIWCRefreshToken == "" || auth.RejectedSIWCRefreshToken != authSIWCTokenField(auth, "refresh_token") { + return auth.Unavailable + } + if authOuterAccessToken(auth) == "" { + return true + } + expiry, hasExpiry := auth.OuterExpirationTime() + return hasExpiry && !expiry.After(now) +} diff --git a/sdk/cliproxy/auth/codex_refresh_selection_test.go b/sdk/cliproxy/auth/codex_refresh_selection_test.go new file mode 100644 index 00000000000..fb0ee8df7ee --- /dev/null +++ b/sdk/cliproxy/auth/codex_refresh_selection_test.go @@ -0,0 +1,258 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestCodexUnknownSIWCRefreshKeepsHealthyOuterAccess(t *testing.T) { + now := time.Now() + manager := NewManager(nil, nil, nil) + manager.RegisterExecutor(&unauthorizedRefreshExecutor{id: "codex", refreshErr: codex.WrapSIWCRefreshError(&codex.RefreshOutcomeUnknownError{Stage: "transport"})}) + auth := &Auth{ID: t.Name(), Provider: "codex", Status: StatusActive, Metadata: map[string]any{ + "access_token": "outer", "refresh_token": "outer-refresh", "expired": now.Add(time.Hour).Format(time.RFC3339), "use_siwc": true, + "siwc": map[string]any{"access_token": "siwc", "refresh_token": "siwc-refresh", "client_id": "client", "expires_at": now.Add(-time.Hour).Format(time.RFC3339)}, + }} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); !codex.IsRefreshOutcomeUnknown(errRefresh) { + t.Fatalf("refresh outcome = %v, want unknown", errRefresh) + } + base, ok := manager.GetByID(auth.ID) + if !ok || !base.RequiresSignIn || !base.Unavailable || base.RejectedSIWCRefreshToken != "siwc-refresh" { + t.Fatalf("SIWC recovery state was not retained: %#v", base) + } + for _, tc := range []struct { + name string + native bool + mutate func(*Auth) + blocked bool + }{ + {name: "healthy outer", native: true}, + {name: "expired SIWC", blocked: true}, + {name: "disabled", native: true, blocked: true, mutate: func(auth *Auth) { auth.Disabled = true }}, + {name: "expired outer", native: true, blocked: true, mutate: func(auth *Auth) { auth.Metadata["expired"] = now.Add(-time.Hour).Format(time.RFC3339) }}, + {name: "rejected outer", native: true, blocked: true, mutate: func(auth *Auth) { auth.RejectedAccessToken = "outer" }}, + {name: "other cooldown", native: true, blocked: true, mutate: func(auth *Auth) { auth.NextRetryAfter = now.Add(time.Hour) }}, + {name: "credential quota", native: true, blocked: true, mutate: func(auth *Auth) { + auth.Quota = QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(time.Hour)} + }}, + {name: "model cooldown", native: true, blocked: true, mutate: func(auth *Auth) { + auth.ModelStates = map[string]*ModelState{"model": {Status: StatusError, Unavailable: true, NextRetryAfter: now.Add(time.Hour)}} + }}, + {name: "unrelated error", native: true, blocked: true, mutate: func(auth *Auth) { auth.LastError = &Error{Code: "unrelated"} }}, + {name: "explicit unauthorized", native: true, blocked: true, mutate: func(auth *Auth) { auth.LastError = &Error{Code: "unauthorized", HTTPStatus: http.StatusUnauthorized} }}, + {name: "no SIWC fence", native: true, blocked: true, mutate: func(auth *Auth) { auth.RejectedSIWCRefreshToken = "" }}, + {name: "other provider", native: true, blocked: true, mutate: func(auth *Auth) { auth.Provider = "claude" }}, + } { + t.Run(tc.name, func(t *testing.T) { + current := base.Clone() + if tc.mutate != nil { + tc.mutate(current) + } + var headers http.Header + if tc.native { + headers = http.Header{"User-Agent": {"codex_cli_rs/1.0.0"}} + } + blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now, headers) + if blocked != tc.blocked { + t.Fatalf("blocked = %v, want %v", blocked, tc.blocked) + } + }) + } +} + +func TestCodexUnknownSIWCRefreshCacheKeepsHealthyOuterAccess(t *testing.T) { + now := time.Now() + manager := NewManager(nil, nil, nil) + executor := &unauthorizedRefreshExecutor{id: "codex", refreshErr: codex.WrapSIWCRefreshError(&codex.RefreshOutcomeUnknownError{Stage: "transport"})} + manager.RegisterExecutor(executor) + auth := &Auth{ID: t.Name(), Provider: "codex", Status: StatusActive, Metadata: map[string]any{ + "access_token": "outer", "expired": now.Add(time.Hour).Format(time.RFC3339), "use_siwc": true, + "siwc": map[string]any{"access_token": "siwc", "refresh_token": "siwc-refresh", "client_id": "client", "expires_at": now.Add(-time.Hour).Format(time.RFC3339)}, + }} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); !codex.IsRefreshOutcomeUnknown(errRefresh) { + t.Fatalf("refresh outcome = %v, want unknown", errRefresh) + } + manager.refreshAuth(context.Background(), auth.ID) + current, _ := manager.GetByID(auth.ID) + if executor.RefreshCalls() != 1 || current.LastError == nil || current.LastError.Code != "refresh_outcome_unknown" { + t.Fatal("cached recovery state lost its classification or resubmitted the grant") + } + headers := http.Header{"User-Agent": {"codex_cli_rs/1.0.0"}} + if blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now, headers); blocked { + t.Fatal("cached SIWC recovery state blocked healthy outer access") + } +} + +type codexPartialRefreshExecutor struct { + unauthorizedRefreshExecutor + refresh func(*Auth) (*Auth, error) +} + +func (e *codexPartialRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return e.refresh(auth) +} + +func TestCodexPartialRefreshKeepsUnknownSelectedPartitionBlocked(t *testing.T) { + for _, useSIWC := range []bool{false, true} { + name := "outer" + if useSIWC { + name = "siwc" + } + t.Run(name, func(t *testing.T) { + now := time.Now() + auth := &Auth{ID: t.Name(), Provider: "codex", Status: StatusActive, Metadata: map[string]any{ + "access_token": "outer", "refresh_token": "outer-refresh", "expired": now.Add(-time.Hour).Format(time.RFC3339), "use_siwc": useSIWC, + "siwc": map[string]any{"access_token": "siwc", "refresh_token": "siwc-refresh", "client_id": "client", "expires_at": now.Add(-time.Hour).Format(time.RFC3339)}, + }} + manager := NewManager(nil, nil, nil) + manager.RegisterExecutor(&codexPartialRefreshExecutor{ + unauthorizedRefreshExecutor: unauthorizedRefreshExecutor{id: "codex"}, + refresh: func(before *Auth) (*Auth, error) { + updated := before.Clone() + var outerErr, siwcErr error + outcome := &codex.RefreshOutcomeUnknownError{Stage: "transport"} + if useSIWC { + updated.Metadata["access_token"] = "outer-new" + updated.Metadata["expired"] = now.Add(time.Hour).Format(time.RFC3339) + siwcErr = codex.WrapSIWCRefreshError(outcome) + } else { + siwc := updated.Metadata["siwc"].(map[string]any) + siwc["access_token"] = "siwc-new" + siwc["expires_at"] = now.Add(time.Hour).Format(time.RFC3339) + outerErr = outcome + } + updated.CodexRefreshPatch = codex.NewCredentialRefreshPatch(before.Metadata, updated.Metadata, useSIWC, !useSIWC) + updated.CodexRefreshPatch.SetPartitionErrors(outerErr, siwcErr) + return updated, nil + }, + }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); errRefresh != nil { + t.Fatal(errRefresh) + } + current, _ := manager.GetByID(auth.ID) + if !current.Unavailable || !current.RequiresSignIn || current.LastError == nil || current.LastError.Code != "refresh_outcome_unknown" { + t.Fatalf("partial refresh lost the unresolved selected partition: unavailable=%v sign_in=%v last_error=%#v refresh_error=%#v", current.Unavailable, current.RequiresSignIn, current.LastError, current.RefreshError) + } + if blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now); !blocked { + t.Fatal("partial refresh restored an expired selected partition") + } + headers := http.Header{"User-Agent": {"codex_cli_rs/1.0.0"}} + if blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now, headers); blocked == useSIWC { + t.Fatal("native request availability does not match its outer access token") + } + }) + } +} + +func TestCodexMissingRecoveryAccessBlocksBeforeHealthyModelState(t *testing.T) { + now := time.Now() + for _, useSIWC := range []bool{false, true} { + for _, unavailable := range []bool{false, true} { + auth := &Auth{Provider: "codex", Status: StatusError, RequiresSignIn: true, Unavailable: unavailable, + RefreshError: &Error{Code: "refresh_outcome_unknown"}, + Metadata: map[string]any{ + "refresh_token": "outer-refresh", "use_siwc": useSIWC, + "siwc": map[string]any{"refresh_token": "siwc-refresh", "client_id": "client"}, + }, + ModelStates: map[string]*ModelState{"model": {Status: StatusActive}}, + } + if useSIWC { + auth.Metadata["access_token"] = "outer-valid" + auth.Metadata["expired"] = now.Add(time.Hour).Format(time.RFC3339) + auth.RejectedSIWCRefreshToken = "siwc-refresh" + } else { + siwc := auth.Metadata["siwc"].(map[string]any) + siwc["access_token"] = "siwc-valid" + siwc["expires_at"] = now.Add(time.Hour).Format(time.RFC3339) + auth.RejectedRefreshToken = "outer-refresh" + } + blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), auth, "model", now) + if !blocked || !isAuthBlockedByQuotaOrAuthForModelForRequest(context.Background(), auth, "model", now) { + t.Fatalf("healthy model state admitted a missing token: use_siwc=%v unavailable=%v", useSIWC, unavailable) + } + } + } +} + +func TestCodexUnknownDiagnostic401UsesPriorUnauthorizedState(t *testing.T) { + for _, mode := range []string{"refresh", "cached"} { + for _, tokenState := range []string{"expired", "missing"} { + for _, priorUnauthorized := range []bool{false, true} { + priorState := "active" + if priorUnauthorized { + priorState = "unauthorized" + } + t.Run(mode+"/"+tokenState+"/"+priorState, func(t *testing.T) { + now := time.Now() + outcome := &codex.RefreshOutcomeUnknownError{Stage: "transport", Cause: oauthStatusError{code: http.StatusUnauthorized, msg: "transport diagnostic"}} + executor := &unauthorizedRefreshExecutor{id: "codex", refreshErr: codex.WrapSIWCRefreshError(outcome)} + manager := NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + siwc := map[string]any{"refresh_token": "siwc-refresh", "client_id": "client"} + if tokenState == "expired" { + siwc["access_token"] = "siwc-expired" + siwc["expires_at"] = now.Add(-time.Hour).Format(time.RFC3339) + } + auth := &Auth{ID: t.Name(), Provider: "codex", Status: StatusActive, Metadata: map[string]any{ + "access_token": "outer-valid", "expired": now.Add(time.Hour).Format(time.RFC3339), "use_siwc": true, "siwc": siwc, + }} + if mode == "cached" { + auth.RequiresSignIn = true + auth.RejectedSIWCRefreshToken = "siwc-refresh" + auth.RefreshError = &Error{Code: "refresh_outcome_unknown", Message: outcome.Error(), HTTPStatus: http.StatusUnauthorized} + auth.LastError = &Error{Code: "unauthorized", Message: outcome.Error(), HTTPStatus: http.StatusUnauthorized} + } + if priorUnauthorized { + auth.Unavailable = true + auth.Status = StatusError + auth.LastError = &Error{Code: "unauthorized", Message: "access token rejected", HTTPStatus: http.StatusUnauthorized} + } + registered, errRegister := manager.Register(context.Background(), auth) + if errRegister != nil || hasUnauthorizedAuthFailure(registered) != priorUnauthorized { + t.Fatalf("incorrect prior state: err=%v terminal=%v", errRegister, hasUnauthorizedAuthFailure(registered)) + } + wantCalls := 0 + if mode == "cached" { + manager.refreshAuth(context.Background(), auth.ID) + } else { + wantCalls = 1 + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); !codex.IsRefreshOutcomeUnknown(errRefresh) { + t.Fatalf("refresh error = %v, want unknown", errRefresh) + } + } + current, _ := manager.GetByID(auth.ID) + if executor.RefreshCalls() != wantCalls || current.RefreshError == nil || current.RefreshError.Code != "refresh_outcome_unknown" || !current.RequiresSignIn { + t.Fatalf("recovery state or replay protection changed: calls=%d refresh_error=%#v", executor.RefreshCalls(), current.RefreshError) + } + if priorUnauthorized { + if !hasUnauthorizedAuthFailure(current) || current.LastError.Code != "unauthorized" { + t.Fatal("unknown refresh cleared the prior terminal rejection") + } + } else if current.LastError == nil || current.LastError.Code != "refresh_outcome_unknown" || current.LastError.HTTPStatus == http.StatusUnauthorized || hasUnauthorizedAuthFailure(current) { + t.Fatalf("diagnostic 401 created a terminal rejection: %#v", current.LastError) + } + headers := http.Header{"User-Agent": {"codex_cli_rs/1.0.0"}} + if blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now, headers); blocked != priorUnauthorized { + t.Fatalf("native outer blocked=%v, want %v", blocked, priorUnauthorized) + } + if blocked, _, _ := isAuthBlockedForModelForRequest(context.Background(), current, "model", now); !blocked { + t.Fatal("unusable SIWC access was admitted") + } + }) + } + } + } +} diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index ee6be67acc8..8347f38ba28 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -938,12 +938,18 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } now = time.Now() responseHeaders := internallogging.GetResponseHeaders(ctx) - modelState := existingModelState(auth, modelKey) var cooldownRecordsBefore []CooldownStateRecord trackCooldownState := m.cooldownStore != nil if trackCooldownState { cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) } + if m != nil { + _ = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) + if modelKey != "" { + modelKey = m.resolveMarkResultModelKeyLocked(auth, modelKey, result.RouteModel) + } + } + modelState := existingModelState(auth, modelKey) auth.recordRecentRequest(now, result.Success) if result.Success { auth.Success++ @@ -1082,7 +1088,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } else { state.NextRetryAfter = now.Add(30 * time.Minute) } - } else if isTransientRateLimitResultError(result.Error) && statusCode == http.StatusTooManyRequests { + } else if !result.CredentialScope && isTransientRateLimitResultError(result.Error) && statusCode == http.StatusTooManyRequests { // Transient per-turn capacity rejection (issue 7017): cool // only the reporting model until the provider retry hint. // The credential quota is not exhausted, so no quota state @@ -1384,7 +1390,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { var targetModels []string if !credentialScopedResult && modelKey != "" { targetModels = append(targetModels, modelKey) - if routeKey := canonicalModelKey(result.RouteModel); routeKey != "" && routeKey != modelKey { + if routeKey := canonicalModelKey(result.RouteModel); routeKey != "" && !sameCanonicalModelKey(routeKey, modelKey) { targetModels = append(targetModels, routeKey) } } @@ -1504,10 +1510,28 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re func existingModelState(auth *Auth, model string) *ModelState { model = canonicalModelKey(model) - if auth == nil || model == "" { + if auth == nil || model == "" || len(auth.ModelStates) == 0 { + return nil + } + if state, ok := auth.ModelStates[model]; ok && state != nil { + return state + } + lookupKey := canonicalModelLookupKey(model) + if lookupKey == "" { return nil } - return auth.ModelStates[model] + var matched *ModelState + for stateModel, state := range auth.ModelStates { + if state == nil || canonicalModelLookupKey(stateModel) != lookupKey { + continue + } + if matched == nil { + matched = state + } else { + matched = mergeModelState(matched.Clone(), state) + } + } + return matched } func ensureModelState(auth *Auth, model string) *ModelState { @@ -1522,6 +1546,12 @@ func ensureModelState(auth *Auth, model string) *ModelState { if state, ok := auth.ModelStates[model]; ok && state != nil { return state } + lookupKey := canonicalModelLookupKey(model) + for stateModel, state := range auth.ModelStates { + if canonicalModelLookupKey(stateModel) == lookupKey && state != nil { + return state + } + } state := &ModelState{Status: StatusActive} auth.ModelStates[model] = state return state @@ -1531,9 +1561,18 @@ func normalizeModelStates(auth *Auth) bool { if auth == nil || len(auth.ModelStates) == 0 { return false } + keys := make([]string, 0, len(auth.ModelStates)) + for model := range auth.ModelStates { + keys = append(keys, model) + } + if len(keys) > 1 { + sort.Strings(keys) + } normalized := make(map[string]*ModelState, len(auth.ModelStates)) + keyByLookup := make(map[string]string, len(auth.ModelStates)) changed := false - for model, state := range auth.ModelStates { + for _, model := range keys { + state := auth.ModelStates[model] modelKey := canonicalModelKey(model) if modelKey == "" { modelKey = strings.TrimSpace(model) @@ -1541,11 +1580,27 @@ func normalizeModelStates(auth *Auth) bool { if modelKey != model { changed = true } - if existing, ok := normalized[modelKey]; ok { - normalized[modelKey] = mergeModelState(existing, state) + lookupKey := canonicalModelLookupKey(modelKey) + if lookupKey == "" { + lookupKey = strings.ToLower(modelKey) + } + if existingKey, ok := keyByLookup[lookupKey]; ok { + var existingClone *ModelState + if normalized[existingKey] != nil { + existingClone = normalized[existingKey].Clone() + } + merged := mergeModelState(existingClone, state) + targetKey := existingKey + if modelKey == lookupKey && existingKey != lookupKey { + delete(normalized, existingKey) + targetKey = modelKey + keyByLookup[lookupKey] = targetKey + } + normalized[targetKey] = merged changed = true continue } + keyByLookup[lookupKey] = modelKey normalized[modelKey] = state } if changed { @@ -1688,7 +1743,7 @@ func (m *Manager) registryModelsForAuthAndModel(auth *Auth, authID, modelKey str return nil } - targetCanonical := canonicalModelKey(modelKey) + targetCanonical := canonicalModelLookupKey(modelKey) supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) if len(supportedModels) == 0 { return nil @@ -1709,7 +1764,7 @@ func (m *Manager) registryModelsForAuthAndModel(auth *Auth, authID, modelKey str if stateKey == "" { stateKey = canonicalModelKey(regModelID) } - if stateKey == targetCanonical { + if canonicalModelLookupKey(stateKey) == targetCanonical { if _, exists := seen[regModelID]; !exists { seen[regModelID] = struct{}{} out = append(out, regModelID) @@ -1945,12 +2000,12 @@ func errorCodeFromError(err error) string { } func isRequestScopedError(err error) bool { - if err == nil { + if err == nil || isCredentialScopedError(err) { return false } // A preserved terminal response belongs to this request unless the executor // explicitly classifies the rejection as safe to retry on another credential. - if isRequestTerminatedError(err) && !isCredentialScopedError(err) { + if isRequestTerminatedError(err) { return true } requestErr, ok := errors.AsType[cliproxyexecutor.RequestScopedError](err) @@ -1961,7 +2016,7 @@ func isRequestScopedError(err error) bool { // upstream streams that ended without a terminal frame. Unlike a received // HTTP refusal, the generation was already submitted upstream. func isStreamTruncatedError(err error) bool { - if err == nil { + if err == nil || isCredentialScopedError(err) { return false } truncatedErr, ok := errors.AsType[cliproxyexecutor.StreamTruncatedError](err) @@ -1997,6 +2052,18 @@ func resultErrorFromError(err error) *Error { resultErr.HTTPStatus = statusCodeFromError(err) } switch { + case isCredentialScopedError(err): + // A credential-scoped failure always stays on the credential quota path; + // normalize conflicting model-only or cooldown-skipping codes so MarkResult + // cannot downgrade whole-credential exhaustion. + switch strings.ToLower(strings.TrimSpace(resultErr.Code)) { + case transientRateLimitErrorCode, + requestScopedErrorCode, + streamTruncatedErrorCode, + connectionLifecycleErrorCode, + transientTransportErrorCode: + resultErr.Code = credentialQuotaErrorCode + } case isStreamTruncatedError(err): // Preserve the local truncation classification alongside the public // 502 status so MarkResult can skip credential cooldown. @@ -2291,7 +2358,8 @@ var ( func isKnownSafeOAuthErrorCode(code string) bool { switch strings.ToLower(strings.TrimSpace(code)) { - case "refresh_token_reused", + case "refresh_outcome_unknown", + "refresh_token_reused", "refresh_token_expired", "refresh_token_invalidated", "token_revoked", @@ -2351,6 +2419,7 @@ func sanitizeRenewalMessage(raw string) string { func isTerminalRefreshGrantMessage(message string) bool { lower := strings.ToLower(message) for _, token := range []string{ + "refresh_outcome_unknown", "refresh_token_reused", "refresh_token_expired", "refresh_token_invalidated", @@ -2392,6 +2461,7 @@ func isTerminalRefreshError(auth *Auth, err error) bool { func terminalRefreshErrorCode(message string, statusCode int) string { lower := strings.ToLower(message) for _, code := range []string{ + "refresh_outcome_unknown", "refresh_token_reused", "refresh_token_expired", "refresh_token_invalidated", @@ -2546,7 +2616,10 @@ func isCredentialScopedError(err error) bool { IsCredentialScoped() bool } var csp credentialScopedProvider - return errors.As(err, &csp) && csp != nil && csp.IsCredentialScoped() + if errors.As(err, &csp) && csp != nil && csp.IsCredentialScoped() { + return true + } + return strings.EqualFold(errorCodeFromError(err), credentialQuotaErrorCode) } func statusCodeFromResult(err *Error) int { diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index df107aaeb28..b1a8e8c626f 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -97,7 +97,7 @@ func unwrapExecutionBoundaryError(err error) error { } func preferredExecutionAttemptError(fallback, upstream error) error { - if errors.Is(fallback, context.Canceled) || errors.Is(fallback, context.DeadlineExceeded) { + if errors.Is(fallback, context.Canceled) || errors.Is(fallback, context.DeadlineExceeded) || isSessionAffinityRetryError(fallback) { return fallback } if upstream == nil { @@ -118,6 +118,46 @@ func preferredExecutionAttemptError(fallback, upstream error) error { return markUpstreamExecutionAttempt(upstream) } +func (m *Manager) normalizeSessionAffinityNeverTerminalError(err error, providers []string, model string, opts cliproxyexecutor.Options) error { + if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + if isRequestTerminatedError(err) || isRequestStopError(err) || isRequestInvalidError(err) || isSessionAffinityRetryError(err) { + return err + } + var authErr *Error + if errors.As(err, &authErr) && authErr != nil { + switch authErr.Code { + case "session_auth_unavailable", "auth_not_found", "auth_unavailable", "model_cooldown", "provider_not_found", "executor_not_found": + return err + } + } + var cooldownErr *modelCooldownError + if errors.As(err, &cooldownErr) && cooldownErr != nil { + return err + } + affinity, ok := m.Selector().(*SessionAffinitySelector) + if !ok || affinity == nil || !affinity.rebindOnNever { + return err + } + if _, okBound := affinity.boundAuthForRequest(providers, model, opts); !okBound { + return err + } + rerr := resultErrorFromError(err) + if shouldRebindSessionAffinityOnError(rerr) { + return err + } + if shouldSkipCredentialCooldown(rerr) && !isTransientTransportResultError(rerr) && !isConnectionLifecycleResultError(rerr) { + return err + } + return &Error{ + Code: "session_auth_retry", + Message: "bound session credential is transiently unavailable", + Retryable: true, + HTTPStatus: http.StatusServiceUnavailable, + } +} + // Execute performs a non-streaming execution using the configured selector and executor. // It supports multiple providers for the same model and round-robins the starting provider per model. func isOpenAIImageExecutionOptions(opts cliproxyexecutor.Options) bool { @@ -173,6 +213,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 resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { return cliproxyexecutor.Response{}, errCredits @@ -234,7 +275,9 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip } } lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) - return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(lastErr) + lastErr = unwrapExecutionBoundaryError(lastErr) + lastErr = m.normalizeSessionAffinityNeverTerminalError(lastErr, normalized, retryModel, opts) + return cliproxyexecutor.Response{}, lastErr } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -314,6 +357,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) } lastErr = unwrapExecutionBoundaryError(lastErr) + lastErr = m.normalizeSessionAffinityNeverTerminalError(lastErr, normalized, retryModel, opts) if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { return nil, errCredits @@ -1112,6 +1156,11 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) + if selection != nil { + execCtx = coreusage.WithHomeSelection(execCtx, coreusage.HomeSelection{ + Reused: selection.Retained(), Model: routeModel, CredentialPolicy: credentialPolicyFromContext(ctx), + }) + } execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.executionModelCandidatesWithAlias(auth, routeModel) // Home owns dispatched credential availability; only local selections need state filtering. @@ -2288,7 +2337,9 @@ func (m *Manager) VerifyInference(ctx context.Context, id, model string) (*Auth, } sessionID := fmt.Sprintf("verify-%s-%d", auth.ID, time.Now().UTC().UnixNano()) - payload := []byte(fmt.Sprintf(`{"model":%q,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"ping"}]}],"max_output_tokens":1,"store":false}`, verifyModel)) + // Codex upstreams reject integer max_output_tokens below 16, so the minimal + // verification budget must sit at that floor. + payload := []byte(fmt.Sprintf(`{"model":%q,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"ping"}]}],"max_output_tokens":16,"store":false}`, verifyModel)) req := cliproxyexecutor.Request{ Model: verifyModel, Payload: payload, diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index d6474210003..26e85115fa5 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -281,8 +281,16 @@ func isHomeRequestRetryExceededError(err error) bool { return strings.EqualFold(strings.TrimSpace(authErr.Code), homeRequestRetryExceededErrorCode) } +func isSessionAffinityRetryError(err error) bool { + if err == nil { + return false + } + var authErr *Error + return errors.As(err, &authErr) && authErr != nil && authErr.Code == "session_auth_retry" +} + func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick error) bool { - if lastErr == nil || isRequestStopError(errPick) { + if lastErr == nil || isRequestStopError(errPick) || isSessionAffinityRetryError(errPick) { return false } if !homeMode { @@ -1076,11 +1084,17 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o } dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) if opts.Metadata != nil { - if nodeKind, ok := opts.Metadata[cliproxyexecutor.NodeKindMetadataKey].(string); ok && strings.TrimSpace(nodeKind) != "" { + nodeKind, _ := opts.Metadata[cliproxyexecutor.NodeKindMetadataKey].(string) + if isFork, _ := opts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); isFork { + nodeKind = "fork" + } + if nodeKind = strings.TrimSpace(nodeKind); nodeKind != "" { + // Keep the caller headers unchanged when adding Home-only lineage metadata. + dispatchHeaders = dispatchHeaders.Clone() if dispatchHeaders == nil { dispatchHeaders = make(http.Header) } - dispatchHeaders.Set("X-Node-Kind", strings.TrimSpace(nodeKind)) + dispatchHeaders.Set("X-Node-Kind", nodeKind) } } credentialPolicy := credentialPolicyFromContext(ctx) @@ -1340,8 +1354,18 @@ func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, r return nil, nil } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + selectionModel := authSelectionModelFromOptions(opts, routeModel) var candidates []creditsCandidateEntry m.mu.RLock() + if affinity, ok := m.selector.(*SessionAffinitySelector); ok && affinity != nil && affinity.rebindOnNever { + if boundAuthID, okBound := affinity.boundAuthForRequest([]string{"antigravity"}, selectionModel, opts); okBound { + if pinnedAuthID != "" && pinnedAuthID != boundAuthID { + m.mu.RUnlock() + return nil, nil + } + pinnedAuthID = boundAuthID + } + } for _, auth := range m.auths { if auth == nil || auth.Disabled || auth.Status == StatusDisabled { continue diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index f30d811fd92..8bccc582b0d 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -135,6 +135,9 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c clearSchedulerSelectionID(opts.Metadata) execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) execCtx = coreusage.WithSelectionID(execCtx, "") + execCtx = coreusage.WithHomeSelection(execCtx, coreusage.HomeSelection{ + Reused: selection.Retained(), Model: routeModel, CredentialPolicy: credentialPolicyFromContext(ctx), + }) execCtx = newUpstreamAttemptContext(execCtx) if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index 117c15c0d8f..b6bc8a90d6e 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "math" - "net/http" "strings" "sync" "time" @@ -135,6 +134,7 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { auth.CredentialVersion = 1 } auth.Generation = 1 + aliasTableBeforeSave := m.oauthModelAlias.Load() // Serialize this credential, but release the manager lock during store I/O. // Persist failures stay non-fatal, but must not be silent: a restart would // lose the credential that only exists in memory. @@ -145,6 +145,12 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { m.mu.Unlock() return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID) } + if m.oauthModelAlias.Load() != aliasTableBeforeSave { + cooldownStateChanged = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) || cooldownStateChanged + if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } + } auth.ClearFileSnapshot() authClone := auth.Clone() m.auths[auth.ID] = authClone @@ -478,7 +484,10 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo if existing.LastError != nil { auth.LastError = cloneError(existing.LastError) } else if auth.RefreshError != nil { - auth.LastError = &Error{Code: "unauthorized", Message: auth.RefreshError.Message, HTTPStatus: http.StatusUnauthorized} + auth.LastError = refreshAccessError(auth, auth.RefreshError.Message, false, hasUnauthorizedAuthFailure(existing)) + if auth.LastError.Code == "refresh_outcome_unknown" { + auth.StatusMessage = auth.LastError.Message + } } } } @@ -525,7 +534,7 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo } } auth.UpdatedAt = now - cooldownStateChanged = normalizeModelStates(auth) || cooldownStateChanged + cooldownStateChanged = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) || cooldownStateChanged if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged } @@ -557,6 +566,10 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo return nil, false, false, fmt.Errorf("auth %s changed registration during persistence", auth.ID) } mergeAuthSaveDelta(auth, existingBeforeSave, existing, false) + if m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) { + cooldownStateChanged = true + updateAggregatedAvailability(auth, now) + } if (auth.RejectedAccessToken == "" && existing.RejectedAccessToken != "") || (auth.RejectedSIWCAccessToken == "" && existing.RejectedSIWCAccessToken != "") { probe := auth.Clone() probe.CodexRefreshPatch = refreshPatchBeforeSave @@ -743,6 +756,7 @@ func (m *Manager) Load(ctx context.Context) error { return err } m.mu.Lock() + now := time.Now() previousAuths := m.auths m.auths = make(map[string]*Auth, len(items)) if m.authEpochs == nil { @@ -757,6 +771,10 @@ func (m *Manager) Load(ctx context.Context) error { continue } auth.EnsureIndex() + _ = m.migrateAuthAliasModelStatesLocked(auth, nil, now, false) + if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } m.authEpochs[auth.ID] = max(m.authEpochs[auth.ID], auth.RegistrationEpoch) + 1 auth.RegistrationEpoch = m.authEpochs[auth.ID] auth.Generation = 1 diff --git a/sdk/cliproxy/auth/conductor_models.go b/sdk/cliproxy/auth/conductor_models.go index bebb4374f96..dd5ab10169c 100644 --- a/sdk/cliproxy/auth/conductor_models.go +++ b/sdk/cliproxy/auth/conductor_models.go @@ -284,7 +284,7 @@ func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel s } stateModel := executionResultModel(routeModel, upstreamModel, pooled) selectionModel := m.selectionModelForAuth(auth, routeModel) - if canonicalModelKey(selectionModel) == canonicalModelKey(upstreamModel) && strings.TrimSpace(selectionModel) != "" { + if sameCanonicalModelKey(selectionModel, upstreamModel) && strings.TrimSpace(selectionModel) != "" { return strings.TrimSpace(upstreamModel) } return stateModel diff --git a/sdk/cliproxy/auth/conductor_oauth_alias_nofork_test.go b/sdk/cliproxy/auth/conductor_oauth_alias_nofork_test.go index 5f68d06f591..c404d046372 100644 --- a/sdk/cliproxy/auth/conductor_oauth_alias_nofork_test.go +++ b/sdk/cliproxy/auth/conductor_oauth_alias_nofork_test.go @@ -4296,3 +4296,646 @@ func TestManager_Load_MonotonicEpochAssignment_AndRemovalTombstones(t *testing.T } } } + +type stub429CooldownExecutor struct { + id string + retryAfter time.Duration + + mu sync.Mutex + models []string +} + +func (e *stub429CooldownExecutor) Identifier() string { return e.id } + +func (e *stub429CooldownExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.models = append(e.models, req.Model) + e.mu.Unlock() + ra := e.retryAfter + return cliproxyexecutor.Response{}, withTrustedRetryAfter(&Error{ + Code: "rate_limit_exceeded", + Message: "upstream 429 cooldown", + HTTPStatus: http.StatusTooManyRequests, + }, &ra) +} + +func (e *stub429CooldownExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "ExecuteStream not implemented"} +} + +func (e *stub429CooldownExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *stub429CooldownExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "CountTokens not implemented"} +} + +func (e *stub429CooldownExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "HttpRequest not implemented"} +} + +func (e *stub429CooldownExecutor) Calls() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.models)) + copy(out, e.models) + return out +} + +func TestManager_Issue7360_Shape1_MixedCaseAliasAndLowercaseModelIDShareCooldown(t *testing.T) { + const ( + provider = "deepseek-plugin" + mixedCaseModel = "DeepSeek-V4.1-flash-sg" + lowerCaseModel = "deepseek-v4.1-flash-sg" + aliasModel = "DeepSeek V4.1 flash" + authID = "auth-issue7360-shape1" + ) + + manager := NewManager(nil, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: mixedCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{ + {ID: lowerCaseModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + _, errFirst := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errFirst == nil { + t.Fatal("expected first Execute to fail with 429") + } + if calls := executor.Calls(); len(calls) != 1 || calls[0] != mixedCaseModel { + t.Fatalf("after first Execute, calls = %v, want [%q]", calls, mixedCaseModel) + } + + _, errSecond := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: lowerCaseModel}, cliproxyexecutor.Options{}) + if errSecond == nil { + t.Fatal("expected second Execute to fail due to live cooldown") + } + if calls := executor.Calls(); len(calls) != 1 { + t.Fatalf("expected second Execute to be blocked by live cooldown without calling upstream again, got calls = %v", calls) + } + + manager.mu.RLock() + currentAuth := manager.auths[authID] + manager.mu.RUnlock() + if len(currentAuth.ModelStates) != 1 { + t.Fatalf("expected exactly 1 ModelStates cell, got %d: %#v", len(currentAuth.ModelStates), currentAuth.ModelStates) + } + if !reg.IsModelSuspendedForClient(authID, lowerCaseModel) { + t.Fatalf("expected %q to be suspended in registry for %s", lowerCaseModel, authID) + } +} + +func TestManager_Issue7360_Shape2_CooldownRecordedBeforeAliasMappingMigratesToTarget(t *testing.T) { + const ( + provider = "deepseek-plugin" + lowerCaseModel = "deepseek-v4.1-flash-sg" + aliasModel = "DeepSeek V4.1 flash" + ) + + t.Run("execute-before-alias-mapping", func(t *testing.T) { + const authID = "auth-issue7360-shape2-exec" + manager := NewManager(nil, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{ + {ID: lowerCaseModel}, + {ID: aliasModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + _, errFirst := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errFirst == nil { + t.Fatal("expected first Execute to fail with 429") + } + if calls := executor.Calls(); len(calls) != 1 || calls[0] != aliasModel { + t.Fatalf("after first Execute, calls = %v, want [%q]", calls, aliasModel) + } + + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: lowerCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + + _, errSecond := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errSecond == nil { + t.Fatal("expected second Execute to fail due to live cooldown") + } + if calls := executor.Calls(); len(calls) != 1 { + t.Fatalf("expected second Execute to honour cooldown recorded before alias mapping without calling upstream again, got calls = %v", calls) + } + + _, errThird := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: lowerCaseModel}, cliproxyexecutor.Options{}) + if errThird == nil { + t.Fatal("expected third Execute on target model to fail due to live cooldown") + } + if calls := executor.Calls(); len(calls) != 1 { + t.Fatalf("expected third Execute on target model to honour migrated cooldown without calling upstream again, got calls = %v", calls) + } + + manager.mu.RLock() + currentAuth := manager.auths[authID] + manager.mu.RUnlock() + if len(currentAuth.ModelStates) != 1 { + t.Fatalf("expected exactly 1 ModelStates cell after alias migration, got %d: %#v", len(currentAuth.ModelStates), currentAuth.ModelStates) + } + if _, existsStale := currentAuth.ModelStates[aliasModel]; existsStale { + t.Fatalf("expected stale route key %q to be migrated away from ModelStates, got %#v", aliasModel, currentAuth.ModelStates) + } + if state := currentAuth.ModelStates[lowerCaseModel]; state == nil || !state.Unavailable { + t.Fatalf("expected target key %q to hold migrated active cooldown, got %#v", lowerCaseModel, state) + } + }) + + t.Run("preexisting-alias-cooldown-with-only-lowercase-model-registered", func(t *testing.T) { + const authID = "auth-issue7360-shape2-preexisting" + manager := NewManager(nil, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + + retryUntil := time.Now().Add(time.Hour) + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + ModelStates: map[string]*ModelState{ + aliasModel: { + Unavailable: true, + Status: StatusError, + StatusMessage: "quota", + NextRetryAfter: retryUntil, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: retryUntil, + }, + }, + }, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{ + {ID: lowerCaseModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: lowerCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + + _, errSecond := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errSecond == nil { + t.Fatal("expected Execute via alias to fail due to live cooldown") + } + if calls := executor.Calls(); len(calls) != 0 { + t.Fatalf("expected Execute via alias to honour pre-existing cooldown without calling upstream, got calls = %v", calls) + } + + _, errThird := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: lowerCaseModel}, cliproxyexecutor.Options{}) + if errThird == nil { + t.Fatal("expected Execute on target model to fail due to live cooldown") + } + if calls := executor.Calls(); len(calls) != 0 { + t.Fatalf("expected Execute on target model to honour migrated cooldown without calling upstream, got calls = %v", calls) + } + + manager.mu.RLock() + currentAuth := manager.auths[authID] + manager.mu.RUnlock() + if len(currentAuth.ModelStates) != 1 { + t.Fatalf("expected exactly 1 ModelStates cell after alias migration, got %d: %#v", len(currentAuth.ModelStates), currentAuth.ModelStates) + } + if _, existsStale := currentAuth.ModelStates[aliasModel]; existsStale { + t.Fatalf("expected stale route key %q to be migrated away from ModelStates, got %#v", aliasModel, currentAuth.ModelStates) + } + if state := currentAuth.ModelStates[lowerCaseModel]; state == nil || !state.Unavailable { + t.Fatalf("expected target key %q to hold migrated active cooldown, got %#v", lowerCaseModel, state) + } + if !reg.IsModelSuspendedForClient(authID, lowerCaseModel) { + t.Fatalf("expected %q to be suspended in registry for %s", lowerCaseModel, authID) + } + }) + + t.Run("inflight-alias-request-completes-after-alias-mapping", func(t *testing.T) { + const authID = "auth-issue7360-shape2-inflight" + manager := NewManager(nil, nil, nil) + + startedCh := make(chan struct{}) + releaseCh := make(chan struct{}) + var callsMu sync.Mutex + var calls []string + retryAfter := time.Hour + + manager.RegisterExecutor(&customStreamMockExecutor{ + identifier: provider, + mockCustomErrorExecutor: mockCustomErrorExecutor{ + executeFn: func(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + callsMu.Lock() + callNum := len(calls) + 1 + calls = append(calls, req.Model) + callsMu.Unlock() + if callNum == 1 { + close(startedCh) + <-releaseCh + } + return cliproxyexecutor.Response{}, withTrustedRetryAfter(&Error{ + Code: "rate_limit_exceeded", + Message: "upstream 429 cooldown", + HTTPStatus: http.StatusTooManyRequests, + }, &retryAfter) + }, + }, + }) + + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{ + {ID: lowerCaseModel}, + {ID: aliasModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + doneCh := make(chan error, 1) + go func() { + _, errExec := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + doneCh <- errExec + }() + + <-startedCh + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: lowerCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + close(releaseCh) + + if errFirst := <-doneCh; errFirst == nil { + t.Fatal("expected in-flight Execute to fail with 429") + } + + _, errSecond := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: lowerCaseModel}, cliproxyexecutor.Options{}) + if errSecond == nil { + t.Fatal("expected Execute on target model to fail due to live cooldown") + } + _, errThird := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errThird == nil { + t.Fatal("expected Execute via alias to fail due to live cooldown") + } + + callsMu.Lock() + gotCalls := append([]string(nil), calls...) + callsMu.Unlock() + if len(gotCalls) != 1 { + t.Fatalf("expected only 1 upstream call, got %v", gotCalls) + } + + manager.mu.RLock() + currentAuth := manager.auths[authID] + manager.mu.RUnlock() + if len(currentAuth.ModelStates) != 1 { + t.Fatalf("expected exactly 1 ModelStates cell, got %d: %#v", len(currentAuth.ModelStates), currentAuth.ModelStates) + } + if _, existsStale := currentAuth.ModelStates[aliasModel]; existsStale { + t.Fatalf("expected stale route key %q not to exist in ModelStates, got %#v", aliasModel, currentAuth.ModelStates) + } + if state := currentAuth.ModelStates[lowerCaseModel]; state == nil || !state.Unavailable { + t.Fatalf("expected target key %q to hold active cooldown, got %#v", lowerCaseModel, state) + } + }) + + t.Run("chained-alias-with-only-intermediate-target-registered-preserves-cooldown", func(t *testing.T) { + const ( + authID = "auth-issue7360-chained-intermediate" + routeB = "DeepSeek V4.1 flash" + targetA = "deepseek-v4.1-flash-sg" + targetC = "deepseek-v4.1-flash-global" + ) + manager := NewManager(nil, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + + aliases := map[string][]internalconfig.OAuthModelAlias{ + provider: { + {Name: targetA, Alias: routeB, Fork: true}, + {Name: targetC, Alias: targetA, Fork: true}, + }, + } + manager.SetOAuthModelAlias(aliases) + + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{ + {ID: targetA}, + }) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + _, errFirst := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: routeB}, cliproxyexecutor.Options{}) + if errFirst == nil { + t.Fatal("expected first Execute for routeB to fail with 429") + } + if calls := executor.Calls(); len(calls) != 1 || calls[0] != targetA { + t.Fatalf("after first Execute, calls = %v, want [%q]", calls, targetA) + } + + // Re-applying alias configuration or reconciling must not migrate targetA's cooldown to targetC + manager.SetOAuthModelAlias(aliases) + manager.ReconcileRegistryModelStates(context.Background(), authID) + + _, errSecond := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: routeB}, cliproxyexecutor.Options{}) + if errSecond == nil { + t.Fatal("expected second Execute for routeB to be blocked by cooldown on targetA") + } + if calls := executor.Calls(); len(calls) != 1 { + t.Fatalf("expected second Execute for routeB not to call upstream again, got calls = %v", calls) + } + }) + + t.Run("atomic-multi-auth-migration-while-first-auth-save-is-blocked", func(t *testing.T) { + const ( + authID1 = "auth-issue7360-atomic-1" + authID2 = "auth-issue7360-atomic-2" + ) + saveBlockedCh := make(chan struct{}) + saveReleaseCh := make(chan struct{}) + var releaseOnce sync.Once + releaseSave := func() { releaseOnce.Do(func() { close(saveReleaseCh) }) } + t.Cleanup(releaseSave) + var blockSaves sync.Once + var shouldBlockSaves sync.Mutex + blockEnabled := false + + store := &blockingSelectiveAuthStore{ + onSave: func(_ *Auth) { + shouldBlockSaves.Lock() + enabled := blockEnabled + shouldBlockSaves.Unlock() + if !enabled { + return + } + blockSaves.Do(func() { + close(saveBlockedCh) + <-saveReleaseCh + }) + }, + } + + manager := NewManager(store, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + + retryUntil := time.Now().Add(time.Hour) + for _, id := range []string{authID1, authID2} { + auth := &Auth{ + ID: id, + Provider: provider, + Status: StatusActive, + Metadata: map[string]any{"type": provider}, + ModelStates: map[string]*ModelState{ + aliasModel: { + Unavailable: true, + Status: StatusError, + StatusMessage: "quota", + NextRetryAfter: retryUntil, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: retryUntil, + }, + }, + }, + } + if _, errReg := manager.Register(context.Background(), auth); errReg != nil { + t.Fatalf("register %s: %v", id, errReg) + } + } + + reg := registry.GetGlobalRegistry() + for _, id := range []string{authID1, authID2} { + reg.RegisterClient(id, provider, []*registry.ModelInfo{{ID: lowerCaseModel}}) + } + t.Cleanup(func() { + reg.UnregisterClient(authID1) + reg.UnregisterClient(authID2) + }) + manager.RefreshSchedulerEntry(authID1) + manager.RefreshSchedulerEntry(authID2) + + shouldBlockSaves.Lock() + blockEnabled = true + shouldBlockSaves.Unlock() + + aliasDoneCh := make(chan struct{}) + go func() { + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: lowerCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + close(aliasDoneCh) + }() + + <-saveBlockedCh + // While the first credential's store persistence is blocked, both credentials + // must already have migrated their in-memory cooldowns atomically with the alias table. + _, errExecDuringSave := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecDuringSave == nil { + t.Fatal("expected Execute during SetOAuthModelAlias persistence to be blocked by migrated cooldowns on both credentials") + } + if calls := executor.Calls(); len(calls) != 0 { + t.Fatalf("expected 0 upstream calls during blocked persistence, got %v", calls) + } + + releaseSave() + <-aliasDoneCh + }) + + t.Run("concurrent-load-and-alias-publication-preserves-cooldown", func(t *testing.T) { + const authID = "auth-issue7360-load-race" + listBlockedCh := make(chan struct{}) + listReleaseCh := make(chan struct{}) + var releaseOnce sync.Once + releaseList := func() { releaseOnce.Do(func() { close(listReleaseCh) }) } + t.Cleanup(releaseList) + + retryUntil := time.Now().Add(time.Hour) + diskSnapshot := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + Metadata: map[string]any{"type": provider}, + ModelStates: map[string]*ModelState{ + aliasModel: { + Unavailable: true, + Status: StatusError, + StatusMessage: "quota", + NextRetryAfter: retryUntil, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: retryUntil, + }, + }, + }, + } + + var listOnce sync.Once + store := &blockingSelectiveAuthStore{ + onList: func() ([]*Auth, error) { + listOnce.Do(func() { + close(listBlockedCh) + <-listReleaseCh + }) + return []*Auth{diskSnapshot.Clone()}, nil + }, + } + + manager := NewManager(store, nil, nil) + executor := &stub429CooldownExecutor{id: provider, retryAfter: time.Hour} + manager.RegisterExecutor(executor) + if _, errReg := manager.Register(context.Background(), diskSnapshot.Clone()); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, provider, []*registry.ModelInfo{{ID: lowerCaseModel}}) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + manager.RefreshSchedulerEntry(authID) + + loadDoneCh := make(chan error, 1) + go func() { + loadDoneCh <- manager.Load(context.Background()) + }() + + <-listBlockedCh + aliasDoneCh := make(chan struct{}) + go func() { + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: lowerCaseModel, + Alias: aliasModel, + Fork: true, + }}, + }) + close(aliasDoneCh) + }() + + releaseList() + if errLoad := <-loadDoneCh; errLoad != nil { + t.Fatalf("manager.Load failed: %v", errLoad) + } + <-aliasDoneCh + + for _, modelSpelling := range []string{aliasModel, lowerCaseModel} { + _, errExec := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: modelSpelling}, cliproxyexecutor.Options{}) + if errExec == nil { + t.Fatalf("expected Execute(%q) after concurrent Load and SetOAuthModelAlias to be blocked by cooldown", modelSpelling) + } + } + if calls := executor.Calls(); len(calls) != 0 { + t.Fatalf("expected 0 upstream calls after concurrent Load and SetOAuthModelAlias, got %v", calls) + } + }) +} + +type blockingSelectiveAuthStore struct { + onSave func(auth *Auth) + onList func() ([]*Auth, error) +} + +func (s *blockingSelectiveAuthStore) List(context.Context) ([]*Auth, error) { + if s.onList != nil { + return s.onList() + } + return nil, nil +} + +func (s *blockingSelectiveAuthStore) Save(_ context.Context, auth *Auth) (string, error) { + if s.onSave != nil { + s.onSave(auth) + } + if auth != nil { + return auth.ID, nil + } + return "", nil +} + +func (s *blockingSelectiveAuthStore) Delete(context.Context, string) error { + return nil +} diff --git a/sdk/cliproxy/auth/conductor_pinned_test.go b/sdk/cliproxy/auth/conductor_pinned_test.go new file mode 100644 index 00000000000..0f183bc6ebe --- /dev/null +++ b/sdk/cliproxy/auth/conductor_pinned_test.go @@ -0,0 +1,96 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "net/http" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestManagerNonWeightedPinnedSelection(t *testing.T) { + for _, mixed := range []bool{false, true} { + for _, affinity := range []bool{false, true} { + for _, strategy := range []string{"round-robin", "fill-first"} { + t.Run(fmt.Sprintf("mixed=%t/affinity=%t/%s", mixed, affinity, strategy), func(t *testing.T) { + var selector Selector = &RoundRobinSelector{} + if strategy == "fill-first" { + selector = &FillFirstSelector{} + } + if affinity { + sticky := NewSessionAffinitySelector(selector) + t.Cleanup(sticky.Stop) + selector = sticky + } + manager := NewManager(nil, selector, nil) + const provider, model = "pinned-non-weighted", "pinned-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + pinnedID := t.Name() + "-pinned" + for i := range 32 { + auth := &Auth{ + ID: fmt.Sprintf("%s-other-%02d", t.Name(), i), Provider: provider, Status: StatusActive, + Attributes: map[string]string{AttributeAuthKind: AuthKindAPIKey, "priority": "10"}, + } + if i == 0 { + auth.ID = pinnedID + auth.Attributes["priority"] = "0" + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(t.Context()), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + // Normalize the legacy selection signatures for the same pinning checks. + pick := func(ctx context.Context, id string, tried map[string]struct{}) (*Auth, error) { + opts := cliproxyexecutor.Options{ + Headers: http.Header{"Session_id": {"non-weighted-pinned"}}, + Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: id}, + } + if mixed { + auth, _, _, errPick := manager.pickNextMixedLegacy(ctx, []string{provider}, model, opts, tried) + return auth, errPick + } + auth, _, errPick := manager.pickNextLegacy(ctx, provider, model, opts, tried) + return auth, errPick + } + ctx := withRequiredAuthKind(t.Context(), AuthKindAPIKey) + for range 3 { + if auth, errPick := pick(ctx, pinnedID, nil); errPick != nil || auth == nil || auth.ID != pinnedID { + t.Fatalf("lower-priority pin not enforced: auth=%v err=%v", auth, errPick) + } + } + for _, tc := range []struct { + name string + ctx context.Context + id string + tried map[string]struct{} + }{ + {name: "missing pin", ctx: ctx, id: "missing"}, + {name: "already tried", ctx: ctx, id: pinnedID, tried: map[string]struct{}{pinnedID: {}}}, + {name: "ineligible kind", ctx: withRequiredAuthKind(t.Context(), AuthKindOAuth), id: pinnedID}, + } { + auth, errPick := pick(tc.ctx, tc.id, tc.tried) + var authErr *Error + if auth != nil || !errors.As(errPick, &authErr) || authErr.Code != "auth_not_found" { + t.Fatalf("%s: auth=%v err=%v, want auth_not_found", tc.name, auth, errPick) + } + } + if sticky, ok := selector.(*SessionAffinitySelector); ok { + namespace := provider + if mixed { + namespace = "mixed" + } + id := CanonicalSessionID(http.Header{"Session_id": {"non-weighted-pinned"}}, nil, nil) + if binding, found := sticky.cache.Get(namespace + "::" + id + "::" + model); !found || binding != pinnedID { + t.Fatalf("pinning checks changed the sticky binding: %q, found=%t", binding, found) + } + } + }) + } + } + } +} diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index cf28ed43e7f..819dd94785b 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -11,6 +11,7 @@ import ( "sync" "time" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -1259,6 +1260,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc m.mu.RLock() auth := m.auths[id].Clone() var exec ProviderExecutor + store := m.store if auth != nil { // Use the same effective provider key as request execution so OpenAI-compat // auths registered under namespaced keys still resolve for refresh. @@ -1268,6 +1270,19 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if auth == nil || exec == nil { return nil, errors.New("auth or executor not found") } + if strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + if resolver, okResolver := store.(interface { + RefreshJournalFor(*Auth) (codex.RefreshJournal, error) + }); okResolver { + journal, errJournal := resolver.RefreshJournalFor(auth) + if errJournal != nil { + return nil, fmt.Errorf("resolve credential refresh journal: %w", errJournal) + } + ctx = codex.WithRefreshJournal(ctx, journal) + } else if journal, okJournal := store.(codex.RefreshJournal); okJournal { + ctx = codex.WithRefreshJournal(ctx, journal) + } + } if registrationEpoch != 0 && auth.RegistrationEpoch != registrationEpoch { return nil, errors.New("auth registration changed before refresh") } @@ -1331,6 +1346,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.RegistrationEpoch == auth.RegistrationEpoch && current.CredentialVersion == auth.CredentialVersion && authHasUnchangedRejectedRefreshToken(current) { + wasTerminalUnauthorized := hasUnauthorizedAuthFailure(current) current.Generation++ current.UpdatedAt = checkNow current.Unavailable = true @@ -1351,7 +1367,10 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if current.RefreshError != nil && strings.TrimSpace(current.RefreshError.Message) != "" { msg = current.RefreshError.Message } - current.LastError = &Error{Code: "unauthorized", Message: msg, HTTPStatus: http.StatusUnauthorized} + current.LastError = refreshAccessError(current, msg, accessTokenRejected, wasTerminalUnauthorized) + if current.LastError.Code == "refresh_outcome_unknown" { + current.StatusMessage = current.LastError.Message + } m.auths[id] = current if m.scheduler != nil { m.scheduler.upsertAuth(current.Clone()) @@ -1368,7 +1387,22 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc base := auth.Clone() updated, err := exec.Refresh(ctx, base.Clone()) now := time.Now() - if err != nil && errors.Is(err, context.Canceled) { + refreshCanceled := (errors.Is(err, context.Canceled) || errors.Is(err, codex.ErrSIWCRefreshDeferred)) && !codex.IsRefreshOutcomeUnknown(err) + if refreshCanceled { + var partitions interface { + OuterRefreshError() error + SIWCRefreshError() error + } + if errors.As(err, &partitions) { + for _, partitionErr := range []error{partitions.OuterRefreshError(), partitions.SIWCRefreshError()} { + if partitionErr != nil && !errors.Is(partitionErr, context.Canceled) && !errors.Is(partitionErr, codex.ErrSIWCRefreshDeferred) { + refreshCanceled = false + break + } + } + } + } + if refreshCanceled { log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID) m.mu.Lock() if current := m.auths[id]; current != nil { @@ -1536,8 +1570,10 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.NextRefreshAfter = time.Time{} current.NextRetryAfter = time.Time{} current.RefreshFailures = 0 - current.LastError = &Error{Code: "unauthorized", Message: sanitizedRefreshErr.Message, HTTPStatus: http.StatusUnauthorized} - if unauthorized { + current.LastError = refreshAccessError(current, sanitizedRefreshErr.Message, accessTokenRejected, wasTerminalUnauthorized) + if current.LastError.Code == "refresh_outcome_unknown" { + current.StatusMessage = current.LastError.Message + } else if unauthorized { current.StatusMessage = "unauthorized" } else { current.StatusMessage = "unauthorized (refresh token invalid)" @@ -1643,7 +1679,10 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc updated.StatusMessage = "unauthorized (refresh token invalid)" } if updated.LastError == nil && updated.RefreshError != nil { - updated.LastError = &Error{Code: "unauthorized", Message: updated.RefreshError.Message, HTTPStatus: http.StatusUnauthorized} + updated.LastError = refreshAccessError(updated, updated.RefreshError.Message, failedAccessToken != "" && authMatchesTokenPartition(updated, failedAccessToken), hasUnauthorizedAuthFailure(base)) + if updated.LastError.Code == "refresh_outcome_unknown" { + updated.StatusMessage = updated.LastError.Message + } } } updated.UpdatedAt = now diff --git a/sdk/cliproxy/auth/conductor_refresh_outcome_test.go b/sdk/cliproxy/auth/conductor_refresh_outcome_test.go new file mode 100644 index 00000000000..3f7f2c53cac --- /dev/null +++ b/sdk/cliproxy/auth/conductor_refresh_outcome_test.go @@ -0,0 +1,37 @@ +package auth + +import ( + "context" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestUnknownRefreshOutcomeRequiresRecoveryWithoutDiscardingAccess(t *testing.T) { + for _, cause := range []error{nil, context.Canceled} { + ctx := context.Background() + manager := NewManager(nil, nil, nil) + executor := &unauthorizedRefreshExecutor{id: "codex", refreshErr: &codex.RefreshOutcomeUnknownError{Stage: "transport", Cause: cause}} + manager.RegisterExecutor(executor) + expiry := time.Now().Add(time.Hour).Truncate(time.Second) + auth := &Auth{ID: "unknown-outcome", Provider: "codex", Status: StatusActive, Metadata: map[string]any{ + "access_token": "valid-access", "refresh_token": "unresolved-grant", "expired": expiry.Format(time.RFC3339), + }} + if _, errRegister := manager.Register(ctx, auth); errRegister != nil { + t.Fatal(errRegister) + } + manager.refreshAuth(ctx, auth.ID) + updated, _ := manager.GetByID(auth.ID) + if updated == nil || !updated.RequiresSignIn || updated.RefreshError == nil || updated.RefreshError.Code != "refresh_outcome_unknown" { + t.Fatalf("unknown outcome not retained as recovery-required: %+v", updated) + } + if updated.Unavailable || updated.Status != StatusActive { + t.Fatal("valid access token must remain usable") + } + manager.refreshAuth(ctx, auth.ID) + if executor.RefreshCalls() != 1 { + t.Fatalf("scheduled refresh replayed: calls=%d", executor.RefreshCalls()) + } + } +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index b30333bc7e5..e379fb32960 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -237,72 +237,23 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin supportedModels, regEpoch = globalReg.GetModelsAndEpochForClient(authID) candidateAuth := &Auth{ ID: auth.ID, + Prefix: auth.Prefix, Provider: auth.Provider, Attributes: auth.Attributes, + Metadata: auth.Metadata, ModelStates: cloneModelStates(auth.ModelStates), } - candidateChanged := normalizeModelStates(candidateAuth) - - // Historical alias migration: - // Migrate legacy alias state strictly based on registered route keys from supportedModels. - // Two-phase safe migration: - // Phase 1: Collect authoritative target keys and route-to-target mappings. - if len(candidateAuth.ModelStates) > 0 && len(supportedModels) > 0 { - authoritativeTargets := make(map[string]bool, len(supportedModels)) - routeToTarget := make(map[string]string, len(supportedModels)) - for _, sm := range supportedModels { - if sm == nil || strings.TrimSpace(sm.ID) == "" { - continue - } - routeID := strings.TrimSpace(sm.ID) - canonicalRoute := canonicalModelKey(routeID) - targetKey := m.selectionModelKeyForAuth(auth, routeID) - if targetKey == "" { - targetKey = canonicalRoute - } - if targetKey != "" { - authoritativeTargets[targetKey] = true - } - if canonicalRoute != "" { - routeToTarget[canonicalRoute] = targetKey - } - } + candidateChanged := m.migrateAuthAliasModelStatesLocked(candidateAuth, supportedModels, now, true) - // Phase 2: For each registered routeKey, migrate legacy state to target ONLY if - // routeKey != target AND routeKey is not itself an authoritative target for any route. - for _, sm := range supportedModels { - if sm == nil || strings.TrimSpace(sm.ID) == "" { - continue - } - routeID := strings.TrimSpace(sm.ID) - canonicalRoute := canonicalModelKey(routeID) - targetKey := routeToTarget[canonicalRoute] - if targetKey == "" || canonicalRoute == "" { - continue - } - if canonicalRoute != targetKey && !authoritativeTargets[canonicalRoute] { - aliasKeys := []string{canonicalRoute} - if routeID != canonicalRoute { - aliasKeys = append(aliasKeys, routeID) - } - for _, aliasKey := range aliasKeys { - if state, hasAliasState := candidateAuth.ModelStates[aliasKey]; hasAliasState { - if state != nil && (!modelStateIsClean(state) || isModelStateActiveCooldown(state, now)) { - if existingTarget, exists := candidateAuth.ModelStates[targetKey]; exists && existingTarget != nil { - candidateAuth.ModelStates[targetKey] = mergeModelState(existingTarget, state) - } else { - candidateAuth.ModelStates[targetKey] = state.Clone() - } - } - delete(candidateAuth.ModelStates, aliasKey) - candidateChanged = true - } - } + supported := make(map[string]struct{}, len(supportedModels)) + if len(supportedModels) > 0 { + authoritativeTargets, _ := m.buildAliasMigrationMapsLocked(candidateAuth, supportedModels, true) + for targetLookup := range authoritativeTargets { + if targetLookup != "" { + supported[targetLookup] = struct{}{} } } } - - supported := make(map[string]struct{}, len(supportedModels)) for _, model := range supportedModels { if model == nil || strings.TrimSpace(model.ID) == "" { continue @@ -311,17 +262,17 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin if stateKey == "" { stateKey = canonicalModelKey(model.ID) } - if stateKey != "" { - supported[stateKey] = struct{}{} + if lookupKey := canonicalModelLookupKey(stateKey); lookupKey != "" { + supported[lookupKey] = struct{}{} } } for modelKey, state := range candidateAuth.ModelStates { - baseModel := canonicalModelKey(modelKey) - if baseModel == "" { - baseModel = strings.TrimSpace(modelKey) + lookupKey := canonicalModelLookupKey(modelKey) + if lookupKey == "" { + lookupKey = strings.ToLower(strings.TrimSpace(modelKey)) } - if _, isSupported := supported[baseModel]; !isSupported { + if _, isSupported := supported[lookupKey]; !isSupported { // Drop state for models that disappeared from the current registry // snapshot and are not reachable via any alias route. delete(candidateAuth.ModelStates, modelKey) @@ -394,6 +345,356 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin } } +func compactModelSpellingKey(model string) string { + key := canonicalModelLookupKey(model) + if key == "" { + return "" + } + var b strings.Builder + b.Grow(len(key)) + for i := 0; i < len(key); i++ { + ch := key[i] + if (ch >= 'a' && ch <= 'z') || (ch >= '0' && ch <= '9') { + b.WriteByte(ch) + } + } + return b.String() +} + +func isSameUnderlyingModelSpelling(a, b string) bool { + ca := compactModelSpellingKey(a) + cb := compactModelSpellingKey(b) + if ca == "" || cb == "" { + return false + } + return ca == cb || strings.HasPrefix(ca, cb) || strings.HasPrefix(cb, ca) +} + +func (m *Manager) buildAliasMigrationMapsLocked(auth *Auth, supportedModels []*registry.ModelInfo, fromExplicitReconcile bool) (map[string]string, map[string]string) { + if m == nil || auth == nil { + return nil, nil + } + if !fromExplicitReconcile && len(supportedModels) == 0 && auth.ID != "" { + supportedModels = registry.GetGlobalRegistry().GetModelsForClient(auth.ID) + } + + allRegistered := make(map[string]struct{}, len(supportedModels)) + nativeRegistered := make(map[string]struct{}, len(supportedModels)) + for _, sm := range supportedModels { + if sm == nil || strings.TrimSpace(sm.ID) == "" { + continue + } + key := canonicalModelLookupKey(sm.ID) + if key == "" { + continue + } + allRegistered[key] = struct{}{} + if strings.TrimSpace(sm.MetadataModelID) != "" && !sameCanonicalModelKey(sm.ID, sm.MetadataModelID) { + continue + } + nativeRegistered[key] = struct{}{} + } + + authoritativeTargets := make(map[string]string, len(supportedModels)) + routeToTarget := make(map[string]string, len(supportedModels)) + for _, sm := range supportedModels { + if sm == nil || strings.TrimSpace(sm.ID) == "" { + continue + } + routeID := strings.TrimSpace(sm.ID) + canonicalRoute := canonicalModelKey(routeID) + routeLookup := canonicalModelLookupKey(canonicalRoute) + targetKey := m.selectionModelKeyForAuth(auth, routeID) + if targetKey == "" { + targetKey = canonicalRoute + } + targetLookup := canonicalModelLookupKey(targetKey) + if !fromExplicitReconcile && routeLookup != "" && targetLookup != "" && routeLookup != targetLookup { + _, routeIsNative := nativeRegistered[routeLookup] + _, targetIsNative := nativeRegistered[targetLookup] + if routeIsNative && targetIsNative && !isSameUnderlyingModelSpelling(canonicalRoute, targetKey) { + authoritativeTargets[routeLookup] = canonicalRoute + routeToTarget[routeLookup] = canonicalRoute + continue + } + } + if targetLookup != "" { + authoritativeTargets[targetLookup] = targetKey + } + if routeLookup != "" && targetKey != "" { + routeToTarget[routeLookup] = targetKey + } + } + for _, route := range m.configuredOAuthAliasRoutes(auth) { + aliasLookup := canonicalModelLookupKey(route.aliasKey) + targetLookup := canonicalModelLookupKey(route.targetKey) + if aliasLookup == "" || targetLookup == "" || aliasLookup == targetLookup { + continue + } + if !fromExplicitReconcile { + _, routeIsNative := nativeRegistered[aliasLookup] + _, targetIsNative := nativeRegistered[targetLookup] + if routeIsNative && targetIsNative && !isSameUnderlyingModelSpelling(route.aliasKey, route.targetKey) { + continue + } + } + reachable := len(supportedModels) == 0 + if !reachable { + _, aliasRegistered := allRegistered[aliasLookup] + _, targetRegistered := allRegistered[targetLookup] + reachable = aliasRegistered || targetRegistered + } + if !reachable { + continue + } + if _, exists := authoritativeTargets[targetLookup]; !exists { + authoritativeTargets[targetLookup] = route.targetKey + } + if _, exists := routeToTarget[aliasLookup]; !exists { + routeToTarget[aliasLookup] = route.targetKey + } + } + return authoritativeTargets, routeToTarget +} + +func (m *Manager) resolveMarkResultModelKeyLocked(auth *Auth, modelKey, routeModel string) string { + canonicalModel := canonicalModelKey(modelKey) + if canonicalModel == "" || m == nil || auth == nil { + return canonicalModel + } + canonicalRoute := canonicalModelKey(routeModel) + if canonicalRoute != "" && !sameCanonicalModelKey(canonicalModel, canonicalRoute) { + if routeTarget := m.selectionModelKeyForAuth(auth, canonicalRoute); sameCanonicalModelKey(canonicalModel, routeTarget) { + return canonicalModel + } + } + lookupKey := canonicalModelLookupKey(canonicalModel) + authoritativeTargets, routeToTarget := m.buildAliasMigrationMapsLocked(auth, nil, false) + if authTarget, isAuthTarget := authoritativeTargets[lookupKey]; isAuthTarget && authTarget != "" { + return authTarget + } + targetKey := routeToTarget[lookupKey] + if targetKey == "" { + targetKey = m.selectionModelKeyForAuth(auth, canonicalModel) + } + if targetKey != "" && !sameCanonicalModelKey(canonicalModel, targetKey) { + if authTarget, isAuthTarget := authoritativeTargets[canonicalModelLookupKey(targetKey)]; isAuthTarget && authTarget != "" { + targetKey = authTarget + } + return targetKey + } + return canonicalModel +} + +func (m *Manager) migrateAuthAliasModelStatesLocked(auth *Auth, supportedModels []*registry.ModelInfo, now time.Time, fromExplicitReconcile bool) bool { + if auth == nil { + return false + } + changed := normalizeModelStates(auth) + if m == nil || len(auth.ModelStates) == 0 { + return changed + } + authoritativeTargets, routeToTarget := m.buildAliasMigrationMapsLocked(auth, supportedModels, fromExplicitReconcile) + + existingKeys := make([]string, 0, len(auth.ModelStates)) + for stateKey := range auth.ModelStates { + existingKeys = append(existingKeys, stateKey) + } + if len(existingKeys) > 1 { + sort.Strings(existingKeys) + } + + var working *Auth + for _, stateKey := range existingKeys { + sourceStates := auth.ModelStates + if working != nil { + sourceStates = working.ModelStates + } + state, exists := sourceStates[stateKey] + if !exists { + continue + } + canonicalRoute := canonicalModelKey(stateKey) + routeLookup := canonicalModelLookupKey(canonicalRoute) + if routeLookup == "" { + continue + } + if _, isAuthTarget := authoritativeTargets[routeLookup]; isAuthTarget { + continue + } + targetKey := routeToTarget[routeLookup] + if targetKey == "" { + targetKey = m.selectionModelKeyForAuth(auth, stateKey) + } + if targetKey == "" || sameCanonicalModelKey(canonicalRoute, targetKey) { + continue + } + if authTarget, isAuthTarget := authoritativeTargets[canonicalModelLookupKey(targetKey)]; isAuthTarget && authTarget != "" { + targetKey = authTarget + } + if working == nil { + working = &Auth{ModelStates: cloneModelStates(auth.ModelStates)} + } + if state != nil && (!modelStateIsClean(state) || isModelStateActiveCooldown(state, now)) { + mergeOrSetModelState(working, targetKey, state) + } + delete(working.ModelStates, stateKey) + changed = true + } + if working != nil { + if len(working.ModelStates) == 0 { + auth.ModelStates = nil + } else { + auth.ModelStates = working.ModelStates + } + } + return changed +} + +func mergeOrSetModelState(auth *Auth, targetKey string, state *ModelState) { + if auth == nil || targetKey == "" || state == nil { + return + } + if auth.ModelStates == nil { + auth.ModelStates = make(map[string]*ModelState) + } + if existingTarget, exists := auth.ModelStates[targetKey]; exists && existingTarget != nil { + auth.ModelStates[targetKey] = mergeModelState(existingTarget, state) + return + } + targetLookup := canonicalModelLookupKey(targetKey) + for existingKey, existingState := range auth.ModelStates { + if canonicalModelLookupKey(existingKey) != targetLookup { + continue + } + if existingState != nil { + auth.ModelStates[existingKey] = mergeModelState(existingState, state) + } else { + auth.ModelStates[existingKey] = state.Clone() + } + return + } + auth.ModelStates[targetKey] = state.Clone() +} + +func (m *Manager) publishOAuthModelAliasAndReconcileStates(ctx context.Context, table *oauthModelAliasTable) { + if m == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + if table == nil { + table = &oauthModelAliasTable{} + } + + globalReg := registry.GetGlobalRegistry() + type reconciledAuthItem struct { + authID string + snapshot *Auth + supportedModels []*registry.ModelInfo + regEpoch uint64 + changed bool + } + + loadGateAcquired := m.authLoadGate.Acquire(ctx, 1) == nil + now := time.Now() + cooldownStateChanged := false + + m.mu.Lock() + m.oauthModelAlias.Store(table) + trackCooldownState := m.cooldownStore != nil + items := make([]reconciledAuthItem, 0, len(m.auths)) + for authID, auth := range m.auths { + if auth == nil || authID == "" { + continue + } + var cooldownRecordsBefore []CooldownStateRecord + if trackCooldownState { + cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) + } + supportedModels, regEpoch := globalReg.GetModelsAndEpochForClient(authID) + candidateChanged := m.migrateAuthAliasModelStatesLocked(auth, supportedModels, now, false) + if candidateChanged { + updateAggregatedAvailability(auth, now) + if !hasModelError(auth, now) && !hasUnauthorizedAuthFailure(auth) { + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + } + auth.Generation++ + auth.UpdatedAt = now + } + if trackCooldownState { + cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) + if candidateChanged || !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) { + cooldownStateChanged = true + } + } + snapshot := auth.Clone() + if m.scheduler != nil { + m.scheduler.upsertAuth(snapshot.Clone()) + } + items = append(items, reconciledAuthItem{ + authID: authID, + snapshot: snapshot, + supportedModels: supportedModels, + regEpoch: regEpoch, + changed: candidateChanged, + }) + } + if len(items) > 0 { + m.structuralEpoch.Add(1) + } + m.mu.Unlock() + if loadGateAcquired { + m.authLoadGate.Release(1) + } + + for _, item := range items { + if len(item.supportedModels) > 0 && item.snapshot != nil { + projections := make([]registry.ClientModelProjection, 0, len(item.supportedModels)) + for _, sm := range item.supportedModels { + if sm == nil || strings.TrimSpace(sm.ID) == "" { + continue + } + projections = append(projections, m.clientModelProjectionForAuth(item.snapshot, sm.ID, now)) + } + globalReg.ApplyClientModelProjections(item.authID, item.regEpoch, item.snapshot.Generation, projections) + } + if item.changed { + releaseMutation := m.lockAuthMutation(item.authID) + m.mu.Lock() + if current, ok := m.auths[item.authID]; ok && current != nil { + if errPersist := m.persistLocked(context.Background(), current); errPersist != nil { + logEntryWithRequestID(ctx).WithField("auth_id", current.ID).Warnf("failed to persist auth changes during alias state reconciliation: %v", errPersist) + } + postNow := time.Now() + if m.migrateAuthAliasModelStatesLocked(current, nil, postNow, false) { + updateAggregatedAvailability(current, postNow) + if !hasModelError(current, postNow) && !hasUnauthorizedAuthFailure(current) { + current.LastError = nil + current.StatusMessage = "" + current.Status = StatusActive + } + current.Generation++ + current.UpdatedAt = postNow + if m.scheduler != nil { + m.scheduler.upsertAuth(current.Clone()) + } + m.structuralEpoch.Add(1) + cooldownStateChanged = true + } + } + m.mu.Unlock() + releaseMutation() + } + } + if cooldownStateChanged { + m.persistCooldownStates(context.Background()) + } +} + func cloneModelStates(states map[string]*ModelState) map[string]*ModelState { if len(states) == 0 { return nil @@ -634,6 +935,11 @@ func (m *Manager) availableAuthsForRouteModelForRequest(ctx context.Context, aut // or scheduler additionally receives lower priority tiers. func (m *Manager) availableAuthsForSelector(ctx context.Context, selector Selector, auths []*Auth, provider, routeModel string, now time.Time, schedulerAcross bool, clientHeaders ...http.Header) (priorityAuths, selectorAuths []*Auth, err error) { _, sessionAffinity := selector.(*SessionAffinitySelector) + if _, weighted := selector.(*WeightedRoundRobinSelector); weighted { + // Exclude zero-weight credentials before priority selection, just as the + // scheduler predicate does, including when only a lower tier has weight. + auths = positiveWeightAuths(auths) + } if !sessionAffinity && !schedulerAcross { priorityAuths, err = m.availableAuthsForRouteModelForRequest(ctx, auths, provider, routeModel, now, false, clientHeaders...) @@ -688,7 +994,7 @@ func selectorContextForAvailableAuths(ctx context.Context, selector Selector, ro func (m *Manager) transientlyUnavailableSessionAffinityAuthIDs(ctx context.Context, selector Selector, affinityCandidates []*Auth, routeModel string, now time.Time, clientHeaders ...http.Header) map[string]struct{} { affinity, ok := selector.(*SessionAffinitySelector) - if !ok || affinity == nil || !affinity.rebindOnQuotaOnly || len(affinityCandidates) == 0 { + if !ok || affinity == nil || (!affinity.rebindOnQuotaOnly && !affinity.rebindOnNever) || len(affinityCandidates) == 0 { return nil } if _, weighted := affinity.fallback.(*WeightedRoundRobinSelector); weighted { @@ -762,15 +1068,7 @@ func latestCandidateErrorForModel(auths []*Auth, selectionModelFunc func(*Auth) var modelErr error var modelTime time.Time if len(candidate.ModelStates) > 0 { - if state, ok := candidate.ModelStates[checkModel]; ok && state != nil { - if state.LastError != nil { - modelErr = state.LastError - modelTime = state.UpdatedAt - } else if strings.TrimSpace(state.StatusMessage) != "" { - modelErr = errors.New(state.StatusMessage) - modelTime = state.UpdatedAt - } - } else if state, ok := candidate.ModelStates[canonicalModelKey(checkModel)]; ok && state != nil { + if state := existingModelState(candidate, checkModel); state != nil { if state.LastError != nil { modelErr = state.LastError modelTime = state.UpdatedAt @@ -1069,7 +1367,7 @@ func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, au return true } selectionKey := m.selectionModelKeyForAuth(auth, routeModel) - return selectionKey != "" && selectionKey != routeKey && registryRef.ClientSupportsModel(auth.ID, selectionKey) + return selectionKey != "" && !sameCanonicalModelKey(selectionKey, routeKey) && registryRef.ClientSupportsModel(auth.ID, selectionKey) } func (m *Manager) normalizeProviders(providers []string) []string { @@ -1211,11 +1509,11 @@ func retryRoundAvailabilityForAuthForRequest(ctx context.Context, auth *Auth, mo return credentialRetryRoundStateEligible(auth.LastError, true), next } - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) if modelKey != "" && len(auth.ModelStates) > 0 { matchedBlocked := false for stateModel, state := range auth.ModelStates { - if state == nil || canonicalModelKey(stateModel) != modelKey { + if state == nil || canonicalModelLookupKey(stateModel) != modelKey { continue } if state.Status == StatusDisabled { @@ -1255,7 +1553,7 @@ func hasLiveQuotaCooldown(auth *Auth, model string, now time.Time) bool { if auth.Quota.Reason == "credential_quota" && quotaLive(auth.Quota) { return true } - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) if modelKey == "" { return quotaLive(auth.Quota) } @@ -1263,7 +1561,7 @@ func hasLiveQuotaCooldown(auth *Auth, model string, now time.Time) bool { if state == nil { continue } - if canonicalModelKey(stateModel) == modelKey && quotaLive(state.Quota) { + if canonicalModelLookupKey(stateModel) == modelKey && quotaLive(state.Quota) { return true } } @@ -1277,11 +1575,8 @@ func hasTransientZeroWaitEligibility(auth *Auth, model string) bool { if auth == nil { return false } - modelKey := canonicalModelKey(model) - if modelKey != "" && len(auth.ModelStates) > 0 { - if state := auth.ModelStates[modelKey]; state != nil { - return isTransientRateLimitResultError(state.LastError) - } + if state := existingModelState(auth, model); state != nil { + return isTransientRateLimitResultError(state.LastError) } return isTransientRateLimitResultError(auth.LastError) } @@ -1523,6 +1818,21 @@ func (m *Manager) shouldRetryAfterErrorWithAttempted(ctx context.Context, opts c } eligibility := authSelectionEligibilityForRequest(ctx, opts) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + if affinity, ok := m.Selector().(*SessionAffinitySelector); ok && affinity != nil && affinity.rebindOnNever { + var authErr *Error + if errors.As(err, &authErr) && authErr != nil && authErr.Code == "session_auth_unavailable" { + return 0, false + } + if boundAuthID, okBound := affinity.boundAuthForRequest(providers, model, opts); okBound { + if pinnedAuthID != "" && pinnedAuthID != boundAuthID { + return 0, false + } + if shouldRebindSessionAffinityOnError(resultErrorFromError(err)) { + return 0, false + } + pinnedAuthID = boundAuthID + } + } if !isRequestRetryRoundError(err) || !m.retryAllowedForRequest(ctx, attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry, opts.Headers) { return 0, false } @@ -1917,7 +2227,7 @@ func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) boo if auth == nil || strings.TrimSpace(routeModel) == "" { return false } - return m.selectionModelKeyForAuth(auth, routeModel) != canonicalModelKey(routeModel) + return !sameCanonicalModelKey(m.selectionModelForAuth(auth, routeModel), routeModel) } func requestAwareCredSelectionRequired(ctx context.Context, auth *Auth, clientHeaders ...http.Header) bool { @@ -1962,8 +2272,13 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op registryRef := registry.GetGlobalRegistry() targetKey := canonicalSchedulingProvider(provider) affinitySelector, hasAffinity := selector.(*SessionAffinitySelector) - trackTransientAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnQuotaOnly - var affinityCandidates []*Auth + trackTransientAffinity := hasAffinity && affinitySelector != nil && (affinitySelector.rebindOnQuotaOnly || affinitySelector.rebindOnNever) + _, weighted := selector.(*WeightedRoundRobinSelector) + if hasAffinity && affinitySelector != nil { + _, weighted = affinitySelector.fallback.(*WeightedRoundRobinSelector) + } + trackPinnedAccounting := pinnedAuthID != "" && weighted + var accountingCandidates, affinityCandidates []*Auth if trackTransientAffinity { affinityCandidates = make([]*Auth, 0, len(m.auths)) } @@ -1971,7 +2286,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if candidate == nil || canonicalSchedulingProvider(executorKeyFromAuth(candidate)) != targetKey || candidate.Disabled { continue } - if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + if !trackPinnedAccounting && pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } if !eligibility.allows(candidate) { @@ -1980,24 +2295,49 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { continue } + _, used := tried[candidate.ID] + // Pinning restricts selection, not the eligible pool used for accounting. + if trackPinnedAccounting && !used { + accountingCandidates = append(accountingCandidates, candidate) + } + if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + continue + } if trackTransientAffinity { affinityCandidates = append(affinityCandidates, candidate) } - if _, used := tried[candidate.ID]; used { + if used { continue } candidates = append(candidates, candidate) } - if len(candidates) == 0 { + canInspectUnavailableAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnNever && pluginScheduler == nil + if len(candidates) == 0 && !canInspectUnavailableAffinity { m.mu.RUnlock() return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"} } now := time.Now() - available, selectorAuths, errAvailable := m.availableAuthsForSelector(ctx, selector, candidates, provider, model, now, schedulerAcross, opts.Headers) - if errAvailable != nil { - m.mu.RUnlock() - m.warnLogAuthUnavailable(ctx, []string{provider}, model, opts, tried, errAvailable) - return nil, nil, errAvailable + var ( + available []*Auth + selectorAuths []*Auth + errAvailable error + ) + if len(candidates) > 0 { + available, selectorAuths, errAvailable = m.availableAuthsForSelector(ctx, selector, candidates, provider, model, now, schedulerAcross, opts.Headers) + if errAvailable != nil && !canInspectUnavailableAffinity { + m.mu.RUnlock() + m.warnLogAuthUnavailable(ctx, []string{provider}, model, opts, tried, errAvailable) + return nil, nil, errAvailable + } + } + var accountingAuths []*Auth + if trackPinnedAccounting && len(candidates) > 0 && errAvailable == nil { + fullAvailable, _, errAccounting := m.availableAuthsForSelector(ctx, selector, accountingCandidates, provider, model, now, true, opts.Headers) + if errAccounting != nil { + m.mu.RUnlock() + return nil, nil, errAccounting + } + accountingAuths = fullAvailable } transientAuthIDs := m.transientlyUnavailableSessionAffinityAuthIDs(ctx, selector, affinityCandidates, model, now, opts.Headers) m.mu.RUnlock() @@ -2010,8 +2350,18 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if !handled { selectorCtx := selectorContextForAvailableAuths(ctx, selector, model) selectorCtx = withSessionAffinityTransientAuthIDs(selectorCtx, transientAuthIDs) + if accountingAuths != nil { + selectorCtx = context.WithValue(selectorCtx, weightedAccountingPoolKey{}, accountingAuths) + } selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { + if isAuthNotFoundError(errPick) { + if errAvailable != nil { + errPick = errAvailable + } else if len(candidates) == 0 { + errPick = &Error{Code: "auth_not_found", Message: "no auth available"} + } + } if isBuiltInSelector(selector) { errPick = restoreModelCooldownErrorModel(errPick, model) } @@ -2367,8 +2717,13 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } registryRef := registry.GetGlobalRegistry() affinitySelector, hasAffinity := selector.(*SessionAffinitySelector) - trackTransientAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnQuotaOnly - var affinityCandidates []*Auth + trackTransientAffinity := hasAffinity && affinitySelector != nil && (affinitySelector.rebindOnQuotaOnly || affinitySelector.rebindOnNever) + _, weighted := selector.(*WeightedRoundRobinSelector) + if hasAffinity && affinitySelector != nil { + _, weighted = affinitySelector.fallback.(*WeightedRoundRobinSelector) + } + trackPinnedAccounting := pinnedAuthID != "" && weighted + var accountingCandidates, affinityCandidates []*Auth if trackTransientAffinity { affinityCandidates = make([]*Auth, 0, len(m.auths)) } @@ -2376,7 +2731,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if candidate == nil || candidate.Disabled { continue } - if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + if !trackPinnedAccounting && pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } if !eligibility.allows(candidate) { @@ -2395,24 +2750,49 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { continue } + _, used := tried[candidate.ID] + // Pinning restricts selection, not the eligible pool used for accounting. + if trackPinnedAccounting && !used { + accountingCandidates = append(accountingCandidates, candidate) + } + if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + continue + } if trackTransientAffinity { affinityCandidates = append(affinityCandidates, candidate) } - if _, used := tried[candidate.ID]; used { + if used { continue } candidates = append(candidates, candidate) } - if len(candidates) == 0 { + canInspectUnavailableAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnNever && pluginScheduler == nil + if len(candidates) == 0 && !canInspectUnavailableAffinity { m.mu.RUnlock() return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} } now := time.Now() - available, selectorAuths, errAvailable := m.availableAuthsForSelector(ctx, selector, candidates, "mixed", model, now, schedulerAcross, opts.Headers) - if errAvailable != nil { - m.mu.RUnlock() - m.warnLogAuthUnavailable(ctx, providers, model, opts, tried, errAvailable) - return nil, nil, "", errAvailable + var ( + available []*Auth + selectorAuths []*Auth + errAvailable error + ) + if len(candidates) > 0 { + available, selectorAuths, errAvailable = m.availableAuthsForSelector(ctx, selector, candidates, "mixed", model, now, schedulerAcross, opts.Headers) + if errAvailable != nil && !canInspectUnavailableAffinity { + m.mu.RUnlock() + m.warnLogAuthUnavailable(ctx, providers, model, opts, tried, errAvailable) + return nil, nil, "", errAvailable + } + } + var accountingAuths []*Auth + if trackPinnedAccounting && len(candidates) > 0 && errAvailable == nil { + fullAvailable, _, errAccounting := m.availableAuthsForSelector(ctx, selector, accountingCandidates, "mixed", model, now, true, opts.Headers) + if errAccounting != nil { + m.mu.RUnlock() + return nil, nil, "", errAccounting + } + accountingAuths = fullAvailable } transientAuthIDs := m.transientlyUnavailableSessionAffinityAuthIDs(ctx, selector, affinityCandidates, model, now, opts.Headers) m.mu.RUnlock() @@ -2425,8 +2805,36 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if !handled { selectorCtx := selectorContextForAvailableAuths(ctx, selector, model) selectorCtx = withSessionAffinityTransientAuthIDs(selectorCtx, transientAuthIDs) + if accountingAuths != nil { + selectorCtx = context.WithValue(selectorCtx, weightedAccountingPoolKey{}, accountingAuths) + } + if _, bareWeighted := selector.(*WeightedRoundRobinSelector); bareWeighted { + // Match the scheduler's single-provider transport preference after + // excluding providers without an executor. + singleProvider := "" + for providerKey := range providerSet { + if _, okExecutor := m.Executor(providerKey); !okExecutor { + continue + } + if singleProvider != "" { + singleProvider = "" + break + } + singleProvider = providerKey + } + if singleProvider != "" { + selectorCtx = context.WithValue(selectorCtx, weightedTransportProviderKey{}, singleProvider) + } + } selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { + if isAuthNotFoundError(errPick) { + if errAvailable != nil { + errPick = errAvailable + } else if len(candidates) == 0 { + errPick = &Error{Code: "auth_not_found", Message: "no auth available"} + } + } if isBuiltInSelector(selector) { errPick = restoreModelCooldownErrorModel(errPick, model) } @@ -2543,6 +2951,14 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s return authCopy, executor, providerKey, nil } +func isAuthNotFoundError(err error) bool { + if err == nil { + return false + } + var authErr *Error + return errors.As(err, &authErr) && authErr != nil && authErr.Code == "auth_not_found" +} + func isAuthUnavailableError(err error) bool { if err == nil { return false @@ -2565,9 +2981,7 @@ func authCoolingSummary(auth *Auth, model string, next time.Time, now time.Time) ident := formatAuthIdentity(auth, auth.Provider) reason := "" if model != "" && len(auth.ModelStates) > 0 { - if state, ok := auth.ModelStates[model]; ok && state != nil { - reason = cooldownReason(state.StatusMessage, state.Quota, state.LastError) - } else if state, ok := auth.ModelStates[canonicalModelKey(model)]; ok && state != nil { + if state := existingModelState(auth, model); state != nil { reason = cooldownReason(state.StatusMessage, state.Quota, state.LastError) } } diff --git a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go index 9b3272b99f7..98fc9320994 100644 --- a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go +++ b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go @@ -1082,7 +1082,7 @@ func TestManager_CodexTerminalRenewalRecoveryLifecycle_Issue5783(t *testing.T) { t.Fatal("VerifyInference must fail when stream ends with response.incomplete") } - // Operator-triggered VerifyInference pins to primary, uses an isolated session and minimal 1-token prompt, and marks InferenceVerified = true upon response.completed. + // Operator-triggered VerifyInference pins to primary, uses an isolated session and a minimal 16-token (Codex floor) prompt, and marks InferenceVerified = true upon response.completed. executor.mu.Lock() executor.streamPayloads = nil executor.streamCalls = nil @@ -1102,8 +1102,8 @@ func TestManager_CodexTerminalRenewalRecoveryLifecycle_Issue5783(t *testing.T) { if gotSession := verifyOpts.Headers.Get("X-Session-Id"); gotSession == "" || len(gotSession) < 7 || gotSession[:7] != "verify-" { t.Fatalf("VerifyInference X-Session-Id = %q, want isolated verify-* session", gotSession) } - if !containsSubstr(string(verifyReq.Payload), `"max_output_tokens":1`) || !containsSubstr(string(verifyReq.Payload), `"store":false`) { - t.Fatalf("VerifyInference payload = %s, want minimal 1-token non-stored prompt", string(verifyReq.Payload)) + if !containsSubstr(string(verifyReq.Payload), `"max_output_tokens":16`) || !containsSubstr(string(verifyReq.Payload), `"store":false`) { + t.Fatalf("VerifyInference payload = %s, want minimal 16-token non-stored prompt", string(verifyReq.Payload)) } if !afterVerify.InferenceVerified || afterVerify.RecoveryStageAt(now) != RecoveryStageInferenceVerified { t.Fatalf("expected InferenceVerified=true and stage=inference_verified after VerifyInference, got %v / %q", afterVerify.InferenceVerified, afterVerify.RecoveryStageAt(now)) diff --git a/sdk/cliproxy/auth/conductor_weighted_pinned_routing_test.go b/sdk/cliproxy/auth/conductor_weighted_pinned_routing_test.go new file mode 100644 index 00000000000..4406b509c1d --- /dev/null +++ b/sdk/cliproxy/auth/conductor_weighted_pinned_routing_test.go @@ -0,0 +1,121 @@ +package auth + +import ( + "context" + "fmt" + "maps" + "net/http" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestManagerWeightedPinsShareCreditsWithoutAffinity(t *testing.T) { + for _, scenario := range []string{"single", "mixed-single", "mixed-providers", "lower-priority", "websocket", "http-pin", "mixed-single-websocket", "xai-websocket", "mixed-single-xai"} { + t.Run(scenario, func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + manager := NewManager(nil, weighted, nil) + const model = "weighted-pin-routing" + providers := []string{"codex"} + if scenario == "xai-websocket" || scenario == "mixed-single-xai" { + providers[0] = "xai" + } + mixed := scenario == "mixed-single" || scenario == "mixed-providers" || scenario == "mixed-single-websocket" || scenario == "mixed-single-xai" + if scenario == "mixed-providers" { + providers = append(providers, "openai") + } + for _, provider := range providers { + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + } + a := &Auth{ID: t.Name() + "-a", Provider: providers[0], Attributes: map[string]string{AttributeWeight: "3"}} + b := &Auth{ID: t.Name() + "-b", Provider: providers[len(providers)-1], Attributes: map[string]string{AttributeWeight: "1"}} + auths := []*Auth{a, b} + // A zero-weight higher tier must not hide eligible positive-weight peers. + auths = append(auths, &Auth{ID: t.Name() + "-zero", Provider: providers[0], Attributes: map[string]string{AttributeWeight: "0", "priority": "100"}}) + if scenario == "lower-priority" { + auths = append(auths, &Auth{ID: t.Name() + "-higher", Provider: providers[0], Attributes: map[string]string{"priority": "1"}}) + } + ctx := t.Context() + if scenario == "websocket" || scenario == "http-pin" || scenario == "mixed-single-websocket" || scenario == "xai-websocket" || scenario == "mixed-single-xai" { + ctx = cliproxyexecutor.WithDownstreamWebsocket(ctx) + a.Attributes["websockets"] = "true" + b.Attributes["websockets"] = "true" + otherWebsocket := "false" + if scenario == "http-pin" { + a.Attributes["websockets"] = "false" + b.Attributes["websockets"] = "false" + otherWebsocket = "true" + } + auths = append(auths, &Auth{ID: t.Name() + "-other-transport", Provider: providers[0], Attributes: map[string]string{"websockets": otherWebsocket}}) + } + for _, auth := range auths { + auth.Status = StatusActive + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + // Normalize single/mixed entrypoints while retaining their real routing decisions. + pick := func(ctx context.Context, opts cliproxyexecutor.Options) *Auth { + t.Helper() + var auth *Auth + var errPick error + if mixed { + auth, _, _, errPick = manager.pickNextMixed(ctx, append([]string{"missing-executor"}, providers...), model, opts, nil) + } else { + auth, _, errPick = manager.pickNext(ctx, providers[0], model, opts, nil) + } + if errPick != nil || auth == nil { + t.Fatalf("selection failed: auth=%v err=%v", auth, errPick) + } + return auth + } + for range 4 { + if auth := pick(ctx, cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: b.ID}}); auth.ID != b.ID { + t.Fatalf("pin escaped: got %q, want %q", auth.ID, b.ID) + } + } + namespace := providers[0] + if mixed { + namespace = "mixed" + } + state := weightedStateForTest(t, weighted, namespace+":"+model, a.ID, b.ID) + if !maps.Equal(state.currentForTest(), map[string]int64{a.ID: 12, b.ID: -12}) || state.inFlight != 0 { + t.Fatalf("wrong pin accounting: current=%v inFlight=%d", state.currentForTest(), state.inFlight) + } + if scenario == "lower-priority" || scenario == "http-pin" { + // Traffic in a different priority/transport pool must not repay this pool. + for range 4 { + pick(ctx, cliproxyexecutor.Options{}) + } + if !maps.Equal(state.currentForTest(), map[string]int64{a.ID: 12, b.ID: -12}) { + t.Fatalf("unrelated pool consumed pin debt: %v", state.currentForTest()) + } + if scenario == "lower-priority" { + manager.Remove(WithSkipPersist(ctx), t.Name()+"-higher") + } else { + manager.Remove(WithSkipPersist(ctx), t.Name()+"-other-transport") + } + } + for i := range 12 { + // Headers must not enable sticky bindings when affinity is disabled. + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}} + if auth := pick(ctx, opts); auth.ID != a.ID { + t.Fatalf("request %d failed to repay pinned debt: got %q, want %q", i, auth.ID, a.ID) + } + } + if !maps.Equal(state.currentForTest(), map[string]int64{a.ID: 0, b.ID: 0}) || len(state.pendingForTest()) != 0 || state.inFlight != 0 { + t.Fatalf("unsettled compensation: current=%v pending=%v inFlight=%d", state.currentForTest(), state.pendingForTest(), state.inFlight) + } + counts := map[string]int{} + for range 40 { + counts[pick(ctx, cliproxyexecutor.Options{}).ID]++ + } + if counts[a.ID] != 30 || counts[b.ID] != 10 { + t.Fatalf("ordinary routing lost 3:1 ratio after repayment: %v", counts) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_weighted_pinned_test.go b/sdk/cliproxy/auth/conductor_weighted_pinned_test.go new file mode 100644 index 00000000000..c3a213872cc --- /dev/null +++ b/sdk/cliproxy/auth/conductor_weighted_pinned_test.go @@ -0,0 +1,172 @@ +package auth + +import ( + "fmt" + "maps" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" +) + +func TestManagerWeightedPinnedSessionCompensation(t *testing.T) { + for _, mixed := range []bool{false, true} { + for _, mode := range []string{"sticky", "cold", "no-session", "lcp", "no-affinity", "no-affinity-legacy"} { + t.Run(fmt.Sprintf("mixed=%t/%s", mixed, mode), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + manager := NewManager(nil, selector, nil) + if mode == "no-affinity" || mode == "no-affinity-legacy" { + manager.SetSelector(weighted) + } + const provider, model = "weighted-pinned", "weighted-pinned-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + ctx := withRequiredAuthKind(t.Context(), AuthKindAPIKey) + aID, bID := t.Name()+"-a", t.Name()+"-b" + triedID := t.Name() + "-tried" + excluded := []*Auth{ + {ID: t.Name() + "-disabled", Disabled: true}, + {ID: t.Name() + "-zero", Attributes: map[string]string{AttributeWeight: "0"}}, + {ID: t.Name() + "-cooldown", ModelStates: map[string]*ModelState{ + model: {Unavailable: true, NextRetryAfter: time.Now().Add(time.Hour)}, + }}, + {ID: t.Name() + "-wrong-model"}, + {ID: t.Name() + "-wrong-kind", Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}}, + {ID: t.Name() + "-wrong-provider", Provider: "other-provider"}, + {ID: t.Name() + "-no-executor", Provider: "no-executor"}, + {ID: t.Name() + "-lower-priority", Attributes: map[string]string{"priority": "0"}}, + {ID: triedID}, + } + auths := append([]*Auth{ + {ID: aID, Attributes: map[string]string{AttributeWeight: "3", "priority": "1"}}, + {ID: bID, Attributes: map[string]string{AttributeWeight: "1", "priority": "1"}}, + }, excluded...) + for _, auth := range auths { + if auth.Provider == "" { + auth.Provider = provider + } + auth.Status = StatusActive + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + if _, ok := auth.Attributes[AttributeWeight]; !ok { + auth.Attributes[AttributeWeight] = "100" + } + if _, ok := auth.Attributes["priority"]; !ok { + auth.Attributes["priority"] = "2" + } + if _, ok := auth.Attributes[AttributeAuthKind]; !ok { + auth.Attributes[AttributeAuthKind] = AuthKindAPIKey + } + registeredModel := model + if auth.ID == t.Name()+"-wrong-model" { + registeredModel = "other-model" + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: registeredModel}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + namespace := provider + if mixed { + namespace = "mixed" + } + stickyHeaders := http.Header{"Session_id": {"pinned-session"}} + id := CanonicalSessionID(stickyHeaders, nil, nil) + bindingKey := namespace + "::" + id + "::" + model + if mode == "sticky" { + selector.cache.Set(bindingKey, bID) + } + tried := map[string]struct{}{triedID: {}} + // Exercise the production entrypoints as well as direct legacy selection. + pick := func(opts cliproxyexecutor.Options) (*Auth, error) { + if mode == "no-affinity" { + if mixed { + auth, _, _, errPick := manager.pickNextMixed(ctx, []string{provider, "no-executor"}, model, opts, tried) + return auth, errPick + } + auth, _, errPick := manager.pickNext(ctx, provider, model, opts, tried) + return auth, errPick + } + if mixed { + auth, _, _, errPick := manager.pickNextMixedLegacy(ctx, []string{provider, "no-executor"}, model, opts, tried) + return auth, errPick + } + auth, _, errPick := manager.pickNextLegacy(ctx, provider, model, opts, tried) + return auth, errPick + } + for i := range 4 { + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: bID}} + switch mode { + case "sticky", "cold": + opts.Headers = stickyHeaders + case "lcp": + opts.SourceFormat = sdktranslator.FormatOpenAI + opts.OriginalRequest = []byte(`{"messages":[{"role":"user","content":"pinned conversation"}]}`) + opts.Metadata[cliproxyexecutor.CallerScopeMetadataKey] = "pinned-caller" + } + if auth, errPick := pick(opts); errPick != nil || auth == nil || auth.ID != bID { + t.Fatalf("pinned request %d: auth=%v err=%v, want %q", i, auth, errPick, bID) + } + } + state := weightedStateForTest(t, weighted, namespace+":"+model, aID, bID) + if state == nil || !maps.Equal(state.currentForTest(), map[string]int64{aID: 12, bID: -12}) || len(state.pendingForTest()) != 0 { + t.Fatalf("pinned traffic did not consume the eligible pool's credit: %#v", state) + } + counts := map[string]int{} + for i := range 12 { + auth, errPick := pick(cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}}) + if errPick != nil || auth == nil { + t.Fatalf("fresh request %d: auth=%v err=%v", i, auth, errPick) + } + counts[auth.ID]++ + } + if counts[aID] != 12 || counts[bID] != 0 { + t.Fatalf("fresh distribution = %v, want a=12 b=0", counts) + } + if !maps.Equal(state.currentForTest(), map[string]int64{aID: 0, bID: 0}) || len(state.pendingForTest()) != 0 { + t.Fatalf("compensation did not balance: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + if mode == "sticky" || mode == "cold" { + if binding, ok := selector.cache.Get(bindingKey); !ok || binding != bID { + t.Fatalf("compensation changed pinned binding: %q, found=%t", binding, ok) + } + } + // A broader accounting pool must never make an invalid pin selectable. + for _, pinnedID := range []string{excluded[0].ID, excluded[1].ID, excluded[2].ID, excluded[3].ID, excluded[4].ID, excluded[5].ID, excluded[6].ID, triedID, "missing"} { + opts := cliproxyexecutor.Options{ + Headers: stickyHeaders, + Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: pinnedID}, + } + if auth, errPick := pick(opts); errPick == nil || auth != nil { + t.Fatalf("invalid pin %q escaped pinning checks: auth=%v err=%v", pinnedID, auth, errPick) + } + if !maps.Equal(state.currentForTest(), map[string]int64{aID: 0, bID: 0}) || len(state.pendingForTest()) != 0 { + t.Fatalf("invalid pin %q changed accounting: current=%v pending=%v", pinnedID, state.currentForTest(), state.pendingForTest()) + } + } + // A valid lower-priority pin still wins without charging another tier. + lowerID := t.Name() + "-lower-priority" + lowerOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session_id": {"lower-pinned-session"}}, + Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: lowerID}, + } + if auth, errPick := pick(lowerOpts); errPick != nil || auth == nil || auth.ID != lowerID { + t.Fatalf("lower-priority pin was not enforced: auth=%v err=%v", auth, errPick) + } + if !maps.Equal(state.currentForTest(), map[string]int64{aID: 0, bID: 0}) || len(state.pendingForTest()) != 0 { + t.Fatalf("lower-priority pin changed higher-tier credits: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + lowerState := weightedStateForTest(t, weighted, namespace+":"+model, lowerID) + if !maps.Equal(lowerState.currentForTest(), map[string]int64{lowerID: 0}) || len(lowerState.pendingForTest()) != 0 { + t.Fatalf("lower-priority pin accounted the wrong pool: %#v", lowerState) + } + }) + } + } +} diff --git a/sdk/cliproxy/auth/cooldown_view.go b/sdk/cliproxy/auth/cooldown_view.go index 48324fe0493..9ad5fc2968a 100644 --- a/sdk/cliproxy/auth/cooldown_view.go +++ b/sdk/cliproxy/auth/cooldown_view.go @@ -51,31 +51,38 @@ func CooldownSnapshotForAuth(auth *Auth, now time.Time) []CooldownView { for _, key := range keys { state := auth.ModelStates[key] model := canonicalModelKey(key) - if state == nil || model == "" { + lookupKey := canonicalModelLookupKey(model) + if state == nil || model == "" || lookupKey == "" { continue } blocked, reason, next := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now) if !blocked || !next.After(now) { continue } - if previous, ok := byModel[model]; ok { + modelName := model + if previous, ok := byModel[lookupKey]; ok { + if previous.view.ModelKey != "" { + modelName = previous.view.ModelKey + } preferQuotaTie := next.Equal(previous.view.RetryAt) && reason == blockReasonCooldown && previous.reason != blockReasonCooldown if !next.After(previous.view.RetryAt) && !preferQuotaTie { continue } } - byModel[model] = modelCooldown{ - view: newCooldownView("model", model, next, now, state.Quota, state.StatusMessage, state.LastError), + byModel[lookupKey] = modelCooldown{ + view: newCooldownView("model", modelName, next, now, state.Quota, state.StatusMessage, state.LastError), reason: reason, } } models := make([]string, 0, len(byModel)) - for model := range byModel { - models = append(models, model) - } - sort.Strings(models) - for _, model := range models { - views = append(views, byModel[model].view) + for lookupKey := range byModel { + models = append(models, lookupKey) + } + sort.Slice(models, func(i, j int) bool { + return byModel[models[i]].view.ModelKey < byModel[models[j]].view.ModelKey + }) + for _, lookupKey := range models { + views = append(views, byModel[lookupKey].view) } return views } diff --git a/sdk/cliproxy/auth/error_events.go b/sdk/cliproxy/auth/error_events.go index aa85220c0fe..5f18053b322 100644 --- a/sdk/cliproxy/auth/error_events.go +++ b/sdk/cliproxy/auth/error_events.go @@ -104,7 +104,7 @@ func errorEventModelStatusFrom(model string, authSnapshot *Auth) *errorEventMode if model == "" || authSnapshot == nil || authSnapshot.ModelStates == nil { return nil } - state := authSnapshot.ModelStates[model] + state := existingModelState(authSnapshot, model) if state == nil { return nil } diff --git a/sdk/cliproxy/auth/errors.go b/sdk/cliproxy/auth/errors.go index b54e14b78b2..9f94371485d 100644 --- a/sdk/cliproxy/auth/errors.go +++ b/sdk/cliproxy/auth/errors.go @@ -47,6 +47,12 @@ const ErrorCodeTransientRateLimit = "transient_rate_limit" const transientRateLimitErrorCode = ErrorCodeTransientRateLimit +// ErrorCodeCredentialQuota marks a rate-limit/quota refusal that applies to the +// entire credential across all models rather than only the reporting model. +const ErrorCodeCredentialQuota = "credential_quota" + +const credentialQuotaErrorCode = ErrorCodeCredentialQuota + // Error describes an authentication related failure in a provider agnostic format. type Error struct { // Code is a short machine readable identifier. diff --git a/sdk/cliproxy/auth/home_session_kind_test.go b/sdk/cliproxy/auth/home_session_kind_test.go new file mode 100644 index 00000000000..76203a8bf14 --- /dev/null +++ b/sdk/cliproxy/auth/home_session_kind_test.go @@ -0,0 +1,59 @@ +package auth + +import ( + "context" + "net/http" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +type sessionKindCaptureDispatcher struct { + fixtureHomeDispatcher + headers http.Header + sessionID string + parentSessionID string +} + +func (d *sessionKindCaptureDispatcher) RPopAuthWithSessionHierarchy(_ context.Context, _ string, sessionID, parentSessionID string, headers http.Header, _ int, _ string, _ *int, _ []string, _ string) ([]byte, error) { + d.headers = headers.Clone() + d.sessionID = sessionID + d.parentSessionID = parentSessionID + return d.payload, nil +} + +func TestHomeDispatchPreservesForkAndCompactionKind(t *testing.T) { + for _, tc := range []struct { + name string + body string + metadata map[string]any + want string + }{ + {name: "body fork", body: `{"thread_id":"child","metadata":{"forked_from_thread_id":"parent"}}`, want: "fork"}, + {name: "compaction", body: `{"session_id":"child","parent_session_id":"parent"}`, metadata: map[string]any{cliproxyexecutor.NodeKindMetadataKey: "compaction"}, want: "compaction"}, + {name: "ordinary child", body: `{"session_id":"child","parent_session_id":"parent"}`}, + } { + t.Run(tc.name, func(t *testing.T) { + dispatcher := &sessionKindCaptureDispatcher{fixtureHomeDispatcher: fixtureHomeDispatcher{payload: []byte(`{"auth":{"id":"kind-auth","provider":"kind-provider"}}`)}} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.RegisterExecutor(schedulerTestExecutor{provider: "kind-provider"}) + headers := http.Header{"User-Agent": {"test-client"}} + opts := cliproxyexecutor.Options{Headers: headers, OriginalRequest: []byte(tc.body), Metadata: tc.metadata} + opts.EnsureMetadata() + selection, errSelection := manager.pickHomeDispatchSelection(t.Context(), "model", opts) + if errSelection != nil { + t.Fatal(errSelection) + } + selection.End("test_complete") + if dispatcher.sessionID == "" || dispatcher.parentSessionID == "" || dispatcher.sessionID == dispatcher.parentSessionID { + t.Fatalf("lost session hierarchy: child=%q parent=%q", dispatcher.sessionID, dispatcher.parentSessionID) + } + if got := dispatcher.headers.Get("X-Node-Kind"); got != tc.want { + t.Fatalf("node kind=%q, want %q", got, tc.want) + } + if got := headers.Get("X-Node-Kind"); got != "" { + t.Fatalf("caller headers mutated: %q", got) + } + }) + } +} diff --git a/sdk/cliproxy/auth/home_usage_accounting_test.go b/sdk/cliproxy/auth/home_usage_accounting_test.go new file mode 100644 index 00000000000..e84a7141eeb --- /dev/null +++ b/sdk/cliproxy/auth/home_usage_accounting_test.go @@ -0,0 +1,76 @@ +package auth + +import ( + "context" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" +) + +type homeUsageAccountingExecutor struct { + retainingHomeExecutionExecutor + observed []*coreusage.HomeSelection +} + +func (e *homeUsageAccountingExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.observed = append(e.observed, coreusage.HomeSelectionFromContext(ctx)) + return e.retainingHomeExecutionExecutor.Execute(ctx, auth, req, opts) +} + +func (e *homeUsageAccountingExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.observed = append(e.observed, coreusage.HomeSelectionFromContext(ctx)) + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed","response":{"id":"test"}}`)} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func TestHomeUsageMarksOnlyReusedExecutions(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "execute", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &homeUsageAccountingExecutor{} + manager.RegisterExecutor(executor) + t.Cleanup(func() { manager.CloseExecutionSession("usage-session") }) + ctx := cliproxyexecutor.WithDownstreamWebsocket(t.Context()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "usage-session"}} + for _, model := range []string{"model-a", "model-a", "model-a", "model-b"} { + if stream { + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: model}, opts) + if errExecute != nil { + t.Fatal(errExecute) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + } else if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: model}, opts); errExecute != nil { + t.Fatal(errExecute) + } + } + if dispatcher.calls.Load() != 2 || len(executor.observed) != 4 { + t.Fatalf("dispatch=%d executions=%d", dispatcher.calls.Load(), len(executor.observed)) + } + for i, want := range []bool{false, true, true, false} { + got := executor.observed[i] + wantModel := "model-a" + if i == 3 { + wantModel = "model-b" + } + if got == nil || got.Reused != want || got.Model != wantModel { + t.Fatalf("turn %d: got=%+v want reused=%t model=%s", i, got, want, wantModel) + } + } + }) + } +} diff --git a/sdk/cliproxy/auth/oauth_model_alias.go b/sdk/cliproxy/auth/oauth_model_alias.go index c9e1af41914..279433cf377 100644 --- a/sdk/cliproxy/auth/oauth_model_alias.go +++ b/sdk/cliproxy/auth/oauth_model_alias.go @@ -89,7 +89,55 @@ func (m *Manager) SetOAuthModelAlias(aliases map[string][]internalconfig.OAuthMo if table == nil { table = &oauthModelAliasTable{} } - m.oauthModelAlias.Store(table) + m.publishOAuthModelAliasAndReconcileStates(nil, table) +} + +type oauthResolvedAliasRoute struct { + aliasKey string + targetKey string +} + +func (m *Manager) configuredOAuthAliasRoutes(auth *Auth) []oauthResolvedAliasRoute { + if m == nil || auth == nil { + return nil + } + channel := modelAliasChannel(auth) + if channel == "" { + return nil + } + seenAliases := make(map[string]struct{}) + var routes []oauthResolvedAliasRoute + addAlias := func(alias string) { + alias = strings.TrimSpace(alias) + if alias == "" { + return + } + aliasLookup := canonicalModelLookupKey(alias) + if aliasLookup == "" { + return + } + if _, exists := seenAliases[aliasLookup]; exists { + return + } + seenAliases[aliasLookup] = struct{}{} + aliasKey := canonicalModelKey(alias) + if targetKey := m.selectionModelKeyForAuth(auth, alias); targetKey != "" && !sameCanonicalModelKey(aliasKey, targetKey) { + routes = append(routes, oauthResolvedAliasRoute{ + aliasKey: aliasKey, + targetKey: targetKey, + }) + } + } + for _, entry := range OAuthModelAliasesFromAttributes(authAttributes(auth)) { + addAlias(entry.Alias) + } + raw := m.oauthModelAlias.Load() + if table, _ := raw.(*oauthModelAliasTable); table != nil && table.reverse != nil { + for _, entry := range table.reverse[channel] { + addAlias(entry.configAlias) + } + } + return routes } // applyOAuthModelAlias resolves the upstream model from OAuth model alias. diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index 752a52d7d57..b62e332d143 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -45,6 +45,7 @@ type scheduledGenerationMeta struct { type authScheduler struct { mu sync.Mutex strategy schedulerStrategy + weightedSelector *WeightedRoundRobinSelector providers map[string]*providerScheduler authProviders map[string]string authGenerations map[string]scheduledGenerationMeta @@ -153,8 +154,10 @@ 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) return &authScheduler{ strategy: selectorStrategy(selector), + weightedSelector: weighted, providers: make(map[string]*providerScheduler), authProviders: make(map[string]string), authGenerations: make(map[string]scheduledGenerationMeta), @@ -185,6 +188,7 @@ func (s *authScheduler) setSelector(selector Selector) { s.mu.Lock() defer s.mu.Unlock() s.strategy = selectorStrategy(selector) + s.weightedSelector, _ = selector.(*WeightedRoundRobinSelector) clear(s.mixedCursors) clear(s.mixedWeightedStates) } @@ -403,7 +407,7 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } providerKey := canonicalSchedulingProvider(provider) - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) eligibility := authSelectionEligibilityForRequest(ctx, opts) preferWebsocket := cliproxyexecutor.DownstreamWebsocket(ctx) && providerPrefersWebsocketTransport(providerKey) && pinnedAuthID == "" @@ -422,6 +426,34 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } predicate := scheduledAuthPredicate(eligibility, tried, pinnedAuthID, strategy == schedulerStrategyWeightedRoundRobin) + if strategy == schedulerStrategyWeightedRoundRobin { + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + shard.promoteExpiredLocked(time.Now()) + priority, ready := shard.highestReadyPriorityLocked(preferWebsocket, predicate) + if !ready { + return nil, shard.unavailableErrorLocked(provider, model, predicate) + } + // A pin controls selection, not the peers against which traffic is charged. + accountingPredicate := scheduledAuthPredicate(eligibility, tried, "", true) + 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 + } else { + entries = bucket.ws.flat + } + } + namespace := providerKey + if opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] == "mixed" { + namespace = "mixed" + } + route := namespace + ":" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) + return s.pickWeightedPoolLocked(route, entries, accountingPredicate, pinnedAuthID, nonWebsocketOnly) + } if picked := shard.pickReadyLocked(preferWebsocket, strategy, predicate); picked != nil { return picked, nil } @@ -465,14 +497,14 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) eligibility := authSelectionEligibilityForRequest(ctx, opts) - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) s.mu.Lock() defer s.mu.Unlock() if strategy == schedulerStrategyCurrent { strategy = s.strategy } - if pinnedAuthID != "" { + if pinnedAuthID != "" && strategy != schedulerStrategyWeightedRoundRobin { providerKey := s.authProviders[pinnedAuthID] if providerKey == "" || !containsProvider(normalized, providerKey) { return nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} @@ -513,6 +545,22 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s hasCandidate = true } } + if pinnedAuthID != "" { + providerKey := s.authProviders[pinnedAuthID] + if providerKey == "" || !containsProvider(normalized, providerKey) { + return nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} + } + shard := s.providers[providerKey].ensureModelLocked(modelKey, now) + if shard == nil { + return nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} + } + pinned := shard.entries[pinnedAuthID] + if pinned == nil || pinned.state != scheduledStateReady || !predicate(pinned) { + pinnedPredicate := scheduledAuthPredicate(eligibility, tried, pinnedAuthID, true) + return nil, "", shard.unavailableErrorLocked("mixed", model, pinnedPredicate) + } + bestPriority, hasCandidate = pinned.meta.priority, true + } if !hasCandidate { return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } @@ -552,20 +600,15 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s } return entries[i].auth.ID < entries[j].auth.ID }) - if s.mixedWeightedStates == nil { - s.mixedWeightedStates = make(map[string]*smoothWeightedState) + if ctx != nil && ctx.Err() != nil { + return nil, "", ctx.Err() } - state := s.mixedWeightedStates[cursorKey] - if state == nil { - state = &smoothWeightedState{} - s.mixedWeightedStates[cursorKey] = state - } - state.prepare(scheduledWeightVectorMatching(entries, predicate)) - picked := pickSmoothWeightedScheduled(entries, state.current, predicate) - if picked != nil && picked.meta != nil { - return picked.auth, picked.meta.providerKey, nil + route := "mixed:" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) + picked, errPick := s.pickWeightedPoolLocked(route, entries, predicate, pinnedAuthID, false) + if errPick != nil { + return nil, "", errPick } - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) + return picked, s.authProviders[picked.ID], nil } weights := make([]int, len(normalized)) @@ -644,7 +687,7 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st if providerState == nil { continue } - shard := providerState.ensureModelLocked(canonicalModelKey(model), now) + shard := providerState.ensureModelLocked(canonicalModelLookupKey(model), now) if shard == nil { continue } @@ -1077,7 +1120,7 @@ func supportedModelSetForAuth(authID string) map[string]struct{} { if model == nil { continue } - modelKey := canonicalModelKey(model.ID) + modelKey := canonicalModelLookupKey(model.ID) if modelKey == "" { continue } @@ -1130,7 +1173,7 @@ func (p *providerScheduler) updateTargetShardsLocked(meta *scheduledAuthMeta, ta seen := make(map[string]struct{}, len(targetModels)) for _, m := range targetModels { - modelKey := canonicalModelKey(m) + modelKey := canonicalModelLookupKey(m) if _, ok := seen[modelKey]; ok { continue } @@ -1166,7 +1209,7 @@ func (p *providerScheduler) ensureModelLocked(modelKey string, now time.Time) *m if p == nil { return nil } - modelKey = canonicalModelKey(modelKey) + modelKey = canonicalModelLookupKey(modelKey) if shard, ok := p.modelShards[modelKey]; ok && shard != nil { shard.promoteExpiredLocked(now) return shard @@ -1188,7 +1231,7 @@ func (p *providerScheduler) ensureModelLocked(modelKey string, now time.Time) *m // supportsModel reports whether the auth metadata currently supports modelKey. func (m *scheduledAuthMeta) supportsModel(modelKey string) bool { - modelKey = canonicalModelKey(modelKey) + modelKey = canonicalModelLookupKey(modelKey) if modelKey == "" { return true } @@ -1461,7 +1504,6 @@ func (m *modelScheduler) candidateErrorsLocked(model string, predicate func(*sch if m == nil { return nil, time.Time{}, "", nil, time.Time{}, "" } - modelKey := canonicalModelKey(model) for _, entry := range m.entries { if predicate != nil && !predicate(entry) { continue @@ -1473,15 +1515,7 @@ func (m *modelScheduler) candidateErrorsLocked(model string, predicate func(*sch var curModelErr error var curModelTime time.Time if len(auth.ModelStates) > 0 { - if state, ok := auth.ModelStates[model]; ok && state != nil { - if state.LastError != nil { - curModelErr = state.LastError - curModelTime = state.UpdatedAt - } else if strings.TrimSpace(state.StatusMessage) != "" { - curModelErr = errors.New(state.StatusMessage) - curModelTime = state.UpdatedAt - } - } else if state, ok := auth.ModelStates[modelKey]; ok && state != nil { + if state := existingModelState(auth, model); state != nil { if state.LastError != nil { curModelErr = state.LastError curModelTime = state.UpdatedAt diff --git a/sdk/cliproxy/auth/scheduler_benchmark_test.go b/sdk/cliproxy/auth/scheduler_benchmark_test.go index 85a9014353d..91817fd600e 100644 --- a/sdk/cliproxy/auth/scheduler_benchmark_test.go +++ b/sdk/cliproxy/auth/scheduler_benchmark_test.go @@ -214,3 +214,41 @@ func BenchmarkManagerPickNextAndMarkResult1000(b *testing.B) { manager.MarkResult(ctx, Result{AuthID: auth.ID, Provider: "gemini", Model: model, Success: true}) } } + +func BenchmarkManagerPinnedLegacy1000(b *testing.B) { + for _, mixed := range []bool{false, true} { + b.Run(fmt.Sprintf("mixed=%t", mixed), func(b *testing.B) { + manager, providers, model := benchmarkManagerSetup(b, 1000, mixed, false) + selector := NewSessionAffinitySelector(&RoundRobinSelector{}) + b.Cleanup(selector.Stop) + manager.SetSelector(selector) + ctx := context.Background() + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.PinnedAuthMetadataKey: "bench-gemini-0000", + }} + // Warm up index assignment and selection metadata outside the timed loop. + if mixed { + if _, _, _, errWarm := manager.pickNextMixedLegacy(ctx, providers, model, opts, nil); errWarm != nil { + b.Fatal(errWarm) + } + } else if _, _, errWarm := manager.pickNextLegacy(ctx, providers[0], model, opts, nil); errWarm != nil { + b.Fatal(errWarm) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if mixed { + auth, exec, provider, errPick := manager.pickNextMixedLegacy(ctx, providers, model, opts, nil) + if errPick != nil || auth == nil || auth.ID != "bench-gemini-0000" || exec == nil || provider != providers[0] { + b.Fatalf("pinned mixed pick failed: auth=%v exec=%v provider=%q err=%v", auth, exec, provider, errPick) + } + } else { + auth, exec, errPick := manager.pickNextLegacy(ctx, providers[0], model, opts, nil) + if errPick != nil || auth == nil || auth.ID != "bench-gemini-0000" || exec == nil { + b.Fatalf("pinned pick failed: auth=%v exec=%v err=%v", auth, exec, errPick) + } + } + } + }) + } +} diff --git a/sdk/cliproxy/auth/scheduler_weighted.go b/sdk/cliproxy/auth/scheduler_weighted.go new file mode 100644 index 00000000000..4b37cc5402c --- /dev/null +++ b/sdk/cliproxy/auth/scheduler_weighted.go @@ -0,0 +1,34 @@ +package auth + +// pickWeightedPoolLocked consumes the scheduler's ID-sorted, immutable ready +// 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) { + if s.weightedSelector == nil { + return nil, &Error{Code: "auth_unavailable", Message: "weighted selector is not configured"} + } + auths := make([]*Auth, 0, len(entries)) + ids := make([]string, 0, len(entries)) + weights := make([]int64, 0, len(entries)) + for _, entry := range entries { + if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { + continue + } + if predicate != nil && !predicate(entry) { + continue + } + if nonWebsocketOnly && entry.meta.websocketEnabled { + continue + } + auths = append(auths, entry.auth) + ids = append(ids, entry.auth.ID) + weights = append(weights, entry.meta.weight) + } + picked, reservation := s.weightedSelector.reservePool(route, auths, weights, ids, pinnedAuthID) + if picked == nil { + return nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} + } + reservation.finish(true) + return picked, nil +} diff --git a/sdk/cliproxy/auth/scheduler_weighted_routing_test.go b/sdk/cliproxy/auth/scheduler_weighted_routing_test.go new file mode 100644 index 00000000000..cd9e5f543b3 --- /dev/null +++ b/sdk/cliproxy/auth/scheduler_weighted_routing_test.go @@ -0,0 +1,161 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestWeightedWebsocketPreferenceAcrossPriorities(t *testing.T) { + for _, provider := range []string{"codex", "xai"} { + for _, mixed := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/mixed=%t", provider, mixed), func(t *testing.T) { + manager := NewManager(nil, &WeightedRoundRobinSelector{}, nil) + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + ctx := cliproxyexecutor.WithDownstreamWebsocket(t.Context()) + const model = "weighted-ws-priority" + for _, auth := range []*Auth{ + {ID: t.Name() + "-http", Provider: provider, Attributes: map[string]string{"priority": "10"}}, + {ID: t.Name() + "-ws", Provider: provider, Attributes: map[string]string{"priority": "0", "websockets": "true"}}, + } { + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + for _, routeModel := range []string{"", model} { + var got *Auth + var errPick error + if mixed { + got, _, _, errPick = manager.pickNextMixed(ctx, []string{provider}, routeModel, cliproxyexecutor.Options{}, nil) + } else { + got, _, errPick = manager.pickNext(ctx, provider, routeModel, cliproxyexecutor.Options{}, nil) + } + if errPick != nil || got == nil { + t.Fatalf("model=%q: selection failed: auth=%v err=%v", routeModel, got, errPick) + } + if got.ID != t.Name()+"-ws" { + t.Fatalf("model=%q: got %q, want lower-priority WebSocket auth", routeModel, got.ID) + } + } + }) + } + } +} + +func TestWeightedSchedulerAndLegacyShareCredits(t *testing.T) { + for _, mixed := range []bool{false, true} { + for _, pinViaLegacy := range []bool{false, true} { + t.Run(fmt.Sprintf("mixed=%t/pinViaLegacy=%t", mixed, pinViaLegacy), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + manager := NewManager(nil, weighted, nil) + const provider, model = "weighted-shared", "weighted-shared-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + for _, auth := range []*Auth{ + {ID: t.Name() + "-a", Provider: provider, Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: t.Name() + "-b", Provider: provider, Attributes: map[string]string{AttributeWeight: "1"}}, + } { + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(t.Context()), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + if !manager.useSchedulerFastPath() { + t.Fatal("bare weighted routing must retain the scheduler fast path") + } + // Normalize the four entrypoints to verify accounting across boundaries. + pick := func(ctx context.Context, legacy bool, opts cliproxyexecutor.Options) (*Auth, error) { + if mixed { + if legacy { + auth, _, _, errPick := manager.pickNextMixedLegacy(ctx, []string{provider}, model, opts, nil) + return auth, errPick + } + auth, _, _, errPick := manager.pickNextMixed(ctx, []string{provider}, model, opts, nil) + return auth, errPick + } + if legacy { + auth, _, errPick := manager.pickNextLegacy(ctx, provider, model, opts, nil) + return auth, errPick + } + auth, _, errPick := manager.pickNext(ctx, provider, model, opts, nil) + return auth, errPick + } + for range 4 { + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: t.Name() + "-b"}} + auth, errPick := pick(t.Context(), pinViaLegacy, opts) + if errPick != nil || auth == nil || auth.ID != t.Name()+"-b" { + t.Fatalf("pin failed: auth=%v err=%v", auth, errPick) + } + } + canceled, cancel := context.WithCancel(t.Context()) + cancel() + if auth, errPick := pick(canceled, false, cliproxyexecutor.Options{}); auth != nil || !errors.Is(errPick, context.Canceled) { + t.Fatalf("canceled selection: auth=%v err=%v", auth, errPick) + } + for i := range 12 { + auth, errPick := pick(t.Context(), !pinViaLegacy, cliproxyexecutor.Options{}) + if errPick != nil || auth == nil || auth.ID != t.Name()+"-a" { + t.Fatalf("repayment %d used a different ledger: auth=%v err=%v", i, auth, errPick) + } + } + }) + } + } +} + +func TestWeightedSchedulerSeesPendingReservations(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "pending-a", Provider: "weighted-pending"}, {ID: "pending-b", Provider: "weighted-pending"}} + scheduler := newSchedulerForTest(weighted, auths...) + _, reservation := weighted.reserve("weighted-pending:", auths, auths[0].ID) + t.Cleanup(func() { reservation.finish(false) }) + picked, errPick := scheduler.pickSingle(t.Context(), "weighted-pending", "", cliproxyexecutor.Options{}, nil) + if errPick != nil || picked == nil || picked.ID != auths[1].ID { + t.Fatalf("scheduler ignored pending credit: auth=%v err=%v", picked, errPick) + } + reservation.finish(false) + picked, errPick = scheduler.pickSingle(t.Context(), "weighted-pending", "", cliproxyexecutor.Options{}, nil) + if errPick != nil || picked == nil || picked.ID != auths[0].ID { + t.Fatalf("rollback lost committed credit: auth=%v err=%v", picked, errPick) + } + state := weightedStateForTest(t, weighted, "weighted-pending:", auths[0].ID, auths[1].ID) + if state.inFlight != 0 || len(state.pendingForTest()) != 0 || state.currentForTest()[auths[0].ID] != 0 || state.currentForTest()[auths[1].ID] != 0 { + t.Fatalf("pool did not settle: current=%v pending=%v inFlight=%d", state.currentForTest(), state.pendingForTest(), state.inFlight) + } +} + +func BenchmarkManagerWeighted1000(b *testing.B) { + for _, mixed := range []bool{false, true} { + b.Run(fmt.Sprintf("mixed=%t", mixed), func(b *testing.B) { + manager, providers, model := benchmarkManagerSetup(b, 1000, mixed, false) + manager.SetSelector(&WeightedRoundRobinSelector{}) + ctx := b.Context() + opts := cliproxyexecutor.Options{} + // Warm both the shard indexes and the pool state outside the timed loop. + if mixed { + if _, _, _, errWarm := manager.pickNextMixed(ctx, providers, model, opts, nil); errWarm != nil { + b.Fatal(errWarm) + } + } else if _, _, errWarm := manager.pickNext(ctx, providers[0], model, opts, nil); errWarm != nil { + b.Fatal(errWarm) + } + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + if mixed { + if _, _, _, errPick := manager.pickNextMixed(ctx, providers, model, opts, nil); errPick != nil { + b.Fatal(errPick) + } + } else if _, _, errPick := manager.pickNext(ctx, providers[0], model, opts, nil); errPick != nil { + b.Fatal(errPick) + } + } + }) + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 9d88187de83..24c76560e03 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -39,9 +39,10 @@ type RoundRobinSelector struct { // WeightedRoundRobinSelector provides smooth weighted round-robin selection. type WeightedRoundRobinSelector struct { - mu sync.Mutex - states map[string]*smoothWeightedState - maxKeys int + mu sync.Mutex + states map[[32]byte]*weightedPoolState + stateCount int + maxKeys int } type smoothWeightedState struct { @@ -411,6 +412,18 @@ func canonicalModelKey(model string) string { return modelName } +func canonicalModelLookupKey(model string) string { + return strings.ToLower(canonicalModelKey(model)) +} + +func sameCanonicalModelKey(left, right string) bool { + leftKey := canonicalModelLookupKey(left) + if leftKey == "" { + return false + } + return leftKey == canonicalModelLookupKey(right) +} + func authWebsocketsEnabled(auth *Auth) bool { if auth == nil { return false @@ -628,7 +641,7 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o return nil, err } available = preferCodexWebsocketAuths(ctx, provider, available) - key := provider + ":" + canonicalModelKey(model) + key := provider + ":" + canonicalModelLookupKey(model) s.mu.Lock() defer s.mu.Unlock() if s.lastPicked == nil { @@ -680,38 +693,11 @@ func positiveWeightAuths(auths []*Auth) []*Auth { // Pick selects the next available auth using smooth weighted round-robin. func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { - _ = opts - available, errAvailable := getSelectorAvailableAuths(ctx, positiveWeightAuths(auths), provider, model, time.Now()) - if errAvailable != nil { - return nil, errAvailable - } - available = preferCodexWebsocketAuths(ctx, provider, available) - stateModel := weightedSelectorStateModel(ctx, model) - key := provider + ":" + canonicalModelKey(stateModel) - - s.mu.Lock() - defer s.mu.Unlock() - if s.states == nil { - s.states = make(map[string]*smoothWeightedState) - } - limit := s.maxKeys - if limit <= 0 { - limit = 4096 - } - if _, ok := s.states[key]; !ok && len(s.states) >= limit { - s.states = make(map[string]*smoothWeightedState) - } - state := s.states[key] - if state == nil { - state = &smoothWeightedState{} - s.states[key] = state - } - weights := authWeightVector(available) - state.prepare(weights) - picked := pickSmoothWeightedAuth(available, state.current) - if picked == nil { - return nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} + picked, reservation, errPick := s.reserveForRequest(ctx, provider, model, opts, auths, nil) + if errPick != nil { + return nil, errPick } + reservation.finish(true) return picked, nil } @@ -784,29 +770,6 @@ func authWeightVector(auths []*Auth) map[string]int64 { return weights } -func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { - var picked *Auth - var pickedCurrent int64 - var totalWeight int64 - for _, auth := range auths { - weight := authWeight(auth) - if auth == nil || weight <= 0 { - continue - } - current[auth.ID] = saturatingAddInt64(current[auth.ID], weight) - totalWeight = saturatingAddInt64(totalWeight, weight) - if picked == nil || current[auth.ID] > pickedCurrent { - picked = auth - pickedCurrent = current[auth.ID] - } - } - if picked == nil { - return nil - } - current[picked.ID] = saturatingAddInt64(current[picked.ID], -totalWeight) - return picked -} - func saturatingAddInt64(value, delta int64) int64 { if delta > 0 && value > math.MaxInt64-delta { return math.MaxInt64 @@ -846,6 +809,9 @@ func isAuthBlockedForModelForRequest(ctx context.Context, auth *Auth, model stri if auth.Disabled || auth.Status == StatusDisabled { return true, blockReasonDisabled, time.Time{} } + if codexRefreshMissingAccessForRequest(ctx, auth, clientHeaders...) { + return true, blockReasonOther, time.Time{} + } if hasUnauthorizedAuthFailure(auth) { return true, blockReasonOther, time.Time{} } @@ -855,15 +821,16 @@ func isAuthBlockedForModelForRequest(ctx context.Context, auth *Auth, model stri if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { return true, blockReasonCooldown, auth.Quota.NextRecoverAt } + unavailable := codexRefreshUnavailableForRequest(ctx, auth, now, clientHeaders...) if model != "" { if len(auth.ModelStates) > 0 { - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) matched := false blocked := false blockedReason := blockReasonNone nextRetry := time.Time{} for stateModel, state := range auth.ModelStates { - if state == nil || canonicalModelKey(stateModel) != modelKey { + if state == nil || canonicalModelLookupKey(stateModel) != modelKey { continue } matched = true @@ -888,17 +855,17 @@ func isAuthBlockedForModelForRequest(ctx context.Context, auth *Auth, model stri } return false, blockReasonNone, time.Time{} } - return availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) + return availabilityBlock(unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) } quotaExceeded := auth.Quota.Exceeded // When model is empty and the credential has individual model states, auth.Quota.Exceeded // is an aggregate of single-model quota cooldowns (reason "quota"). As long as the credential // itself is not unavailable (not all models failed) and not under a credential-wide quota, // do not treat individual model cooldowns as blocking the entire credential. - if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !auth.Unavailable { + if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !unavailable { quotaExceeded = false } - return availabilityBlock(auth.Unavailable, quotaExceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) + return availabilityBlock(unavailable, quotaExceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) } func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextRecoverAt, now time.Time) (bool, blockReason, time.Time) { @@ -934,6 +901,7 @@ type SessionAffinitySelector struct { matcher *cliproxysession.MerklePrefixMatcher subagentAffinity bool rebindOnQuotaOnly bool + rebindOnNever bool } const ( @@ -942,12 +910,17 @@ const ( // SessionAffinityRebindOnQuota preserves the existing session binding across transient failures // and only rebinds on quota, auth, model-support, or permanent unavailability. SessionAffinityRebindOnQuota = "quota" + // SessionAffinityRebindOnNever preserves the existing session binding across all failures + // and refuses to fail over a bound session to another credential. + SessionAffinityRebindOnNever = "never" ) func normalizeSessionAffinityRebindOn(value string) string { switch strings.ToLower(strings.TrimSpace(value)) { case SessionAffinityRebindOnQuota: return SessionAffinityRebindOnQuota + case SessionAffinityRebindOnNever: + return SessionAffinityRebindOnNever default: return SessionAffinityRebindOnAny } @@ -999,6 +972,9 @@ func isAuthBlockedByQuotaOrAuthForModelForRequest(ctx context.Context, auth *Aut if auth.Disabled || auth.Status == StatusDisabled { return true } + if codexRefreshMissingAccessForRequest(ctx, auth, clientHeaders...) { + return true + } if hasUnauthorizedAuthFailure(auth) { return true } @@ -1009,9 +985,9 @@ func isAuthBlockedByQuotaOrAuthForModelForRequest(ctx context.Context, auth *Aut return true } if model != "" && len(auth.ModelStates) > 0 { - modelKey := canonicalModelKey(model) + modelKey := canonicalModelLookupKey(model) for stateModel, state := range auth.ModelStates { - if state == nil || canonicalModelKey(stateModel) != modelKey { + if state == nil || canonicalModelLookupKey(stateModel) != modelKey { continue } if state.Status == StatusDisabled { @@ -1024,10 +1000,11 @@ func isAuthBlockedByQuotaOrAuthForModelForRequest(ctx context.Context, auth *Aut return false } quota := auth.Quota - if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !auth.Unavailable { + unavailable := codexRefreshUnavailableForRequest(ctx, auth, now, clientHeaders...) + if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !unavailable { quota.Exceeded = false } - return isStateBlockedByQuotaOrAuth(auth.Unavailable, quota, auth.NextRetryAfter, auth.LastError, now) + return isStateBlockedByQuotaOrAuth(unavailable, quota, auth.NextRetryAfter, auth.LastError, now) } func isStateBlockedByQuotaOrAuth(unavailable bool, quota QuotaState, nextRetryAfter time.Time, lastErr *Error, now time.Time) bool { @@ -1045,7 +1022,7 @@ func isStateBlockedByQuotaOrAuth(unavailable bool, quota QuotaState, nextRetryAf } func (s *SessionAffinitySelector) shouldPreserveBindingWhenUnavailable(ctx context.Context, candidates []*Auth, model, authID string, now time.Time) bool { - if s == nil || !s.rebindOnQuotaOnly || authID == "" { + if s == nil || (!s.rebindOnQuotaOnly && !s.rebindOnNever) || authID == "" { return false } if ctx != nil { @@ -1063,6 +1040,23 @@ func (s *SessionAffinitySelector) shouldPreserveBindingWhenUnavailable(ctx conte return false } +func (s *SessionAffinitySelector) unavailableBoundAuthError(ctx context.Context, candidates []*Auth, model, authID string, now time.Time) *Error { + if s.shouldPreserveBindingWhenUnavailable(ctx, candidates, model, authID, now) { + return &Error{ + Code: "session_auth_retry", + Message: "bound session credential is transiently unavailable", + Retryable: true, + HTTPStatus: http.StatusServiceUnavailable, + } + } + return &Error{ + Code: "session_auth_unavailable", + Message: "bound session credential is unavailable", + Retryable: false, + HTTPStatus: http.StatusTooManyRequests, + } +} + // SessionAffinityConfig configures the session affinity selector. type SessionAffinityConfig struct { Fallback Selector @@ -1091,13 +1085,14 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff if cfg.SubagentAffinity != nil { subagentAffinity = *cfg.SubagentAffinity } - rebindOnQuotaOnly := normalizeSessionAffinityRebindOn(cfg.RebindOn) == SessionAffinityRebindOnQuota + rebindMode := normalizeSessionAffinityRebindOn(cfg.RebindOn) return &SessionAffinitySelector{ fallback: cfg.Fallback, cache: NewSessionCache(cfg.TTL), matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), subagentAffinity: subagentAffinity, - rebindOnQuotaOnly: rebindOnQuotaOnly, + rebindOnQuotaOnly: rebindMode == SessionAffinityRebindOnQuota, + rebindOnNever: rebindMode == SessionAffinityRebindOnNever, } } @@ -1120,7 +1115,16 @@ func (s *SessionAffinitySelector) Trees() *cliproxysession.InMemorySessionTreeSt // Note: The cache key includes provider, session ID, and model to handle cases where // a session uses multiple models (e.g., gemini-2.5-pro and gemini-3-flash-preview) // that may be supported by different auth credentials, and to avoid cross-provider conflicts. -func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { +func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (selected *Auth, errPick error) { + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + var reservation *weightedReservation + // Local bindings complete inside Pick; failed selections release speculative credit. + defer func() { reservation.finish(selected != nil && errPick == nil) }() entry := selectorLogEntry(ctx) if opts.Metadata == nil { opts.Metadata = make(map[string]any) @@ -1132,7 +1136,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri // consulted when no header, body, or execution-session identity is present. explicitID, explicitFallbackID := extractExplicitSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) if explicitID == "" { - if auth, handled, errLCP := s.pickLCP(ctx, provider, model, opts, auths, entry); handled || errLCP != nil { + if auth, handled, errLCP := s.pickLCP(ctx, provider, model, opts, auths, entry, &reservation); handled || errLCP != nil { return auth, errLCP } } else if opts.Metadata != nil { @@ -1174,13 +1178,13 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri return nil, errAvailable } entry.Debugf("session-affinity: no session ID extracted, falling back to default selector | provider=%s model=%s", provider, model) - return s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + return s.pickWithAffinityReservation(ctx, provider, model, opts, fallbackAuths, nil, &reservation) } // A single availability pass serves both lookups: the bound credential is validated against // every priority tier, while the fallback selector keeps seeing only the highest tier. available, err := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, now) - if err != nil { + if err != nil && !s.rebindOnNever { return nil, err } fallbackAuths := highestPriorityAuths(available) @@ -1224,11 +1228,20 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth.ID == cachedAuthID { bind(auth.ID, true) entry.Infof("session-affinity: cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) - return auth, nil + return s.pickWithAffinityReservation(ctx, provider, model, opts, available, auth, &reservation) } } + if s.rebindOnNever { + bind(cachedAuthID, false) + errBound := s.unavailableBoundAuthError(ctx, availabilityCandidates, model, cachedAuthID, now) + entry.Infof("session-affinity: cache hit but bound auth unavailable, refusing failover | session=%s bound_auth=%s code=%s provider=%s model=%s", truncateSessionID(primaryID), cachedAuthID, errBound.Code, provider, model) + return nil, errBound + } + if err != nil { + return nil, err + } // Cached auth not available, reselect via fallback selector for even distribution - auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + auth, err := s.pickWithAffinityReservation(ctx, provider, model, opts, fallbackAuths, nil, &reservation) if err != nil { return nil, err } @@ -1262,7 +1275,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth.ID == parentAuthID { bind(auth.ID, true) entry.Infof("session-affinity: guardian inherited parent binding | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) - return auth, nil + return s.pickWithAffinityReservation(ctx, provider, model, opts, available, auth, &reservation) } } } @@ -1279,12 +1292,18 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } else { entry.Infof("session-affinity: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) } - return auth, nil + return s.pickWithAffinityReservation(ctx, provider, model, opts, available, auth, &reservation) } } } + if !isSubagent && !isFork && s.rebindOnNever { + bind(cachedAuthID, false) + errBound := s.unavailableBoundAuthError(ctx, availabilityCandidates, model, cachedAuthID, now) + entry.Infof("session-affinity: fallback cache hit but bound auth unavailable, refusing failover | session=%s fallback=%s bound_auth=%s code=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), cachedAuthID, errBound.Code, provider, model) + return nil, errBound + } if !isSubagent && !isFork && s.shouldPreserveBindingWhenUnavailable(ctx, availabilityCandidates, model, cachedAuthID, now) { - auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + auth, err := s.pickWithAffinityReservation(ctx, provider, model, opts, fallbackAuths, nil, &reservation) if err != nil { return nil, err } @@ -1298,7 +1317,10 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } } - auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + if err != nil { + return nil, err + } + auth, err := s.pickWithAffinityReservation(ctx, provider, model, opts, fallbackAuths, nil, &reservation) if err != nil { return nil, err } @@ -1314,7 +1336,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri return auth, nil } -func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth, entry *log.Entry) (*Auth, bool, error) { +func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth, entry *log.Entry, reservation **weightedReservation) (*Auth, bool, error) { if s == nil || s.matcher == nil { return nil, false, nil } @@ -1343,7 +1365,7 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s } now := time.Now() available, errAvailable := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, now) - if errAvailable != nil { + if errAvailable != nil && !s.rebindOnNever { return nil, true, errAvailable } @@ -1387,15 +1409,47 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s } entry.Infof("session-affinity: LCP cache hit | session=%s prefix=%d auth=%s provider=%s model=%s", truncateSessionID(match.SessionID), match.PrefixLength, auth.ID, provider, model) } - return auth, true, nil + picked, errPick := s.pickWithAffinityReservation(ctx, provider, model, opts, available, auth, reservation) + return picked, true, errPick + } + if !match.IsFork && s.rebindOnNever { + if bindRes := s.matcher.BindFingerprintsWithContext(namespace, fingerprints, tailFingerprints, envDigest, minPrefixLength, match.AuthID); bindRes.SessionID != "" { + opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] = bindRes.SessionID + opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] = bindRes.SessionID + if bindRes.ParentSessionID != "" { + opts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = bindRes.ParentSessionID + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + } + if bindRes.AccessNumber > 0 && opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.LCPAccessGenerationMetadataKey] = bindRes.AccessNumber + } + if bindRes.IsCompaction { + if opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.IsCompactionMetadataKey] = true + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + opts.Metadata[cliproxyexecutor.NodeKindMetadataKey] = "compaction" + } + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + delete(opts.Metadata, cliproxyexecutor.IsCompactionMetadataKey) + delete(opts.Metadata, cliproxyexecutor.NodeKindMetadataKey) + } + } + errBound := s.unavailableBoundAuthError(ctx, availabilityCandidates, model, match.AuthID, now) + entry.Infof("session-affinity: LCP cache hit but bound auth unavailable, refusing failover | session=%s bound_auth=%s code=%s provider=%s model=%s", truncateSessionID(match.SessionID), match.AuthID, errBound.Code, provider, model) + return nil, true, errBound } if !match.IsFork && s.shouldPreserveBindingWhenUnavailable(ctx, availabilityCandidates, model, match.AuthID, now) { preservedAuthID = match.AuthID } } + if errAvailable != nil { + return nil, true, errAvailable + } fallbackAuths := highestPriorityAuths(available) - auth, errPick := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + auth, errPick := s.pickWithAffinityReservation(ctx, provider, model, opts, fallbackAuths, nil, reservation) if errPick != nil { return nil, true, errPick } @@ -1706,9 +1760,9 @@ func (s *SessionAffinitySelector) OnResult(res Result) { // selected credential is unhealthy, so preserve both explicit and LCP bindings. return } - if !res.Success && s.rebindOnQuotaOnly && !shouldRebindSessionAffinityOnError(res.Error) { - // Under session-affinity-rebind-on: quota, transient failures fail over for - // the current request without unbinding the session from its bound credential. + if !res.Success && (s.rebindOnNever || (s.rebindOnQuotaOnly && !shouldRebindSessionAffinityOnError(res.Error))) { + // Under session-affinity-rebind-on: never (or quota on transient failures), + // keep the session bound to its original credential. return } @@ -1788,6 +1842,75 @@ func (s *SessionAffinitySelector) OnResult(res Result) { } } +func (s *SessionAffinitySelector) boundAuthForRequest(providers []string, model string, opts cliproxyexecutor.Options) (string, bool) { + if s == nil { + return "", false + } + modelKey := canonicalModelKey(model) + candidateProviders := make([]string, 0, len(providers)+1) + candidateProviders = append(candidateProviders, "mixed") + for _, p := range providers { + p = strings.TrimSpace(p) + if p != "" && p != "mixed" { + candidateProviders = append(candidateProviders, p) + } + } + + explicitID, explicitFallbackID := extractExplicitSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) + if explicitID == "" && s.matcher != nil { + turns := cliproxysession.ExtractCanonicalTurns(opts.SourceFormat, opts.OriginalRequest) + if len(turns) > 0 { + fingerprints, minPrefixLength, tailFingerprints, envDigest := s.matcher.PrepareExt(turns) + if len(fingerprints) > 0 && minPrefixLength > 0 && minPrefixLength <= len(fingerprints) { + for _, candProvider := range candidateProviders { + if namespace := lcpAffinityNamespace(candProvider, model, opts.Metadata); namespace != "" { + if match, ok := s.matcher.MatchFingerprintsWithContext(namespace, fingerprints, tailFingerprints, envDigest, minPrefixLength); ok && !match.IsFork && match.AuthID != "" { + return match.AuthID, true + } + } + } + } + } + } + + if s.cache == nil { + return "", false + } + primaryID, fallbackID := explicitID, explicitFallbackID + if primaryID == "" { + primaryID, fallbackID = extractSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) + } + if primaryID != "" { + primaryID = cliproxysession.BoundSessionIdentity(primaryID) + } + if fallbackID != "" { + fallbackID = cliproxysession.BoundSessionIdentity(fallbackID) + } + isFork := false + if opts.Metadata != nil { + if forkFlag, ok := opts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); ok && forkFlag { + isFork = true + } + } + isSubagent := !isFork && isSubagentSession(primaryID, fallbackID) + + if primaryID != "" { + for _, candProvider := range candidateProviders { + if authID, ok := s.cache.Get(candProvider + "::" + primaryID + "::" + modelKey); ok && authID != "" { + return authID, true + } + } + } + if fallbackID != "" && fallbackID != primaryID && !isSubagent && !isFork { + for _, candProvider := range candidateProviders { + if authID, ok := s.cache.Get(candProvider + "::" + fallbackID + "::" + modelKey); ok && authID != "" { + return authID, true + } + } + } + return "", false +} + // isGuardianRequest deliberately matches the header, not the configurable review model. func isGuardianRequest(headers http.Header) bool { for key, values := range headers { diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 18ed81f30f4..37678684b25 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -205,16 +205,14 @@ func TestAuthWeight_InvalidAndOverflowValuesAreExcluded(t *testing.T) { } } -func TestPickSmoothWeightedAuth_SaturatesCorruptState(t *testing.T) { +func TestSmoothWeightedStateCharge_SaturatesAndClampsCorruptState(t *testing.T) { t.Parallel() - current := map[string]int64{"a": math.MaxInt64, "b": math.MinInt64} - picked := pickSmoothWeightedAuth([]*Auth{{ID: "a"}, {ID: "b"}}, current) - if picked == nil { - t.Fatal("pickSmoothWeightedAuth() returned nil") - } - if current["a"] != math.MaxInt64-2 || current["b"] != math.MinInt64+1 { - t.Fatalf("current state = %#v, want saturated arithmetic", current) + state := weightedPoolForTest("route", map[string]int64{"a": 1, "b": 1}, map[string]int64{"a": math.MaxInt64, "b": math.MinInt64}) + state.charge(0, false) + limit := maxSmoothWeightedCreditRounds * 2 + if state.currentForTest()["a"] != limit || state.currentForTest()["b"] != -limit { + t.Fatalf("current state = %#v, want saturated and clamped arithmetic", state.currentForTest()) } } diff --git a/sdk/cliproxy/auth/selector_weighted.go b/sdk/cliproxy/auth/selector_weighted.go new file mode 100644 index 00000000000..a984fa79d84 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted.go @@ -0,0 +1,397 @@ +package auth + +import ( + "context" + "crypto/sha256" + "encoding/binary" + "math" + "sort" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// weightedAccountingPoolKey carries the manager-validated pre-pin pool. It affects +// accounting only; the selector still receives pin-filtered selection candidates. +type weightedAccountingPoolKey struct{} + +// weightedTransportProviderKey preserves single-provider transport preference on +// the mixed entrypoint without changing its weighted accounting namespace. +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 +} + +// pickWithAffinityReservation normalizes weighted candidates and defers accounting +// until the local binding completes. Other fallback selectors keep their behavior. +func (s *SessionAffinitySelector) pickWithAffinityReservation(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth, bound *Auth, reservation **weightedReservation) (*Auth, error) { + weighted, ok := s.fallback.(*WeightedRoundRobinSelector) + if !ok { + if bound != nil { + return bound, nil + } + return s.fallback.Pick(ctx, provider, model, opts, auths) + } + picked, pending, errPick := weighted.reserveForRequest(ctx, provider, model, opts, auths, bound) + *reservation = pending + return picked, errPick +} + +// reserveForRequest shares candidate normalization and pinned accounting between +// bare weighted routing and session affinity, which owns deferred confirmation. +func (s *WeightedRoundRobinSelector) reserveForRequest(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth, bound *Auth) (*Auth, *weightedReservation, error) { + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return nil, nil, errContext + } + if accountingAuths, ok := ctx.Value(weightedAccountingPoolKey{}).([]*Auth); ok { + pinnedID := pinnedAuthIDFromMetadata(opts.Metadata) + if bound == nil { + for _, auth := range auths { + if auth != nil && auth.ID == pinnedID { + bound = auth + break + } + } + } + if bound == nil || bound.ID != pinnedID || authWeight(bound) <= 0 { + return nil, nil, &Error{Code: "auth_unavailable", Message: "pinned auth is not available for weighted accounting"} + } + auths = positiveWeightAuths(accountingAuths) + } + transportProvider := provider + if singleProvider, ok := ctx.Value(weightedTransportProviderKey{}).(string); ok { + transportProvider = singleProvider + } + if providerPrefersWebsocketTransport(transportProvider) { + // Reuse the transport filter for every provider preferred by the scheduler. + transportProvider = "codex" + } + var available []*Auth + authID := "" + if bound == nil { + var errAvailable error + available, errAvailable = getSelectorAvailableAuths(ctx, positiveWeightAuths(auths), provider, model, time.Now()) + if errAvailable != nil { + return nil, nil, errAvailable + } + available = preferCodexWebsocketAuths(ctx, transportProvider, available) + } else { + authID = bound.ID + // A binding may bypass priority or transport preference, but its traffic + // must only charge peers that compete with it, never another priority tier. + priority := authPriority(bound) + for _, auth := range auths { + if authPriority(auth) == priority { + available = append(available, auth) + } + } + preferred := preferCodexWebsocketAuths(ctx, transportProvider, available) + if len(preferred) < len(available) && !authWebsocketsEnabled(bound) { + // Keep an HTTP-bound session outside the preferred WebSocket pool. + nonWebsocket := available[:0] + for _, auth := range available { + if !authWebsocketsEnabled(auth) { + nonWebsocket = append(nonWebsocket, auth) + } + } + available = nonWebsocket + } else { + available = preferred + } + } + key := provider + ":" + canonicalModelLookupKey(weightedSelectorStateModel(ctx, model)) + picked, pending := s.reserve(key, available, authID) + if picked == nil { + return nil, nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} + } + return picked, pending, nil +} + +// 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 +} + +// 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. +} + +func sortedWeightedIDs(weights map[string]int64) []string { + ids := make([]string, 0, len(weights)) + for id := range weights { + ids = append(ids, id) + } + sort.Strings(ids) + return ids +} + +// weightedPoolStateKey isolates competing sets without retaining concatenated IDs. +// Weights stay out of the key so configuration changes invalidate reservations. +func weightedPoolStateKey(route string, weights map[string]int64) [32]byte { + return weightedPoolDigest(route, sortedWeightedIDs(weights)) +} + +func weightedPoolDigest(route string, ids []string) [32]byte { + size := 8 + len(route) + for _, id := range ids { + size += 8 + len(id) + } + identity := make([]byte, 0, size) + identity = binary.LittleEndian.AppendUint64(identity, uint64(len(route))) + identity = append(identity, route...) + for _, id := range ids { + identity = binary.LittleEndian.AppendUint64(identity, uint64(len(id))) + identity = append(identity, id...) + } + return sha256.Sum256(identity) +} + +// matches verifies the full identity even on a digest hit. Hash collisions must +// never combine the quota or configured weights of unrelated candidate pools. +func (s *weightedPoolState) matches(route string, ids []string) bool { + if s.route != route || len(s.members) != len(ids) { + return false + } + for i, id := range ids { + if s.members[i].id != id { + return false + } + } + return true +} + +func (s *weightedPoolState) prepare(weights []int64) { + 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] + s.members[i].current = 0 + s.members[i].pending = 0 + } + return + } + } +} + +// 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) { + 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) +} + +// 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) + s.mu.Lock() + defer s.mu.Unlock() + if s.states == nil { + s.states = make(map[[32]byte]*weightedPoolState) + } + state := s.states[key] + for state != nil && !state.matches(route, ids) { + 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 + } + 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]} + } + s.states[key] = state + s.stateCount++ + } else { + state.prepare(weights) + } + var picked *Auth + var pickedCurrent, total int64 + pickedMember := -1 + for _, auth := range available { + if auth == nil { + continue + } + index := sort.SearchStrings(ids, auth.ID) + if index == len(ids) || ids[index] != auth.ID || weights[index] <= 0 { + continue + } + total = saturatingAddInt64(total, weights[index]) + if authID != "" && auth.ID != authID { + continue + } + member := &state.members[index] + current := saturatingAddInt64(member.current, member.pending) + current = saturatingAddInt64(current, member.weight) + if picked == nil || current > pickedCurrent { + picked, pickedCurrent, pickedMember = auth, current, index + } + } + if picked == nil { + return nil, nil + } + // Bound traffic beyond the compensation window still keeps its credential, + // but must not create speculative credit that can starve concurrent picks. + if authID != "" && !state.boundChargeFits(pickedMember, total) { + return picked, nil + } + state.adjustPending(pickedMember, total, false) + state.inFlight++ + return picked, &weightedReservation{ + selector: s, state: state, generation: state.generation, + member: pickedMember, total: total, bound: authID != "", + } +} + +// finish confirms or releases a reservation exactly once under the selector lock. +func (r *weightedReservation) finish(committed bool) { + if r == nil { + return + } + r.selector.mu.Lock() + defer r.selector.mu.Unlock() + if r.finished { + return + } + r.finished = true + if r.state.generation != r.generation { + return + } + r.state.adjustPending(r.member, r.total, true) + r.state.inFlight-- + if committed { + r.state.charge(r.member, r.bound) + } else if r.state.inFlight == 0 { + // A final rollback can remove the offset to already committed credit. + r.state.boundCredits(r.total) + } +} + +// adjustPending must be called with the selector mutex held. Pending deltas are +// not clamped, so releasing them cannot lose credit at the compensation bound. +func (s *weightedPoolState) adjustPending(selected int, total int64, release bool) { + for i := range s.members { + member := &s.members[i] + delta := member.weight + if i == selected { + delta -= total + } + if release { + delta = -delta + } + member.pending = saturatingAddInt64(member.pending, delta) + } +} + +// maxSmoothWeightedCreditRounds bounds compensation after a long sticky session. +const maxSmoothWeightedCreditRounds int64 = 32 + +// boundChargeFits admits the whole sticky delta or none. Both reservation and +// commit include outstanding deltas: committed credit alone can transiently +// exceed the cap while successful reservations settle in a different order. +func (s *weightedPoolState) boundChargeFits(selected int, totalWeight int64) bool { + limit := int64(math.MaxInt64) + if totalWeight <= math.MaxInt64/maxSmoothWeightedCreditRounds { + limit = totalWeight * maxSmoothWeightedCreditRounds + } + for i, member := range s.members { + delta := member.weight + if i == selected { + delta -= totalWeight + } + current := saturatingAddInt64(member.current, member.pending) + next := saturatingAddInt64(current, delta) + if next < -limit || next > limit { + return false + } + } + return true +} + +func (s *weightedPoolState) charge(selected int, bound bool) { + var totalWeight int64 + for _, member := range s.members { + totalWeight = saturatingAddInt64(totalWeight, member.weight) + } + if !bound || s.boundChargeFits(selected, totalWeight) { + // At the cap, forced traffic cannot erase a small peer's earnings. + // Pending repayments still count until they actually roll back. + for i := range s.members { + member := &s.members[i] + member.current = saturatingAddInt64(member.current, member.weight) + } + member := &s.members[selected] + member.current = saturatingAddInt64(member.current, -totalWeight) + } + if s.inFlight == 0 { + s.boundCredits(totalWeight) + } +} + +// boundCredits runs only once the pool settles. Net pending deltas can be zero +// even with live reservations, so the in-flight count, not len(pending), decides +// when clipping cannot irreversibly discard another reservation's offset. +func (s *weightedPoolState) boundCredits(totalWeight int64) { + limit := int64(math.MaxInt64) + if totalWeight <= math.MaxInt64/maxSmoothWeightedCreditRounds { + limit = totalWeight * maxSmoothWeightedCreditRounds + } + var minCredit, maxCredit int64 = math.MaxInt64, math.MinInt64 + for _, member := range s.members { + minCredit = min(minCredit, member.current) + maxCredit = max(maxCredit, member.current) + } + // A common offset keeps leading credits distinct at the compensation bound. + // Independent clamps can repeatedly tie the leaders and starve the later + // candidate in deterministic selection order. + minOffset := saturatingAddInt64(maxCredit, -limit) + maxOffset := saturatingAddInt64(minCredit, limit) + var offset int64 + if minOffset > 0 || minOffset > maxOffset { + offset = minOffset + } else if maxOffset < 0 { + offset = maxOffset + } + // If the full spread cannot fit, cap only the trailing debt after rebasing + // the leader to the upper bound. Do not flatten leading candidates' credits. + for i := range s.members { + member := &s.members[i] + member.current = max(-limit, saturatingAddInt64(member.current, -offset)) + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_affinity_test.go b/sdk/cliproxy/auth/selector_weighted_affinity_test.go new file mode 100644 index 00000000000..5736e0c8210 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_affinity_test.go @@ -0,0 +1,436 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "net/http" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + cliproxysession "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/session" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" +) + +func TestWeightedSessionAffinityCompensatesStickyRequests(t *testing.T) { + for _, source := range []string{"sticky", "parent", "guardian", "fork"} { + t.Run(source, func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{ + {ID: "a", Provider: "codex", Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: "b", Provider: "codex"}, + } + parent := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"busy"}}} + parentID := CanonicalSessionID(parent.Headers, nil, nil) + parentModel := "model" + if source == "guardian" { + parentModel = "parent-model" + } + key := "codex::" + parentID + "::" + parentModel + selector.cache.setSessionAliases("b", "codex", parentID, key) + for i := range 4 { + opts := parent + if source != "sticky" { + opts = cliproxyexecutor.Options{Headers: http.Header{ + "Session_id": {fmt.Sprintf("child-%d", i)}, "X-Codex-Parent-Thread-Id": {"busy"}, + }} + if source == "guardian" { + opts.Headers.Set("X-Openai-Subagent", "guardian") + } + if source == "fork" { + opts.OriginalRequest = []byte(`{"forked_from_thread_id":"busy"}`) + } + } + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, auths) + if errPick != nil || picked == nil || picked.ID != "b" { + t.Fatalf("sticky request %d: pick=%v err=%v", i, picked, errPick) + } + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + if state.currentForTest()["a"] != 12 || state.currentForTest()["b"] != -12 || len(state.pendingForTest()) != 0 { + t.Fatalf("sticky credits: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + for i := range 12 { + // Requests without an identity and cold sessions share the same credits. + opts := cliproxyexecutor.Options{} + if i%2 == 0 { + opts.Headers = http.Header{"Session_id": {fmt.Sprintf("new-%d", i)}} + } + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, auths) + if errPick != nil || picked == nil || picked.ID != "a" { + t.Fatalf("compensation %d: pick=%v err=%v", i, picked, errPick) + } + } + if state.currentForTest()["a"] != 0 || state.currentForTest()["b"] != 0 || len(state.pendingForTest()) != 0 { + t.Fatalf("compensated credits: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + if binding, ok := selector.cache.Get(key); !ok || binding != "b" { + t.Fatalf("balancing moved the parent binding: %q, found=%t", binding, ok) + } + }) + } +} + +func TestWeightedSessionAffinityChangingPoolsDoNotStarveEqualWeights(t *testing.T) { + for _, newSession := range []bool{false, true} { + t.Run(fmt.Sprintf("new-session=%t", newSession), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{ + {ID: "a", Attributes: map[string]string{"priority": "0"}}, + {ID: "b", Attributes: map[string]string{"priority": "1"}}, + {ID: "c", Attributes: map[string]string{"priority": "1"}}, + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"lower-priority"}}} + id := CanonicalSessionID(sticky.Headers, nil, nil) + key := "codex::" + id + "::model" + selector.cache.Set(key, "a") + counts := map[string]int{} + const warmup, requests = 100, 1000 + for i := range requests { + picked, errPick := selector.Pick(t.Context(), "codex", "model", sticky, auths) + if errPick != nil || picked == nil || picked.ID != "a" { + t.Fatalf("sticky request %d: pick=%v err=%v", i, picked, errPick) + } + opts := cliproxyexecutor.Options{} + if newSession { + opts.Headers = http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}} + } + picked, errPick = selector.Pick(t.Context(), "codex", "model", opts, auths) + if errPick != nil || picked == nil || (picked.ID != "b" && picked.ID != "c") { + t.Fatalf("fresh request %d: pick=%v err=%v", i, picked, errPick) + } + if i >= warmup { + counts[picked.ID]++ + } + } + if counts["b"] != (requests-warmup)/2 || counts["c"] != (requests-warmup)/2 { + t.Fatalf("changing pools starved equal-weight candidates after warmup: %v, want b=450 c=450", counts) + } + if binding, ok := selector.cache.Get(key); !ok || binding != "a" { + t.Fatalf("balancing moved the sticky binding: %q, found=%t", binding, ok) + } + if state := weightedStateForTest(t, weighted, "codex:model", "b", "c"); len(state.pendingForTest()) != 0 { + t.Fatalf("pending credits leaked: %v", state.pendingForTest()) + } + }) + } +} + +func TestWeightedSessionAffinityColdSelectionsCountOnce(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + counts := map[string]int{} + for i := range 40 { + picked, errPick := selector.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{ + Headers: http.Header{"Session_id": {fmt.Sprintf("session-%d", i)}}, + }, auths) + if errPick != nil || picked == nil { + t.Fatalf("pick=%v err=%v", picked, errPick) + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + if i == 0 && (state.currentForTest()["a"] != -1 || state.currentForTest()["b"] != 1) { + t.Fatalf("first selection charged more than once: %v", state.currentForTest()) + } + counts[picked.ID]++ + } + if counts["a"] != 30 || counts["b"] != 10 { + t.Fatalf("cold distribution = %v, want 30:10", counts) + } +} + +func TestWeightedSessionAffinityBoundsCompensation(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"long"}}} + id := CanonicalSessionID(opts.Headers, nil, nil) + selector.cache.Set("codex::"+id+"::model", "b") + for range 1000 { + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, auths) + if errPick != nil || picked == nil || picked.ID != "b" { + t.Fatalf("long session: pick=%v err=%v", picked, errPick) + } + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + limit := maxSmoothWeightedCreditRounds * 4 + // A bound charge is admitted as a whole, stopping within one delta of the cap. + wantCredit := limit / 3 * 3 + if state.currentForTest()["a"] != wantCredit || state.currentForTest()["b"] != -wantCredit { + t.Fatalf("bounded credits = %v, want +/- %d", state.currentForTest(), wantCredit) + } + for range wantCredit { + if _, errPick := weighted.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths); errPick != nil { + t.Fatal(errPick) + } + } + if state.currentForTest()["a"] != 0 || state.currentForTest()["b"] != 0 { + t.Fatalf("compensation did not finish: %v", state.currentForTest()) + } +} + +func TestWeightedSessionAffinityLCPCredits(t *testing.T) { + for _, source := range []string{"hit", "fork", "compaction"} { + t.Run(source, func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + root := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"step 1"},` + + `{"role":"assistant","content":"ack 1"},{"role":"user","content":"step 2"},` + + `{"role":"assistant","content":"ack 2"},{"role":"user","content":"step 3"}]}`), + Metadata: map[string]any{cliproxyexecutor.CallerScopeMetadataKey: "caller"}, + } + // Seed the matcher without consuming weighted credits. + turns := cliproxysession.ExtractCanonicalTurns(root.SourceFormat, root.OriginalRequest) + selector.matcher.Bind(lcpAffinityNamespace("openai", "model", root.Metadata), turns, "b") + opts := root + if source == "fork" { + opts.OriginalRequest = []byte(`{"messages":[{"role":"user","content":"step 1"},` + + `{"role":"assistant","content":"ack 1"},{"role":"user","content":"branch"}]}`) + } else if source == "compaction" { + opts.OriginalRequest = []byte(`{"messages":[{"role":"user","content":"Steps 1 and 2 completed"},` + + `{"role":"assistant","content":"ack 2"},{"role":"user","content":"step 3"},` + + `{"role":"user","content":"step 4"}]}`) + } + picked, errPick := selector.Pick(t.Context(), "openai", "model", opts, auths) + if errPick != nil || picked == nil || picked.ID != "b" { + t.Fatalf("LCP %s: pick=%v err=%v", source, picked, errPick) + } + if source == "fork" && opts.Metadata[cliproxyexecutor.IsForkMetadataKey] != true { + t.Fatal("expected LCP fork") + } + if source == "compaction" && opts.Metadata[cliproxyexecutor.IsCompactionMetadataKey] != true { + t.Fatal("expected LCP compaction") + } + state := weightedStateForTest(t, weighted, "openai:model", "a", "b") + if state.currentForTest()["a"] != 3 || state.currentForTest()["b"] != -3 || len(state.pendingForTest()) != 0 { + t.Fatalf("LCP credits: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + }) + } +} + +func TestWeightedReservationRollbackPreservesConcurrentCharges(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + key := "codex:model" + for range 1000 { + _, reservation := weighted.reserve(key, auths, "b") + reservation.finish(true) + } + state := weightedStateForTest(t, weighted, key, "a", "b") + _, rejected := weighted.reserve(key, auths, "b") + _, accepted := weighted.reserve(key, auths, "a") + accepted.finish(true) + rejected.finish(false) + rejected.finish(true) + accepted.finish(false) + limit := maxSmoothWeightedCreditRounds * 4 / 3 * 3 + if state.currentForTest()["a"] != limit-1 || state.currentForTest()["b"] != -limit+1 || len(state.pendingForTest()) != 0 { + t.Fatalf("rollback lost another charge: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } +} + +func TestWeightedReservationResetInvalidatesOldGeneration(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + key := "codex:model" + _, old := weighted.reserve(key, auths, "") + auths[0].Attributes[AttributeWeight] = "1" + _, latest := weighted.reserve(key, auths, "") + old.finish(true) + state := weightedStateForTest(t, weighted, key, "a", "b") + if state.currentForTest()["a"] != 0 || state.currentForTest()["b"] != 0 || state.pendingForTest()["a"] != -1 || state.pendingForTest()["b"] != 1 { + t.Fatalf("stale reservation changed new generation: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + latest.finish(true) + if state.currentForTest()["a"] != -1 || state.currentForTest()["b"] != 1 || len(state.pendingForTest()) != 0 { + t.Fatalf("new generation accounting: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } +} + +func TestWeightedReservationConcurrentSelectionsSeePendingCredits(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + const requests = 40 + selected := make(chan *Auth, requests) + reservations := make(chan *weightedReservation, requests) + var ready sync.WaitGroup + ready.Add(requests) + for range requests { + go func() { + defer ready.Done() + picked, reservation := weighted.reserve("codex:model", auths, "") + selected <- picked + reservations <- reservation + }() + } + ready.Wait() + close(selected) + close(reservations) + counts := map[string]int{} + for auth := range selected { + if auth == nil { + t.Fatal("concurrent reservation returned no auth") + } + counts[auth.ID]++ + } + if counts["a"] != 30 || counts["b"] != 10 { + t.Fatalf("pending allocations ignored: %v", counts) + } + for reservation := range reservations { + reservation.finish(true) + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + if state.currentForTest()["a"] != 0 || state.currentForTest()["b"] != 0 || len(state.pendingForTest()) != 0 { + t.Fatalf("concurrent accounting: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } +} + +func TestWeightedSessionAffinityFailureDoesNotCharge(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: weighted, RebindOn: SessionAffinityRebindOnNever, + }) + t.Cleanup(selector.Stop) + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"bound"}}} + id := CanonicalSessionID(opts.Headers, nil, nil) + selector.cache.Set("codex::"+id+"::model", "missing") + if auth, errPick := selector.Pick(t.Context(), "codex", "model", opts, []*Auth{{ID: "a"}, {ID: "b"}}); errPick == nil || auth != nil { + t.Fatalf("expected unavailable binding: auth=%v err=%v", auth, errPick) + } + ctx, cancel := context.WithCancel(t.Context()) + cancel() + if auth, errPick := selector.Pick(ctx, "codex", "model", opts, []*Auth{{ID: "a"}, {ID: "b"}}); !errors.Is(errPick, context.Canceled) || auth != nil { + t.Fatalf("expected cancellation: auth=%v err=%v", auth, errPick) + } + if len(weighted.states) != 0 { + t.Fatalf("failed selection changed weighted state: %v", weighted.states) + } +} + +func TestWeightedSessionAffinityEligiblePoolAndModelIsolation(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{ + {ID: "a", Attributes: map[string]string{AttributeWeight: "3", "priority": "1"}}, + {ID: "b", Attributes: map[string]string{"priority": "1"}}, + {ID: "lower", Attributes: map[string]string{AttributeWeight: "100"}}, + {ID: "zero", Attributes: map[string]string{AttributeWeight: "0", "priority": "2"}}, + {ID: "cooldown", Attributes: map[string]string{"priority": "2"}, ModelStates: map[string]*ModelState{ + "upstream": {Unavailable: true, NextRetryAfter: time.Now().Add(time.Hour)}, + }}, + } + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"alias"}}} + id := CanonicalSessionID(opts.Headers, nil, nil) + selector.cache.Set("codex::"+id+"::upstream", "b") + ctx := withWeightedSelectorStateModel(t.Context(), weighted, "requested(high)") + picked, errPick := selector.Pick(ctx, "codex", "upstream", opts, auths) + if errPick != nil || picked == nil || picked.ID != "b" { + t.Fatalf("sticky pick=%v err=%v", picked, errPick) + } + state := weightedStateForTest(t, weighted, "codex:requested", "a", "b") + if state == nil || len(weighted.states) != 1 || len(state.currentForTest()) != 2 || state.currentForTest()["a"] != 3 || state.currentForTest()["b"] != -3 { + t.Fatalf("wrong model key or eligible pool: %v", weighted.states) + } + picked, errPick = weighted.Pick(ctx, "codex", "", cliproxyexecutor.Options{}, auths[:2]) + if errPick != nil || picked == nil || picked.ID != "a" || state.currentForTest()["a"] != 2 || state.currentForTest()["b"] != -2 { + t.Fatalf("fallback did not share credits: pick=%v err=%v credits=%v", picked, errPick, state.currentForTest()) + } + for _, pair := range [][2]string{{"codex", "other-model"}, {"other-provider", "requested"}} { + if _, errOther := selector.Pick(t.Context(), pair[0], pair[1], cliproxyexecutor.Options{}, auths[:2]); errOther != nil { + t.Fatal(errOther) + } + } + if state.currentForTest()["a"] != 2 || state.currentForTest()["b"] != -2 { + t.Fatal("another provider or model changed credits") + } +} + +func TestWeightedSessionAffinityReselectionCountsOnce(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"reselect"}}} + id := CanonicalSessionID(opts.Headers, nil, nil) + key := "codex::" + id + "::model" + selector.cache.Set(key, "removed") + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, + []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}}) + if errPick != nil || picked == nil || picked.ID != "a" { + t.Fatalf("reselection: pick=%v err=%v", picked, errPick) + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + if state.currentForTest()["a"] != -1 || state.currentForTest()["b"] != 1 || len(state.pendingForTest()) != 0 { + t.Fatalf("reselection counted incorrectly: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + if binding, ok := selector.cache.Get(key); !ok || binding != "a" { + t.Fatalf("reselection did not bind: %q, found=%t", binding, ok) + } +} + +func TestManagerWeightedSessionAffinityCompensation(t *testing.T) { + for _, mixed := range []bool{false, true} { + t.Run(fmt.Sprintf("mixed=%t", mixed), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + manager := NewManager(nil, selector, nil) + provider, model := "weighted-affinity", "weighted-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + aID, bID := t.Name()+"-a", t.Name()+"-b" + for _, auth := range []*Auth{ + {ID: aID, Provider: provider, Status: StatusActive, Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: bID, Provider: provider, Status: StatusActive}, + } { + if _, errRegister := manager.Register(WithSkipPersist(t.Context()), auth); errRegister != nil { + t.Fatal(errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + namespace := provider + if mixed { + namespace = "mixed" + } + parentHeaders := http.Header{"Session_id": {"manager-sticky"}} + id := CanonicalSessionID(parentHeaders, nil, nil) + selector.cache.Set(namespace+"::"+id+"::"+model, bID) + for i := range 16 { + opts := cliproxyexecutor.Options{Headers: parentHeaders} + want := bID + if i >= 4 { + opts.Headers = http.Header{"Session_id": {fmt.Sprintf("new-%d", i)}} + want = aID + } + var auth *Auth + var errPick error + if mixed { + auth, _, _, errPick = manager.pickNextMixed(t.Context(), []string{provider}, model, opts, nil) + } else { + auth, _, errPick = manager.pickNext(t.Context(), provider, model, opts, nil) + } + if errPick != nil || auth == nil || auth.ID != want { + t.Fatalf("request %d: auth=%v err=%v, want %q", i, auth, errPick, want) + } + } + state := weightedStateForTest(t, weighted, namespace+":"+model, aID, bID) + if state == nil || state.currentForTest()[aID] != 0 || state.currentForTest()[bID] != 0 || len(state.pendingForTest()) != 0 { + t.Fatalf("manager credits did not balance: %#v", state) + } + }) + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_commit_order_test.go b/sdk/cliproxy/auth/selector_weighted_commit_order_test.go new file mode 100644 index 00000000000..58d18499a9d --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_commit_order_test.go @@ -0,0 +1,199 @@ +package auth + +import ( + "fmt" + "maps" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestWeightedReservationFinalRollbackBoundsSettledCredits(t *testing.T) { + for _, rollbackFirst := range []bool{false, true} { + t.Run(fmt.Sprintf("rollback-first=%t", rollbackFirst), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a"}, {ID: "b"}} + reservations := map[string][]*weightedReservation{} + t.Cleanup(func() { + for _, group := range reservations { + for _, reservation := range group { + reservation.finish(false) + } + } + }) + for range 256 { + picked, reservation := weighted.reserve("codex:model", auths, "") + if picked == nil || reservation == nil { + t.Fatal("missing cold reservation") + } + reservations[picked.ID] = append(reservations[picked.ID], reservation) + } + ids := []string{"a", "b"} + if rollbackFirst { + ids = []string{"b", "a"} + } + for _, id := range ids { + for _, reservation := range reservations[id] { + reservation.finish(id == "a") + // Repeated finishing must not decrement the counter again. + reservation.finish(false) + } + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + limit := maxSmoothWeightedCreditRounds * 2 + if state.inFlight != 0 || len(state.pendingForTest()) != 0 || !maps.Equal(state.currentForTest(), map[string]int64{"a": -limit, "b": limit}) { + t.Fatalf("settled credit did not reach its bound: current=%v pending=%v in-flight=%d", state.currentForTest(), state.pendingForTest(), state.inFlight) + } + }) + } +} + +func TestWeightedReservationUnequalWeightsOutOfOrderCommits(t *testing.T) { + for _, first := range []string{"a", "b"} { + t.Run(first, func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + reservations := map[string][]*weightedReservation{} + t.Cleanup(func() { + for _, group := range reservations { + for _, reservation := range group { + reservation.finish(false) + } + } + }) + for range 512 { + picked, reservation := weighted.reserve("codex:model", auths, "") + if picked == nil || reservation == nil { + t.Fatal("missing cold reservation") + } + reservations[picked.ID] = append(reservations[picked.ID], reservation) + } + if len(reservations["a"]) != 384 || len(reservations["b"]) != 128 { + t.Fatalf("reservations did not follow 3:1: a=%d b=%d", len(reservations["a"]), len(reservations["b"])) + } + last := "b" + if first == "b" { + last = "a" + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + for _, reservation := range reservations[first] { + reservation.finish(true) + } + // New arrivals must see the remaining offsets before the batch settles. + for _, phase := range []string{"pending", "settled"} { + if phase == "settled" { + for _, reservation := range reservations[last] { + reservation.finish(true) + } + } + counts := map[string]int{} + for range 64 { + picked, errPick := weighted.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil || picked == nil { + t.Fatalf("Pick failed: auth=%v err=%v", picked, errPick) + } + counts[picked.ID]++ + } + if counts["a"] != 48 || counts["b"] != 16 { + t.Fatalf("%s distribution = %v, want 48:16", phase, counts) + } + } + if state.inFlight != 0 || len(state.pendingForTest()) != 0 || !maps.Equal(state.currentForTest(), map[string]int64{"a": 0, "b": 0}) { + t.Fatalf("settled 3:1 batch drifted: %#v", state) + } + }) + } +} + +func TestWeightedReservationResetDoesNotReleaseNewGeneration(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a"}, {ID: "b"}} + var old []*weightedReservation + t.Cleanup(func() { + for _, reservation := range old { + reservation.finish(false) + } + }) + for range 256 { + _, reservation := weighted.reserve("codex:model", auths, "") + old = append(old, reservation) + } + auths[0].Attributes = map[string]string{AttributeWeight: "3"} + _, latest := weighted.reserve("codex:model", auths, "") + if latest == nil { + t.Fatal("missing new-generation reservation") + } + t.Cleanup(func() { latest.finish(false) }) + for _, reservation := range old { + reservation.finish(true) + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + if state.inFlight != 1 || state.hasCurrentForTest() || !maps.Equal(state.pendingForTest(), map[string]int64{"a": -1, "b": 1}) { + t.Fatalf("old generation altered new accounting: %#v", state) + } + latest.finish(true) + if state.inFlight != 0 || len(state.pendingForTest()) != 0 || !maps.Equal(state.currentForTest(), map[string]int64{"a": -1, "b": 1}) { + t.Fatalf("new generation failed to settle: %#v", state) + } +} + +func TestWeightedReservationOutOfOrderBalancedCommits(t *testing.T) { + for _, bound := range []bool{false, true} { + for _, first := range []string{"a", "b"} { + t.Run(fmt.Sprintf("bound=%t/first=%s", bound, first), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a"}, {ID: "b"}} + reservations := map[string][]*weightedReservation{} + t.Cleanup(func() { + for _, group := range reservations { + for _, reservation := range group { + reservation.finish(false) + } + } + }) + for i := range 256 { + boundID := "" + if bound { + boundID = auths[i%2].ID + } + picked, reservation := weighted.reserve("codex:model", auths, boundID) + if picked == nil || reservation == nil { + t.Fatalf("reserve %d returned pick=%v reservation=%v", i, picked, reservation) + } + reservations[picked.ID] = append(reservations[picked.ID], reservation) + } + if len(reservations["a"]) != 128 || len(reservations["b"]) != 128 { + t.Fatalf("unbalanced reservations: a=%d b=%d", len(reservations["a"]), len(reservations["b"])) + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + // A zero net pending vector does not mean all reservations settled. + if state.inFlight != 256 || len(state.pendingForTest()) != 0 { + t.Fatalf("in-flight=%d pending=%v, want 256 live balanced reservations", state.inFlight, state.pendingForTest()) + } + last := "b" + if first == "b" { + last = "a" + } + for _, id := range []string{first, last} { + for _, reservation := range reservations[id] { + reservation.finish(true) + } + } + if !maps.Equal(state.currentForTest(), map[string]int64{"a": 0, "b": 0}) || len(state.pendingForTest()) != 0 || state.inFlight != 0 { + t.Errorf("commit order distorted balanced credit: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } + counts := map[string]int{} + for range 64 { + picked, errPick := weighted.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil || picked == nil { + t.Fatalf("Pick: auth=%v err=%v", picked, errPick) + } + counts[picked.ID]++ + } + if counts["a"] != 32 || counts["b"] != 32 { + t.Fatalf("post-commit distribution = %v, want a=32 b=32", counts) + } + }) + } + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_memory_test.go b/sdk/cliproxy/auth/selector_weighted_memory_test.go new file mode 100644 index 00000000000..eec967e2883 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_memory_test.go @@ -0,0 +1,71 @@ +package auth + +import ( + "fmt" + "runtime" + "testing" +) + +func BenchmarkWeightedPoolReserve1000(b *testing.B) { + for _, bound := range []bool{false, true} { + b.Run(fmt.Sprintf("bound=%t", bound), func(b *testing.B) { + auths := make([]*Auth, 1000) + for i := range auths { + auths[i] = &Auth{ID: fmt.Sprintf("%036d", i)} + } + selector := &WeightedRoundRobinSelector{} + _, warmup := selector.reserve("codex:model", auths, "") + warmup.finish(true) + b.ReportAllocs() + b.ResetTimer() + for i := range b.N { + boundID := "" + if bound { + boundID = auths[i%len(auths)].ID + } + _, reservation := selector.reserve("codex:model", auths, boundID) + reservation.finish(true) + } + }) + } +} + +// BenchmarkWeightedPoolRetainedMemory measures live cache growth separately from +// total heap and allocation churn. Each retry subset omits a different credential. +func BenchmarkWeightedPoolRetainedMemory(b *testing.B) { + const members = 1000 + auths := make([]*Auth, members) + for i := range auths { + auths[i] = &Auth{ID: fmt.Sprintf("%036d", i)} + } + subset := make([]*Auth, 0, members-1) + var retained, heap, allocated int64 + b.ResetTimer() + for range b.N { + selector := &WeightedRoundRobinSelector{} + _, warmup := selector.reserve("codex:model", auths, "") + warmup.finish(true) + runtime.GC() + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + for excluded := range auths { + subset = append(subset[:0], auths[:excluded]...) + subset = append(subset, auths[excluded+1:]...) + _, reservation := selector.reserve("codex:model", subset, "") + reservation.finish(true) + } + runtime.GC() + runtime.ReadMemStats(&after) + if len(selector.states) != members+1 { + b.Fatalf("cached pools = %d, want %d", len(selector.states), members+1) + } + runtime.KeepAlive(selector) + runtime.KeepAlive(auths) + retained += int64(after.HeapAlloc) - int64(before.HeapAlloc) + heap += int64(after.HeapAlloc) + allocated += int64(after.TotalAlloc - before.TotalAlloc) + } + b.ReportMetric(float64(retained)/float64(b.N), "retained-B/op") + b.ReportMetric(float64(heap)/float64(b.N), "heap-total-B/op") + b.ReportMetric(float64(allocated)/float64(b.N), "alloc-total-B/op") +} diff --git a/sdk/cliproxy/auth/selector_weighted_model_key_test.go b/sdk/cliproxy/auth/selector_weighted_model_key_test.go new file mode 100644 index 00000000000..2f988316bad --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_model_key_test.go @@ -0,0 +1,39 @@ +package auth + +import ( + "fmt" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestWeightedSelectorsShareEquivalentModelKeys(t *testing.T) { + for _, firstAffinity := range []bool{false, true} { + for _, secondAffinity := range []bool{false, true} { + for _, routeOverride := range []bool{false, true} { + t.Run(fmt.Sprintf("affinity=%t/%t/route=%t", firstAffinity, secondAffinity, routeOverride), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + affinity := NewSessionAffinitySelector(weighted) + t.Cleanup(affinity.Stop) + auths := []*Auth{{ID: "a"}, {ID: "b"}} + for i, useAffinity := range []bool{firstAffinity, secondAffinity} { + ctx := t.Context() + model := []string{" Model(high) ", "model(low)"}[i] + if routeOverride { + ctx = withWeightedSelectorStateModel(ctx, weighted, model) + model = []string{"upstream-a", "upstream-b"}[i] + } + var selector Selector = weighted + if useAffinity { + selector = affinity + } + picked, errPick := selector.Pick(ctx, "codex", model, cliproxyexecutor.Options{}, auths) + if errPick != nil || picked == nil || picked.ID != auths[i].ID { + t.Fatalf("request %d: pick=%v err=%v, want auth %s", i, picked, errPick, auths[i].ID) + } + } + }) + } + } + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_overload_test.go b/sdk/cliproxy/auth/selector_weighted_overload_test.go new file mode 100644 index 00000000000..f7639a53689 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_overload_test.go @@ -0,0 +1,455 @@ +package auth + +import ( + "fmt" + "maps" + "net/http" + "strconv" + "sync" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestWeightedOverloadExactCompensation(t *testing.T) { + for _, weights := range [][]int{{1, 1}, {3, 1}, {1, 3}, {1, 5, 1}, {2, 3, 5}, {1, 2, 3, 4}} { + for boundIndex := range weights { + t.Run(fmt.Sprintf("weights=%v/bound=%d", weights, boundIndex), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := make([]*Auth, len(weights)) + total := 0 + for i, weight := range weights { + auths[i] = &Auth{ID: fmt.Sprintf("auth-%d", i), Attributes: map[string]string{AttributeWeight: strconv.Itoa(weight)}} + total += weight + } + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"sticky"}}} + key := "codex::" + CanonicalSessionID(opts.Headers, nil, nil) + "::model" + selector.cache.Set(key, auths[boundIndex].ID) + counts := map[string]int{} + for i := range total * 4 { + request := opts + sticky := i < weights[boundIndex]*4 + if !sticky { + request = cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}} + } + picked, errPick := selector.Pick(t.Context(), "codex", "model", request, auths) + if errPick != nil || picked == nil { + t.Fatalf("selection failed: pick=%v err=%v", picked, errPick) + } + if sticky && picked.ID != auths[boundIndex].ID { + t.Fatal("balancing moved an existing binding") + } + counts[picked.ID]++ + } + for i, weight := range weights { + if counts[auths[i].ID] != weight*4 { + t.Errorf("total allocation = %v, want weight vector %v times 4", counts, weights) + } + } + for _, state := range weighted.states { + for id, credit := range state.currentForTest() { + if credit != 0 { + t.Errorf("exact compensation left %q credit=%d", id, credit) + } + } + if len(state.pendingForTest()) != 0 { + t.Errorf("pending credit leaked: %v", state.pendingForTest()) + } + } + }) + } + } +} + +func TestWeightedOverloadSustainedStickyTrafficDoesNotStarveOthers(t *testing.T) { + for _, weights := range [][]int{{1, 1, 1}, {1, 3, 1}, {1, 5, 1}, {1, 1, 5}, {1, 9, 1}, {1, 1, 9}, {1, 5, 2, 1}} { + t.Run(fmt.Sprint(weights), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := make([]*Auth, len(weights)) + for i, weight := range weights { + auths[i] = &Auth{ID: fmt.Sprintf("auth-%d", i), Attributes: map[string]string{AttributeWeight: strconv.Itoa(weight)}} + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"busy"}}} + key := "codex::" + CanonicalSessionID(sticky.Headers, nil, nil) + "::model" + selector.cache.Set(key, auths[0].ID) + counts := map[string]int{} + for i := range 3000 { + picked, errPick := selector.Pick(t.Context(), "codex", "model", sticky, auths) + if errPick != nil || picked == nil || picked.ID != auths[0].ID { + t.Fatalf("binding changed: pick=%v err=%v", picked, errPick) + } + fresh := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}} + picked, errPick = selector.Pick(t.Context(), "codex", "model", fresh, auths) + if errPick != nil || picked == nil { + t.Fatalf("fresh selection failed: pick=%v err=%v", picked, errPick) + } + if i >= 1000 { + counts[picked.ID]++ + } + } + ids := make([]string, len(auths)) + for i, auth := range auths { + ids[i] = auth.ID + } + state := weightedStateForTest(t, weighted, "codex:model", ids...) + t.Logf("fresh distribution after warmup: %v; committed=%v", counts, state.currentForTest()) + // Forced traffic makes global ratios infeasible; the remaining capacity + // should be shared by unbound peers rather than monopolized by one. + peerWeight := 0 + for _, weight := range weights[1:] { + peerWeight += weight + } + for i, auth := range auths[1:] { + want := 2000 * weights[i+1] / peerWeight + if counts[auth.ID] < want-2 || counts[auth.ID] > want+2 { + t.Errorf("peer %q got %d fresh requests, want approximately %d; distribution=%v", auth.ID, counts[auth.ID], want, counts) + } + } + }) + } +} + +func TestWeightedOverloadFeasibleSustainedTraffic(t *testing.T) { + for _, weights := range [][]int{{3, 1}, {1, 5, 1}, {1, 9, 1}, {1, 5, 2, 1}} { + t.Run(fmt.Sprint(weights), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := make([]*Auth, len(weights)) + total := 0 + for i, weight := range weights { + auths[i] = &Auth{ID: fmt.Sprintf("auth-%d", i), Attributes: map[string]string{AttributeWeight: strconv.Itoa(weight)}} + total += weight + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"feasible"}}} + key := "codex::" + CanonicalSessionID(sticky.Headers, nil, nil) + "::model" + selector.cache.Set(key, auths[0].ID) + counts := map[string]int{} + // Forced traffic is below A's target share, so exact aggregate ratios + // remain feasible without moving established bindings. + for i := range total * 2000 { + opts := cliproxyexecutor.Options{} + bound := i%(total*2) == 0 + if bound { + opts = sticky + } + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, auths) + if errPick != nil || picked == nil || (bound && picked.ID != auths[0].ID) { + t.Fatalf("request %d: pick=%v err=%v", i, picked, errPick) + } + if i >= total*1000 { + counts[picked.ID]++ + } + } + for i, weight := range weights { + if counts[auths[i].ID] != weight*1000 { + t.Errorf("feasible distribution drifted: counts=%v weights=%v", counts, weights) + } + } + t.Logf("feasible distribution: %v", counts) + }) + } +} + +func TestWeightedOverloadRecoveryAfterStickyTrafficStops(t *testing.T) { + for _, weights := range [][]int{{1, 5, 1}, {1, 9, 1}, {1, 5, 2, 1}} { + t.Run(fmt.Sprint(weights), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := make([]*Auth, len(weights)) + total := 0 + for i, weight := range weights { + auths[i] = &Auth{ID: fmt.Sprintf("auth-%d", i), Attributes: map[string]string{AttributeWeight: strconv.Itoa(weight)}} + total += weight + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"busy"}}} + key := "codex::" + CanonicalSessionID(sticky.Headers, nil, nil) + "::model" + selector.cache.Set(key, auths[0].ID) + for range 3000 { + if picked, errPick := selector.Pick(t.Context(), "codex", "model", sticky, auths); errPick != nil || picked == nil || picked.ID != auths[0].ID { + t.Fatalf("sticky selection failed: pick=%v err=%v", picked, errPick) + } + if _, errPick := selector.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths); errPick != nil { + t.Fatal(errPick) + } + } + counts := map[string]int{} + for i := range 1000 + total*100 { + picked, errPick := selector.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil || picked == nil { + t.Fatalf("recovery selection failed: pick=%v err=%v", picked, errPick) + } + if i >= 1000 { + counts[picked.ID]++ + } + } + for i, weight := range weights { + if counts[auths[i].ID] != weight*100 { + t.Errorf("allocation did not recover after sticky traffic stopped: counts=%v weights=%v", counts, weights) + } + } + t.Logf("recovered distribution: %v", counts) + }) + } +} + +func TestWeightedOverloadPendingStickyReservationsDoNotStarveOthers(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a"}, {ID: "b", Attributes: map[string]string{AttributeWeight: "5"}}, {ID: "c"}} + const key = "codex:model" + const inFlight = 200 + reservations := make([]*weightedReservation, inFlight) + t.Cleanup(func() { + for _, reservation := range reservations { + reservation.finish(false) + } + }) + for i := range reservations { + picked, reservation := weighted.reserve(key, auths, "a") + if picked == nil || picked.ID != "a" { + t.Fatal("pending sticky request lost its binding") + } + reservations[i] = reservation + } + counts := map[string]int{} + for i := range 3000 { + // Maintain a fixed backlog without sleeps or timing-dependent coordination. + reservations[i%inFlight].finish(true) + _, reservations[i%inFlight] = weighted.reserve(key, auths, "a") + picked, reservation := weighted.reserve(key, auths, "") + reservation.finish(true) + if i >= 1000 { + counts[picked.ID]++ + } + } + if counts["c"] == 0 { + t.Fatalf("speculative sticky traffic starved c: %v", counts) + } + for _, reservation := range reservations { + reservation.finish(true) + } + state := weightedStateForTest(t, weighted, key, "a", "b", "c") + if len(state.pendingForTest()) != 0 { + t.Fatalf("settled reservations leaked credit: %v", state.pendingForTest()) + } +} + +func TestWeightedOverloadCommitAfterSpeculativeRepaymentRollsBack(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + const key = "codex:model" + for range 1000 { + _, reservation := weighted.reserve(key, auths, "b") + reservation.finish(true) + } + state := weightedStateForTest(t, weighted, key, "a", "b") + before := maps.Clone(state.currentForTest()) + _, first := weighted.reserve(key, auths, "") + _, second := weighted.reserve(key, auths, "") + t.Cleanup(func() { first.finish(false); second.finish(false) }) + // Two pending cold selections make room for a bound delta, but neither + // repayment is real yet. Its later commit must recheck the committed cap. + picked, bound := weighted.reserve(key, auths, "b") + if picked == nil || picked.ID != "b" || bound == nil { + t.Fatal("expected bound reservation funded by pending repayments") + } + t.Cleanup(func() { bound.finish(false) }) + first.finish(false) + second.finish(false) + bound.finish(true) + bound.finish(true) + bound.finish(false) + if !maps.Equal(state.currentForTest(), before) || len(state.pendingForTest()) != 0 { + t.Fatalf("late commit crossed the cap or leaked pending credit: before=%v current=%v pending=%v", before, state.currentForTest(), state.pendingForTest()) + } + // An over-limit bound selection succeeds without reserving any credit. + picked, skipped := weighted.reserve(key, auths, "b") + if picked == nil || picked.ID != "b" || skipped != nil || !maps.Equal(state.currentForTest(), before) || len(state.pendingForTest()) != 0 { + t.Fatal("saturated bound selection should preserve binding and all credits") + } +} + +func TestWeightedOverloadMultipleBusyBindings(t *testing.T) { + for _, busyPerFresh := range []int{1, 3, 10} { + t.Run(fmt.Sprint(busyPerFresh), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{ + {ID: "a"}, {ID: "b", Attributes: map[string]string{AttributeWeight: "5"}}, {ID: "c"}, {ID: "d"}, + } + boundOpts := make([]cliproxyexecutor.Options, 2) + for i, id := range []string{"a", "d"} { + boundOpts[i].Headers = http.Header{"Session_id": {"busy-" + id}} + sessionID := CanonicalSessionID(boundOpts[i].Headers, nil, nil) + selector.cache.Set("codex::"+sessionID+"::model", id) + } + windowCounts := map[string]int{} + for i := range 1400 { + for j := range busyPerFresh { + index := (i + j) % len(boundOpts) + picked, errPick := selector.Pick(t.Context(), "codex", "model", boundOpts[index], auths) + if errPick != nil || picked == nil || picked.ID != []string{"a", "d"}[index] { + t.Fatalf("busy binding changed: auth=%v err=%v", picked, errPick) + } + } + picked, errPick := selector.Pick(t.Context(), "codex", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil || picked == nil { + t.Fatalf("fresh pick failed: auth=%v err=%v", picked, errPick) + } + if i >= 1000 { + windowCounts[picked.ID]++ + if (i+1)%100 == 0 { + if windowCounts["b"] == 0 || windowCounts["c"] == 0 { + t.Fatalf("busy=%d unbound peer starved in a 100-request window: %v", busyPerFresh, windowCounts) + } + clear(windowCounts) + } + } + } + }) + } +} + +func TestWeightedOverloadSaturatedChargeLeavesPeersUnchanged(t *testing.T) { + for _, credit := range []map[string]int64{ + {"a": -222, "b": 111, "c": 111}, + {"a": -112, "b": 222, "c": -110}, + } { + state := weightedPoolForTest("route", map[string]int64{"a": 1, "b": 5, "c": 1}, credit) + state.charge(0, true) + if !maps.Equal(state.currentForTest(), credit) { + t.Fatalf("saturated sticky charge changed peers: got=%v want=%v", state.currentForTest(), credit) + } + } +} + +func TestManagerWeightedOverloadStickyAndPinnedSessions(t *testing.T) { + for _, mixed := range []bool{false, true} { + for _, pinned := range []bool{false, true} { + t.Run(fmt.Sprintf("mixed=%t/pinned=%t", mixed, pinned), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + manager := NewManager(nil, selector, nil) + const provider, model = "sticky-overload", "sticky-overload-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + aID, bID, cID := t.Name()+"-a", t.Name()+"-b", t.Name()+"-c" + for _, auth := range []*Auth{ + {ID: aID, Provider: provider, Status: StatusActive}, + {ID: bID, Provider: provider, Status: StatusActive, Attributes: map[string]string{AttributeWeight: "5"}}, + {ID: cID, Provider: provider, Status: StatusActive}, + } { + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(t.Context()), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + namespace := provider + if mixed { + namespace = "mixed" + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"busy"}}, Metadata: map[string]any{}} + key := namespace + "::" + CanonicalSessionID(sticky.Headers, nil, nil) + "::" + model + selector.cache.Set(key, aID) + if pinned { + sticky.Metadata[cliproxyexecutor.PinnedAuthMetadataKey] = aID + } + counts := map[string]int{} + for i := range 1600 { + for _, bound := range []bool{true, false} { + opts := sticky + if !bound { + opts = cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("new-%d", i)}}} + } + var auth *Auth + var errPick error + if mixed { + auth, _, _, errPick = manager.pickNextMixedLegacy(t.Context(), []string{provider}, model, opts, nil) + } else { + auth, _, errPick = manager.pickNextLegacy(t.Context(), provider, model, opts, nil) + } + if errPick != nil || auth == nil || (bound && auth.ID != aID) { + t.Fatalf("request %d bound=%t: auth=%v err=%v", i, bound, auth, errPick) + } + if !bound && i >= 1000 { + counts[auth.ID]++ + } + } + } + if counts[bID] != 500 || counts[cID] != 100 { + t.Fatalf("manager overload allocation = %v, want b=500 c=100", counts) + } + if binding, ok := selector.cache.Get(key); !ok || binding != aID { + t.Fatalf("manager moved busy binding: %q, found=%t", binding, ok) + } + state := weightedStateForTest(t, weighted, namespace+":"+model, aID, bID, cID) + if len(state.pendingForTest()) != 0 { + t.Fatalf("manager left speculative credit: %v", state.pendingForTest()) + } + }) + } + } +} + +func TestWeightedOverloadConcurrentStickyAccounting(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}} + headers := http.Header{"Session_id": {"sticky-concurrent"}} + key := "codex::" + CanonicalSessionID(headers, nil, nil) + "::model" + selector.cache.Set(key, "b") + type selection struct { + id string + err error + sticky bool + } + results := make(chan selection, 64) + start := make(chan struct{}) + var done sync.WaitGroup + for i := range 64 { + done.Add(1) + go func() { + defer done.Done() + isSticky := i < 8 + opts := cliproxyexecutor.Options{Headers: headers} + if !isSticky { + opts.Headers = http.Header{"Session_id": {fmt.Sprintf("new-%d", i)}} + } + <-start + picked, errPick := selector.Pick(t.Context(), "codex", "model", opts, auths) + id := "" + if picked != nil { + id = picked.ID + } + results <- selection{id: id, err: errPick, sticky: isSticky} + }() + } + close(start) + done.Wait() + close(results) + counts := map[string]int{} + for result := range results { + if result.err != nil || result.id == "" || (result.sticky && result.id != "b") { + t.Fatalf("invalid concurrent result: %+v", result) + } + counts[result.id]++ + } + state := weightedStateForTest(t, weighted, "codex:model", "a", "b") + // All 64 calls share one pool. Accounting must match actual accepted picks, + // regardless of whether the final arrivals happened to be sticky requests. + if state.currentForTest()["a"] != int64(64*3-counts["a"]*4) || state.currentForTest()["b"] != int64(64-counts["b"]*4) || len(state.pendingForTest()) != 0 { + t.Fatalf("concurrent accounting drifted: counts=%v committed=%v pending=%v", counts, state.currentForTest(), state.pendingForTest()) + } + if bound, ok := selector.cache.Get(key); !ok || bound != "b" { + t.Fatalf("concurrency moved sticky binding: %q, found=%t", bound, ok) + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_pool_test.go b/sdk/cliproxy/auth/selector_weighted_pool_test.go new file mode 100644 index 00000000000..5fb73ffa9ca --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_pool_test.go @@ -0,0 +1,287 @@ +package auth + +import ( + "fmt" + "maps" + "math/rand/v2" + "net/http" + "strconv" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// weightedStateForTest checks that the requested accounting pool was created. +func weightedStateForTest(t *testing.T, selector *WeightedRoundRobinSelector, route string, ids ...string) *weightedPoolState { + t.Helper() + weights := make(map[string]int64, len(ids)) + for _, id := range ids { + weights[id] = 1 + } + state := selector.states[weightedPoolStateKey(route, weights)] + ids = sortedWeightedIDs(weights) + for state != nil && !state.matches(route, ids) { + state = state.next + } + if state == nil { + t.Fatalf("missing weighted state for %q and %v", route, ids) + } + return state +} + +func TestWeightedPoolsResumeCreditsAndIsolatePendingReservations(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "b"}, {ID: "c"}} + key := "codex:model" + _, pendingFull := weighted.reserve(key, auths, "b") + t.Cleanup(func() { pendingFull.finish(false) }) + _, pendingSubset := weighted.reserve(key, auths[:2], "b") + t.Cleanup(func() { pendingSubset.finish(false) }) + pendingSubset.finish(true) + subset := weightedStateForTest(t, weighted, key, "a", "b") + full := weightedStateForTest(t, weighted, key, "a", "b", "c") + if !maps.Equal(subset.currentForTest(), map[string]int64{"a": 3, "b": -3}) || full.hasCurrentForTest() { + t.Fatalf("commit crossed pool boundary: subset=%v full=%v", subset.currentForTest(), full.currentForTest()) + } + if !maps.Equal(full.pendingForTest(), map[string]int64{"a": 3, "b": -4, "c": 1}) { + t.Fatalf("another pool consumed pending credit: %v", full.pendingForTest()) + } + pendingFull.finish(false) + for range 100 { + _, reservation := weighted.reserve(key, auths, "b") + reservation.finish(true) + } + // A returning pool resumes its debt, independent of full-pool traffic. + for range 3 { + picked, reservation := weighted.reserve(key, []*Auth{auths[1], auths[0]}, "") + if picked == nil || picked.ID != "a" { + t.Fatalf("returning subset lost credit: pick=%v credits=%v", picked, subset.currentForTest()) + } + reservation.finish(true) + } + if !maps.Equal(subset.currentForTest(), map[string]int64{"a": 0, "b": 0}) || len(subset.pendingForTest()) != 0 || len(full.pendingForTest()) != 0 { + t.Fatalf("returning subset failed to balance: %#v", subset) + } +} + +func TestWeightedPoolInterleavingMatchesIndependentSelectors(t *testing.T) { + shared := &WeightedRoundRobinSelector{} + auths := []*Auth{ + {ID: "a", Attributes: map[string]string{AttributeWeight: "1"}}, + {ID: "b", Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: "c", Attributes: map[string]string{AttributeWeight: "2"}}, + {ID: "d", Attributes: map[string]string{AttributeWeight: "5"}}, + } + var independent [16]WeightedRoundRobinSelector + rng := rand.New(rand.NewPCG(17, 29)) + for step := range 5000 { + mask := 1 + rng.IntN(15) + var pool []*Auth + for i, auth := range auths { + if mask&(1< weighted.maxKeys || old.state == current.state { + t.Fatal("pool eviction failed to bound or replace state") + } + old.finish(true) + if current.state.hasCurrentForTest() || !maps.Equal(current.state.pendingForTest(), map[string]int64{"a": -1, "b": 1}) { + t.Fatal("finishing an evicted reservation affected the new pool") + } + current.finish(true) +} + +func TestManagerWeightedPriorityPoolIsolation(t *testing.T) { + for _, mixed := range []bool{false, true} { + for _, pinned := range []bool{false, true} { + t.Run(fmt.Sprintf("mixed=%t/pinned=%t", mixed, pinned), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + manager := NewManager(nil, selector, nil) + const provider, model = "pool-isolation", "pool-model" + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + auths := []*Auth{ + {ID: t.Name() + "-a", Attributes: map[string]string{"priority": "0"}}, + {ID: t.Name() + "-b", Attributes: map[string]string{"priority": "1", AttributeWeight: "3"}}, + {ID: t.Name() + "-c", Attributes: map[string]string{"priority": "1", AttributeWeight: "1"}}, + } + for _, auth := range auths { + auth.Provider, auth.Status = provider, StatusActive + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(WithSkipPersist(t.Context()), auth); errRegister != nil { + t.Fatal(errRegister) + } + } + namespace := provider + if mixed { + namespace = "mixed" + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"priority-bound"}}, Metadata: map[string]any{}} + id := CanonicalSessionID(sticky.Headers, nil, nil) + selector.cache.Set(namespace+"::"+id+"::"+model, auths[0].ID) + if pinned { + sticky.Metadata[cliproxyexecutor.PinnedAuthMetadataKey] = auths[0].ID + } + counts := map[string]int{} + for i := range 1200 { + for _, bound := range []bool{true, false} { + opts := sticky + if !bound { + opts = cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}} + } + var auth *Auth + var errPick error + if mixed { + auth, _, _, errPick = manager.pickNextMixedLegacy(t.Context(), []string{provider}, model, opts, nil) + } else { + auth, _, errPick = manager.pickNextLegacy(t.Context(), provider, model, opts, nil) + } + if errPick != nil || auth == nil || (bound && auth.ID != auths[0].ID) || (!bound && auth.ID == auths[0].ID) { + t.Fatalf("request %d bound=%t: auth=%v err=%v", i, bound, auth, errPick) + } + if !bound && i >= 200 { + counts[auth.ID]++ + } + } + } + if counts[auths[1].ID] != 750 || counts[auths[2].ID] != 250 { + t.Fatalf("manager cross-priority distribution = %v, want 750:250", counts) + } + }) + } + } +} + +func TestWeightedSessionAffinityPoolIsolation(t *testing.T) { + for _, reason := range []string{"priority", "websocket", "retry-exclusion", "cooldown"} { + for _, weights := range [][2]int{{1, 1}, {3, 1}, {1, 3}, {5, 2}} { + t.Run(fmt.Sprintf("%s/%d:%d", reason, weights[0], weights[1]), func(t *testing.T) { + weighted := &WeightedRoundRobinSelector{} + selector := NewSessionAffinitySelector(weighted) + t.Cleanup(selector.Stop) + auths := []*Auth{ + {ID: "a", Attributes: map[string]string{"priority": "0"}}, + {ID: "b", Attributes: map[string]string{AttributeWeight: strconv.Itoa(weights[0])}}, + {ID: "c", Attributes: map[string]string{AttributeWeight: strconv.Itoa(weights[1])}}, + } + ctx := t.Context() + switch reason { + case "priority": + auths[1].Attributes["priority"] = "1" + auths[2].Attributes["priority"] = "1" + case "websocket": + ctx = cliproxyexecutor.WithDownstreamWebsocket(ctx) + auths[1].Attributes["websockets"] = "true" + auths[2].Attributes["websockets"] = "true" + } + sticky := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {"bound-a"}}} + id := CanonicalSessionID(sticky.Headers, nil, nil) + key := "codex::" + id + "::model" + selector.cache.Set(key, "a") + const warmup, measured = 1000, 1400 + counts := map[string]int{} + for i := range warmup + measured { + picked, errPick := selector.Pick(ctx, "codex", "model", sticky, auths) + if errPick != nil || picked == nil || picked.ID != "a" { + t.Fatalf("sticky request %d: pick=%v err=%v", i, picked, errPick) + } + freshAuths := auths + if reason == "retry-exclusion" { + freshAuths = auths[1:] + } else if reason == "cooldown" { + auths[0].Unavailable = true + auths[0].NextRetryAfter = time.Now().Add(time.Hour) + } + opts := cliproxyexecutor.Options{Headers: http.Header{"Session_id": {fmt.Sprintf("fresh-%d", i)}}} + picked, errPick = selector.Pick(ctx, "codex", "model", opts, freshAuths) + if errPick != nil || picked == nil || (picked.ID != "b" && picked.ID != "c") { + t.Fatalf("fresh request %d: pick=%v err=%v", i, picked, errPick) + } + if i >= warmup { + counts[picked.ID]++ + } + auths[0].Unavailable = false + auths[0].NextRetryAfter = time.Time{} + } + wantB := measured * weights[0] / (weights[0] + weights[1]) + if counts["b"] != wantB || counts["c"] != measured-wantB { + t.Fatalf("cross-pool traffic distorted cold selections: %v, want b=%d c=%d", counts, wantB, measured-wantB) + } + if bound, ok := selector.cache.Get(key); !ok || bound != "a" { + t.Fatalf("binding moved: %q, found=%t", bound, ok) + } + for _, state := range weighted.states { + if len(state.pendingForTest()) != 0 { + t.Fatalf("pending credits leaked: %v", state.pendingForTest()) + } + } + }) + } + } +} diff --git a/sdk/cliproxy/auth/selector_weighted_storage_test.go b/sdk/cliproxy/auth/selector_weighted_storage_test.go new file mode 100644 index 00000000000..0d937b90381 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_storage_test.go @@ -0,0 +1,88 @@ +package auth + +import ( + "maps" + "testing" +) + +// These snapshots keep assertions independent of the compact production layout. +func (s *weightedPoolState) currentForTest() map[string]int64 { + current := make(map[string]int64, len(s.members)) + for _, member := range s.members { + current[member.id] = member.current + } + return current +} + +func (s *weightedPoolState) pendingForTest() map[string]int64 { + pending := make(map[string]int64) + for _, member := range s.members { + if member.pending != 0 { + pending[member.id] = member.pending + } + } + return pending +} + +func (s *weightedPoolState) hasCurrentForTest() bool { + for _, member := range s.members { + if member.current != 0 { + return true + } + } + return false +} + +// TestWeightedPoolDigestCollisionIsolation forces digest hits for different exact +// identities, including a matching state behind a collision-chain head. +func TestWeightedPoolDigestCollisionIsolation(t *testing.T) { + for _, unrelated := range []*weightedPoolState{ + weightedPoolForTest("other-route", map[string]int64{"a": 1, "b": 1}, map[string]int64{"a": 7, "b": -7}), + weightedPoolForTest("route", map[string]int64{"a": 1, "c": 1}, map[string]int64{"a": 7, "c": -7}), + weightedPoolForTest("route", map[string]int64{"a": 1}, map[string]int64{"a": 7}), + } { + t.Run(unrelated.route+"/"+unrelated.members[len(unrelated.members)-1].id, func(t *testing.T) { + weights := map[string]int64{"a": 1, "b": 1} + key := weightedPoolStateKey("route", weights) + selector := &WeightedRoundRobinSelector{ + states: map[[32]byte]*weightedPoolState{key: unrelated}, stateCount: 1, maxKeys: 2, + } + before := unrelated.currentForTest() + picked, first := selector.reserve("route", []*Auth{{ID: "b"}, {ID: "a"}}, "b") + if picked == nil || picked.ID != "b" || first == nil || first.state == unrelated { + t.Fatal("digest collision reused an unrelated pool") + } + t.Cleanup(func() { first.finish(false) }) + matching := first.state + // Put the nonmatching state first to exercise a full chain lookup. + matching.next, unrelated.next, selector.states[key] = nil, matching, unrelated + _, second := selector.reserve("route", []*Auth{{ID: "a"}, {ID: "b"}}, "a") + if second == nil || second.state != matching || selector.stateCount != 2 { + t.Fatal("collision lookup failed to reuse the matching pool") + } + t.Cleanup(func() { second.finish(false) }) + first.finish(true) + second.finish(true) + if !maps.Equal(before, unrelated.currentForTest()) || unrelated.inFlight != 0 || len(unrelated.pendingForTest()) != 0 { + t.Fatal("collision accounting changed the unrelated pool") + } + if matching.hasCurrentForTest() || matching.inFlight != 0 || len(matching.pendingForTest()) != 0 { + t.Fatal("matching pool did not balance") + } + // Capacity still counts pools, not just distinct digest buckets. + _, third := selector.reserve("third-route", []*Auth{{ID: "a"}}, "") + third.finish(true) + if selector.stateCount != 1 || len(selector.states) != 1 { + t.Fatal("colliding pools bypassed the existing capacity limit") + } + }) + } +} + +func weightedPoolForTest(route string, weights, current map[string]int64) *weightedPoolState { + state := &weightedPoolState{route: route} + for _, id := range sortedWeightedIDs(weights) { + state.members = append(state.members, weightedPoolMember{id: id, weight: weights[id], current: current[id]}) + } + return state +} diff --git a/sdk/cliproxy/auth/selector_weighted_test.go b/sdk/cliproxy/auth/selector_weighted_test.go new file mode 100644 index 00000000000..31352e30170 --- /dev/null +++ b/sdk/cliproxy/auth/selector_weighted_test.go @@ -0,0 +1,79 @@ +package auth + +import ( + "maps" + "sort" + "testing" +) + +func TestSmoothWeightedStateChargeRebasesBeforeBounding(t *testing.T) { + for _, tc := range []struct { + name string + current map[string]int64 + weights map[string]int64 + authID string + want map[string]int64 + }{ + { + name: "smaller pool preserves leading difference", + current: map[string]int64{"a": -96, "b": 96, "c": 95}, + weights: map[string]int64{"b": 1, "c": 1}, + authID: "b", + want: map[string]int64{"b": 63, "c": 64}, + }, + { + name: "negative credits shift together", + current: map[string]int64{"a": 96, "b": -64, "c": -63}, + weights: map[string]int64{"b": 1, "c": 1}, + authID: "b", + want: map[string]int64{"b": -64, "c": -61}, + }, + { + name: "excessive spread clips debt without tying leaders", + current: map[string]int64{"a": -96, "b": 96, "c": 95}, + weights: map[string]int64{"a": 1, "b": 1, "c": 1}, + authID: "a", + want: map[string]int64{"a": -96, "b": 96, "c": 95}, + }, + { + name: "bounded credits do not shift", + current: map[string]int64{"b": 0, "c": 0}, + weights: map[string]int64{"b": 1, "c": 1}, + authID: "b", + want: map[string]int64{"b": -1, "c": 1}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + state := weightedPoolForTest("route", tc.weights, tc.current) + state.charge(sort.Search(len(state.members), func(i int) bool { return state.members[i].id >= tc.authID }), false) + if !maps.Equal(state.currentForTest(), tc.want) { + t.Fatalf("rebased credits = %v, want %v", state.currentForTest(), tc.want) + } + }) + } +} + +func TestWeightedReservationRollbackPreservesConcurrentCredits(t *testing.T) { + key := "codex:model" + weights := map[string]int64{"a": 1, "b": 1, "c": 1} + state := weightedPoolForTest(key, weights, map[string]int64{"a": -94, "b": 94, "c": 93}) + weighted := &WeightedRoundRobinSelector{states: map[[32]byte]*weightedPoolState{weightedPoolStateKey(key, weights): state}, stateCount: 1} + auths := []*Auth{{ID: "a"}, {ID: "b"}, {ID: "c"}} + _, rejected := weighted.reserve(key, auths, "a") + t.Cleanup(func() { rejected.finish(false) }) + _, accepted := weighted.reserve(key, auths, "b") + t.Cleanup(func() { accepted.finish(false) }) + accepted.finish(true) + want := map[string]int64{"a": -93, "b": 92, "c": 94} + if !maps.Equal(state.currentForTest(), want) { + t.Fatalf("committed credits = %v, want %v", state.currentForTest(), want) + } + if !maps.Equal(state.pendingForTest(), map[string]int64{"a": -2, "b": 1, "c": 1}) { + t.Fatalf("commit changed another request's reservation: %v", state.pendingForTest()) + } + rejected.finish(false) + rejected.finish(true) + if !maps.Equal(state.currentForTest(), want) || len(state.pendingForTest()) != 0 { + t.Fatalf("rollback changed committed credits: current=%v pending=%v", state.currentForTest(), state.pendingForTest()) + } +} diff --git a/sdk/cliproxy/auth/session_affinity_priority_test.go b/sdk/cliproxy/auth/session_affinity_priority_test.go index c6cd1fb3358..649bd329562 100644 --- a/sdk/cliproxy/auth/session_affinity_priority_test.go +++ b/sdk/cliproxy/auth/session_affinity_priority_test.go @@ -2,10 +2,13 @@ package auth import ( "context" + "errors" "net/http" + "strings" "testing" "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" @@ -768,3 +771,643 @@ func TestManagerExecuteAndStreamRebindOnQuotaEndToEnd(t *testing.T) { }) } } + +func TestManagerSessionAffinityRebindOnNeverRefusesFailoverAndPreservesBinding(t *testing.T) { + for _, testCase := range []struct { + name string + providerSuffix string + pick func(*Manager, context.Context, string, string, cliproxyexecutor.Options, map[string]struct{}) (*Auth, error) + }{ + { + name: "single provider", + providerSuffix: "single", + pick: func(manager *Manager, ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, error) { + auth, _, errPick := manager.pickNext(ctx, provider, model, opts, tried) + return auth, errPick + }, + }, + { + name: "mixed provider", + providerSuffix: "mixed", + pick: func(manager *Manager, ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, error) { + auth, _, _, errPick := manager.pickNextMixed(ctx, []string{provider}, model, opts, tried) + return auth, errPick + }, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + ctx := context.Background() + provider := "affinity-rebind-never-" + testCase.providerSuffix + model := "affinity-rebind-never-model" + highID := provider + "-high" + lowID := provider + "-low" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + RebindOn: SessionAffinityRebindOnNever, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + + for _, auth := range []*Auth{ + {ID: highID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "1"}}, + {ID: lowID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "0"}}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.DerivedSessionIDMetadataKey: "never-session", + }} + otherOpts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.DerivedSessionIDMetadataKey: "never-other-session", + }} + + got, errPick := testCase.pick(manager, ctx, provider, model, opts, nil) + if errPick != nil || got == nil || got.ID != highID { + t.Fatalf("cold binding = %v (err=%v), want %q", got, errPick, highID) + } + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: true, + Options: opts, + }) + + // Transient failure on highID must refuse failover with 503 session_auth_retry and keep binding. + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadGateway, Message: "502 overloaded"}, + Options: opts, + }) + if got, errPick = testCase.pick(manager, ctx, provider, model, opts, nil); errPick == nil { + t.Fatalf("transient unavailable pick succeeded with %v, want 503 session_auth_retry error", got) + } else { + var authErr *Error + if !errors.As(errPick, &authErr) || authErr == nil || authErr.Code != "session_auth_retry" || authErr.HTTPStatus != http.StatusServiceUnavailable { + t.Fatalf("transient unavailable pick err = %v, want 503 session_auth_retry", errPick) + } + } + + // A new unbound session while highID is cooling down must still pick lowID. + gotOther, errOther := testCase.pick(manager, ctx, provider, model, otherOpts, nil) + if errOther != nil || gotOther == nil || gotOther.ID != lowID { + t.Fatalf("new session pick while highID unavailable = %v (err=%v), want %q", gotOther, errOther, lowID) + } + + expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + got, errPick = testCase.pick(manager, ctx, provider, model, opts, nil) + if errPick != nil || got == nil || got.ID != highID { + t.Fatalf("pick after transient recovery = %v (err=%v), want %q", got, errPick, highID) + } + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: true, + Options: opts, + }) + + // Quota/auth failures (429, 403) on highID must refuse failover with 429 session_auth_unavailable and keep binding. + for _, blockErr := range []*Error{ + {HTTPStatus: http.StatusTooManyRequests, Message: "usage limit reached"}, + {HTTPStatus: http.StatusForbidden, Message: "forbidden"}, + } { + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: false, + Error: blockErr, + Options: opts, + }) + if got, errPick = testCase.pick(manager, ctx, provider, model, opts, nil); errPick == nil { + t.Fatalf("blocked (%v) pick succeeded with %v, want 429 session_auth_unavailable error", blockErr, got) + } else { + var authErr *Error + if !errors.As(errPick, &authErr) || authErr == nil || authErr.Code != "session_auth_unavailable" || authErr.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("blocked (%v) pick err = %v, want 429 session_auth_unavailable", blockErr, errPick) + } + } + expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + got, errPick = testCase.pick(manager, ctx, provider, model, opts, nil) + if errPick != nil || got == nil || got.ID != highID { + t.Fatalf("pick after blocked (%v) recovery = %v (err=%v), want %q", blockErr, got, errPick, highID) + } + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: true, + Options: opts, + }) + } + }) + } +} + +func TestManagerSessionAffinityRebindOnNeverLCPAndFallbackKeyPreserveBinding(t *testing.T) { + ctx := context.Background() + provider := "openai" + model := "gpt-5-never" + highID := "openai-never-high" + lowID := "openai-never-low" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + RebindOn: SessionAffinityRebindOnNever, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + + for _, auth := range []*Auth{ + {ID: highID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "1"}}, + {ID: lowID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "0"}}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + makeTurnOpts := func(payload string) cliproxyexecutor.Options { + return cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(payload), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "caller-lcp-never", + }, + } + } + + // 1. LCP cold binding on highID. + turn1 := makeTurnOpts(`{"messages":[{"role":"user","content":"turn 1 hello"}]}`) + got1, _, errPick := manager.pickNext(ctx, provider, model, turn1, nil) + if errPick != nil || got1 == nil || got1.ID != highID { + t.Fatalf("turn 1 pick = %v (err=%v), want %q", got1, errPick, highID) + } + manager.MarkResult(ctx, Result{AuthID: highID, Provider: provider, Model: model, Success: true, Options: turn1}) + + // 2. LCP turn 2 fails on highID with 429 quota -> failover attempt must refuse with 429 session_auth_unavailable and preserve binding. + turn2 := makeTurnOpts(`{"messages":[{"role":"user","content":"turn 1 hello"},{"role":"assistant","content":"hi"},{"role":"user","content":"turn 2 question"}]}`) + got2, _, errPick := manager.pickNext(ctx, provider, model, turn2, nil) + if errPick != nil || got2 == nil || got2.ID != highID { + t.Fatalf("turn 2 pick = %v (err=%v), want %q", got2, errPick, highID) + } + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, + Options: turn2, + }) + + turn2Retry := makeTurnOpts(`{"messages":[{"role":"user","content":"turn 1 hello"},{"role":"assistant","content":"hi"},{"role":"user","content":"turn 2 question"}]}`) + if gotFailover, _, errFailover := manager.pickNext(ctx, provider, model, turn2Retry, nil); errFailover == nil { + t.Fatalf("LCP quota unavailable pick succeeded with %v, want 429 session_auth_unavailable", gotFailover) + } else { + var authErr *Error + if !errors.As(errFailover, &authErr) || authErr == nil || authErr.Code != "session_auth_unavailable" || authErr.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("LCP quota unavailable pick err = %v, want 429 session_auth_unavailable", errFailover) + } + } + + expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + turn2Recovered := makeTurnOpts(`{"messages":[{"role":"user","content":"turn 1 hello"},{"role":"assistant","content":"hi"},{"role":"user","content":"turn 2 question"}]}`) + got2Rec, _, errPick := manager.pickNext(ctx, provider, model, turn2Recovered, nil) + if errPick != nil || got2Rec == nil || got2Rec.ID != highID { + t.Fatalf("LCP turn 2 after recovery = %v (err=%v), want %q", got2Rec, errPick, highID) + } + manager.MarkResult(ctx, Result{AuthID: highID, Provider: provider, Model: model, Success: true, Options: turn2Recovered}) + + // 3. Fallback-key path (non-subagent, non-fork: conversation -> prompt_cache_key + conversation) + // refuses failover when bound auth is transiently unavailable. + fallbackInitOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"conversation":"conv-never-1"}`), + } + gotFB1, _, errPick := manager.pickNext(ctx, provider, model, fallbackInitOpts, nil) + if errPick != nil || gotFB1 == nil || gotFB1.ID != highID { + t.Fatalf("fallback init pick = %v (err=%v), want %q", gotFB1, errPick, highID) + } + manager.MarkResult(ctx, Result{AuthID: highID, Provider: provider, Model: model, Success: true, Options: fallbackInitOpts}) + + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadGateway, Message: "502 overloaded"}, + Options: fallbackInitOpts, + }) + fallbackLookupOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"prompt_cache_key":"pck-never-1","conversation":"conv-never-1"}`), + } + if gotFB2, _, errFB2 := manager.pickNext(ctx, provider, model, fallbackLookupOpts, nil); errFB2 == nil { + t.Fatalf("fallback key unavailable pick succeeded with %v, want 503 session_auth_retry", gotFB2) + } else { + var authErr *Error + if !errors.As(errFB2, &authErr) || authErr == nil || authErr.Code != "session_auth_retry" || authErr.HTTPStatus != http.StatusServiceUnavailable { + t.Fatalf("fallback key unavailable pick err = %v, want 503 session_auth_retry", errFB2) + } + } + expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + gotFBRec, _, errPick := manager.pickNext(ctx, provider, model, fallbackLookupOpts, nil) + if errPick != nil || gotFBRec == nil || gotFBRec.ID != highID { + t.Fatalf("fallback key pick after recovery = %v (err=%v), want %q", gotFBRec, errPick, highID) + } +} + +func TestManagerExecuteAndStreamRebindOnNeverEndToEnd(t *testing.T) { + ctx := context.Background() + provider := "exec-rebind-never" + model := "exec-rebind-never-model" + authAID := provider + "-a" + authBID := provider + "-b" + + var failAuthErr map[string]func(call int) error + calls := make(map[string]int) + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(2, 30*time.Second, 0) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + RebindOn: SessionAffinityRebindOnNever, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(scriptedFailoverExecutor{ + provider: provider, + execute: func(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + calls[auth.ID]++ + if fn := failAuthErr[auth.ID]; fn != nil { + if errFail := fn(calls[auth.ID]); errFail != nil { + return cliproxyexecutor.Response{}, errFail + } + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil + }, + stream: func(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + calls[auth.ID]++ + if fn := failAuthErr[auth.ID]; fn != nil { + if errFail := fn(calls[auth.ID]); errFail != nil { + return nil, errFail + } + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + }, + }) + + for _, auth := range []*Auth{ + {ID: authAID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "1"}}, + {ID: authBID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "0"}}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + req := cliproxyexecutor.Request{Model: model} + sessionAOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"never-e2e-session-a"}}, + } + sessionBOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"never-e2e-session-b"}}, + } + + // 1. Bind session A to authAID. + resp, errExec := manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec != nil || string(resp.Payload) != authAID { + t.Fatalf("cold Execute = %q (err=%v), want %q", string(resp.Payload), errExec, authAID) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls = %d, want 0", calls[authBID]) + } + + // 2. Make authA return 429 usage limit reached -> returns 429 to client without calling authB or running extra retry rounds. + calls[authAID] = 0 + failAuthErr = map[string]func(int) error{ + authAID: func(_ int) error { + return &Error{HTTPStatus: http.StatusTooManyRequests, Message: "usage_limit_reached: quota exhausted"} + }, + } + _, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests || !strings.Contains(errExec.Error(), "usage_limit_reached") { + t.Fatalf("Execute on 429 = %v, want 429 usage_limit_reached", errExec) + } + if calls[authAID] != 1 { + t.Fatalf("authA calls on 429 = %d, want 1 (no extra retry rounds on quota error)", calls[authAID]) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls after authA 429 = %d, want 0", calls[authBID]) + } + + // 3. Next request on session A while authA is in 429 cooldown returns 429 session_auth_unavailable without calling authA or authB. + _, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests { + t.Fatalf("subsequent Execute on cooled session A = %v, want 429", errExec) + } + if calls[authAID] != 1 || calls[authBID] != 0 { + t.Fatalf("calls after cooled session A = (A=%d, B=%d), want (1, 0)", calls[authAID], calls[authBID]) + } + + // 4. A request with a new session ID (session B) goes to authB. + respB, errExecB := manager.Execute(ctx, []string{provider}, req, sessionBOpts) + if errExecB != nil || string(respB.Payload) != authBID { + t.Fatalf("new session B Execute = %q (err=%v), want %q", string(respB.Payload), errExecB, authBID) + } + if calls[authBID] != 1 { + t.Fatalf("authB calls = %d, want 1", calls[authBID]) + } + + // 5. Expire authA cooldown -> session A still bound to authA. + failAuthErr = nil + expireSessionAffinityPriorityModelCooldown(t, manager, authAID, model) + resp, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec != nil || string(resp.Payload) != authAID { + t.Fatalf("Execute after authA quota recovery = %q (err=%v), want %q", string(resp.Payload), errExec, authAID) + } + + // 6. Transient 502 error on authA (with transient cooldown disabled) is retried on authA through retry rounds and never sent to authB. + prevTransientSec := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(-1) + defer SetTransientErrorCooldownSeconds(int(prevTransientSec)) + + calls[authAID] = 0 + calls[authBID] = 0 + failAuthErr = map[string]func(int) error{ + authAID: func(call int) error { + if call == 1 { + return &Error{HTTPStatus: http.StatusBadGateway, Message: "502 bad gateway"} + } + return nil + }, + } + resp, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec != nil || string(resp.Payload) != authAID { + t.Fatalf("Execute 502 retry on authA = %q (err=%v), want %q", string(resp.Payload), errExec, authAID) + } + if calls[authAID] != 2 || calls[authBID] != 0 { + t.Fatalf("calls after 502 retry = (A=%d, B=%d), want (2, 0)", calls[authAID], calls[authBID]) + } + + // 7. When authA fails with 502 across all retry rounds (or cooldown exceeds maxWait), Execute returns 503 session_auth_retry and never calls authB. + SetTransientErrorCooldownSeconds(-1) + manager.SetRetryConfig(2, 30*time.Second, 0) + calls[authAID] = 0 + calls[authBID] = 0 + failAuthErr = map[string]func(int) error{ + authAID: func(_ int) error { + return &Error{HTTPStatus: http.StatusBadGateway, Message: "502 overloaded"} + }, + } + _, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusServiceUnavailable { + t.Fatalf("Execute exhausted 502 retry rounds = %v, want 503 session_auth_retry", errExec) + } + var exhaustedRetryErr *Error + if !errors.As(errExec, &exhaustedRetryErr) || exhaustedRetryErr == nil || exhaustedRetryErr.Code != "session_auth_retry" { + t.Fatalf("Execute exhausted 502 error = %v, want session_auth_retry", errExec) + } + if calls[authAID] != 3 || calls[authBID] != 0 { + t.Fatalf("calls after exhausted 502 retry rounds = (A=%d, B=%d), want (3, 0)", calls[authAID], calls[authBID]) + } + + // 8. When authA is in transient 502 cooldown exceeding maxWait, both initial and subsequent requests on session A return 503 session_auth_retry and never call authB. + SetTransientErrorCooldownSeconds(0) + manager.SetRetryConfig(2, 0, 0) + calls[authAID] = 0 + calls[authBID] = 0 + _, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusServiceUnavailable { + t.Fatalf("Execute initial 502 with cooldown = %v, want 503 session_auth_retry", errExec) + } + if calls[authAID] != 1 || calls[authBID] != 0 { + t.Fatalf("calls after initial 502 = (A=%d, B=%d), want (1, 0)", calls[authAID], calls[authBID]) + } + + _, errExec = manager.Execute(ctx, []string{provider}, req, sessionAOpts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusServiceUnavailable { + t.Fatalf("subsequent Execute while authA in 502 cooldown = %v, want 503", errExec) + } + var retryAuthErr *Error + if !errors.As(errExec, &retryAuthErr) || retryAuthErr == nil || retryAuthErr.Code != "session_auth_retry" { + t.Fatalf("subsequent Execute error = %v, want session_auth_retry", errExec) + } + if calls[authAID] != 1 || calls[authBID] != 0 { + t.Fatalf("calls after cooled 502 session A = (A=%d, B=%d), want (1, 0)", calls[authAID], calls[authBID]) + } + + _, errStream := manager.ExecuteStream(ctx, []string{provider}, req, sessionAOpts) + if errStream == nil || statusCodeFromError(errStream) != http.StatusServiceUnavailable { + t.Fatalf("subsequent ExecuteStream while authA in 502 cooldown = %v, want 503", errStream) + } + if calls[authAID] != 1 || calls[authBID] != 0 { + t.Fatalf("calls after cooled 502 ExecuteStream = (A=%d, B=%d), want (1, 0)", calls[authAID], calls[authBID]) + } +} + +func TestManagerSessionAffinityRebindOnNeverBlocksAntigravityCreditsCrossAuthFallback(t *testing.T) { + ctx := context.Background() + provider := "antigravity" + model := "claude-sonnet-4-6" + authAID := "antigravity-never-a" + authBID := "antigravity-never-b" + + calls := make(map[string]int) + var failAuthA bool + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true}, + }) + manager.SetRetryConfig(0, 0, 0) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + RebindOn: SessionAffinityRebindOnNever, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(scriptedFailoverExecutor{ + provider: provider, + execute: func(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + calls[auth.ID]++ + if auth.ID == authAID && failAuthA { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota exhausted"} + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil + }, + stream: func(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + calls[auth.ID]++ + if auth.ID == authAID && failAuthA { + return nil, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota exhausted"} + } + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID)} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + }, + }) + + for _, auth := range []*Auth{ + {ID: authAID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "1"}}, + {ID: authBID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "0"}}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + SetAntigravityCreditsHint(authAID, AntigravityCreditsHint{Known: true, Available: false, UpdatedAt: time.Now()}) + SetAntigravityCreditsHint(authBID, AntigravityCreditsHint{Known: true, Available: true, UpdatedAt: time.Now()}) + t.Cleanup(func() { + SetAntigravityCreditsHint(authAID, AntigravityCreditsHint{}) + SetAntigravityCreditsHint(authBID, AntigravityCreditsHint{}) + }) + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"antigravity-never-session"}}, + } + + // 1. Cold Execute binds session to authA. + resp, errExec := manager.Execute(ctx, []string{provider}, req, opts) + if errExec != nil || string(resp.Payload) != authAID { + t.Fatalf("cold Execute = %q (err=%v), want %q", string(resp.Payload), errExec, authAID) + } + + // 2. authA returns 429 and has no credits; authB has available credits, but rebind-on: never must block switching to authB. + failAuthA = true + calls[authAID] = 0 + calls[authBID] = 0 + if _, errExec = manager.Execute(ctx, []string{provider}, req, opts); errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests { + t.Fatalf("Execute on 429 = %v, want 429", errExec) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls after Execute 429 = %d, want 0", calls[authBID]) + } + + if _, errStream := manager.ExecuteStream(ctx, []string{provider}, req, opts); errStream == nil || statusCodeFromError(errStream) != http.StatusTooManyRequests { + t.Fatalf("ExecuteStream on cooled authA = %v, want 429", errStream) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls after ExecuteStream 429 = %d, want 0", calls[authBID]) + } + + // 3. Explicit PinnedAuthMetadataKey pointing to authB on the same session must also be rejected and never call authB via credits fallback. + pinnedBOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"antigravity-never-session"}}, + Metadata: map[string]any{ + cliproxyexecutor.PinnedAuthMetadataKey: authBID, + }, + } + if _, errExec = manager.Execute(ctx, []string{provider}, req, pinnedBOpts); errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests { + t.Fatalf("Execute with conflicting PinnedAuthMetadataKey = %v, want 429", errExec) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls after conflicting pinned Execute = %d, want 0", calls[authBID]) + } + if _, errStream := manager.ExecuteStream(ctx, []string{provider}, req, pinnedBOpts); errStream == nil || statusCodeFromError(errStream) != http.StatusTooManyRequests { + t.Fatalf("ExecuteStream with conflicting PinnedAuthMetadataKey = %v, want 429", errStream) + } + if calls[authBID] != 0 { + t.Fatalf("authB calls after conflicting pinned ExecuteStream = %d, want 0", calls[authBID]) + } +} + +func TestManagerSessionAffinityRebindOnNeverSingleAuthDisabledOrModelRemovedReturns429(t *testing.T) { + ctx := context.Background() + provider := "single-never-empty-cand" + model := "single-never-model" + authID := provider + "-only" + + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 0) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + RebindOn: SessionAffinityRebindOnNever, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(scriptedFailoverExecutor{provider: provider}) + + auth := &Auth{ID: authID, Provider: provider, Status: StatusActive} + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"single-never-session"}}, + } + + // 1. Bind session to authID. + resp, errExec := manager.Execute(ctx, []string{provider}, req, opts) + if errExec != nil || string(resp.Payload) != authID { + t.Fatalf("cold Execute = %q (err=%v), want %q", string(resp.Payload), errExec, authID) + } + + // 2. Disable the only credential -> candidate list is empty, bound session must return 429 session_auth_unavailable (not 503). + disabled := auth.Clone() + disabled.Disabled = true + disabled.Status = StatusDisabled + if _, errUpdate := manager.Update(WithSkipPersist(ctx), disabled); errUpdate != nil { + t.Fatalf("Update(disabled): %v", errUpdate) + } + _, errExec = manager.Execute(ctx, []string{provider}, req, opts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests { + t.Fatalf("Execute when only bound auth is disabled = %v, want 429 session_auth_unavailable", errExec) + } + var authErr *Error + if !errors.As(errExec, &authErr) || authErr == nil || authErr.Code != "session_auth_unavailable" { + t.Fatalf("Execute disabled bound auth error = %v, want session_auth_unavailable", errExec) + } + + // An unbound session when candidates is empty must still receive auth_not_found. + unboundOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"single-never-unbound-session"}}, + } + if _, errUnbound := manager.Execute(ctx, []string{provider}, req, unboundOpts); !isAuthNotFoundError(errUnbound) { + t.Fatalf("Execute unbound session with empty candidates = %v, want auth_not_found", errUnbound) + } + + // 3. Re-enable auth, but remove model support from registry -> bound session must also return 429 session_auth_unavailable. + enabled := auth.Clone() + enabled.Disabled = false + enabled.Status = StatusActive + if _, errUpdate := manager.Update(WithSkipPersist(ctx), enabled); errUpdate != nil { + t.Fatalf("Update(enabled): %v", errUpdate) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: "other-model"}}) + _, errExec = manager.Execute(ctx, []string{provider}, req, opts) + if errExec == nil || statusCodeFromError(errExec) != http.StatusTooManyRequests { + t.Fatalf("Execute when bound auth loses model support = %v, want 429 session_auth_unavailable", errExec) + } + if !errors.As(errExec, &authErr) || authErr == nil || authErr.Code != "session_auth_unavailable" { + t.Fatalf("Execute model-removed bound auth error = %v, want session_auth_unavailable", errExec) + } +} diff --git a/sdk/cliproxy/auth/siwc_refresh_cancellation_test.go b/sdk/cliproxy/auth/siwc_refresh_cancellation_test.go new file mode 100644 index 00000000000..9fa0ce60f22 --- /dev/null +++ b/sdk/cliproxy/auth/siwc_refresh_cancellation_test.go @@ -0,0 +1,57 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +func TestSIWCWakeCancellationPreservesOuterRefreshFailure(t *testing.T) { + for _, terminalOuter := range []bool{false, true} { + t.Run(fmt.Sprintf("terminalOuter=%t", terminalOuter), func(t *testing.T) { + var outerErr error + if terminalOuter { + outerErr = errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused"}}`) + } + deferred := codex.WrapSIWCRefreshError(fmt.Errorf("SIWC renewal deferred until interactive macOS wake: %w", codex.ErrSIWCRefreshDeferred)) + executor := &issue6819OuterOnlyRefreshExecutor{id: "codex", refreshErr: codex.NewPartitionRefreshError(outerErr, deferred)} + manager := NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &Auth{ID: "wake-cancellation", Provider: "codex", Metadata: map[string]any{ + "access_token": "outer-access", "refresh_token": "outer-refresh", "use_siwc": true, + "expired": time.Now().Add(time.Hour).Format(time.RFC3339), + "siwc": map[string]any{ + "access_token": "siwc-access", "refresh_token": "siwc-refresh", "client_id": "oaiapp_cancellation", + "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }, + }} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), auth.ID); errRefresh == nil { + t.Fatal("expected deferred refresh") + } + current, ok := manager.GetByID(auth.ID) + if !ok { + t.Fatal("credential disappeared") + } + if terminalOuter { + if !current.RequiresSignIn || current.RejectedRefreshToken != "outer-refresh" { + t.Fatal("SIWC cancellation hid the terminal outer refresh rejection") + } + } else if current.RequiresSignIn || current.RefreshError != nil || current.NextRefreshAfter.IsZero() { + t.Fatal("pure SIWC cancellation did not use the cancellation rescheduler") + } + if current.RejectedSIWCRefreshToken != "" { + t.Fatal("wake deferral marked the SIWC grant as rejected") + } + if current.CodexSIWCStatus(time.Now()) != SIWCStatusValid { + t.Fatal("outer refresh rejection invalidated the working SIWC access token") + } + }) + } +} diff --git a/sdk/cliproxy/auth/siwc_renewal_window_test.go b/sdk/cliproxy/auth/siwc_renewal_window_test.go new file mode 100644 index 00000000000..3f91d8701c6 --- /dev/null +++ b/sdk/cliproxy/auth/siwc_renewal_window_test.go @@ -0,0 +1,66 @@ +package auth + +import ( + "testing" + "time" +) + +func TestSIWCRenewalUsesRealExpiryAndServerWindow(t *testing.T) { + expiry := time.Date(2026, 10, 29, 13, 8, 9, 0, time.UTC) + for _, tc := range []struct { + name string + earliest any + wantDue string + }{ + {"real expiry without double margin", nil, "2026-10-29T13:03:09Z"}, + {"early server window", "2026-10-29T12:58:09Z", "2026-10-29T13:03:09Z"}, + {"server RFC3339 window", "2026-10-29T13:07:09Z", "2026-10-29T13:07:09Z"}, + {"server Unix window", int64(1793279229), "2026-10-29T13:07:09Z"}, + {"server JSON Unix window", float64(1793279229), "2026-10-29T13:07:09Z"}, + {"malformed optional window", "invalid", "2026-10-29T13:03:09Z"}, + } { + t.Run(tc.name, func(t *testing.T) { + siwc := map[string]any{ + "access_token": "access", "refresh_token": "refresh", "client_id": "oaiapp_test", + "last_refresh": "2026-10-29T12:08:09Z", "expired": "2026-10-29T13:05:09Z", + "expires_at": "2026-10-29T13:08:09Z", "earliest_refresh_at": tc.earliest, + } + auth := &Auth{Provider: "codex", Metadata: map[string]any{"siwc": siwc}} + wantDue, errParse := time.Parse(time.RFC3339, tc.wantDue) + if errParse != nil { + t.Fatal(errParse) + } + due, has, need := CodexSIWCNeedsRefresh(auth, expiry.Add(-8*time.Minute)) + if !has || need || !due.Equal(wantDue) { + t.Fatalf("due = %v, has = %v, need = %v; want %v, true, false", due, has, need, wantDue) + } + for _, now := range []time.Time{wantDue, expiry.Add(3 * time.Hour)} { + if _, has, need := CodexSIWCNeedsRefresh(auth, now); !has || !need { + t.Fatalf("at %v: has = %v, need = %v; want renewal", now, has, need) + } + } + }) + } +} + +func TestSIWCRenewalServerWindowWithoutUsableAccessToken(t *testing.T) { + now := time.Date(2026, 10, 29, 12, 0, 0, 0, time.UTC) + earliest := now.Add(time.Minute) + for _, accessToken := range []string{"", "rejected", "opaque-without-expiry"} { + t.Run(accessToken, func(t *testing.T) { + auth := &Auth{RejectedSIWCAccessToken: "rejected", Metadata: map[string]any{ + "siwc": map[string]any{ + "access_token": accessToken, "refresh_token": "refresh", "client_id": "oaiapp_window", + "earliest_refresh_at": "2026-10-29T12:01:00Z", + }, + }} + due, has, need := CodexSIWCNeedsRefresh(auth, now) + if !has || need || !due.Equal(earliest) { + t.Fatalf("due = %v, has = %v, need = %v; want %v, true, false", due, has, need, earliest) + } + if _, has, need := CodexSIWCNeedsRefresh(auth, earliest); !has || !need { + t.Fatal("renewal was not due at the server window") + } + }) + } +} diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 1d43d746c48..7cf8d91aa7c 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -1094,6 +1094,15 @@ const CodexOuterRefreshLead = 96 * time.Hour // a Sign in with ChatGPT (SIWC) token, matching pi's DEFAULT_OAUTH_MINIMUM_VALIDITY_MS (5 minutes). const CodexSIWCRefreshLead = 5 * time.Minute +// SIWCEarliestRefreshTime returns the server's lower bound for SIWC token renewal. +func (a *Auth) SIWCEarliestRefreshTime() (time.Time, bool) { + if a == nil || a.Metadata == nil { + return time.Time{}, false + } + siwc, _ := a.Metadata["siwc"].(map[string]any) + return lookupMetadataTime(siwc, "earliest_refresh_at") +} + // SIWCStatus represents the read-only validity state of a Codex SIWC credential. type SIWCStatus string @@ -1311,28 +1320,26 @@ func CodexSIWCNeedsRefresh(a *Auth, now time.Time) (dueAt time.Time, hasRefresh siwcToken, _ := siwcMap["access_token"].(string) siwcToken = strings.TrimSpace(siwcToken) - if siwcToken == "" { - return now, true, true - } - if a.isSIWCAccessTokenRejected(siwcToken) { - return now, true, true - } - expiry, hasExpiry := a.SIWCExpirationTime() - if hasExpiry && !expiry.IsZero() { - due := expiry.Add(-CodexSIWCRefreshLead) - if !due.After(now) { - return now, true, true + due := now + if siwcToken != "" && !a.isSIWCAccessTokenRejected(siwcToken) { + // The legacy expired field already includes a margin; prefer the real expiry. + expiry, hasExpiry := a.SIWCRealExpirationTime() + if !hasExpiry { + expiry, hasExpiry = a.SIWCExpirationTime() } - return due, true, false - } - if lastRefresh, okLast := lookupMetadataTime(siwcMap, "last_refresh", "lastRefresh"); okLast && !lastRefresh.IsZero() { - due := lastRefresh.Add(CodexSIWCRefreshLead) - if !due.After(now) { - return now, true, true + if hasExpiry && !expiry.IsZero() { + due = expiry.Add(-CodexSIWCRefreshLead) + } else if lastRefresh, okLast := lookupMetadataTime(siwcMap, "last_refresh", "lastRefresh"); okLast && !lastRefresh.IsZero() { + due = lastRefresh.Add(CodexSIWCRefreshLead) } - return due, true, false } - return now, true, true + if earliest, okEarliest := a.SIWCEarliestRefreshTime(); okEarliest && earliest.After(due) { + due = earliest + } + if !due.After(now) { + return now, true, true + } + return due, true, false } // OuterAccessTokenExpirationTime returns the expiration time of the outer access_token, diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index cbda28e2716..e0e2d8d07d7 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -92,8 +92,13 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { if state.sessionAffinity && cfg.Routing.SessionAffinitySubagents != nil { state.sessionAffinitySubagents = *cfg.Routing.SessionAffinitySubagents } - if state.sessionAffinity && strings.EqualFold(strings.TrimSpace(cfg.Routing.SessionAffinityRebindOn), coreauth.SessionAffinityRebindOnQuota) { - state.sessionAffinityRebindOn = coreauth.SessionAffinityRebindOnQuota + if state.sessionAffinity { + switch strings.ToLower(strings.TrimSpace(cfg.Routing.SessionAffinityRebindOn)) { + case coreauth.SessionAffinityRebindOnQuota: + state.sessionAffinityRebindOn = coreauth.SessionAffinityRebindOnQuota + case coreauth.SessionAffinityRebindOnNever: + state.sessionAffinityRebindOn = coreauth.SessionAffinityRebindOnNever + } } return state } diff --git a/sdk/cliproxy/service_executionregistry_test.go b/sdk/cliproxy/service_executionregistry_test.go index 38f7f785d09..e168c9e93bd 100644 --- a/sdk/cliproxy/service_executionregistry_test.go +++ b/sdk/cliproxy/service_executionregistry_test.go @@ -2966,6 +2966,46 @@ func TestServiceApplyConfigRuntimeSessionAffinityRebindOnChangeRecreatesSelector if changedSelector == initialSelector { t.Fatal("changing SessionAffinityRebindOn should recreate selector") } + + neverCfg := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + SessionAffinityRebindOn: "never", + }}) + if !service.applyConfigRuntime(context.Background(), neverCfg, false) { + t.Fatal("never config runtime apply failed") + } + neverSelector := manager.Selector() + if neverSelector == changedSelector { + t.Fatal("changing SessionAffinityRebindOn to never should recreate selector") + } + + sameNeverCfg := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + SessionAffinityRebindOn: " NEVER ", + }}) + if !service.applyConfigRuntime(context.Background(), sameNeverCfg, false) { + t.Fatal("same never config runtime apply failed") + } + if manager.Selector() != neverSelector { + t.Fatal("equivalent never SessionAffinityRebindOn should preserve selector") + } + + backToAnyCfg := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + SessionAffinityRebindOn: "any", + }}) + if !service.applyConfigRuntime(context.Background(), backToAnyCfg, false) { + t.Fatal("back to any config runtime apply failed") + } + if manager.Selector() == neverSelector { + t.Fatal("changing SessionAffinityRebindOn from never to any should recreate selector") + } } func TestServiceSerializesHomeAndWatcherConfigRuntimeApply(t *testing.T) { diff --git a/sdk/cliproxy/usage/home_selection.go b/sdk/cliproxy/usage/home_selection.go new file mode 100644 index 00000000000..c197e51865c --- /dev/null +++ b/sdk/cliproxy/usage/home_selection.go @@ -0,0 +1,37 @@ +package usage + +import "context" + +// HomeSelection identifies executions associated with a Home dispatch. Dispatch +// charges the first execution; reused executions are charged by primary usage. +// Additional-model records are billing details, not extra scheduling requests. +// Model is the scheduling route, not the upstream model. +type HomeSelection struct { + Reused bool `json:"reused"` + Additional bool `json:"additional,omitempty"` + Model string `json:"model"` + CredentialPolicy string `json:"credential_policy,omitempty"` +} + +type homeSelectionContextKey struct{} + +// WithHomeSelection captures immutable, attempt-scoped Home accounting metadata. +func WithHomeSelection(ctx context.Context, selection HomeSelection) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, homeSelectionContextKey{}, selection) +} + +// HomeSelectionFromContext returns a copy so follow-up responses cannot change +// the accounting metadata retained by the first response's usage reporter. +func HomeSelectionFromContext(ctx context.Context) *HomeSelection { + if ctx == nil { + return nil + } + selection, ok := ctx.Value(homeSelectionContextKey{}).(HomeSelection) + if !ok { + return nil + } + return &selection +} diff --git a/sdk/cliproxy/usage/home_selection_test.go b/sdk/cliproxy/usage/home_selection_test.go new file mode 100644 index 00000000000..036030efaea --- /dev/null +++ b/sdk/cliproxy/usage/home_selection_test.go @@ -0,0 +1,22 @@ +package usage + +import "testing" + +func TestHomeSelectionContextAndPublishSnapshot(t *testing.T) { + ctx := WithHomeSelection(t.Context(), HomeSelection{Model: "route"}) + followup := HomeSelectionFromContext(ctx) + followup.Reused = true + if HomeSelectionFromContext(ctx).Reused { + t.Fatal("follow-up mutated the first turn") + } + mgr := NewManager(10) + t.Cleanup(mgr.Stop) + plugin := &captureUsagePlugin{records: make(chan Record, 2)} + mgr.Register(plugin) + mgr.Publish(ctx, Record{}) + mgr.Publish(ctx, Record{HomeSelection: followup}) + first, second := <-plugin.records, <-plugin.records + if first.HomeSelection == nil || first.HomeSelection.Reused || second.HomeSelection == nil || !second.HomeSelection.Reused { + t.Fatalf("lost attempt metadata: first=%+v second=%+v", first.HomeSelection, second.HomeSelection) + } +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 9b38c25c007..4ddaebc04d7 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -30,7 +30,9 @@ type Record struct { TraceID string // SelectionID identifies the scheduler selection decision that produced this usage record when available. SelectionID string - Provider string + // HomeSelection carries deferred accounting for reused Home dispatches. + HomeSelection *HomeSelection + Provider string // BaseURL stores the configured upstream base URL when available. BaseURL string // ExecutorType stores the concrete executor type that handled the request. @@ -737,6 +739,9 @@ func (m *Manager) Publish(ctx context.Context, record Record) { record.TraceID = trID } } + if record.HomeSelection == nil { + record.HomeSelection = HomeSelectionFromContext(ctx) + } if strings.TrimSpace(record.SelectionID) == "" { if selID := SelectionIDFromContext(ctx); selID != "" { record.SelectionID = selID diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index 2d7b2654a5f..49a8c972a1d 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -1,6 +1,9 @@ package pluginabi -import "encoding/json" +import ( + "encoding/json" + "strings" +) const ( // ABIVersion tracks the native C ABI shape (native plugin exports). @@ -149,12 +152,65 @@ const ( MethodHostRoutingResetCooldown = "host.routing.reset_cooldown" ) +// ErrorCode constants define the machine-readable Error.Code values recognized +// by CPA's routing and cooldown pipeline. +const ( + // ErrorCodeCredentialQuota marks a 429 refusal as exhausting or limiting the + // entire credential across all of its models (sets Result.CredentialScope and + // records quota reason "credential_quota"). Plugins may also set + // Error.CredentialScoped = true while keeping a provider-specific Code. + ErrorCodeCredentialQuota = "credential_quota" + // ErrorCodeTransientRateLimit marks a 429 refusal as a transient per-model + // capacity/admission limit that cools only the reporting model until + // RetryAfterSeconds without recording quota exhaustion or advancing backoff. + ErrorCodeTransientRateLimit = "transient_rate_limit" + // ErrorCodeRequestScoped marks a failure as caused by the request itself: + // CPA skips credential/model cooldown and stops credential rotation. + ErrorCodeRequestScoped = "request_scoped" + // ErrorCodeForceCooldown forces cooldown scheduling even when cooling is + // disabled for the credential, falling back to the default transient cooldown + // when no status-specific deadline applies. + ErrorCodeForceCooldown = "force_cooldown" + // ErrorCodeStreamTruncated marks an accepted upstream stream that terminated + // without a completion frame: CPA skips credential cooldown and does not + // replay the request on another credential. + ErrorCodeStreamTruncated = "stream_truncated" + // ErrorCodeConnectionLifecycle marks a client disconnect or session-close + // failure that must skip credential cooldown. + ErrorCodeConnectionLifecycle = "connection_lifecycle" + // ErrorCodeTransientTransport marks a pre-HTTP dial/TLS/DNS/reset failure + // that skips credential cooldown and remains eligible for request retries. + ErrorCodeTransientTransport = "transient_transport" + // ErrorCodeModelNotFound marks the requested model as unavailable on the + // selected credential so CPA cools only that model for 12h (or RetryAfterSeconds). + ErrorCodeModelNotFound = "model_not_found" + // ErrorCodeNotImplemented marks an optional capability (such as count_tokens) + // as unsupported so CPA keeps the credential and model available. + ErrorCodeNotImplemented = "not_implemented" +) + type Envelope struct { OK bool `json:"ok"` Result json.RawMessage `json:"result,omitempty"` Error *Error `json:"error,omitempty"` } +// Error carries a structured plugin failure across RPC envelopes and stream +// error_detail callbacks (host.stream.emit and host.stream.close). +// +// Routing and cooldown behavior is controlled by HTTPStatus, RetryAfterSeconds, +// CredentialScoped, and the recognized ErrorCode* values in Code: +// - CredentialScoped: true (or Code: ErrorCodeCredentialQuota) on a 429 marks +// the entire credential as rate-limited across all of its models +// (reason "credential_quota") until RetryAfterSeconds or the quota backoff. +// - Code: ErrorCodeTransientRateLimit on a 429 cools only the reporting model +// until RetryAfterSeconds without marking quota exceeded. +// - Code: ErrorCodeRequestScoped, ErrorCodeStreamTruncated, +// ErrorCodeConnectionLifecycle, or ErrorCodeTransientTransport skips +// credential/model cooldown. +// - Code: ErrorCodeForceCooldown enforces cooldown even when cooling is disabled. +// - Code: ErrorCodeModelNotFound cools only the requested model as unsupported. +// - Code: ErrorCodeNotImplemented keeps count_tokens failures availability-neutral. type Error struct { Code string `json:"code"` Message string `json:"message"` @@ -168,6 +224,10 @@ type Error struct { // for correctness. The same pair is accepted as error_detail on host.stream.emit // and host.stream.close. RetryAfterSeconds int `json:"retry_after_seconds,omitempty"` + // CredentialScoped reports whether the failure applies to the whole credential + // across all models (for example, an account-wide 429 quota or spend limit) + // rather than only the requested model. + CredentialScoped bool `json:"credential_scoped,omitempty"` } // Error implements the error interface for Error. @@ -194,6 +254,38 @@ func (e *Error) StatusCode() int { return e.HTTPStatus } +// IsCredentialScoped reports whether the Error applies to the entire credential. +func (e *Error) IsCredentialScoped() bool { + if e == nil { + return false + } + return e.CredentialScoped || strings.EqualFold(strings.TrimSpace(e.Code), ErrorCodeCredentialQuota) +} + +// IsRequestScoped reports whether the Error is tied to the request rather than the credential. +func (e *Error) IsRequestScoped() bool { + if e == nil || e.IsCredentialScoped() { + return false + } + return strings.EqualFold(strings.TrimSpace(e.Code), ErrorCodeRequestScoped) +} + +// IsTransientRateLimit reports whether the Error is a transient per-model capacity 429. +func (e *Error) IsTransientRateLimit() bool { + if e == nil || e.IsCredentialScoped() { + return false + } + return strings.EqualFold(strings.TrimSpace(e.Code), ErrorCodeTransientRateLimit) +} + +// IsStreamTruncated reports whether the Error represents a truncated accepted upstream stream. +func (e *Error) IsStreamTruncated() bool { + if e == nil || e.IsCredentialScoped() { + return false + } + return strings.EqualFold(strings.TrimSpace(e.Code), ErrorCodeStreamTruncated) +} + // NewError creates an Error instance with an optional HTTP status code. func NewError(code, message string, httpStatus ...int) *Error { status := 0 diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index d79c4c04349..99ad87fb8fa 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -208,3 +208,87 @@ func TestErrorInterfaceAndConstructor(t *testing.T) { t.Fatalf("default envelope JSON unexpectedly contained http_status: %s", string(defaultEnvBytes)) } } + +func TestErrorCredentialScopedAndCodes(t *testing.T) { + var _ interface{ IsCredentialScoped() bool } = (*Error)(nil) + + var nilErr *Error + if nilErr.IsCredentialScoped() { + t.Fatal("nil Error reported IsCredentialScoped() = true") + } + + modelScoped := NewError("rate_limited", "per-model rate limit", 429) + if modelScoped.IsCredentialScoped() { + t.Fatal("model-scoped Error reported IsCredentialScoped() = true") + } + + flagScoped := &Error{ + Code: "rate_limited", + Message: "account limited, resets at 2026-01-01T00:00:00Z", + HTTPStatus: 429, + RetryAfterSeconds: 120, + CredentialScoped: true, + } + if !flagScoped.IsCredentialScoped() { + t.Fatal("Error with CredentialScoped=true reported IsCredentialScoped() = false") + } + + raw, errMarshal := json.Marshal(Envelope{OK: false, Error: flagScoped}) + if errMarshal != nil { + t.Fatalf("marshal credential-scoped envelope: %v", errMarshal) + } + if !strings.Contains(string(raw), `"credential_scoped":true`) { + t.Fatalf("marshaled envelope missing credential_scoped: %s", string(raw)) + } + var decoded Envelope + if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil { + t.Fatalf("unmarshal credential-scoped envelope: %v", errUnmarshal) + } + if decoded.Error == nil || !decoded.Error.CredentialScoped || !decoded.Error.IsCredentialScoped() { + t.Fatalf("decoded error = %+v, want CredentialScoped=true", decoded.Error) + } + + codeScoped := NewError(ErrorCodeCredentialQuota, "daily quota spent", 429) + if !codeScoped.IsCredentialScoped() { + t.Fatal("Error with ErrorCodeCredentialQuota reported IsCredentialScoped() = false") + } + if codeScoped.IsRequestScoped() || codeScoped.IsTransientRateLimit() { + t.Fatal("credential-scoped Error reported request-scoped or transient-rate-limit") + } + + reqScoped := NewError(ErrorCodeRequestScoped, "invalid prompt", 400) + if !reqScoped.IsRequestScoped() { + t.Fatal("Error with ErrorCodeRequestScoped reported IsRequestScoped() = false") + } + + transientRL := NewError(ErrorCodeTransientRateLimit, "turn rate limit", 429) + if !transientRL.IsTransientRateLimit() { + t.Fatal("Error with ErrorCodeTransientRateLimit reported IsTransientRateLimit() = false") + } + + truncated := NewError(ErrorCodeStreamTruncated, "stream ended early", 502) + if !truncated.IsStreamTruncated() { + t.Fatal("Error with ErrorCodeStreamTruncated reported IsStreamTruncated() = false") + } + truncatedCredScoped := &Error{ + Code: ErrorCodeStreamTruncated, + Message: "stream ended early", + HTTPStatus: 429, + CredentialScoped: true, + } + if !truncatedCredScoped.IsCredentialScoped() || truncatedCredScoped.IsStreamTruncated() { + t.Fatal("credential-scoped Error with stream_truncated code must prioritize IsCredentialScoped()") + } + + if ErrorCodeCredentialQuota != "credential_quota" || + ErrorCodeTransientRateLimit != "transient_rate_limit" || + ErrorCodeRequestScoped != "request_scoped" || + ErrorCodeForceCooldown != "force_cooldown" || + ErrorCodeStreamTruncated != "stream_truncated" || + ErrorCodeConnectionLifecycle != "connection_lifecycle" || + ErrorCodeTransientTransport != "transient_transport" || + ErrorCodeModelNotFound != "model_not_found" || + ErrorCodeNotImplemented != "not_implemented" { + t.Fatal("unexpected routing/cooldown error code constant values") + } +} diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 718961c7639..9e84fa29f33 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -1732,6 +1732,29 @@ type UsageDetail struct { CacheCreationTokens int64 // TotalTokens is the total token count. TotalTokens int64 + + // TextInputTokens is the optional text input token count. + TextInputTokens *int64 + // ImageInputTokens is the optional image input token count. + ImageInputTokens *int64 + // AudioInputTokens is the optional audio input token count. + AudioInputTokens *int64 + // VideoInputTokens is the optional video input token count. + VideoInputTokens *int64 + // FileInputTokens is the optional file input token count. + FileInputTokens *int64 + // CachedTextTokens is the optional cached text token count. + CachedTextTokens *int64 + // CachedImageTokens is the optional cached image token count. + CachedImageTokens *int64 + // CachedAudioTokens is the optional cached audio token count. + CachedAudioTokens *int64 + // TextOutputTokens is the optional text output token count. + TextOutputTokens *int64 + // ImageOutputTokens is the optional image output token count. + ImageOutputTokens *int64 + // AudioOutputTokens is the optional audio output token count. + AudioOutputTokens *int64 } // QuotaProvider surfaces credential quota, balance, and billing information for management clients. diff --git a/sdk/pluginapi/types_test.go b/sdk/pluginapi/types_test.go index 3db66fa9c1a..e40fa445cc9 100644 --- a/sdk/pluginapi/types_test.go +++ b/sdk/pluginapi/types_test.go @@ -869,6 +869,78 @@ func TestUsageRecordRequestIDAndTraceID(t *testing.T) { } } +func TestUsageDetailModalityTokenCounters(t *testing.T) { + int64Ptr := func(v int64) *int64 { return &v } + record := UsageRecord{ + Provider: "openai", + Model: "gpt-image-2.5-flare", + Detail: UsageDetail{ + InputTokens: 80, + OutputTokens: 1056, + CachedTokens: 40, + CacheReadTokens: 40, + TotalTokens: 1136, + TextInputTokens: int64Ptr(30), + ImageInputTokens: int64Ptr(50), + AudioInputTokens: int64Ptr(0), + VideoInputTokens: nil, + FileInputTokens: int64Ptr(12), + CachedTextTokens: int64Ptr(25), + CachedImageTokens: int64Ptr(15), + CachedAudioTokens: int64Ptr(0), + TextOutputTokens: int64Ptr(0), + ImageOutputTokens: int64Ptr(1056), + AudioOutputTokens: nil, + }, + } + + raw, errMarshal := json.Marshal(record) + if errMarshal != nil { + t.Fatalf("Marshal(UsageRecord) error = %v", errMarshal) + } + var decoded UsageRecord + if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil { + t.Fatalf("Unmarshal(UsageRecord) error = %v", errUnmarshal) + } + got := decoded.Detail + if got.InputTokens != 80 || got.OutputTokens != 1056 || got.TotalTokens != 1136 { + t.Fatalf("aggregate Detail = %+v", got) + } + if got.TextInputTokens == nil || *got.TextInputTokens != 30 { + t.Fatalf("TextInputTokens = %v, want 30", got.TextInputTokens) + } + if got.ImageInputTokens == nil || *got.ImageInputTokens != 50 { + t.Fatalf("ImageInputTokens = %v, want 50", got.ImageInputTokens) + } + if got.AudioInputTokens == nil || *got.AudioInputTokens != 0 { + t.Fatalf("AudioInputTokens = %v, want explicit 0", got.AudioInputTokens) + } + if got.VideoInputTokens != nil { + t.Fatalf("VideoInputTokens = %v, want nil", got.VideoInputTokens) + } + if got.FileInputTokens == nil || *got.FileInputTokens != 12 { + t.Fatalf("FileInputTokens = %v, want 12", got.FileInputTokens) + } + if got.CachedTextTokens == nil || *got.CachedTextTokens != 25 { + t.Fatalf("CachedTextTokens = %v, want 25", got.CachedTextTokens) + } + if got.CachedImageTokens == nil || *got.CachedImageTokens != 15 { + t.Fatalf("CachedImageTokens = %v, want 15", got.CachedImageTokens) + } + if got.CachedAudioTokens == nil || *got.CachedAudioTokens != 0 { + t.Fatalf("CachedAudioTokens = %v, want explicit 0", got.CachedAudioTokens) + } + if got.TextOutputTokens == nil || *got.TextOutputTokens != 0 { + t.Fatalf("TextOutputTokens = %v, want explicit 0", got.TextOutputTokens) + } + if got.ImageOutputTokens == nil || *got.ImageOutputTokens != 1056 { + t.Fatalf("ImageOutputTokens = %v, want 1056", got.ImageOutputTokens) + } + if got.AudioOutputTokens != nil { + t.Fatalf("AudioOutputTokens = %v, want nil", got.AudioOutputTokens) + } +} + func TestQuotaPayloadJSON(t *testing.T) { // Test camelCase input camelJSON := []byte(`{ diff --git a/test/codex_refresh_recovery_test.go b/test/codex_refresh_recovery_test.go new file mode 100644 index 00000000000..508166a3b3e --- /dev/null +++ b/test/codex_refresh_recovery_test.go @@ -0,0 +1,154 @@ +package test + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" + _ "github.com/router-for-me/CLIProxyAPI/v8/internal/translator" + coreauth "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" +) + +type codexRecoveryTransport struct { + target *url.URL + base http.RoundTripper +} + +func (t codexRecoveryTransport) RoundTrip(req *http.Request) (*http.Response, error) { + local := req.Clone(req.Context()) + local.URL.Scheme = t.target.Scheme + local.URL.Host = t.target.Host + return t.base.RoundTrip(local) +} + +func TestCodexSIWCRecoveryPreservesRequestPartition(t *testing.T) { + for _, tc := range []struct { + name string + missingSIWC bool + siwcStatus int + nativeAllowed bool + }{ + {name: "expired SIWC", siwcStatus: http.StatusServiceUnavailable, nativeAllowed: true}, + {name: "missing SIWC with healthy model state", missingSIWC: true, siwcStatus: http.StatusServiceUnavailable, nativeAllowed: true}, + {name: "explicit SIWC rejection", siwcStatus: http.StatusBadRequest}, + } { + t.Run(tc.name, func(t *testing.T) { + const model = "gpt-5.4" + var outerCalls, siwcCalls, inferenceCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch req.URL.Path { + case "/oauth/token": + outerCalls.Add(1) + _, _ = io.WriteString(w, `{"access_token":"outer-rotated","refresh_token":"outer-next","expires_in":86400}`) + case "/api/accounts/oauth/token": + siwcCalls.Add(1) + w.WriteHeader(tc.siwcStatus) + body := `{"error":"server_error"}` + if tc.siwcStatus == http.StatusBadRequest { + body = `{"error":"invalid_grant","code":"refresh_token_reused"}` + } + _, _ = io.WriteString(w, body) + default: + inferenceCalls.Add(1) + if got := req.Header.Get("Authorization"); got != "Bearer outer-rotated" { + t.Errorf("inference used %q, want the refreshed outer token", got) + w.WriteHeader(http.StatusUnauthorized) + return + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"recovery-response\",\"status\":\"completed\",\"model\":\"gpt-5.4\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n") + } + })) + t.Cleanup(server.Close) + target, errParse := url.Parse(server.URL) + if errParse != nil { + t.Fatal(errParse) + } + transport := &http.Transport{} + t.Cleanup(transport.CloseIdleConnections) + originalTransport := http.DefaultTransport + http.DefaultTransport = codexRecoveryTransport{target: target, base: transport} + t.Cleanup(func() { http.DefaultTransport = originalTransport }) + + cfg := &config.Config{AuthDir: t.TempDir()} + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(runtimeexecutor.NewCodexExecutor(cfg)) + authID := t.Name() + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + now := time.Now() + siwc := map[string]any{"refresh_token": authID + "-siwc", "client_id": "issued-client"} + if !tc.missingSIWC { + siwc["access_token"] = "siwc-expired" + siwc["expires_at"] = now.Add(-time.Hour).Format(time.RFC3339) + } + auth := &coreauth.Auth{ID: authID, Provider: "codex", Status: coreauth.StatusActive, + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{ + "access_token": "outer-before", "refresh_token": authID + "-outer", "expired": now.Add(time.Hour).Format(time.RFC3339), + "use_siwc": true, "siwc": siwc, + }, + } + if tc.missingSIWC { + auth.ModelStates = map[string]*coreauth.ModelState{model: {Status: coreauth.StatusActive}} + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + if _, errRefresh := manager.ForceRefreshAuth(context.Background(), authID); errRefresh != nil { + t.Fatal(errRefresh) + } + if outerCalls.Load() != 1 || siwcCalls.Load() != 1 { + t.Fatalf("refresh HTTP calls = outer %d / SIWC %d, want 1 / 1", outerCalls.Load(), siwcCalls.Load()) + } + recovered, _ := manager.GetByID(authID) + if !recovered.RequiresSignIn || recovered.RefreshError == nil { + t.Fatal("the unresolved or rejected SIWC grant lost its recovery requirement") + } + request := cliproxyexecutor.Request{Model: model, Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)} + assertSIWCBlocked := func() { + t.Helper() + _, errExecute := manager.Execute(context.Background(), []string{"codex"}, request, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + var authErr *coreauth.Error + if !errors.As(errExecute, &authErr) || authErr.Code != "auth_unavailable" { + t.Fatalf("unusable SIWC credential reached execution: %v", errExecute) + } + } + assertSIWCBlocked() + _, errNative := manager.Execute(context.Background(), []string{"codex"}, request, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Headers: http.Header{"User-Agent": {"codex_cli_rs/1.0.0"}}, + }) + if (errNative == nil) != tc.nativeAllowed { + t.Fatalf("native outer inference: error=%v, want allowed=%v", errNative, tc.nativeAllowed) + } + assertSIWCBlocked() + _, _ = manager.ForceRefreshAuth(context.Background(), authID) + if siwcCalls.Load() != 1 { + t.Fatal("a later refresh resubmitted the fenced SIWC grant") + } + assertSIWCBlocked() + wantInference := int32(0) + if tc.nativeAllowed { + wantInference = 1 + } + if inferenceCalls.Load() != wantInference { + t.Fatalf("inference calls = %d, want %d", inferenceCalls.Load(), wantInference) + } + }) + } +}