mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 17:36:47 +08:00
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
This commit is contained in:
Luis Pater
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user