From 5771abbc81fd2f1860d698b54fe3e63757a77d19 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 20 Jun 2026 13:49:40 +0800 Subject: [PATCH] feat(management): add `ResetQuota` endpoint and auth manager quota reset functionality - Introduced a new `/reset-quota` API endpoint in the management handler to clear quota and cooldown state for auth records. - Implemented `ResetQuota` method in the auth manager to handle runtime and registry state resets for affected models. - Added tests to validate quota reset behavior, including proper state updates and registry consistency. - Refactored utility functions to support deduplication and registered models handling in quota resets. Closes: #3866 --- internal/api/handlers/management/quota.go | 53 ++++++- .../api/handlers/management/quota_test.go | 134 ++++++++++++++++++ internal/api/server.go | 1 + sdk/cliproxy/auth/conductor.go | 114 +++++++++++++++ .../auth/conductor_availability_test.go | 74 ++++++++++ 5 files changed, 375 insertions(+), 1 deletion(-) create mode 100644 internal/api/handlers/management/quota_test.go diff --git a/internal/api/handlers/management/quota.go b/internal/api/handlers/management/quota.go index c7efd217..a87a05ef 100644 --- a/internal/api/handlers/management/quota.go +++ b/internal/api/handlers/management/quota.go @@ -1,6 +1,12 @@ package management -import "github.com/gin-gonic/gin" +import ( + "fmt" + "net/http" + "strings" + + "github.com/gin-gonic/gin" +) // Quota exceeded toggles func (h *Handler) GetSwitchProject(c *gin.Context) { @@ -16,3 +22,48 @@ func (h *Handler) GetSwitchPreviewModel(c *gin.Context) { func (h *Handler) PutSwitchPreviewModel(c *gin.Context) { h.updateBoolField(c, func(v bool) { h.cfg.QuotaExceeded.SwitchPreviewModel = v }) } + +// ResetQuota clears quota/cooldown routing state for one auth index. +func (h *Handler) ResetQuota(c *gin.Context) { + if h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + + var req struct { + AuthIndex string `json:"auth_index"` + } + if errBindJSON := c.ShouldBindJSON(&req); errBindJSON != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + + authIndex := strings.TrimSpace(req.AuthIndex) + if authIndex == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "auth_index is required"}) + return + } + + auth := h.authByIndex(authIndex) + if auth == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth not found"}) + return + } + + updated, models, errReset := h.authManager.ResetQuota(c.Request.Context(), auth.ID) + if errReset != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to reset quota: %v", errReset)}) + return + } + if updated == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth not found"}) + return + } + updated.EnsureIndex() + + c.JSON(http.StatusOK, gin.H{ + "status": "ok", + "auth_index": updated.Index, + "models": models, + }) +} diff --git a/internal/api/handlers/management/quota_test.go b/internal/api/handlers/management/quota_test.go new file mode 100644 index 00000000..aee9b1d8 --- /dev/null +++ b/internal/api/handlers/management/quota_test.go @@ -0,0 +1,134 @@ +package management + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestResetQuota_UsesAuthIndex(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + manager := coreauth.NewManager(nil, nil, nil) + next := time.Now().Add(time.Hour) + auth := &coreauth.Auth{ + ID: "reset-auth-id", + FileName: "reset-auth-file.json", + Provider: "claude", + Status: coreauth.StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: next, + Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2}, + ModelStates: map[string]*coreauth.ModelState{ + "claude-reset-model": { + Status: coreauth.StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: next, + Quota: coreauth.QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2}, + }, + }, + } + authIndex := auth.EnsureIndex() + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("failed to register auth record: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager) + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPost, "/v0/management/reset-quota", strings.NewReader(`{"auth_index":"`+authIndex+`"}`)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + h.ResetQuota(ctx) + + if rec.Code != http.StatusOK { + t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String()) + } + + var payload map[string]any + if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil { + t.Fatalf("failed to decode response: %v", errUnmarshal) + } + if payload["auth_index"] != authIndex { + t.Fatalf("auth_index = %#v, want %q", payload["auth_index"], authIndex) + } + + updated, ok := manager.GetByID("reset-auth-id") + if !ok || updated == nil { + t.Fatalf("expected auth record to exist after reset") + } + if updated.Status != coreauth.StatusActive || updated.StatusMessage != "" || updated.Unavailable || !updated.NextRetryAfter.IsZero() { + t.Fatalf("updated auth state = status %q message %q unavailable %v next %v", updated.Status, updated.StatusMessage, updated.Unavailable, updated.NextRetryAfter) + } + if updated.Quota.Exceeded || updated.Quota.Reason != "" || !updated.Quota.NextRecoverAt.IsZero() || updated.Quota.BackoffLevel != 0 { + t.Fatalf("updated auth quota = %+v, want cleared", updated.Quota) + } + state := updated.ModelStates["claude-reset-model"] + if state == nil { + t.Fatalf("expected model state to remain") + } + if state.Status != coreauth.StatusActive || state.StatusMessage != "" || state.Unavailable || !state.NextRetryAfter.IsZero() { + t.Fatalf("updated model state = status %q message %q unavailable %v next %v", state.Status, state.StatusMessage, state.Unavailable, state.NextRetryAfter) + } + if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 { + t.Fatalf("updated model quota = %+v, want cleared", state.Quota) + } +} + +func TestResetQuota_DoesNotAcceptAuthIDOrFileName(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ + ID: "reset-auth-id-only", + FileName: "reset-auth-file-only.json", + Provider: "claude", + Status: coreauth.StatusError, + } + authIndex := auth.EnsureIndex() + if authIndex == auth.ID || authIndex == auth.FileName { + t.Fatalf("test auth_index unexpectedly matches id or file name: %q", authIndex) + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("failed to register auth record: %v", errRegister) + } + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager) + + tests := []struct { + name string + body string + wantCode int + }{ + {name: "auth_id field ignored", body: `{"auth_id":"reset-auth-id-only"}`, wantCode: http.StatusBadRequest}, + {name: "id field ignored", body: `{"id":"reset-auth-id-only"}`, wantCode: http.StatusBadRequest}, + {name: "file name is not an index", body: `{"auth_index":"reset-auth-file-only.json"}`, wantCode: http.StatusNotFound}, + {name: "auth id is not an index", body: `{"auth_index":"reset-auth-id-only"}`, wantCode: http.StatusNotFound}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + req := httptest.NewRequest(http.MethodPost, "/v0/management/reset-quota", strings.NewReader(tt.body)) + req.Header.Set("Content-Type", "application/json") + ctx.Request = req + h.ResetQuota(ctx) + + if rec.Code != tt.wantCode { + t.Fatalf("status = %d, want %d with body %s", rec.Code, tt.wantCode, rec.Body.String()) + } + }) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index e31955eb..bddd963b 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -645,6 +645,7 @@ func (s *Server) registerManagementRoutes() { mgmt.GET("/quota-exceeded/switch-preview-model", s.mgmt.GetSwitchPreviewModel) mgmt.PUT("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) mgmt.PATCH("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) + mgmt.POST("/reset-quota", s.mgmt.ResetQuota) mgmt.GET("/api-keys", s.mgmt.GetAPIKeys) mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 44302da1..b8d0a430 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" "net/http" "path/filepath" @@ -686,6 +687,119 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { return changed } +func dedupeStrings(values []string) []string { + if len(values) < 2 { + return values + } + seen := make(map[string]struct{}, len(values)) + out := values[:0] + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + out = append(out, value) + } + return out +} + +// ResetQuota clears quota/cooldown state for an auth and resumes registry routing. +func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []string, error) { + if m == nil { + return nil, nil, nil + } + authID = strings.TrimSpace(authID) + if authID == "" { + return nil, nil, fmt.Errorf("auth id is required") + } + + now := time.Now() + var snapshot *Auth + models := make([]string, 0) + registeredModels := modelsForRegisteredAuth(authID) + cooldownStateChanged := false + + m.mu.Lock() + auth, ok := m.auths[authID] + if !ok || auth == nil { + m.mu.Unlock() + return nil, nil, nil + } + + var cooldownRecordsBefore []CooldownStateRecord + trackCooldownState := m.cooldownStore != nil + if trackCooldownState { + cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) + } + + for modelKey, state := range auth.ModelStates { + if strings.TrimSpace(modelKey) == "" { + continue + } + models = append(models, modelKey) + if state != nil { + resetModelState(state, now) + } + } + if clearCooldownStateForAuth(auth, now) { + if len(models) == 0 { + models = append(models, registeredModels...) + } + } else if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } + + if len(models) == 0 { + models = append(models, registeredModels...) + } + models = dedupeStrings(models) + + if !auth.Disabled && auth.Status != StatusDisabled && !hasModelError(auth, now) { + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + } + auth.UpdatedAt = now + if errPersist := m.persist(ctx, auth); errPersist != nil { + m.mu.Unlock() + return nil, nil, errPersist + } + snapshot = auth.Clone() + if trackCooldownState { + cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) + cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) + } + m.mu.Unlock() + + for _, modelKey := range models { + registry.GetGlobalRegistry().ClearModelQuotaExceeded(authID, modelKey) + registry.GetGlobalRegistry().ResumeClientModel(authID, modelKey) + } + if m.scheduler != nil && snapshot != nil { + m.scheduler.upsertAuth(snapshot) + } + if snapshot != nil && cooldownStateChanged { + m.persistCooldownStates(ctx) + } + return snapshot, models, nil +} + +func modelsForRegisteredAuth(authID string) []string { + supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) + models := make([]string, 0, len(supportedModels)) + for _, supportedModel := range supportedModels { + if supportedModel == nil || strings.TrimSpace(supportedModel.ID) == "" { + continue + } + models = append(models, supportedModel.ID) + } + return models +} + func (m *Manager) persistCooldownStates(ctx context.Context) { if m == nil { return diff --git a/sdk/cliproxy/auth/conductor_availability_test.go b/sdk/cliproxy/auth/conductor_availability_test.go index 831df3b0..7e07cc07 100644 --- a/sdk/cliproxy/auth/conductor_availability_test.go +++ b/sdk/cliproxy/auth/conductor_availability_test.go @@ -4,6 +4,8 @@ import ( "context" "testing" "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" ) func TestUpdateAggregatedAvailability_UnavailableWithoutNextRetryDoesNotBlockAuth(t *testing.T) { @@ -102,3 +104,75 @@ func TestManager_AvailableProvidersAndHasProviderAuth_ExcludeDisabled(t *testing t.Errorf("HasProviderAuth(codex) = true, want false (only StatusDisabled auth registered)") } } + +func TestManager_ResetQuotaClearsRuntimeAndRegistryState(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := context.Background() + authID := "reset-quota-auth" + model := "reset-quota-model" + next := time.Now().Add(time.Hour) + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(authID) + }) + + if _, errRegister := manager.Register(ctx, &Auth{ + ID: authID, + Provider: "claude", + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2}, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: next, BackoffLevel: 2}, + UpdatedAt: next, + }, + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + reg.SetModelQuotaExceeded(authID, model) + reg.SuspendClientModel(authID, model, "quota") + if count := reg.GetModelCount(model); count != 0 { + t.Fatalf("registry model count before reset = %d, want 0", count) + } + + updated, models, errReset := manager.ResetQuota(ctx, authID) + if errReset != nil { + t.Fatalf("ResetQuota() error = %v", errReset) + } + if updated == nil { + t.Fatalf("ResetQuota() updated auth is nil") + } + if len(models) != 1 || models[0] != model { + t.Fatalf("ResetQuota() models = %v, want [%s]", models, model) + } + if updated.Status != StatusActive || updated.StatusMessage != "" || updated.Unavailable || !updated.NextRetryAfter.IsZero() { + t.Fatalf("updated auth state = status %q message %q unavailable %v next %v", updated.Status, updated.StatusMessage, updated.Unavailable, updated.NextRetryAfter) + } + if updated.Quota.Exceeded || updated.Quota.Reason != "" || !updated.Quota.NextRecoverAt.IsZero() || updated.Quota.BackoffLevel != 0 { + t.Fatalf("updated auth quota = %+v, want cleared", updated.Quota) + } + state := updated.ModelStates[model] + if state == nil { + t.Fatalf("updated model state missing") + } + if state.Status != StatusActive || state.StatusMessage != "" || state.Unavailable || !state.NextRetryAfter.IsZero() { + t.Fatalf("updated model state = status %q message %q unavailable %v next %v", state.Status, state.StatusMessage, state.Unavailable, state.NextRetryAfter) + } + if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 { + t.Fatalf("updated model quota = %+v, want cleared", state.Quota) + } + if count := reg.GetModelCount(model); count != 1 { + t.Fatalf("registry model count after reset = %d, want 1", count) + } +}