From ca601db05d858472b25e7c56580555ab62b90e68 Mon Sep 17 00:00:00 2001 From: sususu98 <33882693+sususu98@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:15:37 +0800 Subject: [PATCH] feat: observe upstream provider quota signals (#5211) Codex and Claude already emit credential-level quota watermarks on ordinary responses. CPA used to drop them. Keep the latest watermark in memory and return it from the management auth-file API. Hard rule: this is observation only. It must not change scheduling, cooldown selection, or auth-file persistence. Snapshot, not accumulation - QuotaState now has ObservedAt and a bounded Signals map. MarkResult fills them from the response headers already recorded on the request. - Signals is the current response, not a union of earlier ones. Retry-After and "limit reached" only appear on the response that produced them; merging across responses would keep an expired value forever. - A response with no quota header (transport failure, 5xx, unrelated endpoint) leaves the previous snapshot in place. - ObservedAt is the time of the current snapshot. It advances even when the values did not change, so a consumer can tell a fresh reading from a stale one. - When two model states merge, keep the newer snapshot. Do not union keys captured at different times. What is observed, and what is not - One predicate, ProviderSupportsQuotaObservation, decides the provider set. - Keep Codex and Claude. Drop Kimi, xAI/Grok, Antigravity, and the Gemini family (gemini/vertex/aistudio): their ordinary headers are not a reliable credential-level remaining quota. - Count-tokens reuses the credential but is not generation traffic. ExecuteCount sets SkipQuotaObservation so those headers cannot replace the last generation snapshot. Cooldown and success/failure accounting still run. Cooldown must not overwrite the last snapshot - Observation writes only ObservedAt and Signals. - Cooldown writes only Exceeded, Reason, NextRecoverAt, and BackoffLevel, through applyCooldownFields. Never assign a fresh QuotaState{...} over a live value: that would zero the snapshot on 429, Cloudflare, credential- scope sibling updates, and cooldown clears. - If a credential-quota cooldown is still active, MarkResult still observes an already-present model state. It does not create scheduler state just to record a watermark. - .cds files persist cooldownFieldsOf(Quota) only. Restore keeps the newer ObservedAt, so reloading cooldown cannot clobber a newer in-memory snapshot. - cooldownQuotaEqual still ignores observation fields, so a watermark change cannot by itself persist cooldown or move the scheduler. - The management payload omits every cooldown field, so it cannot be mistaken for scheduler state or wired back into scheduling. - Manual ResetQuota still clears the full QuotaState. Codex websocket events - Codex WS reports quota as codex.rate_limits frames, not HTTP headers. ParseCodexQuotaEventHeaders turns one event into the same bounded header shape, and MergeResponseHeaders folds it into the request-scoped holder. additional_rate_limits is accepted as an object (websocket) or an array (/wham/usage). - Parse only through AppendCodexAPIWebsocketResponse. The shared AppendAPIWebsocketResponse is also used by xAI, and xAI error frames really do carry x-ratelimit-* headers. Parsing every frame as Codex quota would forge Codex headers into another provider's request log. - Also capture code_review_rate_limits. - A malformed active-limit name drops only that one header, not the window watermarks parsed from the same event. - The type discriminator scans a bounded frame prefix, not every byte of every frame. - HTTP namespaces an extra limit by short name (x-codex-bengalfox-*); WS namespaces it by limit name (GPT-5.3-Codex-Spark). The two paths cannot emit the same header names. The X-Codex-Additional- prefix marks the WS origin, and snapshot replacement keeps the two spellings from piling up. Hardening - Reject observed values with control characters. These strings reach the plain-text request log, and Limit-Name is upstream-controlled, so CR/LF could forge a header line. - When the header cap is hit, keep plan/credits/primary ahead of additional-limit namespaces, then sort names so truncation is deterministic. - QuotaState.Clone deep-copies Signals and is used by Auth.Clone and ModelState.Clone. - Token stores still serialize credential metadata only, so observation adds no auth-file writes. --- .../api/handlers/management/auth_files.go | 44 ++ .../management/auth_files_quota_test.go | 68 ++ .../auth_files_recent_requests_test.go | 42 ++ internal/logging/requestmeta.go | 35 + internal/logging/requestmeta_test.go | 34 + .../executor/codex_websockets_execute.go | 2 +- .../executor/codex_websockets_stream.go | 4 +- .../runtime/executor/helps/codex_quota.go | 356 +++++++++ .../executor/helps/codex_quota_test.go | 340 +++++++++ .../runtime/executor/helps/logging_helpers.go | 7 + sdk/cliproxy/auth/conductor.go | 4 + sdk/cliproxy/auth/conductor_cooldown.go | 53 +- sdk/cliproxy/auth/conductor_execution.go | 21 +- .../auth/conductor_execution_quota_test.go | 133 ++++ sdk/cliproxy/auth/conductor_home_execution.go | 10 +- sdk/cliproxy/auth/conductor_stream.go | 9 +- sdk/cliproxy/auth/quota_signals.go | 227 ++++++ sdk/cliproxy/auth/quota_signals_test.go | 685 ++++++++++++++++++ sdk/cliproxy/auth/types.go | 23 + 19 files changed, 2072 insertions(+), 25 deletions(-) create mode 100644 internal/api/handlers/management/auth_files_quota_test.go create mode 100644 internal/logging/requestmeta_test.go create mode 100644 internal/runtime/executor/helps/codex_quota.go create mode 100644 internal/runtime/executor/helps/codex_quota_test.go create mode 100644 sdk/cliproxy/auth/conductor_execution_quota_test.go create mode 100644 sdk/cliproxy/auth/quota_signals.go create mode 100644 sdk/cliproxy/auth/quota_signals_test.go diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index f42b6811..258893cd 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -344,6 +344,10 @@ func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth) gin.H { entry["success"] = auth.Success entry["failed"] = auth.Failed entry["recent_requests"] = auth.RecentRequestsSnapshot(time.Now()) + entry["quota"] = quotaObservationPayloadForProvider(auth.Provider, auth.Quota) + if modelQuotas := modelQuotaObservationPayload(auth.Provider, auth.ModelStates); len(modelQuotas) > 0 { + entry["model_quotas"] = modelQuotas + } if email := authEmail(auth); email != "" { entry["email"] = email } @@ -441,6 +445,46 @@ func authFileRequestRetryFromJSON(data []byte) (int, bool) { return (&coreauth.Auth{Metadata: metadata}).RequestRetryOverride() } +// quotaObservationPayload exposes only passive provider observations. Cooldown +// fields are intentionally excluded so this management response cannot be +// mistaken for scheduler state or influence scheduling behavior. +func quotaObservationPayloadForProvider(provider string, quota coreauth.QuotaState) gin.H { + if !coreauth.ProviderSupportsQuotaObservation(provider) { + return quotaObservationPayload(coreauth.QuotaState{}) + } + return quotaObservationPayload(quota) +} + +func quotaObservationPayload(quota coreauth.QuotaState) gin.H { + observed := gin.H{} + if !quota.ObservedAt.IsZero() { + observed["observed_at"] = quota.ObservedAt + } + signals := make(map[string]string, len(quota.Signals)) + for key, value := range quota.Signals { + signals[key] = value + } + observed["signals"] = signals + return observed +} + +func modelQuotaObservationPayload(provider string, states map[string]*coreauth.ModelState) gin.H { + if !coreauth.ProviderSupportsQuotaObservation(provider) { + return gin.H{} + } + observations := gin.H{} + for model, state := range states { + if state == nil { + continue + } + if state.Quota.ObservedAt.IsZero() && len(state.Quota.Signals) == 0 { + continue + } + observations[model] = quotaObservationPayloadForProvider(provider, state.Quota) + } + return observations +} + func authWeightValue(auth *coreauth.Auth) (int64, bool) { if auth == nil { return 0, false diff --git a/internal/api/handlers/management/auth_files_quota_test.go b/internal/api/handlers/management/auth_files_quota_test.go new file mode 100644 index 00000000..8717a51f --- /dev/null +++ b/internal/api/handlers/management/auth_files_quota_test.go @@ -0,0 +1,68 @@ +package management + +import ( + "testing" + "time" + + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestModelQuotaObservationPayloadOmitsUnsupportedProviders(t *testing.T) { + states := map[string]*coreauth.ModelState{ + "grok-4": { + Quota: coreauth.QuotaState{ + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Ratelimit-Remaining-Requests": "1"}, + }, + }, + } + if got := modelQuotaObservationPayload("grok", states); len(got) != 0 { + t.Fatalf("unsupported provider returned model observations: %#v", got) + } + for _, provider := range []string{"gemini", "gemini-interactions", "openai", "openai-compatibility", "plugin-provider"} { + if got := modelQuotaObservationPayload(provider, states); len(got) != 0 { + t.Fatalf("provider %q returned model observations: %#v", provider, got) + } + } +} + +func TestModelQuotaObservationPayloadSkipsNilAndEmptyStates(t *testing.T) { + states := map[string]*coreauth.ModelState{ + "nil": nil, + "empty": &coreauth.ModelState{}, + "observed": &coreauth.ModelState{Quota: coreauth.QuotaState{ + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Plan-Type": "pro"}, + }}, + } + got := modelQuotaObservationPayload("codex", states) + if len(got) != 1 { + t.Fatalf("model observations = %#v, want only observed state", got) + } + if _, ok := got["observed"]; !ok { + t.Fatalf("observed model quota missing: %#v", got) + } +} + +func TestQuotaObservationPayloadExcludesCooldownState(t *testing.T) { + payload := quotaObservationPayload(coreauth.QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: time.Unix(20, 0), + BackoffLevel: 3, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Plan-Type": "pro"}, + }) + if _, ok := payload["exceeded"]; ok { + t.Fatalf("cooldown exceeded leaked: %#v", payload) + } + if _, ok := payload["reason"]; ok { + t.Fatalf("cooldown reason leaked: %#v", payload) + } + if _, ok := payload["next_recover_at"]; ok { + t.Fatalf("cooldown recovery leaked: %#v", payload) + } + if _, ok := payload["backoff_level"]; ok { + t.Fatalf("cooldown backoff leaked: %#v", payload) + } +} diff --git a/internal/api/handlers/management/auth_files_recent_requests_test.go b/internal/api/handlers/management/auth_files_recent_requests_test.go index f3c5107c..7ae8fcd1 100644 --- a/internal/api/handlers/management/auth_files_recent_requests_test.go +++ b/internal/api/handlers/management/auth_files_recent_requests_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" @@ -25,6 +26,22 @@ func TestListAuthFiles_IncludesRecentRequestsBuckets(t *testing.T) { Metadata: map[string]any{ "type": "codex", }, + Quota: coreauth.QuotaState{ + ObservedAt: time.Date(2026, 8, 22, 0, 0, 0, 0, time.UTC), + Signals: map[string]string{ + "X-Codex-Primary-Used-Percent": "58", + }, + }, + ModelStates: map[string]*coreauth.ModelState{ + "gpt-5": { + Quota: coreauth.QuotaState{ + ObservedAt: time.Date(2026, 8, 22, 0, 1, 0, 0, time.UTC), + Signals: map[string]string{ + "Retry-After": "120", + }, + }, + }, + }, } if _, errRegister := manager.Register(context.Background(), record); errRegister != nil { t.Fatalf("failed to register auth record: %v", errRegister) @@ -68,6 +85,31 @@ func TestListAuthFiles_IncludesRecentRequestsBuckets(t *testing.T) { t.Fatalf("expected failed number, got %#v", fileEntry["failed"]) } + quota, ok := fileEntry["quota"].(map[string]any) + if !ok { + t.Fatalf("expected quota observation object, got %#v", fileEntry["quota"]) + } + if _, ok := quota["observed_at"].(string); !ok { + t.Fatalf("expected quota observed_at string, got %#v", quota["observed_at"]) + } + quotaSignals, ok := quota["signals"].(map[string]any) + if !ok || quotaSignals["X-Codex-Primary-Used-Percent"] != "58" { + t.Fatalf("expected auth quota signals, got %#v", quota["signals"]) + } + + modelQuotas, ok := fileEntry["model_quotas"].(map[string]any) + if !ok { + t.Fatalf("expected model_quotas object, got %#v", fileEntry["model_quotas"]) + } + modelQuota, ok := modelQuotas["gpt-5"].(map[string]any) + if !ok { + t.Fatalf("expected gpt-5 quota observation, got %#v", modelQuotas["gpt-5"]) + } + modelSignals, ok := modelQuota["signals"].(map[string]any) + if !ok || modelSignals["Retry-After"] != "120" { + t.Fatalf("expected model quota signals, got %#v", modelQuota["signals"]) + } + recentRaw, ok := fileEntry["recent_requests"].([]any) if !ok { t.Fatalf("expected recent_requests array, got %#v", fileEntry["recent_requests"]) diff --git a/internal/logging/requestmeta.go b/internal/logging/requestmeta.go index 576bf5db..f2fc32fa 100644 --- a/internal/logging/requestmeta.go +++ b/internal/logging/requestmeta.go @@ -84,6 +84,17 @@ func WithResponseHeadersHolder(ctx context.Context) context.Context { return context.WithValue(ctx, responseHeadersKey{}, &responseHeadersHolder{}) } +// WithFreshResponseHeadersHolder starts an isolated upstream response attempt. +// Unlike WithResponseHeadersHolder, it always shadows any holder inherited from +// the parent request so a later retry cannot observe headers from an earlier +// credential or model attempt. +func WithFreshResponseHeadersHolder(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, responseHeadersKey{}, &responseHeadersHolder{}) +} + func SetResponseStatus(ctx context.Context, status int) { if ctx == nil || status <= 0 { return @@ -108,6 +119,30 @@ func SetResponseHeaders(ctx context.Context, headers http.Header) { holder.headers = cloneHTTPHeader(headers) } +// MergeResponseHeaders adds headers observed after the initial HTTP response, +// such as quota metadata delivered in a websocket event. +func MergeResponseHeaders(ctx context.Context, headers http.Header) { + if ctx == nil || len(headers) == 0 { + return + } + holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder) + if !ok || holder == nil { + return + } + holder.mu.Lock() + defer holder.mu.Unlock() + if holder.headers == nil { + holder.headers = make(http.Header, len(headers)) + } + for key, values := range headers { + canonicalKey := http.CanonicalHeaderKey(key) + if canonicalKey == "" { + continue + } + holder.headers[canonicalKey] = append([]string(nil), values...) + } +} + func GetResponseStatus(ctx context.Context) int { if ctx == nil { return 0 diff --git a/internal/logging/requestmeta_test.go b/internal/logging/requestmeta_test.go new file mode 100644 index 00000000..8825cdc9 --- /dev/null +++ b/internal/logging/requestmeta_test.go @@ -0,0 +1,34 @@ +package logging + +import ( + "context" + "net/http" + "testing" +) + +func TestWithFreshResponseHeadersHolderIsolatesAttempts(t *testing.T) { + requestCtx := WithResponseHeadersHolder(context.Background()) + SetResponseHeaders(requestCtx, http.Header{"X-Upstream-Attempt": []string{"request"}}) + + firstAttemptCtx := WithFreshResponseHeadersHolder(requestCtx) + if headers := GetResponseHeaders(firstAttemptCtx); len(headers) != 0 { + t.Fatalf("fresh attempt inherited request headers: %#v", headers) + } + SetResponseHeaders(firstAttemptCtx, http.Header{"X-Upstream-Attempt": []string{"first"}}) + + secondAttemptCtx := WithFreshResponseHeadersHolder(firstAttemptCtx) + if headers := GetResponseHeaders(secondAttemptCtx); len(headers) != 0 { + t.Fatalf("second attempt inherited first attempt headers: %#v", headers) + } + SetResponseHeaders(secondAttemptCtx, http.Header{"X-Upstream-Attempt": []string{"second"}}) + + if got := GetResponseHeaders(requestCtx).Get("X-Upstream-Attempt"); got != "request" { + t.Fatalf("request holder = %q, want request", got) + } + if got := GetResponseHeaders(firstAttemptCtx).Get("X-Upstream-Attempt"); got != "first" { + t.Fatalf("first attempt holder = %q, want first", got) + } + if got := GetResponseHeaders(secondAttemptCtx).Get("X-Upstream-Attempt"); got != "second" { + t.Fatalf("second attempt holder = %q, want second", got) + } +} diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index 72bace8c..ef887272 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -285,7 +285,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) - helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 52279bc4..5b859009 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -327,7 +327,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) - helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { @@ -538,7 +538,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) - helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { diff --git a/internal/runtime/executor/helps/codex_quota.go b/internal/runtime/executor/helps/codex_quota.go new file mode 100644 index 00000000..cc0c1c2b --- /dev/null +++ b/internal/runtime/executor/helps/codex_quota.go @@ -0,0 +1,356 @@ +package helps + +import ( + "net/http" + "strings" + + "github.com/tidwall/gjson" +) + +const ( + codexRateLimitsEventType = "codex.rate_limits" + codexQuotaAdditionalHeaderKey = "X-Codex-Additional-" + maxCodexAdditionalRateLimits = 8 +) + +// ParseCodexQuotaEventHeaders converts one Codex websocket quota event into the +// same bounded header representation used by HTTP usage observations. Codex +// sends additional_rate_limits as an object on the websocket path, while the +// /wham/usage probe exposes the equivalent data as an array; both forms are +// normalized into namespaced X-Codex headers here. +// +// The websocket payload identifies each additional limit only by limit name +// ("GPT-5.3-Codex-Spark"), while HTTP responses namespace the same limit by a +// short name (x-codex-bengalfox-*). The two forms therefore cannot produce +// identical header names; the X-Codex-Additional- prefix marks the websocket +// origin, and the snapshot-replacement semantics in QuotaState keep the two +// spellings from accumulating side by side. +func ParseCodexQuotaEventHeaders(payload []byte) http.Header { + if len(payload) == 0 { + return nil + } + + eventKind := codexQuotaEventKind(payload) + if eventKind == codexQuotaEventOther { + return nil + } + root := gjson.ParseBytes(payload) + if eventKind == codexQuotaEventError { + return parseCodexQuotaHeadersObject(root.Get("headers")) + } + + headers := make(http.Header) + hasQuotaData := false + + baseRateLimits := firstCodexQuotaResult(root, "rate_limits", "rateLimit") + if addCodexQuotaRateLimitHeaders(headers, "X-Codex-", baseRateLimits) { + // The main rate-limit object has no user-facing name in the websocket + // payload, so its fields retain the HTTP response header names. + hasQuotaData = true + } + + additional := firstCodexQuotaResult(root, "additional_rate_limits", "additionalRateLimits") + additionalCount := 0 + if additional.Exists() && (additional.IsObject() || additional.IsArray()) { + additional.ForEach(func(key, value gjson.Result) bool { + if additionalCount >= maxCodexAdditionalRateLimits { + return false + } + limitName := strings.TrimSpace(key.String()) + if additional.IsArray() { + limitName = firstCodexQuotaResultString(value, "limit_name", "limitName", "name") + } + if limitName == "" { + return true + } + identifier := normalizeCodexQuotaHeaderIdentifier(limitName) + if identifier == "" { + return true + } + rateInfo := firstCodexQuotaResult(value, "rate_limit", "rateLimit") + if !rateInfo.Exists() { + rateInfo = value + } + prefix := codexQuotaAdditionalHeaderKey + identifier + "-" + if addCodexQuotaRateLimitHeaders(headers, prefix, rateInfo) { + // The limit name is upstream-controlled and lands in the + // plain-text request log, so reject control characters here the + // same way the plan type is validated below. + if validCodexQuotaEventText(limitName) { + headers.Set(prefix+"Limit-Name", limitName) + } + hasQuotaData = true + additionalCount++ + } + return true + }) + } + + codeReview := firstCodexQuotaResult(root, "code_review_rate_limits", "codeReviewRateLimits") + if addCodexQuotaRateLimitHeaders(headers, "X-Codex-Code-Review-", codeReview) { + hasQuotaData = true + } + + credits := firstCodexQuotaResult(root, "credits") + if credits.Exists() && credits.IsObject() { + hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Has-Credits", credits, "has_credits", "hasCredits") || hasQuotaData + hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Unlimited", credits, "unlimited") || hasQuotaData + hasQuotaData = setCodexQuotaScalarHeader(headers, "X-Codex-Credits-Balance", credits, "balance") || hasQuotaData + } + + if !hasQuotaData { + return nil + } + // A malformed active-limit name only invalidates that one header; the + // window percentages collected above stay valid and must not be discarded. + activeLimit := firstCodexQuotaResultString(root, "metered_limit_name", "meteredLimitName", "limit_name", "limitName") + if validCodexQuotaEventIdentifier(activeLimit) { + headers.Set("X-Codex-Active-Limit", activeLimit) + } + if planType := firstCodexQuotaResultString(root, "plan_type", "planType"); validCodexQuotaEventText(planType) { + headers.Set("X-Codex-Plan-Type", planType) + } + return headers +} + +const ( + codexQuotaEventOther = iota + codexQuotaEventError + codexQuotaEventRateLimits +) + +// codexQuotaEventKind performs an allocation-free discriminator check for the +// common non-quota websocket event path. Codex frames carry the top-level +// type field within their opening bytes (after sequence_number at most), so +// only a bounded prefix is scanned instead of the whole frame; full JSON +// parsing is reserved for quota and error events. +func codexQuotaEventKind(payload []byte) int { + const scanLimit = 256 + end := len(payload) + if end > scanLimit { + end = scanLimit + } + for i := 0; i+6 < end; i++ { + if payload[i] != '"' || payload[i+1] != 't' || payload[i+2] != 'y' || payload[i+3] != 'p' || payload[i+4] != 'e' || payload[i+5] != '"' { + continue + } + j := i + 6 + for j < end && (payload[j] == ' ' || payload[j] == '\t' || payload[j] == '\r' || payload[j] == '\n') { + j++ + } + if j >= end || payload[j] != ':' { + continue + } + j++ + for j < end && (payload[j] == ' ' || payload[j] == '\t' || payload[j] == '\r' || payload[j] == '\n') { + j++ + } + if j+7 <= end && payload[j] == '"' && payload[j+1] == 'e' && payload[j+2] == 'r' && payload[j+3] == 'r' && payload[j+4] == 'o' && payload[j+5] == 'r' && payload[j+6] == '"' { + return codexQuotaEventError + } + const rateLimits = `"codex.rate_limits"` + if j+len(rateLimits) <= end { + matched := true + for offset := range rateLimits { + if payload[j+offset] != rateLimits[offset] { + matched = false + break + } + } + if matched { + return codexQuotaEventRateLimits + } + } + } + return codexQuotaEventOther +} + +func addCodexQuotaRateLimitHeaders(headers http.Header, prefix string, rateInfo gjson.Result) bool { + if !rateInfo.Exists() || !rateInfo.IsObject() { + return false + } + + changed := false + if setCodexQuotaScalarHeaderFromResult(headers, prefix+"Allowed", rateInfo, "allowed") { + changed = true + } + if setCodexQuotaScalarHeaderFromResult(headers, prefix+"Limit-Reached", rateInfo, "limit_reached", "limitReached") { + changed = true + } + for _, windowName := range []string{"primary", "secondary"} { + window := firstCodexQuotaResult(rateInfo, windowName) + if !window.Exists() || !window.IsObject() { + continue + } + used := firstCodexQuotaResult(window, "used_percent", "usedPercent") + minutes := firstCodexQuotaResult(window, "window_minutes", "windowMinutes") + resetAfter := firstCodexQuotaResult(window, "reset_after_seconds", "resetAfterSeconds") + resetAt := firstCodexQuotaResult(window, "reset_at", "resetAt") + hasResetAfter := resetAfter.Exists() && resetAfter.Int() >= 0 + hasResetAt := resetAt.Exists() && resetAt.Int() > 0 + if !used.Exists() || !minutes.Exists() || + used.Float() < 0 || used.Float() > 100 || minutes.Int() <= 0 || + (!hasResetAfter && !hasResetAt) { + continue + } + windowPrefix := prefix + strings.ToUpper(windowName[:1]) + windowName[1:] + "-" + setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Used-Percent", window, "used_percent", "usedPercent") + setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Window-Minutes", window, "window_minutes", "windowMinutes") + if hasResetAfter { + setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Reset-After-Seconds", window, "reset_after_seconds", "resetAfterSeconds") + } + if hasResetAt { + setCodexQuotaScalarHeaderFromResult(headers, windowPrefix+"Reset-At", window, "reset_at", "resetAt") + } + changed = true + } + return changed +} + +func parseCodexQuotaHeadersObject(headersNode gjson.Result) http.Header { + if !headersNode.Exists() || !headersNode.IsObject() { + return nil + } + headers := make(http.Header) + headersNode.ForEach(func(key, value gjson.Result) bool { + name := http.CanonicalHeaderKey(strings.TrimSpace(key.String())) + if !isCodexQuotaHeaderName(name) { + return true + } + if raw := codexQuotaScalarValue(value); raw != "" { + headers.Set(name, raw) + } + return true + }) + if len(headers) == 0 { + return nil + } + return headers +} + +func isCodexQuotaHeaderName(name string) bool { + lower := strings.ToLower(strings.TrimSpace(name)) + if lower == "retry-after" || strings.HasPrefix(lower, "x-ratelimit-") { + return true + } + if lower == "x-codex-active-limit" || lower == "x-codex-plan-type" || + strings.HasPrefix(lower, "x-codex-credits-") { + return true + } + if !strings.HasPrefix(lower, "x-codex-") { + return false + } + for _, marker := range []string{ + "-allowed", + "-limit-reached", + "-limit-name", + "-used-percent", + "-window-minutes", + "-reset-after-seconds", + "-reset-at", + "-over-secondary-limit-percent", + } { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +func setCodexQuotaScalarHeader(headers http.Header, name string, object gjson.Result, paths ...string) bool { + return setCodexQuotaScalarHeaderFromResult(headers, name, object, paths...) +} + +func setCodexQuotaScalarHeaderFromResult(headers http.Header, name string, object gjson.Result, paths ...string) bool { + value := firstCodexQuotaResult(object, paths...) + if !value.Exists() { + return false + } + raw := codexQuotaScalarValue(value) + if raw == "" { + return false + } + headers.Set(name, raw) + return true +} + +func codexQuotaScalarValue(value gjson.Result) string { + switch value.Type { + case gjson.String: + return strings.TrimSpace(value.String()) + case gjson.Number, gjson.True, gjson.False: + return strings.TrimSpace(value.Raw) + default: + return "" + } +} + +func firstCodexQuotaResult(object gjson.Result, paths ...string) gjson.Result { + for _, path := range paths { + value := object.Get(path) + if value.Exists() && value.Type != gjson.Null { + return value + } + } + return gjson.Result{} +} + +func firstCodexQuotaResultString(object gjson.Result, paths ...string) string { + for _, path := range paths { + value := firstCodexQuotaResult(object, path) + if raw := codexQuotaScalarValue(value); raw != "" { + return raw + } + } + return "" +} + +func normalizeCodexQuotaHeaderIdentifier(value string) string { + value = strings.TrimSpace(value) + if value == "" || len(value) > 128 { + return "" + } + var builder strings.Builder + builder.Grow(len(value)) + lastDash := false + for _, char := range value { + switch { + case (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') || char == '.' || char == '_': + builder.WriteRune(char) + lastDash = false + case char == '-': + builder.WriteRune(char) + lastDash = false + default: + if !lastDash { + builder.WriteByte('-') + lastDash = true + } + } + } + identifier := strings.Trim(builder.String(), "-_.") + if identifier == "" || !validCodexQuotaEventIdentifier(identifier) { + return "" + } + return identifier +} + +func validCodexQuotaEventIdentifier(value string) bool { + value = strings.TrimSpace(value) + if value == "" || len(value) > 256 { + return false + } + for _, char := range value { + if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || + (char >= '0' && char <= '9') || char == '_' || char == '-' || char == '.' { + continue + } + return false + } + return true +} + +func validCodexQuotaEventText(value string) bool { + value = strings.TrimSpace(value) + return value != "" && len(value) <= 256 && !strings.ContainsAny(value, "\r\n") +} diff --git a/internal/runtime/executor/helps/codex_quota_test.go b/internal/runtime/executor/helps/codex_quota_test.go new file mode 100644 index 00000000..b1f24cce --- /dev/null +++ b/internal/runtime/executor/helps/codex_quota_test.go @@ -0,0 +1,340 @@ +package helps + +import ( + "context" + "net/http" + "strings" + "testing" + + internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" +) + +func TestParseCodexQuotaEventHeadersPreservesActiveLimit(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "plan_type":"pro", + "metered_limit_name":"codex_bengalfox", + "rate_limits":{ + "primary":{"used_percent":2,"window_minutes":10080,"reset_at":1782951970}, + "secondary":null + } + }`)) + if headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" || + headers.Get("X-Codex-Primary-Used-Percent") != "2" || + headers.Get("X-Codex-Primary-Window-Minutes") != "10080" || + headers.Get("X-Codex-Primary-Reset-At") != "1782951970" || + headers.Get("X-Codex-Plan-Type") != "pro" { + t.Fatalf("unexpected quota headers: %#v", headers) + } +} + +func TestParseCodexQuotaEventHeadersPreservesAdditionalLimitsAndCredits(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "plan_type":"pro", + "rate_limits":{ + "allowed":true, + "limit_reached":false, + "primary":{"used_percent":48,"window_minutes":10080,"reset_after_seconds":523210,"reset_at":1786677299}, + "secondary":null + }, + "additional_rate_limits":{ + "GPT-5.3-Codex-Spark":{ + "allowed":true, + "limit_reached":false, + "primary":{"used_percent":3,"window_minutes":300,"reset_after_seconds":10148,"reset_at":1787231961}, + "secondary":{"used_percent":63,"window_minutes":10080,"reset_after_seconds":75420,"reset_at":1787290791} + } + }, + "credits":{"has_credits":false,"unlimited":false,"balance":"0"} + }`)) + for key, want := range map[string]string{ + "X-Codex-Primary-Used-Percent": "48", + "X-Codex-Primary-Window-Minutes": "10080", + "X-Codex-Primary-Reset-After-Seconds": "523210", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Used-Percent": "3", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Window-Minutes": "300", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Secondary-Used-Percent": "63", + "X-Codex-Credits-Has-Credits": "false", + "X-Codex-Credits-Unlimited": "false", + "X-Codex-Credits-Balance": "0", + } { + if got := headers.Get(key); got != want { + t.Fatalf("header %s = %q, want %q; all headers = %#v", key, got, want, headers) + } + } +} + +func TestParseCodexQuotaEventHeadersReadsQuotaHeadersFromErrorFrames(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"error", + "status_code":429, + "headers":{ + "X-Codex-Primary-Used-Percent":"100", + "X-Codex-Primary-Window-Minutes":"10080", + "X-Codex-Primary-Reset-After-Seconds":"437380", + "X-Codex-Credits-Balance":"0", + "X-Codex-Turn-State":"must-not-be-observed", + "Set-Cookie":"must-not-be-observed" + } + }`)) + if headers.Get("X-Codex-Primary-Used-Percent") != "100" || + headers.Get("X-Codex-Primary-Reset-After-Seconds") != "437380" || + headers.Get("X-Codex-Credits-Balance") != "0" { + t.Fatalf("unexpected error quota headers: %#v", headers) + } + if headers.Get("Set-Cookie") != "" || headers.Get("X-Codex-Turn-State") != "" { + t.Fatalf("non-quota error header leaked into quota observation: %#v", headers) + } +} + +func TestParseCodexQuotaEventHeadersDoesNotInventActiveLimitAndRejectsIncompleteWindows(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "rate_limits":{"primary":{"used_percent":42,"window_minutes":300,"reset_after_seconds":60}} + }`)) + if headers.Get("X-Codex-Active-Limit") != "" || headers.Get("X-Codex-Primary-Reset-After-Seconds") != "60" { + t.Fatalf("unexpected quota headers: %#v", headers) + } + if got := ParseCodexQuotaEventHeaders([]byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`)); got != nil { + t.Fatalf("incomplete quota event produced headers: %#v", got) + } + + headers = ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":42,"window_minutes":300}, + "secondary":{"used_percent":17,"window_minutes":10080,"reset_after_seconds":60} + } + }`)) + if headers.Get("X-Codex-Primary-Used-Percent") != "" || + headers.Get("X-Codex-Secondary-Used-Percent") != "17" { + t.Fatalf("partially valid quota event produced incomplete headers: %#v", headers) + } +} + +func TestParseCodexQuotaEventHeadersSupportsAdditionalArrayAndCamelCase(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "planType":"pro", + "meteredLimitName":"premium", + "rateLimit":{ + "allowed":true, + "limitReached":false, + "primary":{"usedPercent":71,"windowMinutes":10080,"resetAfterSeconds":60} + }, + "additionalRateLimits":[{ + "limitName":"GPT-5.3-Codex-Spark", + "rateLimit":{ + "primary":{"usedPercent":13,"windowMinutes":300,"resetAt":1787399817}, + "secondary":{"usedPercent":51,"windowMinutes":10080,"resetAfterSeconds":120} + } + }] + }`)) + for key, want := range map[string]string{ + "X-Codex-Active-Limit": "premium", + "X-Codex-Plan-Type": "pro", + "X-Codex-Primary-Used-Percent": "71", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Used-Percent": "13", + "X-Codex-Additional-GPT-5.3-Codex-Spark-Secondary-Reset-After-Seconds": "120", + } { + if got := headers.Get(key); got != want { + t.Fatalf("header %s = %q, want %q; all headers = %#v", key, got, want, headers) + } + } +} + +func TestParseCodexQuotaEventHeadersKeepsOnlySafeErrorHeaders(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"error", + "headers":{ + "X-Ratelimit-Remaining-Requests":"7", + "Retry-After":"60", + "X-Codex-Primary-Over-Secondary-Limit-Percent":"0", + "X-Codex-Turn-State":"secret", + "Authorization":"Bearer secret" + } + }`)) + for key, want := range map[string]string{ + "X-Ratelimit-Remaining-Requests": "7", + "Retry-After": "60", + "X-Codex-Primary-Over-Secondary-Limit-Percent": "0", + } { + if got := headers.Get(key); got != want { + t.Fatalf("header %s = %q, want %q", key, got, want) + } + } + if headers.Get("X-Codex-Turn-State") != "" || headers.Get("Authorization") != "" { + t.Fatalf("unsafe error headers leaked: %#v", headers) + } +} + +func TestParseCodexQuotaEventHeadersNonQuotaEventDoesNotAllocate(t *testing.T) { + payload := []byte(`{"type":"response.output_text.delta","delta":"hello"}`) + allocs := testing.AllocsPerRun(1000, func() { + if headers := ParseCodexQuotaEventHeaders(payload); headers != nil { + t.Fatalf("non-quota event produced headers: %#v", headers) + } + }) + if allocs != 0 { + t.Fatalf("non-quota event allocations = %v, want 0", allocs) + } +} + +// Large non-quota frames must stay allocation-free: only the bounded frame +// prefix is scanned for the type discriminator. +func TestParseCodexQuotaEventHeadersLargeNonQuotaFrameDoesNotAllocate(t *testing.T) { + big := strings.Repeat("x", 4096) + payload := []byte(`{"sequence_number":7,"type":"response.output_text.delta","delta":"` + big + `"}`) + allocs := testing.AllocsPerRun(100, func() { + if headers := ParseCodexQuotaEventHeaders(payload); headers != nil { + t.Fatalf("non-quota event produced headers: %#v", headers) + } + }) + if allocs != 0 { + t.Fatalf("large non-quota frame allocations = %v, want 0", allocs) + } +} + +// Quota markers live in the frame prefix, so large quota events must still be parsed. +func TestParseCodexQuotaEventHeadersLargeRateLimitsFrameIsParsed(t *testing.T) { + padding := strings.Repeat("p", 4096) + payload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42,"limit_percent":100,"window_minutes":10080,"reset_after_seconds":3600,"used_minutes":120},"secondary":{"used_percent":1,"limit_percent":100,"window_minutes":300,"reset_after_seconds":60,"used_minutes":2},"active_limit":"primary","credits":{"total_credits":"10","available_credits":"5","pct_remaining":"50"}},"padding":"` + padding + `"}`) + headers := ParseCodexQuotaEventHeaders(payload) + if headers == nil { + t.Fatal("large rate_limits frame produced no headers") + } + if got := headers.Get("X-Codex-Primary-Used-Percent"); got != "42" { + t.Fatalf("X-Codex-Primary-Used-Percent = %q, want 42", got) + } +} + +func TestParseCodexQuotaEventHeadersRejectsInvalidAndEmptyEvents(t *testing.T) { + for _, payload := range []string{ + ``, + `{"type":"response.completed"}`, + `{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":101,"window_minutes":300,"reset_after_seconds":60}}}`, + } { + if headers := ParseCodexQuotaEventHeaders([]byte(payload)); headers != nil { + t.Fatalf("payload %q produced unexpected headers: %#v", payload, headers) + } + } +} + +// A malformed active-limit name must not discard the window watermarks that were +// parsed successfully from the same event. +func TestParseCodexQuotaEventHeadersKeepsWindowsWhenActiveLimitIsInvalid(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "metered_limit_name":"bad limit", + "rate_limits":{"primary":{"used_percent":1,"window_minutes":300,"reset_after_seconds":60}} + }`)) + if headers == nil { + t.Fatal("valid window watermarks were discarded") + } + if got := headers.Get("X-Codex-Primary-Used-Percent"); got != "1" { + t.Fatalf("X-Codex-Primary-Used-Percent = %q, want \"1\"", got) + } + if got := headers.Get("X-Codex-Active-Limit"); got != "" { + t.Fatalf("malformed active limit was retained: %q", got) + } +} + +// Real websocket events namespace additional limits by limit name and carry a +// code_review_rate_limits sibling; both must survive parsing. +func TestParseCodexQuotaEventHeadersObservedProductionEvent(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "plan_type":"pro", + "rate_limits":{"allowed":true,"limit_reached":false,"primary":{"used_percent":81,"window_minutes":10080,"reset_after_seconds":137160,"reset_at":1787588999},"secondary":null}, + "code_review_rate_limits":{"primary":{"used_percent":4,"window_minutes":300,"reset_after_seconds":900}}, + "additional_rate_limits":{"GPT-5.3-Codex-Spark":{"allowed":true,"limit_reached":false,"primary":{"used_percent":0,"window_minutes":300,"reset_after_seconds":18000}}}, + "credits":{"has_credits":false,"unlimited":false,"balance":"0"}, + "promo":null + }`)) + if headers == nil { + t.Fatal("observed production event produced no headers") + } + for key, want := range map[string]string{ + "X-Codex-Plan-Type": "pro", + "X-Codex-Primary-Used-Percent": "81", + "X-Codex-Primary-Reset-At": "1787588999", + "X-Codex-Code-Review-Primary-Used-Percent": "4", + "X-Codex-Additional-Gpt-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark", + "X-Codex-Additional-Gpt-5.3-Codex-Spark-Primary-Used-Percent": "0", + "X-Codex-Credits-Balance": "0", + "X-Codex-Credits-Has-Credits": "false", + } { + if got := headers.Get(key); got != want { + t.Fatalf("header %s = %q, want %q", key, got, want) + } + } + // secondary is null upstream and must not be invented. + if got := headers.Get("X-Codex-Secondary-Used-Percent"); got != "" { + t.Fatalf("null secondary window was materialized: %q", got) + } +} + +// An upstream-controlled limit name containing control characters must not reach +// the plain-text request log. +func TestParseCodexQuotaEventHeadersRejectsControlCharactersInLimitName(t *testing.T) { + headers := ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "additional_rate_limits":[{"limit_name":"evil\r\nX-Injected: 1","primary":{"used_percent":3,"window_minutes":300,"reset_after_seconds":60}}] + }`)) + if headers == nil { + t.Fatal("expected window watermarks to survive") + } + for key, values := range headers { + for _, value := range values { + if strings.ContainsAny(value, "\r\n") { + t.Fatalf("header %s carried control characters: %q", key, value) + } + } + } +} + +// Codex quota parsing must stay on the codex-specific websocket entry point. +// The generic entry point is shared with other providers (xAI), whose frames +// must never be reinterpreted as codex quota data: an xAI error frame really +// does carry x-ratelimit-* headers, and merging them here would forge codex +// quota headers into an unrelated provider's request log. +func TestGenericWebsocketEntryPointDoesNotObserveCodexQuota(t *testing.T) { + frame := []byte(`{"type":"error","headers":{"x-ratelimit-remaining-requests":"0","retry-after":"60"}}`) + + genericCtx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(genericCtx, http.Header{"X-Request-Id": []string{"xai-1"}}) + AppendAPIWebsocketResponse(genericCtx, nil, frame) + generic := internallogging.GetResponseHeaders(genericCtx) + if got := generic.Get("X-Ratelimit-Remaining-Requests"); got != "" { + t.Fatalf("generic websocket entry point observed codex quota: %q", got) + } + if got := generic.Get("Retry-After"); got != "" { + t.Fatalf("generic websocket entry point observed retry-after: %q", got) + } + if got := generic.Get("X-Request-Id"); got != "xai-1" { + t.Fatalf("generic websocket entry point disturbed response headers: %q", got) + } + + codexCtx := internallogging.WithResponseHeadersHolder(context.Background()) + AppendCodexAPIWebsocketResponse(codexCtx, nil, frame) + if got := internallogging.GetResponseHeaders(codexCtx).Get("X-Ratelimit-Remaining-Requests"); got != "0" { + t.Fatalf("codex websocket entry point lost quota observation: %q", got) + } +} + +func TestMergeWebsocketQuotaHeaders(t *testing.T) { + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{"X-Request-Id": []string{"req-1"}}) + internallogging.MergeResponseHeaders(ctx, ParseCodexQuotaEventHeaders([]byte(`{ + "type":"codex.rate_limits", + "rate_limits":{"primary":{"used_percent":2,"window_minutes":10080,"reset_at":1782951970}}, + "metered_limit_name":"codex_bengalfox" + }`))) + headers := internallogging.GetResponseHeaders(ctx) + if headers.Get("X-Request-Id") != "req-1" || headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" { + t.Fatalf("merged response headers = %#v", headers) + } +} diff --git a/internal/runtime/executor/helps/logging_helpers.go b/internal/runtime/executor/helps/logging_helpers.go index e1fc2c17..da94bf30 100644 --- a/internal/runtime/executor/helps/logging_helpers.go +++ b/internal/runtime/executor/helps/logging_helpers.go @@ -399,6 +399,13 @@ func AppendAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload appendAPIWebsocketTimeline(ginCtx, []byte(builder.String())) } +// AppendCodexAPIWebsocketResponse stores a codex upstream websocket response frame and merges any +// quota event headers carried by the frame into the request log. +func AppendCodexAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) { + logging.MergeResponseHeaders(ctx, ParseCodexQuotaEventHeaders(payload)) + AppendAPIWebsocketResponse(ctx, cfg, payload) +} + // RecordAPIWebsocketError stores an upstream websocket error event in Gin context. func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) { if !requestLogCaptureEnabled(cfg) || err == nil { diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index a4f2f6ac..87e53078 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -60,6 +60,10 @@ type Result struct { Error *Error // Options carries execution request options (headers, metadata, etc.) for result tracking. Options cliproxyexecutor.Options + // SkipQuotaObservation reports that this result must not replace the last + // observed watermark. Count-tokens requests reuse the credential but are not + // generation traffic; their response headers are not a generation snapshot. + SkipQuotaObservation bool } // Selector chooses an auth candidate for execution. diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 815ce70e..0263e08d 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -15,6 +15,7 @@ import ( "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -355,7 +356,8 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti auth.Unavailable = true auth.Status = StatusError auth.NextRetryAfter = record.NextRetryAfter - auth.Quota = quota + applyCooldownFields(&auth.Quota, quota) + auth.Quota = mergeQuotaObservation(auth.Quota, quota) auth.UpdatedAt = updatedAt if reason != "" { auth.StatusMessage = reason @@ -386,7 +388,7 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { if auth.Unavailable || !auth.NextRetryAfter.IsZero() || auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.IsZero() { auth.Unavailable = false auth.NextRetryAfter = time.Time{} - auth.Quota = QuotaState{} + applyCooldownFields(&auth.Quota, QuotaState{}) auth.UpdatedAt = now changed = true } @@ -397,7 +399,7 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() { state.Unavailable = false state.NextRetryAfter = time.Time{} - state.Quota = QuotaState{} + applyCooldownFields(&state.Quota, QuotaState{}) state.UpdatedAt = now changed = true } @@ -657,7 +659,7 @@ func authCooldownStateRecord(auth *Auth, now time.Time) (CooldownStateRecord, bo Status: "cooling", NextRetryAfter: auth.NextRetryAfter, Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError), - Quota: auth.Quota, + Quota: cooldownFieldsOf(auth.Quota), LastError: cloneError(auth.LastError), UpdatedAt: auth.UpdatedAt, }, true @@ -676,7 +678,7 @@ func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now t Status: "cooling", NextRetryAfter: state.NextRetryAfter, Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError), - Quota: state.Quota, + Quota: cooldownFieldsOf(state.Quota), LastError: cloneError(state.LastError), UpdatedAt: state.UpdatedAt, }, true @@ -718,6 +720,8 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { now := time.Now() + responseHeaders := internallogging.GetResponseHeaders(ctx) + modelState := existingModelState(auth, modelKey) var cooldownRecordsBefore []CooldownStateRecord trackCooldownState := m.cooldownStore != nil if trackCooldownState { @@ -735,6 +739,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { // Retain active credential-scoped cooldown } else if modelKey != "" { state := ensureModelState(auth, modelKey) + modelState = state resetModelState(state, now) updateAggregatedAvailability(auth, now) if !hasModelError(auth, now) { @@ -756,6 +761,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { disableCooling = false } state := ensureModelState(auth, modelKey) + modelState = state state.Unavailable = true state.Status = StatusError state.UpdatedAt = now @@ -779,12 +785,12 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if auth.LastError != nil { auth.StatusMessage = "cloudflare challenge" } - state.Quota = QuotaState{ + applyCooldownFields(&state.Quota, QuotaState{ Exceeded: true, Reason: "cloudflare challenge", NextRecoverAt: next, BackoffLevel: backoffLevel, - } + }) } else if isInvalidGrantResultError(result.Error) { if disableCooling { state.NextRetryAfter = time.Time{} @@ -836,12 +842,12 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } state.NextRetryAfter = next - state.Quota = QuotaState{ + applyCooldownFields(&state.Quota, QuotaState{ Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: backoffLevel, - } + }) if !disableCooling { suspendReason = "quota" shouldSuspendModel = true @@ -857,12 +863,12 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { otherNext = otherState.Quota.NextRecoverAt } otherState.NextRetryAfter = otherNext - otherState.Quota = QuotaState{ + applyCooldownFields(&otherState.Quota, QuotaState{ Exceeded: true, Reason: "credential_quota", NextRecoverAt: otherNext, BackoffLevel: backoffLevel, - } + }) } } auth.Unavailable = true @@ -905,6 +911,13 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } + if !result.SkipQuotaObservation { + auth.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) + if modelState != nil { + modelState.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) + } + } + _ = m.persist(ctx, auth) authSnapshot = auth.Clone() if trackCooldownState { @@ -993,6 +1006,14 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re m.publishErrorEvent(result, authSnapshot) } +func existingModelState(auth *Auth, model string) *ModelState { + model = canonicalModelKey(model) + if auth == nil || model == "" { + return nil + } + return auth.ModelStates[model] +} + func ensureModelState(auth *Auth, model string) *ModelState { model = canonicalModelKey(model) if auth == nil || model == "" { @@ -1065,6 +1086,8 @@ func mergeModelState(target, source *ModelState) *ModelState { }, UpdatedAt: target.UpdatedAt, } + merged.Quota = mergeQuotaObservation(merged.Quota, fallback.Quota) + merged.Quota = mergeQuotaObservation(merged.Quota, preferred.Quota) if source.NextRetryAfter.After(merged.NextRetryAfter) { merged.NextRetryAfter = source.NextRetryAfter } @@ -1104,7 +1127,7 @@ func resetModelState(state *ModelState, now time.Time) { state.StatusMessage = "" state.NextRetryAfter = time.Time{} state.LastError = nil - state.Quota = QuotaState{} + applyCooldownFields(&state.Quota, QuotaState{}) state.UpdatedAt = now } @@ -1210,7 +1233,7 @@ func clearAggregatedAvailability(auth *Auth) { } auth.Unavailable = false auth.NextRetryAfter = time.Time{} - auth.Quota = QuotaState{} + applyCooldownFields(&auth.Quota, QuotaState{}) } func hasModelError(auth *Auth, now time.Time) bool { @@ -1894,12 +1917,12 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati if isCloudflareChallengeResultError(resultErr) { auth.StatusMessage = "cloudflare challenge" next, backoffLevel := nextCloudflareCooldown(auth.Quota.BackoffLevel, disableCooling, now) - auth.Quota = QuotaState{ + applyCooldownFields(&auth.Quota, QuotaState{ Exceeded: true, Reason: "cloudflare challenge", NextRecoverAt: next, BackoffLevel: backoffLevel, - } + }) auth.NextRetryAfter = next return } diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index c98508ab..c7dcebe2 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -23,6 +23,10 @@ import ( log "github.com/sirupsen/logrus" ) +func newUpstreamAttemptContext(ctx context.Context) context.Context { + return logging.WithFreshResponseHeadersHolder(ctx) +} + func claudeOAuthRequestCancellation(ctx context.Context, auth *Auth, err error) error { if auth == nil || !strings.EqualFold(strings.TrimSpace(auth.Provider), "claude") || !strings.EqualFold(strings.TrimSpace(auth.Attributes["auth_kind"]), "oauth") { return nil @@ -349,6 +353,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) if len(models) == 0 { @@ -369,6 +374,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req var authErr error didRefreshOnUnauthorized := false for _, upstreamModel := range models { + execCtx = newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -391,9 +397,11 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + refreshCtx := newUpstreamAttemptContext(execCtx) + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { auth = refreshed didRefreshOnUnauthorized = true + execCtx = newUpstreamAttemptContext(execCtx) startRetry := time.Now() resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts) durationRetry := time.Since(startRetry) @@ -523,6 +531,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) if len(models) == 0 { @@ -535,7 +544,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts, SkipQuotaObservation: true} m.MarkResult(execCtx, result) lastErr = errPrepare continue @@ -543,6 +552,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, var authErr error didRefreshOnUnauthorized := false for _, upstreamModel := range models { + execCtx = newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -565,9 +575,11 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + refreshCtx := newUpstreamAttemptContext(execCtx) + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { auth = refreshed didRefreshOnUnauthorized = true + execCtx = newUpstreamAttemptContext(execCtx) startRetry := time.Now() resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts) durationRetry := time.Since(startRetry) @@ -584,7 +596,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil, Options: execOpts, SkipQuotaObservation: true} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -809,6 +821,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } // Enrich before auth preparation so prepare-stage usage records observe the client request. execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) if selection != nil && aliasResult.ForceMapping && responseAlias != "" { aliasResult.OriginalAlias = responseAlias diff --git a/sdk/cliproxy/auth/conductor_execution_quota_test.go b/sdk/cliproxy/auth/conductor_execution_quota_test.go new file mode 100644 index 00000000..9edbf124 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_execution_quota_test.go @@ -0,0 +1,133 @@ +package auth + +import ( + "context" + "net/http" + "strings" + "testing" + + internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type quotaAttemptIsolationSelector struct{} + +func (quotaAttemptIsolationSelector) Pick(_ context.Context, _, _ string, _ cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + var selected *Auth + for _, auth := range auths { + if auth != nil && (selected == nil || auth.ID < selected.ID) { + selected = auth + } + } + return selected, nil +} + +type quotaAttemptIsolationExecutor struct{} + +func (*quotaAttemptIsolationExecutor) Identifier() string { return "codex" } + +func (*quotaAttemptIsolationExecutor) ShouldPrepareRequestAuth(auth *Auth) bool { + return auth != nil && strings.HasSuffix(auth.ID, "-b") +} + +func (*quotaAttemptIsolationExecutor) PrepareRequestAuth(context.Context, *Auth) (*Auth, error) { + return nil, &Error{HTTPStatus: http.StatusInternalServerError, Message: "prepare failed before upstream response"} +} + +func (*quotaAttemptIsolationExecutor) Execute(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + setQuotaAttemptIsolationHeaders(ctx) + return cliproxyexecutor.Response{}, quotaAttemptIsolationError() +} + +func (*quotaAttemptIsolationExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + setQuotaAttemptIsolationHeaders(ctx) + return nil, quotaAttemptIsolationError() +} + +func (*quotaAttemptIsolationExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (*quotaAttemptIsolationExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func (*quotaAttemptIsolationExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "not implemented"} +} + +func setQuotaAttemptIsolationHeaders(ctx context.Context) { + internallogging.SetResponseHeaders(ctx, http.Header{ + "X-Codex-Plan-Type": []string{"pro"}, + "X-Codex-Primary-Used-Percent": []string{"91"}, + "X-Codex-Primary-Window-Minutes": []string{"10080"}, + "X-Codex-Primary-Reset-After-Seconds": []string{"3600"}, + }) +} + +func quotaAttemptIsolationError() error { + return &Error{HTTPStatus: http.StatusInternalServerError, Message: "first upstream attempt failed"} +} + +func TestExecutionAttemptsDoNotReuseQuotaResponseHeaders(t *testing.T) { + tests := []struct { + name string + run func(*Manager, context.Context, string) error + }{ + { + name: "non-stream", + run: func(manager *Manager, ctx context.Context, model string) error { + _, errExecute := manager.Execute(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + run: func(manager *Manager, ctx context.Context, model string) error { + _, errExecute := manager.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewManager(nil, quotaAttemptIsolationSelector{}, nil) + manager.RegisterExecutor("aAttemptIsolationExecutor{}) + model := "gpt-quota-attempt-isolation-" + test.name + firstID := "quota-attempt-" + test.name + "-a" + secondID := "quota-attempt-" + test.name + "-b" + for _, id := range []string{firstID, secondID} { + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: id, + Provider: "codex", + Status: StatusActive, + }); errRegister != nil { + t.Fatalf("Register(%s) error = %v", id, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(id, "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) }) + } + + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + if errExecute := test.run(manager, ctx, model); errExecute == nil { + t.Fatal("execution error = nil, want terminal prepare error") + } + + first, okFirst := manager.GetByID(firstID) + second, okSecond := manager.GetByID(secondID) + if !okFirst || first == nil || !okSecond || second == nil { + t.Fatalf("auth lookup failed: first=%#v second=%#v", first, second) + } + if got := first.Quota.Signals["X-Codex-Primary-Used-Percent"]; got != "91" { + t.Fatalf("first attempt observation = %q, want 91; quota=%#v", got, first.Quota) + } + if len(second.Quota.Signals) != 0 || !second.Quota.ObservedAt.IsZero() { + t.Fatalf("pre-response failure inherited earlier attempt headers: %#v", second.Quota) + } + if headers := internallogging.GetResponseHeaders(ctx); len(headers) != 0 { + t.Fatalf("attempt headers leaked into request holder: %#v", headers) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index d965712a..9be8f52f 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -122,6 +122,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } // Enrich before auth preparation so prepare-stage usage records observe the client request. execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = newUpstreamAttemptContext(execCtx) if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) @@ -156,6 +157,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } didRefreshOnUnauthorized := false for _, upstreamModel := range models { + execCtx = newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel @@ -222,7 +224,8 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } } if errExecute != nil { - if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(execCtx, selection.Executor, refreshAuth, errExecute, didRefreshOnUnauthorized, true); errRefresh != nil { + refreshCtx := newUpstreamAttemptContext(execCtx) + if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, selection.Executor, refreshAuth, errExecute, didRefreshOnUnauthorized, true); errRefresh != nil { errExecute = errRefresh warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute) } else if okRefresh { @@ -231,6 +234,11 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c didRefreshOnUnauthorized = true publishSelectedAuthMetadata(opts.Metadata, preparedAuth) setEffectiveAuth(preparedAuth) + execCtx = newUpstreamAttemptContext(execCtx) + executorCtx = execCtx + if countTokens { + executorCtx = withAccessTokenFingerprintObserver(execCtx, setEffectiveAuth) + } startHomeRetry := time.Now() response, errExecute = execute() durationHomeRetry := time.Since(startHomeRetry) diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 6efb2bd6..cca3c3da 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -216,6 +216,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi _, didRefreshOnUnauthorized = unauthorizedRefreshTried[auth.ID] } for idx, execModel := range execModels { + ctx = newUpstreamAttemptContext(ctx) resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel @@ -245,7 +246,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if allowRetry { alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(errStream) - refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, errStream, alreadyTried, ephemeralResult) + refreshCtx := newUpstreamAttemptContext(ctx) + refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, executor, auth, errStream, alreadyTried, ephemeralResult) if willAttemptHomeRefresh { didRefreshOnUnauthorized = true if unauthorizedRefreshTried != nil { @@ -260,6 +262,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true + ctx = newUpstreamAttemptContext(ctx) startRetry := time.Now() streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) durationRetry := time.Since(startRetry) @@ -321,7 +324,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if allowRetry { alreadyTried := didRefreshOnUnauthorized willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(bootstrapErr) - refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, executor, auth, bootstrapErr, alreadyTried, ephemeralResult) + refreshCtx := newUpstreamAttemptContext(ctx) + refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, executor, auth, bootstrapErr, alreadyTried, ephemeralResult) if willAttemptHomeRefresh { didRefreshOnUnauthorized = true if unauthorizedRefreshTried != nil { @@ -339,6 +343,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true + ctx = newUpstreamAttemptContext(ctx) startRetry := time.Now() retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) retryStream, retryErr = validateStreamResult(retryStream, retryErr) diff --git a/sdk/cliproxy/auth/quota_signals.go b/sdk/cliproxy/auth/quota_signals.go new file mode 100644 index 00000000..6928c842 --- /dev/null +++ b/sdk/cliproxy/auth/quota_signals.go @@ -0,0 +1,227 @@ +package auth + +import ( + "net/http" + "sort" + "strings" + "time" +) + +const ( + maxQuotaSignalHeaders = 64 + maxQuotaSignalValue = 512 +) + +// ProviderSupportsQuotaObservation reports whether the named provider emits a +// passive credential-level quota snapshot understood by collectQuotaSignals. +func ProviderSupportsQuotaObservation(provider string) bool { + switch strings.ToLower(strings.TrimSpace(provider)) { + case "claude", "codex": + return true + default: + return false + } +} + +// ObserveResponseHeadersForProvider replaces the passive quota snapshot with the +// signals carried by the current upstream response. +// +// The snapshot is replaced rather than merged: a watermark such as Retry-After +// or a "limit reached" flag only appears on the response that produced it, so +// accumulating signals across responses would leave an expired value visible +// indefinitely. Responses that carry no quota signal at all (transport +// failures, 5xx, unrelated endpoints) leave the previous snapshot untouched. +// +// This function only ever touches ObservedAt and Signals. Cooldown and +// scheduling fields are never read or written here. +func (q *QuotaState) ObserveResponseHeadersForProvider(provider string, headers http.Header, observedAt time.Time) bool { + if q == nil { + return false + } + if !ProviderSupportsQuotaObservation(provider) { + return q.ClearObservationSignals() + } + next := collectQuotaSignals(provider, headers) + if len(next) == 0 { + return false + } + if observedAt.IsZero() { + observedAt = time.Now() + } + q.Signals = next + q.ObservedAt = observedAt + return true +} + +// ClearObservationSignals removes only passive observation data. It leaves +// cooldown and scheduler state untouched. +func (q *QuotaState) ClearObservationSignals() bool { + if q == nil || (len(q.Signals) == 0 && q.ObservedAt.IsZero()) { + return false + } + q.Signals = nil + q.ObservedAt = time.Time{} + return true +} + +// cooldownFieldsOf copies only scheduler cooldown fields. Observation data is +// omitted so cooldown persistence and restore cannot replace a live snapshot. +func cooldownFieldsOf(q QuotaState) QuotaState { + return QuotaState{ + Exceeded: q.Exceeded, + Reason: q.Reason, + NextRecoverAt: q.NextRecoverAt, + BackoffLevel: q.BackoffLevel, + } +} + +// applyCooldownFields writes only scheduler cooldown fields. ObservedAt and +// Signals are left untouched so a cooldown transition cannot erase the last +// upstream watermark. +func applyCooldownFields(dst *QuotaState, cooldown QuotaState) { + if dst == nil { + return + } + dst.Exceeded = cooldown.Exceeded + dst.Reason = cooldown.Reason + dst.NextRecoverAt = cooldown.NextRecoverAt + dst.BackoffLevel = cooldown.BackoffLevel +} + +// collectQuotaSignals builds the bounded snapshot for a single response. +func collectQuotaSignals(provider string, headers http.Header) map[string]string { + if len(headers) == 0 { + return nil + } + names := make([]string, 0, len(headers)) + values := make(map[string]string, len(headers)) + for key, headerValues := range headers { + canonicalKey := http.CanonicalHeaderKey(strings.TrimSpace(key)) + if !isQuotaSignalHeaderForProvider(provider, canonicalKey) || len(headerValues) == 0 { + continue + } + value := strings.TrimSpace(headerValues[len(headerValues)-1]) + if !validQuotaSignalValue(value) { + continue + } + if _, exists := values[canonicalKey]; !exists { + names = append(names, canonicalKey) + } + values[canonicalKey] = value + } + if len(names) == 0 { + return nil + } + // Rank then sort so truncation is deterministic and keeps credential-level + // watermarks (plan, credits, primary) ahead of additional-limit namespaces. + sort.Slice(names, func(i, j int) bool { + ri, rj := quotaSignalRetentionRank(names[i]), quotaSignalRetentionRank(names[j]) + if ri != rj { + return ri < rj + } + return names[i] < names[j] + }) + if len(names) > maxQuotaSignalHeaders { + names = names[:maxQuotaSignalHeaders] + } + signals := make(map[string]string, len(names)) + for _, name := range names { + signals[name] = values[name] + } + return signals +} + +// validQuotaSignalValue rejects empty, oversized, and control-character values. +// Observed values reach the plain-text upstream request log, so a value +// containing CR/LF could otherwise forge a header line there. +func validQuotaSignalValue(value string) bool { + if value == "" || len(value) > maxQuotaSignalValue { + return false + } + for _, char := range value { + if char < 0x20 || char == 0x7f { + return false + } + } + return true +} + +func quotaSignalRetentionRank(name string) int { + lower := strings.ToLower(strings.TrimSpace(name)) + switch { + case lower == "retry-after", strings.HasPrefix(lower, "anthropic-ratelimit-unified-"): + return 0 + case lower == "x-codex-plan-type", lower == "x-codex-active-limit", strings.HasPrefix(lower, "x-codex-credits-"): + return 1 + case lower == "x-codex-allowed", lower == "x-codex-limit-reached", + strings.HasPrefix(lower, "x-codex-primary-"), strings.HasPrefix(lower, "x-codex-secondary-"): + return 2 + case strings.HasPrefix(lower, "x-codex-code-review-"): + return 3 + case strings.HasPrefix(lower, "x-codex-additional-"): + return 5 + case strings.HasPrefix(lower, "x-codex-"): + return 4 + default: + return 6 + } +} + +func isQuotaSignalHeaderForProvider(provider, name string) bool { + provider = strings.ToLower(strings.TrimSpace(provider)) + name = strings.ToLower(strings.TrimSpace(name)) + if name == "retry-after" { + return provider == "claude" || provider == "codex" + } + if strings.HasPrefix(name, "anthropic-ratelimit-unified-") { + return provider == "claude" + } + if strings.HasPrefix(name, "x-ratelimit-") { + // Observed Codex responses do not carry x-ratelimit-* headers; the only + // upstream seen emitting them is Grok, which is excluded from quota + // observation. The rule is kept so a future Codex rollout is captured + // without another change, but it is expected to be inert today. + return provider == "codex" + } + if !strings.HasPrefix(name, "x-codex-") { + return false + } + if provider != "codex" { + return false + } + if name == "x-codex-active-limit" || name == "x-codex-plan-type" || + strings.HasPrefix(name, "x-codex-credits-") { + return true + } + // Codex namespaces each additional limit by a short name on the HTTP path + // (x-codex-bengalfox-primary-used-percent) and by limit name on the + // websocket path (x-codex-additional--primary-used-percent), so the + // suffix is matched instead of an exhaustive header list. + for _, marker := range []string{ + "-allowed", + "-limit-reached", + "-limit-name", + "-used-percent", + "-window-minutes", + "-reset-after-seconds", + "-reset-at", + "-over-secondary-limit-percent", + } { + if strings.Contains(name, marker) { + return true + } + } + return false +} + +// mergeQuotaObservation keeps the newest observation snapshot instead of +// unioning signals captured at different times, so merging an older snapshot +// can never resurrect a stale watermark. +func mergeQuotaObservation(target, source QuotaState) QuotaState { + if source.ObservedAt.IsZero() || source.ObservedAt.Before(target.ObservedAt) { + return target + } + target.ObservedAt = source.ObservedAt + target.Signals = source.Clone().Signals + return target +} diff --git a/sdk/cliproxy/auth/quota_signals_test.go b/sdk/cliproxy/auth/quota_signals_test.go new file mode 100644 index 00000000..239d1009 --- /dev/null +++ b/sdk/cliproxy/auth/quota_signals_test.go @@ -0,0 +1,685 @@ +package auth + +import ( + "context" + "fmt" + "net/http" + "reflect" + "strconv" + "strings" + "testing" + "time" + + internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" +) + +func TestQuotaStateObserveResponseHeadersKeepsProviderScopedSignals(t *testing.T) { + observedAt := time.Unix(123, 0) + var quota QuotaState + if !quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "X-Codex-Active-Limit": []string{"codex_bengalfox"}, + "X-Codex-Primary-Used-Percent": []string{"2"}, + "X-Codex-Turn-State": []string{"opaque-state"}, + "X-Codex-Safety-Buffering-Enabled": []string{"true"}, + "Retry-After": []string{"120"}, + "Authorization": []string{"Bearer secret"}, + }, observedAt) { + t.Fatal("ObserveResponseHeadersForProvider() reported no change") + } + if !quota.ObservedAt.Equal(observedAt) { + t.Fatalf("ObservedAt = %v, want %v", quota.ObservedAt, observedAt) + } + if quota.Signals["X-Codex-Active-Limit"] != "codex_bengalfox" || quota.Signals["Retry-After"] != "120" { + t.Fatalf("quota signals = %#v", quota.Signals) + } + if _, ok := quota.Signals["Authorization"]; ok { + t.Fatal("authorization header was retained as a quota signal") + } + if _, ok := quota.Signals["X-Codex-Turn-State"]; ok { + t.Fatal("non-quota Codex response header was retained as a quota signal") + } + if _, ok := quota.Signals["X-Codex-Safety-Buffering-Enabled"]; ok { + t.Fatal("Codex safety-buffering metadata was retained as a quota signal") + } +} + +func TestQuotaStateObserveResponseHeadersBoundsAndCanonicalizesValues(t *testing.T) { + var quota QuotaState + longValue := strings.Repeat("x", maxQuotaSignalValue+1) + if quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "X-Codex-Empty": []string{""}, + "X-Codex-Long": []string{longValue}, + }, time.Unix(123, 0)) { + t.Fatal("invalid-only headers reported a change") + } + if quota.Signals != nil || !quota.ObservedAt.IsZero() { + t.Fatalf("invalid signals were retained: %#v", quota) + } + if !quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "x-codex-plan-type": []string{"pro"}, + }, time.Unix(124, 0)) { + t.Fatal("canonical valid header reported no change") + } + if quota.Signals["X-Codex-Plan-Type"] != "pro" { + t.Fatalf("canonical signal = %#v", quota.Signals) + } +} + +func TestQuotaStateObserveResponseHeadersRetainsMeasuredClaudeAndCodexWatermarks(t *testing.T) { + var claudeQuota QuotaState + if !claudeQuota.ObserveResponseHeadersForProvider("claude", http.Header{ + "Anthropic-Ratelimit-Unified-5h-Status": []string{"allowed"}, + "Anthropic-Ratelimit-Unified-5h-Utilization": []string{"0.0"}, + "Anthropic-Ratelimit-Unified-5h-Reset": []string{"1787296800"}, + "Anthropic-Ratelimit-Unified-7d-Status": []string{"allowed"}, + "Anthropic-Ratelimit-Unified-7d-Utilization": []string{"0.53"}, + "Anthropic-Ratelimit-Unified-7d-Reset": []string{"1787695200"}, + "Anthropic-Ratelimit-Unified-Fallback-Percentage": []string{"0.5"}, + "Anthropic-Ratelimit-Unified-Overage-Disabled-Reason": []string{"member_zero_credit_limit"}, + "Anthropic-Ratelimit-Unified-Overage-Status": []string{"rejected"}, + "Anthropic-Ratelimit-Unified-Representative-Claim": []string{"five_hour"}, + "Anthropic-Ratelimit-Unified-Reset": []string{"1787296800"}, + "Anthropic-Ratelimit-Unified-Status": []string{"allowed"}, + "Anthropic-Workspace-Id": []string{"workspace-must-not-be-quota-signal"}, + }, time.Unix(1787279282, 0)) { + t.Fatal("Claude observation reported no change") + } + for key, want := range map[string]string{ + "Anthropic-Ratelimit-Unified-5h-Reset": "1787296800", + "Anthropic-Ratelimit-Unified-5h-Status": "allowed", + "Anthropic-Ratelimit-Unified-5h-Utilization": "0.0", + "Anthropic-Ratelimit-Unified-7d-Reset": "1787695200", + "Anthropic-Ratelimit-Unified-7d-Status": "allowed", + "Anthropic-Ratelimit-Unified-7d-Utilization": "0.53", + "Anthropic-Ratelimit-Unified-Fallback-Percentage": "0.5", + "Anthropic-Ratelimit-Unified-Overage-Disabled-Reason": "member_zero_credit_limit", + "Anthropic-Ratelimit-Unified-Overage-Status": "rejected", + "Anthropic-Ratelimit-Unified-Representative-Claim": "five_hour", + "Anthropic-Ratelimit-Unified-Reset": "1787296800", + "Anthropic-Ratelimit-Unified-Status": "allowed", + } { + if got := claudeQuota.Signals[key]; got != want { + t.Fatalf("Claude quota signal %s = %q, want %q", key, got, want) + } + } + if _, ok := claudeQuota.Signals["Anthropic-Workspace-Id"]; ok { + t.Fatal("workspace identity header was retained as a quota signal") + } + + var codexQuota QuotaState + if !codexQuota.ObserveResponseHeadersForProvider("codex", http.Header{ + "X-Codex-Plan-Type": []string{"pro"}, + "X-Codex-Primary-Used-Percent": []string{"51"}, + "X-Codex-Primary-Window-Minutes": []string{"10080"}, + "X-Codex-Primary-Reset-After-Seconds": []string{"309718"}, + "X-Codex-Primary-Reset-At": []string{"1787588999"}, + "X-Codex-Bengalfox-Limit-Name": []string{"GPT-5.3-Codex-Spark"}, + "X-Codex-Bengalfox-Secondary-Used-Percent": []string{"35"}, + "X-Codex-Credits-Has-Credits": []string{"False"}, + }, time.Unix(1787279282, 0)) { + t.Fatal("Codex observation reported no change") + } + for key, want := range map[string]string{ + "X-Codex-Plan-Type": "pro", + "X-Codex-Primary-Used-Percent": "51", + "X-Codex-Bengalfox-Limit-Name": "GPT-5.3-Codex-Spark", + "X-Codex-Bengalfox-Secondary-Used-Percent": "35", + "X-Codex-Credits-Has-Credits": "False", + } { + if got := codexQuota.Signals[key]; got != want { + t.Fatalf("Codex quota signal %s = %q, want %q", key, got, want) + } + } +} + +func TestQuotaStateObserveResponseHeadersDropsKimiGrokAndAntigravitySignals(t *testing.T) { + for _, provider := range []string{"kimi", "xai", "grok", "antigravity", "gemini", "vertex", "aistudio"} { + quota := QuotaState{ + ObservedAt: time.Unix(1786082736, 0), + Signals: map[string]string{"old": "value"}, + } + changed := quota.ObserveResponseHeadersForProvider(provider, http.Header{ + "X-Ratelimit-Remaining-Requests": []string{"0"}, + "X-Ratelimit-Remaining-Tokens": []string{"0"}, + "Retry-After": []string{"60"}, + "Anthropic-Ratelimit-Unified-5h-Utilization": []string{"0.5"}, + }, time.Now()) + if !changed || !quota.ObservedAt.IsZero() || len(quota.Signals) != 0 { + t.Fatalf("provider %s retained observation signals: changed=%v quota=%#v", provider, changed, quota) + } + } +} + +func TestCooldownEqualityIgnoresObservationSignals(t *testing.T) { + base := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: time.Unix(1787588999, 0), + BackoffLevel: 2, + } + observed := base.Clone() + observed.ObservedAt = time.Unix(1787279282, 0) + observed.Signals = map[string]string{ + "X-Codex-Primary-Used-Percent": "51", + "X-Codex-Primary-Reset-At": "1787588999", + } + if !cooldownQuotaEqual(base, observed) { + t.Fatal("observation-only quota signals changed cooldown equality") + } +} + +func TestManagerMarkResultRecordsResponseQuotaSignalsInMemory(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth, errRegister := manager.Register(context.Background(), &Auth{ + ID: "quota-signal-auth", + Provider: "codex", + }) + if errRegister != nil || auth == nil { + t.Fatalf("Register() auth=%#v err=%v", auth, errRegister) + } + + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{ + "X-Codex-Active-Limit": []string{"codex_bengalfox"}, + "X-Codex-Primary-Used-Percent": []string{"2"}, + "X-Codex-Primary-Window-Minutes": []string{"10080"}, + "X-Codex-Primary-Reset-At": []string{"1782951970"}, + }) + manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: "codex", + Model: "gpt-5.3-codex", + Success: true, + }) + + updated, ok := manager.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatal("auth not found after MarkResult") + } + if updated.Quota.Signals["X-Codex-Active-Limit"] != "codex_bengalfox" || + updated.Quota.Signals["X-Codex-Primary-Used-Percent"] != "2" { + t.Fatalf("in-memory quota signals = %#v", updated.Quota.Signals) + } +} + +func TestMarkResultCountTokensDoesNotReplaceObservation(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth, errRegister := manager.Register(context.Background(), &Auth{ + ID: "quota-count-tokens-auth", + Provider: "claude", + Quota: QuotaState{ + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"Anthropic-Ratelimit-Unified-Status": "allowed"}, + }, + }) + if errRegister != nil || auth == nil { + t.Fatalf("Register() auth=%#v err=%v", auth, errRegister) + } + + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{ + "Anthropic-Ratelimit-Unified-Status": []string{"rejected"}, + }) + manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: "claude", + Model: "claude-opus-4-6", + Success: true, + SkipQuotaObservation: true, + }) + + updated, ok := manager.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatal("auth not found after MarkResult") + } + if updated.Quota.Signals["Anthropic-Ratelimit-Unified-Status"] != "allowed" || !updated.Quota.ObservedAt.Equal(time.Unix(10, 0)) { + t.Fatalf("count_tokens replaced the last generation snapshot: %#v", updated.Quota) + } +} + +func TestResetModelStatePreservesObservationSignals(t *testing.T) { + state := &ModelState{ + Status: StatusError, + Unavailable: true, + StatusMessage: "quota", + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: time.Unix(20, 0), + BackoffLevel: 2, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Active-Limit": "premium"}, + }, + } + resetModelState(state, time.Unix(30, 0)) + if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 { + t.Fatalf("cooldown state was not reset: %#v", state.Quota) + } + if !state.Quota.ObservedAt.Equal(time.Unix(10, 0)) || state.Quota.Signals["X-Codex-Active-Limit"] != "premium" { + t.Fatalf("observation signals were lost during reset: %#v", state.Quota) + } +} + +func TestMergeModelStateKeepsNewestObservationSnapshot(t *testing.T) { + target := &ModelState{UpdatedAt: time.Unix(20, 0), Quota: QuotaState{ + ObservedAt: time.Unix(20, 0), Signals: map[string]string{"X-Codex-Plan-Type": "pro"}, + }} + source := &ModelState{UpdatedAt: time.Unix(30, 0), Quota: QuotaState{ + ObservedAt: time.Unix(30, 0), Signals: map[string]string{"X-Codex-Active-Limit": "codex_bengalfox"}, + }} + mergeModelState(target, source) + if !target.Quota.ObservedAt.Equal(time.Unix(30, 0)) { + t.Fatalf("merged ObservedAt = %v, want newest snapshot time", target.Quota.ObservedAt) + } + if target.Quota.Signals["X-Codex-Active-Limit"] != "codex_bengalfox" { + t.Fatalf("newest snapshot was lost: %#v", target.Quota.Signals) + } + // Unioning snapshots taken at different times would resurrect the older + // watermark, so the stale key must be gone. + if _, ok := target.Quota.Signals["X-Codex-Plan-Type"]; ok { + t.Fatalf("stale snapshot key survived the merge: %#v", target.Quota.Signals) + } +} + +// A watermark such as Retry-After only appears on the response that produced it. +// Later responses must not keep advertising it. +func TestObserveResponseHeadersReplacesStaleWatermarks(t *testing.T) { + var quota QuotaState + if !quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "Retry-After": []string{"120"}, + "X-Codex-Primary-Used-Percent": []string{"99"}, + }, time.Unix(100, 0)) { + t.Fatal("initial observation reported no change") + } + if quota.Signals["Retry-After"] != "120" { + t.Fatalf("initial snapshot = %#v", quota.Signals) + } + + if !quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "X-Codex-Primary-Used-Percent": []string{"5"}, + }, time.Unix(200, 0)) { + t.Fatal("second observation reported no change") + } + if _, ok := quota.Signals["Retry-After"]; ok { + t.Fatalf("expired Retry-After survived a later response: %#v", quota.Signals) + } + if quota.Signals["X-Codex-Primary-Used-Percent"] != "5" { + t.Fatalf("snapshot not refreshed: %#v", quota.Signals) + } + if !quota.ObservedAt.Equal(time.Unix(200, 0)) { + t.Fatalf("ObservedAt = %v, want the latest observation time", quota.ObservedAt) + } +} + +// Responses without any quota header (transport failures, 5xx) must not erase +// the last known snapshot. +func TestObserveResponseHeadersKeepsSnapshotWhenResponseCarriesNoSignal(t *testing.T) { + quota := QuotaState{ + ObservedAt: time.Unix(100, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "5"}, + } + if quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "Content-Type": []string{"application/json"}, + }, time.Unix(200, 0)) { + t.Fatal("signal-free response reported a change") + } + if quota.Signals["X-Codex-Primary-Used-Percent"] != "5" || !quota.ObservedAt.Equal(time.Unix(100, 0)) { + t.Fatalf("snapshot was disturbed: %#v", quota) + } +} + +// ObservedAt must advance even when the watermark values are unchanged, +// otherwise consumers cannot tell a fresh reading from a stale one. +func TestObserveResponseHeadersAdvancesObservedAtOnRepeatedValues(t *testing.T) { + var quota QuotaState + headers := http.Header{"X-Codex-Primary-Used-Percent": []string{"5"}} + quota.ObserveResponseHeadersForProvider("codex", headers, time.Unix(100, 0)) + quota.ObserveResponseHeadersForProvider("codex", headers, time.Unix(200, 0)) + if !quota.ObservedAt.Equal(time.Unix(200, 0)) { + t.Fatalf("ObservedAt = %v, want 200", quota.ObservedAt) + } +} + +// Observed values reach the plain-text upstream request log, so control +// characters must never be stored. +func TestObserveResponseHeadersRejectsControlCharacterValues(t *testing.T) { + var quota QuotaState + if quota.ObserveResponseHeadersForProvider("codex", http.Header{ + "X-Codex-Bengalfox-Limit-Name": []string{"evil\r\nX-Injected: 1"}, + }, time.Unix(100, 0)) { + t.Fatal("control-character value was accepted") + } + if len(quota.Signals) != 0 { + t.Fatalf("control-character value was stored: %#v", quota.Signals) + } +} + +// Truncation at the header cap must not depend on map iteration order. +func TestObserveResponseHeadersTruncatesDeterministically(t *testing.T) { + headers := make(http.Header, maxQuotaSignalHeaders*2) + for i := 0; i < maxQuotaSignalHeaders*2; i++ { + headers.Set(fmt.Sprintf("X-Codex-L%03d-Primary-Used-Percent", i), strconv.Itoa(i)) + } + var first QuotaState + first.ObserveResponseHeadersForProvider("codex", headers, time.Unix(100, 0)) + if len(first.Signals) != maxQuotaSignalHeaders { + t.Fatalf("snapshot size = %d, want %d", len(first.Signals), maxQuotaSignalHeaders) + } + for attempt := 0; attempt < 5; attempt++ { + var next QuotaState + next.ObserveResponseHeadersForProvider("codex", headers, time.Unix(100, 0)) + if !reflect.DeepEqual(first.Signals, next.Signals) { + t.Fatal("truncated snapshot varied between identical observations") + } + } +} + +func TestProviderSupportsQuotaObservation(t *testing.T) { + for _, provider := range []string{ + "", "kimi", "xai", "grok", "antigravity", "gemini", "gemini-interactions", + "vertex", "aistudio", "openai", "openai-compatibility", "third-party-plugin", + "XAI", " Grok ", " Gemini ", + } { + if ProviderSupportsQuotaObservation(provider) { + t.Fatalf("provider %q unexpectedly supports quota observation", provider) + } + } + for _, provider := range []string{"codex", "claude", "CODEX", " Claude "} { + if !ProviderSupportsQuotaObservation(provider) { + t.Fatalf("provider %q unexpectedly excluded from quota observation", provider) + } + } +} + +func TestQuotaStateCloneCopiesSignals(t *testing.T) { + original := QuotaState{Signals: map[string]string{"X-Codex-Plan-Type": "pro"}} + clone := original.Clone() + clone.Signals["X-Codex-Plan-Type"] = "team" + if original.Signals["X-Codex-Plan-Type"] != "pro" { + t.Fatalf("mutating cloned quota changed original: %#v", original.Signals) + } +} + +func TestApplyCooldownFieldsPreservesObservation(t *testing.T) { + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: time.Unix(20, 0), + BackoffLevel: 1, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "51"}, + } + applyCooldownFields("a, QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: time.Unix(40, 0), + BackoffLevel: 2, + }) + if !quota.Exceeded || quota.Reason != "credential_quota" || quota.BackoffLevel != 2 || !quota.NextRecoverAt.Equal(time.Unix(40, 0)) { + t.Fatalf("cooldown fields were not applied: %#v", quota) + } + if !quota.ObservedAt.Equal(time.Unix(10, 0)) || quota.Signals["X-Codex-Primary-Used-Percent"] != "51" { + t.Fatalf("cooldown overwrote the last observation: %#v", quota) + } +} + +func TestClearCooldownStateForAuthPreservesObservation(t *testing.T) { + auth := &Auth{ + Unavailable: true, + NextRetryAfter: time.Unix(40, 0), + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: time.Unix(40, 0), + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "51"}, + }, + ModelStates: map[string]*ModelState{ + "gpt-5.3-codex": { + Unavailable: true, + NextRetryAfter: time.Unix(40, 0), + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: time.Unix(40, 0), + ObservedAt: time.Unix(11, 0), + Signals: map[string]string{"X-Codex-Plan-Type": "pro"}, + }, + }, + }, + } + if !clearCooldownStateForAuth(auth, time.Unix(50, 0)) { + t.Fatal("clearCooldownStateForAuth reported no change") + } + if auth.Unavailable || auth.Quota.Exceeded || auth.Quota.Reason != "" { + t.Fatalf("cooldown was not cleared: %#v", auth.Quota) + } + if !auth.Quota.ObservedAt.Equal(time.Unix(10, 0)) || auth.Quota.Signals["X-Codex-Primary-Used-Percent"] != "51" { + t.Fatalf("clearing cooldown overwrote auth observation: %#v", auth.Quota) + } + state := auth.ModelStates["gpt-5.3-codex"] + if state.Unavailable || state.Quota.Exceeded { + t.Fatalf("model cooldown was not cleared: %#v", state.Quota) + } + if !state.Quota.ObservedAt.Equal(time.Unix(11, 0)) || state.Quota.Signals["X-Codex-Plan-Type"] != "pro" { + t.Fatalf("clearing cooldown overwrote model observation: %#v", state.Quota) + } +} + +func TestCooldownStateRecordOmitsObservation(t *testing.T) { + now := time.Unix(50, 0) + auth := &Auth{ + ID: "auth-1", + Unavailable: true, + NextRetryAfter: now.Add(time.Hour), + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: now.Add(time.Hour), + BackoffLevel: 2, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "51"}, + }, + } + record, ok := authCooldownStateRecord(auth, now) + if !ok { + t.Fatal("expected a cooldown record") + } + if !record.Quota.Exceeded || record.Quota.Reason != "quota" || record.Quota.BackoffLevel != 2 { + t.Fatalf("cooldown fields missing from record: %#v", record.Quota) + } + if !record.Quota.ObservedAt.IsZero() || len(record.Quota.Signals) != 0 { + t.Fatalf("observation leaked into cooldown persistence: %#v", record.Quota) + } +} + +func TestMarkResultQuotaFailureDoesNotEraseSiblingObservation(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth, errRegister := manager.Register(context.Background(), &Auth{ + ID: "quota-sibling-auth", + Provider: "codex", + ModelStates: map[string]*ModelState{ + "gpt-5.3-codex": { + Status: StatusActive, + Quota: QuotaState{ + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "40"}, + }, + }, + "gpt-5.4": { + Status: StatusActive, + Quota: QuotaState{ + ObservedAt: time.Unix(11, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "41"}, + }, + }, + }, + }) + if errRegister != nil || auth == nil { + t.Fatalf("Register() auth=%#v err=%v", auth, errRegister) + } + + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{ + "Retry-After": []string{"120"}, + "X-Codex-Primary-Used-Percent": []string{"99"}, + }) + manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: "codex", + Model: "gpt-5.3-codex", + Success: false, + CredentialScope: true, + Error: &Error{HTTPStatus: 429, Message: "quota"}, + }) + + updated, ok := manager.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatal("auth not found after MarkResult") + } + current := updated.ModelStates["gpt-5.3-codex"] + if current == nil || !current.Quota.Exceeded || current.Quota.Reason != "quota" { + t.Fatalf("current model cooldown missing: %#v", current) + } + if current.Quota.Signals["X-Codex-Primary-Used-Percent"] != "99" || current.Quota.Signals["Retry-After"] != "120" { + t.Fatalf("current model observation was not refreshed: %#v", current.Quota.Signals) + } + sibling := updated.ModelStates["gpt-5.4"] + if sibling == nil || !sibling.Quota.Exceeded || sibling.Quota.Reason != "credential_quota" { + t.Fatalf("sibling cooldown missing: %#v", sibling) + } + if sibling.Quota.Signals["X-Codex-Primary-Used-Percent"] != "41" { + t.Fatalf("credential-scope cooldown erased sibling observation: %#v", sibling.Quota.Signals) + } + if _, ok := sibling.Quota.Signals["Retry-After"]; ok { + t.Fatalf("sibling observation was replaced by the current request: %#v", sibling.Quota.Signals) + } +} + +func TestMarkResultRetainedCredentialQuotaStillObservesModel(t *testing.T) { + manager := NewManager(nil, nil, nil) + recoverAt := time.Now().Add(time.Hour) + auth, errRegister := manager.Register(context.Background(), &Auth{ + ID: "quota-retain-auth", + Provider: "codex", + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: recoverAt, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "10"}, + }, + ModelStates: map[string]*ModelState{ + "gpt-5.3-codex": { + Status: StatusError, + Unavailable: true, + NextRetryAfter: recoverAt, + Quota: QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: recoverAt, + ObservedAt: time.Unix(10, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "10"}, + }, + }, + }, + }) + if errRegister != nil || auth == nil { + t.Fatalf("Register() auth=%#v err=%v", auth, errRegister) + } + + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, http.Header{ + "X-Codex-Primary-Used-Percent": []string{"20"}, + }) + manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: "codex", + Model: "gpt-5.3-codex", + Success: true, + }) + + updated, ok := manager.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatal("auth not found after MarkResult") + } + if !updated.Quota.Exceeded || updated.Quota.Reason != "credential_quota" { + t.Fatalf("retained cooldown was disturbed: %#v", updated.Quota) + } + if updated.Quota.Signals["X-Codex-Primary-Used-Percent"] != "20" { + t.Fatalf("auth observation was not refreshed: %#v", updated.Quota.Signals) + } + state := updated.ModelStates["gpt-5.3-codex"] + if state == nil || !state.Quota.Exceeded || state.Quota.Reason != "credential_quota" { + t.Fatalf("model cooldown was disturbed: %#v", state) + } + if state.Quota.Signals["X-Codex-Primary-Used-Percent"] != "20" { + t.Fatalf("retained cooldown prevented model observation: %#v", state.Quota.Signals) + } +} + +func TestRestoreCooldownRecordDoesNotOverwriteNewerObservation(t *testing.T) { + nextRetry := time.Now().Add(time.Hour).UTC().Truncate(time.Second) + store := &recordingCooldownStateStore{ + load: []CooldownStateRecord{{ + Provider: "codex", + AuthID: "auth-restore-obs", + Status: "cooling", + NextRetryAfter: nextRetry, + Reason: "quota", + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + ObservedAt: time.Unix(5, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "1"}, + }, + UpdatedAt: nextRetry.Add(-time.Minute), + }}, + } + manager := NewManager(nil, nil, nil) + manager.SetCooldownStateStore(store) + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), &Auth{ + ID: "auth-restore-obs", + Provider: "codex", + Quota: QuotaState{ + ObservedAt: time.Unix(50, 0), + Signals: map[string]string{"X-Codex-Primary-Used-Percent": "77"}, + }, + }); errRegister != nil { + t.Fatalf("Register() returned error: %v", errRegister) + } + if errRestore := manager.RestoreCooldownStates(context.Background()); errRestore != nil { + t.Fatalf("RestoreCooldownStates() returned error: %v", errRestore) + } + auth, ok := manager.GetByID("auth-restore-obs") + if !ok || auth == nil { + t.Fatal("restored auth was not found") + } + if !auth.Unavailable || !auth.Quota.Exceeded || auth.Quota.Reason != "quota" { + t.Fatalf("cooldown was not restored: %#v", auth.Quota) + } + if auth.Quota.Signals["X-Codex-Primary-Used-Percent"] != "77" { + t.Fatalf("restoring cooldown overwrote a newer observation: %#v", auth.Quota.Signals) + } +} + +func TestObserveResponseHeadersKeepsPrimaryWhenTruncatingAdditional(t *testing.T) { + headers := make(http.Header, maxQuotaSignalHeaders+8) + headers.Set("X-Codex-Plan-Type", "pro") + headers.Set("X-Codex-Primary-Used-Percent", "81") + headers.Set("X-Codex-Credits-Balance", "0") + for i := 0; i < maxQuotaSignalHeaders; i++ { + headers.Set(fmt.Sprintf("X-Codex-Additional-L%03d-Primary-Used-Percent", i), strconv.Itoa(i)) + } + var quota QuotaState + quota.ObserveResponseHeadersForProvider("codex", headers, time.Unix(100, 0)) + if quota.Signals["X-Codex-Plan-Type"] != "pro" || + quota.Signals["X-Codex-Primary-Used-Percent"] != "81" || + quota.Signals["X-Codex-Credits-Balance"] != "0" { + t.Fatalf("credential-level watermarks were truncated: %#v", quota.Signals) + } + if len(quota.Signals) != maxQuotaSignalHeaders { + t.Fatalf("snapshot size = %d, want %d", len(quota.Signals), maxQuotaSignalHeaders) + } +} diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 0a9099e0..e8d14c3a 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -174,6 +174,27 @@ type QuotaState struct { NextRecoverAt time.Time `json:"next_recover_at"` // BackoffLevel stores the progressive cooldown exponent used for rate limits. BackoffLevel int `json:"backoff_level,omitempty"` + // ObservedAt is the time the current Signals snapshot was observed. + ObservedAt time.Time `json:"observed_at,omitempty"` + // Signals stores bounded, provider-specific quota watermark values observed + // from upstream response headers or websocket quota events. It is a snapshot + // of one upstream response, not an accumulation across responses, so an + // expired watermark cannot linger after the response that produced it. + // Cooldown transitions must use applyCooldownFields so they cannot replace + // this snapshot. + Signals map[string]string `json:"signals,omitempty"` +} + +// Clone returns an independent copy of the quota state. +func (q QuotaState) Clone() QuotaState { + copyQuota := q + if len(q.Signals) > 0 { + copyQuota.Signals = make(map[string]string, len(q.Signals)) + for key, value := range q.Signals { + copyQuota.Signals[key] = value + } + } + return copyQuota } // ModelState captures the execution state for a specific model under an auth entry. @@ -264,6 +285,7 @@ func (a *Auth) Clone() *Auth { return nil } copyAuth := *a + copyAuth.Quota = a.Quota.Clone() if len(a.Attributes) > 0 { copyAuth.Attributes = make(map[string]string, len(a.Attributes)) for key, value := range a.Attributes { @@ -406,6 +428,7 @@ func (m *ModelState) Clone() *ModelState { return nil } copyState := *m + copyState.Quota = m.Quota.Clone() if m.LastError != nil { copyState.LastError = &Error{ Code: m.LastError.Code, -- 2.51.2