diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index c599ba39b..d69ddc149 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -372,7 +372,7 @@ func (m *Manager) predictedHomeConcurrencyModel(auth *Auth, routeModel string) ( requestedModel := rewriteModelForAuth(routeModel, auth) aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) - if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) != 0 { + if pool := m.resolveAPIKeyUpstreamModelPool(auth, upstreamModel); len(pool) != 0 { if len(pool) != 1 { return "", false } diff --git a/sdk/cliproxy/auth/conductor_models.go b/sdk/cliproxy/auth/conductor_models.go index 50788157d..b827494a7 100644 --- a/sdk/cliproxy/auth/conductor_models.go +++ b/sdk/cliproxy/auth/conductor_models.go @@ -88,6 +88,22 @@ func openAICompatModelPoolKey(auth *Auth, requestedModel string) string { return strings.ToLower(strings.TrimSpace(auth.ID)) + "|" + openAICompatProviderKey(auth) + "|" + strings.ToLower(base) } +func apiKeyModelPoolKey(auth *Auth, requestedModel string) string { + base := strings.TrimSpace(thinking.ParseSuffix(requestedModel).ModelName) + if base == "" { + base = strings.TrimSpace(requestedModel) + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + return strings.ToLower(strings.TrimSpace(auth.ID)) + "|" + provider + "|" + strings.ToLower(base) +} + +func modelPoolKey(auth *Auth, requestedModel string) string { + if isConfiguredOpenAICompatAuth(auth) { + return openAICompatModelPoolKey(auth, requestedModel) + } + return apiKeyModelPoolKey(auth, requestedModel) +} + func (m *Manager) nextModelPoolOffset(key string, size int) int { if m == nil || size <= 1 { return 0 @@ -156,6 +172,28 @@ func resolveOpenAICompatUpstreamModelPool(cfg *internalconfig.Config, auth *Auth return resolveModelAliasPoolFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) } +func (m *Manager) resolveAPIKeyUpstreamModelPool(auth *Auth, requestedModel string) []string { + return resolveAPIKeyUpstreamModelPool(m.loadAPIKeyModelRouting().config, auth, requestedModel) +} + +func resolveAPIKeyUpstreamModelPool(cfg *internalconfig.Config, auth *Auth, requestedModel string) []string { + if !isConfiguredModelRoutingAuth(auth) { + return nil + } + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return nil + } + if cfg == nil { + cfg = &internalconfig.Config{} + } + models := configuredModelAliasEntries(cfg, auth) + if len(models) == 0 { + return nil + } + return resolveModelAliasPoolFromConfigModels(requestedModel, models) +} + func preserveRequestedModelSuffix(requestedModel, resolved string) string { return preserveResolvedModelSuffix(resolved, thinking.ParseSuffix(requestedModel)) } @@ -168,11 +206,11 @@ func (m *Manager) executionModelCandidates(auth *Auth, routeModel string) []stri } requestedModel := rewriteModelForAuth(routeModel, auth) requestedModel = m.applyOAuthModelAlias(auth, requestedModel) - if pool := m.resolveOpenAICompatUpstreamModelPool(auth, requestedModel); len(pool) > 0 { + if pool := m.resolveAPIKeyUpstreamModelPool(auth, requestedModel); len(pool) > 0 { if len(pool) == 1 { return pool } - offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, requestedModel), len(pool)) + offset := m.nextModelPoolOffset(modelPoolKey(auth, requestedModel), len(pool)) return rotateStrings(pool, offset) } resolved := m.applyAPIKeyModelAlias(auth, requestedModel) @@ -289,11 +327,11 @@ func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel strin } } if len(candidates) == 0 { - if pool := resolveOpenAICompatUpstreamModelPool(routing.config, auth, upstreamModel); len(pool) > 0 { + if pool := resolveAPIKeyUpstreamModelPool(routing.config, auth, upstreamModel); len(pool) > 0 { if len(pool) == 1 { candidates = pool } else { - offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, upstreamModel), len(pool)) + offset := m.nextModelPoolOffset(modelPoolKey(auth, upstreamModel), len(pool)) candidates = rotateStrings(pool, offset) } } else { diff --git a/sdk/cliproxy/auth/conductor_models_test.go b/sdk/cliproxy/auth/conductor_models_test.go new file mode 100644 index 000000000..57b155be0 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_models_test.go @@ -0,0 +1,243 @@ +package auth + +import ( + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestExecutionModelCandidates_APIKeyAliasPoolRotates(t *testing.T) { + cfg := &internalconfig.Config{ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "test-key", + Prefix: "tenant", + Models: []internalconfig.ClaudeModel{ + {Name: "claude-sonnet-4", Alias: "public"}, + {Name: "claude-sonnet-3.5", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-claude-pool", + Provider: "claude", + Prefix: "tenant", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "test-key", + AttributeSource: "config:claude[0]", + AttributeConfigIndex: "0", + }, + } + + first := manager.executionModelCandidates(auth, "tenant/public") + if len(first) != 2 || first[0] != "claude-sonnet-4" || first[1] != "claude-sonnet-3.5" { + t.Fatalf("first candidates = %v, want [claude-sonnet-4 claude-sonnet-3.5]", first) + } + + second := manager.executionModelCandidates(auth, "tenant/public") + if len(second) != 2 || second[0] != "claude-sonnet-3.5" || second[1] != "claude-sonnet-4" { + t.Fatalf("second candidates = %v, want [claude-sonnet-3.5 claude-sonnet-4]", second) + } + + third := manager.executionModelCandidates(auth, "tenant/public") + if len(third) != 2 || third[0] != "claude-sonnet-4" || third[1] != "claude-sonnet-3.5" { + t.Fatalf("third candidates = %v, want [claude-sonnet-4 claude-sonnet-3.5]", third) + } +} + +func TestExecutionModelCandidates_APIKeyAliasPoolRotatesWithSuffix(t *testing.T) { + cfg := &internalconfig.Config{ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "test-key", + Prefix: "tenant", + Models: []internalconfig.ClaudeModel{ + {Name: "claude-sonnet-4", Alias: "public"}, + {Name: "claude-sonnet-3.5", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-claude-pool-suffix", + Provider: "claude", + Prefix: "tenant", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "test-key", + AttributeSource: "config:claude[0]", + AttributeConfigIndex: "0", + }, + } + + first := manager.executionModelCandidates(auth, "tenant/public(8192)") + want := []string{"claude-sonnet-4(8192)", "claude-sonnet-3.5(8192)"} + if len(first) != 2 || first[0] != want[0] || first[1] != want[1] { + t.Fatalf("first candidates = %v, want %v", first, want) + } +} + +func TestPreparedExecutionModels_APIKeyPoolSkipsBlockedMembers(t *testing.T) { + cfg := &internalconfig.Config{ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "test-key", + Prefix: "tenant", + Models: []internalconfig.ClaudeModel{ + {Name: "claude-sonnet-4", Alias: "public"}, + {Name: "claude-sonnet-3.5", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-claude-pool-blocked", + Provider: "claude", + Prefix: "tenant", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "test-key", + AttributeSource: "config:claude[0]", + AttributeConfigIndex: "0", + }, + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": { + Unavailable: true, + NextRetryAfter: time.Now().Add(time.Hour), + }, + }, + } + + models, pooled := manager.preparedExecutionModels(auth, "tenant/public") + if !pooled { + t.Fatalf("pooled = false, want true") + } + if len(models) != 1 || models[0] != "claude-sonnet-3.5" { + t.Fatalf("filtered models = %v, want [claude-sonnet-3.5]", models) + } +} + +func TestExecutionModelCandidates_APIKeySingleModelUnchanged(t *testing.T) { + cfg := &internalconfig.Config{GeminiKey: []internalconfig.GeminiKey{{ + APIKey: "gemini-key", + Prefix: "team", + Models: []internalconfig.GeminiModel{ + {Name: "gemini-2.5-pro", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-gemini-single", + Provider: "gemini", + Prefix: "team", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "gemini-key", + AttributeSource: "config:gemini[0]", + AttributeConfigIndex: "0", + }, + } + + got := manager.executionModelCandidates(auth, "team/public") + if len(got) != 1 || got[0] != "gemini-2.5-pro" { + t.Fatalf("single model candidates = %v, want [gemini-2.5-pro]", got) + } +} + +func TestExecutionModelCandidates_APIKeyPoolForCodex(t *testing.T) { + cfg := &internalconfig.Config{CodexKey: []internalconfig.CodexKey{{ + APIKey: "codex-key", + Prefix: "team", + Models: []internalconfig.CodexModel{ + {Name: "deepseek-v4", Alias: "fast"}, + {Name: "gpt-5.4", Alias: "fast"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-codex-pool", + Provider: "codex", + Prefix: "team", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "codex-key", + AttributeSource: "config:codex[0]", + AttributeConfigIndex: "0", + }, + } + + first := manager.executionModelCandidates(auth, "team/fast") + if len(first) != 2 || first[0] != "deepseek-v4" || first[1] != "gpt-5.4" { + t.Fatalf("first codex candidates = %v, want [deepseek-v4 gpt-5.4]", first) + } +} + +func TestPredictedHomeConcurrencyModel_APIKeyPoolRejectsMulti(t *testing.T) { + cfg := &internalconfig.Config{ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "test-key", + Prefix: "tenant", + Models: []internalconfig.ClaudeModel{ + {Name: "claude-sonnet-4", Alias: "public"}, + {Name: "claude-sonnet-3.5", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-claude-pool-home", + Provider: "claude", + Prefix: "tenant", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "test-key", + AttributeSource: "config:claude[0]", + AttributeConfigIndex: "0", + }, + } + + model, ok := manager.predictedHomeConcurrencyModel(auth, "tenant/public") + if ok { + t.Fatalf("predictedHomeConcurrencyModel for multi-model pool = (%q, true), want empty", model) + } +} + +func TestPredictedHomeConcurrencyModel_APIKeyPoolAcceptsSingle(t *testing.T) { + cfg := &internalconfig.Config{ClaudeKey: []internalconfig.ClaudeKey{{ + APIKey: "test-key", + Prefix: "tenant", + Models: []internalconfig.ClaudeModel{ + {Name: "claude-sonnet-4", Alias: "public"}, + }, + }}} + + manager := NewManager(nil, nil, nil) + manager.SetConfig(cfg) + + auth := &Auth{ + ID: "auth-claude-single-home", + Provider: "claude", + Prefix: "tenant", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + AttributeAPIKey: "test-key", + AttributeSource: "config:claude[0]", + AttributeConfigIndex: "0", + }, + } + + model, ok := manager.predictedHomeConcurrencyModel(auth, "tenant/public") + if !ok || model != "claude-sonnet-4" { + t.Fatalf("predictedHomeConcurrencyModel for single-model pool = (%q, %t), want (claude-sonnet-4, true)", model, ok) + } +}