History

feat(cooldowns): add cooldown snapshot feature for management auth files

Supra4E8C committed 2 months ago

Commit 1ca975d1 parent f702bc14 files · +603 −4Browse files

Unified · Split

4 files
  1. internal/api/handlers/management/auth_files.go +11 −4
  2. internal/api/handlers/management/auth_files_cooldown_test.go +189 −0
  3. sdk/cliproxy/auth/cooldown_view.go +171 −0
  4. sdk/cliproxy/auth/cooldown_view_test.go +232 −0
internal/api/handlers/management/auth_files.go+11−4
@@ -106,12 +106,18 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { · 105 unmodified lines
106106 quotaSupportedProviders = host.QuotaSupportedProvidersSet(c.Request.Context())
107107 }
108108 auths := h.authManager.List()
109+ observedAt := time.Now().UTC()
110+ cooldownsKnown := !h.authManager.HomeEnabled()
109111 files := make([]gin.H, 0, len(auths))
110112 for _, auth := range auths {
111113 if !matchesAuthFileLookup(auth, nameFilter, authIndexFilter) {
112114 continue
113115 }
114116 if entry := h.buildAuthFileEntry(auth, quotaSupportedProviders); entry != nil {
117+ entry["cooldowns"] = nil
118+ if cooldownsKnown {
119+ entry["cooldowns"] = coreauth.CooldownSnapshotForAuth(auth, observedAt)
120+ }
115121 files = append(files, entry)
116122 }
117123 }
@@ -120,7 +126,7 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { · 2 unmodified lines
120126 nameJ, _ := files[j]["name"].(string)
121127 return strings.ToLower(nameI) < strings.ToLower(nameJ)
122128 })
123− c.JSON(200, gin.H{"files": files})
129+ c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
124130 }
125131
126132 func lockedAuthIndex(auth *coreauth.Auth) string {
@@ -222,6 +228,7 @@ func (h *Handler) GetAuthFileModels(c *gin.Context) { · 95 unmodified lines
222228
223229 // List auth files from disk when the auth manager is unavailable.
224230 func (h *Handler) listAuthFilesFromDisk(c *gin.Context) {
231+ observedAt := time.Now().UTC()
225232 nameFilter := strings.TrimSpace(c.Query("name"))
226233 authIndexFilter := strings.TrimSpace(c.Query("auth_index"))
227234 entries, err := os.ReadDir(h.cfg.AuthDir)
@@ -231,7 +238,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 3 unmodified lines
231238 }
232239 files := make([]gin.H, 0)
233240 if authIndexFilter != "" {
234− c.JSON(200, gin.H{"files": files})
241+ c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
235242 return
236243 }
237244 for _, e := range entries {
@@ -246,7 +253,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 8 unmodified lines
246253 continue
247254 }
248255 if info, errInfo := e.Info(); errInfo == nil {
249− fileData := gin.H{"name": name, "size": info.Size(), "modtime": info.ModTime()}
256+ fileData := gin.H{"name": name, "size": info.Size(), "modtime": info.ModTime(), "cooldowns": nil}
250257
251258 // Read file to get type field
252259 full := filepath.Join(h.cfg.AuthDir, name)
@@ -307,7 +314,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 54 unmodified lines
307314 files = append(files, fileData)
308315 }
309316 }
310− c.JSON(200, gin.H{"files": files})
317+ c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
311318 }
312319
313320 func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth, quotaSupported ...map[string]struct{}) gin.H {
internal/api/handlers/management/auth_files_cooldown_test.go+189−0
@@ -0,0 +1,189 @@
1+package management
2+
3+import (
4+ "encoding/json"
5+ "net/http"
6+ "net/http/httptest"
7+ "net/url"
8+ "os"
9+ "path/filepath"
10+ "reflect"
11+ "testing"
12+ "time"
13+
14+ "github.com/gin-gonic/gin"
15+ "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
16+ coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
17+)
18+
19+type authFilesCooldownResponse struct {
20+ ObservedAt time.Time `json:"observed_at"`
21+ Files []struct {
22+ ID string `json:"id"`
23+ AuthIndex string `json:"auth_index"`
24+ Name string `json:"name"`
25+ Status string `json:"status"`
26+ Unavailable bool `json:"unavailable"`
27+ NextRetryAfter time.Time `json:"next_retry_after"`
28+ Cooldowns json.RawMessage `json:"cooldowns"`
29+ Quota map[string]any `json:"quota"`
30+ ModelQuotas map[string]map[string]any `json:"model_quotas"`
31+ } `json:"files"`
32+}
33+
34+func requestAuthFilesCooldowns(t *testing.T, h *Handler, query string) authFilesCooldownResponse {
35+ t.Helper()
36+ rec := httptest.NewRecorder()
37+ ctx, _ := gin.CreateTestContext(rec)
38+ ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files"+query, nil)
39+ h.ListAuthFiles(ctx)
40+ if rec.Code != http.StatusOK {
41+ t.Fatalf("status = %d: %s", rec.Code, rec.Body.String())
42+ }
43+ var payload authFilesCooldownResponse
44+ if errDecode := json.Unmarshal(rec.Body.Bytes(), &payload); errDecode != nil {
45+ t.Fatal(errDecode)
46+ }
47+ if payload.ObservedAt.IsZero() || payload.ObservedAt.Location() != time.UTC {
48+ t.Fatalf("invalid observed_at: %v", payload.ObservedAt)
49+ }
50+ return payload
51+}
52+
53+func TestListAuthFilesCooldownsSnapshot(t *testing.T) {
54+ t.Setenv("MANAGEMENT_PASSWORD", "")
55+ now := time.Now().UTC()
56+ next := now.Add(time.Hour)
57+ manager := coreauth.NewManager(nil, nil, nil)
58+ cfg := &config.Config{AuthDir: t.TempDir()}
59+ manager.SetConfig(cfg)
60+ for _, id := range []string{"a", "b"} {
61+ registerAuthForLookupTest(t, manager, &coreauth.Auth{
62+ ID: id, Index: "index-" + id, Provider: "codex", Status: coreauth.StatusError,
63+ Unavailable: true, NextRetryAfter: next,
64+ Attributes: map[string]string{"runtime_only": "true"},
65+ Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, ObservedAt: now, Signals: map[string]string{"x-codex-primary-used-percent": "90"}},
66+ ModelStates: map[string]*coreauth.ModelState{
67+ "model-a": {Unavailable: true, NextRetryAfter: next, Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 6, ObservedAt: now, Signals: map[string]string{"x-codex-primary-used-percent": "90"}}, LastError: &coreauth.Error{HTTPStatus: 429, Message: "private upstream body"}},
68+ "expired": {Status: coreauth.StatusError, Unavailable: true, NextRetryAfter: now.Add(-time.Hour), Quota: coreauth.QuotaState{Exceeded: true, NextRecoverAt: now.Add(-time.Hour), BackoffLevel: 9}},
69+ },
70+ })
71+ }
72+ beforeA, _ := manager.GetByID("a")
73+ beforeB, _ := manager.GetByID("b")
74+ h := NewHandlerWithoutConfigFilePath(cfg, manager)
75+ for range 2 {
76+ payload := requestAuthFilesCooldowns(t, h, "")
77+ if len(payload.Files) != 2 {
78+ t.Fatalf("files = %+v", payload.Files)
79+ }
80+ for i, file := range payload.Files {
81+ if file.ID != []string{"a", "b"}[i] || file.AuthIndex != "index-"+file.ID {
82+ t.Fatalf("identity/order changed: %+v", file)
83+ }
84+ if file.Status != string(coreauth.StatusError) || !file.Unavailable || !file.NextRetryAfter.Equal(next) {
85+ t.Fatalf("existing state changed: %+v", file)
86+ }
87+ var views []coreauth.CooldownView
88+ if errDecode := json.Unmarshal(file.Cooldowns, &views); errDecode != nil {
89+ t.Fatal(errDecode)
90+ }
91+ if len(views) != 1 || views[0].Scope != "model" || views[0].ModelKey != "model-a" || views[0].Reason != "quota" || views[0].HTTPStatus != 429 || views[0].BackoffLevel == nil || *views[0].BackoffLevel != 6 {
92+ t.Fatalf("unexpected cooldowns: %s", file.Cooldowns)
93+ }
94+ remaining := next.Sub(payload.ObservedAt)
95+ wantSeconds := int64(remaining / time.Second)
96+ if remaining%time.Second != 0 {
97+ wantSeconds++
98+ }
99+ if views[0].RemainingSeconds != wantSeconds || !views[0].RetryAt.Equal(next) {
100+ t.Fatalf("inconsistent time basis: %+v", views[0])
101+ }
102+ for _, quota := range []map[string]any{file.Quota, file.ModelQuotas["model-a"]} {
103+ if len(quota) != 2 || quota["signals"] == nil || quota["observed_at"] == nil {
104+ t.Fatalf("quota observation changed: %+v", quota)
105+ }
106+ }
107+ var fields []map[string]any
108+ if errDecode := json.Unmarshal(file.Cooldowns, &fields); errDecode != nil {
109+ t.Fatal(errDecode)
110+ }
111+ if len(fields[0]) != 7 {
112+ t.Fatalf("unexpected fields: %+v", fields[0])
113+ }
114+ }
115+ }
116+ afterA, _ := manager.GetByID("a")
117+ afterB, _ := manager.GetByID("b")
118+ if !reflect.DeepEqual(beforeA, afterA) || !reflect.DeepEqual(beforeB, afterB) {
119+ t.Fatal("GET mutated auth state")
120+ }
121+}
122+
123+func TestListAuthFilesCooldownsCredentialKindsAndFilters(t *testing.T) {
124+ t.Setenv("MANAGEMENT_PASSWORD", "")
125+ cfg := &config.Config{AuthDir: t.TempDir()}
126+ manager := coreauth.NewManager(nil, nil, nil)
127+ manager.SetConfig(cfg)
128+ path := filepath.Join(cfg.AuthDir, "shared.json")
129+ if errWrite := os.WriteFile(path, []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
130+ t.Fatal(errWrite)
131+ }
132+ for _, id := range []string{"file", "virtual", "runtime"} {
133+ auth := &coreauth.Auth{ID: id, Index: "index-" + id, FileName: "shared.json", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"path": path}}
134+ if id == "virtual" {
135+ coreauth.MarkPluginVirtualAuth(auth, path, 0)
136+ }
137+ if id == "runtime" {
138+ auth.Attributes = map[string]string{"runtime_only": "true"}
139+ }
140+ registerAuthForLookupTest(t, manager, auth)
141+ }
142+ h := NewHandlerWithoutConfigFilePath(cfg, manager)
143+ payload := requestAuthFilesCooldowns(t, h, "?name=shared.json")
144+ if len(payload.Files) != 3 {
145+ t.Fatalf("files = %+v", payload.Files)
146+ }
147+ for _, file := range payload.Files {
148+ if string(file.Cooldowns) != "[]" {
149+ t.Fatalf("known empty cooldowns = %s", file.Cooldowns)
150+ }
151+ filtered := requestAuthFilesCooldowns(t, h, "?name=shared.json&auth_index="+url.QueryEscape(file.AuthIndex))
152+ if len(filtered.Files) != 1 || filtered.Files[0].ID != file.ID || string(filtered.Files[0].Cooldowns) != "[]" {
153+ t.Fatalf("filter mismatch: %+v", filtered.Files)
154+ }
155+ }
156+ if missing := requestAuthFilesCooldowns(t, h, "?auth_index=missing"); len(missing.Files) != 0 {
157+ t.Fatal("unknown index matched")
158+ }
159+}
160+
161+func TestListAuthFilesCooldownsUnknown(t *testing.T) {
162+ t.Setenv("MANAGEMENT_PASSWORD", "")
163+ for _, mode := range []string{"disk", "home"} {
164+ t.Run(mode, func(t *testing.T) {
165+ cfg := &config.Config{AuthDir: t.TempDir()}
166+ var manager *coreauth.Manager
167+ if mode == "disk" {
168+ if errWrite := os.WriteFile(filepath.Join(cfg.AuthDir, "a.json"), []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
169+ t.Fatal(errWrite)
170+ }
171+ } else {
172+ cfg.Home.Enabled = true
173+ manager = coreauth.NewManager(nil, nil, nil)
174+ manager.SetConfig(cfg)
175+ registerAuthForLookupTest(t, manager, &coreauth.Auth{ID: "a", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"runtime_only": "true"}})
176+ }
177+ h := NewHandlerWithoutConfigFilePath(cfg, manager)
178+ payload := requestAuthFilesCooldowns(t, h, "")
179+ if len(payload.Files) != 1 || string(payload.Files[0].Cooldowns) != "null" {
180+ t.Fatalf("unknown cooldowns = %+v", payload.Files)
181+ }
182+ if mode == "disk" {
183+ if filtered := requestAuthFilesCooldowns(t, h, "?auth_index=missing"); len(filtered.Files) != 0 {
184+ t.Fatal("disk fallback matched index")
185+ }
186+ }
187+ })
188+ }
189+}
sdk/cliproxy/auth/cooldown_view.go+171−0
@@ -0,0 +1,171 @@
1+package auth
2+
3+import (
4+ "sort"
5+ "strings"
6+ "time"
7+)
8+
9+// CooldownView describes an unexpired local retry restriction, not overall
10+// credential availability. It contains no credential metadata or raw errors.
11+type CooldownView struct {
12+ Scope string `json:"scope"`
13+ ModelKey string `json:"model_key,omitempty"`
14+ Reason string `json:"reason"`
15+ RetryAt time.Time `json:"retry_at"`
16+ RemainingSeconds int64 `json:"remaining_seconds"`
17+ BackoffLevel *int `json:"backoff_level,omitempty"`
18+ HTTPStatus int `json:"http_status,omitempty"`
19+}
20+
21+// CooldownSnapshotForAuth projects a detached auth snapshot without mutating it.
22+// It reports timers even when another restriction (such as disablement or an
23+// expired token) also prevents execution. An empty result does not imply that
24+// the credential is usable. Callers must handle unavailable/remote state separately.
25+func CooldownSnapshotForAuth(auth *Auth, now time.Time) []CooldownView {
26+ views := make([]CooldownView, 0)
27+ if auth == nil {
28+ return views
29+ }
30+ // Match the explicit credential-wide gate in isAuthBlockedForModel. Other
31+ // auth-level fields can be model aggregates and must not become global gates.
32+ if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) {
33+ views = append(views, newCooldownView("credential", "", auth.Quota.NextRecoverAt, now, auth.Quota, auth.StatusMessage, auth.LastError))
34+ } else if len(auth.ModelStates) == 0 {
35+ if blocked, _, next := availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now); blocked && next.After(now) {
36+ views = append(views, newCooldownView("credential", "", next, now, auth.Quota, auth.StatusMessage, auth.LastError))
37+ }
38+ }
39+
40+ // Sorting source keys makes ties deterministic after selector precedence.
41+ keys := make([]string, 0, len(auth.ModelStates))
42+ for key := range auth.ModelStates {
43+ keys = append(keys, key)
44+ }
45+ sort.Strings(keys)
46+ type modelCooldown struct {
47+ view CooldownView
48+ reason blockReason
49+ }
50+ byModel := make(map[string]modelCooldown)
51+ for _, key := range keys {
52+ state := auth.ModelStates[key]
53+ model := canonicalModelKey(key)
54+ if state == nil || model == "" {
55+ continue
56+ }
57+ blocked, reason, next := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now)
58+ if !blocked || !next.After(now) {
59+ continue
60+ }
61+ if previous, ok := byModel[model]; ok {
62+ preferQuotaTie := next.Equal(previous.view.RetryAt) && reason == blockReasonCooldown && previous.reason != blockReasonCooldown
63+ if !next.After(previous.view.RetryAt) && !preferQuotaTie {
64+ continue
65+ }
66+ }
67+ byModel[model] = modelCooldown{
68+ view: newCooldownView("model", model, next, now, state.Quota, state.StatusMessage, state.LastError),
69+ reason: reason,
70+ }
71+ }
72+ models := make([]string, 0, len(byModel))
73+ for model := range byModel {
74+ models = append(models, model)
75+ }
76+ sort.Strings(models)
77+ for _, model := range models {
78+ views = append(views, byModel[model].view)
79+ }
80+ return views
81+}
82+
83+func newCooldownView(scope, model string, next, now time.Time, quota QuotaState, statusMessage string, lastErr *Error) CooldownView {
84+ remaining := next.Sub(now)
85+ seconds := int64(remaining / time.Second)
86+ if remaining%time.Second != 0 {
87+ seconds++
88+ }
89+ view := CooldownView{
90+ Scope: scope, ModelKey: model, Reason: "unknown",
91+ RetryAt: next.UTC(), RemainingSeconds: seconds,
92+ }
93+ // A shorter quota window must not label a longer non-quota retry timer.
94+ // Zero recovery times can occur in legacy quota state with only a retry time.
95+ if quota.Exceeded && (quota.NextRecoverAt.IsZero() || !quota.NextRecoverAt.Before(next)) {
96+ switch quota.Reason {
97+ case "credential_quota", "quota":
98+ view.Reason = quota.Reason
99+ case "cloudflare challenge":
100+ view.Reason = "cloudflare_challenge"
101+ }
102+ }
103+ propagatedQuota := quota.Exceeded && quota.Reason == "credential_quota"
104+ if view.Reason == "credential_quota" {
105+ // The credential-wide gate takes precedence over stale sibling errors.
106+ return view
107+ }
108+ if (view.Reason == "quota" || view.Reason == "cloudflare_challenge") && quota.BackoffLevel >= 0 {
109+ level := quota.BackoffLevel
110+ view.BackoffLevel = &level
111+ }
112+ errorReason := cooldownErrorReason(lastErr)
113+ if view.Reason == "unknown" {
114+ view.Reason = errorReason
115+ }
116+ if view.Reason == "unknown" {
117+ view.Reason = cooldownStatusReason(statusMessage)
118+ }
119+ // Propagation does not replace sibling errors. Avoid attributing a stale
120+ // error status to that quota failure, even if a longer retry timer survives.
121+ if !propagatedQuota && lastErr != nil && lastErr.HTTPStatus >= 400 && lastErr.HTTPStatus <= 599 && errorReason == view.Reason {
122+ view.HTTPStatus = lastErr.HTTPStatus
123+ }
124+ return view
125+}
126+
127+func cooldownErrorReason(err *Error) string {
128+ switch {
129+ case isModelSupportResultError(err):
130+ return "model_not_supported"
131+ case isCloudflareChallengeResultError(err):
132+ return "cloudflare_challenge"
133+ case isInvalidGrantResultError(err):
134+ return "invalid_grant"
135+ }
136+ switch statusCodeFromResult(err) {
137+ case 401:
138+ return "unauthorized"
139+ case 402, 403:
140+ return "payment_required"
141+ case 404:
142+ return "not_found"
143+ case 429:
144+ return "quota"
145+ case 408, 500, 502, 503, 504, 520, 521, 522, 523, 524, 525, 526:
146+ return "transient_error"
147+ }
148+ if err != nil {
149+ return cooldownStatusReason(err.Code)
150+ }
151+ return "unknown"
152+}
153+
154+func cooldownStatusReason(message string) string {
155+ // Only exact known markers can become public reason codes. Never return
156+ // arbitrary status messages, error codes, or upstream response bodies.
157+ switch strings.TrimSpace(message) {
158+ case "quota", "quota exhausted":
159+ return "quota"
160+ case "cloudflare challenge":
161+ return "cloudflare_challenge"
162+ case "invalid_grant", "unauthorized", "payment_required", "not_found":
163+ return strings.TrimSpace(message)
164+ case "model_not_supported":
165+ return "model_not_supported"
166+ case "transient upstream error":
167+ return "transient_error"
168+ default:
169+ return "unknown"
170+ }
171+}
sdk/cliproxy/auth/cooldown_view_test.go+232−0
@@ -0,0 +1,232 @@
1+package auth
2+
3+import (
4+ "encoding/json"
5+ "reflect"
6+ "strings"
7+ "testing"
8+ "time"
9+)
10+
11+func TestCooldownSnapshotForAuthScopes(t *testing.T) {
12+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
13+ modelState := func(after time.Duration) *ModelState {
14+ return &ModelState{Unavailable: true, NextRetryAfter: now.Add(after), Quota: QuotaState{
15+ Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(after), BackoffLevel: 6,
16+ }, LastError: &Error{HTTPStatus: 429}}
17+ }
18+ for _, tt := range []struct {
19+ name string
20+ auth *Auth
21+ want []string
22+ }{
23+ {name: "nil", want: []string{}},
24+ {name: "empty", auth: &Auth{}, want: []string{}},
25+ {name: "aggregate is not credential cooldown", auth: &Auth{
26+ Unavailable: true, NextRetryAfter: now.Add(time.Minute),
27+ Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(time.Minute)},
28+ ModelStates: map[string]*ModelState{"b": modelState(time.Minute), "a": modelState(32 * time.Second), "ready": {}},
29+ }, want: []string{"model:a", "model:b"}},
30+ {name: "credential quota and longer model coexist", auth: &Auth{
31+ NextRetryAfter: now.Add(time.Hour),
32+ Quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(20 * time.Second), BackoffLevel: 9},
33+ ModelStates: map[string]*ModelState{"a": modelState(time.Minute)},
34+ }, want: []string{"credential:", "model:a"}},
35+ {name: "credential fallback", auth: &Auth{
36+ Unavailable: true, NextRetryAfter: now.Add(time.Minute), LastError: &Error{HTTPStatus: 503},
37+ }, want: []string{"credential:"}},
38+ {name: "nil model states still suppress fallback like selector", auth: &Auth{
39+ Unavailable: true, NextRetryAfter: now.Add(time.Minute), ModelStates: map[string]*ModelState{"a": nil},
40+ }, want: []string{}},
41+ {name: "disabled credential retains timer", auth: &Auth{
42+ Disabled: true, Status: StatusDisabled, Unavailable: true, NextRetryAfter: now.Add(time.Minute),
43+ }, want: []string{"credential:"}},
44+ {name: "expired token does not erase timer", auth: &Auth{
45+ Metadata: map[string]any{"expired": now.Add(-time.Hour).Format(time.RFC3339)},
46+ ModelStates: map[string]*ModelState{"a": modelState(time.Minute)},
47+ }, want: []string{"model:a"}},
48+ {name: "disabled without timer", auth: &Auth{Disabled: true, ModelStates: map[string]*ModelState{"a": {Status: StatusDisabled}}}, want: []string{}},
49+ {name: "forced timer survives disable cooling override", auth: &Auth{
50+ Metadata: map[string]any{"disable_cooling": true}, Unavailable: true, NextRetryAfter: now.Add(time.Minute),
51+ LastError: &Error{Code: ErrorCodeForceCooldown},
52+ }, want: []string{"credential:"}},
53+ } {
54+ t.Run(tt.name, func(t *testing.T) {
55+ got := CooldownSnapshotForAuth(tt.auth, now)
56+ if got == nil {
57+ t.Fatal("expected non-nil slice")
58+ }
59+ keys := make([]string, 0, len(got))
60+ for _, view := range got {
61+ keys = append(keys, view.Scope+":"+view.ModelKey)
62+ }
63+ if !reflect.DeepEqual(keys, tt.want) {
64+ t.Fatalf("scopes = %v, want %v", keys, tt.want)
65+ }
66+ if tt.name == "credential quota and longer model coexist" {
67+ if got[0].RemainingSeconds != 20 || got[0].BackoffLevel != nil || got[0].HTTPStatus != 0 {
68+ t.Fatalf("credential gate did not use its own deadline/diagnostics: %+v", got[0])
69+ }
70+ }
71+ })
72+ }
73+}
74+
75+func TestCooldownSnapshotForAuthTimeBoundaries(t *testing.T) {
76+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
77+ for _, tt := range []struct {
78+ name string
79+ state ModelState
80+ seconds int64
81+ }{
82+ {name: "just before expiry", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(time.Nanosecond)}, seconds: 1},
83+ {name: "fraction rounds up", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(1500 * time.Millisecond)}, seconds: 2},
84+ {name: "exact expiry", state: ModelState{Unavailable: true, NextRetryAfter: now}},
85+ {name: "past expiry", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(-time.Nanosecond)}},
86+ {name: "historical error and backoff", state: ModelState{Status: StatusError, Quota: QuotaState{BackoffLevel: 6}}},
87+ {name: "no deadline", state: ModelState{Unavailable: true, Quota: QuotaState{Exceeded: true}}},
88+ {name: "inactive future timestamp", state: ModelState{NextRetryAfter: now.Add(time.Minute)}},
89+ {name: "later quota time", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(time.Second), Quota: QuotaState{Exceeded: true, NextRecoverAt: now.Add(3 * time.Second)}}, seconds: 3},
90+ {name: "later retry time", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(4 * time.Second), Quota: QuotaState{Exceeded: true, NextRecoverAt: now.Add(time.Second)}}, seconds: 4},
91+ {name: "quota only", state: ModelState{Quota: QuotaState{Exceeded: true, NextRecoverAt: now.Add(5 * time.Second)}}, seconds: 5},
92+ {name: "expired quota and retry", state: ModelState{Unavailable: true, NextRetryAfter: now, Quota: QuotaState{Exceeded: true, NextRecoverAt: now, BackoffLevel: 6}}},
93+ {name: "retry hint independent of backoff", state: ModelState{Unavailable: true, NextRetryAfter: now.Add(97 * time.Second), Quota: QuotaState{Exceeded: true, Reason: "quota", BackoffLevel: 2}}, seconds: 97},
94+ } {
95+ t.Run(tt.name, func(t *testing.T) {
96+ got := CooldownSnapshotForAuth(&Auth{ModelStates: map[string]*ModelState{"a": &tt.state}}, now)
97+ if tt.seconds == 0 {
98+ if len(got) != 0 {
99+ t.Fatalf("unexpected cooldown: %+v", got)
100+ }
101+ return
102+ }
103+ if len(got) != 1 || got[0].RemainingSeconds != tt.seconds {
104+ t.Fatalf("views = %+v, want remaining %d", got, tt.seconds)
105+ }
106+ })
107+ }
108+}
109+
110+func TestCooldownSnapshotForAuthReasons(t *testing.T) {
111+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
112+ for _, tt := range []struct {
113+ name string
114+ quota QuotaState
115+ err *Error
116+ message string
117+ reason string
118+ status int
119+ backoff bool
120+ }{
121+ {name: "quota", quota: QuotaState{Exceeded: true, Reason: "quota", BackoffLevel: 6}, err: &Error{HTTPStatus: 429}, reason: "quota", status: 429, backoff: true},
122+ {name: "propagated quota hides old error", quota: QuotaState{Exceeded: true, Reason: "credential_quota", BackoffLevel: 6}, err: &Error{HTTPStatus: 401}, reason: "credential_quota"},
123+ {name: "quota hides unrelated error", quota: QuotaState{Exceeded: true, Reason: "quota"}, err: &Error{HTTPStatus: 503}, reason: "quota", backoff: true},
124+ {name: "challenge", quota: QuotaState{Exceeded: true, Reason: "cloudflare challenge"}, err: &Error{HTTPStatus: 403, Message: "cf-mitigated: challenge"}, reason: "cloudflare_challenge", status: 403, backoff: true},
125+ {name: "model unsupported", err: &Error{HTTPStatus: 400, Message: "model not supported"}, reason: "model_not_supported", status: 400},
126+ {name: "invalid grant", err: &Error{HTTPStatus: 400, Message: "invalid_grant"}, reason: "invalid_grant", status: 400},
127+ {name: "unauthorized", err: &Error{HTTPStatus: 401}, reason: "unauthorized", status: 401},
128+ {name: "payment", err: &Error{HTTPStatus: 402}, reason: "payment_required", status: 402},
129+ {name: "forbidden", err: &Error{HTTPStatus: 403}, reason: "payment_required", status: 403},
130+ {name: "not found", err: &Error{HTTPStatus: 404}, reason: "not_found", status: 404},
131+ {name: "gateway not challenge", err: &Error{HTTPStatus: 520, Message: "cloudflare challenge"}, reason: "transient_error", status: 520},
132+ {name: "shorter active quota does not label longer retry", quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(20 * time.Second), BackoffLevel: 6}, err: &Error{HTTPStatus: 503}, reason: "transient_error", status: 503},
133+ {name: "longer quota supplies deadline", quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(2 * time.Minute), BackoffLevel: 6}, err: &Error{HTTPStatus: 503}, reason: "quota", backoff: true},
134+ {name: "shorter propagated quota preserves longer retry reason", quota: QuotaState{Exceeded: true, Reason: "credential_quota", NextRecoverAt: now.Add(20 * time.Second), BackoffLevel: 6}, err: &Error{HTTPStatus: 401}, reason: "unauthorized"},
135+ {name: "expired quota does not label new failure", quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now, BackoffLevel: 6}, err: &Error{HTTPStatus: 503}, reason: "transient_error", status: 503},
136+ {name: "known marker", message: "transient upstream error", reason: "transient_error"},
137+ {name: "unknown is sanitized", quota: QuotaState{Exceeded: true, Reason: "secret-quota"}, err: &Error{Code: "secret-code", Message: "secret-body"}, message: "secret-message", reason: "unknown"},
138+ {name: "unknown HTTP error", err: &Error{HTTPStatus: 418}, reason: "unknown", status: 418},
139+ } {
140+ t.Run(tt.name, func(t *testing.T) {
141+ state := &ModelState{Unavailable: true, NextRetryAfter: now.Add(time.Minute), Quota: tt.quota, LastError: tt.err, StatusMessage: tt.message}
142+ got := CooldownSnapshotForAuth(&Auth{ModelStates: map[string]*ModelState{"a": state}}, now)
143+ if len(got) != 1 {
144+ t.Fatalf("views = %+v", got)
145+ }
146+ view := got[0]
147+ if view.Reason != tt.reason || view.HTTPStatus != tt.status || (view.BackoffLevel != nil) != tt.backoff {
148+ t.Fatalf("view = %+v, want reason=%s status=%d backoff=%v", view, tt.reason, tt.status, tt.backoff)
149+ }
150+ encoded, errMarshal := json.Marshal(view)
151+ if errMarshal != nil {
152+ t.Fatal(errMarshal)
153+ }
154+ if strings.Contains(string(encoded), "secret") {
155+ t.Fatalf("raw data leaked: %s", encoded)
156+ }
157+ })
158+ }
159+}
160+
161+func TestCooldownSnapshotForAuthDeduplicatesWithoutMutation(t *testing.T) {
162+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.FixedZone("offset", 3600))
163+ auth := &Auth{
164+ Metadata: map[string]any{"access_token": "secret-token"},
165+ ModelStates: map[string]*ModelState{
166+ " model-a(high) ": {Unavailable: true, NextRetryAfter: now.Add(time.Minute), LastError: &Error{HTTPStatus: 503}},
167+ "model-a": {Unavailable: true, NextRetryAfter: now.Add(32 * time.Second), Quota: QuotaState{Exceeded: true, Reason: "quota", BackoffLevel: 6}, LastError: &Error{HTTPStatus: 429}},
168+ "model-b": {Unavailable: true, NextRetryAfter: now.Add(time.Minute), Quota: QuotaState{Exceeded: true, Reason: "quota", BackoffLevel: 6}},
169+ " ": {Unavailable: true, NextRetryAfter: now.Add(time.Minute)},
170+ "nil": nil,
171+ },
172+ }
173+ before := auth.Clone()
174+ got := CooldownSnapshotForAuth(auth, now)
175+ if len(got) != 2 || got[0].ModelKey != "model-a" || got[0].Reason != "transient_error" || got[0].HTTPStatus != 503 || got[0].BackoffLevel != nil || got[0].RemainingSeconds != 60 {
176+ t.Fatalf("deduplicated views = %+v", got)
177+ }
178+ if got[0].RetryAt.Location() != time.UTC {
179+ t.Fatal("retry_at is not UTC")
180+ }
181+ for range 20 {
182+ if again := CooldownSnapshotForAuth(auth, now); !reflect.DeepEqual(again, got) {
183+ t.Fatal("unstable snapshot")
184+ }
185+ }
186+ *got[1].BackoffLevel = 999
187+ if !reflect.DeepEqual(auth, before) {
188+ t.Fatal("projection or returned view mutated auth")
189+ }
190+}
191+
192+func TestCooldownSnapshotForAuthEqualDeadlineIsStable(t *testing.T) {
193+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
194+ for _, quotaKey := range []string{"a(high)", "a(low)"} {
195+ t.Run(quotaKey, func(t *testing.T) {
196+ otherKey := "a(low)"
197+ if quotaKey == otherKey {
198+ otherKey = "a(high)"
199+ }
200+ auth := &Auth{ModelStates: map[string]*ModelState{
201+ quotaKey: {Unavailable: true, NextRetryAfter: now.Add(time.Minute), LastError: &Error{HTTPStatus: 429}, Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(time.Minute)}},
202+ otherKey: {Unavailable: true, NextRetryAfter: now.Add(time.Minute), LastError: &Error{HTTPStatus: 503}},
203+ }}
204+ blocked, reason, next := isAuthBlockedForModel(auth, "a", now)
205+ if !blocked || reason != blockReasonCooldown {
206+ t.Fatal("precondition: selector must prefer quota on equal deadlines")
207+ }
208+ for range 20 {
209+ got := CooldownSnapshotForAuth(auth, now)
210+ if len(got) != 1 || got[0].Reason != "quota" || !got[0].RetryAt.Equal(next) || got[0].HTTPStatus != 429 {
211+ t.Fatalf("unstable equal-deadline selection: %+v", got)
212+ }
213+ }
214+ })
215+ }
216+}
217+
218+func TestCooldownSnapshotForAuthLongerRetryReasonSurvivesQuotaExpiry(t *testing.T) {
219+ now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC)
220+ auth := &Auth{ModelStates: map[string]*ModelState{
221+ "a": {
222+ Unavailable: true, NextRetryAfter: now.Add(time.Hour), LastError: &Error{HTTPStatus: 503},
223+ Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: now.Add(5 * time.Minute), BackoffLevel: 6},
224+ },
225+ }}
226+ for _, observed := range []time.Time{now, now.Add(5 * time.Minute), now.Add(6 * time.Minute)} {
227+ got := CooldownSnapshotForAuth(auth, observed)
228+ if len(got) != 1 || got[0].Reason != "transient_error" || got[0].HTTPStatus != 503 || got[0].BackoffLevel != nil || !got[0].RetryAt.Equal(now.Add(time.Hour)) {
229+ t.Fatalf("longer retry diagnostics changed at %v: %+v", observed, got)
230+ }
231+ }
232+}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t