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
106 quotaSupportedProviders = host.QuotaSupportedProvidersSet(c.Request.Context())106 quotaSupportedProviders = host.QuotaSupportedProvidersSet(c.Request.Context())
107 }107 }
108 auths := h.authManager.List()108 auths := h.authManager.List()
109 observedAt := time.Now().UTC()
110 cooldownsKnown := !h.authManager.HomeEnabled()
109 files := make([]gin.H, 0, len(auths))111 files := make([]gin.H, 0, len(auths))
110 for _, auth := range auths {112 for _, auth := range auths {
111 if !matchesAuthFileLookup(auth, nameFilter, authIndexFilter) {113 if !matchesAuthFileLookup(auth, nameFilter, authIndexFilter) {
112 continue114 continue
113 }115 }
114 if entry := h.buildAuthFileEntry(auth, quotaSupportedProviders); entry != nil {116 if entry := h.buildAuthFileEntry(auth, quotaSupportedProviders); entry != nil {
117 entry["cooldowns"] = nil
118 if cooldownsKnown {
119 entry["cooldowns"] = coreauth.CooldownSnapshotForAuth(auth, observedAt)
120 }
115 files = append(files, entry)121 files = append(files, entry)
116 }122 }
117 }123 }
@@ -120,7 +126,7 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { · 2 unmodified lines
120 nameJ, _ := files[j]["name"].(string)126 nameJ, _ := files[j]["name"].(string)
121 return strings.ToLower(nameI) < strings.ToLower(nameJ)127 return strings.ToLower(nameI) < strings.ToLower(nameJ)
122 })128 })
123 c.JSON(200, gin.H{"files": files})129 c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
124}130}
125 131
126func lockedAuthIndex(auth *coreauth.Auth) string {132func lockedAuthIndex(auth *coreauth.Auth) string {
@@ -222,6 +228,7 @@ func (h *Handler) GetAuthFileModels(c *gin.Context) { · 95 unmodified lines
222 228
223// List auth files from disk when the auth manager is unavailable.229// List auth files from disk when the auth manager is unavailable.
224func (h *Handler) listAuthFilesFromDisk(c *gin.Context) {230func (h *Handler) listAuthFilesFromDisk(c *gin.Context) {
231 observedAt := time.Now().UTC()
225 nameFilter := strings.TrimSpace(c.Query("name"))232 nameFilter := strings.TrimSpace(c.Query("name"))
226 authIndexFilter := strings.TrimSpace(c.Query("auth_index"))233 authIndexFilter := strings.TrimSpace(c.Query("auth_index"))
227 entries, err := os.ReadDir(h.cfg.AuthDir)234 entries, err := os.ReadDir(h.cfg.AuthDir)
@@ -231,7 +238,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 3 unmodified lines
231 }238 }
232 files := make([]gin.H, 0)239 files := make([]gin.H, 0)
233 if authIndexFilter != "" {240 if authIndexFilter != "" {
234 c.JSON(200, gin.H{"files": files})241 c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
235 return242 return
236 }243 }
237 for _, e := range entries {244 for _, e := range entries {
@@ -246,7 +253,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 8 unmodified lines
246 continue253 continue
247 }254 }
248 if info, errInfo := e.Info(); errInfo == nil {255 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}
250 257
251 // Read file to get type field258 // Read file to get type field
252 full := filepath.Join(h.cfg.AuthDir, name)259 full := filepath.Join(h.cfg.AuthDir, name)
@@ -307,7 +314,7 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context) { · 54 unmodified lines
307 files = append(files, fileData)314 files = append(files, fileData)
308 }315 }
309 }316 }
310 c.JSON(200, gin.H{"files": files})317 c.JSON(200, gin.H{"observed_at": observedAt, "files": files})
311}318}
312 319
313func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth, quotaSupported ...map[string]struct{}) gin.H {320func (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 @@
1package management
2
3import (
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
19type 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
34func 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
53func 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
123func 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
161func 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 @@
1package auth
2
3import (
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.
11type 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.
25func 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
83func 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
127func 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
154func 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 @@
1package auth
2
3import (
4 "encoding/json"
5 "reflect"
6 "strings"
7 "testing"
8 "time"
9)
10
11func 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
75func 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
110func 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
161func 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
192func 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
218func 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