diff --git a/AGENTS.md b/AGENTS.md index 8a40ce91d..81c02bc81 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -56,7 +56,7 @@ go build -o test-output ./cmd/server && rm test-output # Verify compile (REQUIRE - Wrap defer errors: `defer func() { if err := f.Close(); err != nil { log.Errorf(...) } }()` - Use logrus structured logging; avoid leaking secrets/tokens in logs - Avoid panics in HTTP handlers; prefer logged errors and meaningful HTTP status codes -- Timeouts are allowed only during credential acquisition; after an upstream connection is established, do not set timeouts for any subsequent network behavior. Intentional exceptions that must remain allowed are the Codex websocket liveness deadlines in `internal/runtime/executor/codex_websockets_executor.go`, the wsrelay session deadlines in `internal/wsrelay/session.go`, the management APICall timeout in `internal/api/handlers/management/api_tools.go`, and the `cmd/fetch_antigravity_models` utility timeouts +- Timeouts are allowed only during credential acquisition; after an upstream connection is established, do not set timeouts for any subsequent network behavior. Intentional exceptions that must remain allowed are the Codex websocket liveness deadlines in `internal/runtime/executor/codex_websockets_executor.go`, the wsrelay session deadlines in `internal/wsrelay/session.go`, the management APICall timeout in `internal/api/handlers/management/api_tools.go`, the quota usage read timeout in `sdk/cliproxy/service_quota.go`, and the `cmd/fetch_antigravity_models` utility timeouts - Avoid wall-clock `time.Sleep` in TTL, expiration, ordering, or cache-eviction unit tests due to platform timer granularity (e.g. Windows default timer resolution of ~15.6ms) and CI jitter under load; prefer controllable clocks (`nowFunc` / mock clock), explicit timestamp manipulation, or deterministic synchronization primitives. - Note: if modifying features that involve CLIProxyAPIHome, check if corresponding updates are needed in the CLIProxyAPIHome repository. - Endpoints under the `/v0/management` base URL are deprecated and no longer maintained. For any feature changes, do not modify endpoints under `/v0/management` unless necessary to fix compilation errors. diff --git a/internal/runtime/executor/helps/quota_usage.go b/internal/runtime/executor/helps/quota_usage.go new file mode 100644 index 000000000..30753d6ad --- /dev/null +++ b/internal/runtime/executor/helps/quota_usage.go @@ -0,0 +1,69 @@ +package helps + +import ( + "net/http" + "strconv" + "strings" + "time" + + "github.com/tidwall/gjson" +) + +// ParseCodexUsageHeaders converts a /backend-api/wham/usage response into the +// X-Codex primary and secondary window headers that Codex responses carry. +// Each non-null window's fields are copied as given, like the proxied-response +// path, because a usage read replaces the whole observation and a dropped +// window would erase its exhaustion. It returns nil when no window is present, +// so a snapshot holding only the limit-reached flag never erases the windows. +func ParseCodexUsageHeaders(body []byte) http.Header { + rateLimit := gjson.GetBytes(body, "rate_limit") + headers := make(http.Header) + for _, name := range []string{"Primary", "Secondary"} { + window := rateLimit.Get(strings.ToLower(name) + "_window") + if !window.IsObject() { + continue + } + prefix := "X-Codex-" + name + "-" + if used := window.Get("used_percent"); used.Type == gjson.Number { + headers.Set(prefix+"Used-Percent", used.Raw) + } + if seconds := window.Get("limit_window_seconds"); seconds.Type == gjson.Number { + headers.Set(prefix+"Window-Minutes", strconv.FormatInt(seconds.Int()/60, 10)) + } + if resetAfter := window.Get("reset_after_seconds"); resetAfter.Type == gjson.Number { + headers.Set(prefix+"Reset-After-Seconds", resetAfter.Raw) + } + if resetAt := window.Get("reset_at"); resetAt.Type == gjson.Number { + headers.Set(prefix+"Reset-At", resetAt.Raw) + } + } + if len(headers) == 0 { + return nil + } + setCodexQuotaScalarHeaderFromResult(headers, "X-Codex-Limit-Reached", rateLimit, "limit_reached") + return headers +} + +// ParseClaudeUsageHeaders converts an /api/oauth/usage response into the +// Anthropic unified 5h and 7d headers that Claude responses carry. Utilization +// arrives as a percent and leaves as a fraction; resets_at arrives as RFC 3339 +// and leaves as Unix seconds. It returns nil when the response holds no usable +// window. +func ParseClaudeUsageHeaders(body []byte) http.Header { + headers := make(http.Header) + for name, path := range map[string]string{"5h": "five_hour", "7d": "seven_day"} { + window := gjson.GetBytes(body, path) + utilization := window.Get("utilization") + resetAt, errParse := time.Parse(time.RFC3339, window.Get("resets_at").String()) + if utilization.Type != gjson.Number || errParse != nil { + continue + } + prefix := "Anthropic-Ratelimit-Unified-" + name + "-" + headers.Set(prefix+"Utilization", strconv.FormatFloat(utilization.Float()/100, 'f', -1, 64)) + headers.Set(prefix+"Reset", strconv.FormatInt(resetAt.Unix(), 10)) + } + if len(headers) == 0 { + return nil + } + return headers +} diff --git a/internal/runtime/executor/helps/quota_usage_test.go b/internal/runtime/executor/helps/quota_usage_test.go new file mode 100644 index 000000000..2cf12b9aa --- /dev/null +++ b/internal/runtime/executor/helps/quota_usage_test.go @@ -0,0 +1,65 @@ +package helps + +import ( + "net/http" + "reflect" + "testing" +) + +// Recorded usage responses with identities removed. +const ( + recordedCodexUsage = `{"plan_type":"pro","rate_limit":{"allowed":true,"limit_reached":false,"primary_window":{"used_percent":9,"limit_window_seconds":604800,"reset_after_seconds":505219,"reset_at":1791580412},"secondary_window":null},"code_review_rate_limit":null,"additional_rate_limits":[{"limit_name":"gpt-reserve","metered_feature":"base_model_inference","rate_limit":{"allowed":true,"limit_reached":false,"primary_window":{"used_percent":0,"limit_window_seconds":604800,"reset_after_seconds":604800,"reset_at":1791679994},"secondary_window":null}}],"credits":{"unlimited":false,"overage_limit_reached":false},"rate_limit_reached_type":null}` + recordedClaudeUsage = `{"five_hour":{"utilization":14.0,"resets_at":"2026-10-04T02:40:00.347417+00:00"},"seven_day":{"utilization":31.0,"resets_at":"2026-10-08T23:00:00.347435+00:00"},"iguana_necktie":{"utilization":5.3976432,"resets_at":"2026-11-05T07:59:00+00:00"},"extra_usage":null,"spend":null,"member_dashboard_available":false}` +) + +func TestParseCodexUsageHeadersMapsRecordedResponse(t *testing.T) { + got := ParseCodexUsageHeaders([]byte(recordedCodexUsage)) + want := http.Header{} + want.Set("X-Codex-Primary-Used-Percent", "9") + want.Set("X-Codex-Primary-Window-Minutes", "10080") + want.Set("X-Codex-Primary-Reset-After-Seconds", "505219") + want.Set("X-Codex-Primary-Reset-At", "1791580412") + want.Set("X-Codex-Limit-Reached", "false") + // The null secondary window and the additional limits map to nothing. + if !reflect.DeepEqual(got, want) { + t.Fatalf("ParseCodexUsageHeaders() = %v, want %v", got, want) + } +} + +func TestParseCodexUsageHeadersKeepsExhaustedWindowAsGiven(t *testing.T) { + body := `{"rate_limit":{"limit_reached":true,"primary_window":{"used_percent":100.4,"reset_after_seconds":3600,"reset_at":1791580412},"secondary_window":null}}` + got := ParseCodexUsageHeaders([]byte(body)) + want := http.Header{} + want.Set("X-Codex-Primary-Used-Percent", "100.4") + want.Set("X-Codex-Primary-Reset-After-Seconds", "3600") + want.Set("X-Codex-Primary-Reset-At", "1791580412") + want.Set("X-Codex-Limit-Reached", "true") + if !reflect.DeepEqual(got, want) { + t.Fatalf("ParseCodexUsageHeaders() = %v, want %v", got, want) + } +} + +func TestParseCodexUsageHeadersWithoutWindowReturnsNil(t *testing.T) { + body := `{"rate_limit":{"limit_reached":true,"primary_window":null,"secondary_window":null}}` + if got := ParseCodexUsageHeaders([]byte(body)); got != nil { + t.Fatalf("ParseCodexUsageHeaders() = %v, want nil", got) + } +} + +func TestParseClaudeUsageHeadersMapsRecordedResponse(t *testing.T) { + got := ParseClaudeUsageHeaders([]byte(recordedClaudeUsage)) + want := http.Header{} + want.Set("Anthropic-Ratelimit-Unified-5h-Utilization", "0.14") + want.Set("Anthropic-Ratelimit-Unified-5h-Reset", "1791081600") + want.Set("Anthropic-Ratelimit-Unified-7d-Utilization", "0.31") + want.Set("Anthropic-Ratelimit-Unified-7d-Reset", "1791500400") + if !reflect.DeepEqual(got, want) { + t.Fatalf("ParseClaudeUsageHeaders() = %v, want %v", got, want) + } +} + +func TestParseClaudeUsageHeadersWithoutWindowReturnsNil(t *testing.T) { + if got := ParseClaudeUsageHeaders([]byte(`{"five_hour":null,"seven_day":null}`)); got != nil { + t.Fatalf("ParseClaudeUsageHeaders() = %v, want nil", got) + } +} diff --git a/sdk/cliproxy/auth/quota_signals.go b/sdk/cliproxy/auth/quota_signals.go index fea368e53..6ab3d8472 100644 --- a/sdk/cliproxy/auth/quota_signals.go +++ b/sdk/cliproxy/auth/quota_signals.go @@ -1,6 +1,7 @@ package auth import ( + "context" "net/http" "sort" "strings" @@ -53,6 +54,25 @@ func (q *QuotaState) ObserveResponseHeadersForProvider(provider string, headers return true } +// ObserveQuotaHeaders records a quota snapshot read outside the request path, +// such as a usage endpoint, on the registered credential. It touches only the +// credential-wide observation: counters, availability, cooldown fields and the +// auth record on disk stay unchanged. A snapshot without quota signals leaves +// the previous observation in place. +func (m *Manager) ObserveQuotaHeaders(authID string, headers http.Header) bool { + if m == nil { + return false + } + m.mu.Lock() + auth := m.auths[authID] + observed := auth != nil && auth.Quota.ObserveResponseHeadersForProvider(auth.Provider, headers, time.Now()) + m.mu.Unlock() + if observed { + m.persistCooldownStates(context.Background()) + } + return observed +} + // ClearObservationSignals removes only passive observation data. It leaves // cooldown and scheduler state untouched. func (q *QuotaState) ClearObservationSignals() bool { diff --git a/sdk/cliproxy/auth/quota_signals_test.go b/sdk/cliproxy/auth/quota_signals_test.go index 2a2b08dd2..318da7b28 100644 --- a/sdk/cliproxy/auth/quota_signals_test.go +++ b/sdk/cliproxy/auth/quota_signals_test.go @@ -683,3 +683,50 @@ func TestObserveResponseHeadersKeepsPrimaryWhenTruncatingAdditional(t *testing.T t.Fatalf("snapshot size = %d, want %d", len(quota.Signals), maxQuotaSignalHeaders) } } + +func TestObserveQuotaHeadersKeepsCooldownAndAuthRecord(t *testing.T) { + store := &countingStore{} + cooldownStore := &recordingCooldownStateStore{} + manager := NewManager(store, nil, nil) + manager.SetCooldownStateStore(cooldownStore) + recoverAt := time.Now().Add(time.Hour).Truncate(time.Second) + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), &Auth{ + ID: "usage-read-auth", Provider: "codex", Status: StatusError, Unavailable: true, NextRetryAfter: recoverAt, + Quota: QuotaState{ + Exceeded: true, Reason: "credential_quota", NextRecoverAt: recoverAt, BackoffLevel: 2, + ObservedAt: time.Now().Add(-time.Hour), Signals: map[string]string{"X-Codex-Primary-Used-Percent": "100"}, + }, + }); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + before, _ := manager.GetByID("usage-read-auth") + if !before.Quota.Exceeded || !before.Unavailable { + t.Fatalf("registered auth lost its cooldown: %+v", before) + } + + if manager.ObserveQuotaHeaders("usage-read-auth", nil) { + t.Fatal("ObserveQuotaHeaders(nil) reported an observation") + } + if unchanged, _ := manager.GetByID("usage-read-auth"); !reflect.DeepEqual(unchanged.Quota, before.Quota) { + t.Fatalf("empty read changed quota: %+v, want %+v", unchanged.Quota, before.Quota) + } + + if !manager.ObserveQuotaHeaders("usage-read-auth", http.Header{"X-Codex-Primary-Used-Percent": []string{"9"}}) { + t.Fatal("ObserveQuotaHeaders() did not record the observation") + } + after, _ := manager.GetByID("usage-read-auth") + if after.Quota.Signals["X-Codex-Primary-Used-Percent"] != "9" || !after.Quota.ObservedAt.After(before.Quota.ObservedAt) { + t.Fatalf("observation not recorded: %+v", after.Quota) + } + if !reflect.DeepEqual(cooldownFieldsOf(after.Quota), cooldownFieldsOf(before.Quota)) || after.Unavailable != before.Unavailable || + after.Status != before.Status || !after.NextRetryAfter.Equal(before.NextRetryAfter) || + after.Generation != before.Generation { + t.Fatalf("read changed cooldown or availability: before %+v, after %+v", before, after) + } + if got := store.saveCount.Load(); got != 0 { + t.Fatalf("auth store saves = %d, want 0", got) + } + if got := cooldownStore.saveCount.Load(); got != 1 { + t.Fatalf("cooldown state saves = %d, want 1", got) + } +} diff --git a/sdk/cliproxy/service_lifecycle.go b/sdk/cliproxy/service_lifecycle.go index 7f313edb6..dd1bd1c9f 100644 --- a/sdk/cliproxy/service_lifecycle.go +++ b/sdk/cliproxy/service_lifecycle.go @@ -91,6 +91,7 @@ func (s *Service) Run(ctx context.Context) error { interval := 15 * time.Minute s.coreManager.StartAutoRefresh(ctx, interval) log.Infof("core auth auto-refresh started (interval=%s)", interval) + s.startQuotaUsageReads(ctx) } if !homeEnabled { diff --git a/sdk/cliproxy/service_quota.go b/sdk/cliproxy/service_quota.go new file mode 100644 index 000000000..1588329ec --- /dev/null +++ b/sdk/cliproxy/service_quota.go @@ -0,0 +1,99 @@ +package cliproxy + +import ( + "context" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +const ( + quotaUsageReadInterval = 5 * time.Minute + // quotaUsageReadTimeout bounds one read: the executor HTTP client has no timeout, and + // reads run serially, so one hung read would otherwise stall every later read. + quotaUsageReadTimeout = 30 * time.Second + codexUsageURL = "https://chatgpt.com/backend-api/wham/usage" + claudeUsageURL = "https://api.anthropic.com/api/oauth/usage" +) + +// startQuotaUsageReads reads the usage endpoint of every Codex and Claude OAuth credential +// now and then every quotaUsageReadInterval until ctx ends. The result feeds the quota +// observation that soonest-reset ranks on, so the selector sees quota consumed outside the +// proxy and credentials the proxy has not used yet. +func (s *Service) startQuotaUsageReads(ctx context.Context) { + go func() { + ticker := time.NewTicker(quotaUsageReadInterval) + defer ticker.Stop() + for { + for _, auth := range s.coreManager.List() { + if ctx.Err() != nil { + return + } + s.readQuotaUsage(ctx, auth.ID) + } + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } + }() +} + +// readQuotaUsage reads one credential's usage endpoint and records the result. A failed or +// rejected read leaves the previous observation untouched. +func (s *Service) readQuotaUsage(ctx context.Context, authID string) { + auth, ok := s.coreManager.GetByID(authID) + if !ok || auth.AuthKind() != coreauth.AuthKindOAuth || !auth.HasValidAccessToken(time.Now()) { + return + } + headers := make(http.Header) + var usageURL string + var parse func([]byte) http.Header + switch strings.ToLower(strings.TrimSpace(auth.Provider)) { + case "codex": + usageURL, parse = codexUsageURL, helps.ParseCodexUsageHeaders + if accountID, _ := auth.Metadata["account_id"].(string); accountID != "" { + headers.Set("Chatgpt-Account-Id", accountID) + } + case "claude": + usageURL, parse = claudeUsageURL, helps.ParseClaudeUsageHeaders + headers.Set("Anthropic-Beta", "oauth-2025-04-20") + default: + return + } + readCtx, cancel := context.WithTimeout(ctx, quotaUsageReadTimeout) + defer cancel() + req, errRequest := http.NewRequestWithContext(readCtx, http.MethodGet, usageURL, nil) + if errRequest != nil { + return + } + req.Header = headers + resp, errDo := s.coreManager.HttpRequest(readCtx, auth, req) + if errDo != nil { + log.Warnf("quota usage read failed for auth %s: %v", authID, errDo) + return + } + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("quota usage read: close response body: %v", errClose) + } + }() + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + log.Warnf("quota usage read for auth %s returned status %d", authID, resp.StatusCode) + return + } + body, errRead := io.ReadAll(resp.Body) + if errRead != nil { + log.Warnf("quota usage read failed for auth %s: %v", authID, errRead) + return + } + if observed := parse(body); observed != nil { + s.coreManager.ObserveQuotaHeaders(authID, observed) + } +} diff --git a/sdk/cliproxy/service_quota_test.go b/sdk/cliproxy/service_quota_test.go new file mode 100644 index 000000000..db0b7f2be --- /dev/null +++ b/sdk/cliproxy/service_quota_test.go @@ -0,0 +1,305 @@ +package cliproxy + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "reflect" + "sort" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +type quotaTestStore struct { + auths []*coreauth.Auth + saves atomic.Int32 +} + +func (s *quotaTestStore) List(context.Context) ([]*coreauth.Auth, error) { + out := make([]*coreauth.Auth, 0, len(s.auths)) + for _, auth := range s.auths { + out = append(out, auth.Clone()) + } + return out, nil +} + +func (s *quotaTestStore) Save(_ context.Context, auth *coreauth.Auth) (string, error) { + s.saves.Add(1) + return auth.ID, nil +} + +func (s *quotaTestStore) Delete(context.Context, string) error { return nil } + +type quotaTestCooldownStore struct { + mu sync.Mutex + saves int + records []coreauth.CooldownStateRecord +} + +func (s *quotaTestCooldownStore) Load(context.Context) ([]coreauth.CooldownStateRecord, error) { + return nil, nil +} + +func (s *quotaTestCooldownStore) Save(_ context.Context, records []coreauth.CooldownStateRecord) error { + s.mu.Lock() + defer s.mu.Unlock() + s.saves++ + s.records = records + return nil +} + +type quotaTestRoundTripper func(*http.Request) (*http.Response, error) + +func (f quotaTestRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +type quotaTestRTProvider struct{ rt http.RoundTripper } + +func (p quotaTestRTProvider) RoundTripperFor(*coreauth.Auth) http.RoundTripper { return p.rt } + +func quotaTestResponse(req *http.Request, status int, contentType, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{contentType}}, + Body: io.NopCloser(strings.NewReader(body)), + Request: req, + } +} + +func codexWeeklyUsage(resetAt time.Time) string { + return fmt.Sprintf(`{"plan_type":"pro","rate_limit":{"allowed":true,"limit_reached":false,"primary_window":{"used_percent":9,"limit_window_seconds":604800,"reset_after_seconds":%d,"reset_at":%d},"secondary_window":null},"additional_rate_limits":[]}`, + int64(time.Until(resetAt).Seconds()), resetAt.Unix()) +} + +func codexWeeklySignals(resetAt time.Time) coreauth.QuotaState { + return coreauth.QuotaState{ObservedAt: time.Now(), Signals: map[string]string{ + "X-Codex-Primary-Used-Percent": "50", + "X-Codex-Primary-Window-Minutes": "10080", + "X-Codex-Primary-Reset-At": strconv.FormatInt(resetAt.Unix(), 10), + }} +} + +// TestQuotaUsageSweepRoutesNewSessionToUnusedCredential runs the service with three Codex +// OAuth credentials and one Claude OAuth credential. codex-a has the soonest observed weekly +// reset and its usage read is rejected; codex-c was never used, so without a usage read it +// sorts last. The sweep shows codex-c resetting soonest, so a new session goes to codex-c. +func TestQuotaUsageSweepRoutesNewSessionToUnusedCredential(t *testing.T) { + // A Codex catalog model, so the request routes whether or not the service has replaced + // the registrations below with its own catalog registration yet. + const model = "gpt-5.5" + now := time.Now() + resetA, resetB, resetBRead, resetC := now.Add(48*time.Hour), now.Add(72*time.Hour), now.Add(60*time.Hour), now.Add(24*time.Hour) + // Tokens expire well past every refresh lead and carry no refresh token, so the + // auto-refresh loop has nothing due and no test path can reach a real token endpoint. + expiry := now.Add(10 * 24 * time.Hour).Format(time.RFC3339) + codexAuth := func(id, token string, quota coreauth.QuotaState) *coreauth.Auth { + return &coreauth.Auth{ID: id, Provider: "codex", Status: coreauth.StatusActive, Quota: quota, + Metadata: map[string]any{"type": "codex", "access_token": token, "account_id": "account-" + token, + "expired": expiry}} + } + store := "aTestStore{auths: []*coreauth.Auth{ + codexAuth("codex-a", "tok-a", codexWeeklySignals(resetA)), + codexAuth("codex-b", "tok-b", codexWeeklySignals(resetB)), + codexAuth("codex-c", "tok-c", coreauth.QuotaState{}), + {ID: "claude-d", Provider: "claude", Status: coreauth.StatusActive, + Metadata: map[string]any{"type": "claude", "access_token": "tok-d", "expired": expiry}}, + }} + for _, auth := range store.auths { + id, provider := auth.ID, auth.Provider + registry.GetGlobalRegistry().RegisterClient(id, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) }) + } + + var mu sync.Mutex + var usageReads, responses, unexpected []string + rt := quotaTestRoundTripper(func(req *http.Request) (*http.Response, error) { + token := strings.TrimPrefix(req.Header.Get("Authorization"), "Bearer ") + route := req.Method + " " + req.URL.Host + req.URL.Path + mu.Lock() + defer mu.Unlock() + switch route { + case "GET chatgpt.com/backend-api/wham/usage": + usageReads = append(usageReads, token) + if req.Header.Get("Chatgpt-Account-Id") != "account-"+token { + unexpected = append(unexpected, "missing account header for "+token) + } + switch token { + case "tok-a": + return quotaTestResponse(req, http.StatusTooManyRequests, "application/json", `{"error":"rate limited"}`), nil + case "tok-b": + return quotaTestResponse(req, http.StatusOK, "application/json", codexWeeklyUsage(resetBRead)), nil + case "tok-c": + return quotaTestResponse(req, http.StatusOK, "application/json", codexWeeklyUsage(resetC)), nil + } + case "GET api.anthropic.com/api/oauth/usage": + usageReads = append(usageReads, token) + if req.Header.Get("Anthropic-Beta") != "oauth-2025-04-20" { + unexpected = append(unexpected, "missing anthropic-beta header") + } + return quotaTestResponse(req, http.StatusOK, "application/json", + `{"five_hour":{"utilization":14.0,"resets_at":"2026-10-04T02:40:00.347417+00:00"},"seven_day":{"utilization":31.0,"resets_at":"2026-10-08T23:00:00.347435+00:00"}}`), nil + case "POST chatgpt.com/backend-api/codex/responses": + responses = append(responses, token) + completed := fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp-1","object":"response","status":"completed","model":%q,"output":[{"type":"message","id":"msg-1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok","annotations":[]}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, model) + return quotaTestResponse(req, http.StatusOK, "text/event-stream", "data: "+completed+"\n\n"), nil + } + unexpected = append(unexpected, route) + return nil, errors.New("unexpected upstream request: " + route) + }) + + ln, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + port := ln.Addr().(*net.TCPAddr).Port + if errClose := ln.Close(); errClose != nil { + t.Fatalf("close listener: %v", errClose) + } + authDir := t.TempDir() + cfgPath := filepath.Join(authDir, "config.yaml") + if errWrite := os.WriteFile(cfgPath, []byte("{}"), 0o644); errWrite != nil { + t.Fatalf("write config: %v", errWrite) + } + cfg := &config.Config{Host: "127.0.0.1", Port: port, AuthDir: authDir, SaveCooldownStatus: true} + cfg.APIKeys = []string{"proxy-key"} + cfg.Routing.Strategy = "soonest-reset" + cooldownStore := "aTestCooldownStore{} + manager := coreauth.NewManager(store, newRoutingSelector(normalizedRoutingRuntimeState(cfg)), nil) + service, errBuild := NewBuilder(). + WithConfig(cfg). + WithConfigPath(cfgPath). + WithCoreAuthManager(manager). + WithCooldownStateStore(cooldownStore). + WithWatcherFactory(func(string, string, func(*config.Config)) (*WatcherWrapper, error) { + return &WatcherWrapper{}, nil + }). + Build() + if errBuild != nil { + t.Fatalf("build service: %v", errBuild) + } + // Execution reads the transport from the round-tripper provider; the sweep, which runs + // outside any request, reads it from the service context. + manager.SetRoundTripperProvider(quotaTestRTProvider{rt: rt}) + runCtx, cancel := context.WithCancel(context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(rt))) + runDone := make(chan error, 1) + go func() { runDone <- service.Run(runCtx) }() + defer func() { + cancel() + select { + case <-runDone: + case <-time.After(10 * time.Second): + t.Error("timed out waiting for service.Run to exit") + } + }() + + resetOf := func(id string) string { + auth, _ := manager.GetByID(id) + if auth == nil { + return "" + } + return auth.Quota.Signals["X-Codex-Primary-Reset-At"] + } + swept := func() bool { + mu.Lock() + reads := len(usageReads) + mu.Unlock() + claude, _ := manager.GetByID("claude-d") + return reads == 4 && claude != nil && claude.Quota.Signals["Anthropic-Ratelimit-Unified-7d-Reset"] != "" && + resetOf("codex-b") == strconv.FormatInt(resetBRead.Unix(), 10) && + resetOf("codex-c") == strconv.FormatInt(resetC.Unix(), 10) + } + for deadline := time.Now().Add(5 * time.Second); !swept() && time.Now().Before(deadline); { + time.Sleep(10 * time.Millisecond) + } + if !swept() { + t.Fatal("the sweep did not read and record every credential's usage within 5s") + } + + mu.Lock() + reads := append([]string(nil), usageReads...) + sort.Strings(reads) + if want := []string{"tok-a", "tok-b", "tok-c", "tok-d"}; !reflect.DeepEqual(reads, want) { + t.Errorf("usage reads = %v, want one per credential %v", usageReads, want) + } + if len(unexpected) != 0 { + t.Errorf("unexpected upstream traffic during the sweep (refreshes included): %v", unexpected) + } + mu.Unlock() + if saves := store.saves.Load(); saves != 0 { + t.Errorf("token store saves during the sweep = %d, want 0", saves) + } + cooldownStore.mu.Lock() + cdsSaves, records := cooldownStore.saves, cooldownStore.records + cooldownStore.mu.Unlock() + if cdsSaves == 0 { + t.Error("cooldown state store was not written after the sweep") + } + savedReset := map[string]string{} + for _, record := range records { + if record.Model == "" { + savedReset[record.AuthID] = record.Quota.Signals["X-Codex-Primary-Reset-At"] + } + } + if got, want := savedReset["codex-c"], strconv.FormatInt(resetC.Unix(), 10); got != want { + t.Errorf("saved codex-c reset = %q, want the read %q", got, want) + } + if got, want := savedReset["codex-a"], strconv.FormatInt(resetA.Unix(), 10); got != want { + t.Errorf("saved codex-a reset after a rejected read = %q, want previous %q", got, want) + } + if got, want := resetOf("codex-a"), strconv.FormatInt(resetA.Unix(), 10); got != want { + t.Errorf("codex-a reset after a rejected read = %q, want previous %q", got, want) + } + for _, id := range []string{"codex-a", "codex-b", "codex-c", "claude-d"} { + auth, _ := manager.GetByID(id) + if token, _ := auth.Metadata["access_token"].(string); token != "tok-"+id[len(id)-1:] { + t.Errorf("%s access token changed during the sweep: %q", id, token) + } + // A refresh attempt sets NextRefreshAfter when it starts and LastRefreshedAt on success. + if !auth.NextRefreshAfter.IsZero() || !auth.LastRefreshedAt.IsZero() { + t.Errorf("%s was refreshed during the sweep", id) + } + } + + var resp *http.Response + var errDo error + for deadline := time.Now().Add(5 * time.Second); time.Now().Before(deadline); time.Sleep(20 * time.Millisecond) { + req, errRequest := http.NewRequest(http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/v1/responses", port), + strings.NewReader(fmt.Sprintf(`{"model":%q,"input":"hello"}`, model))) + if errRequest != nil { + t.Fatalf("new request: %v", errRequest) + } + req.Header.Set("Authorization", "Bearer proxy-key") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Session-Id", fmt.Sprintf("session-%d", time.Now().UnixNano())) + if resp, errDo = http.DefaultClient.Do(req); errDo == nil { + break + } + } + if errDo != nil { + t.Fatalf("proxy request: %v", errDo) + } + body, _ := io.ReadAll(resp.Body) + if errClose := resp.Body.Close(); errClose != nil { + t.Errorf("close response body: %v", errClose) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("proxy status = %d, body = %s", resp.StatusCode, body) + } + mu.Lock() + defer mu.Unlock() + if len(responses) != 1 || responses[0] != "tok-c" { + t.Fatalf("new session went to %v, want [tok-c]: the never-used credential resets soonest", responses) + } +}