diff --git a/config.example.yaml b/config.example.yaml index 83959ea7..fedefab3 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -200,7 +200,10 @@ quota-exceeded: # Routing strategy for selecting credentials when multiple match. routing: - strategy: "round-robin" # round-robin (default), fill-first + strategy: "round-robin" # round-robin (default), weighted-round-robin, fill-first + # weighted-round-robin uses each credential's integer weight (default 1, maximum 1,000,000). + # Non-positive weights exclude the credential while this strategy is active. + # For OAuth/file credentials, add a top-level numeric "weight" field to the auth JSON. # Enable universal session-sticky routing for all clients. # Explicit Claude Code, Codex, OpenCode, and pi session headers are preferred, # followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, @@ -278,6 +281,7 @@ nonstream-keepalive-interval: 0 # Gemini API keys # gemini-api-key: # - api-key: "AIzaSy...01" +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/gemini-3-pro-preview" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://generativelanguage.googleapis.com" @@ -301,6 +305,7 @@ nonstream-keepalive-interval: 0 # send Gemini generateContent/streamGenerateContent requests when the client enters through the interactions API. # interactions-api-key: # - api-key: "AIzaSy...03" +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "native" # optional: require calls like "native/gemini-3-pro-preview" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://generativelanguage.googleapis.com" @@ -317,6 +322,7 @@ nonstream-keepalive-interval: 0 # Codex API keys # codex-api-key: # - api-key: "sk-atSM..." +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/gpt-5-codex" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://www.example.com" # use the custom codex API endpoint @@ -339,6 +345,7 @@ nonstream-keepalive-interval: 0 # Uses the native xAI executor, including its Responses namespace-tool handling. # xai-api-key: # - api-key: "xai-..." +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "xai" # optional: require calls like "xai/grok-4.5" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://api.x.ai/v1" # xAI-compatible Responses API endpoint @@ -360,6 +367,7 @@ nonstream-keepalive-interval: 0 # claude-api-key: # - api-key: "sk-atSM..." # use the official claude API key, no need to set the base url # - api-key: "sk-atSM..." +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/claude-sonnet-latest" to target this credential # disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling # base-url: "https://www.example.com" # use the custom claude API endpoint @@ -429,6 +437,7 @@ nonstream-keepalive-interval: 0 # X-Custom-Header: "custom-value" # api-key-entries: # - api-key: "sk-or-v1-...b780" +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential # - api-key: "sk-or-v1-...b781" # without proxy-url @@ -456,6 +465,7 @@ nonstream-keepalive-interval: 0 # Vertex API keys (Vertex-compatible endpoints, base-url is optional) # vertex-api-key: # - api-key: "vk-123..." # x-goog-api-key header +# weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/vertex-pro" to target this credential # base-url: "https://example.com/api" # optional, e.g. https://zenmux.ai/api; falls back to Google Vertex when omitted # proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override diff --git a/internal/api/handlers/management/auth_files_crud.go b/internal/api/handlers/management/auth_files_crud.go index c7416334..496930ff 100644 --- a/internal/api/handlers/management/auth_files_crud.go +++ b/internal/api/handlers/management/auth_files_crud.go @@ -498,7 +498,11 @@ func (h *Handler) buildAuthFromFileData(path string, data []byte) (*coreauth.Aut Now: time.Now(), IDGenerator: synthesizer.NewStableIDGenerator(), } - if generated := synthesizer.SynthesizeAuthFile(sctx, path, data); len(generated) > 0 && generated[0] != nil { + generated, errSynthesize := synthesizer.SynthesizeAuthFile(sctx, path, data) + if errSynthesize != nil { + return nil, fmt.Errorf("invalid auth file: %w", errSynthesize) + } + if len(generated) > 0 && generated[0] != nil { auth = generated[0].Clone() } } diff --git a/internal/api/handlers/management/auth_files_fields.go b/internal/api/handlers/management/auth_files_fields.go index a1d633d6..3dc4411e 100644 --- a/internal/api/handlers/management/auth_files_fields.go +++ b/internal/api/handlers/management/auth_files_fields.go @@ -15,6 +15,7 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) @@ -272,7 +273,25 @@ func (h *Handler) PatchAuthFileFields(c *gin.Context) { targetAuth.Metadata = make(map[string]any) } - if fieldPath == "headers" { + if fieldPath == coreauth.AttributeWeight { + if value == nil { + delete(targetAuth.Metadata, coreauth.AttributeWeight) + } else { + if _, okNumber := value.(json.Number); !okNumber { + c.JSON(http.StatusBadRequest, gin.H{"error": "weight must be an integer"}) + return + } + weight, errWeight := credentialweight.ParseValue(value) + if errWeight != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": errWeight.Error()}) + return + } + targetAuth.Metadata[coreauth.AttributeWeight] = weight + } + } else if rootAuthFileField(fieldPath) == coreauth.AttributeWeight { + c.JSON(http.StatusBadRequest, gin.H{"error": "weight does not support nested fields"}) + return + } else if fieldPath == "headers" { applyAuthFileHeadersPatch(targetAuth, value) } else if errSet := setAuthFileMetadataValue(targetAuth.Metadata, fieldPath, value); errSet != nil { c.JSON(http.StatusBadRequest, gin.H{"error": errSet.Error()}) @@ -429,6 +448,9 @@ func syncAuthFileMetadataFields(auth *coreauth.Auth, touchedRoots map[string]str if _, ok := touchedRoots["priority"]; ok { syncAuthFilePriorityAttribute(auth) } + if _, ok := touchedRoots[coreauth.AttributeWeight]; ok { + syncAuthFileWeightAttribute(auth) + } if _, ok := touchedRoots["note"]; ok { syncAuthFileNoteAttribute(auth) } @@ -476,6 +498,21 @@ func syncAuthFilePriorityAttribute(auth *coreauth.Auth) { auth.Attributes["priority"] = strconv.Itoa(priority) } +func syncAuthFileWeightAttribute(auth *coreauth.Auth) { + if auth == nil { + return + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + weight, errWeight := credentialweight.ParseValue(auth.Metadata[coreauth.AttributeWeight]) + if errWeight != nil { + delete(auth.Attributes, coreauth.AttributeWeight) + return + } + auth.Attributes[coreauth.AttributeWeight] = strconv.FormatInt(weight, 10) +} + func authFileIntValue(value any) (int, bool) { switch typed := value.(type) { case int: diff --git a/internal/api/handlers/management/auth_files_patch_fields_test.go b/internal/api/handlers/management/auth_files_patch_fields_test.go index e01f1d5c..c37d3757 100644 --- a/internal/api/handlers/management/auth_files_patch_fields_test.go +++ b/internal/api/handlers/management/auth_files_patch_fields_test.go @@ -276,3 +276,96 @@ func TestPatchAuthFileFields_ArbitraryFieldsPersistToFile(t *testing.T) { t.Fatalf("fgh.ijk = %#v, want true", got) } } + +func TestPatchAuthFileFields_WeightPersistsAndSyncsRuntime(t *testing.T) { + t.Setenv("MANAGEMENT_PASSWORD", "") + + authDir := t.TempDir() + fileName := "weighted.json" + filePath := filepath.Join(authDir, fileName) + store := fileauth.NewFileTokenStore() + store.SetBaseDir(authDir) + manager := coreauth.NewManager(store, nil, nil) + record := &coreauth.Auth{ + ID: fileName, + FileName: fileName, + Provider: "codex", + Attributes: map[string]string{"path": filePath}, + Metadata: map[string]any{"type": "codex"}, + } + if _, errRegister := manager.Register(context.Background(), record); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager) + + patch := func(weight string) *httptest.ResponseRecorder { + t.Helper() + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + body := `{"name":"weighted.json","weight":` + weight + `}` + ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + h.PatchAuthFileFields(ctx) + return rec + } + + if rec := patch("7"); rec.Code != http.StatusOK { + t.Fatalf("update status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + updated, ok := manager.GetByID(fileName) + if !ok || updated.Attributes[coreauth.AttributeWeight] != "7" { + t.Fatalf("runtime weight = %#v, want 7", updated) + } + raw, errRead := os.ReadFile(filePath) + if errRead != nil { + t.Fatalf("ReadFile() error = %v", errRead) + } + var persisted map[string]any + if errUnmarshal := json.Unmarshal(raw, &persisted); errUnmarshal != nil { + t.Fatalf("Unmarshal() error = %v", errUnmarshal) + } + if persisted["weight"] != float64(7) { + t.Fatalf("persisted weight = %#v, want 7", persisted["weight"]) + } + + if rec := patch("null"); rec.Code != http.StatusOK { + t.Fatalf("reset status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + updated, _ = manager.GetByID(fileName) + if _, exists := updated.Attributes[coreauth.AttributeWeight]; exists { + t.Fatal("runtime weight remains after reset") + } + raw, errRead = os.ReadFile(filePath) + if errRead != nil { + t.Fatalf("ReadFile() after reset error = %v", errRead) + } + persisted = nil + if errUnmarshal := json.Unmarshal(raw, &persisted); errUnmarshal != nil { + t.Fatalf("Unmarshal() after reset error = %v", errUnmarshal) + } + if _, exists := persisted["weight"]; exists { + t.Fatal("persisted weight remains after reset") + } +} + +func TestPatchAuthFileFields_RejectsInvalidWeights(t *testing.T) { + store := &memoryAuthStore{} + manager := coreauth.NewManager(store, nil, nil) + record := &coreauth.Auth{ID: "auth.json", FileName: "auth.json", Provider: "codex", Metadata: map[string]any{"type": "codex"}} + if _, errRegister := manager.Register(context.Background(), record); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + h := NewHandlerWithoutConfigFilePath(&config.Config{}, manager) + + for _, weight := range []string{"1.5", "1000001", "9223372036854775808", `"7"`} { + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + body := `{"name":"auth.json","weight":` + weight + `}` + ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/auth-files/fields", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + h.PatchAuthFileFields(ctx) + if rec.Code != http.StatusBadRequest { + t.Fatalf("weight %s status = %d, want 400; body=%s", weight, rec.Code, rec.Body.String()) + } + } +} diff --git a/internal/api/handlers/management/config_basic.go b/internal/api/handlers/management/config_basic.go index a0818aa8..a7af35b9 100644 --- a/internal/api/handlers/management/config_basic.go +++ b/internal/api/handlers/management/config_basic.go @@ -284,6 +284,8 @@ func normalizeRoutingStrategy(strategy string) (string, bool) { switch normalized { case "", "round-robin", "roundrobin", "rr": return "round-robin", true + case "weighted-round-robin", "weightedroundrobin", "wrr": + return "weighted-round-robin", true case "fill-first", "fillfirst", "ff": return "fill-first", true default: diff --git a/internal/api/handlers/management/config_basic_weight_test.go b/internal/api/handlers/management/config_basic_weight_test.go new file mode 100644 index 00000000..427690da --- /dev/null +++ b/internal/api/handlers/management/config_basic_weight_test.go @@ -0,0 +1,12 @@ +package management + +import "testing" + +func TestNormalizeRoutingStrategyWeightedRoundRobin(t *testing.T) { + for _, input := range []string{"weighted-round-robin", "weightedroundrobin", "wrr"} { + got, ok := normalizeRoutingStrategy(input) + if !ok || got != "weighted-round-robin" { + t.Fatalf("normalizeRoutingStrategy(%q) = %q, %v; want weighted-round-robin, true", input, got, ok) + } + } +} diff --git a/internal/api/handlers/management/config_lists.go b/internal/api/handlers/management/config_lists.go index b4138127..541184ff 100644 --- a/internal/api/handlers/management/config_lists.go +++ b/internal/api/handlers/management/config_lists.go @@ -1,6 +1,7 @@ package management import ( + "bytes" "encoding/json" "fmt" "strings" @@ -9,6 +10,32 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/config" ) +func parseCredentialWeightPatch(raw json.RawMessage) (*int, error) { + if len(raw) == 0 { + return nil, fmt.Errorf("weight is missing") + } + if bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return nil, nil + } + var weight int + decoder := json.NewDecoder(bytes.NewReader(raw)) + if errDecode := decoder.Decode(&weight); errDecode != nil { + return nil, fmt.Errorf("weight must be an integer") + } + if errValidate := config.ValidateCredentialWeight(&weight); errValidate != nil { + return nil, errValidate + } + return &weight, nil +} + +func rejectInvalidCredentialWeight(c *gin.Context, field string, weight *int) bool { + if errValidate := config.ValidateCredentialWeight(weight); errValidate != nil { + c.JSON(400, gin.H{"error": fmt.Sprintf("%s: %v", field, errValidate)}) + return true + } + return false +} + // Generic helpers for list[string] func (h *Handler) putStringList(c *gin.Context, set func([]string), after func()) { data, err := c.GetRawData() @@ -139,6 +166,11 @@ func (h *Handler) PutGeminiKeys(c *gin.Context) { } arr = obj.Items } + for index := range arr { + if rejectInvalidCredentialWeight(c, fmt.Sprintf("gemini-api-key[%d].weight", index), arr[index].Weight) { + return + } + } h.mu.Lock() defer h.mu.Unlock() h.cfg.GeminiKey = append([]config.GeminiKey(nil), arr...) @@ -148,6 +180,7 @@ func (h *Handler) PutGeminiKeys(c *gin.Context) { func (h *Handler) PatchGeminiKey(c *gin.Context) { type geminiKeyPatch struct { APIKey *string `json:"api-key"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` @@ -197,6 +230,14 @@ func (h *Handler) PatchGeminiKey(c *gin.Context) { } entry.APIKey = trimmed } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } @@ -298,6 +339,11 @@ func (h *Handler) PutInteractionsKeys(c *gin.Context) { } arr = obj.Items } + for index := range arr { + if rejectInvalidCredentialWeight(c, fmt.Sprintf("interactions-api-key[%d].weight", index), arr[index].Weight) { + return + } + } h.mu.Lock() defer h.mu.Unlock() h.cfg.InteractionsKey = append([]config.GeminiKey(nil), arr...) @@ -307,6 +353,7 @@ func (h *Handler) PutInteractionsKeys(c *gin.Context) { func (h *Handler) PatchInteractionsKey(c *gin.Context) { type geminiKeyPatch struct { APIKey *string `json:"api-key"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` @@ -357,6 +404,14 @@ func (h *Handler) PatchInteractionsKey(c *gin.Context) { } entry.APIKey = trimmed } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } @@ -459,6 +514,9 @@ func (h *Handler) PutClaudeKeys(c *gin.Context) { } for i := range arr { normalizeClaudeKey(&arr[i]) + if rejectInvalidCredentialWeight(c, fmt.Sprintf("claude-api-key[%d].weight", i), arr[i].Weight) { + return + } } h.mu.Lock() defer h.mu.Unlock() @@ -469,6 +527,7 @@ func (h *Handler) PutClaudeKeys(c *gin.Context) { func (h *Handler) PatchClaudeKey(c *gin.Context) { type claudeKeyPatch struct { APIKey *string `json:"api-key"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` @@ -511,6 +570,14 @@ func (h *Handler) PatchClaudeKey(c *gin.Context) { if body.Value.APIKey != nil { entry.APIKey = strings.TrimSpace(*body.Value.APIKey) } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } @@ -615,9 +682,16 @@ func (h *Handler) PutOpenAICompat(c *gin.Context) { filtered := make([]config.OpenAICompatibility, 0, len(arr)) for i := range arr { normalizeOpenAICompatibilityEntry(&arr[i]) - if strings.TrimSpace(arr[i].BaseURL) != "" { - filtered = append(filtered, arr[i]) + if strings.TrimSpace(arr[i].BaseURL) == "" { + continue + } + for keyIndex := range arr[i].APIKeyEntries { + field := fmt.Sprintf("openai-compatibility[%d].api-key-entries[%d].weight", i, keyIndex) + if rejectInvalidCredentialWeight(c, field, arr[i].APIKeyEntries[keyIndex].Weight) { + return + } } + filtered = append(filtered, arr[i]) } h.mu.Lock() defer h.mu.Unlock() @@ -690,6 +764,12 @@ func (h *Handler) PatchOpenAICompat(c *gin.Context) { entry.BaseURL = trimmed } if body.Value.APIKeyEntries != nil { + for keyIndex := range *body.Value.APIKeyEntries { + weight := (*body.Value.APIKeyEntries)[keyIndex].Weight + if rejectInvalidCredentialWeight(c, fmt.Sprintf("api-key-entries[%d].weight", keyIndex), weight) { + return + } + } entry.APIKeyEntries = append([]config.OpenAICompatibilityAPIKey(nil), (*body.Value.APIKeyEntries)...) } if body.Value.Models != nil { @@ -759,6 +839,9 @@ func (h *Handler) PutVertexCompatKeys(c *gin.Context) { c.JSON(400, gin.H{"error": fmt.Sprintf("vertex-api-key[%d].api-key is required", i)}) return } + if rejectInvalidCredentialWeight(c, fmt.Sprintf("vertex-api-key[%d].weight", i), arr[i].Weight) { + return + } } h.mu.Lock() defer h.mu.Unlock() @@ -769,6 +852,7 @@ func (h *Handler) PutVertexCompatKeys(c *gin.Context) { func (h *Handler) PatchVertexCompatKey(c *gin.Context) { type vertexCompatPatch struct { APIKey *string `json:"api-key"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` @@ -819,6 +903,14 @@ func (h *Handler) PatchVertexCompatKey(c *gin.Context) { } entry.APIKey = trimmed } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } @@ -1114,6 +1206,9 @@ func (h *Handler) PutCodexKeys(c *gin.Context) { if entry.BaseURL == "" { continue } + if rejectInvalidCredentialWeight(c, fmt.Sprintf("codex-api-key[%d].weight", i), entry.Weight) { + return + } filtered = append(filtered, entry) } h.mu.Lock() @@ -1125,6 +1220,7 @@ func (h *Handler) PutCodexKeys(c *gin.Context) { func (h *Handler) PatchCodexKey(c *gin.Context) { type codexKeyPatch struct { APIKey *string `json:"api-key"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` ProxyURL *string `json:"proxy-url"` @@ -1166,6 +1262,14 @@ func (h *Handler) PatchCodexKey(c *gin.Context) { if body.Value.APIKey != nil { entry.APIKey = strings.TrimSpace(*body.Value.APIKey) } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } @@ -1279,6 +1383,9 @@ func (h *Handler) PutXAIKeys(c *gin.Context) { if entry.BaseURL == "" { continue } + if rejectInvalidCredentialWeight(c, fmt.Sprintf("xai-api-key[%d].weight", i), entry.Weight) { + return + } filtered = append(filtered, entry) } h.mu.Lock() @@ -1292,6 +1399,7 @@ func (h *Handler) PatchXAIKey(c *gin.Context) { type xaiKeyPatch struct { APIKey *string `json:"api-key"` Priority *int `json:"priority"` + Weight json.RawMessage `json:"weight"` Prefix *string `json:"prefix"` BaseURL *string `json:"base-url"` Websockets *bool `json:"websockets"` @@ -1338,6 +1446,14 @@ func (h *Handler) PatchXAIKey(c *gin.Context) { if body.Value.Priority != nil { entry.Priority = *body.Value.Priority } + if len(body.Value.Weight) > 0 { + weight, errWeight := parseCredentialWeightPatch(body.Value.Weight) + if errWeight != nil { + c.JSON(400, gin.H{"error": errWeight.Error()}) + return + } + entry.Weight = weight + } if body.Value.Prefix != nil { entry.Prefix = strings.TrimSpace(*body.Value.Prefix) } diff --git a/internal/api/handlers/management/config_weight_test.go b/internal/api/handlers/management/config_weight_test.go new file mode 100644 index 00000000..6442dd8f --- /dev/null +++ b/internal/api/handlers/management/config_weight_test.go @@ -0,0 +1,104 @@ +package management + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestPatchAPIKeyWeightForEveryFamily(t *testing.T) { + tests := []struct { + name string + setup func(*config.Config) + patch func(*Handler, *gin.Context) + get func(*config.Config) *int + }{ + {name: "gemini", setup: func(cfg *config.Config) { cfg.GeminiKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchGeminiKey, get: func(cfg *config.Config) *int { return cfg.GeminiKey[0].Weight }}, + {name: "interactions", setup: func(cfg *config.Config) { cfg.InteractionsKey = []config.GeminiKey{{APIKey: "key"}} }, patch: (*Handler).PatchInteractionsKey, get: func(cfg *config.Config) *int { return cfg.InteractionsKey[0].Weight }}, + {name: "claude", setup: func(cfg *config.Config) { cfg.ClaudeKey = []config.ClaudeKey{{APIKey: "key"}} }, patch: (*Handler).PatchClaudeKey, get: func(cfg *config.Config) *int { return cfg.ClaudeKey[0].Weight }}, + {name: "vertex", setup: func(cfg *config.Config) { + cfg.VertexCompatAPIKey = []config.VertexCompatKey{{APIKey: "key", BaseURL: "https://example.com"}} + }, patch: (*Handler).PatchVertexCompatKey, get: func(cfg *config.Config) *int { return cfg.VertexCompatAPIKey[0].Weight }}, + {name: "codex", setup: func(cfg *config.Config) { + cfg.CodexKey = []config.CodexKey{{APIKey: "key", BaseURL: "https://example.com"}} + }, patch: (*Handler).PatchCodexKey, get: func(cfg *config.Config) *int { return cfg.CodexKey[0].Weight }}, + {name: "xai", setup: func(cfg *config.Config) { + cfg.XAIKey = []config.XAIKey{{APIKey: "key", BaseURL: "https://example.com"}} + }, patch: (*Handler).PatchXAIKey, get: func(cfg *config.Config) *int { return cfg.XAIKey[0].Weight }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := &config.Config{} + test.setup(cfg) + h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)} + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/key", strings.NewReader(`{"index":0,"value":{"weight":7}}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + test.patch(h, ctx) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + if weight := test.get(cfg); weight == nil || *weight != 7 { + t.Fatalf("weight = %v, want 7", weight) + } + }) + } +} + +func TestPatchAPIKeyWeightResetAndStrictValidation(t *testing.T) { + initial := 5 + cfg := &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key", Weight: &initial}}} + h := &Handler{cfg: cfg, configFilePath: writeTestConfigFile(t)} + + patch := func(raw string) *httptest.ResponseRecorder { + t.Helper() + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + body := fmt.Sprintf(`{"index":0,"value":{"weight":%s}}`, raw) + ctx.Request = httptest.NewRequest(http.MethodPatch, "/v0/management/gemini-api-key", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + h.PatchGeminiKey(ctx) + return rec + } + + for _, invalid := range []string{"1.5", "1000001", "9223372036854775808", `"7"`} { + rec := patch(invalid) + if rec.Code != http.StatusBadRequest { + t.Fatalf("weight %s status = %d, want 400; body=%s", invalid, rec.Code, rec.Body.String()) + } + if cfg.GeminiKey[0].Weight == nil || *cfg.GeminiKey[0].Weight != initial { + t.Fatalf("invalid weight %s changed config", invalid) + } + } + + if rec := patch("null"); rec.Code != http.StatusOK { + t.Fatalf("reset status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + if cfg.GeminiKey[0].Weight != nil { + t.Fatalf("reset weight = %v, want nil default", cfg.GeminiKey[0].Weight) + } +} + +func TestPutAPIKeyWeightRejectsAboveMaximum(t *testing.T) { + h := &Handler{cfg: &config.Config{}, configFilePath: writeTestConfigFile(t)} + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodPut, "/v0/management/gemini-api-key", strings.NewReader(`[{"api-key":"key","weight":1000001}]`)) + ctx.Request.Header.Set("Content-Type", "application/json") + h.PutGeminiKeys(ctx) + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body=%s", rec.Code, rec.Body.String()) + } + if len(h.cfg.GeminiKey) != 0 { + t.Fatal("invalid PUT changed config") + } +} diff --git a/internal/config/config_load.go b/internal/config/config_load.go index 61570320..c5e6beaf 100644 --- a/internal/config/config_load.go +++ b/internal/config/config_load.go @@ -51,6 +51,15 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { return cfg, nil } + if errValidate := validateCredentialWeightYAML(data); errValidate != nil { + if optional { + cfgOptional := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} + cfgOptional.NormalizePluginsConfig() + return cfgOptional, nil + } + return nil, errValidate + } + // Unmarshal the YAML data into the Config struct. var cfg Config // Set defaults before unmarshal so that absent keys keep defaults. @@ -86,6 +95,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { if errValidate := cfg.Codex.LiveMediaRelay.Validate(); errValidate != nil { return nil, errValidate } + if errValidate := cfg.ValidateCredentialWeights(); errValidate != nil { + return nil, errValidate + } // Hash remote management key if plaintext is detected (nested) // We consider a value to be already hashed if it looks like a bcrypt hash ($2a$, $2b$, or $2y$ prefix). diff --git a/internal/config/config_types.go b/internal/config/config_types.go index c0eaa1d6..fd7f6f69 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -203,7 +203,7 @@ type QuotaExceeded struct { // RoutingConfig configures how credentials are selected for requests. type RoutingConfig struct { // Strategy selects the credential selection strategy. - // Supported values: "round-robin" (default), "fill-first". + // Supported values: "round-robin" (default), "weighted-round-robin", "fill-first". Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` // SessionAffinity enables universal session-sticky routing for all clients. @@ -317,6 +317,10 @@ type ClaudeKey struct { // Higher values are preferred; defaults to 0. Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + // Weight controls proportional selection under weighted-round-robin. + // An omitted value defaults to 1; non-positive values exclude this credential; maximum 1,000,000. + Weight *int `yaml:"weight,omitempty" json:"weight,omitempty"` + // Prefix optionally namespaces models for this credential (e.g., "teamA/claude-sonnet-4"). Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` @@ -388,6 +392,10 @@ type CodexKey struct { // Higher values are preferred; defaults to 0. Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + // Weight controls proportional selection under weighted-round-robin. + // An omitted value defaults to 1; non-positive values exclude this credential; maximum 1,000,000. + Weight *int `yaml:"weight,omitempty" json:"weight,omitempty"` + // Prefix optionally namespaces models for this credential (e.g., "teamA/gpt-5-codex"). Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` @@ -457,6 +465,10 @@ type GeminiKey struct { // Higher values are preferred; defaults to 0. Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + // Weight controls proportional selection under weighted-round-robin. + // An omitted value defaults to 1; non-positive values exclude this credential; maximum 1,000,000. + Weight *int `yaml:"weight,omitempty" json:"weight,omitempty"` + // Prefix optionally namespaces models for this credential (e.g., "teamA/gemini-3-pro-preview"). Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` @@ -543,6 +555,10 @@ type OpenAICompatibilityAPIKey struct { // APIKey is the authentication key for accessing the external API services. APIKey string `yaml:"api-key" json:"api-key"` + // Weight controls proportional selection under weighted-round-robin. + // An omitted value defaults to 1; non-positive values exclude this credential; maximum 1,000,000. + Weight *int `yaml:"weight,omitempty" json:"weight,omitempty"` + // ProxyURL overrides the global proxy setting for this API key if provided. ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` } diff --git a/internal/config/config_yaml.go b/internal/config/config_yaml.go index 69f4490b..d5dfd693 100644 --- a/internal/config/config_yaml.go +++ b/internal/config/config_yaml.go @@ -311,6 +311,11 @@ func appendPath(path []string, key string) []string { // represents a known default value that should not be written to the config file. // This prevents non-zero defaults from polluting the config. func isKnownDefaultValue(path []string, node *yaml.Node) bool { + // Weight is pointer-backed, so an explicit zero is meaningful and must be preserved. + if len(path) > 0 && path[len(path)-1] == "weight" && node != nil && node.Kind == yaml.ScalarNode && node.Tag == "!!int" { + return false + } + // First check if it's a zero value if isZeroValueNode(node) { return true diff --git a/internal/config/parse.go b/internal/config/parse.go index 5c1f59fa..ba6af9f9 100644 --- a/internal/config/parse.go +++ b/internal/config/parse.go @@ -16,6 +16,10 @@ func ParseConfigBytes(data []byte) (*Config, error) { return nil, fmt.Errorf("config payload is empty") } + if errValidate := validateCredentialWeightYAML(data); errValidate != nil { + return nil, errValidate + } + var cfg Config // Keep defaults aligned with LoadConfigOptional. cfg.Host = "" // Default empty: binds to all interfaces (IPv4 + IPv6) @@ -42,6 +46,9 @@ func ParseConfigBytes(data []byte) (*Config, error) { if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { return nil, errValidate } + if errValidate := cfg.ValidateCredentialWeights(); errValidate != nil { + return nil, errValidate + } // Hash remote management key if plaintext is detected (nested), but do NOT persist. if cfg.RemoteManagement.SecretKey != "" && !looksLikeBcrypt(cfg.RemoteManagement.SecretKey) { diff --git a/internal/config/vertex_compat.go b/internal/config/vertex_compat.go index 2d3d9014..4a73f98c 100644 --- a/internal/config/vertex_compat.go +++ b/internal/config/vertex_compat.go @@ -17,6 +17,10 @@ type VertexCompatKey struct { // Higher values are preferred; defaults to 0. Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + // Weight controls proportional selection under weighted-round-robin. + // An omitted value defaults to 1; non-positive values exclude this credential; maximum 1,000,000. + Weight *int `yaml:"weight,omitempty" json:"weight,omitempty"` + // Prefix optionally namespaces model aliases for this credential (e.g., "teamA/vertex-pro"). Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` diff --git a/internal/config/weight.go b/internal/config/weight.go new file mode 100644 index 00000000..e67ff727 --- /dev/null +++ b/internal/config/weight.go @@ -0,0 +1,153 @@ +package config + +import ( + "fmt" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight" + "gopkg.in/yaml.v3" +) + +// MaxCredentialWeight is the largest positive credential routing weight. +const MaxCredentialWeight = int(credentialweight.Max) + +// ValidateCredentialWeight validates one optional config credential weight. +func ValidateCredentialWeight(weight *int) error { + if weight == nil { + return nil + } + _, errNormalize := credentialweight.Normalize(int64(*weight)) + return errNormalize +} + +func validateCredentialWeightYAML(data []byte) error { + var document yaml.Node + if errUnmarshal := yaml.Unmarshal(data, &document); errUnmarshal != nil { + return nil + } + if len(document.Content) == 0 { + return nil + } + root := document.Content[0] + families := map[string]struct{}{ + "gemini-api-key": {}, "interactions-api-key": {}, "claude-api-key": {}, + "vertex-api-key": {}, "codex-api-key": {}, "xai-api-key": {}, + } + for index := 0; root != nil && root.Kind == yaml.MappingNode && index+1 < len(root.Content); index += 2 { + name := root.Content[index].Value + value := root.Content[index+1] + if _, ok := families[name]; ok { + if errValidate := validateWeightSequenceNode(value, name); errValidate != nil { + return errValidate + } + continue + } + if name == "openai-compatibility" { + if errValidate := validateOpenAICompatibilityWeightNodes(value); errValidate != nil { + return errValidate + } + } + } + return nil +} + +func validateWeightSequenceNode(sequence *yaml.Node, path string) error { + if sequence == nil || sequence.Kind != yaml.SequenceNode { + return nil + } + for index, item := range sequence.Content { + if errValidate := validateWeightMappingNode(item, fmt.Sprintf("%s[%d]", path, index)); errValidate != nil { + return errValidate + } + } + return nil +} + +func validateWeightMappingNode(mapping *yaml.Node, path string) error { + if mapping == nil || mapping.Kind != yaml.MappingNode { + return nil + } + for index := 0; index+1 < len(mapping.Content); index += 2 { + if mapping.Content[index].Value != "weight" { + continue + } + value := mapping.Content[index+1] + if value.Kind != yaml.ScalarNode || value.Tag != "!!int" { + return fmt.Errorf("%s.weight: weight must be an integer", path) + } + var weight int64 + if errDecode := value.Decode(&weight); errDecode != nil { + return fmt.Errorf("%s.weight: weight must be an integer", path) + } + if _, errNormalize := credentialweight.Normalize(weight); errNormalize != nil { + return fmt.Errorf("%s.weight: %w", path, errNormalize) + } + } + return nil +} + +func validateOpenAICompatibilityWeightNodes(sequence *yaml.Node) error { + if sequence == nil || sequence.Kind != yaml.SequenceNode { + return nil + } + for providerIndex, provider := range sequence.Content { + if provider == nil || provider.Kind != yaml.MappingNode { + continue + } + for index := 0; index+1 < len(provider.Content); index += 2 { + if provider.Content[index].Value != "api-key-entries" { + continue + } + path := fmt.Sprintf("openai-compatibility[%d].api-key-entries", providerIndex) + if errValidate := validateWeightSequenceNode(provider.Content[index+1], path); errValidate != nil { + return errValidate + } + } + } + return nil +} + +// ValidateCredentialWeights validates weights for every API-key family. +func (cfg *Config) ValidateCredentialWeights() error { + if cfg == nil { + return nil + } + for index := range cfg.GeminiKey { + if errValidate := ValidateCredentialWeight(cfg.GeminiKey[index].Weight); errValidate != nil { + return fmt.Errorf("gemini-api-key[%d].weight: %w", index, errValidate) + } + } + for index := range cfg.InteractionsKey { + if errValidate := ValidateCredentialWeight(cfg.InteractionsKey[index].Weight); errValidate != nil { + return fmt.Errorf("interactions-api-key[%d].weight: %w", index, errValidate) + } + } + for index := range cfg.ClaudeKey { + if errValidate := ValidateCredentialWeight(cfg.ClaudeKey[index].Weight); errValidate != nil { + return fmt.Errorf("claude-api-key[%d].weight: %w", index, errValidate) + } + } + for index := range cfg.VertexCompatAPIKey { + if errValidate := ValidateCredentialWeight(cfg.VertexCompatAPIKey[index].Weight); errValidate != nil { + return fmt.Errorf("vertex-api-key[%d].weight: %w", index, errValidate) + } + } + for index := range cfg.CodexKey { + if errValidate := ValidateCredentialWeight(cfg.CodexKey[index].Weight); errValidate != nil { + return fmt.Errorf("codex-api-key[%d].weight: %w", index, errValidate) + } + } + for index := range cfg.XAIKey { + if errValidate := ValidateCredentialWeight(cfg.XAIKey[index].Weight); errValidate != nil { + return fmt.Errorf("xai-api-key[%d].weight: %w", index, errValidate) + } + } + for providerIndex := range cfg.OpenAICompatibility { + for keyIndex := range cfg.OpenAICompatibility[providerIndex].APIKeyEntries { + weight := cfg.OpenAICompatibility[providerIndex].APIKeyEntries[keyIndex].Weight + if errValidate := ValidateCredentialWeight(weight); errValidate != nil { + return fmt.Errorf("openai-compatibility[%d].api-key-entries[%d].weight: %w", providerIndex, keyIndex, errValidate) + } + } + } + return nil +} diff --git a/internal/config/weight_test.go b/internal/config/weight_test.go new file mode 100644 index 00000000..d3a508bf --- /dev/null +++ b/internal/config/weight_test.go @@ -0,0 +1,62 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestAPIKeyWeightValidation(t *testing.T) { + tests := []struct { + name string + weight string + valid bool + }{ + {name: "negative excludes", weight: "-1", valid: true}, + {name: "maximum", weight: "1000000", valid: true}, + {name: "fraction", weight: "1.5", valid: false}, + {name: "above maximum", weight: "1000001", valid: false}, + {name: "integer overflow", weight: "9223372036854775808", valid: false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, errParse := ParseConfigBytes([]byte("gemini-api-key:\n - api-key: key\n weight: " + test.weight + "\n")) + if (errParse == nil) != test.valid { + t.Fatalf("ParseConfigBytes(weight=%s) error = %v, want valid=%v", test.weight, errParse, test.valid) + } + }) + } +} + +func TestAPIKeyWeightParsingAndZeroPersistence(t *testing.T) { + cfg, errParse := ParseConfigBytes([]byte(`xai-api-key: + - api-key: key + base-url: https://api.x.ai/v1 + weight: 0 +`)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + if len(cfg.XAIKey) != 1 || cfg.XAIKey[0].Weight == nil || *cfg.XAIKey[0].Weight != 0 { + t.Fatalf("parsed weight = %#v, want explicit zero", cfg.XAIKey) + } + + configPath := filepath.Join(t.TempDir(), "config.yaml") + if errWrite := os.WriteFile(configPath, []byte(`xai-api-key: + - api-key: key + base-url: https://api.x.ai/v1 +`), 0644); errWrite != nil { + t.Fatalf("WriteFile() error = %v", errWrite) + } + if errSave := SaveConfigPreserveComments(configPath, cfg); errSave != nil { + t.Fatalf("SaveConfigPreserveComments() error = %v", errSave) + } + saved, errRead := os.ReadFile(configPath) + if errRead != nil { + t.Fatalf("ReadFile() error = %v", errRead) + } + if !strings.Contains(string(saved), "weight: 0") { + t.Fatalf("saved config does not preserve explicit zero weight:\n%s", saved) + } +} diff --git a/internal/config/xai_api_key_test.go b/internal/config/xai_api_key_test.go index 6e2e729d..355c9d09 100644 --- a/internal/config/xai_api_key_test.go +++ b/internal/config/xai_api_key_test.go @@ -26,6 +26,7 @@ func TestParseConfigBytesXAIAPIKeyMatchesCodexShape(t *testing.T) { cfg, errParse := ParseConfigBytes([]byte(`xai-api-key: - api-key: " xai-key " priority: 3 + weight: 5 prefix: " team-xai " base-url: " https://api.x.ai/v1 " websockets: true @@ -56,6 +57,9 @@ func TestParseConfigBytesXAIAPIKeyMatchesCodexShape(t *testing.T) { if entry.Priority != 3 { t.Fatalf("priority = %d, want 3", entry.Priority) } + if entry.Weight == nil || *entry.Weight != 5 { + t.Fatalf("weight = %v, want 5", entry.Weight) + } if entry.Prefix != "team-xai" { t.Fatalf("prefix = %q, want team-xai", entry.Prefix) } diff --git a/internal/credentialweight/weight.go b/internal/credentialweight/weight.go new file mode 100644 index 00000000..0a2de932 --- /dev/null +++ b/internal/credentialweight/weight.go @@ -0,0 +1,100 @@ +// Package credentialweight defines shared credential weight validation and parsing. +package credentialweight + +import ( + "encoding/json" + "fmt" + "math" + "strconv" + "strings" +) + +const ( + // Default is used when a credential does not define a weight. + Default int64 = 1 + // Max bounds scheduler arithmetic while allowing practical proportional routing. + Max int64 = 1_000_000 +) + +// Normalize validates and normalizes an explicit weight. Non-positive values are +// valid and normalize to zero, which excludes the credential from weighted routing. +func Normalize(weight int64) (int64, error) { + if weight <= 0 { + return 0, nil + } + if weight > Max { + return 0, fmt.Errorf("weight must not exceed %d", Max) + } + return weight, nil +} + +// ParseString parses a scheduler attribute. An empty value uses the default weight. +func ParseString(raw string) (int64, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return Default, nil + } + weight, errParse := strconv.ParseInt(raw, 10, 64) + if errParse != nil { + return 0, fmt.Errorf("weight must be an integer: %w", errParse) + } + return Normalize(weight) +} + +// ParseValue parses a JSON-compatible auth-file metadata value. +func ParseValue(value any) (int64, error) { + switch typed := value.(type) { + case int: + return Normalize(int64(typed)) + case int8: + return Normalize(int64(typed)) + case int16: + return Normalize(int64(typed)) + case int32: + return Normalize(int64(typed)) + case int64: + return Normalize(typed) + case uint: + if uint64(typed) > uint64(Max) { + return 0, fmt.Errorf("weight must not exceed %d", Max) + } + return int64(typed), nil + case uint8: + return int64(typed), nil + case uint16: + return int64(typed), nil + case uint32: + if uint64(typed) > uint64(Max) { + return 0, fmt.Errorf("weight must not exceed %d", Max) + } + return int64(typed), nil + case uint64: + if typed > uint64(Max) { + return 0, fmt.Errorf("weight must not exceed %d", Max) + } + return int64(typed), nil + case float64: + if math.IsNaN(typed) || math.IsInf(typed, 0) || math.Trunc(typed) != typed { + return 0, fmt.Errorf("weight must be an integer") + } + if typed <= 0 { + return 0, nil + } + if typed > float64(Max) { + return 0, fmt.Errorf("weight must not exceed %d", Max) + } + return int64(typed), nil + case float32: + return ParseValue(float64(typed)) + case json.Number: + weight, errParse := typed.Int64() + if errParse != nil { + return 0, fmt.Errorf("weight must be an integer: %w", errParse) + } + return Normalize(weight) + case string: + return ParseString(typed) + default: + return 0, fmt.Errorf("weight must be an integer") + } +} diff --git a/internal/credentialweight/weight_test.go b/internal/credentialweight/weight_test.go new file mode 100644 index 00000000..37a50755 --- /dev/null +++ b/internal/credentialweight/weight_test.go @@ -0,0 +1,33 @@ +package credentialweight + +import ( + "encoding/json" + "testing" +) + +func TestParseValueValidation(t *testing.T) { + tests := []struct { + name string + value any + want int64 + wantErr bool + }{ + {name: "default string", value: "", want: Default}, + {name: "negative excluded", value: json.Number("-5"), want: 0}, + {name: "fraction rejected", value: json.Number("1.5"), wantErr: true}, + {name: "maximum", value: json.Number("1000000"), want: Max}, + {name: "above maximum", value: json.Number("1000001"), wantErr: true}, + {name: "int64 overflow", value: json.Number("9223372036854775808"), wantErr: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, errParse := ParseValue(test.value) + if (errParse != nil) != test.wantErr { + t.Fatalf("ParseValue(%v) error = %v, wantErr=%v", test.value, errParse, test.wantErr) + } + if !test.wantErr && got != test.want { + t.Fatalf("ParseValue(%v) = %d, want %d", test.value, got, test.want) + } + }) + } +} diff --git a/internal/pluginhost/auth_callbacks.go b/internal/pluginhost/auth_callbacks.go index 3573999a..559de6b8 100644 --- a/internal/pluginhost/auth_callbacks.go +++ b/internal/pluginhost/auth_callbacks.go @@ -351,6 +351,9 @@ func (h *Host) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, auth.Runtime = existing.Runtime } } + if errWeight := coreauth.ValidateAuthWeight(auth); errWeight != nil { + return nil, fmt.Errorf("invalid auth weight: %w", errWeight) + } coreauth.ApplyCustomHeadersFromMetadata(auth) return auth, nil } diff --git a/internal/pluginhost/auth_callbacks_test.go b/internal/pluginhost/auth_callbacks_test.go index 2a1b325e..cc46404d 100644 --- a/internal/pluginhost/auth_callbacks_test.go +++ b/internal/pluginhost/auth_callbacks_test.go @@ -211,6 +211,34 @@ func TestHostAuthGetRuntimeCallbackReturnsRuntimeInfo(t *testing.T) { } } +func TestHostAuthSaveCallbackRejectsInvalidWeightBeforePersistence(t *testing.T) { + for _, rawWeight := range []string{`1.5`, `1000001`, `9223372036854775808`, `"invalid"`} { + t.Run(rawWeight, func(t *testing.T) { + authDir := t.TempDir() + host := New() + host.runtimeConfig = &config.Config{AuthDir: authDir} + host.SetAuthManager(coreauth.NewManager(nil, nil, nil)) + + req, errMarshal := json.Marshal(pluginapi.HostAuthSaveRequest{ + Name: "invalid.json", + JSON: json.RawMessage(`{"type":"demo","weight":` + rawWeight + `}`), + }) + if errMarshal != nil { + t.Fatalf("marshal request: %v", errMarshal) + } + if _, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAuthSave, req); errCall == nil { + t.Fatal("host.auth.save accepted an invalid weight") + } + if _, errStat := os.Stat(filepath.Join(authDir, "invalid.json")); !os.IsNotExist(errStat) { + t.Fatalf("invalid auth file was persisted: %v", errStat) + } + if auths := host.currentAuthManager().List(); len(auths) != 0 { + t.Fatalf("invalid auth was registered: %#v", auths) + } + }) + } +} + func TestHostAuthSaveCallbackWritesPhysicalFile(t *testing.T) { authDir := t.TempDir() host := New() diff --git a/internal/store/gitstore.go b/internal/store/gitstore.go index 14a3dfb6..222130a6 100644 --- a/internal/store/gitstore.go +++ b/internal/store/gitstore.go @@ -265,6 +265,9 @@ func (s *GitTokenStore) Save(_ context.Context, auth *cliproxyauth.Auth) (string if auth == nil { return "", fmt.Errorf("auth filestore: auth is nil") } + if errWeight := cliproxyauth.ValidateAuthWeight(auth); errWeight != nil { + return "", fmt.Errorf("auth filestore: %w", errWeight) + } path, err := s.resolveAuthPath(auth) if err != nil { @@ -532,6 +535,9 @@ func (s *GitTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, if err = json.Unmarshal(data, &metadata); err != nil { return nil, fmt.Errorf("unmarshal auth json: %w", err) } + if errWeight := cliproxyauth.ValidateAuthWeight(&cliproxyauth.Auth{Metadata: metadata}); errWeight != nil { + return nil, errWeight + } provider, _ := metadata["type"].(string) if provider == "" { provider = "unknown" diff --git a/internal/store/objectstore.go b/internal/store/objectstore.go index dff9211c..58ec6f25 100644 --- a/internal/store/objectstore.go +++ b/internal/store/objectstore.go @@ -160,6 +160,9 @@ func (s *ObjectTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (s if auth == nil { return "", fmt.Errorf("object store: auth is nil") } + if errWeight := cliproxyauth.ValidateAuthWeight(auth); errWeight != nil { + return "", fmt.Errorf("object store: %w", errWeight) + } path, err := s.resolveAuthPath(auth) if err != nil { @@ -574,6 +577,9 @@ func (s *ObjectTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Aut if err = json.Unmarshal(data, &metadata); err != nil { return nil, fmt.Errorf("unmarshal auth json: %w", err) } + if errWeight := cliproxyauth.ValidateAuthWeight(&cliproxyauth.Auth{Metadata: metadata}); errWeight != nil { + return nil, errWeight + } provider := strings.TrimSpace(valueAsString(metadata["type"])) if provider == "" { provider = "unknown" diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index 46e7515d..c9a22174 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -211,6 +211,9 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri if auth == nil { return "", fmt.Errorf("postgres store: auth is nil") } + if errWeight := cliproxyauth.ValidateAuthWeight(auth); errWeight != nil { + return "", fmt.Errorf("postgres store: %w", errWeight) + } path, err := s.resolveAuthPath(auth) if err != nil { @@ -319,6 +322,10 @@ func (s *PostgresStore) List(ctx context.Context) ([]*cliproxyauth.Auth, error) log.WithError(err).Warnf("postgres store: skipping auth %s with invalid json", id) continue } + if errWeight := cliproxyauth.ValidateAuthWeight(&cliproxyauth.Auth{Metadata: metadata}); errWeight != nil { + log.WithError(errWeight).Warnf("postgres store: skipping auth %s with invalid weight", id) + continue + } provider := strings.TrimSpace(valueAsString(metadata["type"])) if provider == "" { provider = "unknown" diff --git a/internal/watcher/clients.go b/internal/watcher/clients.go index ec964125..3ef58b55 100644 --- a/internal/watcher/clients.go +++ b/internal/watcher/clients.go @@ -119,7 +119,10 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string IDGenerator: synthesizer.NewStableIDGenerator(), PluginAuthParser: parser, } - if generated := synthesizer.SynthesizeAuthFile(ctx, fullPath, data); len(generated) > 0 { + generated, errSynthesize := synthesizer.SynthesizeAuthFile(ctx, fullPath, data) + if errSynthesize != nil { + log.WithError(errSynthesize).Warnf("skipping auth file %s", name) + } else if len(generated) > 0 { if pathAuths := authSliceToMap(generated); len(pathAuths) > 0 { newFileAuthsByPath[normalizedPath] = authIDSet(pathAuths) } @@ -250,7 +253,10 @@ func (w *Watcher) addOrUpdateClientLocked(path string) { IDGenerator: synthesizer.NewStableIDGenerator(), PluginAuthParser: parser, } - generated := synthesizer.SynthesizeAuthFile(sctx, path, data) + generated, errSynthesize := synthesizer.SynthesizeAuthFile(sctx, path, data) + if errSynthesize != nil { + log.WithError(errSynthesize).Warnf("skipping auth file %s", filepath.Base(path)) + } newByID := authSliceToMap(generated) w.clientsMutex.Lock() if len(newByID) > 0 { @@ -261,7 +267,9 @@ func (w *Watcher) addOrUpdateClientLocked(path string) { updates := w.computePerPathUpdatesLocked(oldByID, newByID) w.clientsMutex.Unlock() - w.persistAuthAsync(fmt.Sprintf("Sync auth %s", filepath.Base(path)), path) + if errSynthesize == nil { + w.persistAuthAsync(fmt.Sprintf("Sync auth %s", filepath.Base(path)), path) + } w.dispatchAuthUpdates(updates) redisqueue.NotifyUsageRefresh() } diff --git a/internal/watcher/synthesizer/config.go b/internal/watcher/synthesizer/config.go index 83e83d93..3b003a9c 100644 --- a/internal/watcher/synthesizer/config.go +++ b/internal/watcher/synthesizer/config.go @@ -21,12 +21,26 @@ func NewConfigSynthesizer() *ConfigSynthesizer { return &ConfigSynthesizer{} } +func addWeightToAttrs(weight *int, attrs map[string]string) { + if weight == nil { + return + } + normalized := *weight + if normalized <= 0 { + normalized = 0 + } + attrs[coreauth.AttributeWeight] = strconv.Itoa(normalized) +} + // Synthesize generates Auth entries from config API keys. func (s *ConfigSynthesizer) Synthesize(ctx *SynthesisContext) ([]*coreauth.Auth, error) { out := make([]*coreauth.Auth, 0, 32) if ctx == nil || ctx.Config == nil { return out, nil } + if errValidate := ctx.Config.ValidateCredentialWeights(); errValidate != nil { + return nil, fmt.Errorf("synthesize config API key auths: %w", errValidate) + } // Gemini API Keys out = append(out, s.synthesizeGeminiKeys(ctx)...) @@ -83,6 +97,7 @@ func (s *ConfigSynthesizer) synthesizeGeminiKeyEntries(ctx *SynthesisContext, en if entry.Priority != 0 { attrs["priority"] = strconv.Itoa(entry.Priority) } + addWeightToAttrs(entry.Weight, attrs) if base != "" { attrs["base_url"] = base } @@ -138,6 +153,7 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea if ck.Priority != 0 { attrs["priority"] = strconv.Itoa(ck.Priority) } + addWeightToAttrs(ck.Weight, attrs) if base != "" { attrs["base_url"] = base } @@ -206,6 +222,7 @@ func (s *ConfigSynthesizer) synthesizeCodexStyleKeys(ctx *SynthesisContext, entr if entry.Priority != 0 { attrs["priority"] = strconv.Itoa(entry.Priority) } + addWeightToAttrs(entry.Weight, attrs) if baseURL != "" { attrs["base_url"] = baseURL } @@ -279,6 +296,7 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor if compat.Priority != 0 { attrs["priority"] = strconv.Itoa(compat.Priority) } + addWeightToAttrs(entry.Weight, attrs) if key != "" { attrs["api_key"] = key } @@ -370,6 +388,7 @@ func (s *ConfigSynthesizer) synthesizeVertexCompat(ctx *SynthesisContext) []*cor if compat.Priority != 0 { attrs["priority"] = strconv.Itoa(compat.Priority) } + addWeightToAttrs(compat.Weight, attrs) if key != "" { attrs["api_key"] = key } diff --git a/internal/watcher/synthesizer/config_test.go b/internal/watcher/synthesizer/config_test.go index d06619ed..2ce96079 100644 --- a/internal/watcher/synthesizer/config_test.go +++ b/internal/watcher/synthesizer/config_test.go @@ -1,6 +1,8 @@ package synthesizer import ( + "strconv" + "strings" "testing" "time" @@ -743,6 +745,146 @@ func TestConfigSynthesizer_IDStability(t *testing.T) { } } +func TestConfigSynthesizer_RejectsInvalidWeightsForAllAPIKeyTypes(t *testing.T) { + invalidWeight := config.MaxCredentialWeight + 1 + tests := []struct { + name string + cfg *config.Config + wantPath string + }{ + { + name: "gemini", + cfg: &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "gemini-api-key[0].weight", + }, + { + name: "interactions", + cfg: &config.Config{InteractionsKey: []config.GeminiKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "interactions-api-key[0].weight", + }, + { + name: "claude", + cfg: &config.Config{ClaudeKey: []config.ClaudeKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "claude-api-key[0].weight", + }, + { + name: "codex", + cfg: &config.Config{CodexKey: []config.CodexKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "codex-api-key[0].weight", + }, + { + name: "xai", + cfg: &config.Config{XAIKey: []config.XAIKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "xai-api-key[0].weight", + }, + { + name: "openai compatibility", + cfg: &config.Config{OpenAICompatibility: []config.OpenAICompatibility{{ + APIKeyEntries: []config.OpenAICompatibilityAPIKey{{APIKey: "key", Weight: &invalidWeight}}, + }}}, + wantPath: "openai-compatibility[0].api-key-entries[0].weight", + }, + { + name: "vertex", + cfg: &config.Config{VertexCompatAPIKey: []config.VertexCompatKey{{APIKey: "key", Weight: &invalidWeight}}}, + wantPath: "vertex-api-key[0].weight", + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + auths, errSynthesize := NewConfigSynthesizer().Synthesize(&SynthesisContext{ + Config: testCase.cfg, + Now: time.Now(), + IDGenerator: NewStableIDGenerator(), + }) + if errSynthesize == nil { + t.Fatal("Synthesize() accepted an invalid credential weight") + } + if auths != nil { + t.Fatalf("Synthesize() auths = %#v, want nil", auths) + } + if !strings.Contains(errSynthesize.Error(), "synthesize config API key auths: "+testCase.wantPath) { + t.Fatalf("Synthesize() error = %q, want contextual path %q", errSynthesize, testCase.wantPath) + } + }) + } +} + +func TestConfigSynthesizer_OmittedWeightRemainsUnset(t *testing.T) { + auths, errSynthesize := NewConfigSynthesizer().Synthesize(&SynthesisContext{ + Config: &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key"}}}, + Now: time.Now(), + IDGenerator: NewStableIDGenerator(), + }) + if errSynthesize != nil { + t.Fatalf("Synthesize() error = %v", errSynthesize) + } + if len(auths) != 1 { + t.Fatalf("auth count = %d, want 1", len(auths)) + } + if _, exists := auths[0].Attributes[coreauth.AttributeWeight]; exists { + t.Fatal("omitted weight was added to synthesized attributes") + } +} + +func TestConfigSynthesizer_NormalizesNonPositiveWeightToZero(t *testing.T) { + weight := -5 + auths, errSynthesize := NewConfigSynthesizer().Synthesize(&SynthesisContext{ + Config: &config.Config{GeminiKey: []config.GeminiKey{{APIKey: "key", Weight: &weight}}}, + Now: time.Now(), + IDGenerator: NewStableIDGenerator(), + }) + if errSynthesize != nil { + t.Fatalf("Synthesize() error = %v", errSynthesize) + } + if len(auths) != 1 { + t.Fatalf("auth count = %d, want 1", len(auths)) + } + if gotWeight := auths[0].Attributes[coreauth.AttributeWeight]; gotWeight != "0" { + t.Fatalf("weight = %q, want 0", gotWeight) + } +} + +func TestConfigSynthesizer_PropagatesWeightsForAllAPIKeyTypes(t *testing.T) { + weight := func(value int) *int { return &value } + synth := NewConfigSynthesizer() + ctx := &SynthesisContext{ + Config: &config.Config{ + GeminiKey: []config.GeminiKey{{APIKey: "gemini", Weight: weight(1)}}, + InteractionsKey: []config.GeminiKey{{APIKey: "interactions", Weight: weight(2)}}, + ClaudeKey: []config.ClaudeKey{{APIKey: "claude", Weight: weight(3)}}, + CodexKey: []config.CodexKey{{APIKey: "codex", Weight: weight(4)}}, + XAIKey: []config.XAIKey{{APIKey: "xai", Weight: weight(5)}}, + OpenAICompatibility: []config.OpenAICompatibility{{ + Name: "compat", + BaseURL: "https://compat.example.com", + APIKeyEntries: []config.OpenAICompatibilityAPIKey{{ + APIKey: "compat", + Weight: weight(6), + }}, + }}, + VertexCompatAPIKey: []config.VertexCompatKey{{APIKey: "vertex", Weight: weight(7)}}, + }, + Now: time.Now(), + IDGenerator: NewStableIDGenerator(), + } + + auths, errSynthesize := synth.Synthesize(ctx) + if errSynthesize != nil { + t.Fatalf("Synthesize() error = %v", errSynthesize) + } + if len(auths) != 7 { + t.Fatalf("auth count = %d, want 7", len(auths)) + } + for index, auth := range auths { + wantWeight := strconv.Itoa(index + 1) + if gotWeight := auth.Attributes[coreauth.AttributeWeight]; gotWeight != wantWeight { + t.Fatalf("auth[%d] weight = %q, want %q", index, gotWeight, wantWeight) + } + } +} + func TestConfigSynthesizer_AllProviders(t *testing.T) { synth := NewConfigSynthesizer() ctx := &SynthesisContext{ diff --git a/internal/watcher/synthesizer/file.go b/internal/watcher/synthesizer/file.go index 2b19759c..dad21bd1 100644 --- a/internal/watcher/synthesizer/file.go +++ b/internal/watcher/synthesizer/file.go @@ -3,6 +3,7 @@ package synthesizer import ( "context" "encoding/json" + "fmt" "os" "path/filepath" "runtime" @@ -13,6 +14,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + log "github.com/sirupsen/logrus" ) // FileSynthesizer generates Auth entries from OAuth JSON files. @@ -50,7 +52,11 @@ func (s *FileSynthesizer) Synthesize(ctx *SynthesisContext) ([]*coreauth.Auth, e if errRead != nil || len(data) == 0 { continue } - auths := synthesizeFileAuths(ctx, full, data) + auths, errSynthesize := synthesizeFileAuths(ctx, full, data) + if errSynthesize != nil { + log.WithError(errSynthesize).Warnf("skipping auth file %s", name) + continue + } if len(auths) == 0 { continue } @@ -61,19 +67,22 @@ func (s *FileSynthesizer) Synthesize(ctx *SynthesisContext) ([]*coreauth.Auth, e // SynthesizeAuthFile generates Auth entries for one auth JSON file payload. // It shares exactly the same mapping behavior as FileSynthesizer.Synthesize. -func SynthesizeAuthFile(ctx *SynthesisContext, fullPath string, data []byte) []*coreauth.Auth { +func SynthesizeAuthFile(ctx *SynthesisContext, fullPath string, data []byte) ([]*coreauth.Auth, error) { return synthesizeFileAuths(ctx, fullPath, data) } -func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) []*coreauth.Auth { +func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) ([]*coreauth.Auth, error) { if ctx == nil || len(data) == 0 { - return nil + return nil, nil } now := ctx.Now cfg := ctx.Config var metadata map[string]any if errUnmarshal := json.Unmarshal(data, &metadata); errUnmarshal != nil { - return nil + return nil, nil + } + if errWeight := coreauth.ValidateAuthWeight(&coreauth.Auth{Metadata: metadata}); errWeight != nil { + return nil, fmt.Errorf("invalid weight in %s: %w", filepath.Base(fullPath), errWeight) } t, _ := metadata["type"].(string) provider := strings.ToLower(strings.TrimSpace(t)) @@ -90,7 +99,7 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) [] if errParse == nil && handled { auths = compactPluginAuths(auths) if len(auths) == 0 { - return nil + return nil, nil } perAccountExcluded := extractExcludedModelsFromMetadata(metadata) perAccountModelAliases := extractOAuthModelAliasesFromMetadata(metadata) @@ -118,15 +127,18 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) [] } auth.Metadata["disabled"] = true } + if errWeight := coreauth.ApplyAuthWeightMetadata(auth, metadata); errWeight != nil { + return nil, fmt.Errorf("invalid plugin auth weight in %s: %w", filepath.Base(fullPath), errWeight) + } coreauth.SetOAuthModelAliasesAttribute(auth, perAccountModelAliases) ApplyAuthExcludedModelsMeta(auth, cfg, perAccountExcluded, "oauth") coreauth.ApplyCustomHeadersFromMetadata(auth) } - return auths + return auths, nil } } if provider == "" || provider == "gemini-cli" { - return nil + return nil, nil } label := provider if email, _ := metadata["email"].(string); email != "" { @@ -196,6 +208,9 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) [] } } } + if errWeight := coreauth.ApplyAuthWeightMetadata(a, metadata); errWeight != nil { + return nil, fmt.Errorf("invalid auth weight in %s: %w", filepath.Base(fullPath), errWeight) + } // Read note from auth file. if rawNote, ok := metadata["note"]; ok { if note, isStr := rawNote.(string); isStr { @@ -217,7 +232,7 @@ func synthesizeFileAuths(ctx *SynthesisContext, fullPath string, data []byte) [] } } } - return []*coreauth.Auth{a} + return []*coreauth.Auth{a}, nil } func parsePluginFileAuths(parser PluginAuthParser, req pluginapi.AuthParseRequest) ([]*coreauth.Auth, bool, error) { @@ -243,6 +258,9 @@ func compactPluginAuths(auths []*coreauth.Auth) []*coreauth.Auth { if auth == nil { continue } + if errWeight := coreauth.ValidateAuthWeight(auth); errWeight != nil { + continue + } out = append(out, auth) } return out diff --git a/internal/watcher/synthesizer/file_test.go b/internal/watcher/synthesizer/file_test.go index caac1c13..20026080 100644 --- a/internal/watcher/synthesizer/file_test.go +++ b/internal/watcher/synthesizer/file_test.go @@ -202,7 +202,10 @@ func TestSynthesizeAuthFileExpandsPluginMultiAuths(t *testing.T) { }), } - auths := SynthesizeAuthFile(ctx, fullPath, raw) + auths, errSynthesize := SynthesizeAuthFile(ctx, fullPath, raw) + if errSynthesize != nil { + t.Fatalf("SynthesizeAuthFile() error = %v", errSynthesize) + } if len(auths) != 2 { t.Fatalf("SynthesizeAuthFile() len = %d, want two plugin auths", len(auths)) } @@ -231,6 +234,30 @@ func TestSynthesizeAuthFileExpandsPluginMultiAuths(t *testing.T) { } } +func TestSynthesizeAuthFileSkipsInvalidPluginAuthWeight(t *testing.T) { + tempDir := t.TempDir() + fullPath := filepath.Join(tempDir, "plugin.json") + ctx := &SynthesisContext{ + Config: &config.Config{}, + AuthDir: tempDir, + Now: time.Date(2026, 6, 21, 0, 0, 0, 0, time.UTC), + PluginAuthParser: multiAuthParserFunc(func(context.Context, pluginapi.AuthParseRequest) ([]*coreauth.Auth, bool, error) { + return []*coreauth.Auth{ + {ID: "invalid", Provider: "plugin", Attributes: map[string]string{coreauth.AttributeWeight: "1.5"}}, + {ID: "valid", Provider: "plugin", Attributes: map[string]string{coreauth.AttributeWeight: "0"}}, + }, true, nil + }), + } + + auths, errSynthesize := SynthesizeAuthFile(ctx, fullPath, []byte(`{"type":"plugin"}`)) + if errSynthesize != nil { + t.Fatalf("SynthesizeAuthFile() error = %v", errSynthesize) + } + if len(auths) != 1 || auths[0].ID != "valid" { + t.Fatalf("SynthesizeAuthFile() auths = %#v, want only valid zero-weight auth", auths) + } +} + func TestSynthesizeAuthFileAppliesSourceDisabledToPluginMultiAuths(t *testing.T) { tempDir := t.TempDir() fullPath := filepath.Join(tempDir, "geminicli.json") @@ -248,7 +275,10 @@ func TestSynthesizeAuthFileAppliesSourceDisabledToPluginMultiAuths(t *testing.T) }), } - auths := SynthesizeAuthFile(ctx, fullPath, raw) + auths, errSynthesize := SynthesizeAuthFile(ctx, fullPath, raw) + if errSynthesize != nil { + t.Fatalf("SynthesizeAuthFile() error = %v", errSynthesize) + } if len(auths) != 2 { t.Fatalf("SynthesizeAuthFile() len = %d, want two plugin auths", len(auths)) } @@ -276,7 +306,10 @@ func TestSynthesizeAuthFilePluginHandledEmptySuppressesBuiltin(t *testing.T) { }), } - auths := SynthesizeAuthFile(ctx, fullPath, raw) + auths, errSynthesize := SynthesizeAuthFile(ctx, fullPath, raw) + if errSynthesize != nil { + t.Fatalf("SynthesizeAuthFile() error = %v", errSynthesize) + } if len(auths) != 0 { t.Fatalf("SynthesizeAuthFile() len = %d, want plugin-handled empty result", len(auths)) } @@ -505,6 +538,63 @@ func TestFileSynthesizer_Synthesize_PriorityParsing(t *testing.T) { } } +func TestFileSynthesizer_Synthesize_WeightParsing(t *testing.T) { + tests := []struct { + name string + weight any + want string + valid bool + }{ + {name: "number", weight: 5, want: "5", valid: true}, + {name: "numeric string", weight: " 3 ", want: "3", valid: true}, + {name: "zero excludes", weight: 0, want: "0", valid: true}, + {name: "negative excludes", weight: -5, want: "0", valid: true}, + {name: "maximum", weight: 1000000, want: "1000000", valid: true}, + {name: "fraction rejected", weight: 1.5}, + {name: "above maximum rejected", weight: 1000001}, + {name: "overflow rejected", weight: "9223372036854775808"}, + {name: "invalid string", weight: "heavy"}, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + tempDir := t.TempDir() + data, errMarshal := json.Marshal(map[string]any{"type": "claude", "weight": testCase.weight}) + if errMarshal != nil { + t.Fatalf("json.Marshal() error = %v", errMarshal) + } + if errWrite := os.WriteFile(filepath.Join(tempDir, "auth.json"), data, 0644); errWrite != nil { + t.Fatalf("WriteFile() error = %v", errWrite) + } + ctx := &SynthesisContext{ + Config: &config.Config{}, + AuthDir: tempDir, + Now: time.Now(), + IDGenerator: NewStableIDGenerator(), + } + auths, errSynthesize := NewFileSynthesizer().Synthesize(ctx) + if errSynthesize != nil { + t.Fatalf("Synthesize() error = %v", errSynthesize) + } + if !testCase.valid { + if len(auths) != 0 { + t.Fatalf("auth count = %d, want invalid credential skipped", len(auths)) + } + if _, errDirect := SynthesizeAuthFile(ctx, filepath.Join(tempDir, "auth.json"), data); errDirect == nil { + t.Fatal("SynthesizeAuthFile() error = nil, want weight validation error") + } + return + } + if len(auths) != 1 { + t.Fatalf("auth count = %d, want 1", len(auths)) + } + if gotWeight := auths[0].Attributes[coreauth.AttributeWeight]; gotWeight != testCase.want { + t.Fatalf("weight = %q, want %q", gotWeight, testCase.want) + } + }) + } +} + func TestFileSynthesizer_Synthesize_OAuthExcludedModelsMerged(t *testing.T) { tempDir := t.TempDir() authData := map[string]any{ diff --git a/sdk/api/handlers/handlers_errors.go b/sdk/api/handlers/handlers_errors.go index 1f555e22..1dc05913 100644 --- a/sdk/api/handlers/handlers_errors.go +++ b/sdk/api/handlers/handlers_errors.go @@ -25,6 +25,15 @@ func statusFromError(err error) int { return 0 } +func isAuthSelectionUnavailable(err error) bool { + var authErr *coreauth.Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + code := strings.TrimSpace(authErr.Code) + return code == "auth_not_found" || code == "auth_unavailable" +} + func enrichAuthSelectionError(err error, providers []string, model string) error { if err == nil { return nil diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index d0861764..4daa2e98 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -461,7 +461,12 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context bootstrapRetries++ retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) if retryErr != nil { - bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel)) + originalBootstrapErr := executionErrorMessage(bootstrapStreamErr) + if isAuthSelectionUnavailable(retryErr) && originalBootstrapErr.StatusCode >= http.StatusInternalServerError { + bootstrapErr = originalBootstrapErr + } else { + bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel)) + } break } if retryResult == nil { diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go index 90c6316f..8abef89e 100644 --- a/sdk/auth/filestore.go +++ b/sdk/auth/filestore.go @@ -77,6 +77,9 @@ func (s *FileTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (str if auth == nil { return "", fmt.Errorf("auth filestore: auth is nil") } + if errWeight := cliproxyauth.ValidateAuthWeight(auth); errWeight != nil { + return "", fmt.Errorf("auth filestore: %w", errWeight) + } path, err := s.resolveAuthPath(auth) if err != nil { @@ -233,6 +236,9 @@ func (s *FileTokenStore) readAuthFiles(path, baseDir string) ([]*cliproxyauth.Au if err = json.Unmarshal(data, &metadata); err != nil { return nil, fmt.Errorf("unmarshal auth json: %w", err) } + if errWeight := cliproxyauth.ValidateAuthWeight(&cliproxyauth.Auth{Metadata: metadata}); errWeight != nil { + return nil, errWeight + } provider, _ := metadata["type"].(string) provider = strings.TrimSpace(provider) if strings.EqualFold(provider, "gemini") { @@ -278,6 +284,9 @@ func (s *FileTokenStore) readAuthFiles(path, baseDir string) ([]*cliproxyauth.Au } auth.Metadata["disabled"] = true } + if errWeight := cliproxyauth.ApplyAuthWeightMetadata(auth, metadata); errWeight != nil { + return nil, errWeight + } cliproxyauth.ApplyCustomHeadersFromMetadata(auth) } return auths, nil @@ -373,6 +382,9 @@ func compactPluginAuths(auths []*cliproxyauth.Auth) []*cliproxyauth.Auth { if auth == nil { continue } + if errWeight := cliproxyauth.ValidateAuthWeight(auth); errWeight != nil { + continue + } out = append(out, auth) } return out diff --git a/sdk/auth/filestore_test.go b/sdk/auth/filestore_test.go index add3e9b5..e638ffee 100644 --- a/sdk/auth/filestore_test.go +++ b/sdk/auth/filestore_test.go @@ -146,10 +146,61 @@ func TestFileTokenStoreSaveExistingMetadataSetsFileAttributes(t *testing.T) { } } +func TestFileTokenStoreSaveRejectsInvalidWeight(t *testing.T) { + baseDir := t.TempDir() + store := NewFileTokenStore() + store.SetBaseDir(baseDir) + auth := &cliproxyauth.Auth{ + ID: "invalid.json", + FileName: "invalid.json", + Metadata: map[string]any{ + "type": "test", + cliproxyauth.AttributeWeight: 1.5, + }, + } + + if _, errSave := store.Save(context.Background(), auth); errSave == nil { + t.Fatal("Save() accepted an invalid weight") + } + if _, errStat := os.Stat(filepath.Join(baseDir, auth.FileName)); !os.IsNotExist(errStat) { + t.Fatalf("invalid auth file was persisted: %v", errStat) + } +} + +func TestFileTokenStoreListSkipsInvalidPluginSourceWeight(t *testing.T) { + baseDir := t.TempDir() + path := filepath.Join(baseDir, "plugin.json") + if errWrite := os.WriteFile(path, []byte(`{"type":"plugin","weight":"invalid"}`), 0o600); errWrite != nil { + t.Fatalf("write auth file: %v", errWrite) + } + + parserCalled := false + RegisterPluginAuthParser(fileStoreMultiAuthParserFunc(func(context.Context, pluginapi.AuthParseRequest) ([]*cliproxyauth.Auth, bool, error) { + parserCalled = true + return []*cliproxyauth.Auth{{ID: "plugin.json", Provider: "plugin"}}, true, nil + })) + t.Cleanup(func() { + RegisterPluginAuthParser(nil) + }) + + store := NewFileTokenStore() + store.SetBaseDir(baseDir) + auths, errList := store.List(context.Background()) + if errList != nil { + t.Fatalf("List() error = %v", errList) + } + if parserCalled { + t.Fatal("plugin parser was called for an invalid persisted source") + } + if len(auths) != 0 { + t.Fatalf("List() returned invalid plugin auths: %#v", auths) + } +} + func TestFileTokenStoreListExpandsPluginMultiAuths(t *testing.T) { baseDir := t.TempDir() path := filepath.Join(baseDir, "geminicli.json") - if errWrite := os.WriteFile(path, []byte(`{"type":"gemini-cli","headers":{"X-Test":"value"}}`), 0o600); errWrite != nil { + if errWrite := os.WriteFile(path, []byte(`{"type":"gemini-cli","weight":3,"headers":{"X-Test":"value"}}`), 0o600); errWrite != nil { t.Fatalf("write auth file: %v", errWrite) } @@ -211,6 +262,9 @@ func TestFileTokenStoreListExpandsPluginMultiAuths(t *testing.T) { if gotHeader := auth.Attributes["header:X-Test"]; gotHeader != "value" { t.Fatalf("header:X-Test = %q, want value", gotHeader) } + if gotWeight := auth.Attributes[cliproxyauth.AttributeWeight]; gotWeight != "3" { + t.Fatalf("weight = %q, want 3", gotWeight) + } } if gotProject := auths[1].Metadata["project_id"]; gotProject != "project-a" { t.Fatalf("project_id = %#v, want project-a", gotProject) diff --git a/sdk/cliproxy/auth/classification.go b/sdk/cliproxy/auth/classification.go index b8c71718..f39864bd 100644 --- a/sdk/cliproxy/auth/classification.go +++ b/sdk/cliproxy/auth/classification.go @@ -19,6 +19,7 @@ const ( AttributeRuntimeOnly = "runtime_only" AttributeSource = "source" AttributeSourceBackend = "source_backend" + AttributeWeight = "weight" ) // AuthKind returns the credential kind using explicit metadata first and legacy diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 6d24ce46..7e143ed3 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -86,6 +86,13 @@ func nextTransientErrorRetryAfter(now time.Time) time.Time { return now.Add(time.Duration(seconds) * time.Second) } +func recoverableFailureRetryAfter(now time.Time, disableCooling bool) time.Time { + if disableCooling { + return time.Time{} + } + return nextTransientErrorRetryAfter(now) +} + // SetConfig updates the runtime config snapshot used by request-time helpers. // Callers should provide the latest config on reload so per-credential alias mapping stays in sync. func (m *Manager) SetConfig(cfg *internalconfig.Config) { @@ -820,16 +827,18 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { setModelQuota = true } case 408, 500, 502, 503, 504: - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - state.NextRetryAfter = nextTransientErrorRetryAfter(now) - } + state.NextRetryAfter = recoverableFailureRetryAfter(now, disableCooling) + state.Unavailable = !state.NextRetryAfter.IsZero() default: - state.NextRetryAfter = time.Time{} + state.NextRetryAfter = recoverableFailureRetryAfter(now, disableCooling) + state.Unavailable = !state.NextRetryAfter.IsZero() } } + if disableCooling && state.NextRetryAfter.IsZero() && state.Quota.NextRecoverAt.IsZero() { + state.Unavailable = false + state.Quota.Exceeded = false + } auth.Status = StatusError auth.UpdatedAt = now updateAggregatedAvailability(auth, now) @@ -1571,6 +1580,12 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati if isRequestScopedResultError(resultErr) { return } + defer func() { + if disableCooling && auth.NextRetryAfter.IsZero() && auth.Quota.NextRecoverAt.IsZero() { + auth.Unavailable = false + auth.Quota.Exceeded = false + } + }() auth.Unavailable = true auth.Status = StatusError auth.UpdatedAt = now @@ -1640,15 +1655,14 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati auth.NextRetryAfter = next case 408, 500, 502, 503, 504: auth.StatusMessage = "transient upstream error" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = nextTransientErrorRetryAfter(now) - } + auth.NextRetryAfter = recoverableFailureRetryAfter(now, disableCooling) + auth.Unavailable = !auth.NextRetryAfter.IsZero() default: if auth.StatusMessage == "" { auth.StatusMessage = "request failed" } + auth.NextRetryAfter = recoverableFailureRetryAfter(now, disableCooling) + auth.Unavailable = !auth.NextRetryAfter.IsZero() } } diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index 109f2f87..6e1d2535 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -2,6 +2,7 @@ package auth import ( "context" + "fmt" "strings" "time" @@ -68,6 +69,9 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { if auth == nil { return nil, nil } + if errWeight := ValidateAuthWeight(auth); errWeight != nil { + return nil, fmt.Errorf("register auth: %w", errWeight) + } if auth.ID == "" { auth.ID = uuid.NewString() } @@ -101,6 +105,9 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { if auth == nil || auth.ID == "" { return nil, nil } + if errWeight := ValidateAuthWeight(auth); errWeight != nil { + return nil, fmt.Errorf("update auth: %w", errWeight) + } m.mu.Lock() existing, ok := m.auths[auth.ID] if !ok || existing == nil { @@ -222,6 +229,9 @@ func (m *Manager) Load(ctx context.Context) error { if auth == nil || auth.ID == "" { continue } + if errWeight := ValidateAuthWeight(auth); errWeight != nil { + continue + } auth.EnsureIndex() m.auths[auth.ID] = auth.Clone() } @@ -239,6 +249,9 @@ func (m *Manager) persist(ctx context.Context, auth *Auth) error { if m.store == nil || auth == nil { return nil } + if errWeight := ValidateAuthWeight(auth); errWeight != nil { + return fmt.Errorf("persist auth: %w", errWeight) + } if shouldSkipPersist(ctx) { return nil } diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 97f1a35a..81e41b38 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -42,13 +42,42 @@ func (m *Manager) hasPluginScheduler() bool { func isBuiltInSelector(selector Selector) bool { switch selector.(type) { - case *RoundRobinSelector, *FillFirstSelector: + case *RoundRobinSelector, *WeightedRoundRobinSelector, *FillFirstSelector: return true default: return false } } +type requiredAuthKindContextKey struct{} + +type authSelectionEligibility struct { + requiredKind string + disallowFreeAuth bool +} + +func withRequiredAuthKind(ctx context.Context, requiredKind string) context.Context { + return context.WithValue(ctx, requiredAuthKindContextKey{}, requiredKind) +} + +func authSelectionEligibilityForRequest(ctx context.Context, opts cliproxyexecutor.Options) authSelectionEligibility { + eligibility := authSelectionEligibility{disallowFreeAuth: disallowFreeAuthFromMetadata(opts.Metadata)} + if ctx != nil { + eligibility.requiredKind, _ = ctx.Value(requiredAuthKindContextKey{}).(string) + } + return eligibility +} + +func (e authSelectionEligibility) allows(auth *Auth) bool { + if auth == nil { + return false + } + if e.requiredKind != "" && auth.AuthKind() != e.requiredKind { + return false + } + return !e.disallowFreeAuth || !isFreeCodexAuth(auth) +} + func (m *Manager) syncSchedulerFromSnapshot(auths []*Auth) { if m == nil || m.scheduler == nil { return @@ -444,38 +473,28 @@ func (m *Manager) pickViaBuiltinScheduler(ctx context.Context, strategy schedule return nil, false, nil } providerKey := strings.ToLower(strings.TrimSpace(provider)) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - var selected *Auth - var errPick error - if providerKey == "mixed" { + var selected *Auth + var errPick error + if providerKey == "mixed" { + selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) - } - } else { - selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) - } - } - if errPick != nil { - return nil, true, errPick } - if selected == nil { - return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue + } else { + selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) } - return selected, true, nil } + if errPick != nil { + return nil, true, errPick + } + if selected == nil { + return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + return selected, true, nil } func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}, candidates []*Auth) (*Auth, bool, error) { @@ -932,7 +951,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + eligibility := authSelectionEligibilityForRequest(ctx, opts) m.mu.RLock() selector := m.selector @@ -959,7 +978,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } - if disallowFreeAuth && isFreeCodexAuth(candidate) { + if !eligibility.allows(candidate) { continue } if _, used := tried[candidate.ID]; used { @@ -987,7 +1006,8 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op return nil, nil, errPick } if !handled { - selected, errPick = selector.Pick(ctx, provider, selectionArgForSelector(selector, model), opts, available) + selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) + selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, available) if errPick != nil { return nil, nil, errPick } @@ -1034,30 +1054,18 @@ func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, require return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} } - tried := make(map[string]struct{}) - for { - selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, tried) - if errPick != nil { - return nil, errPick - } - if selected == nil { - return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if selected.AuthKind() == requiredKind { - if m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} - } - return selected, nil - } - authID := strings.TrimSpace(selected.ID) - if authID == "" { - return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} - } - if _, alreadyTried := tried[authID]; alreadyTried { - return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} - } - tried[authID] = struct{}{} + selectionCtx := withRequiredAuthKind(ctx, requiredKind) + selected, _, errPick := m.pickNextLegacy(selectionCtx, provider, model, opts, nil) + if errPick != nil { + return nil, errPick + } + if selected == nil { + return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} } + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + return selected, nil } // SelectHomeAuthByKind selects a Home dispatch while retaining its execution scope. @@ -1115,12 +1123,16 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli if m.hasPluginScheduler() || !m.useSchedulerFastPath() { return m.pickNextLegacy(ctx, provider, model, opts, tried) } + eligibility := authSelectionEligibilityForRequest(ctx, opts) if strings.TrimSpace(model) != "" { m.mu.RLock() for _, candidate := range m.auths { if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled { continue } + if !eligibility.allows(candidate) { + continue + } if _, used := tried[candidate.ID]; used { continue } @@ -1135,37 +1147,27 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli if !okExecutor { return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) - } - if errPick != nil { - return nil, nil, errPick - } - if selected == nil { - return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() + selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) + } + if errPick != nil { + return nil, nil, errPick + } + if selected == nil { + return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() } - return authCopy, executor, nil + m.mu.Unlock() } + return authCopy, executor, nil } func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { @@ -1174,7 +1176,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + eligibility := authSelectionEligibilityForRequest(ctx, opts) providerSet := make(map[string]struct{}, len(providers)) for _, provider := range providers { @@ -1208,7 +1210,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if pinnedAuthID != "" && candidate.ID != pinnedAuthID { continue } - if disallowFreeAuth && isFreeCodexAuth(candidate) { + if !eligibility.allows(candidate) { continue } providerKey := executorKeyFromAuth(candidate) @@ -1246,7 +1248,8 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m return nil, nil, "", errPick } if !handled { - selected, errPick = selector.Pick(ctx, "mixed", selectionArgForSelector(selector, model), opts, available) + selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) + selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, available) if errPick != nil { return nil, nil, "", errPick } @@ -1299,6 +1302,7 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s if len(eligibleProviders) == 0 { return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} } + eligibility := authSelectionEligibilityForRequest(ctx, opts) if strings.TrimSpace(model) != "" { providerSet := make(map[string]struct{}, len(eligibleProviders)) for _, providerKey := range eligibleProviders { @@ -1312,6 +1316,9 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s if _, ok := providerSet[executorKeyFromAuth(candidate)]; !ok { continue } + if !eligibility.allows(candidate) { + continue + } if _, used := tried[candidate.ID]; used { continue } @@ -1323,39 +1330,29 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s m.mu.RUnlock() } - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - } - if errPick != nil { - return nil, nil, "", errPick - } - if selected == nil { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue - } - executor, okExecutor := m.Executor(providerKey) - if !okExecutor { - return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() + selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) + } + if errPick != nil { + return nil, nil, "", errPick + } + if selected == nil { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + executor, okExecutor := m.Executor(providerKey) + if !okExecutor { + return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() } - return authCopy, executor, providerKey, nil + m.mu.Unlock() } + return authCopy, executor, providerKey, nil } diff --git a/sdk/cliproxy/auth/conductor_weight_validation_test.go b/sdk/cliproxy/auth/conductor_weight_validation_test.go new file mode 100644 index 00000000..75971edf --- /dev/null +++ b/sdk/cliproxy/auth/conductor_weight_validation_test.go @@ -0,0 +1,93 @@ +package auth + +import ( + "context" + "encoding/json" + "testing" +) + +type weightValidationStore struct { + auths []*Auth + saveCount int +} + +func (s *weightValidationStore) List(context.Context) ([]*Auth, error) { + return s.auths, nil +} + +func (s *weightValidationStore) Save(context.Context, *Auth) (string, error) { + s.saveCount++ + return "", nil +} + +func (s *weightValidationStore) Delete(context.Context, string) error { + return nil +} + +func TestManagerLoadSkipsInvalidExplicitWeights(t *testing.T) { + store := &weightValidationStore{auths: []*Auth{ + {ID: "omitted", Provider: "test"}, + {ID: "zero", Provider: "test", Metadata: map[string]any{AttributeWeight: json.Number("0")}}, + {ID: "fraction", Provider: "test", Metadata: map[string]any{AttributeWeight: json.Number("1.5")}}, + {ID: "overflow", Provider: "test", Attributes: map[string]string{AttributeWeight: "9223372036854775808"}}, + }} + manager := NewManager(store, nil, nil) + + if errLoad := manager.Load(context.Background()); errLoad != nil { + t.Fatalf("Load() error = %v", errLoad) + } + if _, ok := manager.GetByID("omitted"); !ok { + t.Fatal("omitted weight auth was not loaded") + } + if _, ok := manager.GetByID("zero"); !ok { + t.Fatal("zero weight auth was not loaded") + } + for _, id := range []string{"fraction", "overflow"} { + if _, ok := manager.GetByID(id); ok { + t.Fatalf("invalid auth %q remained active after Load()", id) + } + } +} + +func TestManagerRegisterAndUpdateRejectInvalidExplicitWeights(t *testing.T) { + store := &weightValidationStore{} + manager := NewManager(store, nil, nil) + ctx := context.Background() + + invalid := &Auth{ + ID: "invalid", + Provider: "test", + Metadata: map[string]any{AttributeWeight: "nonnumeric"}, + } + if _, errRegister := manager.Register(ctx, invalid); errRegister == nil { + t.Fatal("Register() accepted an invalid weight") + } + if _, ok := manager.GetByID(invalid.ID); ok { + t.Fatal("invalid registered auth became active") + } + if store.saveCount != 0 { + t.Fatalf("invalid Register() save count = %d, want 0", store.saveCount) + } + + valid := &Auth{ + ID: "valid", + Provider: "test", + Attributes: map[string]string{AttributeWeight: "2"}, + Metadata: map[string]any{"type": "test"}, + } + if _, errRegister := manager.Register(ctx, valid); errRegister != nil { + t.Fatalf("Register(valid) error = %v", errRegister) + } + invalidUpdate := valid.Clone() + invalidUpdate.Attributes[AttributeWeight] = "1000001" + if _, errUpdate := manager.Update(ctx, invalidUpdate); errUpdate == nil { + t.Fatal("Update() accepted an invalid weight") + } + current, ok := manager.GetByID(valid.ID) + if !ok || current.Attributes[AttributeWeight] != "2" { + t.Fatalf("invalid Update() changed active auth: %#v", current) + } + if store.saveCount != 1 { + t.Fatalf("save count = %d, want only the valid Register() save", store.saveCount) + } +} diff --git a/sdk/cliproxy/auth/cooldown_backoff_test.go b/sdk/cliproxy/auth/cooldown_backoff_test.go index b1d77ebc..73a7bdcf 100644 --- a/sdk/cliproxy/auth/cooldown_backoff_test.go +++ b/sdk/cliproxy/auth/cooldown_backoff_test.go @@ -5,6 +5,9 @@ import ( "net/http" "testing" "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) func withQuotaCooldownEnabled(t *testing.T) { @@ -152,6 +155,118 @@ func TestApplyAuthFailureStateQuotaBackoffOncePerWindow(t *testing.T) { } } +func TestRecoverableUnknownFailuresHaveFiniteCooldown(t *testing.T) { + withQuotaCooldownEnabled(t) + previousTransient := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(0) + t.Cleanup(func() { transientErrorCooldownSeconds.Store(previousTransient) }) + + testCases := []struct { + name string + model string + resultErr *Error + }{ + {name: "model failure without error details", model: "gpt-5"}, + {name: "auth transport failure without status", resultErr: &Error{Message: "connection reset"}}, + } + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-unknown-" + testCase.name, Provider: "codex"} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register returned error: %v", errRegister) + } + + manager.MarkResult(context.Background(), Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Model: testCase.model, + Success: false, + Error: testCase.resultErr, + }) + + updated, ok := manager.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatal("expected auth after failure") + } + var nextRetryAfter time.Time + if testCase.model == "" { + nextRetryAfter = updated.NextRetryAfter + } else { + state := updated.ModelStates[testCase.model] + if state == nil { + t.Fatalf("expected model state for %q", testCase.model) + } + nextRetryAfter = state.NextRetryAfter + } + if nextRetryAfter.IsZero() { + t.Fatal("recoverable failure has no retry deadline") + } + if blocked, _, _ := isAuthBlockedForModel(updated, testCase.model, time.Now()); !blocked { + t.Fatal("auth was not blocked during recoverable failure cooldown") + } + if blocked, _, _ := isAuthBlockedForModel(updated, testCase.model, nextRetryAfter.Add(time.Nanosecond)); blocked { + t.Fatal("auth did not automatically recover after retry deadline") + } + }) + } +} + +func TestSchedulerPromotesUnknownFailureAfterRetryDeadline(t *testing.T) { + withQuotaCooldownEnabled(t) + previousTransient := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(0) + t.Cleanup(func() { transientErrorCooldownSeconds.Store(previousTransient) }) + + const ( + provider = "gemini" + model = "scheduler-unknown-recovery-model" + authID = "scheduler-unknown-recovery-auth" + ) + modelRegistry := registry.GetGlobalRegistry() + modelRegistry.RegisterClient(authID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { modelRegistry.UnregisterClient(authID) }) + + manager := NewManager(nil, &RoundRobinSelector{}, nil) + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), &Auth{ID: authID, Provider: provider}); errRegister != nil { + t.Fatalf("Register returned error: %v", errRegister) + } + if _, errPick := manager.scheduler.pickSingle(context.Background(), provider, model, cliproxyexecutor.Options{}, nil); errPick != nil { + t.Fatalf("initial scheduler pick returned error: %v", errPick) + } + + manager.MarkResult(context.Background(), Result{ + AuthID: authID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{Message: "transport closed"}, + }) + + manager.scheduler.mu.Lock() + defer manager.scheduler.mu.Unlock() + providerScheduler := manager.scheduler.providers[provider] + if providerScheduler == nil { + t.Fatalf("scheduler provider %q is missing", provider) + } + shard := providerScheduler.modelShards[model] + if shard == nil { + t.Fatalf("scheduler model shard %q is missing", model) + } + entry := shard.entries[authID] + if entry == nil { + t.Fatalf("scheduler auth %q is missing", authID) + } + if entry.state != scheduledStateBlocked || entry.nextRetryAt.IsZero() { + t.Fatalf("scheduler entry state = %v, retry = %v; want finite blocked state", entry.state, entry.nextRetryAt) + } + + shard.promoteExpiredLocked(entry.nextRetryAt.Add(time.Nanosecond)) + if entry.state != scheduledStateReady { + t.Fatalf("scheduler entry state after deadline = %v, want ready", entry.state) + } +} + func TestJitteredCooldownWaitBounds(t *testing.T) { cases := []struct { wait time.Duration diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index 8c864221..8bec6123 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -15,10 +15,11 @@ import ( type schedulerStrategy int const ( - schedulerStrategyCurrent schedulerStrategy = -1 - schedulerStrategyCustom schedulerStrategy = 0 - schedulerStrategyRoundRobin schedulerStrategy = 1 - schedulerStrategyFillFirst schedulerStrategy = 2 + schedulerStrategyCurrent schedulerStrategy = -1 + schedulerStrategyCustom schedulerStrategy = 0 + schedulerStrategyRoundRobin schedulerStrategy = 1 + schedulerStrategyFillFirst schedulerStrategy = 2 + schedulerStrategyWeightedRoundRobin schedulerStrategy = 3 ) // scheduledState describes how an auth currently participates in a model shard. @@ -33,11 +34,12 @@ const ( // authScheduler keeps the incremental provider/model scheduling state used by Manager. type authScheduler struct { - mu sync.Mutex - strategy schedulerStrategy - providers map[string]*providerScheduler - authProviders map[string]string - mixedCursors map[string]int + mu sync.Mutex + strategy schedulerStrategy + providers map[string]*providerScheduler + authProviders map[string]string + mixedCursors map[string]int + mixedWeightedStates map[string]*smoothWeightedState } // providerScheduler stores auth metadata and model shards for a single provider. @@ -52,6 +54,7 @@ type scheduledAuthMeta struct { auth *Auth providerKey string priority int + weight int64 websocketEnabled bool supportedModelSet map[string]struct{} } @@ -81,15 +84,17 @@ type readyBucket struct { // readyView holds the selection order for flat round-robin traversal. type readyView struct { - flat []*scheduledAuth - cursor int + flat []*scheduledAuth + cursor int + weightedState smoothWeightedState } // cooldownQueue is the blocked auth collection ordered by next retry time during rebuilds. type cooldownQueue []*scheduledAuth type readyViewCursorState struct { - cursor int + cursor int + weightedState smoothWeightedState } type readyBucketCursorState struct { @@ -98,7 +103,20 @@ type readyBucketCursorState struct { } func snapshotReadyViewCursors(view readyView) readyViewCursorState { - return readyViewCursorState{cursor: view.cursor} + state := readyViewCursorState{cursor: view.cursor} + if len(view.weightedState.current) > 0 { + state.weightedState.current = make(map[string]int64, len(view.weightedState.current)) + for authID, current := range view.weightedState.current { + state.weightedState.current[authID] = current + } + } + if len(view.weightedState.weights) > 0 { + state.weightedState.weights = make(map[string]int64, len(view.weightedState.weights)) + for authID, weight := range view.weightedState.weights { + state.weightedState.weights[authID] = weight + } + } + return state } func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { @@ -108,6 +126,12 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { if len(view.flat) > 0 { view.cursor = normalizeCursor(state.cursor, len(view.flat)) } + weights := scheduledWeightVector(view.flat) + if len(state.weightedState.current) == 0 || !weightVectorsEqual(state.weightedState.weights, weights) { + return + } + view.weightedState.current = state.weightedState.current + view.weightedState.weights = weights } func normalizeCursor(cursor, size int) int { @@ -124,10 +148,11 @@ func normalizeCursor(cursor, size int) int { // newAuthScheduler constructs an empty scheduler configured for the supplied selector strategy. func newAuthScheduler(selector Selector) *authScheduler { return &authScheduler{ - strategy: selectorStrategy(selector), - providers: make(map[string]*providerScheduler), - authProviders: make(map[string]string), - mixedCursors: make(map[string]int), + strategy: selectorStrategy(selector), + providers: make(map[string]*providerScheduler), + authProviders: make(map[string]string), + mixedCursors: make(map[string]int), + mixedWeightedStates: make(map[string]*smoothWeightedState), } } @@ -136,6 +161,8 @@ func selectorStrategy(selector Selector) schedulerStrategy { switch selector.(type) { case *FillFirstSelector: return schedulerStrategyFillFirst + case *WeightedRoundRobinSelector: + return schedulerStrategyWeightedRoundRobin case nil, *RoundRobinSelector: return schedulerStrategyRoundRobin default: @@ -152,6 +179,7 @@ func (s *authScheduler) setSelector(selector Selector) { defer s.mu.Unlock() s.strategy = selectorStrategy(selector) clear(s.mixedCursors) + clear(s.mixedWeightedStates) } // rebuild recreates the complete scheduler state from an auth snapshot. @@ -164,6 +192,7 @@ func (s *authScheduler) rebuild(auths []*Auth) { s.providers = make(map[string]*providerScheduler) s.authProviders = make(map[string]string) s.mixedCursors = make(map[string]int) + s.mixedWeightedStates = make(map[string]*smoothWeightedState) now := time.Now() for _, auth := range auths { s.upsertAuthLocked(auth, now) @@ -206,6 +235,7 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo providerKey := strings.ToLower(strings.TrimSpace(provider)) modelKey := canonicalModelKey(model) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + eligibility := authSelectionEligibilityForRequest(ctx, opts) preferWebsocket := cliproxyexecutor.DownstreamWebsocket(ctx) && providerPrefersWebsocketTransport(providerKey) && pinnedAuthID == "" s.mu.Lock() @@ -221,20 +251,7 @@ func (s *authScheduler) pickSingleWithStrategy(ctx context.Context, provider, mo if shard == nil { return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } - predicate := func(entry *scheduledAuth) bool { - if entry == nil || entry.auth == nil { - return false - } - if pinnedAuthID != "" && entry.auth.ID != pinnedAuthID { - return false - } - if len(tried) > 0 { - if _, ok := tried[entry.auth.ID]; ok { - return false - } - } - return true - } + predicate := scheduledAuthPredicate(eligibility, tried, pinnedAuthID, strategy == schedulerStrategyWeightedRoundRobin) if picked := shard.pickReadyLocked(preferWebsocket, strategy, predicate); picked != nil { return picked, nil } @@ -277,6 +294,7 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s return picked, providerKey, nil } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + eligibility := authSelectionEligibilityForRequest(ctx, opts) modelKey := canonicalModelKey(model) s.mu.Lock() @@ -294,23 +312,14 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s return nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} } shard := providerState.ensureModelLocked(modelKey, time.Now()) - predicate := func(entry *scheduledAuth) bool { - if entry == nil || entry.auth == nil || entry.auth.ID != pinnedAuthID { - return false - } - if len(tried) == 0 { - return true - } - _, ok := tried[pinnedAuthID] - return !ok - } + predicate := scheduledAuthPredicate(eligibility, tried, pinnedAuthID, strategy == schedulerStrategyWeightedRoundRobin) if picked := shard.pickReadyLocked(false, strategy, predicate); picked != nil { return picked, providerKey, nil } return nil, "", shard.unavailableErrorLocked("mixed", model, predicate) } - predicate := triedPredicate(tried) + predicate := scheduledAuthPredicate(eligibility, tried, "", strategy == schedulerStrategyWeightedRoundRobin) candidateShards := make([]*modelScheduler, len(normalized)) bestPriority := 0 hasCandidate := false @@ -335,7 +344,7 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s } } if !hasCandidate { - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, tried) + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } if strategy == schedulerStrategyFillFirst { @@ -349,10 +358,46 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s return picked, providerKey, nil } } - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, tried) + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } cursorKey := strings.Join(normalized, ",") + ":" + modelKey + if strategy == schedulerStrategyWeightedRoundRobin { + entries := make([]*scheduledAuth, 0) + for _, shard := range candidateShards { + if shard == nil { + continue + } + bucket := shard.readyByPriority[bestPriority] + if bucket != nil { + entries = append(entries, bucket.all.flat...) + } + } + sort.Slice(entries, func(i, j int) bool { + if entries[i] == nil || entries[i].auth == nil { + return false + } + if entries[j] == nil || entries[j].auth == nil { + return true + } + return entries[i].auth.ID < entries[j].auth.ID + }) + if s.mixedWeightedStates == nil { + s.mixedWeightedStates = make(map[string]*smoothWeightedState) + } + state := s.mixedWeightedStates[cursorKey] + if state == nil { + state = &smoothWeightedState{} + s.mixedWeightedStates[cursorKey] = state + } + state.prepare(scheduledWeightVectorMatching(entries, predicate)) + picked := pickSmoothWeightedScheduled(entries, state.current, predicate) + if picked != nil && picked.meta != nil { + return picked.auth, picked.meta.providerKey, nil + } + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) + } + weights := make([]int, len(normalized)) segmentStarts := make([]int, len(normalized)) segmentEnds := make([]int, len(normalized)) @@ -360,13 +405,13 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s for providerIndex, shard := range candidateShards { segmentStarts[providerIndex] = totalWeight if shard != nil { - weights[providerIndex] = shard.readyCountAtPriorityLocked(false, bestPriority) + weights[providerIndex] = shard.readyCountAtPriorityLocked(false, bestPriority, predicate) } totalWeight += weights[providerIndex] segmentEnds[providerIndex] = totalWeight } if totalWeight == 0 { - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, tried) + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } startSlot := s.mixedCursors[cursorKey] % totalWeight @@ -381,7 +426,7 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s } } if startProviderIndex < 0 { - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, tried) + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } slot := startSlot @@ -405,11 +450,11 @@ func (s *authScheduler) pickMixedWithStrategy(ctx context.Context, providers []s s.mixedCursors[cursorKey] = slot + 1 return picked, providerKey, nil } - return nil, "", s.mixedUnavailableErrorLocked(normalized, model, tried) + return nil, "", s.mixedUnavailableErrorLocked(normalized, model, predicate) } // mixedUnavailableErrorLocked synthesizes the mixed-provider cooldown or unavailable error. -func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model string, tried map[string]struct{}) error { +func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model string, predicate func(*scheduledAuth) bool) error { now := time.Now() total := 0 cooldownCount := 0 @@ -423,7 +468,7 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st if shard == nil { continue } - localTotal, localCooldownCount, localEarliest := shard.availabilitySummaryLocked(triedPredicate(tried)) + localTotal, localCooldownCount, localEarliest := shard.availabilitySummaryLocked(predicate) total += localTotal cooldownCount += localCooldownCount if !localEarliest.IsZero() && (earliest.IsZero() || localEarliest.Before(earliest)) { @@ -443,17 +488,24 @@ func (s *authScheduler) mixedUnavailableErrorLocked(providers []string, model st return &Error{Code: "auth_unavailable", Message: "no auth available"} } -// triedPredicate builds a filter that excludes auths already attempted for the current request. -func triedPredicate(tried map[string]struct{}) func(*scheduledAuth) bool { - if len(tried) == 0 { - return func(entry *scheduledAuth) bool { return entry != nil && entry.auth != nil } - } +// scheduledAuthPredicate filters request-ineligible auths before scheduler state advances. +func scheduledAuthPredicate(eligibility authSelectionEligibility, tried map[string]struct{}, pinnedAuthID string, requirePositiveWeight bool) func(*scheduledAuth) bool { return func(entry *scheduledAuth) bool { - if entry == nil || entry.auth == nil { + if entry == nil || entry.auth == nil || !eligibility.allows(entry.auth) { + return false + } + if requirePositiveWeight && (entry.meta == nil || entry.meta.weight <= 0) { + return false + } + if pinnedAuthID != "" && entry.auth.ID != pinnedAuthID { return false } - _, ok := tried[entry.auth.ID] - return !ok + if len(tried) > 0 { + if _, ok := tried[entry.auth.ID]; ok { + return false + } + } + return true } } @@ -543,6 +595,7 @@ func buildScheduledAuthMeta(auth *Auth) *scheduledAuthMeta { auth: auth, providerKey: providerKey, priority: authPriority(auth), + weight: authWeight(auth), websocketEnabled: authWebsocketsEnabled(auth), supportedModelSet: supportedModelSetForAuth(auth.ID), } @@ -789,9 +842,12 @@ func (m *modelScheduler) pickReadyAtPriorityLocked(preferWebsocket bool, priorit view = &bucket.ws } var picked *scheduledAuth - if strategy == schedulerStrategyFillFirst { + switch strategy { + case schedulerStrategyFillFirst: picked = view.pickFirst(predicate) - } else { + case schedulerStrategyWeightedRoundRobin: + picked = view.pickWeighted(predicate) + default: picked = view.pickRoundRobin(predicate) } if picked == nil || picked.auth == nil { @@ -800,7 +856,7 @@ func (m *modelScheduler) pickReadyAtPriorityLocked(preferWebsocket bool, priorit return picked.auth } -func (m *modelScheduler) readyCountAtPriorityLocked(preferWebsocket bool, priority int) int { +func (m *modelScheduler) readyCountAtPriorityLocked(preferWebsocket bool, priority int, predicate func(*scheduledAuth) bool) int { if m == nil { return 0 } @@ -808,10 +864,17 @@ func (m *modelScheduler) readyCountAtPriorityLocked(preferWebsocket bool, priori if bucket == nil { return 0 } - if preferWebsocket && len(bucket.ws.flat) > 0 { - return len(bucket.ws.flat) + view := &bucket.all + if preferWebsocket && bucket.ws.pickFirst(predicate) != nil { + view = &bucket.ws + } + count := 0 + for _, entry := range view.flat { + if predicate == nil || predicate(entry) { + count++ + } } - return len(bucket.all.flat) + return count } // unavailableErrorLocked returns the correct unavailable or cooldown error for the shard. @@ -974,3 +1037,71 @@ func (v *readyView) pickRoundRobin(predicate func(*scheduledAuth) bool) *schedul } return nil } + +// pickWeighted returns the next ready entry using smooth weighted round-robin. +func (v *readyView) pickWeighted(predicate func(*scheduledAuth) bool) *scheduledAuth { + if v == nil || len(v.flat) == 0 { + return nil + } + v.weightedState.prepare(scheduledWeightVectorMatching(v.flat, predicate)) + return pickSmoothWeightedScheduled(v.flat, v.weightedState.current, predicate) +} + +func scheduledWeightVector(entries []*scheduledAuth) map[string]int64 { + return scheduledWeightVectorMatching(entries, nil) +} + +func scheduledWeightVectorMatching(entries []*scheduledAuth, predicate func(*scheduledAuth) bool) map[string]int64 { + weights := make(map[string]int64, len(entries)) + for _, entry := range entries { + if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { + continue + } + if predicate != nil && !predicate(entry) { + continue + } + weights[entry.auth.ID] = entry.meta.weight + } + return weights +} + +func pickSmoothWeightedScheduled(entries []*scheduledAuth, current map[string]int64, predicate func(*scheduledAuth) bool) *scheduledAuth { + active := make(map[string]struct{}, len(entries)) + for _, entry := range entries { + if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { + continue + } + if predicate != nil && !predicate(entry) { + continue + } + active[entry.auth.ID] = struct{}{} + } + for authID := range current { + if _, ok := active[authID]; !ok { + delete(current, authID) + } + } + + var picked *scheduledAuth + var pickedCurrent int64 + var totalWeight int64 + for _, entry := range entries { + if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { + continue + } + if predicate != nil && !predicate(entry) { + continue + } + current[entry.auth.ID] = saturatingAddInt64(current[entry.auth.ID], entry.meta.weight) + totalWeight = saturatingAddInt64(totalWeight, entry.meta.weight) + if picked == nil || current[entry.auth.ID] > pickedCurrent { + picked = entry + pickedCurrent = current[entry.auth.ID] + } + } + if picked == nil { + return nil + } + current[picked.auth.ID] = saturatingAddInt64(current[picked.auth.ID], -totalWeight) + return picked +} diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 1e6ae7c2..3e3156d1 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -20,6 +20,22 @@ type schedulerTestExecutor struct { provider string } +type schedulerLoadStore struct { + auths []*Auth +} + +func (s *schedulerLoadStore) List(context.Context) ([]*Auth, error) { + return s.auths, nil +} + +func (s *schedulerLoadStore) Save(context.Context, *Auth) (string, error) { + return "", nil +} + +func (s *schedulerLoadStore) Delete(context.Context, string) error { + return nil +} + func (e schedulerTestExecutor) Identifier() string { if e.provider != "" { return e.provider @@ -153,6 +169,165 @@ func TestSchedulerPick_RoundRobinHighestPriority(t *testing.T) { } } +func TestSchedulerPick_WeightedRoundRobin(t *testing.T) { + t.Parallel() + + scheduler := newSchedulerForTest( + &WeightedRoundRobinSelector{}, + &Auth{ID: "a", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "5"}}, + &Auth{ID: "b", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "3"}}, + &Auth{ID: "c", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "2"}}, + ) + + counts := make(map[string]int) + for index := 0; index < 100; index++ { + got, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickSingle() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + want := map[string]int{"a": 50, "b": 30, "c": 20} + for authID, wantCount := range want { + if counts[authID] != wantCount { + t.Fatalf("auth %q picks = %d, want %d", authID, counts[authID], wantCount) + } + } +} + +func TestManagerLoad_WeightedRoundRobinUsesPersistedMetadataWeight(t *testing.T) { + t.Parallel() + + manager := NewManager(&schedulerLoadStore{auths: []*Auth{ + {ID: "a", Provider: "gemini", Metadata: map[string]any{AttributeWeight: float64(5)}}, + {ID: "b", Provider: "gemini", Metadata: map[string]any{AttributeWeight: float64(1)}}, + }}, &WeightedRoundRobinSelector{}, nil) + if errLoad := manager.Load(context.Background()); errLoad != nil { + t.Fatalf("Load() error = %v", errLoad) + } + + counts := make(map[string]int) + for index := 0; index < 60; index++ { + got, errPick := manager.scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickSingle() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 50 || counts["b"] != 10 { + t.Fatalf("metadata-weighted picks = %#v, want a:b=50:10", counts) + } +} + +func TestSchedulerPick_WeightedRoundRobinResetsCreditsWhenWeightsChange(t *testing.T) { + t.Parallel() + + authA := &Auth{ID: "a", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "1000000"}} + authB := &Auth{ID: "b", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "1"}} + scheduler := newSchedulerForTest(&WeightedRoundRobinSelector{}, authA, authB) + for index := 0; index < 1000; index++ { + if _, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil); errPick != nil { + t.Fatalf("warmup pickSingle() #%d error = %v", index, errPick) + } + } + + authA.Attributes[AttributeWeight] = "1" + scheduler.upsertAuth(authA) + counts := make(map[string]int) + for index := 0; index < 20; index++ { + got, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickSingle() after weight change #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 10 || counts["b"] != 10 { + t.Fatalf("picks after weight change = %#v, want a:b=10:10", counts) + } +} + +func TestSchedulerPick_WeightedWebsocketResetsCreditsWhenWeightsChange(t *testing.T) { + t.Parallel() + + authA := &Auth{ID: "a", Provider: "codex", Attributes: map[string]string{AttributeWeight: "1000000", "websockets": "true"}} + authB := &Auth{ID: "b", Provider: "codex", Attributes: map[string]string{AttributeWeight: "1", "websockets": "true"}} + scheduler := newSchedulerForTest(&WeightedRoundRobinSelector{}, authA, authB) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + for index := 0; index < 1000; index++ { + if _, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil); errPick != nil { + t.Fatalf("warmup websocket pickSingle() #%d error = %v", index, errPick) + } + } + + authA.Attributes[AttributeWeight] = "1" + scheduler.upsertAuth(authA) + counts := make(map[string]int) + for index := 0; index < 20; index++ { + got, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("websocket pickSingle() after weight change #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 10 || counts["b"] != 10 { + t.Fatalf("websocket picks after weight change = %#v, want a:b=10:10", counts) + } +} + +func TestManagerLegacyWeightedRoundRobinKeepsIndependentAliasPrefixedModelState(t *testing.T) { + manager := NewManager(nil, &WeightedRoundRobinSelector{}, nil) + manager.executors["gemini"] = schedulerTestExecutor{} + manager.SetPluginScheduler(&fakePluginScheduler{}) + + auths := []*Auth{ + {ID: "a-heavy", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: "a-light", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "1"}}, + {ID: "b-light", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "1"}}, + {ID: "b-heavy", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "3"}}, + } + for _, auth := range auths { + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register(%s) error = %v", auth.ID, errRegister) + } + } + registerSchedulerModels(t, "gemini", "team-a/shared", "a-heavy", "a-light") + registerSchedulerModels(t, "gemini", "team-b/shared", "b-light", "b-heavy") + + counts := make(map[string]int) + for index := 0; index < 40; index++ { + for _, model := range []string{"team-a/shared", "team-b/shared"} { + got, _, errPick := manager.pickNext(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickNext(%q) #%d error = %v", model, index, errPick) + } + counts[got.ID]++ + } + } + want := map[string]int{"a-heavy": 30, "a-light": 10, "b-light": 10, "b-heavy": 30} + for authID, wantCount := range want { + if counts[authID] != wantCount { + t.Fatalf("auth %q picks = %d, want %d; all=%#v", authID, counts[authID], wantCount, counts) + } + } +} + +func TestSchedulerPick_WeightedRoundRobinSkipsNonPositiveWeightPriorityTier(t *testing.T) { + t.Parallel() + + scheduler := newSchedulerForTest( + &WeightedRoundRobinSelector{}, + &Auth{ID: "excluded", Provider: "gemini", Attributes: map[string]string{"priority": "10", AttributeWeight: "0"}}, + &Auth{ID: "available", Provider: "gemini", Attributes: map[string]string{"priority": "0", AttributeWeight: "1"}}, + ) + got, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickSingle() error = %v", errPick) + } + if got == nil || got.ID != "available" { + t.Fatalf("pickSingle() auth = %#v, want available", got) + } +} + func TestSchedulerPick_FillFirstSticksToFirstReady(t *testing.T) { t.Parallel() @@ -316,6 +491,63 @@ func TestSchedulerPick_MixedProvidersUsesWeightedProviderRotationOverReadyCandid } } +func TestSchedulerPick_MixedProvidersWeightedRoundRobin(t *testing.T) { + t.Parallel() + + scheduler := newSchedulerForTest( + &WeightedRoundRobinSelector{}, + &Auth{ID: "gemini-a", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "5"}}, + &Auth{ID: "claude-b", Provider: "claude", Attributes: map[string]string{AttributeWeight: "3"}}, + &Auth{ID: "claude-c", Provider: "claude", Attributes: map[string]string{AttributeWeight: "2"}}, + ) + + counts := make(map[string]int) + for index := 0; index < 100; index++ { + got, provider, errPick := scheduler.pickMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickMixed() #%d error = %v", index, errPick) + } + if got == nil || provider == "" { + t.Fatalf("pickMixed() #%d returned auth=%v provider=%q", index, got, provider) + } + counts[got.ID]++ + } + want := map[string]int{"gemini-a": 50, "claude-b": 30, "claude-c": 20} + for authID, wantCount := range want { + if counts[authID] != wantCount { + t.Fatalf("auth %q picks = %d, want %d", authID, counts[authID], wantCount) + } + } +} + +func TestSchedulerPick_MixedProvidersResetsCreditsWhenWeightsChange(t *testing.T) { + t.Parallel() + + authA := &Auth{ID: "gemini-a", Provider: "gemini", Attributes: map[string]string{AttributeWeight: "1000000"}} + authB := &Auth{ID: "claude-b", Provider: "claude", Attributes: map[string]string{AttributeWeight: "1"}} + scheduler := newSchedulerForTest(&WeightedRoundRobinSelector{}, authA, authB) + providers := []string{"gemini", "claude"} + for index := 0; index < 1000; index++ { + if _, _, errPick := scheduler.pickMixed(context.Background(), providers, "", cliproxyexecutor.Options{}, nil); errPick != nil { + t.Fatalf("warmup pickMixed() #%d error = %v", index, errPick) + } + } + + authA.Attributes[AttributeWeight] = "1" + scheduler.upsertAuth(authA) + counts := make(map[string]int) + for index := 0; index < 20; index++ { + got, _, errPick := scheduler.pickMixed(context.Background(), providers, "", cliproxyexecutor.Options{}, nil) + if errPick != nil { + t.Fatalf("pickMixed() after weight change #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts[authA.ID] != 10 || counts[authB.ID] != 10 { + t.Fatalf("mixed picks after weight change = %#v, want 10 each", counts) + } +} + func TestSchedulerPick_MixedProvidersPrefersHighestPriorityTier(t *testing.T) { t.Parallel() @@ -481,8 +713,113 @@ func TestManagerSelectAuthByKindSkipsAPIKey(t *testing.T) { if selected == nil || selected.ID != "codex-oauth" { t.Fatalf("SelectAuthByKind() auth = %#v, want codex-oauth", selected) } - if scheduler.calls != 2 { - t.Fatalf("scheduler.calls = %d, want 2", scheduler.calls) + if scheduler.calls != 1 { + t.Fatalf("scheduler.calls = %d, want 1", scheduler.calls) + } + if len(scheduler.requests) != 1 || len(scheduler.requests[0].Candidates) != 1 || scheduler.requests[0].Candidates[0].ID != "codex-oauth" { + t.Fatalf("scheduler candidates = %#v, want only codex-oauth", scheduler.requests) + } +} + +func TestManagerSelectAuthByKindWeightedRoundRobinIgnoresIneligibleAPIKeyWeight(t *testing.T) { + manager := NewManager(nil, &WeightedRoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + for _, candidate := range []*Auth{ + {ID: "api-high", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "test-key", AttributeWeight: "100"}}, + {ID: "oauth-heavy", Provider: "codex", Attributes: map[string]string{AttributeWeight: "5"}, Metadata: map[string]any{"access_token": "heavy-token"}}, + {ID: "oauth-light", Provider: "codex", Attributes: map[string]string{AttributeWeight: "1"}, Metadata: map[string]any{"access_token": "light-token"}}, + } { + if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil { + t.Fatalf("Register(%s) error = %v", candidate.ID, errRegister) + } + } + + counts := make(map[string]int) + for index := 0; index < 600; index++ { + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectAuthByKind() #%d error = %v", index, errSelect) + } + counts[selected.ID]++ + } + if counts["oauth-heavy"] != 500 || counts["oauth-light"] != 100 || counts["api-high"] != 0 { + t.Fatalf("weighted OAuth picks = %#v, want oauth-heavy:oauth-light=500:100 and no API key", counts) + } +} + +func TestManagerWeightedRoundRobinDisallowFreeAuthIgnoresFreeWeight(t *testing.T) { + tests := []struct { + name string + mixed bool + }{ + {name: "single provider"}, + {name: "mixed providers", mixed: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + manager := NewManager(nil, &WeightedRoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + lightProvider := "codex" + if tt.mixed { + lightProvider = "gemini" + manager.executors["gemini"] = schedulerTestExecutor{provider: "gemini"} + } + for _, candidate := range []*Auth{ + {ID: "free-high", Provider: "codex", Attributes: map[string]string{"plan_type": "free", AttributeWeight: "100"}, Metadata: map[string]any{"access_token": "free-token"}}, + {ID: "paid-heavy", Provider: "codex", Attributes: map[string]string{"plan_type": "plus", AttributeWeight: "5"}, Metadata: map[string]any{"access_token": "heavy-token"}}, + {ID: "paid-light", Provider: lightProvider, Attributes: map[string]string{"plan_type": "plus", AttributeWeight: "1"}, Metadata: map[string]any{"access_token": "light-token"}}, + } { + if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil { + t.Fatalf("Register(%s) error = %v", candidate.ID, errRegister) + } + } + + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.DisallowFreeAuthMetadataKey: true}} + counts := make(map[string]int) + for index := 0; index < 600; index++ { + var selected *Auth + var errPick error + if tt.mixed { + selected, _, _, errPick = manager.pickNextMixed(context.Background(), []string{"codex", "gemini"}, "", opts, nil) + } else { + selected, _, errPick = manager.pickNext(context.Background(), "codex", "", opts, nil) + } + if errPick != nil { + t.Fatalf("weighted pick #%d error = %v", index, errPick) + } + counts[selected.ID]++ + } + if counts["paid-heavy"] != 500 || counts["paid-light"] != 100 || counts["free-high"] != 0 { + t.Fatalf("weighted non-free picks = %#v, want paid-heavy:paid-light=500:100 and no free auth", counts) + } + }) + } +} + +func TestManagerSelectAuthByKindRoundRobinKeepsEligibleRotation(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + manager.executors["codex"] = schedulerTestExecutor{} + for _, candidate := range []*Auth{ + {ID: "api-key", Provider: "codex", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, + {ID: "oauth-a", Provider: "codex", Metadata: map[string]any{"access_token": "token-a"}}, + {ID: "oauth-b", Provider: "codex", Metadata: map[string]any{"access_token": "token-b"}}, + } { + if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil { + t.Fatalf("Register(%s) error = %v", candidate.ID, errRegister) + } + } + + counts := make(map[string]int) + for index := 0; index < 6; index++ { + selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectAuthByKind() #%d error = %v", index, errSelect) + } + counts[selected.ID]++ + } + if counts["oauth-a"] != 3 || counts["oauth-b"] != 3 || counts["api-key"] != 0 { + t.Fatalf("round-robin OAuth picks = %#v, want three picks per OAuth auth and no API key", counts) } } diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 4be06257..00656f42 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -16,6 +16,7 @@ import ( log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" + "github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -29,6 +30,36 @@ type RoundRobinSelector struct { maxKeys int } +// WeightedRoundRobinSelector provides smooth weighted round-robin selection. +type WeightedRoundRobinSelector struct { + mu sync.Mutex + states map[string]*smoothWeightedState + maxKeys int +} + +type smoothWeightedState struct { + current map[string]int64 + weights map[string]int64 +} + +type weightedSelectorStateModelKey struct{} + +func withWeightedSelectorStateModel(ctx context.Context, selector Selector, routeModel string) context.Context { + if _, ok := selector.(*WeightedRoundRobinSelector); !ok || strings.TrimSpace(routeModel) == "" { + return ctx + } + return context.WithValue(ctx, weightedSelectorStateModelKey{}, routeModel) +} + +func weightedSelectorStateModel(ctx context.Context, availabilityModel string) string { + if ctx != nil { + if routeModel, ok := ctx.Value(weightedSelectorStateModelKey{}).(string); ok && strings.TrimSpace(routeModel) != "" { + return routeModel + } + } + return availabilityModel +} + // FillFirstSelector selects the first available credential (deterministic ordering). // This "burns" one account before moving to the next, which can help stagger // rolling-window subscription caps (e.g. chat message limits). @@ -127,6 +158,27 @@ func authPriority(auth *Auth) int { return parsed } +func authWeight(auth *Auth) int64 { + if auth == nil { + return credentialweight.Default + } + if rawWeight, ok := auth.Attributes[AttributeWeight]; ok && strings.TrimSpace(rawWeight) != "" { + weight, errParse := credentialweight.ParseString(rawWeight) + if errParse != nil { + return 0 + } + return weight + } + if rawWeight, ok := auth.Metadata[AttributeWeight]; ok { + weight, errParse := credentialweight.ParseValue(rawWeight) + if errParse != nil { + return 0 + } + return weight + } + return credentialweight.Default +} + func canonicalModelKey(model string) string { model = strings.TrimSpace(model) if model == "" { @@ -290,6 +342,125 @@ func (s *RoundRobinSelector) ensureCursorKey(key string, limit int) { } } +func positiveWeightAuths(auths []*Auth) []*Auth { + weightedCandidates := make([]*Auth, 0, len(auths)) + for _, auth := range auths { + if authWeight(auth) > 0 { + weightedCandidates = append(weightedCandidates, auth) + } + } + return weightedCandidates +} + +// Pick selects the next available auth using smooth weighted round-robin. +func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + _ = opts + available, errAvailable := getAvailableAuths(positiveWeightAuths(auths), provider, model, time.Now()) + if errAvailable != nil { + return nil, errAvailable + } + available = preferCodexWebsocketAuths(ctx, provider, available) + stateModel := weightedSelectorStateModel(ctx, model) + key := provider + ":" + canonicalModelKey(stateModel) + + s.mu.Lock() + defer s.mu.Unlock() + if s.states == nil { + s.states = make(map[string]*smoothWeightedState) + } + limit := s.maxKeys + if limit <= 0 { + limit = 4096 + } + if _, ok := s.states[key]; !ok && len(s.states) >= limit { + s.states = make(map[string]*smoothWeightedState) + } + state := s.states[key] + if state == nil { + state = &smoothWeightedState{} + s.states[key] = state + } + weights := authWeightVector(available) + state.prepare(weights) + picked := pickSmoothWeightedAuth(available, state.current) + if picked == nil { + return nil, &Error{Code: "auth_unavailable", Message: "no auth available with positive weight"} + } + return picked, nil +} + +func (s *smoothWeightedState) prepare(weights map[string]int64) { + if s.current == nil || !weightVectorsEqual(s.weights, weights) { + s.current = make(map[string]int64) + } + s.weights = weights +} + +func weightVectorsEqual(left, right map[string]int64) bool { + if len(left) != len(right) { + return false + } + for authID, weight := range left { + if right[authID] != weight { + return false + } + } + return true +} + +func authWeightVector(auths []*Auth) map[string]int64 { + weights := make(map[string]int64, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + if weight := authWeight(auth); weight > 0 { + weights[auth.ID] = weight + } + } + return weights +} + +func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { + active := make(map[string]struct{}, len(auths)) + var picked *Auth + var pickedCurrent int64 + var totalWeight int64 + for _, auth := range auths { + weight := authWeight(auth) + if auth == nil || weight <= 0 { + continue + } + active[auth.ID] = struct{}{} + current[auth.ID] = saturatingAddInt64(current[auth.ID], weight) + totalWeight = saturatingAddInt64(totalWeight, weight) + if picked == nil || current[auth.ID] > pickedCurrent { + picked = auth + pickedCurrent = current[auth.ID] + } + } + for authID := range current { + if _, ok := active[authID]; !ok { + delete(current, authID) + } + } + if picked == nil { + return nil + } + current[picked.ID] = saturatingAddInt64(current[picked.ID], -totalWeight) + return picked +} + +func saturatingAddInt64(value, delta int64) int64 { + if delta > 0 && value > math.MaxInt64-delta { + return math.MaxInt64 + } + if delta < 0 && value < math.MinInt64-delta { + return math.MinInt64 + } + return value + delta +} + // Pick selects the first available auth for the provider in a deterministic manner. func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { _ = opts @@ -322,43 +493,38 @@ func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, block if state.Status == StatusDisabled { return true, blockReasonDisabled, time.Time{} } - if state.Unavailable { - if state.NextRetryAfter.IsZero() { - return false, blockReasonNone, time.Time{} - } - if state.NextRetryAfter.After(now) { - next := state.NextRetryAfter - if !state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.After(now) { - next = state.Quota.NextRecoverAt - } - if next.Before(now) { - next = now - } - if state.Quota.Exceeded { - return true, blockReasonCooldown, next - } - return true, blockReasonOther, next - } - } - return false, blockReasonNone, time.Time{} + return availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now) } + // Auth-level availability can aggregate failures from other models. + return false, blockReasonNone, time.Time{} } + return availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) + } + return availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) +} + +func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextRecoverAt, now time.Time) (bool, blockReason, time.Time) { + if !unavailable && !quotaExceeded { return false, blockReasonNone, time.Time{} } - if auth.Unavailable && auth.NextRetryAfter.After(now) { - next := auth.NextRetryAfter - if !auth.Quota.NextRecoverAt.IsZero() && auth.Quota.NextRecoverAt.After(now) { - next = auth.Quota.NextRecoverAt - } - if next.Before(now) { - next = now + + hasRecoveryTime := !nextRetryAfter.IsZero() || !nextRecoverAt.IsZero() + var next time.Time + for _, candidate := range []time.Time{nextRetryAfter, nextRecoverAt} { + if candidate.After(now) && (next.IsZero() || candidate.After(next)) { + next = candidate } - if auth.Quota.Exceeded { + } + if !next.IsZero() { + if quotaExceeded { return true, blockReasonCooldown, next } return true, blockReasonOther, next } - return false, blockReasonNone, time.Time{} + if hasRecoveryTime { + return false, blockReasonNone, time.Time{} + } + return true, blockReasonOther, time.Time{} } // SessionAffinitySelector wraps another selector with session-sticky behavior. @@ -413,7 +579,11 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } now := time.Now() - available, err := getAvailableAuths(auths, provider, model, now) + availabilityCandidates := auths + if _, weighted := s.fallback.(*WeightedRoundRobinSelector); weighted { + availabilityCandidates = positiveWeightAuths(auths) + } + available, err := getAvailableAuths(availabilityCandidates, provider, model, now) if err != nil { return nil, err } diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index cb1bfb40..ac4cc56d 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "math" "net/http" "strings" "sync" @@ -63,6 +64,217 @@ func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) { } } +func TestWeightedRoundRobinSelectorPick_DistributesAndSkipsNonPositiveWeights(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + auths := []*Auth{ + {ID: "a", Attributes: map[string]string{AttributeWeight: "5"}}, + {ID: "b", Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: "c", Attributes: map[string]string{AttributeWeight: "2"}}, + {ID: "disabled-by-weight", Attributes: map[string]string{AttributeWeight: "0"}}, + } + + counts := make(map[string]int) + for index := 0; index < 100; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil { + t.Fatalf("Pick() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + want := map[string]int{"a": 50, "b": 30, "c": 20} + for authID, wantCount := range want { + if counts[authID] != wantCount { + t.Fatalf("auth %q picks = %d, want %d", authID, counts[authID], wantCount) + } + } + if counts["disabled-by-weight"] != 0 { + t.Fatalf("non-positive weight auth picks = %d, want 0", counts["disabled-by-weight"]) + } +} + +func TestWeightedRoundRobinSelectorPick_ResetsCreditsWhenWeightsChange(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + authA := &Auth{ID: "a", Attributes: map[string]string{AttributeWeight: "1000000"}} + authB := &Auth{ID: "b", Attributes: map[string]string{AttributeWeight: "1"}} + auths := []*Auth{authA, authB} + for index := 0; index < 1000; index++ { + if _, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths); errPick != nil { + t.Fatalf("warmup Pick() #%d error = %v", index, errPick) + } + } + + authA.Attributes[AttributeWeight] = "1" + counts := make(map[string]int) + for index := 0; index < 20; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil { + t.Fatalf("Pick() after weight change #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 10 || counts["b"] != 10 { + t.Fatalf("picks after weight change = %#v, want a:b=10:10", counts) + } +} + +func TestWeightedRoundRobinSelectorPick_RebalancesWhenHighestWeightUnavailable(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + auths := []*Auth{ + {ID: "a", Disabled: true, Attributes: map[string]string{AttributeWeight: "5"}}, + {ID: "b", Attributes: map[string]string{AttributeWeight: "3"}}, + {ID: "c", Attributes: map[string]string{AttributeWeight: "2"}}, + } + counts := make(map[string]int) + for index := 0; index < 100; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil { + t.Fatalf("Pick() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 0 || counts["b"] != 60 || counts["c"] != 40 { + t.Fatalf("weighted failover counts = %#v, want b:c=60:40 with a skipped", counts) + } +} + +func TestWeightedRoundRobinSelectorPick_SkipsUnavailableAndQuotaExceededWithoutRecovery(t *testing.T) { + t.Parallel() + + model := "test-model" + selector := &WeightedRoundRobinSelector{} + auths := []*Auth{ + { + ID: "model-unavailable", + ModelStates: map[string]*ModelState{ + model: {Unavailable: true}, + }, + }, + {ID: "quota-exceeded", Quota: QuotaState{Exceeded: true}}, + {ID: "available"}, + } + + gotModel, errModel := selector.Pick(context.Background(), "gemini", model, cliproxyexecutor.Options{}, auths) + if errModel != nil || gotModel == nil || gotModel.ID != "available" { + t.Fatalf("model Pick() = %#v, %v; want available", gotModel, errModel) + } + for index := 0; index < 4; index++ { + gotAuth, errAuth := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) + if errAuth != nil || gotAuth == nil { + t.Fatalf("auth Pick() #%d = %#v, %v; want available auth", index, gotAuth, errAuth) + } + if gotAuth.ID == "quota-exceeded" { + t.Fatalf("auth Pick() #%d selected quota-exceeded credential", index) + } + } +} + +func TestAuthWeight_MetadataFallbackAndAttributePrecedence(t *testing.T) { + t.Parallel() + + if got := authWeight(&Auth{Metadata: map[string]any{AttributeWeight: float64(7)}}); got != 7 { + t.Fatalf("authWeight(metadata) = %d, want 7", got) + } + if got := authWeight(&Auth{ + Attributes: map[string]string{AttributeWeight: "3"}, + Metadata: map[string]any{AttributeWeight: float64(7)}, + }); got != 3 { + t.Fatalf("authWeight(attribute and metadata) = %d, want attribute weight 3", got) + } +} + +func TestAuthWeight_InvalidAndOverflowValuesAreExcluded(t *testing.T) { + t.Parallel() + + for _, raw := range []string{"1.5", "1000001", "9223372036854775807", "9223372036854775808"} { + auth := &Auth{Attributes: map[string]string{AttributeWeight: raw}} + if got := authWeight(auth); got != 0 { + t.Fatalf("authWeight(%q) = %d, want 0", raw, got) + } + } + if got := authWeight(&Auth{Metadata: map[string]any{AttributeWeight: 1.5}}); got != 0 { + t.Fatalf("authWeight(invalid metadata) = %d, want 0", got) + } + if got := authWeight(&Auth{Attributes: map[string]string{AttributeWeight: "-1"}}); got != 0 { + t.Fatalf("authWeight(-1) = %d, want 0", got) + } +} + +func TestPickSmoothWeightedAuth_SaturatesCorruptState(t *testing.T) { + t.Parallel() + + current := map[string]int64{"a": math.MaxInt64, "b": math.MinInt64} + picked := pickSmoothWeightedAuth([]*Auth{{ID: "a"}, {ID: "b"}}, current) + if picked == nil { + t.Fatal("pickSmoothWeightedAuth() returned nil") + } + if current["a"] != math.MaxInt64-2 || current["b"] != math.MinInt64+1 { + t.Fatalf("current state = %#v, want saturated arithmetic", current) + } +} + +func TestWeightedRoundRobinSelectorPick_RecoveredAuthReturnsWithoutAccumulatedCredit(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + authA := &Auth{ID: "a", Attributes: map[string]string{AttributeWeight: "5"}} + authB := &Auth{ID: "b", Attributes: map[string]string{AttributeWeight: "1"}} + auths := []*Auth{authA, authB} + + for index := 0; index < 6; index++ { + if _, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths); errPick != nil { + t.Fatalf("warmup Pick() #%d error = %v", index, errPick) + } + } + authA.Unavailable = true + authA.NextRetryAfter = time.Now().Add(time.Hour) + for index := 0; index < 6; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil || got == nil || got.ID != "b" { + t.Fatalf("unavailable Pick() #%d = %#v, %v; want b", index, got, errPick) + } + } + authA.Unavailable = false + authA.NextRetryAfter = time.Time{} + + counts := make(map[string]int) + for index := 0; index < 6; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil { + t.Fatalf("recovered Pick() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["a"] != 5 || counts["b"] != 1 { + t.Fatalf("recovered picks = %#v, want a:b=5:1", counts) + } +} + +func TestWeightedRoundRobinSelectorPick_DefaultWeightIsOne(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + auths := []*Auth{{ID: "a"}, {ID: "b"}, {ID: "c"}} + counts := make(map[string]int) + for index := 0; index < 30; index++ { + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) + if errPick != nil { + t.Fatalf("Pick() #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + for _, authID := range []string{"a", "b", "c"} { + if counts[authID] != 10 { + t.Fatalf("auth %q picks = %d, want 10", authID, counts[authID]) + } + } +} + func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) { t.Parallel() @@ -285,7 +497,7 @@ func TestSelectorPick_AllCooldownReturnsModelCooldownError(t *testing.T) { }) } -func TestIsAuthBlockedForModel_UnavailableWithoutNextRetryIsNotBlocked(t *testing.T) { +func TestIsAuthBlockedForModel_UnavailableWithoutNextRetryIsBlocked(t *testing.T) { t.Parallel() now := time.Now() @@ -304,17 +516,48 @@ func TestIsAuthBlockedForModel_UnavailableWithoutNextRetryIsNotBlocked(t *testin } blocked, reason, next := isAuthBlockedForModel(auth, model, now) - if blocked { - t.Fatalf("blocked = true, want false") + if !blocked { + t.Fatalf("blocked = false, want true") } - if reason != blockReasonNone { - t.Fatalf("reason = %v, want %v", reason, blockReasonNone) + if reason != blockReasonOther { + t.Fatalf("reason = %v, want %v", reason, blockReasonOther) } if !next.IsZero() { t.Fatalf("next = %v, want zero", next) } } +func TestIsAuthBlockedForModel_AuthQuotaExceededWithoutRecoveryIsBlocked(t *testing.T) { + t.Parallel() + + auth := &Auth{ID: "a", Quota: QuotaState{Exceeded: true}} + for _, model := range []string{"", "test-model"} { + blocked, reason, next := isAuthBlockedForModel(auth, model, time.Now()) + if !blocked || reason != blockReasonOther || !next.IsZero() { + t.Fatalf("isAuthBlockedForModel(%q) = %v, %v, %v; want true, other, zero", model, blocked, reason, next) + } + } +} + +func TestIsAuthBlockedForModel_ExpiredRecoveryIsAvailable(t *testing.T) { + t.Parallel() + + now := time.Now() + auth := &Auth{ + ID: "a", + Unavailable: true, + NextRetryAfter: now.Add(-time.Minute), + Quota: QuotaState{ + Exceeded: true, + NextRecoverAt: now.Add(-time.Second), + }, + } + blocked, reason, next := isAuthBlockedForModel(auth, "", now) + if blocked || reason != blockReasonNone || !next.IsZero() { + t.Fatalf("isAuthBlockedForModel() = %v, %v, %v; want false, none, zero", blocked, reason, next) + } +} + func TestFillFirstSelectorPick_ThinkingSuffixFallsBackToBaseModelState(t *testing.T) { t.Parallel() @@ -499,6 +742,75 @@ func TestSessionAffinitySelector_SameSessionSameAuth(t *testing.T) { } } +func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) + defer selector.Stop() + + authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "1"}} + authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "1"}} + auths := []*Auth{authA, authB} + opts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"metadata":{"user_id":"user_xxx_account__session_weight-change"}}`)} + + first, errFirst := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errFirst != nil { + t.Fatalf("first Pick() error = %v", errFirst) + } + if first.ID != authA.ID { + t.Fatalf("first Pick() auth.ID = %q, want %q", first.ID, authA.ID) + } + + authA.Attributes[AttributeWeight] = "0" + second, errSecond := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errSecond != nil { + t.Fatalf("Pick() after weight update error = %v", errSecond) + } + if second.ID != authB.ID { + t.Fatalf("Pick() after weight update auth.ID = %q, want %q", second.ID, authB.ID) + } + + authA.Attributes[AttributeWeight] = "10" + third, errThird := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errThird != nil { + t.Fatalf("Pick() after rebind error = %v", errThird) + } + if third.ID != authB.ID { + t.Fatalf("Pick() after rebind auth.ID = %q, want sticky auth %q", third.ID, authB.ID) + } +} + +func TestSessionAffinitySelector_WeightedNewSessionsResetAfterWeightChange(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) + defer selector.Stop() + authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "1000000"}} + authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "1"}} + auths := []*Auth{authA, authB} + pickSession := func(index int) *Auth { + t.Helper() + opts := cliproxyexecutor.Options{OriginalRequest: []byte(fmt.Sprintf(`{"session_id":"session-%d"}`, index))} + picked, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errPick != nil { + t.Fatalf("Pick(session-%d) error = %v", index, errPick) + } + return picked + } + for index := 0; index < 1000; index++ { + pickSession(index) + } + + authA.Attributes[AttributeWeight] = "1" + counts := make(map[string]int) + for index := 1000; index < 1020; index++ { + counts[pickSession(index).ID]++ + } + if counts[authA.ID] != 10 || counts[authB.ID] != 10 { + t.Fatalf("new session picks after weight change = %#v, want 10 each", counts) + } +} + func TestSessionAffinitySelector_NoSessionFallback(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/auth/weight.go b/sdk/cliproxy/auth/weight.go new file mode 100644 index 00000000..471bf9a6 --- /dev/null +++ b/sdk/cliproxy/auth/weight.go @@ -0,0 +1,49 @@ +package auth + +import ( + "fmt" + "strconv" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/credentialweight" +) + +// ValidateAuthWeight validates every explicit credential weight source. +func ValidateAuthWeight(auth *Auth) error { + if auth == nil { + return nil + } + if rawWeight, ok := auth.Attributes[AttributeWeight]; ok { + if _, errParse := credentialweight.ParseString(rawWeight); errParse != nil { + return fmt.Errorf("invalid attributes weight: %w", errParse) + } + } + if rawWeight, ok := auth.Metadata[AttributeWeight]; ok { + if _, errParse := credentialweight.ParseValue(rawWeight); errParse != nil { + return fmt.Errorf("invalid metadata weight: %w", errParse) + } + } + return nil +} + +// ApplyAuthWeightMetadata validates the auth and applies a source metadata weight. +func ApplyAuthWeightMetadata(auth *Auth, metadata map[string]any) error { + if errWeight := ValidateAuthWeight(auth); errWeight != nil { + return errWeight + } + if auth == nil || metadata == nil { + return nil + } + rawWeight, ok := metadata[AttributeWeight] + if !ok { + return nil + } + weight, errParse := credentialweight.ParseValue(rawWeight) + if errParse != nil { + return fmt.Errorf("invalid metadata weight: %w", errParse) + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Attributes[AttributeWeight] = strconv.FormatInt(weight, 10) + return nil +} diff --git a/sdk/cliproxy/auth/weight_test.go b/sdk/cliproxy/auth/weight_test.go new file mode 100644 index 00000000..ddba0cb1 --- /dev/null +++ b/sdk/cliproxy/auth/weight_test.go @@ -0,0 +1,43 @@ +package auth + +import ( + "encoding/json" + "testing" +) + +func TestValidateAuthWeight(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + auth *Auth + wantErr bool + }{ + {name: "omitted", auth: &Auth{}}, + {name: "positive attribute", auth: &Auth{Attributes: map[string]string{AttributeWeight: "7"}}}, + {name: "zero metadata", auth: &Auth{Metadata: map[string]any{AttributeWeight: json.Number("0")}}}, + {name: "negative attribute", auth: &Auth{Attributes: map[string]string{AttributeWeight: "-2"}}}, + {name: "fraction metadata", auth: &Auth{Metadata: map[string]any{AttributeWeight: json.Number("1.5")}}, wantErr: true}, + {name: "above maximum attribute", auth: &Auth{Attributes: map[string]string{AttributeWeight: "1000001"}}, wantErr: true}, + {name: "overflow metadata", auth: &Auth{Metadata: map[string]any{AttributeWeight: json.Number("9223372036854775808")}}, wantErr: true}, + {name: "nonnumeric attribute", auth: &Auth{Attributes: map[string]string{AttributeWeight: "invalid"}}, wantErr: true}, + { + name: "valid attribute does not hide invalid metadata", + auth: &Auth{ + Attributes: map[string]string{AttributeWeight: "2"}, + Metadata: map[string]any{AttributeWeight: 1.5}, + }, + wantErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + errValidate := ValidateAuthWeight(test.auth) + if (errValidate != nil) != test.wantErr { + t.Fatalf("ValidateAuthWeight() error = %v, wantErr = %v", errValidate, test.wantErr) + } + }) + } +} diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go index 1081b6a5..bc1a6853 100644 --- a/sdk/cliproxy/builder.go +++ b/sdk/cliproxy/builder.go @@ -194,6 +194,9 @@ func (b *Builder) Build() (*Service, error) { if b.configPath == "" { return nil, fmt.Errorf("cliproxy: configuration path is required") } + if errValidate := b.cfg.ValidateCredentialWeights(); errValidate != nil { + return nil, fmt.Errorf("cliproxy: validate credential weights: %w", errValidate) + } b.cfg.NormalizePluginsConfig() if errResolvePluginsDir := b.cfg.ResolvePluginsDir(); errResolvePluginsDir != nil && b.cfg.Plugins.Enabled { return nil, fmt.Errorf("cliproxy: %w", errResolvePluginsDir) diff --git a/sdk/cliproxy/builder_weight_validation_test.go b/sdk/cliproxy/builder_weight_validation_test.go new file mode 100644 index 00000000..7e505a9d --- /dev/null +++ b/sdk/cliproxy/builder_weight_validation_test.go @@ -0,0 +1,32 @@ +package cliproxy + +import ( + "strings" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestBuilderBuildRejectsInvalidWithConfigCredentialWeight(t *testing.T) { + invalidWeight := internalconfig.MaxCredentialWeight + 1 + cfg := &internalconfig.Config{ + ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "claude-key", + Weight: &invalidWeight, + }}, + } + + service, errBuild := NewBuilder(). + WithConfig(cfg). + WithConfigPath(t.TempDir() + "/config.yaml"). + Build() + if errBuild == nil { + t.Fatal("Build() accepted an invalid credential weight") + } + if service != nil { + t.Fatal("Build() returned a service for an invalid credential weight") + } + if !strings.Contains(errBuild.Error(), "cliproxy: validate credential weights: claude-api-key[0].weight") { + t.Fatalf("Build() error = %q, want contextual credential weight path", errBuild) + } +} diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index c0e74eab..4b0f12bd 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -40,6 +40,8 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { } switch strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) { + case "weighted-round-robin", "weightedroundrobin", "wrr": + state.strategy = "weighted-round-robin" case "fill-first", "fillfirst", "ff": state.strategy = "fill-first" } @@ -54,9 +56,12 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { func newRoutingSelector(state routingRuntimeState) coreauth.Selector { var selector coreauth.Selector - if state.strategy == "fill-first" { + switch state.strategy { + case "weighted-round-robin": + selector = &coreauth.WeightedRoundRobinSelector{} + case "fill-first": selector = &coreauth.FillFirstSelector{} - } else { + default: selector = &coreauth.RoundRobinSelector{} } if state.sessionAffinity { @@ -94,6 +99,10 @@ func (s *Service) commitConfigUpdate(newCfg *config.Config) configCommit { if newCfg == nil { return configCommit{} } + if errValidate := newCfg.ValidateCredentialWeights(); errValidate != nil { + log.WithError(errValidate).Warn("rejected config update with invalid credential weights") + return configCommit{} + } s.cfgMu.Lock() s.cfg = newCfg diff --git a/sdk/cliproxy/service_config_weight_test.go b/sdk/cliproxy/service_config_weight_test.go new file mode 100644 index 00000000..e4a0bd7a --- /dev/null +++ b/sdk/cliproxy/service_config_weight_test.go @@ -0,0 +1,42 @@ +package cliproxy + +import ( + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestWeightedRoundRobinRoutingSelector(t *testing.T) { + state := normalizedRoutingRuntimeState(&internalconfig.Config{ + Routing: internalconfig.RoutingConfig{Strategy: "wrr"}, + }) + if state.strategy != "weighted-round-robin" { + t.Fatalf("strategy = %q, want weighted-round-robin", state.strategy) + } + if _, ok := newRoutingSelector(state).(*coreauth.WeightedRoundRobinSelector); !ok { + t.Fatalf("selector type = %T, want *auth.WeightedRoundRobinSelector", newRoutingSelector(state)) + } +} + +func TestServiceRejectsInvalidCredentialWeightConfigCommit(t *testing.T) { + originalCfg := &internalconfig.Config{} + service := &Service{cfg: originalCfg} + invalidWeight := internalconfig.MaxCredentialWeight + 1 + newCfg := &internalconfig.Config{ + VertexCompatAPIKey: []internalconfig.VertexCompatKey{{ + APIKey: "vertex-key", + Weight: &invalidWeight, + }}, + } + + if service.applyConfigUpdateWithAuthSynthesis(nil, newCfg, true) { + t.Fatal("hot config application accepted an invalid credential weight") + } + if service.cfg != originalCfg { + t.Fatal("invalid hot config replaced the active config") + } + if service.configSequence != 0 { + t.Fatalf("config sequence = %d, want 0", service.configSequence) + } +}