From 2fb67d24552169823f8197b9bc916cbe0402062d Mon Sep 17 00:00:00 2001 From: Eric Litman Date: Sat, 3 Oct 2026 21:17:56 -0400 Subject: [PATCH 1/3] feat(auth): read Codex and Claude usage on a timer for soonest-reset (MASTRA-787) The service reads each Codex and Claude OAuth credential's usage endpoint at startup and every five minutes, maps the response onto the quota signals the soonest-reset selector already ranks on, and records it on the credential without touching counters, cooldowns or the auth file. The selector now sees quota consumed outside the proxy and credentials the proxy has never used. Co-Authored-By: Claude Opus 5.5 --- .../executor/helps/claude_ratelimit.go | 25 ++ .../runtime/executor/helps/codex_quota.go | 36 +++ .../executor/helps/quota_usage_test.go | 52 ++++ sdk/cliproxy/auth/quota_signals.go | 25 ++ sdk/cliproxy/auth/quota_signals_test.go | 47 +++ sdk/cliproxy/service.go | 3 + sdk/cliproxy/service_lifecycle.go | 6 + sdk/cliproxy/service_quota.go | 99 ++++++ sdk/cliproxy/service_quota_test.go | 284 ++++++++++++++++++ 9 files changed, 577 insertions(+) create mode 100644 internal/runtime/executor/helps/quota_usage_test.go create mode 100644 sdk/cliproxy/service_quota.go create mode 100644 sdk/cliproxy/service_quota_test.go diff --git a/internal/runtime/executor/helps/claude_ratelimit.go b/internal/runtime/executor/helps/claude_ratelimit.go index 12264fcd3..764dd235f 100644 --- a/internal/runtime/executor/helps/claude_ratelimit.go +++ b/internal/runtime/executor/helps/claude_ratelimit.go @@ -10,6 +10,7 @@ import ( "time" log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" ) const ( @@ -306,3 +307,27 @@ func randomClaudeFuzzDuration(minSec, maxSec int) time.Duration { } return time.Duration(minSec+int(nBig.Int64())) * time.Second } + +// 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/codex_quota.go b/internal/runtime/executor/helps/codex_quota.go index cc0c1c2b4..ddce8a479 100644 --- a/internal/runtime/executor/helps/codex_quota.go +++ b/internal/runtime/executor/helps/codex_quota.go @@ -2,6 +2,7 @@ package helps import ( "net/http" + "strconv" "strings" "github.com/tidwall/gjson" @@ -113,6 +114,41 @@ func ParseCodexQuotaEventHeaders(payload []byte) http.Header { return headers } +// ParseCodexUsageHeaders converts a /backend-api/wham/usage response into the +// X-Codex primary and secondary window headers that Codex responses carry. It +// returns nil when the response holds no usable window. +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") + used := window.Get("used_percent") + seconds := window.Get("limit_window_seconds").Int() + resetAfter := window.Get("reset_after_seconds") + resetAt := window.Get("reset_at") + hasResetAfter := resetAfter.Type == gjson.Number && resetAfter.Int() >= 0 + hasResetAt := resetAt.Type == gjson.Number && resetAt.Int() > 0 + if used.Type != gjson.Number || used.Float() < 0 || used.Float() > 100 || seconds < 60 || + (!hasResetAfter && !hasResetAt) { + continue + } + prefix := "X-Codex-" + name + "-" + headers.Set(prefix+"Used-Percent", used.Raw) + headers.Set(prefix+"Window-Minutes", strconv.FormatInt(seconds/60, 10)) + if hasResetAfter { + headers.Set(prefix+"Reset-After-Seconds", resetAfter.Raw) + } + if hasResetAt { + headers.Set(prefix+"Reset-At", resetAt.Raw) + } + } + if len(headers) == 0 { + return nil + } + setCodexQuotaScalarHeaderFromResult(headers, "X-Codex-Limit-Reached", rateLimit, "limit_reached") + return headers +} + const ( codexQuotaEventOther = iota codexQuotaEventError 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..5acfe2ef3 --- /dev/null +++ b/internal/runtime/executor/helps/quota_usage_test.go @@ -0,0 +1,52 @@ +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 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..36fe53397 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,30 @@ 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] + if auth == nil || !auth.Quota.ObserveResponseHeadersForProvider(auth.Provider, headers, time.Now()) { + m.mu.Unlock() + return false + } + snapshot := auth.Clone() + m.mu.Unlock() + if m.scheduler != nil { + m.scheduler.upsertAuthResult(snapshot, nil, true) + } + m.persistCooldownStates(context.Background()) + return true +} + // 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.go b/sdk/cliproxy/service.go index a8d9603ab..c90ef0e0c 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -81,6 +81,9 @@ type Service struct { // watcherCancel cancels the watcher context. watcherCancel context.CancelFunc + // quotaReadsCancel stops the quota usage reads. + quotaReadsCancel context.CancelFunc + // authUpdates channel for authentication updates. authUpdates chan watcher.AuthUpdate diff --git a/sdk/cliproxy/service_lifecycle.go b/sdk/cliproxy/service_lifecycle.go index 7f313edb6..d8a201471 100644 --- a/sdk/cliproxy/service_lifecycle.go +++ b/sdk/cliproxy/service_lifecycle.go @@ -91,6 +91,9 @@ 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) + quotaCtx, quotaCancel := context.WithCancel(ctx) + s.quotaReadsCancel = quotaCancel + s.startQuotaUsageReads(quotaCtx) } if !homeEnabled { @@ -286,6 +289,9 @@ func (s *Service) Shutdown(ctx context.Context) error { if s.watcherCancel != nil { s.watcherCancel() } + if s.quotaReadsCancel != nil { + s.quotaReadsCancel() + } if s.coreManager != nil { s.coreManager.StopAutoRefresh() } 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..c78772ccf --- /dev/null +++ b/sdk/cliproxy/service_quota_test.go @@ -0,0 +1,284 @@ +package cliproxy + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "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) + } + + mu.Lock() + t.Logf("usage reads: %v", usageReads) + 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 := cooldownStore.saves + cooldownStore.mu.Unlock() + if cdsSaves == 0 { + t.Error("cooldown state store was not written after the sweep") + } + 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) + } +} From 51f8a2057f50242037b4927690eb620229d20633 Mon Sep 17 00:00:00 2001 From: Eric Litman Date: Sat, 3 Oct 2026 21:21:28 -0400 Subject: [PATCH 2/3] docs(agents): allow the quota usage read timeout (MASTRA-787) The usage reads run serially on one timer, and the executor HTTP client has no timeout, so the per-read deadline is what keeps one hung read from stopping every later observation. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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. From 93889015188bb1bc14a515da5810e2a943c74050 Mon Sep 17 00:00:00 2001 From: Eric Litman Date: Sat, 3 Oct 2026 21:49:39 -0400 Subject: [PATCH 3/3] refactor(auth): subtract review findings from the usage reads (MASTRA-787) The Codex usage parser now copies each non-null window as given instead of dropping windows that fail range and reset checks. ObserveQuotaHeaders no longer pushes a clone to the scheduler, which reads no observation fields. The usage read loop runs on Run's own context, so the quotaReadsCancel field and its Shutdown branch are gone. Both usage parsers move to helps/quota_usage.go, and codex_quota.go and claude_ratelimit.go return to their upstream content. The end-to-end sweep test now fails when the sweep does not finish, checks the exact usage reads, and checks the saved cooldown records. Co-Authored-By: Claude Opus 5.5 --- .../executor/helps/claude_ratelimit.go | 25 ------- .../runtime/executor/helps/codex_quota.go | 36 ---------- .../runtime/executor/helps/quota_usage.go | 69 +++++++++++++++++++ .../executor/helps/quota_usage_test.go | 13 ++++ sdk/cliproxy/auth/quota_signals.go | 13 ++-- sdk/cliproxy/service.go | 3 - sdk/cliproxy/service_lifecycle.go | 7 +- sdk/cliproxy/service_quota_test.go | 25 ++++++- 8 files changed, 110 insertions(+), 81 deletions(-) create mode 100644 internal/runtime/executor/helps/quota_usage.go diff --git a/internal/runtime/executor/helps/claude_ratelimit.go b/internal/runtime/executor/helps/claude_ratelimit.go index 764dd235f..12264fcd3 100644 --- a/internal/runtime/executor/helps/claude_ratelimit.go +++ b/internal/runtime/executor/helps/claude_ratelimit.go @@ -10,7 +10,6 @@ import ( "time" log "github.com/sirupsen/logrus" - "github.com/tidwall/gjson" ) const ( @@ -307,27 +306,3 @@ func randomClaudeFuzzDuration(minSec, maxSec int) time.Duration { } return time.Duration(minSec+int(nBig.Int64())) * time.Second } - -// 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/codex_quota.go b/internal/runtime/executor/helps/codex_quota.go index ddce8a479..cc0c1c2b4 100644 --- a/internal/runtime/executor/helps/codex_quota.go +++ b/internal/runtime/executor/helps/codex_quota.go @@ -2,7 +2,6 @@ package helps import ( "net/http" - "strconv" "strings" "github.com/tidwall/gjson" @@ -114,41 +113,6 @@ func ParseCodexQuotaEventHeaders(payload []byte) http.Header { return headers } -// ParseCodexUsageHeaders converts a /backend-api/wham/usage response into the -// X-Codex primary and secondary window headers that Codex responses carry. It -// returns nil when the response holds no usable window. -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") - used := window.Get("used_percent") - seconds := window.Get("limit_window_seconds").Int() - resetAfter := window.Get("reset_after_seconds") - resetAt := window.Get("reset_at") - hasResetAfter := resetAfter.Type == gjson.Number && resetAfter.Int() >= 0 - hasResetAt := resetAt.Type == gjson.Number && resetAt.Int() > 0 - if used.Type != gjson.Number || used.Float() < 0 || used.Float() > 100 || seconds < 60 || - (!hasResetAfter && !hasResetAt) { - continue - } - prefix := "X-Codex-" + name + "-" - headers.Set(prefix+"Used-Percent", used.Raw) - headers.Set(prefix+"Window-Minutes", strconv.FormatInt(seconds/60, 10)) - if hasResetAfter { - headers.Set(prefix+"Reset-After-Seconds", resetAfter.Raw) - } - if hasResetAt { - headers.Set(prefix+"Reset-At", resetAt.Raw) - } - } - if len(headers) == 0 { - return nil - } - setCodexQuotaScalarHeaderFromResult(headers, "X-Codex-Limit-Reached", rateLimit, "limit_reached") - return headers -} - const ( codexQuotaEventOther = iota codexQuotaEventError 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 index 5acfe2ef3..2cf12b9aa 100644 --- a/internal/runtime/executor/helps/quota_usage_test.go +++ b/internal/runtime/executor/helps/quota_usage_test.go @@ -26,6 +26,19 @@ func TestParseCodexUsageHeadersMapsRecordedResponse(t *testing.T) { } } +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 { diff --git a/sdk/cliproxy/auth/quota_signals.go b/sdk/cliproxy/auth/quota_signals.go index 36fe53397..6ab3d8472 100644 --- a/sdk/cliproxy/auth/quota_signals.go +++ b/sdk/cliproxy/auth/quota_signals.go @@ -65,17 +65,12 @@ func (m *Manager) ObserveQuotaHeaders(authID string, headers http.Header) bool { } m.mu.Lock() auth := m.auths[authID] - if auth == nil || !auth.Quota.ObserveResponseHeadersForProvider(auth.Provider, headers, time.Now()) { - m.mu.Unlock() - return false - } - snapshot := auth.Clone() + observed := auth != nil && auth.Quota.ObserveResponseHeadersForProvider(auth.Provider, headers, time.Now()) m.mu.Unlock() - if m.scheduler != nil { - m.scheduler.upsertAuthResult(snapshot, nil, true) + if observed { + m.persistCooldownStates(context.Background()) } - m.persistCooldownStates(context.Background()) - return true + return observed } // ClearObservationSignals removes only passive observation data. It leaves diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index c90ef0e0c..a8d9603ab 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -81,9 +81,6 @@ type Service struct { // watcherCancel cancels the watcher context. watcherCancel context.CancelFunc - // quotaReadsCancel stops the quota usage reads. - quotaReadsCancel context.CancelFunc - // authUpdates channel for authentication updates. authUpdates chan watcher.AuthUpdate diff --git a/sdk/cliproxy/service_lifecycle.go b/sdk/cliproxy/service_lifecycle.go index d8a201471..dd1bd1c9f 100644 --- a/sdk/cliproxy/service_lifecycle.go +++ b/sdk/cliproxy/service_lifecycle.go @@ -91,9 +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) - quotaCtx, quotaCancel := context.WithCancel(ctx) - s.quotaReadsCancel = quotaCancel - s.startQuotaUsageReads(quotaCtx) + s.startQuotaUsageReads(ctx) } if !homeEnabled { @@ -289,9 +287,6 @@ func (s *Service) Shutdown(ctx context.Context) error { if s.watcherCancel != nil { s.watcherCancel() } - if s.quotaReadsCancel != nil { - s.quotaReadsCancel() - } if s.coreManager != nil { s.coreManager.StopAutoRefresh() } diff --git a/sdk/cliproxy/service_quota_test.go b/sdk/cliproxy/service_quota_test.go index c78772ccf..db0b7f2be 100644 --- a/sdk/cliproxy/service_quota_test.go +++ b/sdk/cliproxy/service_quota_test.go @@ -9,6 +9,8 @@ import ( "net/http" "os" "path/filepath" + "reflect" + "sort" "strconv" "strings" "sync" @@ -221,9 +223,16 @@ func TestQuotaUsageSweepRoutesNewSessionToUnusedCredential(t *testing.T) { 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() - t.Logf("usage reads: %v", usageReads) + 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) } @@ -232,11 +241,23 @@ func TestQuotaUsageSweepRoutesNewSessionToUnusedCredential(t *testing.T) { t.Errorf("token store saves during the sweep = %d, want 0", saves) } cooldownStore.mu.Lock() - cdsSaves := cooldownStore.saves + 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) }