fix(auth): rewrite API key alias response models (#4002)

This commit is contained in:
sususu98
2026-06-25 22:32:48 +08:00
committed by GitHub
parent c4cf0fd324
commit eb2e1e33d4
8 changed files with 435 additions and 12 deletions
+2
View File
@@ -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)
+24 -8
View File
@@ -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,
+6 -2
View File
@@ -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)
}
}
+81 -2
View File
@@ -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)
}
})
}
}
+49
View File
@@ -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"}