diff --git a/README.md b/README.md index 5146b87f304..39f30075fd2 100644 --- a/README.md +++ b/README.md @@ -110,11 +110,11 @@ You can access the following providers locally and with multiple CLI accounts th Kimi - Kimi series models (Kimi K3, K2.8 Preview, etc.). Kimi K3 is Moonshot AI’s most capable model and the world’s first open 3T-class model. With 2.8 trillion parameters, native vision, and a 1-million-token context window, K3 is built for long-horizon coding, knowledge work, and reasoning. CLIProxyAPI supports Kimi through OAuth or compatible API interfaces. Try a Kimi Code plan (中文站 | Global), or get an API key from the Kimi Open Platform (中文站 | Global). Thanks to Kimi for supporting CLIProxyAPI and the open-source community! + Kimi series models (Kimi K3.1 Pro, K3.1 Pro Highspeed, K3, K2.8 Preview, etc.). Kimi K3 is Moonshot AI’s previous-generation flagship and the world’s first open 3T-class model. With 2.8 trillion parameters, native vision, and a 1-million-token context window, K3 is built for long-horizon coding, knowledge work, and reasoning. CLIProxyAPI supports Kimi through OAuth or compatible API interfaces. Try a Kimi Code plan (中文站 | Global), or get an API key from the Kimi Open Platform (中文站 | Global). Thanks to Kimi for supporting CLIProxyAPI and the open-source community! OpenAI - OpenAI GPT-6 models (GPT-6 Astra, GPT-6.1 Sol, GPT-6 Luna), including access through Codex OAuth. Astra is the flagship for complex reasoning and coding, Sol balances capability and cost, and Luna handles focused, high-volume work. + OpenAI GPT-6 models (GPT-6 Astra, GPT-6.1 Sol, GPT-6.1 Terra, GPT-6 Luna), including access through Codex OAuth. Astra is the flagship for complex reasoning and coding, Sol balances capability and cost, Terra focuses on visual reasoning and computer use, and Luna handles focused, high-volume work. Anthropic @@ -404,6 +404,14 @@ A Cloudflare-based LLM gateway and proxy server: CLIProxyAPI deployed to Cloudfl 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`. +### [Ebb](https://github.com/lucasnsaraujo/ebb) + +Native macOS menu bar app for CLIProxyAPI pools. It shows how much of the pool's Claude, Codex and other providers' 5-hour and weekly limits is left and when it comes back, and sets each login's routing weight through the Management API so the accounts drain evenly. + +### [Nia Live Demo](https://github.com/nia-teams/nia-live-demo) + +Real-time full-duplex voice companion with a live, interactive SVG/HTML canvas. It combines two CLIProxyAPI capabilities: the Codex GPT-Live 1.1 media relay (`gpt-live-1-codex`, WebRTC + Sideband) for voice with instant barge-in, and `ultrafast` Responses (`gpt-6.1-sol` / `gpt-6-astra`) for drawing and editing the canvas while you talk. + > [!NOTE] > If you developed a project based on CLIProxyAPI, please open a PR to add it to this list. diff --git a/README_CN.md b/README_CN.md index 2d2f04002df..4a44e757312 100644 --- a/README_CN.md +++ b/README_CN.md @@ -16,11 +16,11 @@ CLIProxyAPI 是一个为 CLI 提供 OpenAI/Gemini/Claude/Codex/Grok 兼容 API Kimi - Kimi 系列模型（Kimi K3、K2.8 Preview 等）。Kimi K3 是 Moonshot AI 迄今能力最强的模型，也是全球首个开源 3T 级模型。K3 拥有 2.8T 参数、原生视觉能力与 100 万 Token 上下文，面向长周期编码、知识工作与推理任务。CLIProxyAPI 支持通过 OAuth 或兼容 API 接入 Kimi。立即体验 Kimi Code 订阅（中文站｜Global），或前往 Kimi 开放平台（中文站｜Global）获取 API Key。感谢 Kimi 对 CLIProxyAPI 及开源社区的支持！ + Kimi 系列模型（Kimi K3.1 Pro、K3.1 Pro Highspeed、K3、K2.8 Preview 等）。Kimi K3 是 Moonshot AI 的上一代旗舰模型，也是全球首个开源 3T 级模型。K3 拥有 2.8T 参数、原生视觉能力与 100 万 Token 上下文，面向长周期编码、知识工作与推理任务。CLIProxyAPI 支持通过 OAuth 或兼容 API 接入 Kimi。立即体验 Kimi Code 订阅（中文站｜Global），或前往 Kimi 开放平台（中文站｜Global）获取 API Key。感谢 Kimi 对 CLIProxyAPI 及开源社区的支持！ OpenAI - OpenAI GPT-6 系列模型（GPT-6 Astra、GPT-6.1 Sol、GPT-6 Luna），也支持通过 Codex OAuth 接入。Astra 适合复杂推理与编程，Sol 兼顾能力与成本，Luna 适合高吞吐量的明确任务。 + OpenAI GPT-6 系列模型（GPT-6 Astra、GPT-6.1 Sol、GPT-6.1 Terra、GPT-6 Luna），也支持通过 Codex OAuth 接入。Astra 适合复杂推理与编程，Sol 兼顾能力与成本，Terra 专注于视觉推理与计算机操作，Luna 适合高吞吐量的明确任务。 Anthropic @@ -300,6 +300,14 @@ Claude Code 状态栏，按当前 CPA 实例展示 Codex / Grok / Antigravity / [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`。 +### [Ebb](https://github.com/lucasnsaraujo/ebb) + +面向 CLIProxyAPI 账号池的原生 macOS 菜单栏应用。显示账号池中 Claude、Codex 等服务的 5 小时与每周额度还剩多少、何时恢复，并通过 Management API 设置每个登录的路由权重，让各账号均匀消耗。 + +### [Nia Live Demo](https://github.com/nia-teams/nia-live-demo) + +全双工实时语音伙伴，带一块边聊边画、可点击交互的 SVG/HTML 画布。同时依赖 CLIProxyAPI 的两项能力：Codex GPT-Live 1.1 媒体中继（`gpt-live-1-codex`，WebRTC + Sideband）负责语音和随时打断；`ultrafast` 档 Responses（`gpt-6.1-sol` / `gpt-6-astra`）负责在对话过程中实时生成和修改画布。 + > [!NOTE] > 如果你开发了基于 CLIProxyAPI 的项目，请提交一个 PR（拉取请求）将其添加到此列表中。 diff --git a/README_JA.md b/README_JA.md index 5cd6b3b7f26..7c7af13f269 100644 --- a/README_JA.md +++ b/README_JA.md @@ -16,11 +16,11 @@ CLIProxyAPI は、CLI向けのOpenAI/Gemini/Claude/Codex/Grok互換APIインタ Kimi - Kimiシリーズモデル（Kimi K3、K2.8 Previewなど）。Kimi K3は、Moonshot AIで最も高性能なモデルであり、世界初のオープンな3兆パラメータ級モデルです。2.8兆のパラメータ、ネイティブな視覚機能、100万トークンのコンテキストウィンドウを備え、長期間にわたるコーディング、知識作業、推論向けに構築されています。CLIProxyAPIはOAuthまたは互換APIインターフェース経由でKimiをサポートします。Kimi Code プラン（中文站 | Global）を試すか、Kimi Open Platform（中文站 | Global）でAPIキーを取得してください。CLIProxyAPIとオープンソースコミュニティを支援してくださるKimiに感謝します！ + Kimiシリーズモデル（Kimi K3.1 Pro、K3.1 Pro Highspeed、K3、K2.8 Previewなど）。Kimi K3は、Moonshot AIの前世代のフラッグシップモデルであり、世界初のオープンな3兆パラメータ級モデルです。2.8兆のパラメータ、ネイティブな視覚機能、100万トークンのコンテキストウィンドウを備え、長期間にわたるコーディング、知識作業、推論向けに構築されています。CLIProxyAPIはOAuthまたは互換APIインターフェース経由でKimiをサポートします。Kimi Code プラン（中文站 | Global）を試すか、Kimi Open Platform（中文站 | Global）でAPIキーを取得してください。CLIProxyAPIとオープンソースコミュニティを支援してくださるKimiに感謝します！ OpenAI - OpenAI GPT-6シリーズモデル（GPT-6 Astra、GPT-6.1 Sol、GPT-6 Luna）。Codex OAuth経由でも利用できます。Astraは複雑な推論とコーディング、Solは性能とコストのバランス、Lunaは大量の明確なタスクに適しています。 + OpenAI GPT-6シリーズモデル（GPT-6 Astra、GPT-6.1 Sol、GPT-6.1 Terra、GPT-6 Luna）。Codex OAuth経由でも利用できます。Astraは複雑な推論とコーディング、Solは性能とコストのバランス、Terraは視覚的推論とコンピューター操作、Lunaは大量の明確なタスクに適しています。 Anthropic @@ -299,6 +299,14 @@ Cloudflare ベースの LLM ゲートウェイ兼プロキシサーバー。CLIP [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`。 +### [Ebb](https://github.com/lucasnsaraujo/ebb) + +CLIProxyAPI のアカウントプール向けのネイティブ macOS メニューバーアプリ。プール全体で Claude、Codex などの 5 時間枠と週間枠がどれだけ残っているか、いつ戻るかを表示し、Management API で各ログインのルーティング重みを設定してアカウントを均等に消費させます。 + +### [Nia Live Demo](https://github.com/nia-teams/nia-live-demo) + +話しながらリアルタイムで描画・操作できるインタラクティブな SVG/HTML キャンバスを備えた、全二重音声コンパニオン。CLIProxyAPI の 2 つの機能を組み合わせています：音声対話と即時割り込みを担う Codex GPT-Live 1.1 メディアリレー（`gpt-live-1-codex`、WebRTC + Sideband）と、会話中にキャンバスをリアルタイムで生成・編集する `ultrafast` Responses（`gpt-6.1-sol` / `gpt-6-astra`）です。 + > [!NOTE] > CLIProxyAPIをベースにプロジェクトを開発した場合は、PRを送ってこのリストに追加してください。 diff --git a/config.example.yaml b/config.example.yaml index 880d013146a..27dca93332c 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -36,7 +36,8 @@ server: # Negotiated compression of downstream HTTP responses (including SSE streams). # Responses are compressed only when the client advertises zstd or gzip in - # Accept-Encoding; WebSocket upgrades are never wrapped. + # Accept-Encoding; WebSocket upgrades are never wrapped. Encoder memory is + # retained for the lifetime of an open stream and depends on prior pool usage. response-compression: # Default: true. enabled: true diff --git a/docs/management-api-v8.md b/docs/management-api-v8.md index afdc17e1445..f91304442d9 100644 --- a/docs/management-api-v8.md +++ b/docs/management-api-v8.md @@ -104,8 +104,9 @@ retain the corresponding business operation's fields. | `/credentials/status` | PATCH | Change credential status. | | `/credentials/fields` | PATCH | Change credential fields. | | `/credentials/refresh` | POST | Refresh credentials. | +| `/credentials/verify` | POST | Run an operator-triggered, account-pinned inference verification using an isolated session and minimal 1-token prompt. | | `/oauth/import?provider=vertex` | POST | Import a Vertex service account using a multipart `file` upload. | -| `/oauth/auth-url?provider=` | GET | Start built-in or plugin OAuth. | +| `/oauth/auth-url?provider=` | GET, POST | Start built-in or plugin OAuth (supports `name`/`auth_index` for targeted credential recovery and optional `verify` opt-in). | | `/oauth/siwc/bind` | POST | Start SIWC binding to an existing file-backed Codex OAuth credential. | | `/oauth/siwc/status` | GET | Read local SIWC credential status (`off`, `valid`, or `invalid`). | | `/oauth/status?state=` | GET | Get login status. | @@ -128,9 +129,16 @@ GET /v8/management/oauth/auth-url?provider=codex Authorization: Bearer ``` -Keep any provider-specific login parameters in the query string. The login -response includes the authorization URL and a session `state`. Use that state -for status queries and cancellation. +Keep any provider-specific login parameters in the query string. To target recovery +at a specific existing credential, pass `name` (or `auth_index`) via query parameter +or POST JSON body; authenticating a different account will be rejected rather than +overwriting or clearing the target credential's recovery state. By default, post-login +inference verification is not executed automatically to avoid unexpected token +consumption; operators may explicitly opt in by passing `verify: true` (and optional +`verify_model`) when starting the targeted login, or trigger verification explicitly +after adoption via `POST /v8/management/credentials/verify` (`{"name":"","model":""}`). +The login response includes the authorization URL and a session `state` (plus `target_auth_id` +when targeted). Use that state for status queries and cancellation. Callbacks accept `provider`, `state`, `code`, and `error` as GET query parameters or POST JSON fields. POST also accepts `redirect_url` containing the complete @@ -220,6 +228,18 @@ For Codex entries, `GET /credentials` includes these non-secret summary fields: | `siwc_status` | string | Local SIWC status: `off`, `valid`, or `invalid`. | | `siwc_bound` | boolean | SIWC access token, refresh token, and issued client ID are present; not a live validity check. | | `siwc_expires_at` | optional RFC3339 string | Stored SIWC expiration, when available. | +| `access_token_expires_at` | optional RFC3339 string | Active access token expiration deadline, when available. | +| `next_refresh_after` | optional RFC3339 string | Next scheduled automatic refresh check, when in the future. | +| `requires_sign_in` | boolean | Automatic renewal failed with a terminal grant rejection and independent sign-in is required. | +| `renewal_state` | string | Renewal lifecycle state: `sign_in_required`, `expired_sign_in_required`, `retry_scheduled`, `awaiting_first_renewal`, or `renewal_verified`. | +| `recovery_stage` | string | Post-login runtime adoption and verification stage: `sign_in_required`, `quota_blocked`, `loaded`, `inference_verified`, or `renewal_verified`. | +| `adopted` | boolean | Whether the current credential version is loaded in the running auth manager. | +| `inference_verified` | boolean | Whether the current credential version has completed an inference response. | +| `renewal_verified` | boolean | Whether the current credential version has completed an automatic refresh cycle. | +| `refresh_error` | optional object | Sanitized last refresh failure (`code`, `message`, `http_status`) with no tokens, codes, or URLs. | +| `next_action` | optional string | Required operator recovery action (`sign_in` when `requires_sign_in` is true). | +| `recovery_hint` | optional string | Operator guidance explaining instance-scoped OAuth recovery when `requires_sign_in` is true. | +| `renewal_event` | optional object | Deduplicated structured event (`renewal_sign_in_required`) for terminal renewal failures. | The v0 callback and credential-list contracts are unchanged. SIWC activation, unbinding, and standalone SIWC credential creation are not separate management diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index 4ebda82bcee..f59aaf0d15a 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -558,6 +558,9 @@ func reconcileAuthFileCooldownState(auth *coreauth.Auth, now time.Time) (unavail if !nextRetry.IsZero() && !nextRetry.After(now) { nextRetry = time.Time{} } + if statusMessage == "" && auth.RequiresSignIn { + statusMessage = "unauthorized (refresh token invalid)" + } return true, coreauth.StatusError, statusMessage, nextRetry } @@ -786,6 +789,62 @@ func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth, quota authFileQu if requestRetry, ok := auth.RequestRetryOverride(); ok { entry["request_retry"] = requestRetry } + if quota.v8 { + 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) + } + if !auth.NextRefreshAfter.IsZero() && auth.NextRefreshAfter.After(nowUTC) { + entry["next_refresh_after"] = auth.NextRefreshAfter.UTC().Format(time.RFC3339) + } + entry["requires_sign_in"] = auth.RequiresSignIn + if statusMessage != "" { + entry["status_message"] = coreauth.SanitizeRenewalMessage(statusMessage) + } + entry["renewal_state"] = auth.RenewalStateAt(nowUTC) + recoveryStage := auth.RecoveryStageAt(nowUTC) + entry["recovery_stage"] = recoveryStage + entry["quota_blocked"] = recoveryStage == coreauth.RecoveryStageQuotaBlocked + entry["adopted"] = !auth.AdoptedAt.IsZero() || auth.CredentialVersion > 0 + if !auth.AdoptedAt.IsZero() { + entry["adopted_at"] = auth.AdoptedAt.UTC().Format(time.RFC3339) + } + entry["inference_verified"] = auth.InferenceVerified + if !auth.LastInferenceVerifiedAt.IsZero() { + entry["last_inference_verified_at"] = auth.LastInferenceVerifiedAt.UTC().Format(time.RFC3339) + } + entry["renewal_verified"] = auth.RenewalVerified + if !auth.LastRenewalVerifiedAt.IsZero() { + entry["last_renewal_verified_at"] = auth.LastRenewalVerifiedAt.UTC().Format(time.RFC3339) + } + if auth.RefreshError != nil { + entry["refresh_error"] = gin.H{ + "code": auth.RefreshError.Code, + "message": auth.RefreshError.Message, + "http_status": auth.RefreshError.HTTPStatus, + } + } + if auth.RequiresSignIn { + entry["next_action"] = "sign_in" + entry["recovery_hint"] = "Re-authenticate this account on this proxy instance via OAuth login; signing in on another machine or a client's separate direct login does not update this instance." + event := gin.H{ + "type": "renewal_sign_in_required", + "auth_id": auth.ID, + "provider": strings.TrimSpace(auth.Provider), + "next_action": "sign_in", + } + if !auth.RenewalEventAt.IsZero() { + event["occurred_at"] = auth.RenewalEventAt.UTC().Format(time.RFC3339) + } + if email := authEmail(auth); email != "" { + event["account"] = email + } + if exp, okExp := auth.AccessTokenExpirationTime(); okExp && !exp.IsZero() && exp.Unix() > 0 { + event["access_token_expires_at"] = exp.UTC().Format(time.RFC3339) + } + entry["renewal_event"] = event + } + } if quota.v8 && strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { entry["siwc_bind_supported"] = siwcBindEligible(auth) entry["use_siwc"] = coreauth.IsCodexUseSIWC(auth) diff --git a/internal/api/handlers/management/auth_files_cooldown_test.go b/internal/api/handlers/management/auth_files_cooldown_test.go index 05d79c29eff..6d291bacac5 100644 --- a/internal/api/handlers/management/auth_files_cooldown_test.go +++ b/internal/api/handlers/management/auth_files_cooldown_test.go @@ -1,19 +1,24 @@ package management import ( + "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "net/url" "os" "path/filepath" "reflect" + "strings" "testing" "time" "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" ) type authFilesCooldownResponse struct { @@ -761,3 +766,540 @@ func TestListAuthFiles_SparseModelState_BlockedModelDoesNotMakeAuthUnavailable(t t.Errorf("expected status=%q, got %q", coreauth.StatusActive, file.Status) } } + +func TestListAuthFiles_V8CodexTerminalRenewalRecoveryWorkflow_Issue5783(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + now := time.Now().UTC() + futureExpiry := now.Add(24 * time.Hour).Truncate(time.Second) + + manager := coreauth.NewManager(nil, nil, nil) + authDir := t.TempDir() + cfg := &config.Config{AuthDir: authDir} + manager.SetConfig(cfg) + + exec := &failingRefreshManagementExecutor{ + provider: "codex", + err: fmt.Errorf( + `token refresh failed with status 401: {"error":{"message":"Your refresh token has already been used to generate a new access token. Please try signing in again. https://auth.openai.com/authorize?code=secret-code&state=secret-state refresh_token=rt-secret-123 Authorization: Bearer opaque-bearer-secret","code":"refresh_token_reused"}}`, + ), + } + manager.RegisterExecutor(exec) + + credPath := filepath.Join(authDir, "codex-a.json") + if errWrite := os.WriteFile(credPath, []byte(`{"type":"codex","email":"a@example.com"}`), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + + registerAuthForLookupTest(t, manager, &coreauth.Auth{ + ID: "codex-a.json", + Index: "idx-codex-a", + FileName: "codex-a.json", + Provider: "codex", + Status: coreauth.StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "a@example.com", + "account_id": "acc-a", + "access_token": "valid-access-token-secret", + "refresh_token": "rt-secret-123", + "expired": futureExpiry.Format(time.RFC3339), + }, + Attributes: map[string]string{"path": credPath, "source": credPath}, + }) + + // Trigger terminal refresh failure via ForceRefreshAuth while access_token is still valid. + _, _ = manager.ForceRefreshAuth(context.Background(), "codex-a.json") + + // Ordinary successful completed inference must NOT mask the terminal renewal failure. + manager.MarkResult(context.Background(), coreauth.Result{ + AuthID: "codex-a.json", + Provider: "codex", + Model: "gpt-5.5", + Success: true, + CompletedResponse: true, + }) + + h := NewHandlerWithoutConfigFilePath(cfg, manager) + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v8/management/credentials?name=codex-a.json", nil) + ginCtx.Set(ConfigV8ContextKey, true) + h.ListAuthFiles(ginCtx) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200: %s", rec.Code, rec.Body.String()) + } + + var resp struct { + Files []map[string]any `json:"files"` + } + if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &resp); errUnmarshal != nil { + t.Fatalf("unmarshal response: %v", errUnmarshal) + } + if len(resp.Files) != 1 { + t.Fatalf("len(files) = %d, want 1", len(resp.Files)) + } + entry := resp.Files[0] + if entry["requires_sign_in"] != true { + t.Fatalf("requires_sign_in = %v, want true (entry=%#v)", entry["requires_sign_in"], entry) + } + if entry["renewal_state"] != coreauth.RenewalStateSignInRequired { + t.Fatalf("renewal_state = %v, want %s", entry["renewal_state"], coreauth.RenewalStateSignInRequired) + } + if entry["recovery_stage"] != coreauth.RecoveryStageSignInRequired { + t.Fatalf("recovery_stage = %v, want %s", entry["recovery_stage"], coreauth.RecoveryStageSignInRequired) + } + if entry["access_token_expires_at"] != futureExpiry.Format(time.RFC3339) { + t.Fatalf("access_token_expires_at = %v, want %s", entry["access_token_expires_at"], futureExpiry.Format(time.RFC3339)) + } + if entry["next_action"] != "sign_in" { + t.Fatalf("next_action = %v, want sign_in", entry["next_action"]) + } + if entry["renewal_event"] == nil { + t.Fatalf("expected deduplicated renewal_event in entry: %#v", entry) + } + rawBody := rec.Body.String() + for _, forbidden := range []string{"valid-access-token-secret", "rt-secret-123", "secret-code", "secret-state", "opaque-bearer-secret", "https://auth.openai.com"} { + if strings.Contains(rawBody, forbidden) { + t.Fatalf("response leaked sensitive value %q in %s", forbidden, rawBody) + } + } + + // Targeted OAuth recovery via StartOAuthV8: cancelling or wrong provider does not mark codex-a.json repaired. + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ConfigV8ContextKey, true) + c.Next() + }) + router.POST("/v8/management/oauth/auth-url", h.StartOAuthV8) + router.GET("/v8/management/oauth/status", h.GetAuthStatus) + router.DELETE("/v8/management/oauth/session", h.CancelAuthSession) + router.POST("/v8/management/credentials/verify", h.VerifyAuthFileV8) + + wMismatch := httptest.NewRecorder() + reqMismatch := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"claude","name":"codex-a.json"}`)) + reqMismatch.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wMismatch, reqMismatch) + if wMismatch.Code != http.StatusBadRequest { + t.Fatalf("expected 400 for provider mismatch, got %d: %s", wMismatch.Code, wMismatch.Body.String()) + } + + wStart := httptest.NewRecorder() + reqStart := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json"}`)) + reqStart.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wStart, reqStart) + if wStart.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 status = %d, want 200: %s", wStart.Code, wStart.Body.String()) + } + var startResp map[string]any + if errUnmarshal := json.Unmarshal(wStart.Body.Bytes(), &startResp); errUnmarshal != nil { + t.Fatalf("unmarshal StartOAuthV8: %v", errUnmarshal) + } + if startResp["target_auth_id"] != "codex-a.json" { + t.Fatalf("target_auth_id = %v, want codex-a.json", startResp["target_auth_id"]) + } + state, _ := startResp["state"].(string) + + // Cancel login session and confirm codex-a.json still requires sign-in. + wCancel := httptest.NewRecorder() + router.ServeHTTP(wCancel, httptest.NewRequest(http.MethodDelete, "/v8/management/oauth/session?state="+state, nil)) + if wCancel.Code != http.StatusOK { + t.Fatalf("CancelAuthSession status = %d, want 200", wCancel.Code) + } + stillBroken, _ := manager.GetByID("codex-a.json") + if !stillBroken.RequiresSignIn { + t.Fatal("cancelling OAuth session must not clear RequiresSignIn") + } + + // Targeted recovery with a DIFFERENT account (other@example.com) must fail and not mark codex-a.json repaired. + mockSvc := &targetedRecoveryTestCodexOAuthService{ + authDir: authDir, + email: "other@example.com", + accountID: "acc-other", + } + origFactory := newCodexOAuthService + newCodexOAuthService = func(*config.Config) codexOAuthService { return mockSvc } + t.Cleanup(func() { newCodexOAuthService = origFactory }) + + wDiffStart := httptest.NewRecorder() + reqDiffStart := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json"}`)) + reqDiffStart.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wDiffStart, reqDiffStart) + if wDiffStart.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 (diff account) status = %d: %s", wDiffStart.Code, wDiffStart.Body.String()) + } + var diffStartResp struct { + State string `json:"state"` + } + if errUnmarshal := json.Unmarshal(wDiffStart.Body.Bytes(), &diffStartResp); errUnmarshal != nil { + t.Fatalf("unmarshal diffStartResp: %v", errUnmarshal) + } + diffStatus := pollOAuthStatusUntilDone(t, router, diffStartResp.State) + if diffStatus["status"] != "error" || !strings.Contains(fmt.Sprint(diffStatus["error"]), "does not match target credential") { + t.Fatalf("expected account mismatch error, got %#v", diffStatus) + } + afterMismatch, _ := manager.GetByID("codex-a.json") + if !afterMismatch.RequiresSignIn { + t.Fatal("authenticating a different account must not mark target credential repaired") + } + + // Targeted recovery with missing account_id or missing email must also be rejected (fail-closed). + for _, missingCase := range []struct { + name string + email string + accountID string + }{ + {name: "missing_account_id", email: "a@example.com", accountID: ""}, + {name: "missing_email", email: "", accountID: "acc-a"}, + } { + mockSvc.email = missingCase.email + mockSvc.accountID = missingCase.accountID + wMissing := httptest.NewRecorder() + reqMissing := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json"}`)) + reqMissing.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wMissing, reqMissing) + if wMissing.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 (%s) status = %d: %s", missingCase.name, wMissing.Code, wMissing.Body.String()) + } + var missingStartResp struct { + State string `json:"state"` + } + if errUnmarshal := json.Unmarshal(wMissing.Body.Bytes(), &missingStartResp); errUnmarshal != nil { + t.Fatalf("unmarshal missingStartResp (%s): %v", missingCase.name, errUnmarshal) + } + missingStatus := pollOAuthStatusUntilDone(t, router, missingStartResp.State) + if missingStatus["status"] != "error" || !strings.Contains(fmt.Sprint(missingStatus["error"]), "does not match target credential") { + t.Fatalf("%s: expected mismatch error, got %#v", missingCase.name, missingStatus) + } + afterMissing, _ := manager.GetByID("codex-a.json") + if !afterMissing.RequiresSignIn { + t.Fatalf("%s: missing identity field must not mark target credential repaired", missingCase.name) + } + } + + // If saving to storage succeeds but runtime adoption fails, GetAuthStatus must report saved=true, adopted=false, recovery_stage=credential_saved. + mockSvc.email = "a@example.com" + mockSvc.accountID = "acc-a" + h.SetPostAuthPersistHook(func(context.Context, *coreauth.Auth) error { + return fmt.Errorf("simulated runtime adoption failure") + }) + wUnadoptedStart := httptest.NewRecorder() + reqUnadoptedStart := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json"}`)) + reqUnadoptedStart.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wUnadoptedStart, reqUnadoptedStart) + if wUnadoptedStart.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 (unadopted) status = %d: %s", wUnadoptedStart.Code, wUnadoptedStart.Body.String()) + } + var unadoptedStartResp struct { + State string `json:"state"` + } + if errUnmarshal := json.Unmarshal(wUnadoptedStart.Body.Bytes(), &unadoptedStartResp); errUnmarshal != nil { + t.Fatalf("unmarshal unadoptedStartResp: %v", errUnmarshal) + } + unadoptedStatus := pollOAuthStatusUntilDone(t, router, unadoptedStartResp.State) + if unadoptedStatus["status"] != "ok" || unadoptedStatus["saved"] != true || unadoptedStatus["adopted"] != false || unadoptedStatus["recovery_stage"] != coreauth.RecoveryStageCredentialSaved { + t.Fatalf("expected saved=true adopted=false recovery_stage=credential_saved when runtime adoption fails, got %#v", unadoptedStatus) + } + // An unrelated runtime update or replacement with a different grant must NOT falsely satisfy adoption for unadoptedStartResp.State. + unrelatedReplacement := stillBroken.Clone() + unrelatedReplacement.Metadata["access_token"] = "unrelated-access-token" + unrelatedReplacement.Metadata["refresh_token"] = "unrelated-refresh-token" + if _, errUnrelated := manager.Update(coreauth.WithSkipPersist(context.Background()), unrelatedReplacement); errUnrelated != nil { + t.Fatalf("Update unrelatedReplacement: %v", errUnrelated) + } + unrelatedStatus := pollOAuthStatusUntilDone(t, router, unadoptedStartResp.State) + if unrelatedStatus["adopted"] != false || unrelatedStatus["recovery_stage"] != coreauth.RecoveryStageCredentialSaved { + t.Fatalf("unrelated credential update must not mark OAuth session adopted, got %#v", unrelatedStatus) + } + // When a delayed watcher/reload subsequently adopts the actual saved file into manager AND an automatic refresh rotates + // the tokens BEFORE GetAuthStatus is polled again, polling GetAuthStatus must still report adopted=true and renewal_verified=true. + if errDelayedReg := h.registerAuthFromFile(context.Background(), credPath, nil); errDelayedReg != nil { + t.Fatalf("delayed registerAuthFromFile error: %v", errDelayedReg) + } + exec.err = nil + exec.rotateTokens = true + if _, errDelayedRot := manager.ForceRefreshAuth(coreauth.WithSkipPersist(context.Background()), "codex-a.json"); errDelayedRot != nil { + t.Fatalf("ForceRefreshAuth after delayed adoption error: %v", errDelayedRot) + } + exec.rotateTokens = false + delayedStatus := pollOAuthStatusUntilDone(t, router, unadoptedStartResp.State) + if delayedStatus["status"] != "ok" || delayedStatus["saved"] != true || delayedStatus["adopted"] != true || delayedStatus["renewal_verified"] != true || delayedStatus["recovery_stage"] != coreauth.RecoveryStageRenewalVerified { + t.Fatalf("expected delayed adoption + rotation to transition GetAuthStatus to adopted=true renewal_verified=true recovery_stage=renewal_verified, got %#v", delayedStatus) + } + h.SetPostAuthPersistHook(nil) + + // Targeted recovery with the MATCHING account (a@example.com / acc-a) saves fresh tokens, adopts into manager, and clears RequiresSignIn. + wMatchStart := httptest.NewRecorder() + reqMatchStart := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json"}`)) + reqMatchStart.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wMatchStart, reqMatchStart) + if wMatchStart.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 (matching account) status = %d: %s", wMatchStart.Code, wMatchStart.Body.String()) + } + var matchStartResp struct { + State string `json:"state"` + } + if errUnmarshal := json.Unmarshal(wMatchStart.Body.Bytes(), &matchStartResp); errUnmarshal != nil { + t.Fatalf("unmarshal matchStartResp: %v", errUnmarshal) + } + matchStatus := pollOAuthStatusUntilDone(t, router, matchStartResp.State) + if matchStatus["status"] != "ok" || matchStatus["auth_id"] != "codex-a.json" || matchStatus["saved"] != true || matchStatus["adopted"] != true { + t.Fatalf("expected saved & adopted ok status for codex-a.json, got %#v", matchStatus) + } + if matchStatus["requires_sign_in"] != false || matchStatus["renewal_state"] != coreauth.RenewalStateAwaitingFirstRenewal || matchStatus["recovery_stage"] != coreauth.RecoveryStageLoaded || matchStatus["inference_verified"] != false || matchStatus["renewal_verified"] != false { + t.Fatalf("unexpected post-recovery status fields: %#v", matchStatus) + } + repaired, _ := manager.GetByID("codex-a.json") + if repaired.RequiresSignIn || repaired.RefreshError != nil { + t.Fatalf("expected codex-a.json repaired in manager, got RequiresSignIn=%v RefreshError=%+v", repaired.RequiresSignIn, repaired.RefreshError) + } + + // Operator-triggered POST /v8/management/credentials/verify: + // 1) response.incomplete must fail (HTTP 502) and NOT mark inference_verified=true. + exec.streamPayloads = [][]byte{[]byte(`data: {"type":"response.incomplete","response":{"status":"incomplete"}}` + "\n\n")} + wVerifyIncomp := httptest.NewRecorder() + reqVerifyIncomp := httptest.NewRequest(http.MethodPost, "/v8/management/credentials/verify", strings.NewReader(`{"name":"codex-a.json","model":"gpt-5.5"}`)) + reqVerifyIncomp.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wVerifyIncomp, reqVerifyIncomp) + if wVerifyIncomp.Code != http.StatusBadGateway { + t.Fatalf("VerifyAuthFileV8 (incomplete) status = %d, want 502: %s", wVerifyIncomp.Code, wVerifyIncomp.Body.String()) + } + + // 1b) 503 transient failure must NOT be misclassified as quota_blocked. + exec.streamPayloads = nil + exec.streamErr = &coreauth.Error{Code: "service_unavailable", Message: "upstream 503", HTTPStatus: http.StatusServiceUnavailable, Retryable: true} + wVerify503 := httptest.NewRecorder() + reqVerify503 := httptest.NewRequest(http.MethodPost, "/v8/management/credentials/verify", strings.NewReader(`{"name":"codex-a.json","model":"gpt-5.5"}`)) + reqVerify503.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wVerify503, reqVerify503) + if wVerify503.Code != http.StatusServiceUnavailable { + t.Fatalf("VerifyAuthFileV8 (503) status = %d, want 503: %s", wVerify503.Code, wVerify503.Body.String()) + } + var verify503Resp map[string]any + if errUnmarshal := json.Unmarshal(wVerify503.Body.Bytes(), &verify503Resp); errUnmarshal != nil { + t.Fatalf("unmarshal verify503Resp: %v", errUnmarshal) + } + if verify503Resp["quota_blocked"] != false || verify503Resp["recovery_stage"] == coreauth.RecoveryStageQuotaBlocked { + t.Fatalf("503 transient failure must not be classified as quota_blocked, got %#v", verify503Resp) + } + _, _, _ = manager.ResetQuota(context.Background(), "codex-a.json") + + // 2) 429 quota error after valid auth must transition recovery_stage to quota_blocked and sanitize any secrets in the error message. + exec.streamPayloads = nil + exec.streamErr = &coreauth.Error{ + Code: "rate_limit_exceeded", + Message: "You have hit your usage limit https://chatgpt.com/backend-api/codex?code=secret-verify-code&state=secret-verify-state Bearer secret-verify-bearer", + HTTPStatus: http.StatusTooManyRequests, + } + wVerifyQuota := httptest.NewRecorder() + reqVerifyQuota := httptest.NewRequest(http.MethodPost, "/v8/management/credentials/verify", strings.NewReader(`{"name":"codex-a.json","model":"gpt-5.5"}`)) + reqVerifyQuota.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wVerifyQuota, reqVerifyQuota) + if wVerifyQuota.Code != http.StatusTooManyRequests { + t.Fatalf("VerifyAuthFileV8 (quota) status = %d, want 429: %s", wVerifyQuota.Code, wVerifyQuota.Body.String()) + } + for _, secret := range []string{"secret-verify-code", "secret-verify-state", "secret-verify-bearer", "https://chatgpt.com"} { + if strings.Contains(wVerifyQuota.Body.String(), secret) { + t.Fatalf("VerifyAuthFileV8 response leaked secret %q: %s", secret, wVerifyQuota.Body.String()) + } + } + var quotaVerifyResp map[string]any + if errUnmarshal := json.Unmarshal(wVerifyQuota.Body.Bytes(), "aVerifyResp); errUnmarshal != nil { + t.Fatalf("unmarshal quotaVerifyResp: %v", errUnmarshal) + } + if quotaVerifyResp["recovery_stage"] != coreauth.RecoveryStageQuotaBlocked || quotaVerifyResp["quota_blocked"] != true { + t.Fatalf("expected recovery_stage=quota_blocked and quota_blocked=true, got %#v", quotaVerifyResp) + } + _, _, _ = manager.ResetQuota(context.Background(), "codex-a.json") + + // 3) Completed stream transitions recovery_stage to inference_verified. + exec.streamErr = nil + exec.streamPayloads = [][]byte{[]byte(`data: {"type":"response.completed"}` + "\n\n")} + wVerifyOK := httptest.NewRecorder() + reqVerifyOK := httptest.NewRequest(http.MethodPost, "/v8/management/credentials/verify", strings.NewReader(`{"name":"codex-a.json","model":"gpt-5.5"}`)) + reqVerifyOK.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wVerifyOK, reqVerifyOK) + if wVerifyOK.Code != http.StatusOK { + t.Fatalf("VerifyAuthFileV8 (ok) status = %d, want 200: %s", wVerifyOK.Code, wVerifyOK.Body.String()) + } + var okVerifyResp map[string]any + if errUnmarshal := json.Unmarshal(wVerifyOK.Body.Bytes(), &okVerifyResp); errUnmarshal != nil { + t.Fatalf("unmarshal okVerifyResp: %v", errUnmarshal) + } + if okVerifyResp["verified"] != true || okVerifyResp["inference_verified"] != true || okVerifyResp["recovery_stage"] != coreauth.RecoveryStageInferenceVerified { + t.Fatalf("expected verified=true inference_verified=true recovery_stage=inference_verified, got %#v", okVerifyResp) + } + + // 3b) After a subsequent automatic refresh rotates tokens on the adopted credential, GetAuthStatus for matchStartResp.State + // must still report adopted=true and transition to renewal_verified=true / recovery_stage=renewal_verified. + exec.err = nil + exec.rotateTokens = true + if _, errForceRefresh := manager.ForceRefreshAuth(coreauth.WithSkipPersist(context.Background()), "codex-a.json"); errForceRefresh != nil { + t.Fatalf("ForceRefreshAuth after recovery error: %v", errForceRefresh) + } + rotatedStatus := pollOAuthStatusUntilDone(t, router, matchStartResp.State) + if rotatedStatus["adopted"] != true || rotatedStatus["renewal_verified"] != true || rotatedStatus["recovery_stage"] != coreauth.RecoveryStageRenewalVerified { + t.Fatalf("expected adopted=true renewal_verified=true recovery_stage=renewal_verified after token rotation, got %#v", rotatedStatus) + } + + // 3c) If the adopted credential is subsequently replaced via Update by an unrelated credential + // and that replacement credential is refreshed, the original OAuth session must NOT claim the replacement's renewal. + currentAdopted, _ := manager.GetByID("codex-a.json") + postRotationUnrelated := currentAdopted.Clone() + postRotationUnrelated.Metadata["access_token"] = "unrelated-access-token" + postRotationUnrelated.Metadata["refresh_token"] = "unrelated-refresh-token" + if _, errUpdUnrelated := manager.Update(coreauth.WithSkipPersist(context.Background()), postRotationUnrelated); errUpdUnrelated != nil { + t.Fatalf("Update postRotationUnrelated error: %v", errUpdUnrelated) + } + if _, errForceUnrelated := manager.ForceRefreshAuth(coreauth.WithSkipPersist(context.Background()), "codex-a.json"); errForceUnrelated != nil { + t.Fatalf("ForceRefreshAuth unrelatedReplacement error: %v", errForceUnrelated) + } + replacedAfterRotationStatus := pollOAuthStatusUntilDone(t, router, matchStartResp.State) + if replacedAfterRotationStatus["adopted"] != false || strings.TrimSpace(fmt.Sprint(replacedAfterRotationStatus["adoption_error"])) == "" { + t.Fatalf("expected adopted=false and adoption_error after unrelated credential replacement + refresh, got %#v", replacedAfterRotationStatus) + } + + // 4) Opt-in post-login verification (verify=true on StartOAuthV8) sanitizes verification_error in GetAuthStatus. + exec.streamPayloads = nil + exec.streamErr = &coreauth.Error{ + Code: "rate_limit_exceeded", + Message: "Quota error https://chatgpt.com/backend-api/codex?code=secret-postlogin-code&state=secret-postlogin-state Bearer secret-postlogin-bearer", + HTTPStatus: http.StatusTooManyRequests, + } + wVerifyOptIn := httptest.NewRecorder() + reqVerifyOptIn := httptest.NewRequest(http.MethodPost, "/v8/management/oauth/auth-url", strings.NewReader(`{"provider":"codex","name":"codex-a.json","verify":true,"verify_model":"gpt-5.5"}`)) + reqVerifyOptIn.Header.Set("Content-Type", "application/json") + router.ServeHTTP(wVerifyOptIn, reqVerifyOptIn) + if wVerifyOptIn.Code != http.StatusOK { + t.Fatalf("StartOAuthV8 (verify=true) status = %d: %s", wVerifyOptIn.Code, wVerifyOptIn.Body.String()) + } + var verifyOptInStartResp struct { + State string `json:"state"` + } + if errUnmarshal := json.Unmarshal(wVerifyOptIn.Body.Bytes(), &verifyOptInStartResp); errUnmarshal != nil { + t.Fatalf("unmarshal verifyOptInStartResp: %v", errUnmarshal) + } + optInStatus := pollOAuthStatusUntilDone(t, router, verifyOptInStartResp.State) + if optInStatus["status"] != "ok" || optInStatus["adopted"] != true || optInStatus["quota_blocked"] != true { + t.Fatalf("expected adopted=true quota_blocked=true after opt-in verify quota error, got %#v", optInStatus) + } + verifyErrText := fmt.Sprint(optInStatus["verification_error"]) + if verifyErrText == "" { + t.Fatalf("expected sanitized verification_error in status response, got %#v", optInStatus) + } + for _, secret := range []string{"secret-postlogin-code", "secret-postlogin-state", "secret-postlogin-bearer", "https://chatgpt.com"} { + if strings.Contains(verifyErrText, secret) { + t.Fatalf("GetAuthStatus verification_error leaked secret %q: %s", secret, verifyErrText) + } + } +} + +type targetedRecoveryTestCodexOAuthService struct { + authDir string + email string + accountID string +} + +func (s *targetedRecoveryTestCodexOAuthService) GenerateAuthURL(state string, _ *codex.PKCECodes) (string, error) { + waitFile := filepath.Join(s.authDir, fmt.Sprintf(".oauth-codex-%s.oauth", state)) + payload, _ := json.Marshal(map[string]string{"state": state, "code": "code-" + state}) + _ = os.WriteFile(waitFile, payload, 0o600) + return "https://auth.example.test/authorize?state=" + state, nil +} + +func (s *targetedRecoveryTestCodexOAuthService) ExchangeCodeForTokens(_ context.Context, code string, _ *codex.PKCECodes) (*codex.CodexAuthBundle, error) { + now := time.Now().UTC() + return &codex.CodexAuthBundle{ + TokenData: codex.CodexTokenData{ + IDToken: makeOAuthTestJWT("plus"), + AccessToken: "fresh-access-" + code, + RefreshToken: "fresh-refresh-" + code, + AccountID: s.accountID, + Email: s.email, + PlanType: "plus", + Expire: now.Add(48 * time.Hour).Format(time.RFC3339), + }, + LastRefresh: now.Format(time.RFC3339), + }, nil +} + +func (s *targetedRecoveryTestCodexOAuthService) CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage { + return &codex.CodexTokenStorage{ + IDToken: bundle.TokenData.IDToken, + AccessToken: bundle.TokenData.AccessToken, + RefreshToken: bundle.TokenData.RefreshToken, + AccountID: bundle.TokenData.AccountID, + LastRefresh: bundle.LastRefresh, + Email: bundle.TokenData.Email, + Expire: bundle.TokenData.Expire, + PlanType: bundle.TokenData.PlanType, + } +} + +func pollOAuthStatusUntilDone(t *testing.T, router http.Handler, state string) map[string]any { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v8/management/oauth/status?state="+state, nil)) + if w.Code != http.StatusOK { + t.Fatalf("oauth status HTTP %d: %s", w.Code, w.Body.String()) + } + var out map[string]any + if errUnmarshal := json.Unmarshal(w.Body.Bytes(), &out); errUnmarshal != nil { + t.Fatalf("unmarshal oauth status: %v", errUnmarshal) + } + if out["status"] != "wait" { + return out + } + } + t.Fatalf("timed out waiting for oauth session %s", state) + return nil +} + +type failingRefreshManagementExecutor struct { + provider string + err error + rotateTokens bool + streamErr error + streamPayloads [][]byte +} + +func (e *failingRefreshManagementExecutor) Identifier() string { return e.provider } + +func (e *failingRefreshManagementExecutor) Execute(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *failingRefreshManagementExecutor) ExecuteStream(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.streamErr != nil { + return nil, e.streamErr + } + ch := make(chan cliproxyexecutor.StreamChunk, len(e.streamPayloads)) + for _, p := range e.streamPayloads { + ch <- cliproxyexecutor.StreamChunk{Payload: p} + } + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *failingRefreshManagementExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + if e.err != nil { + return nil, e.err + } + if e.rotateTokens && auth != nil { + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["access_token"] = "rotated-access-token" + auth.Metadata["refresh_token"] = "rotated-refresh-token" + } + return auth, nil +} + +func (e *failingRefreshManagementExecutor) CountTokens(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *failingRefreshManagementExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, nil +} diff --git a/internal/api/handlers/management/auth_files_crud.go b/internal/api/handlers/management/auth_files_crud.go index 30fee8dec6c..f1d29069884 100644 --- a/internal/api/handlers/management/auth_files_crud.go +++ b/internal/api/handlers/management/auth_files_crud.go @@ -587,10 +587,12 @@ func (h *Handler) buildAuthFromFileData(path string, data []byte) (*coreauth.Aut if h != nil && h.authManager != nil { if existing, ok := h.authManager.GetByID(authID); ok { auth.CreatedAt = existing.CreatedAt - if !hasLastRefresh { - auth.LastRefreshedAt = existing.LastRefreshedAt + if !coreauth.CredentialsChanged(existing, auth) { + if !hasLastRefresh { + auth.LastRefreshedAt = existing.LastRefreshedAt + } + auth.NextRefreshAfter = existing.NextRefreshAfter } - auth.NextRefreshAfter = existing.NextRefreshAfter auth.Runtime = existing.Runtime } } diff --git a/internal/api/handlers/management/auth_files_fields.go b/internal/api/handlers/management/auth_files_fields.go index 39020d47b6b..6352f5ca607 100644 --- a/internal/api/handlers/management/auth_files_fields.go +++ b/internal/api/handlers/management/auth_files_fields.go @@ -1049,6 +1049,10 @@ func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (s if errHook := h.postAuthPersistHook(ctx, persistedRecord); errHook != nil { return savedPath, fmt.Errorf("post-auth persist hook failed: %w", errHook) } + } else if h.authManager != nil && strings.TrimSpace(savedPath) != "" { + if errReg := h.registerAuthFromFile(ctx, savedPath, nil); errReg != nil { + return savedPath, fmt.Errorf("register persisted auth failed: %w", errReg) + } } return savedPath, nil } diff --git a/internal/api/handlers/management/auth_files_provider_oauth.go b/internal/api/handlers/management/auth_files_provider_oauth.go index dffe87337b6..52a70422745 100644 --- a/internal/api/handlers/management/auth_files_provider_oauth.go +++ b/internal/api/handlers/management/auth_files_provider_oauth.go @@ -241,6 +241,10 @@ func (h *Handler) RequestCodexToken(c *gin.Context) { } RegisterOAuthSessionWithProxy(state, "codex", sessionProxyURL) + sessionTarget, hasSessionTarget := oauthSessionTargetFromContext(c) + if hasSessionTarget { + SetOAuthSessionTarget(state, sessionTarget) + } isWebUI := isWebUIRequest(c) var forwarder *callbackForwarder @@ -334,6 +338,16 @@ func (h *Handler) RequestCodexToken(c *gin.Context) { tokenStorage := openaiAuth.CreateTokenStorage(bundle) tokenStorage.PlanType = planType fileName := codex.CredentialFileName(tokenStorage.Email, planType, hashAccountID, true) + if target, okTarget := GetOAuthSessionTarget(state); okTarget && target.AuthID != "" { + if !matchesOAuthRecoveryTarget(target, tokenStorage.Email, tokenStorage.AccountID, fileName) { + SetOAuthSessionError(state, "Authenticated account does not match target credential") + log.Warnf("Codex OAuth recovery rejected for %s: authenticated account does not match target credential", target.AuthID) + return + } + if target.FileName != "" { + fileName = target.FileName + } + } record := &coreauth.Auth{ ID: fileName, Provider: "codex", @@ -348,25 +362,125 @@ func (h *Handler) RequestCodexToken(c *gin.Context) { "plan_type": planType, }, } + target, hasTarget := GetOAuthSessionTarget(state) + if hasTarget && target.AuthID != "" { + record.ID = target.AuthID + } + var hadPrevAuth bool + var prevVersion uint64 + var prevEpoch uint64 + if h.authManager != nil { + if existing, okExisting := h.authManager.GetByID(record.ID); okExisting && existing != nil { + hadPrevAuth = true + prevVersion = existing.CredentialVersion + prevEpoch = existing.RegistrationEpoch + } + } applySessionProxyToAuthRecord(record, sessionProxyURL) if errGuard := claimOAuthSessionCommit(state, "codex"); errGuard != nil { return } savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { + if errSave != nil && strings.TrimSpace(savedPath) == "" { FailOAuthSessionCommit(state, "Failed to save authentication tokens") log.Errorf("Failed to save authentication tokens: %v", errSave) return } + savedGrantHash := oauthGrantFingerprint(tokenStorage.AccessToken, tokenStorage.RefreshToken, tokenStorage.Email, tokenStorage.AccountID) + var adopted bool + var adoptedVersion uint64 + var adoptedEpoch uint64 + var adoptErr string + if errSave != nil { + adoptErr = "Credential saved to storage, but runtime adoption failed" + log.Errorf("Codex credential saved to %s, but runtime adoption failed: %v", savedPath, errSave) + } else if h.authManager != nil { + if runtimeAuth, okRuntime := h.authManager.GetByID(record.ID); okRuntime && runtimeAuth != nil { + if authGrantFingerprint(runtimeAuth) == savedGrantHash { + adopted = true + adoptedVersion = runtimeAuth.CredentialVersion + adoptedEpoch = runtimeAuth.RegistrationEpoch + } else { + adoptErr = "Credential saved to storage, but runtime did not adopt the new credential version" + } + } else { + adoptErr = "Credential saved to storage, but was not registered in active runtime" + } + } + var verifyErr string + if adopted && hasTarget && target.VerifyAfterLogin && h.authManager != nil { + if _, errVerify := h.authManager.VerifyInference(ctx, record.ID, target.VerifyModel); errVerify != nil { + verifyErr = coreauth.SanitizeRenewalMessage(errVerify.Error()) + } + } fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) if bundle.APIKey != "" { fmt.Println("API key obtained and saved") } fmt.Println("You can now use Codex services through this CLI") - CompleteOAuthSession(state) + CompleteOAuthSessionWithOutcome(state, record.ID, savedGrantHash, strings.TrimSpace(savedPath) != "", adopted, hadPrevAuth, prevVersion, prevEpoch, adoptedVersion, adoptedEpoch, adoptErr, verifyErr) }() - c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) + respPayload := gin.H{"status": "ok", "url": authURL, "state": state} + if hasSessionTarget { + respPayload["target_auth_id"] = sessionTarget.AuthID + } + c.JSON(200, respPayload) +} + +func matchesOAuthRecoveryTarget(target oauthRecoveryTarget, email, accountID, candidateFileName string) bool { + targetAccountID := strings.TrimSpace(target.AccountID) + targetEmail := strings.TrimSpace(target.Email) + gotAccountID := strings.TrimSpace(accountID) + gotEmail := strings.TrimSpace(email) + + if targetAccountID != "" { + if gotAccountID == "" || !strings.EqualFold(gotAccountID, targetAccountID) { + return false + } + } + if targetEmail != "" { + if gotEmail == "" || !strings.EqualFold(gotEmail, targetEmail) { + return false + } + } + if targetAccountID == "" && targetEmail == "" { + if gotAccountID == "" && gotEmail == "" { + return false + } + if target.FileName != "" && candidateFileName != target.FileName && candidateFileName != target.AuthID { + return false + } + } + return true +} + +func oauthGrantFingerprint(accessToken, refreshToken, email, accountID string) string { + return coreauth.OAuthGrantFingerprint(accessToken, refreshToken, email, accountID) +} + +func authGrantFingerprint(auth *coreauth.Auth) string { + if auth == nil { + return "" + } + var accessToken, refreshToken, accountID string + if storage, ok := auth.Storage.(*codex.CodexTokenStorage); ok && storage != nil { + accessToken = strings.TrimSpace(storage.AccessToken) + refreshToken = strings.TrimSpace(storage.RefreshToken) + accountID = strings.TrimSpace(storage.AccountID) + } + if auth.Metadata != nil { + if accessToken == "" { + accessToken, _ = auth.Metadata["access_token"].(string) + } + if refreshToken == "" { + refreshToken, _ = auth.Metadata["refresh_token"].(string) + } + if accountID == "" { + accountID, _ = auth.Metadata["account_id"].(string) + } + } + return oauthGrantFingerprint(accessToken, refreshToken, authEmail(auth), accountID) } func (h *Handler) RequestAntigravityToken(c *gin.Context) { @@ -998,7 +1112,60 @@ func (h *Handler) GetAuthStatus(c *gin.Context) { return } if session.Completed { - c.JSON(http.StatusOK, gin.H{"status": "ok"}) + resp := gin.H{"status": "ok"} + if session.AuthID != "" { + resp["auth_id"] = session.AuthID + resp["saved"] = session.Saved + var adoptedAuth *coreauth.Auth + var adoptErr string + if h != nil && h.authManager != nil { + if candidate, okAuth := h.authManager.GetByID(session.AuthID); okAuth && candidate != nil { + if session.SavedGrantHash != "" { + if authGrantFingerprint(candidate) == session.SavedGrantHash || + candidate.AdoptedGrantHash == session.SavedGrantHash { + adoptedAuth = candidate + if !session.Adopted { + oauthSessions.MarkAdopted(state, candidate.CredentialVersion, candidate.RegistrationEpoch) + session.Adopted = true + session.AdoptedVersion = candidate.CredentialVersion + session.AdoptedEpoch = candidate.RegistrationEpoch + session.AdoptError = "" + } + } else if session.Adopted { + adoptErr = "Credential was replaced in active runtime after login" + } + } else if session.Adopted { + if (session.AdoptedVersion == 0 && session.AdoptedEpoch == 0) || + (candidate.RegistrationEpoch == session.AdoptedEpoch && candidate.CredentialVersion == session.AdoptedVersion) { + adoptedAuth = candidate + } + } + } + } + if adoptErr == "" { + adoptErr = session.AdoptError + } + if adoptedAuth != nil { + nowUTC := time.Now().UTC() + resp["adopted"] = true + resp["recovery_stage"] = adoptedAuth.RecoveryStageAt(nowUTC) + resp["renewal_state"] = adoptedAuth.RenewalStateAt(nowUTC) + resp["requires_sign_in"] = adoptedAuth.RequiresSignIn + resp["inference_verified"] = adoptedAuth.InferenceVerified + resp["renewal_verified"] = adoptedAuth.RenewalVerified + resp["quota_blocked"] = adoptedAuth.RecoveryStageAt(nowUTC) == coreauth.RecoveryStageQuotaBlocked + } else if session.Saved { + resp["adopted"] = false + resp["recovery_stage"] = coreauth.RecoveryStageCredentialSaved + if adoptErr != "" { + resp["adoption_error"] = adoptErr + } + } + if session.VerifyError != "" { + resp["verification_error"] = session.VerifyError + } + } + c.JSON(http.StatusOK, resp) return } if session.Status != "" { diff --git a/internal/api/handlers/management/auth_files_v8.go b/internal/api/handlers/management/auth_files_v8.go index 15096d7a223..365523c0655 100644 --- a/internal/api/handlers/management/auth_files_v8.go +++ b/internal/api/handlers/management/auth_files_v8.go @@ -20,8 +20,21 @@ const ( oauthSessionDomainContextKey = "management.oauth.domain" oauthSessionIsWebUIContextKey = "management.oauth.is_webui" oauthSessionMetadataContextKey = "management.oauth.metadata" + oauthSessionTargetContextKey = "management.oauth.target" ) +func oauthSessionTargetFromContext(c *gin.Context) (oauthRecoveryTarget, bool) { + if c == nil { + return oauthRecoveryTarget{}, false + } + val, ok := c.Get(oauthSessionTargetContextKey) + if !ok { + return oauthRecoveryTarget{}, false + } + target, okTarget := val.(oauthRecoveryTarget) + return target, okTarget && target.AuthID != "" +} + func oauthSessionProxyURLFromContext(c *gin.Context) string { if c == nil { return "" @@ -125,6 +138,10 @@ func (h *Handler) StartOAuthV8(c *gin.Context) { } provider := strings.TrimSpace(c.Query("provider")) + targetName := firstNonEmpty(strings.TrimSpace(c.Query("name")), strings.TrimSpace(c.Query("target_auth_id")), strings.TrimSpace(c.Query("auth_id"))) + targetAuthIndex := strings.TrimSpace(c.Query("auth_index")) + verifyAfterLogin := strings.EqualFold(strings.TrimSpace(c.Query("verify")), "true") || strings.TrimSpace(c.Query("verify")) == "1" + verifyModel := strings.TrimSpace(c.Query("verify_model")) if c.Request != nil && c.Request.Body != nil { rawBody, errRead := io.ReadAll(c.Request.Body) if errRead != nil { @@ -157,6 +174,28 @@ func (h *Handler) StartOAuthV8(c *gin.Context) { provider = bodyProvider } } + for _, key := range []string{"name", "target_auth_id", "auth_id"} { + if rawVal, exists := bodyMap[key]; exists && rawVal != nil { + strVal, ok := rawVal.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid " + key}) + return + } + if trimmed := strings.TrimSpace(strVal); trimmed != "" && targetName == "" { + targetName = trimmed + } + } + } + if rawIdx, exists := bodyMap["auth_index"]; exists && rawIdx != nil { + idxStr, ok := rawIdx.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid auth_index"}) + return + } + if trimmed := strings.TrimSpace(idxStr); trimmed != "" && targetAuthIndex == "" { + targetAuthIndex = trimmed + } + } var rawProxySnake, rawProxyKebab string if v, exists := bodyMap["proxy_url"]; exists && v != nil { @@ -232,10 +271,40 @@ func (h *Handler) StartOAuthV8(c *gin.Context) { } } + if rawVerify, exists := bodyMap["verify"]; exists && rawVerify != nil { + switch v := rawVerify.(type) { + case bool: + verifyAfterLogin = v + case string: + switch strings.ToLower(strings.TrimSpace(v)) { + case "1", "true", "yes", "on": + verifyAfterLogin = true + case "0", "false", "no", "off", "": + verifyAfterLogin = false + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid verify"}) + return + } + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid verify"}) + return + } + } + if rawVerifyModel, exists := bodyMap["verify_model"]; exists && rawVerifyModel != nil { + vmStr, ok := rawVerifyModel.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid verify_model"}) + return + } + if trimmed := strings.TrimSpace(vmStr); trimmed != "" { + verifyModel = trimmed + } + } + extraMeta := make(map[string]any) for k, v := range bodyMap { switch k { - case "provider", "proxy_url", "proxy-url", "is_webui": + case "provider", "proxy_url", "proxy-url", "is_webui", "name", "target_auth_id", "auth_id", "auth_index", "verify", "verify_model": continue default: extraMeta[k] = v @@ -247,6 +316,34 @@ func (h *Handler) StartOAuthV8(c *gin.Context) { } } + if targetName != "" || targetAuthIndex != "" { + targetAuth, okTarget := h.lookupAuthFile(targetName, targetAuthIndex) + if !okTarget || targetAuth == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) + return + } + if provider != "" && !strings.EqualFold(strings.TrimSpace(targetAuth.Provider), strings.TrimSpace(provider)) { + c.JSON(http.StatusBadRequest, gin.H{"error": "provider does not match target credential"}) + return + } + if provider == "" { + provider = strings.TrimSpace(targetAuth.Provider) + } + targetAccountID := "" + if targetAuth.Metadata != nil { + targetAccountID, _ = targetAuth.Metadata["account_id"].(string) + } + c.Set(oauthSessionTargetContextKey, oauthRecoveryTarget{ + AuthID: targetAuth.ID, + FileName: firstNonEmpty(strings.TrimSpace(targetAuth.FileName), targetAuth.ID), + Provider: strings.TrimSpace(targetAuth.Provider), + Email: authEmail(targetAuth), + AccountID: strings.TrimSpace(targetAccountID), + VerifyAfterLogin: verifyAfterLogin, + VerifyModel: verifyModel, + }) + } + c.Set(oauthSessionProviderContextKey, provider) switch strings.ToLower(strings.TrimSpace(provider)) { @@ -286,3 +383,112 @@ func (h *Handler) ImportOAuthV8(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "provider_not_found"}) } } + +// VerifyAuthFileV8 runs an operator-triggered, account-pinned inference verification +// using an isolated session and minimal non-user prompt/token budget. +func (h *Handler) VerifyAuthFileV8(c *gin.Context) { + if h == nil || h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "auth manager is not available"}) + return + } + targetName := firstNonEmpty(strings.TrimSpace(c.Query("name")), strings.TrimSpace(c.Query("target_auth_id")), strings.TrimSpace(c.Query("auth_id"))) + targetAuthIndex := strings.TrimSpace(c.Query("auth_index")) + model := strings.TrimSpace(c.Query("model")) + + if c.Request != nil && c.Request.Body != nil { + rawBody, errRead := io.ReadAll(c.Request.Body) + if errRead != nil { + c.JSON(clienterror.HTTPStatusFromErrorOr(errRead, http.StatusBadRequest), gin.H{"error": "invalid request body"}) + return + } + if len(bytes.TrimSpace(rawBody)) > 0 { + var bodyMap map[string]any + if errUnmarshal := json.Unmarshal(rawBody, &bodyMap); errUnmarshal != nil || bodyMap == nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + for _, key := range []string{"name", "target_auth_id", "auth_id"} { + if rawVal, exists := bodyMap[key]; exists && rawVal != nil { + strVal, ok := rawVal.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid " + key}) + return + } + if trimmed := strings.TrimSpace(strVal); trimmed != "" && targetName == "" { + targetName = trimmed + } + } + } + if rawIdx, exists := bodyMap["auth_index"]; exists && rawIdx != nil { + idxStr, ok := rawIdx.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid auth_index"}) + return + } + if trimmed := strings.TrimSpace(idxStr); trimmed != "" && targetAuthIndex == "" { + targetAuthIndex = trimmed + } + } + if rawModel, exists := bodyMap["model"]; exists && rawModel != nil { + modelStr, ok := rawModel.(string) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid model"}) + return + } + if trimmed := strings.TrimSpace(modelStr); trimmed != "" && model == "" { + model = trimmed + } + } + } + } + + if targetName == "" && targetAuthIndex == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "name or auth_index is required"}) + return + } + targetAuth, okTarget := h.lookupAuthFile(targetName, targetAuthIndex) + if !okTarget || targetAuth == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) + return + } + + updated, errVerify := h.authManager.VerifyInference(c.Request.Context(), targetAuth.ID, model) + var entry gin.H + if updated != nil { + entry = h.buildAuthFileEntry(updated, authFileQuotaContext{v8: true}) + } + if errVerify != nil { + status := clienterror.HTTPStatusFromErrorOr(errVerify, http.StatusBadGateway) + if status < 400 { + status = http.StatusBadGateway + } + if statusCoder, okCode := errVerify.(interface{ StatusCode() int }); okCode && statusCoder != nil && statusCoder.StatusCode() >= 400 { + status = statusCoder.StatusCode() + } + resp := gin.H{ + "status": "error", + "verified": false, + "error": coreauth.SanitizeRenewalMessage(errVerify.Error()), + } + if entry != nil { + resp["recovery_stage"] = entry["recovery_stage"] + resp["renewal_state"] = entry["renewal_state"] + resp["quota_blocked"] = entry["quota_blocked"] + resp["credential"] = entry + } + c.JSON(status, resp) + return + } + verified := updated != nil && updated.InferenceVerified + resp := gin.H{ + "status": "ok", + "verified": verified, + } + if entry != nil { + resp["recovery_stage"] = entry["recovery_stage"] + resp["renewal_state"] = entry["renewal_state"] + resp["inference_verified"] = entry["inference_verified"] + resp["credential"] = entry + } + c.JSON(http.StatusOK, resp) +} diff --git a/internal/api/handlers/management/oauth_sessions.go b/internal/api/handlers/management/oauth_sessions.go index 57563ab48ea..62559a1dc83 100644 --- a/internal/api/handlers/management/oauth_sessions.go +++ b/internal/api/handlers/management/oauth_sessions.go @@ -37,18 +37,40 @@ var ( // ErrSIWCCallbackRequiresV8 prevents legacy callback sinks from accepting SIWC codes. var ErrSIWCCallbackRequiresV8 = fmt.Errorf("%w: SIWC callbacks require /v8/management/oauth/callback", errOAuthSessionNotPending) +type oauthRecoveryTarget struct { + AuthID string + FileName string + Provider string + Email string + AccountID string + VerifyAfterLogin bool + VerifyModel string +} + type oauthSession struct { - Provider string - Status string - Source string - ProxyURL string - Metadata map[string]any - SIWC *siwcBindSession - Committing bool - SIWCCommitting bool - Completed bool - CreatedAt time.Time - ExpiresAt time.Time + Provider string + Status string + Source string + ProxyURL string + Metadata map[string]any + SIWC *siwcBindSession + Target *oauthRecoveryTarget + AuthID string + SavedGrantHash string + Saved bool + Adopted bool + BaselineExisted bool + BaselineVersion uint64 + BaselineEpoch uint64 + AdoptedVersion uint64 + AdoptedEpoch uint64 + AdoptError string + VerifyError string + Committing bool + SIWCCommitting bool + Completed bool + CreatedAt time.Time + ExpiresAt time.Time } type oauthSessionStore struct { @@ -235,7 +257,16 @@ func (s *oauthSessionStore) FailPendingPoll(state, message string) (errStatus st } func (s *oauthSessionStore) Complete(state string) { + s.CompleteWithAuth(state, "") +} + +func (s *oauthSessionStore) CompleteWithAuth(state, authID string) { + s.CompleteWithOutcome(state, authID, "", strings.TrimSpace(authID) != "", strings.TrimSpace(authID) != "", false, 0, 0, 0, 0, "", "") +} + +func (s *oauthSessionStore) CompleteWithOutcome(state, authID, savedGrantHash string, saved, adopted, baselineExisted bool, baselineVersion, baselineEpoch, adoptedVersion, adoptedEpoch uint64, adoptErr, verifyErr string) { state = strings.TrimSpace(state) + authID = strings.TrimSpace(authID) if state == "" { return } @@ -252,6 +283,17 @@ func (s *oauthSessionStore) Complete(state string) { session.Status = "" session.ProxyURL = "" session.Metadata = nil + session.AuthID = authID + session.SavedGrantHash = strings.TrimSpace(savedGrantHash) + session.Saved = saved + session.Adopted = adopted + session.BaselineExisted = baselineExisted + session.BaselineVersion = baselineVersion + session.BaselineEpoch = baselineEpoch + session.AdoptedVersion = adoptedVersion + session.AdoptedEpoch = adoptedEpoch + session.AdoptError = strings.TrimSpace(adoptErr) + session.VerifyError = strings.TrimSpace(verifyErr) session.Committing = false session.SIWCCommitting = false session.Completed = true @@ -259,6 +301,52 @@ func (s *oauthSessionStore) Complete(state string) { s.sessions[state] = session } +func (s *oauthSessionStore) SetTarget(state string, target oauthRecoveryTarget) { + state = strings.TrimSpace(state) + if state == "" { + return + } + s.mu.Lock() + defer s.mu.Unlock() + session, ok := s.sessions[state] + if !ok || session.Completed { + return + } + copyTarget := target + session.Target = ©Target + s.sessions[state] = session +} + +func (s *oauthSessionStore) MarkAdopted(state string, adoptedVersion, adoptedEpoch uint64) { + state = strings.TrimSpace(state) + if state == "" { + return + } + s.mu.Lock() + defer s.mu.Unlock() + session, ok := s.sessions[state] + if !ok || !session.Completed { + return + } + session.Adopted = true + session.AdoptedVersion = adoptedVersion + session.AdoptedEpoch = adoptedEpoch + session.AdoptError = "" + s.sessions[state] = session +} + +func (s *oauthSessionStore) GetTarget(state string) (oauthRecoveryTarget, bool) { + state = strings.TrimSpace(state) + if state == "" { + return oauthRecoveryTarget{}, false + } + session, ok := s.Get(state) + if !ok || session.Target == nil { + return oauthRecoveryTarget{}, false + } + return *session.Target, true +} + func (s *oauthSessionStore) CompleteProvider(provider string, source string) int { provider = strings.ToLower(strings.TrimSpace(provider)) if provider == "" { @@ -439,6 +527,22 @@ func FailOAuthSessionCommit(state, message string) { oauthSessions.FailCommit(st func CompleteOAuthSession(state string) { oauthSessions.Complete(state) } +func CompleteOAuthSessionWithAuth(state, authID string) { + oauthSessions.CompleteWithAuth(state, authID) +} + +func CompleteOAuthSessionWithOutcome(state, authID, savedGrantHash string, saved, adopted, baselineExisted bool, baselineVersion, baselineEpoch, adoptedVersion, adoptedEpoch uint64, adoptErr, verifyErr string) { + oauthSessions.CompleteWithOutcome(state, authID, savedGrantHash, saved, adopted, baselineExisted, baselineVersion, baselineEpoch, adoptedVersion, adoptedEpoch, adoptErr, verifyErr) +} + +func SetOAuthSessionTarget(state string, target oauthRecoveryTarget) { + oauthSessions.SetTarget(state, target) +} + +func GetOAuthSessionTarget(state string) (oauthRecoveryTarget, bool) { + return oauthSessions.GetTarget(state) +} + func CompleteOAuthSessionsByProvider(provider string) int { return oauthSessions.CompleteProvider(provider, oauthSessionSourceBuiltin) } diff --git a/internal/api/response_compression.go b/internal/api/response_compression.go index 4a5e213c916..6fc57071773 100644 --- a/internal/api/response_compression.go +++ b/internal/api/response_compression.go @@ -10,6 +10,7 @@ import ( "net/http" "strconv" "strings" + "sync" "syscall" "github.com/gin-gonic/gin" @@ -164,6 +165,41 @@ func (w *responseCompressionCommitWriter) FlushError() error { type responseCompressor interface { io.WriteCloser Flush() error + Reset(io.Writer) +} + +type responseCompressorPool interface { + Get() any + Put(any) +} + +// responseCompressorPools reuse encoders across responses. The default zstd encoder keeps an +// 8 MiB window and its buffers (about 18 MB per stream); a 1 MiB window with lower encoder +// memory loses little ratio on API responses, which are far smaller than the window. +var responseCompressorPools = map[string]responseCompressorPool{ + "zstd": &sync.Pool{New: func() any { + encoder, errNewWriter := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1), zstd.WithWindowSize(1<<20), zstd.WithLowerEncoderMem(true)) + if errNewWriter != nil { + log.Errorf("create response compression encoder: %v", errNewWriter) + return nil + } + return encoder + }}, + "gzip": &sync.Pool{New: func() any { return gzip.NewWriter(io.Discard) }}, +} + +// newResponseCompressor takes a pooled encoder writing to dst, or nil when none can be created. +func newResponseCompressor(encoding string, dst io.Writer) responseCompressor { + pool := responseCompressorPools[encoding] + if pool == nil { + return nil + } + compressor, _ := pool.Get().(responseCompressor) + if compressor == nil { + return nil + } + compressor.Reset(dst) + return compressor } // responseCompressionWriter keeps headers uncommitted until the encoding is known. @@ -307,17 +343,10 @@ func (w *responseCompressionWriter) decide(force bool, size int) { return } } - var compressor responseCompressor - if w.encoding == "gzip" { - compressor = gzip.NewWriter(w.ResponseWriter) - } else { - encoder, errNewWriter := zstd.NewWriter(w.ResponseWriter, zstd.WithEncoderConcurrency(1)) - if errNewWriter != nil { - log.Errorf("create response compression encoder: %v", errNewWriter) - w.decided = true - return - } - compressor = encoder + compressor := newResponseCompressor(w.encoding, w.ResponseWriter) + if compressor == nil { + w.decided = true + return } w.compressor = compressor w.decided = true @@ -373,5 +402,13 @@ func (w *responseCompressionWriter) finish() { log.Errorf("close response compression encoder: %v", errClose) } } + // Drop the response writer reference before pooling the encoder, and fail later writes + // instead of sending them uncompressed under Content-Encoding. + w.compressor.Reset(io.Discard) + responseCompressorPools[w.encoding].Put(w.compressor) + w.compressor = nil + if w.writeErr == nil { + w.writeErr = io.ErrClosedPipe + } } } diff --git a/internal/api/response_compression_bench_test.go b/internal/api/response_compression_bench_test.go new file mode 100644 index 00000000000..b6bb9924e6c --- /dev/null +++ b/internal/api/response_compression_bench_test.go @@ -0,0 +1,159 @@ +package api + +import ( + "fmt" + "io" + "math/rand" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +// benchSSEEvents returns Responses-style SSE deltas whose text mixes common words and random +// identifiers, so they compress like real model output. +func benchSSEEvents(n int) []string { + r := rand.New(rand.NewSource(1)) + words := strings.Fields("the function return error context request response model input output tool call result value index token cache stream event message content type text file path line") + events := make([]string, n) + for i := range events { + var delta strings.Builder + for j, k := 0, 4+r.Intn(12); j < k; j++ { + delta.WriteString(words[r.Intn(len(words))]) + if r.Intn(10) == 0 { + fmt.Fprintf(&delta, "_%x", r.Int31()) + } + delta.WriteByte(' ') + } + events[i] = fmt.Sprintf("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"sequence_number\":%d,\"item_id\":\"msg_%x\",\"delta\":%q}\n\n", i, r.Int63(), delta.String()) + } + return events +} + +func benchCompressionEngine(events []string, hold <-chan struct{}, started *sync.WaitGroup) *gin.Engine { + engine := gin.New() + engine.Use(responseCompressionMiddleware(&config.Config{})) + engine.GET("/sse", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + for i, event := range events { + _, _ = c.Writer.WriteString(event) + c.Writer.Flush() + if i == 0 && hold != nil { + started.Done() + <-hold + } + } + }) + return engine +} + +// BenchmarkResponseCompressionSSE serves whole compressed SSE responses concurrently. +func BenchmarkResponseCompressionSSE(b *testing.B) { + gin.SetMode(gin.ReleaseMode) + for _, n := range []int{20, 500} { + events := benchSSEEvents(n) + for _, encoding := range []string{"zstd", "gzip"} { + b.Run(fmt.Sprintf("%s/%devents", encoding, n), func(b *testing.B) { + engine := benchCompressionEngine(events, nil, nil) + var wire int64 + var mu sync.Mutex + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + req := httptest.NewRequest(http.MethodGet, "/sse", nil) + req.Header.Set("Accept-Encoding", encoding) + rec := httptest.NewRecorder() + engine.ServeHTTP(rec, req) + if rec.Header().Get("Content-Encoding") != encoding { + b.Fatal("response not compressed") + } + mu.Lock() + wire = int64(rec.Body.Len()) + mu.Unlock() + } + }) + b.ReportMetric(float64(wire), "wire-B") + }) + } + } +} + +// BenchmarkResponseCompressionOpenStreams measures open SSE streams with fresh encoders. +func BenchmarkResponseCompressionOpenStreams(b *testing.B) { + benchmarkResponseCompressionOpenStreams(b, false) +} + +// BenchmarkResponseCompressionOpenStreamsAfterShortResponse measures streams whose encoders +// have already compressed a short non-streaming response. +func BenchmarkResponseCompressionOpenStreamsAfterShortResponse(b *testing.B) { + benchmarkResponseCompressionOpenStreams(b, true) +} + +func benchmarkResponseCompressionOpenStreams(b *testing.B, afterShortResponse bool) { + gin.SetMode(gin.ReleaseMode) + const streams = 256 + events := benchSSEEvents(2) + shortResponse := []byte(`{"message":"` + strings.Repeat("response data ", 150) + `"}`) + for _, encoding := range []string{"zstd", "gzip"} { + b.Run(encoding, func(b *testing.B) { + if afterShortResponse { + original := responseCompressorPools[encoding] + // Prime each encoder before it is handed to a stream, independently of pool eviction. + responseCompressorPools[encoding] = &sync.Pool{New: func() any { + compressor, ok := original.Get().(responseCompressor) + if !ok || compressor == nil { + b.Errorf("create %s response compressor", encoding) + return nil + } + compressor.Reset(io.Discard) + if n, errWrite := compressor.Write(shortResponse); errWrite != nil || n != len(shortResponse) { + b.Errorf("prime %s response compressor: wrote %d bytes: %v", encoding, n, errWrite) + return nil + } + if errClose := compressor.Close(); errClose != nil { + b.Errorf("close priming %s response compressor: %v", encoding, errClose) + return nil + } + compressor.Reset(io.Discard) + return compressor + }} + defer func() { responseCompressorPools[encoding] = original }() + } + var perStream float64 + for i := 0; i < b.N; i++ { + hold := make(chan struct{}) + var started, done sync.WaitGroup + engine := benchCompressionEngine(events, hold, &started) + // Two cycles empty sync.Pool: the first moves pooled encoders to the victim cache, + // the second frees them. Otherwise the baseline already counts them. + runtime.GC() + runtime.GC() + var before, during runtime.MemStats + runtime.ReadMemStats(&before) + started.Add(streams) + for s := 0; s < streams; s++ { + done.Add(1) + go func() { + defer done.Done() + req := httptest.NewRequest(http.MethodGet, "/sse", nil) + req.Header.Set("Accept-Encoding", encoding) + engine.ServeHTTP(httptest.NewRecorder(), req) + }() + } + started.Wait() + runtime.GC() + runtime.ReadMemStats(&during) + perStream += float64(int64(during.HeapInuse)-int64(before.HeapInuse)) / streams + close(hold) + done.Wait() + } + b.ReportMetric(perStream/float64(b.N)/1024, "heap-KiB/stream") + }) + } +} diff --git a/internal/api/response_compression_test.go b/internal/api/response_compression_test.go index c5413718bef..067c7f957d7 100644 --- a/internal/api/response_compression_test.go +++ b/internal/api/response_compression_test.go @@ -5,6 +5,7 @@ import ( "bytes" "compress/gzip" "context" + "errors" "io" "net/http" "net/http/httptest" @@ -307,3 +308,129 @@ func TestResponseCompressionSSEIncremental(t *testing.T) { }) } } + +type deterministicResponseCompressorPool struct { + newCompressor func() any + item any +} + +func (p *deterministicResponseCompressorPool) Get() any { + if p.item == nil { + return p.newCompressor() + } + item := p.item + p.item = nil + return item +} + +func (p *deterministicResponseCompressorPool) Put(item any) { + p.item = item +} + +func TestResponseCompressionReusesPooledEncoders(t *testing.T) { + bodies := []string{ + "{\"message\":\"" + strings.Repeat("first response ", 200) + "\"}", + "{\"message\":\"" + strings.Repeat("second, different response ", 150) + "\"}", + } + engine := gin.New() + engine.Use(responseCompressionMiddleware(&config.Config{})) + body := "" + engine.GET("/", func(c *gin.Context) { + c.Data(http.StatusOK, "application/json", []byte(body)) + }) + for _, encoding := range []string{"zstd", "gzip"} { + original := responseCompressorPools[encoding] + created := 0 + responseCompressorPools[encoding] = &deterministicResponseCompressorPool{ + newCompressor: func() any { created++; return original.Get() }, + } + t.Cleanup(func() { responseCompressorPools[encoding] = original }) + const rounds = 16 + for round := 0; round < rounds; round++ { + if round%2 == 1 { + // A failed client write must not poison the encoder handed to the next response. + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Header("Content-Type", "application/json") + writer := &responseCompressionWriter{ResponseWriter: &compressionErrorWriter{ResponseWriter: context.Writer}, encoding: encoding, minBytes: 1} + _, _ = writer.WriteString(bodies[0]) + writer.finish() + continue + } + body = bodies[(round/2)%len(bodies)] + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.Header.Set("Accept-Encoding", encoding) + rec := httptest.NewRecorder() + engine.ServeHTTP(rec, req) + if got := rec.Header().Get("Content-Encoding"); got != encoding { + t.Fatalf("%s round %d: Content-Encoding = %q", encoding, round, got) + } + if got := string(compressionTestBody(t, encoding, rec.Body.Bytes())); got != body { + t.Fatalf("%s round %d: decoded body mismatch", encoding, round) + } + } + if created != 1 { + t.Fatalf("%s: created %d encoders for %d responses, want 1", encoding, created, rounds) + } + } +} + +func TestResponseCompressionReusesEncoderAfterWriteFailure(t *testing.T) { + for _, encoding := range []string{"zstd", "gzip"} { + t.Run(encoding, func(t *testing.T) { + original := responseCompressorPools[encoding] + responseCompressorPools[encoding] = &deterministicResponseCompressorPool{newCompressor: original.Get} + t.Cleanup(func() { responseCompressorPools[encoding] = original }) + + failedContext, _ := gin.CreateTestContext(httptest.NewRecorder()) + failedContext.Header("Content-Type", "text/event-stream") + failedWriter := &responseCompressionWriter{ + ResponseWriter: &compressionErrorWriter{ResponseWriter: failedContext.Writer}, + encoding: encoding, + minBytes: 1, + } + _, errWrite := failedWriter.WriteString("data: failed\n\n") + if errWrite == nil { + errWrite = failedWriter.FlushError() + } + if !errors.Is(errWrite, io.ErrClosedPipe) { + t.Fatalf("failed response error = %v, want %v", errWrite, io.ErrClosedPipe) + } + failedCompressor := failedWriter.compressor + if failedCompressor == nil { + t.Fatal("failed response has no compressor") + } + failedWriter.finish() + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Header("Content-Type", "application/json") + writer := &responseCompressionWriter{ResponseWriter: context.Writer, encoding: encoding, minBytes: 1} + body := `{"message":"response after failure"}` + if _, errWrite = writer.WriteString(body); errWrite != nil { + t.Fatal(errWrite) + } + if writer.compressor != failedCompressor { + t.Fatal("the encoder from the failed response was not reused") + } + writer.finish() + if got := string(compressionTestBody(t, encoding, recorder.Body.Bytes())); got != body { + t.Fatalf("decoded body = %q, want %q", got, body) + } + }) + } +} + +func TestResponseCompressionWriteAfterFinishFails(t *testing.T) { + for _, encoding := range []string{"zstd", "gzip"} { + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Header("Content-Type", "text/event-stream") + writer := &responseCompressionWriter{ResponseWriter: context.Writer, encoding: encoding, minBytes: 1} + if _, errWrite := writer.WriteString("data: one\n\n"); errWrite != nil { + t.Fatal(errWrite) + } + writer.finish() + if _, errWrite := writer.WriteString("data: late\n\n"); errWrite == nil { + t.Fatalf("%s: a write after finish must fail, not go out uncompressed", encoding) + } + } +} diff --git a/internal/api/server_management_v8.go b/internal/api/server_management_v8.go index 47cbee0246a..e73ad893132 100644 --- a/internal/api/server_management_v8.go +++ b/internal/api/server_management_v8.go @@ -50,6 +50,7 @@ func (s *Server) registerManagementV8Routes() { v8.PATCH("/credentials/status", s.mgmt.PatchAuthFileStatus) v8.PATCH("/credentials/fields", s.mgmt.PatchAuthFileFields) v8.POST("/credentials/refresh", s.mgmt.RefreshAuthFiles) + v8.POST("/credentials/verify", s.mgmt.VerifyAuthFileV8) v8.POST("/oauth/import", s.mgmt.ImportOAuthV8) v8.GET("/oauth/auth-url", s.mgmt.StartOAuthV8) v8.POST("/oauth/auth-url", s.mgmt.StartOAuthV8) diff --git a/internal/api/server_management_v8_test.go b/internal/api/server_management_v8_test.go index da91cd74db0..b59ad2b038d 100644 --- a/internal/api/server_management_v8_test.go +++ b/internal/api/server_management_v8_test.go @@ -88,7 +88,7 @@ func TestManagementV8IndependentContract(t *testing.T) { for _, route := range []string{ "GET /v0/management/debug", "PUT /v0/management/request-retry", "GET /v0/management/auth-files", "GET /v8/management/observability/logs", "GET /v8/management/observability/usage/queue", - "GET /v8/management/credentials", "POST /v8/management/credentials", + "GET /v8/management/credentials", "POST /v8/management/credentials", "POST /v8/management/credentials/verify", "GET /v8/management/oauth/auth-url", "POST /v8/management/oauth/auth-url", "POST /v8/management/oauth/import", "POST /v8/management/oauth/callback", "POST /v8/management/routing/cooldown/reset", "POST /v8/management/routing/model-definitions/refresh", diff --git a/internal/api/server_multi_agent_config_test.go b/internal/api/server_multi_agent_config_test.go index 176bfed478d..51f920b767d 100644 --- a/internal/api/server_multi_agent_config_test.go +++ b/internal/api/server_multi_agent_config_test.go @@ -36,6 +36,11 @@ func TestModelsMultiAgentClientConfigAndReload(t *testing.T) { raw string want bool }{ + // The synthetic model is synthesized from the catalog's default + // template, which carries multi_agent_version="v2". Runtime code + // preserves that template value unless the opt-in rewrites it, + // so every case observes "v2"; tc.want asserts the published + // client setting after reload, not the catalog field. {"omitted", "{}", false}, {"enabled", "client: {codex: {optimize-multi-agent-v2: true}}", true}, {"disabled", "client: {codex: {optimize-multi-agent-v2: false}}", false}, @@ -72,10 +77,9 @@ func TestModelsMultiAgentClientConfigAndReload(t *testing.T) { if entry["slug"] != modelID { continue } - var want any - if tc.want { - want = "v2" - } + // Synthesized from the default template: its v2 value is + // preserved regardless of the opt-in setting. + want := any("v2") if value, present := entry["multi_agent_version"]; !present || value != want { t.Fatalf("multi_agent_version = %#v (present %t), want %#v", value, present, want) } diff --git a/internal/api/server_routes.go b/internal/api/server_routes.go index 206e3332830..4de317bde8d 100644 --- a/internal/api/server_routes.go +++ b/internal/api/server_routes.go @@ -567,11 +567,14 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { } // AttachWebsocketRoute registers a websocket upgrade handler on the primary Gin engine. -// The handler is served as-is without additional middleware beyond the standard stack already configured. +// Home mode does not support AI Studio and must not expose its websocket gateway. func (s *Server) AttachWebsocketRoute(path string, handler http.Handler) { if s == nil || s.engine == nil || handler == nil { return } + if cfg := s.getConfig(); cfg != nil && cfg.Home.Enabled { + return + } trimmed := strings.TrimSpace(path) if trimmed == "" { trimmed = "/v1/ws" diff --git a/internal/api/server_websocket_route_test.go b/internal/api/server_websocket_route_test.go new file mode 100644 index 00000000000..392c153b211 --- /dev/null +++ b/internal/api/server_websocket_route_test.go @@ -0,0 +1,68 @@ +package api + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func TestAttachWebsocketRouteDisabledInHomeMode(t *testing.T) { + for _, homeEnabled := range []bool{false, true} { + for _, wsAuth := range []bool{false, true} { + t.Run(fmt.Sprintf("home=%t/ws-auth=%t", homeEnabled, wsAuth), func(t *testing.T) { + cfg := &config.Config{ + SDKConfig: config.SDKConfig{APIKeys: []string{"test-key"}}, + Home: config.HomeConfig{Enabled: homeEnabled}, + WebsocketAuth: wsAuth, + } + if homeEnabled { + // Home mode clears local API keys during startup. + cfg.APIKeys = nil + } + server := newTestServerWithConfig(t, cfg) + // Isolate route registration from the Home heartbeat readiness gate. + server.engine = gin.New() + handlerCalled := false + server.AttachWebsocketRoute("/v1/ws", http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerCalled = true + w.WriteHeader(http.StatusNoContent) + })) + + for _, route := range server.engine.Routes() { + if homeEnabled && route.Path == "/v1/ws" { + t.Fatal("websocket route must not be registered in Home mode") + } + } + + for _, token := range []string{"", "test-key", "invalid-key"} { + handlerCalled = false + req := httptest.NewRequest(http.MethodGet, "/v1/ws", nil) + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Upgrade", "websocket") + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + recorder := httptest.NewRecorder() + server.engine.ServeHTTP(recorder, req) + + wantStatus := http.StatusNoContent + if homeEnabled { + wantStatus = http.StatusNotFound + } else if wsAuth && token != "test-key" { + wantStatus = http.StatusUnauthorized + } + if recorder.Code != wantStatus { + t.Errorf("token=%q: status = %d, want %d", token, recorder.Code, wantStatus) + } + if wantCalled := wantStatus == http.StatusNoContent; handlerCalled != wantCalled { + t.Errorf("token=%q: handler called = %t, want %t", token, handlerCalled, wantCalled) + } + } + }) + } + } +} diff --git a/internal/auth/codex/openai_auth.go b/internal/auth/codex/openai_auth.go index bc9ce24a088..d5246acec50 100644 --- a/internal/auth/codex/openai_auth.go +++ b/internal/auth/codex/openai_auth.go @@ -348,7 +348,22 @@ func isNonRetryableRefreshErr(err error) bool { return false } raw := strings.ToLower(err.Error()) - return strings.Contains(raw, "refresh_token_reused") + for _, token := range []string{ + "refresh_token_reused", + "refresh_token_expired", + "refresh_token_invalidated", + "token_revoked", + "invalid_grant", + "invalid_client", + "status 401", + "401 unauthorized", + "request failed (401)", + } { + if strings.Contains(raw, token) { + return true + } + } + return false } // UpdateTokenStorage updates an existing CodexTokenStorage with new token data. diff --git a/internal/auth/codex/openai_auth_test.go b/internal/auth/codex/openai_auth_test.go index b0c86678a87..aeaeb5e1c70 100644 --- a/internal/auth/codex/openai_auth_test.go +++ b/internal/auth/codex/openai_auth_test.go @@ -304,3 +304,74 @@ func TestCreateAndUpdateTokenStorage_PlanType(t *testing.T) { t.Fatalf("updated storage.PlanType = %q, want free", storage.PlanType) } } + +func TestRefreshTokensWithRetry_TerminalGrantErrorsOnlyAttemptOnce(t *testing.T) { + resetCodexRefreshGroupForTest() + t.Cleanup(resetCodexRefreshGroupForTest) + + cases := []struct { + name string + statusCode int + body string + wantCalls int32 + }{ + { + name: "refresh_token_reused", + statusCode: http.StatusUnauthorized, + body: `{"error":{"message":"Your refresh token has already been used","code":"refresh_token_reused"}}`, + wantCalls: 1, + }, + { + name: "invalid_grant", + statusCode: http.StatusBadRequest, + body: `{"error":"invalid_grant","error_description":"Refresh token not found or invalid"}`, + wantCalls: 1, + }, + { + name: "refresh_token_expired", + statusCode: http.StatusUnauthorized, + body: `{"error":{"message":"Refresh token expired","code":"refresh_token_expired"}}`, + wantCalls: 1, + }, + { + name: "refresh_token_invalidated", + statusCode: http.StatusUnauthorized, + body: `{"error":{"message":"Refresh token invalidated","code":"refresh_token_invalidated"}}`, + wantCalls: 1, + }, + { + name: "transient_503_retries", + statusCode: http.StatusServiceUnavailable, + body: `{"error":"service_unavailable"}`, + wantCalls: 2, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resetCodexRefreshGroupForTest() + var calls int32 + auth := &CodexAuth{ + httpClient: &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + atomic.AddInt32(&calls, 1) + return &http.Response{ + StatusCode: tc.statusCode, + Body: io.NopCloser(strings.NewReader(tc.body)), + Header: make(http.Header), + Request: req, + }, nil + }), + }, + } + + _, errRefresh := auth.RefreshTokensWithRetry(context.Background(), "rt-"+tc.name, 2) + if errRefresh == nil { + t.Fatalf("expected error for %s", tc.name) + } + if got := atomic.LoadInt32(&calls); got != tc.wantCalls { + t.Fatalf("%s: RefreshTokensWithRetry attempts = %d, want %d", tc.name, got, tc.wantCalls) + } + }) + } +} diff --git a/internal/auth/codex/refresh_patch.go b/internal/auth/codex/refresh_patch.go index e3ce17f52be..2e831dbfddf 100644 --- a/internal/auth/codex/refresh_patch.go +++ b/internal/auth/codex/refresh_patch.go @@ -11,23 +11,34 @@ var codexSIWCCredentialKeys = []string{ // Its original token identities (both in-memory and on-disk baselines) act as versions, // independent of timestamp precision. type CredentialRefreshPatch struct { - outerBefore map[string]any - outerDiskBefore map[string]any - outerAfter map[string]any - siwcBefore map[string]any - siwcDiskBefore map[string]any - siwcAfter map[string]any + outerBefore map[string]any + outerDiskBefore map[string]any + outerAfter map[string]any + siwcBefore map[string]any + siwcDiskBefore map[string]any + siwcAfter map[string]any + baselineOuterRefresh string + baselineSIWCRefresh string + upstreamRefreshed bool + outerErr error + siwcErr error } // NewCredentialRefreshPatch records only credential partitions actually refreshed. func NewCredentialRefreshPatch(before, after map[string]any, outer, siwc bool) *CredentialRefreshPatch { - return NewCredentialRefreshPatchWithDiskBaseline(before, nil, after, outer, siwc) + patch := NewCredentialRefreshPatchWithDiskBaseline(before, nil, after, outer, siwc) + patch.upstreamRefreshed = outer || siwc + return patch } // NewCredentialRefreshPatchWithDiskBaseline records both the in-memory baseline (before) // and the on-disk baseline (diskBefore) observed under the file lock before refreshing to after. func NewCredentialRefreshPatchWithDiskBaseline(before, diskBefore, after map[string]any, outer, siwc bool) *CredentialRefreshPatch { - patch := &CredentialRefreshPatch{} + patch := &CredentialRefreshPatch{ + upstreamRefreshed: outer || siwc, + baselineOuterRefresh: lookupRefreshTokenField(diskBefore, before, false), + baselineSIWCRefresh: lookupRefreshTokenField(diskBefore, before, true), + } if outer { patch.outerBefore = credentialFields(before, codexOuterCredentialKeys) if len(diskBefore) > 0 { @@ -48,6 +59,26 @@ func NewCredentialRefreshPatchWithDiskBaseline(before, diskBefore, after map[str return patch } +func lookupRefreshTokenField(diskBefore, memBefore map[string]any, siwc bool) string { + for _, src := range []map[string]any{diskBefore, memBefore} { + if len(src) == 0 { + continue + } + target := src + if siwc { + nested, ok := src["siwc"].(map[string]any) + if !ok || len(nested) == 0 { + continue + } + target = nested + } + if rt, ok := target["refresh_token"].(string); ok && rt != "" { + return rt + } + } + return "" +} + func credentialFields(metadata map[string]any, keys []string) map[string]any { fields := make(map[string]any) for _, key := range keys { @@ -119,6 +150,128 @@ func (p *CredentialRefreshPatch) RefreshedSIWC() bool { return p != nil && p.siwcAfter != nil } +// SetUpstreamRefreshed records whether an upstream OAuth token refresh actually succeeded +// (as opposed to only adopting rotated credentials already present on disk). +func (p *CredentialRefreshPatch) SetUpstreamRefreshed(refreshed bool) { + if p == nil { + return + } + p.upstreamRefreshed = refreshed +} + +// UpstreamRefreshed reports whether p includes an upstream OAuth token refresh. +func (p *CredentialRefreshPatch) UpstreamRefreshed() bool { + if p == nil { + return false + } + return p.upstreamRefreshed +} + +// SetPartitionErrors records any per-partition refresh failures that occurred alongside a partial refresh. +func (p *CredentialRefreshPatch) SetPartitionErrors(outerErr, siwcErr error) { + if p == nil { + return + } + p.outerErr = outerErr + p.siwcErr = siwcErr +} + +// OuterError returns any error encountered while refreshing the outer OAuth partition. +func (p *CredentialRefreshPatch) OuterError() error { + if p == nil { + return nil + } + return p.outerErr +} + +// SIWCError returns any error encountered while refreshing the SIWC partition. +func (p *CredentialRefreshPatch) SIWCError() error { + if p == nil { + return nil + } + return p.siwcErr +} + +// FailedOuterRefreshToken returns the baseline outer refresh token that was submitted when OuterError occurred. +func (p *CredentialRefreshPatch) FailedOuterRefreshToken() string { + if p == nil || p.outerErr == nil { + return "" + } + return p.baselineOuterRefresh +} + +// FailedSIWCRefreshToken returns the baseline SIWC refresh token that was submitted when SIWCError occurred. +func (p *CredentialRefreshPatch) FailedSIWCRefreshToken() string { + if p == nil || p.siwcErr == nil { + return "" + } + return p.baselineSIWCRefresh +} + +// PartitionRefreshError carries per-partition refresh errors when a credential refresh fails. +type PartitionRefreshError struct { + outerErr error + siwcErr error +} + +// NewPartitionRefreshError returns a combined error carrying both outer and SIWC refresh errors when present. +func NewPartitionRefreshError(outerErr, siwcErr error) error { + if outerErr == nil && siwcErr == nil { + return nil + } + if outerErr != nil && siwcErr == nil { + return outerErr + } + if siwcErr != nil && outerErr == nil { + return siwcErr + } + return &PartitionRefreshError{outerErr: outerErr, siwcErr: siwcErr} +} + +func (e *PartitionRefreshError) Error() string { + if e == nil { + return "" + } + if e.outerErr != nil && e.siwcErr != nil { + return e.outerErr.Error() + "; " + e.siwcErr.Error() + } + if e.outerErr != nil { + return e.outerErr.Error() + } + if e.siwcErr != nil { + return e.siwcErr.Error() + } + return "" +} + +func (e *PartitionRefreshError) Unwrap() []error { + if e == nil { + return nil + } + var errs []error + if e.outerErr != nil { + errs = append(errs, e.outerErr) + } + if e.siwcErr != nil { + errs = append(errs, e.siwcErr) + } + return errs +} + +func (e *PartitionRefreshError) OuterRefreshError() error { + if e == nil { + return nil + } + return e.outerErr +} + +func (e *PartitionRefreshError) SIWCRefreshError() error { + if e == nil { + return nil + } + return e.siwcErr +} + // CanApplyTo reports whether at least one refreshed partition in p still matches disk's version. func (p *CredentialRefreshPatch) CanApplyTo(disk map[string]any) bool { if p == nil { diff --git a/internal/auth/codex/siwc_auth.go b/internal/auth/codex/siwc_auth.go index 9752516522b..3b0ffd5bea5 100644 --- a/internal/auth/codex/siwc_auth.go +++ b/internal/auth/codex/siwc_auth.go @@ -5,6 +5,7 @@ import ( "crypto/rand" "encoding/base64" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -202,15 +203,51 @@ func (o *CodexAuth) ExchangeSIWCCodeForTokens(ctx context.Context, code, issuedC return o.requestSIWCToken(ctx, form, issuedClientID, deviceID, true) } +// SIWCRefreshError marks an error as originating from a Sign in with ChatGPT (SIWC) token refresh. +type SIWCRefreshError struct { + Err error +} + +func (e *SIWCRefreshError) Error() string { + if e == nil || e.Err == nil { + return "" + } + return e.Err.Error() +} + +func (e *SIWCRefreshError) Unwrap() error { + if e == nil { + return nil + } + return e.Err +} + +// IsSIWCRefreshError reports that this error came from the SIWC refresh partition. +func (e *SIWCRefreshError) IsSIWCRefreshError() bool { + return true +} + +// WrapSIWCRefreshError wraps err in a SIWCRefreshError if non-nil. +func WrapSIWCRefreshError(err error) error { + if err == nil { + return nil + } + var existing *SIWCRefreshError + if errors.As(err, &existing) { + return err + } + return &SIWCRefreshError{Err: err} +} + // RefreshSIWCTokens refreshes a Sign in with ChatGPT access token using its refresh token and issued client ID. func (o *CodexAuth) RefreshSIWCTokens(ctx context.Context, refreshToken, clientID, deviceID string) (*CodexSIWCTokenData, error) { refreshToken = strings.TrimSpace(refreshToken) if refreshToken == "" { - return nil, fmt.Errorf("SIWC refresh token is required") + return nil, WrapSIWCRefreshError(fmt.Errorf("SIWC refresh token is required")) } clientID = strings.TrimSpace(clientID) if clientID == "" { - return nil, fmt.Errorf("Stored OpenAI OAuth credential does not contain an issued client ID; reconnect ChatGPT") + return nil, WrapSIWCRefreshError(fmt.Errorf("Stored OpenAI OAuth credential does not contain an issued client ID; reconnect ChatGPT")) } if ctx == nil { ctx = context.Background() @@ -229,11 +266,11 @@ func (o *CodexAuth) RefreshSIWCTokens(ctx context.Context, refreshToken, clientI return o.requestSIWCToken(refreshCtx, form, clientID, deviceID, false) }) if errSF != nil { - return nil, errSF + return nil, WrapSIWCRefreshError(errSF) } tokenData, ok := result.(*CodexSIWCTokenData) if !ok || tokenData == nil { - return nil, fmt.Errorf("SIWC token refresh failed: invalid single-flight result") + return nil, WrapSIWCRefreshError(fmt.Errorf("SIWC token refresh failed: invalid single-flight result")) } return tokenData, nil } diff --git a/internal/client/codex/models/models_test.go b/internal/client/codex/models/models_test.go index d3b40311818..843784790ea 100644 --- a/internal/client/codex/models/models_test.go +++ b/internal/client/codex/models/models_test.go @@ -476,7 +476,14 @@ func TestApplyCodexClientModelMetadataPreservesMultiAgentVersionWhenDisabled(t * func TestCodexClientModelsResponseAppliesMaxContextLengthOverride(t *testing.T) { const wantOverride = 1048576 - const wantDefault = 272000 + // Non-template models inherit context windows from the catalog's + // default template (currently the gpt-5.5 entry). + _, defaultTemplate, errLoadTemplates := loadCodexClientModelTemplates() + if errLoadTemplates != nil || defaultTemplate == nil { + t.Fatalf("load codex client model templates: %v", errLoadTemplates) + } + wantDefaultContext := intModelValue(defaultTemplate, "context_window") + wantDefaultMaxContext := intModelValue(defaultTemplate, "max_context_window") resp := BuildResponse([]map[string]any{ {"id": "deepseek-v4-flash", "max_context_length": wantOverride}, @@ -494,22 +501,23 @@ func TestCodexClientModelsResponseAppliesMaxContextLengthOverride(t *testing.T) } for _, testCase := range []struct { - slug string - want int + slug string + wantContext int + wantMaxContext int }{ - {slug: "deepseek-v4-flash", want: wantOverride}, - {slug: "deepseek-v4-pro", want: wantDefault}, - {slug: "gpt-5.5", want: wantOverride}, + {slug: "deepseek-v4-flash", wantContext: wantOverride, wantMaxContext: wantOverride}, + {slug: "deepseek-v4-pro", wantContext: wantDefaultContext, wantMaxContext: wantDefaultMaxContext}, + {slug: "gpt-5.5", wantContext: wantOverride, wantMaxContext: wantOverride}, } { entry := bySlug[testCase.slug] if entry == nil { t.Fatalf("missing model %q", testCase.slug) } - if got := intModelValue(entry, "context_window"); got != testCase.want { - t.Errorf("%s context_window = %d, want %d", testCase.slug, got, testCase.want) + if got := intModelValue(entry, "context_window"); got != testCase.wantContext { + t.Errorf("%s context_window = %d, want %d", testCase.slug, got, testCase.wantContext) } - if got := intModelValue(entry, "max_context_window"); got != testCase.want { - t.Errorf("%s max_context_window = %d, want %d", testCase.slug, got, testCase.want) + if got := intModelValue(entry, "max_context_window"); got != testCase.wantMaxContext { + t.Errorf("%s max_context_window = %d, want %d", testCase.slug, got, testCase.wantMaxContext) } } } @@ -931,8 +939,14 @@ func TestCodexClientModelsResponse_UnrelatedCustomProviderRetainsFallback(t *tes if got := stringModelValue(entry, "slug"); got != "my-unrelated-custom-model" { t.Errorf("slug = %q, want my-unrelated-custom-model", got) } - if got := stringModelValue(entry, "comp_hash"); got != "2911" { - t.Errorf("comp_hash = %q, want 2911 (from generic gpt-5.5 fallback)", got) + // The fallback clones the catalog's default template (currently the + // gpt-5.5 entry); comp_hash is inherited from it verbatim. + _, defaultTemplate, errLoadTemplates := loadCodexClientModelTemplates() + if errLoadTemplates != nil || defaultTemplate == nil { + t.Fatalf("load codex client model templates: %v", errLoadTemplates) + } + if got, want := stringModelValue(entry, "comp_hash"), stringModelValue(defaultTemplate, "comp_hash"); got != want { + t.Errorf("comp_hash = %q, want %q (from generic fallback template)", got, want) } if got, _ := entry["supports_search_tool"].(bool); got { t.Errorf("supports_search_tool = %v, want false for non-template fallback", got) @@ -1556,7 +1570,7 @@ func TestCodexClientModelsResponse_DevinDisplayName(t *testing.T) { "id": "devin/gpt-6-astra", "display_name": "GPT-6 Astra", }, - // 2. Template Devin model without display_name (inherits template "GPT-5.5") + // 2. Template Devin model without display_name (inherits the template's display_name) { "id": "devin/gpt-5.5", }, @@ -1636,15 +1650,21 @@ func TestCodexClientModelsResponse_DevinDisplayName(t *testing.T) { bySlug[slug] = m } + templates, _, errLoadTemplates := loadCodexClientModelTemplates() + if errLoadTemplates != nil || templates["gpt-5.5"] == nil { + t.Fatalf("load codex client model templates: %v", errLoadTemplates) + } + templateDisplayName := stringModelValue(templates["gpt-5.5"], "display_name") + testCases := []struct { slug string wantDisplayName string }{ {"devin/gpt-6-astra", "GPT-6 Astra (Devin)"}, - {"devin/gpt-5.5", "GPT-5.5 (Devin)"}, + {"devin/gpt-5.5", templateDisplayName + " (Devin)"}, {"devin/swe-2", "SWE-2 (Devin)"}, {"gpt-6-astra", "GPT 6.0 Astra"}, - {"gpt-5.5", "GPT-5.5"}, + {"gpt-5.5", templateDisplayName}, {"devin/swe-1-7", "SWE-1.7 (Devin)"}, {"custom-devin-by-type", "Custom Model (Devin)"}, {"custom-devin-by-owned", "Cognition Model (Devin)"}, diff --git a/internal/discovery/discovery_test.go b/internal/discovery/discovery_test.go index 5c82f8f9f96..3a5ee98ec1d 100644 --- a/internal/discovery/discovery_test.go +++ b/internal/discovery/discovery_test.go @@ -509,8 +509,11 @@ func TestAdvertiserAndBrowser_Integration(t *testing.T) { browser := NewZeroconfBrowser(ifaces...) results, err := browser.Browse(browseCtx, DefaultServiceType, DefaultDomain) if err != nil { - if strings.Contains(err.Error(), "failed to join any of these interfaces") { - t.Skipf("live multicast is unavailable on this host: %v", err) + // Some platforms (notably windows/amd64) cannot open mDNS sockets or + // join multicast interfaces; the live browse cannot run there. + if strings.Contains(err.Error(), "open mDNS sockets") || + strings.Contains(err.Error(), "failed to join any of these interfaces") { + t.Skipf("skipping live multicast browse: %v", err) } t.Fatalf("browse failed: %v", err) } diff --git a/internal/logging/request_capture.go b/internal/logging/request_capture.go new file mode 100644 index 00000000000..6a97ea4ae8c --- /dev/null +++ b/internal/logging/request_capture.go @@ -0,0 +1,25 @@ +package logging + +import "context" + +type deferredAPIRequestCaptureDisabledKey struct{} + +// WithDeferredAPIRequestCaptureDisabled skips error-only upstream request capture +// for this context and its descendants. Full request logging is unaffected. +// Connection-scoped contexts keep this policy even when asynchronous callbacks +// outlive the handler and its pooled Gin context is reused. +func WithDeferredAPIRequestCaptureDisabled(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, deferredAPIRequestCaptureDisabledKey{}, true) +} + +// DeferredAPIRequestCaptureDisabled reports whether error-only capture is disabled. +func DeferredAPIRequestCaptureDisabled(ctx context.Context) bool { + if ctx == nil { + return false + } + disabled, _ := ctx.Value(deferredAPIRequestCaptureDisabledKey{}).(bool) + return disabled +} diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 7410b88e03f..164a3b48322 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -227,6 +227,16 @@ func cloneRegistryModels(in []*registry.ModelInfo) []*registry.ModelInfo { return out } +func normalizeModelRegistrationResponse(resp pluginapi.ModelRegistrationResponse) []pluginapi.ProviderModels { + if len(resp.Registrations) > 0 { + return resp.Registrations + } + return []pluginapi.ProviderModels{{ + Provider: resp.Provider, + Models: resp.Models, + }} +} + // RegisterModels collects static catalogs for executor and per-auth registration. // Catalog-only plugins must not advertise availability independently of credentials. func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) { @@ -237,7 +247,9 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) snap := h.Snapshot() records := h.activeRecordsFromSnapshot(snap) nextProviders := make(map[string]string) + nextProviderSets := make(map[string][]string) nextModelRegistrations := make(map[string]pluginModelRegistration) + nextPluginModelRegistrations := make(map[string][]pluginModelRegistration) for _, record := range records { modelProvider := record.plugin.Capabilities.ModelProvider registrar := record.plugin.Capabilities.ModelRegistrar @@ -253,8 +265,9 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) modelResp, errStaticModels := h.callModelProviderStaticModels(ctx, record, modelProvider) errRegisterModels = errStaticModels resp = pluginapi.ModelRegistrationResponse{ - Provider: modelResp.Provider, - Models: modelResp.Models, + Provider: modelResp.Provider, + Models: modelResp.Models, + Registrations: modelResp.Registrations, } } else { resp, errRegisterModels = h.callModelRegistrar(ctx, record, registrar) @@ -264,32 +277,57 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) continue } - provider := strings.ToLower(strings.TrimSpace(resp.Provider)) - if provider == "" || len(resp.Models) == 0 { - continue - } + entries := normalizeModelRegistrationResponse(resp) + var providersForPlugin []string + seenProviders := make(map[string]struct{}, len(entries)) + var regsForPlugin []pluginModelRegistration + regIndexByProvider := make(map[string]int, len(entries)) + for _, entry := range entries { + provider := strings.ToLower(strings.TrimSpace(entry.Provider)) + if provider == "" || len(entry.Models) == 0 { + continue + } - models := make([]*registry.ModelInfo, 0, len(resp.Models)) - for _, item := range resp.Models { - model := pluginModelInfoToRegistryModelInfo(item) - if model == nil || strings.TrimSpace(model.ID) == "" { + models := make([]*registry.ModelInfo, 0, len(entry.Models)) + for _, item := range entry.Models { + model := pluginModelInfoToRegistryModelInfo(item) + if model == nil || strings.TrimSpace(model.ID) == "" { + continue + } + model.ID = strings.TrimSpace(model.ID) + models = append(models, model) + } + if len(models) == 0 { continue } - model.ID = strings.TrimSpace(model.ID) - models = append(models, model) + + if _, exists := seenProviders[provider]; !exists { + seenProviders[provider] = struct{}{} + providersForPlugin = append(providersForPlugin, provider) + } + if idx, exists := regIndexByProvider[provider]; exists { + merged := map[string][]*registry.ModelInfo{provider: regsForPlugin[idx].models} + appendModelsForProvider(merged, provider, models) + regsForPlugin[idx].models = merged[provider] + continue + } + regIndexByProvider[provider] = len(regsForPlugin) + regsForPlugin = append(regsForPlugin, pluginModelRegistration{ + pluginID: record.id, + provider: provider, + priority: record.priority, + models: cloneRegistryModels(models), + hasExecutor: record.plugin.Capabilities.Executor != nil, + }) } - if len(models) == 0 { + if len(regsForPlugin) == 0 { continue } - nextModelRegistrations[record.id] = pluginModelRegistration{ - pluginID: record.id, - provider: provider, - priority: record.priority, - models: cloneRegistryModels(models), - hasExecutor: record.plugin.Capabilities.Executor != nil, - } - nextProviders[record.id] = provider + nextModelRegistrations[record.id] = regsForPlugin[0] + nextPluginModelRegistrations[record.id] = regsForPlugin + nextProviders[record.id] = providersForPlugin[0] + nextProviderSets[record.id] = providersForPlugin } h.mu.Lock() defer h.mu.Unlock() @@ -297,7 +335,9 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) return } h.modelProviders = nextProviders + h.modelProviderSets = nextProviderSets h.modelRegistrations = nextModelRegistrations + h.pluginModelRegistrations = nextPluginModelRegistrations } // HasAuthModelProvider reports whether an active plugin declares per-auth model @@ -327,11 +367,20 @@ func (h *Host) authModelProviderMatches(record capabilityRecord, provider string identifier, ok := h.callAuthProviderIdentifier(record.id, caps.AuthProvider) return ok && normalizeProviderID(identifier) == provider } - candidate := normalizeProviderID(h.modelProvider(record.id)) - if candidate == "" && caps.Executor != nil { - candidate, _ = h.executorProvider(record, caps.Executor) + candidates := h.modelProvidersForPlugin(record.id) + if len(candidates) > 0 { + for _, candidate := range candidates { + if normalizeProviderID(candidate) == provider { + return true + } + } + return false + } + if caps.Executor != nil { + candidate, _ := h.executorProvider(record, caps.Executor) + return candidate == provider } - return candidate == provider + return false } func (h *Host) ModelsForAuth(ctx context.Context, auth *coreauth.Auth) AuthModelResult { @@ -465,7 +514,10 @@ func (h *Host) callModelRegistrar(ctx context.Context, record capabilityRecord, err = fmt.Errorf("model registrar panic: %v", recovered) } }() - return registrar.RegisterModels(ctx, pluginapi.ModelRegistrationRequest{Plugin: record.meta}) + return registrar.RegisterModels(ctx, pluginapi.ModelRegistrationRequest{ + Plugin: record.meta, + MultiProviderModels: true, + }) } func (h *Host) callModelProviderStaticModels(ctx context.Context, record capabilityRecord, provider pluginapi.ModelProvider) (resp pluginapi.ModelResponse, err error) { @@ -480,8 +532,9 @@ func (h *Host) callModelProviderStaticModels(ctx context.Context, record capabil } }() return provider.StaticModels(ctx, pluginapi.StaticModelRequest{ - Plugin: record.meta, - Host: h.hostConfigSummary(), + Plugin: record.meta, + Host: h.hostConfigSummary(), + MultiProviderModels: true, }) } diff --git a/internal/pluginhost/adapters_auth.go b/internal/pluginhost/adapters_auth.go index f55d29842d2..1a901980953 100644 --- a/internal/pluginhost/adapters_auth.go +++ b/internal/pluginhost/adapters_auth.go @@ -12,6 +12,7 @@ import ( "github.com/klauspost/compress/zstd" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" sdkaccess "github.com/router-for-me/CLIProxyAPI/v8/sdk/access" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) @@ -113,6 +114,7 @@ func (h *Host) RegisterFrontendAuthProviders() { path: record.path, version: record.version, provider: provider, + omitBody: frontendAuthOmitsBody(record), } key := strings.TrimSpace(adapter.Identifier()) if key == "" { @@ -168,6 +170,28 @@ type accessAdapter struct { path string version string provider pluginapi.FrontendAuthProvider + omitBody bool +} + +func frontendAuthOmitsBody(record capabilityRecord) bool { + if record.plugin.SchemaVersion < pluginabi.SchemaVersionFrontendAuthOmitBody { + return false + } + caps := record.plugin.Capabilities + if caps.FrontendAuthProvider == nil { + return false + } + return frontendAuthProviderWantsOmitBody(caps) +} + +func frontendAuthRequestBody(provider pluginapi.FrontendAuthProvider, body []byte) []byte { + if len(body) == 0 { + return nil + } + if _, isRPC := provider.(rpcFrontendAuthProvider); isRPC { + return body + } + return bytes.Clone(body) } func (a *accessAdapter) Identifier() (identifier string) { @@ -202,55 +226,59 @@ func (a *accessAdapter) Authenticate(ctx context.Context, r *http.Request) (resu } }() - body, errReadAll := readAndRestoreRequestBody(r) - if errReadAll != nil { - var maxBytesErr *http.MaxBytesError - if errors.As(errReadAll, &maxBytesErr) { - return nil, sdkaccess.NewRequestEntityTooLargeError(errReadAll) - } - return nil, sdkaccess.NewInternalAuthError("failed to read plugin auth request body", errReadAll) - } headers := cloneHeader(r.Header) + var authBody []byte var decodedBody []byte - cfg := a.host.currentRuntimeConfig() - // With image generation disabled the image handlers answer 404 before reading the body. - imagesOff := cfg != nil && cfg.DisableImageGeneration == config.DisableImageGenerationAll && strings.HasPrefix(r.URL.Path, "/v1/images/") - if encoding := strings.TrimSpace(headers.Get("Content-Encoding")); encoding != "" && !strings.EqualFold(encoding, "identity") && strings.Count(encoding, ",") < maxAuthContentCodings && !imagesOff && handlerDecodesBody(r) { - // Decoding runs before authentication, so it is bounded; a slot is held until the plugin call finishes. - var releaseSlot func() - select { - case authDecodeSlots <- struct{}{}: - ctx, releaseSlot = withPluginCallFinish(ctx, func() { <-authDecodeSlots }) - defer releaseSlot() - case <-ctx.Done(): - return nil, sdkaccess.NewInternalAuthError("plugin auth request canceled", ctx.Err()) - } - // The server middleware limits only compressed bytes; ReadRequestBody applies the configured - // limit to decoded bytes after authentication, so apply it here too. - limit := int64(maxAuthDecodedBodyBytes) - if cfg != nil && cfg.MaxRequestBodyBytes > 0 && cfg.MaxRequestBodyBytes < limit { - limit = cfg.MaxRequestBodyBytes - } - if (r.URL.Path == "/v1/audio/speech" || r.URL.Path == "/v1/tts") && limit > maxAuthSpeechBodyBytes { - limit = maxAuthSpeechBodyBytes - } - decoded, errDecode := decodeAuthBody(body, encoding, limit) - if errors.Is(errDecode, errAuthBodyTooLarge) { - return nil, sdkaccess.NewRequestEntityTooLargeError(errDecode) + if !a.omitBody { + body, errReadAll := readAndRestoreRequestBody(r) + if errReadAll != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(errReadAll, &maxBytesErr) { + return nil, sdkaccess.NewRequestEntityTooLargeError(errReadAll) + } + return nil, sdkaccess.NewInternalAuthError("failed to read plugin auth request body", errReadAll) } - // Any other failure leaves DecodedBody unset; the handler falls back to the raw body or rejects it. - if errDecode == nil { - decodedBody = decoded - } else { - releaseSlot() + cfg := a.host.currentRuntimeConfig() + // With image generation disabled the image handlers answer 404 before reading the body. + imagesOff := cfg != nil && cfg.DisableImageGeneration == config.DisableImageGenerationAll && strings.HasPrefix(r.URL.Path, "/v1/images/") + if encoding := strings.TrimSpace(headers.Get("Content-Encoding")); encoding != "" && !strings.EqualFold(encoding, "identity") && strings.Count(encoding, ",") < maxAuthContentCodings && !imagesOff && handlerDecodesBody(r) { + // Decoding runs before authentication, so it is bounded; a slot is held until the plugin call finishes. + var releaseSlot func() + select { + case authDecodeSlots <- struct{}{}: + ctx, releaseSlot = withPluginCallFinish(ctx, func() { <-authDecodeSlots }) + defer releaseSlot() + case <-ctx.Done(): + return nil, sdkaccess.NewInternalAuthError("plugin auth request canceled", ctx.Err()) + } + // The server middleware limits only compressed bytes; ReadRequestBody applies the configured + // limit to decoded bytes after authentication, so apply it here too. + limit := int64(maxAuthDecodedBodyBytes) + if cfg != nil && cfg.MaxRequestBodyBytes > 0 && cfg.MaxRequestBodyBytes < limit { + limit = cfg.MaxRequestBodyBytes + } + if (r.URL.Path == "/v1/audio/speech" || r.URL.Path == "/v1/tts") && limit > maxAuthSpeechBodyBytes { + limit = maxAuthSpeechBodyBytes + } + decoded, errDecode := decodeAuthBody(body, encoding, limit) + if errors.Is(errDecode, errAuthBodyTooLarge) { + return nil, sdkaccess.NewRequestEntityTooLargeError(errDecode) + } + // Any other failure leaves DecodedBody unset; the handler falls back to the raw body or rejects it. + if errDecode == nil { + decodedBody = decoded + } else { + releaseSlot() + } } + authBody = frontendAuthRequestBody(a.provider, body) } resp, errAuthenticate := a.provider.Authenticate(ctx, pluginapi.FrontendAuthRequest{ Method: r.Method, Path: r.URL.Path, Headers: headers, Query: cloneValues(r.URL.Query()), - Body: bytes.Clone(body), + Body: authBody, DecodedBody: decodedBody, }) if errAuthenticate != nil { diff --git a/internal/pluginhost/adapters_auth_body_test.go b/internal/pluginhost/adapters_auth_body_test.go index f8f34858812..96e06b5d89a 100644 --- a/internal/pluginhost/adapters_auth_body_test.go +++ b/internal/pluginhost/adapters_auth_body_test.go @@ -193,3 +193,224 @@ func TestFrontendAuthRequestDecodedBodyOverRPC(t *testing.T) { t.Fatalf("old plugin type failed on %s: %v", raw, errUnmarshal) } } + +type countingAuthBodyReader struct { + reader io.Reader + readBytes int +} + +func (r *countingAuthBodyReader) Read(p []byte) (int, error) { + n, errRead := r.reader.Read(p) + r.readBytes += n + return n, errRead +} + +type omitBodyFrontendAuthProvider struct { + frontendAuthProviderFunc +} + +func (omitBodyFrontendAuthProvider) FrontendAuthProviderWantsOmitBody() bool { + return true +} + +func TestAccessAdapterFrontendAuthOmitBodySkipsReadingAndDecodingBody(t *testing.T) { + const ( + omitKey = "plugin:auth-omit:quota" + aliasKey = "plugin:auth-alias:quota" + bodyKey = "plugin:auth-body:quota" + legacyKey = "plugin:auth-legacy:quota" + ) + for _, key := range []string{omitKey, aliasKey, bodyKey, legacyKey} { + sdkaccess.UnregisterProvider(key) + defer sdkaccess.UnregisterProvider(key) + } + + var omitSeen, aliasSeen, bodySeen, legacySeen pluginapi.FrontendAuthRequest + var omitCalls, aliasCalls, bodyCalls, legacyCalls int + newProvider := func(got *pluginapi.FrontendAuthRequest, calls *int) frontendAuthProviderFunc { + return frontendAuthProviderFunc{ + identifier: "quota", + authenticate: func(_ context.Context, req pluginapi.FrontendAuthRequest) (pluginapi.FrontendAuthResponse, error) { + *calls++ + *got = req + return pluginapi.FrontendAuthResponse{Authenticated: true, Principal: "user-1"}, nil + }, + } + } + + host := newHostWithRecords( + capabilityRecord{ + id: "auth-omit", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, + Capabilities: pluginapi.Capabilities{ + FrontendAuthProvider: newProvider(&omitSeen, &omitCalls), + FrontendAuthProviderOmitBody: true, + }, + }, + }, + capabilityRecord{ + id: "auth-alias", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, + Capabilities: pluginapi.Capabilities{ + FrontendAuthProvider: newProvider(&aliasSeen, &aliasCalls), + FrontendAuthOmitBody: true, + }, + }, + }, + capabilityRecord{ + id: "auth-body", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, + Capabilities: pluginapi.Capabilities{ + FrontendAuthProvider: newProvider(&bodySeen, &bodyCalls), + }, + }, + }, + capabilityRecord{ + id: "auth-legacy", + plugin: pluginapi.Plugin{ + SchemaVersion: 6, + Capabilities: pluginapi.Capabilities{ + FrontendAuthProvider: newProvider(&legacySeen, &legacyCalls), + FrontendAuthProviderOmitBody: true, + FrontendAuthOmitBody: true, + }, + }, + }, + ) + host.RegisterFrontendAuthProviders() + + findProvider := func(t *testing.T, key string) sdkaccess.Provider { + t.Helper() + for _, candidate := range sdkaccess.RegisteredProviders() { + if candidate.Identifier() == key { + return candidate + } + } + t.Fatalf("provider %q was not registered", key) + return nil + } + + plain := []byte(`{"model":"gpt-5","input":"large payload"}`) + compressed := zstdCompress(t, plain) + + for _, tc := range []struct { + name string + key string + seen *pluginapi.FrontendAuthRequest + call *int + }{ + {name: "FrontendAuthProviderOmitBody", key: omitKey, seen: &omitSeen, call: &omitCalls}, + {name: "FrontendAuthOmitBody", key: aliasKey, seen: &aliasSeen, call: &aliasCalls}, + } { + t.Run(tc.name, func(t *testing.T) { + savedSlots := authDecodeSlots + authDecodeSlots = make(chan struct{}, 1) + authDecodeSlots <- struct{}{} + defer func() { authDecodeSlots = savedSlots }() + + provider := findProvider(t, tc.key) + counter := &countingAuthBodyReader{reader: bytes.NewReader(compressed)} + req := httptest.NewRequest(http.MethodPost, "/v1/responses?stream=true", io.NopCloser(counter)) + req.Header.Set("Authorization", "Bearer sk-test") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Content-Encoding", "zstd") + + res, authErr := provider.Authenticate(context.Background(), req) + if authErr != nil { + t.Fatalf("Authenticate() error = %v, want nil", authErr) + } + if res == nil || res.Principal != "user-1" { + t.Fatalf("Authenticate() result = %#v, want principal user-1", res) + } + if *tc.call != 1 { + t.Fatalf("Authenticate() calls = %d, want 1", *tc.call) + } + if counter.readBytes != 0 { + t.Fatalf("bytes read from r.Body during Authenticate = %d, want 0", counter.readBytes) + } + if tc.seen.Body != nil || tc.seen.DecodedBody != nil { + t.Fatalf("seen Body=%q DecodedBody=%q, want both nil", tc.seen.Body, tc.seen.DecodedBody) + } + if tc.seen.Headers.Get("Authorization") != "Bearer sk-test" || tc.seen.Query.Get("stream") != "true" { + t.Fatalf("seen headers/query = %v / %v, want preserved metadata", tc.seen.Headers, tc.seen.Query) + } + rest, errReadAll := io.ReadAll(req.Body) + if errReadAll != nil || !bytes.Equal(rest, compressed) { + t.Fatalf("downstream body = %q (err=%v), want untouched compressed payload", rest, errReadAll) + } + if counter.readBytes != len(compressed) { + t.Fatalf("bytes read after downstream ReadAll = %d, want %d", counter.readBytes, len(compressed)) + } + + // A failing body reader must not cause Authenticate to fail when omitBody is active. + failReq, errNewReq := http.NewRequest(http.MethodPost, "http://example.test/v1/responses", failingReadCloser{}) + if errNewReq != nil { + t.Fatalf("NewRequest() error = %v", errNewReq) + } + if _, authErrFail := provider.Authenticate(context.Background(), failReq); authErrFail != nil { + t.Fatalf("Authenticate() with failing body reader error = %v, want nil", authErrFail) + } + }) + } + + for _, tc := range []struct { + name string + key string + seen *pluginapi.FrontendAuthRequest + call *int + }{ + {name: "schema 7 without opt-out", key: bodyKey, seen: &bodySeen, call: &bodyCalls}, + {name: "schema 6 legacy ignores opt-out", key: legacyKey, seen: &legacySeen, call: &legacyCalls}, + } { + t.Run(tc.name, func(t *testing.T) { + provider := findProvider(t, tc.key) + counter := &countingAuthBodyReader{reader: bytes.NewReader(compressed)} + req := httptest.NewRequest(http.MethodPost, "/v1/responses", io.NopCloser(counter)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Content-Encoding", "zstd") + + if _, authErr := provider.Authenticate(context.Background(), req); authErr != nil { + t.Fatalf("Authenticate() error = %v, want nil", authErr) + } + if *tc.call != 1 { + t.Fatalf("Authenticate() calls = %d, want 1", *tc.call) + } + if counter.readBytes != len(compressed) { + t.Fatalf("bytes read = %d, want %d", counter.readBytes, len(compressed)) + } + if !bytes.Equal(tc.seen.Body, compressed) || !bytes.Equal(tc.seen.DecodedBody, plain) { + t.Fatalf("seen Body=%q DecodedBody=%q, want preserved", tc.seen.Body, tc.seen.DecodedBody) + } + }) + } +} + +func TestFrontendAuthRequestBodySkipsCloneForRPCProvider(t *testing.T) { + body := []byte("original-request-body") + + rpcProvider := rpcFrontendAuthProvider{rpcPluginAdapter: &rpcPluginAdapter{client: &capturePluginClient{}}} + gotRPC := frontendAuthRequestBody(rpcProvider, body) + if len(gotRPC) != len(body) || &gotRPC[0] != &body[0] { + t.Fatal("frontendAuthRequestBody cloned body for rpcFrontendAuthProvider, want shared backing slice") + } + + inProcProvider := frontendAuthProviderFunc{identifier: "in-proc"} + gotInProc := frontendAuthRequestBody(inProcProvider, body) + if !bytes.Equal(gotInProc, body) { + t.Fatalf("frontendAuthRequestBody(inProc) = %q, want %q", gotInProc, body) + } + if &gotInProc[0] == &body[0] { + t.Fatal("frontendAuthRequestBody did not clone body for in-process provider") + } + gotInProc[0] = 'X' + if body[0] != 'o' { + t.Fatalf("mutating in-process request body mutated original buffer: %q", body) + } + + if gotNil := frontendAuthRequestBody(rpcProvider, nil); gotNil != nil { + t.Fatalf("frontendAuthRequestBody(nil) = %v, want nil", gotNil) + } +} diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index 46b719d0a26..593a0840f5e 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -229,10 +229,30 @@ func (h *Host) snapshotModelRegistrations() []pluginModelRegistration { } h.mu.Lock() defer h.mu.Unlock() - registrations := make([]pluginModelRegistration, 0, len(h.modelRegistrations)) - for _, registration := range h.modelRegistrations { - registration.models = cloneRegistryModels(registration.models) - registrations = append(registrations, registration) + return h.snapshotModelRegistrationsLocked() +} + +func (h *Host) snapshotModelRegistrationsLocked() []pluginModelRegistration { + if h == nil { + return nil + } + pluginIDs := make([]string, 0, len(h.modelRegistrations)+len(h.pluginModelRegistrations)) + seenPlugins := make(map[string]struct{}, len(h.modelRegistrations)+len(h.pluginModelRegistrations)) + for pluginID := range h.modelRegistrations { + seenPlugins[pluginID] = struct{}{} + pluginIDs = append(pluginIDs, pluginID) + } + for pluginID := range h.pluginModelRegistrations { + if _, exists := seenPlugins[pluginID]; !exists { + seenPlugins[pluginID] = struct{}{} + pluginIDs = append(pluginIDs, pluginID) + } + } + sort.Strings(pluginIDs) + + registrations := make([]pluginModelRegistration, 0, len(pluginIDs)) + for _, pluginID := range pluginIDs { + registrations = append(registrations, h.pluginRegistrationsLocked(pluginID)...) } sort.SliceStable(registrations, func(i, j int) bool { if registrations[i].priority == registrations[j].priority { @@ -243,6 +263,29 @@ func (h *Host) snapshotModelRegistrations() []pluginModelRegistration { return registrations } +func (h *Host) pluginRegistrationsLocked(pluginID string) []pluginModelRegistration { + regs := h.pluginModelRegistrations[pluginID] + if len(regs) == 1 { + if single, ok := h.modelRegistrations[pluginID]; ok && (single.pluginID != "" || single.provider != "" || len(single.models) > 0) { + single.models = cloneRegistryModels(single.models) + return []pluginModelRegistration{single} + } + } + if len(regs) > 0 { + out := make([]pluginModelRegistration, 0, len(regs)) + for _, reg := range regs { + reg.models = cloneRegistryModels(reg.models) + out = append(out, reg) + } + return out + } + if single, ok := h.modelRegistrations[pluginID]; ok && (single.pluginID != "" || single.provider != "" || len(single.models) > 0) { + single.models = cloneRegistryModels(single.models) + return []pluginModelRegistration{single} + } + return nil +} + func (h *Host) modelRegistration(pluginID string) pluginModelRegistration { if h == nil { return pluginModelRegistration{} @@ -384,6 +427,37 @@ func (h *Host) modelProvider(pluginID string) string { return h.modelProviders[pluginID] } +func (h *Host) modelProvidersForPlugin(pluginID string) []string { + if h == nil { + return nil + } + h.mu.Lock() + defer h.mu.Unlock() + list := h.modelProviderSets[pluginID] + primary := strings.TrimSpace(h.modelProviders[pluginID]) + if len(list) == 1 && primary != "" { + return []string{primary} + } + out := make([]string, 0, len(list)+1) + seen := make(map[string]struct{}, len(list)+1) + if primary != "" { + seen[primary] = struct{}{} + out = append(out, primary) + } + for _, item := range list { + item = strings.TrimSpace(item) + if item == "" { + continue + } + if _, exists := seen[item]; exists { + continue + } + seen[item] = struct{}{} + out = append(out, item) + } + return out +} + type executorAdapter struct { host *Host pluginID string diff --git a/internal/pluginhost/adapters_interceptors.go b/internal/pluginhost/adapters_interceptors.go index 657dd5fc115..a415118eaf7 100644 --- a/internal/pluginhost/adapters_interceptors.go +++ b/internal/pluginhost/adapters_interceptors.go @@ -3,10 +3,13 @@ package pluginhost import ( "bytes" "context" + "encoding/json" + "fmt" "io" "net/http" "net/url" "reflect" + "sort" "strings" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" @@ -35,6 +38,25 @@ func (h *Host) callRequestInterceptor(ctx context.Context, record capabilityReco return resp, true } +func (h *Host) callRequestBodyPatcher(ctx context.Context, record capabilityRecord, method string, patcher pluginapi.RequestBodyPatcher, body []byte) (out pluginapi.RequestBodyPatchResponse, ok bool) { + if h == nil || patcher == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return pluginapi.RequestBodyPatchResponse{}, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, method, recovered) + out = pluginapi.RequestBodyPatchResponse{} + ok = false + } + }() + resp, errPatch := patcher.PatchRequestBody(ctx, body) + if errPatch != nil { + log.Warnf("pluginhost: request body patcher %s failed: %v", record.id, errPatch) + return pluginapi.RequestBodyPatchResponse{}, false + } + return resp, true +} + func (h *Host) callResponseInterceptor(ctx context.Context, record capabilityRecord, interceptor pluginapi.ResponseInterceptor, req pluginapi.ResponseInterceptRequest) (out pluginapi.ResponseInterceptResponse, ok bool) { if h == nil || interceptor == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { return pluginapi.ResponseInterceptResponse{}, false @@ -114,18 +136,45 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc skipPluginID = strings.TrimSpace(skipPluginID) currentBase := req.Body bodyModified := false + patchMethod := "RequestBodyPatcher.PatchRequestBodyAfterAuth" + if beforeAuth { + patchMethod = "RequestBodyPatcher.PatchRequestBodyBeforeAuth" + } for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) || record.id == skipPluginID { + continue + } + patcher := requestRecordBodyPatcher(record, beforeAuth) interceptor := record.plugin.Capabilities.RequestInterceptor - if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { + if patcher == nil && interceptor == nil { continue } + if patcher != nil { + if len(currentBase) > 0 { + patchInput := requestPatcherInputBody(patcher, currentBase) + if patchResp, ok := h.callRequestBodyPatcher(ctx, record, patchMethod, patcher, patchInput); ok && len(patchResp.Edits) > 0 { + patched, modified, errPatch := applyRequestBodyEdits(currentBase, patchResp.Edits) + if errPatch != nil { + log.Warnf("pluginhost: request body patcher %s returned invalid edits: %v", record.id, errPatch) + } else if modified { + currentBase = patched + bodyModified = true + } + } + } + // When a plugin opts into raw body patching for this phase without also requesting + // metadata-only request interception (omit_body), skip the legacy full-body RPC call. + if interceptor == nil || !requestInterceptorOmitsBody(record, beforeAuth) { + continue + } + } nextReq := req nextReq.Headers = cloneHeader(current.Headers) if requestInterceptorOmitsBody(record, beforeAuth) { nextReq.Body = nil } else if len(currentBase) > 0 { - nextReq.Body = bytes.Clone(currentBase) + nextReq.Body = requestInterceptorInputBody(interceptor, currentBase) } else { nextReq.Body = nil } @@ -150,9 +199,20 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc return invoke(interceptor, callCtx, callReq) }, nextReq); ok { current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.ClearHeaders) > 0 { + current.ClearHeaders = append(current.ClearHeaders, resp.ClearHeaders...) + } if len(resp.Body) > 0 { - currentBase = bytes.Clone(resp.Body) + currentBase = requestInterceptorOutputBody(interceptor, resp.Body) bodyModified = true + } else if edits := requestInterceptResponseEdits(resp); len(edits) > 0 && len(currentBase) > 0 { + patched, modified, errPatch := applyRequestBodyEdits(currentBase, edits) + if errPatch != nil { + log.Warnf("pluginhost: request interceptor %s returned invalid edits: %v", record.id, errPatch) + } else if modified { + currentBase = patched + bodyModified = true + } } if strings.TrimSpace(resp.Path) != "" { current.Path = strings.TrimSpace(resp.Path) @@ -166,6 +226,19 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc } } } + if len(current.ClearHeaders) > 0 && len(current.Headers) > 0 { + filtered := current.ClearHeaders[:0] + for _, key := range current.ClearHeaders { + if current.Headers.Values(key) == nil { + filtered = append(filtered, key) + } + } + if len(filtered) == 0 { + current.ClearHeaders = nil + } else { + current.ClearHeaders = filtered + } + } if bodyModified { current.Body = currentBase } @@ -408,13 +481,172 @@ func (h *Host) HasRequestInterceptors() bool { if h.isPluginFused(record.id) { continue } - if record.plugin.Capabilities.RequestInterceptor != nil { + if recordHasRequestInterceptor(record) { return true } } return false } +func recordHasRequestInterceptor(record capabilityRecord) bool { + caps := record.plugin.Capabilities + if caps.RequestInterceptor != nil { + return true + } + return requestRecordBodyPatcher(record, true) != nil || requestRecordBodyPatcher(record, false) != nil +} + +type funcBodyPatcher func(context.Context, []byte) (pluginapi.RequestBodyPatchResponse, error) + +func (f funcBodyPatcher) PatchRequestBody(ctx context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + return f(ctx, body) +} + +func requestRecordBodyPatcher(record capabilityRecord, beforeAuth bool) pluginapi.RequestBodyPatcher { + if record.plugin.SchemaVersion < pluginabi.SchemaVersionRequestInterceptorBodyPatch { + return nil + } + caps := record.plugin.Capabilities + if beforeAuth { + if caps.RequestBodyPatcherBefore != nil { + return caps.RequestBodyPatcherBefore + } + if caps.RequestInterceptorPatchBodyBefore && caps.RequestInterceptor != nil { + if p, ok := caps.RequestInterceptor.(interface { + PatchRequestBodyBeforeAuth(context.Context, []byte) (pluginapi.RequestBodyPatchResponse, error) + }); ok && p != nil { + return funcBodyPatcher(p.PatchRequestBodyBeforeAuth) + } + if p, ok := caps.RequestInterceptor.(pluginapi.RequestBodyPatcher); ok && p != nil { + return p + } + } + return nil + } + if caps.RequestBodyPatcherAfter != nil { + return caps.RequestBodyPatcherAfter + } + if caps.RequestInterceptorPatchBodyAfter && caps.RequestInterceptor != nil { + if p, ok := caps.RequestInterceptor.(interface { + PatchRequestBodyAfterAuth(context.Context, []byte) (pluginapi.RequestBodyPatchResponse, error) + }); ok && p != nil { + return funcBodyPatcher(p.PatchRequestBodyAfterAuth) + } + if p, ok := caps.RequestInterceptor.(pluginapi.RequestBodyPatcher); ok && p != nil { + return p + } + } + return nil +} + +func requestPatcherInputBody(patcher pluginapi.RequestBodyPatcher, body []byte) []byte { + if _, isRPC := patcher.(rpcRequestBodyPatcher); isRPC { + return body + } + return bytes.Clone(body) +} + +func requestInterceptorInputBody(interceptor pluginapi.RequestInterceptor, body []byte) []byte { + if _, isRPC := interceptor.(*rpcPluginAdapter); isRPC { + return body + } + return bytes.Clone(body) +} + +func requestInterceptorOutputBody(interceptor pluginapi.RequestInterceptor, body []byte) []byte { + if _, isRPC := interceptor.(*rpcPluginAdapter); isRPC { + return body + } + return bytes.Clone(body) +} + +func requestInterceptResponseEdits(resp pluginapi.RequestInterceptResponse) []pluginapi.RequestBodyEdit { + if len(resp.BodyEdits) > 0 { + return resp.BodyEdits + } + return resp.Edits +} + +func requestBodyEditReplacementLen(edit pluginapi.RequestBodyEdit) int { + if len(edit.ReplacementBytes) > 0 { + return len(edit.ReplacementBytes) + } + return len(edit.Replacement) +} + +func appendRequestBodyEditReplacement(dst []byte, edit pluginapi.RequestBodyEdit) []byte { + if len(edit.ReplacementBytes) > 0 { + return append(dst, edit.ReplacementBytes...) + } + return append(dst, edit.Replacement...) +} + +func applyRequestBodyEdits(base []byte, edits []pluginapi.RequestBodyEdit) ([]byte, bool, error) { + if len(edits) == 0 { + return base, false, nil + } + type indexedEdit struct { + edit pluginapi.RequestBodyEdit + } + sorted := make([]indexedEdit, len(edits)) + for i, edit := range edits { + sorted[i] = indexedEdit{edit: edit} + } + if len(sorted) > 1 { + sort.SliceStable(sorted, func(i, j int) bool { + if sorted[i].edit.Start != sorted[j].edit.Start { + return sorted[i].edit.Start < sorted[j].edit.Start + } + return sorted[i].edit.End < sorted[j].edit.End + }) + } + baseLen := len(base) + outLen := baseLen + prevEnd := 0 + anyChange := false + for i, item := range sorted { + e := item.edit + if e.Start < 0 || e.End < e.Start || e.End > baseLen { + return base, false, fmt.Errorf("edit[%d] range [%d:%d] out of bounds for body length %d", i, e.Start, e.End, baseLen) + } + if i > 0 && e.Start < prevEnd { + return base, false, fmt.Errorf("edit[%d] range [%d:%d] overlaps previous end %d", i, e.Start, e.End, prevEnd) + } + replLen := requestBodyEditReplacementLen(e) + if e.Start != e.End || replLen > 0 { + if len(e.ReplacementBytes) > 0 { + if !bytes.Equal(base[e.Start:e.End], e.ReplacementBytes) { + anyChange = true + } + } else if string(base[e.Start:e.End]) != e.Replacement { + anyChange = true + } + } + outLen += replLen - (e.End - e.Start) + prevEnd = e.End + } + if !anyChange { + return base, false, nil + } + out := make([]byte, 0, outLen) + cursor := 0 + for _, item := range sorted { + e := item.edit + if e.Start > cursor { + out = append(out, base[cursor:e.Start]...) + } + out = appendRequestBodyEditReplacement(out, e) + cursor = e.End + } + if cursor < baseLen { + out = append(out, base[cursor:]...) + } + if !json.Valid(out) { + return base, false, fmt.Errorf("patched request body is not valid JSON") + } + return out, true, nil +} + func readAndRestoreRequestBody(r *http.Request) ([]byte, error) { if r == nil || r.Body == nil { return nil, nil diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index e62f64dbee4..db951931ad1 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -4522,3 +4522,339 @@ func TestUsageAdapterPropagatesRequestIDTraceIDAndSelectionID(t *testing.T) { t.Fatalf("captured.SelectionID from context fallback = %q, want %q", captured.SelectionID, "sel-from-ctx") } } + +type funcRequestBodyPatcher func(context.Context, []byte) (pluginapi.RequestBodyPatchResponse, error) + +func (f funcRequestBodyPatcher) PatchRequestBody(ctx context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + return f(ctx, body) +} + +func TestInterceptRequestBodyPatchValidationAndSequentialChaining(t *testing.T) { + rawInput := []byte(`{"model":"gpt-5","input":[{"role":"developer","content":[{"type":"input_text","text":"\n UTC\n 2026-04-19\n"}]}]}`) + + // Plugin 1: replaces UTC -> Asia/Shanghai (+10 bytes), + // shifting the byte offset of 2026-04-19 in the resulting body. + plugin1 := normalizeTestCapabilityRecord(capabilityRecord{ + id: "01-tz-plugin", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionRequestInterceptorBodyPatch, + Capabilities: pluginapi.Capabilities{ + RequestBodyPatcherBefore: funcRequestBodyPatcher(func(_ context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + target := []byte("UTC") + idx := bytes.Index(body, target) + if idx < 0 { + return pluginapi.RequestBodyPatchResponse{}, nil + } + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{ + { + Start: idx, + End: idx + len(target), + Replacement: "Asia/Shanghai", + }, + }, + }, nil + }), + }, + }, + }) + + // Plugin 2: receives Plugin 1's patched body and returns byte-range edits relative to the body it received. + plugin2 := normalizeTestCapabilityRecord(capabilityRecord{ + id: "02-date-plugin", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionRequestInterceptorBodyPatch, + Capabilities: pluginapi.Capabilities{ + RequestBodyPatcherBefore: funcRequestBodyPatcher(func(_ context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + if !bytes.Contains(body, []byte("Asia/Shanghai")) { + t.Fatalf("plugin2 did not receive plugin1's patched body: %s", body) + } + target := []byte("2026-04-19") + idx := bytes.Index(body, target) + if idx < 0 { + return pluginapi.RequestBodyPatchResponse{}, nil + } + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{ + { + Start: idx, + End: idx + len(target), + Replacement: "2026-04-20", + }, + }, + }, nil + }), + }, + }, + }) + + // Plugin 3: returns invalid/overlapping/out-of-bounds edits or invalid JSON; host must reject and preserve Plugin 2's valid body. + var invalidCase int + plugin3 := normalizeTestCapabilityRecord(capabilityRecord{ + id: "03-invalid-edits-plugin", + plugin: pluginapi.Plugin{ + SchemaVersion: pluginabi.SchemaVersionRequestInterceptorBodyPatch, + Capabilities: pluginapi.Capabilities{ + RequestBodyPatcherBefore: funcRequestBodyPatcher(func(_ context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + switch invalidCase { + case 0: // Out of bounds + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{{Start: 0, End: len(body) + 5, Replacement: "{}"}}, + }, nil + case 1: // Overlapping ranges + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{ + {Start: 2, End: 8, Replacement: `"a"`}, + {Start: 5, End: 10, Replacement: `"b"`}, + }, + }, nil + case 2: // Invalid JSON after splice + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{{Start: 0, End: 1, Replacement: "not-json"}}, + }, nil + default: + return pluginapi.RequestBodyPatchResponse{}, nil + } + }), + }, + }, + }) + + host := newHostWithRecords(plugin1, plugin2, plugin3) + wantChained := `{"model":"gpt-5","input":[{"role":"developer","content":[{"type":"input_text","text":"\n Asia/Shanghai\n 2026-04-20\n"}]}]}` + + for invalidCase = 0; invalidCase <= 2; invalidCase++ { + resp := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{ + Model: "gpt-5", + Body: rawInput, + }) + if string(resp.Body) != wantChained { + t.Fatalf("case %d: chained patched body = %s, want %s", invalidCase, resp.Body, wantChained) + } + } +} + +func TestRegisterModelsMultiProviderRegistrationsAndUnload(t *testing.T) { + modelRegistry := newFakeModelRegistry() + manager := newFakeExecutorManager() + host := New() + t.Cleanup(host.ShutdownAll) + + host.RegisterPluginForTest("catalog-multi", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + ModelRegistrar: modelRegistrarFunc(func(ctx context.Context, req pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) { + if !req.MultiProviderModels { + t.Fatalf("req.MultiProviderModels = false, want true") + } + return pluginapi.ModelRegistrationResponse{ + Provider: "compat-direct", + Models: []pluginapi.ModelInfo{{ID: "legacy-fallback-model"}}, + Registrations: []pluginapi.ProviderModels{ + { + Provider: " Compat-Direct ", + Models: []pluginapi.ModelInfo{{ID: " model-direct "}}, + }, + { + Provider: " Compat-VPN ", + Models: []pluginapi.ModelInfo{{ID: " model-vpn "}}, + }, + }, + }, nil + }), + }, + }) + host.RegisterPluginForTest("auth-multi", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + ModelProvider: modelProviderFunc{ + staticModels: func(ctx context.Context, req pluginapi.StaticModelRequest) (pluginapi.ModelResponse, error) { + if !req.MultiProviderModels { + t.Fatalf("req.MultiProviderModels = false, want true") + } + return pluginapi.ModelResponse{ + Registrations: []pluginapi.ProviderModels{ + { + Provider: " Compat-Auth-Direct ", + Models: []pluginapi.ModelInfo{{ID: " static-auth-direct "}}, + }, + { + Provider: " Compat-Auth-VPN ", + Models: []pluginapi.ModelInfo{{ID: " static-auth-vpn "}}, + }, + }, + }, nil + }, + modelsForAuth: func(ctx context.Context, req pluginapi.AuthModelRequest) (pluginapi.ModelResponse, error) { + switch strings.ToLower(strings.TrimSpace(req.AuthProvider)) { + case "compat-auth-direct": + return pluginapi.ModelResponse{ + Provider: "compat-auth-direct", + Models: []pluginapi.ModelInfo{{ID: "auth-model-direct"}}, + }, nil + case "compat-auth-vpn": + return pluginapi.ModelResponse{ + Provider: "compat-auth-vpn", + Models: []pluginapi.ModelInfo{{ID: "auth-model-vpn"}}, + }, nil + default: + return pluginapi.ModelResponse{}, nil + } + }, + }, + }, + }) + + host.RegisterModels(context.Background(), modelRegistry) + host.RegisterExecutors(manager, modelRegistry) + + if models := host.ModelsForProvider("compat-direct"); len(models) != 1 || models[0].ID != "model-direct" { + t.Fatalf("compat-direct models = %#v, want [model-direct]", models) + } + if models := host.ModelsForProvider("compat-vpn"); len(models) != 1 || models[0].ID != "model-vpn" { + t.Fatalf("compat-vpn models = %#v, want [model-vpn]", models) + } + + if !host.HasAuthModelProvider("compat-auth-direct") || !host.HasAuthModelProvider("compat-auth-vpn") { + t.Fatalf("HasAuthModelProvider = (%v, %v), want (true, true)", host.HasAuthModelProvider("compat-auth-direct"), host.HasAuthModelProvider("compat-auth-vpn")) + } + directAuthRes := host.ModelsForAuth(context.Background(), &coreauth.Auth{ID: "auth-1", Provider: "compat-auth-direct"}) + if !directAuthRes.Handled || len(directAuthRes.Models) != 1 || directAuthRes.Models[0].ID != "auth-model-direct" { + t.Fatalf("ModelsForAuth(compat-auth-direct) = %#v, want handled with auth-model-direct", directAuthRes) + } + vpnAuthRes := host.ModelsForAuth(context.Background(), &coreauth.Auth{ID: "auth-2", Provider: "compat-auth-vpn"}) + if !vpnAuthRes.Handled || len(vpnAuthRes.Models) != 1 || vpnAuthRes.Models[0].ID != "auth-model-vpn" { + t.Fatalf("ModelsForAuth(compat-auth-vpn) = %#v, want handled with auth-model-vpn", vpnAuthRes) + } + + if !host.UnloadPlugin("catalog-multi") { + t.Fatal("UnloadPlugin(catalog-multi) = false, want true") + } + if models := host.ModelsForProvider("compat-direct"); len(models) != 0 { + t.Fatalf("compat-direct models after unload = %#v, want empty", models) + } + if models := host.ModelsForProvider("compat-vpn"); len(models) != 0 { + t.Fatalf("compat-vpn models after unload = %#v, want empty", models) + } + + if !host.UnloadPlugin("auth-multi") { + t.Fatal("UnloadPlugin(auth-multi) = false, want true") + } + if host.HasAuthModelProvider("compat-auth-direct") || host.HasAuthModelProvider("compat-auth-vpn") { + t.Fatalf("HasAuthModelProvider after unload = (%v, %v), want (false, false)", host.HasAuthModelProvider("compat-auth-direct"), host.HasAuthModelProvider("compat-auth-vpn")) + } +} + +func TestRegisterModelsLegacyEmptyStaticModelsPreservesSingleProviderBehavior(t *testing.T) { + modelRegistry := newFakeModelRegistry() + manager := newFakeExecutorManager() + registrarCalled := false + host := newHostWithRecords( + capabilityRecord{ + id: "dual-empty-static", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + ModelProvider: modelProviderFunc{ + staticModels: func(context.Context, pluginapi.StaticModelRequest) (pluginapi.ModelResponse, error) { + return pluginapi.ModelResponse{}, nil + }, + modelsForAuth: func(context.Context, pluginapi.AuthModelRequest) (pluginapi.ModelResponse, error) { + return pluginapi.ModelResponse{ + Provider: "legacy-provider", + Models: []pluginapi.ModelInfo{{ID: "unexpected"}}, + }, nil + }, + }, + ModelRegistrar: modelRegistrarFunc(func(context.Context, pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) { + registrarCalled = true + return pluginapi.ModelRegistrationResponse{ + Provider: "legacy-provider", + Models: []pluginapi.ModelInfo{{ID: "legacy-model"}}, + }, nil + }), + }}, + }, + capabilityRecord{ + id: "empty-models-provider", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + ModelProvider: modelProviderFunc{ + staticModels: func(context.Context, pluginapi.StaticModelRequest) (pluginapi.ModelResponse, error) { + return pluginapi.ModelResponse{Provider: "gemini"}, nil + }, + modelsForAuth: func(context.Context, pluginapi.AuthModelRequest) (pluginapi.ModelResponse, error) { + return pluginapi.ModelResponse{Provider: "gemini"}, nil + }, + }, + }}, + }, + ) + + host.RegisterModels(context.Background(), modelRegistry) + host.RegisterExecutors(manager, modelRegistry) + + if registrarCalled { + t.Fatal("ModelRegistrar.RegisterModels was called despite ModelProvider being present") + } + if models := host.ModelsForProvider("legacy-provider"); len(models) != 0 { + t.Fatalf("legacy-provider models = %#v, want empty when StaticModels returns empty", models) + } + if host.HasAuthModelProvider("legacy-provider") { + t.Fatal("HasAuthModelProvider(legacy-provider) = true, want false when StaticModels returns empty") + } + if host.HasAuthModelProvider("gemini") { + t.Fatal("HasAuthModelProvider(gemini) = true, want false when StaticModels returns empty models") + } + if res := host.ModelsForAuth(context.Background(), &coreauth.Auth{ID: "g1", Provider: "gemini"}); res.Handled { + t.Fatalf("ModelsForAuth(gemini) = %#v, want unhandled when StaticModels had no models", res) + } +} + +func TestRegisterModelsStaticModelsRegistrationsMergeAndSharedProviderUnload(t *testing.T) { + modelRegistry := newFakeModelRegistry() + manager := newFakeExecutorManager() + host := New() + t.Cleanup(host.ShutdownAll) + + host.RegisterPluginForTest("alpha", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + ModelProvider: modelProviderFunc{ + staticModels: func(ctx context.Context, req pluginapi.StaticModelRequest) (pluginapi.ModelResponse, error) { + if !req.MultiProviderModels { + t.Fatalf("req.MultiProviderModels = false, want true") + } + return pluginapi.ModelResponse{ + Registrations: []pluginapi.ProviderModels{ + {Provider: "shared-provider", Models: []pluginapi.ModelInfo{{ID: "alpha-1"}}}, + {Provider: "shared-provider", Models: []pluginapi.ModelInfo{{ID: "alpha-1"}, {ID: "alpha-2"}}}, + {Provider: "alpha-only", Models: []pluginapi.ModelInfo{{ID: "alpha-only-1"}}}, + }, + }, nil + }, + }, + }, + }) + host.RegisterPluginForTest("bravo", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + ModelRegistrar: staticModelRegistrar("shared-provider", "bravo-1"), + }, + }) + + host.RegisterModels(context.Background(), modelRegistry) + host.RegisterExecutors(manager, modelRegistry) + + shared := host.ModelsForProvider("shared-provider") + if len(shared) != 3 || shared[0].ID != "alpha-1" || shared[1].ID != "alpha-2" || shared[2].ID != "bravo-1" { + t.Fatalf("shared-provider models = %#v, want [alpha-1 alpha-2 bravo-1]", shared) + } + if alphaOnly := host.ModelsForProvider("alpha-only"); len(alphaOnly) != 1 || alphaOnly[0].ID != "alpha-only-1" { + t.Fatalf("alpha-only models = %#v, want [alpha-only-1]", alphaOnly) + } + + if !host.UnloadPlugin("alpha") { + t.Fatal("UnloadPlugin(alpha) = false, want true") + } + sharedAfterUnload := host.ModelsForProvider("shared-provider") + if len(sharedAfterUnload) != 1 || sharedAfterUnload[0].ID != "bravo-1" { + t.Fatalf("shared-provider models after unloading alpha = %#v, want [bravo-1]", sharedAfterUnload) + } + if alphaOnlyAfterUnload := host.ModelsForProvider("alpha-only"); len(alphaOnlyAfterUnload) != 0 { + t.Fatalf("alpha-only models after unloading alpha = %#v, want empty", alphaOnlyAfterUnload) + } +} diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index d48e543fd2c..d1c8b5420d4 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -61,69 +61,73 @@ type pluginLoadResult struct { } type Host struct { - applyMu chan struct{} - mu sync.Mutex - loader pluginLoader - loaded map[string]*loadedPlugin - retired map[string][]*loadedPlugin - loading map[string]*pluginLoadRequest - fused map[string]string - pluginFileVersions map[string]string - activePluginVersions map[string]string - activePluginPaths map[string]string - cleanupFilesPending bool - runtimeConfig *config.Config - authManager *coreauth.Manager - modelExecutor modelExecutor - modelClientIDs map[string]struct{} // Reserved for future use; catalog-only plugins do not populate it. - executorModelClientIDs map[string]struct{} - modelProviders map[string]string - modelRegistrations map[string]pluginModelRegistration - providerModels map[string][]*registryModelInfo - executorProviders map[string]struct{} - accessProviderKeys map[string]struct{} - commandLineFlags map[string]commandLineFlagRecord - commandLineHits map[string]struct{} - managementRoutes map[string]managementRouteRecord - resourceRoutes map[string]resourceRouteRecord - streams *streamBridge - httpStreams *hostHTTPStreamBridge - httpOperations *hostHTTPOperationBridge - modelStreams *modelStreamBridge - callbackContexts *callbackContextRegistry - snapshot atomic.Value + applyMu chan struct{} + mu sync.Mutex + loader pluginLoader + loaded map[string]*loadedPlugin + retired map[string][]*loadedPlugin + loading map[string]*pluginLoadRequest + fused map[string]string + pluginFileVersions map[string]string + activePluginVersions map[string]string + activePluginPaths map[string]string + cleanupFilesPending bool + runtimeConfig *config.Config + authManager *coreauth.Manager + modelExecutor modelExecutor + modelClientIDs map[string]struct{} // Reserved for future use; catalog-only plugins do not populate it. + executorModelClientIDs map[string]struct{} + modelProviders map[string]string + modelProviderSets map[string][]string + modelRegistrations map[string]pluginModelRegistration + pluginModelRegistrations map[string][]pluginModelRegistration + providerModels map[string][]*registryModelInfo + executorProviders map[string]struct{} + accessProviderKeys map[string]struct{} + commandLineFlags map[string]commandLineFlagRecord + commandLineHits map[string]struct{} + managementRoutes map[string]managementRouteRecord + resourceRoutes map[string]resourceRouteRecord + streams *streamBridge + httpStreams *hostHTTPStreamBridge + httpOperations *hostHTTPOperationBridge + modelStreams *modelStreamBridge + callbackContexts *callbackContextRegistry + snapshot atomic.Value // loadCleanupDoneHook, set by tests, runs after a late load cleanup dropped its load token. loadCleanupDoneHook func() } func New() *Host { h := &Host{ - applyMu: make(chan struct{}, 1), - loader: defaultPluginLoader(), - loaded: make(map[string]*loadedPlugin), - retired: make(map[string][]*loadedPlugin), - loading: make(map[string]*pluginLoadRequest), - fused: make(map[string]string), - pluginFileVersions: make(map[string]string), - activePluginVersions: make(map[string]string), - activePluginPaths: make(map[string]string), - cleanupFilesPending: true, - modelClientIDs: make(map[string]struct{}), - executorModelClientIDs: make(map[string]struct{}), - modelProviders: make(map[string]string), - modelRegistrations: make(map[string]pluginModelRegistration), - providerModels: make(map[string][]*registryModelInfo), - executorProviders: make(map[string]struct{}), - accessProviderKeys: make(map[string]struct{}), - commandLineFlags: make(map[string]commandLineFlagRecord), - commandLineHits: make(map[string]struct{}), - managementRoutes: make(map[string]managementRouteRecord), - resourceRoutes: make(map[string]resourceRouteRecord), - streams: newStreamBridge(), - httpStreams: newHostHTTPStreamBridge(), - httpOperations: newHostHTTPOperationBridge(), - modelStreams: newModelStreamBridge(), - callbackContexts: newCallbackContextRegistry(), + applyMu: make(chan struct{}, 1), + loader: defaultPluginLoader(), + loaded: make(map[string]*loadedPlugin), + retired: make(map[string][]*loadedPlugin), + loading: make(map[string]*pluginLoadRequest), + fused: make(map[string]string), + pluginFileVersions: make(map[string]string), + activePluginVersions: make(map[string]string), + activePluginPaths: make(map[string]string), + cleanupFilesPending: true, + modelClientIDs: make(map[string]struct{}), + executorModelClientIDs: make(map[string]struct{}), + modelProviders: make(map[string]string), + modelProviderSets: make(map[string][]string), + modelRegistrations: make(map[string]pluginModelRegistration), + pluginModelRegistrations: make(map[string][]pluginModelRegistration), + providerModels: make(map[string][]*registryModelInfo), + executorProviders: make(map[string]struct{}), + accessProviderKeys: make(map[string]struct{}), + commandLineFlags: make(map[string]commandLineFlagRecord), + commandLineHits: make(map[string]struct{}), + managementRoutes: make(map[string]managementRouteRecord), + resourceRoutes: make(map[string]resourceRouteRecord), + streams: newStreamBridge(), + httpStreams: newHostHTTPStreamBridge(), + httpOperations: newHostHTTPOperationBridge(), + modelStreams: newModelStreamBridge(), + callbackContexts: newCallbackContextRegistry(), } h.snapshot.Store(emptySnapshot()) return h @@ -841,7 +845,9 @@ func (h *Host) ShutdownAllContext(ctx context.Context) { h.modelClientIDs = make(map[string]struct{}) h.executorModelClientIDs = make(map[string]struct{}) h.modelProviders = make(map[string]string) + h.modelProviderSets = make(map[string][]string) h.modelRegistrations = make(map[string]pluginModelRegistration) + h.pluginModelRegistrations = make(map[string][]pluginModelRegistration) h.providerModels = make(map[string][]*registryModelInfo) h.executorProviders = make(map[string]struct{}) h.commandLineFlags = make(map[string]commandLineFlagRecord) @@ -992,11 +998,40 @@ func (h *Host) removePluginRuntimeStateLocked(id string) { delete(h.commandLineHits, name) } } - if registration, ok := h.modelRegistrations[id]; ok { - delete(h.providerModels, registration.provider) + affectedProviders := make(map[string]struct{}) + if registration, ok := h.modelRegistrations[id]; ok && registration.provider != "" { + affectedProviders[registration.provider] = struct{}{} + } + for _, registration := range h.pluginModelRegistrations[id] { + if registration.provider != "" { + affectedProviders[registration.provider] = struct{}{} + } + } + if provider := normalizeProviderID(h.modelProviders[id]); provider != "" { + affectedProviders[provider] = struct{}{} + } + for _, provider := range h.modelProviderSets[id] { + if normalized := normalizeProviderID(provider); normalized != "" { + affectedProviders[normalized] = struct{}{} + } } delete(h.modelProviders, id) + delete(h.modelProviderSets, id) delete(h.modelRegistrations, id) + delete(h.pluginModelRegistrations, id) + for provider := range affectedProviders { + delete(h.providerModels, provider) + } + if len(affectedProviders) > 0 { + for _, registration := range h.snapshotModelRegistrationsLocked() { + if registration.hasExecutor { + continue + } + if _, ok := affectedProviders[registration.provider]; ok { + appendModelsForProvider(h.providerModels, registration.provider, registration.models) + } + } + } } func (h *Host) rebuildActivePluginMapsLocked(records []capabilityRecord) { @@ -1213,6 +1248,8 @@ func validPlugin(plugin pluginapi.Plugin) bool { caps.RequestTranslator != nil || caps.RequestNormalizer != nil || caps.RequestInterceptor != nil || + caps.RequestBodyPatcherBefore != nil || + caps.RequestBodyPatcherAfter != nil || caps.RequestLifecyclePlugin != nil || caps.ResponseTranslator != nil || caps.ResponseBeforeTranslator != nil || diff --git a/internal/pluginhost/quota_provider.go b/internal/pluginhost/quota_provider.go index f366560d12b..a596dc7d12c 100644 --- a/internal/pluginhost/quota_provider.go +++ b/internal/pluginhost/quota_provider.go @@ -24,6 +24,14 @@ func (h *Host) RegisterPluginForTest(id string, plugin pluginapi.Plugin) { } records := append(h.activeRecords(), record) h.mu.Lock() + h.loaded[id] = &loadedPlugin{ + id: id, + name: plugin.Metadata.Name, + path: record.path, + version: record.version, + plugin: plugin, + registered: true, + } h.rebuildActivePluginMapsLocked(records) h.snapshot.Store(&Snapshot{enabled: true, records: records, quotaSupportedProviders: make(map[string][]string)}) h.mu.Unlock() diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 4aaf3c66413..495bb63f72f 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -86,6 +86,15 @@ type rpcResponseNormalizer struct { method string } +type rpcRequestBodyPatcher struct { + *rpcPluginAdapter + method string +} + +func (p rpcRequestBodyPatcher) PatchRequestBody(ctx context.Context, body []byte) (pluginapi.RequestBodyPatchResponse, error) { + return callPluginRaw[pluginapi.RequestBodyPatchResponse](ctx, p.client, p.method, body) +} + func registerRPCPlugin(ctx context.Context, host *Host, id string, client pluginClient, method string, configYAML []byte) (pluginapi.Plugin, error) { if client == nil { return pluginapi.Plugin{}, fmt.Errorf("plugin client is nil") @@ -136,7 +145,10 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin } } if resp.Capabilities.FrontendAuthProvider { + omitBody := resp.Capabilities.FrontendAuthProviderOmitBody || resp.Capabilities.FrontendAuthOmitBody plugin.Capabilities.FrontendAuthProvider = rpcFrontendAuthProvider{rpcPluginAdapter: adapter} + plugin.Capabilities.FrontendAuthProviderOmitBody = omitBody + plugin.Capabilities.FrontendAuthOmitBody = omitBody } if resp.Capabilities.Scheduler { plugin.Capabilities.Scheduler = adapter @@ -159,6 +171,22 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin plugin.Capabilities.RequestInterceptorOmitBodyBefore = resp.Capabilities.RequestInterceptorOmitBodyBefore plugin.Capabilities.RequestInterceptorOmitBodyAfter = resp.Capabilities.RequestInterceptorOmitBodyAfter } + if schemaVersion >= pluginabi.SchemaVersionRequestInterceptorBodyPatch { + if resp.Capabilities.RequestBodyPatcherBefore || resp.Capabilities.RequestInterceptorPatchBodyBefore { + plugin.Capabilities.RequestBodyPatcherBefore = rpcRequestBodyPatcher{ + rpcPluginAdapter: adapter, + method: pluginabi.MethodRequestPatchBodyBefore, + } + plugin.Capabilities.RequestInterceptorPatchBodyBefore = true + } + if resp.Capabilities.RequestBodyPatcherAfter || resp.Capabilities.RequestInterceptorPatchBodyAfter { + plugin.Capabilities.RequestBodyPatcherAfter = rpcRequestBodyPatcher{ + rpcPluginAdapter: adapter, + method: pluginabi.MethodRequestPatchBodyAfter, + } + plugin.Capabilities.RequestInterceptorPatchBodyAfter = true + } + } if resp.Capabilities.RequestLifecyclePlugin { plugin.Capabilities.RequestLifecyclePlugin = adapter } @@ -214,6 +242,11 @@ func callPlugin[T any](ctx context.Context, client pluginClient, method string, if errMarshal != nil { return zero, fmt.Errorf("marshal plugin request %s: %w", method, errMarshal) } + return callPluginRaw[T](ctx, client, method, rawRequest) +} + +func callPluginRaw[T any](ctx context.Context, client pluginClient, method string, rawRequest []byte) (T, error) { + var zero T rawResp, errCall := client.Call(ctx, method, rawRequest) if errCall != nil { return zero, errCall diff --git a/internal/pluginhost/rpc_schema.go b/internal/pluginhost/rpc_schema.go index f2887029b5a..ca56a965595 100644 --- a/internal/pluginhost/rpc_schema.go +++ b/internal/pluginhost/rpc_schema.go @@ -19,35 +19,41 @@ type rpcRegistration struct { } type rpcCapabilities struct { - ModelRegistrar bool `json:"model_registrar"` - ModelProvider bool `json:"model_provider"` - AuthProvider bool `json:"auth_provider"` - FrontendAuthProvider bool `json:"frontend_auth_provider"` - FrontendAuthProviderExclusive bool `json:"frontend_auth_provider_exclusive"` - Scheduler bool `json:"scheduler"` - SchedulerAcrossPriorities bool `json:"scheduler_across_priorities,omitempty"` - ModelRouter bool `json:"model_router"` - Executor bool `json:"executor"` - ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"` - ExecutorInputFormats []string `json:"executor_input_formats,omitempty"` - ExecutorOutputFormats []string `json:"executor_output_formats,omitempty"` - RequestTranslator bool `json:"request_translator"` - RequestNormalizer bool `json:"request_normalizer"` - RequestInterceptor bool `json:"request_interceptor"` - RequestInterceptorOmitBodyBefore bool `json:"request_interceptor_omit_body_before,omitempty"` - RequestInterceptorOmitBodyAfter bool `json:"request_interceptor_omit_body_after,omitempty"` - RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` - ResponseTranslator bool `json:"response_translator"` - ResponseBeforeTranslator bool `json:"response_before_translator"` - ResponseAfterTranslator bool `json:"response_after_translator"` - ResponseInterceptor bool `json:"response_interceptor"` - StreamChunkInterceptor bool `json:"response_stream_interceptor"` - WebSocketResponseObserver bool `json:"websocket_response_observer"` - ThinkingApplier bool `json:"thinking_applier"` - UsagePlugin bool `json:"usage_plugin"` - CommandLinePlugin bool `json:"command_line_plugin"` - ManagementAPI bool `json:"management_api"` - QuotaProvider bool `json:"quota_provider"` + ModelRegistrar bool `json:"model_registrar"` + ModelProvider bool `json:"model_provider"` + AuthProvider bool `json:"auth_provider"` + FrontendAuthProvider bool `json:"frontend_auth_provider"` + FrontendAuthProviderExclusive bool `json:"frontend_auth_provider_exclusive"` + FrontendAuthProviderOmitBody bool `json:"frontend_auth_provider_omit_body,omitempty"` + FrontendAuthOmitBody bool `json:"frontend_auth_omit_body,omitempty"` + Scheduler bool `json:"scheduler"` + SchedulerAcrossPriorities bool `json:"scheduler_across_priorities,omitempty"` + ModelRouter bool `json:"model_router"` + Executor bool `json:"executor"` + ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"` + ExecutorInputFormats []string `json:"executor_input_formats,omitempty"` + ExecutorOutputFormats []string `json:"executor_output_formats,omitempty"` + RequestTranslator bool `json:"request_translator"` + RequestNormalizer bool `json:"request_normalizer"` + RequestInterceptor bool `json:"request_interceptor"` + RequestInterceptorOmitBodyBefore bool `json:"request_interceptor_omit_body_before,omitempty"` + RequestInterceptorOmitBodyAfter bool `json:"request_interceptor_omit_body_after,omitempty"` + RequestBodyPatcherBefore bool `json:"request_body_patcher_before,omitempty"` + RequestBodyPatcherAfter bool `json:"request_body_patcher_after,omitempty"` + RequestInterceptorPatchBodyBefore bool `json:"request_interceptor_patch_body_before,omitempty"` + RequestInterceptorPatchBodyAfter bool `json:"request_interceptor_patch_body_after,omitempty"` + RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` + ResponseTranslator bool `json:"response_translator"` + ResponseBeforeTranslator bool `json:"response_before_translator"` + ResponseAfterTranslator bool `json:"response_after_translator"` + ResponseInterceptor bool `json:"response_interceptor"` + StreamChunkInterceptor bool `json:"response_stream_interceptor"` + WebSocketResponseObserver bool `json:"websocket_response_observer"` + ThinkingApplier bool `json:"thinking_applier"` + UsagePlugin bool `json:"usage_plugin"` + CommandLinePlugin bool `json:"command_line_plugin"` + ManagementAPI bool `json:"management_api"` + QuotaProvider bool `json:"quota_provider"` } type rpcIdentifierResponse struct { @@ -147,6 +153,22 @@ type rpcQuotaResetRequest struct { type rpcEmptyResponse struct{} +func frontendAuthProviderWantsOmitBody(caps pluginapi.Capabilities) bool { + if caps.FrontendAuthProvider == nil { + return false + } + if caps.FrontendAuthProviderOmitBody || caps.FrontendAuthOmitBody { + return true + } + if opt, ok := caps.FrontendAuthProvider.(interface{ FrontendAuthProviderWantsOmitBody() bool }); ok && opt != nil { + return opt.FrontendAuthProviderWantsOmitBody() + } + if opt, ok := caps.FrontendAuthProvider.(interface{ FrontendAuthWantsOmitBody() bool }); ok && opt != nil { + return opt.FrontendAuthWantsOmitBody() + } + return false +} + func requestInterceptorWantsOmitBody(caps pluginapi.Capabilities, before bool) bool { if caps.RequestInterceptor == nil { return false @@ -163,6 +185,22 @@ func requestInterceptorWantsOmitBody(caps pluginapi.Capabilities, before bool) b return false } +func requestInterceptorWantsPatchBody(caps pluginapi.Capabilities, before bool) bool { + if before { + if caps.RequestBodyPatcherBefore != nil || caps.RequestInterceptorPatchBodyBefore { + return true + } + } else { + if caps.RequestBodyPatcherAfter != nil || caps.RequestInterceptorPatchBodyAfter { + return true + } + } + if opt, ok := caps.RequestInterceptor.(interface{ RequestInterceptorWantsPatchBody(bool) bool }); ok && opt != nil { + return opt.RequestInterceptorWantsPatchBody(before) + } + return false +} + func schedulerWantsAcrossPriorities(caps pluginapi.Capabilities) bool { if caps.Scheduler == nil { return false @@ -179,35 +217,41 @@ func schedulerWantsAcrossPriorities(caps pluginapi.Capabilities) bool { func rpcCapabilitiesFromPlugin(plugin pluginapi.Plugin) rpcCapabilities { caps := plugin.Capabilities return rpcCapabilities{ - ModelRegistrar: caps.ModelRegistrar != nil, - ModelProvider: caps.ModelProvider != nil, - AuthProvider: caps.AuthProvider != nil, - FrontendAuthProvider: caps.FrontendAuthProvider != nil, - FrontendAuthProviderExclusive: caps.FrontendAuthProvider != nil && caps.FrontendAuthProviderExclusive, - Scheduler: caps.Scheduler != nil, - SchedulerAcrossPriorities: schedulerWantsAcrossPriorities(caps), - ModelRouter: caps.ModelRouter != nil, - Executor: caps.Executor != nil, - ExecutorModelScope: normalizedExecutorModelScope(caps), - ExecutorInputFormats: append([]string(nil), caps.ExecutorInputFormats...), - ExecutorOutputFormats: append([]string(nil), caps.ExecutorOutputFormats...), - RequestTranslator: caps.RequestTranslator != nil, - RequestNormalizer: caps.RequestNormalizer != nil, - RequestInterceptor: caps.RequestInterceptor != nil, - RequestInterceptorOmitBodyBefore: requestInterceptorWantsOmitBody(caps, true), - RequestInterceptorOmitBodyAfter: requestInterceptorWantsOmitBody(caps, false), - RequestLifecyclePlugin: caps.RequestLifecyclePlugin != nil, - ResponseTranslator: caps.ResponseTranslator != nil, - ResponseBeforeTranslator: caps.ResponseBeforeTranslator != nil, - ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, - ResponseInterceptor: caps.ResponseInterceptor != nil, - StreamChunkInterceptor: caps.StreamChunkInterceptor != nil, - WebSocketResponseObserver: caps.WebSocketResponseObserver != nil, - ThinkingApplier: caps.ThinkingApplier != nil, - UsagePlugin: caps.UsagePlugin != nil, - CommandLinePlugin: caps.CommandLinePlugin != nil, - ManagementAPI: caps.ManagementAPI != nil, - QuotaProvider: caps.QuotaProvider != nil, + ModelRegistrar: caps.ModelRegistrar != nil, + ModelProvider: caps.ModelProvider != nil, + AuthProvider: caps.AuthProvider != nil, + FrontendAuthProvider: caps.FrontendAuthProvider != nil, + FrontendAuthProviderExclusive: caps.FrontendAuthProvider != nil && caps.FrontendAuthProviderExclusive, + FrontendAuthProviderOmitBody: frontendAuthProviderWantsOmitBody(caps), + FrontendAuthOmitBody: frontendAuthProviderWantsOmitBody(caps), + Scheduler: caps.Scheduler != nil, + SchedulerAcrossPriorities: schedulerWantsAcrossPriorities(caps), + ModelRouter: caps.ModelRouter != nil, + Executor: caps.Executor != nil, + ExecutorModelScope: normalizedExecutorModelScope(caps), + ExecutorInputFormats: append([]string(nil), caps.ExecutorInputFormats...), + ExecutorOutputFormats: append([]string(nil), caps.ExecutorOutputFormats...), + RequestTranslator: caps.RequestTranslator != nil, + RequestNormalizer: caps.RequestNormalizer != nil, + RequestInterceptor: caps.RequestInterceptor != nil, + RequestInterceptorOmitBodyBefore: requestInterceptorWantsOmitBody(caps, true), + RequestInterceptorOmitBodyAfter: requestInterceptorWantsOmitBody(caps, false), + RequestBodyPatcherBefore: requestInterceptorWantsPatchBody(caps, true), + RequestBodyPatcherAfter: requestInterceptorWantsPatchBody(caps, false), + RequestInterceptorPatchBodyBefore: requestInterceptorWantsPatchBody(caps, true), + RequestInterceptorPatchBodyAfter: requestInterceptorWantsPatchBody(caps, false), + RequestLifecyclePlugin: caps.RequestLifecyclePlugin != nil, + ResponseTranslator: caps.ResponseTranslator != nil, + ResponseBeforeTranslator: caps.ResponseBeforeTranslator != nil, + ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, + ResponseInterceptor: caps.ResponseInterceptor != nil, + StreamChunkInterceptor: caps.StreamChunkInterceptor != nil, + WebSocketResponseObserver: caps.WebSocketResponseObserver != nil, + ThinkingApplier: caps.ThinkingApplier != nil, + UsagePlugin: caps.UsagePlugin != nil, + CommandLinePlugin: caps.CommandLinePlugin != nil, + ManagementAPI: caps.ManagementAPI != nil, + QuotaProvider: caps.QuotaProvider != nil, } } diff --git a/internal/pluginhost/rpc_schema_test.go b/internal/pluginhost/rpc_schema_test.go index 8251d5168f1..cf1350c793a 100644 --- a/internal/pluginhost/rpc_schema_test.go +++ b/internal/pluginhost/rpc_schema_test.go @@ -1,12 +1,16 @@ package pluginhost import ( + "bytes" "context" "encoding/json" + "io" + "net/http" "reflect" "strings" "testing" + sdkaccess "github.com/router-for-me/CLIProxyAPI/v8/sdk/access" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) @@ -41,6 +45,176 @@ func TestRPCCapabilitiesIncludeFrontendAuthProviderExclusive(t *testing.T) { if decoded["frontend_auth_provider_exclusive"] != true { t.Fatalf("frontend_auth_provider_exclusive = %#v, want true", decoded["frontend_auth_provider_exclusive"]) } + if decoded["frontend_auth_provider_omit_body"] != nil && decoded["frontend_auth_provider_omit_body"] != false { + t.Fatalf("frontend_auth_provider_omit_body = %#v, want omitted or false", decoded["frontend_auth_provider_omit_body"]) + } + if decoded["frontend_auth_omit_body"] != nil && decoded["frontend_auth_omit_body"] != false { + t.Fatalf("frontend_auth_omit_body = %#v, want omitted or false", decoded["frontend_auth_omit_body"]) + } +} + +func TestRPCCapabilitiesIncludeFrontendAuthProviderOmitBody(t *testing.T) { + for _, tc := range []struct { + name string + caps pluginapi.Capabilities + }{ + { + name: "FrontendAuthProviderOmitBody", + caps: pluginapi.Capabilities{ + FrontendAuthProvider: frontendAuthProviderFunc{identifier: "omit-auth"}, + FrontendAuthProviderOmitBody: true, + }, + }, + { + name: "FrontendAuthOmitBody", + caps: pluginapi.Capabilities{ + FrontendAuthProvider: frontendAuthProviderFunc{identifier: "omit-auth"}, + FrontendAuthOmitBody: true, + }, + }, + { + name: "method hook", + caps: pluginapi.Capabilities{ + FrontendAuthProvider: omitBodyFrontendAuthProvider{ + frontendAuthProviderFunc: frontendAuthProviderFunc{identifier: "omit-auth"}, + }, + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + plugin := pluginapi.Plugin{Capabilities: tc.caps} + caps := rpcCapabilitiesFromPlugin(plugin) + if !caps.FrontendAuthProvider { + t.Fatal("FrontendAuthProvider = false, want true") + } + if !caps.FrontendAuthProviderOmitBody || !caps.FrontendAuthOmitBody { + t.Fatalf("omit body flags = provider:%v short:%v, want both true", caps.FrontendAuthProviderOmitBody, caps.FrontendAuthOmitBody) + } + raw, errMarshal := json.Marshal(caps) + if errMarshal != nil { + t.Fatalf("Marshal() error = %v", errMarshal) + } + var decoded map[string]any + if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil { + t.Fatalf("Unmarshal() error = %v", errUnmarshal) + } + if decoded["frontend_auth_provider_omit_body"] != true { + t.Fatalf("frontend_auth_provider_omit_body = %#v, want true", decoded["frontend_auth_provider_omit_body"]) + } + if decoded["frontend_auth_omit_body"] != true { + t.Fatalf("frontend_auth_omit_body = %#v, want true", decoded["frontend_auth_omit_body"]) + } + }) + } +} + +func TestRPCFrontendAuthOmitsBodyOnTheWireForSchema7(t *testing.T) { + for _, tc := range []struct { + name string + schemaVersion uint32 + omitBody bool + shortFlagOnly bool + wantBodyLen int + }{ + {name: "schema 7 omit body", schemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, omitBody: true, wantBodyLen: 0}, + {name: "schema 7 short flag omit body", schemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, omitBody: true, shortFlagOnly: true, wantBodyLen: 0}, + {name: "schema 7 default includes body", schemaVersion: pluginabi.SchemaVersionFrontendAuthOmitBody, omitBody: false, wantBodyLen: len("request")}, + {name: "schema 6 ignores omit body", schemaVersion: 6, omitBody: true, wantBodyLen: len("request")}, + } { + t.Run(tc.name, func(t *testing.T) { + const key = "plugin:auth-wire:quota" + sdkaccess.UnregisterProvider(key) + defer sdkaccess.UnregisterProvider(key) + + var wireGot pluginapi.FrontendAuthRequest + var wireCalls int + plugin := validTestPlugin("auth-wire") + plugin.Capabilities.FrontendAuthProvider = frontendAuthProviderFunc{ + identifier: "quota", + authenticate: func(_ context.Context, req pluginapi.FrontendAuthRequest) (pluginapi.FrontendAuthResponse, error) { + wireCalls++ + wireGot = req + return pluginapi.FrontendAuthResponse{Authenticated: true, Principal: "wire-user"}, nil + }, + } + if tc.omitBody && !tc.shortFlagOnly { + plugin.Capabilities.FrontendAuthProviderOmitBody = true + } + lookup := newTestSymbolLookup(&testPlugin{registerResult: plugin}) + lookup.schemaVersion = tc.schemaVersion + + var client pluginClient = lookup + if tc.shortFlagOnly { + client = &shortFrontendAuthOmitBodyClient{inner: lookup} + } + + registered, errRegister := registerRPCPlugin(context.Background(), nil, "auth-wire", client, pluginabi.MethodPluginRegister, nil) + if errRegister != nil { + t.Fatalf("registerRPCPlugin() error = %v", errRegister) + } + if registered.SchemaVersion != tc.schemaVersion { + t.Fatalf("registered SchemaVersion = %d, want %d", registered.SchemaVersion, tc.schemaVersion) + } + if registered.Capabilities.FrontendAuthProviderOmitBody != tc.omitBody || registered.Capabilities.FrontendAuthOmitBody != tc.omitBody { + t.Fatalf("registered omit flags = provider:%v short:%v, want %v", registered.Capabilities.FrontendAuthProviderOmitBody, registered.Capabilities.FrontendAuthOmitBody, tc.omitBody) + } + + host := newHostWithRecords(capabilityRecord{id: "auth-wire", plugin: registered}) + host.RegisterFrontendAuthProviders() + + var provider sdkaccess.Provider + for _, candidate := range sdkaccess.RegisteredProviders() { + if candidate.Identifier() == key { + provider = candidate + break + } + } + if provider == nil { + t.Fatalf("provider %q was not registered", key) + } + + counter := &countingAuthBodyReader{reader: strings.NewReader("request")} + req, errNewReq := http.NewRequest(http.MethodPost, "http://example.test/v1/responses", io.NopCloser(counter)) + if errNewReq != nil { + t.Fatalf("NewRequest() error = %v", errNewReq) + } + req.Header.Set("Authorization", "Bearer wire-token") + res, authErr := provider.Authenticate(context.Background(), req) + if authErr != nil { + t.Fatalf("Authenticate() error = %v, want nil", authErr) + } + if res == nil || res.Principal != "wire-user" { + t.Fatalf("Authenticate() result = %#v, want wire-user", res) + } + if wireCalls != 1 { + t.Fatalf("wire calls = %d, want 1", wireCalls) + } + if len(wireGot.Body) != tc.wantBodyLen { + t.Fatalf("wire body len = %d, want %d", len(wireGot.Body), tc.wantBodyLen) + } + if counter.readBytes != tc.wantBodyLen { + t.Fatalf("counter.readBytes = %d, want %d", counter.readBytes, tc.wantBodyLen) + } + }) + } +} + +type shortFrontendAuthOmitBodyClient struct { + inner pluginClient +} + +func (c *shortFrontendAuthOmitBodyClient) Call(ctx context.Context, method string, request []byte) ([]byte, error) { + if method == pluginabi.MethodPluginRegister || method == pluginabi.MethodPluginReconfigure { + if _, errCall := c.inner.Call(ctx, method, request); errCall != nil { + return nil, errCall + } + return []byte(`{"ok":true,"result":{"schema_version":7,"metadata":{"Name":"auth-wire","Version":"1.0.0","Author":"test","GitHubRepository":"https://github.com/router-for-me/CLIProxyAPI"},"capabilities":{"frontend_auth_provider":true,"frontend_auth_omit_body":true}}}`), nil + } + return c.inner.Call(ctx, method, request) +} + +func (c *shortFrontendAuthOmitBodyClient) Shutdown() { + c.inner.Shutdown() } func TestRPCCapabilitiesIncludeScheduler(t *testing.T) { @@ -638,3 +812,142 @@ func TestRPCSchedulerAndUsageSelectionIDCompatibility(t *testing.T) { t.Fatalf("legacyUsageRecord = %+v", legacy) } } + +type recordingRawPatchClient struct { + schemaVersion uint32 + methods []string + rawRequests [][]byte + rawRequestPtr []uintptr + patchFunc func(method string, rawBody []byte) (pluginapi.RequestBodyPatchResponse, error) +} + +func (c *recordingRawPatchClient) Call(_ context.Context, method string, request []byte) ([]byte, error) { + c.methods = append(c.methods, method) + c.rawRequests = append(c.rawRequests, bytes.Clone(request)) + if len(request) > 0 { + c.rawRequestPtr = append(c.rawRequestPtr, uintptr(reflect.ValueOf(request).Pointer())) + } else { + c.rawRequestPtr = append(c.rawRequestPtr, 0) + } + switch method { + case pluginabi.MethodPluginRegister: + return marshalRPCResult(rpcRegistration{ + Metadata: pluginapi.Metadata{ + Name: "Timezone Patcher", + Version: "1.0.0", + Author: "test", + GitHubRepository: "https://github.com/router-for-me/CLIProxyAPI", + }, + SchemaVersion: c.schemaVersion, + Capabilities: rpcCapabilities{ + RequestInterceptorPatchBodyBefore: true, + RequestInterceptorPatchBodyAfter: true, + }, + }) + case pluginabi.MethodRequestPatchBodyBefore, pluginabi.MethodRequestPatchBodyAfter: + resp, err := c.patchFunc(method, request) + if err != nil { + return nil, err + } + return marshalRPCResult(resp) + case pluginabi.MethodRequestInterceptBefore, pluginabi.MethodRequestInterceptAfter: + var req pluginapi.RequestInterceptRequest + if err := json.Unmarshal(request, &req); err != nil { + return nil, err + } + return marshalRPCResult(pluginapi.RequestInterceptResponse{ + Body: bytes.ReplaceAll(req.Body, []byte("UTC"), []byte("Legacy")), + }) + default: + return marshalRPCResult(struct{}{}) + } +} + +func (c *recordingRawPatchClient) Shutdown() {} +func (c *recordingRawPatchClient) Close() error { return nil } + +func TestRPCRequestBodyPatchRawCABIAndSchemaVersionGate(t *testing.T) { + rawInput := []byte(`{"model":"gpt-5","input":[{"role":"developer","content":[{"type":"input_text","text":"\n 2026-04-19\n UTC\n"}]}]}`) + + clientV8 := &recordingRawPatchClient{ + schemaVersion: pluginabi.SchemaVersionRequestInterceptorBodyPatch, + patchFunc: func(method string, rawBody []byte) (pluginapi.RequestBodyPatchResponse, error) { + if method == pluginabi.MethodRequestPatchBodyBefore { + target := []byte("UTC") + idx := bytes.Index(rawBody, target) + if idx < 0 { + return pluginapi.RequestBodyPatchResponse{}, nil + } + return pluginapi.RequestBodyPatchResponse{ + Edits: []pluginapi.RequestBodyEdit{ + { + Start: idx, + End: idx + len(target), + Replacement: "Asia/Shanghai", + }, + }, + }, nil + } + return pluginapi.RequestBodyPatchResponse{}, nil + }, + } + + host := New() + pluginV8, errReg := registerRPCPlugin(context.Background(), host, "tz-patcher", clientV8, pluginabi.MethodPluginRegister, nil) + if errReg != nil { + t.Fatalf("registerRPCPlugin(v8) error = %v", errReg) + } + if !validPlugin(pluginV8) { + t.Fatalf("validPlugin(patch-only pluginV8) = false, want true") + } + setHostSnapshotForTest(host, true, capabilityRecord{id: "tz-patcher", plugin: pluginV8}) + + if !host.HasRequestInterceptors() { + t.Fatal("HasRequestInterceptors() = false, want true for patch-only plugin") + } + + resp := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{ + Model: "gpt-5", + Body: rawInput, + }) + if len(clientV8.methods) != 2 || clientV8.methods[1] != pluginabi.MethodRequestPatchBodyBefore { + t.Fatalf("called methods = %v, want [%s %s]", clientV8.methods, pluginabi.MethodPluginRegister, pluginabi.MethodRequestPatchBodyBefore) + } + // Crucial requirement 1: C ABI receives the raw request body directly, not a JSON/Base64 envelope, + // and without a pre-call bytes.Clone on the host side. + if !bytes.Equal(clientV8.rawRequests[1], rawInput) { + t.Fatalf("raw C ABI request = %s, want exact rawInput %s", clientV8.rawRequests[1], rawInput) + } + if gotPtr, wantPtr := clientV8.rawRequestPtr[1], uintptr(reflect.ValueOf(rawInput).Pointer()); gotPtr != wantPtr { + t.Fatalf("raw C ABI request slice pointer = %#x, want original rawInput pointer %#x (zero pre-call copy)", gotPtr, wantPtr) + } + wantPatched := `{"model":"gpt-5","input":[{"role":"developer","content":[{"type":"input_text","text":"\n 2026-04-19\n Asia/Shanghai\n"}]}]}` + if string(resp.Body) != wantPatched { + t.Fatalf("patched body = %s, want %s", resp.Body, wantPatched) + } + + // Crucial requirement 3: When plugin returns empty edits, host reuses original body (resp.Body == nil). + afterResp := host.InterceptRequestAfterAuth(context.Background(), pluginapi.RequestInterceptRequest{ + Model: "gpt-5", + Body: rawInput, + }) + if afterResp.Body != nil { + t.Fatalf("no-op patch returned non-nil Body = %s, want nil to reuse original body", afterResp.Body) + } + + // Schema v7 must ignore request_interceptor_patch_body_before. + clientV7 := &recordingRawPatchClient{ + schemaVersion: pluginabi.SchemaVersionRequestInterceptorOmitBody, + patchFunc: func(string, []byte) (pluginapi.RequestBodyPatchResponse, error) { + t.Fatal("patchFunc must not be called on schema v7") + return pluginapi.RequestBodyPatchResponse{}, nil + }, + } + pluginV7, errRegV7 := registerRPCPlugin(context.Background(), host, "tz-patcher-v7", clientV7, pluginabi.MethodPluginRegister, nil) + if errRegV7 != nil { + t.Fatalf("registerRPCPlugin(v7) error = %v", errRegV7) + } + if pluginV7.Capabilities.RequestBodyPatcherBefore != nil || pluginV7.Capabilities.RequestInterceptorPatchBodyBefore { + t.Fatalf("schema v7 kept patch capability: %+v", pluginV7.Capabilities) + } +} diff --git a/internal/pluginhost/test_helpers_test.go b/internal/pluginhost/test_helpers_test.go index fb199f722d8..b359918e4b6 100644 --- a/internal/pluginhost/test_helpers_test.go +++ b/internal/pluginhost/test_helpers_test.go @@ -97,6 +97,26 @@ func (l *testSymbolLookup) Call(ctx context.Context, method string, request []by return nil, errIntercept } return marshalRPCResult(resp) + case pluginabi.MethodRequestPatchBodyBefore: + patcher := requestRecordBodyPatcher(capabilityRecord{plugin: l.active}, true) + if patcher == nil { + return nil, fmt.Errorf("missing request body patcher before") + } + resp, errPatch := patcher.PatchRequestBody(ctx, request) + if errPatch != nil { + return nil, errPatch + } + return marshalRPCResult(resp) + case pluginabi.MethodRequestPatchBodyAfter: + patcher := requestRecordBodyPatcher(capabilityRecord{plugin: l.active}, false) + if patcher == nil { + return nil, fmt.Errorf("missing request body patcher after") + } + resp, errPatch := patcher.PatchRequestBody(ctx, request) + if errPatch != nil { + return nil, errPatch + } + return marshalRPCResult(resp) case pluginabi.MethodRequestComplete: if l.active.Capabilities.RequestLifecyclePlugin == nil { return nil, fmt.Errorf("missing request lifecycle plugin") @@ -140,6 +160,24 @@ func (l *testSymbolLookup) Call(ctx context.Context, method string, request []by return nil, fmt.Errorf("missing auth provider") } return marshalRPCResult(rpcIdentifierResponse{Identifier: l.active.Capabilities.AuthProvider.Identifier()}) + case pluginabi.MethodFrontendAuthIdentifier: + if l.active.Capabilities.FrontendAuthProvider == nil { + return nil, fmt.Errorf("missing frontend auth provider") + } + return marshalRPCResult(rpcIdentifierResponse{Identifier: l.active.Capabilities.FrontendAuthProvider.Identifier()}) + case pluginabi.MethodFrontendAuthAuthenticate: + if l.active.Capabilities.FrontendAuthProvider == nil { + return nil, fmt.Errorf("missing frontend auth provider") + } + var req pluginapi.FrontendAuthRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, errUnmarshal + } + resp, errAuth := l.active.Capabilities.FrontendAuthProvider.Authenticate(ctx, req) + if errAuth != nil { + return nil, errAuth + } + return marshalRPCResult(resp) case pluginabi.MethodSchedulerPick: if l.active.Capabilities.Scheduler == nil { return nil, fmt.Errorf("missing scheduler") diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index c9d0a944f40..c922a73c01b 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -105,14 +105,29 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec usageDetail := coreusage.EnsureTokenBreakdownForProvider(record.Detail, record.Provider, record.ExecutorType) tokens := tokenStats{ - InputTokens: usageDetail.InputTokens, - OutputTokens: usageDetail.OutputTokens, - ReasoningTokens: usageDetail.ReasoningTokens, - CachedTokens: usageDetail.CachedTokens, - CacheReadTokens: usageDetail.CacheReadTokens, - CacheReadTokensPresent: true, - CacheCreationTokens: usageDetail.CacheCreationTokens, - TotalTokens: usageDetail.TotalTokens, + InputTokens: usageDetail.InputTokens, + OutputTokens: usageDetail.OutputTokens, + ReasoningTokens: usageDetail.ReasoningTokens, + CachedTokens: usageDetail.CachedTokens, + CacheReadTokens: usageDetail.CacheReadTokens, + CacheReadTokensPresent: true, + CacheCreationTokens: usageDetail.CacheCreationTokens, + CacheCreation5mTokens: coreusage.CloneInt64Ptr(usageDetail.CacheCreation5mTokens), + CacheCreation30mTokens: coreusage.CloneInt64Ptr(usageDetail.CacheCreation30mTokens), + CacheCreation1hTokens: coreusage.CloneInt64Ptr(usageDetail.CacheCreation1hTokens), + CacheCreation12hTokens: coreusage.CloneInt64Ptr(usageDetail.CacheCreation12hTokens), + TextInputTokens: coreusage.CloneInt64Ptr(usageDetail.TextInputTokens), + ImageInputTokens: coreusage.CloneInt64Ptr(usageDetail.ImageInputTokens), + AudioInputTokens: coreusage.CloneInt64Ptr(usageDetail.AudioInputTokens), + VideoInputTokens: coreusage.CloneInt64Ptr(usageDetail.VideoInputTokens), + FileInputTokens: coreusage.CloneInt64Ptr(usageDetail.FileInputTokens), + CachedTextTokens: coreusage.CloneInt64Ptr(usageDetail.CachedTextTokens), + TextOutputTokens: coreusage.CloneInt64Ptr(usageDetail.TextOutputTokens), + ImageOutputTokens: coreusage.CloneInt64Ptr(usageDetail.ImageOutputTokens), + AudioOutputTokens: coreusage.CloneInt64Ptr(usageDetail.AudioOutputTokens), + AcceptedPredictionTokens: coreusage.CloneInt64Ptr(usageDetail.AcceptedPredictionTokens), + RejectedPredictionTokens: coreusage.CloneInt64Ptr(usageDetail.RejectedPredictionTokens), + TotalTokens: usageDetail.TotalTokens, } failed := record.Failed @@ -146,31 +161,43 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } payload, err := json.Marshal(queuedUsageDetail{ - requestDetail: detail, - AccountingVersion: coreusage.TokenAccountingSchemaVersion, - TokenBreakdown: usageDetail.TokenBreakdown, - Provider: provider, - ExecutorType: executorType, - Model: modelName, - Alias: aliasName, - Endpoint: resolveEndpoint(ctx), - AuthType: authType, - APIKey: apiKey, - IsNativeKey: isNativeKey, - AccessProvider: accessProvider, - RequestID: requestID, - ExecutionID: executionID, - TraceID: traceID, - SessionID: sessionID, - ParentSessionID: parentSessionID, - NodeKind: strings.TrimSpace(clientRequestMetadata.NodeKind), - IsFork: clientRequestMetadata.IsFork, - IsCompaction: clientRequestMetadata.IsCompaction, - ReasoningEffort: reasoningEffort, - ServiceTier: serviceTier, - EffectiveServiceTier: effectiveServiceTier, - ResponseServiceTier: responseServiceTier, - ResponseModel: responseModel, + requestDetail: detail, + AccountingVersion: coreusage.TokenAccountingSchemaVersion, + TokenBreakdown: usageDetail.TokenBreakdown, + Provider: provider, + ExecutorType: executorType, + Model: modelName, + Alias: aliasName, + Endpoint: resolveEndpoint(ctx), + AuthType: authType, + APIKey: apiKey, + IsNativeKey: isNativeKey, + AccessProvider: accessProvider, + RequestID: requestID, + ExecutionID: executionID, + TraceID: traceID, + SessionID: sessionID, + ParentSessionID: parentSessionID, + NodeKind: strings.TrimSpace(clientRequestMetadata.NodeKind), + IsFork: clientRequestMetadata.IsFork, + IsCompaction: clientRequestMetadata.IsCompaction, + ReasoningEffort: reasoningEffort, + ServiceTier: serviceTier, + EffectiveServiceTier: effectiveServiceTier, + ResponseServiceTier: responseServiceTier, + UpstreamServiceTier: strings.TrimSpace(usageDetail.UpstreamServiceTier), + ResponseModel: responseModel, + Speed: strings.TrimSpace(usageDetail.Speed), + InferenceGeo: strings.TrimSpace(usageDetail.InferenceGeo), + CacheDetails: usageDetail.CacheDetails.Clone(), + WebSearchRequests: coreusage.CloneInt64Ptr(usageDetail.WebSearchRequests), + ImageCount: coreusage.CloneInt64Ptr(usageDetail.ImageCount), + RequestedImageDetails: usageDetail.RequestedImageDetails.Clone(), + ImageDetails: usageDetail.ImageDetails.Clone(), + MediaCounts: usageDetail.MediaCounts.Clone(), + ToolUsageDetails: usageDetail.ToolUsageDetails.Clone(), + Iterations: usageDetail.IterationSlice(), + CostDetails: usageDetail.CostDetails.Clone(), }) if err != nil { return @@ -180,30 +207,42 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec type queuedUsageDetail struct { requestDetail - AccountingVersion int `json:"accounting_version"` - TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` - Provider string `json:"provider"` - ExecutorType string `json:"executor_type"` - Model string `json:"model"` - Alias string `json:"alias"` - Endpoint string `json:"endpoint"` - AuthType string `json:"auth_type"` - APIKey string `json:"api_key"` - IsNativeKey bool `json:"is_native_key"` - AccessProvider string `json:"access_provider,omitempty"` - RequestID string `json:"request_id"` - ExecutionID string `json:"execution_id,omitempty"` - TraceID string `json:"trace_id,omitempty"` - SessionID string `json:"session_id,omitempty"` - ParentSessionID string `json:"parent_session_id,omitempty"` - NodeKind string `json:"node_kind,omitempty"` - IsFork bool `json:"is_fork,omitempty"` - IsCompaction bool `json:"is_compaction,omitempty"` - ReasoningEffort string `json:"reasoning_effort"` - ServiceTier string `json:"service_tier"` - EffectiveServiceTier string `json:"effective_service_tier,omitempty"` - ResponseServiceTier string `json:"response_service_tier,omitempty"` - ResponseModel string `json:"response_model,omitempty"` + AccountingVersion int `json:"accounting_version"` + TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` + Provider string `json:"provider"` + ExecutorType string `json:"executor_type"` + Model string `json:"model"` + Alias string `json:"alias"` + Endpoint string `json:"endpoint"` + AuthType string `json:"auth_type"` + APIKey string `json:"api_key"` + IsNativeKey bool `json:"is_native_key"` + AccessProvider string `json:"access_provider,omitempty"` + RequestID string `json:"request_id"` + ExecutionID string `json:"execution_id,omitempty"` + TraceID string `json:"trace_id,omitempty"` + SessionID string `json:"session_id,omitempty"` + ParentSessionID string `json:"parent_session_id,omitempty"` + NodeKind string `json:"node_kind,omitempty"` + IsFork bool `json:"is_fork,omitempty"` + IsCompaction bool `json:"is_compaction,omitempty"` + ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` + EffectiveServiceTier string `json:"effective_service_tier,omitempty"` + ResponseServiceTier string `json:"response_service_tier,omitempty"` + UpstreamServiceTier string `json:"upstream_service_tier,omitempty"` + ResponseModel string `json:"response_model,omitempty"` + Speed string `json:"speed,omitempty"` + InferenceGeo string `json:"inference_geo,omitempty"` + CacheDetails *coreusage.CacheDetails `json:"cache_details,omitempty"` + WebSearchRequests *int64 `json:"web_search_requests,omitempty"` + ImageCount *int64 `json:"image_count,omitempty"` + RequestedImageDetails *coreusage.ImageDetails `json:"requested_image_details,omitempty"` + ImageDetails *coreusage.ImageDetails `json:"image_details,omitempty"` + MediaCounts *coreusage.MediaCounts `json:"media_counts,omitempty"` + ToolUsageDetails *coreusage.ToolUsageDetails `json:"tool_usage_details,omitempty"` + Iterations []coreusage.IterationUsage `json:"iterations,omitempty"` + CostDetails *coreusage.CostDetails `json:"cost_details,omitempty"` } type requestDetail struct { @@ -226,14 +265,29 @@ type requestDetail struct { } type tokenStats struct { - InputTokens int64 `json:"input_tokens"` - OutputTokens int64 `json:"output_tokens"` - ReasoningTokens int64 `json:"reasoning_tokens"` - CachedTokens int64 `json:"cached_tokens"` - CacheReadTokens int64 `json:"cache_read_tokens"` - CacheReadTokensPresent bool `json:"cache_read_tokens_present"` - CacheCreationTokens int64 `json:"cache_creation_tokens"` - TotalTokens int64 `json:"total_tokens"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + ReasoningTokens int64 `json:"reasoning_tokens"` + CachedTokens int64 `json:"cached_tokens"` + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheReadTokensPresent bool `json:"cache_read_tokens_present"` + CacheCreationTokens int64 `json:"cache_creation_tokens"` + CacheCreation5mTokens *int64 `json:"cache_creation_5m_tokens,omitempty"` + CacheCreation30mTokens *int64 `json:"cache_creation_30m_tokens,omitempty"` + CacheCreation1hTokens *int64 `json:"cache_creation_1h_tokens,omitempty"` + CacheCreation12hTokens *int64 `json:"cache_creation_12h_tokens,omitempty"` + TextInputTokens *int64 `json:"text_input_tokens,omitempty"` + ImageInputTokens *int64 `json:"image_input_tokens,omitempty"` + AudioInputTokens *int64 `json:"audio_input_tokens,omitempty"` + VideoInputTokens *int64 `json:"video_input_tokens,omitempty"` + FileInputTokens *int64 `json:"file_input_tokens,omitempty"` + CachedTextTokens *int64 `json:"cached_text_tokens,omitempty"` + TextOutputTokens *int64 `json:"text_output_tokens,omitempty"` + ImageOutputTokens *int64 `json:"image_output_tokens,omitempty"` + AudioOutputTokens *int64 `json:"audio_output_tokens,omitempty"` + AcceptedPredictionTokens *int64 `json:"accepted_prediction_tokens,omitempty"` + RejectedPredictionTokens *int64 `json:"rejected_prediction_tokens,omitempty"` + TotalTokens int64 `json:"total_tokens"` } type failDetail struct { diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 41c3a5f0816..0b34177fc43 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -907,3 +907,78 @@ func TestUsageQueuePluginEmitsFailCode(t *testing.T) { } }) } + +func TestUsageQueuePluginPublishesAnthropicStructuredUsageDetails(t *testing.T) { + withEnabledQueue(t, func() { + ctx := internallogging.WithRequestID(context.Background(), "claude-req-1") + ctx = internallogging.WithEndpoint(ctx, "POST /v1/messages") + ctx = internallogging.WithResponseStatusHolder(ctx) + internallogging.SetResponseStatus(ctx, http.StatusOK) + + plugin := &usageQueuePlugin{} + plugin.HandleUsage(ctx, coreusage.Record{ + Provider: "claude", + ExecutorType: "ClaudeExecutor", + Model: "claude-opus-5-5", + ResponseServiceTier: "standard", + Detail: coreusage.Detail{ + InputTokens: 6, + OutputTokens: 483, + ReasoningTokens: 45, + CachedTokens: 3036, + CacheReadTokens: 3036, + CacheCreationTokens: 8846, + TotalTokens: 12371, + CacheCreation5mTokens: coreusage.Int64Ptr(5810), + CacheCreation1hTokens: coreusage.Int64Ptr(3036), + CacheDetails: &coreusage.CacheDetails{ + Ephemeral5mInputTokens: coreusage.Int64Ptr(5810), + Ephemeral1hInputTokens: coreusage.Int64Ptr(3036), + ByTTL: map[string]int64{"5m": 5810, "1h": 3036}, + }, + WebSearchRequests: coreusage.Int64Ptr(1), + ToolUsageDetails: &coreusage.ToolUsageDetails{ + WebSearchCalls: coreusage.Int64Ptr(1), + WebFetchCalls: coreusage.Int64Ptr(0), + }, + Iterations: coreusage.IterationsPtr([]coreusage.IterationUsage{ + { + Type: "message", + InputTokens: coreusage.Int64Ptr(4), + OutputTokens: coreusage.Int64Ptr(53), + CacheReadInputTokens: coreusage.Int64Ptr(0), + CacheCreationInputTokens: coreusage.Int64Ptr(3036), + CacheCreation5mTokens: coreusage.Int64Ptr(0), + CacheCreation1hTokens: coreusage.Int64Ptr(3036), + }, + { + Type: "message", + InputTokens: coreusage.Int64Ptr(2), + OutputTokens: coreusage.Int64Ptr(430), + CacheReadInputTokens: coreusage.Int64Ptr(3036), + CacheCreationInputTokens: coreusage.Int64Ptr(5810), + CacheCreation5mTokens: coreusage.Int64Ptr(5810), + CacheCreation1hTokens: coreusage.Int64Ptr(0), + }, + }), + Speed: "standard", + InferenceGeo: "global", + }, + }) + + payload := popSinglePayload(t) + tokens := requireTokensPayload(t, payload) + requireIntField(t, tokens, "cache_creation_5m_tokens", 5810) + requireIntField(t, tokens, "cache_creation_1h_tokens", 3036) + requireMissingField(t, tokens, "cache_creation_30m_tokens") + requireIntField(t, payload, "web_search_requests", 1) + requireStringField(t, payload, "speed", "standard") + requireStringField(t, payload, "inference_geo", "global") + if _, ok := payload["iterations"]; !ok { + t.Fatal("expected iterations in payload") + } + if _, ok := payload["tool_usage_details"]; !ok { + t.Fatal("expected tool_usage_details in payload") + } + }) +} diff --git a/internal/registry/model_definitions_test.go b/internal/registry/model_definitions_test.go index ee12e616ba6..b5a73d9494e 100644 --- a/internal/registry/model_definitions_test.go +++ b/internal/registry/model_definitions_test.go @@ -7,6 +7,8 @@ import ( ) func TestCodexConfigurationUpdateCapability(t *testing.T) { + // Counterexample: present in every tier without support_configuration_update. + const incapableModelID = "gpt-5.6-luna" for _, tier := range []struct { name string models func() []*ModelInfo @@ -24,13 +26,13 @@ func TestCodexConfigurationUpdateCapability(t *testing.T) { byID[info.ID] = info } } - for _, id := range append([]string{"gpt-5.5"}, tier.capable...) { + for _, id := range append([]string{incapableModelID}, tier.capable...) { info := byID[id] if info == nil { t.Errorf("%s is missing", id) continue } - want := id != "gpt-5.5" + want := id != incapableModelID if info.SupportConfigurationUpdate != want { t.Errorf("%s configuration_update = %v, want %v", id, info.SupportConfigurationUpdate, want) } diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 0c0a84319d6..7f82e7dcf39 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -202,6 +202,14 @@ type ModelRegistry struct { mutex *sync.RWMutex // availableModelsCache stores per-handler snapshots for GetAvailableModels. availableModelsCache map[string]availableModelsCacheEntry + // webSearchClientCounts maps model ID to the number of clients advertising + // web search for it. It makes hasClientSupportingWebSearchLocked O(1) + // instead of scanning every client, which is quadratic when the whole + // registry is re-registered (thousands of clients x models). + webSearchClientCounts map[string]int + // webSearchProviderCounts maps model ID to provider -> client count for + // clients advertising web search, mirroring webSearchClientCounts. + webSearchProviderCounts map[string]map[string]int // generation tracks changes to model registrations and availability. generation uint64 // registrationEpoch tracks monotonic client registration and deregistration structural changes. @@ -218,14 +226,15 @@ var registryOnce sync.Once func GetGlobalRegistry() *ModelRegistry { registryOnce.Do(func() { globalRegistry = &ModelRegistry{ - models: make(map[string]*ModelRegistration), - clientModels: make(map[string][]string), - clientModelInfos: make(map[string]map[string]*ModelInfo), - clientProviders: make(map[string]string), - clientEpochs: make(map[string]uint64), - clientGenerations: make(map[string]uint64), - availableModelsCache: make(map[string]availableModelsCacheEntry), - mutex: &sync.RWMutex{}, + models: make(map[string]*ModelRegistration), + clientModels: make(map[string][]string), + clientModelInfos: make(map[string]map[string]*ModelInfo), + clientProviders: make(map[string]string), + clientEpochs: make(map[string]uint64), + clientGenerations: make(map[string]uint64), + availableModelsCache: make(map[string]availableModelsCacheEntry), + webSearchClientCounts: map[string]int{}, + mutex: &sync.RWMutex{}, } }) return globalRegistry @@ -522,6 +531,7 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo if len(uniqueModelIDs) == 0 { // No models supplied; unregister existing client state if present. + // unregisterClientInternal retracts the client's web-search counts. r.unregisterClientInternal(clientID) delete(r.clientModels, clientID) delete(r.clientModelInfos, clientID) @@ -545,7 +555,9 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo now := time.Now() if !hadExisting { - // Pure addition path. + // Pure addition normally has no prior metadata. Clear orphaned metadata + // and refresh its models before publishing the new catalog. + r.clearOrphanedClientModelInfosLocked(clientID) for _, modelID := range rawModelIDs { model := newModels[modelID] r.addModelRegistration(modelID, provider, model, now, clientID) @@ -562,6 +574,7 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo } else { delete(r.clientProviders, clientID) } + r.addClientWebSearchCountsLocked(clientID) r.invalidateAvailableModelsCacheLocked() r.triggerModelsRegistered(provider, clientID, models) log.Debugf("Registered client %s from provider %s with %d models", clientID, clientProvider, len(rawModelIDs)) @@ -701,6 +714,9 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo } // Update client bookkeeping. + // Retract this client's previous web-search contributions before rewriting + // clientModelInfos/clientProviders, then re-add them from the new state. + r.dropClientWebSearchCountsLocked(clientID) if len(rawModelIDs) > 0 { r.clientModels[clientID] = append([]string(nil), rawModelIDs...) } @@ -715,6 +731,7 @@ func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, mo } else { delete(r.clientProviders, clientID) } + r.addClientWebSearchCountsLocked(clientID) r.invalidateAvailableModelsCacheLocked() r.triggerModelsRegistered(provider, clientID, models) @@ -902,6 +919,7 @@ func (r *ModelRegistry) unregisterClientInternal(clientID string) { models, exists := r.clientModels[clientID] provider, hasProvider := r.clientProviders[clientID] if !exists { + r.clearOrphanedClientModelInfosLocked(clientID) if hasProvider { delete(r.clientProviders, clientID) } @@ -950,6 +968,11 @@ func (r *ModelRegistry) unregisterClientInternal(clientID string) { } } + // Retract this client's web-search contributions while clientModelInfos and + // clientProviders still describe it. This must run after the per-model loop + // above, which relies on excludeClientID to skip this client. + r.dropClientWebSearchCountsLocked(clientID) + delete(r.clientModels, clientID) delete(r.clientModelInfos, clientID) if hasProvider { @@ -1151,9 +1174,18 @@ func (r *ModelRegistry) ApplyClientModelCapabilities(clientID string, expectedEp } provider := r.clientProviders[clientID] + // The callback mutates capabilities in place, which can add or clear this + // client's web-search contributions. Retract the current ones first so the + // index stays in sync with clientInfos. + r.dropClientWebSearchCountsLocked(clientID) for id, info := range clientInfos { if info != nil { mutate(id, info) + // Reinstate this client's contribution before querying the index so + // the aggregate reflects the post-mutation state, including itself. + if info.SupportsWebSearch { + r.incrementWebSearchCountLocked(id, provider) + } if reg, okReg := r.models[id]; okReg && reg != nil { hasWebSearch := r.hasClientSupportingWebSearchLocked(id, "", "") if reg.Info != nil { @@ -1169,19 +1201,131 @@ func (r *ModelRegistry) ApplyClientModelCapabilities(clientID string, expectedEp return true } +// hasClientSupportingWebSearchLocked excludes the given ID even when it is empty, +// matching the original scan without walking the client set. func (r *ModelRegistry) hasClientSupportingWebSearchLocked(modelID, provider, excludeClientID string) bool { - for cID, infos := range r.clientModelInfos { - if cID == excludeClientID { - continue + if provider != "" { + count := r.webSearchProviderCounts[modelID][provider] + if count <= 0 { + return false } - if provider != "" && r.clientProviders[cID] != provider { - continue + if r.clientProviders[excludeClientID] == provider { + if info := r.clientModelInfos[excludeClientID][modelID]; info != nil && info.SupportsWebSearch { + count-- + } } - if info, ok := infos[modelID]; ok && info != nil && info.SupportsWebSearch { - return true + return count > 0 + } + + count := r.webSearchClientCounts[modelID] + if count <= 0 { + return false + } + if info := r.clientModelInfos[excludeClientID][modelID]; info != nil && info.SupportsWebSearch { + count-- + } + return count > 0 +} + +// incrementWebSearchCountLocked records one client's web-search contribution for modelID. +func (r *ModelRegistry) incrementWebSearchCountLocked(modelID, provider string) { + if modelID == "" { + return + } + if r.webSearchClientCounts == nil { + r.webSearchClientCounts = make(map[string]int) + } + r.webSearchClientCounts[modelID]++ + if provider == "" { + return + } + if r.webSearchProviderCounts == nil { + r.webSearchProviderCounts = make(map[string]map[string]int) + } + counts := r.webSearchProviderCounts[modelID] + if counts == nil { + counts = make(map[string]int) + r.webSearchProviderCounts[modelID] = counts + } + counts[provider]++ +} + +// decrementWebSearchCountLocked removes one client's web-search contribution for modelID. +func (r *ModelRegistry) decrementWebSearchCountLocked(modelID, provider string) { + if modelID == "" { + return + } + if count := r.webSearchClientCounts[modelID]; count > 1 { + r.webSearchClientCounts[modelID] = count - 1 + } else { + delete(r.webSearchClientCounts, modelID) + } + if provider == "" { + return + } + counts := r.webSearchProviderCounts[modelID] + if counts == nil { + return + } + if count := counts[provider]; count > 1 { + counts[provider] = count - 1 + } else { + delete(counts, provider) + } + if len(counts) == 0 { + delete(r.webSearchProviderCounts, modelID) + } +} + +// clearOrphanedClientModelInfosLocked refreshes capabilities before retracting +// and discarding metadata for a client without clientModels. The caller must +// retain clientProviders until this returns; model multiplicities are unknown. +func (r *ModelRegistry) clearOrphanedClientModelInfosLocked(clientID string) { + provider, hasProvider := r.clientProviders[clientID] + for modelID := range r.clientModelInfos[clientID] { + if registration := r.models[modelID]; registration != nil { + if registration.Info != nil { + registration.Info.SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, "", clientID) + } + if hasProvider && registration.InfoByProvider[provider] != nil { + registration.InfoByProvider[provider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, provider, clientID) + } + } + } + r.dropClientWebSearchCountsLocked(clientID) + delete(r.clientModelInfos, clientID) +} + +// dropClientWebSearchCountsLocked removes every web-search contribution currently +// recorded for clientID. It must run before clientModelInfos/clientProviders are +// rewritten for that client. +func (r *ModelRegistry) dropClientWebSearchCountsLocked(clientID string) { + infos := r.clientModelInfos[clientID] + if len(infos) == 0 { + return + } + provider := r.clientProviders[clientID] + for modelID, info := range infos { + if info != nil && info.SupportsWebSearch { + r.decrementWebSearchCountLocked(modelID, provider) + } + } +} + +// addClientWebSearchCountsLocked records every web-search contribution currently +// present for clientID. It must run after clientModelInfos/clientProviders are +// rewritten for that client. +func (r *ModelRegistry) addClientWebSearchCountsLocked(clientID string) { + infos := r.clientModelInfos[clientID] + if len(infos) == 0 { + return + } + provider := r.clientProviders[clientID] + for modelID, info := range infos { + if info != nil && info.SupportsWebSearch { + r.incrementWebSearchCountLocked(modelID, provider) } } - return false } // SuspendClientModel marks a client's model as temporarily unavailable until explicitly resumed. diff --git a/internal/registry/web_search_index_test.go b/internal/registry/web_search_index_test.go new file mode 100644 index 00000000000..e4353179d1b --- /dev/null +++ b/internal/registry/web_search_index_test.go @@ -0,0 +1,396 @@ +package registry + +import ( + "fmt" + "sync" + "testing" +) + +// referenceHasClientSupportingWebSearch recomputes the web-search aggregate with +// the original full scan. The index must agree with it in every state the +// registry can reach, so tests compare the fast path against this oracle +// instead of only asserting hand-picked expectations. +func referenceHasClientSupportingWebSearch(r *ModelRegistry, modelID, provider, excludeClientID string) bool { + for clientID, infos := range r.clientModelInfos { + if clientID == excludeClientID { + continue + } + if provider != "" && r.clientProviders[clientID] != provider { + continue + } + info := infos[modelID] + if info != nil && info.SupportsWebSearch { + return true + } + } + return false +} + +// assertWebSearchIndexMatchesScan verifies the incremental index against the +// full scan for every client/model/provider combination currently registered. +func assertWebSearchIndexMatchesScan(t *testing.T, r *ModelRegistry, context string) { + t.Helper() + r.mutex.RLock() + defer r.mutex.RUnlock() + + providers := make(map[string]struct{}) + modelIDs := make(map[string]struct{}) + for clientID, infos := range r.clientModelInfos { + providers[r.clientProviders[clientID]] = struct{}{} + for modelID := range infos { + modelIDs[modelID] = struct{}{} + } + } + + for modelID := range modelIDs { + for provider := range providers { + for _, exclude := range []string{"", "__none__"} { + want := referenceHasClientSupportingWebSearch(r, modelID, provider, exclude) + got := r.hasClientSupportingWebSearchLocked(modelID, provider, exclude) + if got != want { + t.Fatalf("%s: model=%q provider=%q exclude=%q: index=%v scan=%v", context, modelID, provider, exclude, got, want) + } + } + } + // Cross-check per-client exclusion, which the index answers by + // subtracting the excluded client instead of skipping it. + for clientID := range r.clientModelInfos { + want := referenceHasClientSupportingWebSearch(r, modelID, "", clientID) + got := r.hasClientSupportingWebSearchLocked(modelID, "", clientID) + if got != want { + t.Fatalf("%s: model=%q exclude=%q: index=%v scan=%v", context, modelID, clientID, got, want) + } + } + + // The published Info flags are what callers actually read, so verify + // them against the scan too: a self-consistent index would still be a + // bug if the stored aggregate had been computed from a stale one. + registration := r.models[modelID] + if registration == nil || registration.Info == nil { + continue + } + if want, got := referenceHasClientSupportingWebSearch(r, modelID, "", ""), registration.Info.SupportsWebSearch; got != want { + t.Fatalf("%s: model=%q Info.SupportsWebSearch=%v scan=%v", context, modelID, got, want) + } + for provider, info := range registration.InfoByProvider { + if info == nil { + continue + } + if want, got := referenceHasClientSupportingWebSearch(r, modelID, provider, ""), info.SupportsWebSearch; got != want { + t.Fatalf("%s: model=%q provider=%q InfoByProvider.SupportsWebSearch=%v scan=%v", context, modelID, provider, got, want) + } + } + } +} + +func newWebSearchIndexRegistry() *ModelRegistry { + reg := newTestModelRegistry() + reg.RegisterClient("a", "codex", []*ModelInfo{ + {ID: "solo", SupportsWebSearch: true}, + {ID: "shared", SupportsWebSearch: true}, + }) + reg.RegisterClient("b", "codex", []*ModelInfo{ + {ID: "shared", SupportsWebSearch: true}, + }) + reg.RegisterClient("c", "gemini", []*ModelInfo{ + {ID: "shared", SupportsWebSearch: true}, + {ID: "off"}, + }) + return reg +} + +func TestWebSearchIndexMatchesScanAcrossMutations(t *testing.T) { + reg := newWebSearchIndexRegistry() + assertWebSearchIndexMatchesScan(t, reg, "initial") + + // Re-registration may add or drop contributions; both paths must retract + // the old counts before writing the new ones. + reg.RegisterClient("b", "codex", []*ModelInfo{{ID: "shared"}}) + assertWebSearchIndexMatchesScan(t, reg, "re-register drops web search") + + reg.RegisterClient("b", "codex", []*ModelInfo{{ID: "shared", SupportsWebSearch: true}}) + assertWebSearchIndexMatchesScan(t, reg, "re-register restores web search") + + reg.RegisterClient("d", "codex", []*ModelInfo{{ID: "solo", SupportsWebSearch: true}}) + assertWebSearchIndexMatchesScan(t, reg, "second supporter of solo") + + reg.ApplyClientModelCapabilities("d", reg.ClientRegistrationEpoch("d"), func(_ string, info *ModelInfo) { + info.SupportsWebSearch = false + }) + assertWebSearchIndexMatchesScan(t, reg, "capabilities clear web search") + + reg.ApplyClientModelCapabilities("a", reg.ClientRegistrationEpoch("a"), func(_ string, info *ModelInfo) { + info.SupportsWebSearch = true + }) + assertWebSearchIndexMatchesScan(t, reg, "capabilities set web search") + + reg.UnregisterClient("a") + assertWebSearchIndexMatchesScan(t, reg, "unregister sole supporter") + + reg.UnregisterClient("d") + reg.UnregisterClient("b") + reg.UnregisterClient("c") + assertWebSearchIndexMatchesScan(t, reg, "unregister all") +} + +func TestWebSearchIndexDropsOrphanCountsOnPureAddition(t *testing.T) { + // The pure-addition path assumes the client has no prior bookkeeping, but + // clientModelInfos can outlive clientModels (any code path that rewrites one + // without the other). A stale entry still contributes to the index, so + // registering the client again must retract it instead of counting it twice. + reg := newTestModelRegistry() + reg.mutex.Lock() + reg.clientModelInfos["ghost"] = map[string]*ModelInfo{ + "orphaned": {ID: "orphaned", SupportsWebSearch: true}, + } + reg.clientProviders["ghost"] = "codex" + // The orphan is seeded in the shape any buggy path would leave behind: the + // index still agrees with clientModelInfos, but clientModels no longer + // carries the client, so the next registration takes the pure-addition path. + reg.incrementWebSearchCountLocked("orphaned", "codex") + reg.mutex.Unlock() + assertWebSearchIndexMatchesScan(t, reg, "orphan seeded") + + // Registering the same client with a catalog that advertises no web search + // must leave a clean index, not the orphaned contribution plus nothing. + reg.RegisterClient("ghost", "codex", []*ModelInfo{{ID: "other"}}) + assertWebSearchIndexMatchesScan(t, reg, "orphan replaced by pure addition") + if reg.hasClientSupportingWebSearchLocked("orphaned", "", "") { + t.Fatalf("index kept the orphaned web-search contribution after re-registration") + } + if reg.hasClientSupportingWebSearchLocked("orphaned", "codex", "") { + t.Fatalf("provider index kept the orphaned web-search contribution after re-registration") + } +} + +func assertPublishedWebSearchSupport(t *testing.T, reg *ModelRegistry, modelID, provider string, want bool) { + t.Helper() + for _, view := range []string{"", provider} { + info := reg.GetModelInfo(modelID, view) + if info == nil { + t.Errorf("GetModelInfo(%q, %q) returned nil", modelID, view) + continue + } + if info.SupportsWebSearch != want { + t.Errorf("GetModelInfo(%q, %q).SupportsWebSearch = %v, want %v", modelID, view, info.SupportsWebSearch, want) + } + } +} + +func TestWebSearchIndexEmptyClientExclusion(t *testing.T) { + for _, clientID := range []string{"", "supporter"} { + t.Run(fmt.Sprintf("client=%q", clientID), func(t *testing.T) { + reg := newTestModelRegistry() + reg.RegisterClient(clientID, "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + reg.RegisterClient("observer", "antigravity", []*ModelInfo{{ID: "m"}}) + + reg.mutex.RLock() + defer reg.mutex.RUnlock() + for _, provider := range []string{"", "antigravity", "gemini"} { + for _, exclude := range []string{"", "supporter", "observer", "missing"} { + want := provider != "gemini" && clientID != exclude + if got := referenceHasClientSupportingWebSearch(reg, "m", provider, exclude); got != want { + t.Errorf("provider=%q exclude=%q: scan=%v, want %v", provider, exclude, got, want) + } + if got := reg.hasClientSupportingWebSearchLocked("m", provider, exclude); got != want { + t.Errorf("provider=%q exclude=%q: index=%v, want %v", provider, exclude, got, want) + } + } + } + }) + } +} + +func TestWebSearchIndexPublishedViewsAfterSupporterUnregister(t *testing.T) { + for _, clientID := range []string{"", "supporter"} { + t.Run(fmt.Sprintf("client=%q", clientID), func(t *testing.T) { + reg := newTestModelRegistry() + reg.RegisterClient(clientID, "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + reg.RegisterClient("observer", "antigravity", []*ModelInfo{{ID: "m"}}) + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true) + + reg.UnregisterClient(clientID) + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", false) + assertWebSearchIndexMatchesScan(t, reg, "supporter unregistered") + }) + } +} + +func TestWebSearchIndexPreservesRealSupporterOnOrphanAddition(t *testing.T) { + reg := newTestModelRegistry() + reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + // Artificially leave orphaned metadata and index entries. Normal registry + // operations keep these maps together; this exercises defensive cleanup. + reg.mutex.Lock() + delete(reg.clientModels, "ghost") + reg.mutex.Unlock() + assertWebSearchIndexMatchesScan(t, reg, "orphan seeded beside real supporter") + + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m"}}) + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true) + assertWebSearchIndexMatchesScan(t, reg, "orphan replaced beside real supporter") +} + +func TestWebSearchIndexRefreshesOrphanedModelsOnPureAddition(t *testing.T) { + for _, realSupportsWebSearch := range []bool{false, true} { + t.Run(fmt.Sprintf("real_supports_web_search=%v", realSupportsWebSearch), func(t *testing.T) { + reg := newTestModelRegistry() + reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: realSupportsWebSearch}}) + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + // Artificially orphan the old catalog to exercise defensive cleanup. + reg.mutex.Lock() + delete(reg.clientModels, "ghost") + reg.mutex.Unlock() + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", true) + assertWebSearchIndexMatchesScan(t, reg, "orphan seeded before catalog replacement") + + // The old model is absent from the new catalog, so additions cannot + // refresh its published capability as a side effect. + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "n"}}) + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch) + assertPublishedWebSearchSupport(t, reg, "n", "antigravity", false) + assertWebSearchIndexMatchesScan(t, reg, "orphan replaced with a different catalog") + }) + } +} + +func TestWebSearchIndexCleansOrphanOnUnregister(t *testing.T) { + for _, realSupportsWebSearch := range []bool{false, true} { + t.Run(fmt.Sprintf("real_supports_web_search=%v", realSupportsWebSearch), func(t *testing.T) { + reg := newTestModelRegistry() + reg.RegisterClient("real", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: realSupportsWebSearch}}) + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m", SupportsWebSearch: true}}) + // This inconsistent state is constructed only to test orphan cleanup. + reg.mutex.Lock() + delete(reg.clientModels, "ghost") + reg.mutex.Unlock() + + reg.UnregisterClient("ghost") + reg.mutex.RLock() + if _, exists := reg.clientModelInfos["ghost"]; exists { + t.Error("unregister left orphaned client model metadata") + } + reg.mutex.RUnlock() + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch) + + // Repeated cleanup must not retract the real client's contribution. + reg.UnregisterClient("ghost") + reg.mutex.RLock() + wantCount := 0 + if realSupportsWebSearch { + wantCount = 1 + } + if got := reg.webSearchClientCounts["m"]; got != wantCount { + t.Errorf("global count after repeated orphan cleanup = %d, want %d", got, wantCount) + } + if got := reg.webSearchProviderCounts["m"]["antigravity"]; got != wantCount { + t.Errorf("provider count after repeated orphan cleanup = %d, want %d", got, wantCount) + } + reg.mutex.RUnlock() + assertWebSearchIndexMatchesScan(t, reg, "orphan unregistered twice") + + reg.RegisterClient("ghost", "antigravity", []*ModelInfo{{ID: "m"}}) + assertPublishedWebSearchSupport(t, reg, "m", "antigravity", realSupportsWebSearch) + assertWebSearchIndexMatchesScan(t, reg, "orphan re-registered after cleanup") + }) + } +} + +func TestWebSearchIndexTracksSoleSupporterAfterUnregister(t *testing.T) { + reg := newWebSearchIndexRegistry() + if !reg.hasClientSupportingWebSearchLocked("solo", "", "") { + t.Fatalf("expected solo to be supported before its only supporter left") + } + reg.UnregisterClient("a") + if reg.hasClientSupportingWebSearchLocked("solo", "", "") { + t.Fatalf("sole supporter left, but the aggregate still reports web search") + } + // The published flag must follow the index, otherwise callers keep + // advertising web search for a model no client can serve. + if reg.models["solo"] != nil && reg.models["solo"].Info != nil && reg.models["solo"].Info.SupportsWebSearch { + t.Fatalf("published Info.SupportsWebSearch stayed true after the sole supporter left") + } +} + +func TestWebSearchIndexToleratesNilIndexMaps(t *testing.T) { + // The index maps are nil in a partially constructed registry (tests and any + // future literal). The helpers must lazily initialize instead of panicking + // on assignment to a nil map. + reg := &ModelRegistry{mutex: &sync.RWMutex{}} + if reg.hasClientSupportingWebSearchLocked("m", "", "") { + t.Fatalf("registry without index must not report web-search support") + } + reg.incrementWebSearchCountLocked("m", "codex") + if !reg.hasClientSupportingWebSearchLocked("m", "codex", "") { + t.Fatalf("expected the increment to be visible after lazy initialization") + } + reg.decrementWebSearchCountLocked("m", "codex") + if reg.hasClientSupportingWebSearchLocked("m", "codex", "") { + t.Fatalf("expected the decrement to clear the count") + } +} + +func TestWebSearchIndexProviderChangeRetractsOldProvider(t *testing.T) { + // Re-registering the same client under a different provider must retract + // the old provider's contribution. The old-provider aggregate is recomputed + // while this client's counts are still indexed, so a wrong drop/add ordering + // leaves a phantom count for the provider the client no longer serves. + reg := newTestModelRegistry() + reg.RegisterClient("x", "codex", []*ModelInfo{ + {ID: "m", SupportsWebSearch: true}, + }) + if !reg.hasClientSupportingWebSearchLocked("m", "codex", "") { + t.Fatalf("expected codex to support web search before the provider change") + } + + reg.RegisterClient("x", "gemini", []*ModelInfo{ + {ID: "m", SupportsWebSearch: true}, + }) + if reg.hasClientSupportingWebSearchLocked("m", "codex", "") { + t.Fatalf("codex must no longer report web search after the client moved to gemini") + } + if !reg.hasClientSupportingWebSearchLocked("m", "gemini", "") { + t.Fatalf("gemini must report web search after the provider change") + } + if got := reg.webSearchProviderCounts["m"]["codex"]; got != 0 { + t.Fatalf("stale codex provider count = %d, want 0", got) + } + if got := reg.webSearchProviderCounts["m"]["gemini"]; got != 1 { + t.Fatalf("gemini provider count = %d, want 1", got) + } + assertWebSearchIndexMatchesScan(t, reg, "provider change") +} + +// BenchmarkRegisterClientWebSearchAggregate measures whole-registry +// registration, the hot path that blocked startup: every registered model +// recomputes the web-search aggregate while the registry write lock is held. +// Sizes are reported separately because the pre-fix implementation is quadratic +// in the client count and only the growth curve exposes that. +func BenchmarkRegisterClientWebSearchAggregate(b *testing.B) { + const modelsPerClient = 20 + + infos := make([]*ModelInfo, 0, modelsPerClient) + for i := 0; i < modelsPerClient; i++ { + // Only the first model advertises web search, and it must be set on + // SupportsWebSearch: the aggregate is keyed off that field, not off + // NativeCapabilities.WebSearch. Leaving it false makes every lookup miss + // the client set, which silently benchmarks a different code path. + infos = append(infos, &ModelInfo{ + ID: fmt.Sprintf("model-%d", i), + SupportsWebSearch: i == 0, + }) + } + + for _, clients := range []int{250, 500, 1000} { + b.Run(fmt.Sprintf("clients=%d", clients), func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + reg := newTestModelRegistry() + for c := 0; c < clients; c++ { + reg.RegisterClient(fmt.Sprintf("client-%d", c), "codex", infos) + } + } + }) + } +} diff --git a/internal/runtime/executor/aistudio_executor.go b/internal/runtime/executor/aistudio_executor.go index 7969fb1ca3c..f2dc2b6b332 100644 --- a/internal/runtime/executor/aistudio_executor.go +++ b/internal/runtime/executor/aistudio_executor.go @@ -201,7 +201,9 @@ func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, reporter.PublishFailureWithDetail(ctx, helps.ParseGeminiUsage(wsResp.Body), err) return cliproxyexecutor.Response{}, err } - reporter.Publish(ctx, helps.ParseGeminiUsage(wsResp.Body)) + gemDetail := helps.ParseGeminiUsage(wsResp.Body) + helps.ApplyGeminiRequestImageConfig(&gemDetail, body.payload) + reporter.Publish(ctx, gemDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } @@ -328,6 +330,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth } claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, body.toFormat, responseFormat, originalRequest) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(body.payload) defer streamUsage.Publish(ctx, reporter) var streamErr error defer func() { @@ -336,7 +339,36 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth var param any 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) { + reporter.ObserveResponseModel(frame) + streamUsage.ObserveGeminiStream(frame) + filtered := helps.FilterSSEUsageMetadata(frame) + lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, filtered, ¶m, claudeInputTokens) + helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) + helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { + return statusErr{code: http.StatusBadGateway, msg: msg} + }) + for i := range lines { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: + case <-ctx.Done(): + return false + } + } + if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) { + return false + } + } + } + return true + } finishStream := func() { + if !processRelayFrames(relayFrames.Finish()) { + return + } if helps.EndApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) { return } @@ -385,29 +417,8 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth if len(event.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) - // A relay chunk may hold several SSE frames; the translators read one frame per call (gjson parses the first - // JSON object only), so each frame goes through the whole per-frame path: model, usage, translation, errors. - for _, frame := range splitRelayFrames(event.Payload) { - reporter.ObserveResponseModel(frame) - if detail, ok := helps.ParseGeminiStreamUsage(frame); ok { - streamUsage.Observe(detail, true) - } - filtered := helps.FilterSSEUsageMetadata(frame) - lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translatedReq, filtered, ¶m, claudeInputTokens) - helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) - helps.CaptureTranslatedClaudeStreamError(responseFormat, lines, &streamErr, func(msg string) error { - return statusErr{code: http.StatusBadGateway, msg: msg} - }) - for i := range lines { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: - case <-ctx.Done(): - return false - } - } - if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) { - return false - } + if !processRelayFrames(relayFrames.Feed(event.Payload)) { + return false } break } diff --git a/internal/runtime/executor/aistudio_fragmented_test.go b/internal/runtime/executor/aistudio_fragmented_test.go new file mode 100644 index 00000000000..5d2d1f4e92a --- /dev/null +++ b/internal/runtime/executor/aistudio_fragmented_test.go @@ -0,0 +1,48 @@ +package executor + +import ( + "strings" + "testing" + + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" +) + +func TestAIStudioStreamReassemblesFragmentedFrames(t *testing.T) { + terminal := strings.Replace(multiFrameB, `"text":"BBB"`, `"text":"BBB","thoughtSignature":"`+strings.Repeat("signature", 512)+`"`, 1) + frame := "data: " + terminal + "\n\n" + crlf := "data: " + terminal + "\r\n\r\n" + for _, format := range []sdktranslator.Format{ + sdktranslator.FormatOpenAI, + sdktranslator.FormatOpenAIResponse, + sdktranslator.FormatClaude, + } { + t.Run(format.String(), func(t *testing.T) { + for name, fragments := range map[string][]string{ + "large signature": {frame[:1181], frame[1181:2362], frame[2362:3543], frame[3543:]}, + "split data prefix": {frame[:2], frame[2:4], frame[4:]}, + "split crlf": {crlf[:len(crlf)-3], crlf[len(crlf)-3 : len(crlf)-1], crlf[len(crlf)-1:]}, + "unterminated tail": {frame[:1181], strings.TrimSuffix(frame[1181:], "\n\n")}, + } { + t.Run(name, func(t *testing.T) { + chunks := append([]string{"data: " + multiFrameA + "\n\n"}, fragments...) + got, errGot := streamThroughRelay(t, format, chunks) + if errGot != nil { + t.Fatal(errGot) + } + if !strings.Contains(got, "AAA") || !strings.Contains(got, "BBB") { + t.Errorf("fragmented stream lost content: %s", got) + } + terminalMarker := `"finish_reason": "stop"` + if format == sdktranslator.FormatOpenAIResponse { + terminalMarker = "event: response.completed" + } else if format == sdktranslator.FormatClaude { + terminalMarker = "event: message_stop" + } + if !strings.Contains(got, terminalMarker) { + t.Errorf("fragmented stream lacks %q: %s", terminalMarker, got) + } + }) + } + }) + } +} diff --git a/internal/runtime/executor/antigravity_bench_test.go b/internal/runtime/executor/antigravity_bench_test.go new file mode 100644 index 00000000000..d9ede8e1bd5 --- /dev/null +++ b/internal/runtime/executor/antigravity_bench_test.go @@ -0,0 +1,46 @@ +package executor + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + antigravitygemini "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/antigravity/gemini" +) + +func buildSynthetic4KAntigravitySSE() ([]byte, []byte) { + // ~8MB base64 payload representing a 4K JPEG/PNG inlineData response. + largeB64 := strings.Repeat("QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVowMTIzNDU2Nzg5", 160000) + var line1 bytes.Buffer + line1.WriteString(`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_abc","inlineData":{"mimeType":"image/jpeg","data":"`) + line1.WriteString(largeB64) + line1.WriteString(`"}}]}}],"usageMetadata":{"promptTokenCount":15,"candidatesTokenCount":2000,"totalTokenCount":2015},"modelVersion":"gemini-3.1-flash-image","responseId":"resp_1"},"traceId":"trace_1","metadata":{}}`) + line2 := []byte(`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":15,"candidatesTokenCount":2000,"totalTokenCount":2015},"modelVersion":"gemini-3.1-flash-image","responseId":"resp_1"},"traceId":"trace_1","metadata":{}}`) + return line1.Bytes(), line2 +} + +func BenchmarkAntigravity4KImage_NonStreamPipeline(b *testing.B) { + line1, line2 := buildSynthetic4KAntigravitySSE() + reqPayload := []byte(`{"request":{"generationConfig":{"responseModalities":["TEXT","IMAGE"],"imageConfig":{"aspectRatio":"16:9","imageSize":"4K"}}}}`) + exec := &AntigravityExecutor{} + b.SetBytes(int64(len(line1) + len(line2))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rawPayloads := make([][]byte, 0, 2) + for _, rawLine := range [][]byte{line1, line2} { + if payload := helps.JSONPayload(rawLine); payload != nil { + rawPayloads = append(rawPayloads, payload) + } + } + assembled := exec.convertStreamChunksToNonStream(rawPayloads) + converted := antigravitygemini.ConvertAntigravityResponseToGeminiNonStream(context.Background(), "gemini-3.1-flash-image", reqPayload, reqPayload, assembled, nil) + detail := helps.ParseAntigravityUsage(assembled) + helps.ApplyGeminiRequestImageConfig(&detail, reqPayload) + if len(converted) == 0 || detail.ImageCount == nil || *detail.ImageCount != 1 { + b.Fatal("unexpected result") + } + } +} diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go index e1eba61172b..235c193f69a 100644 --- a/internal/runtime/executor/antigravity_executor_execute.go +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -14,6 +14,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" "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" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" @@ -202,7 +203,9 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au if helps.ApplyPatchTranslationError(param) != nil || len(converted) == 0 { return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} } - reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) + antiDetail := helps.ParseAntigravityUsage(bodyBytes) + helps.ApplyGeminiRequestImageConfig(&antiDetail, requestPayload) + reporter.Publish(ctx, antiDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { converted = helps.EnsureResponsesUsageDetails(converted) } @@ -418,57 +421,35 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * clearAntigravityCreditsFailureState(auth) } replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) - out := make(chan cliproxyexecutor.StreamChunk) - go func(resp *http.Response) { - defer close(out) - defer func() { - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - }() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, streamScannerBuffer) - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if replayAccumulator != nil { - replayAccumulator.ObserveSSELine(line) - } - - // Filter usage metadata for all models - // Only retain usage statistics in the terminal chunk - line = helps.FilterSSEUsageMetadata(line) - - payload := helps.JSONPayload(line) - if payload == nil { - continue - } - reporter.ObserveResponseModel(payload) - - out <- cliproxyexecutor.StreamChunk{Payload: payload} - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - out <- cliproxyexecutor.StreamChunk{Err: errScan} - } else { - if replayAccumulator != nil { - replayAccumulator.Commit(ctx) - } - } - }(httpResp) - - var buffer bytes.Buffer - for chunk := range out { - if chunk.Err != nil { - return resp, chunk.Err - } - if len(chunk.Payload) > 0 { - _, _ = buffer.Write(chunk.Payload) - _, _ = buffer.Write([]byte("\n")) + var rawPayloads [][]byte + scanner := bufio.NewScanner(httpResp.Body) + scanner.Buffer(nil, streamScannerBuffer) + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } + payload := helps.JSONPayload(line) + if payload == nil { + continue } + reporter.ObserveResponseModel(payload) + rawPayloads = append(rawPayloads, bytes.Clone(payload)) } - resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} + scanErr := scanner.Err() + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if scanErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, scanErr) + reporter.PublishFailure(ctx, scanErr) + return resp, scanErr + } + if replayAccumulator != nil { + replayAccumulator.Commit(ctx) + } + resp = cliproxyexecutor.Response{Payload: e.convertStreamChunksToNonStream(rawPayloads)} resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload) reporter.ObserveResponseModel(resp.Payload) @@ -477,7 +458,9 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * if helps.ApplyPatchTranslationError(param) != nil || len(converted) == 0 { return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} } - reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload)) + antiDetail := helps.ParseAntigravityUsage(resp.Payload) + helps.ApplyGeminiRequestImageConfig(&antiDetail, requestPayload) + reporter.Publish(ctx, antiDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { converted = helps.EnsureResponsesUsageDetails(converted) } @@ -487,18 +470,43 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * return resp, nil } +const antigravityPartsPlaceholder = `"__cpa_parts_placeholder__"` + func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { - responseTemplate := "" - var traceID string - var finishReason string - var modelVersion string - var responseID string - var role string - var usageRaw string - parts := make([]map[string]interface{}, 0) - var pendingKind string - var pendingText strings.Builder - var pendingThoughtSig string + var chunks [][]byte + remaining := stream + for len(remaining) > 0 { + var line []byte + if nl := bytes.IndexByte(remaining, '\n'); nl >= 0 { + line = remaining[:nl] + remaining = remaining[nl+1:] + } else { + line = remaining + remaining = nil + } + if trimmed := bytes.TrimSpace(line); len(trimmed) > 0 { + chunks = append(chunks, trimmed) + } + } + return e.convertStreamChunksToNonStream(chunks) +} + +func (e *AntigravityExecutor) convertStreamChunksToNonStream(chunks [][]byte) []byte { + var ( + responseNodeRaw string + responsePartsRaw string + responsePartsIndex int + traceID string + finishReason string + modelVersion string + responseID string + role string + usageRaw string + parts []string + pendingKind string + pendingText strings.Builder + pendingThoughtSig string + ) flushPending := func() { if pendingKind == "" { @@ -513,7 +521,8 @@ func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { pendingThoughtSig = "" return } - parts = append(parts, map[string]interface{}{"text": text}) + encoded, _ := json.Marshal(map[string]any{"text": text}) + parts = append(parts, string(encoded)) case "thought": if strings.TrimSpace(text) == "" && pendingThoughtSig == "" { pendingKind = "" @@ -521,46 +530,45 @@ func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { pendingThoughtSig = "" return } - part := map[string]interface{}{"thought": true} - part["text"] = text + part := map[string]any{"thought": true, "text": text} if pendingThoughtSig != "" { part["thoughtSignature"] = pendingThoughtSig } - parts = append(parts, part) + encoded, _ := json.Marshal(part) + parts = append(parts, string(encoded)) } pendingKind = "" pendingText.Reset() pendingThoughtSig = "" } - normalizePart := func(partResult gjson.Result) map[string]interface{} { - var m map[string]interface{} - _ = json.Unmarshal([]byte(partResult.Raw), &m) - if m == nil { - m = map[string]interface{}{} - } - sig := partResult.Get("thoughtSignature").String() - if sig == "" { - sig = partResult.Get("thought_signature").String() + normalizeRawPart := func(partResult gjson.Result, hasSnakeSig, hasSnakeInline bool, sig string, inlineRaw string) string { + if !hasSnakeSig && !hasSnakeInline { + return partResult.Raw } - if sig != "" { - m["thoughtSignature"] = sig - delete(m, "thought_signature") + raw := []byte(partResult.Raw) + if hasSnakeSig { + if sig != "" { + raw, _ = sjson.SetBytes(raw, "thoughtSignature", sig) + } + raw, _ = sjson.DeleteBytes(raw, "thought_signature") } - if inlineData, ok := m["inline_data"]; ok { - m["inlineData"] = inlineData - delete(m, "inline_data") + if hasSnakeInline { + if inlineRaw != "" { + raw, _ = sjson.SetRawBytes(raw, "inlineData", []byte(inlineRaw)) + } + raw, _ = sjson.DeleteBytes(raw, "inline_data") } - return m + return string(raw) } - for _, line := range bytes.Split(stream, []byte("\n")) { + for _, line := range chunks { trimmed := bytes.TrimSpace(line) if len(trimmed) == 0 || !gjson.ValidBytes(trimmed) { continue } - root := gjson.ParseBytes(trimmed) + root := util.GetGJSONBytesNoCopy(trimmed, "@this") responseNode := root.Get("response") if !responseNode.Exists() { if root.Get("candidates").Exists() { @@ -569,50 +577,113 @@ func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { continue } } - responseTemplate = responseNode.Raw + responseNodeRaw = responseNode.Raw + responsePartsRaw = "" + responsePartsIndex = -1 if traceResult := root.Get("traceId"); traceResult.Exists() && traceResult.String() != "" { traceID = traceResult.String() } - if roleResult := responseNode.Get("candidates.0.content.role"); roleResult.Exists() { - role = roleResult.String() - } - - if finishResult := responseNode.Get("candidates.0.finishReason"); finishResult.Exists() && finishResult.String() != "" { - finishReason = finishResult.String() - } - - if modelResult := responseNode.Get("modelVersion"); modelResult.Exists() && modelResult.String() != "" { - modelVersion = modelResult.String() - } - if responseIDResult := responseNode.Get("responseId"); responseIDResult.Exists() && responseIDResult.String() != "" { - responseID = responseIDResult.String() - } - if usageResult := responseNode.Get("usageMetadata"); usageResult.Exists() { - usageRaw = usageResult.Raw - } else if usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() { - usageRaw = usageMetadataResult.Raw + var partsResult gjson.Result + responseNode.ForEach(func(rk, rv gjson.Result) bool { + switch rk.String() { + case "modelVersion": + if rv.String() != "" { + modelVersion = rv.String() + } + case "responseId": + if rv.String() != "" { + responseID = rv.String() + } + case "usageMetadata", "cpaUsageMetadata": + if rk.String() == "usageMetadata" || usageRaw == "" { + usageRaw = rv.Raw + } + case "candidates": + if rv.IsArray() { + cand0 := rv.Get("0") + if cand0.IsObject() { + cand0.ForEach(func(ck, cv gjson.Result) bool { + switch ck.String() { + case "finishReason": + if cv.String() != "" { + finishReason = cv.String() + } + case "content": + if cv.IsObject() { + cv.ForEach(func(ctk, ctv gjson.Result) bool { + switch ctk.String() { + case "role": + role = ctv.String() + case "parts": + partsResult = ctv + responsePartsRaw = ctv.Raw + responsePartsIndex = ctv.Index - responseNode.Index + } + return true + }) + } + } + return true + }) + } + } + } + return true + }) + if usageRaw == "" { + if usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() { + usageRaw = usageMetadataResult.Raw + } } - if partsResult := responseNode.Get("candidates.0.content.parts"); partsResult.IsArray() { - for _, part := range partsResult.Array() { - hasFunctionCall := part.Get("functionCall").Exists() - hasInlineData := part.Get("inlineData").Exists() || part.Get("inline_data").Exists() - sig := part.Get("thoughtSignature").String() - if sig == "" { - sig = part.Get("thought_signature").String() - } - text := part.Get("text").String() - thought := part.Get("thought").Bool() + if partsResult.IsArray() { + partsResult.ForEach(func(_, part gjson.Result) bool { + var ( + hasFunctionCall bool + hasInlineData bool + hasSnakeInline bool + hasSnakeSig bool + hasText bool + sig string + inlineRaw string + text string + thought bool + ) + part.ForEach(func(pk, pv gjson.Result) bool { + switch pk.String() { + case "functionCall": + hasFunctionCall = true + case "inlineData": + hasInlineData = true + case "inline_data": + hasInlineData = true + hasSnakeInline = true + inlineRaw = pv.Raw + case "thoughtSignature": + sig = pv.String() + case "thought_signature": + hasSnakeSig = true + if sig == "" { + sig = pv.String() + } + case "text": + hasText = true + text = pv.String() + case "thought": + thought = pv.Bool() + } + return true + }) if hasFunctionCall || hasInlineData { flushPending() - parts = append(parts, normalizePart(part)) - continue + parts = append(parts, normalizeRawPart(part, hasSnakeSig, hasSnakeInline, sig, inlineRaw)) + return true } - if thought || part.Get("text").Exists() { + if thought || hasText { kind := "text" if thought { kind = "thought" @@ -625,57 +696,97 @@ func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { if kind == "thought" && sig != "" { pendingThoughtSig = sig } - continue + return true } flushPending() - parts = append(parts, normalizePart(part)) - } + parts = append(parts, normalizeRawPart(part, hasSnakeSig, hasSnakeInline, sig, inlineRaw)) + return true + }) } } flushPending() - if responseTemplate == "" { - responseTemplate = `{"candidates":[{"content":{"role":"model","parts":[]}}]}` + var templateBytes []byte + if responseNodeRaw == "" { + templateBytes = []byte(`{"candidates":[{"content":{"role":"model","parts":` + antigravityPartsPlaceholder + `}}]}`) + } else if responsePartsRaw != "" { + if idx := responsePartsIndex; idx >= 0 && idx+len(responsePartsRaw) <= len(responseNodeRaw) { + templateBytes = make([]byte, 0, len(responseNodeRaw)-len(responsePartsRaw)+len(antigravityPartsPlaceholder)) + templateBytes = append(templateBytes, responseNodeRaw[:idx]...) + templateBytes = append(templateBytes, antigravityPartsPlaceholder...) + templateBytes = append(templateBytes, responseNodeRaw[idx+len(responsePartsRaw):]...) + } else { + templateBytes, _ = sjson.SetRawBytes([]byte(responseNodeRaw), "candidates.0.content.parts", []byte(antigravityPartsPlaceholder)) + } + } else { + templateBytes, _ = sjson.SetRawBytes([]byte(responseNodeRaw), "candidates.0.content.parts", []byte(antigravityPartsPlaceholder)) } - partsJSON, _ := json.Marshal(parts) - updatedTemplate, _ := sjson.SetRawBytes([]byte(responseTemplate), "candidates.0.content.parts", partsJSON) - responseTemplate = string(updatedTemplate) if role != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.content.role", role) - responseTemplate = string(updatedTemplate) + templateBytes, _ = sjson.SetBytes(templateBytes, "candidates.0.content.role", role) } if finishReason != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.finishReason", finishReason) - responseTemplate = string(updatedTemplate) + templateBytes, _ = sjson.SetBytes(templateBytes, "candidates.0.finishReason", finishReason) } if modelVersion != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "modelVersion", modelVersion) - responseTemplate = string(updatedTemplate) + templateBytes, _ = sjson.SetBytes(templateBytes, "modelVersion", modelVersion) } if responseID != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "responseId", responseID) - responseTemplate = string(updatedTemplate) + templateBytes, _ = sjson.SetBytes(templateBytes, "responseId", responseID) } if usageRaw != "" { - updatedTemplate, _ = sjson.SetRawBytes([]byte(responseTemplate), "usageMetadata", []byte(usageRaw)) - responseTemplate = string(updatedTemplate) - } else if !gjson.Get(responseTemplate, "usageMetadata").Exists() { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.promptTokenCount", 0) - responseTemplate = string(updatedTemplate) - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.candidatesTokenCount", 0) - responseTemplate = string(updatedTemplate) - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.totalTokenCount", 0) - responseTemplate = string(updatedTemplate) - } - - output := `{"response":{},"traceId":""}` - updatedOutput, _ := sjson.SetRawBytes([]byte(output), "response", []byte(responseTemplate)) - output = string(updatedOutput) - if traceID != "" { - updatedOutput, _ = sjson.SetBytes([]byte(output), "traceId", traceID) - output = string(updatedOutput) - } - return []byte(output) + templateBytes, _ = sjson.SetRawBytes(templateBytes, "usageMetadata", []byte(usageRaw)) + } else if !gjson.GetBytes(templateBytes, "usageMetadata").Exists() { + templateBytes, _ = sjson.SetRawBytes(templateBytes, "usageMetadata", []byte(`{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}`)) + } + templateBytes, _ = sjson.DeleteBytes(templateBytes, "cpaUsageMetadata") + + partsTotalLen := 2 + for i, p := range parts { + if i > 0 { + partsTotalLen++ + } + partsTotalLen += len(p) + } + + traceJSON, _ := json.Marshal(traceID) + placeholder := gjson.GetBytes(templateBytes, "candidates.0.content.parts") + placeholderIdx := placeholder.Index + if placeholder.Raw != antigravityPartsPlaceholder || placeholderIdx <= 0 { + partsJSON := make([]byte, 0, partsTotalLen) + partsJSON = append(partsJSON, '[') + for i, p := range parts { + if i > 0 { + partsJSON = append(partsJSON, ',') + } + partsJSON = append(partsJSON, p...) + } + partsJSON = append(partsJSON, ']') + templateBytes, _ = sjson.SetRawBytes(templateBytes, "candidates.0.content.parts", partsJSON) + output := make([]byte, 0, len(templateBytes)+len(traceJSON)+28) + output = append(output, `{"response":`...) + output = append(output, templateBytes...) + output = append(output, `,"traceId":`...) + output = append(output, traceJSON...) + output = append(output, '}') + return output + } + + output := make([]byte, 0, len(templateBytes)-len(antigravityPartsPlaceholder)+partsTotalLen+len(traceJSON)+28) + output = append(output, `{"response":`...) + output = append(output, templateBytes[:placeholderIdx]...) + output = append(output, '[') + for i, p := range parts { + if i > 0 { + output = append(output, ',') + } + output = append(output, p...) + } + output = append(output, ']') + output = append(output, templateBytes[placeholderIdx+len(antigravityPartsPlaceholder):]...) + output = append(output, `,"traceId":`...) + output = append(output, traceJSON...) + output = append(output, '}') + return output } diff --git a/internal/runtime/executor/antigravity_executor_finish_reason_test.go b/internal/runtime/executor/antigravity_executor_finish_reason_test.go index c54065f65bf..3cbcccd46de 100644 --- a/internal/runtime/executor/antigravity_executor_finish_reason_test.go +++ b/internal/runtime/executor/antigravity_executor_finish_reason_test.go @@ -323,3 +323,26 @@ func TestAntigravityStreamSynthesizesOpenAIToolCallsFinishReason(t *testing.T) { t.Fatalf("expected exactly one finish_reason chunk, got %d", finishReasonCount) } } + +func TestAntigravityAssemblyReplacesOnlyContentParts(t *testing.T) { + for _, metadata := range []string{ + `"groundingMetadata":{"webSearchQueries":[]}`, + `"marker":"__cpa_parts_placeholder__"`, + } { + t.Run(metadata, func(t *testing.T) { + raw := []byte(`{"response":{"candidates":[{` + metadata + `,"content":{"role":"model","parts":[]},"finishReason":"STOP"}]}}`) + exec := &AntigravityExecutor{} + first := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}}`) + out := exec.convertStreamChunksToNonStream([][]byte{first, raw}) + if !gjson.ValidBytes(out) || gjson.GetBytes(out, "response.candidates.0.content.parts.0.text").String() != "hello" { + t.Fatalf("assembled parts replaced the wrong value: %s", out) + } + before := gjson.GetBytes(raw, "response.candidates.0") + after := gjson.GetBytes(out, "response.candidates.0") + if before.Get("groundingMetadata").Raw != after.Get("groundingMetadata").Raw || + before.Get("marker").Raw != after.Get("marker").Raw { + t.Fatalf("assembly changed unrelated metadata: %s", out) + } + }) + } +} diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index bba3a9615f9..4b819b7b338 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -213,6 +213,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya scanner.Buffer(nil, streamScannerBuffer) claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(requestPayload) var streamErr error var pendingJSON []byte var replayCommitted bool @@ -268,7 +269,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya } // Capture accounting before the client-facing filter renames usage. - streamUsage.Observe(helps.ParseAntigravityStreamUsage(line)) + streamUsage.ObserveAntigravityStream(line) // Filter usage metadata for all models // Only retain usage statistics in the terminal chunk diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go index 8a9c808a787..f7cc78b59bc 100644 --- a/internal/runtime/executor/claude_executor_execute.go +++ b/internal/runtime/executor/claude_executor_execute.go @@ -370,7 +370,14 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r e.rememberClaudeOAuthToolAliases(body, oauthToolNamesReverseMap, msgID) } lines := bytes.Split(data, []byte("\n")) + var eventBuf bytes.Buffer for i, line := range lines { + eventBuf.Write(line) + eventBuf.WriteByte('\n') + if len(bytes.TrimSpace(line)) == 0 { + streamUsage.ObserveClaudeSSEEvent(eventBuf.Bytes()) + eventBuf.Reset() + } reporter.ObserveResponseModel(line) streamUsage.ObserveClaudeStream(line) restoredLine, errRestore := restoreClaudeOAuthToolNamesFromStreamLine(line, oauthToolNamesReverseMap) @@ -382,6 +389,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r } lines[i] = restoredLine } + if eventBuf.Len() > 0 { + streamUsage.ObserveClaudeSSEEvent(eventBuf.Bytes()) + eventBuf.Reset() + } data = bytes.Join(lines, []byte("\n")) } else { messageID := claudeMessageIDFromResponse(data) diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index c0ec95ca933..bb441d2c7bb 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -4,6 +4,7 @@ import ( "bufio" "bytes" "context" + "errors" "fmt" "io" "net/http" @@ -352,16 +353,25 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A defer streamUsage.Publish(ctx, reporter) emitCancellation := func(cause error) bool { cancelErr := newClaudeOAuthCancellationError(ctx, fp.OAuthCancellation, cause) - if cancelErr == nil { - return false + if cancelErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, cancelErr) + streamUsage.PublishFailure(ctx, reporter, cancelErr) + select { + case out <- cliproxyexecutor.StreamChunk{Err: cancelErr}: + default: + } + return true } - helps.RecordAPIResponseError(ctx, e.cfg, cancelErr) - streamUsage.PublishFailure(ctx, reporter, cancelErr) - select { - case out <- cliproxyexecutor.StreamChunk{Err: cancelErr}: - default: + c := cause + if ctx != nil && ctx.Err() != nil { + c = ctx.Err() + } + if errors.Is(c, context.Canceled) { + helps.RecordAPIResponseError(ctx, e.cfg, c) + streamUsage.PublishFailure(ctx, reporter, c) + return true } - return true + return false } emitResponseError := func(errResponse error) { errResponse = wrapClaudeFastRequestError(fastRequest, httpResp.StatusCode, errResponse) @@ -386,6 +396,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return true } observeClaudeSSEEvent(event.Bytes(), &upstreamMessageID, &upstreamCompleted, &upstreamStreamErr) + streamUsage.ObserveClaudeSSEEvent(event.Bytes()) cloned := bytes.Clone(event.Bytes()) event.Reset() select { @@ -466,6 +477,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A eventBuf.WriteByte('\n') if len(bytes.TrimSpace(line)) == 0 { observeClaudeSSEEvent(eventBuf.Bytes(), &upstreamMessageID, &upstreamCompleted, &upstreamStreamErr) + streamUsage.ObserveClaudeSSEEvent(eventBuf.Bytes()) eventBuf.Reset() } helps.AppendAPIResponseChunk(ctx, e.cfg, line) @@ -498,7 +510,14 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: case <-ctx.Done(): - emitCancellation(ctx.Err()) + if eventBuf.Len() > 0 { + observeClaudeSSEEvent(eventBuf.Bytes(), &upstreamMessageID, &upstreamCompleted, &upstreamStreamErr) + streamUsage.ObserveClaudeSSEEvent(eventBuf.Bytes()) + eventBuf.Reset() + } + if !upstreamCompleted { + emitCancellation(ctx.Err()) + } return } } @@ -511,6 +530,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } if eventBuf.Len() > 0 { observeClaudeSSEEvent(eventBuf.Bytes(), &upstreamMessageID, &upstreamCompleted, &upstreamStreamErr) + streamUsage.ObserveClaudeSSEEvent(eventBuf.Bytes()) eventBuf.Reset() } if helps.EndApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) { diff --git a/internal/runtime/executor/claude_executor_stream_truncation_test.go b/internal/runtime/executor/claude_executor_stream_truncation_test.go index 8a21bbb6e89..2fadc1580b0 100644 --- a/internal/runtime/executor/claude_executor_stream_truncation_test.go +++ b/internal/runtime/executor/claude_executor_stream_truncation_test.go @@ -935,3 +935,95 @@ func TestTranslatedClaudeStreamTruncationPreservesBufferedUsageOnFailure(t *test } }) } + +func TestClaudeExecutorStreamCancel499RecordsLastPacketUsage(t *testing.T) { + for _, mode := range []struct { + name string + format sdktranslator.Format + }{ + {name: "passthrough", format: sdktranslator.FormatClaude}, + {name: "translated", format: sdktranslator.FormatOpenAIResponse}, + } { + t.Run(mode.name, func(t *testing.T) { + deltaFlushed := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + _, _ = fmt.Fprint(w, "event: message_start\n") + _, _ = fmt.Fprint(w, `data: {"type":"message_start","message":{"id":"msg_cancel_499","type":"message","role":"assistant","content":[],"model":"claude-opus-5-5","stop_reason":null,"usage":{"input_tokens":4,"cache_creation_input_tokens":3036,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":3036},"output_tokens":1,"service_tier":"standard","inference_geo":"global","speed":"standard"}}}`+"\n\n") + _, _ = fmt.Fprint(w, "event: content_block_start\n") + _, _ = fmt.Fprint(w, `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`+"\n\n") + _, _ = fmt.Fprint(w, "event: content_block_delta\n") + _, _ = fmt.Fprint(w, `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`+"\n\n") + _, _ = fmt.Fprint(w, "event: message_delta\n") + _, _ = fmt.Fprint(w, `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":6,"cache_creation_input_tokens":8846,"cache_read_input_tokens":3036,"output_tokens":120,"output_tokens_details":{"thinking_tokens":30},"server_tool_use":{"web_search_requests":1,"web_fetch_requests":0},"iterations":[{"type":"message","input_tokens":4,"output_tokens":20,"cache_read_input_tokens":0,"cache_creation_input_tokens":3036,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":3036}},{"type":"message","input_tokens":2,"output_tokens":100,"cache_read_input_tokens":3036,"cache_creation_input_tokens":5810,"cache_creation":{"ephemeral_5m_input_tokens":5810,"ephemeral_1h_input_tokens":0}}]}}`+"\n\n") + if flusher != nil { + flusher.Flush() + } + close(deltaFlushed) + <-r.Context().Done() + })) + defer server.Close() + + pluginName := "test-claude-cancel-499-" + mode.name + plugin := &captureClaudeUsagePlugin{records: make(chan usage.Record, 4)} + usage.RegisterNamedPlugin(pluginName, plugin) + t.Cleanup(func() { + usage.RegisterNamedPlugin(pluginName, noopClaudeUsagePlugin{}) + }) + + exec := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-opus-5-5", + Payload: []byte(`{"model":"claude-opus-5-5","messages":[{"role":"user","content":"hi"}],"stream":true}`), + }, cliproxyexecutor.Options{ + SourceFormat: mode.format, + ResponseFormat: mode.format, + Stream: true, + }) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + <-deltaFlushed + // Drain until the message_delta chunk has been processed, then cancel before message_stop. + for chunk := range result.Chunks { + if strings.Contains(string(chunk.Payload), "message_delta") || strings.Contains(string(chunk.Payload), "partial") { + cancel() + break + } + } + for range result.Chunks { + } + + select { + case record := <-plugin.records: + if !record.Failed || record.Fail.StatusCode != 499 { + t.Fatalf("expected failed=true status=499 on canceled stream, got failed=%v fail=%+v", record.Failed, record.Fail) + } + if record.Detail.InputTokens != 6 || record.Detail.OutputTokens != 120 || record.Detail.ReasoningTokens != 30 || record.Detail.CacheReadTokens != 3036 || record.Detail.CacheCreationTokens != 8846 { + t.Fatalf("unexpected last-packet token counts on 499 cancel: %+v", record.Detail) + } + if record.Detail.CacheCreation5mTokens == nil || *record.Detail.CacheCreation5mTokens != 5810 || record.Detail.CacheCreation1hTokens == nil || *record.Detail.CacheCreation1hTokens != 3036 { + t.Fatalf("expected 5m=5810 1h=3036 on 499 cancel, got 5m=%v 1h=%v", record.Detail.CacheCreation5mTokens, record.Detail.CacheCreation1hTokens) + } + if record.Detail.WebSearchRequests == nil || *record.Detail.WebSearchRequests != 1 { + t.Fatalf("expected WebSearchRequests=1 on 499 cancel, got %v", record.Detail.WebSearchRequests) + } + if len(record.Detail.IterationSlice()) != 2 { + t.Fatalf("expected 2 iterations on 499 cancel, got %d", len(record.Detail.IterationSlice())) + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for 499 canceled usage record") + } + }) + } +} diff --git a/internal/runtime/executor/codex_executor_auth.go b/internal/runtime/executor/codex_executor_auth.go index a4e594393c7..dd1477d3455 100644 --- a/internal/runtime/executor/codex_executor_auth.go +++ b/internal/runtime/executor/codex_executor_auth.go @@ -292,10 +292,17 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au needSIWC = hasSIWC } } + if !diskRotatedOuter && auth.RejectedRefreshToken != "" && outerRefreshToken == strings.TrimSpace(auth.RejectedRefreshToken) { + needOuter = false + } + if !diskRotatedSIWC && auth.RejectedSIWCRefreshToken != "" && siwcRefreshToken == strings.TrimSpace(auth.RejectedSIWCRefreshToken) { + needSIWC = false + } if !needOuter && !needSIWC { if diskRotatedOuter || diskRotatedSIWC { patch := codexauth.NewCredentialRefreshPatchWithDiskBaseline(beforeRefresh.Metadata, diskBaseline, auth.Metadata, diskRotatedOuter, diskRotatedSIWC) + patch.SetUpstreamRefreshed(false) auth.CodexRefreshPatch = patch if storage, ok := auth.Storage.(*codexauth.CodexTokenStorage); ok && storage != nil { clonedStorage := *storage @@ -439,11 +446,8 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au } if !outerRefreshed && !siwcRefreshed && !diskRotatedOuter && !diskRotatedSIWC { - if outerErr != nil { - return nil, outerErr - } - if siwcErr != nil { - return nil, siwcErr + if errCombined := codexauth.NewPartitionRefreshError(outerErr, siwcErr); errCombined != nil { + return nil, errCombined } return nil, fmt.Errorf("codex executor: refresh failed") } @@ -455,6 +459,8 @@ func (e *CodexExecutor) refreshLocked(ctx context.Context, auth *cliproxyauth.Au } patch := codexauth.NewCredentialRefreshPatchWithDiskBaseline(beforeRefresh.Metadata, diskBaseline, auth.Metadata, outerRefreshed || diskRotatedOuter, siwcRefreshed || diskRotatedSIWC) + patch.SetUpstreamRefreshed(outerRefreshed || siwcRefreshed) + patch.SetPartitionErrors(outerErr, siwcErr) auth.CodexRefreshPatch = patch if storage, ok := auth.Storage.(*codexauth.CodexTokenStorage); ok && storage != nil { clonedStorage := *storage diff --git a/internal/runtime/executor/codex_executor_compact_test.go b/internal/runtime/executor/codex_executor_compact_test.go index c374ed98f19..1ca9b367a0e 100644 --- a/internal/runtime/executor/codex_executor_compact_test.go +++ b/internal/runtime/executor/codex_executor_compact_test.go @@ -1,13 +1,18 @@ package executor import ( + "bytes" "context" "io" "net/http" "net/http/httptest" + "strings" "testing" + "time" + "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "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" @@ -433,3 +438,376 @@ func TestCodexExecutorReplayedWebSearchCallHistoryEnsuresWebSearchTool_Issue6770 } }) } + +func TestCodexExecutorExpandsRecognizedCPACompactionCapsules_Issue7118(t *testing.T) { + summaryText := "Claude compacted summary before switching to GPT." + capsule, errSeal := helps.SealAntigravityCompaction(summaryText, "claude-sonnet-4-6") + if errSeal != nil { + t.Fatalf("SealAntigravityCompaction() error = %v", errSeal) + } + + mixedPayload := `{"model":"gpt-6.1-sol","input":[` + + `{"id":"cmp_ag_compact_1","type":"compaction","encrypted_content":"` + capsule + `"},` + + `{"id":"cmp_native_1","type":"compaction","encrypted_content":"gAAAAABnative-openai-compaction"},` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"continue on GPT"}]}` + + `]}` + + assertExpandedAndNativePreserved := func(t *testing.T, body []byte) { + t.Helper() + if bytes.Contains(body, []byte("cpa-ag-compact-v1:")) || bytes.Contains(body, []byte("cmp_ag_compact_1")) { + t.Fatalf("CPA compaction capsule leaked to Codex upstream: %s", body) + } + input := gjson.GetBytes(body, "input").Array() + if len(input) != 3 { + t.Fatalf("input len = %d, want 3; body=%s", len(input), body) + } + if input[0].Get("type").String() != "message" || input[0].Get("role").String() != "developer" { + t.Fatalf("input[0] = %s, want developer message", input[0].Raw) + } + if !strings.Contains(input[0].Get("content.0.text").String(), summaryText) { + t.Fatalf("input[0] text = %q, want summary %q", input[0].Get("content.0.text").String(), summaryText) + } + if input[1].Get("type").String() != "compaction" || input[1].Get("encrypted_content").String() != "gAAAAABnative-openai-compaction" { + t.Fatalf("native OpenAI compaction item not preserved at input[1]: %s", input[1].Raw) + } + } + + t.Run("ExecuteStream expands CPA capsule and preserves native OpenAI compaction", func(t *testing.T) { + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + gotBody = body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"ok\"}]}],\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12}}}\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }, + } + + streamResult, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: []byte(mixedPayload), + Stream: true, + }) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + } + assertExpandedAndNativePreserved(t, gotBody) + }) + + t.Run("Execute and executeCompact expand CPA capsule and preserve native OpenAI compaction", func(t *testing.T) { + var gotResponsesBody []byte + var gotCompactBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if r.URL.Path == "/responses/compact" { + gotCompactBody = body + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"resp_compact_1","object":"response.compaction","usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}`)) + return + } + gotResponsesBody = body + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12}}}\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }, + } + + _, errExec := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: []byte(mixedPayload), + }) + if errExec != nil { + t.Fatalf("Execute() error = %v", errExec) + } + assertExpandedAndNativePreserved(t, gotResponsesBody) + + _, errCompact := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: []byte(mixedPayload), + Alt: "responses/compact", + }) + if errCompact != nil { + t.Fatalf("Execute(responses/compact) error = %v", errCompact) + } + assertExpandedAndNativePreserved(t, gotCompactBody) + }) + + t.Run("prepareCodexWebsocketStream and CountTokens expand CPA capsule", func(t *testing.T) { + wsExecutor := NewCodexWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": "http://127.0.0.1:1", + "api_key": "test", + }, + } + prepared, errPrep := wsExecutor.prepareCodexWebsocketStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: []byte(mixedPayload), + Stream: true, + }) + if errPrep != nil { + t.Fatalf("prepareCodexWebsocketStream() error = %v", errPrep) + } + assertExpandedAndNativePreserved(t, prepared.clientBody) + + executor := NewCodexExecutor(&config.Config{}) + withoutCapsulePayload := `{"model":"gpt-6.1-sol","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"continue on GPT"}]}]}` + baseCountResp, errBaseCount := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(withoutCapsulePayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + }) + if errBaseCount != nil { + t.Fatalf("CountTokens(base) error = %v", errBaseCount) + } + expandedCountResp, errExpandedCount := executor.CountTokens(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + }) + if errExpandedCount != nil { + t.Fatalf("CountTokens(expanded) error = %v", errExpandedCount) + } + if gjson.GetBytes(expandedCountResp.Payload, "response.usage.input_tokens").Int() <= gjson.GetBytes(baseCountResp.Payload, "response.usage.input_tokens").Int() { + t.Fatalf("expected expanded capsule token count (%s) > base token count (%s)", expandedCountResp.Payload, baseCountResp.Payload) + } + }) + + t.Run("damaged CPA capsule returns 400 without forwarding to upstream", func(t *testing.T) { + upstreamHits := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamHits++ + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + wsExecutor := NewCodexWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "codex", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }, + } + + assertBadRequest := func(t *testing.T, name string, err error) { + t.Helper() + if err == nil { + t.Fatalf("%s with damaged capsule succeeded, want 400 error", name) + } + var se statusErr + if statusCode := err.(interface{ StatusCode() int }).StatusCode(); statusCode != http.StatusBadRequest { + t.Fatalf("%s status = %d (%v), want %d", name, statusCode, se, http.StatusBadRequest) + } + } + + corruptedPayloads := []string{ + `{"model":"gpt-6.1-sol","input":[{"id":"cmp_ag_compact_bad","type":"compaction","encrypted_content":"cpa-ag-compact-v1:corrupted"},{"type":"message","role":"user","content":"hi"}]}`, + `{"model":"gpt-6.1-sol","input":[{"id":"cmp_ag_compact_bad","type":"compaction","encrypted_content":"truncated"},{"type":"message","role":"user","content":"hi"}]}`, + } + + for _, badPayload := range corruptedPayloads { + req := cliproxyexecutor.Request{Model: "gpt-6.1-sol", Payload: []byte(badPayload)} + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: []byte(badPayload)} + + _, errExec := executor.Execute(context.Background(), auth, req, opts) + assertBadRequest(t, "Execute()", errExec) + + _, errStream := executor.ExecuteStream(context.Background(), auth, req, opts) + assertBadRequest(t, "ExecuteStream()", errStream) + + compactOpts := opts + compactOpts.Alt = "responses/compact" + _, errCompact := executor.Execute(context.Background(), auth, req, compactOpts) + assertBadRequest(t, "Execute(responses/compact)", errCompact) + + _, errTokens := executor.CountTokens(context.Background(), auth, req, opts) + assertBadRequest(t, "CountTokens()", errTokens) + + _, errWSExec := wsExecutor.Execute(context.Background(), auth, req, opts) + assertBadRequest(t, "wsExecutor.Execute()", errWSExec) + + _, errWSStream := wsExecutor.ExecuteStream(context.Background(), auth, req, opts) + assertBadRequest(t, "wsExecutor.ExecuteStream()", errWSStream) + } + if upstreamHits != 0 { + t.Fatalf("upstreamHits = %d, want 0 when CPA capsule is damaged", upstreamHits) + } + }) + + t.Run("websocket wire frames and duplex follow-up frames expand CPA capsule before PayloadFinalizer", func(t *testing.T) { + receivedFrames := make(chan []byte, 4) + serverDone := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(serverDone) + conn, errUpgrade := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { _ = conn.Close() }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + + // 1. Initial response.create + _, first, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read initial create: %v", errRead) + return + } + receivedFrames <- first + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"r1","output":[]}}`)) + + // 2. Duplex response.steer + _, second, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read steer: %v", errRead) + return + } + receivedFrames <- second + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.steer.accepted","steer":{"id":"s1","previous_response_id":"r1"}}`)) + _ = conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"r1","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`)) + + // 3. Duplex follow-up response.create (Wait for damaged steer to abort before reaching upstream) + _, third, errRead := conn.ReadMessage() + if errRead == nil { + receivedFrames <- third + } + })) + defer server.Close() + + cfg := &config.Config{ + Codex: config.CodexConfig{ResponseSteering: true}, + Payload: config.PayloadConfig{ + Override: []config.PayloadRule{ + { + Models: []config.PayloadModelRule{{Name: "gpt-6.1-sol", Protocol: "codex"}}, + Params: map[string]any{"input.0.content.0.text": "User rule override after CPA expansion"}, + }, + }, + }, + } + wsExec := NewCodexWebsocketsExecutor(cfg) + wsExec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ + ID: "ws-compact-auth", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "test-key", + "base_url": server.URL, + "websockets": "true", + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + inputCh := make(chan cliproxyexecutor.WebsocketInput, 4) + ctx = cliproxyexecutor.WithWebsocketInput(cliproxyexecutor.WithDownstreamWebsocket(ctx), inputCh) + + streamRes, errStream := wsExec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-6.1-sol", + Payload: []byte(mixedPayload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: []byte(mixedPayload), + Stream: true, + Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "duplex-compact-session"}, + }) + if errStream != nil { + t.Fatalf("wsExec.ExecuteStream() error = %v", errStream) + } + + steerPayload := `{"type":"response.steer","previous_response_id":"r1","input":[` + + `{"id":"cmp_ag_compact_2","type":"compaction","encrypted_content":"` + capsule + `"},` + + `{"id":"cmp_native_2","type":"compaction","encrypted_content":"gAAAAABnative-openai-compaction"},` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"steer on GPT"}]}` + + `]}` + badSteerPayload := `{"type":"response.steer","previous_response_id":"r1","input":[` + + `{"id":"cmp_ag_compact_bad","type":"compaction","encrypted_content":"cpa-ag-compact-v1:corrupted"}` + + `]}` + + sawStreamErr := false + for chunk := range streamRes.Chunks { + if chunk.Err != nil { + sawStreamErr = true + break + } + evType := gjson.GetBytes(chunk.Payload, "type").String() + switch evType { + case "response.created": + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(steerPayload)} + case "response.completed": + inputCh <- cliproxyexecutor.WebsocketInput{Payload: []byte(badSteerPayload)} + } + } + if !sawStreamErr { + t.Fatal("expected stream error when duplex response.steer carries damaged CPA capsule") + } + cancel() + <-serverDone + close(receivedFrames) + + var frames [][]byte + for f := range receivedFrames { + frames = append(frames, f) + } + if len(frames) != 2 { + t.Fatalf("upstream received %d websocket frames, want 2 (initial create + valid steer; damaged steer must be rejected locally)", len(frames)) + } + for idx, frame := range frames { + if bytes.Contains(frame, []byte("cpa-ag-compact-v1:")) || bytes.Contains(frame, []byte("cmp_ag_compact_")) { + t.Fatalf("frame[%d] leaked CPA compaction capsule: %s", idx, frame) + } + input := gjson.GetBytes(frame, "input").Array() + if len(input) != 3 { + t.Fatalf("frame[%d] input len = %d, want 3; frame=%s", idx, len(input), frame) + } + if input[0].Get("type").String() != "message" || input[0].Get("role").String() != "developer" { + t.Fatalf("frame[%d] input[0] = %s, want developer message", idx, input[0].Raw) + } + // Verify PayloadFinalizer ran after capsule expansion and overrode input.0.content.0.text. + if gotText := input[0].Get("content.0.text").String(); gotText != "User rule override after CPA expansion" { + t.Fatalf("frame[%d] input[0].content.0.text = %q, want %q (PayloadFinalizer must run after capsule expansion)", idx, gotText, "User rule override after CPA expansion") + } + if input[1].Get("type").String() != "compaction" || input[1].Get("encrypted_content").String() != "gAAAAABnative-openai-compaction" { + t.Fatalf("frame[%d] native compaction item not preserved at input[1]: %s", idx, input[1].Raw) + } + } + }) +} diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go index cbaf2c9c718..2c3fbb36e90 100644 --- a/internal/runtime/executor/codex_executor_execute.go +++ b/internal/runtime/executor/codex_executor_execute.go @@ -20,6 +20,9 @@ import ( func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { ctx = helps.EnsureSessionContext(ctx, opts, req.Payload) + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } if opts.Alt == "responses/compact" { return e.executeCompact(ctx, auth, req, opts) } @@ -205,14 +208,14 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 { return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage} } - if detail, ok := helps.ParseCodexUsage(eventData); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(completedData, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } - publishCodexImageToolUsage(ctx, reporter, body, eventData) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, completedData, outputItemsByIndex, outputItemsFallback) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } - resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} + resp = cliproxyexecutor.Response{Payload: out, Incomplete: eventType == "response.incomplete", Headers: httpResp.Header.Clone()} return resp, nil } if errRead != nil { @@ -228,6 +231,9 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re } func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } baseModel := thinking.ParseSuffix(req.Model).ModelName apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) diff --git a/internal/runtime/executor/codex_executor_request.go b/internal/runtime/executor/codex_executor_request.go index b03948d7ec2..3f9bb2d9c18 100644 --- a/internal/runtime/executor/codex_executor_request.go +++ b/internal/runtime/executor/codex_executor_request.go @@ -726,7 +726,11 @@ func normalizeCodexParallelToolCallsForTools(body []byte) []byte { } func publishCodexImageToolUsage(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte) { - detail, ok := helps.ParseCodexImageToolUsage(completedData) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, completedData, nil, nil) +} + +func publishCodexImageToolUsageWithOutputItems(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte, itemsByIndex map[int64][]byte, fallback [][]byte) { + detail, ok := helps.ParseCodexImageToolUsageWithOutputItems(body, completedData, itemsByIndex, fallback) if !ok { return } diff --git a/internal/runtime/executor/codex_executor_sse_events_test.go b/internal/runtime/executor/codex_executor_sse_events_test.go index 558bfdb56a0..28309d07cdd 100644 --- a/internal/runtime/executor/codex_executor_sse_events_test.go +++ b/internal/runtime/executor/codex_executor_sse_events_test.go @@ -12,6 +12,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" @@ -454,3 +455,152 @@ func TestCodexExecuteStreamEventExceedsByteBudgetWhileAssemblingReleasesStream(t t.Fatal("expected in-stream delivery of overload error after byte budget exhausted") } } + +func TestCodexVerifyInferenceTranslatesPayloadAndRequiresCompletedEvent_Issue5783(t *testing.T) { + var capturedBody []byte + var sseResponse string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var buf strings.Builder + _, _ = fmt.Fscan(r.Body, &buf) + _ = r.Body.Close() + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(sseResponse)) + })) + defer server.Close() + + rt := roundTripFunc(func(req *http.Request) (*http.Response, error) { + if req.Body != nil { + var b strings.Builder + buf := make([]byte, 4096) + for { + n, errRead := req.Body.Read(buf) + if n > 0 { + b.Write(buf[:n]) + } + if errRead != nil { + break + } + } + capturedBody = []byte(b.String()) + } + return codexAllocationTransport{stream: sseResponse}.RoundTrip(req) + }) + + cfg := &config.Config{} + executor := NewCodexExecutor(cfg) + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + + ctx := cliproxyauth.WithSkipPersist(context.WithValue(t.Context(), "cliproxy.roundtripper", rt)) + auth := &cliproxyauth.Auth{ + ID: "codex-real-verify.json", + FileName: "codex-real-verify.json", + Provider: "codex", + Status: cliproxyauth.StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "verify@example.com", + "access_token": "at-verify", + "refresh_token": "rt-verify", + "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.4"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + + // 1. Output text containing "response.completed" followed by response.incomplete must fail verification. + sseResponse = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\",\"model\":\"gpt-5.4\"}}\n\n" + + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"response.completed [DONE]\"}\n\n" + + "event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"r1\",\"model\":\"gpt-5.4\",\"status\":\"incomplete\"}}\n\n" + if _, errVerifyIncomp := manager.VerifyInference(ctx, auth.ID, "gpt-5.4"); errVerifyIncomp == nil { + t.Fatal("VerifyInference should fail when stream ends with response.incomplete despite output text mentioning response.completed") + } + afterIncomp, _ := manager.GetByID(auth.ID) + if afterIncomp.InferenceVerified { + t.Fatal("InferenceVerified must remain false after incomplete stream") + } + + // 1b. Chat Completions (FormatOpenAI) stream translated from upstream response.incomplete (finish_reason="length" + [DONE]) must NOT mark InferenceVerified=true. + chatStreamRes, errChatStream := manager.ExecuteStream( + ctx, + []string{"codex"}, + cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`), + }, + cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + }, + ) + if errChatStream != nil { + t.Fatalf("ExecuteStream (Chat Completions incomplete) error: %v", errChatStream) + } + for range chatStreamRes.Chunks { + } + afterChatIncomp, _ := manager.GetByID(auth.ID) + if afterChatIncomp.InferenceVerified { + t.Fatal("Chat Completions stream translated from response.incomplete must not mark InferenceVerified=true") + } + + // 1c. With StreamBootstrapBuffering enabled, an immediate response.incomplete without incomplete_details.reason + // (which Chat Completions translates with finish_reason="stop") must STILL be marked Incomplete and NOT set InferenceVerified=true. + cfg.Codex.StreamBootstrapBuffering = true + sseResponse = "event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"r_buf\",\"model\":\"gpt-5.4\",\"status\":\"incomplete\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"partial\"}]}]}}\n\n" + bufferedChatRes, errBufChat := manager.ExecuteStream( + ctx, + []string{"codex"}, + cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`), + }, + cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + }, + ) + if errBufChat != nil { + t.Fatalf("ExecuteStream (buffered Chat Completions incomplete) error: %v", errBufChat) + } + for range bufferedChatRes.Chunks { + } + afterBufChatIncomp, _ := manager.GetByID(auth.ID) + if afterBufChatIncomp.InferenceVerified { + t.Fatal("bootstrap-buffered response.incomplete without incomplete_details.reason must not mark InferenceVerified=true") + } + cfg.Codex.StreamBootstrapBuffering = false + + // 2. Real completed stream succeeds and translated upstream body strips max_output_tokens 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" + verifiedAuth, errVerifyOK := manager.VerifyInference(ctx, auth.ID, "gpt-5.4") + if errVerifyOK != nil { + t.Fatalf("VerifyInference error: %v", errVerifyOK) + } + if !verifiedAuth.InferenceVerified { + t.Fatal("expected InferenceVerified=true after completed stream") + } + var upstreamJSON map[string]any + if errUnmarshal := json.Unmarshal(capturedBody, &upstreamJSON); errUnmarshal != nil { + t.Fatalf("unmarshal upstream body %s: %v", string(capturedBody), errUnmarshal) + } + if upstreamJSON["stream"] != true || upstreamJSON["store"] != false { + t.Fatalf("expected translated Codex body with stream=true and store=false, got: %s", string(capturedBody)) + } + if _, hasMaxOutput := upstreamJSON["max_output_tokens"]; hasMaxOutput { + t.Fatalf("translated Codex body must strip unsupported max_output_tokens, got: %s", string(capturedBody)) + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index 5b7084525f0..d65c94c875c 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -25,6 +25,9 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} } + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } if isCodexOpenAIImageRequest(opts) { return e.executeOpenAIImageStream(ctx, auth, req, opts) } @@ -204,6 +207,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au var initialChunks [][]byte streamStarted := false immediateTerminal := false + immediateIncomplete := false // bootstrapTerminalErr holds a non-overload terminal failure seen while buffering. It is // delivered as an in-stream chunk after the buffered handshake so downstream behaviour stays // identical to the unbuffered path instead of silently turning into a credential failover. @@ -287,13 +291,14 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au collectCodexOutputItemDone(data, outputItemsByIndex, &outputItemsFallback) case "response.completed", "response.incomplete", "response.done": terminalSuccess = true + immediateIncomplete = eventType == "response.incomplete" data = normalizeCodexWebsocketCompletion(data) - if detail, ok := helps.ParseCodexUsage(data); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(data, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } - publishCodexImageToolUsage(ctx, reporter, body, data) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, data, outputItemsByIndex, outputItemsFallback) if !preserveNativeOutput { data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) } @@ -379,7 +384,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } chanCapacity := len(bufferedChunks) + len(initialChunks) - if bootstrapTerminalErr != nil { + if bootstrapTerminalErr != nil || (immediateIncomplete && len(initialChunks) == 0) { chanCapacity++ } out := make(chan cliproxyexecutor.StreamChunk, chanCapacity) @@ -387,7 +392,10 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au out <- cliproxyexecutor.StreamChunk{Payload: chunk} } for _, chunk := range initialChunks { - out <- cliproxyexecutor.StreamChunk{Payload: chunk} + out <- cliproxyexecutor.StreamChunk{Payload: chunk, Incomplete: immediateIncomplete} + } + if immediateIncomplete && len(initialChunks) == 0 { + out <- cliproxyexecutor.StreamChunk{Incomplete: true} } if bootstrapTerminalErr != nil { // Buffered handshake payloads are flushed first so the conductor observes a committed @@ -424,6 +432,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au line := scanner.Bytes() var translatedLine []byte terminalSuccess := false + isIncompleteEvent := false if transformed, ok := grokbuild.TransformKeepaliveSSELine(line, isGrokClient); ok { translatedLine = transformed @@ -469,13 +478,14 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au collectCodexOutputItemDone(data, outputItemsByIndex, &outputItemsFallback) case "response.completed", "response.incomplete", "response.done": terminalSuccess = true + isIncompleteEvent = eventType == "response.incomplete" data = normalizeCodexWebsocketCompletion(data) - if detail, ok := helps.ParseCodexUsage(data); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(data, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } - publishCodexImageToolUsage(ctx, reporter, body, data) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, data, outputItemsByIndex, outputItemsFallback) if !preserveNativeOutput { data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) } @@ -491,7 +501,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) for i := range chunks { select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i], Incomplete: isIncompleteEvent}: if len(chunks[i]) > 0 { emittedCount++ } @@ -500,6 +510,12 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } } if terminalSuccess { + if isIncompleteEvent && len(chunks) == 0 { + select { + case out <- cliproxyexecutor.StreamChunk{Incomplete: true}: + case <-ctx.Done(): + } + } return } } diff --git a/internal/runtime/executor/codex_executor_tokens.go b/internal/runtime/executor/codex_executor_tokens.go index 5d8f70be68b..4719e1a9da0 100644 --- a/internal/runtime/executor/codex_executor_tokens.go +++ b/internal/runtime/executor/codex_executor_tokens.go @@ -3,6 +3,7 @@ package executor import ( "context" "fmt" + "net/http" "strings" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" @@ -16,6 +17,9 @@ import ( ) func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } baseModel := thinking.ParseSuffix(req.Model).ModelName from := opts.SourceFormat diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 7a689a102fa..abea9f74b7e 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -168,10 +168,10 @@ func (e *CodexExecutor) executeOpenAIImage(ctx context.Context, auth *cliproxyau case "response.output_item.done": collectCodexOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) case "response.completed": - if detail, ok := helps.ParseCodexUsage(eventData); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } - publishCodexImageToolUsage(ctx, reporter, body, eventData) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, eventData, outputItemsByIndex, outputItemsFallback) results, createdAt, usageRaw, firstMeta, errExtract := codexExtractImageResults(eventData, outputItemsByIndex, outputItemsFallback) if errExtract != nil { return resp, errExtract @@ -307,10 +307,10 @@ func (e *CodexExecutor) executeOpenAIImageStream(ctx context.Context, auth *clip return } case "response.completed": - if detail, ok := helps.ParseCodexUsage(eventData); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } - publishCodexImageToolUsage(ctx, reporter, body, eventData) + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, eventData, outputItemsByIndex, outputItemsFallback) results, _, usageRaw, _, errExtract := codexExtractImageResults(eventData, outputItemsByIndex, outputItemsFallback) if errExtract != nil { sendError(errExtract) @@ -405,7 +405,7 @@ func (e *CodexExecutor) executeDirectOpenAIImage(ctx context.Context, auth *clip return resp, err } - reporter.Publish(ctx, helps.ParseOpenAIUsage(data)) + reporter.Publish(ctx, helps.ParseOpenAIImageUsage(body, data, contentType)) reporter.EnsurePublished(ctx) return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil } @@ -478,10 +478,13 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth go func() { defer close(out) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(body, contentType) + usageObserver := helps.NewOpenAIStreamUsageObserver(&streamUsage, streamScannerBuffer) defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("codex executor: close response body error: %v", errClose) } + usageObserver.Finish() streamUsage.Publish(ctx, reporter) reporter.EnsurePublished(ctx) }() @@ -492,9 +495,7 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth if n > 0 { chunk := bytes.Clone(buffer[:n]) helps.AppendAPIResponseChunk(ctx, e.cfg, chunk) - for _, line := range bytes.Split(chunk, []byte("\n")) { - streamUsage.ObserveOpenAIStream(bytes.TrimSpace(line)) - } + usageObserver.Feed(chunk) select { case out <- cliproxyexecutor.StreamChunk{Payload: chunk}: case <-ctx.Done(): @@ -504,7 +505,10 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth if errRead != nil { if errRead != io.EOF { helps.RecordAPIResponseError(ctx, e.cfg, errRead) - reporter.PublishFailure(ctx, errRead) + usageObserver.Finish() + if !streamUsage.PublishFailure(ctx, reporter, errRead) { + reporter.PublishFailure(ctx, errRead) + } select { case out <- cliproxyexecutor.StreamChunk{Err: errRead}: case <-ctx.Done(): diff --git a/internal/runtime/executor/codex_openai_images_test.go b/internal/runtime/executor/codex_openai_images_test.go index e138e5d343e..ae50a90caee 100644 --- a/internal/runtime/executor/codex_openai_images_test.go +++ b/internal/runtime/executor/codex_openai_images_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" 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" "github.com/tidwall/gjson" ) @@ -605,3 +606,47 @@ func TestCodexExecutorDirectOpenAIImage25Models(t *testing.T) { }) } } + +func TestCodexDirectOpenAIImageStreamAccountsLargeCompletedEvents(t *testing.T) { + for _, size := range []int{16, 100000} { + t.Run(fmt.Sprint(size), func(t *testing.T) { + event := `data: {"type":"image_generation.completed","model":"gpt-image-2","b64_json":"` + strings.Repeat("A", size) + `","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}` + "\n\n" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + if _, err := fmt.Fprint(w, event); err != nil { + t.Errorf("write upstream event: %v", err) + } + })) + defer server.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + res, err := NewCodexExecutor(&config.Config{}).ExecuteStream(ctx, newCodexOpenAIImageTestAuth(server.URL), cliproxyexecutor.Request{ + Model: "gpt-image-2", + Payload: []byte(`{"model":"gpt-image-2","prompt":"test"}`), + }, codexOpenAIImageTestOptions(codexImagesGenerationsPath, true)) + if err != nil { + t.Fatal(err) + } + var got bytes.Buffer + for chunk := range res.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + got.Write(chunk.Payload) + } + if got.String() != event { + t.Fatal("forwarded image SSE bytes changed") + } + record := capture.await(t) + if record.Detail.TotalTokens != 30 || record.Detail.ImageCount == nil || *record.Detail.ImageCount != 1 { + t.Fatalf("image completion usage missing: %+v", record.Detail) + } + }) + } +} diff --git a/internal/runtime/executor/codex_websockets_active_response_test.go b/internal/runtime/executor/codex_websockets_active_response_test.go new file mode 100644 index 00000000000..53fe4ebc53a --- /dev/null +++ b/internal/runtime/executor/codex_websockets_active_response_test.go @@ -0,0 +1,150 @@ +package executor + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// On a duplex socket a rejected queued create can fail while another response +// is still running. That failure must not end the live response's interrupt +// window, the running response's own terminal event must, and an unknown id +// must be rejected even while a different response is active. +func TestInterruptExecutionSessionTracksActiveResponseIdentity(t *testing.T) { + received := make(chan string, 8) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := (&websocket.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 upstream: %v", errClose) + } + }() + // Hold the completion until the active-r1 interrupt has arrived, so + // the read loop cannot process it before the checks below run. + for _, payload := range []string{ + `{"type":"response.created","response":{"id":"r1"}}`, + // A queued create rejected before its own response.created. + `{"type":"response.failed","response":{"id":"queued-1"}}`, + } { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Errorf("write: %v", errWrite) + return + } + } + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + received <- string(payload) + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"r1","output":[]}}`)); errWrite != nil { + t.Errorf("write completion: %v", errWrite) + return + } + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + received <- string(payload) + if string(payload) == "barrier" { + return + } + } + })) + defer upstream.Close() + + client, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(upstream.URL, "http"), nil) + if errDial != nil { + t.Fatal(errDial) + } + exec := NewCodexWebsocketsExecutor(nil) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + const sessionID = "active-response-identity" + defer exec.CloseExecutionSession(sessionID) + sess := exec.getOrCreateSession(sessionID) + sess.conn = client + sess.connCloser = newWebsocketConnectionCloser(client) + sess.configureConn(client) + readCh := sess.activate(client) + go exec.readUpstreamLoop(sess, client) + read := func() codexWebsocketRead { + t.Helper() + select { + case event := <-readCh: + if event.err != nil { + t.Fatal(event.err) + } + return event + case <-time.After(3 * time.Second): + t.Fatal("upstream event timed out") + return codexWebsocketRead{} + } + } + interrupt := func(payload string) error { + t.Helper() + return exec.InterruptExecutionSession(context.Background(), sessionID, []byte(payload)) + } + + read() // response.created r1 + read() // response.failed queued-1 + + // The queued rejection must not close r1's interrupt window. + if errInterrupt := interrupt(`{"type":"response.interrupt","response_id":"r1"}`); errInterrupt != nil { + t.Fatalf("interrupt active r1 after queued failure = %v, want nil (forwarded)", errInterrupt) + } + // An unknown id is not the running response: reject outright, without + // inviting a local HTTP fallback against the live websocket turn. + if errUnknown := interrupt(`{"type":"response.interrupt","response_id":"unknown"}`); !errors.Is(errUnknown, cliproxyexecutor.ErrUpstreamWebsocketInterruptRejected) { + t.Fatalf("unknown interrupt during r1 = %v, want ErrUpstreamWebsocketInterruptRejected", errUnknown) + } + + read() // response.completed r1 + + // The terminal id itself stays idempotent, but an id that never ran must + // be rejected while teardown of the active channel is still pending. + if errTerminal := interrupt(`{"type":"response.interrupt","response_id":"r1"}`); errTerminal != nil { + t.Fatalf("interrupt completed r1 = %v, want nil (already terminal)", errTerminal) + } + if errUnknown := interrupt(`{"type":"response.interrupt","response_id":"unknown-2"}`); !errors.Is(errUnknown, cliproxyexecutor.ErrUpstreamWebsocketInterruptRejected) { + t.Fatalf("unknown interrupt after completion = %v, want ErrUpstreamWebsocketInterruptRejected", errUnknown) + } + + // Only the active-r1 interrupt may have reached the socket. + if errWrite := client.WriteMessage(websocket.TextMessage, []byte("barrier")); errWrite != nil { + t.Fatal(errWrite) + } + var got []string +loop: + for { + select { + case payload := <-received: + got = append(got, payload) + if payload == "barrier" { + break loop + } + case <-time.After(3 * time.Second): + t.Fatalf("upstream barrier timed out, received %v", got) + } + } + want := []string{`{"type":"response.interrupt","response_id":"r1"}`, "barrier"} + if len(got) != len(want) { + t.Fatalf("upstream received %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("upstream received %v, want %v", got, want) + } + } + sess.clearActive(client, readCh) +} diff --git a/internal/runtime/executor/codex_websockets_duplex.go b/internal/runtime/executor/codex_websockets_duplex.go index 41daa3e03ee..6c70a8bc8ee 100644 --- a/internal/runtime/executor/codex_websockets_duplex.go +++ b/internal/runtime/executor/codex_websockets_duplex.go @@ -16,6 +16,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" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -182,6 +183,117 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } return true, true } + finalBeforeHeaders := func(current, intercepted http.Header, clear []string) http.Header { + if intercepted != nil { + if len(intercepted) == 0 { + return nil + } + return intercepted.Clone() + } + if len(current) == 0 { + return nil + } + out := current.Clone() + for _, key := range clear { + out.Del(key) + } + if len(out) == 0 { + return nil + } + return out + } + mergeAfterHeaders := func(current, updates http.Header, clear []string) http.Header { + if updates == nil && len(clear) == 0 { + if len(current) == 0 { + return nil + } + return current.Clone() + } + out := http.Header{} + if len(current) > 0 { + out = current.Clone() + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + if len(out) == 0 { + return nil + } + return out + } + rejectTerminated := func(resp cliproxyexecutor.RequestAfterAuthInterceptResponse) bool { + status := resp.StatusCode + if status < 400 || status > 599 { + status = http.StatusForbidden + } + msg := "request terminated by plugin" + if len(resp.ResponseBody) > 0 { + if parsed := strings.TrimSpace(gjson.GetBytes(resp.ResponseBody, "error.message").String()); parsed != "" { + msg = parsed + } else if trimmed := strings.TrimSpace(string(resp.ResponseBody)); trimmed != "" { + msg = trimmed + } + } + payload, _ := json.Marshal(map[string]any{ + "type": "error", "status": status, + "error": map[string]string{"type": "invalid_request_error", "message": msg}, + }) + select { + case out <- cliproxyexecutor.StreamChunk{Payload: payload}: + consumedPayload = nil + return true + case <-streamCtx.Done(): + return false + } + } + interceptTurnPayload := func(payload []byte, baseHeaders http.Header) ([]byte, http.Header, bool, bool) { + if opts.RequestBeforeAuthInterceptor == nil && opts.RequestAfterAuthInterceptor == nil { + return payload, baseHeaders, true, true + } + curHeaders := baseHeaders + if curHeaders != nil { + curHeaders = curHeaders.Clone() + } + interceptReq := cliproxyexecutor.RequestAfterAuthInterceptRequest{ + SourceFormat: opts.SourceFormat, + Model: req.Model, + RequestedModel: boundModel, + Stream: true, + Headers: curHeaders, + Body: payload, + Metadata: opts.Metadata, + } + if opts.RequestBeforeAuthInterceptor != nil { + resp := opts.RequestBeforeAuthInterceptor(streamCtx, interceptReq) + curHeaders = finalBeforeHeaders(curHeaders, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + payload = resp.Body + interceptReq.Body = payload + } + interceptReq.Headers = curHeaders + if resp.Terminate { + return payload, curHeaders, false, rejectTerminated(resp) + } + } + if opts.RequestAfterAuthInterceptor != nil { + interceptReq.ToFormat = sdktranslator.FromString("codex") + resp := opts.RequestAfterAuthInterceptor(streamCtx, interceptReq) + curHeaders = mergeAfterHeaders(curHeaders, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + payload = resp.Body + } + if resp.Terminate { + return payload, curHeaders, false, rejectTerminated(resp) + } + } + return payload, curHeaders, true, true + } processCreatePayload := func(payload []byte) bool { consumedPayload = payload if allowed, keepOpen := authorizePayload(payload, false); !allowed { @@ -212,13 +324,19 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( fail(cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()) return false } + var turnHeaders http.Header + var allowed, keepOpen bool + payload, turnHeaders, allowed, keepOpen = interceptTurnPayload(payload, opts.Headers) + if !allowed { + return keepOpen + } if model == "" { model = req.Model if model == "" { model = boundModel } - payload, _ = sjson.SetBytes(payload, "model", model) } + payload, _ = sjson.SetBytes(payload, "model", model) if isAppend && !gjson.GetBytes(payload, "instructions").Exists() { metadataMu.Lock() parentSettings := responseSettings[prevID] @@ -247,6 +365,7 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( nextReq, nextOpts := req, opts nextReq.Payload = payload nextOpts.OriginalRequest = payload + nextOpts.Headers = turnHeaders prepared, errPrepare := e.prepareCodexWebsocketStream(streamCtx, auth, nextReq, nextOpts) if errPrepare != nil { fail(errPrepare) @@ -342,13 +461,30 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( } continue } + var steerHeaders http.Header + var allowed, keepOpen bool + payload, steerHeaders, allowed, keepOpen = interceptTurnPayload(payload, opts.Headers) + if !allowed { + if !keepOpen { + return + } + continue + } // Steering carries business input but must not inherit response.create defaults. steerReq := req steerReq.Payload = payload - if codexUseSIWCForRequest(ctx, auth, opts.Headers) { + steerOpts := opts + steerOpts.OriginalRequest = payload + steerOpts.Headers = steerHeaders + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&steerReq, &steerOpts); errExpand != nil { + fail(statusErr{code: http.StatusBadRequest, msg: errExpand.Error()}) + return + } + payload = steerReq.Payload + if codexUseSIWCForRequest(ctx, auth, steerHeaders) { payload = thinking.StripConfigurationUpdates(payload) } - payload = helps.NewPayloadFinalizer(e.cfg, "codex-websockets", thinking.ParseSuffix(req.Model).ModelName, "codex", "", payload, steerReq, opts)(payload) + payload = helps.NewPayloadFinalizer(e.cfg, "codex-websockets", thinking.ParseSuffix(req.Model).ModelName, "codex", "", payload, steerReq, steerOpts)(payload) payload, _ = sjson.SetBytes(payload, "type", "response.steer") parent := gjson.GetBytes(payload, "previous_response_id").String() metadataMu.Lock() @@ -739,13 +875,14 @@ func (e *CodexWebsocketsExecutor) streamCodexDuplex( if eventType != "response.incomplete" { cacheCodexReasoningReplayFromCompleted(current.replayScope, payload) } - if detail, ok := helps.ParseCodexUsage(payload); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(payload, outputItems, outputFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, current.clientBody, payload, outputItems, outputFallback) } - if !send(cliproxyexecutor.StreamChunk{Payload: helps.EnsureResponsesUsageDetails(payload)}) { + if !send(cliproxyexecutor.StreamChunk{Payload: helps.EnsureResponsesUsageDetails(payload), Incomplete: eventType == "response.incomplete"}) { return } if establishing { diff --git a/internal/runtime/executor/codex_websockets_error_terminal_test.go b/internal/runtime/executor/codex_websockets_error_terminal_test.go new file mode 100644 index 00000000000..1dc2c00a7e3 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_error_terminal_test.go @@ -0,0 +1,181 @@ +package executor + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// Failure events terminate responses too, in the shapes duplex attribution +// accepts (top-level response_id, or id-less while a response runs), for +// both error and response.failed. The session must close the matching +// interrupt window for them, must keep the running response's window open +// for other responses' failures, and must let the next response.created +// reopen it. +func TestInterruptExecutionSessionErrorTerminalShapes(t *testing.T) { + received := make(chan string, 8) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := (&websocket.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 upstream: %v", errClose) + } + }() + write := func(payload string) bool { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Errorf("write: %v", errWrite) + return false + } + return true + } + readFrame := func() bool { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return false + } + received <- string(payload) + return true + } + // r1 runs; an error for another response must not end it. Reading + // the forwarded r1 interrupt barriers before r1's own error. + if !write(`{"type":"response.created","response":{"id":"r1"}}`) { + return + } + if !write(`{"type":"error","response_id":"queued-1","error":{"message":"rejected"}}`) { + return + } + if !readFrame() { + return + } + // r1 ends with a top-level response_id error. The rejected checks + // write nothing upstream, so after r2 is created the next frame + // read here must be the r2 interrupt; reading it before r2's + // id-less failure is sent keeps the checks deterministic. + if !write(`{"type":"error","response_id":"r1","error":{"message":"boom"}}`) { + return + } + if !write(`{"type":"response.created","response":{"id":"r2"}}`) { + return + } + if !readFrame() { + return + } + if !write(`{"type":"response.failed","response":{"error":{"message":"boom"}}}`) { + return + } + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + received <- string(payload) + if string(payload) == "barrier" { + return + } + } + })) + defer upstream.Close() + + client, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(upstream.URL, "http"), nil) + if errDial != nil { + t.Fatal(errDial) + } + exec := NewCodexWebsocketsExecutor(nil) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + const sessionID = "error-terminal-shapes" + defer exec.CloseExecutionSession(sessionID) + sess := exec.getOrCreateSession(sessionID) + sess.conn = client + sess.connCloser = newWebsocketConnectionCloser(client) + sess.configureConn(client) + readCh := sess.activate(client) + go exec.readUpstreamLoop(sess, client) + read := func() codexWebsocketRead { + t.Helper() + select { + case event := <-readCh: + if event.err != nil { + t.Fatal(event.err) + } + return event + case <-time.After(3 * time.Second): + t.Fatal("upstream event timed out") + return codexWebsocketRead{} + } + } + interrupt := func(payload string) error { + t.Helper() + return exec.InterruptExecutionSession(context.Background(), sessionID, []byte(payload)) + } + wantRejected := func(payload, what string) { + t.Helper() + if errInterrupt := interrupt(payload); !errors.Is(errInterrupt, cliproxyexecutor.ErrUpstreamWebsocketInterruptRejected) { + t.Fatalf("%s = %v, want ErrUpstreamWebsocketInterruptRejected", what, errInterrupt) + } + } + + read() // response.created r1 + read() // error queued-1 + if errInterrupt := interrupt(`{"type":"response.interrupt","response_id":"r1"}`); errInterrupt != nil { + t.Fatalf("interrupt r1 after queued error = %v, want nil (forwarded)", errInterrupt) + } + read() // error r1 + wantRejected(`{"type":"response.interrupt","response_id":"unknown"}`, "unknown interrupt after r1 error") + wantRejected(`{"type":"response.interrupt","response_id":"r1"}`, "interrupt r1 after its error") + + read() // response.created r2 + if errInterrupt := interrupt(`{"type":"response.interrupt","response_id":"r2"}`); errInterrupt != nil { + t.Fatalf("interrupt r2 = %v, want nil (forwarded)", errInterrupt) + } + read() // id-less response.failed + wantRejected(`{"type":"response.interrupt","response_id":"unknown-2"}`, "unknown interrupt after id-less failure") + // r2 itself is already terminal: idempotent via the recorded + // response.failed terminal id, and never forwarded (proven by the + // upstream barrier below). Returning nil here is what distinguishes + // "window closed" from "still active" — the latter would forward. + if errInterrupt := interrupt(`{"type":"response.interrupt","response_id":"r2"}`); errInterrupt != nil { + t.Fatalf("interrupt r2 after id-less failure = %v, want nil (already terminal)", errInterrupt) + } + + if errWrite := client.WriteMessage(websocket.TextMessage, []byte("barrier")); errWrite != nil { + t.Fatal(errWrite) + } + var got []string +loop: + for { + select { + case payload := <-received: + got = append(got, payload) + if payload == "barrier" { + break loop + } + case <-time.After(3 * time.Second): + t.Fatalf("upstream barrier timed out, received %v", got) + } + } + want := []string{ + `{"type":"response.interrupt","response_id":"r1"}`, + `{"type":"response.interrupt","response_id":"r2"}`, + "barrier", + } + if len(got) != len(want) { + t.Fatalf("upstream received %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("upstream received %v, want %v", got, want) + } + } + sess.clearActive(client, readCh) +} diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index 1637c5d1a87..a8087033835 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -22,6 +22,9 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if ctx == nil { ctx = context.Background() } + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return resp, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } if opts.Alt == "responses/compact" { return e.CodexExecutor.executeCompact(ctx, auth, req, opts) } @@ -370,17 +373,18 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if eventType != "response.incomplete" { cacheCodexReasoningReplayFromCompleted(replayScope, payload) } - if detail, ok := helps.ParseCodexUsage(payload); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(payload, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, body, payload, outputItemsByIndex, outputItemsFallback) var param any out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, payload, ¶m) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } - resp = cliproxyexecutor.Response{Payload: out} + resp = cliproxyexecutor.Response{Payload: out, Incomplete: eventType == "response.incomplete"} return resp, nil } } diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 65ec499c690..b7a2c832b4b 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -227,9 +227,17 @@ func (e *CodexWebsocketsExecutor) InterruptExecutionSession(ctx context.Context, return cliproxyexecutor.ErrNoActiveUpstreamWebsocket } sess.activeMu.Lock() - _, alreadyTerminal := sess.terminalResponseIDs[gjson.GetBytes(payload, "response_id").String()] + targetID := gjson.GetBytes(payload, "response_id").String() + _, alreadyTerminal := sess.terminalResponseIDs[targetID] alreadyTerminal = alreadyTerminal && sess.responseStateConn == conn - active := sess.activeConn == conn && sess.activeCh != nil + // activeResponseID/Terminal track the running response at read-loop + // granularity: only its own id forwards while it runs, and its terminal + // event closes the window before teardown clears the channel. Before + // response.created the id is unknown, so the live channel alone + // authorizes forwarding and upstream validates the frame. + liveChannel := sess.activeConn == conn && sess.activeCh != nil + active := liveChannel && !sess.activeResponseTerminal && + (sess.activeResponseID == "" || sess.activeResponseID == targetID) var disconnectErr error if sess.terminalConn == conn { disconnectErr = sess.terminalErr @@ -247,9 +255,15 @@ func (e *CodexWebsocketsExecutor) InterruptExecutionSession(ctx context.Context, } // A retained socket from an earlier turn is not the current upstream. // HTTP turns must fall through to local cancellation instead. - if !active { + if !liveChannel { return cliproxyexecutor.ErrNoActiveUpstreamWebsocket } + // The socket owns this turn, but the frame does not target its running + // response. Reject outright: a local HTTP fallback would fabricate an + // incomplete event and cancel the healthy websocket turn. + if !active { + return cliproxyexecutor.ErrUpstreamWebsocketInterruptRejected + } if !authEnabled { return fmt.Errorf("websocket credential is no longer enabled") } diff --git a/internal/runtime/executor/codex_websockets_late_interrupt_unknown_test.go b/internal/runtime/executor/codex_websockets_late_interrupt_unknown_test.go new file mode 100644 index 00000000000..09611bc836c --- /dev/null +++ b/internal/runtime/executor/codex_websockets_late_interrupt_unknown_test.go @@ -0,0 +1,110 @@ +package executor + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +// An unknown-id interrupt arriving after the upstream terminal event was +// delivered, but before the stream goroutine cleared the active channel, +// must be rejected outright instead of being written to the still-active +// socket or falling back to a local cancellation. The test does not call +// clearActive on purpose: it reproduces the window between completion +// delivery and turn teardown. +func TestInterruptExecutionSessionUnknownAfterTerminalDelivery(t *testing.T) { + interruptCount := make(chan int, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := (&websocket.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 upstream: %v", errClose) + } + }() + for _, payload := range []string{ + `{"type":"response.created","response":{"id":"r1"}}`, + `{"type":"response.completed","response":{"id":"r1","output":[]}}`, + } { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Errorf("write: %v", errWrite) + return + } + } + count := 0 + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + if string(payload) == "barrier" { + interruptCount <- count + return + } + count++ + } + })) + defer upstream.Close() + + client, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(upstream.URL, "http"), nil) + if errDial != nil { + t.Fatal(errDial) + } + exec := NewCodexWebsocketsExecutor(nil) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + const sessionID = "unknown-after-terminal" + defer exec.CloseExecutionSession(sessionID) + sess := exec.getOrCreateSession(sessionID) + sess.conn = client + sess.connCloser = newWebsocketConnectionCloser(client) + sess.configureConn(client) + readCh := sess.activate(client) + go exec.readUpstreamLoop(sess, client) + read := func() codexWebsocketRead { + t.Helper() + select { + case event := <-readCh: + if event.err != nil { + t.Fatal(event.err) + } + return event + case <-time.After(3 * time.Second): + t.Fatal("upstream event timed out") + return codexWebsocketRead{} + } + } + read() + read() + + // Deliberately no clearActive here: the stream goroutine may not have + // run its deferred teardown yet, but the client has already seen + // response.completed and now sends an unknown interrupt. + errInterrupt := exec.InterruptExecutionSession(context.Background(), sessionID, []byte(`{"type":"response.interrupt","response_id":"unknown"}`)) + if !errors.Is(errInterrupt, cliproxyexecutor.ErrUpstreamWebsocketInterruptRejected) { + t.Fatalf("unknown interrupt after completion = %v, want ErrUpstreamWebsocketInterruptRejected", errInterrupt) + } + // A transport-order barrier avoids sleeps or assumptions about whether + // the upstream reader has already consumed a wrongly forwarded frame. + if errWrite := client.WriteMessage(websocket.TextMessage, []byte("barrier")); errWrite != nil { + t.Fatal(errWrite) + } + select { + case count := <-interruptCount: + if count != 0 { + t.Fatalf("forwarded %d interrupt(s) upstream after completion", count) + } + case <-time.After(3 * time.Second): + t.Fatal("upstream barrier timed out") + } + sess.clearActive(client, readCh) +} diff --git a/internal/runtime/executor/codex_websockets_session.go b/internal/runtime/executor/codex_websockets_session.go index cd22affd2cb..633ee1d8c38 100644 --- a/internal/runtime/executor/codex_websockets_session.go +++ b/internal/runtime/executor/codex_websockets_session.go @@ -84,6 +84,16 @@ type codexWebsocketSession struct { responseStateConn *websocket.Conn terminalResponseIDs map[string]struct{} + // activeResponseID is the newest response known to run, or to have run, + // on responseStateConn. The read loop sets it on response.created and on + // the terminal event of that same response, before the payload is + // delivered downstream. activeResponseTerminal then tells an interrupt + // snapshot that the turn already ended even while its read channel is + // still awaiting teardown. A duplex socket may also reject a queued + // create while another response runs; only the matching id ends the + // active response. An empty id means no response.created was seen yet. + activeResponseID string + activeResponseTerminal bool readerConn *websocket.Conn @@ -122,6 +132,8 @@ func (s *codexWebsocketSession) setActiveLocked(conn *websocket.Conn, ch chan co } s.activeConn = conn s.activeCh = ch + s.activeResponseID = "" + s.activeResponseTerminal = false if conn != nil && ch != nil { activeCtx, activeCancel := context.WithCancel(context.Background()) s.activeDone = activeCtx.Done() @@ -206,6 +218,8 @@ func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexW } s.activeConn = nil s.activeCh = nil + s.activeResponseID = "" + s.activeResponseTerminal = false if s.activeCancel != nil { s.activeCancel() } @@ -839,15 +853,53 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, if eventType != "" { sess.setLastEventType(conn, eventType) } - // Record terminal IDs before delivery: downstream may still be - // processing this turn after its active read channel is cleared. - responseID := gjson.GetBytes(payload, "response.id") - if isTerminalEvent(eventType) && eventType != "error" && responseID.Type == gjson.String && responseID.String() != "" { - sess.activeMu.Lock() - if sess.responseStateConn == conn { - sess.terminalResponseIDs[responseID.String()] = struct{}{} + // Track the running response before delivery: downstream may + // still be processing this turn after its active read channel is + // cleared, and interrupts arriving after the completion must not + // forward while teardown is still pending. Terminal events only + // end their own response; a duplex socket can reject a queued + // create while another response is still running. + // + // Failure events (error, response.failed) use the duplex + // attribution shapes: an explicit response id (response.id or + // top-level response_id) ends only that response, while an + // id-less failure ends the running one. Success terminals only + // carry response.id, and one without an id is left unattributed, + // as before. + responseID := gjson.GetBytes(payload, "response.id").String() + if eventType == "response.created" { + if responseID != "" { + sess.activeMu.Lock() + if sess.responseStateConn == conn && sess.activeConn == conn && sess.activeCh != nil { + sess.activeResponseID = responseID + sess.activeResponseTerminal = false + } + sess.activeMu.Unlock() + } + } else if isTerminalEvent(eventType) { + failure := eventType == "error" || eventType == "response.failed" + if failure && responseID == "" { + responseID = gjson.GetBytes(payload, "response_id").String() + } + if responseID != "" || failure { + sess.activeMu.Lock() + if sess.responseStateConn == conn { + terminalID := responseID + if terminalID == "" { + // Id-less failure: the running response ended. + terminalID = sess.activeResponseID + } + if terminalID != "" && eventType != "error" { + // Match the pre-existing terminal-ID recording, + // which never covered error events. + sess.terminalResponseIDs[terminalID] = struct{}{} + } + if terminalID != "" && sess.activeResponseID == terminalID { + sess.activeResponseTerminal = true + } + } + sess.activeMu.Unlock() } - sess.activeMu.Unlock() } } diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 91989d7ddb3..4aa68f321e3 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -303,6 +303,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr bufferedBytes := 0 var initialChunks [][]byte immediateTerminal := false + immediateIncomplete := false // bootstrapTerminalErr holds a non-overload terminal failure seen while buffering. It is // delivered as an in-stream chunk after the buffered handshake so downstream behaviour stays // identical to the unbuffered path instead of silently turning into a credential failover. @@ -513,11 +514,12 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if eventType != "response.incomplete" { cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) } - if detail, ok := helps.ParseCodexUsage(completedPayload); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(completedPayload, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, clientBody, completedPayload, outputItemsByIndex, outputItemsFallback) } var currentChunks [][]byte @@ -555,13 +557,14 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr initialChunks = currentChunks if isTerminalEvent { immediateTerminal = true + immediateIncomplete = eventType == "response.incomplete" } break } } chanCapacity := len(bufferedChunks) + len(initialChunks) - if bootstrapTerminalErr != nil { + if bootstrapTerminalErr != nil || (immediateIncomplete && len(initialChunks) == 0) { chanCapacity++ } out := make(chan cliproxyexecutor.StreamChunk, chanCapacity) @@ -569,7 +572,10 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr out <- cliproxyexecutor.StreamChunk{Payload: chunk} } for _, chunk := range initialChunks { - out <- cliproxyexecutor.StreamChunk{Payload: chunk} + out <- cliproxyexecutor.StreamChunk{Payload: chunk, Incomplete: immediateIncomplete} + } + if immediateIncomplete && len(initialChunks) == 0 { + out <- cliproxyexecutor.StreamChunk{Incomplete: true} } if bootstrapTerminalErr != nil { if isEphemeralSession { @@ -771,11 +777,12 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if eventType != "response.incomplete" { cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) } - if detail, ok := helps.ParseCodexUsage(completedPayload); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(completedPayload, outputItemsByIndex, outputItemsFallback); ok { reporter.Publish(ctx, detail) } else { reporter.EnsurePublished(ctx) } + publishCodexImageToolUsageWithOutputItems(ctx, reporter, clientBody, completedPayload, outputItemsByIndex, outputItemsFallback) } if cliproxyexecutor.DownstreamWebsocket(ctx) { @@ -783,7 +790,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr payload = completedPayload } downstreamPayload := helps.EnsureResponsesUsageDetails(payload) - if !send(cliproxyexecutor.StreamChunk{Payload: downstreamPayload}) { + if !send(cliproxyexecutor.StreamChunk{Payload: downstreamPayload, Incomplete: eventType == "response.incomplete"}) { terminateReason = "context_done" terminateErr = ctx.Err() return @@ -800,16 +807,24 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr payload = completedPayload } eventType = gjson.GetBytes(payload, "type").String() + isIncomplete := eventType == "response.incomplete" line := encodeCodexWebsocketAsSSE(payload) chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens) for i := range chunks { - if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { + if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i], Incomplete: isIncomplete}) { terminateReason = "context_done" terminateErr = ctx.Err() return } emittedPayload = true } + if isIncomplete && len(chunks) == 0 { + if !send(cliproxyexecutor.StreamChunk{Incomplete: true}) { + terminateReason = "context_done" + terminateErr = ctx.Err() + return + } + } if isTerminalEvent || eventType == "response.completed" || eventType == "response.done" || eventType == "response.incomplete" { return } @@ -837,6 +852,9 @@ type codexWebsocketPrepared struct { } func (e *CodexWebsocketsExecutor) prepareCodexWebsocketStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*codexWebsocketPrepared, error) { + if errExpand := helps.ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + return nil, statusErr{code: http.StatusBadRequest, msg: errExpand.Error()} + } baseModel := thinking.ParseSuffix(req.Model).ModelName apiKey, baseURL := codexCredsForRequest(ctx, auth, opts.Headers) if errSIWC := ensureCodexSIWCCredentialsForRequest(ctx, auth, apiKey, opts.Headers); errSIWC != nil { diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index e3a7764acbe..32954d8756c 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -256,7 +256,9 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r reporter.PublishFailureWithDetail(ctx, helps.ParseGeminiUsage(data), err) return cliproxyexecutor.Response{}, err } - reporter.Publish(ctx, helps.ParseGeminiUsage(data)) + gemDetail := helps.ParseGeminiUsage(data) + helps.ApplyGeminiRequestImageConfig(&gemDetail, body) + reporter.Publish(ctx, gemDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } @@ -377,6 +379,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A scanner.Buffer(nil, streamScannerBuffer) claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(body) var terminalDelivered bool defer func() { if errContext := ctx.Err(); errContext != nil && !terminalDelivered { @@ -400,9 +403,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if !sawFinishReason && geminiStreamHasFinishSignal(rawPayload) { sawFinishReason = true } - if detail, ok := helps.ParseGeminiStreamUsage(rawPayload); ok { - streamUsage.Observe(detail, true) - } + streamUsage.ObserveGeminiStream(rawPayload) if blockMsg, blocked := geminiPromptBlockErrorMessage(rawPayload); blocked { blockErr := statusErr{code: http.StatusBadGateway, msg: blockMsg} helps.RecordAPIResponseError(ctx, e.cfg, blockErr) diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index efca220f97a..467e2f6e740 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -568,7 +568,9 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au reporter.PublishFailureWithDetail(ctx, helps.ParseGeminiUsage(data), err) return cliproxyexecutor.Response{}, err } - reporter.Publish(ctx, helps.ParseGeminiUsage(data)) + gemDetail := helps.ParseGeminiUsage(data) + helps.ApplyGeminiRequestImageConfig(&gemDetail, body) + reporter.Publish(ctx, gemDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } @@ -710,7 +712,9 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip reporter.PublishFailureWithDetail(ctx, helps.ParseGeminiUsage(data), err) return cliproxyexecutor.Response{}, err } - reporter.Publish(ctx, helps.ParseGeminiUsage(data)) + gemDetail := helps.ParseGeminiUsage(data) + helps.ApplyGeminiRequestImageConfig(&gemDetail, body) + reporter.Publish(ctx, gemDetail) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } @@ -836,6 +840,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte scanner.Buffer(nil, streamScannerBuffer) claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(body) defer streamUsage.Publish(ctx, reporter) var streamErr error defer func() { @@ -853,9 +858,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte if !sawFinishReason && geminiStreamHasFinishSignal(rawPayload) { sawFinishReason = true } - if detail, ok := helps.ParseGeminiStreamUsage(line); ok { - streamUsage.Observe(detail, true) - } + streamUsage.ObserveGeminiStream(line) if blockMsg, blocked := geminiPromptBlockErrorMessage(rawPayload); blocked { blockErr := statusErr{code: http.StatusBadGateway, msg: blockMsg} helps.RecordAPIResponseError(ctx, e.cfg, blockErr) @@ -1057,6 +1060,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth scanner.Buffer(nil, streamScannerBuffer) claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(body) defer streamUsage.Publish(ctx, reporter) var streamErr error defer func() { @@ -1074,9 +1078,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth if !sawFinishReason && geminiStreamHasFinishSignal(rawPayload) { sawFinishReason = true } - if detail, ok := helps.ParseGeminiStreamUsage(line); ok { - streamUsage.Observe(detail, true) - } + streamUsage.ObserveGeminiStream(line) if blockMsg, blocked := geminiPromptBlockErrorMessage(rawPayload); blocked { blockErr := statusErr{code: http.StatusBadGateway, msg: blockMsg} helps.RecordAPIResponseError(ctx, e.cfg, blockErr) diff --git a/internal/runtime/executor/helps/antigravity_compaction.go b/internal/runtime/executor/helps/antigravity_compaction.go index a550c934c23..c5c749a04c3 100644 --- a/internal/runtime/executor/helps/antigravity_compaction.go +++ b/internal/runtime/executor/helps/antigravity_compaction.go @@ -21,6 +21,7 @@ import ( const ( antigravityCompactionCapsulePrefix = "cpa-ag-compact-v1:" + antigravityCompactionItemIDPrefix = "cmp_ag_compact_" fixedCompactionKeySecret = "CLIProxyAPI" ) @@ -35,6 +36,30 @@ func RecognizedAntigravityCompactionCapsule(encryptedContent string) bool { return strings.HasPrefix(strings.TrimSpace(encryptedContent), antigravityCompactionCapsulePrefix) } +func isRecognizedCPACompactionItem(item gjson.Result) bool { + if item.Get("type").String() != "compaction" { + return false + } + if RecognizedAntigravityCompactionCapsule(item.Get("encrypted_content").String()) { + return true + } + return strings.HasPrefix(strings.TrimSpace(item.Get("id").String()), antigravityCompactionItemIDPrefix) +} + +// HasRecognizedResponsesCompactionItem checks whether input contains a CPA-minted compaction item. +func HasRecognizedResponsesCompactionItem(payload []byte) bool { + input := gjson.GetBytes(payload, "input") + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + if isRecognizedCPACompactionItem(item) { + return true + } + } + return false +} + // HasResponsesCompactionTrigger checks whether input contains a compaction_trigger item. func HasResponsesCompactionTrigger(payload []byte) bool { input := gjson.GetBytes(payload, "input") @@ -89,7 +114,7 @@ func DropForeignResponsesCompactionItems(payload []byte, provider string) []byte dropped := 0 items := make([]string, 0, len(input.Array())) for _, item := range input.Array() { - if item.Get("type").String() == "compaction" && !RecognizedAntigravityCompactionCapsule(item.Get("encrypted_content").String()) { + if item.Get("type").String() == "compaction" && !isRecognizedCPACompactionItem(item) { dropped++ continue } @@ -133,6 +158,32 @@ func ExpandResponsesCompactionCapsules(req *cliproxyexecutor.Request, opts *clip return nil } +// ExpandRecognizedResponsesCompactionCapsules expands recognized CPA compaction capsules in req and opts while preserving foreign compaction items. +func ExpandRecognizedResponsesCompactionCapsules(req *cliproxyexecutor.Request, opts *cliproxyexecutor.Options) error { + sameOriginal := req != nil && opts != nil && len(opts.OriginalRequest) > 0 && bytes.Equal(req.Payload, opts.OriginalRequest) + if req != nil && HasRecognizedResponsesCompactionItem(req.Payload) { + expanded, errExpand := ExpandRecognizedAntigravityCompactionCapsules(req.Payload) + if errExpand != nil { + return errExpand + } + req.Payload = expanded + } + if opts != nil && len(opts.OriginalRequest) > 0 { + if sameOriginal { + opts.OriginalRequest = req.Payload + return nil + } + if HasRecognizedResponsesCompactionItem(opts.OriginalRequest) { + expanded, errExpand := ExpandRecognizedAntigravityCompactionCapsules(opts.OriginalRequest) + if errExpand != nil { + return errExpand + } + opts.OriginalRequest = expanded + } + } + return nil +} + // PrepareAntigravityCompactionSummaryPayload prepares a payload for non-stream Antigravity summary generation. func PrepareAntigravityCompactionSummaryPayload(payload []byte, modelName string) []byte { out := payload @@ -219,6 +270,7 @@ func SealAntigravityCompaction(summary, modelName string) (string, error) { // UnsealAntigravityCompaction decrypts and validates an opaque capsule. func UnsealAntigravityCompaction(encryptedContent string) (string, error) { + encryptedContent = strings.TrimSpace(encryptedContent) if !strings.HasPrefix(encryptedContent, antigravityCompactionCapsulePrefix) { return "", fmt.Errorf("unrecognized compaction capsule format") } @@ -291,6 +343,39 @@ func ExpandAntigravityCompactionCapsules(payload []byte) ([]byte, error) { return sjson.SetRawBytes(payload, "input", []byte(newInput)) } +// ExpandRecognizedAntigravityCompactionCapsules expands recognized CPA compaction items in input into developer context while leaving non-CPA compaction items untouched. +func ExpandRecognizedAntigravityCompactionCapsules(payload []byte) ([]byte, error) { + input := gjson.GetBytes(payload, "input") + if !input.IsArray() { + return payload, nil + } + + var expandedItems []string + changed := false + for _, item := range input.Array() { + if isRecognizedCPACompactionItem(item) { + encryptedContent := item.Get("encrypted_content").String() + summary, errUnseal := UnsealAntigravityCompaction(encryptedContent) + if errUnseal != nil { + return nil, fmt.Errorf("invalid compaction capsule: %w", errUnseal) + } + devMsg := []byte(`{"type":"message","role":"developer","content":[{"type":"input_text"}]}`) + devMsg, _ = sjson.SetBytes(devMsg, "content.0.text", "Context summary from previous turns:\n"+summary) + expandedItems = append(expandedItems, string(devMsg)) + changed = true + continue + } + expandedItems = append(expandedItems, item.Raw) + } + + if !changed { + return payload, nil + } + + newInput := "[" + strings.Join(expandedItems, ",") + "]" + return sjson.SetRawBytes(payload, "input", []byte(newInput)) +} + // ExtractAntigravitySummaryText extracts summary text from Antigravity or Responses response. func ExtractAntigravitySummaryText(respPayload []byte) (string, error) { // 1. OpenAI Responses format: iterate through output array and skip reasoning diff --git a/internal/runtime/executor/helps/antigravity_compaction_test.go b/internal/runtime/executor/helps/antigravity_compaction_test.go index 203a5cf9106..f958e277043 100644 --- a/internal/runtime/executor/helps/antigravity_compaction_test.go +++ b/internal/runtime/executor/helps/antigravity_compaction_test.go @@ -322,3 +322,35 @@ func TestExpandResponsesCompactionCapsulesAndBuildFromSummary(t *testing.T) { t.Fatal("expected error for incomplete summary response, got nil") } } + +func TestExpandRecognizedResponsesCompactionCapsules_PreservesForeign(t *testing.T) { + capsule, errSeal := SealAntigravityCompaction("Sealed CPA summary.", "claude-sonnet-4-6") + if errSeal != nil { + t.Fatalf("SealAntigravityCompaction failed: %v", errSeal) + } + + rawReq := []byte(`{"input":[{"id":"cmp_ag_compact_1","type":"compaction","encrypted_content":"` + capsule + `"},{"id":"cmp_openai_1","type":"compaction","encrypted_content":"gAAAAABopenai-native"},{"type":"message","role":"user","content":"next"}]}`) + req := cliproxyexecutor.Request{Payload: rawReq} + opts := cliproxyexecutor.Options{OriginalRequest: rawReq} + if errExpand := ExpandRecognizedResponsesCompactionCapsules(&req, &opts); errExpand != nil { + t.Fatalf("ExpandRecognizedResponsesCompactionCapsules failed: %v", errExpand) + } + for _, target := range [][]byte{req.Payload, opts.OriginalRequest} { + items := gjson.GetBytes(target, "input").Array() + if len(items) != 3 { + t.Fatalf("input len = %d, want 3; payload=%s", len(items), target) + } + if items[0].Get("type").String() != "message" || items[0].Get("role").String() != "developer" || !strings.Contains(items[0].Get("content.0.text").String(), "Sealed CPA summary.") { + t.Fatalf("items[0] not expanded: %s", items[0].Raw) + } + if items[1].Get("type").String() != "compaction" || items[1].Get("encrypted_content").String() != "gAAAAABopenai-native" { + t.Fatalf("native compaction item not preserved at items[1]: %s", items[1].Raw) + } + } + + badReq := cliproxyexecutor.Request{Payload: []byte(`{"input":[{"id":"cmp_ag_compact_1","type":"compaction","encrypted_content":"cpa-ag-compact-v1:invalid"}]}`)} + badOpts := cliproxyexecutor.Options{} + if errBad := ExpandRecognizedResponsesCompactionCapsules(&badReq, &badOpts); errBad == nil { + t.Fatal("expected error for corrupted CPA compaction capsule, got nil") + } +} diff --git a/internal/runtime/executor/helps/codex_input_ids.go b/internal/runtime/executor/helps/codex_input_ids.go index 07a925c93be..694b5201198 100644 --- a/internal/runtime/executor/helps/codex_input_ids.go +++ b/internal/runtime/executor/helps/codex_input_ids.go @@ -18,9 +18,12 @@ const ( codexFunctionCallItemIDPrefix = "fc" codexCustomToolCallItemIDPrefix = "ctc" codexCustomToolCallOutputItemIDPrefix = "ctco" + codexToolSearchCallItemIDPrefix = "tsc" codexInputItemIDOccupied uint8 = 1 << 0 codexInputItemIDPreserved uint8 = 1 << 1 + codexInputItemIDSeen uint8 = 1 << 2 + codexInputItemIDDuplicate uint8 = 1 << 3 ) // SanitizeCodexInputItemIDs normalizes supported input item IDs for Codex, removes encrypted @@ -45,15 +48,17 @@ func SanitizeCodexInputItemIDs(body []byte) []byte { originalID := itemID.String() id := normalizeCodexInputItemID(item, originalID) state := idStates[id] + if state&codexInputItemIDSeen != 0 { + state |= codexInputItemIDDuplicate + } + state |= codexInputItemIDSeen if id == originalID { state |= codexInputItemIDPreserved } if len([]rune(id)) <= codexInputItemIDLimit { state |= codexInputItemIDOccupied } - if state != 0 { - idStates[id] = state - } + idStates[id] = state } var mapped map[string]string @@ -71,9 +76,18 @@ func SanitizeCodexInputItemIDs(body []byte) []byte { if itemID.Type == gjson.String { originalID := itemID.String() id := normalizeCodexInputItemID(item, originalID) - if id != originalID && idStates[id]&codexInputItemIDPreserved != 0 { - collisionID, ok := collisionMapped[id] - if !ok { + if id != originalID && idStates[id]&(codexInputItemIDPreserved|codexInputItemIDDuplicate) != 0 { + collisionKey := id + "\x00" + originalID + if collisionID, ok := collisionMapped[collisionKey]; ok { + id = collisionID + } else if idStates[id]&codexInputItemIDPreserved == 0 { + if collisionMapped == nil { + collisionMapped = make(map[string]string) + } + collisionMapped[collisionKey] = id + idStates[id] |= codexInputItemIDPreserved + } else { + var collisionID string for attempt := 0; ; attempt++ { collisionID = codexInputItemIDWithHashSuffix(id, attempt) if idStates[collisionID]&codexInputItemIDOccupied != 0 { @@ -82,12 +96,12 @@ func SanitizeCodexInputItemIDs(body []byte) []byte { if collisionMapped == nil { collisionMapped = make(map[string]string) } - collisionMapped[id] = collisionID - idStates[collisionID] |= codexInputItemIDOccupied + collisionMapped[collisionKey] = collisionID + idStates[collisionID] |= codexInputItemIDOccupied | codexInputItemIDPreserved break } + id = collisionID } - id = collisionID } if len([]rune(id)) > codexInputItemIDLimit { shortened, ok := mapped[id] @@ -142,6 +156,11 @@ func normalizeCodexInputItemID(item gjson.Result, id string) string { prefix = codexCustomToolCallItemIDPrefix case "custom_tool_call_output": prefix = codexCustomToolCallOutputItemIDPrefix + case "tool_search_call": + if after, ok := strings.CutPrefix(id, "ts_"); ok && after != "" { + return codexToolSearchCallItemIDPrefix + "_" + after + } + prefix = codexToolSearchCallItemIDPrefix default: return id } diff --git a/internal/runtime/executor/helps/codex_input_ids_test.go b/internal/runtime/executor/helps/codex_input_ids_test.go index c6264f860fe..8dd7bc7c77b 100644 --- a/internal/runtime/executor/helps/codex_input_ids_test.go +++ b/internal/runtime/executor/helps/codex_input_ids_test.go @@ -60,14 +60,19 @@ func TestSanitizeCodexInputItemIDsNormalizesResponseItemIDs(t *testing.T) { reasoningID = "item_reasoning" functionCallID = "item_function_call" functionCallOutputID = "item_function_call_output" + toolSearchCallID = "item_tool_search_call" + legacyToolSearchID = "ts_call_4000860" ) body := []byte(`{"input":[` + `{"type":"message","id":"` + messageID + `"},` + `{"type":"reasoning","id":"` + reasoningID + `"},` + `{"type":"function_call","id":"` + functionCallID + `","call_id":"call-1"},` + `{"type":"function_call_output","id":"` + functionCallOutputID + `","call_id":"call-1"},` + + `{"type":"tool_search_call","id":"` + toolSearchCallID + `","call_id":"call-search-1","execution":"client"},` + + `{"type":"tool_search_call","id":"` + legacyToolSearchID + `","call_id":"call_4000860","execution":"client"},` + `{"type":"reasoning","id":"rs-existing"},` + `{"type":"function_call","id":"fc-existing","call_id":"call-2"},` + + `{"type":"tool_search_call","id":"tsc-existing","call_id":"call-search-2","execution":"client"},` + `{"type":"message","id":"msg-existing"}` + `]}`) @@ -77,8 +82,11 @@ func TestSanitizeCodexInputItemIDsNormalizesResponseItemIDs(t *testing.T) { "rs_" + reasoningID, "fc_" + functionCallID, functionCallOutputID, + "tsc_" + toolSearchCallID, + "tsc_call_4000860", "rs-existing", "fc-existing", + "tsc-existing", "msg-existing", } @@ -105,6 +113,7 @@ func TestSanitizeCodexInputItemIDsAvoidsNormalizationCollisions(t *testing.T) { {name: "function call", itemType: "function_call", prefix: "fc_"}, {name: "custom tool call", itemType: "custom_tool_call", prefix: "ctc_"}, {name: "custom tool call output", itemType: "custom_tool_call_output", prefix: "ctco_"}, + {name: "tool search call", itemType: "tool_search_call", prefix: "tsc_"}, } { for _, idCase := range []struct { name string @@ -159,6 +168,62 @@ func TestSanitizeCodexInputItemIDsAvoidsNormalizationCollisions(t *testing.T) { } } +func TestSanitizeCodexInputItemIDsNormalizesLegacyToolSearchCallIDs(t *testing.T) { + original := []byte(`{"input":[` + + `{"type":"reasoning","id":"rs_signed","encrypted_content":"gAAAA-opaque-signature","summary":[]},` + + `{"type":"tool_search_call","id":"ts_call_4000860","call_id":"call_4000860","execution":"client","status":"completed","arguments":{"query":"agents"}},` + + `{"type":"tool_search_output","id":"tso_call_4000860","call_id":"call_4000860","execution":"client","status":"completed","tools":[]}` + + `]}`) + snapshot := string(original) + + got := SanitizeCodexInputItemIDs(original) + + if string(original) != snapshot { + t.Fatalf("SanitizeCodexInputItemIDs mutated caller slice: got %s, want %s", original, snapshot) + } + if actual := gjson.GetBytes(got, "input.0.id").String(); actual != "rs_signed" { + t.Fatalf("reasoning id = %q, want rs_signed", actual) + } + if actual := gjson.GetBytes(got, "input.0.encrypted_content").String(); actual != "gAAAA-opaque-signature" { + t.Fatalf("encrypted_content = %q, want gAAAA-opaque-signature", actual) + } + if actual := gjson.GetBytes(got, "input.1.id").String(); actual != "tsc_call_4000860" { + t.Fatalf("tool_search_call id = %q, want tsc_call_4000860", actual) + } + if actual := gjson.GetBytes(got, "input.1.call_id").String(); actual != "call_4000860" { + t.Fatalf("tool_search_call call_id = %q, want call_4000860", actual) + } + if actual := gjson.GetBytes(got, "input.2.call_id").String(); actual != "call_4000860" { + t.Fatalf("tool_search_output call_id = %q, want call_4000860", actual) + } + + overlongBase := strings.Repeat("x", codexInputItemIDLimit) + for _, ids := range [][2]string{ + {"ts_call_4000860", "tsc_call_4000860"}, + {"tsc_call_4000860", "ts_call_4000860"}, + {"ts_call_4000860", "call_4000860"}, + {"call_4000860", "ts_call_4000860"}, + {"ts_" + overlongBase, overlongBase}, + {overlongBase, "ts_" + overlongBase}, + } { + collisionBody := []byte(fmt.Sprintf(`{"input":[{"type":"tool_search_call","id":%q},{"type":"tool_search_call","id":%q}]}`, ids[0], ids[1])) + first := SanitizeCodexInputItemIDs(collisionBody) + second := SanitizeCodexInputItemIDs(collisionBody) + again := SanitizeCodexInputItemIDs(first) + id0 := gjson.GetBytes(first, "input.0.id").String() + id1 := gjson.GetBytes(first, "input.1.id").String() + if id0 == id1 { + t.Fatalf("distinct tool_search_call IDs %v collided: %q; payload=%s", ids, id0, first) + } + if !strings.HasPrefix(id0, "tsc_") || !strings.HasPrefix(id1, "tsc_") { + t.Fatalf("normalized IDs must start with tsc_: got (%q, %q)", id0, id1) + } + if string(first) != string(second) || string(first) != string(again) { + t.Fatalf("legacy collision normalization must be deterministic and idempotent: first=%s second=%s again=%s", first, second, again) + } + } +} + func TestSanitizeCodexInputItemIDsNormalizesCustomToolCallIDs(t *testing.T) { const invalidID = "item_44e13caebc1ddf25f1337cbe" body := []byte(`{"input":[{"type":"custom_tool_call","id":"` + invalidID + `","call_id":"call-1","name":"lookup","input":"{}"}]}`) diff --git a/internal/runtime/executor/helps/devin_apply_patch_diagnostic.go b/internal/runtime/executor/helps/devin_apply_patch_diagnostic.go index 7b46f0ca98d..7088d95d39e 100644 --- a/internal/runtime/executor/helps/devin_apply_patch_diagnostic.go +++ b/internal/runtime/executor/helps/devin_apply_patch_diagnostic.go @@ -2,6 +2,7 @@ package helps import ( "fmt" + "strconv" "strings" "unicode/utf8" @@ -229,6 +230,11 @@ func devinPatchReasonCode(reason string) string { switch { case strings.HasSuffix(reason, "in string literal"): return "invalid control character in apply_patch input" + case strings.HasPrefix(reason, "decode apply_patch input value: invalid character ") && strings.HasSuffix(reason, " in string"): + quoted := strings.TrimSuffix(strings.TrimPrefix(reason, "decode apply_patch input value: invalid character "), " in string") + if character, errUnquote := strconv.Unquote(quoted); errUnquote == nil && len(character) == 1 && character[0] < 0x20 { + return "invalid control character in apply_patch input" + } case strings.Contains(reason, "in \\u hexadecimal character escape"): return "invalid Unicode escape in apply_patch input" case strings.Contains(reason, "in string escape code"): diff --git a/internal/runtime/executor/helps/devin_apply_patch_diagnostic_test.go b/internal/runtime/executor/helps/devin_apply_patch_diagnostic_test.go index a24d42a29ff..72441febf87 100644 --- a/internal/runtime/executor/helps/devin_apply_patch_diagnostic_test.go +++ b/internal/runtime/executor/helps/devin_apply_patch_diagnostic_test.go @@ -2,9 +2,11 @@ package helps import ( "errors" + "fmt" "strings" "testing" + applypatch "github.com/router-for-me/CLIProxyAPI/v8/internal/client/codex/apply-patch" "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/common" ) @@ -58,6 +60,31 @@ func TestDevinApplyPatchDiagnosticsUsesTranslatorReason(t *testing.T) { } } +func TestDevinApplyPatchDiagnosticsControlCharacterReasons(t *testing.T) { + for c := byte(0); c < 0x20; c++ { + t.Run(fmt.Sprintf("control_%02x", c), func(t *testing.T) { + _, err := applypatch.UnwrapInput("{\"input\":\"secret" + string(c) + "secret\"}") + if err == nil { + t.Fatal("expected invalid JSON") + } + state := &common.ApplyPatchErrorState{} + state.SetToolInputError(err) + got := (*DevinApplyPatchDiagnostics)(nil).Render(state, "") + if !strings.Contains(got, `reason="invalid control character in apply_patch input"`) || strings.Contains(got, "secret") { + t.Fatalf("unexpected diagnostic: %s", got) + } + }) + } + for _, reason := range []string{ + `decode apply_patch input value: invalid character 'x' in string`, + `decode apply_patch input value: invalid character 'secret' in string`, + } { + if got := devinPatchReasonCode(reason); got != "other" { + t.Errorf("reason %q classified as %q", reason, got) + } + } +} + func TestDevinApplyPatchDiagnosticsBounded(t *testing.T) { if got := (*DevinApplyPatchDiagnostics)(nil).Render(nil, ""); got != ApplyPatchUpstreamErrorMessage { t.Fatalf("nil diagnostic = %q", got) diff --git a/internal/runtime/executor/helps/gemini_stream_tool_usage.go b/internal/runtime/executor/helps/gemini_stream_tool_usage.go new file mode 100644 index 00000000000..08e283880a4 --- /dev/null +++ b/internal/runtime/executor/helps/gemini_stream_tool_usage.go @@ -0,0 +1,77 @@ +package helps + +import ( + "fmt" + "strings" + + translatorcommon "github.com/router-for-me/CLIProxyAPI/v8/internal/translator/common" + "github.com/tidwall/gjson" +) + +type geminiFunctionCallCounter struct { + activeByCandidate map[int64]*translatorcommon.VertexStreamingToolCall + seenCallIDs map[string]struct{} +} + +// observePart counts logical Gemini/Vertex function calls, deduplicating streamed +// partialArgs continuations, repeated call IDs, and empty separator objects. +func (c *geminiFunctionCallCounter) observePart(candidateIdx int64, call gjson.Result) int64 { + if c == nil || !call.Exists() || !call.IsObject() { + return 0 + } + rawName := strings.TrimSpace(call.Get("name").String()) + upstreamID := strings.TrimSpace(call.Get("id").String()) + + if active := c.activeByCandidate[candidateIdx]; active != nil { + ownedCall := gjson.Parse(strings.Clone(call.Raw)) + if active.CanContinueWith(ownedCall) { + _, done := active.Assembler.PushFunctionCall(ownedCall) + if active.Assembler.Err() != nil || done { + delete(c.activeByCandidate, candidateIdx) + } + return 0 + } + } + + delete(c.activeByCandidate, candidateIdx) + + if translatorcommon.IsVertexStreamingFunctionCall(call) { + if rawName == "" { + return 0 + } + ownCall := gjson.Parse(strings.Clone(call.Raw)) + clonedName := strings.Clone(rawName) + clonedID := strings.Clone(upstreamID) + active := translatorcommon.NewVertexStreamingToolCall(clonedName, clonedName, clonedID, clonedID, 0) + _, done := active.Assembler.PushFunctionCall(ownCall) + if active.Assembler.Err() != nil { + return 0 + } + if !done { + if c.activeByCandidate == nil { + c.activeByCandidate = make(map[int64]*translatorcommon.VertexStreamingToolCall) + } + c.activeByCandidate[candidateIdx] = active + } + return c.countLogicalCall(candidateIdx, clonedID) + } + + if rawName == "" && upstreamID == "" && !call.Get("args").Exists() { + return 0 + } + return c.countLogicalCall(candidateIdx, strings.Clone(upstreamID)) +} + +func (c *geminiFunctionCallCounter) countLogicalCall(candidateIdx int64, clonedID string) int64 { + if clonedID != "" { + key := fmt.Sprintf("%d:%s", candidateIdx, clonedID) + if c.seenCallIDs == nil { + c.seenCallIDs = make(map[string]struct{}) + } + if _, exists := c.seenCallIDs[key]; exists { + return 0 + } + c.seenCallIDs[key] = struct{}{} + } + return 1 +} diff --git a/internal/runtime/executor/helps/logging_helpers.go b/internal/runtime/executor/helps/logging_helpers.go index e1db2978eef..cdf1264d16d 100644 --- a/internal/runtime/executor/helps/logging_helpers.go +++ b/internal/runtime/executor/helps/logging_helpers.go @@ -66,6 +66,10 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ if cfg == nil || cfg.CommercialMode { return } + // Check the immutable policy before accessing the pooled request context. + if !cfg.RequestLog && logging.DeferredAPIRequestCaptureDisabled(ctx) { + return + } ginCtx := ginContextFrom(ctx) if ginCtx == nil { return diff --git a/internal/runtime/executor/helps/logging_helpers_test.go b/internal/runtime/executor/helps/logging_helpers_test.go index c9c9f21ecaf..607ecea9cd0 100644 --- a/internal/runtime/executor/helps/logging_helpers_test.go +++ b/internal/runtime/executor/helps/logging_helpers_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" "testing" "github.com/gin-gonic/gin" @@ -71,6 +72,91 @@ func TestRecordAPIRequestClonesDeferredBodyWhenRequestLogDisabled(t *testing.T) } } +func TestRecordAPIRequestDeferredCapturePolicy(t *testing.T) { + for _, tt := range []struct { + name string + disabled bool + fullLog bool + wantDefer bool + }{ + {name: "HTTP with upgrade headers", wantDefer: true}, + {name: "established WebSocket", disabled: true}, + {name: "WebSocket with full logging", disabled: true, fullLog: true}, + } { + t.Run(tt.name, func(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + ginCtx.Request.Header.Set("Connection", "Upgrade") + ginCtx.Request.Header.Set("Upgrade", "websocket") + ctx := context.WithValue(context.Background(), "gin", ginCtx) + if tt.disabled { + ctx = logging.WithDeferredAPIRequestCaptureDisabled(ctx) + } + cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: tt.fullLog}} + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{Method: http.MethodPost, Body: []byte("request-payload")}) + value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey) + if exists != tt.wantDefer { + t.Fatalf("deferred capture present = %t, want %t", exists, tt.wantDefer) + } + if tt.wantDefer { + requests, _ := value.([]logging.DeferredAPIRequest) + if len(requests) != 1 || !strings.Contains(string(requests[0]()), "request-payload") { + t.Fatal("ordinary HTTP lost its error-only request body") + } + } + if tt.fullLog { + logged, _ := ginCtx.Get(apiRequestKey) + if body, _ := logged.([]byte); !strings.Contains(string(body), "request-payload") { + t.Fatal("full logging lost the upstream request body") + } + } + }) + } +} + +func TestRecordAPIRequestDisabledCaptureDoesNotRetainLargeTurns(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx := logging.WithDeferredAPIRequestCaptureDisabled(context.WithValue(context.Background(), "gin", ginCtx)) + body := []byte(strings.Repeat("x", 1<<20)) + for range 40 { + RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{Body: body}) + } + for _, key := range []string{logging.DeferredAPIRequestContextKey, deferredAPIRequestBytesKey} { + if _, exists := ginCtx.Get(key); exists { + t.Fatalf("disabled capture retained %s after 40 MiB of upstream requests", key) + } + } +} + +func TestRecordAPIRequestAsyncCapturePolicySurvivesGinContextReuse(t *testing.T) { + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + parent := context.WithValue(context.Background(), "gin", ginCtx) + socketCtx, cancel := context.WithCancel(logging.WithDeferredAPIRequestCaptureDisabled(parent)) + cancel() + callbackCtx := context.WithoutCancel(socketCtx) + + // A completed socket can leave asynchronous plugin callbacks holding the + // same Gin pointer that is now serving an unrelated HTTP request. + ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + RecordAPIRequest(parent, &config.Config{}, UpstreamRequestLog{Body: []byte("new-http-request")}) + var callbacks sync.WaitGroup + for range 8 { + callbacks.Go(func() { + RecordAPIRequest(callbackCtx, &config.Config{}, UpstreamRequestLog{Body: []byte("old-websocket-request")}) + }) + } + callbacks.Wait() + value, _ := ginCtx.Get(logging.DeferredAPIRequestContextKey) + requests, _ := value.([]logging.DeferredAPIRequest) + if len(requests) != 1 || !strings.Contains(string(requests[0]()), "new-http-request") { + t.Fatal("old WebSocket callbacks changed the new HTTP request's deferred log") + } + used, _ := ginCtx.Get(deferredAPIRequestBytesKey) + if used != len("new-http-request") { + t.Fatalf("deferred bytes = %v, want only the new HTTP body", used) + } +} + func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing.T) { ctx := logging.WithResponseHeadersHolder(context.Background()) headers := http.Header{} diff --git a/internal/runtime/executor/helps/openai_stream_usage_observer.go b/internal/runtime/executor/helps/openai_stream_usage_observer.go new file mode 100644 index 00000000000..1e35ce60c56 --- /dev/null +++ b/internal/runtime/executor/helps/openai_stream_usage_observer.go @@ -0,0 +1,111 @@ +package helps + +import "bytes" + +// OpenAIStreamUsageObserver assembles SSE events for accounting while callers +// continue forwarding the original network chunks unchanged. +type OpenAIStreamUsageObserver struct { + usage *StreamUsageBuffer + line []byte + data [][]byte + dataBytes int + maxBytes int + lineOverflow bool + eventOverflow bool +} + +// NewOpenAIStreamUsageObserver bounds both incomplete lines and assembled events. +func NewOpenAIStreamUsageObserver(usage *StreamUsageBuffer, maxBytes int) *OpenAIStreamUsageObserver { + if maxBytes <= 0 { + maxBytes = defaultMaxStreamModelBufferBound + } + return &OpenAIStreamUsageObserver{usage: usage, maxBytes: maxBytes} +} + +// Feed consumes arbitrary network chunks, retaining incomplete lines across reads. +func (o *OpenAIStreamUsageObserver) Feed(chunk []byte) { + if o == nil || o.usage == nil { + return + } + for len(chunk) > 0 { + idx := bytes.IndexByte(chunk, '\n') + if idx < 0 { + if !o.lineOverflow { + if len(o.line)+len(chunk) > o.maxBytes { + o.dropEvent() + o.line = nil + o.lineOverflow = true + } else { + o.line = append(o.line, chunk...) + } + } + return + } + part := chunk[:idx] + chunk = chunk[idx+1:] + if o.lineOverflow || len(o.line)+len(part) > o.maxBytes { + o.dropEvent() + o.lineOverflow = false + o.line = nil + continue + } + if len(o.line) > 0 { + o.line = append(o.line, part...) + o.observeLine(o.line) + o.line = o.line[:0] + } else { + o.observeLine(part) + } + } +} + +// Finish observes a trailing event even if the upstream omitted its final newline. +func (o *OpenAIStreamUsageObserver) Finish() { + if o == nil || o.usage == nil { + return + } + if !o.lineOverflow && len(o.line) > 0 { + o.observeLine(o.line) + } + o.flushEvent() + o.line = nil + o.lineOverflow = false +} + +func (o *OpenAIStreamUsageObserver) observeLine(line []byte) { + line = bytes.TrimSuffix(line, []byte("\r")) + if len(line) == 0 { + o.flushEvent() + return + } + if o.eventOverflow || !bytes.HasPrefix(line, []byte("data:")) { + return + } + payload := bytes.TrimPrefix(line[len("data:"):], []byte(" ")) + cost := len(payload) + streamEventLineOverhead + if o.dataBytes+cost > o.maxBytes || len(o.data) >= defaultMaxLinesPerStreamEvent { + o.dropEvent() + return + } + o.data = append(o.data, bytes.Clone(payload)) + o.dataBytes += cost +} + +func (o *OpenAIStreamUsageObserver) dropEvent() { + o.eventOverflow = true + o.data = nil + o.dataBytes = 0 +} + +func (o *OpenAIStreamUsageObserver) flushEvent() { + if !o.eventOverflow && len(o.data) > 0 { + if len(o.data) == 1 { + o.usage.ObserveOpenAIStream(o.data[0]) + } else { + o.usage.ObserveOpenAIStream(bytes.Join(o.data, []byte("\n"))) + } + } + o.data = nil + o.dataBytes = 0 + o.eventOverflow = false +} diff --git a/internal/runtime/executor/helps/openai_stream_usage_observer_test.go b/internal/runtime/executor/helps/openai_stream_usage_observer_test.go new file mode 100644 index 00000000000..8963285dc66 --- /dev/null +++ b/internal/runtime/executor/helps/openai_stream_usage_observer_test.go @@ -0,0 +1,58 @@ +package helps + +import ( + "fmt" + "strings" + "testing" +) + +func TestOpenAIStreamUsageObserverChunkBoundaries(t *testing.T) { + const limit = 256 * 1024 + image := `{"type":"image_generation.completed","b64_json":"` + strings.Repeat("A", 100000) + `","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}` + for _, stream := range []string{ + "event: done\ndata: " + image + "\n\n", + ": keepalive\r\nevent: done\r\ndata: " + image + "\r\n\r\n", + "data: " + strings.Replace(image, `,"usage":`, ",\ndata: \"usage\":", 1) + "\n\n", + "data: " + image, + } { + for _, chunkSize := range []int{1, 7, 32768, len(stream)} { + t.Run(fmt.Sprintf("stream_bytes=%d/chunk=%d", len(stream), chunkSize), func(t *testing.T) { + var b StreamUsageBuffer + observer := NewOpenAIStreamUsageObserver(&b, limit) + for remaining := []byte(stream); len(remaining) > 0; { + n := min(chunkSize, len(remaining)) + observer.Feed(remaining[:n]) + remaining = remaining[n:] + } + observer.Finish() + observer.Finish() + d, ok := b.Detail() + if !ok || d.TotalTokens != 30 || d.ImageCount == nil || *d.ImageCount != 1 || b.openAIStreamImages != 1 { + t.Fatalf("chunk framing lost or duplicated image usage: %+v", d) + } + }) + } + } +} + +func TestOpenAIStreamUsageObserverBoundsAndRecovery(t *testing.T) { + const limit = 1024 + oversize := "data: " + strings.Repeat("A", limit*2) + "\n\n" + manyLines := strings.Repeat("data: "+strings.Repeat("B", 100)+"\n", 20) + "\n" + valid := "data: {\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}\n\n" + for _, input := range []string{oversize, manyLines} { + var b StreamUsageBuffer + observer := NewOpenAIStreamUsageObserver(&b, limit) + for _, c := range []byte(input + valid) { + observer.Feed([]byte{c}) + if len(observer.line) > limit || observer.dataBytes > limit { + t.Fatal("observer exceeded its buffer bound") + } + } + observer.Finish() + d, ok := b.Detail() + if !ok || d.TotalTokens != 3 { + t.Fatalf("observer did not recover after oversized event: %+v", d) + } + } +} diff --git a/internal/runtime/executor/helps/plugin_executor_usage.go b/internal/runtime/executor/helps/plugin_executor_usage.go index 27dcf04a5a6..89a6680fb35 100644 --- a/internal/runtime/executor/helps/plugin_executor_usage.go +++ b/internal/runtime/executor/helps/plugin_executor_usage.go @@ -34,11 +34,11 @@ func ParsePluginExecutorResponseUsage(protocol string, payload []byte) usage.Det // response.usage must not hide the completed object's token counts. func parseResponsesPluginExecutorUsage(payload []byte) usage.Detail { detail, ok := ParseCodexUsage(payload) - if ok && hasNonZeroTokenUsage(detail) { + if ok && (hasNonZeroTokenUsage(detail) || hasStructuredUsageDetail(detail)) { return detail } openAIDetail := ParseOpenAIUsage(payload) - if hasNonZeroTokenUsage(openAIDetail) { + if hasNonZeroTokenUsage(openAIDetail) || hasStructuredUsageDetail(openAIDetail) { if openAIDetail.ResponseServiceTier == "" { openAIDetail.ResponseServiceTier = detail.ResponseServiceTier } @@ -58,15 +58,11 @@ func ObservePluginExecutorStreamUsage(protocol string, payload []byte, buffer *S switch strings.ToLower(strings.TrimSpace(protocol)) { case "claude": IterateStreamLines(payload, func(line []byte) { - if detail, ok := parseClaudeStreamLine(line); ok { - ObserveMergedStreamUsage(buffer, detail) - } + buffer.ObserveClaudeStream(line) }) case "gemini": IterateStreamLines(payload, func(line []byte) { - if detail, ok := ParseGeminiStreamUsage(line); ok { - buffer.Observe(detail, ok) - } + buffer.ObserveGeminiStream(line) }) case "interactions", "interactions-response": IterateStreamLines(payload, func(line []byte) { @@ -76,16 +72,17 @@ func ObservePluginExecutorStreamUsage(protocol string, payload []byte, buffer *S }) case "antigravity": IterateStreamLines(payload, func(line []byte) { - if detail, ok := ParseAntigravityStreamUsage(line); ok { - buffer.Observe(detail, ok) - } + buffer.ObserveAntigravityStream(line) }) case "codex", "openai-response": IterateStreamLines(payload, func(line []byte) { if jsonBytes := ExtractStreamJSONPayload(line); len(jsonBytes) > 0 { + if bytes.Contains(jsonBytes, responsesOutputItemDoneMarker) && buffer.ObserveResponsesStream(jsonBytes) { + return + } // A service tier without response.usage is not token usage. Fall through // so a completed Responses object's top-level usage is still recorded. - if detail, ok := ParseCodexUsage(jsonBytes); ok && hasNonZeroTokenUsage(detail) { + if detail, ok := ParseCodexUsageWithOutputItems(jsonBytes, buffer.responsesOutputByIndex, buffer.responsesOutputFallback); ok && (hasNonZeroTokenUsage(detail) || hasStructuredUsageDetail(detail)) { buffer.Observe(detail, true) return } @@ -121,14 +118,17 @@ func parseClaudePayloadUsage(payload []byte) usage.Detail { if len(payload) == 0 || !gjson.ValidBytes(payload) { return usage.Detail{} } - usageNode := gjson.GetBytes(payload, "usage") + root := gjson.ParseBytes(payload) + usageNode := root.Get("usage") if !usageNode.Exists() { - usageNode = gjson.GetBytes(payload, "message.usage") + usageNode = root.Get("message.usage") } if !usageNode.Exists() { return usage.Detail{} } - return ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) + detail := parseClaudeUsageNode(usageNode) + applyClaudeContentBlocksUsage(&detail, root.Get("content")) + return detail } func parseClaudeStreamLine(line []byte) (usage.Detail, bool) { @@ -143,8 +143,7 @@ func parseClaudeStreamLine(line []byte) (usage.Detail, bool) { if !usageNode.Exists() { return usage.Detail{}, false } - detail := ParseClaudeUsage([]byte(`{"usage":` + usageNode.Raw + `}`)) - return detail, true + return parseClaudeUsageNode(usageNode), true } // ObserveMergedStreamUsage updates buffer with merged usage details. @@ -198,6 +197,7 @@ func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { calculatedTotal, okTotal := safeUsageTokenSum(merged.InputTokens, merged.OutputTokens, merged.ReasoningTokens) if !okTotal { merged.TokenBreakdown = invalidUsageTokenBreakdown(merged.TotalTokens) + mergeStreamStructuredFields(&merged, existing, update) return merged } if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal { @@ -211,6 +211,7 @@ func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { merged.ReasoningTokens, merged.TotalTokens, ) + mergeStreamStructuredFields(&merged, existing, update) return merged } @@ -218,6 +219,7 @@ func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { calculatedTotal, okTotal := safeUsageTokenSum(merged.InputTokens, merged.OutputTokens) if !okTotal { merged.TokenBreakdown = invalidUsageTokenBreakdown(merged.TotalTokens) + mergeStreamStructuredFields(&merged, existing, update) return merged } if merged.TotalTokens == 0 || merged.TotalTokens < calculatedTotal { @@ -231,6 +233,7 @@ func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { merged.ReasoningTokens, merged.TotalTokens, ) + mergeStreamStructuredFields(&merged, existing, update) return merged } @@ -254,9 +257,186 @@ func MergeStreamUsageDetail(existing, update usage.Detail) usage.Detail { merged.ReasoningTokens, merged.TotalTokens, ) + mergeStreamStructuredFields(&merged, existing, update) return merged } +func mergeStreamStructuredFields(merged *usage.Detail, existing, update usage.Detail) { + if merged == nil { + return + } + if merged.UpstreamServiceTier == "" { + merged.UpstreamServiceTier = existing.UpstreamServiceTier + } + if merged.Speed == "" { + merged.Speed = existing.Speed + } + if merged.InferenceGeo == "" { + merged.InferenceGeo = existing.InferenceGeo + } + if len(merged.IterationSlice()) == 0 && len(existing.IterationSlice()) > 0 { + merged.Iterations = usage.IterationsPtr(existing.IterationSlice()) + } else if len(merged.IterationSlice()) > 0 { + merged.Iterations = usage.IterationsPtr(merged.IterationSlice()) + } + merged.ToolUsageDetails = mergeToolUsageDetails(existing.ToolUsageDetails, update.ToolUsageDetails) + if merged.WebSearchRequests == nil && existing.WebSearchRequests != nil { + merged.WebSearchRequests = usage.CloneInt64Ptr(existing.WebSearchRequests) + } else if merged.WebSearchRequests != nil && *merged.WebSearchRequests == 0 && existing.WebSearchRequests != nil && *existing.WebSearchRequests > 0 { + merged.WebSearchRequests = usage.CloneInt64Ptr(existing.WebSearchRequests) + } else if merged.WebSearchRequests != nil { + merged.WebSearchRequests = usage.CloneInt64Ptr(merged.WebSearchRequests) + } else if merged.ToolUsageDetails != nil && merged.ToolUsageDetails.WebSearchCalls != nil { + merged.WebSearchRequests = usage.CloneInt64Ptr(merged.ToolUsageDetails.WebSearchCalls) + } + merged.TextInputTokens = mergeInt64Ptr(existing.TextInputTokens, update.TextInputTokens) + merged.ImageInputTokens = mergeInt64Ptr(existing.ImageInputTokens, update.ImageInputTokens) + merged.AudioInputTokens = mergeInt64Ptr(existing.AudioInputTokens, update.AudioInputTokens) + merged.VideoInputTokens = mergeInt64Ptr(existing.VideoInputTokens, update.VideoInputTokens) + merged.FileInputTokens = mergeInt64Ptr(existing.FileInputTokens, update.FileInputTokens) + merged.CachedTextTokens = mergeInt64Ptr(existing.CachedTextTokens, update.CachedTextTokens) + merged.TextOutputTokens = mergeInt64Ptr(existing.TextOutputTokens, update.TextOutputTokens) + merged.ImageOutputTokens = mergeInt64Ptr(existing.ImageOutputTokens, update.ImageOutputTokens) + merged.AudioOutputTokens = mergeInt64Ptr(existing.AudioOutputTokens, update.AudioOutputTokens) + merged.AcceptedPredictionTokens = mergeInt64Ptr(existing.AcceptedPredictionTokens, update.AcceptedPredictionTokens) + merged.RejectedPredictionTokens = mergeInt64Ptr(existing.RejectedPredictionTokens, update.RejectedPredictionTokens) + merged.ImageCount = mergeInt64Ptr(existing.ImageCount, update.ImageCount) + if update.RequestedImageDetails != nil { + merged.RequestedImageDetails = update.RequestedImageDetails.Clone() + } else if existing.RequestedImageDetails != nil { + merged.RequestedImageDetails = existing.RequestedImageDetails.Clone() + } + if update.ImageDetails != nil { + merged.ImageDetails = update.ImageDetails.Clone() + } else if existing.ImageDetails != nil { + merged.ImageDetails = existing.ImageDetails.Clone() + } + if update.MediaCounts != nil { + merged.MediaCounts = update.MediaCounts.Clone() + } else if existing.MediaCounts != nil { + merged.MediaCounts = existing.MediaCounts.Clone() + } + merged.CostDetails = mergeCostDetails(existing.CostDetails, update.CostDetails) + cacheDetails := resolveMergedCacheDetails(existing, update, *merged) + applyCacheDetailsToUsageDetail(merged, cacheDetails) +} + +func mergeInt64Ptr(existing, update *int64) *int64 { + if update != nil { + if *update == 0 && existing != nil && *existing > 0 { + return usage.CloneInt64Ptr(existing) + } + return usage.CloneInt64Ptr(update) + } + return usage.CloneInt64Ptr(existing) +} + +func resolveMergedCacheDetails(existing, update, merged usage.Detail) *usage.CacheDetails { + existingCC := existing.CacheDetails + updateCC := update.CacheDetails + if updateCC != nil { + if update.CacheCreationTokens == 0 && existing.CacheCreationTokens > 0 && existingCC != nil { + if sum, ok := updateCC.SumTokens(); !ok || sum == 0 { + return existingCC.Clone() + } + } + if sum, ok := updateCC.SumTokens(); ok && sum == merged.CacheCreationTokens { + return updateCC.Clone() + } + if iterations := merged.IterationSlice(); len(iterations) > 0 { + if derived := deriveClaudeCacheCreationFromIterations(iterations, merged.CacheCreationTokens); derived != nil { + return derived + } + } + return updateCC.Clone() + } + if iterations := merged.IterationSlice(); len(iterations) > 0 { + if derived := deriveClaudeCacheCreationFromIterations(iterations, merged.CacheCreationTokens); derived != nil { + return derived + } + } + if existingCC != nil && merged.CacheCreationTokens == existing.CacheCreationTokens { + return existingCC.Clone() + } + return nil +} + +func mergeToolUsageDetails(existing, update *usage.ToolUsageDetails) *usage.ToolUsageDetails { + if existing == nil { + return update.Clone() + } + if update == nil { + return existing.Clone() + } + out := update.Clone() + mergeCounter := func(cur, prev *int64) *int64 { + if cur == nil { + return usage.CloneInt64Ptr(prev) + } + if *cur == 0 && prev != nil && *prev > 0 { + return usage.CloneInt64Ptr(prev) + } + return cur + } + out.WebSearchCalls = mergeCounter(out.WebSearchCalls, existing.WebSearchCalls) + out.WebSearchQueries = mergeCounter(out.WebSearchQueries, existing.WebSearchQueries) + out.WebFetchCalls = mergeCounter(out.WebFetchCalls, existing.WebFetchCalls) + out.NumSearchResults = mergeCounter(out.NumSearchResults, existing.NumSearchResults) + out.NumSourcesUsed = mergeCounter(out.NumSourcesUsed, existing.NumSourcesUsed) + out.NumServerSideToolsUsed = mergeCounter(out.NumServerSideToolsUsed, existing.NumServerSideToolsUsed) + out.ToolCallsRequested = mergeCounter(out.ToolCallsRequested, existing.ToolCallsRequested) + out.ToolCallsExecuted = mergeCounter(out.ToolCallsExecuted, existing.ToolCallsExecuted) + out.XSearchCalls = mergeCounter(out.XSearchCalls, existing.XSearchCalls) + out.XPostsFetched = mergeCounter(out.XPostsFetched, existing.XPostsFetched) + out.XUsersFetched = mergeCounter(out.XUsersFetched, existing.XUsersFetched) + out.CodeExecutionCalls = mergeCounter(out.CodeExecutionCalls, existing.CodeExecutionCalls) + out.FileSearchCalls = mergeCounter(out.FileSearchCalls, existing.FileSearchCalls) + out.MCPCalls = mergeCounter(out.MCPCalls, existing.MCPCalls) + out.DocumentSearchCalls = mergeCounter(out.DocumentSearchCalls, existing.DocumentSearchCalls) + out.ImageGenerationCalls = mergeCounter(out.ImageGenerationCalls, existing.ImageGenerationCalls) + if len(existing.Extra) > 0 { + if out.Extra == nil { + out.Extra = make(map[string]int64, len(existing.Extra)) + } + for k, prev := range existing.Extra { + if cur, ok := out.Extra[k]; !ok || (cur == 0 && prev > 0) { + out.Extra[k] = prev + } + } + } + return out +} + +func mergeCostDetails(existing, update *usage.CostDetails) *usage.CostDetails { + if existing == nil { + return update.Clone() + } + if update == nil { + return existing.Clone() + } + out := update.Clone() + e := existing.Clone() + if out.Cost == nil { + out.Cost = e.Cost + } + if out.CostInUSDTicks == nil { + out.CostInUSDTicks = e.CostInUSDTicks + } + if out.IsBYOK == nil { + out.IsBYOK = e.IsBYOK + } + if out.UpstreamInferenceCost == nil { + out.UpstreamInferenceCost = e.UpstreamInferenceCost + } + if out.UpstreamInferencePromptCost == nil { + out.UpstreamInferencePromptCost = e.UpstreamInferencePromptCost + } + if out.UpstreamInferenceCompletionsCost == nil { + out.UpstreamInferenceCompletionsCost = e.UpstreamInferenceCompletionsCost + } + return out +} + func detailInputIncludesCache(detail usage.Detail) (bool, bool) { cache, okCache := safeUsageTokenSum(detail.CacheReadTokens, detail.CacheCreationTokens) if !okCache || (cache == 0 && detail.CachedTokens < 0) { diff --git a/internal/runtime/executor/helps/relay_sse_buffer.go b/internal/runtime/executor/helps/relay_sse_buffer.go new file mode 100644 index 00000000000..12e0c916edc --- /dev/null +++ b/internal/runtime/executor/helps/relay_sse_buffer.go @@ -0,0 +1,69 @@ +package helps + +import "bytes" + +// RelaySSEBuffer reassembles SSE frames across arbitrary relay chunk boundaries. +// Bare JSON relay messages remain whole, including pretty-printed JSON. +type RelaySSEBuffer struct { + pending []byte + sse bool + scan int + lineStart int +} + +// Feed returns complete frames and retains an unfinished SSE frame for the next chunk. +func (b *RelaySSEBuffer) Feed(chunk []byte) [][]byte { + b.pending = append(b.pending, chunk...) + if !b.sse { + trimmed := bytes.TrimLeft(b.pending, " \t\r\n") + if len(trimmed) == 0 { + return nil + } + partialPrefix := false + for _, prefix := range []string{"data:", "event:", "id:", "retry:", ":"} { + if bytes.HasPrefix(trimmed, []byte(prefix)) { + b.sse = true + break + } + partialPrefix = partialPrefix || bytes.HasPrefix([]byte(prefix), trimmed) + } + if !b.sse { + if partialPrefix { + return nil + } + return b.Finish() + } + } + + var frames [][]byte + for b.scan < len(b.pending) { + offset := bytes.IndexByte(b.pending[b.scan:], '\n') + if offset < 0 { + b.scan = len(b.pending) + break + } + end := b.scan + offset + blank := end == b.lineStart || (end == b.lineStart+1 && b.pending[b.lineStart] == '\r') + b.scan, b.lineStart = end+1, end+1 + if blank { + frames = append(frames, bytes.Clone(b.pending[:end+1])) + b.pending = b.pending[end+1:] + b.scan, b.lineStart = 0, 0 + } + } + if len(b.pending) == 0 { + b.pending = nil + } + return frames +} + +// Finish flushes a final frame from relays that omit its trailing blank line. +func (b *RelaySSEBuffer) Finish() [][]byte { + if len(b.pending) == 0 { + return nil + } + frame := b.pending + b.pending = nil + b.scan, b.lineStart = 0, 0 + return [][]byte{frame} +} diff --git a/internal/runtime/executor/helps/relay_sse_buffer_test.go b/internal/runtime/executor/helps/relay_sse_buffer_test.go new file mode 100644 index 00000000000..e1da6859d06 --- /dev/null +++ b/internal/runtime/executor/helps/relay_sse_buffer_test.go @@ -0,0 +1,59 @@ +package helps + +import ( + "bytes" + "testing" +) + +func TestRelaySSEBufferChunkBoundaries(t *testing.T) { + for _, newline := range []string{"\n", "\r\n"} { + first := "event: message" + newline + "data: {\"text\":\"first\"}" + newline + newline + second := ": keep-alive" + newline + newline + last := "data: {\"finishReason\":\"STOP\"}" + newline + newline + wire := []byte(first + second + last) + for split := 0; split <= len(wire); split++ { + var buffer RelaySSEBuffer + frames := buffer.Feed(wire[:split]) + frames = append(frames, buffer.Feed(wire[split:])...) + if len(frames) != 3 || string(frames[0]) != first || string(frames[1]) != second || string(frames[2]) != last { + t.Fatalf("newline %q, split %d: got frames %q", newline, split, frames) + } + if tail := buffer.Finish(); len(tail) != 0 { + t.Fatalf("newline %q, split %d: unexpected tail %q", newline, split, tail) + } + } + var buffer RelaySSEBuffer + for i := range len(first) { + frames := buffer.Feed([]byte{first[i]}) + if i < len(first)-1 && len(frames) != 0 { + t.Fatalf("incomplete frame emitted at byte %d: %q", i, frames) + } + if i == len(first)-1 && (len(frames) != 1 || string(frames[0]) != first) { + t.Fatalf("one-byte chunks lost frame: %q", frames) + } + } + } +} + +func TestRelaySSEBufferFinalFrameAndBareJSON(t *testing.T) { + var buffer RelaySSEBuffer + if frames := buffer.Feed([]byte("data: {\"text\":")); len(frames) != 0 { + t.Fatalf("incomplete frame emitted: %q", frames) + } + if frames := buffer.Feed([]byte("\"tail\"}")); len(frames) != 0 { + t.Fatalf("unterminated frame emitted: %q", frames) + } + if tail := buffer.Finish(); len(tail) != 1 || string(tail[0]) != "data: {\"text\":\"tail\"}" { + t.Fatalf("lost unterminated final frame: %q", tail) + } + if tail := buffer.Finish(); len(tail) != 0 { + t.Fatalf("final frame emitted twice: %q", tail) + } + for _, payload := range []string{"{\n \"text\": \"value\"\n\n}", "[DONE]"} { + var raw RelaySSEBuffer + frames := raw.Feed([]byte(payload)) + if len(frames) != 1 || !bytes.Equal(frames[0], []byte(payload)) { + t.Fatalf("bare payload changed: %q", frames) + } + } +} diff --git a/internal/runtime/executor/helps/response_decode.go b/internal/runtime/executor/helps/response_decode.go index 09373dacfec..6f1419cc719 100644 --- a/internal/runtime/executor/helps/response_decode.go +++ b/internal/runtime/executor/helps/response_decode.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "strings" + "sync" "github.com/andybalholm/brotli" "github.com/klauspost/compress/zstd" @@ -58,43 +59,9 @@ func DecodeResponseBody(body io.ReadCloser, header http.Header) (io.ReadCloser, // br and deflate have none and are left as-is. // The bufio wrapper preserves unread bytes so callers always see the full // stream regardless of whether decompression was applied. - pb := &peekableBody{Reader: bufio.NewReader(body), closer: body} - magic, peekErr := pb.Peek(4) - if peekErr == nil || (peekErr == io.EOF && len(magic) >= 2) { - switch { - case len(magic) >= 2 && magic[0] == 0x1f && magic[1] == 0x8b: - gzipReader, gzErr := gzip.NewReader(pb) - if gzErr != nil { - if errClose := pb.Close(); errClose != nil { - log.Errorf("response decoder: close body error: %v", errClose) - } - return nil, fmt.Errorf("magic-byte gzip: failed to create reader: %w", gzErr) - } - return &compositeReadCloser{ - Reader: gzipReader, - closers: []func() error{ - gzipReader.Close, - pb.Close, - }, - }, nil - case len(magic) >= 4 && magic[0] == 0x28 && magic[1] == 0xb5 && magic[2] == 0x2f && magic[3] == 0xfd: - decoder, zdErr := zstd.NewReader(pb) - if zdErr != nil { - if errClose := pb.Close(); errClose != nil { - log.Errorf("response decoder: close body error: %v", errClose) - } - return nil, fmt.Errorf("magic-byte zstd: failed to create reader: %w", zdErr) - } - return &compositeReadCloser{ - Reader: decoder, - closers: []func() error{ - func() error { decoder.Close(); return nil }, - pb.Close, - }, - }, nil - } - } - return pb, nil + // Detection waits for the first Read: a streaming upstream may send headers long + // before its first byte, and peeking here would block the caller until then. + return &sniffedResponseBody{body: &peekableBody{Reader: bufio.NewReader(body), closer: body}}, nil } encodings := strings.Split(contentEncoding, ",") reader := io.Reader(body) @@ -133,7 +100,7 @@ func DecodeResponseBody(body io.ReadCloser, header http.Header) (io.ReadCloser, case "br": reader = brotli.NewReader(reader) case "zstd": - decoder, errZstd := zstd.NewReader(reader) + decoder, errZstd := newResponseZstdReader(reader) if errZstd != nil { cleanup() return nil, fmt.Errorf("failed to create zstd reader: %w", errZstd) @@ -179,3 +146,78 @@ func isZlibHeader(header []byte) bool { cmf, flg := header[0], header[1] return cmf&0x0f == 8 && cmf>>4 <= 7 && (uint16(cmf)<<8|uint16(flg))%31 == 0 } + +// newResponseZstdReader decodes one response synchronously. The default decoder hands blocks to up +// to four asynchronous block decoders, which costs allocations and handoff latency and gains +// nothing on bodies that are read as they arrive. +func newResponseZstdReader(reader io.Reader) (*zstd.Decoder, error) { + return zstd.NewReader(reader, zstd.WithDecoderConcurrency(1)) +} + +// sniffedResponseBody decodes a body that arrived without Content-Encoding but starts with +// gzip or zstd magic bytes. Detection runs on the first Read and waits for more than one byte +// only when the first byte can start a gzip or zstd frame. +type sniffedResponseBody struct { + body *peekableBody + once sync.Once + mu sync.Mutex + reader io.Reader + decoder func() error + err error +} + +func (s *sniffedResponseBody) detect() { + reader := io.Reader(s.body) + var closeDecoder func() error + defer func() { + s.mu.Lock() + s.reader, s.decoder = reader, closeDecoder + s.mu.Unlock() + }() + first, _ := s.body.Peek(1) + if len(first) == 0 || (first[0] != 0x1f && first[0] != 0x28) { + return + } + magic, errPeek := s.body.Peek(4) + if errPeek != nil && !(errPeek == io.EOF && len(magic) >= 2) { + return + } + switch { + case len(magic) >= 2 && magic[0] == 0x1f && magic[1] == 0x8b: + gzipReader, errGzip := gzip.NewReader(s.body) + if errGzip != nil { + s.err = fmt.Errorf("magic-byte gzip: failed to create reader: %w", errGzip) + return + } + reader, closeDecoder = gzipReader, gzipReader.Close + case len(magic) >= 4 && magic[0] == 0x28 && magic[1] == 0xb5 && magic[2] == 0x2f && magic[3] == 0xfd: + decoder, errZstd := newResponseZstdReader(s.body) + if errZstd != nil { + s.err = fmt.Errorf("magic-byte zstd: failed to create reader: %w", errZstd) + return + } + reader, closeDecoder = decoder, func() error { decoder.Close(); return nil } + } +} + +func (s *sniffedResponseBody) Read(p []byte) (int, error) { + s.once.Do(s.detect) + if s.err != nil { + return 0, s.err + } + return s.reader.Read(p) +} + +func (s *sniffedResponseBody) Close() error { + s.mu.Lock() + closeDecoder := s.decoder + s.mu.Unlock() + var errDecoder error + if closeDecoder != nil { + errDecoder = closeDecoder() + } + if errClose := s.body.Close(); errClose != nil { + return errClose + } + return errDecoder +} diff --git a/internal/runtime/executor/helps/response_decode_bench_test.go b/internal/runtime/executor/helps/response_decode_bench_test.go new file mode 100644 index 00000000000..f7b5444e9aa --- /dev/null +++ b/internal/runtime/executor/helps/response_decode_bench_test.go @@ -0,0 +1,117 @@ +package helps + +import ( + "bytes" + "fmt" + "io" + "math/rand" + "net/http" + "runtime" + "strings" + "testing" + "time" + + "github.com/klauspost/compress/zstd" +) + +// benchZstdSSE returns a zstd-encoded body of about size bytes of Responses-style SSE deltas. +func benchZstdSSE(b *testing.B, size int) []byte { + b.Helper() + r := rand.New(rand.NewSource(1)) + words := strings.Fields("the function return error context request response model input output tool call result value index token cache stream event message content type text") + var plain strings.Builder + for i := 0; plain.Len() < size; i++ { + var delta strings.Builder + for j, n := 0, 4+r.Intn(12); j < n; j++ { + delta.WriteString(words[r.Intn(len(words))]) + delta.WriteByte(' ') + } + fmt.Fprintf(&plain, "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"sequence_number\":%d,\"delta\":%q}\n\n", i, delta.String()) + } + encoder, errNewWriter := zstd.NewWriter(nil) + if errNewWriter != nil { + b.Fatal(errNewWriter) + } + defer func() { _ = encoder.Close() }() + return encoder.EncodeAll([]byte(plain.String()), nil) +} + +var benchZstdHeader = http.Header{"Content-Encoding": {"zstd"}} + +// BenchmarkDecodeResponseBodyZstd decodes whole zstd upstream responses concurrently. +func BenchmarkDecodeResponseBodyZstd(b *testing.B) { + for _, size := range []int{64 << 10, 2 << 20} { + wire := benchZstdSSE(b, size) + b.Run(fmt.Sprintf("%dKiB", size>>10), func(b *testing.B) { + b.SetBytes(int64(size)) + b.ReportAllocs() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + body, errDecode := DecodeResponseBody(io.NopCloser(bytes.NewReader(wire)), benchZstdHeader) + if errDecode != nil { + b.Fatal(errDecode) + } + if _, errCopy := io.Copy(io.Discard, body); errCopy != nil { + b.Fatal(errCopy) + } + _ = body.Close() + } + }) + }) + } +} + +// BenchmarkDecodeResponseBodyOpenStreams holds 256 zstd responses open after their first read, +// like streams waiting on the model, and reports the heap held per open stream. +func BenchmarkDecodeResponseBodyOpenStreams(b *testing.B) { + const streams = 256 + wire := benchZstdSSE(b, 256<<10) + var heap float64 + for i := 0; i < b.N; i++ { + runtime.GC() + var before, during runtime.MemStats + runtime.ReadMemStats(&before) + bodies := make([]io.ReadCloser, streams) + buf := make([]byte, 512) + for s := range bodies { + body, errDecode := DecodeResponseBody(io.NopCloser(bytes.NewReader(wire)), benchZstdHeader) + if errDecode != nil { + b.Fatal(errDecode) + } + if _, errRead := body.Read(buf); errRead != nil { + b.Fatal(errRead) + } + bodies[s] = body + } + runtime.GC() + runtime.ReadMemStats(&during) + heap += float64(int64(during.HeapInuse)-int64(before.HeapInuse)) / streams / 1024 + for _, body := range bodies { + _ = body.Close() + } + } + b.ReportMetric(heap/float64(b.N), "heap-KiB/stream") +} + +// BenchmarkDecodeResponseBodyUnlabeledFirstByte measures how long DecodeResponseBody blocks +// for a response without Content-Encoding whose first byte arrives 20 ms after the headers. +func BenchmarkDecodeResponseBodyUnlabeledFirstByte(b *testing.B) { + var blocked time.Duration + for i := 0; i < b.N; i++ { + reader, writer := io.Pipe() + go func() { + time.Sleep(20 * time.Millisecond) + _, _ = writer.Write([]byte("data: {}\n\n")) + _ = writer.Close() + }() + began := time.Now() + body, errDecode := DecodeResponseBody(reader, http.Header{}) + blocked += time.Since(began) + if errDecode != nil { + b.Fatal(errDecode) + } + _, _ = io.Copy(io.Discard, body) + _ = body.Close() + } + b.ReportMetric(float64(blocked.Milliseconds())/float64(b.N), "blocked-ms") +} diff --git a/internal/runtime/executor/helps/response_decode_test.go b/internal/runtime/executor/helps/response_decode_test.go index 999351bad1d..bc08169a659 100644 --- a/internal/runtime/executor/helps/response_decode_test.go +++ b/internal/runtime/executor/helps/response_decode_test.go @@ -10,6 +10,7 @@ import ( "net/http" "reflect" "testing" + "time" "github.com/andybalholm/brotli" "github.com/klauspost/compress/zstd" @@ -131,3 +132,56 @@ func TestDecodeResponseBodyCompositeClose(t *testing.T) { t.Fatalf("close must attempt all closers and return first error: %v, order=%v", errClose, order) } } + +func TestDecodeResponseBodyUnlabeledDoesNotWaitForFirstByte(t *testing.T) { + reader, writer := io.Pipe() + returned := make(chan io.ReadCloser, 1) + go func() { + body, errDecode := DecodeResponseBody(reader, http.Header{}) + if errDecode != nil { + t.Error(errDecode) + } + returned <- body + }() + var body io.ReadCloser + select { + case body = <-returned: + case <-time.After(2 * time.Second): + t.Fatal("DecodeResponseBody blocked until the upstream sent its first byte") + } + go func() { + _, _ = writer.Write([]byte("data: {}\n\n")) + _ = writer.Close() + }() + got, errRead := io.ReadAll(body) + if errRead != nil || string(got) != "data: {}\n\n" { + t.Fatalf("read %q, %v", got, errRead) + } + if errClose := body.Close(); errClose != nil { + t.Fatal(errClose) + } +} + +func TestDecodeResponseBodySniffsSplitMagicBytes(t *testing.T) { + plain := []byte("data: response.completed\n\n") + for _, encoding := range []string{"gzip", "zstd"} { + wire := encodeResponseDecodeTestBody(t, encoding, plain) + reader, writer := io.Pipe() + go func() { + // The magic bytes arrive one write at a time. + for i := range wire { + _, _ = writer.Write(wire[i : i+1]) + } + _ = writer.Close() + }() + body, errDecode := DecodeResponseBody(reader, http.Header{}) + if errDecode != nil { + t.Fatal(errDecode) + } + got, errRead := io.ReadAll(body) + if errRead != nil || !bytes.Equal(got, plain) { + t.Fatalf("%s: read %q, %v", encoding, got, errRead) + } + _ = body.Close() + } +} diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 75a5f989747..a037cba6137 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -6,12 +6,15 @@ import ( "errors" "fmt" "io" + "mime" + "mime/multipart" "net/http" "reflect" "strings" "sync" "sync/atomic" "time" + "unsafe" "github.com/gin-gonic/gin" "github.com/google/uuid" @@ -21,6 +24,7 @@ import ( 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" + log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -917,34 +921,276 @@ func resolveUsageAuthType(auth *cliproxyauth.Auth) string { // StreamUsageBuffer keeps the latest usage detail observed in a stream. type StreamUsageBuffer struct { - detail usage.Detail - ok bool - responseModel string + detail usage.Detail + ok bool + responseModel string + seenClaudeToolUseIDs map[string]struct{} + seenClaudeToolResultIDs map[string]struct{} + claudeWebSearchQueries int64 + claudeCompletedSearches int64 + claudeSearchResultsCount int64 + claudeHasCompletedSearch bool + claudeReportedWebSearchCalls *int64 + seenOpenAIToolCallKeys map[string]struct{} + openAIToolCallsRequested int64 + openAIExplicitToolCalls bool + geminiCallCounter geminiFunctionCallCounter + openAIStreamImages int64 + openAIExplicitImageCount bool + openAIStreamImageDetails *usage.ImageDetails + responsesOutputByIndex map[int64][]byte + responsesOutputFallback [][]byte + geminiToolDetails *usage.ToolUsageDetails + geminiWebSearchRequests *int64 + geminiImageCount int64 + geminiImageDetails *usage.ImageDetails + requestedImageDetails *usage.ImageDetails } var ( openAIStreamUsageMarker = []byte(`"usage"`) + openAIStreamCacheMarker = []byte(`"prompt_cache_`) openAIStreamServiceTierMarker = []byte(`"service_tier"`) + openAIStreamToolCallsMarker = []byte(`"tool_calls"`) + openAIStreamImageDoneMarker = []byte(`.completed"`) + responsesOutputItemDoneMarker = []byte(`"response.output_item.done"`) + claudeStreamBlockStartMarker = []byte(`"content_block_start"`) + geminiUsageMetaMarker = []byte(`"usageMetadata"`) + geminiUsageMetaSnakeMarker = []byte(`"usage_metadata"`) + geminiGroundingMetaMarker = []byte(`"groundingMetadata"`) + geminiGroundingSnakeMarker = []byte(`"grounding_metadata"`) + geminiURLContextMetaMarker = []byte(`"urlContextMetadata"`) + geminiURLContextSnakeMarker = []byte(`"url_context_metadata"`) + geminiExecCodeMarker = []byte(`"executableCode"`) + geminiExecCodeSnakeMarker = []byte(`"executable_code"`) + geminiFunctionCallMarker = []byte(`"functionCall"`) + geminiFunctionCallSnakeMarker = []byte(`"function_call"`) + geminiInlineDataMarker = []byte(`"inlineData"`) + geminiInlineDataSnakeMarker = []byte(`"inline_data"`) + geminiPredictionsMarker = []byte(`"predictions"`) ) +// unsafeBytesToString converts a byte slice to a string without heap allocation. +// The returned string must only be used for read-only gjson traversal within the +// call frame, and any string values stored into long-lived structs are cloned. +func unsafeBytesToString(b []byte) string { + if len(b) == 0 { + return "" + } + return unsafe.String(unsafe.SliceData(b), len(b)) +} + +func cloneTrimmedString(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return "" + } + return strings.Clone(s) +} + +func hasGeminiCandidateArtifactMarker(payload []byte) bool { + return bytes.Contains(payload, geminiGroundingMetaMarker) || + bytes.Contains(payload, geminiGroundingSnakeMarker) || + bytes.Contains(payload, geminiURLContextMetaMarker) || + bytes.Contains(payload, geminiURLContextSnakeMarker) || + bytes.Contains(payload, geminiExecCodeMarker) || + bytes.Contains(payload, geminiExecCodeSnakeMarker) || + bytes.Contains(payload, geminiFunctionCallMarker) || + bytes.Contains(payload, geminiFunctionCallSnakeMarker) || + bytes.Contains(payload, geminiInlineDataMarker) || + bytes.Contains(payload, geminiInlineDataSnakeMarker) || + bytes.Contains(payload, geminiPredictionsMarker) +} + // Observe records detail when ok is true, allowing the final stream usage to win. func (b *StreamUsageBuffer) Observe(detail usage.Detail, ok bool) { if b == nil || !ok { return } responseServiceTier := strings.TrimSpace(detail.ResponseServiceTier) - if responseServiceTier == "" || hasNonZeroTokenUsage(detail) { + if responseServiceTier == "" || hasNonZeroTokenUsage(detail) || hasStructuredUsageDetail(detail) { preservedTier := b.detail.ResponseServiceTier + preservedUpstreamTier := b.detail.UpstreamServiceTier b.detail = detail if b.detail.ResponseServiceTier == "" { b.detail.ResponseServiceTier = preservedTier } + if b.detail.UpstreamServiceTier == "" { + b.detail.UpstreamServiceTier = preservedUpstreamTier + } } else { b.detail.ResponseServiceTier = responseServiceTier + if b.detail.UpstreamServiceTier == "" { + b.detail.UpstreamServiceTier = detail.UpstreamServiceTier + } } + b.applyGeminiStreamObservations() b.ok = true } +func (b *StreamUsageBuffer) rememberGeminiCandidateArtifacts(detail *usage.Detail) { + if b == nil || detail == nil { + return + } + if detail.ToolUsageDetails != nil { + // Content parts are incremental executions; grounding metadata is a snapshot. + previous := b.geminiToolDetails + b.geminiToolDetails = mergeToolUsageDetails(previous, detail.ToolUsageDetails) + if previous != nil { + addCounter := func(prev, delta *int64) *int64 { + if prev == nil || delta == nil { + return mergeInt64Ptr(prev, delta) + } + if sum, ok := safeUsageTokenSum(*prev, *delta); ok { + return usage.Int64Ptr(sum) + } + return nil + } + b.geminiToolDetails.ToolCallsRequested = addCounter( + previous.ToolCallsRequested, detail.ToolUsageDetails.ToolCallsRequested, + ) + b.geminiToolDetails.CodeExecutionCalls = addCounter( + previous.CodeExecutionCalls, detail.ToolUsageDetails.CodeExecutionCalls, + ) + } + } + if detail.WebSearchRequests != nil { + b.geminiWebSearchRequests = mergeInt64Ptr(b.geminiWebSearchRequests, detail.WebSearchRequests) + } + if detail.ImageCount != nil && *detail.ImageCount > 0 { + b.geminiImageCount += *detail.ImageCount + } + if detail.ImageDetails != nil { + b.geminiImageDetails = detail.ImageDetails.Clone() + } +} + +func (b *StreamUsageBuffer) applyGeminiStreamObservations() { + if b == nil { + return + } + if b.geminiToolDetails != nil { + b.detail.ToolUsageDetails = mergeToolUsageDetails(b.detail.ToolUsageDetails, b.geminiToolDetails) + } + if b.geminiWebSearchRequests != nil { + b.detail.WebSearchRequests = usage.CloneInt64Ptr(b.geminiWebSearchRequests) + } + if b.geminiImageCount > 0 { + b.detail.ImageCount = usage.Int64Ptr(b.geminiImageCount) + } + if b.geminiImageDetails != nil && b.detail.ImageDetails == nil { + b.detail.ImageDetails = b.geminiImageDetails.Clone() + } + if b.geminiImageCount > 0 && b.requestedImageDetails != nil && b.detail.RequestedImageDetails == nil { + b.detail.RequestedImageDetails = b.requestedImageDetails.Clone() + } +} + +// SetRequestPayload records request-side image/media configuration on the stream buffer +// so it can be attached as RequestedImageDetails if the stream produces images. +func (b *StreamUsageBuffer) SetRequestPayload(requestPayload []byte, contentType ...string) { + if b == nil || len(requestPayload) == 0 { + return + } + b.requestedImageDetails = ParseRequestImageDetails(requestPayload, contentType...) +} + +// ObserveGeminiStream records and merges usage and candidate tool/media artifacts from a Gemini SSE line. +func (b *StreamUsageBuffer) ObserveGeminiStream(line []byte) { + b.observeGeminiFamilyStream(line, false) +} + +// ObserveAntigravityStream records and merges usage and candidate tool/media artifacts from an Antigravity SSE line. +func (b *StreamUsageBuffer) ObserveAntigravityStream(line []byte) { + b.observeGeminiFamilyStream(line, true) +} + +func (b *StreamUsageBuffer) observeGeminiFamilyStream(line []byte, allowResponseWrapper bool) { + if b == nil { + return + } + payload := jsonPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return + } + root := gjson.Parse(unsafeBytesToString(payload)) + var ( + node gjson.Result + candidatesNode gjson.Result + predsNode gjson.Result + ) + root.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "usageMetadata", "usage_metadata": + if !node.Exists() { + node = val + } + case "candidates": + candidatesNode = val + case "predictions": + predsNode = val + case "response": + if allowResponseWrapper && val.IsObject() { + val.ForEach(func(rk, rv gjson.Result) bool { + switch rk.String() { + case "usageMetadata", "usage_metadata": + if !node.Exists() { + node = rv + } + case "candidates": + if !candidatesNode.Exists() { + candidatesNode = rv + } + } + return true + }) + } + } + return true + }) + if candidatesNode.Exists() || predsNode.Exists() { + var artifactDetail usage.Detail + applyGeminiNodesArtifactsWithCounter(&artifactDetail, candidatesNode, predsNode, &b.geminiCallCounter) + b.rememberGeminiCandidateArtifacts(&artifactDetail) + if !node.Exists() && (artifactDetail.ToolUsageDetails != nil || artifactDetail.ImageCount != nil) { + b.applyGeminiStreamObservations() + b.ok = true + return + } + } + if !node.Exists() { + return + } + detail := parseGeminiFamilyUsageDetail(node) + if !allowResponseWrapper && !hasNonZeroTokenUsage(detail) { + return + } + if b.ok { + update := detail + mergeStreamStructuredFields(&detail, b.detail, update) + // Explicit zero modality counts are authoritative, not absent fields. + for _, pair := range []struct { + dst **int64 + src *int64 + }{ + {&detail.TextInputTokens, update.TextInputTokens}, + {&detail.ImageInputTokens, update.ImageInputTokens}, + {&detail.AudioInputTokens, update.AudioInputTokens}, + {&detail.VideoInputTokens, update.VideoInputTokens}, + {&detail.FileInputTokens, update.FileInputTokens}, + {&detail.CachedTextTokens, update.CachedTextTokens}, + {&detail.TextOutputTokens, update.TextOutputTokens}, + {&detail.ImageOutputTokens, update.ImageOutputTokens}, + {&detail.AudioOutputTokens, update.AudioOutputTokens}, + } { + if pair.src != nil { + *pair.dst = usage.CloneInt64Ptr(pair.src) + } + } + } + b.Observe(detail, true) +} + // ObserveOpenAIStream records response-tier state and the latest usage from an // OpenAI-style stream while avoiding JSON parsing for irrelevant chunks. func (b *StreamUsageBuffer) ObserveOpenAIStream(line []byte) { @@ -957,33 +1203,145 @@ func (b *StreamUsageBuffer) ObserveOpenAIStream(line []byte) { } hasUsageCandidate := bytes.Contains(payload, openAIStreamUsageMarker) + hasCacheCandidate := bytes.Contains(payload, openAIStreamCacheMarker) needTier := b.detail.ResponseServiceTier == "" || hasUsageCandidate hasTierCandidate := needTier && bytes.Contains(payload, openAIStreamServiceTierMarker) + hasToolCallsCandidate := bytes.Contains(payload, openAIStreamToolCallsMarker) + hasImageCompletedCandidate := bytes.Contains(payload, openAIStreamImageDoneMarker) if b.responseModel == "" { if model, _ := extractGenericResponseModelEvent(payload); model != "" { b.responseModel = model } } - if !hasUsageCandidate && !hasTierCandidate { + if !hasUsageCandidate && !hasCacheCandidate && !hasTierCandidate && !hasToolCallsCandidate && !hasImageCompletedCandidate { return } if !gjson.ValidBytes(payload) { return } + if hasToolCallsCandidate { + b.observeOpenAIStreamToolCalls(payload) + } + if hasImageCompletedCandidate { + b.observeOpenAIImageStreamCompleted(payload) + } + detail := usage.Detail{} usageOK := false - if hasUsageCandidate { - usageNode := gjson.GetBytes(payload, "usage") - if hasOpenAIStyleUsageTokenFields(usageNode) { - detail = parseOpenAIStyleUsageNode(usageNode) - usageOK = true + if hasUsageCandidate || hasCacheCandidate { + root := gjson.Parse(unsafeBytesToString(payload)) + usageNode := root.Get("usage") + detail, usageOK = parseOpenAIStyleUsageWithCache(root, usageNode) + if detail.ToolUsageDetails != nil && detail.ToolUsageDetails.ToolCallsRequested != nil { + b.openAIExplicitToolCalls = true + } + if detail.ImageCount != nil { + b.openAIExplicitImageCount = true + } + if usageOK && b.ok { + // Metadata-only frames do not replace the last token snapshot. + if !hasOpenAIStyleUsageTokenFields(usageNode) { + detail.InputTokens = b.detail.InputTokens + detail.OutputTokens = b.detail.OutputTokens + detail.ReasoningTokens = b.detail.ReasoningTokens + detail.CachedTokens = b.detail.CachedTokens + detail.CacheReadTokens = b.detail.CacheReadTokens + detail.CacheCreationTokens = b.detail.CacheCreationTokens + detail.TotalTokens = b.detail.TotalTokens + detail.TokenBreakdown = b.detail.TokenBreakdown + } + mergeStreamStructuredFields(&detail, b.detail, detail) } } if hasTierCandidate { detail.ResponseServiceTier = extractResponseServiceTierFromValidJSON(payload) + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = detail.ResponseServiceTier + } } b.Observe(detail, usageOK || detail.ResponseServiceTier != "") + b.applyOpenAIStreamObservations() +} + +func (b *StreamUsageBuffer) observeOpenAIStreamToolCalls(payload []byte) { + root := gjson.ParseBytes(payload) + for choiceIdx, choice := range root.Get("choices").Array() { + if idx := choice.Get("index"); idx.Type == gjson.Number { + choiceIdx = int(idx.Int()) + } + tcNode := choice.Get("delta.tool_calls") + if !tcNode.IsArray() { + tcNode = choice.Get("message.tool_calls") + } + if !tcNode.IsArray() { + continue + } + for itemIdx, tc := range tcNode.Array() { + key := "" + if idxNode := tc.Get("index"); idxNode.Type == gjson.Number { + key = fmt.Sprintf("%d:%d", choiceIdx, idxNode.Int()) + } else if id := strings.TrimSpace(tc.Get("id").String()); id != "" { + key = fmt.Sprintf("%d:id:%s", choiceIdx, id) + } else { + key = fmt.Sprintf("%d:pos:%d", choiceIdx, itemIdx) + } + if b.seenOpenAIToolCallKeys == nil { + b.seenOpenAIToolCallKeys = make(map[string]struct{}) + } + if _, exists := b.seenOpenAIToolCallKeys[key]; exists { + continue + } + b.seenOpenAIToolCallKeys[key] = struct{}{} + b.openAIToolCallsRequested++ + } + } +} + +func (b *StreamUsageBuffer) observeOpenAIImageStreamCompleted(payload []byte) { + root := gjson.ParseBytes(payload) + evType := strings.TrimSpace(root.Get("type").String()) + if evType != "image_generation.completed" && evType != "image_edit.completed" { + return + } + b.openAIStreamImages++ + details := &usage.ImageDetails{ + Size: strings.TrimSpace(root.Get("size").String()), + Quality: strings.TrimSpace(root.Get("quality").String()), + Background: strings.TrimSpace(root.Get("background").String()), + OutputFormat: strings.TrimSpace(root.Get("output_format").String()), + MimeType: strings.TrimSpace(root.Get("mime_type").String()), + } + if !isEmptyImageDetails(details) { + b.openAIStreamImageDetails = details + } +} + +func (b *StreamUsageBuffer) applyOpenAIStreamObservations() { + if b == nil { + return + } + if b.openAIToolCallsRequested > 0 { + if b.detail.ToolUsageDetails == nil { + b.detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if !b.openAIExplicitToolCalls { + b.detail.ToolUsageDetails.ToolCallsRequested = usage.Int64Ptr(b.openAIToolCallsRequested) + } + b.ok = true + } + if b.openAIStreamImages > 0 { + if !b.openAIExplicitImageCount || b.detail.ImageCount == nil { + b.detail.ImageCount = usage.Int64Ptr(b.openAIStreamImages) + } + if b.openAIStreamImageDetails != nil { + b.detail.ImageDetails = b.openAIStreamImageDetails.Clone() + } + if b.detail.RequestedImageDetails == nil && b.requestedImageDetails != nil { + b.detail.RequestedImageDetails = b.requestedImageDetails.Clone() + } + b.ok = true + } } // ObserveClaudeStream records and merges usage from a Claude SSE line. @@ -991,16 +1349,129 @@ func (b *StreamUsageBuffer) ObserveClaudeStream(line []byte) { if b == nil { return } - if b.responseModel == "" { - if payload := jsonPayload(line); len(payload) > 0 { - if model, _ := extractClaudeResponseModelEvent(payload); model != "" { - b.responseModel = model - } + payload := jsonPayload(line) + if b.responseModel == "" && len(payload) > 0 { + if model, _ := extractClaudeResponseModelEvent(payload); model != "" { + b.responseModel = model } } + if len(payload) > 0 { + b.observeClaudeStreamToolBlocks(payload) + } if detail, ok := ParseClaudeStreamUsage(line); ok { + if detail.ToolUsageDetails != nil && detail.ToolUsageDetails.WebSearchCalls != nil && *detail.ToolUsageDetails.WebSearchCalls > 0 { + b.claudeReportedWebSearchCalls = usage.CloneInt64Ptr(detail.ToolUsageDetails.WebSearchCalls) + } ObserveMergedStreamUsage(b, detail) + b.applyClaudeStreamToolObservations() + } +} + +// ObserveClaudeSSEEvent extracts joined multi-line SSE data payloads from rawEvent +// and records any model or usage details when a frame was split across lines. +func (b *StreamUsageBuffer) ObserveClaudeSSEEvent(rawEvent []byte) { + if b == nil || len(rawEvent) == 0 { + return + } + var dataLines [][]byte + for _, line := range bytes.Split(rawEvent, []byte("\n")) { + trimmed := bytes.TrimSpace(line) + if !bytes.HasPrefix(trimmed, []byte("data:")) { + continue + } + payload := bytes.TrimSpace(trimmed[len("data:"):]) + if len(payload) > 0 { + dataLines = append(dataLines, payload) + } + } + if len(dataLines) <= 1 { + return } + joined := bytes.TrimSpace(bytes.Join(dataLines, []byte("\n"))) + if len(joined) == 0 || !gjson.ValidBytes(joined) { + return + } + b.ObserveClaudeStream(joined) +} + +func (b *StreamUsageBuffer) observeClaudeStreamToolBlocks(payload []byte) { + if b == nil || !bytes.Contains(payload, claudeStreamBlockStartMarker) || !gjson.ValidBytes(payload) { + return + } + root := gjson.Parse(unsafeBytesToString(payload)) + if root.Get("type").String() != "content_block_start" { + return + } + cb := root.Get("content_block") + if !cb.Exists() || !cb.IsObject() { + return + } + switch cb.Get("type").String() { + case "server_tool_use": + if strings.TrimSpace(cb.Get("name").String()) != "web_search" { + return + } + id := strings.TrimSpace(cb.Get("id").String()) + if id != "" { + if b.seenClaudeToolUseIDs == nil { + b.seenClaudeToolUseIDs = make(map[string]struct{}) + } + if _, exists := b.seenClaudeToolUseIDs[id]; exists { + return + } + b.seenClaudeToolUseIDs[strings.Clone(id)] = struct{}{} + } + b.claudeWebSearchQueries++ + b.applyClaudeStreamToolObservations() + case "web_search_tool_result": + content := cb.Get("content") + if !content.Exists() || !content.IsArray() { + return + } + toolUseID := strings.TrimSpace(cb.Get("tool_use_id").String()) + if toolUseID != "" { + if b.seenClaudeToolResultIDs == nil { + b.seenClaudeToolResultIDs = make(map[string]struct{}) + } + if _, exists := b.seenClaudeToolResultIDs[toolUseID]; exists { + return + } + b.seenClaudeToolResultIDs[strings.Clone(toolUseID)] = struct{}{} + } + b.claudeHasCompletedSearch = true + b.claudeCompletedSearches++ + for _, item := range content.Array() { + if item.Get("type").String() == "web_search_result" { + b.claudeSearchResultsCount++ + } + } + b.applyClaudeStreamToolObservations() + } +} + +func (b *StreamUsageBuffer) applyClaudeStreamToolObservations() { + if b == nil { + return + } + if b.claudeWebSearchQueries == 0 && !b.claudeHasCompletedSearch { + return + } + if b.detail.ToolUsageDetails == nil { + b.detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if b.claudeWebSearchQueries > 0 { + b.detail.ToolUsageDetails.WebSearchQueries = usage.Int64Ptr(b.claudeWebSearchQueries) + } + if b.claudeHasCompletedSearch { + b.detail.ToolUsageDetails.NumSearchResults = usage.Int64Ptr(b.claudeSearchResultsCount) + calls := b.claudeCompletedSearches + if b.claudeReportedWebSearchCalls != nil { + calls = *b.claudeReportedWebSearchCalls + } + b.detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(calls) + b.detail.WebSearchRequests = usage.Int64Ptr(calls) + } + b.ok = true } // Publish emits the latest observed usage detail, if any. @@ -1043,39 +1514,276 @@ func (b *StreamUsageBuffer) ResponseModel() string { return b.responseModel } +// ObserveResponsesStream collects response.output_item.done items across a +// Responses SSE/WebSocket stream and parses terminal response.usage payloads +// with those collected items so tool usage is preserved when response.output is empty. +func (b *StreamUsageBuffer) ObserveResponsesStream(line []byte) bool { + if b == nil { + return false + } + payload := jsonPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return false + } + if b.responseModel == "" { + if model, _ := extractGenericResponseModelEvent(payload); model != "" { + b.responseModel = model + } + } + if bytes.Contains(payload, responsesOutputItemDoneMarker) { + root := gjson.Parse(unsafeBytesToString(payload)) + if strings.TrimSpace(root.Get("type").String()) == "response.output_item.done" { + if item := root.Get("item"); item.Exists() && item.IsObject() { + rawItem := bytes.Clone([]byte(item.Raw)) + if idxNode := root.Get("output_index"); idxNode.Type == gjson.Number && idxNode.Int() >= 0 { + if b.responsesOutputByIndex == nil { + b.responsesOutputByIndex = make(map[int64][]byte) + } + b.responsesOutputByIndex[idxNode.Int()] = rawItem + } else { + b.responsesOutputFallback = append(b.responsesOutputFallback, rawItem) + } + } + return true + } + } + detail, ok := ParseCodexUsageWithOutputItems(payload, b.responsesOutputByIndex, b.responsesOutputFallback) + if !ok || (!hasNonZeroTokenUsage(detail) && !hasStructuredUsageDetail(detail) && detail.ResponseServiceTier == "") { + return false + } + b.Observe(detail, true) + return true +} + func ParseCodexUsage(data []byte) (usage.Detail, bool) { - responseServiceTier := extractResponseServiceTier(data) - usageNode := gjson.ParseBytes(data).Get("response.usage") - if !hasOpenAIStyleUsageTokenFields(usageNode) { + return ParseCodexUsageWithOutputItems(data, nil, nil) +} + +func ParseCodexUsageWithOutputItems(data []byte, itemsByIndex map[int64][]byte, fallback [][]byte) (usage.Detail, bool) { + if len(data) == 0 { + return usage.Detail{}, false + } + var responseServiceTier string + if bytes.Contains(data, openAIStreamServiceTierMarker) { + responseServiceTier = extractResponseServiceTier(data) + } + root := gjson.Parse(unsafeBytesToString(data)) + usageNode := root.Get("response.usage") + detail, ok := parseOpenAIStyleUsageWithCache(root, usageNode) + if !ok { if responseServiceTier == "" { return usage.Detail{}, false } - return usage.Detail{ResponseServiceTier: responseServiceTier}, true + return usage.Detail{ + ResponseServiceTier: responseServiceTier, + UpstreamServiceTier: responseServiceTier, + }, true } - detail := parseOpenAIStyleUsageNode(usageNode) detail.ResponseServiceTier = responseServiceTier + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = responseServiceTier + } + if wsReq := root.Get("response.tool_usage.web_search.num_requests"); wsReq.Type == gjson.Number { + v := wsReq.Int() + if detail.ToolUsageDetails == nil { + detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if detail.ToolUsageDetails.WebSearchCalls == nil || (v > 0 && *detail.ToolUsageDetails.WebSearchCalls == 0) { + detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(v) + } + if detail.WebSearchRequests == nil || (v > 0 && *detail.WebSearchRequests == 0) { + detail.WebSearchRequests = usage.Int64Ptr(v) + } + } + outputNode := root.Get("response.output") + if outputNode.Exists() && outputNode.IsArray() && len(outputNode.Array()) > 0 { + applyResponsesOutputItemsUsage(&detail, outputNode) + } else if len(itemsByIndex) > 0 || len(fallback) > 0 { + applyResponsesCollectedOutputItemsUsage(&detail, itemsByIndex, fallback) + } + // When response.tool_usage.image_gen is present, CodexExecutor publishes a dedicated + // additional usage record for the image model (e.g. gpt-image-2) carrying ImageCount + // and ImageDetails; clear them on the orchestrator model record to avoid double-counting. + if hasOpenAIStyleUsageTokenFields(root.Get("response.tool_usage.image_gen")) { + detail.ImageCount = nil + detail.ImageDetails = nil + } return detail, true } +func applyResponsesCacheRetention(detail *usage.Detail, root, usageNode gjson.Result) { + if detail == nil { + return + } + retention := strings.TrimSpace(firstExistingUsageNode( + root, + "response.prompt_cache_retention", + "prompt_cache_retention", + ).String()) + if retention == "" { + retention = strings.TrimSpace(firstExistingUsageNode( + usageNode, + "prompt_cache_retention", + "prompt_cache_options.retention", + ).String()) + } + options := firstExistingUsageNode(root, "response.prompt_cache_options", "prompt_cache_options") + mode := cloneTrimmedString(options.Get("mode").String()) + ttl := cloneTrimmedString(options.Get("ttl").String()) + buckets := firstExistingUsageNode(usageNode, "input_tokens_details.cache_write_tokens_details", "prompt_tokens_details.cache_write_tokens_details") + short := buckets.Get("ttl_30m_tokens") + long := buckets.Get("ttl_12h_tokens") + if retention == "" && mode == "" && ttl == "" && short.Type != gjson.Number && long.Type != gjson.Number { + return + } + if detail.CacheDetails == nil { + detail.CacheDetails = &usage.CacheDetails{} + } + detail.CacheDetails.Retention = strings.Clone(retention) + detail.CacheDetails.Mode = mode + detail.CacheDetails.TTL = ttl + if short.Type == gjson.Number { + detail.CacheCreation30mTokens = usage.Int64Ptr(short.Int()) + detail.CacheDetails.Ephemeral30mInputTokens = usage.Int64Ptr(short.Int()) + } + if long.Type == gjson.Number { + detail.CacheCreation12hTokens = usage.Int64Ptr(long.Int()) + detail.CacheDetails.Ephemeral12hInputTokens = usage.Int64Ptr(long.Int()) + } +} + func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) { - usageNode := gjson.ParseBytes(data).Get("response.tool_usage.image_gen") + return ParseCodexImageToolUsageWithOutputItems(nil, data, nil, nil) +} + +func ParseCodexImageToolUsageWithOutputItems(requestBody, data []byte, itemsByIndex map[int64][]byte, fallback [][]byte) (usage.Detail, bool) { + if len(data) == 0 { + return usage.Detail{}, false + } + root := gjson.Parse(unsafeBytesToString(data)) + usageNode := root.Get("response.tool_usage.image_gen") if !hasOpenAIStyleUsageTokenFields(usageNode) { return usage.Detail{}, false } - return parseOpenAIStyleUsageNode(usageNode), true + detail := parseOpenAIStyleUsageNode(usageNode) + applyCodexImageToolOutputMetadata(&detail, requestBody, root.Get("response.output"), itemsByIndex, fallback) + return detail, true +} + +func applyCodexImageToolOutputMetadata(detail *usage.Detail, requestBody []byte, outputNode gjson.Result, itemsByIndex map[int64][]byte, fallback [][]byte) { + if detail == nil { + return + } + var items []gjson.Result + if outputNode.Exists() && outputNode.IsArray() && len(outputNode.Array()) > 0 { + items = outputNode.Array() + } else { + items = parseCollectedOutputItems(itemsByIndex, fallback) + } + var count int64 + var respMeta *usage.ImageDetails + for _, item := range items { + if strings.TrimSpace(item.Get("type").String()) != "image_generation_call" { + continue + } + if status := strings.TrimSpace(item.Get("status").String()); strings.EqualFold(status, "failed") || strings.EqualFold(status, "in_progress") { + continue + } + count++ + if respMeta == nil { + meta := &usage.ImageDetails{ + Size: cloneTrimmedString(item.Get("size").String()), + Resolution: cloneTrimmedString(item.Get("resolution").String()), + AspectRatio: cloneTrimmedString(item.Get("aspect_ratio").String()), + Quality: cloneTrimmedString(item.Get("quality").String()), + Background: cloneTrimmedString(item.Get("background").String()), + OutputFormat: cloneTrimmedString(item.Get("output_format").String()), + MimeType: cloneTrimmedString(firstExistingUsageNode(item, "mime_type", "mimeType").String()), + InputFidelity: cloneTrimmedString(item.Get("input_fidelity").String()), + } + if !isEmptyImageDetails(meta) { + respMeta = meta + } + } + } + if count > 0 { + detail.ImageCount = usage.Int64Ptr(count) + } + detail.ImageDetails = respMeta + detail.RequestedImageDetails = ParseRequestImageDetails(requestBody) } func ParseOpenAIUsage(data []byte) usage.Detail { - responseServiceTier := extractResponseServiceTier(data) - usageNode := gjson.ParseBytes(data).Get("usage") - if !hasOpenAIStyleUsageTokenFields(usageNode) { - return usage.Detail{ResponseServiceTier: responseServiceTier} + if len(data) == 0 { + return usage.Detail{} } - detail := parseOpenAIStyleUsageNode(usageNode) + var responseServiceTier string + if bytes.Contains(data, openAIStreamServiceTierMarker) { + responseServiceTier = extractResponseServiceTier(data) + } + root := gjson.Parse(unsafeBytesToString(data)) + usageNode := root.Get("usage") + detail, _ := parseOpenAIStyleUsageWithCache(root, usageNode) detail.ResponseServiceTier = responseServiceTier + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = responseServiceTier + } + applyOpenAIChatToolCallsUsage(&detail, root) + if outputNode := root.Get("output"); outputNode.Exists() && outputNode.IsArray() { + applyResponsesOutputItemsUsage(&detail, outputNode) + } return detail } +func applyOpenAIChatToolCallsUsage(detail *usage.Detail, root gjson.Result) { + if detail == nil || !root.Exists() || !root.IsObject() { + return + } + var requested int64 + for _, choice := range root.Get("choices").Array() { + if tc := choice.Get("message.tool_calls"); tc.IsArray() { + requested += int64(len(tc.Array())) + } else if tc := choice.Get("delta.tool_calls"); tc.IsArray() { + requested += int64(len(tc.Array())) + } + } + if requested > 0 { + if detail.ToolUsageDetails == nil { + detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if detail.ToolUsageDetails.ToolCallsRequested == nil { + detail.ToolUsageDetails.ToolCallsRequested = usage.Int64Ptr(requested) + } + } +} + +// parseOpenAIStyleUsageWithCache keeps structured-only usage and cache metadata +// consistent across Chat, Responses, and streaming entry points. +func parseOpenAIStyleUsageWithCache(root, usageNode gjson.Result) (usage.Detail, bool) { + var detail usage.Detail + hasTokens := hasOpenAIStyleUsageTokenFields(usageNode) + if hasTokens { + detail = parseOpenAIStyleUsageNode(usageNode) + } else { + populateOpenAIStyleStructuredFields(&detail, usageNode) + } + applyResponsesCacheRetention(&detail, root, usageNode) + return detail, hasTokens || hasStructuredUsageDetail(detail) +} + +func hasStructuredUsageDetail(d usage.Detail) bool { + return d.CacheDetails != nil || d.ToolUsageDetails != nil || d.CostDetails != nil || + d.WebSearchRequests != nil || d.ImageCount != nil || d.ImageDetails != nil || + d.RequestedImageDetails != nil || d.MediaCounts != nil || d.Iterations != nil || + d.CacheCreation5mTokens != nil || d.CacheCreation30mTokens != nil || + d.CacheCreation1hTokens != nil || d.CacheCreation12hTokens != nil || + d.TextInputTokens != nil || d.ImageInputTokens != nil || d.AudioInputTokens != nil || + d.VideoInputTokens != nil || d.FileInputTokens != nil || d.CachedTextTokens != nil || + d.TextOutputTokens != nil || d.ImageOutputTokens != nil || d.AudioOutputTokens != nil || + d.AcceptedPredictionTokens != nil || d.RejectedPredictionTokens != nil || + d.Speed != "" || d.InferenceGeo != "" +} + func hasOpenAIStyleUsageTokenFields(usageNode gjson.Result) bool { if !usageNode.Exists() || !usageNode.IsObject() { return false @@ -1100,50 +1808,177 @@ func hasOpenAIStyleUsageBucketFields(usageNode gjson.Result) bool { } func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { - inputNode := usageNode.Get("prompt_tokens") - if !inputNode.Exists() { - inputNode = usageNode.Get("input_tokens") - } - outputNode := usageNode.Get("completion_tokens") - if !outputNode.Exists() { - outputNode = usageNode.Get("output_tokens") - } - detail := usage.Detail{ - InputTokens: inputNode.Int(), - OutputTokens: outputNode.Int(), - TotalTokens: usageNode.Get("total_tokens").Int(), - } - cached := usageNode.Get("prompt_tokens_details.cached_tokens") - if !cached.Exists() { - cached = usageNode.Get("input_tokens_details.cached_tokens") - } - if cached.Exists() { - detail.CachedTokens = cached.Int() - detail.CacheReadTokens = cached.Int() + if !usageNode.Exists() || !usageNode.IsObject() { + return usage.Detail{} } - cacheCreation := firstExistingUsageNode( - usageNode, - "input_tokens_details.cache_creation_tokens", - "input_tokens_details.cache_write_tokens", - "prompt_tokens_details.cache_creation_tokens", - "prompt_tokens_details.cache_write_tokens", + var ( + detail usage.Detail + hasInput bool + hasOutput bool + hasBucketFields bool + inputDetails gjson.Result + promptDetails gjson.Result + outputDetails gjson.Result + compDetails gjson.Result + flatReasoning gjson.Result ) - if cacheCreation.Exists() { - detail.CacheCreationTokens = cacheCreation.Int() - } - reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens") - if !reasoning.Exists() { - reasoning = usageNode.Get("output_tokens_details.reasoning_tokens") + usageNode.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "prompt_tokens": + detail.InputTokens = val.Int() + hasInput = true + hasBucketFields = true + case "input_tokens": + if !hasInput { + detail.InputTokens = val.Int() + hasInput = true + } + hasBucketFields = true + case "completion_tokens": + detail.OutputTokens = val.Int() + hasOutput = true + hasBucketFields = true + case "output_tokens": + if !hasOutput { + detail.OutputTokens = val.Int() + hasOutput = true + } + hasBucketFields = true + case "total_tokens": + detail.TotalTokens = val.Int() + case "reasoning_tokens": + flatReasoning = val + hasBucketFields = true + case "prompt_tokens_details": + promptDetails = val + case "input_tokens_details": + inputDetails = val + case "completion_tokens_details": + compDetails = val + case "output_tokens_details": + outputDetails = val + } + return true + }) + + var hasCached, hasCacheCreation, hasReasoning bool + scanInDetails := func(node gjson.Result) { + if !node.Exists() || !node.IsObject() { + return + } + node.ForEach(func(k, v gjson.Result) bool { + switch k.String() { + case "cached_tokens": + if !hasCached { + detail.CachedTokens = v.Int() + detail.CacheReadTokens = v.Int() + hasCached = true + } + hasBucketFields = true + case "cache_creation_tokens", "cache_write_tokens": + if !hasCacheCreation { + detail.CacheCreationTokens = v.Int() + hasCacheCreation = true + } + hasBucketFields = true + case "text_tokens": + if v.Type == gjson.Number && detail.TextInputTokens == nil { + detail.TextInputTokens = usage.Int64Ptr(v.Int()) + } + case "image_tokens", "prompt_image_tokens": + if v.Type == gjson.Number && detail.ImageInputTokens == nil { + detail.ImageInputTokens = usage.Int64Ptr(v.Int()) + } + case "audio_tokens": + if v.Type == gjson.Number && detail.AudioInputTokens == nil { + detail.AudioInputTokens = usage.Int64Ptr(v.Int()) + } + case "video_tokens": + if v.Type == gjson.Number && detail.VideoInputTokens == nil { + detail.VideoInputTokens = usage.Int64Ptr(v.Int()) + } + case "file_tokens": + if v.Type == gjson.Number && detail.FileInputTokens == nil { + detail.FileInputTokens = usage.Int64Ptr(v.Int()) + } + case "cached_prompt_text_tokens": + if v.Type == gjson.Number && detail.CachedTextTokens == nil { + detail.CachedTextTokens = usage.Int64Ptr(v.Int()) + } + case "cached_tokens_details": + if detail.CachedTextTokens == nil && v.IsObject() { + if tv := v.Get("text_tokens"); tv.Type == gjson.Number { + detail.CachedTextTokens = usage.Int64Ptr(tv.Int()) + } + } + } + return true + }) + } + // Match original precedence: prompt_tokens_details first for cached_tokens, input_tokens_details first for structured fields & cache_creation_tokens. + if promptDetails.Exists() { + if c := promptDetails.Get("cached_tokens"); c.Exists() { + detail.CachedTokens = c.Int() + detail.CacheReadTokens = c.Int() + hasCached = true + hasBucketFields = true + } } - if !reasoning.Exists() { - // Flat dialect (e.g. SGLang): usage.reasoning_tokens at top level. - reasoning = usageNode.Get("reasoning_tokens") + scanInDetails(inputDetails) + scanInDetails(promptDetails) + + scanOutDetails := func(node gjson.Result) { + if !node.Exists() || !node.IsObject() { + return + } + node.ForEach(func(k, v gjson.Result) bool { + switch k.String() { + case "reasoning_tokens": + if !hasReasoning { + detail.ReasoningTokens = v.Int() + hasReasoning = true + } + hasBucketFields = true + case "text_tokens": + if v.Type == gjson.Number && detail.TextOutputTokens == nil { + detail.TextOutputTokens = usage.Int64Ptr(v.Int()) + } + case "image_tokens": + if v.Type == gjson.Number && detail.ImageOutputTokens == nil { + detail.ImageOutputTokens = usage.Int64Ptr(v.Int()) + } + case "audio_tokens": + if v.Type == gjson.Number && detail.AudioOutputTokens == nil { + detail.AudioOutputTokens = usage.Int64Ptr(v.Int()) + } + case "accepted_prediction_tokens": + if v.Type == gjson.Number && detail.AcceptedPredictionTokens == nil { + detail.AcceptedPredictionTokens = usage.Int64Ptr(v.Int()) + } + case "rejected_prediction_tokens": + if v.Type == gjson.Number && detail.RejectedPredictionTokens == nil { + detail.RejectedPredictionTokens = usage.Int64Ptr(v.Int()) + } + } + return true + }) + } + if compDetails.Exists() { + if r := compDetails.Get("reasoning_tokens"); r.Exists() { + detail.ReasoningTokens = r.Int() + hasReasoning = true + hasBucketFields = true + } } - if reasoning.Exists() { - detail.ReasoningTokens = reasoning.Int() + scanOutDetails(outputDetails) + scanOutDetails(compDetails) + if !hasReasoning && flatReasoning.Exists() { + detail.ReasoningTokens = flatReasoning.Int() + hasReasoning = true } - if hasOpenAIStyleUsageBucketFields(usageNode) { - if inputNode.Exists() && outputNode.Exists() { + + if hasBucketFields { + if hasInput && hasOutput { detail.TokenBreakdown = usage.NewSubsetTokenBreakdown( detail.InputTokens, detail.CacheReadTokens, @@ -1155,12 +1990,12 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { } else { cacheReadTokens := detail.CacheReadTokens cacheCreationTokens := detail.CacheCreationTokens - if !inputNode.Exists() { + if !hasInput { cacheReadTokens = 0 cacheCreationTokens = 0 } reasoningTokens := detail.ReasoningTokens - if !outputNode.Exists() { + if !hasOutput { reasoningTokens = 0 } detail.TokenBreakdown = usage.NewPartialSubsetTokenBreakdown( @@ -1172,122 +2007,1225 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { detail.TotalTokens, ) } - } else { - detail.TokenBreakdown = usage.NewUnclassifiedTokenBreakdown(detail.TotalTokens) + } else { + detail.TokenBreakdown = usage.NewUnclassifiedTokenBreakdown(detail.TotalTokens) + } + if detail.TotalTokens == 0 { + detail.TotalTokens = detail.TokenBreakdown.TotalTokens + } + populateOpenAIStyleStructuredFields(&detail, usageNode) + return detail +} + +func populateOpenAIStyleStructuredFields(detail *usage.Detail, usageNode gjson.Result) { + if detail == nil || !usageNode.Exists() || !usageNode.IsObject() { + return + } + var ( + toolDetails *usage.ToolUsageDetails + costNode gjson.Result + ticksNode gjson.Result + byokNode gjson.Result + costDetails gjson.Result + hasCost bool + ) + usageNode.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "service_tier": + if detail.UpstreamServiceTier == "" && val.Type == gjson.String { + detail.UpstreamServiceTier = cloneTrimmedString(val.String()) + } + case "attribution": + if detail.UpstreamServiceTier == "" && val.IsObject() { + if t := val.Get("tier"); t.Type == gjson.String { + detail.UpstreamServiceTier = cloneTrimmedString(t.String()) + } + } + case "server_side_tool_usage_details", "server_side_tool_usage", "server_tool_use_details", "server_tool_use": + if parsed := parseOpenAIToolUsageDetailsNode(val); parsed != nil { + toolDetails = mergeToolUsageDetails(toolDetails, parsed) + } + case "num_sources_used": + if val.Type == gjson.Number { + if toolDetails == nil { + toolDetails = &usage.ToolUsageDetails{} + } + toolDetails.NumSourcesUsed = usage.Int64Ptr(val.Int()) + } + case "num_server_side_tools_used": + if val.Type == gjson.Number { + if toolDetails == nil { + toolDetails = &usage.ToolUsageDetails{} + } + toolDetails.NumServerSideToolsUsed = usage.Int64Ptr(val.Int()) + } + case "cost": + costNode = val + hasCost = true + case "cost_in_usd_ticks": + ticksNode = val + hasCost = true + case "is_byok": + byokNode = val + hasCost = true + case "cost_details": + costDetails = val + hasCost = true + } + return true + }) + detail.ToolUsageDetails = toolDetails + if toolDetails != nil { + if toolDetails.WebSearchCalls != nil { + detail.WebSearchRequests = usage.CloneInt64Ptr(toolDetails.WebSearchCalls) + } + if detail.ImageCount == nil && toolDetails.ImageGenerationCalls != nil && *toolDetails.ImageGenerationCalls > 0 { + detail.ImageCount = usage.CloneInt64Ptr(toolDetails.ImageGenerationCalls) + } + } + if hasCost { + _ = costNode + _ = ticksNode + _ = byokNode + _ = costDetails + detail.CostDetails = parseClaudeCostNode(usageNode) + } +} + +func parseCollectedOutputItems(itemsByIndex map[int64][]byte, fallback [][]byte) []gjson.Result { + if len(itemsByIndex) == 0 && len(fallback) == 0 { + return nil + } + items := make([]gjson.Result, 0, len(itemsByIndex)+len(fallback)) + if len(itemsByIndex) > 0 { + indexes := make([]int64, 0, len(itemsByIndex)) + for idx := range itemsByIndex { + indexes = append(indexes, idx) + } + for i := 0; i < len(indexes); i++ { + for j := i + 1; j < len(indexes); j++ { + if indexes[j] < indexes[i] { + indexes[i], indexes[j] = indexes[j], indexes[i] + } + } + } + for _, idx := range indexes { + if raw := itemsByIndex[idx]; len(raw) > 0 && gjson.ValidBytes(raw) { + items = append(items, gjson.ParseBytes(raw)) + } + } + } + for _, raw := range fallback { + if len(raw) > 0 && gjson.ValidBytes(raw) { + items = append(items, gjson.ParseBytes(raw)) + } + } + return items +} + +func applyResponsesCollectedOutputItemsUsage(detail *usage.Detail, itemsByIndex map[int64][]byte, fallback [][]byte) { + if detail == nil { + return + } + items := parseCollectedOutputItems(itemsByIndex, fallback) + if len(items) == 0 { + return + } + applyResponsesOutputItemSliceUsage(detail, items) +} + +func applyResponsesOutputItemsUsage(detail *usage.Detail, outputNode gjson.Result) { + if detail == nil || !outputNode.Exists() || !outputNode.IsArray() { + return + } + applyResponsesOutputItemSliceUsage(detail, outputNode.Array()) +} + +func applyResponsesOutputItemSliceUsage(detail *usage.Detail, items []gjson.Result) { + if detail == nil || len(items) == 0 { + return + } + seenIDs := make(map[string]struct{}) + var ( + webSearchCalls int64 + webSearchQueries int64 + webSearchSources int64 + imageSearchCalls int64 + xSearchCalls int64 + codeExecCalls int64 + fileSearchCalls int64 + documentSearchCalls int64 + mcpCalls int64 + imageGenCalls int64 + toolCallsRequested int64 + firstImageMeta *usage.ImageDetails + ) + for _, item := range items { + if !item.IsObject() { + continue + } + if status := strings.TrimSpace(item.Get("status").String()); strings.EqualFold(status, "failed") || strings.EqualFold(status, "in_progress") { + continue + } + if id := strings.TrimSpace(firstExistingUsageNode(item, "id", "call_id").String()); id != "" { + if _, exists := seenIDs[id]; exists { + continue + } + seenIDs[id] = struct{}{} + } + switch strings.TrimSpace(item.Get("type").String()) { + case "web_search_call": + webSearchCalls++ + if strings.HasPrefix(strings.TrimSpace(item.Get("id").String()), "is_") { + imageSearchCalls++ + } + action := item.Get("action") + if queries := action.Get("queries"); queries.IsArray() && len(queries.Array()) > 0 { + webSearchQueries += int64(len(queries.Array())) + } else if q := strings.TrimSpace(action.Get("query").String()); q != "" || strings.EqualFold(strings.TrimSpace(action.Get("type").String()), "search") { + webSearchQueries++ + } + if sources := action.Get("sources"); sources.IsArray() { + webSearchSources += int64(len(sources.Array())) + } + case "file_search_call": + fileSearchCalls++ + case "code_interpreter_call": + codeExecCalls++ + case "mcp_call": + mcpCalls++ + case "image_generation_call": + imageGenCalls++ + if firstImageMeta == nil { + meta := &usage.ImageDetails{ + Size: cloneTrimmedString(item.Get("size").String()), + Resolution: cloneTrimmedString(item.Get("resolution").String()), + AspectRatio: cloneTrimmedString(item.Get("aspect_ratio").String()), + Quality: cloneTrimmedString(item.Get("quality").String()), + Background: cloneTrimmedString(item.Get("background").String()), + OutputFormat: cloneTrimmedString(item.Get("output_format").String()), + MimeType: cloneTrimmedString(firstExistingUsageNode(item, "mime_type", "mimeType").String()), + InputFidelity: cloneTrimmedString(item.Get("input_fidelity").String()), + } + if !isEmptyImageDetails(meta) { + firstImageMeta = meta + } + } + case "function_call", "local_shell_call", "shell_call", "computer_call", "apply_patch_call": + toolCallsRequested++ + case "custom_tool_call": + callID := strings.TrimSpace(item.Get("call_id").String()) + switch { + case strings.HasPrefix(callID, "xs_"): + xSearchCalls++ + case strings.HasPrefix(callID, "as_"): + documentSearchCalls++ + default: + toolCallsRequested++ + } + } + } + if webSearchCalls == 0 && xSearchCalls == 0 && codeExecCalls == 0 && fileSearchCalls == 0 && documentSearchCalls == 0 && mcpCalls == 0 && imageGenCalls == 0 && toolCallsRequested == 0 { + return + } + if detail.ToolUsageDetails == nil { + detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if webSearchCalls > 0 && (detail.ToolUsageDetails.WebSearchCalls == nil || *detail.ToolUsageDetails.WebSearchCalls == 0) { + detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(webSearchCalls) + if detail.WebSearchRequests == nil || *detail.WebSearchRequests == 0 { + detail.WebSearchRequests = usage.Int64Ptr(webSearchCalls) + } + } else if imageSearchCalls > 0 && detail.ToolUsageDetails.WebSearchCalls != nil && *detail.ToolUsageDetails.WebSearchCalls == 0 { + // xAI Image Search (id prefix "is_") is billed at the standard Web Search rate + // even when server_side_tool_usage_details.web_search_calls reports 0. + detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(imageSearchCalls) + detail.WebSearchRequests = usage.Int64Ptr(imageSearchCalls) + } + if imageSearchCalls > 0 { + if detail.ToolUsageDetails.Extra == nil { + detail.ToolUsageDetails.Extra = make(map[string]int64) + } + detail.ToolUsageDetails.Extra["image_search_calls"] = imageSearchCalls + } + if xSearchCalls > 0 && detail.ToolUsageDetails.XSearchCalls == nil { + detail.ToolUsageDetails.XSearchCalls = usage.Int64Ptr(xSearchCalls) + } + if codeExecCalls > 0 && detail.ToolUsageDetails.CodeExecutionCalls == nil { + detail.ToolUsageDetails.CodeExecutionCalls = usage.Int64Ptr(codeExecCalls) + } + if fileSearchCalls > 0 && detail.ToolUsageDetails.FileSearchCalls == nil { + detail.ToolUsageDetails.FileSearchCalls = usage.Int64Ptr(fileSearchCalls) + } + if documentSearchCalls > 0 && detail.ToolUsageDetails.DocumentSearchCalls == nil { + detail.ToolUsageDetails.DocumentSearchCalls = usage.Int64Ptr(documentSearchCalls) + } + if mcpCalls > 0 && detail.ToolUsageDetails.MCPCalls == nil { + detail.ToolUsageDetails.MCPCalls = usage.Int64Ptr(mcpCalls) + } + if webSearchQueries > 0 && detail.ToolUsageDetails.WebSearchQueries == nil { + detail.ToolUsageDetails.WebSearchQueries = usage.Int64Ptr(webSearchQueries) + } + if webSearchSources > 0 && detail.ToolUsageDetails.NumSourcesUsed == nil { + detail.ToolUsageDetails.NumSourcesUsed = usage.Int64Ptr(webSearchSources) + } + if imageGenCalls > 0 { + if detail.ToolUsageDetails.ImageGenerationCalls == nil { + detail.ToolUsageDetails.ImageGenerationCalls = usage.Int64Ptr(imageGenCalls) + } + if detail.ImageCount == nil { + detail.ImageCount = usage.Int64Ptr(imageGenCalls) + } + if detail.ImageDetails == nil && firstImageMeta != nil { + detail.ImageDetails = firstImageMeta + } + } + if toolCallsRequested > 0 && detail.ToolUsageDetails.ToolCallsRequested == nil { + detail.ToolUsageDetails.ToolCallsRequested = usage.Int64Ptr(toolCallsRequested) + } +} + +func isEmptyImageDetails(d *usage.ImageDetails) bool { + if d == nil { + return true + } + return d.Size == "" && d.Resolution == "" && d.ImageSize == "" && d.AspectRatio == "" && d.Quality == "" && d.Background == "" && d.OutputFormat == "" && d.ResponseFormat == "" && d.MimeType == "" && d.InputFidelity == "" +} + +// ParseRequestImageDetails extracts requested image/video parameters strictly from a request payload, +// preserving each upstream protocol's native field names. +func ParseRequestImageDetails(requestPayload []byte, contentType ...string) *usage.ImageDetails { + if len(contentType) > 0 { + mediaType, params, errParse := mime.ParseMediaType(contentType[0]) + if errParse == nil && mediaType == "multipart/form-data" { + return parseMultipartImageDetails(requestPayload, params["boundary"]) + } + } + if len(requestPayload) == 0 || !gjson.ValidBytes(requestPayload) { + return nil + } + root := gjson.ParseBytes(requestPayload) + pick := func(paths ...string) string { + for _, p := range paths { + if v := root.Get(p); v.Type == gjson.String { + if s := strings.TrimSpace(v.String()); s != "" { + return s + } + } + } + return "" + } + out := &usage.ImageDetails{ + Size: pick("size"), + Resolution: pick("resolution", "image_config.resolution"), + ImageSize: pick("image_size", "imageSize", "image_config.image_size", "image_config.imageSize", "generationConfig.imageConfig.imageSize", "generation_config.image_config.image_size", "request.generationConfig.imageConfig.imageSize", "request.generation_config.image_config.image_size", "parameters.sampleImageSize"), + AspectRatio: pick("aspect_ratio", "aspectRatio", "image_config.aspect_ratio", "image_config.aspectRatio", "generationConfig.imageConfig.aspectRatio", "generation_config.image_config.aspect_ratio", "request.generationConfig.imageConfig.aspectRatio", "request.generation_config.image_config.aspect_ratio", "parameters.aspectRatio"), + Quality: pick("quality"), + Background: pick("background"), + OutputFormat: pick("output_format"), + ResponseFormat: pick("response_format"), + MimeType: pick("mime_type", "mimeType", "parameters.outputOptions.mimeType"), + InputFidelity: pick("input_fidelity"), + } + if isEmptyImageDetails(out) { + for _, tool := range root.Get("tools").Array() { + if strings.TrimSpace(tool.Get("type").String()) != "image_generation" { + continue + } + out = &usage.ImageDetails{ + Size: strings.TrimSpace(tool.Get("size").String()), + Resolution: strings.TrimSpace(tool.Get("resolution").String()), + ImageSize: strings.TrimSpace(firstExistingUsageNode(tool, "image_size", "imageSize").String()), + AspectRatio: strings.TrimSpace(firstExistingUsageNode(tool, "aspect_ratio", "aspectRatio").String()), + Quality: strings.TrimSpace(tool.Get("quality").String()), + Background: strings.TrimSpace(tool.Get("background").String()), + OutputFormat: strings.TrimSpace(tool.Get("output_format").String()), + MimeType: strings.TrimSpace(firstExistingUsageNode(tool, "mime_type", "mimeType").String()), + InputFidelity: strings.TrimSpace(tool.Get("input_fidelity").String()), + } + break + } + } + if isEmptyImageDetails(out) { + return nil + } + return out +} + +// parseMultipartImageDetails reads only small scalar attributes. File parts and +// unrelated fields are drained without allocating another copy or temporary files. +func parseMultipartImageDetails(payload []byte, boundary string) *usage.ImageDetails { + if boundary == "" { + return nil + } + out := &usage.ImageDetails{} + fields := map[string]*string{ + "size": &out.Size, "resolution": &out.Resolution, "image_size": &out.ImageSize, + "aspect_ratio": &out.AspectRatio, "quality": &out.Quality, + "background": &out.Background, "output_format": &out.OutputFormat, + "response_format": &out.ResponseFormat, "mime_type": &out.MimeType, + "input_fidelity": &out.InputFidelity, + } + reader := multipart.NewReader(bytes.NewReader(payload), boundary) + for { + part, errNext := reader.NextPart() + if errNext == io.EOF { + break + } + if errNext != nil { + log.Debugf("usage: read multipart image attributes: %v", errNext) + return nil + } + dst := fields[part.FormName()] + var value []byte + var errRead error + const maxAttributeBytes = 4096 + if dst != nil && part.FileName() == "" { + value, errRead = io.ReadAll(io.LimitReader(part, maxAttributeBytes+1)) + } + errClose := part.Close() + if errRead != nil || errClose != nil { + log.Debugf("usage: read multipart image attribute: read=%v close=%v", errRead, errClose) + return nil + } + if dst != nil && len(value) <= maxAttributeBytes && len(value) > 0 { + *dst = strings.TrimSpace(string(value)) + } + } + if isEmptyImageDetails(out) { + return nil + } + return out +} + +// ParseOpenAIImageUsage extracts token usage, generated image count, and separate +// request vs upstream response image parameters from an OpenAI or xAI image request/response. +func ParseOpenAIImageUsage(requestPayload, responsePayload []byte, contentType ...string) usage.Detail { + var detail usage.Detail + detail.RequestedImageDetails = ParseRequestImageDetails(requestPayload, contentType...) + if len(responsePayload) == 0 || !gjson.ValidBytes(responsePayload) { + return detail + } + root := gjson.Parse(unsafeBytesToString(responsePayload)) + var ( + usageNode gjson.Result + serviceTier string + imageCount int64 + imgDetails usage.ImageDetails + firstItemMeta usage.ImageDetails + ) + root.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "usage": + usageNode = val + case "service_tier": + if val.Type == gjson.String { + serviceTier = cloneTrimmedString(val.String()) + } + case "data": + if val.IsArray() { + arr := val.Array() + imageCount = int64(len(arr)) + if len(arr) > 0 && arr[0].IsObject() { + arr[0].ForEach(func(dk, dv gjson.Result) bool { + if dv.Type != gjson.String { + return true + } + switch dk.String() { + case "b64_json", "url", "revised_prompt": + return true + case "size": + firstItemMeta.Size = cloneTrimmedString(dv.String()) + case "resolution": + firstItemMeta.Resolution = cloneTrimmedString(dv.String()) + case "image_size", "imageSize": + if firstItemMeta.ImageSize == "" { + firstItemMeta.ImageSize = cloneTrimmedString(dv.String()) + } + case "aspect_ratio", "aspectRatio": + if firstItemMeta.AspectRatio == "" { + firstItemMeta.AspectRatio = cloneTrimmedString(dv.String()) + } + case "quality": + firstItemMeta.Quality = cloneTrimmedString(dv.String()) + case "background": + firstItemMeta.Background = cloneTrimmedString(dv.String()) + case "output_format": + firstItemMeta.OutputFormat = cloneTrimmedString(dv.String()) + case "mime_type", "mimeType": + if firstItemMeta.MimeType == "" { + firstItemMeta.MimeType = cloneTrimmedString(dv.String()) + } + case "input_fidelity": + firstItemMeta.InputFidelity = cloneTrimmedString(dv.String()) + } + return true + }) + } + } + case "size": + if val.Type == gjson.String { + imgDetails.Size = cloneTrimmedString(val.String()) + } + case "resolution": + if val.Type == gjson.String { + imgDetails.Resolution = cloneTrimmedString(val.String()) + } + case "image_size", "imageSize": + if val.Type == gjson.String && imgDetails.ImageSize == "" { + imgDetails.ImageSize = cloneTrimmedString(val.String()) + } + case "aspect_ratio", "aspectRatio": + if val.Type == gjson.String && imgDetails.AspectRatio == "" { + imgDetails.AspectRatio = cloneTrimmedString(val.String()) + } + case "quality": + if val.Type == gjson.String { + imgDetails.Quality = cloneTrimmedString(val.String()) + } + case "background": + if val.Type == gjson.String { + imgDetails.Background = cloneTrimmedString(val.String()) + } + case "output_format": + if val.Type == gjson.String { + imgDetails.OutputFormat = cloneTrimmedString(val.String()) + } + case "mime_type", "mimeType": + if val.Type == gjson.String && imgDetails.MimeType == "" { + imgDetails.MimeType = cloneTrimmedString(val.String()) + } + case "input_fidelity": + if val.Type == gjson.String { + imgDetails.InputFidelity = cloneTrimmedString(val.String()) + } + } + return true + }) + if usageNode.Exists() { + if hasOpenAIStyleUsageTokenFields(usageNode) { + detail = parseOpenAIStyleUsageNode(usageNode) + } else { + populateOpenAIStyleStructuredFields(&detail, usageNode) + } + detail.RequestedImageDetails = ParseRequestImageDetails(requestPayload, contentType...) + } + if serviceTier != "" { + detail.ResponseServiceTier = serviceTier + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = serviceTier + } + } + if imageCount > 0 { + detail.ImageCount = usage.Int64Ptr(imageCount) + } + if imgDetails.Size == "" { + imgDetails.Size = firstItemMeta.Size + } + if imgDetails.Resolution == "" { + imgDetails.Resolution = firstItemMeta.Resolution + } + if imgDetails.ImageSize == "" { + imgDetails.ImageSize = firstItemMeta.ImageSize + } + if imgDetails.AspectRatio == "" { + imgDetails.AspectRatio = firstItemMeta.AspectRatio + } + if imgDetails.Quality == "" { + imgDetails.Quality = firstItemMeta.Quality + } + if imgDetails.Background == "" { + imgDetails.Background = firstItemMeta.Background + } + if imgDetails.OutputFormat == "" { + imgDetails.OutputFormat = firstItemMeta.OutputFormat + } + if imgDetails.MimeType == "" { + imgDetails.MimeType = firstItemMeta.MimeType + } + if imgDetails.InputFidelity == "" { + imgDetails.InputFidelity = firstItemMeta.InputFidelity + } + if !isEmptyImageDetails(&imgDetails) { + detail.ImageDetails = &imgDetails + } + return detail +} + +// ParseXAIImageUsage extracts token usage, generated image count, and image parameters +// from an xAI image generation/edit request and response. +func ParseXAIImageUsage(requestPayload, responsePayload []byte) usage.Detail { + return ParseOpenAIImageUsage(requestPayload, responsePayload) +} + +// ParseXAIVideoUsage extracts token usage and video parameters from an xAI video response. +func ParseXAIVideoUsage(requestPayload, responsePayload []byte) usage.Detail { + detail := ParseOpenAIUsage(responsePayload) + root := gjson.ParseBytes(responsePayload) + reqRoot := gjson.ParseBytes(requestPayload) + // Only attach video media attributes when a completed video output or usage is present, + // not on an initial async job creation response that only returns {"request_id":"..."}. + hasCompletedVideo := root.Get("usage").Exists() || root.Get("video").Exists() || root.Get("data").Exists() || root.Get("url").Exists() || strings.EqualFold(strings.TrimSpace(root.Get("status").String()), "completed") + if !hasCompletedVideo { + return detail + } + pickRespString := func(paths ...string) string { + for _, p := range paths { + if v := root.Get(p); v.Type == gjson.String { + if s := strings.TrimSpace(v.String()); s != "" { + return s + } + } + } + return "" + } + imgDetails := &usage.ImageDetails{ + Resolution: pickRespString("resolution", "video.resolution", "data.0.resolution"), + AspectRatio: pickRespString("aspect_ratio", "video.aspect_ratio", "data.0.aspect_ratio"), + } + if !isEmptyImageDetails(imgDetails) { + detail.ImageDetails = imgDetails + } + detail.RequestedImageDetails = ParseRequestImageDetails(requestPayload) + if dur := firstExistingUsageNode(root, "duration", "video.duration"); dur.Type == gjson.Number && dur.Float() > 0 { + detail.MediaCounts = &usage.MediaCounts{ + CompletionMedia: usage.Int64Ptr(1), + AudioSeconds: usage.Float64Ptr(dur.Float()), + } + } else if dur := reqRoot.Get("duration"); dur.Type == gjson.Number && dur.Float() > 0 { + detail.MediaCounts = &usage.MediaCounts{ + CompletionMedia: usage.Int64Ptr(1), + AudioSeconds: usage.Float64Ptr(dur.Float()), + } + } + return detail +} + +func parseOpenAIToolUsageDetailsNode(node gjson.Result) *usage.ToolUsageDetails { + if !node.Exists() || !node.IsObject() { + return nil + } + out := &usage.ToolUsageDetails{} + hasAny := false + node.ForEach(func(key, value gjson.Result) bool { + if value.Type != gjson.Number { + return true + } + k := strings.TrimSpace(key.String()) + if k == "" { + return true + } + v := value.Int() + switch k { + case "web_search_calls", "web_search_requests", "SERVER_SIDE_TOOL_WEB_SEARCH": + out.WebSearchCalls = usage.Int64Ptr(v) + case "web_search_queries": + out.WebSearchQueries = usage.Int64Ptr(v) + case "web_fetch_calls", "web_fetch_requests": + out.WebFetchCalls = usage.Int64Ptr(v) + case "num_search_results": + out.NumSearchResults = usage.Int64Ptr(v) + case "num_sources_used": + out.NumSourcesUsed = usage.Int64Ptr(v) + case "num_server_side_tools_used": + out.NumServerSideToolsUsed = usage.Int64Ptr(v) + case "tool_calls_requested": + out.ToolCallsRequested = usage.Int64Ptr(v) + case "tool_calls_executed": + out.ToolCallsExecuted = usage.Int64Ptr(v) + case "x_search_calls", "x_search_requests", "SERVER_SIDE_TOOL_X_SEARCH": + out.XSearchCalls = usage.Int64Ptr(v) + case "x_posts_fetched": + out.XPostsFetched = usage.Int64Ptr(v) + case "x_users_fetched": + out.XUsersFetched = usage.Int64Ptr(v) + case "code_interpreter_calls", "code_execution_calls", "code_execution_requests", "SERVER_SIDE_TOOL_CODE_EXECUTION": + out.CodeExecutionCalls = usage.Int64Ptr(v) + case "file_search_calls", "file_search_requests", "SERVER_SIDE_TOOL_COLLECTIONS_SEARCH": + out.FileSearchCalls = usage.Int64Ptr(v) + case "mcp_calls", "mcp_requests", "SERVER_SIDE_TOOL_MCP": + out.MCPCalls = usage.Int64Ptr(v) + case "document_search_calls", "document_search_requests", "attachment_search_calls", "attachment_search_requests", "SERVER_SIDE_TOOL_ATTACHMENT_SEARCH", "SERVER_SIDE_TOOL_DOCUMENT_SEARCH": + out.DocumentSearchCalls = usage.Int64Ptr(v) + case "image_generation_calls", "image_generation_requests": + out.ImageGenerationCalls = usage.Int64Ptr(v) + default: + if out.Extra == nil { + out.Extra = make(map[string]int64) + } + out.Extra[strings.Clone(k)] = v + } + hasAny = true + return true + }) + if !hasAny { + return nil + } + return out +} + +func ParseOpenAIStreamUsage(line []byte) (usage.Detail, bool) { + payload := jsonPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return usage.Detail{}, false + } + root := gjson.Parse(unsafeBytesToString(payload)) + detail, ok := parseOpenAIStyleUsageWithCache(root, root.Get("usage")) + detail.ResponseServiceTier = extractResponseServiceTierFromValidJSON(payload) + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = detail.ResponseServiceTier + } + return detail, ok || detail.ResponseServiceTier != "" +} + +func ParseClaudeUsage(data []byte) usage.Detail { + root := gjson.ParseBytes(data) + usageNode := root.Get("usage") + if !usageNode.Exists() { + return usage.Detail{} + } + detail := parseClaudeUsageNode(usageNode) + applyClaudeContentBlocksUsage(&detail, root.Get("content")) + return detail +} + +func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) { + payload := jsonPayload(line) + if len(payload) == 0 || !bytes.Contains(payload, openAIStreamUsageMarker) || !gjson.ValidBytes(payload) { + return usage.Detail{}, false + } + root := gjson.Parse(unsafeBytesToString(payload)) + usageNode := root.Get("usage") + if !usageNode.Exists() { + usageNode = root.Get("message.usage") + } + if !usageNode.Exists() { + return usage.Detail{}, false + } + return parseClaudeUsageNode(usageNode), true +} + +func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail { + cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int() + cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int() + rawOutputTokens := usageNode.Get("output_tokens").Int() + // Anthropic reports thinking as a subset of output_tokens. Prefer the official + // nested field, then fall back to legacy aliases used by some gateways. + reasoningNode := firstExistingUsageNode( + usageNode, + "output_tokens_details.thinking_tokens", + "output_tokens_details.reasoning_tokens", + "thinking_tokens", + ) + reasoningTokens := reasoningNode.Int() + if reasoningTokens < 0 { + reasoningTokens = 0 + } + nonReasoningOutput := rawOutputTokens + if reasoningTokens > 0 && reasoningTokens <= rawOutputTokens { + nonReasoningOutput = rawOutputTokens - reasoningTokens + } else if reasoningTokens > rawOutputTokens { + // Keep Detail.OutputTokens authoritative for keeper subset checks and + // avoid inventing extra non-reasoning output when the upstream payload + // is inconsistent. + nonReasoningOutput = 0 + } + detail := usage.Detail{ + InputTokens: usageNode.Get("input_tokens").Int(), + OutputTokens: rawOutputTokens, + ReasoningTokens: reasoningTokens, + CachedTokens: cacheReadTokens, + CacheReadTokens: cacheReadTokens, + CacheCreationTokens: cacheCreationTokens, + } + // Anthropic reports the speed a response was served and billed at in + // usage.speed ("fast" or "standard"): fast mode costs more, and a + // fast request on a model without it runs, and is billed, at standard. + // Carried as the response tier so usage sinks can price it. + if speed := cloneTrimmedString(usageNode.Get("speed").String()); speed != "" { + detail.ResponseServiceTier = speed + } + if detail.CachedTokens == 0 { + detail.CachedTokens = detail.CacheCreationTokens + } + // raw output_tokens already includes thinking; cache fields are independent + // from input_tokens in the Messages API. + detail.TotalTokens = detail.InputTokens + rawOutputTokens + detail.CacheReadTokens + detail.CacheCreationTokens + detail.TokenBreakdown = usage.NewIndependentTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + nonReasoningOutput, + detail.ReasoningTokens, + detail.TotalTokens, + ) + populateClaudeUsageStructuredFields(&detail, usageNode) + return detail +} + +func populateClaudeUsageStructuredFields(detail *usage.Detail, usageNode gjson.Result) { + if detail == nil || !usageNode.Exists() || !usageNode.IsObject() { + return + } + if v := usageNode.Get("service_tier"); v.Type == gjson.String { + detail.UpstreamServiceTier = cloneTrimmedString(v.String()) + } + if v := usageNode.Get("speed"); v.Type == gjson.String { + detail.Speed = cloneTrimmedString(v.String()) + } + if v := usageNode.Get("inference_geo"); v.Type == gjson.String { + detail.InferenceGeo = cloneTrimmedString(v.String()) + } + iterations := parseClaudeIterationsNode(usageNode.Get("iterations")) + detail.Iterations = usage.IterationsPtr(iterations) + cacheDetails := parseClaudeCacheCreationNode(usageNode.Get("cache_creation")) + if cacheDetails == nil && len(iterations) > 0 && + (usageNode.Get("cache_creation_input_tokens").Type == gjson.Number || detail.CacheCreationTokens > 0) { + cacheDetails = deriveClaudeCacheCreationFromIterations(iterations, detail.CacheCreationTokens) + } + applyCacheDetailsToUsageDetail(detail, cacheDetails) + detail.ToolUsageDetails = parseClaudeServerToolUseNode(usageNode.Get("server_tool_use")) + if detail.ToolUsageDetails != nil && detail.ToolUsageDetails.WebSearchCalls != nil { + detail.WebSearchRequests = usage.CloneInt64Ptr(detail.ToolUsageDetails.WebSearchCalls) + } + detail.CostDetails = parseClaudeCostNode(usageNode) +} + +func applyClaudeContentBlocksUsage(detail *usage.Detail, contentNode gjson.Result) { + if detail == nil || !contentNode.Exists() || !contentNode.IsArray() { + return + } + seenToolUseIDs := make(map[string]struct{}) + seenToolResultIDs := make(map[string]struct{}) + var ( + webSearchQueries int64 + completedSearches int64 + searchResults int64 + hasCompleted bool + ) + for _, block := range contentNode.Array() { + if !block.IsObject() { + continue + } + switch block.Get("type").String() { + case "server_tool_use": + if strings.TrimSpace(block.Get("name").String()) != "web_search" { + continue + } + if id := strings.TrimSpace(block.Get("id").String()); id != "" { + if _, ok := seenToolUseIDs[id]; ok { + continue + } + seenToolUseIDs[id] = struct{}{} + } + webSearchQueries++ + case "web_search_tool_result": + content := block.Get("content") + if !content.Exists() || !content.IsArray() { + continue + } + if id := strings.TrimSpace(block.Get("tool_use_id").String()); id != "" { + if _, ok := seenToolResultIDs[id]; ok { + continue + } + seenToolResultIDs[id] = struct{}{} + } + hasCompleted = true + completedSearches++ + for _, item := range content.Array() { + if item.Get("type").String() == "web_search_result" { + searchResults++ + } + } + } + } + if webSearchQueries == 0 && !hasCompleted { + return + } + if detail.ToolUsageDetails == nil { + detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if webSearchQueries > 0 { + detail.ToolUsageDetails.WebSearchQueries = usage.Int64Ptr(webSearchQueries) + } + if hasCompleted { + detail.ToolUsageDetails.NumSearchResults = usage.Int64Ptr(searchResults) + if detail.ToolUsageDetails.WebSearchCalls == nil || (completedSearches > 0 && *detail.ToolUsageDetails.WebSearchCalls == 0) { + detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(completedSearches) + } + if detail.WebSearchRequests == nil || (completedSearches > 0 && *detail.WebSearchRequests == 0) { + detail.WebSearchRequests = usage.Int64Ptr(completedSearches) + } + } +} + +func applyCacheDetailsToUsageDetail(detail *usage.Detail, cacheDetails *usage.CacheDetails) { + if detail == nil { + return + } + detail.CacheDetails = cacheDetails.Clone() + if cacheDetails == nil { + detail.CacheCreation5mTokens = nil + detail.CacheCreation30mTokens = nil + detail.CacheCreation1hTokens = nil + detail.CacheCreation12hTokens = nil + return + } + detail.CacheCreation5mTokens = usage.CloneInt64Ptr(cacheDetails.Ephemeral5mInputTokens) + detail.CacheCreation30mTokens = usage.CloneInt64Ptr(cacheDetails.Ephemeral30mInputTokens) + detail.CacheCreation1hTokens = usage.CloneInt64Ptr(cacheDetails.Ephemeral1hInputTokens) + detail.CacheCreation12hTokens = usage.CloneInt64Ptr(cacheDetails.Ephemeral12hInputTokens) +} + +func parseClaudeCacheCreationNode(node gjson.Result) *usage.CacheDetails { + if !node.Exists() || !node.IsObject() { + return nil + } + out := &usage.CacheDetails{} + hasAny := false + node.ForEach(func(key, value gjson.Result) bool { + k := strings.TrimSpace(key.String()) + switch k { + case "ttl": + if value.Type == gjson.String { + out.TTL = cloneTrimmedString(value.String()) + hasAny = out.TTL != "" || hasAny + } + return true + case "mode": + if value.Type == gjson.String { + out.Mode = cloneTrimmedString(value.String()) + hasAny = out.Mode != "" || hasAny + } + return true + case "retention": + if value.Type == gjson.String { + out.Retention = cloneTrimmedString(value.String()) + hasAny = out.Retention != "" || hasAny + } + return true + } + if value.Type != gjson.Number { + return true + } + v := value.Int() + if v < 0 { + return true + } + ttlKey := "" + switch { + case k == "ephemeral_5m_input_tokens" || k == "5m": + out.Ephemeral5mInputTokens = usage.Int64Ptr(v) + ttlKey = "5m" + case k == "ephemeral_30m_input_tokens" || k == "30m": + out.Ephemeral30mInputTokens = usage.Int64Ptr(v) + ttlKey = "30m" + case k == "ephemeral_1h_input_tokens" || k == "1h": + out.Ephemeral1hInputTokens = usage.Int64Ptr(v) + ttlKey = "1h" + case k == "ephemeral_12h_input_tokens" || k == "12h": + out.Ephemeral12hInputTokens = usage.Int64Ptr(v) + ttlKey = "12h" + case strings.HasPrefix(k, "ephemeral_") && strings.HasSuffix(k, "_input_tokens"): + ttlKey = strings.TrimSuffix(strings.TrimPrefix(k, "ephemeral_"), "_input_tokens") + } + if ttlKey != "" { + if out.ByTTL == nil { + out.ByTTL = make(map[string]int64) + } + out.ByTTL[strings.Clone(ttlKey)] = v + hasAny = true + } + return true + }) + if !hasAny { + return nil + } + return out +} + +func parseClaudeServerToolUseNode(node gjson.Result) *usage.ToolUsageDetails { + if !node.Exists() || !node.IsObject() { + return nil + } + out := &usage.ToolUsageDetails{} + hasAny := false + node.ForEach(func(key, value gjson.Result) bool { + if value.Type != gjson.Number { + return true + } + k := strings.TrimSpace(key.String()) + if k == "" { + return true + } + v := value.Int() + switch k { + case "web_search_requests", "web_search_calls": + out.WebSearchCalls = usage.Int64Ptr(v) + case "web_fetch_requests", "web_fetch_calls": + out.WebFetchCalls = usage.Int64Ptr(v) + case "code_execution_requests", "code_execution_calls", "code_interpreter_calls": + out.CodeExecutionCalls = usage.Int64Ptr(v) + case "tool_calls_requested": + out.ToolCallsRequested = usage.Int64Ptr(v) + case "tool_calls_executed": + out.ToolCallsExecuted = usage.Int64Ptr(v) + case "file_search_requests", "file_search_calls": + out.FileSearchCalls = usage.Int64Ptr(v) + case "mcp_requests", "mcp_calls": + out.MCPCalls = usage.Int64Ptr(v) + default: + if out.Extra == nil { + out.Extra = make(map[string]int64) + } + out.Extra[strings.Clone(k)] = v + } + hasAny = true + return true + }) + if !hasAny { + return nil + } + return out +} + +func parseClaudeIterationsNode(node gjson.Result) []usage.IterationUsage { + if !node.Exists() || !node.IsArray() { + return nil + } + items := node.Array() + if len(items) == 0 { + return nil + } + out := make([]usage.IterationUsage, 0, len(items)) + for _, item := range items { + if !item.IsObject() { + continue + } + it := usage.IterationUsage{} + if v := item.Get("type"); v.Type == gjson.String { + it.Type = cloneTrimmedString(v.String()) + } + if v := item.Get("model"); v.Type == gjson.String { + it.Model = cloneTrimmedString(v.String()) + } + if v := item.Get("input_tokens"); v.Type == gjson.Number { + it.InputTokens = usage.Int64Ptr(v.Int()) + } + if v := item.Get("output_tokens"); v.Type == gjson.Number { + it.OutputTokens = usage.Int64Ptr(v.Int()) + } + if v := firstExistingUsageNode(item, "output_tokens_details.thinking_tokens", "thinking_tokens"); v.Type == gjson.Number { + val := v.Int() + if val < 0 { + val = 0 + } + it.ThinkingTokens = usage.Int64Ptr(val) + it.ReasoningTokens = usage.Int64Ptr(val) + } else if v := item.Get("output_tokens_details.reasoning_tokens"); v.Type == gjson.Number { + val := v.Int() + if val < 0 { + val = 0 + } + it.ReasoningTokens = usage.Int64Ptr(val) + } + if v := item.Get("cache_read_input_tokens"); v.Type == gjson.Number { + it.CacheReadInputTokens = usage.Int64Ptr(v.Int()) + } + if v := item.Get("cache_creation_input_tokens"); v.Type == gjson.Number { + it.CacheCreationInputTokens = usage.Int64Ptr(v.Int()) + } + it.CacheCreation = parseClaudeCacheCreationNode(item.Get("cache_creation")) + if it.CacheCreation != nil { + if it.CacheCreation.Ephemeral5mInputTokens != nil { + it.CacheCreation5mTokens = usage.Int64Ptr(*it.CacheCreation.Ephemeral5mInputTokens) + } + if it.CacheCreation.Ephemeral1hInputTokens != nil { + it.CacheCreation1hTokens = usage.Int64Ptr(*it.CacheCreation.Ephemeral1hInputTokens) + } + } + out = append(out, it) } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.TokenBreakdown.TotalTokens + if len(out) == 0 { + return nil } - return detail + return out } -func ParseOpenAIStreamUsage(line []byte) (usage.Detail, bool) { - payload := jsonPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return usage.Detail{}, false +func deriveClaudeCacheCreationFromIterations(iterations []usage.IterationUsage, totalCacheCreation int64) *usage.CacheDetails { + if len(iterations) == 0 || totalCacheCreation < 0 { + return nil } - responseServiceTier := extractResponseServiceTier(payload) - usageNode := gjson.GetBytes(payload, "usage") - if !hasOpenAIStyleUsageTokenFields(usageNode) { - if responseServiceTier == "" { - return usage.Detail{}, false + byTTL := make(map[string]int64) + var has5m, has30m, has1h, has12h bool + for _, it := range iterations { + if it.Type != "" && it.Type != "message" { + continue + } + if it.CacheCreation == nil { + if it.CacheCreationInputTokens != nil && *it.CacheCreationInputTokens > 0 { + return nil + } + continue + } + if it.CacheCreation.Ephemeral5mInputTokens != nil { + has5m = true + } + if it.CacheCreation.Ephemeral30mInputTokens != nil { + has30m = true + } + if it.CacheCreation.Ephemeral1hInputTokens != nil { + has1h = true + } + if it.CacheCreation.Ephemeral12hInputTokens != nil { + has12h = true + } + for ttl, tokens := range it.CacheCreation.ByTTL { + sum, ok := safeUsageTokenSum(byTTL[ttl], tokens) + if !ok { + return nil + } + byTTL[ttl] = sum } - return usage.Detail{ResponseServiceTier: responseServiceTier}, true } - detail := parseOpenAIStyleUsageNode(usageNode) - detail.ResponseServiceTier = responseServiceTier - return detail, true -} - -func ParseClaudeUsage(data []byte) usage.Detail { - usageNode := gjson.ParseBytes(data).Get("usage") - if !usageNode.Exists() { - return usage.Detail{} + if len(byTTL) == 0 { + return nil } - return parseClaudeUsageNode(usageNode) -} - -func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) { - payload := jsonPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return usage.Detail{}, false + out := &usage.CacheDetails{ByTTL: byTTL} + sum, ok := out.SumTokens() + if !ok || sum != totalCacheCreation { + return nil } - usageNode := gjson.GetBytes(payload, "usage") - if !usageNode.Exists() { - usageNode = gjson.GetBytes(payload, "message.usage") + if has5m { + out.Ephemeral5mInputTokens = usage.Int64Ptr(byTTL["5m"]) } - if !usageNode.Exists() { - return usage.Detail{}, false + if has30m { + out.Ephemeral30mInputTokens = usage.Int64Ptr(byTTL["30m"]) } - return parseClaudeUsageNode(usageNode), true + if has1h { + out.Ephemeral1hInputTokens = usage.Int64Ptr(byTTL["1h"]) + } + if has12h { + out.Ephemeral12hInputTokens = usage.Int64Ptr(byTTL["12h"]) + } + return out } -func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail { - cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int() - cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int() - rawOutputTokens := usageNode.Get("output_tokens").Int() - // Anthropic reports thinking as a subset of output_tokens. Prefer the official - // nested field, then fall back to legacy aliases used by some gateways. - reasoningNode := firstExistingUsageNode( - usageNode, - "output_tokens_details.thinking_tokens", - "output_tokens_details.reasoning_tokens", - "thinking_tokens", - ) - reasoningTokens := reasoningNode.Int() - if reasoningTokens < 0 { - reasoningTokens = 0 +func parseClaudeCostNode(usageNode gjson.Result) *usage.CostDetails { + costNode := usageNode.Get("cost") + byokNode := usageNode.Get("is_byok") + detailsNode := usageNode.Get("cost_details") + out := &usage.CostDetails{} + hasAny := false + if costNode.Type == gjson.Number { + out.Cost = usage.Float64Ptr(costNode.Float()) + hasAny = true } - nonReasoningOutput := rawOutputTokens - if reasoningTokens > 0 && reasoningTokens <= rawOutputTokens { - nonReasoningOutput = rawOutputTokens - reasoningTokens - } else if reasoningTokens > rawOutputTokens { - // Keep Detail.OutputTokens authoritative for keeper subset checks and - // avoid inventing extra non-reasoning output when the upstream payload - // is inconsistent. - nonReasoningOutput = 0 + if ticksNode := usageNode.Get("cost_in_usd_ticks"); ticksNode.Type == gjson.Number { + out.CostInUSDTicks = usage.Int64Ptr(ticksNode.Int()) + hasAny = true } - detail := usage.Detail{ - InputTokens: usageNode.Get("input_tokens").Int(), - OutputTokens: rawOutputTokens, - ReasoningTokens: reasoningTokens, - CachedTokens: cacheReadTokens, - CacheReadTokens: cacheReadTokens, - CacheCreationTokens: cacheCreationTokens, + if byokNode.Type == gjson.True || byokNode.Type == gjson.False { + out.IsBYOK = usage.BoolPtr(byokNode.Bool()) + hasAny = true } - // Anthropic reports the speed a response was served and billed at in - // usage.speed ("fast" or "standard"): fast mode costs more, and a - // fast request on a model without it runs, and is billed, at standard. - // Carried as the response tier so usage sinks can price it. - if speed := strings.TrimSpace(usageNode.Get("speed").String()); speed != "" { - detail.ResponseServiceTier = speed + if detailsNode.Exists() && detailsNode.IsObject() { + if v := detailsNode.Get("upstream_inference_cost"); v.Type == gjson.Number { + out.UpstreamInferenceCost = usage.Float64Ptr(v.Float()) + hasAny = true + } + if v := detailsNode.Get("upstream_inference_prompt_cost"); v.Type == gjson.Number { + out.UpstreamInferencePromptCost = usage.Float64Ptr(v.Float()) + hasAny = true + } + if v := detailsNode.Get("upstream_inference_completions_cost"); v.Type == gjson.Number { + out.UpstreamInferenceCompletionsCost = usage.Float64Ptr(v.Float()) + hasAny = true + } } - if detail.CachedTokens == 0 { - detail.CachedTokens = detail.CacheCreationTokens + if !hasAny { + return nil } - // raw output_tokens already includes thinking; cache fields are independent - // from input_tokens in the Messages API. - detail.TotalTokens = detail.InputTokens + rawOutputTokens + detail.CacheReadTokens + detail.CacheCreationTokens - detail.TokenBreakdown = usage.NewIndependentTokenBreakdown( - detail.InputTokens, - detail.CacheReadTokens, - detail.CacheCreationTokens, - nonReasoningOutput, - detail.ReasoningTokens, - detail.TotalTokens, - ) - return detail + return out } func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail { - cachedTokens := node.Get("cachedContentTokenCount").Int() - toolUseTokens := firstExistingUsageNode(node, "toolUsePromptTokenCount", "tool_use_prompt_token_count").Int() - inputTokens, okInput := safeUsageTokenSum(node.Get("promptTokenCount").Int(), toolUseTokens) + if !node.Exists() || !node.IsObject() { + return usage.Detail{} + } + var ( + promptTokens int64 + toolUseTokens int64 + candidatesTokens int64 + thoughtsTokens int64 + totalTokens int64 + cachedTokens int64 + serviceTier string + trafficType string + promptDetails gjson.Result + toolUseDetails gjson.Result + cacheDetails gjson.Result + candDetails gjson.Result + ) + node.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "promptTokenCount", "prompt_token_count": + if promptTokens == 0 { + promptTokens = val.Int() + } + case "toolUsePromptTokenCount", "tool_use_prompt_token_count": + if toolUseTokens == 0 { + toolUseTokens = val.Int() + } + case "candidatesTokenCount", "candidates_token_count": + if candidatesTokens == 0 { + candidatesTokens = val.Int() + } + case "thoughtsTokenCount", "thoughts_token_count": + if thoughtsTokens == 0 { + thoughtsTokens = val.Int() + } + case "totalTokenCount", "total_token_count": + if totalTokens == 0 { + totalTokens = val.Int() + } + case "cachedContentTokenCount", "cached_content_token_count": + if cachedTokens == 0 { + cachedTokens = val.Int() + } + case "serviceTier", "service_tier": + if serviceTier == "" && val.Type == gjson.String { + serviceTier = cloneTrimmedString(val.String()) + } + case "trafficType", "traffic_type": + if trafficType == "" && val.Type == gjson.String { + trafficType = cloneTrimmedString(val.String()) + } + case "promptTokensDetails", "prompt_tokens_details": + if !promptDetails.Exists() { + promptDetails = val + } + case "toolUsePromptTokensDetails", "tool_use_prompt_tokens_details": + if !toolUseDetails.Exists() { + toolUseDetails = val + } + case "cacheTokensDetails", "cache_tokens_details": + if !cacheDetails.Exists() { + cacheDetails = val + } + case "candidatesTokensDetails", "candidates_tokens_details": + if !candDetails.Exists() { + candDetails = val + } + } + return true + }) + inputTokens, okInput := safeUsageTokenSum(promptTokens, toolUseTokens) detail := usage.Detail{ InputTokens: inputTokens, - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), + OutputTokens: candidatesTokens, + ReasoningTokens: thoughtsTokens, + TotalTokens: totalTokens, CachedTokens: cachedTokens, CacheReadTokens: cachedTokens, } @@ -1312,9 +3250,352 @@ func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail { detail.ReasoningTokens, detail.TotalTokens, ) + if serviceTier != "" { + detail.ResponseServiceTier = serviceTier + if detail.UpstreamServiceTier == "" { + detail.UpstreamServiceTier = serviceTier + } + } + if trafficType != "" { + detail.UpstreamServiceTier = trafficType + if detail.ResponseServiceTier == "" { + detail.ResponseServiceTier = trafficType + } + } + if toolUseTokens > 0 { + detail.ToolUsageDetails = &usage.ToolUsageDetails{ + Extra: map[string]int64{"tool_use_prompt_tokens": toolUseTokens}, + } + } + parseModalityList := func(arr gjson.Result, target func(modality string, count int64)) { + if !arr.Exists() || !arr.IsArray() { + return + } + arr.ForEach(func(_, entry gjson.Result) bool { + mod := strings.ToUpper(strings.TrimSpace(entry.Get("modality").String())) + if mod == "" { + return true + } + cntNode := entry.Get("tokenCount") + if !cntNode.Exists() { + cntNode = entry.Get("token_count") + } + if cntNode.Type == gjson.Number { + target(mod, cntNode.Int()) + } + return true + }) + } + if promptDetails.Exists() || toolUseDetails.Exists() { + var textIn, imgIn, audioIn, videoIn, docIn int64 + var hasTextIn, hasImgIn, hasAudioIn, hasVideoIn, hasDocIn bool + addInputModality := func(mod string, cnt int64) { + switch mod { + case "TEXT": + textIn += cnt + hasTextIn = true + case "IMAGE": + imgIn += cnt + hasImgIn = true + case "AUDIO": + audioIn += cnt + hasAudioIn = true + case "VIDEO": + videoIn += cnt + hasVideoIn = true + case "DOCUMENT": + docIn += cnt + hasDocIn = true + } + } + parseModalityList(promptDetails, addInputModality) + parseModalityList(toolUseDetails, addInputModality) + if hasTextIn { + detail.TextInputTokens = usage.Int64Ptr(textIn) + } + if hasImgIn { + detail.ImageInputTokens = usage.Int64Ptr(imgIn) + } + if hasAudioIn { + detail.AudioInputTokens = usage.Int64Ptr(audioIn) + } + if hasVideoIn { + detail.VideoInputTokens = usage.Int64Ptr(videoIn) + } + if hasDocIn { + detail.FileInputTokens = usage.Int64Ptr(docIn) + } + } + if cacheDetails.Exists() { + var cachedText int64 + var hasCachedText bool + parseModalityList(cacheDetails, func(mod string, cnt int64) { + if mod == "TEXT" { + cachedText += cnt + hasCachedText = true + } + }) + if hasCachedText { + detail.CachedTextTokens = usage.Int64Ptr(cachedText) + } + } + if candDetails.Exists() { + var textOut, imgOut, audioOut int64 + var hasTextOut, hasImgOut, hasAudioOut bool + parseModalityList(candDetails, func(mod string, cnt int64) { + switch mod { + case "TEXT": + textOut += cnt + hasTextOut = true + case "IMAGE": + imgOut += cnt + hasImgOut = true + case "AUDIO": + audioOut += cnt + hasAudioOut = true + } + }) + if hasTextOut { + detail.TextOutputTokens = usage.Int64Ptr(textOut) + } + if hasImgOut { + detail.ImageOutputTokens = usage.Int64Ptr(imgOut) + } + if hasAudioOut { + detail.AudioOutputTokens = usage.Int64Ptr(audioOut) + } + } return detail } +func applyGeminiPayloadArtifacts(detail *usage.Detail, root gjson.Result) { + if detail == nil || !root.Exists() || !root.IsObject() { + return + } + candidates := root.Get("candidates") + if !candidates.Exists() { + candidates = root.Get("response.candidates") + } + var preds gjson.Result + if strings.Contains(root.Raw, `"predictions"`) { + preds = root.Get("predictions") + } + applyGeminiNodesArtifacts(detail, candidates, preds) +} + +func applyGeminiNodesArtifacts(detail *usage.Detail, candidates, preds gjson.Result) { + var callCounter geminiFunctionCallCounter + applyGeminiNodesArtifactsWithCounter(detail, candidates, preds, &callCounter) +} + +func applyGeminiNodesArtifactsWithCounter(detail *usage.Detail, candidates, preds gjson.Result, callCounter *geminiFunctionCallCounter) { + if detail == nil { + return + } + if callCounter == nil { + callCounter = &geminiFunctionCallCounter{} + } + if candidates.Exists() && candidates.IsArray() { + var ( + webSearchGrounded bool + webSearchQueries int64 + webChunks int64 + fileChunks int64 + mapsChunks int64 + urlFetchCount int64 + codeExecCount int64 + fnCallCount int64 + imageOutCount int64 + firstImageMime string + ) + candidates.ForEach(func(candPos, cand gjson.Result) bool { + if !cand.IsObject() { + return true + } + candIdx := candPos.Int() + if idxNode := cand.Get("index"); idxNode.Type == gjson.Number { + candIdx = idxNode.Int() + } + cand.ForEach(func(ck, cv gjson.Result) bool { + switch ck.String() { + case "groundingMetadata", "grounding_metadata": + if !cv.IsObject() { + return true + } + cv.ForEach(func(gk, gv gjson.Result) bool { + switch gk.String() { + case "webSearchQueries", "web_search_queries": + if gv.IsArray() { + gv.ForEach(func(_, qItem gjson.Result) bool { + if strings.TrimSpace(qItem.String()) != "" { + webSearchQueries++ + webSearchGrounded = true + } + return true + }) + } + case "searchEntryPoint", "search_entry_point": + webSearchGrounded = true + case "groundingChunks", "grounding_chunks": + if gv.IsArray() { + gv.ForEach(func(_, ch gjson.Result) bool { + switch { + case ch.Get("web").Exists(): + webChunks++ + webSearchGrounded = true + case ch.Get("retrievedContext").Exists() || ch.Get("retrieved_context").Exists(): + fileChunks++ + case ch.Get("maps").Exists(): + mapsChunks++ + } + return true + }) + } + } + return true + }) + case "urlContextMetadata", "url_context_metadata": + if cv.IsObject() { + if urls := firstExistingUsageNode(cv, "urlMetadata", "url_metadata"); urls.IsArray() { + urlFetchCount += int64(len(urls.Array())) + } + } + case "content": + if parts := cv.Get("parts"); parts.IsArray() { + parts.ForEach(func(_, part gjson.Result) bool { + if !part.IsObject() { + return true + } + part.ForEach(func(pk, pv gjson.Result) bool { + switch pk.String() { + case "executableCode", "executable_code": + codeExecCount++ + case "functionCall", "function_call": + fnCallCount += callCounter.observePart(candIdx, pv) + case "inlineData", "inline_data": + if pv.IsObject() { + mimeNode := pv.Get("mimeType") + if !mimeNode.Exists() { + mimeNode = pv.Get("mime_type") + } + mimeType := cloneTrimmedString(mimeNode.String()) + if len(mimeType) >= 6 && strings.EqualFold(mimeType[:6], "image/") { + imageOutCount++ + if firstImageMime == "" { + firstImageMime = mimeType + } + } + } + } + return true + }) + return true + }) + } + } + return true + }) + return true + }) + if webSearchGrounded || webSearchQueries > 0 || webChunks > 0 || fileChunks > 0 || mapsChunks > 0 || urlFetchCount > 0 || codeExecCount > 0 || fnCallCount > 0 { + if detail.ToolUsageDetails == nil { + detail.ToolUsageDetails = &usage.ToolUsageDetails{} + } + if webSearchGrounded { + calls := int64(1) + if webSearchQueries > 1 { + calls = webSearchQueries + } + detail.ToolUsageDetails.WebSearchCalls = usage.Int64Ptr(calls) + detail.WebSearchRequests = usage.Int64Ptr(calls) + } + if webSearchQueries > 0 { + detail.ToolUsageDetails.WebSearchQueries = usage.Int64Ptr(webSearchQueries) + } + if webChunks > 0 { + detail.ToolUsageDetails.NumSearchResults = usage.Int64Ptr(webChunks) + detail.ToolUsageDetails.NumSourcesUsed = usage.Int64Ptr(webChunks) + } + if fileChunks > 0 { + detail.ToolUsageDetails.FileSearchCalls = usage.Int64Ptr(fileChunks) + } + if urlFetchCount > 0 { + detail.ToolUsageDetails.WebFetchCalls = usage.Int64Ptr(urlFetchCount) + } + if codeExecCount > 0 { + detail.ToolUsageDetails.CodeExecutionCalls = usage.Int64Ptr(codeExecCount) + } + if fnCallCount > 0 { + detail.ToolUsageDetails.ToolCallsRequested = usage.Int64Ptr(fnCallCount) + } + if mapsChunks > 0 { + if detail.ToolUsageDetails.Extra == nil { + detail.ToolUsageDetails.Extra = make(map[string]int64) + } + detail.ToolUsageDetails.Extra["google_maps_chunks"] = mapsChunks + } + } + if imageOutCount > 0 { + detail.ImageCount = usage.Int64Ptr(imageOutCount) + if firstImageMime != "" { + if detail.ImageDetails == nil { + detail.ImageDetails = &usage.ImageDetails{} + } + detail.ImageDetails.MimeType = firstImageMime + } + } + } + if preds.Exists() && preds.IsArray() { + var count int64 + var firstMime string + preds.ForEach(func(_, pred gjson.Result) bool { + if !pred.IsObject() { + return true + } + var hasImg bool + var mime string + pred.ForEach(func(pk, pv gjson.Result) bool { + switch pk.String() { + case "bytesBase64Encoded", "gcsUri": + if pv.Type == gjson.String && len(pv.Raw) > 2 { + hasImg = true + } + case "mimeType": + if pv.Type == gjson.String { + mime = cloneTrimmedString(pv.String()) + } + } + return true + }) + if hasImg { + count++ + if firstMime == "" { + firstMime = mime + } + } + return true + }) + if count > 0 { + detail.ImageCount = usage.Int64Ptr(count) + if firstMime != "" { + if detail.ImageDetails == nil { + detail.ImageDetails = &usage.ImageDetails{} + } + detail.ImageDetails.MimeType = firstMime + } + } + } +} + +// ApplyGeminiRequestImageConfig populates detail.RequestedImageDetails from a Gemini or Vertex +// request payload without overwriting upstream-returned detail.ImageDetails. +func ApplyGeminiRequestImageConfig(detail *usage.Detail, requestPayload []byte) { + if detail == nil || detail.ImageCount == nil || *detail.ImageCount <= 0 { + return + } + detail.RequestedImageDetails = ParseRequestImageDetails(requestPayload) +} + func parseInteractionsUsageDetail(node gjson.Result) usage.Detail { cacheRead := firstExistingUsageNode(node, "cache_read_tokens", "cacheReadTokens") toolUseTokens := firstExistingUsageNode(node, "tool_use_tokens", "total_tool_use_tokens", "toolUseTokens", "totalToolUseTokens").Int() @@ -1410,25 +3691,70 @@ func ParseInteractionsStreamUsage(line []byte) (usage.Detail, bool) { } func ParseGeminiUsage(data []byte) usage.Detail { - usageNode := gjson.ParseBytes(data) - node := usageNode.Get("usageMetadata") - if !node.Exists() { - node = usageNode.Get("usage_metadata") - } - if !node.Exists() { + if len(data) == 0 || !gjson.ValidBytes(data) { return usage.Detail{} } - return parseGeminiFamilyUsageDetail(node) + root := gjson.Parse(unsafeBytesToString(data)) + var ( + node gjson.Result + candidatesNode gjson.Result + predsNode gjson.Result + ) + root.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "usageMetadata", "usage_metadata": + if !node.Exists() { + node = val + } + case "candidates": + candidatesNode = val + case "predictions": + predsNode = val + case "response": + if val.IsObject() { + val.ForEach(func(rk, rv gjson.Result) bool { + switch rk.String() { + case "usageMetadata", "usage_metadata": + if !node.Exists() { + node = rv + } + case "candidates": + if !candidatesNode.Exists() { + candidatesNode = rv + } + } + return true + }) + } + } + return true + }) + var detail usage.Detail + if node.Exists() { + detail = parseGeminiFamilyUsageDetail(node) + } + if candidatesNode.Exists() || predsNode.Exists() { + applyGeminiNodesArtifacts(&detail, candidatesNode, predsNode) + } + return detail } func ParseGeminiStreamUsage(line []byte) (usage.Detail, bool) { payload := jsonPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { + if len(payload) == 0 { return usage.Detail{}, false } - node := gjson.GetBytes(payload, "usageMetadata") - if !node.Exists() { - node = gjson.GetBytes(payload, "usage_metadata") + hasCamel := bytes.Contains(payload, geminiUsageMetaMarker) + hasSnake := !hasCamel && bytes.Contains(payload, geminiUsageMetaSnakeMarker) + if (!hasCamel && !hasSnake) || !gjson.ValidBytes(payload) { + return usage.Detail{}, false + } + root := gjson.Parse(unsafeBytesToString(payload)) + var node gjson.Result + if hasCamel { + node = root.Get("usageMetadata") + } else { + node = root.Get("usage_metadata") } if !node.Exists() { return usage.Detail{}, false @@ -1437,6 +3763,9 @@ func ParseGeminiStreamUsage(line []byte) (usage.Detail, bool) { if !hasNonZeroTokenUsage(detail) { return usage.Detail{}, false } + if hasGeminiCandidateArtifactMarker(payload) { + applyGeminiPayloadArtifacts(&detail, root) + } return detail, true } @@ -1474,36 +3803,35 @@ func invalidUsageTokenBreakdown(total int64) usage.TokenBreakdown { } func ParseAntigravityUsage(data []byte) usage.Detail { - usageNode := gjson.ParseBytes(data) - node := usageNode.Get("response.usageMetadata") - if !node.Exists() { - node = usageNode.Get("usageMetadata") - } - if !node.Exists() { - node = usageNode.Get("usage_metadata") - } - if !node.Exists() { - return usage.Detail{} - } - return parseGeminiFamilyUsageDetail(node) + return ParseGeminiUsage(data) } func ParseAntigravityStreamUsage(line []byte) (usage.Detail, bool) { payload := jsonPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { + if len(payload) == 0 { return usage.Detail{}, false } - node := gjson.GetBytes(payload, "response.usageMetadata") - if !node.Exists() { - node = gjson.GetBytes(payload, "usageMetadata") + hasCamel := bytes.Contains(payload, geminiUsageMetaMarker) + hasSnake := !hasCamel && bytes.Contains(payload, geminiUsageMetaSnakeMarker) + if (!hasCamel && !hasSnake) || !gjson.ValidBytes(payload) { + return usage.Detail{}, false } - if !node.Exists() { - node = gjson.GetBytes(payload, "usage_metadata") + root := gjson.Parse(unsafeBytesToString(payload)) + node := root.Get("response.usageMetadata") + if !node.Exists() && hasCamel { + node = root.Get("usageMetadata") + } + if !node.Exists() && hasSnake { + node = root.Get("usage_metadata") } if !node.Exists() { return usage.Detail{}, false } - return parseGeminiFamilyUsageDetail(node), true + detail := parseGeminiFamilyUsageDetail(node) + if hasGeminiCandidateArtifactMarker(payload) { + applyGeminiPayloadArtifacts(&detail, root) + } + return detail, true } var stopChunkWithoutUsage sync.Map @@ -1551,9 +3879,13 @@ func FilterSSEUsageMetadata(payload []byte) []byte { if !changed { continue } - var rebuilt []byte + rebuiltCap := dataIdx + 5 + if len(cleaned) > 0 { + rebuiltCap += 1 + len(cleaned) + } + rebuilt := make([]byte, 0, rebuiltCap) rebuilt = append(rebuilt, line[:dataIdx]...) - rebuilt = append(rebuilt, []byte("data:")...) + rebuilt = append(rebuilt, "data:"...) if len(cleaned) > 0 { rebuilt = append(rebuilt, ' ') rebuilt = append(rebuilt, cleaned...) @@ -1582,50 +3914,80 @@ func FilterSSEUsageMetadata(payload []byte) []byte { // - Antigravity: response.candidates.0.finishReason func StripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) { jsonBytes := bytes.TrimSpace(rawJSON) - if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) { + if len(jsonBytes) == 0 || !bytes.Contains(jsonBytes, geminiUsageMetaMarker) || !gjson.ValidBytes(jsonBytes) { return rawJSON, false } - // Check for finishReason in both aistudio and antigravity formats - finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason") - if !finishReason.Exists() { - finishReason = gjson.GetBytes(jsonBytes, "response.candidates.0.finishReason") - } - terminalReason := finishReason.Exists() && strings.TrimSpace(finishReason.String()) != "" - - usageMetadata := gjson.GetBytes(jsonBytes, "usageMetadata") - if !usageMetadata.Exists() { - usageMetadata = gjson.GetBytes(jsonBytes, "response.usageMetadata") - } + root := gjson.Parse(unsafeBytesToString(jsonBytes)) + var ( + finishReason string + topUsage gjson.Result + respUsage gjson.Result + topUsageKey gjson.Result + respUsageKey gjson.Result + ) + root.ForEach(func(key, val gjson.Result) bool { + switch key.String() { + case "usageMetadata": + topUsage = val + topUsageKey = key + case "candidates": + if finishReason == "" && val.IsArray() { + finishReason = strings.TrimSpace(val.Get("0.finishReason").String()) + } + case "response": + if val.IsObject() { + val.ForEach(func(rk, rv gjson.Result) bool { + switch rk.String() { + case "usageMetadata": + respUsage = rv + respUsageKey = rk + case "candidates": + if finishReason == "" && rv.IsArray() { + finishReason = strings.TrimSpace(rv.Get("0.finishReason").String()) + } + } + return true + }) + } + } + return true + }) - // Terminal chunk: keep as-is. - if terminalReason { + if finishReason != "" || (!topUsage.Exists() && !respUsage.Exists()) { return rawJSON, false } - // Nothing to strip - if !usageMetadata.Exists() { - return rawJSON, false + // Fast single-allocation key rename when only one usageMetadata exists + // (avoids 2 full-buffer sjson copies on multi-MB 4K image chunks). + if topUsage.Exists() != respUsage.Exists() { + key := topUsageKey + if respUsage.Exists() { + key = respUsageKey + } + if key.Raw == string(geminiUsageMetaMarker) && key.Index > 0 && + key.Index+len(key.Raw) <= len(jsonBytes) { + keyPos := key.Index + 1 // Right after the selected key's opening quote. + out := make([]byte, 0, len(jsonBytes)+3) + out = append(out, jsonBytes[:keyPos]...) + out = append(out, "cpaUsageMetadata"...) + out = append(out, jsonBytes[keyPos+len("usageMetadata"):]...) + return out, true + } } - // Remove usageMetadata from both possible locations cleaned := jsonBytes var changed bool - - if usageMetadata = gjson.GetBytes(cleaned, "usageMetadata"); usageMetadata.Exists() { - // Rename usageMetadata to cpaUsageMetadata - cleaned, _ = sjson.SetRawBytes(cleaned, "cpaUsageMetadata", []byte(usageMetadata.Raw)) + if topUsage.Exists() { + cleaned, _ = sjson.SetRawBytes(cleaned, "cpaUsageMetadata", []byte(topUsage.Raw)) cleaned, _ = sjson.DeleteBytes(cleaned, "usageMetadata") changed = true } - - if usageMetadata = gjson.GetBytes(cleaned, "response.usageMetadata"); usageMetadata.Exists() { - // Rename usageMetadata to cpaUsageMetadata - cleaned, _ = sjson.SetRawBytes(cleaned, "response.cpaUsageMetadata", []byte(usageMetadata.Raw)) + if respUsage.Exists() { + cleaned, _ = sjson.SetRawBytes(cleaned, "response.cpaUsageMetadata", []byte(respUsage.Raw)) cleaned, _ = sjson.DeleteBytes(cleaned, "response.usageMetadata") changed = true } - return cleaned, changed } @@ -1642,8 +4004,10 @@ func hasUsageMetadata(jsonBytes []byte) bool { return false } +var finishReasonMarker = []byte(`"finishReason"`) + func isStopChunkWithoutUsage(jsonBytes []byte) bool { - if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) { + if len(jsonBytes) == 0 || !bytes.Contains(jsonBytes, finishReasonMarker) || bytes.Contains(jsonBytes, geminiUsageMetaMarker) || !gjson.ValidBytes(jsonBytes) { return false } finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason") diff --git a/internal/runtime/executor/helps/usage_helpers_bench_test.go b/internal/runtime/executor/helps/usage_helpers_bench_test.go new file mode 100644 index 00000000000..661d9e3e95e --- /dev/null +++ b/internal/runtime/executor/helps/usage_helpers_bench_test.go @@ -0,0 +1,86 @@ +package helps + +import ( + "bytes" + "strings" + "testing" +) + +func BenchmarkObserveOpenAIStream_TextDelta(b *testing.B) { + deltaChunk := []byte(`data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1700000000,"model":"gpt-5.4","choices":[{"index":0,"delta":{"content":"Hello world"},"finish_reason":null}]}`) + finalChunk := []byte(`data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1700000000,"model":"gpt-5.4","service_tier":"default","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":120,"completion_tokens":45,"total_tokens":165,"prompt_tokens_details":{"cached_tokens":64,"text_tokens":120},"completion_tokens_details":{"reasoning_tokens":10,"text_tokens":35}}}`) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + var buf StreamUsageBuffer + for j := 0; j < 50; j++ { + buf.ObserveOpenAIStream(deltaChunk) + } + buf.ObserveOpenAIStream(finalChunk) + } +} + +func BenchmarkObserveClaudeStream_TextDelta(b *testing.B) { + msgStart := []byte(`data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-opus-4-6","usage":{"input_tokens":200,"cache_creation_input_tokens":100,"cache_read_input_tokens":50,"cache_creation":{"ephemeral_5m_input_tokens":100,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard"}}}`) + deltaChunk := []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello world"}}`) + msgDelta := []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":85,"server_tool_use":{"web_search_requests":1}}}`) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + var buf StreamUsageBuffer + buf.ObserveClaudeStream(msgStart) + for j := 0; j < 50; j++ { + buf.ObserveClaudeStream(deltaChunk) + } + buf.ObserveClaudeStream(msgDelta) + } +} + +func BenchmarkObserveGeminiStream_TextAndUsage(b *testing.B) { + deltaChunk := []byte(`data: {"candidates":[{"content":{"parts":[{"text":"Go 1.24 was released in "}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":22,"candidatesTokenCount":12,"totalTokenCount":114,"promptTokensDetails":[{"modality":"TEXT","tokenCount":22}],"thoughtsTokenCount":80,"serviceTier":"standard"},"modelVersion":"gemini-3.8-flash"}`) + finalChunk := []byte(`data: {"candidates":[{"content":{"parts":[{"text":"2025."}],"role":"model"},"finishReason":"STOP","index":0,"groundingMetadata":{"webSearchQueries":["Go 1.24 release year"],"groundingChunks":[{"web":{"uri":"https://go.dev/doc/devel/release","title":"Release History"}}]}}],"usageMetadata":{"promptTokenCount":103,"candidatesTokenCount":22,"totalTokenCount":286,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"thoughtsTokenCount":161,"serviceTier":"standard"},"modelVersion":"gemini-3.8-flash"}`) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + var buf StreamUsageBuffer + for j := 0; j < 10; j++ { + buf.ObserveGeminiStream(deltaChunk) + } + buf.ObserveGeminiStream(finalChunk) + } +} + +func BenchmarkParseOpenAIImageUsage_LargeBase64(b *testing.B) { + reqPayload := []byte(`{"model":"grok-imagine-image","prompt":"A minimal orange square","response_format":"b64_json","aspect_ratio":"16:9","resolution":"1k"}`) + largeB64 := strings.Repeat("QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVowMTIzNDU2Nzg5", 30000) // ~1.44MB + respPayload := []byte(`{"data":[{"b64_json":"` + largeB64 + `","mime_type":"image/jpeg"}],"usage":{"cost_in_usd_ticks":200000000}}`) + b.SetBytes(int64(len(respPayload))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + detail := ParseOpenAIImageUsage(reqPayload, respPayload) + if detail.ImageCount == nil || *detail.ImageCount != 1 { + b.Fatal("expected image count 1") + } + } +} + +func BenchmarkParseGeminiUsage_LargeBase64(b *testing.B) { + reqPayload := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw"}]}],"generationConfig":{"responseModalities":["TEXT","IMAGE"],"imageConfig":{"aspectRatio":"16:9","imageSize":"1K"}}}`) + largeB64 := strings.Repeat("QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVowMTIzNDU2Nzg5", 30000) // ~1.44MB + var buf bytes.Buffer + buf.WriteString(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/jpeg","data":"`) + buf.WriteString(largeB64) + buf.WriteString(`"}}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":1343,"totalTokenCount":1354,"promptTokensDetails":[{"modality":"TEXT","tokenCount":11}],"candidatesTokensDetails":[{"modality":"IMAGE","tokenCount":1120}],"serviceTier":"standard"}}`) + respPayload := buf.Bytes() + b.SetBytes(int64(len(respPayload))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + detail := ParseGeminiUsage(respPayload) + ApplyGeminiRequestImageConfig(&detail, reqPayload) + if detail.ImageCount == nil || *detail.ImageCount != 1 { + b.Fatal("expected image count 1") + } + } +} diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 83bc0ae1580..48f9fb4cfb9 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -1,6 +1,8 @@ package helps import ( + "bufio" + "bytes" "context" "errors" "fmt" @@ -23,6 +25,7 @@ import ( 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" + "github.com/tidwall/gjson" ) func TestParseOpenAIUsageChatCompletions(t *testing.T) { @@ -1501,3 +1504,577 @@ func TestFailFromErrorsCarriesStreamTruncatedCode(t *testing.T) { t.Fatalf("plain error must not gain truncated classification, got code %q", plain.Code) } } + +func TestParseClaudeUsageExtractsCacheTTLServerToolsIterationsAndCost(t *testing.T) { + var buffer StreamUsageBuffer + lineStart := []byte(`data: {"type":"message_start","message":{"id":"msg_1","model":"claude-opus-5-5","usage":{"input_tokens":4,"cache_creation_input_tokens":3036,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":3036},"output_tokens":16,"service_tier":"standard","inference_geo":"global","speed":"standard"}}}`) + lineToolUse := []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"server_tool_use","id":"srvtoolu_1","name":"web_search","input":{"query":"anthropic"}}}`) + lineToolResult := []byte(`data: {"type":"content_block_start","index":1,"content_block":{"type":"web_search_tool_result","tool_use_id":"srvtoolu_1","content":[{"type":"web_search_result","url":"https://www.anthropic.com","title":"Anthropic"},{"type":"web_search_result","url":"https://docs.anthropic.com","title":"Docs"}]}}`) + lineDelta := []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":6,"cache_creation_input_tokens":8846,"cache_read_input_tokens":3036,"output_tokens":483,"output_tokens_details":{"thinking_tokens":45},"server_tool_use":{"web_search_requests":1,"web_fetch_requests":0},"iterations":[{"input_tokens":4,"output_tokens":53,"cache_read_input_tokens":0,"cache_creation_input_tokens":3036,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":3036},"type":"message","model":null},{"input_tokens":2,"output_tokens":430,"cache_read_input_tokens":3036,"cache_creation_input_tokens":5810,"cache_creation":{"ephemeral_5m_input_tokens":5810,"ephemeral_1h_input_tokens":0},"type":"message","model":null}],"cost":0.035471,"is_byok":false,"cost_details":{"upstream_inference_cost":0.035471,"upstream_inference_prompt_cost":0.023151,"upstream_inference_completions_cost":0.00232}}}`) + + buffer.ObserveClaudeStream(lineStart) + buffer.ObserveClaudeStream(lineToolUse) + buffer.ObserveClaudeStream(lineToolResult) + buffer.ObserveClaudeStream(lineDelta) + + detail, ok := buffer.Detail() + if !ok { + t.Fatal("expected observed stream usage") + } + if detail.InputTokens != 6 || detail.CacheReadTokens != 3036 || detail.CacheCreationTokens != 8846 || detail.OutputTokens != 483 || detail.ReasoningTokens != 45 || detail.TotalTokens != 12371 { + t.Fatalf("unexpected token totals: %+v", detail) + } + if detail.CacheCreation5mTokens == nil || *detail.CacheCreation5mTokens != 5810 { + t.Fatalf("CacheCreation5mTokens = %v, want 5810", detail.CacheCreation5mTokens) + } + if detail.CacheCreation1hTokens == nil || *detail.CacheCreation1hTokens != 3036 { + t.Fatalf("CacheCreation1hTokens = %v, want 3036", detail.CacheCreation1hTokens) + } + if detail.WebSearchRequests == nil || *detail.WebSearchRequests != 1 { + t.Fatalf("WebSearchRequests = %v, want 1", detail.WebSearchRequests) + } + if detail.ToolUsageDetails == nil || detail.ToolUsageDetails.WebSearchCalls == nil || *detail.ToolUsageDetails.WebSearchCalls != 1 || detail.ToolUsageDetails.WebFetchCalls == nil || *detail.ToolUsageDetails.WebFetchCalls != 0 || detail.ToolUsageDetails.WebSearchQueries == nil || *detail.ToolUsageDetails.WebSearchQueries != 1 || detail.ToolUsageDetails.NumSearchResults == nil || *detail.ToolUsageDetails.NumSearchResults != 2 { + t.Fatalf("ToolUsageDetails = %+v, want web_search=1 queries=1 results=2 web_fetch=0", detail.ToolUsageDetails) + } + iterations := detail.IterationSlice() + if len(iterations) != 2 { + t.Fatalf("len(iterations) = %d, want 2", len(iterations)) + } + if iterations[0].CacheCreation1hTokens == nil || *iterations[0].CacheCreation1hTokens != 3036 || iterations[1].CacheCreation5mTokens == nil || *iterations[1].CacheCreation5mTokens != 5810 { + t.Fatalf("iterations = %+v", iterations) + } + if detail.InferenceGeo != "global" || detail.Speed != "standard" || detail.UpstreamServiceTier != "standard" { + t.Fatalf("geo=%q speed=%q tier=%q", detail.InferenceGeo, detail.Speed, detail.UpstreamServiceTier) + } + if detail.CostDetails == nil || detail.CostDetails.Cost == nil || *detail.CostDetails.Cost != 0.035471 || detail.CostDetails.IsBYOK == nil || *detail.CostDetails.IsBYOK != false { + t.Fatalf("CostDetails = %+v", detail.CostDetails) + } +} + +func TestParseGeminiAndAntigravityUsageWithModalitiesAndTools(t *testing.T) { + payload := []byte(`{"candidates":[{"content":{"parts":[{"executableCode":{"language":"PYTHON","code":"print(1)"}},{"functionCall":{"name":"lookup","args":{"q":"a"}}},{"inlineData":{"mimeType":"image/png","data":"AA=="}}]},"groundingMetadata":{"webSearchQueries":["gemini pricing","vertex grounding"],"groundingChunks":[{"web":{"uri":"https://a.example"}},{"web":{"uri":"https://b.example"}},{"retrievedContext":{"uri":"gs://b/doc"}}]},"urlContextMetadata":{"urlMetadata":[{"retrievedUrl":"https://c.example"}]}}],"usageMetadata":{"promptTokenCount":1000,"cachedContentTokenCount":400,"candidatesTokenCount":300,"thoughtsTokenCount":120,"toolUsePromptTokenCount":50,"totalTokenCount":1470,"trafficType":"ON_DEMAND_PRIORITY","promptTokensDetails":[{"modality":"TEXT","tokenCount":800},{"modality":"IMAGE","tokenCount":150},{"modality":"DOCUMENT","tokenCount":50}],"cacheTokensDetails":[{"modality":"TEXT","tokenCount":400}],"candidatesTokensDetails":[{"modality":"TEXT","tokenCount":100},{"modality":"IMAGE","tokenCount":200}]}}`) + detail := ParseGeminiUsage(payload) + ApplyGeminiRequestImageConfig(&detail, []byte(`{"generationConfig":{"imageConfig":{"aspectRatio":"16:9","imageSize":"2K"}}}`)) + if detail.InputTokens != 1050 || detail.OutputTokens != 300 || detail.ReasoningTokens != 120 || detail.TotalTokens != 1470 { + t.Fatalf("unexpected Gemini token totals: %+v", detail) + } + if detail.TextInputTokens == nil || *detail.TextInputTokens != 800 || detail.ImageInputTokens == nil || *detail.ImageInputTokens != 150 || detail.FileInputTokens == nil || *detail.FileInputTokens != 50 || detail.CachedTextTokens == nil || *detail.CachedTextTokens != 400 { + t.Fatalf("unexpected Gemini input modalities: %+v", detail) + } + if detail.TextOutputTokens == nil || *detail.TextOutputTokens != 100 || detail.ImageOutputTokens == nil || *detail.ImageOutputTokens != 200 { + t.Fatalf("unexpected Gemini output modalities: %+v", detail) + } + if detail.UpstreamServiceTier != "ON_DEMAND_PRIORITY" || detail.ResponseServiceTier != "ON_DEMAND_PRIORITY" { + t.Fatalf("unexpected Gemini trafficType tier: upstream=%q response=%q", detail.UpstreamServiceTier, detail.ResponseServiceTier) + } + if detail.WebSearchRequests == nil || *detail.WebSearchRequests != 2 || detail.ToolUsageDetails == nil || + detail.ToolUsageDetails.WebSearchCalls == nil || *detail.ToolUsageDetails.WebSearchCalls != 2 || + detail.ToolUsageDetails.WebSearchQueries == nil || *detail.ToolUsageDetails.WebSearchQueries != 2 || + detail.ToolUsageDetails.NumSearchResults == nil || *detail.ToolUsageDetails.NumSearchResults != 2 || + detail.ToolUsageDetails.FileSearchCalls == nil || *detail.ToolUsageDetails.FileSearchCalls != 1 || + detail.ToolUsageDetails.WebFetchCalls == nil || *detail.ToolUsageDetails.WebFetchCalls != 1 || + detail.ToolUsageDetails.CodeExecutionCalls == nil || *detail.ToolUsageDetails.CodeExecutionCalls != 1 || + detail.ToolUsageDetails.ToolCallsRequested == nil || *detail.ToolUsageDetails.ToolCallsRequested != 1 || + detail.ToolUsageDetails.Extra["tool_use_prompt_tokens"] != 50 { + t.Fatalf("unexpected Gemini ToolUsageDetails: %+v", detail.ToolUsageDetails) + } + if detail.ImageCount == nil || *detail.ImageCount != 1 || detail.ImageDetails == nil || detail.ImageDetails.MimeType != "image/png" { + t.Fatalf("unexpected Gemini response ImageCount/ImageDetails: count=%+v details=%+v", detail.ImageCount, detail.ImageDetails) + } + if detail.RequestedImageDetails == nil || detail.RequestedImageDetails.AspectRatio != "16:9" || detail.RequestedImageDetails.ImageSize != "2K" { + t.Fatalf("unexpected Gemini RequestedImageDetails: %+v", detail.RequestedImageDetails) + } + + var streamBuf StreamUsageBuffer + streamBuf.ObserveAntigravityStream([]byte(`data: {"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"tool_a","args":{}}}]},"groundingMetadata":{"webSearchQueries":["q1"],"groundingChunks":[{"web":{"uri":"https://x.example"}}]}}]}}`)) + streamBuf.ObserveAntigravityStream([]byte(`data: {"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":120,"candidatesTokenCount":30,"totalTokenCount":150}}}`)) + streamDetail, ok := streamBuf.Detail() + if !ok { + t.Fatal("expected streamBuf.Detail() ok = true") + } + if streamDetail.TotalTokens != 150 || streamDetail.WebSearchRequests == nil || *streamDetail.WebSearchRequests != 1 || streamDetail.ToolUsageDetails == nil || streamDetail.ToolUsageDetails.ToolCallsRequested == nil || *streamDetail.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("unexpected Antigravity stream merged usage: %+v tools=%+v", streamDetail, streamDetail.ToolUsageDetails) + } +} + +func TestParseCodexUsageWithCollectedOutputItemsAndImageToolUsage(t *testing.T) { + completed := []byte(`{"type":"response.completed","response":{"id":"resp_codex_1","service_tier":"priority","prompt_cache_retention":"24h","output":[],"usage":{"input_tokens":1200,"input_tokens_details":{"cached_tokens":1024,"text_tokens":1100,"image_tokens":100,"cached_tokens_details":{"text_tokens":1024}},"output_tokens":350,"output_tokens_details":{"reasoning_tokens":200,"text_tokens":150},"total_tokens":1550,"attribution":{"tier":"priority"}},"tool_usage":{"image_gen":{"input_tokens":80,"input_tokens_details":{"text_tokens":30,"image_tokens":50},"output_tokens":1056,"output_tokens_details":{"image_tokens":1056,"text_tokens":0},"total_tokens":1136}}}}`) + itemsByIndex := map[int64][]byte{ + 0: []byte(`{"id":"ws_1","type":"web_search_call","status":"completed","action":{"type":"search","query":"openai responses usage","sources":[{"url":"https://a.example"},{"url":"https://b.example"}]}}`), + 1: []byte(`{"id":"ci_1","type":"code_interpreter_call","status":"completed"}`), + 2: []byte(`{"id":"ig_1","type":"image_generation_call","status":"completed","size":"1024x1024","quality":"high","background":"transparent","output_format":"png","result":"AA=="}`), + 3: []byte(`{"id":"fc_1","type":"function_call","status":"completed","name":"run_check","call_id":"call_1","arguments":"{}"}`), + } + detail, ok := ParseCodexUsageWithOutputItems(completed, itemsByIndex, nil) + if !ok { + t.Fatal("expected ParseCodexUsageWithOutputItems ok = true") + } + if detail.ResponseServiceTier != "priority" || detail.UpstreamServiceTier != "priority" { + t.Fatalf("unexpected service tier: response=%q upstream=%q", detail.ResponseServiceTier, detail.UpstreamServiceTier) + } + if detail.CacheDetails == nil || detail.CacheDetails.Retention != "24h" { + t.Fatalf("unexpected CacheDetails: %+v", detail.CacheDetails) + } + if detail.CachedTextTokens == nil || *detail.CachedTextTokens != 1024 { + t.Fatalf("unexpected CachedTextTokens: %+v", detail.CachedTextTokens) + } + if detail.WebSearchRequests == nil || *detail.WebSearchRequests != 1 { + t.Fatalf("unexpected WebSearchRequests: %+v", detail.WebSearchRequests) + } + if detail.ToolUsageDetails == nil || + detail.ToolUsageDetails.WebSearchCalls == nil || *detail.ToolUsageDetails.WebSearchCalls != 1 || + detail.ToolUsageDetails.WebSearchQueries == nil || *detail.ToolUsageDetails.WebSearchQueries != 1 || + detail.ToolUsageDetails.NumSourcesUsed == nil || *detail.ToolUsageDetails.NumSourcesUsed != 2 || + detail.ToolUsageDetails.CodeExecutionCalls == nil || *detail.ToolUsageDetails.CodeExecutionCalls != 1 || + detail.ToolUsageDetails.ImageGenerationCalls == nil || *detail.ToolUsageDetails.ImageGenerationCalls != 1 || + detail.ToolUsageDetails.ToolCallsRequested == nil || *detail.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("unexpected ToolUsageDetails: %+v", detail.ToolUsageDetails) + } + if detail.ImageCount != nil || detail.ImageDetails != nil { + t.Fatalf("expected orchestrator record ImageCount/ImageDetails to be nil when tool_usage.image_gen is present, got count=%+v details=%+v", detail.ImageCount, detail.ImageDetails) + } + + imgToolDetail, okImg := ParseCodexImageToolUsageWithOutputItems([]byte(`{"tools":[{"type":"image_generation","input_fidelity":"high"}]}`), completed, itemsByIndex, nil) + if !okImg { + t.Fatal("expected ParseCodexImageToolUsageWithOutputItems ok = true") + } + if imgToolDetail.TotalTokens != 1136 || imgToolDetail.TextInputTokens == nil || *imgToolDetail.TextInputTokens != 30 || imgToolDetail.ImageInputTokens == nil || *imgToolDetail.ImageInputTokens != 50 || imgToolDetail.ImageOutputTokens == nil || *imgToolDetail.ImageOutputTokens != 1056 { + t.Fatalf("unexpected image tool token breakdown: %+v", imgToolDetail) + } + if imgToolDetail.ImageCount == nil || *imgToolDetail.ImageCount != 1 || imgToolDetail.ImageDetails == nil || imgToolDetail.ImageDetails.Size != "1024x1024" || imgToolDetail.ImageDetails.InputFidelity != "" { + t.Fatalf("unexpected image tool response metadata: count=%+v details=%+v", imgToolDetail.ImageCount, imgToolDetail.ImageDetails) + } + if imgToolDetail.RequestedImageDetails == nil || imgToolDetail.RequestedImageDetails.InputFidelity != "high" { + t.Fatalf("unexpected image tool requested metadata: %+v", imgToolDetail.RequestedImageDetails) + } +} + +func TestStreamUsageBufferObserveOpenAIStreamToolCallsAndImageCompleted(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"f1","arguments":""}}]}}]}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}},{"index":1,"id":"call_2","type":"function","function":{"name":"f2","arguments":"{}"}}]}}]}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"id":"chatcmpl-1","service_tier":"default","choices":[],"usage":{"prompt_tokens":50,"completion_tokens":20,"total_tokens":70}}`)) + detail, ok := buffer.Detail() + if !ok { + t.Fatal("expected buffer.Detail() ok = true") + } + if detail.TotalTokens != 70 || detail.ToolUsageDetails == nil || detail.ToolUsageDetails.ToolCallsRequested == nil || *detail.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("unexpected OpenAI stream tool call usage: %+v, tools=%+v", detail, detail.ToolUsageDetails) + } + + var imgBuffer StreamUsageBuffer + imgBuffer.ObserveOpenAIStream([]byte(`data: {"type":"image_generation.completed","b64_json":"AA==","size":"1024x1536","quality":"medium","background":"opaque","output_format":"webp","usage":{"input_tokens":25,"input_tokens_details":{"text_tokens":25,"image_tokens":0},"output_tokens":1500,"output_tokens_details":{"image_tokens":1500,"text_tokens":0},"total_tokens":1525}}`)) + imgDetail, okImg := imgBuffer.Detail() + if !okImg { + t.Fatal("expected imgBuffer.Detail() ok = true") + } + if imgDetail.TotalTokens != 1525 || imgDetail.ImageCount == nil || *imgDetail.ImageCount != 1 || imgDetail.ImageDetails == nil || imgDetail.ImageDetails.Size != "1024x1536" || imgDetail.ImageDetails.Quality != "medium" || imgDetail.ImageDetails.OutputFormat != "webp" { + t.Fatalf("unexpected OpenAI image stream usage: %+v, details=%+v", imgDetail, imgDetail.ImageDetails) + } +} + +func TestParseCodexUsageExtractsXAIServerSideToolUsageAndModalities(t *testing.T) { + data := []byte(`{"type":"response.completed","response":{"id":"resp_xai_1","usage":{"input_tokens":3804,"input_tokens_details":{"cached_tokens":3072,"text_tokens":3500,"image_tokens":304,"cached_prompt_text_tokens":3072},"output_tokens":889,"output_tokens_details":{"reasoning_tokens":828,"text_tokens":61,"accepted_prediction_tokens":5,"rejected_prediction_tokens":2},"total_tokens":4693,"num_sources_used":0,"num_server_side_tools_used":1,"cost_in_usd_ticks":76246640,"server_side_tool_usage_details":{"web_search_calls":0,"x_search_calls":1,"x_posts_fetched":1,"x_users_fetched":0,"code_interpreter_calls":0,"file_search_calls":0,"mcp_calls":0,"document_search_calls":0,"image_generation_calls":0}}}}`) + detail, ok := ParseCodexUsage(data) + if !ok { + t.Fatal("expected ParseCodexUsage ok = true") + } + if detail.TextInputTokens == nil || *detail.TextInputTokens != 3500 || detail.ImageInputTokens == nil || *detail.ImageInputTokens != 304 || detail.CachedTextTokens == nil || *detail.CachedTextTokens != 3072 { + t.Fatalf("unexpected input modality tokens: %+v", detail) + } + if detail.TextOutputTokens == nil || *detail.TextOutputTokens != 61 || detail.AcceptedPredictionTokens == nil || *detail.AcceptedPredictionTokens != 5 || detail.RejectedPredictionTokens == nil || *detail.RejectedPredictionTokens != 2 { + t.Fatalf("unexpected output modality tokens: %+v", detail) + } + if detail.ToolUsageDetails == nil || detail.ToolUsageDetails.XSearchCalls == nil || *detail.ToolUsageDetails.XSearchCalls != 1 || detail.ToolUsageDetails.XPostsFetched == nil || *detail.ToolUsageDetails.XPostsFetched != 1 || detail.ToolUsageDetails.NumServerSideToolsUsed == nil || *detail.ToolUsageDetails.NumServerSideToolsUsed != 1 { + t.Fatalf("unexpected ToolUsageDetails: %+v", detail.ToolUsageDetails) + } + if detail.CostDetails == nil || detail.CostDetails.CostInUSDTicks == nil || *detail.CostDetails.CostInUSDTicks != 76246640 { + t.Fatalf("unexpected CostDetails: %+v", detail.CostDetails) + } +} + +func TestParseXAIImageAndVideoUsage(t *testing.T) { + imgDetail := ParseXAIImageUsage( + []byte(`{"model":"grok-imagine-image","prompt":"cat","aspect_ratio":"16:9","resolution":"2k","quality":"medium"}`), + []byte(`{"data":[{"b64_json":"AA=="},{"b64_json":"BB=="}],"usage":{"prompt_tokens":12,"completion_tokens":100,"total_tokens":112}}`), + ) + if imgDetail.TotalTokens != 112 || imgDetail.ImageCount == nil || *imgDetail.ImageCount != 2 { + t.Fatalf("unexpected image usage: %+v", imgDetail) + } + if imgDetail.RequestedImageDetails == nil || imgDetail.RequestedImageDetails.AspectRatio != "16:9" || imgDetail.RequestedImageDetails.Resolution != "2k" || imgDetail.RequestedImageDetails.Quality != "medium" { + t.Fatalf("unexpected requested image details: %+v", imgDetail.RequestedImageDetails) + } + + vidDetail := ParseXAIVideoUsage( + []byte(`{"model":"grok-imagine-video","prompt":"wave","aspect_ratio":"16:9","resolution":"720p","duration":5}`), + []byte(`{"status":"completed","url":"https://example.com/v.mp4","video":{"resolution":"720p","aspect_ratio":"16:9"},"usage":{"prompt_tokens":20,"completion_tokens":400,"total_tokens":420}}`), + ) + if vidDetail.TotalTokens != 420 || vidDetail.ImageDetails == nil || vidDetail.ImageDetails.Resolution != "720p" || vidDetail.ImageDetails.AspectRatio != "16:9" || vidDetail.RequestedImageDetails == nil || vidDetail.RequestedImageDetails.Resolution != "720p" { + t.Fatalf("unexpected video usage: %+v", vidDetail) + } + if vidDetail.MediaCounts == nil || vidDetail.MediaCounts.CompletionMedia == nil || *vidDetail.MediaCounts.CompletionMedia != 1 || vidDetail.MediaCounts.AudioSeconds == nil || *vidDetail.MediaCounts.AudioSeconds != 5 { + t.Fatalf("unexpected video media counts: %+v", vidDetail.MediaCounts) + } +} + +func TestStreamUsageBufferObserveClaudeStream_PreservesSingleTurnCacheCreationTTL(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveClaudeStream([]byte(`data: {"type":"message_start","message":{"id":"msg_1","model":"claude-haiku-5-5","usage":{"input_tokens":16,"cache_creation_input_tokens":45045,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":45045},"output_tokens":4,"speed":"standard"}}}`)) + buffer.ObserveClaudeStream([]byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":4}}`)) + + detail, ok := buffer.Detail() + if !ok { + t.Fatal("expected observed stream usage") + } + if detail.CacheCreation5mTokens == nil || *detail.CacheCreation5mTokens != 0 { + t.Fatalf("CacheCreation5mTokens = %v, want explicit 0", detail.CacheCreation5mTokens) + } + if detail.CacheCreation1hTokens == nil || *detail.CacheCreation1hTokens != 45045 { + t.Fatalf("CacheCreation1hTokens = %v, want 45045", detail.CacheCreation1hTokens) + } +} + +func TestClaudeStreamUsageOwnsStrings(t *testing.T) { + line := []byte(`data: {"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1,"speed":"standard","service_tier":"priority","inference_geo":"global","cache_creation":{"mode":"explicit","ttl":"2h","retention":"24h","ephemeral_2h_input_tokens":5},"server_tool_use":{"custom_tool_requests":2},"iterations":[{"type":"message","model":"claude-sonnet"}]}}}`) + var b StreamUsageBuffer + b.ObserveClaudeStream(line) + for i := range line { + line[i] = 'X' + } + d, ok := b.Detail() + if !ok || d.ResponseServiceTier != "standard" || d.Speed != "standard" || + d.UpstreamServiceTier != "priority" || d.InferenceGeo != "global" { + t.Fatalf("usage strings changed after input reuse: %+v", d) + } + iterations := d.IterationSlice() + if len(iterations) != 1 || iterations[0].Type != "message" || iterations[0].Model != "claude-sonnet" { + t.Fatalf("iteration strings changed after input reuse: %+v", iterations) + } + if d.CacheDetails == nil || d.CacheDetails.Mode != "explicit" || d.CacheDetails.TTL != "2h" || + d.CacheDetails.Retention != "24h" || d.CacheDetails.ByTTL["2h"] != 5 { + t.Fatalf("cache strings changed after input reuse: %+v", d.CacheDetails) + } + if d.ToolUsageDetails == nil || d.ToolUsageDetails.Extra["custom_tool_requests"] != 2 { + t.Fatalf("tool counter key changed after input reuse: %+v", d.ToolUsageDetails) + } +} + +func TestClaudeStreamUsageSurvivesScannerReuse(t *testing.T) { + start := `data: {"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1,"speed":"standard","service_tier":"priority","inference_geo":"global"}}}` + delta := `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"` + strings.Repeat("A", 100) + `"}}` + stream := start + "\n" + strings.Repeat(delta+"\n", 100) + `data: {"type":"message_delta","usage":{"output_tokens":10}}` + "\n" + scanner := bufio.NewScanner(strings.NewReader(stream)) + var b StreamUsageBuffer + for scanner.Scan() { + b.ObserveClaudeStream(scanner.Bytes()) + } + if err := scanner.Err(); err != nil { + t.Fatal(err) + } + d, ok := b.Detail() + if !ok || d.InputTokens != 10 || d.OutputTokens != 10 || d.Speed != "standard" || + d.UpstreamServiceTier != "priority" || d.InferenceGeo != "global" { + t.Fatalf("scanner reuse corrupted usage: %+v", d) + } +} + +func TestClaudeStreamToolDeduplicationOwnsIDsAndCountsMultipleSearches(t *testing.T) { + use1 := []byte(`data: {"type":"content_block_start","content_block":{"type":"server_tool_use","id":"tool_1","name":"web_search"}}`) + result1 := []byte(`data: {"type":"content_block_start","content_block":{"type":"web_search_tool_result","tool_use_id":"tool_1","content":[]}}`) + result2 := []byte(`data: {"type":"content_block_start","content_block":{"type":"web_search_tool_result","tool_use_id":"tool_2","content":[]}}`) + var b StreamUsageBuffer + for i := 0; i < 2; i++ { + useCopy, resultCopy := bytes.Clone(use1), bytes.Clone(result1) + b.ObserveClaudeStream(useCopy) + b.ObserveClaudeStream(resultCopy) + for j := range useCopy { + useCopy[j] = 'X' + } + for j := range resultCopy { + resultCopy[j] = 'X' + } + } + b.ObserveClaudeStream(result2) + d, ok := b.Detail() + if !ok || d.ToolUsageDetails == nil || d.ToolUsageDetails.WebSearchQueries == nil || + *d.ToolUsageDetails.WebSearchQueries != 1 || d.ToolUsageDetails.WebSearchCalls == nil || + *d.ToolUsageDetails.WebSearchCalls != 2 || d.WebSearchRequests == nil || *d.WebSearchRequests != 2 { + t.Fatalf("unexpected Claude stream search counts: detail=%+v tools=%+v", d, d.ToolUsageDetails) + } +} + +func TestGeminiStreamUsageAccumulatesToolExecutions(t *testing.T) { + for _, wrapped := range []bool{false, true} { + t.Run(fmt.Sprintf("wrapped=%v", wrapped), func(t *testing.T) { + var b StreamUsageBuffer + observe := func(payload string) { + if wrapped { + b.ObserveAntigravityStream([]byte(`data: {"response":` + payload + `}`)) + } else { + b.ObserveGeminiStream([]byte("data: " + payload)) + } + } + for _, name := range []string{"first", "second"} { + observe(fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"functionCall":{"name":%q,"args":{}}},{"executableCode":{"language":"PYTHON","code":"print(1)"}},{"inlineData":{"mimeType":"image/png","data":"AA=="}}]},"groundingMetadata":{"webSearchQueries":["q1","q2"]}}]}`, name)) + d, ok := b.Detail() + if !ok || d.ToolUsageDetails == nil || d.ToolUsageDetails.ToolCallsRequested == nil || + d.ToolUsageDetails.CodeExecutionCalls == nil { + t.Fatalf("missing incremental tool usage: %+v", d) + } + want := int64(1) + if name == "second" { + want = 2 + } + if *d.ToolUsageDetails.ToolCallsRequested != want || *d.ToolUsageDetails.CodeExecutionCalls != want || + d.ImageCount == nil || *d.ImageCount != want { + t.Fatalf("incremental counts: tools=%+v images=%v, want %d", d.ToolUsageDetails, d.ImageCount, want) + } + } + observe(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12,"toolUsePromptTokenCount":1}}`) + d, _ := b.Detail() + if d.TotalTokens != 12 || *d.ToolUsageDetails.ToolCallsRequested != 2 || + *d.ToolUsageDetails.CodeExecutionCalls != 2 || *d.ImageCount != 2 || + d.WebSearchRequests == nil || *d.WebSearchRequests != 2 || + d.ToolUsageDetails.Extra["tool_use_prompt_tokens"] != 1 { + t.Fatalf("terminal usage lost counts or double-counted grounding: %+v tools=%+v", d, d.ToolUsageDetails) + } + }) + } +} + +func TestStripUsageMetadataPreservesToolArguments(t *testing.T) { + for _, wrapped := range []bool{false, true} { + t.Run(fmt.Sprintf("wrapped=%v", wrapped), func(t *testing.T) { + raw := `{"usageMetadata":{"promptTokenCount":5},"candidates":[{"content":{"parts":[{"functionCall":{"name":"save","args":{"usageMetadata":{"value":1}}}}]}}]}` + prefix := "" + if wrapped { + raw = `{"response":` + raw + `}` + prefix = "response." + } + out, changed := StripUsageMetadataFromJSON([]byte(raw)) + if !changed || !gjson.ValidBytes(out) || gjson.GetBytes(out, prefix+"usageMetadata").Exists() || + gjson.GetBytes(out, prefix+"cpaUsageMetadata.promptTokenCount").Int() != 5 || + gjson.GetBytes(out, prefix+"candidates.0.content.parts.0.functionCall.args.usageMetadata.value").Int() != 1 { + t.Fatalf("metadata rename changed the wrong field: %s", out) + } + }) + } +} + +func TestCodexUsageOwnsImageMetadataStrings(t *testing.T) { + for _, imageTool := range []bool{false, true} { + t.Run(fmt.Sprintf("image_tool=%v", imageTool), func(t *testing.T) { + raw := []byte(`{"response":{"usage":{"input_tokens":1,"output_tokens":1,"server_tool_use":{"custom_tool_requests":3}},"prompt_cache_retention":"24h","tool_usage":{"image_gen":{"input_tokens":1,"output_tokens":2}},"output":[{"type":"image_generation_call","status":"completed","size":"1024x1024","quality":"high"}]}}`) + var detailJSON []byte + if !imageTool { + detailJSON = bytes.ReplaceAll(raw, []byte(`"image_gen"`), []byte(`"other_tool"`)) + } else { + detailJSON = raw + } + d, ok := ParseCodexUsage(detailJSON) + if imageTool { + d, ok = ParseCodexImageToolUsage(detailJSON) + } + for i := range detailJSON { + detailJSON[i] = 'X' + } + if !ok || d.ImageDetails == nil || d.ImageDetails.Size != "1024x1024" || d.ImageDetails.Quality != "high" { + t.Fatalf("image metadata aliases input buffer: %+v", d.ImageDetails) + } + if !imageTool && (d.CacheDetails == nil || d.CacheDetails.Retention != "24h" || + d.ToolUsageDetails == nil || d.ToolUsageDetails.Extra["custom_tool_requests"] != 3) { + t.Fatalf("cache metadata or tool keys alias input buffer: %+v", d) + } + }) + } +} + +func TestResponsesCacheWriteLifetimesAreReportedNotInferred(t *testing.T) { + tests := []struct { + name, body string + short, long *int64 + }{ + {name: "mixed", body: `{"response":{"prompt_cache_options":{"mode":"explicit","ttl":"12h"},"usage":{"input_tokens":100,"output_tokens":1,"total_tokens":101,"input_tokens_details":{"cache_write_tokens":70,"cache_write_tokens_details":{"ttl_30m_tokens":30,"ttl_12h_tokens":40}}}}}`, short: usage.Int64Ptr(30), long: usage.Int64Ptr(40)}, + {name: "explicit zero", body: `{"response":{"usage":{"input_tokens":100,"output_tokens":1,"input_tokens_details":{"cache_write_tokens":70,"cache_write_tokens_details":{"ttl_30m_tokens":0,"ttl_12h_tokens":70}}}}}`, short: usage.Int64Ptr(0), long: usage.Int64Ptr(70)}, + {name: "ttl does not imply bucket", body: `{"response":{"prompt_cache_options":{"mode":"explicit","ttl":"12h"},"usage":{"input_tokens":100,"output_tokens":1,"input_tokens_details":{"cache_write_tokens":70}}}}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d, ok := ParseCodexUsage([]byte(tt.body)) + if !ok { + t.Fatal("missing usage") + } + check := func(name string, got, want *int64) { + t.Helper() + if (got == nil) != (want == nil) || (got != nil && *got != *want) { + t.Errorf("%s got %v want %v", name, got, want) + } + } + check("30m", d.CacheCreation30mTokens, tt.short) + check("12h", d.CacheCreation12hTokens, tt.long) + if d.CacheCreationTokens != 70 { + t.Fatalf("cache total=%d", d.CacheCreationTokens) + } + }) + } +} + +func TestOpenAIStreamToolCallsAcrossChoicesAndIncrementalChunks(t *testing.T) { + var b StreamUsageBuffer + b.ObserveOpenAIStream([]byte(`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"a"}]}}]}`)) + b.ObserveOpenAIStream([]byte(`data: {"choices":[{"index":1,"delta":{"tool_calls":[{"index":0,"id":"b"}]}}]}`)) + d, ok := b.Detail() + if !ok || d.ToolUsageDetails == nil || d.ToolUsageDetails.ToolCallsRequested == nil || *d.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("partial tool calls = %+v, want 2", d.ToolUsageDetails) + } + b.ObserveOpenAIStream([]byte(`data: {"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`)) + d, _ = b.Detail() + if d.ToolUsageDetails == nil || d.ToolUsageDetails.ToolCallsRequested == nil || *d.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("terminal tool calls = %+v, want 2", d.ToolUsageDetails) + } + + var inc StreamUsageBuffer + inc.ObserveOpenAIStream([]byte(`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0}]}}]}`)) + inc.ObserveOpenAIStream([]byte(`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1}]}}]}`)) + dInc, _ := inc.Detail() + if dInc.ToolUsageDetails == nil || dInc.ToolUsageDetails.ToolCallsRequested == nil || *dInc.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("incremental tool calls = %+v, want 2", dInc.ToolUsageDetails) + } +} + +func TestOpenAIStreamMultipleImageCompletionsWithoutUsage(t *testing.T) { + var b StreamUsageBuffer + b.ObserveOpenAIStream([]byte(`data: {"type":"image_generation.completed","b64_json":"AA==","size":"1024x1024"}`)) + b.ObserveOpenAIStream([]byte(`data: {"type":"image_generation.completed","b64_json":"BB==","size":"1024x1536"}`)) + d, ok := b.Detail() + if !ok || d.ImageCount == nil || *d.ImageCount != 2 || d.ImageDetails == nil || d.ImageDetails.Size != "1024x1536" { + t.Fatalf("image stream detail = %+v details=%+v, want ImageCount=2 and latest size", d, d.ImageDetails) + } +} + +func TestGeminiPartialArgsDeduplicationAndScannerReuse(t *testing.T) { + combined := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","partialArgs":[{"jsonPath":"$.city","stringValue":"Par","willContinue":true}],"willContinue":true}},{"functionCall":{"partialArgs":[{"jsonPath":"$.city","stringValue":"is"}],"willContinue":true}},{"functionCall":{}},{"functionCall":{"name":"get_weather","partialArgs":[{"jsonPath":"$.city","stringValue":"Lyon"}],"willContinue":false}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}`) + d := ParseGeminiUsage(combined) + if d.ToolUsageDetails == nil || d.ToolUsageDetails.ToolCallsRequested == nil || *d.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("direct ParseGeminiUsage ToolCallsRequested = %+v, want 2", d.ToolUsageDetails) + } + + var b StreamUsageBuffer + for _, chunk := range []string{ + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","partialArgs":[{"jsonPath":"$.city","stringValue":"Tok","willContinue":true}],"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"partialArgs":[{"jsonPath":"$.city","stringValue":"yo","willContinue":true}],"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"partialArgs":[{"jsonPath":"$.city"},{"jsonPath":"$.days","numberValue":0}],"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"name":"get_weather","partialArgs":[{"jsonPath":"$.city","stringValue":"Os","willContinue":true}],"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"partialArgs":[{"jsonPath":"$.city","stringValue":"aka"}]}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}`, + } { + b.ObserveGeminiStream([]byte(chunk)) + } + dStream, ok := b.Detail() + if !ok || dStream.ToolUsageDetails == nil || dStream.ToolUsageDetails.ToolCallsRequested == nil || *dStream.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("stream ToolCallsRequested = %+v, want 2", dStream.ToolUsageDetails) + } + + first := `data: {"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"name":"weather","partialArgs":[{"jsonPath":"$.city","stringValue":"Paris"}],"willContinue":true}}]}}]}` + filler := `data: {"candidates":[{"content":{"parts":[{"text":"` + strings.Repeat("X", 120) + `"}]}}]}` + last := `data: {"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"name":"weather","partialArgs":[{"jsonPath":"$.city","stringValue":"London"}],"willContinue":false}}]}}]}` + scanner := bufio.NewScanner(strings.NewReader(first + "\n" + strings.Repeat(filler+"\n", 100) + last + "\n")) + var scannerBuf StreamUsageBuffer + for scanner.Scan() { + scannerBuf.ObserveGeminiStream(scanner.Bytes()) + } + if err := scanner.Err(); err != nil { + t.Fatal(err) + } + dScanner, ok := scannerBuf.Detail() + if !ok || dScanner.ToolUsageDetails == nil || dScanner.ToolUsageDetails.ToolCallsRequested == nil || *dScanner.ToolUsageDetails.ToolCallsRequested != 2 { + t.Fatalf("scanner reuse ToolCallsRequested = %+v, want 2", dScanner.ToolUsageDetails) + } +} + +func TestGeminiStreamUsageMergePreservesEarlierModalitiesAndTier(t *testing.T) { + for _, wrapped := range []bool{false, true} { + t.Run(map[bool]string{false: "gemini", true: "antigravity"}[wrapped], func(t *testing.T) { + var b StreamUsageBuffer + observe := func(s string) { + if wrapped { + b.ObserveAntigravityStream([]byte(`data: {"response":` + s + `}`)) + } else { + b.ObserveGeminiStream([]byte("data: " + s)) + } + } + observe(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1,"totalTokenCount":11,"trafficType":"ON_DEMAND","promptTokensDetails":[{"modality":"IMAGE","tokenCount":10}]}}`) + observe(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`) + d, _ := b.Detail() + if d.ImageInputTokens == nil || *d.ImageInputTokens != 10 || d.UpstreamServiceTier != "ON_DEMAND" { + t.Fatalf("lost structured usage: %+v", d) + } + }) + } +} + +func TestCodexCustomToolCallWithoutXSearchPrefixCountsAsRequestedToolCall(t *testing.T) { + d, ok := ParseCodexUsage([]byte(`{"response":{"usage":{"input_tokens":1,"output_tokens":1},"output":[{"type":"custom_tool_call","status":"completed","name":"x_keyword_search","call_id":"call_local_1"}]}}`)) + if !ok || d.ToolUsageDetails == nil || d.ToolUsageDetails.ToolCallsRequested == nil || *d.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("misclassified client custom_tool_call: %+v", d.ToolUsageDetails) + } +} + +func TestClaudeStreamProvisionalZeroWebSearchDoesNotFreezeCompletedSearches(t *testing.T) { + for _, initial := range []string{"", `,"server_tool_use":{"web_search_requests":0}`} { + t.Run(fmt.Sprintf("initial_zero=%v", initial != ""), func(t *testing.T) { + var b StreamUsageBuffer + b.ObserveClaudeStream([]byte(fmt.Sprintf(`data: {"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1%s}}}`, initial))) + for i := 1; i <= 2; i++ { + b.ObserveClaudeStream([]byte(fmt.Sprintf(`data: {"type":"content_block_start","content_block":{"type":"web_search_tool_result","tool_use_id":"tool_%d","content":[]}}`, i))) + } + b.ObserveClaudeStream([]byte(`data: {"type":"message_delta","usage":{"output_tokens":20}}`)) + d, ok := b.Detail() + if !ok || d.WebSearchRequests == nil || *d.WebSearchRequests != 2 || + d.ToolUsageDetails == nil || d.ToolUsageDetails.WebSearchCalls == nil || *d.ToolUsageDetails.WebSearchCalls != 2 { + t.Fatalf("expected 2 completed searches, got detail=%+v tools=%+v", d, d.ToolUsageDetails) + } + }) + } +} + +func TestOpenAIStreamPreservesCacheTTLBucketsAndStructuredOnlyUsage(t *testing.T) { + var cacheBuf StreamUsageBuffer + cacheBuf.ObserveOpenAIStream([]byte(`data: {"usage":{"prompt_tokens":100,"completion_tokens":1,"total_tokens":101,"prompt_tokens_details":{"cache_write_tokens":70,"cache_write_tokens_details":{"ttl_30m_tokens":30,"ttl_12h_tokens":40}}}}`)) + cacheDetail, ok := cacheBuf.Detail() + if !ok || cacheDetail.CacheCreation30mTokens == nil || *cacheDetail.CacheCreation30mTokens != 30 || + cacheDetail.CacheCreation12hTokens == nil || *cacheDetail.CacheCreation12hTokens != 40 { + t.Fatalf("stream lost cache TTL buckets: %+v", cacheDetail) + } + + var costBuf StreamUsageBuffer + costBuf.ObserveOpenAIStream([]byte(`data: {"usage":{"cost":0.25,"cost_in_usd_ticks":2500000000,"server_tool_use":{"web_search_requests":2}}}`)) + costDetail, ok := costBuf.Detail() + if !ok || costDetail.CostDetails == nil || costDetail.CostDetails.Cost == nil || *costDetail.CostDetails.Cost != 0.25 || + costDetail.WebSearchRequests == nil || *costDetail.WebSearchRequests != 2 { + t.Fatalf("stream lost structured-only usage: %+v", costDetail) + } +} + +func TestObserveResponsesStreamPreservesOutputItemDoneToolUsageWhenCompletedOutputEmpty(t *testing.T) { + var b StreamUsageBuffer + b.ObserveResponsesStream([]byte(`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"ct_1","type":"custom_tool_call","status":"completed","call_id":"xs_1","name":"x_keyword_search","input":"{}"}}`)) + b.ObserveResponsesStream([]byte(`data: {"type":"response.output_item.done","output_index":1,"item":{"id":"fc_1","type":"function_call","status":"completed","call_id":"call_1","name":"lookup","arguments":"{}"}}`)) + b.ObserveResponsesStream([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","output":[],"usage":{"input_tokens":20,"output_tokens":10,"total_tokens":30}}}`)) + d, ok := b.Detail() + if !ok || d.ToolUsageDetails == nil || + d.ToolUsageDetails.XSearchCalls == nil || *d.ToolUsageDetails.XSearchCalls != 1 || + d.ToolUsageDetails.ToolCallsRequested == nil || *d.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("ObserveResponsesStream lost output_item.done tool usage: %+v", d.ToolUsageDetails) + } + + var pluginBuf StreamUsageBuffer + ObservePluginExecutorStreamUsage("openai-response", []byte(strings.Join([]string{ + `data: {"type":"response.output_item.done","output_index":0,"item":{"id":"fc_1","type":"function_call","status":"completed","call_id":"call_1","name":"lookup","arguments":"{}"}}`, + "", + `data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}`, + }, "\n")), &pluginBuf) + pd, pok := pluginBuf.Detail() + if !pok || pd.ToolUsageDetails == nil || pd.ToolUsageDetails.ToolCallsRequested == nil || *pd.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("ObservePluginExecutorStreamUsage lost output_item.done tool usage: %+v", pd.ToolUsageDetails) + } +} diff --git a/internal/runtime/executor/meta_executor_execute.go b/internal/runtime/executor/meta_executor_execute.go index d8af8539673..d3b1ed75e74 100644 --- a/internal/runtime/executor/meta_executor_execute.go +++ b/internal/runtime/executor/meta_executor_execute.go @@ -214,9 +214,7 @@ func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxye } filteringSubscriptionUsage = false } - if detail, ok := helps.ParseCodexUsage(eventData); ok { - upstreamUsage.Observe(detail, true) - } + upstreamUsage.ObserveResponsesStream(eventData) if errEvent := metaStreamEventError(eventData); errEvent != nil { return metaCompletedTranslation{}, errEvent } @@ -232,6 +230,7 @@ func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxye xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) case "response.completed", "response.incomplete": completedData := patchCodexCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + upstreamUsage.ObserveResponsesStream(completedData) var param any out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, ¶m) if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 { diff --git a/internal/runtime/executor/meta_executor_stream.go b/internal/runtime/executor/meta_executor_stream.go index 904f30650a2..e5f1255aa86 100644 --- a/internal/runtime/executor/meta_executor_stream.go +++ b/internal/runtime/executor/meta_executor_stream.go @@ -213,7 +213,7 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut case "response.output_item.done": xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) case "response.completed", "response.incomplete": - if detail, ok := helps.ParseCodexUsage(eventData); ok { + if detail, ok := helps.ParseCodexUsageWithOutputItems(eventData, outputItemsByIndex, outputItemsFallback); ok { streamUsage.Observe(detail, true) } eventData = patchCodexCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index d270f621da1..b9caea42209 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -330,7 +330,7 @@ func (e *OpenAICompatExecutor) executeImages(ctx context.Context, auth *cliproxy } reporter.ObserveResponseModel(body) - reporter.Publish(ctx, helps.ParseOpenAIUsage(body)) + reporter.Publish(ctx, helps.ParseOpenAIImageUsage(payload, body, contentType)) reporter.EnsurePublished(ctx) resp = cliproxyexecutor.Response{Payload: body, Headers: httpResp.Header.Clone()} return resp, nil @@ -483,7 +483,9 @@ func (e *OpenAICompatExecutor) executeStreamWithCredentials(ctx context.Context, loggedErr = statusErr{code: streamErr.code, msg: "upstream stream returned an error payload"} } helps.RecordAPIResponseError(ctx, e.cfg, loggedErr) - reporter.PublishFailure(ctx, loggedErr) + if !streamUsage.PublishFailure(ctx, reporter, loggedErr) { + reporter.PublishFailure(ctx, loggedErr) + } select { case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: case <-ctx.Done(): @@ -533,6 +535,9 @@ func (e *OpenAICompatExecutor) executeStreamWithCredentials(ctx context.Context, } // Observe complete upstream frames before translation or downstream backpressure. + if !isDone && len(dataLines) > 1 { + streamUsage.ObserveOpenAIStream(dataPayload) + } helps.ObserveChatTokenEvent(reporter, dataPayload) streamLine := append([]byte("data: "), dataPayload...) chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, helps.ApplyPatchOriginalRequest(req, opts), translated, streamLine, ¶m, claudeInputTokens) @@ -612,7 +617,9 @@ func (e *OpenAICompatExecutor) executeStreamWithCredentials(ctx context.Context, } if errScan != nil && !seenDone { helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) + if !streamUsage.PublishFailure(ctx, reporter, errScan) { + reporter.PublishFailure(ctx, errScan) + } select { case out <- cliproxyexecutor.StreamChunk{Err: errScan}: case <-ctx.Done(): @@ -621,7 +628,7 @@ func (e *OpenAICompatExecutor) executeStreamWithCredentials(ctx context.Context, // Without a translator-confirmed terminal state, a clean Responses or Claude EOF // without [DONE] remains a failed stream instead of completing it. if responseFormat == sdktranslator.FormatOpenAIResponse || (responseFormat == sdktranslator.FormatClaude && !sawFinishReason) { - streamErr := statusErr{code: http.StatusBadGateway, msg: "upstream stream closed before [DONE]"} + streamErr := statusErr{code: http.StatusBadGateway, msg: "upstream stream closed before [DONE]", streamTruncated: true} helps.RecordAPIResponseError(ctx, e.cfg, streamErr) if !helps.RecordApplyPatchStreamFailureWithUsage(ctx, param, reporter, &streamUsage, streamErr) { if !streamUsage.PublishFailure(ctx, reporter, streamErr) { @@ -740,12 +747,17 @@ func (e *OpenAICompatExecutor) executeImagesStream(ctx context.Context, auth *cl out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) + var streamUsage helps.StreamUsageBuffer + streamUsage.SetRequestPayload(payload, contentType) + usageObserver := helps.NewOpenAIStreamUsageObserver(&streamUsage, streamScannerBuffer) observer := helps.NewStreamResponseModelObserver(reporter) defer func() { observer.Finish() if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("openai compat executor: close response body error: %v", errClose) } + usageObserver.Finish() + streamUsage.Publish(ctx, reporter) reporter.EnsurePublished(ctx) }() buffer := make([]byte, 32*1024) @@ -755,6 +767,7 @@ func (e *OpenAICompatExecutor) executeImagesStream(ctx context.Context, auth *cl chunk := bytes.Clone(buffer[:n]) helps.AppendAPIResponseChunk(ctx, e.cfg, chunk) observer.Feed(chunk) + usageObserver.Feed(chunk) select { case out <- cliproxyexecutor.StreamChunk{Payload: chunk}: case <-ctx.Done(): @@ -764,7 +777,11 @@ func (e *OpenAICompatExecutor) executeImagesStream(ctx context.Context, auth *cl if errRead != nil { if errRead != io.EOF { helps.RecordAPIResponseError(ctx, e.cfg, errRead) - reporter.PublishFailure(ctx, errRead) + observer.Finish() + usageObserver.Finish() + if !streamUsage.PublishFailure(ctx, reporter, errRead) { + reporter.PublishFailure(ctx, errRead) + } select { case out <- cliproxyexecutor.StreamChunk{Err: errRead}: case <-ctx.Done(): diff --git a/internal/runtime/executor/openai_compat_executor_compact_test.go b/internal/runtime/executor/openai_compat_executor_compact_test.go index ce5922f1be0..0f9ba8dd5a5 100644 --- a/internal/runtime/executor/openai_compat_executor_compact_test.go +++ b/internal/runtime/executor/openai_compat_executor_compact_test.go @@ -11,9 +11,13 @@ import ( "net/http/httptest" "net/textproto" "strings" + "sync/atomic" "testing" + "time" "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" @@ -954,6 +958,144 @@ func TestOpenAICompatExecutorResponsesStreamFailsOnEOFWithoutDone(t *testing.T) if !strings.Contains(streamErr.Error(), "closed before [DONE]") { t.Fatalf("stream error does not explain the missing terminal marker: %v", streamErr) } + truncatedErr, okTruncated := streamErr.(cliproxyexecutor.StreamTruncatedError) + if !okTruncated || !truncatedErr.IsStreamTruncated() { + t.Fatalf("stream error = %T (%v), want IsStreamTruncated() = true", streamErr, streamErr) + } +} + +func TestOpenAICompatMissingDoneDoesNotCooldownPool(t *testing.T) { + const provider = "openai-compatible-minimax" + const model = "MiniMax-M3" + const partialChunk = `data: {"id":"chatcmpl_1","object":"chat.completion.chunk","created":1773896263,"model":"MiniMax-M3","choices":[{"index":0,"delta":{"role":"assistant","content":"partial"},"finish_reason":null}]}` + "\n\n" + + cases := []struct { + name string + status int + body string + requests int + wantCalls int32 + wantCooled int + wantLastErrorText string + forbidAuthUnavail bool + }{ + { + name: "partial_stream_without_done_across_requests", + status: http.StatusOK, + body: partialChunk, + requests: 4, + wantCalls: 4, + wantCooled: 0, + wantLastErrorText: "upstream stream closed before [DONE]", + forbidAuthUnavail: true, + }, + { + name: "empty_200_stream_without_done", + status: http.StatusOK, + body: "", + requests: 1, + wantCalls: 1, + wantCooled: 0, + wantLastErrorText: "upstream stream closed before [DONE]", + forbidAuthUnavail: true, + }, + { + name: "real_http_502_control", + status: http.StatusBadGateway, + body: `{"error":{"message":"mock upstream gateway failure","type":"server_error"}}`, + requests: 2, + wantCalls: 3, + wantCooled: 3, + wantLastErrorText: "auth_unavailable: no auth available", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(tc.status) + _, _ = io.WriteString(w, tc.body) + })) + defer server.Close() + + cfg := &config.Config{} + manager := cliproxyauth.NewManager(nil, &cliproxyauth.RoundRobinSelector{}, nil) + manager.SetConfig(cfg) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(NewOpenAICompatExecutor(provider, cfg)) + + for i := 0; i < 3; i++ { + id := fmt.Sprintf("minimax-compat-%s-%d", tc.name, i) + registry.GetGlobalRegistry().RegisterClient(id, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) }) + _, errRegister := manager.Register(context.Background(), &cliproxyauth.Auth{ + ID: id, + Provider: provider, + Status: cliproxyauth.StatusActive, + Attributes: map[string]string{ + "api_key": fmt.Sprintf("minimax-key-%d", i), + "base_url": server.URL + "/v1", + "compat_name": "minimax", + "provider_key": provider, + }, + }) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + } + + requestPayload := []byte(`{"model":"MiniMax-M3","input":"hi","stream":true}`) + var lastObservedErr error + for reqIdx := 0; reqIdx < tc.requests; reqIdx++ { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + result, errStream := manager.ExecuteStream(ctx, []string{provider}, cliproxyexecutor.Request{ + Model: model, + Payload: requestPayload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: requestPayload, + Stream: true, + }) + observedErr := errStream + if result != nil { + for chunk := range result.Chunks { + if chunk.Err != nil { + observedErr = chunk.Err + } + } + } + cancel() + if observedErr == nil { + t.Fatalf("request %d succeeded unexpectedly", reqIdx+1) + } + lastObservedErr = observedErr + } + + cooled := 0 + for _, auth := range manager.List() { + if state := auth.ModelStates[model]; state != nil && state.Unavailable && state.NextRetryAfter.After(time.Now()) { + cooled++ + } + } + + if calls.Load() != tc.wantCalls { + t.Errorf("upstream calls = %d, want %d", calls.Load(), tc.wantCalls) + } + if cooled != tc.wantCooled { + t.Errorf("cooled credentials = %d, want %d", cooled, tc.wantCooled) + } + if lastObservedErr == nil || !strings.Contains(lastObservedErr.Error(), tc.wantLastErrorText) { + t.Errorf("last error = %v, want substring %q", lastObservedErr, tc.wantLastErrorText) + } + if tc.forbidAuthUnavail && lastObservedErr != nil && strings.Contains(lastObservedErr.Error(), "auth_unavailable") { + t.Errorf("last error unexpectedly degraded to auth_unavailable: %v", lastObservedErr) + } + }) + } } func TestOpenAICompatExecutorResponsesStreamPreservesUpstreamDataError(t *testing.T) { diff --git a/internal/runtime/executor/openai_compat_executor_images_test.go b/internal/runtime/executor/openai_compat_executor_images_test.go index 9fbb4730f5f..e4ae11448d7 100644 --- a/internal/runtime/executor/openai_compat_executor_images_test.go +++ b/internal/runtime/executor/openai_compat_executor_images_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "fmt" + "mime/multipart" "net/http" "net/http/httptest" "strings" @@ -184,3 +185,321 @@ func TestOpenAICompatExecutor_ImageEndpointPath_HonorsOverriddenRequestPath_Issu t.Fatalf("upstream path = %q, want /images/generations", requestedPath) } } + +func TestOpenAICompatImageStreamAccountsLargeCompletedEvents(t *testing.T) { + for _, size := range []int{16, 100000} { + t.Run(fmt.Sprint(size), func(t *testing.T) { + event := `data: {"type":"image_generation.completed","model":"gpt-image-2","b64_json":"` + strings.Repeat("A", size) + `","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}` + "\n\n" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + if _, err := fmt.Fprint(w, event); err != nil { + t.Errorf("write upstream event: %v", err) + } + })) + defer server.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + exec := NewOpenAICompatExecutor("openai-compatibility", &config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test"}, + } + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + res, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-image-2", + Payload: []byte(`{"model":"gpt-image-2","prompt":"test"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-image"), + Stream: true, + }) + if err != nil { + t.Fatal(err) + } + var got bytes.Buffer + for chunk := range res.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + got.Write(chunk.Payload) + } + if got.String() != event { + t.Fatal("forwarded image SSE bytes changed") + } + record := capture.await(t) + if record.Detail.TotalTokens != 30 || record.Detail.ImageCount == nil || *record.Detail.ImageCount != 1 { + t.Fatalf("image completion usage missing: %+v", record.Detail) + } + }) + } +} + +func TestOpenAICompatAndCodexImageStreamReadFailurePreservesObservedUsage(t *testing.T) { + for _, provider := range []string{"codex", "compat"} { + t.Run(provider, func(t *testing.T) { + event := `data: {"type":"image_generation.completed","model":"gpt-image-2","b64_json":"AA==","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}` + "\n\n" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Content-Length", fmt.Sprint(len(event)+100)) + if _, err := fmt.Fprint(w, event); err != nil { + t.Errorf("write upstream event: %v", err) + } + })) + defer server.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + req := cliproxyexecutor.Request{ + Model: "gpt-image-2", + Payload: []byte(`{"model":"gpt-image-2","prompt":"test"}`), + } + var ( + res *cliproxyexecutor.StreamResult + err error + ) + if provider == "codex" { + res, err = NewCodexExecutor(&config.Config{}).ExecuteStream(ctx, newCodexOpenAIImageTestAuth(server.URL), req, codexOpenAIImageTestOptions(codexImagesGenerationsPath, true)) + } else { + res, err = NewOpenAICompatExecutor("openai-compatibility", &config.Config{}).ExecuteStream(ctx, &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test"}, + }, req, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-image"), + Stream: true, + }) + } + if err != nil { + t.Fatal(err) + } + var got bytes.Buffer + var streamErr error + for chunk := range res.Chunks { + got.Write(chunk.Payload) + if chunk.Err != nil { + streamErr = chunk.Err + } + } + if got.String() != event || streamErr == nil { + t.Fatalf("unexpected stream output=%q err=%v", got.String(), streamErr) + } + record := capture.await(t) + if !record.Failed || record.Detail.TotalTokens != 30 || record.Detail.ImageCount == nil || *record.Detail.ImageCount != 1 { + t.Fatalf("read failure lost observed image usage: failed=%v detail=%+v", record.Failed, record.Detail) + } + }) + } +} + +func TestOpenAICompatMultipartImageEditExtractsRequestedImageDetails(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%v", stream), func(t *testing.T) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + for k, v := range map[string]string{ + "model": "gpt-image-2", + "prompt": "test", + "size": "1024x1536", + "quality": "high", + } { + if err := writer.WriteField(k, v); err != nil { + t.Fatal(err) + } + } + part, err := writer.CreateFormFile("image", "image.png") + if err != nil { + t.Fatal(err) + } + if _, err = part.Write([]byte("PNG")); err != nil { + t.Fatal(err) + } + if err = writer.Close(); err != nil { + t.Fatal(err) + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if errParse := r.ParseMultipartForm(1 << 20); errParse != nil { + t.Errorf("parse multipart: %v", errParse) + return + } + defer func() { + if errRemove := r.MultipartForm.RemoveAll(); errRemove != nil { + t.Errorf("remove multipart form: %v", errRemove) + } + }() + if stream { + w.Header().Set("Content-Type", "text/event-stream") + if _, errWrite := fmt.Fprint(w, `data: {"type":"image_edit.completed","b64_json":"AA==","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}`+"\n\n"); errWrite != nil { + t.Errorf("write stream: %v", errWrite) + } + return + } + w.Header().Set("Content-Type", "application/json") + if _, errWrite := fmt.Fprint(w, `{"data":[{"b64_json":"AA=="}],"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}`); errWrite != nil { + t.Errorf("write json: %v", errWrite) + } + })) + defer server.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + exec := NewOpenAICompatExecutor("openai-compatibility", &config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{"base_url": server.URL, "api_key": "test"}, + } + req := cliproxyexecutor.Request{Model: "gpt-image-2", Payload: body.Bytes()} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-image"), + Stream: stream, + Headers: http.Header{"Content-Type": []string{writer.FormDataContentType()}}, + Metadata: map[string]any{cliproxyexecutor.RequestPathMetadataKey: "/v1/images/edits"}, + } + if stream { + res, errExec := exec.ExecuteStream(ctx, auth, req, opts) + if errExec != nil { + t.Fatal(errExec) + } + for chunk := range res.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + } else { + if _, errExec := exec.Execute(ctx, auth, req, opts); errExec != nil { + t.Fatal(errExec) + } + } + record := capture.await(t) + if record.Detail.TotalTokens != 30 || record.Detail.RequestedImageDetails == nil || + record.Detail.RequestedImageDetails.Size != "1024x1536" || + record.Detail.RequestedImageDetails.Quality != "high" { + t.Fatalf("unexpected multipart image usage: %+v requested=%+v", record.Detail, record.Detail.RequestedImageDetails) + } + }) + } +} + +func TestOpenAICompatMultilineSSEAndXAIOutputItemDoneUsage(t *testing.T) { + t.Run("openai_compat_multiline_sse_usage", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],"usage":`, + `data: {"prompt_tokens":12,"completion_tokens":8,"total_tokens":20}}`, + "", + "data: [DONE]", + "", + }, "\n")) + })) + defer srv.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + exec := NewOpenAICompatExecutor("openai-compatibility", &config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "openai-compatibility", + Attributes: map[string]string{"base_url": srv.URL, "api_key": "test"}, + } + res, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-4o", + Payload: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"stream":true}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Stream: true, + }) + if err != nil { + t.Fatal(err) + } + for range res.Chunks { + } + rec := capture.await(t) + if rec.Detail.TotalTokens != 20 { + t.Fatalf("multiline SSE lost usage: %+v", rec.Detail) + } + }) + + t.Run("xai_output_item_done_with_empty_completed_output", func(t *testing.T) { + sseBody := strings.Join([]string{ + "event: response.output_item.done", + `data: {"type":"response.output_item.done","output_index":0,"item":{"id":"ct_1","type":"custom_tool_call","status":"completed","call_id":"xs_1","name":"x_keyword_search","input":"{}"}}`, + "", + "event: response.output_item.done", + `data: {"type":"response.output_item.done","output_index":1,"item":{"id":"fc_1","type":"function_call","status":"completed","call_id":"call_1","name":"lookup","arguments":"{}"}}`, + "", + "event: response.completed", + `data: {"type":"response.completed","response":{"id":"resp_1","model":"grok-4","status":"completed","output":[],"usage":{"input_tokens":20,"output_tokens":10,"total_tokens":30}}}`, + "", + }, "\n") + + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%v", stream), func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, sseBody) + })) + defer srv.Close() + + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan coreusage.Record, 4)} + coreusage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) + }) + ctx := coreusage.WithRequestedModelAlias(context.Background(), alias) + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": srv.URL, "api_key": "xai-test"}, + } + req := cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","input":"hi","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Stream: stream, + } + if stream { + res, err := exec.ExecuteStream(ctx, auth, req, opts) + if err != nil { + t.Fatal(err) + } + for chunk := range res.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + } else { + if _, err := exec.Execute(ctx, auth, req, opts); err != nil { + t.Fatal(err) + } + } + rec := capture.await(t) + if rec.Detail.ToolUsageDetails == nil || + rec.Detail.ToolUsageDetails.XSearchCalls == nil || *rec.Detail.ToolUsageDetails.XSearchCalls != 1 || + rec.Detail.ToolUsageDetails.ToolCallsRequested == nil || *rec.Detail.ToolUsageDetails.ToolCallsRequested != 1 { + t.Fatalf("lost output_item.done tool usage: %+v", rec.Detail.ToolUsageDetails) + } + }) + } + }) +} diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go index f92d436b731..2b6078738da 100644 --- a/internal/runtime/executor/xai_executor_execute.go +++ b/internal/runtime/executor/xai_executor_execute.go @@ -97,9 +97,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req continue } eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) - if detail, ok := helps.ParseCodexUsage(eventData); ok { - upstreamUsage.Observe(detail, true) - } + upstreamUsage.ObserveResponsesStream(eventData) prepared.applyPatch.RememberDispatcherEvent(eventData) eventData = namespaceRestorer.restore(eventData) if prepared.webSearchAlias != "" { @@ -136,8 +134,10 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req upstreamUsage.PublishFailure(ctx, reporter, errTranslation) return cliproxyexecutor.Response{}, errTranslation } - if detail, ok := helps.ParseCodexUsage(eventData); ok { - reporter.Publish(ctx, detail) + if !upstreamUsage.Publish(ctx, reporter) { + if detail, ok := helps.ParseCodexUsageWithOutputItems(completedData, outputItemsByIndex, outputItemsFallback); ok { + reporter.Publish(ctx, detail) + } } if prepared.responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) diff --git a/internal/runtime/executor/xai_executor_media.go b/internal/runtime/executor/xai_executor_media.go index 0d00c515949..f6a8e7ac3cc 100644 --- a/internal/runtime/executor/xai_executor_media.go +++ b/internal/runtime/executor/xai_executor_media.go @@ -74,7 +74,7 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth } reporter.ObserveResponseModel(data) - reporter.EnsurePublished(ctx) + reporter.Publish(ctx, helps.ParseXAIImageUsage(payload, data)) return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil } @@ -153,6 +153,6 @@ func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth } reporter.ObserveResponseModel(data) - reporter.EnsurePublished(ctx) + reporter.Publish(ctx, helps.ParseXAIVideoUsage(payload, data)) return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil } diff --git a/internal/runtime/executor/xai_executor_speech.go b/internal/runtime/executor/xai_executor_speech.go index 58c60d00a74..3b51a0b658d 100644 --- a/internal/runtime/executor/xai_executor_speech.go +++ b/internal/runtime/executor/xai_executor_speech.go @@ -129,6 +129,6 @@ func (e *XAIExecutor) executeSpeech(ctx context.Context, auth *cliproxyauth.Auth return resp, err } - reporter.EnsurePublished(ctx) + reporter.Publish(ctx, helps.ParseOpenAIUsage(data)) return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil } diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go index 715fbcf5b6f..e4d7148ad44 100644 --- a/internal/runtime/executor/xai_executor_stream.go +++ b/internal/runtime/executor/xai_executor_stream.go @@ -155,6 +155,9 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):])) hasPendingEventLine := pendingEventLine != nil for i, eventData := range eventDataList { + if rawEventType := gjson.GetBytes(eventData, "type").String(); rawEventType == "response.output_item.done" || rawEventType == "response.completed" || rawEventType == "response.incomplete" || rawEventType == "response.done" { + streamUsage.ObserveResponsesStream(eventData) + } prepared.applyPatch.RememberDispatcherEvent(eventData) eventData = namespaceRestorer.restore(eventData) if prepared.webSearchAlias != "" { @@ -169,11 +172,6 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth } reporter.ObserveResponseModel(eventData) normalizedEventName := gjson.GetBytes(eventData, "type").String() - if normalizedEventName == "response.completed" || normalizedEventName == "response.incomplete" { - if detail, ok := helps.ParseCodexUsage(eventData); ok { - streamUsage.Observe(detail, true) - } - } if hasPendingEventLine { eventLine := []byte("event: " + normalizedEventName) diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 886b5283302..1c14c2323e6 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -3608,8 +3608,8 @@ func TestXAIExecutorExecuteImagesUsesImagesEndpointAndPublishesUsage(t *testing. if record.Failed { t.Fatalf("failed = true, want false; failure=%+v", record.Fail) } - if record.Detail != (usage.Detail{}) { - t.Fatalf("detail = %+v, want zero token usage", record.Detail) + if record.Detail.TotalTokens != 0 || record.Detail.ImageCount == nil || *record.Detail.ImageCount != 1 { + t.Fatalf("detail = %+v, want zero token usage and ImageCount=1", record.Detail) } assertXAIUsageRecordTTFT(t, record.TTFT) assertNoAdditionalXAIUsageRecord(t, plugin.records) @@ -4048,8 +4048,8 @@ func TestXAIExecutorExecuteVideosCreate(t *testing.T) { if record.Failed { t.Fatalf("failed = true, want false; failure=%+v", record.Fail) } - if record.Detail != (usage.Detail{}) { - t.Fatalf("detail = %+v, want zero token usage", record.Detail) + if record.Detail.TotalTokens != 0 || record.Detail.ImageDetails != nil || record.Detail.MediaCounts != nil { + t.Fatalf("detail = %+v, want zero token usage on async video creation", record.Detail) } assertXAIUsageRecordTTFT(t, record.TTFT) assertNoAdditionalXAIUsageRecord(t, plugin.records) @@ -7935,3 +7935,52 @@ func TestApplyPatchResponsesExecutorLocalFailures(t *testing.T) { } } } + +func TestXAIWebsocketUsagePreservesFilteredSearch(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close: %v", errClose) + } + }() + if _, _, err = conn.ReadMessage(); err != nil { + t.Errorf("read: %v", err) + return + } + event := `{"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[{"id":"ctc_1","type":"custom_tool_call","call_id":"xs_call-1","name":"x_keyword_search","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + if err = conn.WriteMessage(websocket.TextMessage, []byte(event)); err != nil { + t.Errorf("write: %v", err) + } + })) + defer server.Close() + alias := t.Name() + capture := &multiProviderUsageCapture{alias: alias, records: make(chan usage.Record, 4)} + usage.RegisterNamedPlugin(t.Name(), capture) + t.Cleanup(func() { usage.RegisterNamedPlugin(t.Name(), multiProviderNoopUsagePlugin{}) }) + ctx := usage.WithRequestedModelAlias(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), alias) + exec := NewXAIWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Provider: "xai", Attributes: map[string]string{"base_url": server.URL, "websockets": "true"}, Metadata: map[string]any{"access_token": "test"}} + result, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "grok-4.5", Payload: []byte(`{"model":"grok-4.5","input":"search","tools":[{"type":"x_search"}]}`)}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Stream: true}) + if err != nil { + t.Fatal(err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + select { + case record := <-capture.records: + if record.Detail.TotalTokens != 2 || record.Detail.ToolUsageDetails == nil || record.Detail.ToolUsageDetails.XSearchCalls == nil || *record.Detail.ToolUsageDetails.XSearchCalls != 1 { + t.Fatalf("lost raw upstream usage: %+v tools=%+v", record.Detail, record.Detail.ToolUsageDetails) + } + case <-time.After(3 * time.Second): + t.Fatal("usage not published") + } +} diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index cca1478d6f7..276ba250198 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -808,6 +808,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } for _, payload := range xaiNormalizeReasoningSummaryDataEvents(payload) { + if rawEventType := gjson.GetBytes(payload, "type").String(); rawEventType == "response.output_item.done" || rawEventType == "response.completed" || rawEventType == "response.incomplete" || rawEventType == "response.done" { + streamUsage.ObserveResponsesStream(payload) + } prepared.applyPatch.RememberDispatcherEvent(payload) payload = namespaceRestorer.restore(payload) if prepared.webSearchAlias != "" { @@ -852,9 +855,6 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox xaiCollectOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) case "response.completed": logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload) - if detail, ok := helps.ParseCodexUsage(payload); ok { - streamUsage.Observe(detail, true) - } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) @@ -864,9 +864,6 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } case "response.done": logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload) - if detail, ok := helps.ParseCodexUsage(payload); ok { - streamUsage.Observe(detail, true) - } if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true diff --git a/internal/translator/antigravity/claude/antigravity_claude_request.go b/internal/translator/antigravity/claude/antigravity_claude_request.go index 57352bbe997..5d328469453 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_request.go +++ b/internal/translator/antigravity/claude/antigravity_claude_request.go @@ -657,10 +657,15 @@ func convertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _ functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "id", toolCallID) functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "name", util.MapSanitizedFunctionName(functionNameMap, funcName)) + responsePath := "response.result" + if contentResult.Get("is_error").Type == gjson.True { + responsePath = "response.error" + } + responseData := "" if functionResponseResult.Type == gjson.String { responseData = functionResponseResult.String() - functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "response.result", responseData) + functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, responsePath, responseData) } else if functionResponseResult.IsArray() { frResults := functionResponseResult.Array() nonImageItems := make([][]byte, 0, len(frResults)) @@ -685,11 +690,11 @@ func convertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _ } if len(nonImageItems) == 1 { - functionResponseJSON = translatorcommon.SetGeminiFunctionResponseRaw(functionResponseJSON, "response.result", string(nonImageItems[0])) + functionResponseJSON = translatorcommon.SetGeminiFunctionResponseRaw(functionResponseJSON, responsePath, string(nonImageItems[0])) } else if len(nonImageItems) > 1 { - functionResponseJSON = translatorcommon.SetGeminiFunctionResponseRaw(functionResponseJSON, "response.result", string(translatorcommon.JoinRawArray(nonImageItems))) + functionResponseJSON = translatorcommon.SetGeminiFunctionResponseRaw(functionResponseJSON, responsePath, string(translatorcommon.JoinRawArray(nonImageItems))) } else { - functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "response.result", "") + functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, responsePath, "") } // Place image data inside functionResponse.parts as inlineData @@ -712,16 +717,16 @@ func convertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _ imagePartJSON := []byte(`{}`) imagePartJSON, _ = sjson.SetRawBytes(imagePartJSON, "inlineData", inlineDataJSON) functionResponseJSON, _ = sjson.SetRawBytes(functionResponseJSON, "parts", translatorcommon.JoinRawArray([][]byte{imagePartJSON})) - functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "response.result", "") + functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, responsePath, "") } else { - functionResponseJSON = translatorcommon.SetGeminiFunctionResponseResult(functionResponseJSON, "response.result", functionResponseResult) + functionResponseJSON = translatorcommon.SetGeminiFunctionResponseResult(functionResponseJSON, responsePath, functionResponseResult) } } else if functionResponseResult.Raw != "" { - functionResponseJSON = translatorcommon.SetGeminiFunctionResponseResult(functionResponseJSON, "response.result", functionResponseResult) + functionResponseJSON = translatorcommon.SetGeminiFunctionResponseResult(functionResponseJSON, responsePath, functionResponseResult) } else { // Content field is missing entirely — .Raw is empty which // causes sjson.SetRaw to produce invalid JSON (e.g. "result":}). - functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, "response.result", "") + functionResponseJSON, _ = sjson.SetBytes(functionResponseJSON, responsePath, "") } partJSON := []byte(`{}`) diff --git a/internal/translator/antigravity/claude/antigravity_claude_request_test.go b/internal/translator/antigravity/claude/antigravity_claude_request_test.go index 4a60074f4aa..bbb232f0748 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_request_test.go +++ b/internal/translator/antigravity/claude/antigravity_claude_request_test.go @@ -3924,3 +3924,108 @@ func TestConvertClaudeRequestToAntigravity_FunctionResponseJSONRef(t *testing.T) t.Fatalf("expected result to contain ref target, got %q", result.String()) } } + +func TestConvertClaudeRequestToAntigravity_PreservesToolResultErrors(t *testing.T) { + tests := []struct { + name string + fields string + wantResponse string + wantImage bool + }{ + { + name: "success_without_flag", + fields: `"content":"permission denied"`, + wantResponse: `{"result":"permission denied"}`, + }, + { + name: "explicit_success", + fields: `"is_error":false,"content":{"ok":true}`, + wantResponse: `{"result":{"ok":true}}`, + }, + { + name: "non_bool_is_error_string_treated_as_success", + fields: `"is_error":"true","content":"permission denied"`, + wantResponse: `{"result":"permission denied"}`, + }, + { + name: "non_bool_is_error_number_treated_as_success", + fields: `"is_error":1,"content":"permission denied"`, + wantResponse: `{"result":"permission denied"}`, + }, + { + name: "failed_string", + fields: `"is_error":true,"content":"permission denied"`, + wantResponse: `{"error":"permission denied"}`, + }, + { + name: "failed_structured_result", + fields: `"is_error":true,"content":{"code":"EACCES","retryable":false}`, + wantResponse: `{"error":{"code":"EACCES","retryable":false}}`, + }, + { + name: "failed_text_and_image", + fields: `"is_error":true,"content":[{"type":"text","text":"permission denied"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}}]`, + wantResponse: `{"error":{"type":"text","text":"permission denied"}}`, + wantImage: true, + }, + { + name: "failed_image_only_array", + fields: `"is_error":true,"content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}}]`, + wantResponse: `{"error":""}`, + wantImage: true, + }, + { + name: "failed_single_image_object", + fields: `"is_error":true,"content":{"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}}`, + wantResponse: `{"error":""}`, + wantImage: true, + }, + { + name: "failed_multiple_blocks", + fields: `"is_error":true,"content":[{"type":"text","text":"first"},{"type":"text","text":"second"}]`, + wantResponse: `{"error":[{"type":"text","text":"first"},{"type":"text","text":"second"}]}`, + }, + { + name: "failed_without_content", + fields: `"is_error":true`, + wantResponse: `{"error":""}`, + }, + { + name: "failed_schema_ref_remains_opaque", + fields: `"is_error":true,"content":{"$ref":"#/definitions/Error"}`, + wantResponse: `{"error":"{\"$ref\":\"#/definitions/Error\"}"}`, + }, + } + + for _, stream := range []bool{false, true} { + for _, tt := range tests { + t.Run(fmt.Sprintf("stream=%t/%s", stream, tt.name), func(t *testing.T) { + input := []byte(fmt.Sprintf(`{"model":"claude-sonnet-4-5","messages":[ + {"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1",%s}]} + ]}`, tt.fields)) + output, errTranslate := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5", input, stream) + if errTranslate != nil { + t.Fatal(errTranslate) + } + if !gjson.ValidBytes(output) { + t.Fatalf("invalid translated JSON: %s", output) + } + response := gjson.GetBytes(output, "request.contents.1.parts.0.functionResponse") + if response.Get("id").String() != "call_1" || response.Get("name").String() != "Read" { + t.Fatalf("tool result lost its call identity: %s", response.Raw) + } + if got := response.Get("response").Raw; got != tt.wantResponse { + t.Fatalf("functionResponse.response = %s, want %s", got, tt.wantResponse) + } + if tt.wantImage { + if got := response.Get("parts.0.inlineData").Raw; got != `{"mimeType":"image/png","data":"aGVsbG8="}` { + t.Fatalf("functionResponse.parts[0].inlineData = %s", got) + } + } else if response.Get("parts").Exists() { + t.Fatalf("unexpected functionResponse.parts: %s", response.Get("parts").Raw) + } + }) + } + } +} diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_response.go b/internal/translator/antigravity/gemini/antigravity_gemini_response.go index e7a686ae068..c0f1550a6d1 100644 --- a/internal/translator/antigravity/gemini/antigravity_gemini_response.go +++ b/internal/translator/antigravity/gemini/antigravity_gemini_response.go @@ -164,9 +164,13 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR if alt, ok := ctx.Value("alt").(string); ok { var chunk []byte if alt == "" { - responseResult := gjson.GetBytes(rawJSON, "response") + responseResult := util.GetGJSONBytesNoCopy(rawJSON, "response") if responseResult.Exists() { - chunk = []byte(responseResult.Raw) + if responseResult.Index > 0 && responseResult.Index+len(responseResult.Raw) <= len(rawJSON) { + chunk = rawJSON[responseResult.Index : responseResult.Index+len(responseResult.Raw)] + } else { + chunk = []byte(responseResult.Raw) + } chunk = restoreUsageMetadata(chunk) chunk = restoreGeminiFunctionNames(chunk, originalRequestRawJSON) } @@ -201,24 +205,34 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR // // Returns: // - []byte: A Gemini-compatible JSON response containing the response data. +var cpaUsageMetadataMarker = []byte(`"cpaUsageMetadata"`) + func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) []byte { var chunk []byte - responseResult := gjson.GetBytes(rawJSON, "response") + responseResult := util.GetGJSONBytesNoCopy(rawJSON, "response") if responseResult.Exists() { - chunk = restoreUsageMetadata([]byte(responseResult.Raw)) + if responseResult.Index > 0 && responseResult.Index+len(responseResult.Raw) <= len(rawJSON) { + chunk = rawJSON[responseResult.Index : responseResult.Index+len(responseResult.Raw)] + } else { + chunk = []byte(responseResult.Raw) + } + chunk = restoreUsageMetadata(chunk) chunk = restoreGeminiFunctionNames(chunk, originalRequestRawJSON) } else { chunk = restoreGeminiFunctionNames(rawJSON, originalRequestRawJSON) } - candidates := gjson.GetBytes(chunk, "candidates") + candidates := util.GetGJSONBytesNoCopy(chunk, "candidates") if candidates.IsArray() { // Default every candidate, not just the first one, so a multi-candidate // response cannot leave some of them unterminated. - for i, candidate := range candidates.Array() { + idx := 0 + candidates.ForEach(func(_, candidate gjson.Result) bool { if candidate.Get("finishReason").String() == "" { - chunk, _ = sjson.SetBytes(chunk, fmt.Sprintf("candidates.%d.finishReason", i), "STOP") + chunk, _ = sjson.SetBytes(chunk, fmt.Sprintf("candidates.%d.finishReason", idx), "STOP") } - } + idx++ + return true + }) } return chunk } @@ -258,7 +272,24 @@ func GeminiTokenCount(ctx context.Context, count int64) []byte { // to preserve usage data while hiding it from clients that don't expect it. // When returning standard Gemini API format, we must restore the original name. func restoreUsageMetadata(chunk []byte) []byte { - if cpaUsage := gjson.GetBytes(chunk, "cpaUsageMetadata"); cpaUsage.Exists() { + if !bytes.Contains(chunk, cpaUsageMetadataMarker) { + return chunk + } + if cpaUsage := util.GetGJSONBytesNoCopy(chunk, "cpaUsageMetadata"); cpaUsage.Exists() { + if cpaUsage.Index > 0 && cpaUsage.Index <= len(chunk) { + searchStart := cpaUsage.Index - 36 + if searchStart < 0 { + searchStart = 0 + } + if rel := bytes.LastIndex(chunk[searchStart:cpaUsage.Index], cpaUsageMetadataMarker); rel >= 0 { + keyPos := searchStart + rel + 1 + out := make([]byte, 0, len(chunk)-3) + out = append(out, chunk[:keyPos]...) + out = append(out, "usageMetadata"...) + out = append(out, chunk[keyPos+len("cpaUsageMetadata"):]...) + return out + } + } chunk, _ = sjson.SetRawBytes(chunk, "usageMetadata", []byte(cpaUsage.Raw)) chunk, _ = sjson.DeleteBytes(chunk, "cpaUsageMetadata") } diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go b/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go index ad73fd59f00..be40fc778b1 100644 --- a/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go +++ b/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go @@ -239,3 +239,16 @@ func TestConvertAntigravityResponseToGeminiNonStream_DefaultsMissingFinishReason t.Fatalf("ConvertAntigravityResponseToGeminiNonStream finishReason = %q, want STOP", fr) } } + +func TestAntigravityGeminiResponseAllowsWhitespace(t *testing.T) { + body := `{"candidates":[{"content":{"parts":[{"text":"ok"}]},"finishReason":"STOP"}]}` + for _, separator := range []string{"", " ", "\n\t"} { + t.Run(fmt.Sprintf("separator=%q", separator), func(t *testing.T) { + raw := []byte(`{"response":` + separator + body + `}`) + out := ConvertAntigravityResponseToGeminiNonStream(context.Background(), "", nil, nil, raw, nil) + if !gjson.ValidBytes(out) || string(out) != body { + t.Fatalf("incorrect response slice: %s", out) + } + }) + } +} diff --git a/internal/translator/antigravity/openai/responses/antigravity_openai-responses_tool_search_test.go b/internal/translator/antigravity/openai/responses/antigravity_openai-responses_tool_search_test.go index 7299796a08e..c698e95ef16 100644 --- a/internal/translator/antigravity/openai/responses/antigravity_openai-responses_tool_search_test.go +++ b/internal/translator/antigravity/openai/responses/antigravity_openai-responses_tool_search_test.go @@ -53,18 +53,25 @@ func TestConvertAntigravityResponseToOpenAIResponses_ToolSearchCall(t *testing.T chunk := []byte(`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"tool_search","args":{"query":"fixture"}}}]},"finishReason":"STOP"}],"responseId":"resp_ag_tool_search"}}`) var param any - var done gjson.Result - for _, output := range ConvertAntigravityResponseToOpenAIResponses(context.Background(), "gemini-3.8-flash-high", originalRequest, nil, chunk, ¶m) { - payload := string(output) - if idx := strings.Index(payload, "data:"); idx >= 0 { - data := gjson.Parse(payload[idx+len("data:"):]) - if data.Get("type").String() == "response.output_item.done" { - done = data + var added, done, completed gjson.Result + for _, frame := range [][]byte{chunk, []byte("data: [DONE]")} { + for _, output := range ConvertAntigravityResponseToOpenAIResponses(context.Background(), "gemini-3.8-flash-high", originalRequest, nil, frame, ¶m) { + payload := string(output) + if idx := strings.Index(payload, "data:"); idx >= 0 { + data := gjson.Parse(payload[idx+len("data:"):]) + switch data.Get("type").String() { + case "response.output_item.added": + added = data + case "response.output_item.done": + done = data + case "response.completed": + completed = data + } } } } - if !done.Exists() { - t.Fatalf("missing response.output_item.done event") + if !added.Exists() || !done.Exists() || !completed.Exists() { + t.Fatalf("missing tool_search_call lifecycle events: added=%v done=%v completed=%v", added.Exists(), done.Exists(), completed.Exists()) } if got := done.Get("item.type").String(); got != "tool_search_call" { t.Fatalf("item type = %q, want tool_search_call; raw: %s", got, done.Raw) @@ -75,4 +82,33 @@ func TestConvertAntigravityResponseToOpenAIResponses_ToolSearchCall(t *testing.T if got := done.Get("item.execution").String(); got != "client" { t.Fatalf("item execution = %q", got) } + callID := done.Get("item.call_id").String() + if callID == "" || added.Get("item.call_id").String() != callID || completed.Get("response.output.0.call_id").String() != callID { + t.Fatalf("call_id must be non-empty and stable across events") + } + wantItemID := "tsc_" + callID + for _, test := range []struct { + label string + item gjson.Result + }{ + {label: "added", item: added.Get("item")}, + {label: "done", item: done.Get("item")}, + {label: "completed", item: completed.Get("response.output.0")}, + } { + if got := test.item.Get("id").String(); got != wantItemID { + t.Fatalf("%s id = %q, want %q; raw: %s", test.label, got, wantItemID, test.item.Raw) + } + } + + nonStreamRaw := []byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"tool_search","args":{"query":"fixture"}}}]},"finishReason":"STOP"}],"responseId":"resp_ag_tool_search_nonstream"}}`) + nonStreamOut := ConvertAntigravityResponseToOpenAIResponsesNonStream(context.Background(), "gemini-3.8-flash-high", originalRequest, nil, nonStreamRaw, nil) + nonStreamItem := gjson.GetBytes(nonStreamOut, "output.0") + if got := nonStreamItem.Get("type").String(); got != "tool_search_call" { + t.Fatalf("non-stream item type = %q, want tool_search_call; raw: %s", got, nonStreamOut) + } + if nonStreamCallID := nonStreamItem.Get("call_id").String(); nonStreamCallID == "" { + t.Fatalf("non-stream call_id must be non-empty: %s", nonStreamItem.Raw) + } else if got := nonStreamItem.Get("id").String(); got != "tsc_"+nonStreamCallID { + t.Fatalf("non-stream id = %q, want %q; raw: %s", got, "tsc_"+nonStreamCallID, nonStreamItem.Raw) + } } diff --git a/internal/translator/codex/claude/codex_claude_request.go b/internal/translator/codex/claude/codex_claude_request.go index 504a5b11617..5842e509cf8 100644 --- a/internal/translator/codex/claude/codex_claude_request.go +++ b/internal/translator/codex/claude/codex_claude_request.go @@ -59,10 +59,15 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, template := []byte(`{"model":"","instructions":"","input":[]}`) rootResult := gjson.ParseBytes(rawJSON) - toolNameMap := buildReverseMapFromClaudeOriginalToShort(rawJSON) + mayHaveToolAddition := claudePayloadMayContainToolAddition(rawJSON) + toolNameMap := buildReverseMapFromClaudeOriginalToShortFromResult(rootResult, mayHaveToolAddition) signatureTargetModel := translatorcommon.ResolveCodexReplayTargetModel(modelName, rawJSON) template, _ = sjson.SetBytes(template, "model", modelName) inputItems := translatorcommon.NewRawArrayItems(rootResult.Get("messages.#").Int()) + toolsResult := rootResult.Get("tools") + var declaredToolsByName map[string]gjson.Result + declaredToolsIndexed := false + hasAdditionalTools := false // Process system messages and convert them to input content format. systemsResult := rootResult.Get("system") @@ -109,7 +114,8 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, messageResult := messageResults[i] messageRole := messageResult.Get("role").String() if messageRole == "system" { - if reminderText, ok := translatorcommon.ClaudeMessageSystemReminderText(messageResult.Get("content")); ok { + systemContent := messageResult.Get("content") + if reminderText, ok := translatorcommon.ClaudeMessageSystemReminderText(systemContent); ok { message := []byte(`{"type":"message","role":"user","content":[{"type":"input_text","text":""}]}`) message, _ = sjson.SetBytes(message, "content.0.text", reminderText) if len(pendingToolUseIDs) > 0 { @@ -118,6 +124,14 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, inputItems = append(inputItems, message) } } + if additionalTools, ok := convertClaudeSystemAdditionalTools(systemContent, toolsResult, &declaredToolsByName, &declaredToolsIndexed, toolNameMap); ok { + hasAdditionalTools = true + if len(pendingToolUseIDs) > 0 { + pendingSystemReminders = append(pendingSystemReminders, additionalTools) + } else { + inputItems = append(inputItems, additionalTools) + } + } continue } @@ -349,12 +363,13 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, } // Convert tools declarations to the expected format for the Codex API. - toolsResult := rootResult.Get("tools") hasWebSearchTool := false var toolItems [][]byte - if toolsResult.IsArray() { + if toolsResult.IsArray() || hasAdditionalTools { webSearchToolNames := buildClaudeWebSearchToolNameSet(toolsResult) template, _ = sjson.SetRawBytes(template, "tool_choice", convertClaudeToolChoiceToCodex(rootResult.Get("tool_choice"), toolNameMap, webSearchToolNames)) + } + if toolsResult.IsArray() { toolResults := toolsResult.Array() toolItems = make([][]byte, 0, len(toolResults)) for i := 0; i < len(toolResults); i++ { @@ -365,33 +380,7 @@ func convertClaudeRequestToCodex(modelName string, inputRawJSON []byte, _ bool, toolItems = append(toolItems, convertClaudeWebSearchToolToCodex(toolResult)) continue } - tool := []byte(toolResult.Raw) - if toolResult.Get("type").Type != gjson.String || toolResult.Get("type").String() != "function" { - tool, _ = sjson.SetBytes(tool, "type", "function") - } - // Apply shortened name if needed - if v := toolResult.Get("name"); v.Exists() { - originalName := v.String() - name := originalName - if short, ok := toolNameMap[name]; ok { - name = short - } else { - name = shortenNameIfNeeded(name) - } - if v.Type != gjson.String || name != originalName { - tool, _ = sjson.SetBytes(tool, "name", name) - } - } - tool, _ = sjson.SetRawBytes(tool, "parameters", []byte(normalizeToolParameters(toolResult.Get("input_schema").Raw))) - for _, path := range []string{"input_schema", "parameters.$schema", "cache_control", "defer_loading"} { - if gjson.GetBytes(tool, path).Exists() { - tool, _ = sjson.DeleteBytes(tool, path) - } - } - if gjson.GetBytes(tool, "strict").Type != gjson.False { - tool, _ = sjson.SetBytes(tool, "strict", false) - } - toolItems = append(toolItems, tool) + toolItems = append(toolItems, convertClaudeFunctionToolToCodex(toolResult, toolNameMap)) } } @@ -655,6 +644,9 @@ func buildShortNameMap(names []string) map[string]string { } for _, n := range names { + if _, exists := m[n]; exists { + continue + } cand := baseCandidate(n) uniq := makeUnique(cand) used[uniq] = struct{}{} @@ -663,25 +655,216 @@ func buildShortNameMap(names []string) map[string]string { return m } -// buildReverseMapFromClaudeOriginalToShort builds original->short map, used to map tool_use names to short. -func buildReverseMapFromClaudeOriginalToShort(original []byte) map[string]string { - tools := gjson.GetBytes(original, "tools") - m := map[string]string{} - if !tools.IsArray() { - return m +func convertClaudeFunctionToolToCodex(toolResult gjson.Result, toolNameMap map[string]string) []byte { + tool := []byte(toolResult.Raw) + if toolResult.Get("type").Type != gjson.String || toolResult.Get("type").String() != "function" { + tool, _ = sjson.SetBytes(tool, "type", "function") + } + // Apply shortened name if needed + if v := toolResult.Get("name"); v.Exists() { + originalName := v.String() + name := originalName + if short, ok := toolNameMap[name]; ok { + name = short + } else { + name = shortenNameIfNeeded(name) + } + if v.Type != gjson.String || name != originalName { + tool, _ = sjson.SetBytes(tool, "name", name) + } + } + tool, _ = sjson.SetRawBytes(tool, "parameters", []byte(normalizeToolParameters(toolResult.Get("input_schema").Raw))) + for _, path := range []string{"input_schema", "parameters.$schema", "cache_control", "defer_loading"} { + if gjson.GetBytes(tool, path).Exists() { + tool, _ = sjson.DeleteBytes(tool, path) + } + } + if gjson.GetBytes(tool, "strict").Type != gjson.False { + tool, _ = sjson.SetBytes(tool, "strict", false) + } + return tool +} + +func ensureClaudeDeclaredToolsIndexed(toolsResult gjson.Result, declaredToolsByName *map[string]gjson.Result, declaredToolsIndexed *bool) { + if *declaredToolsIndexed { + return + } + *declaredToolsIndexed = true + if !toolsResult.IsArray() { + return + } + toolArr := toolsResult.Array() + if *declaredToolsByName == nil { + *declaredToolsByName = make(map[string]gjson.Result, len(toolArr)) + } + for i := 0; i < len(toolArr); i++ { + if name := toolArr[i].Get("name").String(); name != "" { + if _, exists := (*declaredToolsByName)[name]; !exists { + (*declaredToolsByName)[name] = toolArr[i] + } + } + } +} + +// convertClaudeSystemAdditionalTools converts mid-conversation tool_addition blocks +// inside a role=system message into a Responses additional_tools input item. +func convertClaudeSystemAdditionalTools(content, toolsResult gjson.Result, declaredToolsByName *map[string]gjson.Result, declaredToolsIndexed *bool, toolNameMap map[string]string) ([]byte, bool) { + if !content.IsArray() { + return nil, false + } + parts := content.Array() + var toolItems [][]byte + var indexByName map[string]int + for i := 0; i < len(parts); i++ { + part := parts[i] + if part.Get("type").String() != "tool_addition" { + continue + } + toolSpec := part.Get("tool") + var def gjson.Result + switch toolSpec.Get("type").String() { + case "tool_definition": + def = toolSpec.Get("definition") + if def.IsObject() { + if name := def.Get("name").String(); name != "" { + ensureClaudeDeclaredToolsIndexed(toolsResult, declaredToolsByName, declaredToolsIndexed) + if *declaredToolsByName == nil { + *declaredToolsByName = make(map[string]gjson.Result) + } + (*declaredToolsByName)[name] = def + } + } + case "tool_reference": + ensureClaudeDeclaredToolsIndexed(toolsResult, declaredToolsByName, declaredToolsIndexed) + refName := toolSpec.Get("name").String() + if refName == "" { + refName = toolSpec.Get("tool_name").String() + } + if refName != "" && *declaredToolsByName != nil { + def = (*declaredToolsByName)[refName] + } + default: + if candidate := toolSpec.Get("definition"); candidate.IsObject() { + def = candidate + if name := def.Get("name").String(); name != "" { + ensureClaudeDeclaredToolsIndexed(toolsResult, declaredToolsByName, declaredToolsIndexed) + if *declaredToolsByName == nil { + *declaredToolsByName = make(map[string]gjson.Result) + } + (*declaredToolsByName)[name] = def + } + } + } + if !def.IsObject() || isClaudeWebSearchToolType(def.Get("type").String()) { + continue + } + origName := strings.TrimSpace(def.Get("name").String()) + if origName == "" { + continue + } + converted := convertClaudeFunctionToolToCodex(def, toolNameMap) + if indexByName == nil { + indexByName = make(map[string]int) + } + if idx, exists := indexByName[origName]; exists { + toolItems[idx] = converted + } else { + indexByName[origName] = len(toolItems) + toolItems = append(toolItems, converted) + } + } + if len(toolItems) == 0 { + return nil, false } + item := []byte(`{"type":"additional_tools","role":"developer","tools":[]}`) + item, _ = sjson.SetRawBytes(item, "tools", translatorcommon.JoinRawArray(toolItems)) + return item, true +} + +func collectClaudeDeclaredToolNames(rootResult gjson.Result, hasToolAddition bool) []string { var names []string - arr := tools.Array() - for i := 0; i < len(arr); i++ { - n := arr[i].Get("name").String() - if n != "" { - names = append(names, n) + 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 len(names) > 0 { - m = buildShortNameMap(names) + if !hasToolAddition { + return names } - return m + messages := rootResult.Get("messages") + if !messages.IsArray() { + return names + } + for _, msg := range messages.Array() { + if msg.Get("role").String() != "system" { + continue + } + content := msg.Get("content") + if !content.IsArray() { + continue + } + for _, part := range content.Array() { + if part.Get("type").String() != "tool_addition" { + 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 n := def.Get("name").String(); n != "" { + names = append(names, n) + } + } + } + } + return names +} + +func buildReverseMapFromClaudeOriginalToShortFromResult(rootResult gjson.Result, hasToolAddition bool) map[string]string { + names := collectClaudeDeclaredToolNames(rootResult, hasToolAddition) + if len(names) == 0 { + return map[string]string{} + } + return buildShortNameMap(names) +} + +func claudePayloadMayContainToolAddition(rawJSON []byte) bool { + return bytes.Contains(rawJSON, []byte("tool_addition")) || + bytes.Contains(rawJSON, []byte(`\u`)) || + bytes.Contains(rawJSON, []byte(`\U`)) +} + +// buildReverseMapFromClaudeOriginalToShort builds original->short map, used to map tool_use names to short. +func buildReverseMapFromClaudeOriginalToShort(original []byte) map[string]string { + mayHaveToolAddition := claudePayloadMayContainToolAddition(original) + if !mayHaveToolAddition { + tools := gjson.GetBytes(original, "tools") + m := map[string]string{} + if !tools.IsArray() { + return m + } + var names []string + arr := tools.Array() + for i := 0; i < len(arr); i++ { + n := arr[i].Get("name").String() + if n != "" { + names = append(names, n) + } + } + if len(names) > 0 { + m = buildShortNameMap(names) + } + return m + } + return buildReverseMapFromClaudeOriginalToShortFromResult(gjson.ParseBytes(original), true) } // normalizeToolParameters ensures object schemas contain at least an empty properties map, diff --git a/internal/translator/codex/claude/codex_claude_request_test.go b/internal/translator/codex/claude/codex_claude_request_test.go index cbe9d0f6da3..b23b5e39374 100644 --- a/internal/translator/codex/claude/codex_claude_request_test.go +++ b/internal/translator/codex/claude/codex_claude_request_test.go @@ -1406,3 +1406,154 @@ func TestConvertClaudeRequestToCodex_StripsPatternPropertiesIncompatibleKeys(t * t.Errorf("expected patternProperties key '^[a-z]+$' to be preserved, got: %s", params.Get("patternProperties").Raw) } } + +func TestConvertClaudeRequestToCodex_MidConversationToolAdditions(t *testing.T) { + longMCPTool := "mcp__very_long_server_name_for_testing__deliver_final_subagent_handback_report" + inputJSON := `{ + "model": "gpt-6.1-sol", + "tools": [ + { + "name": "Read", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}} + }, + { + "name": "deferred_lookup", + "description": "Look up deferred notes", + "defer_loading": true, + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} + } + ], + "messages": [ + {"role": "user", "content": "Start task"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "SubagentHandback is now available"}, + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": { + "name": "SubagentHandback", + "description": "Deliver subagent report", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"summary": {"type": "string"}}, + "required": ["summary"] + } + } + } + }, + { + "type": "tool_addition", + "tool": {"type": "tool_reference", "name": "deferred_lookup"} + }, + { + "type": "tool_removal", + "tool": {"type": "tool_reference", "name": "Read"} + } + ] + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "call_1", "name": "Read", "input": {"path": "README.md"}} + ] + }, + { + "role": "system", + "content": [ + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": { + "name": "` + longMCPTool + `", + "description": "Long MCP handback tool", + "input_schema": {"type": "object"} + } + } + } + ] + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "call_1", "content": "file contents"}, + {"type": "text", "text": "Finish up"} + ] + } + ] + }` + + translated, err := ConvertClaudeRequestToCodex("gpt-6.1-sol", []byte(inputJSON), false) + if err != nil { + 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. + topTools := gjson.GetBytes(translated, "tools").Array() + if len(topTools) != 2 { + t.Fatalf("expected 2 top-level tools, got %d: %s", len(topTools), gjson.GetBytes(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()) + } + + // Expected input items: + // 0: message (user: "Start task") + // 1: message (user: "...") + // 2: additional_tools (developer: SubagentHandback, deferred_lookup) + // 3: function_call (call_1) + // 4: function_call_output (call_1) + // 5: additional_tools (developer: shortened longMCPTool, deferred after function_call_output) + // 6: message (user: "Finish up") + inputs := gjson.GetBytes(translated, "input").Array() + if len(inputs) != 7 { + t.Fatalf("expected 7 input items, got %d: %s", len(inputs), gjson.GetBytes(translated, "input").Raw) + } + + addTools1 := inputs[2] + if got := addTools1.Get("type").String(); got != "additional_tools" { + t.Fatalf("inputs[2].type = %q, want additional_tools: %s", got, addTools1.Raw) + } + if got := addTools1.Get("role").String(); got != "developer" { + t.Fatalf("inputs[2].role = %q, want developer: %s", got, addTools1.Raw) + } + addTools1List := addTools1.Get("tools").Array() + if len(addTools1List) != 2 { + t.Fatalf("expected 2 tools in first additional_tools, got %d: %s", len(addTools1List), addTools1.Raw) + } + if addTools1List[0].Get("type").String() != "function" || addTools1List[0].Get("name").String() != "SubagentHandback" { + t.Fatalf("unexpected first additional tool: %s", addTools1List[0].Raw) + } + if addTools1List[0].Get("strict").Type != gjson.False { + t.Fatalf("expected strict=false on SubagentHandback, got: %s", addTools1List[0].Raw) + } + if addTools1List[0].Get("input_schema").Exists() || addTools1List[0].Get("parameters.$schema").Exists() { + t.Fatalf("expected input_schema and parameters.$schema to be stripped: %s", addTools1List[0].Raw) + } + if addTools1List[0].Get("parameters.properties.summary.type").String() != "string" { + t.Fatalf("expected normalized parameters on SubagentHandback, got: %s", addTools1List[0].Raw) + } + if addTools1List[1].Get("type").String() != "function" || addTools1List[1].Get("name").String() != "deferred_lookup" { + t.Fatalf("unexpected second additional tool (resolved tool_reference): %s", addTools1List[1].Raw) + } + if addTools1List[1].Get("defer_loading").Exists() { + t.Fatalf("expected defer_loading to be stripped on resolved tool_reference: %s", addTools1List[1].Raw) + } + + if inputs[3].Get("type").String() != "function_call" || inputs[4].Get("type").String() != "function_call_output" { + t.Fatalf("expected function_call (3) followed by function_call_output (4), got %q and %q", inputs[3].Get("type").String(), inputs[4].Get("type").String()) + } + addTools2 := inputs[5] + if got := addTools2.Get("type").String(); got != "additional_tools" { + t.Fatalf("inputs[5].type = %q, want additional_tools: %s", got, addTools2.Raw) + } + shortLongName := addTools2.Get("tools.0.name").String() + if shortLongName == "" || shortLongName == longMCPTool || len(shortLongName) > 64 { + t.Fatalf("expected shortened MCP tool name <= 64 chars, got %q (len=%d)", shortLongName, len(shortLongName)) + } +} diff --git a/internal/translator/codex/claude/codex_claude_response.go b/internal/translator/codex/claude/codex_claude_response.go index edba12ae66d..5c28be11486 100644 --- a/internal/translator/codex/claude/codex_claude_response.go +++ b/internal/translator/codex/claude/codex_claude_response.go @@ -927,24 +927,10 @@ func setClaudeReasoningUsage(out []byte, usage gjson.Result) []byte { // buildReverseMapFromClaudeOriginalShortToOriginal builds a map[short]original from original Claude request tools. func buildReverseMapFromClaudeOriginalShortToOriginal(original []byte) map[string]string { - tools := gjson.GetBytes(original, "tools") rev := map[string]string{} - if !tools.IsArray() { - return rev - } - var names []string - arr := tools.Array() - for i := 0; i < len(arr); i++ { - n := arr[i].Get("name").String() - if n != "" { - names = append(names, n) - } - } - if len(names) > 0 { - m := buildShortNameMap(names) - for orig, short := range m { - rev[short] = orig - } + m := buildReverseMapFromClaudeOriginalToShort(original) + for orig, short := range m { + rev[short] = orig } return rev } diff --git a/internal/translator/codex/claude/codex_claude_response_test.go b/internal/translator/codex/claude/codex_claude_response_test.go index e827a165e35..8de5dd84c16 100644 --- a/internal/translator/codex/claude/codex_claude_response_test.go +++ b/internal/translator/codex/claude/codex_claude_response_test.go @@ -2045,3 +2045,63 @@ func TestExtractResponsesUsage(t *testing.T) { }) } } + +func TestConvertCodexResponseToClaude_RestoresLongMidConversationToolAdditionName(t *testing.T) { + // Two >64-char MCP tools that share the same post-"__" suffix so the mid-conversation + // tool_addition collides with the top-level tool and receives a "_1" suffix. + topLongMCPTool := "mcp__alpha_very_long_server_name_for_testing__deliver_final_subagent_handback_report" + midLongMCPTool := "mcp__bravo_very_long_server_name_for_testing__deliver_final_subagent_handback_report" + // Use JSON Unicode escape \u005f in "tool\u005faddition" to cover escaped JSON tokens too. + originalRequest := []byte(`{ + "model": "gpt-6.1-sol", + "tools": [{"name": "` + topLongMCPTool + `", "input_schema": {"type": "object"}}], + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "system", + "content": [ + { + "type": "tool\u005faddition", + "tool": { + "type": "tool_definition", + "definition": { + "name": "` + midLongMCPTool + `", + "description": "Long MCP tool added mid-conversation", + "input_schema": {"type": "object"} + } + } + } + ] + } + ] + }`) + + translatedReq, err := ConvertClaudeRequestToCodex("gpt-6.1-sol", originalRequest, false) + if err != nil { + t.Fatalf("ConvertClaudeRequestToCodex returned error: %v", err) + } + topShortName := gjson.GetBytes(translatedReq, "tools.0.name").String() + midShortName := gjson.GetBytes(translatedReq, "input.1.tools.0.name").String() + if topShortName == "" || midShortName == "" || topShortName == midShortName { + t.Fatalf("expected distinct shortened names for colliding MCP tools, got top=%q mid=%q", topShortName, midShortName) + } + + // 1. Streaming response restoration + var param any + addedEvent := []byte(fmt.Sprintf(`data: {"type":"response.output_item.added","output_index":0,"item":{"type":"function_call","id":"fc_1","call_id":"call_1","name":%q,"arguments":""}}`, midShortName)) + chunks := ConvertCodexResponseToClaude(context.Background(), "gpt-6.1-sol", originalRequest, nil, addedEvent, ¶m) + joined := string(bytes.Join(chunks, nil)) + if !strings.Contains(joined, `"name":"`+midLongMCPTool+`"`) { + t.Fatalf("streaming response did not restore long mid-conversation tool name %q, got: %s", midLongMCPTool, joined) + } + + // 2. Non-streaming response restoration + completedEvent := []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_1","model":"gpt-6.1-sol","status":"completed","output":[{"type":"function_call","id":"fc_1","call_id":"call_1","name":%q,"arguments":"{}"},{"type":"function_call","id":"fc_2","call_id":"call_2","name":%q,"arguments":"{}"}],"usage":{"input_tokens":10,"output_tokens":5}}}`, topShortName, midShortName)) + nonStreamOut := ConvertCodexResponseToClaudeNonStream(context.Background(), "gpt-6.1-sol", originalRequest, nil, completedEvent, nil) + if got := gjson.GetBytes(nonStreamOut, "content.0.name").String(); got != topLongMCPTool { + t.Fatalf("non-streaming response content.0.name = %q, want %q; body=%s", got, topLongMCPTool, nonStreamOut) + } + if got := gjson.GetBytes(nonStreamOut, "content.1.name").String(); got != midLongMCPTool { + t.Fatalf("non-streaming response content.1.name = %q, want %q; body=%s", got, midLongMCPTool, nonStreamOut) + } +} diff --git a/internal/translator/gemini/claude/gemini_claude_request.go b/internal/translator/gemini/claude/gemini_claude_request.go index cae76027a10..92e245ab6d4 100644 --- a/internal/translator/gemini/claude/gemini_claude_request.go +++ b/internal/translator/gemini/claude/gemini_claude_request.go @@ -175,14 +175,18 @@ func convertClaudeRequestToGemini(modelName string, inputRawJSON []byte, _ bool, funcName = toolCallID } funcName = util.SanitizeFunctionName(funcName) - toolResult := util.ConvertClaudeToolResultContent(contentResult.Get("content")) - part := []byte(`{"functionResponse":{"name":"","response":{"result":""}}}`) + toolResult := util.ConvertClaudeToolResult(contentResult) + part := []byte(`{"functionResponse":{"name":"","response":{}}}`) part, _ = sjson.SetBytes(part, "functionResponse.id", toolCallID) part, _ = sjson.SetBytes(part, "functionResponse.name", funcName) + resultPath := "functionResponse.response.result" + if toolResult.IsError { + resultPath = "functionResponse.response.error" + } if toolResult.ResultIsRaw { - part = translatorcommon.SetGeminiFunctionResponseRaw(part, "functionResponse.response.result", toolResult.Result) + part = translatorcommon.SetGeminiFunctionResponseRaw(part, resultPath, toolResult.Result) } else { - part, _ = sjson.SetBytes(part, "functionResponse.response.result", toolResult.Result) + part, _ = sjson.SetBytes(part, resultPath, toolResult.Result) } partItems = append(partItems, part) for _, img := range toolResult.Images { diff --git a/internal/translator/gemini/claude/gemini_claude_tool_result_error_test.go b/internal/translator/gemini/claude/gemini_claude_tool_result_error_test.go new file mode 100644 index 00000000000..956995f9569 --- /dev/null +++ b/internal/translator/gemini/claude/gemini_claude_tool_result_error_test.go @@ -0,0 +1,102 @@ +package claude + +import ( + "fmt" + "testing" + + "github.com/tidwall/gjson" +) + +func TestConvertClaudeRequestToGemini_PreservesToolResultErrors(t *testing.T) { + tests := []struct { + name string + fields string + wantResponse string + wantImage bool + }{ + { + name: "success_without_flag", + fields: `"content":"permission denied"`, + wantResponse: `{"result":"permission denied"}`, + }, + { + name: "explicit_success", + fields: `"is_error":false,"content":{"ok":true}`, + wantResponse: `{"result":{"ok":true}}`, + }, + { + name: "failed_string", + fields: `"is_error":true,"content":"permission denied"`, + wantResponse: `{"error":"permission denied"}`, + }, + { + name: "failed_structured_result", + fields: `"is_error":true,"content":{"code":"EACCES","retryable":false}`, + wantResponse: `{"error":{"code":"EACCES","retryable":false}}`, + }, + { + name: "failed_text_and_image", + fields: `"is_error":true,"content":[{"type":"text","text":"permission denied"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}}]`, + wantResponse: `{"error":{"type":"text","text":"permission denied"}}`, + wantImage: true, + }, + { + name: "failed_multiple_blocks", + fields: `"is_error":true,"content":[{"type":"text","text":"first"},{"type":"text","text":"second"}]`, + wantResponse: `{"error":[{"type":"text","text":"first"},{"type":"text","text":"second"}]}`, + }, + { + name: "failed_without_content", + fields: `"is_error":true`, + wantResponse: `{"error":""}`, + }, + { + name: "failed_schema_ref_remains_opaque", + fields: `"is_error":true,"content":{"$ref":"#/definitions/Error"}`, + wantResponse: `{"error":"{\"$ref\":\"#/definitions/Error\"}"}`, + }, + } + for _, entrypoint := range []struct { + name string + translate func(string, []byte, bool) ([]byte, error) + }{ + {"native", ConvertClaudeRequestToGemini}, + {"compat", ConvertClaudeRequestToGeminiWithCompat}, + } { + for _, stream := range []bool{false, true} { + for _, tt := range tests { + t.Run(fmt.Sprintf("%s/stream=%t/%s", entrypoint.name, stream, tt.name), func(t *testing.T) { + input := []byte(fmt.Sprintf(`{"messages":[ + {"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"Read","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1",%s}]} + ]}`, tt.fields)) + output, errTranslate := entrypoint.translate("gemini-3-flash-preview", input, stream) + if errTranslate != nil { + t.Fatal(errTranslate) + } + if !gjson.ValidBytes(output) { + t.Fatalf("invalid translated JSON: %s", output) + } + response := gjson.GetBytes(output, "contents.1.parts.0.functionResponse") + if response.Get("id").String() != "call_1" || response.Get("name").String() != "Read" { + t.Fatalf("tool result lost its call identity: %s", response.Raw) + } + if got := response.Get("response").Raw; got != tt.wantResponse { + t.Fatalf("functionResponse.response = %s, want %s", got, tt.wantResponse) + } + parts := gjson.GetBytes(output, "contents.1.parts").Array() + wantParts := 1 + if tt.wantImage { + wantParts++ + if got := gjson.GetBytes(output, "contents.1.parts.1.inline_data").Raw; got != `{"mime_type":"image/png","data":"aGVsbG8="}` { + t.Fatalf("tool image changed: %s", got) + } + } + if len(parts) != wantParts { + t.Fatalf("tool result parts = %d, want %d", len(parts), wantParts) + } + }) + } + } + } +} diff --git a/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search.go b/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search.go index 46cbb1573db..ff2ab57d2ca 100644 --- a/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search.go +++ b/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search.go @@ -9,7 +9,7 @@ import ( "github.com/tidwall/sjson" ) -const openAIResponsesToolSearchIDPrefix = "ts_" +const openAIResponsesToolSearchIDPrefix = "tsc_" // normalizeOpenAIResponsesToolSearchItems rewrites replayed client-executed // tool_search history into the function_call / function_call_output shape that diff --git a/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search_test.go b/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search_test.go index 96995c075b1..7fad8a625a2 100644 --- a/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search_test.go +++ b/internal/translator/gemini/openai/responses/gemini_openai-responses_tool_search_test.go @@ -169,6 +169,26 @@ func TestConvertGeminiResponseToOpenAIResponses_ToolSearchCallStream(t *testing. if added.Get("item.call_id").String() == "" || added.Get("item.call_id").String() != done.Get("item.call_id").String() || done.Get("item.call_id").String() != completed.Get("response.output.0.call_id").String() { t.Fatalf("call_id must be non-empty and stable across events") } + wantItemID := "tsc_" + added.Get("item.call_id").String() + for _, test := range []struct { + label string + item gjson.Result + }{ + {label: "added", item: added.Get("item")}, + {label: "done", item: done.Get("item")}, + {label: "completed", item: completed.Get("response.output.0")}, + } { + if got := test.item.Get("id").String(); got != wantItemID { + t.Fatalf("%s id = %q, want %q; raw: %s", test.label, got, wantItemID, test.item.Raw) + } + } +} + +func TestOpenAIResponsesToolSearchCallItemID(t *testing.T) { + got := openAIResponsesToolSearchCallItemID("call_fixture") + if got != "tsc_call_fixture" { + t.Fatalf("item ID = %q, want tsc_call_fixture", got) + } } func TestConvertGeminiResponseToOpenAIResponsesNonStream_ToolSearchCall(t *testing.T) { @@ -189,6 +209,11 @@ func TestConvertGeminiResponseToOpenAIResponsesNonStream_ToolSearchCall(t *testi if got := item.Get("arguments.query").String(); got != "fixture tools" { t.Fatalf("non-stream arguments = %s", item.Get("arguments").Raw) } + if callID := item.Get("call_id").String(); callID == "" { + t.Fatalf("non-stream call_id must be non-empty: %s", item.Raw) + } else if got := item.Get("id").String(); got != "tsc_"+callID { + t.Fatalf("non-stream id = %q, want %q; raw: %s", got, "tsc_"+callID, item.Raw) + } if item.Get("name").Exists() { t.Fatalf("non-stream tool_search_call must not carry a name: %s", item.Raw) } diff --git a/internal/translator/openai/claude/openai_claude_request.go b/internal/translator/openai/claude/openai_claude_request.go index e178d8a1efe..00ecf40fd03 100644 --- a/internal/translator/openai/claude/openai_claude_request.go +++ b/internal/translator/openai/claude/openai_claude_request.go @@ -156,6 +156,7 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream } // Process Anthropic messages + var midConversationToolDefs []gjson.Result if messages := root.Get("messages"); messages.Exists() && messages.IsArray() { var pendingToolUseIDs []string var pendingSystemReminders [][]byte @@ -174,6 +175,20 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream messageItems = append(messageItems, msgJSON) } } + if contentResult.IsArray() { + contentResult.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() != "tool_addition" { + return true + } + toolSpec := part.Get("tool") + if toolSpec.Get("type").String() == "tool_definition" || toolSpec.Get("definition").IsObject() { + if def := toolSpec.Get("definition"); def.IsObject() && def.Get("name").String() != "" { + midConversationToolDefs = append(midConversationToolDefs, def) + } + } + return true + }) + } return true } @@ -370,29 +385,37 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream out = translatorcommon.SetRawArrayItems(out, "messages", messageItems) } - // Process tools - convert Anthropic tools to OpenAI functions + // Process tools - convert Anthropic tools (and mid-conversation tool_addition definitions) to OpenAI functions + var toolItems [][]byte + var toolIndexByName map[string]int + if len(midConversationToolDefs) > 0 { + toolIndexByName = make(map[string]int) + } if tools := root.Get("tools"); tools.Exists() && tools.IsArray() { - var toolItems [][]byte tools.ForEach(func(_, tool gjson.Result) bool { - openAIToolJSON := []byte(`{"type":"function","function":{"name":"","description":""}}`) - openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.name", tool.Get("name").String()) - openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.description", tool.Get("description").String()) - - // Convert Anthropic input_schema to OpenAI function parameters - if inputSchema := tool.Get("input_schema"); inputSchema.Exists() && inputSchema.Type != gjson.Null { - openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.parameters", normalizeObjectSchemaProperties(inputSchema.Value())) - } else { - openAIToolJSON, _ = sjson.SetRawBytes(openAIToolJSON, "function.parameters", []byte(`{"type":"object","properties":{}}`)) + openAIToolJSON := convertClaudeToolToOpenAI(tool) + if toolIndexByName != nil { + if name := tool.Get("name").String(); name != "" { + toolIndexByName[name] = len(toolItems) + } } - toolItems = append(toolItems, openAIToolJSON) return true }) - - if len(toolItems) > 0 { - out, _ = sjson.SetRawBytes(out, "tools", translatorcommon.JoinRawArray(toolItems)) + } + for _, def := range midConversationToolDefs { + name := def.Get("name").String() + converted := convertClaudeToolToOpenAI(def) + if idx, exists := toolIndexByName[name]; exists { + toolItems[idx] = converted + } else { + toolIndexByName[name] = len(toolItems) + toolItems = append(toolItems, converted) } } + if len(toolItems) > 0 { + out, _ = sjson.SetRawBytes(out, "tools", translatorcommon.JoinRawArray(toolItems)) + } // Tool choice mapping - convert Anthropic tool_choice to OpenAI format if toolChoice := root.Get("tool_choice"); toolChoice.Exists() && toolChoice.Type != gjson.Null { @@ -435,6 +458,20 @@ func convertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream return out, drops.Err() } +func convertClaudeToolToOpenAI(tool gjson.Result) []byte { + openAIToolJSON := []byte(`{"type":"function","function":{"name":"","description":""}}`) + openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.name", tool.Get("name").String()) + openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.description", tool.Get("description").String()) + + // Convert Anthropic input_schema to OpenAI function parameters + if inputSchema := tool.Get("input_schema"); inputSchema.Exists() && inputSchema.Type != gjson.Null { + openAIToolJSON, _ = sjson.SetBytes(openAIToolJSON, "function.parameters", normalizeObjectSchemaProperties(inputSchema.Value())) + } else { + openAIToolJSON, _ = sjson.SetRawBytes(openAIToolJSON, "function.parameters", []byte(`{"type":"object","properties":{}}`)) + } + return openAIToolJSON +} + func normalizeObjectSchemaProperties(schema any) any { switch value := schema.(type) { case bool: diff --git a/internal/translator/openai/claude/openai_claude_request_test.go b/internal/translator/openai/claude/openai_claude_request_test.go index ccdb853422c..932338a2db1 100644 --- a/internal/translator/openai/claude/openai_claude_request_test.go +++ b/internal/translator/openai/claude/openai_claude_request_test.go @@ -1827,3 +1827,61 @@ func TestClaudeTextWithUncachedFileKeepsText(t *testing.T) { t.Fatalf("body = %s", envelope.Body) } } + +func TestConvertClaudeRequestToOpenAI_MidConversationToolAdditionAppendsToTools(t *testing.T) { + inputJSON := []byte(`{ + "model": "gpt-5", + "tools": [ + { + "name": "Read", + "description": "Read file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}} + } + ], + "messages": [ + {"role": "user", "content": "start"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "SubagentHandback is now available"}, + { + "type": "tool_addition", + "tool": { + "type": "tool_definition", + "definition": { + "name": "SubagentHandback", + "description": "Deliver subagent report", + "input_schema": { + "type": "object", + "properties": {"summary": {"type": "string"}}, + "required": ["summary"] + } + } + } + } + ] + } + ] + }`) + + output, err := ConvertClaudeRequestToOpenAI("gpt-5", inputJSON, false) + if err != nil { + t.Fatalf("ConvertClaudeRequestToOpenAI returned error: %v", err) + } + tools := gjson.GetBytes(output, "tools").Array() + if len(tools) != 2 { + t.Fatalf("expected 2 tools in OpenAI request, got %d: %s", len(tools), gjson.GetBytes(output, "tools").Raw) + } + if got := tools[1].Get("function.name").String(); got != "SubagentHandback" { + t.Fatalf("tools[1].function.name = %q, want SubagentHandback", got) + } + if got := tools[1].Get("function.parameters.properties.summary.type").String(); got != "string" { + t.Fatalf("unexpected tools[1] parameters: %s", tools[1].Get("function.parameters").Raw) + } + + openAIResp := []byte(`{"id":"chatcmpl_1","model":"gpt-5","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"subagenthandback","arguments":"{\"summary\":\"done\"}"}}]}}]}`) + claudeResp := ConvertOpenAIResponseToClaudeNonStream(context.Background(), "gpt-5", inputJSON, output, openAIResp, nil) + if got := gjson.GetBytes(claudeResp, "content.0.name").String(); got != "SubagentHandback" { + t.Fatalf("ConvertOpenAIResponseToClaudeNonStream tool name = %q, want SubagentHandback; resp=%s", got, claudeResp) + } +} diff --git a/internal/tui/auth_tab.go b/internal/tui/auth_tab.go index 7106a091dbf..621f8cd15a9 100644 --- a/internal/tui/auth_tab.go +++ b/internal/tui/auth_tab.go @@ -191,6 +191,12 @@ func (m authTabModel) renderContent() string { for i, f := range m.files { name := getString(f, "name") channel := getString(f, "channel") + if channel == "" { + channel = getString(f, "provider") + } + if channel == "" { + channel = getString(f, "type") + } email := getString(f, "email") disabled := getBool(f, "disabled") @@ -199,6 +205,19 @@ func (m authTabModel) renderContent() string { if disabled { statusIcon = lipgloss.NewStyle().Foreground(colorMuted).Render("○") statusText = T("status_disabled") + } else if getBool(f, "requires_sign_in") { + statusIcon = warningStyle.Render("⚠") + statusText = "sign-in required" + if getBool(f, "unavailable") { + statusIcon = errorStyle.Render("✗") + statusText = "expired (sign-in required)" + } + } else if getBool(f, "unavailable") { + statusIcon = errorStyle.Render("✗") + statusText = getString(f, "status") + if statusText == "" { + statusText = "unavailable" + } } cursor := " " @@ -275,6 +294,11 @@ func (m authTabModel) renderDetail(f map[string]any) string { {"Email", "email", false}, {"Status", "status", false}, {"Status Msg", "status_message", false}, + {"Renewal", "renewal_state", false}, + {"Stage", "recovery_stage", false}, + {"Token Exp", "access_token_expires_at", false}, + {"Next Action", "next_action", false}, + {"Hint", "recovery_hint", false}, {"File Name", "file_name", false}, {"Auth Type", "auth_type", false}, {"Prefix", "prefix", true}, diff --git a/internal/tui/client.go b/internal/tui/client.go index 85f177c032b..7c688e117f2 100644 --- a/internal/tui/client.go +++ b/internal/tui/client.go @@ -163,9 +163,14 @@ func (c *Client) PutConfigYAML(yamlContent string) error { // GetAuthFiles lists auth credential files. // API returns {"files": [...]}. func (c *Client) GetAuthFiles() ([]map[string]any, error) { - wrapper, err := c.getJSON("/v0/management/auth-files") + wrapper, err := c.getJSON("/v8/management/credentials") if err != nil { - return nil, err + if strings.HasPrefix(err.Error(), "HTTP 404:") { + wrapper, err = c.getJSON("/v0/management/auth-files") + } + if err != nil { + return nil, err + } } return extractList(wrapper, "files") } diff --git a/internal/tui/client_test.go b/internal/tui/client_test.go index 7fe0d85a20e..ead3947b7c0 100644 --- a/internal/tui/client_test.go +++ b/internal/tui/client_test.go @@ -97,3 +97,72 @@ func TestNewAppWithBaseURL(t *testing.T) { t.Fatalf("expected app client baseURL https://proxy.example.com, got %q", app.client.baseURL) } } + +func TestClient_GetAuthFilesUsesV8CredentialsAndRendersRecoveryState(t *testing.T) { + var requestedPath string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestedPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "files": []map[string]any{ + { + "name": "codex-a.json", + "provider": "codex", + "email": "a@example.com", + "status": "active", + "disabled": false, + "unavailable": false, + "requires_sign_in": true, + "renewal_state": "sign_in_required", + "recovery_stage": "sign_in_required", + "access_token_expires_at": "2026-10-31T00:00:00Z", + "next_action": "sign_in", + "recovery_hint": "Re-authenticate this account on this proxy instance via OAuth login.", + }, + }, + }) + })) + defer ts.Close() + + client := NewClientWithBaseURL(ts.URL, "secret") + files, err := client.GetAuthFiles() + if err != nil { + t.Fatalf("GetAuthFiles error: %v", err) + } + if requestedPath != "/v8/management/credentials" { + t.Fatalf("requestedPath = %q, want /v8/management/credentials", requestedPath) + } + if len(files) != 1 { + t.Fatalf("len(files) = %d, want 1", len(files)) + } + + tab := newAuthTabModel(client) + tab.SetSize(120, 40) + tab.files = files + tab.expanded = 0 + view := tab.renderContent() + for _, expected := range []string{ + "sign-in required", + "sign_in_required", + "2026-10-31T00:00:00Z", + "sign_in", + "Re-authenticate this account on this proxy instance via OAuth login.", + } { + if !containsString(view, expected) { + t.Fatalf("expected TUI auth tab view to contain %q, got:\n%s", expected, view) + } + } +} + +func containsString(haystack, needle string) bool { + return len(needle) == 0 || (len(haystack) >= len(needle) && (haystack == needle || indexString(haystack, needle) >= 0)) +} + +func indexString(s, substr string) int { + for i := 0; i+len(substr) <= len(s); i++ { + if s[i:i+len(substr)] == substr { + return i + } + } + return -1 +} diff --git a/internal/util/claude_tool_name_perf_test.go b/internal/util/claude_tool_name_perf_test.go index c3a295708fd..e858e678116 100644 --- a/internal/util/claude_tool_name_perf_test.go +++ b/internal/util/claude_tool_name_perf_test.go @@ -16,11 +16,14 @@ func TestUnambiguousClaudeToolNames(t *testing.T) { `{"tools":[{"name":"Read"},{"name":"Read"}]}`, `{"tools":[{"name":" " ,"function":{"name":"Read"}},{"name":"___"},{"name":"读取"}]}`, `{"tools":[{"name":" Read "},{"function":{"name":"Write"}}],"tools":[]}`, + `{"messages":[{"role":"system","content":[{"type":"tool_addition","tool":{"type":"tool_definition","definition":{"name":"SubagentHandback"}}}]}]}`, + `{"messages":[{"role":"system","content":[{"type":"tool\u005faddition","tool":{"type":"tool_definition","definition":{"name":"SubagentHandback"}}}]}]}`, + `{"tools":[{"name":"Read"}],"messages":[{"role":"system","content":[{"type":"tool_addition","tool":{"type":"tool_definition","definition":{"name":"read"}}}]}]}`, } for _, request := range requests { legacy := ToolNameMapFromClaudeRequest([]byte(request)) cached := UnambiguousToolNameMapFromClaudeRequest([]byte(request)) - for _, name := range []string{"", "Read", "read", "READ", "_Read", "__read", " Read ", "Write", "write", "读取", "missing", "___"} { + for _, name := range []string{"", "Read", "read", "READ", "_Read", "__read", " Read ", "Write", "write", "SubagentHandback", "subagenthandback", "读取", "missing", "___"} { want := MapClaudeToolName(legacy, []byte(request), name) if got := MapToolName(cached, name); got != want { t.Errorf("request=%s name=%q got=%q want=%q", request, name, got, want) diff --git a/internal/util/claude_tool_result.go b/internal/util/claude_tool_result.go index 58554853561..dc2c3cd885c 100644 --- a/internal/util/claude_tool_result.go +++ b/internal/util/claude_tool_result.go @@ -13,20 +13,30 @@ type ClaudeToolResultImage struct { Data string } -// ClaudeToolResult is the normalized form of a Claude tool_result `content` field, +// ClaudeToolResult is the normalized form of a Claude tool_result block, // ready to be written into a Gemini-style functionResponse. type ClaudeToolResult struct { - // Result is the value for functionResponse.response.result. + // Result is the value for functionResponse.response.result or .error. Result string // ResultIsRaw reports whether Result holds raw JSON (write with sjson.SetRaw*) // or a plain string (write with sjson.Set*). Writing raw JSON text through // sjson.Set as a string value would double-encode it, so callers must honor // this flag. ResultIsRaw bool + // IsError preserves an explicit tool_result.is_error failure flag. + IsError bool // Images holds base64 image blocks separated out of the content. Images []ClaudeToolResultImage } +// ConvertClaudeToolResult normalizes a complete Claude tool_result block while +// preserving its success/failure semantics. +func ConvertClaudeToolResult(block gjson.Result) ClaudeToolResult { + result := ConvertClaudeToolResultContent(block.Get("content")) + result.IsError = block.Get("is_error").Type == gjson.True + return result +} + // ConvertClaudeToolResultContent normalizes a Claude tool_result `content` field into // a deterministic Gemini functionResponse result plus any extracted images. // diff --git a/internal/util/claude_tool_result_test.go b/internal/util/claude_tool_result_test.go index 6ac24081b67..fc1f309d7fc 100644 --- a/internal/util/claude_tool_result_test.go +++ b/internal/util/claude_tool_result_test.go @@ -6,6 +6,31 @@ import ( "github.com/tidwall/gjson" ) +func TestConvertClaudeToolResult_PreservesErrorFlag(t *testing.T) { + for _, tt := range []struct { + name string + block string + isError bool + }{ + {"omitted", `{"content":"permission denied"}`, false}, + {"false", `{"is_error":false,"content":"permission denied"}`, false}, + {"true", `{"is_error":true,"content":"permission denied"}`, true}, + {"null", `{"is_error":null,"content":"permission denied"}`, false}, + {"string", `{"is_error":"true","content":"permission denied"}`, false}, + {"number", `{"is_error":1,"content":"permission denied"}`, false}, + } { + t.Run(tt.name, func(t *testing.T) { + got := ConvertClaudeToolResult(gjson.Parse(tt.block)) + if got.IsError != tt.isError { + t.Errorf("IsError = %v, want %v", got.IsError, tt.isError) + } + if got.Result != "permission denied" || got.ResultIsRaw || len(got.Images) != 0 { + t.Fatalf("tool result content changed: %+v", got) + } + }) + } +} + func TestConvertClaudeToolResultContent(t *testing.T) { tests := []struct { name string diff --git a/internal/util/nocopy_invariant_test.go b/internal/util/nocopy_invariant_test.go index 9dbfb5ceec2..2f0ae46d450 100644 --- a/internal/util/nocopy_invariant_test.go +++ b/internal/util/nocopy_invariant_test.go @@ -105,6 +105,7 @@ type reviewedInPlaceByteWrite struct { var reviewedInPlaceByteWrites = map[string]reviewedInPlaceByteWrite{ "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"}, "internal/runtime/executor/claude_executor_cloaking.go": {1, "shifts []string headers to prepend a block; no byte of any payload is rewritten"}, "internal/runtime/executor/claude_executor_request.go": {3, "shifts []string headers to insert a part; no byte of any payload is rewritten"}, diff --git a/internal/util/translator.go b/internal/util/translator.go index d605918234e..75b4a5995a5 100644 --- a/internal/util/translator.go +++ b/internal/util/translator.go @@ -263,29 +263,71 @@ func UnambiguousToolNameMapFromClaudeRequest(rawJSON []byte) map[string]string { return toolNameMapFromClaudeRequest(rawJSON, true) } +func claudePayloadMayContainToolAddition(rawJSON []byte) bool { + return bytes.Contains(rawJSON, []byte("tool_addition")) || + bytes.Contains(rawJSON, []byte(`\u`)) || + bytes.Contains(rawJSON, []byte(`\U`)) +} + +func forEachClaudeInlineToolAdditionDefinition(rawJSON []byte, fn func(def gjson.Result)) { + if !claudePayloadMayContainToolAddition(rawJSON) { + return + } + messages := gjson.GetBytes(rawJSON, "messages") + if !messages.IsArray() { + return + } + messages.ForEach(func(_, message gjson.Result) bool { + if message.Get("role").String() != "system" { + return true + } + content := message.Get("content") + if !content.IsArray() { + return true + } + content.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() != "tool_addition" { + return true + } + toolSpec := part.Get("tool") + if toolSpec.Get("type").String() == "tool_definition" || toolSpec.Get("definition").IsObject() { + if def := toolSpec.Get("definition"); def.IsObject() { + fn(def) + } + } + return true + }) + return true + }) +} + func toolNameMapFromClaudeRequest(rawJSON []byte, rejectAmbiguous bool) map[string]string { if len(rawJSON) == 0 || !gjson.ValidBytes(rawJSON) { return nil } tools := gjson.GetBytes(rawJSON, "tools") - if !tools.Exists() || !tools.IsArray() { + hasTools := tools.Exists() && tools.IsArray() + hasToolAddition := claudePayloadMayContainToolAddition(rawJSON) + if !hasTools && !hasToolAddition { return nil } - toolResults := tools.Array() - out := make(map[string]string, len(toolResults)) - tools.ForEach(func(_, tool gjson.Result) bool { + var out map[string]string + recordTool := func(tool gjson.Result) { name := strings.TrimSpace(tool.Get("name").String()) if name == "" { name = strings.TrimSpace(tool.Get("function.name").String()) } if name == "" { - return true + return } key := CanonicalToolName(name) if key == "" { - return true + return + } + if out == nil { + out = make(map[string]string, 4) } if previous, exists := out[key]; !exists { out[key] = name @@ -294,8 +336,17 @@ func toolNameMapFromClaudeRequest(rawJSON []byte, rejectAmbiguous bool) map[stri // an ambiguous canonical key, including exact declared names. out[key] = "" } - return true - }) + } + + if hasTools { + tools.ForEach(func(_, tool gjson.Result) bool { + recordTool(tool) + return true + }) + } + if hasToolAddition { + forEachClaudeInlineToolAdditionDefinition(rawJSON, recordTool) + } if len(out) == 0 { return nil @@ -325,21 +376,39 @@ func MapClaudeToolName(toolNameMap map[string]string, originalRequestRawJSON []b } if len(originalRequestRawJSON) > 0 && gjson.ValidBytes(originalRequestRawJSON) { tools := gjson.GetBytes(originalRequestRawJSON, "tools") - if tools.IsArray() { + hasTools := tools.IsArray() + hasToolAddition := claudePayloadMayContainToolAddition(originalRequestRawJSON) + if hasTools || hasToolAddition { canonical := CanonicalToolName(name) canonicalNames := map[string]bool{} - for _, tool := range tools.Array() { + exactMatch := false + inspectTool := func(tool gjson.Result) { declared := strings.TrimSpace(tool.Get("name").String()) if declared == "" { declared = strings.TrimSpace(tool.Get("function.name").String()) } if declared == name { - return name + exactMatch = true + return } if declared != "" && CanonicalToolName(declared) == canonical { canonicalNames[declared] = true } } + if hasTools { + for _, tool := range tools.Array() { + inspectTool(tool) + if exactMatch { + return name + } + } + } + if hasToolAddition { + forEachClaudeInlineToolAdditionDefinition(originalRequestRawJSON, inspectTool) + if exactMatch { + return name + } + } if len(canonicalNames) > 1 { return name } diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index 34130f76280..742c1573b95 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -334,13 +334,13 @@ func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx co RequestedModel: originalRequestedModel, Stream: opts.Stream, Headers: cloneHeader(opts.Headers), - Body: cloneBytes(req.Payload), + Body: req.Payload, Metadata: opts.Metadata, }, requestID, skipPluginID) opts.Headers = mergeRequestInterceptorHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) if len(resp.Body) > 0 { - req.Payload = cloneBytes(resp.Body) - opts.OriginalRequest = cloneBytes(resp.Body) + req.Payload = resp.Body + opts.OriginalRequest = resp.Body } if path := strings.TrimSpace(resp.Path); path != "" { if opts.Metadata == nil { diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go index 330ab51329b..16203c447b7 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -393,8 +393,8 @@ func (c *requestAfterAuthCapture) record(req coreexecutor.RequestAfterAuthInterc var originalRequest []byte originalRequestReplaced := false if len(resp.Body) > 0 { - body = cloneBytes(resp.Body) - originalRequest = cloneBytes(resp.Body) + body = resp.Body + originalRequest = resp.Body originalRequestReplaced = true } path := strings.TrimSpace(resp.Path) @@ -419,8 +419,8 @@ func (c *requestAfterAuthCapture) apply(req coreexecutor.Request, opts coreexecu return req, opts } if c.originalRequestReplaced { - req.Payload = cloneBytes(c.body) - opts.OriginalRequest = cloneBytes(c.originalRequest) + req.Payload = c.body + opts.OriginalRequest = c.originalRequest } opts.Headers = cloneHeader(c.headers) if c.path != "" { @@ -509,8 +509,8 @@ func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, }, skipPluginID) opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers) if len(resp.Body) > 0 { - req.Payload = cloneBytes(resp.Body) - opts.OriginalRequest = cloneBytes(resp.Body) + req.Payload = resp.Body + opts.OriginalRequest = resp.Body } if strings.TrimSpace(resp.Path) != "" { if opts.Metadata == nil { @@ -526,6 +526,41 @@ func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, return req, opts, nil } +func (h *BaseAPIHandler) requestBeforeAuthInterceptor(requestID, skipPluginID string) coreexecutor.RequestBeforeAuthInterceptor { + if !requestInterceptorsEnabled(h.interceptorHost()) { + return nil + } + return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { + host := h.interceptorHost() + if !requestInterceptorsEnabled(host) { + return coreexecutor.RequestAfterAuthInterceptResponse{} + } + resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{ + RequestID: requestID, + TraceID: logging.GetRequestID(ctx), + SourceFormat: req.SourceFormat.String(), + ToFormat: "", + Model: req.Model, + RequestedModel: req.RequestedModel, + Stream: req.Stream, + Headers: cloneHeader(req.Headers), + Body: req.Body, + CallerScope: callerScopeFromMetadata(req.Metadata), + Metadata: req.Metadata, + }, skipPluginID) + return coreexecutor.RequestAfterAuthInterceptResponse{ + Path: strings.TrimSpace(resp.Path), + Headers: resp.Headers, + Body: resp.Body, + ClearHeaders: resp.ClearHeaders, + Terminate: resp.Terminate, + StatusCode: normalizedTerminationStatus(resp.StatusCode), + ResponseHeaders: resp.ResponseHeaders, + ResponseBody: resp.ResponseBody, + } + } +} + func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor { if !requestInterceptorsEnabled(h.interceptorHost()) { return nil diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index 816de409c45..86a30602308 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -333,16 +333,17 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context afterAuthCapture := &requestAfterAuthCapture{} lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, true, reqMeta, execOptions.SkipInterceptorPluginID) opts := coreexecutor.Options{ - Stream: true, - Alt: alt, - OriginalRequest: originalRequest, - SourceFormat: sdktranslator.FromString(entryProtocol), - ResponseFormat: sdktranslator.FromString(responseProtocol), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), - WebSocketResponseObserver: h.webSocketResponseObserver(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), - ProxyURL: execOptions.ProxyURL, + Stream: true, + Alt: alt, + OriginalRequest: originalRequest, + SourceFormat: sdktranslator.FromString(entryProtocol), + ResponseFormat: sdktranslator.FromString(responseProtocol), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + RequestBeforeAuthInterceptor: h.requestBeforeAuthInterceptor(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + WebSocketResponseObserver: h.webSocketResponseObserver(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + ProxyURL: execOptions.ProxyURL, } opts.Metadata = reqMeta ctx = enrichContextWithSessionHierarchy(ctx, opts.Headers, req.Payload, opts.Metadata) diff --git a/sdk/api/handlers/openai/codex_client_models_test.go b/sdk/api/handlers/openai/codex_client_models_test.go index 84e6258df53..364c4077cb5 100644 --- a/sdk/api/handlers/openai/codex_client_models_test.go +++ b/sdk/api/handlers/openai/codex_client_models_test.go @@ -46,6 +46,10 @@ func TestCodexClientModelsResponseMultiAgentV2FollowsConfig(t *testing.T) { if entry == nil { t.Fatalf("missing synthesized model %q", modelID) } + // The synthesized model inherits multi_agent_version from the + // catalog's default template, which currently carries "v2"; + // runtime code preserves it when the opt-in is disabled and + // rewrites it to "v2" when enabled. value, exists := entry["multi_agent_version"] if tt.enabled { if !exists || value != "v2" { @@ -53,8 +57,8 @@ func TestCodexClientModelsResponseMultiAgentV2FollowsConfig(t *testing.T) { } return } - if !exists || value != nil { - t.Fatalf("multi_agent_version = %#v, want preserved null", value) + if !exists || value != "v2" { + t.Fatalf("multi_agent_version = %#v, want preserved template value v2", value) } }) } diff --git a/sdk/api/handlers/openai/openai_images_handlers.go b/sdk/api/handlers/openai/openai_images_handlers.go index 6445eebdf3c..c01d5c8a350 100644 --- a/sdk/api/handlers/openai/openai_images_handlers.go +++ b/sdk/api/handlers/openai/openai_images_handlers.go @@ -52,7 +52,8 @@ type imageCallResult struct { } type sseFrameAccumulator struct { - pending []byte + pending []byte + syntheticBreaks []int } type xaiImageResult struct { @@ -151,33 +152,34 @@ func (a *sseFrameAccumulator) AddChunk(chunk []byte) [][]byte { var frames [][]byte if responsesSSEStartsNewDataFrame(a.pending, chunk) { - frames = append(frames, bytes.Clone(a.pending)) + frames = append(frames, bytes.Clone(responsesSSERecoverJSONFrame(a.pending, a.syntheticBreaks))) a.pending = a.pending[:0] + a.syntheticBreaks = a.syntheticBreaks[:0] } - if responsesSSENeedsLineBreak(a.pending, chunk) { - a.pending = append(a.pending, '\n') - } - a.pending = append(a.pending, chunk...) + a.pending = responsesSSEAppendChunk(a.pending, &a.syntheticBreaks, chunk) for { frameLen := responsesSSEFrameLen(a.pending) if frameLen == 0 { break } - frames = append(frames, bytes.Clone(a.pending[:frameLen])) + frames = append(frames, bytes.Clone(responsesSSERecoverJSONFrame(a.pending[:frameLen], a.syntheticBreaks))) copy(a.pending, a.pending[frameLen:]) a.pending = a.pending[:len(a.pending)-frameLen] + a.syntheticBreaks = responsesSSEConsumeBreaks(a.syntheticBreaks, frameLen) } if len(bytes.TrimSpace(a.pending)) == 0 { a.pending = a.pending[:0] + a.syntheticBreaks = a.syntheticBreaks[:0] return frames } if len(a.pending) == 0 || !responsesSSECanEmitWithoutDelimiter(a.pending) { return frames } - frames = append(frames, bytes.Clone(a.pending)) + frames = append(frames, bytes.Clone(responsesSSERecoverJSONFrame(a.pending, a.syntheticBreaks))) a.pending = a.pending[:0] + a.syntheticBreaks = a.syntheticBreaks[:0] return frames } @@ -192,19 +194,22 @@ func (a *sseFrameAccumulator) Flush() [][]byte { if frameLen == 0 { break } - frames = append(frames, bytes.Clone(a.pending[:frameLen])) + frames = append(frames, bytes.Clone(responsesSSERecoverJSONFrame(a.pending[:frameLen], a.syntheticBreaks))) copy(a.pending, a.pending[frameLen:]) a.pending = a.pending[:len(a.pending)-frameLen] + a.syntheticBreaks = responsesSSEConsumeBreaks(a.syntheticBreaks, frameLen) } if len(bytes.TrimSpace(a.pending)) == 0 { a.pending = nil return frames } - if responsesSSECanFlushWithoutDelimiter(a.pending) { - frames = append(frames, bytes.Clone(a.pending)) + frame := responsesSSERecoverJSONFrame(a.pending, a.syntheticBreaks) + if responsesSSECanFlushWithoutDelimiter(frame) { + frames = append(frames, bytes.Clone(frame)) } a.pending = nil + a.syntheticBreaks = nil return frames } diff --git a/sdk/api/handlers/openai/openai_responses_framing_metadata_test.go b/sdk/api/handlers/openai/openai_responses_framing_metadata_test.go new file mode 100644 index 00000000000..fbd67a358ef --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_framing_metadata_test.go @@ -0,0 +1,321 @@ +package openai + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/interfaces" + "github.com/tidwall/gjson" +) + +func TestResponsesSSEFramerFragmentedMetadataDataField(t *testing.T) { + payload := `{"type":"response.metadata","metadata":{"marker":"private"}}` + for _, tc := range []struct { + name string + lineEnding string + isCodexClient bool + }{ + {name: "standard-lf", lineEnding: "\n"}, + {name: "codex-lf", lineEnding: "\n", isCodexClient: true}, + {name: "standard-crlf", lineEnding: "\r\n"}, + {name: "codex-crlf", lineEnding: "\r\n", isCodexClient: true}, + } { + t.Run(tc.name, func(t *testing.T) { + prefix := "event: response.metadata" + tc.lineEnding + "data:" + frame := []byte(prefix + " " + payload + tc.lineEnding + tc.lineEnding) + for _, split := range []int{len(prefix), len(prefix) + 1} { + var output bytes.Buffer + framer := &responsesSSEFramer{isCodexClient: tc.isCodexClient} + framer.WriteChunk(&output, frame[:split]) + framer.WriteChunk(&output, frame[split:]) + framer.WriteChunk(&output, []byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n")) + got := output.String() + if tc.isCodexClient { + if !strings.Contains(got, prefix+" "+payload) { + t.Fatalf("split %d: Codex metadata missing or damaged: %q", split, got) + } + } else if strings.Contains(got, "response.metadata") || strings.Contains(got, "private") { + t.Fatalf("split %d: private metadata leaked: %q", split, got) + } + if !strings.Contains(got, "response.completed") { + t.Fatalf("split %d: completion missing: %q", split, got) + } + } + }) + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeFragmentedCompletion(t *testing.T) { + completed := []byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n") + for split := 1; split <= len("event:"); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 3) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- completed[:split] + data <- completed[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") { + t.Fatalf("split %d: empty metadata consumed completion: %q", split, got) + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeLeadingSeparator(t *testing.T) { + for _, prefix := range []string{"\n", "\r\n", ": keepalive\n", "id: event-42\n", "retry: 1000\n", "x-extra: value\n"} { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 2) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- []byte(prefix + "event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n") + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "response.completed") { + t.Fatalf("prefix %q: empty metadata consumed completion: %q", prefix, got) + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeFragmentedFields(t *testing.T) { + for _, prefix := range []string{"id: event-42\n", "retry: 1000\n", "x-extra: value\n"} { + field := []byte(prefix + "event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n") + for split := 1; split < strings.IndexByte(prefix, ':')+1; split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 3) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- field[:split] + data <- field[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "response.completed") { + t.Fatalf("prefix %q split %d: empty metadata consumed completion: %q", prefix, split, got) + } + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeDataOnlyCompletion(t *testing.T) { + completed := []byte("data: {\"type\":\"response.completed\"}\n\n") + for split := 1; split <= len("data:"); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 3) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- completed[:split] + data <- completed[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "response.completed") { + t.Fatalf("split %d: empty metadata consumed data-only completion: %q", split, got) + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeDataThenEvent(t *testing.T) { + created := []byte(`data: {"response":{"id":"resp-1","status":"completed"}}`) + completed := []byte("event: response.completed\n\n") + for split := 1; split <= len("event:"); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 4) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- created + data <- completed[:split] + data <- completed[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") || !strings.Contains(got, "resp-1") { + t.Fatalf("split %d: data-before-event completion lost: %q", split, got) + } + } + + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 2) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- append(append(bytes.Clone(created), '\n'), completed...) + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") || !strings.Contains(got, "resp-1") { + t.Fatalf("combined data-before-event completion lost: %q", got) + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeDataFieldThenEvent(t *testing.T) { + for _, field := range []string{"id: event-42", "retry: 1000", ": keepalive", "x-extra: value"} { + for split := 1; split <= len(field); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 5) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- []byte(`data: {"response":{"id":"resp-1","status":"completed"}}`) + data <- []byte(field[:split]) + data <- []byte(field[split:]) + data <- []byte("event: response.completed\n\n") + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") || !strings.Contains(got, "resp-1") { + t.Fatalf("field %q split %d: data-before-event completion lost: %q", field, split, got) + } + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeDataAndFragmentedEventAfterLineBreak(t *testing.T) { + for _, field := range []string{"", "id: event-42\n", ": keepalive\n"} { + for split := 1; split < len("event:"); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 5) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + data <- []byte("data: {\"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n") + if field != "" { + data <- []byte(field) + } + event := []byte("event: response.completed\n\n") + data <- event[:split] + data <- event[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") || !strings.Contains(got, "resp-1") { + t.Fatalf("field %q split %d: completion lost: %q", field, split, got) + } + } + } +} + +func TestForwardResponsesStreamEmptyMetadataBeforeDataWithFragmentedEvent(t *testing.T) { + for _, field := range []string{"", "id: event-42\n", ": keepalive\n"} { + for split := 1; split < len("event:"); split++ { + h, recorder, c, flusher := newResponsesStreamTestHandler(t) + data := make(chan []byte, 3) + errs := make(chan *interfaces.ErrorMessage) + data <- []byte("event: response.metadata\ndata:") + event := []byte("event: response.completed\n\n") + data <- append([]byte("data: {\"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n"+field), event[:split]...) + data <- event[split:] + close(data) + close(errs) + h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil) + got := recorder.Body.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "event: error") || !strings.Contains(got, "event: response.completed") || !strings.Contains(got, "resp-1") { + t.Fatalf("field %q split %d: completion lost: %q", field, split, got) + } + } + } +} + +func TestResponsesSSEFramerPreservesEmptyEvent(t *testing.T) { + var output bytes.Buffer + framer := &responsesSSEFramer{} + framer.WriteChunk(&output, []byte("event: custom\ndata:")) + if output.Len() != 0 { + t.Fatalf("empty data field was emitted before completion: %q", output.String()) + } + framer.WriteChunk(&output, []byte("\n\n")) + if got, want := output.String(), "event: custom\ndata:\n\n"; got != want { + t.Fatalf("empty data frame = %q, want %q", got, want) + } +} + +func TestResponsesSSEFramerMetadataWithFieldPrefixInJSONString(t *testing.T) { + for _, marker := range []string{"data: private", "event: private", "id: private", "retry: private", ": private", `escaped " data: private`, `escaped \ event: private`} { + for _, isCodexClient := range []bool{false, true} { + encodedMarker, errMarshal := json.Marshal(marker) + if errMarshal != nil { + t.Fatalf("marshal marker: %v", errMarshal) + } + payload := `{"type":"response.metadata","metadata":{"marker":` + string(encodedMarker) + `}}` + frame := []byte("data: " + payload + "\n\n") + field := strings.Index(string(frame), strings.TrimPrefix(marker, `escaped " `)) + if field < 0 { + field = strings.Index(string(frame), "event: private") + } + if field < 0 { + t.Fatalf("marker %q not found in %q", marker, frame) + } + var output bytes.Buffer + framer := &responsesSSEFramer{isCodexClient: isCodexClient} + framer.WriteChunk(&output, frame[:field]) + framer.WriteChunk(&output, frame[field:]) + framer.WriteChunk(&output, []byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n")) + got := output.String() + if isCodexClient { + if !strings.Contains(got, payload) { + t.Fatalf("marker %q: Codex payload corrupted: %q", marker, got) + } + } else if strings.Contains(got, "response.metadata") || strings.Contains(got, "private") { + t.Fatalf("marker %q: metadata leaked: %q", marker, got) + } + if !strings.Contains(got, "response.completed") { + t.Fatalf("marker %q: completion missing: %q", marker, got) + } + } + } +} + +func TestResponsesSSEFramerFiltersMalformedPrivatePayload(t *testing.T) { + var output bytes.Buffer + framer := &responsesSSEFramer{} + framer.WriteChunk(&output, []byte("data: {\"type\":\"response.metadata\",\"metadata\":{\"marker\":\"private\"}\n\n")) + framer.WriteChunk(&output, []byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n")) + if got := output.String(); strings.Contains(got, "response.metadata") || strings.Contains(got, "private") { + t.Fatalf("malformed private payload leaked: %q", got) + } +} + +func TestResponsesSSEFramerKeepsEmptyCommentInMultilineData(t *testing.T) { + var output bytes.Buffer + framer := &responsesSSEFramer{} + for _, chunk := range []string{ + `data: {"type"`, + ":", + `data: :"response.metadata","metadata":{"marker":"private"}}` + "\n\n", + "event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n", + } { + framer.WriteChunk(&output, []byte(chunk)) + } + got := output.String() + if strings.Contains(got, "response.metadata") || strings.Contains(got, "private") || !strings.Contains(got, "response.completed") { + t.Fatalf("metadata with empty comment leaked: %q", got) + } +} + +func TestSSEFrameAccumulatorKeepsEmptyCommentInMultilineData(t *testing.T) { + acc := &sseFrameAccumulator{} + var frames [][]byte + for _, chunk := range []string{ + `data: {"type"`, + ":", + `data: :"response.completed","response":{"output":[]}}` + "\n\n", + } { + frames = append(frames, acc.AddChunk([]byte(chunk))...) + } + if len(frames) != 1 { + t.Fatalf("frames = %q, want one", frames) + } + payload, ok := responsesSSEDataPayload(frames[0]) + if !ok || !json.Valid(payload) || gjson.GetBytes(payload, "type").String() != "response.completed" { + t.Fatalf("empty comment corrupted image data: %q", frames[0]) + } +} diff --git a/sdk/api/handlers/openai/openai_responses_framing_recovery_test.go b/sdk/api/handlers/openai/openai_responses_framing_recovery_test.go new file mode 100644 index 00000000000..8511b432426 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_framing_recovery_test.go @@ -0,0 +1,183 @@ +package openai + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + + "github.com/tidwall/gjson" +) + +func TestResponsesSSEFramerRecoversFragmentedJSON(t *testing.T) { + for _, tc := range []struct { + name string + chunks []string + }{ + { + name: "multiline-json", + chunks: []string{ + `data: {"type"`, + `:"response.metadata",` + "\ndata: " + `"metadata":{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "multiple-colons", + chunks: []string{ + `data: {"type"`, + `:"response.metadata","metadata"`, + `:{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "real-comment", + chunks: []string{ + `data: {"type"`, + `:"response.metadata",` + "\n: keepalive\ndata: " + `"metadata":{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "split-data-field", + chunks: []string{ + `data`, + `: {"type":"response.metadata","metadata":{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "multiple-colons-with-comment", + chunks: []string{ + `data: {"type"`, + `:"response.metadata",` + "\n: keepalive\ndata: " + `"metadata"`, + `:{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "separate-comment-with-split-key", + chunks: []string{ + `data: {"type"`, + `: keepalive`, + `data: :"response.metadata","metadata"`, + `:{"marker":"private"}}` + "\n\n", + }, + }, + { + name: "colon-before-string-field", + chunks: []string{ + `data: {"type":"response.metadata","metadata":{"marker"`, + `:"`, + `data: private"}}` + "\n\n", + }, + }, + } { + for _, isCodexClient := range []bool{false, true} { + t.Run(tc.name, func(t *testing.T) { + var output bytes.Buffer + framer := &responsesSSEFramer{isCodexClient: isCodexClient} + for _, chunk := range tc.chunks { + framer.WriteChunk(&output, []byte(chunk)) + } + framer.WriteChunk(&output, []byte("event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n")) + got := output.String() + if isCodexClient { + frames := strings.Split(got, "\n\n") + payload, ok := responsesSSEDataPayload([]byte(frames[0])) + wantMarker := "private" + if tc.name == "colon-before-string-field" { + wantMarker = "data: private" + } + if !ok || !json.Valid(payload) || gjson.GetBytes(payload, "type").String() != "response.metadata" || gjson.GetBytes(payload, "metadata.marker").String() != wantMarker { + t.Fatalf("Codex metadata was corrupted: %q", got) + } + if strings.Contains(strings.Join(tc.chunks, ""), ": keepalive") && !strings.Contains(got, ": keepalive") { + t.Fatalf("genuine SSE comment was removed: %q", got) + } + } else if strings.Contains(got, "response.metadata") || strings.Contains(got, "private") { + t.Fatalf("private metadata leaked: %q", got) + } + if !strings.Contains(got, "response.completed") { + t.Fatalf("completion missing: %q", got) + } + }) + } + } +} + +func TestSSEFrameAccumulatorRecoversFragmentedJSON(t *testing.T) { + for _, tc := range []struct { + name string + chunks []string + }{ + { + name: "multiline-json", + chunks: []string{ + `data: {"type"`, + `:"response.completed",` + "\ndata: " + `"response":{"output":[]}}` + "\n\n", + }, + }, + { + name: "multiple-colons", + chunks: []string{ + `data: {"type"`, + `:"response.completed","response"`, + `:{"output":[]}}` + "\n\n", + }, + }, + { + name: "real-comment", + chunks: []string{ + `data: {"type"`, + `:"response.completed",` + "\n: keepalive\ndata: " + `"response":{"output":[]}}` + "\n\n", + }, + }, + { + name: "split-data-field", + chunks: []string{ + `data`, + `: {"type":"response.completed","response":{"output":[]}}` + "\n\n", + }, + }, + { + name: "multiple-colons-with-comment", + chunks: []string{ + `data: {"type"`, + `:"response.completed",` + "\n: keepalive\ndata: " + `"response"`, + `:{"output":[]}}` + "\n\n", + }, + }, + { + name: "separate-comment-with-split-key", + chunks: []string{ + `data: {"type"`, + `: keepalive`, + `data: :"response.completed","response"`, + `:{"output":[]}}` + "\n\n", + }, + }, + { + name: "colon-before-string-field", + chunks: []string{ + `data: {"type":"response.completed","metadata":{"marker"`, + `:"`, + `data: private"}}` + "\n\n", + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + acc := &sseFrameAccumulator{} + var frames [][]byte + for _, chunk := range tc.chunks { + frames = append(frames, acc.AddChunk([]byte(chunk))...) + } + if len(frames) != 1 { + t.Fatalf("frames = %q, want one", frames) + } + payload, ok := responsesSSEDataPayload(frames[0]) + if !ok || !json.Valid(payload) || gjson.GetBytes(payload, "type").String() != "response.completed" { + t.Fatalf("image response was corrupted: %q", frames[0]) + } + if strings.Contains(strings.Join(tc.chunks, ""), ": keepalive") && !bytes.Contains(frames[0], []byte(": keepalive")) { + t.Fatalf("genuine SSE comment was removed: %q", frames[0]) + } + }) + } +} diff --git a/sdk/api/handlers/openai/openai_responses_handlers.go b/sdk/api/handlers/openai/openai_responses_handlers.go index d0396925c41..0c0eac68eb9 100644 --- a/sdk/api/handlers/openai/openai_responses_handlers.go +++ b/sdk/api/handlers/openai/openai_responses_handlers.go @@ -64,6 +64,9 @@ func flushResponsesSSE(w http.ResponseWriter) error { type responsesSSEFramer struct { pending []byte + nextFramePrefix []byte + nextFieldSuffix []byte + syntheticBreaks []int outputItems map[int][]byte outputOrder []int unindexedOutputItems [][]byte @@ -81,17 +84,102 @@ func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) { if len(chunk) == 0 || f.terminalEvent != "" { return } + if len(f.nextFieldSuffix) > 0 { + trimmed := bytes.TrimLeft(chunk, " \t\r\n") + _, complete, _ := responsesSSEChunkField(trimmed) + lastLine := bytes.TrimSpace(f.nextFramePrefix[bytes.LastIndexByte(f.nextFramePrefix, '\n')+1:]) + if complete && len(lastLine) > 0 && !bytes.HasPrefix(lastLine, []byte("data:")) { + f.nextFramePrefix = append(f.nextFramePrefix, f.nextFieldSuffix...) + f.nextFieldSuffix = nil + } else { + f.nextFieldSuffix = append(f.nextFieldSuffix, chunk...) + trimmed = bytes.TrimLeft(f.nextFieldSuffix, " \t\r\n") + _, _, partial := responsesSSEChunkField(trimmed) + if partial { + return + } + chunk = f.nextFieldSuffix + f.nextFieldSuffix = nil + } + } + if len(f.nextFramePrefix) > 0 { + payload, found := responsesSSEDataPayload(f.nextFramePrefix) + if found && json.Valid(payload) && gjson.GetBytes(payload, "type").String() == "" && + !responsesSSEHasField(f.nextFramePrefix, []byte("event:")) && responsesSSEFrameLen(f.nextFramePrefix) == 0 { + trimmed := bytes.TrimLeft(chunk, " \t\r\n") + field, complete, partial := responsesSSEChunkField(trimmed) + lastLine := bytes.TrimSpace(f.nextFramePrefix[bytes.LastIndexByte(f.nextFramePrefix, '\n')+1:]) + _, _, lastPartial := responsesSSEChunkField(lastLine) + if complete && lastPartial { + joined := append(bytes.Clone(lastLine), trimmed...) + joinedField, joinedComplete, _ := responsesSSEChunkField(joined) + if joinedComplete && (joinedField == "event" || joinedField == "data" || joinedField == "id" || joinedField == "retry" || + (field != "" && field != "event" && field != "data" && field != "id" && field != "retry")) { + complete = false + } + } + newFieldPrefix := bytes.HasPrefix(lastLine, []byte("data:")) + for _, prefix := range [][]byte{[]byte("event:"), []byte("data:"), []byte("id:"), []byte("retry:")} { + newFieldPrefix = newFieldPrefix || (len(trimmed) < len(prefix) && bytes.HasPrefix(prefix, trimmed)) + } + if partial && newFieldPrefix { + f.nextFieldSuffix = append(f.nextFieldSuffix, chunk...) + return + } + if complete && !bytes.HasSuffix(f.nextFramePrefix, []byte("\n")) && !bytes.HasSuffix(f.nextFramePrefix, []byte("\r")) && + chunk[0] != '\n' && chunk[0] != '\r' { + f.nextFramePrefix = append(f.nextFramePrefix, '\n') + } + } + f.nextFramePrefix = append(f.nextFramePrefix, chunk...) + chunk = f.nextFramePrefix + f.nextFramePrefix = nil + } + if responsesSSEHasEmptyEventData(f.pending) { + trimmed := bytes.TrimLeft(chunk, " \t\r\n") + field, complete, partial := responsesSSEChunkField(trimmed) + if partial { + f.nextFramePrefix = append(f.nextFramePrefix, chunk...) + return + } + newFrame := complete && field != "data" + if complete && field == "data" { + payload, _ := responsesSSEDataPayload(trimmed) + valid := json.Valid(payload) + if !valid && !bytes.Equal(payload, []byte("[DONE]")) && responsesSSEFrameLen(trimmed) == 0 { + f.nextFramePrefix = append(f.nextFramePrefix, chunk...) + return + } + hasEvent := responsesSSEHasField(trimmed, []byte("event:")) + if valid && hasEvent && responsesSSEEventName(trimmed) == "" && responsesSSEFrameLen(trimmed) == 0 { + f.nextFramePrefix = append(f.nextFramePrefix, chunk...) + return + } + newFrame = bytes.Equal(payload, []byte("[DONE]")) || + (valid && (gjson.GetBytes(payload, "type").String() != "" || hasEvent)) + if valid && !newFrame && responsesSSEFrameLen(trimmed) == 0 { + f.nextFramePrefix = append(f.nextFramePrefix, chunk...) + return + } + } + if newFrame { + f.writeFrame(w, f.pending) + f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] + if f.terminalEvent != "" { + return + } + } + } if responsesSSEStartsNewDataFrame(f.pending, chunk) { f.writeFrame(w, f.pending) f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] if f.terminalEvent != "" { return } } - if responsesSSENeedsLineBreak(f.pending, chunk) { - f.pending = append(f.pending, '\n') - } - f.pending = append(f.pending, chunk...) + f.pending = responsesSSEAppendChunk(f.pending, &f.syntheticBreaks, chunk) for { frameLen := responsesSSEFrameLen(f.pending) if frameLen == 0 { @@ -100,6 +188,7 @@ func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) { f.writeFrame(w, f.pending[:frameLen]) copy(f.pending, f.pending[frameLen:]) f.pending = f.pending[:len(f.pending)-frameLen] + f.syntheticBreaks = responsesSSEConsumeBreaks(f.syntheticBreaks, frameLen) if f.terminalEvent != "" { f.pending = f.pending[:0] return @@ -107,16 +196,36 @@ func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) { } if len(bytes.TrimSpace(f.pending)) == 0 { f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] return } - if len(f.pending) == 0 || !responsesSSECanEmitWithoutDelimiter(f.pending) { + if len(f.pending) == 0 || responsesSSEHasEmptyEventData(f.pending) || !responsesSSECanEmitWithoutDelimiter(f.pending) { return } f.writeFrame(w, f.pending) f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] } func (f *responsesSSEFramer) Flush(w io.Writer) { + if len(f.nextFieldSuffix) > 0 { + if responsesSSENeedsLineBreak(f.nextFramePrefix, f.nextFieldSuffix) { + f.nextFramePrefix = append(f.nextFramePrefix, '\n') + } + f.nextFramePrefix = append(f.nextFramePrefix, f.nextFieldSuffix...) + f.nextFieldSuffix = nil + } + if len(f.nextFramePrefix) > 0 { + trimmed := bytes.TrimLeft(f.nextFramePrefix, " \t\r\n") + _, _, partial := responsesSSEChunkField(trimmed) + if partial && !bytes.HasPrefix([]byte("data"), trimmed) { + f.writeFrame(w, f.pending) + f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] + } + f.pending = responsesSSEAppendChunk(f.pending, &f.syntheticBreaks, f.nextFramePrefix) + f.nextFramePrefix = nil + } if len(f.pending) == 0 || f.terminalEvent != "" { return } @@ -124,17 +233,19 @@ func (f *responsesSSEFramer) Flush(w io.Writer) { f.pending = f.pending[:0] return } - if !responsesSSECanFlushWithoutDelimiter(f.pending) { + frame := responsesSSERecoverJSONFrame(f.pending, f.syntheticBreaks) + if !responsesSSECanFlushWithoutDelimiter(frame) { f.pending = f.pending[:0] return } - f.writeFrame(w, f.pending) + f.writeFrame(w, frame) f.pending = f.pending[:0] + f.syntheticBreaks = f.syntheticBreaks[:0] } func (f *responsesSSEFramer) writeFrame(w io.Writer, frame []byte) { if f.writeErr == nil { - f.writeErr = writeResponsesSSEChunk(w, f.repairFrame(frame)) + f.writeErr = writeResponsesSSEChunk(w, f.repairFrame(responsesSSERecoverJSONFrame(frame, f.syntheticBreaks))) } } @@ -161,8 +272,10 @@ func (f *responsesSSEFramer) shouldFilterPrivateEvent(streamEvent, payloadType s return false } - // For standard Responses API clients: filter any codex.* private events. - if strings.HasPrefix(name, "codex.") { + // For standard Responses API clients: filter any codex.* private events, and the Codex + // backend's moderation metadata, which is not a public Responses event: clients that parse + // a closed event list fail on it. + if strings.HasPrefix(name, "codex.") || name == "response.metadata" { return true } @@ -185,14 +298,13 @@ func (f *responsesSSEFramer) repairFrame(frame []byte) []byte { f.dataFrames++ return frame } - if !json.Valid(payload) { - return frame - } - payloadType := gjson.GetBytes(payload, "type").String() if f.shouldFilterPrivateEvent(streamEvent, payloadType) { return nil } + if !json.Valid(payload) { + return frame + } f.dataFrames++ @@ -331,6 +443,73 @@ func responsesSSEDataPayload(frame []byte) ([]byte, bool) { return payload, found } +func responsesSSERecoverJSONFrame(frame []byte, syntheticBreaks []int) []byte { + if len(syntheticBreaks) == 0 { + return frame + } + payload, found := responsesSSEDataPayload(frame) + if found && json.Valid(payload) { + return frame + } + + // Consider only newlines inserted at chunk boundaries when recovering invalid JSON. + positions := make([]int, 0, len(syntheticBreaks)) + for _, pos := range syntheticBreaks { + if pos >= 0 && pos < len(frame) && frame[pos] == '\n' { + positions = append(positions, pos) + } + } + try := func(candidate []byte) []byte { + recovered, ok := responsesSSEDataPayload(candidate) + if !ok || !json.Valid(recovered) { + return nil + } + return responsesSSEFrameWithData(candidate, recovered) + } + for i, pos := range positions { + if i == 8 { + break + } + candidate := make([]byte, 0, len(frame)-1) + candidate = append(candidate, frame[:pos]...) + candidate = append(candidate, frame[pos+1:]...) + if recovered := try(candidate); recovered != nil { + return recovered + } + } + if len(positions) > 1 { + candidate := make([]byte, 0, len(frame)-len(positions)) + start := 0 + for _, pos := range positions { + candidate = append(candidate, frame[start:pos]...) + start = pos + 1 + } + candidate = append(candidate, frame[start:]...) + if recovered := try(candidate); recovered != nil { + return recovered + } + } + return frame +} + +func responsesSSEAppendChunk(pending []byte, syntheticBreaks *[]int, chunk []byte) []byte { + if responsesSSENeedsLineBreak(pending, chunk) { + *syntheticBreaks = append(*syntheticBreaks, len(pending)) + pending = append(pending, '\n') + } + return append(pending, chunk...) +} + +func responsesSSEConsumeBreaks(syntheticBreaks []int, consumed int) []int { + remaining := syntheticBreaks[:0] + for _, pos := range syntheticBreaks { + if pos >= consumed { + remaining = append(remaining, pos-consumed) + } + } + return remaining +} + func responsesSSEFrameWithData(frame, payload []byte) []byte { var out bytes.Buffer for _, line := range bytes.Split(frame, []byte("\n")) { @@ -469,6 +648,29 @@ func responsesSSECanEmitWithoutDelimiter(chunk []byte) bool { return responsesSSEDataLinesValid(trimmed) } +func responsesSSEChunkField(chunk []byte) (name string, complete, partial bool) { + if len(chunk) == 0 { + return "", false, false + } + for i, b := range chunk { + if b == ':' || b == '\n' || b == '\r' { + return string(chunk[:i]), true, false + } + if !(b >= 'a' && b <= 'z') && !(b >= 'A' && b <= 'Z') && !(b >= '0' && b <= '9') && b != '-' && b != '_' && b != '.' { + return "", false, false + } + } + return "", false, true +} + +func responsesSSEHasEmptyEventData(chunk []byte) bool { + if !responsesSSEHasField(chunk, []byte("event:")) { + return false + } + payload, found := responsesSSEDataPayload(chunk) + return found && len(bytes.TrimSpace(payload)) == 0 +} + func responsesSSECanFlushWithoutDelimiter(chunk []byte) bool { trimmed := bytes.TrimSpace(chunk) return len(trimmed) > 0 && responsesSSEHasField(trimmed, []byte("data:")) && responsesSSEDataLinesValid(trimmed) @@ -519,12 +721,36 @@ func responsesSSENeedsLineBreak(pending, chunk []byte) bool { } for _, prefix := range [][]byte{[]byte("data:"), []byte("event:"), []byte("id:"), []byte("retry:"), []byte(":")} { if bytes.HasPrefix(trimmed, prefix) { - return true + return !responsesSSEDataEndsInsideJSONString(pending) } } return false } +func responsesSSEDataEndsInsideJSONString(chunk []byte) bool { + payload, found := responsesSSEDataPayload(chunk) + payload = bytes.TrimSpace(payload) + if !found || len(payload) == 0 || (payload[0] != '{' && payload[0] != '[' && payload[0] != '"') { + return false + } + inString := false + escaped := false + for _, b := range payload { + if escaped { + escaped = false + continue + } + if inString && b == '\\' { + escaped = true + continue + } + if b == '"' { + inString = !inString + } + } + return inString +} + // OpenAIResponsesAPIHandler contains the handlers for OpenAIResponses API endpoints. // It holds a pool of clients to interact with the backend service. type OpenAIResponsesAPIHandler struct { diff --git a/sdk/api/handlers/openai/openai_responses_handlers_stream_test.go b/sdk/api/handlers/openai/openai_responses_handlers_stream_test.go index df67ca282bd..203037b43fe 100644 --- a/sdk/api/handlers/openai/openai_responses_handlers_stream_test.go +++ b/sdk/api/handlers/openai/openai_responses_handlers_stream_test.go @@ -625,6 +625,9 @@ func TestResponsesSSEFramer_FiltersUpstreamPrivateEvents(t *testing.T) { subscriptionUsageFrame := []byte("event: response.subscription_usage\ndata: {\"subscription\":{\"tier\":\"1234567890123456\",\"weekly\":{\"resets_at\":1792368000,\"used_percent\":80},\"window\":{\"resets_at\":1792272519,\"used_percent\":10,\"window_duration_mins\":300}},\"type\":\"response.subscription_usage\"}\n\n") framer.WriteChunk(&output, subscriptionUsageFrame) + moderationFrame := []byte("event: response.metadata\ndata: {\"type\":\"response.metadata\",\"metadata\":{\"openai_chatgpt_moderation_metadata\":{\"prompt\":{}}}}\n\n") + framer.WriteChunk(&output, moderationFrame) + completedFrame := []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n\n") framer.WriteChunk(&output, completedFrame) @@ -632,6 +635,9 @@ func TestResponsesSSEFramer_FiltersUpstreamPrivateEvents(t *testing.T) { if strings.Contains(got, "codex.rate_limits") { t.Fatalf("expected codex.rate_limits to be stripped, got: %q", got) } + if strings.Contains(got, "event: response.metadata") || strings.Contains(got, "openai_chatgpt_moderation_metadata") { + t.Fatalf("expected response.metadata to be stripped, got: %q", got) + } if strings.Contains(got, "codex.response.metadata") { t.Fatalf("expected codex.response.metadata to be stripped, got: %q", got) } @@ -692,6 +698,9 @@ func TestResponsesSSEFramer_PreservesMetadataForCodexClient(t *testing.T) { createdFrame := []byte("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}\n\n") framer.WriteChunk(&output, createdFrame) + moderationFrame := []byte("event: response.metadata\ndata: {\"type\":\"response.metadata\",\"metadata\":{\"openai_chatgpt_moderation_metadata\":{\"prompt\":{}}}}\n\n") + framer.WriteChunk(&output, moderationFrame) + completedFrame := []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"status\":\"completed\"}}\n\n") framer.WriteChunk(&output, completedFrame) @@ -699,6 +708,9 @@ func TestResponsesSSEFramer_PreservesMetadataForCodexClient(t *testing.T) { if !strings.Contains(got, "codex.response.metadata") { t.Fatalf("expected codex.response.metadata to be preserved for official Codex client, got: %q", got) } + if !strings.Contains(got, "event: response.metadata") { + t.Fatalf("expected response.metadata to be preserved for official Codex client, got: %q", got) + } if !strings.Contains(got, "event: response.created") { t.Fatalf("expected response.created in output, got: %q", got) } diff --git a/sdk/api/handlers/openai/openai_responses_steering_auth_test.go b/sdk/api/handlers/openai/openai_responses_steering_auth_test.go index ef3e42b65ad..130d9042b52 100644 --- a/sdk/api/handlers/openai/openai_responses_steering_auth_test.go +++ b/sdk/api/handlers/openai/openai_responses_steering_auth_test.go @@ -1,6 +1,7 @@ package openai import ( + "bytes" "context" "net/http" "net/http/httptest" @@ -464,3 +465,209 @@ func TestResponsesSteeringInterceptorRewrittenBodyBindsClientModel(t *testing.T) t.Fatalf("upstream creates after explicit client alias continuation = %d, want 3", got) } } + +type environmentContextRewriteInterceptorHost struct { + handlers.PluginInterceptorHost + afterSawRemovedHeader atomic.Bool + afterMissingKept atomic.Bool +} + +func (*environmentContextRewriteInterceptorHost) HasStreamInterceptors() bool { return false } + +func (*environmentContextRewriteInterceptorHost) HasRequestInterceptors() bool { return true } + +func (*environmentContextRewriteInterceptorHost) InterceptRequestBeforeAuth(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + var headers http.Header + if req.Headers != nil { + headers = req.Headers.Clone() + headers.Del("X-Strip-In-Before") + } + if !bytes.Contains(req.Body, []byte("UTC")) { + return pluginapi.RequestInterceptResponse{Headers: headers} + } + patched := bytes.ReplaceAll(req.Body, []byte("UTC"), []byte("Asia/Shanghai")) + return pluginapi.RequestInterceptResponse{Headers: headers, Body: patched} +} + +func (h *environmentContextRewriteInterceptorHost) InterceptRequestAfterAuth(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + if req.Headers.Get("X-Strip-In-Before") != "" { + h.afterSawRemovedHeader.Store(true) + } + if req.Headers.Get("X-Keep-Header") != "kept" { + h.afterMissingKept.Store(true) + } + if !bytes.Contains(req.Body, []byte("2026-04-19")) { + return pluginapi.RequestInterceptResponse{ + Headers: http.Header{"X-Codex-Beta-Features": []string{"siwc"}}, + ClearHeaders: []string{"X-Codex-Beta-Features"}, + } + } + patched := bytes.ReplaceAll(req.Body, []byte("2026-04-19"), []byte("2026-04-20")) + return pluginapi.RequestInterceptResponse{ + Headers: http.Header{"X-Codex-Beta-Features": []string{"siwc"}}, + ClearHeaders: []string{"X-Codex-Beta-Features"}, + Body: patched, + } +} + +func TestResponsesSteeringSubsequentWebsocketMessagesRunRequestInterceptors(t *testing.T) { + gin.SetMode(gin.TestMode) + + const ( + authID = "codex-steering-env-rewrite-auth" + modelName = "steering-env-rewrite-model" + ) + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: modelName}}) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + upstreamFrames := make(chan []byte, 8) + upstreamWS := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer func() { _ = conn.Close() }() + for { + _, msg, errRead := conn.ReadMessage() + if errRead != nil { + return + } + upstreamFrames <- bytes.Clone(msg) + switch gjson.GetBytes(msg, "type").String() { + case "response.create", "response.append": + if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"resp_1"}}`)); err != nil { + return + } + if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_1"}}`)); err != nil { + return + } + case "response.steer": + if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.output_text.delta","delta":"steered"}`)); err != nil { + return + } + } + } + })) + defer upstreamWS.Close() + + cfg := &config.Config{} + cfg.Codex.ResponseSteering = true + cfg.CodexResponseSteering = true + + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + manager.RegisterExecutor(runtimeexecutor.NewCodexAutoExecutor(cfg)) + if _, err := manager.Register(context.Background(), &coreauth.Auth{ + ID: "codex-steering-env-rewrite-auth", + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "base_url": upstreamWS.URL, + "api_key": "upstream-token", + "websockets": "true", + }, + }); err != nil { + t.Fatalf("Register auth failed: %v", err) + } + + interceptorHost := &environmentContextRewriteInterceptorHost{} + baseHandler := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager) + baseHandler.SetPluginHost(interceptorHost) + responsesHandler := NewOpenAIResponsesAPIHandler(baseHandler) + + router := gin.New() + router.GET("/v1/responses", responsesHandler.ResponsesWebsocket) + proxyServer := httptest.NewServer(router) + defer proxyServer.Close() + + wsURL := "ws" + strings.TrimPrefix(proxyServer.URL, "http") + "/v1/responses" + dialHeaders := http.Header{ + "X-Strip-In-Before": []string{"should-be-removed"}, + "X-Keep-Header": []string{"kept"}, + } + c, resp, err := websocket.DefaultDialer.Dial(wsURL, dialHeaders) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer func() { _ = c.Close() }() + + assertRewrittenUpstreamFrame := func(step string) { + t.Helper() + select { + case frame := <-upstreamFrames: + if bytes.Contains(frame, []byte("UTC")) { + t.Fatalf("%s upstream frame still contains unpatched timezone: %s", step, frame) + } + if !bytes.Contains(frame, []byte("Asia/Shanghai")) { + t.Fatalf("%s upstream frame missing before-auth patched timezone: %s", step, frame) + } + if bytes.Contains(frame, []byte("2026-04-19")) { + t.Fatalf("%s upstream frame still contains unpatched current_date: %s", step, frame) + } + if !bytes.Contains(frame, []byte("2026-04-20")) { + t.Fatalf("%s upstream frame missing after-auth patched current_date: %s", step, frame) + } + case <-time.After(3 * time.Second): + t.Fatalf("%s timed out waiting for upstream frame", step) + } + } + + // 1. Initial response.create + turn1 := `{"type":"response.create","model":"` + modelName + `","input":[{"role":"developer","content":[{"type":"input_text","text":"\n 2026-04-19\n UTC\n"}]}]}` + if err := c.WriteMessage(websocket.TextMessage, []byte(turn1)); err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + if _, _, errRead := c.ReadMessage(); errRead != nil { + t.Fatalf("initial turn read failed: %v", errRead) + } + } + assertRewrittenUpstreamFrame("initial response.create") + + // 2. Subsequent response.create over the same duplex WebSocket + turn2 := `{"type":"response.create","model":"` + modelName + `","input":[{"role":"developer","content":[{"type":"input_text","text":"\n 2026-04-19\n UTC\n"}]}]}` + if err := c.WriteMessage(websocket.TextMessage, []byte(turn2)); err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + if _, _, errRead := c.ReadMessage(); errRead != nil { + t.Fatalf("subsequent response.create read failed: %v", errRead) + } + } + assertRewrittenUpstreamFrame("subsequent response.create") + + // 3. Subsequent response.append over the same duplex WebSocket + turn3 := `{"type":"response.append","input":[{"role":"developer","content":[{"type":"input_text","text":"\n 2026-04-19\n UTC\n"}]}]}` + if err := c.WriteMessage(websocket.TextMessage, []byte(turn3)); err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + if _, _, errRead := c.ReadMessage(); errRead != nil { + t.Fatalf("subsequent response.append read failed: %v", errRead) + } + } + assertRewrittenUpstreamFrame("subsequent response.append") + + // 4. Subsequent response.steer over the same duplex WebSocket + turn4 := `{"type":"response.steer","input":[{"role":"user","content":[{"type":"input_text","text":"\n 2026-04-19\n UTC\n"}]}]}` + if err := c.WriteMessage(websocket.TextMessage, []byte(turn4)); err != nil { + t.Fatal(err) + } + if _, _, errRead := c.ReadMessage(); errRead != nil { + t.Fatalf("subsequent response.steer read failed: %v", errRead) + } + assertRewrittenUpstreamFrame("subsequent response.steer") + if interceptorHost.afterSawRemovedHeader.Load() { + t.Fatal("after-auth interceptor saw header X-Strip-In-Before that was removed by before-auth interceptor") + } + if interceptorHost.afterMissingKept.Load() { + t.Fatal("after-auth interceptor lost X-Keep-Header across turns") + } +} diff --git a/sdk/api/handlers/openai/openai_responses_steering_integration_test.go b/sdk/api/handlers/openai/openai_responses_steering_integration_test.go index cf13f9f7972..5e9f3cc0844 100644 --- a/sdk/api/handlers/openai/openai_responses_steering_integration_test.go +++ b/sdk/api/handlers/openai/openai_responses_steering_integration_test.go @@ -13,8 +13,10 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/tidwall/gjson" @@ -98,9 +100,22 @@ func TestResponsesSteeringFullDuplexIntegration(t *testing.T) { } registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: "steering-test-model"}}) defer registry.GetGlobalRegistry().UnregisterClient(authID) - h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager)) + base := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager) + base.SetModelRouterHost(&websocketLoggingRouter{ + record: func(ctx context.Context, _ string, body []byte) { + // Model-router host HTTP calls share the duplex connection lifetime. + helps.RecordAPIRequest(ctx, cfg, helps.UpstreamRequestLog{URL: "https://router.invalid/route", Body: body}) + }, + }) + h := NewOpenAIResponsesAPIHandler(base) router := gin.New() - router.GET("/v1/responses", h.ResponsesWebsocket) + deferredCount := make(chan int, 1) + router.GET("/v1/responses", func(c *gin.Context) { + h.ResponsesWebsocket(c) + value, _ := c.Get(logging.DeferredAPIRequestContextKey) + requests, _ := value.([]logging.DeferredAPIRequest) + deferredCount <- len(requests) + }) downstream := httptest.NewServer(router) defer downstream.Close() c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(downstream.URL, "http")+"/v1/responses", nil) @@ -121,12 +136,15 @@ func TestResponsesSteeringFullDuplexIntegration(t *testing.T) { for { _, b, err := c.ReadMessage() if err != nil { - if !strings.HasPrefix(scenario, "disconnect_") { - t.Fatalf("early close: %v", err) + t.Fatalf("downstream read before expected terminal event: %v", err) + } + kind := gjson.GetBytes(b, "type").String() + if kind == "error" { + if !strings.HasPrefix(scenario, "disconnect_") || gjson.GetBytes(b, "error.code").String() != "upstream_websocket_lost" { + t.Fatalf("unexpected downstream error: %s", b) } break } - kind := gjson.GetBytes(b, "type").String() if kind == "response.created" && gjson.GetBytes(b, "response.id").String() == "r1" { send(control1) send(control2) @@ -141,6 +159,9 @@ func TestResponsesSteeringFullDuplexIntegration(t *testing.T) { send([]byte(`{"type":"response.create","previous_response_id":"r1","input":[{"type":"function_call_output","call_id":"call1","output":"found"}]}`)) } if kind == "response.completed" && gjson.GetBytes(b, "response.id").String() == "r2" { + if strings.HasPrefix(scenario, "disconnect_") { + t.Fatal("successor completed instead of reporting upstream disconnect") + } completed = gjson.GetBytes(b, "response.output.0.content.0.text").String() == "STEER_OK" _ = c.Close() break @@ -163,6 +184,16 @@ func TestResponsesSteeringFullDuplexIntegration(t *testing.T) { if connections.Load() != 1 { t.Fatalf("unexpected reconnect/replay: %d", connections.Load()) } + // An upstream disconnect can leave the downstream ready for another turn. + _ = c.Close() + select { + case n := <-deferredCount: + if n != 0 { + t.Fatalf("duplex socket retained %d deferred upstream requests", n) + } + case <-time.After(3 * time.Second): + t.Fatal("duplex handler did not return") + } }) } } diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 329fb6e1c01..71b849f9295 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -19,6 +19,7 @@ import ( "github.com/google/uuid" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v8/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" @@ -354,7 +355,9 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } writer := newResponsesWebsocketWriter(conn) passthroughSessionID := uuid.NewString() - socketCtx, cancelSocket := context.WithCancelCause(c.Request.Context()) + // After the upgrade, error-only HTTP logs cannot consume deferred bodies. + // Apply the policy before routing so plugin callbacks inherit it as well. + socketCtx, cancelSocket := context.WithCancelCause(logging.WithDeferredAPIRequestCaptureDisabled(c.Request.Context())) defer cancelSocket(nil) c.Request = c.Request.WithContext(socketCtx) upgradeRequest := c.Request.Clone(socketCtx) @@ -927,6 +930,9 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { toolCacheTurn: toolCacheTurn, suppressError: replayPinnedAuthFailure, localInterrupt: localInterrupt, + websocketOwned: func() bool { + return attemptedUpstreamMode == responsesWebsocketUpstreamModeWS + }, }, ) if errForward != nil { diff --git a/sdk/api/handlers/openai/openai_responses_websocket_claimed_interrupt_test.go b/sdk/api/handlers/openai/openai_responses_websocket_claimed_interrupt_test.go new file mode 100644 index 00000000000..5f19d03e1f5 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_claimed_interrupt_test.go @@ -0,0 +1,75 @@ +package openai + +import ( + "testing" +) + +// Every frame deliver claims must stay answerable: deactivating or ending the +// turn may stop new claims, but must never silently drop a frame that was +// already accepted. The forward loop relies on the queued frame surviving +// end() so it can be answered instead of timing the client out. +func TestResponsesLocalInterruptRetainsClaimedFrameAfterEnd(t *testing.T) { + local := newResponsesLocalInterrupt() + local.begin(false) + if !local.deliver([]byte(`{"type":"response.interrupt","response_id":"unknown"}`)) { + t.Fatal("frame was not claimed while active") + } + local.deactivate() + local.end() + if local.deliver([]byte(`{"type":"response.interrupt","response_id":"unknown-2"}`)) { + t.Fatal("frame claimed after end") + } + select { + case frame := <-local.framesChan(): + if got := string(frame); got != `{"type":"response.interrupt","response_id":"unknown"}` { + t.Fatalf("retained frame = %s", got) + } + default: + t.Fatal("claimed frame was silently dropped by end()") + } + // The claim must stay spent: nothing is queued twice. + select { + case frame := <-local.framesChan(): + t.Fatalf("unexpected second frame: %s", frame) + default: + } +} + +// A full queue must reject a second frame instead of claiming it: claiming +// without retaining it would leave the client waiting for a local answer +// that only ever addresses the first frame. +func TestResponsesLocalInterruptRejectsWhenQueueFull(t *testing.T) { + local := newResponsesLocalInterrupt() + local.begin(false) + if !local.deliver([]byte(`{"type":"response.interrupt","response_id":"first"}`)) { + t.Fatal("first frame was not claimed") + } + if local.deliver([]byte(`{"type":"response.interrupt","response_id":"second"}`)) { + t.Fatal("second frame claimed while queue was full") + } + frame := local.takeClaimed() + if got := string(frame); got != `{"type":"response.interrupt","response_id":"first"}` { + t.Fatalf("claimed frame = %s", got) + } + if frame := local.takeClaimed(); frame != nil { + t.Fatalf("unexpected second claimed frame: %s", frame) + } +} + +// Turn ownership must be part of claim activation (begin sets both under +// one lock): if begin ran in two steps, a reader could claim a frame for a +// websocket-owned turn in the gap and the forward loop would answer it with +// a fabricated local incomplete. This test pins the single-step API and +// its visible outcome; the one-critical-section guarantee is by construction. +func TestResponsesLocalInterruptWebsocketOwnershipIsAtomic(t *testing.T) { + local := newResponsesLocalInterrupt() + local.begin(true) + if local.deliver([]byte(`{"type":"response.interrupt","response_id":"r1"}`)) { + t.Fatal("frame claimed for a websocket-owned turn") + } + local.endWebsocketTurn() + local.begin(false) + if !local.deliver([]byte(`{"type":"response.interrupt","response_id":"r1"}`)) { + t.Fatal("frame not claimed for an HTTP turn") + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_disconnect_interrupt_test.go b/sdk/api/handlers/openai/openai_responses_websocket_disconnect_interrupt_test.go new file mode 100644 index 00000000000..407099f533e --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_disconnect_interrupt_test.go @@ -0,0 +1,148 @@ +package openai + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + "github.com/tidwall/gjson" +) + +// websocketLossInterruptExecutor runs a turn whose upstream websocket is +// lost mid-flight: InterruptExecutionSession stops seeing a live socket only +// after the test releases the loss, mirroring disconnect teardown winning +// the race against the handler's error delivery. +type websocketLossInterruptExecutor struct { + lost chan struct{} +} + +func (*websocketLossInterruptExecutor) Identifier() string { return "codex" } + +func (*websocketLossInterruptExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (e *websocketLossInterruptExecutor) ExecuteStream(ctx context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 2) + go func() { + defer close(chunks) + chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"type":"response.created","response":{"id":"r-ws"}}`)} + select { + case <-e.lost: + chunks <- coreexecutor.StreamChunk{Err: &coreexecutor.UpstreamWebsocketTurnError{Cause: errors.New("upstream websocket lost")}} + case <-ctx.Done(): + } + }() + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func (*websocketLossInterruptExecutor) Refresh(context.Context, *coreauth.Auth) (*coreauth.Auth, error) { + return nil, errors.New("not implemented") +} + +func (*websocketLossInterruptExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (*websocketLossInterruptExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +func (e *websocketLossInterruptExecutor) InterruptExecutionSession(context.Context, string, []byte) error { + select { + case <-e.lost: + return coreexecutor.ErrNoActiveUpstreamWebsocket + default: + return errors.New("interrupt reached a live websocket") + } +} + +// When the upstream websocket dies mid-turn, an interrupt racing the loss +// delivery must not be answered with a fabricated local incomplete: the turn +// is websocket-owned, so the client only sees the loss result. The upstream +// loss is detected through InterruptExecutionSession itself, which makes the +// interleaving deterministic. +func TestResponsesInterruptRacingWebsocketLossIsNotAnsweredLocally(t *testing.T) { + gin.SetMode(gin.TestMode) + executor := &websocketLossInterruptExecutor{lost: make(chan struct{})} + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(nil) + manager.RegisterExecutor(executor) + const authID = "interrupt-ws-loss-auth" + const model = "interrupt-ws-loss-model" + _, errRegister := manager.Register(context.Background(), &coreauth.Auth{ + ID: authID, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{"api_key": "test-key", "websockets": "true"}, + }) + if errRegister != nil { + t.Fatal(errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: model}}) + defer registry.GetGlobalRegistry().UnregisterClient(authID) + + handler := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(nil, manager)) + router := gin.New() + router.GET("/v1/responses", handler.ResponsesWebsocket) + downstream := httptest.NewServer(router) + defer downstream.Close() + client, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(downstream.URL, "http")+"/v1/responses", nil) + if errDial != nil { + t.Fatal(errDial) + } + defer func() { _ = client.Close() }() + _ = client.SetReadDeadline(time.Now().Add(8 * time.Second)) + send := func(payload string) { + t.Helper() + if errSend := client.WriteMessage(websocket.TextMessage, []byte(payload)); errSend != nil { + t.Fatal(errSend) + } + } + send(fmt.Sprintf(`{"type":"response.create","model":%q,"event_id":"evt-loss","input":[]}`, model)) + _, created, errRead := client.ReadMessage() + if errRead != nil { + t.Fatal(errRead) + } + if got := gjson.GetBytes(created, "type").String(); got != "response.created" { + t.Fatalf("expected response.created, got %s", created) + } + + // The upstream socket is gone before the interrupt is dispatched, but + // the handler has not consumed the loss yet. + close(executor.lost) + send(`{"type":"response.interrupt","response_id":"r-ws"}`) + + // The next downstream message must be the upstream-loss notice, never a + // locally fabricated incomplete for r-ws. If the interrupt had been + // claimed locally, the forward loop would have cancelled the executor + // context and written that incomplete before the loss terminal could + // be delivered, so observing the loss notice first also proves the + // frame was never locally answered. + deadline := time.Now().Add(8 * time.Second) + for time.Now().Before(deadline) { + _, payload, errNext := client.ReadMessage() + if errNext != nil { + t.Fatalf("waiting for upstream loss notice: %v", errNext) + } + eventType := gjson.GetBytes(payload, "type").String() + if eventType == "response.incomplete" { + t.Fatalf("fabricated local incomplete after websocket loss: %s", payload) + } + if eventType == "error" && gjson.GetBytes(payload, "error.event_id").String() == "evt-loss" { + return + } + } + t.Fatal("upstream loss notice was not observed") +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_forward.go b/sdk/api/handlers/openai/openai_responses_websocket_forward.go index b7882d020a2..4202dc11b0e 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_forward.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_forward.go @@ -30,6 +30,9 @@ type responsesWebsocketForwardOptions struct { suppressError func(*interfaces.ErrorMessage) bool keepAliveInterval *time.Duration localInterrupt *responsesLocalInterrupt + // websocketOwned reports whether this turn runs on an upstream + // websocket. While true, no interrupt may claim the local HTTP path. + websocketOwned func() bool } func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( @@ -48,6 +51,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } toolCacheTurn := opts.toolCacheTurn completed := false + // turnEnded tracks every terminal outcome (success or failure) for the + // local-interrupt window; completed stays success-only for stream + // completion bookkeeping. + turnEnded := false responseStarted := false completedOutput := []byte("[]") completedResponseID := "" @@ -74,8 +81,28 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( keepAliveC = keepAliveTicker.C } if opts.localInterrupt != nil { - opts.localInterrupt.begin() - defer opts.localInterrupt.end() + websocketOwned := opts.websocketOwned != nil && opts.websocketOwned() + opts.localInterrupt.begin(websocketOwned) + if websocketOwned { + defer opts.localInterrupt.endWebsocketTurn() + } + defer func() { + opts.localInterrupt.end() + // The reader may have claimed a frame in the gap between the + // last loop iteration and this return. It already logged the + // frame as locally handled, so answer it here instead of + // dropping it; otherwise the client waits for a reply that + // never comes. Frames claimed while the socket or request is + // already gone need no answer. + if c.Request.Context().Err() == nil { + if claimed := opts.localInterrupt.takeClaimed(); claimed != nil { + _, _ = writeResponsesWebsocketError(writer, wsTimelineLog, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New("no active response to interrupt"), + }) + } + } + }() } localFrames := opts.localInterrupt.framesChan() @@ -112,6 +139,20 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( for { select { case interruptPayload := <-localFrames: + if turnEnded { + // The frame was claimed before the terminal payload landed, so + // the reader already told the client it was handled locally. + // The response is over: reject instead of fabricating an + // incomplete event, and keep waiting for the turn to drain. + if _, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New("no active response to interrupt"), + }); errWrite != nil { + cancel(errWrite) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite + } + continue + } return completeResponsesWebsocketLocalInterrupt(writer, wsTimelineLog, cancel, interruptPayload, outputItemsByIndex, outputItemsFallback, pendingToolCallIDs) case <-c.Request.Context().Done(): cancel(c.Request.Context().Err()) @@ -175,15 +216,40 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( keepAliveTicker.Reset(keepAliveInterval) } + // A queued local interrupt belongs to an in-flight HTTP turn: + // consume it before forwarding more upstream payload. Once the + // turn ended the frame can only be stale, so reject it. + if opts.localInterrupt != nil { + if interruptPayload := opts.localInterrupt.takeClaimed(); interruptPayload != nil { + if !turnEnded { + return completeResponsesWebsocketLocalInterrupt(writer, wsTimelineLog, cancel, interruptPayload, outputItemsByIndex, outputItemsFallback, pendingToolCallIDs) + } + // The response is over: reject instead of fabricating an + // incomplete event, then keep forwarding this chunk. + if _, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: errors.New("no active response to interrupt"), + }); errWrite != nil { + cancel(errWrite) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite + } + } + } + payloads := websocketJSONPayloadsFromChunk(chunk) for i := range payloads { eventType := gjson.GetBytes(payloads[i], "type").String() - if eventType == "response.subscription_usage" { + if eventType == "response.subscription_usage" || + (eventType == "response.metadata" && !isCodexResponsesClientRequest(c)) { continue } if eventType == "response.created" { responseStarted = true completed = false + turnEnded = false + if opts.localInterrupt != nil { + opts.localInterrupt.reactivate() + } outputItemsByIndex = make(map[int64][]byte) outputItemsFallback = nil pendingToolCallIDs = make(map[string]struct{}) @@ -217,8 +283,13 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( // reason "interrupted", to acknowledge response.interrupt. Closing // the client socket here would block the follow-up request. completed = true + turnEnded = true completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback) completedResponseID = responseCompletedIDFromPayload(payloads[i]) + } else if eventType == "response.failed" { + // A failed response ends the turn without a completion + // identity; only its local-interrupt window closes. + turnEnded = true } markAPIResponseTimestamp(c) if payloadErrMsg != nil { @@ -249,6 +320,15 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( // websocketPayloadEventType(payloads[i]), // websocketPayloadPreview(payloads[i]), // ) + if turnEnded && opts.localInterrupt != nil && opts.localInterrupt.isActive() { + // Stop claiming local interrupts before the client can see + // the terminal event; a later frame is then rejected by + // the reader instead of being answered or dropped here. + // The isActive guard keeps multi-payload chunks from + // toggling the state back off after response.created + // reactivated it. + opts.localInterrupt.deactivate() + } if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { log.Warnf( "responses websocket: downstream_out write failed id=%s event=%s error=%v", diff --git a/sdk/api/handlers/openai/openai_responses_websocket_input.go b/sdk/api/handlers/openai/openai_responses_websocket_input.go index e1d912c067b..5ea827ebb3d 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_input.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_input.go @@ -22,52 +22,108 @@ import ( // responsesLocalInterrupt cancels an in-flight HTTP turn. A websocket upstream // does not use it; that interrupt is written to the existing socket instead. +// A claimed frame is never silently dropped: the forward loop consumes it, or +// deliver returns false so the reader answers the interrupt itself. +// +// wsTurnActive is the socket-scoped ownership gate: while the current turn +// runs on an upstream websocket, no frame may claim the local path at all, +// not even in the window after that socket was lost. "No live websocket +// channel right now" must never be read as "this turn is HTTP". type responsesLocalInterrupt struct { - mu sync.Mutex - active bool - frames chan []byte + mu sync.Mutex + active bool + wsTurnActive bool + frames chan []byte } func newResponsesLocalInterrupt() *responsesLocalInterrupt { return &responsesLocalInterrupt{frames: make(chan []byte, 1)} } -func (s *responsesLocalInterrupt) begin() { +// begin activates claims for a turn and records its ownership in the same +// critical section: a websocket-owned turn must never expose an instant in +// which claims are active but ownership is not yet visible to deliver. +func (s *responsesLocalInterrupt) begin(websocketOwned bool) { if s == nil { return } s.mu.Lock() s.active = true + s.wsTurnActive = websocketOwned s.mu.Unlock() } -func (s *responsesLocalInterrupt) end() { +// reactivate reopens claims for the next response on the same turn (duplex), +// preserving the ownership recorded by begin. +func (s *responsesLocalInterrupt) reactivate() { if s == nil { return } s.mu.Lock() - s.active = false - select { - case <-s.frames: - default: + s.active = true + s.mu.Unlock() +} + +func (s *responsesLocalInterrupt) endWebsocketTurn() { + if s == nil { + return } + s.mu.Lock() + s.wsTurnActive = false s.mu.Unlock() } +// stopClaims stops accepting new frames without draining the queue: a frame +// already claimed by deliver stays consumable so the forward loop answers it +// instead of silently dropping it when the turn ends. +func (s *responsesLocalInterrupt) stopClaims() { + if s == nil { + return + } + s.mu.Lock() + s.active = false + s.mu.Unlock() +} + +// end stops new claims at turn teardown; claimed frames stay queued for +// takeClaimed so they can still be answered. +func (s *responsesLocalInterrupt) end() { + s.stopClaims() +} + +// deactivate stops new claims once the turn completed. +func (s *responsesLocalInterrupt) deactivate() { + s.stopClaims() +} + +// isActive reports whether new frames may still be claimed for the turn. +func (s *responsesLocalInterrupt) isActive() bool { + if s == nil { + return false + } + s.mu.Lock() + defer s.mu.Unlock() + return s.active +} + +// deliver claims a frame only if it is actually retained: a full queue +// rejects so the reader answers the frame itself instead of the client +// waiting for an answer that only ever addresses the first claim. func (s *responsesLocalInterrupt) deliver(payload []byte) bool { if s == nil { return false } s.mu.Lock() defer s.mu.Unlock() - if !s.active { + if !s.active || s.wsTurnActive { return false } select { case s.frames <- bytes.Clone(payload): + return true default: + return false } - return true } func (s *responsesLocalInterrupt) framesChan() <-chan []byte { @@ -77,6 +133,20 @@ func (s *responsesLocalInterrupt) framesChan() <-chan []byte { return s.frames } +// takeClaimed consumes a queued claimed frame without blocking. The caller +// becomes responsible for answering it exactly once. +func (s *responsesLocalInterrupt) takeClaimed() []byte { + if s == nil { + return nil + } + select { + case frame := <-s.frames: + return frame + default: + return nil + } +} + // readResponsesWebsocketInput is the only downstream reader. Interrupts are // handled immediately so they are not queued behind an active response. // The bounded queue backpressures clients instead of retaining unlimited input. diff --git a/sdk/api/handlers/openai/openai_responses_websocket_late_local_interrupt_test.go b/sdk/api/handlers/openai/openai_responses_websocket_late_local_interrupt_test.go new file mode 100644 index 00000000000..327f7e86780 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_late_local_interrupt_test.go @@ -0,0 +1,156 @@ +package openai + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" + "github.com/tidwall/gjson" +) + +// Once a terminal payload (completion or failure) has been written +// downstream, a local interrupt must no longer be accepted for that turn: +// accepting it either fabricates a response.incomplete or drops the frame +// when the turn ends. The data channel stays open in this test so the +// forward loop is deterministically still in its post-terminal window when +// deliver is attempted. +func TestForwardResponsesWebsocketRejectsLocalInterruptAfterCompletion(t *testing.T) { + testForwardResponsesWebsocketRejectsLocalInterruptAfterTerminal(t, + `{"type":"response.completed","response":{"id":"resp-1","output":[]}}`, wsEventTypeCompleted, true) +} + +func TestForwardResponsesWebsocketRejectsLocalInterruptAfterFailure(t *testing.T) { + testForwardResponsesWebsocketRejectsLocalInterruptAfterTerminal(t, + `{"type":"response.failed","response":{"id":"resp-1","status":"failed","output":[]}}`, "response.failed", false) +} + +func testForwardResponsesWebsocketRejectsLocalInterruptAfterTerminal(t *testing.T, terminalPayload, terminalType string, wantCompleted bool) { + t.Helper() + gin.SetMode(gin.TestMode) + + local := newResponsesLocalInterrupt() + data := make(chan []byte, 1) + errCh := make(chan *interfaces.ErrorMessage) + var streamClosed atomic.Bool + closeStream := func() { + if streamClosed.CompareAndSwap(false, true) { + close(data) + close(errCh) + } + } + // Always end the forward turn, even on assertion failure. + defer closeStream() + data <- []byte(terminalPayload) + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + serverErrCh <- errUpgrade + return + } + defer func() { + if errClose := conn.Close(); errClose != nil { + serverErrCh <- errClose + } + }() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = r + + // A nil handler panics on the stream-closed-before-completed path + // (LoggingAPIResponseError), which the failure variant reaches; use + // a real handler shell without an auth manager. + testHandler := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(nil, nil)) + _, completedResponseID, _, errMsg, errForward := testHandler.forwardResponsesWebsocket( + ctx, + newResponsesWebsocketWriter(conn), + func(...interface{}) {}, + data, + errCh, + newInMemoryWebsocketTimelineLog(), + "session-1", + responsesWebsocketForwardOptions{localInterrupt: local}, + ) + if errForward != nil { + // The test ends the turn by canceling the request context. The + // failure variant instead ends inside the pre-completion close + // path, which reports ErrCloseSent after closing the socket. + if errForward != context.Canceled && !errors.Is(errForward, websocket.ErrCloseSent) { + serverErrCh <- errForward + return + } + } else if wantCompleted { + if errMsg != nil { + serverErrCh <- fmt.Errorf("unexpected websocket error message: %v", errMsg.Error) + return + } + if completedResponseID != "resp-1" { + serverErrCh <- fmt.Errorf("completed response id = %q, want resp-1", completedResponseID) + return + } + } else { + // A failed turn carries no completed identity. + if completedResponseID != "" { + serverErrCh <- fmt.Errorf("completed response id = %q, want empty after failure", completedResponseID) + return + } + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + closed := false + defer func() { + if !closed { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + } + }() + if errDeadline := conn.SetReadDeadline(time.Now().Add(3 * time.Second)); errDeadline != nil { + t.Fatal(errDeadline) + } + + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read terminal message: %v", errRead) + } + if got := gjson.GetBytes(payload, "type").String(); got != terminalType { + t.Fatalf("payload type = %s, want %s", got, terminalType) + } + + // The terminal event was already written, so no local interrupt may be taken. + if local.deliver([]byte(`{"type":"response.interrupt","response_id":"unknown"}`)) { + t.Fatal("local interrupt accepted after terminal event was written downstream") + } + + // The turn is still alive; closing the stream ends it. + closeStream() + if errClose := conn.Close(); errClose != nil { + t.Fatalf("close websocket: %v", errClose) + } + closed = true + select { + case errServer := <-serverErrCh: + if errServer != nil { + t.Fatalf("server error: %v", errServer) + } + case <-time.After(3 * time.Second): + t.Fatal("forward loop did not finish after client close") + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_logging_test.go b/sdk/api/handlers/openai/openai_responses_websocket_logging_test.go new file mode 100644 index 00000000000..df30b8ae111 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_logging_test.go @@ -0,0 +1,215 @@ +package openai + +import ( + "bytes" + "context" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/api/middleware" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" + "github.com/tidwall/gjson" +) + +type websocketLoggingExecutor struct { + websocketDirectCaptureExecutor + record func(context.Context, string, []byte) +} + +func (e *websocketLoggingExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + e.record(ctx, "executor", req.Payload) + e.mu.Lock() + if len(e.payloads) == 2 { + e.failStatus = http.StatusBadGateway + } + e.mu.Unlock() + return e.websocketDirectCaptureExecutor.ExecuteStream(ctx, auth, req, opts) +} + +type websocketLoggingRouter struct { + record func(context.Context, string, []byte) + contexts chan context.Context +} + +func (*websocketLoggingRouter) HasModelRouters() bool { return true } + +func (r *websocketLoggingRouter) RouteModel(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + r.record(ctx, "router", req.Body) + if r.contexts != nil { + r.contexts <- ctx + } + return pluginapi.ModelRouteResponse{}, false +} + +func TestResponsesWebsocketRequestLoggingLifecycle(t *testing.T) { + for _, fullLog := range []bool{false, true} { + name := "error-only" + if fullLog { + name = "full" + } + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: fullLog}} + observations := make(chan int, 16) + record := func(ctx context.Context, source string, body []byte) { + helps.RecordAPIRequest(ctx, cfg, helps.UpstreamRequestLog{URL: "https://example.invalid/" + source, Method: http.MethodPost, Body: body}) + c, _ := ctx.Value("gin").(*gin.Context) + value, _ := c.Get(logging.DeferredAPIRequestContextKey) + requests, _ := value.([]logging.DeferredAPIRequest) + observations <- len(requests) + } + executor := &websocketLoggingExecutor{record: record} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + authID := "websocket-log-" + name + if _, err := manager.Register(context.Background(), &coreauth.Auth{ID: authID, Provider: "codex", Status: coreauth.StatusActive}); err != nil { + t.Fatal(err) + } + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: "logging-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + base := handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager) + modelRouter := &websocketLoggingRouter{record: record, contexts: make(chan context.Context, 8)} + base.SetModelRouterHost(modelRouter) + h := NewOpenAIResponsesAPIHandler(base) + logsDir := t.TempDir() + logger := logging.NewFileRequestLogger(fullLog, logsDir, "", 10) + handlerDone := make(chan struct{}) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Next() + close(handlerDone) + }) + router.Use(middleware.RequestLoggingMiddleware(logger)) + router.GET("/v1/responses", h.ResponsesWebsocket) + server := httptest.NewServer(router) + defer server.Close() + conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http")+"/v1/responses", nil) + if err != nil { + t.Fatal(err) + } + defer func() { _ = conn.Close() }() + _ = conn.SetReadDeadline(time.Now().Add(10 * time.Second)) + send := func(payload string) { + t.Helper() + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Fatal(errWrite) + } + } + readEvent := func(want string) { + t.Helper() + _, data, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatal(errRead) + } + if got := gjson.GetBytes(data, "type").String(); got != want { + t.Fatalf("event = %q, want %q", got, want) + } + } + assertNoCapture := func() { + t.Helper() + select { + case n := <-observations: + if n != 0 { + t.Fatalf("open WebSocket retained %d deferred requests", n) + } + case <-time.After(5 * time.Second): + t.Fatal("request capture was not observed") + } + } + // Routing happens even for a locally answered prewarm request. + send(`{"type":"response.create","model":"logging-model","generate":false}`) + readEvent("response.created") + readEvent("response.completed") + assertNoCapture() + callbackCtx := <-modelRouter.contexts + payload := `{"type":"response.create","model":"logging-model","input":[{"role":"user","content":"` + strings.Repeat("x", 256<<10) + `"}]}` + for range 2 { + send(payload) + readEvent("response.completed") + assertNoCapture() + assertNoCapture() + } + // A lifecycle callback can perform host HTTP work after its turn ended. + go record(context.WithoutCancel(callbackCtx), "late-plugin", []byte("callback-body")) + assertNoCapture() + send(payload) + if _, data, errRead := conn.ReadMessage(); errRead == nil { + if gjson.GetBytes(data, "type").String() != "error" { + t.Fatalf("failed upstream returned %s", data) + } + } else { + var closeErr *websocket.CloseError + if !errors.As(errRead, &closeErr) { + t.Fatalf("upstream failure returned %v, want an error event or close frame", errRead) + } + } + assertNoCapture() + assertNoCapture() + _ = conn.Close() + select { + case <-handlerDone: + case <-time.After(5 * time.Second): + t.Fatal("WebSocket handler did not finish") + } + if got := len(executor.Payloads()); got != 3 { + t.Fatalf("upstream calls = %d, want 3", got) + } + if fullLog { + entries, errReadDir := os.ReadDir(logsDir) + if errReadDir != nil { + t.Fatal(errReadDir) + } + var logged []byte + for _, entry := range entries { + if !strings.HasSuffix(entry.Name(), ".log") { + continue + } + body, errRead := os.ReadFile(filepath.Join(logsDir, entry.Name())) + if errRead != nil { + t.Fatal(errRead) + } + logged = append(logged, body...) + } + for _, source := range []string{"router", "executor", "late-plugin"} { + if !bytes.Contains(logged, []byte("https://example.invalid/"+source)) { + t.Errorf("full request log lost %s requests", source) + } + } + } + }) + } +} + +func TestResponsesWebsocketRejectedUpgradeKeepsErrorCapture(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Request.Header.Set("Connection", "Upgrade") + c.Request.Header.Set("Upgrade", "websocket") + // Missing handshake headers make the upgrade fail before a socket exists. + (&OpenAIResponsesAPIHandler{}).ResponsesWebsocket(c) + if c.Writer.Status() < http.StatusBadRequest { + t.Fatal("invalid WebSocket handshake was accepted") + } + ctx := context.WithValue(c.Request.Context(), "gin", c) + helps.RecordAPIRequest(ctx, &config.Config{}, helps.UpstreamRequestLog{Body: []byte("handshake-error")}) + value, _ := c.Get(logging.DeferredAPIRequestContextKey) + requests, _ := value.([]logging.DeferredAPIRequest) + if len(requests) != 1 || !bytes.Contains(requests[0](), []byte("handshake-error")) { + t.Fatal("failed upgrade lost ordinary HTTP error capture") + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_metadata_test.go b/sdk/api/handlers/openai/openai_responses_websocket_metadata_test.go new file mode 100644 index 00000000000..61a621e0432 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_metadata_test.go @@ -0,0 +1,115 @@ +package openai + +import ( + "context" + "net/http" + "net/http/httptest" + "slices" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdkconfig "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" + "github.com/tidwall/gjson" +) + +type websocketMetadataExecutor struct { + websocketCaptureExecutor +} + +func (*websocketMetadataExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 5) + for _, payload := range []string{ + `data: {"type":"response.created","response":{"id":"resp-1"}}`, + `data: {"type":"codex.response.metadata","headers":{"x-turn-state":"turn-1"}}`, + `data: {"type":"response.metadata","metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{}}}}`, + `data: {"type":"response.subscription_usage","subscription":{"tier":"private"}}`, + `data: {"type":"response.completed","response":{"id":"resp-1","output":[]}}`, + } { + chunks <- coreexecutor.StreamChunk{Payload: []byte(payload)} + } + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func TestResponsesWebsocketFiltersModerationMetadataForStandardClients(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + headers http.Header + wantTypes []string + }{ + { + name: "standard client", + headers: http.Header{"User-Agent": []string{"standard-client/1.0"}}, + wantTypes: []string{"response.created", "codex.response.metadata", "response.completed"}, + }, + { + name: "codex user agent", + headers: http.Header{"User-Agent": []string{"codex_cli_rs/1.0"}}, + wantTypes: []string{"response.created", "codex.response.metadata", "response.metadata", "response.completed"}, + }, + { + name: "codex originator", + headers: http.Header{"Originator": []string{"codex_cli_rs"}}, + wantTypes: []string{"response.created", "codex.response.metadata", "response.metadata", "response.completed"}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + executor := &websocketMetadataExecutor{} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ID: "websocket-metadata-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "websocket-metadata-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)) + router := gin.New() + router.GET("/v1/responses", h.ResponsesWebsocket) + server := httptest.NewServer(router) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses" + conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, tc.headers) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + request := `{"type":"response.create","model":"websocket-metadata-model","input":[]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil { + t.Fatalf("write request: %v", errWrite) + } + + var gotTypes []string + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket: %v; events: %v", errRead, gotTypes) + } + eventType := gjson.GetBytes(payload, "type").String() + gotTypes = append(gotTypes, eventType) + if eventType == "error" { + t.Fatalf("unexpected error event: %s", payload) + } + if eventType == "response.completed" { + break + } + } + if !slices.Equal(gotTypes, tc.wantTypes) { + t.Fatalf("event types = %v, want %v", gotTypes, tc.wantTypes) + } + }) + } +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_unknown_interrupt_test.go b/sdk/api/handlers/openai/openai_responses_websocket_unknown_interrupt_test.go new file mode 100644 index 00000000000..3df1f554c42 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_unknown_interrupt_test.go @@ -0,0 +1,157 @@ +package openai + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/api/handlers" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + "github.com/tidwall/gjson" +) + +// An unknown interrupt while a websocket response is running belongs to the +// live upstream socket: it must be rejected outright, never fall back to a +// local cancellation that fabricates an incomplete event and kills the +// running turn. Only a turn with no active websocket (HTTP fallback) may be +// interrupted locally. +func TestResponsesUnknownInterruptDoesNotCancelWebsocketTurn(t *testing.T) { + gin.SetMode(gin.TestMode) + var connections atomic.Int32 + release := make(chan struct{}) + var releaseOnce atomic.Bool + // Always unblock the upstream goroutine, even on assertion failure. + defer func() { + if releaseOnce.CompareAndSwap(false, true) { + close(release) + } + }() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connections.Add(1) + conn, errUpgrade := (&websocket.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 upstream: %v", errClose) + } + }() + _, create, errRead := conn.ReadMessage() + if errRead != nil { + return + } + if gjson.GetBytes(create, "type").String() != "response.create" { + t.Errorf("expected response.create, got %s", create) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"r1"}}`)); errWrite != nil { + t.Errorf("created: %v", errWrite) + return + } + // Keep r1 running until the test checked the unknown interrupt. + <-release + // The only frame allowed next is the valid r1 interrupt. + _ = conn.SetReadDeadline(time.Now().Add(5 * time.Second)) + _, interrupt, errInterrupt := conn.ReadMessage() + if errInterrupt != nil { + return + } + if gjson.GetBytes(interrupt, "response_id").String() != "r1" { + t.Errorf("upstream received unexpected frame: %s", interrupt) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.incomplete","response":{"id":"r1","status":"incomplete","incomplete_details":{"reason":"interrupted"},"output":[]}}`)); errWrite != nil { + t.Errorf("incomplete: %v", errWrite) + return + } + _, _, _ = conn.ReadMessage() + })) + defer upstream.Close() + + cfg := &config.Config{} + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(cfg) + manager.RegisterExecutor(runtimeexecutor.NewCodexAutoExecutor(cfg)) + const authID = "unknown-interrupt-ws-turn" + const model = "unknown-interrupt-model" + _, errRegister := manager.Register(context.Background(), &coreauth.Auth{ + ID: authID, Provider: "codex", Status: coreauth.StatusActive, + Attributes: map[string]string{"api_key": "synthetic-key", "base_url": upstream.URL, "websockets": "true"}, + }) + if errRegister != nil { + t.Fatal(errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authID, "codex", []*registry.ModelInfo{{ID: model}}) + defer registry.GetGlobalRegistry().UnregisterClient(authID) + handler := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&cfg.SDKConfig, manager)) + router := gin.New() + router.GET("/v1/responses", handler.ResponsesWebsocket) + downstream := httptest.NewServer(router) + defer downstream.Close() + + client, _, errDial := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(downstream.URL, "http")+"/v1/responses", nil) + if errDial != nil { + t.Fatal(errDial) + } + clientClosed := false + defer func() { + if !clientClosed { + if errClose := client.Close(); errClose != nil { + t.Errorf("close client: %v", errClose) + } + } + }() + if errDeadline := client.SetReadDeadline(time.Now().Add(8 * time.Second)); errDeadline != nil { + t.Fatal(errDeadline) + } + send := func(payload string) { + t.Helper() + if errWrite := client.WriteMessage(websocket.TextMessage, []byte(payload)); errWrite != nil { + t.Fatal(errWrite) + } + } + expect := func(event, id string) { + t.Helper() + _, payload, errRead := client.ReadMessage() + if errRead != nil { + t.Fatal(errRead) + } + if gjson.GetBytes(payload, "type").String() != event || gjson.GetBytes(payload, "response.id").String() != id { + t.Fatalf("expected %s for %s, got %s", event, id, payload) + } + } + + send(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[]}`, model)) + expect("response.created", "r1") + + // Unknown target on a live websocket turn: reject, do not cancel r1. + send(`{"type":"response.interrupt","response_id":"unknown"}`) + expect("error", "") + + // r1 must still be alive: a valid interrupt now forwards upstream and + // the upstream terminal event comes back on the same socket. + if releaseOnce.CompareAndSwap(false, true) { + close(release) + } + send(`{"type":"response.interrupt","response_id":"r1"}`) + expect("response.incomplete", "r1") + if got := connections.Load(); got != 1 { + t.Fatalf("connections = %d, want 1", got) + } + if errClose := client.Close(); errClose != nil { + t.Fatal(errClose) + } + clientClosed = true +} diff --git a/sdk/cliproxy/auth/codex_siwc_bind_test.go b/sdk/cliproxy/auth/codex_siwc_bind_test.go index 555742a8238..7a9928c7f7b 100644 --- a/sdk/cliproxy/auth/codex_siwc_bind_test.go +++ b/sdk/cliproxy/auth/codex_siwc_bind_test.go @@ -279,7 +279,15 @@ func TestBindCodexSIWCRejectsNonFileSources(t *testing.T) { func TestBindCodexSIWCTildeAndAttributeOverride(t *testing.T) { m, base, path, _ := setupSIWCBinding(t, "auth.json") - t.Setenv("HOME", filepath.Dir(path)) + homeDir := filepath.Dir(path) + // os.UserHomeDir reads HOME on Unix but %USERPROFILE% on Windows, and + // the runtime resolves "~" through it; pin both so the injection below + // lands on the same directory on every platform. + t.Setenv("HOME", homeDir) + t.Setenv("USERPROFILE", homeDir) + if userHomeDir, errUserHomeDir := os.UserHomeDir(); errUserHomeDir != nil || filepath.Clean(userHomeDir) != filepath.Clean(homeDir) { + t.Fatalf("os.UserHomeDir() = %q, %v; want %q", userHomeDir, errUserHomeDir, homeDir) + } m.runtimeConfig.Store(&config.Config{AuthDir: "~"}) m.mu.Lock() m.auths[base.ID].Attributes["use_siwc"] = "false" diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 3baead92fd7..3f485fcf10d 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -74,6 +74,9 @@ type Result struct { RouteModel string // Success marks whether the execution succeeded. Success bool + // CompletedResponse reports that the upstream execution produced a completed + // generation response (as opposed to an incomplete stream or non-generation call). + CompletedResponse bool // RetryAfter carries a provider supplied retry hint (e.g. 429 retryDelay). RetryAfter *time.Duration // CredentialScope indicates that the failure affects the whole credential across models (e.g. Anthropic 5h/7d unified limits). diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 04e95fd23a8..388ac914f67 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -8,6 +8,7 @@ import ( "io" "net" "net/http" + "regexp" "sort" "strings" "sync/atomic" @@ -968,6 +969,10 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { isAuthActiveCooldown(auth, now) && auth.cooldownRecordedSeq > 0 && result.ExecutionSeq <= auth.cooldownRecordedSeq if result.Success { + if result.CompletedResponse && !result.SkipQuotaObservation && !wasTerminalUnauthorized && !(auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) && !retainModelCooldown && !retainAuthCooldown { + auth.InferenceVerified = true + auth.LastInferenceVerifiedAt = now + } if wasTerminalUnauthorized { if modelKey != "" { state := ensureModelState(auth, modelKey) @@ -2208,7 +2213,7 @@ func isUnauthorizedError(err error) bool { return true } raw := strings.ToLower(err.Error()) - return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized") + return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized") || strings.Contains(raw, "(401)") } func hasUnauthorizedAuthFailure(auth *Auth) bool { @@ -2247,6 +2252,170 @@ func HasDisabledInvalidGrantFailure(auth *Auth) bool { return hasDisabledInvalidGrantFailure(auth) } +var ( + renewalURLPattern = regexp.MustCompile(`(?i)https?://[^\s"'<>\\]+`) + renewalJWTPattern = regexp.MustCompile(`\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+(?:\.[A-Za-z0-9_-]+)?\b`) + renewalSecretKVPattern = regexp.MustCompile(`(?i)\b(access_token|refresh_token|id_token|client_secret|code_verifier|code_challenge|authorization)\s*([=:]|"\s*:\s*")\s*(?:(?:bearer|basic|token)\s+)?([^\s"&',}\\]+)`) + renewalOAuthParamPatten = regexp.MustCompile(`(?i)(?:[?&]|\b)(code|state)=([^\s"&',}\\]+)`) + renewalEscapedParamPattern = regexp.MustCompile(`(?i)(code|state)%3[dD]([^\s"&',}\\%]+)`) + renewalJSONStatePattern = regexp.MustCompile(`(?i)("state"\s*:\s*)"([^"]*)(?:"|$)`) + renewalJSONCodePattern = regexp.MustCompile(`(?i)("code"\s*:\s*)"([^"]*)(?:"|$)`) + renewalBearerPattern = regexp.MustCompile(`(?i)\bBearer\s+[^\s"',}\\]+`) +) + +func isKnownSafeOAuthErrorCode(code string) bool { + switch strings.ToLower(strings.TrimSpace(code)) { + case "refresh_token_reused", + "refresh_token_expired", + "refresh_token_invalidated", + "token_revoked", + "invalid_grant", + "invalid_client", + "invalid_request", + "invalid_request_error", + "unauthorized", + "rate_limit_exceeded", + "insufficient_quota", + "billing_hard_limit_reached", + "account_deactivated", + "context_length_exceeded", + "server_error", + "service_unavailable", + "bad_gateway", + "model_not_found": + return true + default: + return false + } +} + +// SanitizeRenewalMessage strips URLs, bearer tokens, JWTs, OAuth codes/states, and secret fields from diagnostic messages. +func SanitizeRenewalMessage(raw string) string { + return sanitizeRenewalMessage(raw) +} + +func redactRenewalPatterns(in string) string { + out := renewalURLPattern.ReplaceAllString(in, "[redacted-url]") + out = renewalJWTPattern.ReplaceAllString(out, "[redacted-token]") + out = renewalBearerPattern.ReplaceAllString(out, "Bearer [redacted]") + out = renewalSecretKVPattern.ReplaceAllString(out, "${1}=[redacted]") + out = renewalOAuthParamPatten.ReplaceAllString(out, "${1}=[redacted]") + out = renewalEscapedParamPattern.ReplaceAllString(out, "${1}=[redacted]") + out = renewalJSONStatePattern.ReplaceAllString(out, `${1}"[redacted]"`) + out = renewalJSONCodePattern.ReplaceAllStringFunc(out, func(match string) string { + sub := renewalJSONCodePattern.FindStringSubmatch(match) + if len(sub) == 3 && isKnownSafeOAuthErrorCode(sub[2]) && strings.HasSuffix(match, `"`) { + return match + } + if len(sub) >= 2 { + return sub[1] + `"[redacted]"` + } + return `"code":"[redacted]"` + }) + return out +} + +func sanitizeRenewalMessage(raw string) string { + // Redact secrets on the full input BEFORE truncation so secrets straddling the truncation boundary cannot leak. + preRedacted := redactRenewalPatterns(raw) + out := internallogging.SafeDiagnosticForLog(preRedacted) + return redactRenewalPatterns(out) +} + +func isTerminalRefreshGrantMessage(message string) bool { + lower := strings.ToLower(message) + for _, token := range []string{ + "refresh_token_reused", + "refresh_token_expired", + "refresh_token_invalidated", + "token_revoked", + "invalid_client", + } { + if strings.Contains(lower, token) { + return true + } + } + return false +} + +func isTerminalRefreshError(auth *Auth, err error) bool { + if err == nil { + return false + } + var partErr interface { + OuterRefreshError() error + SIWCRefreshError() error + } + if errors.As(err, &partErr) && partErr != nil { + return isTerminalRefreshError(auth, partErr.OuterRefreshError()) || + isTerminalRefreshError(auth, partErr.SIWCRefreshError()) + } + if isUnauthorizedError(err) || isTerminalRefreshGrantMessage(err.Error()) { + return true + } + if auth != nil && strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") && isInvalidGrantError(err) { + return true + } + return false +} + +func terminalRefreshErrorCode(message string, statusCode int) string { + lower := strings.ToLower(message) + for _, code := range []string{ + "refresh_token_reused", + "refresh_token_expired", + "refresh_token_invalidated", + "token_revoked", + "invalid_client", + } { + if strings.Contains(lower, code) { + return code + } + } + if isInvalidGrantErrorMessage(lower) { + return "invalid_grant" + } + if statusCode == http.StatusUnauthorized { + return "unauthorized" + } + return "refresh_failed" +} + +func sanitizeRenewalError(auth *Auth, err error) *Error { + if err == nil { + return nil + } + var partErr interface { + OuterRefreshError() error + SIWCRefreshError() error + } + if errors.As(err, &partErr) && partErr != nil { + outerErr := partErr.OuterRefreshError() + siwcErr := partErr.SIWCRefreshError() + if outerErr != nil && isTerminalRefreshError(auth, outerErr) { + err = outerErr + } else if siwcErr != nil && isTerminalRefreshError(auth, siwcErr) { + err = siwcErr + } + } + raw := err.Error() + statusCode := statusCodeFromError(err) + if statusCode == 0 { + if isUnauthorizedError(err) { + statusCode = http.StatusUnauthorized + } else if isInvalidGrantError(err) || isTerminalRefreshGrantMessage(raw) { + statusCode = http.StatusBadRequest + } + } + terminal := isTerminalRefreshError(auth, err) + return &Error{ + Code: terminalRefreshErrorCode(raw, statusCode), + Message: sanitizeRenewalMessage(raw), + HTTPStatus: statusCode, + Retryable: !terminal, + } +} + func refreshErrorFromError(err error) *Error { if err == nil { return nil @@ -2255,7 +2424,7 @@ func refreshErrorFromError(err error) *Error { if statusCode == 0 && isUnauthorizedError(err) { statusCode = http.StatusUnauthorized } - authErr := &Error{Message: err.Error(), HTTPStatus: statusCode} + authErr := &Error{Message: sanitizeRenewalMessage(err.Error()), HTTPStatus: statusCode} if statusCode == http.StatusUnauthorized { authErr.Code = "unauthorized" authErr.Retryable = false diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 4f8bf2de497..8b2f8980236 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -16,6 +16,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" cliproxysession "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/session" @@ -351,8 +352,8 @@ func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExec }) opts.Headers = mergeRequestHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) if len(resp.Body) > 0 { - req.Payload = bytes.Clone(resp.Body) - opts.OriginalRequest = bytes.Clone(resp.Body) + req.Payload = resp.Body + opts.OriginalRequest = resp.Body } if path := strings.TrimSpace(resp.Path); path != "" { if opts.Metadata == nil { @@ -661,7 +662,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: execOpts, AccessToken: authAccessTokenForRequest(execCtx, auth, execOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, CompletedResponse: errExec == nil && !resp.Incomplete && isCompletedExecutionPayload(resp.Payload), Options: execOpts, AccessToken: authAccessTokenForRequest(execCtx, auth, execOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -2227,3 +2228,223 @@ func syncMetadataSessionToContext(ctx context.Context, metadata map[string]any) ctx = logging.WithClientRequestMetadata(ctx, clientMeta) return util.WithSessionID(ctx, clientMeta.SessionID) } + +// VerifyInference performs an operator-triggered, account-pinned inference verification +// using an isolated session and minimal non-user prompt/token budget. It never falls back +// to another credential and requires a completed response before marking InferenceVerified. +func (m *Manager) VerifyInference(ctx context.Context, id, model string) (*Auth, error) { + if m == nil { + return nil, &Error{Code: "manager_unavailable", Message: "auth manager is not initialized", HTTPStatus: http.StatusServiceUnavailable} + } + id = strings.TrimSpace(id) + if id == "" { + return nil, &Error{Code: "invalid_request", Message: "auth id is required", HTTPStatus: http.StatusBadRequest} + } + auth, ok := m.GetByID(id) + if !ok || auth == nil { + return nil, &Error{Code: "auth_not_found", Message: "credential not found", HTTPStatus: http.StatusNotFound} + } + if auth.Disabled || auth.Status == StatusDisabled { + return auth, &Error{Code: "auth_disabled", Message: "credential is disabled", HTTPStatus: http.StatusConflict} + } + if auth.RequiresSignIn { + return auth, &Error{Code: "sign_in_required", Message: "credential requires sign-in before inference verification", HTTPStatus: http.StatusConflict} + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + m.mu.RLock() + executor := m.executors[provider] + m.mu.RUnlock() + if executor == nil { + return auth, &Error{Code: "executor_not_found", Message: "executor not registered for provider", HTTPStatus: http.StatusBadRequest} + } + executor = executorForAuth(executor, auth) + + routeModel := strings.TrimSpace(model) + if routeModel == "" { + if models := registry.GetGlobalRegistry().GetModelsForClient(auth.ID); len(models) > 0 && models[0] != nil && strings.TrimSpace(models[0].ID) != "" { + routeModel = strings.TrimSpace(models[0].ID) + } else { + routeModel = "gpt-5" + } + } + verifyModel := routeModel + if candidates, _, _, _ := m.executionModelCandidatesWithAlias(auth, routeModel); len(candidates) > 0 && strings.TrimSpace(candidates[0]) != "" { + verifyModel = strings.TrimSpace(candidates[0]) + } + stateModel := m.selectionModelKeyForAuth(auth, routeModel) + if stateModel == "" { + stateModel = canonicalModelKey(routeModel) + } + + 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)) + req := cliproxyexecutor.Request{ + Model: verifyModel, + Payload: payload, + } + opts := cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + OriginalRequest: payload, + Headers: http.Header{"X-Session-Id": {sessionID}}, + Metadata: map[string]any{cliproxyexecutor.CanonicalSessionIDMetadataKey: sessionID}, + } + + execCtx := newUpstreamAttemptContext(ctx) + execCtx = syncMetadataSessionToContext(execCtx, opts.Metadata) + executionSeq := m.nextExecutionSeq() + accessToken := authAccessTokenForRequest(execCtx, auth, opts.Headers) + credVersion := auth.CredentialVersion + regEpoch := auth.RegistrationEpoch + + streamRes, errStream := executor.ExecuteStream(execCtx, auth, req, opts) + if errStream != nil { + rerr := resultErrorFromError(errStream) + if rerr != nil && rerr.HTTPStatus == http.StatusNotImplemented { + opts.Stream = false + resp, errExec := executor.Execute(execCtx, auth, req, opts) + rerrExec := resultErrorFromError(errExec) + if rerrExec != nil { + rerrExec.Message = sanitizeRenewalMessage(rerrExec.Message) + } + if errExec != nil { + m.MarkResult(execCtx, Result{ + AuthID: auth.ID, + Provider: provider, + Model: stateModel, + RouteModel: routeModel, + Success: false, + Error: rerrExec, + RetryAfter: retryAfterFromError(errExec), + CredentialScope: isCredentialScopedError(errExec), + Options: opts, + AccessToken: accessToken, + CredentialVersion: credVersion, + RegistrationEpoch: regEpoch, + ExecutionSeq: executionSeq, + }) + updated, _ := m.GetByID(auth.ID) + return updated, errExec + } + if resp.Incomplete || !isCompletedExecutionPayload(resp.Payload) { + incompleteErr := &Error{Code: "incomplete_response", Message: "verification response did not complete", HTTPStatus: http.StatusBadGateway} + updated, _ := m.GetByID(auth.ID) + return updated, incompleteErr + } + m.MarkResult(execCtx, Result{ + AuthID: auth.ID, + Provider: provider, + Model: stateModel, + RouteModel: routeModel, + Success: true, + CompletedResponse: true, + Options: opts, + AccessToken: accessToken, + CredentialVersion: credVersion, + RegistrationEpoch: regEpoch, + ExecutionSeq: executionSeq, + }) + updated, okUpdated := m.GetByID(auth.ID) + if !okUpdated || updated == nil { + return nil, &Error{Code: "auth_not_found", Message: "credential was removed during verification", HTTPStatus: http.StatusConflict} + } + if updated.CredentialVersion != credVersion || updated.RegistrationEpoch != regEpoch || !updated.InferenceVerified { + return updated, &Error{Code: "credential_superseded", Message: "credential was replaced during verification; re-run verification for the new credential", HTTPStatus: http.StatusConflict} + } + return updated, nil + } + if rerr != nil { + rerr.Message = sanitizeRenewalMessage(rerr.Message) + } + m.MarkResult(execCtx, Result{ + AuthID: auth.ID, + Provider: provider, + Model: stateModel, + RouteModel: routeModel, + Success: false, + Error: rerr, + RetryAfter: retryAfterFromError(errStream), + CredentialScope: isCredentialScopedError(errStream), + Options: opts, + AccessToken: accessToken, + CredentialVersion: credVersion, + RegistrationEpoch: regEpoch, + ExecutionSeq: executionSeq, + }) + updated, _ := m.GetByID(auth.ID) + return updated, errStream + } + + var sawCompleted bool + var sawIncomplete bool + var chunkErr error + if streamRes != nil && streamRes.Chunks != nil { + for chunk := range streamRes.Chunks { + if chunk.Err != nil && chunkErr == nil { + chunkErr = chunk.Err + } + if chunk.Incomplete { + sawIncomplete = true + } + if len(chunk.Payload) > 0 { + comp, incomp := streamPayloadCompletionFlags(chunk.Payload) + if comp { + sawCompleted = true + } + if incomp { + sawIncomplete = true + } + } + } + } + if chunkErr != nil { + rerrChunk := resultErrorFromError(chunkErr) + if rerrChunk != nil { + rerrChunk.Message = sanitizeRenewalMessage(rerrChunk.Message) + } + m.MarkResult(execCtx, Result{ + AuthID: auth.ID, + Provider: provider, + Model: stateModel, + RouteModel: routeModel, + Success: false, + Error: rerrChunk, + RetryAfter: retryAfterFromError(chunkErr), + CredentialScope: isCredentialScopedError(chunkErr), + Options: opts, + AccessToken: accessToken, + CredentialVersion: credVersion, + RegistrationEpoch: regEpoch, + ExecutionSeq: executionSeq, + }) + updated, _ := m.GetByID(auth.ID) + return updated, chunkErr + } + if !sawCompleted || sawIncomplete { + incompleteErr := &Error{Code: "incomplete_response", Message: "verification stream ended without a completed response", HTTPStatus: http.StatusBadGateway} + updated, _ := m.GetByID(auth.ID) + return updated, incompleteErr + } + m.MarkResult(execCtx, Result{ + AuthID: auth.ID, + Provider: provider, + Model: stateModel, + RouteModel: routeModel, + Success: true, + CompletedResponse: true, + Options: opts, + AccessToken: accessToken, + CredentialVersion: credVersion, + RegistrationEpoch: regEpoch, + ExecutionSeq: executionSeq, + }) + updated, okUpdated := m.GetByID(auth.ID) + if !okUpdated || updated == nil { + return nil, &Error{Code: "auth_not_found", Message: "credential was removed during verification", HTTPStatus: http.StatusConflict} + } + if updated.CredentialVersion != credVersion || updated.RegistrationEpoch != regEpoch || !updated.InferenceVerified { + return updated, &Error{Code: "credential_superseded", Message: "credential was replaced during verification; re-run verification for the new credential", HTTPStatus: http.StatusConflict} + } + return updated, nil +} diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index 488eb3cd360..9a54071bc93 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -104,6 +104,12 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { if auth.CreatedAt.IsZero() { auth.CreatedAt = now } + if auth.AdoptedAt.IsZero() { + auth.AdoptedAt = now + } + if auth.AdoptedGrantHash == "" { + auth.AdoptedGrantHash = AuthGrantFingerprint(auth) + } auth.UpdatedAt = now cooldownStateChanged := normalizeModelStates(auth) if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { @@ -392,7 +398,7 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo } else if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { auth.ModelStates = existing.ModelStates } - if isFileSnapshot { + if isFileSnapshot || (credChanged && mode == updateModeReplace && existing.Quota.Exceeded && existing.Quota.NextRecoverAt.After(now) && !auth.Quota.NextRecoverAt.After(now)) { if (auth.Status == "" || auth.Status == StatusActive) && existing.Status != "" { auth.Status = existing.Status } @@ -412,14 +418,26 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo } if credChanged || mode == updateModeRefresh { applyPreservedRejectedTokens(existing, auth, mode == updateModeRefresh) + applyPreservedRenewalState(existing, auth, mode, now) if auth.RejectedAccessToken == "" || authAccessToken(auth) != auth.RejectedAccessToken { auth.RefreshFailures = 0 if credChanged && mode != updateModeRefresh && auth.NextRefreshAfter.Equal(existing.NextRefreshAfter) { - auth.NextRefreshAfter = time.Time{} + if authHasUnchangedRejectedRefreshToken(auth) { + if exp, ok := auth.AccessTokenExpirationTime(); ok && !exp.IsZero() && exp.After(now) { + auth.NextRefreshAfter = exp + } else { + auth.NextRefreshAfter = time.Time{} + } + } else { + auth.NextRefreshAfter = time.Time{} + } } + canServeWithAccessToken := !auth.RequiresSignIn || auth.HasValidAccessToken(now) isAuthOrGrantFailure := hasUnauthorizedAuthFailure(existing) || - (auth.LastError != nil && (isUnauthorizedError(auth.LastError) || isInvalidGrantResultError(auth.LastError))) - if !auth.Quota.NextRecoverAt.After(now) && + existing.RequiresSignIn || + (existing.RefreshError != nil && isTerminalRefreshError(existing, existing.RefreshError)) || + (auth.LastError != nil && (isUnauthorizedError(auth.LastError) || isInvalidGrantResultError(auth.LastError) || isTerminalRefreshError(auth, auth.LastError))) + if canServeWithAccessToken && !auth.Quota.NextRecoverAt.After(now) && (isAuthOrGrantFailure || (auth.Unavailable && !auth.NextRetryAfter.After(now))) { auth.Unavailable = false auth.NextRetryAfter = time.Time{} @@ -428,18 +446,25 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.StatusMessage = "" auth.Status = StatusActive cooldownStateChanged = true - } else if isAuthOrGrantFailure && auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) { + } else if canServeWithAccessToken && isAuthOrGrantFailure && auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) { + auth.Unavailable = true + if !auth.NextRetryAfter.After(now) { + auth.NextRetryAfter = auth.Quota.NextRecoverAt + } auth.LastError = nil auth.StatusMessage = auth.Quota.Reason cooldownStateChanged = true } - resumed := clearUnauthorizedModelStates(auth, now) - if len(resumed) > 0 { - cooldownStateChanged = true + if canServeWithAccessToken { + resumed := clearUnauthorizedModelStates(auth, now) + if len(resumed) > 0 { + cooldownStateChanged = true + } } } } else { applyPreservedRejectedTokens(existing, auth, false) + applyPreservedRenewalState(existing, auth, mode, now) } // A refresh that confirms quota restoration (issue #6690) releases the // credential-wide cooldown early; otherwise the active credential @@ -460,6 +485,20 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.Status = existing.Status } } + } else { + applyPreservedRejectedTokens(existing, auth, mode == updateModeRefresh) + applyPreservedRenewalState(existing, auth, mode, now) + if !credChanged { + auth.LastRefreshedAt = existing.LastRefreshedAt + auth.NextRefreshAfter = existing.NextRefreshAfter + auth.RefreshFailures = existing.RefreshFailures + } else if authHasUnchangedRejectedRefreshToken(auth) { + if exp, ok := auth.AccessTokenExpirationTime(); ok && !exp.IsZero() && exp.After(now) { + auth.NextRefreshAfter = exp + } else { + auth.NextRefreshAfter = time.Time{} + } + } } auth.UpdatedAt = now cooldownStateChanged = normalizeModelStates(auth) || cooldownStateChanged @@ -501,6 +540,39 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo auth.RejectedAccessToken = probe.RejectedAccessToken auth.RejectedSIWCAccessToken = probe.RejectedSIWCAccessToken } + if existing.RequiresSignIn != existingBeforeSave.RequiresSignIn || + !equalErrorPtr(existing.RefreshError, existingBeforeSave.RefreshError) || + existing.RejectedRefreshToken != existingBeforeSave.RejectedRefreshToken || + existing.RejectedSIWCRefreshToken != existingBeforeSave.RejectedSIWCRefreshToken || + existing.RenewalWarnedRefreshToken != existingBeforeSave.RenewalWarnedRefreshToken || + !existing.RenewalEventAt.Equal(existingBeforeSave.RenewalEventAt) || + existing.InferenceVerified != existingBeforeSave.InferenceVerified || + !existing.LastInferenceVerifiedAt.Equal(existingBeforeSave.LastInferenceVerifiedAt) || + existing.RenewalVerified != existingBeforeSave.RenewalVerified || + !existing.LastRenewalVerifiedAt.Equal(existingBeforeSave.LastRenewalVerifiedAt) { + probe := auth.Clone() + probe.CodexRefreshPatch = refreshPatchBeforeSave + applyPreservedRenewalState(existing, probe, mode, now) + auth.RequiresSignIn = probe.RequiresSignIn + auth.RefreshError = probe.RefreshError + auth.RejectedRefreshToken = probe.RejectedRefreshToken + auth.RejectedSIWCRefreshToken = probe.RejectedSIWCRefreshToken + auth.RenewalWarnedRefreshToken = probe.RenewalWarnedRefreshToken + auth.RenewalEventAt = probe.RenewalEventAt + auth.AdoptedAt = probe.AdoptedAt + auth.AdoptedGrantHash = probe.AdoptedGrantHash + auth.InferenceVerified = probe.InferenceVerified + auth.LastInferenceVerifiedAt = probe.LastInferenceVerifiedAt + auth.RenewalVerified = probe.RenewalVerified + auth.LastRenewalVerifiedAt = probe.LastRenewalVerifiedAt + if authHasUnchangedRejectedRefreshToken(auth) && !existingBeforeSave.RequiresSignIn { + if exp, okExp := auth.AccessTokenExpirationTime(); okExp && !exp.IsZero() && exp.After(now) { + if auth.NextRefreshAfter.IsZero() || auth.NextRefreshAfter.Before(exp) { + auth.NextRefreshAfter = exp + } + } + } + } if existing.Generation >= auth.Generation { auth.Generation = existing.Generation + 1 } @@ -518,6 +590,16 @@ func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mo return auth.Clone(), true, cooldownStateChanged, nil } +func equalErrorPtr(a, b *Error) bool { + if a == nil && b == nil { + return true + } + if a == nil || b == nil { + return false + } + return a.Code == b.Code && a.Message == b.Message && a.Retryable == b.Retryable && a.HTTPStatus == b.HTTPStatus +} + // Remove deletes an auth from runtime state without persisting. // Disk and token-store deletion must be handled by the caller. func (m *Manager) Remove(ctx context.Context, id string) { diff --git a/sdk/cliproxy/auth/conductor_persistence.go b/sdk/cliproxy/auth/conductor_persistence.go index a845f21f1da..fef0e160024 100644 --- a/sdk/cliproxy/auth/conductor_persistence.go +++ b/sdk/cliproxy/auth/conductor_persistence.go @@ -241,7 +241,16 @@ func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) { next := reflect.ValueOf(after).Elem() for i := 0; i < dst.NumField(); i++ { field := dst.Type().Field(i) - if field.PkgPath != "" || field.Name == "ID" || field.Name == "RegistrationEpoch" || field.Name == "CredentialVersion" || field.Name == "Generation" || field.Name == "RejectedAccessToken" || field.Name == "RejectedSIWCAccessToken" { + switch field.Name { + case "ID", "RegistrationEpoch", "CredentialVersion", "Generation", + "RejectedAccessToken", "RejectedSIWCAccessToken", + "RejectedRefreshToken", "RejectedSIWCRefreshToken", "RenewalWarnedRefreshToken", + "RequiresSignIn", "RefreshError", "RenewalEventAt", + "AdoptedAt", "InferenceVerified", "LastInferenceVerifiedAt", + "RenewalVerified", "LastRenewalVerifiedAt": + continue + } + if field.PkgPath != "" { continue } d, b, n := dst.Field(i), base.Field(i), next.Field(i) diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index e21998ce51e..b2421f7ac94 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -452,6 +452,21 @@ func authHasRefreshCredential(auth *Auth) bool { // CredentialsChanged reports whether authentication credentials (tokens or API keys) // differ between two Auth records. func CredentialsChanged(existing, incoming *Auth) bool { + if existing == nil || incoming == nil { + return false + } + if secretGrantChanged(existing, incoming) { + return true + } + existingUseSIWC, existingHasUseSIWC := parseBoolAny(existing.Metadata["use_siwc"]) + incomingUseSIWC, incomingHasUseSIWC := parseBoolAny(incoming.Metadata["use_siwc"]) + if existingHasUseSIWC != incomingHasUseSIWC || existingUseSIWC != incomingUseSIWC { + return true + } + return false +} + +func secretGrantChanged(existing, incoming *Auth) bool { if existing == nil || incoming == nil { return false } @@ -495,14 +510,275 @@ func CredentialsChanged(existing, incoming *Auth) bool { authSIWCTokenField(existing, "client_id") != authSIWCTokenField(incoming, "client_id") { return true } - existingUseSIWC, existingHasUseSIWC := parseBoolAny(existing.Metadata["use_siwc"]) - incomingUseSIWC, incomingHasUseSIWC := parseBoolAny(incoming.Metadata["use_siwc"]) - if existingHasUseSIWC != incomingHasUseSIWC || existingUseSIWC != incomingUseSIWC { + return false +} + +func isSIWCRefreshFailure(err error) bool { + if err == nil { + return false + } + var marker interface{ IsSIWCRefreshError() bool } + if errors.As(err, &marker) && marker != nil && marker.IsSIWCRefreshError() { + return true + } + lower := strings.ToLower(err.Error()) + return strings.Contains(lower, "siwc") || strings.Contains(lower, "openai oauth token request failed") +} + +func recordRejectedRefreshTokens(auth *Auth, err error) { + if auth == nil { + return + } + outerRT := authRefreshToken(auth) + siwcRT := authSIWCTokenField(auth, "refresh_token") + var partErr interface { + OuterRefreshError() error + SIWCRefreshError() error + } + if errors.As(err, &partErr) && partErr != nil { + outerErr := partErr.OuterRefreshError() + siwcErr := partErr.SIWCRefreshError() + if outerErr != nil && isTerminalRefreshError(auth, outerErr) && outerRT != "" { + auth.RejectedRefreshToken = outerRT + } + if siwcErr != nil && isTerminalRefreshError(auth, siwcErr) && siwcRT != "" { + auth.RejectedSIWCRefreshToken = siwcRT + } + return + } + if isSIWCRefreshFailure(err) && siwcRT != "" { + auth.RejectedSIWCRefreshToken = siwcRT + return + } + if outerRT != "" { + auth.RejectedRefreshToken = outerRT + return + } + if siwcRT != "" { + auth.RejectedSIWCRefreshToken = siwcRT + } +} + +func authHasAnyRejectedRefreshPartition(auth *Auth) bool { + if auth == nil { + return false + } + if auth.RejectedRefreshToken != "" && authRefreshToken(auth) == auth.RejectedRefreshToken { + return true + } + if auth.RejectedSIWCRefreshToken != "" && authSIWCTokenField(auth, "refresh_token") == auth.RejectedSIWCRefreshToken { return true } return false } +func authHasUnchangedRejectedRefreshToken(auth *Auth) bool { + if auth == nil || !auth.RequiresSignIn { + return false + } + outerRT := authRefreshToken(auth) + siwcRT := authSIWCTokenField(auth, "refresh_token") + hasOuter := outerRT != "" + hasSIWC := siwcRT != "" + if !hasOuter && !hasSIWC { + return true + } + if hasOuter && auth.RejectedRefreshToken != outerRT { + return false + } + if hasSIWC && auth.RejectedSIWCRefreshToken != siwcRT { + return false + } + return true +} + +func applyPreservedRenewalState(existing, updated *Auth, mode updateAuthMode, now time.Time) { + if updated == nil { + return + } + if now.IsZero() { + now = time.Now() + } + if mode == updateModeRefresh { + var outerErr, siwcErr error + var failedOuterRT, failedSIWCRT string + if updated.CodexRefreshPatch != nil { + outerErr = updated.CodexRefreshPatch.OuterError() + siwcErr = updated.CodexRefreshPatch.SIWCError() + failedOuterRT = updated.CodexRefreshPatch.FailedOuterRefreshToken() + failedSIWCRT = updated.CodexRefreshPatch.FailedSIWCRefreshToken() + } + currentOuterRT := authRefreshToken(updated) + currentSIWCRT := authSIWCTokenField(updated, "refresh_token") + outerStillMatchesFailed := failedOuterRT == "" || (currentOuterRT == failedOuterRT && (existing == nil || authRefreshToken(existing) == failedOuterRT)) + siwcStillMatchesFailed := failedSIWCRT == "" || (currentSIWCRT == failedSIWCRT && (existing == nil || authSIWCTokenField(existing, "refresh_token") == failedSIWCRT)) + outerTerminal := outerErr != nil && outerStillMatchesFailed && isTerminalRefreshError(updated, outerErr) + siwcTerminal := siwcErr != nil && siwcStillMatchesFailed && isTerminalRefreshError(updated, siwcErr) + preservedRejectedOuter := existing != nil && existing.RejectedRefreshToken != "" && currentOuterRT == existing.RejectedRefreshToken + preservedRejectedSIWC := existing != nil && existing.RejectedSIWCRefreshToken != "" && currentSIWCRT == existing.RejectedSIWCRefreshToken + if outerTerminal || siwcTerminal || preservedRejectedOuter || preservedRejectedSIWC { + updated.RequiresSignIn = true + updated.RenewalVerified = false + updated.LastRenewalVerifiedAt = time.Time{} + if outerTerminal { + updated.RejectedRefreshToken = authRefreshToken(updated) + updated.RefreshError = sanitizeRenewalError(updated, outerErr) + } else if preservedRejectedOuter { + updated.RejectedRefreshToken = existing.RejectedRefreshToken + } else { + updated.RejectedRefreshToken = "" + } + if siwcTerminal { + updated.RejectedSIWCRefreshToken = authSIWCTokenField(updated, "refresh_token") + updated.RefreshError = sanitizeRenewalError(updated, siwcErr) + } else if preservedRejectedSIWC { + updated.RejectedSIWCRefreshToken = existing.RejectedSIWCRefreshToken + } else { + updated.RejectedSIWCRefreshToken = "" + } + if updated.RefreshError == nil && existing != nil { + updated.RefreshError = cloneError(existing.RefreshError) + } + if existing != nil && !existing.RenewalEventAt.IsZero() { + updated.RenewalEventAt = existing.RenewalEventAt + } else { + updated.RenewalEventAt = now + } + if existing != nil && !existing.AdoptedAt.IsZero() { + updated.AdoptedAt = existing.AdoptedAt + } else if updated.AdoptedAt.IsZero() { + updated.AdoptedAt = now + } + if existing != nil && existing.AdoptedGrantHash != "" { + updated.AdoptedGrantHash = existing.AdoptedGrantHash + } else if updated.AdoptedGrantHash == "" { + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + } + if existing != nil { + updated.InferenceVerified = existing.InferenceVerified + updated.LastInferenceVerifiedAt = existing.LastInferenceVerifiedAt + } + return + } + upstreamVerified := updated.CodexRefreshPatch == nil || updated.CodexRefreshPatch.UpstreamRefreshed() + if !upstreamVerified { + // Disk-only credential adoption during Refresh: treat as a newly loaded grant awaiting its first renewal. + updated.RequiresSignIn = false + updated.RefreshError = nil + updated.RejectedRefreshToken = "" + updated.RejectedSIWCRefreshToken = "" + updated.RenewalWarnedRefreshToken = "" + updated.RenewalEventAt = time.Time{} + updated.AdoptedAt = now + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + updated.InferenceVerified = false + updated.LastInferenceVerifiedAt = time.Time{} + updated.RenewalVerified = false + updated.LastRenewalVerifiedAt = time.Time{} + return + } + updated.RequiresSignIn = false + updated.RefreshError = nil + updated.RejectedRefreshToken = "" + updated.RejectedSIWCRefreshToken = "" + updated.RenewalWarnedRefreshToken = "" + updated.RenewalEventAt = time.Time{} + if existing != nil && !existing.AdoptedAt.IsZero() { + updated.AdoptedAt = existing.AdoptedAt + } else if updated.AdoptedAt.IsZero() { + updated.AdoptedAt = now + } + if existing != nil && existing.AdoptedGrantHash != "" { + updated.AdoptedGrantHash = existing.AdoptedGrantHash + } else if existing != nil { + updated.AdoptedGrantHash = AuthGrantFingerprint(existing) + } else if updated.AdoptedGrantHash == "" { + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + } + if existing != nil { + updated.InferenceVerified = existing.InferenceVerified + updated.LastInferenceVerifiedAt = existing.LastInferenceVerifiedAt + } + updated.RenewalVerified = true + updated.LastRenewalVerifiedAt = now + return + } + if existing == nil { + if updated.AdoptedAt.IsZero() { + updated.AdoptedAt = now + } + if updated.AdoptedGrantHash == "" { + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + } + return + } + if secretGrantChanged(existing, updated) { + rejectedOuterUnchanged := existing.RejectedRefreshToken != "" && authRefreshToken(updated) == existing.RejectedRefreshToken + rejectedSIWCUnchanged := existing.RejectedSIWCRefreshToken != "" && authSIWCTokenField(updated, "refresh_token") == existing.RejectedSIWCRefreshToken + sameRejectedRefresh := existing.RequiresSignIn && (rejectedOuterUnchanged || rejectedSIWCUnchanged) + if !sameRejectedRefresh { + updated.RequiresSignIn = false + updated.RefreshError = nil + updated.RejectedRefreshToken = "" + updated.RejectedSIWCRefreshToken = "" + updated.RenewalWarnedRefreshToken = "" + updated.RenewalEventAt = time.Time{} + updated.AdoptedAt = now + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + updated.InferenceVerified = false + updated.LastInferenceVerifiedAt = time.Time{} + updated.RenewalVerified = false + updated.LastRenewalVerifiedAt = time.Time{} + return + } + if authAccessToken(updated) != authAccessToken(existing) || authOuterAccessToken(updated) != authOuterAccessToken(existing) { + updated.RequiresSignIn = existing.RequiresSignIn + updated.RefreshError = cloneError(existing.RefreshError) + if rejectedOuterUnchanged { + updated.RejectedRefreshToken = existing.RejectedRefreshToken + } else { + updated.RejectedRefreshToken = "" + } + if rejectedSIWCUnchanged { + updated.RejectedSIWCRefreshToken = existing.RejectedSIWCRefreshToken + } else { + updated.RejectedSIWCRefreshToken = "" + } + updated.RenewalWarnedRefreshToken = existing.RenewalWarnedRefreshToken + updated.RenewalEventAt = existing.RenewalEventAt + updated.AdoptedAt = now + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + updated.InferenceVerified = false + updated.LastInferenceVerifiedAt = time.Time{} + updated.RenewalVerified = false + updated.LastRenewalVerifiedAt = time.Time{} + if exp, ok := updated.AccessTokenExpirationTime(); ok && !exp.IsZero() && exp.After(now) { + updated.NextRefreshAfter = exp + } + return + } + } + updated.RequiresSignIn = existing.RequiresSignIn + updated.RefreshError = cloneError(existing.RefreshError) + updated.RejectedRefreshToken = existing.RejectedRefreshToken + updated.RejectedSIWCRefreshToken = existing.RejectedSIWCRefreshToken + updated.RenewalWarnedRefreshToken = existing.RenewalWarnedRefreshToken + updated.RenewalEventAt = existing.RenewalEventAt + updated.AdoptedAt = existing.AdoptedAt + if updated.AdoptedAt.IsZero() { + updated.AdoptedAt = now + } + if existing.AdoptedGrantHash != "" { + updated.AdoptedGrantHash = existing.AdoptedGrantHash + } else { + updated.AdoptedGrantHash = AuthGrantFingerprint(updated) + } + updated.InferenceVerified = existing.InferenceVerified + updated.LastInferenceVerifiedAt = existing.LastInferenceVerifiedAt + updated.RenewalVerified = existing.RenewalVerified + updated.LastRenewalVerifiedAt = existing.LastRenewalVerifiedAt +} + func authSIWCTokenField(auth *Auth, key string) string { if auth == nil || auth.Metadata == nil { return "" @@ -982,6 +1258,52 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc } } + // Do not repeatedly submit an unchanged refresh grant already known to be terminally rejected. + if !forceRefresh && authHasUnchangedRejectedRefreshToken(auth) { + checkNow := time.Now() + accessTokenRejected := failedAccessToken != "" && authMatchesTokenPartition(auth, failedAccessToken) + if !accessTokenRejected && auth.HasValidAccessToken(checkNow) { + if auth.RefreshError != nil { + return nil, auth.RefreshError + } + if auth.LastError != nil { + return nil, auth.LastError + } + return nil, errors.New("refresh token requires sign-in") + } + // Access token has expired (or was rejected by a 401) while refresh grant remains terminally rejected. + m.mu.Lock() + if current := m.auths[id]; current != nil && + current.RegistrationEpoch == auth.RegistrationEpoch && + current.CredentialVersion == auth.CredentialVersion && + authHasUnchangedRejectedRefreshToken(current) { + current.Generation++ + current.UpdatedAt = checkNow + current.Unavailable = true + current.Status = StatusError + current.NextRefreshAfter = time.Time{} + current.NextRetryAfter = time.Time{} + current.RefreshFailures = 0 + current.RequiresSignIn = true + msg := "unauthorized (refresh token invalid)" + if current.RefreshError != nil && strings.TrimSpace(current.RefreshError.Message) != "" { + msg = current.RefreshError.Message + } + current.LastError = &Error{Code: "unauthorized", Message: msg, HTTPStatus: http.StatusUnauthorized} + current.StatusMessage = "unauthorized (refresh token invalid)" + m.auths[id] = current + if m.scheduler != nil { + m.scheduler.upsertAuth(current.Clone()) + } + } + m.mu.Unlock() + m.queueRefreshUnschedule(id) + if auth.RefreshError != nil { + return nil, auth.RefreshError + } + return nil, &Error{Code: "unauthorized", Message: "unauthorized (refresh token invalid)", HTTPStatus: http.StatusUnauthorized} + } + base := auth.Clone() updated, err := exec.Refresh(ctx, base.Clone()) now := time.Now() @@ -1010,7 +1332,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc shouldUnschedule := false m.mu.Lock() if current := m.auths[id]; current != nil { - if base != nil && current.RegistrationEpoch != base.RegistrationEpoch { + if base != nil && (current.RegistrationEpoch != base.RegistrationEpoch || current.CredentialVersion != base.CredentialVersion || secretGrantChanged(base, current)) { m.mu.Unlock() return nil, err } @@ -1018,6 +1340,8 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc m.mu.Unlock() return nil, err } + terminalRefresh := isTerminalRefreshError(current, err) + sanitizedRefreshErr := sanitizeRenewalError(current, err) wasTerminalUnauthorized := hasUnauthorizedAuthFailure(current) if wasTerminalUnauthorized { current.Generation++ @@ -1026,8 +1350,14 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.Status = StatusError current.NextRefreshAfter = time.Time{} current.NextRetryAfter = time.Time{} - if isUnauthorizedError(err) || isInvalidGrantError(err) { - current.LastError = &Error{Code: "unauthorized", Message: err.Error(), HTTPStatus: http.StatusUnauthorized} + current.RefreshError = sanitizedRefreshErr + if terminalRefresh || isInvalidGrantError(err) { + current.RequiresSignIn = true + recordRejectedRefreshTokens(current, err) + if current.RenewalEventAt.IsZero() { + current.RenewalEventAt = now + } + current.LastError = &Error{Code: "unauthorized", Message: sanitizedRefreshErr.Message, HTTPStatus: http.StatusUnauthorized} current.StatusMessage = "unauthorized (refresh token invalid)" } m.auths[id] = current @@ -1041,6 +1371,26 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.Generation++ current.UpdatedAt = now current.LastError = refreshErrorFromError(err) + if terminalRefresh { + current.RefreshError = sanitizedRefreshErr + current.RequiresSignIn = true + current.RenewalVerified = false + current.LastRenewalVerifiedAt = time.Time{} + recordRejectedRefreshTokens(current, err) + if current.RenewalEventAt.IsZero() { + current.RenewalEventAt = now + } + } else if authHasAnyRejectedRefreshPartition(current) { + // Preserve existing terminal sign-in requirement from the other partition + // when a healthy partition encounters a transient refresh failure. + current.RequiresSignIn = true + if current.RefreshError == nil || !isTerminalRefreshError(current, current.RefreshError) { + current.RefreshError = sanitizedRefreshErr + } + } else { + current.RefreshError = sanitizedRefreshErr + current.RequiresSignIn = false + } isDisabled = current.Disabled || current.Status == StatusDisabled hasValidAccessToken := current.HasValidAccessToken(now) @@ -1062,7 +1412,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.StatusMessage = "disabled" } shouldReschedule = true - } else if accessTokenRejected && invalidGrant { + } else if accessTokenRejected && (invalidGrant || terminalRefresh) { // Neither token can recover without a new login. Stop selecting the // credential until its tokens change instead of retrying it after // every cooldown and returning the same 401 to clients. @@ -1071,17 +1421,27 @@ 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: err.Error(), HTTPStatus: http.StatusUnauthorized} - current.StatusMessage = "unauthorized (refresh token invalid)" + current.LastError = &Error{Code: "unauthorized", Message: sanitizedRefreshErr.Message, HTTPStatus: http.StatusUnauthorized} + if unauthorized && !invalidGrant { + current.StatusMessage = "unauthorized" + } else { + current.StatusMessage = "unauthorized (refresh token invalid)" + } shouldUnschedule = true } else if !hasValidAccessToken || accessTokenRejected { current.Unavailable = true current.Status = StatusError - if unauthorized { + if unauthorized || terminalRefresh { current.NextRefreshAfter = time.Time{} current.NextRetryAfter = time.Time{} current.RefreshFailures = 0 - current.StatusMessage = "unauthorized" + current.LastError = &Error{Code: "unauthorized", Message: sanitizedRefreshErr.Message, HTTPStatus: http.StatusUnauthorized} + if unauthorized { + current.StatusMessage = "unauthorized" + } else { + current.StatusMessage = "unauthorized (refresh token invalid)" + } + shouldUnschedule = true } else if invalidGrant { current.RefreshFailures++ current.NextRefreshAfter = now.Add(invalidGrantBackoffDuration(current.RefreshFailures)) @@ -1093,6 +1453,29 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.StatusMessage = "token expired" shouldReschedule = true } + } else if terminalRefresh { + // Access token remains valid, but refresh grant was terminally rejected. + // Preserve working inference and defer refresh check until access_token expiry + // when no other healthy refresh partition remains. + current.RefreshFailures = 0 + if !authHasUnchangedRejectedRefreshToken(current) { + current.NextRefreshAfter = time.Time{} + shouldReschedule = true + } else if exp, ok := current.AccessTokenExpirationTime(); ok && !exp.IsZero() && exp.After(now) { + current.NextRefreshAfter = exp + shouldReschedule = true + } else { + current.NextRefreshAfter = time.Time{} + shouldUnschedule = true + } + warnKey := current.RejectedRefreshToken + "|" + current.RejectedSIWCRefreshToken + if warnKey == "|" { + warnKey = "terminal" + } + if !current.Unavailable && current.RenewalWarnedRefreshToken != warnKey { + current.RenewalWarnedRefreshToken = warnKey + log.Warnf("credential refresh failed for %s (%s): %s; retaining active credential as access token is unexpired", current.Provider, current.ID, sanitizedRefreshErr.Message) + } } else { // Access token remains valid. Preserve current in-flight/cooldown status without overwrite. nextRetry := now.Add(refreshFailureBackoff) @@ -1109,7 +1492,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc shouldReschedule = true if !current.Unavailable { - log.Warnf("credential refresh failed for %s (%s): %s; retaining active credential as access token is unexpired", current.Provider, current.ID, safeErrorDiagnosticForLog(err)) + log.Warnf("credential refresh failed for %s (%s): %s; retaining active credential as access token is unexpired", current.Provider, current.ID, sanitizedRefreshErr.Message) } } m.auths[id] = current @@ -1135,7 +1518,13 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc } updated.LastRefreshedAt = now updated.NextRefreshAfter = time.Time{} - applyPreservedRejectedTokens(auth, updated, true) + applyPreservedRenewalState(base, updated, updateModeRefresh, now) + if authHasUnchangedRejectedRefreshToken(updated) { + if exp, okExp := updated.AccessTokenExpirationTime(); okExp && !exp.IsZero() && exp.After(now) { + updated.NextRefreshAfter = exp + } + } + applyPreservedRejectedTokens(base, updated, true) if updated.RejectedAccessToken == "" || authAccessToken(updated) != updated.RejectedAccessToken { updated.LastError = nil updated.StatusMessage = "" diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 6f8d6cbc856..be5c66cb9b1 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -1,7 +1,9 @@ package auth import ( + "bytes" "context" + "encoding/json" "net/http" "strings" "time" @@ -150,6 +152,8 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re go func() { defer close(out) var failed bool + var sawCompleted bool + var sawIncomplete bool forward := true var rewriter *StreamRewriter if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" { @@ -184,9 +188,19 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return true } } + if chunk.Incomplete { + sawIncomplete = true + } if len(chunk.Payload) == 0 { return true } + completedChunk, incompleteChunk := streamPayloadCompletionFlags(chunk.Payload) + if completedChunk { + sawCompleted = true + } + if incompleteChunk { + sawIncomplete = true + } payload := rewriteForceMappedStreamChunk(rewriter, chunk.Payload) if len(payload) == 0 { return true @@ -223,12 +237,171 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re } } if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) { - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: true, Options: opts, AccessToken: accessToken, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq}, auth, ephemeralResult) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: true, CompletedResponse: sawCompleted && !sawIncomplete, Options: opts, AccessToken: accessToken, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq}, auth, ephemeralResult) } }() return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} } +func streamPayloadCompletionFlags(payload []byte) (completed, incomplete bool) { + if len(payload) == 0 { + return false, false + } + for _, rawLine := range bytes.Split(payload, []byte("\n")) { + line := bytes.TrimSpace(rawLine) + if len(line) == 0 || bytes.HasPrefix(line, []byte(":")) || bytes.HasPrefix(line, []byte("event:")) || bytes.HasPrefix(line, []byte("id:")) || bytes.HasPrefix(line, []byte("retry:")) { + continue + } + data := line + if bytes.HasPrefix(line, []byte("data:")) { + data = bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))) + } + if bytes.Equal(data, []byte("[DONE]")) { + continue + } + if len(data) < 2 || data[0] != '{' || data[len(data)-1] != '}' { + continue + } + comp, incomp := jsonObjectCompletionFlags(data) + if comp { + completed = true + } + if incomp { + incomplete = true + } + } + return completed, incomplete +} + +func jsonObjectCompletionFlags(rawJSON []byte) (completed, incomplete bool) { + var obj map[string]any + if err := json.Unmarshal(rawJSON, &obj); err != nil || len(obj) == 0 { + return false, false + } + if obj["error"] != nil { + return false, true + } + eventType, _ := obj["type"].(string) + eventType = strings.TrimSpace(eventType) + status, _ := obj["status"].(string) + status = strings.ToLower(strings.TrimSpace(status)) + if respObj, ok := obj["response"].(map[string]any); ok { + if respObj["error"] != nil || respObj["incomplete_details"] != nil { + incomplete = true + } + if respStatus, _ := respObj["status"].(string); strings.TrimSpace(respStatus) != "" { + status = strings.ToLower(strings.TrimSpace(respStatus)) + } + } + switch status { + case "incomplete", "failed", "cancelled", "canceled": + return false, true + } + switch eventType { + case "response.incomplete", "response.failed", "response.cancelled", "response.canceled", "error": + return false, true + case "response.completed", "response.done", "message_stop": + if status == "" || status == "completed" { + return true, incomplete + } + return false, true + case "message_delta": + if deltaObj, ok := obj["delta"].(map[string]any); ok { + if sr, okSR := deltaObj["stop_reason"].(string); okSR { + switch strings.ToLower(strings.TrimSpace(sr)) { + case "max_tokens", "StopSequenceMaxTokens": + incomplete = true + case "end_turn", "stop_sequence", "tool_use": + completed = true + } + } + } + } + if choices, ok := obj["choices"].([]any); ok { + for _, item := range choices { + choice, okChoice := item.(map[string]any) + if !okChoice { + continue + } + fr, okFR := choice["finish_reason"].(string) + if !okFR { + continue + } + switch strings.ToLower(strings.TrimSpace(fr)) { + case "length", "content_filter": + incomplete = true + case "stop", "tool_calls", "function_call", "end_turn": + completed = true + } + } + } + if candidates, ok := obj["candidates"].([]any); ok { + for _, item := range candidates { + cand, okCand := item.(map[string]any) + if !okCand { + continue + } + fr, okFR := cand["finishReason"].(string) + if !okFR { + continue + } + switch strings.ToUpper(strings.TrimSpace(fr)) { + case "MAX_TOKENS", "SAFETY", "RECITATION", "OTHER": + incomplete = true + case "STOP": + completed = true + } + } + } + return completed, incomplete +} + +func isCompletedExecutionPayload(payload []byte) bool { + trimmed := bytes.TrimSpace(payload) + if len(trimmed) < 2 || trimmed[0] != '{' || trimmed[len(trimmed)-1] != '}' { + return false + } + var obj map[string]any + if err := json.Unmarshal(trimmed, &obj); err != nil || len(obj) == 0 { + return false + } + if obj["error"] != nil || obj["incomplete_details"] != nil { + return false + } + comp, incomp := jsonObjectCompletionFlags(trimmed) + if incomp { + return false + } + if comp { + return true + } + status, _ := obj["status"].(string) + status = strings.ToLower(strings.TrimSpace(status)) + if respObj, ok := obj["response"].(map[string]any); ok { + if respObj["error"] != nil || respObj["incomplete_details"] != nil { + return false + } + if respStatus, _ := respObj["status"].(string); strings.TrimSpace(respStatus) != "" { + status = strings.ToLower(strings.TrimSpace(respStatus)) + } + } + switch status { + case "incomplete", "failed", "cancelled", "canceled", "in_progress", "queued": + return false + case "completed": + return true + } + if objType, _ := obj["type"].(string); strings.TrimSpace(objType) == "message" { + if stopReason, ok := obj["stop_reason"].(string); ok && strings.TrimSpace(stopReason) != "" { + return true + } + } + if output, ok := obj["output"].([]any); ok && len(output) > 0 && status == "" { + return true + } + return false +} + func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, routing *apiKeyModelRoutingSnapshot, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} diff --git a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go index 21fe3f9d6f1..1ac52ae2acf 100644 --- a/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go +++ b/sdk/cliproxy/auth/conductor_unauthorized_refresh_test.go @@ -10,6 +10,8 @@ import ( "testing" "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" ) @@ -17,16 +19,21 @@ import ( type unauthorizedRefreshExecutor struct { id string - mu sync.Mutex - executeCalls []string - streamCalls []string - refreshCalls int - tokenInvalid map[string]struct{} - refreshFail bool - refreshErr error - refreshTokens map[string]string - onExecute func(authID string) - onRefresh func() + mu sync.Mutex + executeCalls []string + streamCalls []string + refreshCalls int + tokenInvalid map[string]struct{} + refreshFail bool + refreshErr error + streamChunkErr error + streamPayloads [][]byte + lastStreamReq cliproxyexecutor.Request + lastStreamOpts cliproxyexecutor.Options + refreshTokens map[string]string + refreshPatch *codex.CredentialRefreshPatch + onExecute func(authID string) + onRefresh func() } func (e *unauthorizedRefreshExecutor) Identifier() string { return e.id } @@ -50,11 +57,15 @@ func (e *unauthorizedRefreshExecutor) Execute(_ context.Context, auth *Auth, _ c return cliproxyexecutor.Response{Payload: []byte(auth.ID + ":" + token)}, nil } -func (e *unauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { +func (e *unauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { e.mu.Lock() e.streamCalls = append(e.streamCalls, auth.ID) + e.lastStreamReq = req + e.lastStreamOpts = opts token := authAccessToken(auth) _, invalid := e.tokenInvalid[token] + chunkErr := e.streamChunkErr + customPayloads := e.streamPayloads onExec := e.onExecute e.mu.Unlock() if onExec != nil { @@ -66,8 +77,24 @@ func (e *unauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Aut Message: "Your authentication token has been invalidated. Please try signing in again.", } } - ch := make(chan cliproxyexecutor.StreamChunk, 1) + if chunkErr != nil { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.created"}` + "\n\n")} + ch <- cliproxyexecutor.StreamChunk{Err: chunkErr} + close(ch) + return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Auth": {auth.ID}}, Chunks: ch}, nil + } + if len(customPayloads) > 0 { + ch := make(chan cliproxyexecutor.StreamChunk, len(customPayloads)) + for _, p := range customPayloads { + ch <- cliproxyexecutor.StreamChunk{Payload: p} + } + close(ch) + return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Auth": {auth.ID}}, Chunks: ch}, nil + } + ch := make(chan cliproxyexecutor.StreamChunk, 2) ch <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID + ":" + token)} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.completed"}` + "\n\n")} close(ch) return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Auth": {auth.ID}}, Chunks: ch}, nil } @@ -79,6 +106,7 @@ func (e *unauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*A err := e.refreshErr fail := e.refreshFail next := e.refreshTokens[auth.ID] + patch := e.refreshPatch e.mu.Unlock() if onRef != nil { onRef() @@ -96,11 +124,14 @@ func (e *unauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*A next = "refreshed-access-token" } auth.Metadata["access_token"] = next + if patch != nil { + auth.CodexRefreshPatch = patch + } return auth, nil } -func (e *unauthorizedRefreshExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +func (e *unauthorizedRefreshExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(`{"total_tokens":5}`)}, nil } func (e *unauthorizedRefreshExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { @@ -795,3 +826,1006 @@ func reloginSnapshot(t *testing.T, path, token string, expiry time.Time) *Auth { a.BindFileSnapshot(path, data) return a } + +func TestManager_CodexTerminalRenewalRecoveryLifecycle_Issue5783(t *testing.T) { + ctx := WithSkipPersist(context.Background()) + model := "gpt-5.5" + now := time.Now().UTC() + futureExpiry := now.Add(24 * time.Hour).Truncate(time.Second) + + primary := &Auth{ + ID: "aa-codex-primary.json", + FileName: "aa-codex-primary.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "primary@example.com", + "access_token": "valid-access-token", + "refresh_token": "reused-refresh-token", + "expired": futureExpiry.Format(time.RFC3339), + }, + } + backup := &Auth{ + ID: "bb-codex-backup.json", + FileName: "bb-codex-backup.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "backup@example.com", + "access_token": "backup-access-token", + "refresh_token": "backup-refresh-token", + "expired": now.Add(48 * time.Hour).Format(time.RFC3339), + }, + } + + executor := &unauthorizedRefreshExecutor{ + id: "codex", + refreshErr: errors.New( + `token refresh failed with status 401: {"error":{"message":"Your refresh token has already been used to generate a new access token. Please try signing in again.","type":"invalid_request_error","code":"refresh_token_reused"}}`, + ), + } + + m := NewManager(nil, &FillFirstSelector{}, nil) + m.RegisterExecutor(executor) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(primary.ID, "codex", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(backup.ID, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(primary.ID) + reg.UnregisterClient(backup.ID) + }) + + if _, errReg := m.Register(ctx, primary); errReg != nil { + t.Fatalf("register primary: %v", errReg) + } + if _, errReg := m.Register(ctx, backup); errReg != nil { + t.Fatalf("register backup: %v", errReg) + } + + // 1. Proactive refresh fails with terminal refresh_token_reused while access_token is still valid. + m.refreshAuth(ctx, primary.ID) + if got := executor.RefreshCalls(); got != 1 { + t.Fatalf("RefreshCalls = %d, want 1", got) + } + + afterFirstFailure, ok := m.GetByID(primary.ID) + if !ok || afterFirstFailure == nil { + t.Fatal("primary auth missing") + } + if afterFirstFailure.Unavailable || afterFirstFailure.Status != StatusActive { + t.Fatalf("expected unexpired access_token to remain active, got status=%s unavailable=%v", afterFirstFailure.Status, afterFirstFailure.Unavailable) + } + if !afterFirstFailure.RequiresSignIn || afterFirstFailure.RenewalStateAt(now) != RenewalStateSignInRequired { + t.Fatalf("RequiresSignIn=%v RenewalState=%q, want true/%q", afterFirstFailure.RequiresSignIn, afterFirstFailure.RenewalStateAt(now), RenewalStateSignInRequired) + } + if afterFirstFailure.RefreshError == nil || afterFirstFailure.RefreshError.Code != "refresh_token_reused" { + t.Fatalf("RefreshError = %+v, want code refresh_token_reused", afterFirstFailure.RefreshError) + } + // NextRefreshAfter must be deferred to access_token expiry (~24h), not 5m retry storm. + if afterFirstFailure.NextRefreshAfter.Before(futureExpiry.Add(-time.Minute)) { + t.Fatalf("NextRefreshAfter = %v, want deferred until access_token expiry %v", afterFirstFailure.NextRefreshAfter, futureExpiry) + } + + // 2. Repeated unchanged refresh attempts before expiry must NOT call upstream Refresh again. + m.refreshAuth(ctx, primary.ID) + if got := executor.RefreshCalls(); got != 1 { + t.Fatalf("repeated refreshAuth with unchanged rejected grant called Refresh %d times, want 1", got) + } + + // 3. Ordinary inference on primary remains usable while access_token is unexpired, and does not clear RequiresSignIn/RefreshError. + resp, errExec := m.Execute(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errExec != nil { + t.Fatalf("Execute error = %v, want success on primary", errExec) + } + if got := string(resp.Payload); got != primary.ID+":valid-access-token" { + t.Fatalf("Execute payload = %q, want primary response", got) + } + afterExec, _ := m.GetByID(primary.ID) + if !afterExec.RequiresSignIn || afterExec.RefreshError == nil { + t.Fatalf("successful inference cleared renewal failure state: RequiresSignIn=%v RefreshError=%+v", afterExec.RequiresSignIn, afterExec.RefreshError) + } + + // 4. Transient refresh failure on backup follows bounded retry and does NOT require permanent re-authentication. + executor.mu.Lock() + executor.refreshErr = errors.New("token refresh failed with status 503: service unavailable") + executor.mu.Unlock() + m.refreshAuth(ctx, backup.ID) + afterTransient, _ := m.GetByID(backup.ID) + if afterTransient.RequiresSignIn || afterTransient.RenewalStateAt(now) != RenewalStateRetryScheduled { + t.Fatalf("transient failure set RequiresSignIn=%v RenewalState=%q, want false/%q", afterTransient.RequiresSignIn, afterTransient.RenewalStateAt(now), RenewalStateRetryScheduled) + } + if afterTransient.NextRefreshAfter.After(now.Add(10 * time.Minute)) { + t.Fatalf("transient NextRefreshAfter = %v, want ~5m backoff", afterTransient.NextRefreshAfter) + } + + // 5. While access_token is still valid on primary, a genuine 429 quota cooldown is recorded before access expiry. + retryAfter := 30 * time.Minute + m.MarkResult(ctx, Result{ + AuthID: primary.ID, + Provider: "codex", + Success: false, + CredentialScope: true, + RetryAfter: &retryAfter, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota exceeded"}, + }) + + // At access expiry, the unrenewable primary credential becomes unavailable without extra upstream refresh calls, and backup serves traffic. + m.mu.Lock() + if cur := m.auths[primary.ID]; cur != nil { + cur.Metadata["expired"] = now.Add(-time.Minute).Format(time.RFC3339) + m.scheduler.upsertAuth(cur.Clone()) + } + m.mu.Unlock() + callsBeforeExpiry := executor.RefreshCalls() + m.refreshAuth(ctx, primary.ID) + if got := executor.RefreshCalls(); got != callsBeforeExpiry { + t.Fatalf("expiry transition called upstream Refresh %d -> %d for already-rejected grant", callsBeforeExpiry, got) + } + afterExpiry, _ := m.GetByID(primary.ID) + if !afterExpiry.Unavailable || afterExpiry.Status != StatusError || !afterExpiry.RequiresSignIn || afterExpiry.RenewalStateAt(now) != RenewalStateExpiredSignInRequired { + t.Fatalf("expected unavailable expired_sign_in_required at expiry, got unavailable=%v status=%s requires_sign_in=%v state=%q", + afterExpiry.Unavailable, afterExpiry.Status, afterExpiry.RequiresSignIn, afterExpiry.RenewalStateAt(now)) + } + respAtExpiry, errAtExpiry := m.Execute(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errAtExpiry != nil { + t.Fatalf("Execute at primary expiry error = %v, want success on backup", errAtExpiry) + } + if got := string(respAtExpiry.Payload); got != backup.ID+":backup-access-token" { + t.Fatalf("Execute at primary expiry payload = %q, want backup", got) + } + + // 6. Fresh sign-in clears superseded auth failures while preserving active quota cooldown, and distinguishes awaiting_first_renewal from renewal_verified. + relogged := afterExpiry.Clone() + relogged.Metadata["access_token"] = "relogged-access-token" + relogged.Metadata["refresh_token"] = "relogged-refresh-token" + relogged.Metadata["expired"] = now.Add(48 * time.Hour).Format(time.RFC3339) + if _, errUpdate := m.Update(ctx, relogged); errUpdate != nil { + t.Fatalf("Update relogged primary: %v", errUpdate) + } + afterRelogin, _ := m.GetByID(primary.ID) + if afterRelogin.RequiresSignIn || afterRelogin.RefreshError != nil { + t.Fatalf("expected RequiresSignIn/RefreshError cleared on relogin, got %v / %+v", afterRelogin.RequiresSignIn, afterRelogin.RefreshError) + } + if afterRelogin.RenewalVerified || afterRelogin.RenewalStateAt(now) != RenewalStateAwaitingFirstRenewal { + t.Fatalf("expected fresh grant to be awaiting_first_renewal, got RenewalVerified=%v state=%q", afterRelogin.RenewalVerified, afterRelogin.RenewalStateAt(now)) + } + if !afterRelogin.Unavailable || afterRelogin.RecoveryStageAt(now) != RecoveryStageQuotaBlocked { + t.Fatalf("expected active quota cooldown preserved after relogin, got unavailable=%v stage=%q", afterRelogin.Unavailable, afterRelogin.RecoveryStageAt(now)) + } + + // Release quota cooldown and verify streaming completion vs incomplete stream. + _, _, _ = m.ResetQuota(ctx, primary.ID) + afterReset, _ := m.GetByID(primary.ID) + if afterReset.InferenceVerified || afterReset.RecoveryStageAt(now) != RecoveryStageLoaded { + t.Fatalf("expected stage loaded before inference verification, got Verified=%v stage=%q", afterReset.InferenceVerified, afterReset.RecoveryStageAt(now)) + } + + // Another account's completed response does not verify primary. + m.MarkResult(ctx, Result{AuthID: backup.ID, Provider: "codex", Model: model, Success: true, CompletedResponse: true}) + afterOtherSuccess, _ := m.GetByID(primary.ID) + if afterOtherSuccess.InferenceVerified { + t.Fatal("another account's inference success must not verify primary") + } + + // Local ExecuteCount (CountTokens) on primary must NOT verify inference. + if _, errCount := m.ExecuteCount(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}); errCount != nil { + t.Fatalf("ExecuteCount error: %v", errCount) + } + afterCountTokens, _ := m.GetByID(primary.ID) + if afterCountTokens.InferenceVerified { + t.Fatal("CountTokens success must not mark InferenceVerified=true") + } + + // An HTTP 200 stream that terminates without a terminal completion event (chunk error mid-stream) must NOT verify inference. + executor.mu.Lock() + executor.streamChunkErr = &Error{HTTPStatus: http.StatusBadGateway, Message: "incomplete stream without response.completed", Retryable: true} + executor.mu.Unlock() + incompleteStream, errIncomplete := m.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errIncomplete != nil { + t.Fatalf("ExecuteStream bootstrap error: %v", errIncomplete) + } + for range incompleteStream.Chunks { + } + afterIncomplete, _ := m.GetByID(primary.ID) + if afterIncomplete.InferenceVerified { + t.Fatal("HTTP 200 stream without terminal completion must not mark InferenceVerified=true") + } + _, _, _ = m.ResetQuota(ctx, primary.ID) + + // A stream that closes cleanly at EOF after response.created without a completion event must NOT verify inference. + executor.mu.Lock() + executor.streamChunkErr = nil + executor.streamPayloads = [][]byte{[]byte(`data: {"type":"response.created"}` + "\n\n")} + executor.mu.Unlock() + eofStream, errEOFStream := m.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errEOFStream != nil { + t.Fatalf("ExecuteStream (EOF without completed) error: %v", errEOFStream) + } + for range eofStream.Chunks { + } + afterEOFStream, _ := m.GetByID(primary.ID) + if afterEOFStream.InferenceVerified { + t.Fatal("clean EOF stream without terminal completion event must not mark InferenceVerified=true") + } + + // A stream where delta text contains "response.completed" or finish_reason is null must NOT verify inference. + for _, falsePositiveChunk := range [][]byte{ + []byte(`data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}` + "\n\n"), + []byte(`data: {"type":"response.output_text.delta","delta":"response.completed [DONE] message_stop"}` + "\n\n"), + } { + if comp, incomp := streamPayloadCompletionFlags(falsePositiveChunk); comp || incomp { + t.Fatalf("false-positive chunk %s returned completed=%v incomplete=%v, want false/false", string(falsePositiveChunk), comp, incomp) + } + } + if isCompletedExecutionPayload([]byte(`{"object":"list","data":[{"id":"gpt-5"}]}`)) { + t.Fatal("non-generation model list JSON must not be treated as a completed execution payload") + } + + // A stream that finishes with response.incomplete must NOT verify inference. + executor.mu.Lock() + executor.streamPayloads = [][]byte{[]byte(`data: {"type":"response.incomplete","response":{"status":"incomplete"}}` + "\n\n")} + executor.mu.Unlock() + respIncompleteStream, errRespIncomplete := m.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errRespIncomplete != nil { + t.Fatalf("ExecuteStream (response.incomplete) error: %v", errRespIncomplete) + } + for range respIncompleteStream.Chunks { + } + afterRespIncomplete, _ := m.GetByID(primary.ID) + if afterRespIncomplete.InferenceVerified { + t.Fatal("response.incomplete stream must not mark InferenceVerified=true") + } + if _, errVerifyIncomplete := m.VerifyInference(ctx, primary.ID, model); errVerifyIncomplete == nil { + 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. + executor.mu.Lock() + executor.streamPayloads = nil + executor.streamCalls = nil + executor.mu.Unlock() + afterVerify, errVerify := m.VerifyInference(ctx, primary.ID, model) + if errVerify != nil { + t.Fatalf("VerifyInference error: %v", errVerify) + } + executor.mu.Lock() + verifyCalls := append([]string(nil), executor.streamCalls...) + verifyReq := executor.lastStreamReq + verifyOpts := executor.lastStreamOpts + executor.mu.Unlock() + if len(verifyCalls) != 1 || verifyCalls[0] != primary.ID { + t.Fatalf("VerifyInference streamCalls = %v, want [%s]", verifyCalls, primary.ID) + } + 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 !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)) + } + if afterVerify.RenewalVerified { + t.Fatal("RenewalVerified must remain false until first automatic refresh succeeds") + } + + // Subsequent refresh cycle verifies renewal. + executor.mu.Lock() + executor.refreshErr = nil + executor.refreshTokens = map[string]string{primary.ID: "rotated-access-token"} + executor.mu.Unlock() + if _, errForce := m.ForceRefreshAuth(ctx, primary.ID); errForce != nil { + t.Fatalf("ForceRefreshAuth error: %v", errForce) + } + afterRenewal, _ := m.GetByID(primary.ID) + if !afterRenewal.RenewalVerified || afterRenewal.RenewalStateAt(now) != RenewalStateRenewalVerified || afterRenewal.RecoveryStageAt(now) != RecoveryStageRenewalVerified { + t.Fatalf("expected RenewalVerified=true after successful refresh, got verified=%v renewal_state=%q stage=%q", + afterRenewal.RenewalVerified, afterRenewal.RenewalStateAt(now), afterRenewal.RecoveryStageAt(now)) + } +} + +func TestManager_StaleRefreshFailureDoesNotPoisonReloggedCredential_Issue5783(t *testing.T) { + ctx := WithSkipPersist(context.Background()) + now := time.Now().UTC() + futureExpiry := now.Add(24 * time.Hour).Truncate(time.Second) + + auth := &Auth{ + ID: "codex-stale-race.json", + FileName: "codex-stale-race.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "race@example.com", + "access_token": "old-access-token", + "refresh_token": "old-refresh-token", + "expired": futureExpiry.Format(time.RFC3339), + }, + } + + refreshStarted := make(chan struct{}) + releaseRefresh := make(chan struct{}) + var once sync.Once + executor := &unauthorizedRefreshExecutor{ + id: "codex", + onRefresh: func() { + once.Do(func() { + close(refreshStarted) + <-releaseRefresh + }) + }, + refreshErr: errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"Refresh token already used"}}`), + } + + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + if _, errReg := m.Register(ctx, auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + refreshDone := make(chan struct{}) + go func() { + defer close(refreshDone) + m.refreshAuth(ctx, auth.ID) + }() + + <-refreshStarted + + // While the old refresh is in flight, operator completes re-login with a new grant. + relogged := auth.Clone() + relogged.Metadata["access_token"] = "new-access-token" + relogged.Metadata["refresh_token"] = "new-refresh-token" + relogged.Metadata["expired"] = now.Add(48 * time.Hour).Format(time.RFC3339) + if _, errUpdate := m.Update(ctx, relogged); errUpdate != nil { + t.Fatalf("Update relogged auth: %v", errUpdate) + } + + // Release the stale in-flight refresh failure. + close(releaseRefresh) + <-refreshDone + + after, ok := m.GetByID(auth.ID) + if !ok || after == nil { + t.Fatal("auth missing after stale refresh") + } + if after.RequiresSignIn || after.RefreshError != nil || after.RejectedRefreshToken != "" { + t.Fatalf("stale refresh failure poisoned newly relogged credential: RequiresSignIn=%v RefreshError=%+v RejectedRefreshToken=%q", + after.RequiresSignIn, after.RefreshError, after.RejectedRefreshToken) + } + if after.Unavailable || after.Status != StatusActive { + t.Fatalf("expected active credential after stale refresh discard, got status=%s unavailable=%v", after.Status, after.Unavailable) + } + + // Also verify that if a credential is replaced while VerifyInference is in flight, VerifyInference returns a superseded error and does NOT mark the new credential verified. + verifyStarted := make(chan struct{}) + releaseVerify := make(chan struct{}) + var verifyOnce sync.Once + executor.mu.Lock() + executor.onExecute = func(string) { + verifyOnce.Do(func() { + close(verifyStarted) + <-releaseVerify + }) + } + executor.mu.Unlock() + + verifyErrCh := make(chan error, 1) + go func() { + _, errV := m.VerifyInference(ctx, auth.ID, "gpt-5.5") + verifyErrCh <- errV + }() + + <-verifyStarted + reloggedAgain := after.Clone() + reloggedAgain.Metadata["access_token"] = "newer-access-token" + reloggedAgain.Metadata["refresh_token"] = "newer-refresh-token" + if _, errUpdate2 := m.Update(ctx, reloggedAgain); errUpdate2 != nil { + t.Fatalf("Update reloggedAgain: %v", errUpdate2) + } + close(releaseVerify) + + errVerifyRace := <-verifyErrCh + if errVerifyRace == nil || !containsSubstr(errVerifyRace.Error(), "replaced during verification") { + t.Fatalf("expected credential_superseded error when credential is replaced during VerifyInference, got: %v", errVerifyRace) + } + afterVerifyRace, _ := m.GetByID(auth.ID) + if afterVerifyRace.InferenceVerified { + t.Fatal("new credential version must not inherit InferenceVerified from superseded in-flight VerifyInference") + } + + // Also verify that if Store.Save blocks during a non-secret Update while a terminal refresh failure is recorded concurrently, + // the terminal RequiresSignIn and RejectedRefreshToken state is preserved when Store.Save completes. + saveStarted := make(chan struct{}) + releaseSave := make(chan struct{}) + var saveOnce sync.Once + blockingStore := &blockingTestAuthStore{ + onSave: func() { + saveOnce.Do(func() { + close(saveStarted) + <-releaseSave + }) + }, + } + mWithStore := NewManager(blockingStore, nil, nil) + termExec := &unauthorizedRefreshExecutor{ + id: "codex", + refreshErr: errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"reused"}}`), + } + mWithStore.RegisterExecutor(termExec) + persistAuth := &Auth{ + ID: "codex-persist-race.json", + FileName: "codex-persist-race.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "persist@example.com", + "access_token": "persist-at", + "refresh_token": "persist-rt", + "expired": futureExpiry.Format(time.RFC3339), + }, + } + if _, errRegPersist := mWithStore.Register(ctx, persistAuth); errRegPersist != nil { + t.Fatalf("Register persistAuth: %v", errRegPersist) + } + + updateDone := make(chan error, 1) + go func() { + nonSecretEdit := persistAuth.Clone() + nonSecretEdit.Metadata["note"] = "updated-note" + _, errU := mWithStore.Update(context.Background(), nonSecretEdit) + updateDone <- errU + }() + + <-saveStarted + mWithStore.refreshAuth(ctx, persistAuth.ID) + close(releaseSave) + if errU := <-updateDone; errU != nil { + t.Fatalf("concurrent non-secret Update error: %v", errU) + } + + afterPersistRace, _ := mWithStore.GetByID(persistAuth.ID) + if !afterPersistRace.RequiresSignIn || afterPersistRace.RefreshError == nil || afterPersistRace.RejectedRefreshToken != "persist-rt" { + t.Fatalf("concurrent Store.Save overwrote terminal renewal failure: RequiresSignIn=%v RefreshError=%+v Rejected=%q", + afterPersistRace.RequiresSignIn, afterPersistRace.RefreshError, afterPersistRace.RejectedRefreshToken) + } +} + +type blockingTestAuthStore struct { + onSave func() +} + +func (s *blockingTestAuthStore) List(context.Context) ([]*Auth, error) { return nil, nil } +func (s *blockingTestAuthStore) Save(_ context.Context, auth *Auth) (string, error) { + if s.onSave != nil { + s.onSave() + } + return auth.ID, nil +} +func (s *blockingTestAuthStore) Delete(context.Context, string) error { return nil } + +func TestManager_AccessTokenOnlyUpdateAndDiskAdoption_Issue5783(t *testing.T) { + ctx := WithSkipPersist(context.Background()) + now := time.Now().UTC() + futureExpiry := now.Add(12 * time.Hour).Truncate(time.Second) + + auth := &Auth{ + ID: "codex-partial.json", + FileName: "codex-partial.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "partial@example.com", + "access_token": "at-1", + "refresh_token": "dead-rt", + "expired": futureExpiry.Format(time.RFC3339), + }, + } + + executor := &unauthorizedRefreshExecutor{ + id: "codex", + refreshErr: errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"reused"}}`), + } + m := NewManager(nil, nil, nil) + m.RegisterExecutor(executor) + if _, errReg := m.Register(ctx, auth); errReg != nil { + t.Fatalf("register auth: %v", errReg) + } + + m.refreshAuth(ctx, auth.ID) + broken, _ := m.GetByID(auth.ID) + if !broken.RequiresSignIn { + t.Fatal("expected RequiresSignIn=true after terminal refresh error") + } + // Mark old access_token as inference-verified before replacing only access_token. + m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: "codex", Model: "gpt-5.5", Success: true, CompletedResponse: true}) + verifiedOld, _ := m.GetByID(auth.ID) + if !verifiedOld.InferenceVerified { + t.Fatal("expected InferenceVerified=true on old access_token") + } + + // 1. Updating only access_token while keeping the same rejected refresh_token must keep RequiresSignIn=true AND reset InferenceVerified=false. + newExp := now.Add(36 * time.Hour).Truncate(time.Second) + accessOnly := verifiedOld.Clone() + accessOnly.Metadata["access_token"] = "at-2" + accessOnly.Metadata["refresh_token"] = "dead-rt" + accessOnly.Metadata["expired"] = newExp.Format(time.RFC3339) + if _, errUpdate := m.Update(ctx, accessOnly); errUpdate != nil { + t.Fatalf("Update accessOnly: %v", errUpdate) + } + afterAccessOnly, _ := m.GetByID(auth.ID) + if !afterAccessOnly.RequiresSignIn || afterAccessOnly.RefreshError == nil || afterAccessOnly.RejectedRefreshToken != "dead-rt" { + t.Fatalf("updating only access_token cleared rejected refresh state: RequiresSignIn=%v RefreshError=%+v Rejected=%q", + afterAccessOnly.RequiresSignIn, afterAccessOnly.RefreshError, afterAccessOnly.RejectedRefreshToken) + } + if afterAccessOnly.InferenceVerified { + t.Fatal("replacing access_token must reset InferenceVerified=false so new access_token does not inherit old verification") + } + if afterAccessOnly.NextRefreshAfter.Before(newExp.Add(-time.Minute)) { + t.Fatalf("NextRefreshAfter = %v, want deferred to new access_token expiry %v", afterAccessOnly.NextRefreshAfter, newExp) + } + + // 2. A Refresh that only adopts disk-rotated credentials (UpstreamRefreshed=false) clears RequiresSignIn but keeps RenewalVerified=false. + diskPatch := codex.NewCredentialRefreshPatch( + afterAccessOnly.Metadata, + map[string]any{"access_token": "disk-at", "refresh_token": "disk-rt", "expired": newExp.Format(time.RFC3339)}, + true, + false, + ) + diskPatch.SetUpstreamRefreshed(false) + executor.mu.Lock() + executor.refreshErr = nil + executor.refreshTokens = map[string]string{auth.ID: "disk-at"} + executor.refreshPatch = diskPatch + executor.mu.Unlock() + + if _, errForce := m.ForceRefreshAuth(ctx, auth.ID); errForce != nil { + t.Fatalf("ForceRefreshAuth (disk adoption) error: %v", errForce) + } + afterDiskAdopt, _ := m.GetByID(auth.ID) + if afterDiskAdopt.RequiresSignIn || afterDiskAdopt.RefreshError != nil { + t.Fatalf("disk adoption should clear RequiresSignIn, got %v / %+v", afterDiskAdopt.RequiresSignIn, afterDiskAdopt.RefreshError) + } + if afterDiskAdopt.RenewalVerified || afterDiskAdopt.RenewalStateAt(now) != RenewalStateAwaitingFirstRenewal { + t.Fatalf("disk-only adoption must NOT mark RenewalVerified=true, got verified=%v state=%q", + afterDiskAdopt.RenewalVerified, afterDiskAdopt.RenewalStateAt(now)) + } + + // 3. When a credential has both outer refresh_token and a healthy siwc.refresh_token, an outer terminal refresh failure + // must NOT poison the healthy siwc.refresh_token or prevent recovery when only the outer refresh_token is re-authenticated. + siwcAuth := &Auth{ + ID: "codex-siwc-outer.json", + FileName: "codex-siwc-outer.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "siwc@example.com", + "access_token": "outer-at-1", + "refresh_token": "dead-outer-rt", + "expired": futureExpiry.Format(time.RFC3339), + "siwc": map[string]any{ + "access_token": "siwc-at-1", + "refresh_token": "healthy-siwc-rt", + }, + }, + } + if _, errRegSIWC := m.Register(ctx, siwcAuth); errRegSIWC != nil { + t.Fatalf("register siwcAuth: %v", errRegSIWC) + } + executor.mu.Lock() + executor.refreshPatch = nil + executor.refreshErr = errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused"},"code":"authorization-secret-xyz","state":"callback-secret-123"} callback%3Fcode%3Descaped-code-secret%26state%3Descaped-state-secret`) + executor.mu.Unlock() + m.refreshAuth(ctx, siwcAuth.ID) + siwcBroken, _ := m.GetByID(siwcAuth.ID) + if !siwcBroken.RequiresSignIn || siwcBroken.RejectedRefreshToken != "dead-outer-rt" || siwcBroken.RejectedSIWCRefreshToken != "" { + t.Fatalf("outer refresh failure should only mark outer refresh_token rejected, got RequiresSignIn=%v outer=%q siwc=%q", + siwcBroken.RequiresSignIn, siwcBroken.RejectedRefreshToken, siwcBroken.RejectedSIWCRefreshToken) + } + for _, secret := range []string{"authorization-secret-xyz", "callback-secret-123", "escaped-code-secret", "escaped-state-secret"} { + if containsSubstr(siwcBroken.RefreshError.Message, secret) { + t.Fatalf("RefreshError leaked secret %q: %s", secret, siwcBroken.RefreshError.Message) + } + } + authBearerSanitized := SanitizeRenewalMessage(`Authorization: Bearer opaque-secret-plain {"authorization": "Bearer opaque-secret-json"} AUTHORIZATION: BEARER opaque-secret-upper`) + for _, secret := range []string{"opaque-secret-plain", "opaque-secret-json", "opaque-secret-upper"} { + if containsSubstr(authBearerSanitized, secret) { + t.Fatalf("SanitizeRenewalMessage leaked Authorization Bearer secret %q: %s", secret, authBearerSanitized) + } + } + if !containsSubstr(siwcBroken.RefreshError.Message, "refresh_token_reused") { + t.Fatalf("RefreshError should preserve safe error code refresh_token_reused, got: %s", siwcBroken.RefreshError.Message) + } + + reloggedOuter := siwcBroken.Clone() + reloggedOuter.Metadata["access_token"] = "outer-at-2" + reloggedOuter.Metadata["refresh_token"] = "fresh-outer-rt" + reloggedOuter.Metadata["siwc"] = map[string]any{ + "access_token": "siwc-at-1", + "refresh_token": "healthy-siwc-rt", + } + if _, errUpdateOuter := m.Update(ctx, reloggedOuter); errUpdateOuter != nil { + t.Fatalf("Update reloggedOuter: %v", errUpdateOuter) + } + afterOuterRelogin, _ := m.GetByID(siwcAuth.ID) + if afterOuterRelogin.RequiresSignIn || afterOuterRelogin.RefreshError != nil { + t.Fatalf("re-authenticating outer refresh_token while healthy SIWC token is unchanged must clear RequiresSignIn, got %v / %+v", + afterOuterRelogin.RequiresSignIn, afterOuterRelogin.RefreshError) + } + + // 4. Conversely, a real SIWC HTTP 401 refresh error ("OpenAI OAuth token request failed (401): ...") must only mark + // RejectedSIWCRefreshToken and allow recovery when only the SIWC refresh_token is re-bound. + executor.mu.Lock() + executor.refreshErr = codex.WrapSIWCRefreshError(errors.New(`OpenAI OAuth token request failed (401): {"error":{"code":"refresh_token_reused","message":"SIWC refresh token already used"}}`)) + executor.mu.Unlock() + m.refreshAuth(ctx, siwcAuth.ID) + siwcOnlyBroken, _ := m.GetByID(siwcAuth.ID) + if !siwcOnlyBroken.RequiresSignIn || siwcOnlyBroken.RejectedSIWCRefreshToken != "healthy-siwc-rt" || siwcOnlyBroken.RejectedRefreshToken != "" { + t.Fatalf("SIWC refresh failure should only mark SIWC refresh_token rejected, got RequiresSignIn=%v outer=%q siwc=%q", + siwcOnlyBroken.RequiresSignIn, siwcOnlyBroken.RejectedRefreshToken, siwcOnlyBroken.RejectedSIWCRefreshToken) + } + reboundSIWC := siwcOnlyBroken.Clone() + reboundSIWC.Metadata["siwc"] = map[string]any{ + "access_token": "siwc-at-2", + "refresh_token": "fresh-siwc-rt", + "client_id": "cid-siwc", + "expired": time.Now().UTC().Add(45 * time.Minute).Format(time.RFC3339), + } + if _, errUpdateSIWC := m.Update(ctx, reboundSIWC); errUpdateSIWC != nil { + t.Fatalf("Update reboundSIWC: %v", errUpdateSIWC) + } + afterSIWCRebind, _ := m.GetByID(siwcAuth.ID) + if afterSIWCRebind.RequiresSignIn || afterSIWCRebind.RefreshError != nil { + t.Fatalf("re-binding SIWC refresh_token while healthy outer token is unchanged must clear RequiresSignIn, got %v / %+v", + afterSIWCRebind.RequiresSignIn, afterSIWCRebind.RefreshError) + } + + // 4b. A partial refresh where SIWC succeeds but outer refresh fails terminally (recorded on CodexRefreshPatch) + // must NOT clear RequiresSignIn or mark RenewalVerified=true. + partialPatch := codex.NewCredentialRefreshPatch( + afterSIWCRebind.Metadata, + afterSIWCRebind.Metadata, + false, + true, + ) + partialPatch.SetUpstreamRefreshed(true) + partialPatch.SetPartitionErrors( + errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"outer reused"}}`), + nil, + ) + executor.mu.Lock() + executor.refreshErr = nil + executor.refreshPatch = partialPatch + executor.mu.Unlock() + if _, errPartial := m.ForceRefreshAuth(ctx, siwcAuth.ID); errPartial != nil { + t.Fatalf("ForceRefreshAuth (partial refresh) error: %v", errPartial) + } + afterPartial, _ := m.GetByID(siwcAuth.ID) + if !afterPartial.RequiresSignIn || afterPartial.RenewalVerified || afterPartial.RejectedRefreshToken != "fresh-outer-rt" { + t.Fatalf("partial refresh with terminal outer failure must set RequiresSignIn=true RenewalVerified=false RejectedRefreshToken=fresh-outer-rt, got RequiresSignIn=%v RenewalVerified=%v Rejected=%q", + afterPartial.RequiresSignIn, afterPartial.RenewalVerified, afterPartial.RejectedRefreshToken) + } + if !afterPartial.NextRefreshAfter.IsZero() { + t.Fatalf("partial refresh with healthy SIWC partition must not defer NextRefreshAfter to outer expiry without Update, got %v", + afterPartial.NextRefreshAfter) + } + executor.mu.Lock() + executor.refreshPatch = nil + // A subsequent transient 503 failure on the healthy SIWC partition must NOT clear RequiresSignIn + // or overwrite the terminal outer RefreshError. + executor.refreshErr = errors.New(`OpenAI OAuth token request failed (503): service unavailable`) + executor.mu.Unlock() + m.refreshAuth(ctx, siwcAuth.ID) + afterTransientOnHealthy, _ := m.GetByID(siwcAuth.ID) + if !afterTransientOnHealthy.RequiresSignIn || afterTransientOnHealthy.RejectedRefreshToken != "fresh-outer-rt" { + t.Fatalf("transient failure on healthy SIWC partition must preserve RequiresSignIn=true and RejectedRefreshToken=fresh-outer-rt, got RequiresSignIn=%v Rejected=%q", + afterTransientOnHealthy.RequiresSignIn, afterTransientOnHealthy.RejectedRefreshToken) + } + if afterTransientOnHealthy.RefreshError == nil || !isTerminalRefreshError(afterTransientOnHealthy, afterTransientOnHealthy.RefreshError) { + t.Fatalf("transient failure on healthy SIWC partition must preserve terminal RefreshError from rejected outer partition, got %+v", + afterTransientOnHealthy.RefreshError) + } + executor.mu.Lock() + executor.refreshErr = nil + executor.mu.Unlock() + + // 4c. Bidirectional automatic refresh (m.shouldRefresh + m.refreshAuth) when one partition is rejected + // and the other partition is healthy and due for refresh: + // Direction 1: outer is rejected ("fresh-outer-rt"), while SIWC is healthy ("fresh-siwc-rt") and due for refresh. + nowCheck := time.Now().UTC() + siwcDueClone := afterPartial.Clone() + siwcDueClone.Metadata["siwc"] = map[string]any{ + "access_token": "siwc-at-due", + "refresh_token": "fresh-siwc-rt", + "client_id": "cid-siwc", + "expired": nowCheck.Add(5 * time.Minute).Format(time.RFC3339), + } + if _, errUpdSIWCDue := m.Update(ctx, siwcDueClone); errUpdSIWCDue != nil { + t.Fatalf("Update siwcDueClone: %v", errUpdSIWCDue) + } + currentSIWCDue, _ := m.GetByID(siwcAuth.ID) + if !m.shouldRefresh(currentSIWCDue, nowCheck) { + t.Fatal("shouldRefresh must be true when SIWC partition is healthy and due despite rejected outer partition") + } + callsBeforeAutoSIWC := executor.RefreshCalls() + m.refreshAuth(ctx, siwcAuth.ID) + if executor.RefreshCalls() != callsBeforeAutoSIWC+1 { + t.Fatalf("m.refreshAuth must invoke executor for healthy SIWC partition, calls=%d want %d", executor.RefreshCalls(), callsBeforeAutoSIWC+1) + } + + // Direction 2: SIWC is rejected, while outer is healthy and due for refresh. + outerHealthyClone := currentSIWCDue.Clone() + outerHealthyClone.Metadata["refresh_token"] = "brand-new-outer-rt" + outerHealthyClone.Metadata["expired"] = nowCheck.Add(10 * time.Minute).Format(time.RFC3339) + outerHealthyClone.Metadata["siwc"] = map[string]any{ + "access_token": "siwc-at-valid", + "refresh_token": "dead-siwc-rt", + "client_id": "cid-siwc", + "expired": nowCheck.Add(45 * time.Minute).Format(time.RFC3339), + } + if _, errUpdOuter := m.Update(ctx, outerHealthyClone); errUpdOuter != nil { + t.Fatalf("Update outerHealthyClone: %v", errUpdOuter) + } + executor.mu.Lock() + executor.refreshErr = errors.New(`OpenAI OAuth token request failed (401): {"error":{"code":"refresh_token_reused","message":"SIWC reused"}}`) + executor.mu.Unlock() + m.refreshAuth(ctx, siwcAuth.ID) + afterSIWCDead, _ := m.GetByID(siwcAuth.ID) + if afterSIWCDead.RejectedSIWCRefreshToken != "dead-siwc-rt" || afterSIWCDead.RejectedRefreshToken != "" { + t.Fatalf("expected only RejectedSIWCRefreshToken=dead-siwc-rt, got outer=%q siwc=%q", + afterSIWCDead.RejectedRefreshToken, afterSIWCDead.RejectedSIWCRefreshToken) + } + if !m.shouldRefresh(afterSIWCDead, nowCheck) { + t.Fatal("shouldRefresh must be true when outer partition is healthy and due despite rejected SIWC partition") + } + executor.mu.Lock() + executor.refreshErr = nil + executor.mu.Unlock() + callsBeforeAutoOuter := executor.RefreshCalls() + m.refreshAuth(ctx, siwcAuth.ID) + if executor.RefreshCalls() != callsBeforeAutoOuter+1 { + t.Fatalf("m.refreshAuth must invoke executor for healthy outer partition, calls=%d want %d", executor.RefreshCalls(), callsBeforeAutoOuter+1) + } + + // 5. Disable -> watcher file reload -> re-enable transition must preserve terminal renewal state + // when the rejected refresh token is unchanged. + disableAuth := &Auth{ + ID: "codex-disable-reload.json", + FileName: "codex-disable-reload.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "disable-reload@example.com", + "access_token": "valid-at", + "refresh_token": "dead-rt- across-disable", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + }, + } + if _, errRegDis := m.Register(ctx, disableAuth); errRegDis != nil { + t.Fatalf("Register(disableAuth) error: %v", errRegDis) + } + executor.mu.Lock() + executor.refreshErr = errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"reused"}}`) + executor.mu.Unlock() + m.refreshAuth(ctx, disableAuth.ID) + + // Simulate watcher loading a disabled snapshot from disk (runtime fields are zeroed on fresh disk load). + disabledFromDisk := &Auth{ + ID: disableAuth.ID, + FileName: disableAuth.FileName, + Provider: "codex", + Disabled: true, + Status: StatusDisabled, + Metadata: map[string]any{ + "type": "codex", + "email": "disable-reload@example.com", + "access_token": "valid-at", + "refresh_token": "dead-rt- across-disable", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + }, + } + if _, errUpdDis := m.Update(ctx, disabledFromDisk); errUpdDis != nil { + t.Fatalf("Update(disabledFromDisk) error: %v", errUpdDis) + } + // Simulate watcher loading the re-enabled snapshot from disk with the same unchanged refresh token. + reenabledFromDisk := &Auth{ + ID: disableAuth.ID, + FileName: disableAuth.FileName, + Provider: "codex", + Disabled: false, + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "disable-reload@example.com", + "access_token": "valid-at", + "refresh_token": "dead-rt- across-disable", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + }, + } + if _, errUpdRe := m.Update(ctx, reenabledFromDisk); errUpdRe != nil { + t.Fatalf("Update(reenabledFromDisk) error: %v", errUpdRe) + } + afterReenable, _ := m.GetByID(disableAuth.ID) + if !afterReenable.RequiresSignIn || afterReenable.RejectedRefreshToken != "dead-rt- across-disable" || afterReenable.RefreshError == nil { + t.Fatalf("disable/re-enable with unchanged refresh_token must preserve RequiresSignIn and RejectedRefreshToken, got RequiresSignIn=%v Rejected=%q Err=%+v", + afterReenable.RequiresSignIn, afterReenable.RejectedRefreshToken, afterReenable.RefreshError) + } + callsBeforeReenableRefresh := executor.RefreshCalls() + m.refreshAuth(ctx, disableAuth.ID) + if executor.RefreshCalls() != callsBeforeReenableRefresh { + t.Fatalf("m.refreshAuth after re-enabling unchanged rejected token must short-circuit, calls=%d want %d", + executor.RefreshCalls(), callsBeforeReenableRefresh) + } + + // 6. Dual-partition refresh failure where outer fails with transient 503 and SIWC fails with terminal 401 refresh_token_reused + // must still record SIWC's terminal rejection and RequiresSignIn=true. + dualFailAuth := &Auth{ + ID: "codex-dual-fail.json", + FileName: "codex-dual-fail.json", + Provider: "codex", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "dualfail@example.com", + "access_token": "outer-at-valid", + "refresh_token": "outer-rt-transient", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + "use_siwc": true, + "siwc": map[string]any{ + "access_token": "siwc-at-valid", + "refresh_token": "siwc-rt-terminal", + "client_id": "cid-dual", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + }, + }, + } + if _, errRegDual := m.Register(ctx, dualFailAuth); errRegDual != nil { + t.Fatalf("Register(dualFailAuth) error: %v", errRegDual) + } + executor.mu.Lock() + executor.refreshErr = codex.NewPartitionRefreshError( + errors.New(`token refresh failed with status 503: service unavailable`), + codex.WrapSIWCRefreshError(errors.New(`OpenAI OAuth token request failed (401): {"error":{"code":"refresh_token_reused","message":"SIWC reused"}}`)), + ) + executor.mu.Unlock() + m.refreshAuth(ctx, dualFailAuth.ID) + afterDualFail, _ := m.GetByID(dualFailAuth.ID) + if !afterDualFail.RequiresSignIn || afterDualFail.RejectedSIWCRefreshToken != "siwc-rt-terminal" || afterDualFail.RejectedRefreshToken != "" { + t.Fatalf("dual-partition failure (outer 503 + SIWC 401) must set RequiresSignIn=true and RejectedSIWCRefreshToken=siwc-rt-terminal, got RequiresSignIn=%v outer=%q siwc=%q", + afterDualFail.RequiresSignIn, afterDualFail.RejectedRefreshToken, afterDualFail.RejectedSIWCRefreshToken) + } + if afterDualFail.RefreshError == nil || afterDualFail.RefreshError.Code != "refresh_token_reused" { + t.Fatalf("dual-partition failure must preserve terminal SIWC RefreshError code refresh_token_reused, got %+v", afterDualFail.RefreshError) + } + executor.mu.Lock() + executor.refreshErr = nil + executor.mu.Unlock() + + // 7. VerifyInference must resolve model prefix and OAuth model aliases before calling executor. + m.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + "codex": {{Name: "gpt-5.5", Alias: "codex-verified"}}, + }) + aliasedAuth := &Auth{ + ID: "codex-alias-verify.json", + FileName: "codex-alias-verify.json", + Provider: "codex", + Prefix: "team", + Status: StatusActive, + Metadata: map[string]any{ + "type": "codex", + "email": "alias@example.com", + "access_token": "alias-at", + "refresh_token": "alias-rt", + "expired": nowCheck.Add(40 * time.Minute).Format(time.RFC3339), + }, + } + if _, errRegAlias := m.Register(ctx, aliasedAuth); errRegAlias != nil { + t.Fatalf("Register(aliasedAuth) error: %v", errRegAlias) + } + executor.mu.Lock() + executor.streamPayloads = [][]byte{[]byte(`data: {"type":"response.completed","response":{"status":"completed"}}` + "\n\n")} + executor.mu.Unlock() + verifiedAliasAuth, errVerifyAlias := m.VerifyInference(ctx, aliasedAuth.ID, "team/codex-verified") + if errVerifyAlias != nil { + t.Fatalf("VerifyInference with prefix and alias failed: %v", errVerifyAlias) + } + if !verifiedAliasAuth.InferenceVerified { + t.Fatal("expected InferenceVerified=true after VerifyInference with prefix and alias") + } + executor.mu.Lock() + gotModel := executor.lastStreamReq.Model + executor.mu.Unlock() + if gotModel != "gpt-5.5" { + t.Fatalf("VerifyInference must resolve prefix+alias to upstream model gpt-5.5, got %q", gotModel) + } + + // 8. Concurrent relogin of the failed partition before a partial refresh commits must NOT poison + // the newly relogged partition in either direction (outer failed + SIWC refreshed, or SIWC failed + outer refreshed). + // Direction A: outer failed on old token while SIWC refreshed; outer relogged before UpdateRefreshedAuth commits. + raceBaseA, _ := m.GetByID(siwcAuth.ID) + reloggedOuterA := raceBaseA.Clone() + reloggedOuterA.Metadata["access_token"] = "relogged-outer-at-A" + reloggedOuterA.Metadata["refresh_token"] = "relogged-outer-rt-A" + if _, errUpdRaceA := m.Update(ctx, reloggedOuterA); errUpdRaceA != nil { + t.Fatalf("Update reloggedOuterA: %v", errUpdRaceA) + } + stalePartialOuterFail := raceBaseA.Clone() + stalePartialOuterFail.Metadata["siwc"] = map[string]any{ + "access_token": "siwc-rotated-at-A", + "refresh_token": "siwc-rotated-rt-A", + "client_id": "cid-siwc", + "expired": nowCheck.Add(50 * time.Minute).Format(time.RFC3339), + } + patchA := codex.NewCredentialRefreshPatch(raceBaseA.Metadata, stalePartialOuterFail.Metadata, false, true) + patchA.SetUpstreamRefreshed(true) + patchA.SetPartitionErrors( + errors.New(`token refresh failed with status 401: {"error":{"code":"refresh_token_reused","message":"old outer reused"}}`), + nil, + ) + stalePartialOuterFail.CodexRefreshPatch = patchA + mergedA, errMergeA := m.UpdateRefreshedAuth(ctx, raceBaseA, stalePartialOuterFail) + if errMergeA != nil { + t.Fatalf("UpdateRefreshedAuth (Direction A) error: %v", errMergeA) + } + if mergedA.RequiresSignIn || mergedA.RejectedRefreshToken != "" || !mergedA.RenewalVerified { + t.Fatalf("stale outer failure in partial refresh must not poison relogged outer token: RequiresSignIn=%v RejectedOuter=%q RenewalVerified=%v", + mergedA.RequiresSignIn, mergedA.RejectedRefreshToken, mergedA.RenewalVerified) + } + + // Direction B: SIWC failed on old token while outer refreshed; SIWC rebound before UpdateRefreshedAuth commits. + raceBaseB, _ := m.GetByID(siwcAuth.ID) + reboundSIWCB := raceBaseB.Clone() + reboundSIWCB.Metadata["siwc"] = map[string]any{ + "access_token": "rebound-siwc-at-B", + "refresh_token": "rebound-siwc-rt-B", + "client_id": "cid-siwc", + "expired": nowCheck.Add(50 * time.Minute).Format(time.RFC3339), + } + if _, errUpdRaceB := m.Update(ctx, reboundSIWCB); errUpdRaceB != nil { + t.Fatalf("Update reboundSIWCB: %v", errUpdRaceB) + } + stalePartialSIWCFail := raceBaseB.Clone() + stalePartialSIWCFail.Metadata["access_token"] = "outer-rotated-at-B" + stalePartialSIWCFail.Metadata["refresh_token"] = "outer-rotated-rt-B" + patchB := codex.NewCredentialRefreshPatch(raceBaseB.Metadata, stalePartialSIWCFail.Metadata, true, false) + patchB.SetUpstreamRefreshed(true) + patchB.SetPartitionErrors( + nil, + codex.WrapSIWCRefreshError(errors.New(`OpenAI OAuth token request failed (401): {"error":{"code":"refresh_token_reused","message":"old SIWC reused"}}`)), + ) + stalePartialSIWCFail.CodexRefreshPatch = patchB + mergedB, errMergeB := m.UpdateRefreshedAuth(ctx, raceBaseB, stalePartialSIWCFail) + if errMergeB != nil { + t.Fatalf("UpdateRefreshedAuth (Direction B) error: %v", errMergeB) + } + if mergedB.RequiresSignIn || mergedB.RejectedSIWCRefreshToken != "" || !mergedB.RenewalVerified { + t.Fatalf("stale SIWC failure in partial refresh must not poison rebound SIWC token: RequiresSignIn=%v RejectedSIWC=%q RenewalVerified=%v", + mergedB.RequiresSignIn, mergedB.RejectedSIWCRefreshToken, mergedB.RenewalVerified) + } + + // 5. Secrets straddling the ~300-character truncation boundary must be redacted before truncation. + padding := make([]byte, 285) + for i := range padding { + padding[i] = 'x' + } + boundaryRaw := string(padding) + ` {"code":"boundary-secret-code-1234567890","state":"boundary-secret-state-1234567890"}` + sanitizedBoundary := SanitizeRenewalMessage(boundaryRaw) + if containsSubstr(sanitizedBoundary, "boundary-secret-code") || containsSubstr(sanitizedBoundary, "boundary-secret-state") { + t.Fatalf("SanitizeRenewalMessage leaked truncation-boundary secret: %s", sanitizedBoundary) + } +} + +func containsSubstr(s, sub string) bool { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/sdk/cliproxy/auth/cooldown_state.go b/sdk/cliproxy/auth/cooldown_state.go index 7c0797103e1..28c9c91f048 100644 --- a/sdk/cliproxy/auth/cooldown_state.go +++ b/sdk/cliproxy/auth/cooldown_state.go @@ -283,7 +283,7 @@ func (s *FileCooldownStateStore) statePath(record CooldownStateRecord) (string, func (s *FileCooldownStateStore) stateRelativePath(record CooldownStateRecord) string { authFile := strings.TrimSpace(record.AuthFile) if authFile != "" { - if filepath.IsAbs(authFile) && strings.TrimSpace(s.authDir) != "" { + if strings.TrimSpace(s.authDir) != "" { if rel, errRel := filepath.Rel(s.authDir, authFile); errRel == nil && rel != "." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) && rel != ".." { return cdsPathForRel(rel) } diff --git a/sdk/cliproxy/auth/cooldown_state_test.go b/sdk/cliproxy/auth/cooldown_state_test.go index 0506740cf28..ec7dd305471 100644 --- a/sdk/cliproxy/auth/cooldown_state_test.go +++ b/sdk/cliproxy/auth/cooldown_state_test.go @@ -111,6 +111,25 @@ func TestFileCooldownStateStore_StateRelativePath(t *testing.T) { } } +// With a relative auth-dir the file synthesizer records an auth file as "/", so the state path must strip +// the auth dir there too; otherwise the store, rooted at the auth dir, writes "//.cds". +func TestFileCooldownStateStore_StateRelativePathRelativeAuthDir(t *testing.T) { + store := NewFileCooldownStateStoreWithAuthDir("auths", "auths") + cases := []struct { + authFile string + want string + }{ + {filepath.Join("auths", "claude.json"), "claude.cds"}, + {filepath.Join("auths", "team", "xai.json"), filepath.Join("team", "xai.cds")}, + {filepath.Join("other", "xai.json"), filepath.Join("other", "xai.cds")}, + } + for _, tc := range cases { + if got := store.stateRelativePath(CooldownStateRecord{AuthID: "id", AuthFile: tc.authFile}); got != tc.want { + t.Errorf("stateRelativePath(%q) = %q, want %q", tc.authFile, got, tc.want) + } + } +} + func TestFileCooldownStateStore_SaveLoadAndCleanStale(t *testing.T) { authDir := t.TempDir() store := NewFileCooldownStateStoreWithAuthDir(authDir, authDir) diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index ba307249853..48dd5a69395 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -113,6 +113,31 @@ type Auth struct { RejectedAccessToken string `json:"-"` // RejectedSIWCAccessToken tracks the nested SIWC access token rejected by upstream 401 until replaced. RejectedSIWCAccessToken string `json:"-"` + // RejectedRefreshToken tracks the outer refresh token rejected as terminal until replaced. + RejectedRefreshToken string `json:"-"` + // RejectedSIWCRefreshToken tracks the nested SIWC refresh token rejected as terminal until replaced. + RejectedSIWCRefreshToken string `json:"-"` + // RenewalWarnedRefreshToken tracks the rejected refresh token that already emitted a terminal warning/event. + RenewalWarnedRefreshToken string `json:"-"` + // RequiresSignIn flags that automatic credential renewal failed with a terminal grant rejection. + RequiresSignIn bool `json:"-"` + // RefreshError stores the sanitized last refresh failure without being overwritten by request execution results. + RefreshError *Error `json:"-"` + // RenewalEventAt records when the deduplicated terminal renewal failure event was first observed for the current grant. + RenewalEventAt time.Time `json:"-"` + // AdoptedAt records when the current credential version was adopted into runtime state. + AdoptedAt time.Time `json:"-"` + // AdoptedGrantHash records the fingerprint of the credential grant when it was adopted into runtime state, + // preserved across upstream automatic refresh rotations and reset when the credential grant is externally replaced. + AdoptedGrantHash string `json:"-"` + // InferenceVerified reports whether the current credential version completed a real inference response. + InferenceVerified bool `json:"-"` + // LastInferenceVerifiedAt records when the current credential version last completed an inference response. + LastInferenceVerifiedAt time.Time `json:"-"` + // RenewalVerified reports whether the current credential version completed an automatic/proactive refresh cycle. + RenewalVerified bool `json:"-"` + // LastRenewalVerifiedAt records when the current credential version last completed a token refresh. + LastRenewalVerifiedAt time.Time `json:"-"` Success int64 `json:"-"` Failed int64 `json:"-"` @@ -451,6 +476,14 @@ func (a *Auth) Clone() *Auth { } copyAuth := *a copyAuth.Quota = a.Quota.Clone() + if a.LastError != nil { + copyErr := *a.LastError + copyAuth.LastError = ©Err + } + if a.RefreshError != nil { + copyErr := *a.RefreshError + copyAuth.RefreshError = ©Err + } if a.Attributes != nil { copyAuth.Attributes = make(map[string]string, len(a.Attributes)) for key, value := range a.Attributes { @@ -503,6 +536,40 @@ func stableAuthIndex(seed string) string { return hex.EncodeToString(sum[:8]) } +// OAuthGrantFingerprint computes a SHA-256 fingerprint of an OAuth grant's tokens and account identity. +func OAuthGrantFingerprint(accessToken, refreshToken, email, accountID string) string { + accessToken = strings.TrimSpace(accessToken) + refreshToken = strings.TrimSpace(refreshToken) + email = strings.ToLower(strings.TrimSpace(email)) + accountID = strings.ToLower(strings.TrimSpace(accountID)) + if accessToken == "" && refreshToken == "" { + return "" + } + sum := sha256.Sum256([]byte(accessToken + "\x00" + refreshToken + "\x00" + email + "\x00" + accountID)) + return hex.EncodeToString(sum[:]) +} + +// AuthGrantFingerprint computes a SHA-256 fingerprint of the credential's token grant and account identity. +func AuthGrantFingerprint(a *Auth) string { + if a == nil { + return "" + } + accessToken := authOuterAccessToken(a) + if accessToken == "" { + accessToken = authAccessToken(a) + } + refreshToken := authRefreshToken(a) + email := authMetadataString(a, "email") + if email == "" && a.Attributes != nil { + email = strings.TrimSpace(a.Attributes["email"]) + } + accountID := authMetadataString(a, "account_id") + if accountID == "" && a.Attributes != nil { + accountID = strings.TrimSpace(a.Attributes["account_id"]) + } + return OAuthGrantFingerprint(accessToken, refreshToken, email, accountID) +} + func (a *Auth) indexSeed() string { if a == nil { return "" @@ -1062,6 +1129,19 @@ func CodexOuterNeedsRefresh(a *Auth, now time.Time) (dueAt time.Time, hasRefresh } hasRefresh = true + if a.RejectedRefreshToken != "" && authRefreshToken(a) == a.RejectedRefreshToken { + if expiry, hasExpiry := a.OuterExpirationTime(); hasExpiry && !expiry.IsZero() { + if expiry.After(now) { + return expiry, true, false + } + if IsCodexUseSIWC(a) { + return time.Time{}, false, false + } + return now, true, true + } + return time.Time{}, false, false + } + outerToken := authOuterAccessToken(a) if outerToken == "" || (a.RejectedAccessToken != "" && a.RejectedAccessToken == outerToken) { return now, true, true @@ -1150,6 +1230,19 @@ func CodexSIWCNeedsRefresh(a *Auth, now time.Time) (dueAt time.Time, hasRefresh } hasRefresh = true + if a.RejectedSIWCRefreshToken != "" && strings.TrimSpace(refreshToken) == a.RejectedSIWCRefreshToken { + if expiry, hasExpiry := a.SIWCExpirationTime(); hasExpiry && !expiry.IsZero() { + if expiry.After(now) { + return expiry, true, false + } + if !IsCodexUseSIWC(a) { + return time.Time{}, false, false + } + return now, true, true + } + return time.Time{}, false, false + } + siwcToken, _ := siwcMap["access_token"].(string) siwcToken = strings.TrimSpace(siwcToken) if siwcToken == "" { @@ -1215,6 +1308,93 @@ func (a *Auth) HasValidAccessToken(now time.Time) bool { return true } +const ( + RenewalStateSignInRequired = "sign_in_required" + RenewalStateExpiredSignInRequired = "expired_sign_in_required" + RenewalStateRetryScheduled = "retry_scheduled" + RenewalStateRenewalVerified = "renewal_verified" + RenewalStateAwaitingFirstRenewal = "awaiting_first_renewal" + + RecoveryStageSignInRequired = "sign_in_required" + RecoveryStageCredentialSaved = "credential_saved" + RecoveryStageQuotaBlocked = "quota_blocked" + RecoveryStageRenewalVerified = "renewal_verified" + RecoveryStageInferenceVerified = "inference_verified" + RecoveryStageLoaded = "loaded" +) + +// RenewalStateAt returns the structured renewal lifecycle state at the given time. +func (a *Auth) RenewalStateAt(now time.Time) string { + if a == nil { + return "" + } + if now.IsZero() { + now = time.Now() + } + if a.RequiresSignIn { + if !a.HasValidAccessToken(now) || hasUnauthorizedAuthFailure(a) { + return RenewalStateExpiredSignInRequired + } + return RenewalStateSignInRequired + } + if a.RefreshError != nil { + return RenewalStateRetryScheduled + } + if a.RenewalVerified { + return RenewalStateRenewalVerified + } + return RenewalStateAwaitingFirstRenewal +} + +// RecoveryStageAt returns the structured post-login / runtime adoption and verification stage at the given time. +func (a *Auth) RecoveryStageAt(now time.Time) string { + if a == nil { + return "" + } + if now.IsZero() { + now = time.Now() + } + if a.RequiresSignIn { + return RecoveryStageSignInRequired + } + if a.hasActiveQuotaBlock(now) { + return RecoveryStageQuotaBlocked + } + if a.RenewalVerified { + return RecoveryStageRenewalVerified + } + if a.InferenceVerified { + return RecoveryStageInferenceVerified + } + return RecoveryStageLoaded +} + +func (a *Auth) hasActiveQuotaBlock(now time.Time) bool { + if a == nil { + return false + } + if a.Quota.Exceeded && !a.Quota.NextRecoverAt.IsZero() && a.Quota.NextRecoverAt.After(now) { + return true + } + if a.Unavailable && !a.NextRetryAfter.IsZero() && a.NextRetryAfter.After(now) && + a.LastError != nil && a.LastError.HTTPStatus == http.StatusTooManyRequests { + return true + } + for _, ms := range a.ModelStates { + if ms == nil { + continue + } + if ms.Quota.Exceeded && !ms.Quota.NextRecoverAt.IsZero() && ms.Quota.NextRecoverAt.After(now) { + return true + } + if ms.Unavailable && !ms.NextRetryAfter.IsZero() && ms.NextRetryAfter.After(now) && + ms.LastError != nil && ms.LastError.HTTPStatus == http.StatusTooManyRequests { + return true + } + } + return false +} + var ( refreshLeadMu sync.RWMutex refreshLeadFactories = make(map[string]func() *time.Duration) diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 41ba9a4a5d3..a1032861342 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -102,6 +102,9 @@ type Request struct { Metadata map[string]any } +// RequestBeforeAuthInterceptor rewrites a follow-up request on a bound session before after-auth interception. +type RequestBeforeAuthInterceptor func(context.Context, RequestAfterAuthInterceptRequest) RequestAfterAuthInterceptResponse + // RequestAfterAuthInterceptor rewrites a request after credential selection and before executor translation. type RequestAfterAuthInterceptor func(context.Context, RequestAfterAuthInterceptRequest) RequestAfterAuthInterceptResponse @@ -218,6 +221,8 @@ type Options struct { ResponseFormat sdktranslator.Format // Metadata carries extra execution hints shared across selection and executors. Metadata map[string]any + // RequestBeforeAuthInterceptor runs before after-auth interception on bound follow-up turns. + RequestBeforeAuthInterceptor RequestBeforeAuthInterceptor // RequestAfterAuthInterceptor runs after credential selection and before executor translation. RequestAfterAuthInterceptor RequestAfterAuthInterceptor // WebSocketResponseObserver receives upstream WebSocket response events during execution. @@ -252,6 +257,9 @@ func ResponseFormatOrSource(opts Options) sdktranslator.Format { type Response struct { // Payload is the provider response in the executor format. Payload []byte + // Incomplete reports that the upstream execution terminated in an incomplete state + // (e.g., response.incomplete) even if translated into a normal non-stream payload. + Incomplete bool // Metadata exposes optional structured data for translators. Metadata map[string]any // Headers carries upstream HTTP response headers for passthrough to clients. @@ -262,6 +270,9 @@ type Response struct { type StreamChunk struct { // Payload is the raw provider chunk payload. Payload []byte + // Incomplete reports that the upstream stream ended in an incomplete terminal state + // (e.g., response.incomplete) even if protocol translation mapped it to a normal chunk. + Incomplete bool // Err reports any terminal error encountered while producing chunks. Err error } diff --git a/sdk/cliproxy/executor/websocket_input.go b/sdk/cliproxy/executor/websocket_input.go index fe648df80fc..31e2a255af4 100644 --- a/sdk/cliproxy/executor/websocket_input.go +++ b/sdk/cliproxy/executor/websocket_input.go @@ -8,6 +8,13 @@ import ( // ErrNoActiveUpstreamWebsocket means response.interrupt has no live upstream socket. var ErrNoActiveUpstreamWebsocket = errors.New("no active upstream websocket for response.interrupt") +// ErrUpstreamWebsocketInterruptRejected means a live upstream socket owns the +// current turn, but response.interrupt does not target its running response +// (unknown id, or that response already finished). Handlers must reject the +// frame outright and must not fall back to local cancellation, which would +// cancel a healthy websocket turn. +var ErrUpstreamWebsocketInterruptRejected = errors.New("response.interrupt does not target the active upstream response") + // WebsocketInput is a frame from the single downstream reader. Err terminates // the connection; Payload is owned by the receiver and must not be replayed. type WebsocketInput struct { diff --git a/sdk/cliproxy/plugin_static_model_exclusions_test.go b/sdk/cliproxy/plugin_static_model_exclusions_test.go index 08509dba2a8..522ecf033c5 100644 --- a/sdk/cliproxy/plugin_static_model_exclusions_test.go +++ b/sdk/cliproxy/plugin_static_model_exclusions_test.go @@ -1,14 +1,20 @@ package cliproxy import ( + "bytes" "context" + "net/http" + "net/http/httptest" + "sync/atomic" "testing" "github.com/router-for-me/CLIProxyAPI/v8/internal/pluginhost" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" ) const ( @@ -205,3 +211,160 @@ func TestAntigravityModelsForHints_PluginStaticModelsFollowExclusions(t *testing }) } } + +type multiProviderModelRegistrar struct { + registrations []pluginapi.ProviderModels +} + +func (r multiProviderModelRegistrar) RegisterModels(context.Context, pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) { + return pluginapi.ModelRegistrationResponse{ + Registrations: r.registrations, + }, nil +} + +func TestPluginMultiProviderRegistrarPublishesRoutesAndUnloadsBothCompatProviders(t *testing.T) { + var directHits atomic.Int32 + directServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + directHits.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"chatcmpl-direct","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"direct"},"finish_reason":"stop"}]}`)) + })) + defer directServer.Close() + + var vpnHits atomic.Int32 + vpnServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + vpnHits.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"chatcmpl-vpn","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"vpn"},"finish_reason":"stop"}]}`)) + })) + defer vpnServer.Close() + + const ( + compatDirect = "multi-direct" + compatVPN = "multi-vpn" + providerDirect = "openai-compatible-multi-direct" + providerVPN = "openai-compatible-multi-vpn" + modelDirect = "multi-model-direct" + modelVPN = "multi-model-vpn" + ) + + cfg := &config.Config{ + OpenAICompatibility: []config.OpenAICompatibility{ + { + Name: compatDirect, + BaseURL: directServer.URL, + APIKeyEntries: []config.OpenAICompatibilityAPIKey{{APIKey: "direct-key"}}, + }, + { + Name: compatVPN, + BaseURL: vpnServer.URL, + APIKeyEntries: []config.OpenAICompatibilityAPIKey{{APIKey: "vpn-key"}}, + }, + }, + } + + host := pluginhost.New() + t.Cleanup(host.ShutdownAll) + pluginID := "multi-compat-catalog" + host.RegisterPluginForTest(pluginID, pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + ModelRegistrar: multiProviderModelRegistrar{ + registrations: []pluginapi.ProviderModels{ + {Provider: providerDirect, Models: []pluginapi.ModelInfo{{ID: modelDirect}}}, + {Provider: providerVPN, Models: []pluginapi.ModelInfo{{ID: modelVPN}}}, + }, + }, + }, + }) + + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + service := &Service{ + cfg: cfg, + coreManager: manager, + pluginHost: host, + } + + authDirect := &coreauth.Auth{ + ID: t.Name() + "/direct", + Provider: providerDirect, + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "auth_kind": "apikey", + "compat_name": compatDirect, + "provider_key": providerDirect, + "base_url": directServer.URL, + "api_key": "direct-key", + }, + } + authVPN := &coreauth.Auth{ + ID: t.Name() + "/vpn", + Provider: providerVPN, + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "auth_kind": "apikey", + "compat_name": compatVPN, + "provider_key": providerVPN, + "base_url": vpnServer.URL, + "api_key": "vpn-key", + }, + } + if _, errReg := manager.Register(t.Context(), authDirect); errReg != nil { + t.Fatalf("Register(authDirect): %v", errReg) + } + if _, errReg := manager.Register(t.Context(), authVPN); errReg != nil { + t.Fatalf("Register(authVPN): %v", errReg) + } + reg := GlobalModelRegistry() + t.Cleanup(func() { + reg.UnregisterClient(authDirect.ID) + reg.UnregisterClient(authVPN.ID) + }) + + service.syncPluginModelRuntime(t.Context()) + + available := openAIModelIDSet(reg.GetAvailableModels("openai")) + if _, okDirect := available[modelDirect]; !okDirect { + t.Fatalf("/v1/models missing %s; available = %v", modelDirect, available) + } + if _, okVPN := available[modelVPN]; !okVPN { + t.Fatalf("/v1/models missing %s; available = %v", modelVPN, available) + } + + reqOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + } + respDirect, errExec := manager.Execute(t.Context(), registry.GetGlobalRegistry().GetModelProviders(modelDirect), cliproxyexecutor.Request{ + Model: modelDirect, + Payload: []byte(`{"model":"` + modelDirect + `","messages":[{"role":"user","content":"hi"}]}`), + }, reqOpts) + if errExec != nil { + t.Fatalf("Execute(%s) error = %v", modelDirect, errExec) + } + if !bytes.Contains(respDirect.Payload, []byte("chatcmpl-direct")) || directHits.Load() != 1 || vpnHits.Load() != 0 { + t.Fatalf("Execute(%s) payload=%q hits=(direct=%d, vpn=%d), want direct upstream only", modelDirect, string(respDirect.Payload), directHits.Load(), vpnHits.Load()) + } + respVPN, errExec := manager.Execute(t.Context(), registry.GetGlobalRegistry().GetModelProviders(modelVPN), cliproxyexecutor.Request{ + Model: modelVPN, + Payload: []byte(`{"model":"` + modelVPN + `","messages":[{"role":"user","content":"hi"}]}`), + }, reqOpts) + if errExec != nil { + t.Fatalf("Execute(%s) error = %v", modelVPN, errExec) + } + if !bytes.Contains(respVPN.Payload, []byte("chatcmpl-vpn")) || directHits.Load() != 1 || vpnHits.Load() != 1 { + t.Fatalf("Execute(%s) payload=%q hits=(direct=%d, vpn=%d), want vpn upstream", modelVPN, string(respVPN.Payload), directHits.Load(), vpnHits.Load()) + } + + if !host.UnloadPlugin(pluginID) { + t.Fatalf("UnloadPlugin(%s) = false, want true", pluginID) + } + service.syncPluginModelRuntime(t.Context()) + + availableAfterUnload := openAIModelIDSet(reg.GetAvailableModels("openai")) + if _, okDirect := availableAfterUnload[modelDirect]; okDirect { + t.Fatalf("%s still listed after plugin unload", modelDirect) + } + if _, okVPN := availableAfterUnload[modelVPN]; okVPN { + t.Fatalf("%s still listed after plugin unload", modelVPN) + } +} diff --git a/sdk/cliproxy/service_auth.go b/sdk/cliproxy/service_auth.go index 92e2693a86a..acc0d28a0b0 100644 --- a/sdk/cliproxy/service_auth.go +++ b/sdk/cliproxy/service_auth.go @@ -439,8 +439,10 @@ func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth } } if !auth.HasFileSnapshot() && !existing.Disabled && existing.Status != coreauth.StatusDisabled && !auth.Disabled && auth.Status != coreauth.StatusDisabled { - auth.LastRefreshedAt = existing.LastRefreshedAt - auth.NextRefreshAfter = existing.NextRefreshAfter + if !coreauth.CredentialsChanged(existing, auth) { + auth.LastRefreshedAt = existing.LastRefreshedAt + auth.NextRefreshAfter = existing.NextRefreshAfter + } if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { auth.ModelStates = existing.ModelStates } diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index 890af119e08..bc6214f4c14 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -344,7 +344,7 @@ func forceHomeRuntimeConfig(cfg *config.Config) { cfg.UsageStatisticsEnabled = true cfg.DisableCooling = true cfg.SaveCooldownStatus = false - cfg.WebsocketAuth = false + cfg.WebsocketAuth = true cfg.RemoteManagement.AllowRemote = false cfg.RemoteManagement.DisableControlPanel = true cfg.Plugins.StoreAuth = nil diff --git a/sdk/cliproxy/service_stale_state_test.go b/sdk/cliproxy/service_stale_state_test.go index 0a61d295609..8577d2abc71 100644 --- a/sdk/cliproxy/service_stale_state_test.go +++ b/sdk/cliproxy/service_stale_state_test.go @@ -78,6 +78,7 @@ func TestForceHomeRuntimeConfigEnablesUsageStatistics(t *testing.T) { UsageStatisticsEnabled: false, DisableCooling: false, SaveCooldownStatus: true, + WebsocketAuth: false, } forceHomeRuntimeConfig(cfg) @@ -91,6 +92,9 @@ func TestForceHomeRuntimeConfigEnablesUsageStatistics(t *testing.T) { if cfg.SaveCooldownStatus { t.Fatal("expected home runtime config to force cooldown status persistence disabled") } + if !cfg.WebsocketAuth { + t.Fatal("expected home runtime config to force websocket authentication enabled") + } } func TestLifetimeRegistryObservesBarrierFromAppliedHomeConfig(t *testing.T) { diff --git a/sdk/cliproxy/usage/detail.go b/sdk/cliproxy/usage/detail.go new file mode 100644 index 00000000000..ebc8228e6bc --- /dev/null +++ b/sdk/cliproxy/usage/detail.go @@ -0,0 +1,274 @@ +package usage + +import "maps" + +// CacheDetails holds prompt cache configuration and TTL-bucketed write token counts. +type CacheDetails struct { + Mode string `json:"mode,omitempty"` + TTL string `json:"ttl,omitempty"` + Retention string `json:"retention,omitempty"` + Ephemeral5mInputTokens *int64 `json:"ephemeral_5m_input_tokens,omitempty"` + Ephemeral30mInputTokens *int64 `json:"ephemeral_30m_input_tokens,omitempty"` + Ephemeral1hInputTokens *int64 `json:"ephemeral_1h_input_tokens,omitempty"` + Ephemeral12hInputTokens *int64 `json:"ephemeral_12h_input_tokens,omitempty"` + ByTTL map[string]int64 `json:"by_ttl,omitempty"` +} + +// ToolUsageDetails holds structured server-side tool execution counters. +type ToolUsageDetails struct { + WebSearchCalls *int64 `json:"web_search_calls,omitempty"` + WebSearchQueries *int64 `json:"web_search_queries,omitempty"` + WebFetchCalls *int64 `json:"web_fetch_calls,omitempty"` + NumSearchResults *int64 `json:"num_search_results,omitempty"` + NumSourcesUsed *int64 `json:"num_sources_used,omitempty"` + NumServerSideToolsUsed *int64 `json:"num_server_side_tools_used,omitempty"` + ToolCallsRequested *int64 `json:"tool_calls_requested,omitempty"` + ToolCallsExecuted *int64 `json:"tool_calls_executed,omitempty"` + XSearchCalls *int64 `json:"x_search_calls,omitempty"` + XPostsFetched *int64 `json:"x_posts_fetched,omitempty"` + XUsersFetched *int64 `json:"x_users_fetched,omitempty"` + CodeExecutionCalls *int64 `json:"code_execution_calls,omitempty"` + FileSearchCalls *int64 `json:"file_search_calls,omitempty"` + MCPCalls *int64 `json:"mcp_calls,omitempty"` + DocumentSearchCalls *int64 `json:"document_search_calls,omitempty"` + ImageGenerationCalls *int64 `json:"image_generation_calls,omitempty"` + Extra map[string]int64 `json:"extra,omitempty"` +} + +// IterationUsage holds per-iteration usage reported in usage.iterations[]. +type IterationUsage struct { + Type string `json:"type,omitempty"` + Model string `json:"model,omitempty"` + InputTokens *int64 `json:"input_tokens,omitempty"` + OutputTokens *int64 `json:"output_tokens,omitempty"` + ThinkingTokens *int64 `json:"thinking_tokens,omitempty"` + ReasoningTokens *int64 `json:"reasoning_tokens,omitempty"` + CacheReadInputTokens *int64 `json:"cache_read_input_tokens,omitempty"` + CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens,omitempty"` + CacheCreation5mTokens *int64 `json:"cache_creation_5m_tokens,omitempty"` + CacheCreation1hTokens *int64 `json:"cache_creation_1h_tokens,omitempty"` + CacheCreation *CacheDetails `json:"cache_creation,omitempty"` +} + +// ImageDetails holds image/video generation dimension and quality attributes. +type ImageDetails struct { + Size string `json:"size,omitempty"` + Resolution string `json:"resolution,omitempty"` + ImageSize string `json:"image_size,omitempty"` + AspectRatio string `json:"aspect_ratio,omitempty"` + Quality string `json:"quality,omitempty"` + Background string `json:"background,omitempty"` + OutputFormat string `json:"output_format,omitempty"` + ResponseFormat string `json:"response_format,omitempty"` + MimeType string `json:"mime_type,omitempty"` + InputFidelity string `json:"input_fidelity,omitempty"` +} + +// Clone returns a deep copy of ImageDetails. +func (d *ImageDetails) Clone() *ImageDetails { + if d == nil { + return nil + } + out := *d + return &out +} + +// MediaCounts holds prompt/completion media item counts and audio/video duration. +type MediaCounts struct { + PromptMedia *int64 `json:"prompt_media,omitempty"` + CompletionMedia *int64 `json:"completion_media,omitempty"` + PromptAudio *int64 `json:"prompt_audio,omitempty"` + AudioSeconds *float64 `json:"audio_seconds,omitempty"` +} + +// Clone returns a deep copy of MediaCounts. +func (m *MediaCounts) Clone() *MediaCounts { + if m == nil { + return nil + } + return &MediaCounts{ + PromptMedia: cloneInt64Ptr(m.PromptMedia), + CompletionMedia: cloneInt64Ptr(m.CompletionMedia), + PromptAudio: cloneInt64Ptr(m.PromptAudio), + AudioSeconds: cloneFloat64Ptr(m.AudioSeconds), + } +} + +// CostDetails holds optional upstream-reported cost information (e.g. OpenRouter). +type CostDetails struct { + Cost *float64 `json:"cost,omitempty"` + CostInUSDTicks *int64 `json:"cost_in_usd_ticks,omitempty"` + IsBYOK *bool `json:"is_byok,omitempty"` + UpstreamInferenceCost *float64 `json:"upstream_inference_cost,omitempty"` + UpstreamInferencePromptCost *float64 `json:"upstream_inference_prompt_cost,omitempty"` + UpstreamInferenceCompletionsCost *float64 `json:"upstream_inference_completions_cost,omitempty"` +} + +// Int64Ptr returns a pointer to v. +func Int64Ptr(v int64) *int64 { + return &v +} + +// Float64Ptr returns a pointer to v. +func Float64Ptr(v float64) *float64 { + return &v +} + +// BoolPtr returns a pointer to v. +func BoolPtr(v bool) *bool { + return &v +} + +func cloneInt64Ptr(p *int64) *int64 { + if p == nil { + return nil + } + v := *p + return &v +} + +func cloneFloat64Ptr(p *float64) *float64 { + if p == nil { + return nil + } + v := *p + return &v +} + +func cloneBoolPtr(p *bool) *bool { + if p == nil { + return nil + } + v := *p + return &v +} + +// CloneInt64Ptr returns a shallow copy of p when non-nil. +func CloneInt64Ptr(p *int64) *int64 { + return cloneInt64Ptr(p) +} + +// Clone returns a deep copy of CacheDetails. +func (c *CacheDetails) Clone() *CacheDetails { + if c == nil { + return nil + } + out := &CacheDetails{ + Mode: c.Mode, + TTL: c.TTL, + Retention: c.Retention, + Ephemeral5mInputTokens: cloneInt64Ptr(c.Ephemeral5mInputTokens), + Ephemeral30mInputTokens: cloneInt64Ptr(c.Ephemeral30mInputTokens), + Ephemeral1hInputTokens: cloneInt64Ptr(c.Ephemeral1hInputTokens), + Ephemeral12hInputTokens: cloneInt64Ptr(c.Ephemeral12hInputTokens), + } + if len(c.ByTTL) > 0 { + out.ByTTL = maps.Clone(c.ByTTL) + } + return out +} + +// SumTokens returns the sum of TTL-bucketed tokens and whether at least one TTL bucket was present. +func (c *CacheDetails) SumTokens() (int64, bool) { + if c == nil || len(c.ByTTL) == 0 { + return 0, false + } + var sum int64 + for _, v := range c.ByTTL { + if v < 0 || sum > int64(^uint64(0)>>1)-v { + return 0, false + } + sum += v + } + return sum, true +} + +// Clone returns a deep copy of ToolUsageDetails. +func (t *ToolUsageDetails) Clone() *ToolUsageDetails { + if t == nil { + return nil + } + out := &ToolUsageDetails{ + WebSearchCalls: cloneInt64Ptr(t.WebSearchCalls), + WebSearchQueries: cloneInt64Ptr(t.WebSearchQueries), + WebFetchCalls: cloneInt64Ptr(t.WebFetchCalls), + NumSearchResults: cloneInt64Ptr(t.NumSearchResults), + NumSourcesUsed: cloneInt64Ptr(t.NumSourcesUsed), + NumServerSideToolsUsed: cloneInt64Ptr(t.NumServerSideToolsUsed), + ToolCallsRequested: cloneInt64Ptr(t.ToolCallsRequested), + ToolCallsExecuted: cloneInt64Ptr(t.ToolCallsExecuted), + XSearchCalls: cloneInt64Ptr(t.XSearchCalls), + XPostsFetched: cloneInt64Ptr(t.XPostsFetched), + XUsersFetched: cloneInt64Ptr(t.XUsersFetched), + CodeExecutionCalls: cloneInt64Ptr(t.CodeExecutionCalls), + FileSearchCalls: cloneInt64Ptr(t.FileSearchCalls), + MCPCalls: cloneInt64Ptr(t.MCPCalls), + DocumentSearchCalls: cloneInt64Ptr(t.DocumentSearchCalls), + ImageGenerationCalls: cloneInt64Ptr(t.ImageGenerationCalls), + } + if len(t.Extra) > 0 { + out.Extra = maps.Clone(t.Extra) + } + return out +} + +// Clone returns a deep copy of IterationUsage. +func (it IterationUsage) Clone() IterationUsage { + return IterationUsage{ + Type: it.Type, + Model: it.Model, + InputTokens: cloneInt64Ptr(it.InputTokens), + OutputTokens: cloneInt64Ptr(it.OutputTokens), + ThinkingTokens: cloneInt64Ptr(it.ThinkingTokens), + ReasoningTokens: cloneInt64Ptr(it.ReasoningTokens), + CacheReadInputTokens: cloneInt64Ptr(it.CacheReadInputTokens), + CacheCreationInputTokens: cloneInt64Ptr(it.CacheCreationInputTokens), + CacheCreation5mTokens: cloneInt64Ptr(it.CacheCreation5mTokens), + CacheCreation1hTokens: cloneInt64Ptr(it.CacheCreation1hTokens), + CacheCreation: it.CacheCreation.Clone(), + } +} + +// CloneIterations returns a deep copy of a slice of IterationUsage. +func CloneIterations(items []IterationUsage) []IterationUsage { + if len(items) == 0 { + return nil + } + out := make([]IterationUsage, len(items)) + for i := range items { + out[i] = items[i].Clone() + } + return out +} + +// IterationsPtr returns a pointer to a cloned non-empty IterationUsage slice, or nil when empty. +func IterationsPtr(items []IterationUsage) *[]IterationUsage { + cloned := CloneIterations(items) + if len(cloned) == 0 { + return nil + } + return &cloned +} + +// IterationSlice returns a deep copy of Detail.Iterations when present. +func (d Detail) IterationSlice() []IterationUsage { + if d.Iterations == nil { + return nil + } + return CloneIterations(*d.Iterations) +} + +// Clone returns a deep copy of CostDetails. +func (c *CostDetails) Clone() *CostDetails { + if c == nil { + return nil + } + return &CostDetails{ + Cost: cloneFloat64Ptr(c.Cost), + CostInUSDTicks: cloneInt64Ptr(c.CostInUSDTicks), + IsBYOK: cloneBoolPtr(c.IsBYOK), + UpstreamInferenceCost: cloneFloat64Ptr(c.UpstreamInferenceCost), + UpstreamInferencePromptCost: cloneFloat64Ptr(c.UpstreamInferencePromptCost), + UpstreamInferenceCompletionsCost: cloneFloat64Ptr(c.UpstreamInferenceCompletionsCost), + } +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 470b3dadb73..686d78a46ae 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -94,7 +94,7 @@ type Failure struct { Code string } -// Detail holds the token usage breakdown. +// Detail holds the token and structured usage breakdown captured from an upstream response. type Detail struct { InputTokens int64 OutputTokens int64 @@ -105,6 +105,34 @@ type Detail struct { TotalTokens int64 TokenBreakdown TokenBreakdown ResponseServiceTier string + + CacheCreation5mTokens *int64 + CacheCreation30mTokens *int64 + CacheCreation1hTokens *int64 + CacheCreation12hTokens *int64 + TextInputTokens *int64 + ImageInputTokens *int64 + AudioInputTokens *int64 + VideoInputTokens *int64 + FileInputTokens *int64 + CachedTextTokens *int64 + TextOutputTokens *int64 + ImageOutputTokens *int64 + AudioOutputTokens *int64 + AcceptedPredictionTokens *int64 + RejectedPredictionTokens *int64 + CacheDetails *CacheDetails + WebSearchRequests *int64 + ImageCount *int64 + RequestedImageDetails *ImageDetails + ImageDetails *ImageDetails + MediaCounts *MediaCounts + ToolUsageDetails *ToolUsageDetails + Iterations *[]IterationUsage + CostDetails *CostDetails + Speed string + InferenceGeo string + UpstreamServiceTier string } type requestedModelAliasContextKey struct{} diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index 5ec6947949d..7be274ca9a9 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -18,11 +18,20 @@ const ( // Plugins that still require HTML entity escaping on JSON response strings // should keep schema_version < 6. // Version 7 lets request interceptors opt out of receiving the request body - // per phase via request_interceptor_omit_body_before/after capabilities; the - // host then omits RequestInterceptRequest.Body for that phase instead of - // cloning and Base64-encoding the full payload. Plugins that declare the - // flags on schema_version < 7 (or do not declare them) still receive it. - SchemaVersion uint32 = 7 + // per phase via request_interceptor_omit_body_before/after capabilities and + // lets frontend auth providers opt out via frontend_auth_provider_omit_body / + // frontend_auth_omit_body; the host then omits the request body instead of + // reading, cloning, and Base64-encoding the full payload. Plugins that declare + // the flags on schema_version < 7 (or do not declare them) still receive it. + // Version 8 lets request interceptors opt into raw request body delivery and + // localized byte-range patching per phase via request_interceptor_patch_body_before/after + // (or request_body_patcher_before/after) capabilities and the + // request.patch_body_before / request.patch_body_after methods; the host passes + // raw request body bytes directly across the C ABI without JSON/Base64 wrapping + // and splices validated byte-range edits in a single pass. + // Version 9 lets ModelRegistrar.RegisterModels and ModelProvider.StaticModels + // publish model catalogs for multiple providers in one response via Registrations. + SchemaVersion uint32 = 9 // SchemaVersionStreamChunkOmitRequestBody is the first schema version that omits // request bodies on payload stream-chunk interceptor calls. SchemaVersionStreamChunkOmitRequestBody uint32 = 3 @@ -38,6 +47,15 @@ const ( // SchemaVersionRequestInterceptorOmitBody is the first schema version where // request interceptors may opt out of receiving the request body per phase. SchemaVersionRequestInterceptorOmitBody uint32 = 7 + // SchemaVersionFrontendAuthOmitBody is the first schema version where + // frontend auth providers may opt out of receiving the request body. + SchemaVersionFrontendAuthOmitBody uint32 = 7 + // SchemaVersionRequestInterceptorBodyPatch is the first schema version where + // request interceptors may opt into raw request body delivery and byte-range patching. + SchemaVersionRequestInterceptorBodyPatch uint32 = 8 + // SchemaVersionMultiProviderModels is the first schema version where model + // registration responses may publish models for multiple providers via Registrations. + SchemaVersionMultiProviderModels uint32 = 9 ) const ( @@ -74,6 +92,8 @@ const ( MethodRequestNormalize = "request.normalize" MethodRequestInterceptBefore = "request.intercept_before" MethodRequestInterceptAfter = "request.intercept_after" + MethodRequestPatchBodyBefore = "request.patch_body_before" + MethodRequestPatchBodyAfter = "request.patch_body_after" MethodRequestComplete = "request.complete" MethodResponseTranslate = "response.translate" diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index 6954cd37c44..a9f892271da 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -28,8 +28,8 @@ func TestEnvelopeRoundTrip(t *testing.T) { } func TestMethodNamesAreStable(t *testing.T) { - if SchemaVersion != 7 { - t.Fatalf("SchemaVersion = %d, want 7", SchemaVersion) + if SchemaVersion != 9 { + t.Fatalf("SchemaVersion = %d, want 9", SchemaVersion) } if SchemaVersionWebSocketResponseObserver != 4 { t.Fatalf("SchemaVersionWebSocketResponseObserver = %d, want 4", SchemaVersionWebSocketResponseObserver) @@ -46,6 +46,15 @@ func TestMethodNamesAreStable(t *testing.T) { if SchemaVersionRequestInterceptorOmitBody != 7 { t.Fatalf("SchemaVersionRequestInterceptorOmitBody = %d, want 7", SchemaVersionRequestInterceptorOmitBody) } + if SchemaVersionFrontendAuthOmitBody != 7 { + t.Fatalf("SchemaVersionFrontendAuthOmitBody = %d, want 7", SchemaVersionFrontendAuthOmitBody) + } + if SchemaVersionRequestInterceptorBodyPatch != 8 { + t.Fatalf("SchemaVersionRequestInterceptorBodyPatch = %d, want 8", SchemaVersionRequestInterceptorBodyPatch) + } + if SchemaVersionMultiProviderModels != 9 { + t.Fatalf("SchemaVersionMultiProviderModels = %d, want 9", SchemaVersionMultiProviderModels) + } if MethodPluginRegister != "plugin.register" { t.Fatalf("MethodPluginRegister = %q", MethodPluginRegister) } @@ -58,6 +67,12 @@ func TestMethodNamesAreStable(t *testing.T) { if MethodRequestInterceptAfter != "request.intercept_after" { t.Fatalf("MethodRequestInterceptAfter = %q", MethodRequestInterceptAfter) } + if MethodRequestPatchBodyBefore != "request.patch_body_before" { + t.Fatalf("MethodRequestPatchBodyBefore = %q", MethodRequestPatchBodyBefore) + } + if MethodRequestPatchBodyAfter != "request.patch_body_after" { + t.Fatalf("MethodRequestPatchBodyAfter = %q", MethodRequestPatchBodyAfter) + } if MethodRequestComplete != "request.complete" { t.Fatalf("MethodRequestComplete = %q", MethodRequestComplete) } diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 714b6fac194..e7c0da24887 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -82,6 +82,13 @@ type Capabilities struct { FrontendAuthProvider FrontendAuthProvider // FrontendAuthProviderExclusive makes this frontend auth provider the only active request auth provider when selected. FrontendAuthProviderExclusive bool + // FrontendAuthProviderOmitBody opts out of receiving FrontendAuthRequest.Body and + // FrontendAuthRequest.DecodedBody in FrontendAuthProvider.Authenticate. It is honored + // only on schema version 7+; the host then skips reading, decoding, and cloning the + // request body during authentication. + FrontendAuthProviderOmitBody bool + // FrontendAuthOmitBody is an alias for FrontendAuthProviderOmitBody. + FrontendAuthOmitBody bool // Scheduler chooses an auth candidate before the built-in scheduler runs. Scheduler Scheduler // SchedulerAcrossPriorities opts into receiving available candidates across all priority tiers @@ -120,6 +127,20 @@ type Capabilities struct { // in InterceptRequestAfterAuth. It is honored only on schema version 7+; the plugin // may still replace the body blindly via RequestInterceptResponse.Body. RequestInterceptorOmitBodyAfter bool + // RequestBodyPatcherBefore scans raw request bodies in the before-auth phase and + // returns localized byte-range edits via request.patch_body_before (schema v8+). + RequestBodyPatcherBefore RequestBodyPatcher + // RequestBodyPatcherAfter scans raw request bodies in the after-auth phase and + // returns localized byte-range edits via request.patch_body_after (schema v8+). + RequestBodyPatcherAfter RequestBodyPatcher + // RequestInterceptorPatchBodyBefore opts into raw request body delivery and byte-range + // patching in the before-auth phase via request.patch_body_before. It is honored only + // on schema version 8+. + RequestInterceptorPatchBodyBefore bool + // RequestInterceptorPatchBodyAfter opts into raw request body delivery and byte-range + // patching in the after-auth phase via request.patch_body_after. It is honored only + // on schema version 8+. + RequestInterceptorPatchBodyAfter bool // RequestLifecyclePlugin asynchronously receives one terminal event for each request that reached request interception. RequestLifecyclePlugin RequestLifecyclePlugin // ResponseInterceptor rewrites successful non-streaming HTTP execution responses before downstream delivery. @@ -395,6 +416,16 @@ type ModelRegistrar interface { type ModelRegistrationRequest struct { // Plugin is the metadata of the plugin being registered. Plugin Metadata + // MultiProviderModels reports whether the host supports ModelRegistrationResponse.Registrations. + MultiProviderModels bool `json:"multi_provider_models,omitempty"` +} + +// ProviderModels pairs a provider key with the models registered for it. +type ProviderModels struct { + // Provider is the provider key associated with the returned models. + Provider string `json:"provider,omitempty"` + // Models is the complete set of plugin-provided models for Provider. + Models []ModelInfo `json:"models,omitempty"` } // ModelRegistrationResponse returns provider and model metadata to register. @@ -403,6 +434,12 @@ type ModelRegistrationResponse struct { Provider string // Models is the complete set of plugin-provided models. Models []ModelInfo + // Registrations registers model sets for multiple providers in one call. + // When non-empty, each entry is registered under its own Provider for catalog + // and per-auth model discovery (plugins that also declare an Executor still + // bind that executor to the first registered provider); when empty, Provider + // and Models behave as a single registration. + Registrations []ProviderModels `json:"registrations,omitempty"` } // ModelProvider contributes provider-native static and per-auth model metadata. @@ -418,6 +455,8 @@ type StaticModelRequest struct { Plugin Metadata // Host contains relevant host configuration. Host HostConfigSummary + // MultiProviderModels reports whether the host supports ModelResponse.Registrations. + MultiProviderModels bool `json:"multi_provider_models,omitempty"` } // AuthModelRequest carries auth context for provider model discovery. @@ -446,6 +485,9 @@ type ModelResponse struct { Provider string // Models is the complete set of discovered provider models. Models []ModelInfo + // Registrations registers static model sets for multiple providers in one call. + // Used by StaticModels when non-empty; ignored by ModelsForAuth. + Registrations []ProviderModels `json:"registrations,omitempty"` // AuthUpdate contains updated auth data from model discovery when needed. AuthUpdate AuthData } @@ -1218,6 +1260,10 @@ type RequestInterceptResponse struct { Headers http.Header // Body replaces the current request body only when non-empty. Body []byte + // BodyEdits applies localized byte-range replacements to the current request body when Body is empty. + BodyEdits []RequestBodyEdit `json:"body_edits,omitempty"` + // Edits is an alias for BodyEdits. + Edits []RequestBodyEdit `json:"edits,omitempty"` // ClearHeaders explicitly removes current request headers before Headers is applied. ClearHeaders []string // Terminate stops the interceptor chain and prevents the request from reaching an upstream executor. @@ -1230,6 +1276,31 @@ type RequestInterceptResponse struct { ResponseBody []byte } +// RequestBodyPatcher scans a raw request body and returns localized byte-range edits +// without JSON/Base64 request body wrapping. +type RequestBodyPatcher interface { + PatchRequestBody(context.Context, []byte) (RequestBodyPatchResponse, error) +} + +// RequestBodyEdit describes a 0-based [Start:End] byte-range replacement in a UTF-8 JSON request body. +type RequestBodyEdit struct { + // Start is the inclusive 0-based start byte offset in the input body. + Start int `json:"start"` + // End is the exclusive 0-based end byte offset in the input body. + End int `json:"end"` + // Replacement is the UTF-8 replacement text spliced into [Start:End]. + Replacement string `json:"replacement,omitempty"` + // ReplacementBytes optionally carries raw replacement bytes when non-empty. + ReplacementBytes []byte `json:"replacement_bytes,omitempty"` +} + +// RequestBodyPatchResponse returns localized byte-range edits for a request body. +type RequestBodyPatchResponse struct { + // Edits lists byte-range replacements relative to the input body received by this plugin. + // When empty, the host reuses the input body without copying. + Edits []RequestBodyEdit `json:"edits,omitempty"` +} + // RequestCompletionOutcome identifies how an intercepted request ended. type RequestCompletionOutcome string