mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-09 01:46:48 +08:00
fix(auth): rewrite API key alias response models (#4002)
This commit is contained in:
@@ -231,6 +231,7 @@ nonstream-keepalive-interval: 0
|
||||
# models:
|
||||
# - name: "gpt-5-codex" # upstream model name
|
||||
# alias: "codex-latest" # client alias mapped to the upstream model
|
||||
# force-mapping: true # optional: rewrite response model fields back to the alias
|
||||
# excluded-models:
|
||||
# - "gpt-5.1" # exclude specific models (exact match)
|
||||
# - "gpt-5-*" # wildcard matching prefix (e.g. gpt-5-medium, gpt-5-codex)
|
||||
@@ -251,6 +252,7 @@ nonstream-keepalive-interval: 0
|
||||
# models:
|
||||
# - name: "claude-3-5-sonnet-20241022" # upstream model name
|
||||
# alias: "claude-sonnet-latest" # client alias mapped to the upstream model
|
||||
# force-mapping: true # optional: rewrite response model fields back to the alias
|
||||
# excluded-models:
|
||||
# - "claude-opus-4-5-20251101" # exclude specific models (exact match)
|
||||
# - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219)
|
||||
|
||||
@@ -483,10 +483,14 @@ type ClaudeModel struct {
|
||||
|
||||
// Alias is the client-facing model name that maps to Name.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// ForceMapping rewrites upstream response model fields back to Alias.
|
||||
ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"`
|
||||
}
|
||||
|
||||
func (m ClaudeModel) GetName() string { return m.Name }
|
||||
func (m ClaudeModel) GetAlias() string { return m.Alias }
|
||||
func (m ClaudeModel) GetName() string { return m.Name }
|
||||
func (m ClaudeModel) GetAlias() string { return m.Alias }
|
||||
func (m ClaudeModel) GetForceMapping() bool { return m.ForceMapping }
|
||||
|
||||
// CodexKey represents the configuration for a Codex API key,
|
||||
// including the API key itself and an optional base URL for the API endpoint.
|
||||
@@ -534,10 +538,14 @@ type CodexModel struct {
|
||||
|
||||
// Alias is the client-facing model name that maps to Name.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// ForceMapping rewrites upstream response model fields back to Alias.
|
||||
ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"`
|
||||
}
|
||||
|
||||
func (m CodexModel) GetName() string { return m.Name }
|
||||
func (m CodexModel) GetAlias() string { return m.Alias }
|
||||
func (m CodexModel) GetName() string { return m.Name }
|
||||
func (m CodexModel) GetAlias() string { return m.Alias }
|
||||
func (m CodexModel) GetForceMapping() bool { return m.ForceMapping }
|
||||
|
||||
// GeminiKey represents the configuration for a Gemini API key,
|
||||
// including optional overrides for upstream base URL, proxy routing, and headers.
|
||||
@@ -581,10 +589,14 @@ type GeminiModel struct {
|
||||
|
||||
// Alias is the client-facing model name that maps to Name.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// ForceMapping rewrites upstream response model fields back to Alias.
|
||||
ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"`
|
||||
}
|
||||
|
||||
func (m GeminiModel) GetName() string { return m.Name }
|
||||
func (m GeminiModel) GetAlias() string { return m.Alias }
|
||||
func (m GeminiModel) GetName() string { return m.Name }
|
||||
func (m GeminiModel) GetAlias() string { return m.Alias }
|
||||
func (m GeminiModel) GetForceMapping() bool { return m.ForceMapping }
|
||||
|
||||
// OpenAICompatibility represents the configuration for OpenAI API compatibility
|
||||
// with external providers, allowing model aliases to be routed through OpenAI API format.
|
||||
@@ -636,6 +648,9 @@ type OpenAICompatibilityModel struct {
|
||||
// Alias is the model name alias that clients will use to reference this model.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// ForceMapping rewrites upstream response model fields back to Alias.
|
||||
ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"`
|
||||
|
||||
// Image marks this model as callable through /v1/images/generations and /v1/images/edits.
|
||||
Image bool `yaml:"image,omitempty" json:"image,omitempty"`
|
||||
|
||||
@@ -644,8 +659,9 @@ type OpenAICompatibilityModel struct {
|
||||
Thinking *registry.ThinkingSupport `yaml:"thinking,omitempty" json:"thinking,omitempty"`
|
||||
}
|
||||
|
||||
func (m OpenAICompatibilityModel) GetName() string { return m.Name }
|
||||
func (m OpenAICompatibilityModel) GetAlias() string { return m.Alias }
|
||||
func (m OpenAICompatibilityModel) GetName() string { return m.Name }
|
||||
func (m OpenAICompatibilityModel) GetAlias() string { return m.Alias }
|
||||
func (m OpenAICompatibilityModel) GetForceMapping() bool { return m.ForceMapping }
|
||||
|
||||
// LoadConfig reads a YAML configuration file from the given path,
|
||||
// unmarshals it into a Config struct, applies environment variable overrides,
|
||||
|
||||
@@ -50,10 +50,14 @@ type VertexCompatModel struct {
|
||||
|
||||
// Alias is the model name alias that clients will use to reference this model.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// ForceMapping rewrites upstream response model fields back to Alias.
|
||||
ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"`
|
||||
}
|
||||
|
||||
func (m VertexCompatModel) GetName() string { return m.Name }
|
||||
func (m VertexCompatModel) GetAlias() string { return m.Alias }
|
||||
func (m VertexCompatModel) GetName() string { return m.Name }
|
||||
func (m VertexCompatModel) GetAlias() string { return m.Alias }
|
||||
func (m VertexCompatModel) GetForceMapping() bool { return m.ForceMapping }
|
||||
|
||||
// SanitizeVertexCompatKeys deduplicates and normalizes Vertex-compatible API key credentials.
|
||||
func (cfg *Config) SanitizeVertexCompatKeys() {
|
||||
|
||||
@@ -178,3 +178,92 @@ func TestApplyAPIKeyModelAlias(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAPIKeyModelAliasWithResult_ForceMapping(t *testing.T) {
|
||||
cfg := &internalconfig.Config{
|
||||
ClaudeKey: []internalconfig.ClaudeKey{{
|
||||
APIKey: "claude-key",
|
||||
Models: []internalconfig.ClaudeModel{{
|
||||
Name: "glm-5.2",
|
||||
Alias: "claude-sonnet-latest",
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}},
|
||||
}
|
||||
|
||||
mgr := NewManager(nil, nil, nil)
|
||||
mgr.SetConfig(cfg)
|
||||
|
||||
ctx := context.Background()
|
||||
auth := &Auth{ID: "claude-auth", Provider: "claude", Attributes: map[string]string{"api_key": "claude-key"}}
|
||||
if _, err := mgr.Register(ctx, auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
}
|
||||
|
||||
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "claude-sonnet-latest")
|
||||
if result.UpstreamModel != "glm-5.2" || !result.ForceMapping || result.OriginalAlias != "claude-sonnet-latest" {
|
||||
t.Fatalf("resolveAPIKeyModelAliasWithResult() = %+v, want upstream glm-5.2 with force mapping", result)
|
||||
}
|
||||
|
||||
noRewrite := mgr.resolveAPIKeyModelAliasWithResult(auth, "glm-5.2")
|
||||
if noRewrite.UpstreamModel != "glm-5.2" || noRewrite.ForceMapping || noRewrite.OriginalAlias != "" {
|
||||
t.Fatalf("resolveAPIKeyModelAliasWithResult() direct upstream = %+v, want passthrough without rewrite", noRewrite)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAPIKeyModelAliasWithResult_SameBasePreservesSuffix(t *testing.T) {
|
||||
cfg := &internalconfig.Config{
|
||||
GeminiKey: []internalconfig.GeminiKey{{
|
||||
APIKey: "k",
|
||||
Models: []internalconfig.GeminiModel{{
|
||||
Name: "gemini-2.5-pro",
|
||||
Alias: "gemini-2.5-pro(8192)",
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}},
|
||||
}
|
||||
|
||||
mgr := NewManager(nil, nil, nil)
|
||||
mgr.SetConfig(cfg)
|
||||
|
||||
ctx := context.Background()
|
||||
auth := &Auth{ID: "gemini-auth", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}}
|
||||
if _, err := mgr.Register(ctx, auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
}
|
||||
|
||||
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "gemini-2.5-pro(8192)")
|
||||
if result.UpstreamModel != "gemini-2.5-pro(8192)" || !result.ForceMapping || result.OriginalAlias != "gemini-2.5-pro(8192)" {
|
||||
t.Fatalf("resolveAPIKeyModelAliasWithResult() = %+v, want same-base suffix preserved", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAPIKeyModelAliasWithResult_ForceMappingUsesConfigAliasNotRequestSuffix(t *testing.T) {
|
||||
cfg := &internalconfig.Config{
|
||||
CodexKey: []internalconfig.CodexKey{{
|
||||
APIKey: "codex-key",
|
||||
Models: []internalconfig.CodexModel{{
|
||||
Name: "gpt-5.5",
|
||||
Alias: "claude-sonnet-4-5",
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}},
|
||||
}
|
||||
|
||||
mgr := NewManager(nil, nil, nil)
|
||||
mgr.SetConfig(cfg)
|
||||
|
||||
ctx := context.Background()
|
||||
auth := &Auth{ID: "codex-auth", Provider: "codex", Attributes: map[string]string{"api_key": "codex-key"}}
|
||||
if _, err := mgr.Register(ctx, auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
}
|
||||
|
||||
result := mgr.resolveAPIKeyModelAliasWithResult(auth, "claude-sonnet-4-5(high)")
|
||||
if result.UpstreamModel != "gpt-5.5(high)" {
|
||||
t.Fatalf("upstream = %q want gpt-5.5(high)", result.UpstreamModel)
|
||||
}
|
||||
if result.OriginalAlias != "claude-sonnet-4-5" {
|
||||
t.Fatalf("OriginalAlias = %q want claude-sonnet-4-5", result.OriginalAlias)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1209,8 +1209,8 @@ func (m *Manager) preparedExecutionModelsWithAlias(auth *Auth, routeModel string
|
||||
|
||||
func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) {
|
||||
requestedModel := rewriteModelForAuth(routeModel, auth)
|
||||
aliasResult := m.applyOAuthModelAliasWithResult(auth, requestedModel)
|
||||
upstreamModel := aliasResult.UpstreamModel
|
||||
aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel)
|
||||
upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult)
|
||||
|
||||
var candidates []string
|
||||
if auth != nil && auth.Attributes != nil {
|
||||
@@ -1238,6 +1238,84 @@ func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel strin
|
||||
return candidates, pooled, aliasResult
|
||||
}
|
||||
|
||||
func (m *Manager) resolveExecutionAliasResult(auth *Auth, routeModel string) OAuthModelAliasResult {
|
||||
requestedModel := rewriteModelForAuth(routeModel, auth)
|
||||
return m.resolveExecutionAliasResultForRequested(auth, requestedModel)
|
||||
}
|
||||
|
||||
func (m *Manager) resolveExecutionAliasResultForRequested(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
||||
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
||||
return m.resolveAPIKeyModelAliasWithResult(auth, requestedModel)
|
||||
}
|
||||
return m.applyOAuthModelAliasWithResult(auth, requestedModel)
|
||||
}
|
||||
|
||||
func executionAliasPoolModel(auth *Auth, requestedModel string, aliasResult OAuthModelAliasResult) string {
|
||||
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
||||
if strings.TrimSpace(requestedModel) != "" {
|
||||
return requestedModel
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(aliasResult.UpstreamModel) != "" {
|
||||
return aliasResult.UpstreamModel
|
||||
}
|
||||
return requestedModel
|
||||
}
|
||||
|
||||
func (m *Manager) resolveAPIKeyModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
||||
if m == nil || auth == nil {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
requestedModel = strings.TrimSpace(requestedModel)
|
||||
if requestedModel == "" {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
||||
if cfg == nil {
|
||||
cfg = &internalconfig.Config{}
|
||||
}
|
||||
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
var models []modelAliasEntry
|
||||
switch provider {
|
||||
case "gemini":
|
||||
if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil {
|
||||
models = asModelAliasEntries(entry.Models)
|
||||
}
|
||||
case "claude":
|
||||
if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil {
|
||||
models = asModelAliasEntries(entry.Models)
|
||||
}
|
||||
case "codex":
|
||||
if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil {
|
||||
models = asModelAliasEntries(entry.Models)
|
||||
}
|
||||
case "vertex":
|
||||
if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
|
||||
models = asModelAliasEntries(entry.Models)
|
||||
}
|
||||
default:
|
||||
providerKey := ""
|
||||
compatName := ""
|
||||
if auth.Attributes != nil {
|
||||
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
||||
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
||||
}
|
||||
if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
||||
if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil {
|
||||
models = asModelAliasEntries(entry.Models)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(models) == 0 {
|
||||
return OAuthModelAliasResult{UpstreamModel: requestedModel}
|
||||
}
|
||||
result := resolveModelAliasResultFromConfigModels(requestedModel, models)
|
||||
if strings.TrimSpace(result.UpstreamModel) == "" {
|
||||
return OAuthModelAliasResult{UpstreamModel: requestedModel}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (m *Manager) prepareExecutionModels(auth *Auth, routeModel string) []string {
|
||||
models, _ := m.preparedExecutionModels(auth, routeModel)
|
||||
return models
|
||||
@@ -3164,6 +3242,7 @@ func resolveOpenAICompatConfig(cfg *internalconfig.Config, providerKey, compatNa
|
||||
func asModelAliasEntries[T interface {
|
||||
GetName() string
|
||||
GetAlias() string
|
||||
GetForceMapping() bool
|
||||
}](models []T) []modelAliasEntry {
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
|
||||
@@ -553,3 +553,144 @@ func TestManagerExecuteStream_AntigravityCreditsFallbackForceMappingRewritesResp
|
||||
t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel)
|
||||
}
|
||||
}
|
||||
|
||||
func setupAPIKeyForceMappingManager(t *testing.T, provider, upstreamModel, aliasModel string) (*Manager, *forceMappingExecutor) {
|
||||
t.Helper()
|
||||
manager := NewManager(nil, nil, nil)
|
||||
executor := &forceMappingExecutor{id: provider}
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
cfg := &internalconfig.Config{}
|
||||
apiKey := provider + "-key"
|
||||
switch provider {
|
||||
case "claude":
|
||||
cfg.ClaudeKey = []internalconfig.ClaudeKey{{
|
||||
APIKey: apiKey,
|
||||
Models: []internalconfig.ClaudeModel{{
|
||||
Name: upstreamModel,
|
||||
Alias: aliasModel,
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}}
|
||||
case "codex":
|
||||
cfg.CodexKey = []internalconfig.CodexKey{{
|
||||
APIKey: apiKey,
|
||||
Models: []internalconfig.CodexModel{{
|
||||
Name: upstreamModel,
|
||||
Alias: aliasModel,
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}}
|
||||
case "vertex":
|
||||
cfg.VertexCompatAPIKey = []internalconfig.VertexCompatKey{{
|
||||
APIKey: apiKey,
|
||||
Models: []internalconfig.VertexCompatModel{{
|
||||
Name: upstreamModel,
|
||||
Alias: aliasModel,
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}}
|
||||
case "openai-compatibility":
|
||||
cfg.OpenAICompatibility = []internalconfig.OpenAICompatibility{{
|
||||
Name: provider,
|
||||
Models: []internalconfig.OpenAICompatibilityModel{{
|
||||
Name: upstreamModel,
|
||||
Alias: aliasModel,
|
||||
ForceMapping: true,
|
||||
}},
|
||||
}}
|
||||
default:
|
||||
t.Fatalf("unsupported provider %q", provider)
|
||||
}
|
||||
manager.SetConfig(cfg)
|
||||
|
||||
auth := &Auth{
|
||||
ID: provider + "-api-key-force-mapping-auth",
|
||||
Provider: provider,
|
||||
Attributes: map[string]string{"api_key": apiKey},
|
||||
}
|
||||
if provider == "openai-compatibility" {
|
||||
auth.Attributes["compat_name"] = provider
|
||||
auth.Attributes["provider_key"] = provider
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("register auth: %v", errRegister)
|
||||
}
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: aliasModel}, {ID: upstreamModel}})
|
||||
t.Cleanup(func() {
|
||||
reg.UnregisterClient(auth.ID)
|
||||
})
|
||||
manager.RefreshSchedulerEntry(auth.ID)
|
||||
|
||||
return manager, executor
|
||||
}
|
||||
|
||||
func TestManagerExecute_APIKeyAliasForceMappingRewritesResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
provider string
|
||||
upstreamModel string
|
||||
aliasModel string
|
||||
}{
|
||||
{provider: "claude", upstreamModel: "glm-5.2", aliasModel: "claude-sonnet-latest"},
|
||||
{provider: "codex", upstreamModel: "gpt-5.5", aliasModel: "claude-sonnet-4-5"},
|
||||
{provider: "vertex", upstreamModel: "gemini-3-pro", aliasModel: "claude-opus-4-5"},
|
||||
{provider: "openai-compatibility", upstreamModel: "deepseek-v3.1", aliasModel: "claude-opus-4.66"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.provider, func(t *testing.T) {
|
||||
manager, executor := setupAPIKeyForceMappingManager(t, tt.provider, tt.upstreamModel, tt.aliasModel)
|
||||
resp, errExecute := manager.Execute(context.Background(), []string{tt.provider}, cliproxyexecutor.Request{Model: tt.aliasModel}, cliproxyexecutor.Options{})
|
||||
if errExecute != nil {
|
||||
t.Fatalf("execute error = %v, want success", errExecute)
|
||||
}
|
||||
|
||||
gotModels := executor.ExecuteModels()
|
||||
if len(gotModels) != 1 || gotModels[0] != tt.upstreamModel {
|
||||
t.Fatalf("execute models = %v, want [%s]", gotModels, tt.upstreamModel)
|
||||
}
|
||||
if got := string(resp.Payload); !strings.Contains(got, tt.aliasModel) || forceMappingPayloadLeaksUpstream(got, tt.upstreamModel) {
|
||||
t.Fatalf("response payload = %s, want alias %q without upstream %q", got, tt.aliasModel, tt.upstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerExecuteStream_APIKeyAliasForceMappingRewritesResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
provider string
|
||||
upstreamModel string
|
||||
aliasModel string
|
||||
}{
|
||||
{provider: "claude", upstreamModel: "glm-5.2", aliasModel: "claude-sonnet-latest"},
|
||||
{provider: "codex", upstreamModel: "gpt-5.5", aliasModel: "claude-sonnet-4-5"},
|
||||
{provider: "vertex", upstreamModel: "gemini-3-pro", aliasModel: "claude-opus-4-5"},
|
||||
{provider: "openai-compatibility", upstreamModel: "deepseek-v3.1", aliasModel: "claude-opus-4.66"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.provider, func(t *testing.T) {
|
||||
manager, executor := setupAPIKeyForceMappingManager(t, tt.provider, tt.upstreamModel, tt.aliasModel)
|
||||
streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{tt.provider}, cliproxyexecutor.Request{Model: tt.aliasModel}, cliproxyexecutor.Options{})
|
||||
if errExecute != nil {
|
||||
t.Fatalf("execute stream error = %v, want success", errExecute)
|
||||
}
|
||||
|
||||
gotModels := executor.StreamModels()
|
||||
if len(gotModels) != 1 || gotModels[0] != tt.upstreamModel {
|
||||
t.Fatalf("stream models = %v, want [%s]", gotModels, tt.upstreamModel)
|
||||
}
|
||||
|
||||
var payload []byte
|
||||
for chunk := range streamResult.Chunks {
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("unexpected stream error: %v", chunk.Err)
|
||||
}
|
||||
payload = append(payload, chunk.Payload...)
|
||||
}
|
||||
if got := string(payload); !strings.Contains(got, tt.aliasModel) || forceMappingPayloadLeaksUpstream(got, tt.upstreamModel) {
|
||||
t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, tt.aliasModel, tt.upstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ const oauthModelAliasesAttributeKey = "model_aliases"
|
||||
type modelAliasEntry interface {
|
||||
GetName() string
|
||||
GetAlias() string
|
||||
GetForceMapping() bool
|
||||
}
|
||||
|
||||
// oauthModelAliasEntry stores the upstream model name and mapping flags for an alias.
|
||||
@@ -200,6 +201,54 @@ func resolveModelAliasFromConfigModels(requestedModel string, models []modelAlia
|
||||
return ""
|
||||
}
|
||||
|
||||
func resolveModelAliasResultFromConfigModels(requestedModel string, models []modelAliasEntry) OAuthModelAliasResult {
|
||||
requestedModel = strings.TrimSpace(requestedModel)
|
||||
if requestedModel == "" || len(models) == 0 {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
requestResult, candidates := modelAliasLookupCandidates(requestedModel)
|
||||
if len(candidates) == 0 {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
baseModel := requestResult.ModelName
|
||||
if baseModel == "" {
|
||||
baseModel = requestedModel
|
||||
}
|
||||
for i := range models {
|
||||
original := strings.TrimSpace(models[i].GetName())
|
||||
alias := strings.TrimSpace(models[i].GetAlias())
|
||||
if original == "" || alias == "" {
|
||||
continue
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
key := strings.TrimSpace(candidate)
|
||||
if key == "" || !strings.EqualFold(alias, key) {
|
||||
continue
|
||||
}
|
||||
if strings.EqualFold(original, baseModel) {
|
||||
if !models[i].GetForceMapping() {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
return OAuthModelAliasResult{
|
||||
UpstreamModel: preserveResolvedModelSuffix(original, requestResult),
|
||||
ForceMapping: models[i].GetForceMapping(),
|
||||
OriginalAlias: oauthModelAliasForceMappingResponseModel(alias),
|
||||
}
|
||||
}
|
||||
originalAlias := requestedModel
|
||||
if models[i].GetForceMapping() {
|
||||
originalAlias = oauthModelAliasForceMappingResponseModel(alias)
|
||||
}
|
||||
return OAuthModelAliasResult{
|
||||
UpstreamModel: preserveResolvedModelSuffix(original, requestResult),
|
||||
ForceMapping: models[i].GetForceMapping(),
|
||||
OriginalAlias: originalAlias,
|
||||
}
|
||||
}
|
||||
}
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
|
||||
// resolveOAuthUpstreamModel resolves the upstream model name from OAuth model alias.
|
||||
// If an alias exists, returns the original (upstream) model name that corresponds
|
||||
// to the requested alias.
|
||||
|
||||
@@ -21,6 +21,7 @@ type openAICompatPoolExecutor struct {
|
||||
executeModels []string
|
||||
countModels []string
|
||||
streamModels []string
|
||||
executePayloads map[string][]byte
|
||||
executeErrors map[string]error
|
||||
countErrors map[string]error
|
||||
streamFirstErrors map[string]error
|
||||
@@ -35,11 +36,15 @@ func (e *openAICompatPoolExecutor) Execute(ctx context.Context, auth *Auth, req
|
||||
_ = opts
|
||||
e.mu.Lock()
|
||||
e.executeModels = append(e.executeModels, req.Model)
|
||||
payload := append([]byte(nil), e.executePayloads[req.Model]...)
|
||||
err := e.executeErrors[req.Model]
|
||||
e.mu.Unlock()
|
||||
if err != nil {
|
||||
return cliproxyexecutor.Response{}, err
|
||||
}
|
||||
if len(payload) > 0 {
|
||||
return cliproxyexecutor.Response{Payload: payload}, nil
|
||||
}
|
||||
return cliproxyexecutor.Response{Payload: []byte(req.Model)}, nil
|
||||
}
|
||||
|
||||
@@ -281,6 +286,44 @@ func TestManagerExecute_OpenAICompatAliasPoolRotatesWithinAuth(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerExecute_OpenAICompatAliasPoolForceMappingRotatesAndRewritesResponse(t *testing.T) {
|
||||
alias := "claude-opus-4.66"
|
||||
executor := &openAICompatPoolExecutor{
|
||||
id: openAICompatPoolProviderKey,
|
||||
executePayloads: map[string][]byte{
|
||||
"deepseek-v3.1": []byte(`{"model":"deepseek-v3.1"}`),
|
||||
"glm-5": []byte(`{"model":"glm-5"}`),
|
||||
},
|
||||
}
|
||||
m := newOpenAICompatPoolTestManager(t, alias, []internalconfig.OpenAICompatibilityModel{
|
||||
{Name: "deepseek-v3.1", Alias: alias, ForceMapping: true},
|
||||
{Name: "glm-5", Alias: alias, ForceMapping: true},
|
||||
}, executor)
|
||||
|
||||
var payloads []string
|
||||
for i := 0; i < 2; i++ {
|
||||
resp, err := m.Execute(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: alias}, cliproxyexecutor.Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("execute %d: %v", i, err)
|
||||
}
|
||||
payloads = append(payloads, string(resp.Payload))
|
||||
}
|
||||
|
||||
got := executor.ExecuteModels()
|
||||
wantModels := []string{"deepseek-v3.1", "glm-5"}
|
||||
for i := range wantModels {
|
||||
if got[i] != wantModels[i] {
|
||||
t.Fatalf("execute call %d model = %q, want %q", i, got[i], wantModels[i])
|
||||
}
|
||||
}
|
||||
wantPayloads := []string{`{"model":"claude-opus-4.66"}`, `{"model":"claude-opus-4.66"}`}
|
||||
for i := range wantPayloads {
|
||||
if payloads[i] != wantPayloads[i] {
|
||||
t.Fatalf("payload %d = %s, want %s", i, payloads[i], wantPayloads[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerExecute_OpenAICompatAliasPoolStopsOnBadRequest(t *testing.T) {
|
||||
alias := "claude-opus-4.66"
|
||||
invalidErr := &Error{HTTPStatus: http.StatusBadRequest, Message: "invalid_request_error: malformed payload"}
|
||||
|
||||
Reference in New Issue
Block a user