diff --git a/config.example.yaml b/config.example.yaml index c480b2f5..fae425e3 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -355,6 +355,9 @@ nonstream-keepalive-interval: 0 # client-visible names can become ambiguous across providers. For strict backend pinning, use # unique aliases/prefixes or avoid overlapping names. # You can repeat the same name with different aliases to expose multiple client model names. +# Optional per-entry flags: +# fork: true # keep the upstream model and also expose the alias as a separate client-visible model +# force-mapping: true # optional: rewrite upstream response model fields back to the client-visible alias (example below uses antigravity only) # Per-auth OAuth aliases can also be stored in an OAuth auth JSON file as "model-aliases". # They apply only to that selected auth and take precedence over global aliases for the same client-visible alias. # Example auth JSON: @@ -374,8 +377,10 @@ nonstream-keepalive-interval: 0 # - name: "gemini-2.5-pro" # alias: "g2.5p" # antigravity: -# - name: "gemini-3-pro-high" -# alias: "gemini-3-pro-preview" +# - name: "gemini-pro-agent" # upstream Antigravity model id +# alias: "gemini-3.1-pro-preview" # client-visible id (Gemini 3.1 Pro Preview) +# fork: true +# force-mapping: true # claude: # - name: "claude-sonnet-4-5-20250929" # alias: "cs4.5" @@ -436,6 +441,13 @@ nonstream-keepalive-interval: 0 # "generationConfig.responseJsonSchema": "{\"type\":\"object\",\"properties\":{\"answer\":{\"type\":\"string\"}}}" # override: # Override rules always set parameters, overwriting any existing values. # - models: +# - name: "gpt-5.4-fast" +# protocol: "codex" +# - name: "gpt-5.5-fast" +# protocol: "codex" +# params: +# service_tier: priority +# - models: # - name: "gpt-*" # Supports wildcards (e.g., "gpt-*") # protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex, antigravity # params: # JSON path (gjson/sjson syntax) -> value diff --git a/internal/config/config.go b/internal/config/config.go index ffb67e42..bf052833 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -351,6 +351,8 @@ type OAuthModelAlias struct { Name string `yaml:"name" json:"name"` Alias string `yaml:"alias" json:"alias"` Fork bool `yaml:"fork,omitempty" json:"fork,omitempty"` + + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` } // PayloadConfig defines default and override parameter rules applied to provider payloads. @@ -922,7 +924,7 @@ func (cfg *Config) SanitizeOAuthModelAlias() { continue } seenAlias[aliasKey] = struct{}{} - clean = append(clean, OAuthModelAlias{Name: name, Alias: alias, Fork: entry.Fork}) + clean = append(clean, OAuthModelAlias{Name: name, Alias: alias, Fork: entry.Fork, ForceMapping: entry.ForceMapping}) } if len(clean) > 0 { out[channel] = clean diff --git a/internal/watcher/diff/oauth_model_alias.go b/internal/watcher/diff/oauth_model_alias.go index 8c14089b..d95bfd39 100644 --- a/internal/watcher/diff/oauth_model_alias.go +++ b/internal/watcher/diff/oauth_model_alias.go @@ -83,6 +83,9 @@ func summarizeOAuthModelAliasList(list []config.OAuthModelAlias) OAuthModelAlias if alias.Fork { key += "|fork" } + if alias.ForceMapping { + key += "|force-mapping" + } if _, exists := seen[key]; exists { continue } diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index b47ed23d..71090f29 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1202,11 +1202,78 @@ func (m *Manager) preparedExecutionModels(auth *Auth, routeModel string) ([]stri return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled } +func (m *Manager) preparedExecutionModelsWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { + candidates, pooled, aliasResult := m.executionModelCandidatesWithAlias(auth, routeModel) + return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled, aliasResult +} + +func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { + requestedModel := rewriteModelForAuth(routeModel, auth) + aliasResult := m.applyOAuthModelAliasWithResult(auth, requestedModel) + upstreamModel := aliasResult.UpstreamModel + + var candidates []string + if auth != nil && auth.Attributes != nil { + if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { + candidates = []string{homeModel} + } + } + if len(candidates) == 0 { + if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) > 0 { + if len(pool) == 1 { + candidates = pool + } else { + offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, upstreamModel), len(pool)) + candidates = rotateStrings(pool, offset) + } + } else { + resolved := m.applyAPIKeyModelAlias(auth, upstreamModel) + if strings.TrimSpace(resolved) == "" { + resolved = upstreamModel + } + candidates = []string{resolved} + } + } + pooled := len(candidates) > 1 + return candidates, pooled, aliasResult +} + func (m *Manager) prepareExecutionModels(auth *Auth, routeModel string) []string { models, _ := m.preparedExecutionModels(auth, routeModel) return models } +func rewriteForceMappedResponse(resp *cliproxyexecutor.Response, aliasResult OAuthModelAliasResult) { + if resp == nil || !aliasResult.ForceMapping || strings.TrimSpace(aliasResult.OriginalAlias) == "" { + return + } + resp.Payload = rewriteModelInResponse(resp.Payload, aliasResult.OriginalAlias) +} + +func rewriteForceMappedStreamChunk(rewriter *StreamRewriter, payload []byte) []byte { + if rewriter == nil || len(payload) == 0 { + return payload + } + rewritten := rewriter.RewriteChunk(payload) + if len(rewritten) > 0 { + return rewritten + } + if len(rewriter.pendingBuf) > 0 { + return nil + } + if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 { + return lineWise + } + return nil +} + +func finishForceMappedStreamChunks(rewriter *StreamRewriter) []byte { + if rewriter == nil { + return nil + } + return rewriter.Finish() +} + func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) { if len(auths) == 0 { return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} @@ -1591,12 +1658,16 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC } } -func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk) *cliproxyexecutor.StreamResult { +func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) var failed bool forward := true + var rewriter *StreamRewriter + if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" { + rewriter = NewStreamRewriter(StreamRewriteOptions{RewriteModel: aliasResult.OriginalAlias}) + } emit := func(chunk cliproxyexecutor.StreamChunk) bool { if chunk.Err != nil && !failed { failed = true @@ -1609,6 +1680,27 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re if !forward { return false } + if chunk.Err != nil { + if ctx == nil { + out <- chunk + return true + } + select { + case <-ctx.Done(): + forward = false + return false + case out <- chunk: + return true + } + } + if len(chunk.Payload) == 0 { + return true + } + payload := rewriteForceMappedStreamChunk(rewriter, chunk.Payload) + if len(payload) == 0 { + return true + } + chunk.Payload = payload if ctx == nil { out <- chunk return true @@ -1633,6 +1725,12 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return } } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + tailChunk := cliproxyexecutor.StreamChunk{Payload: tail} + if !emit(tailChunk) { + return + } + } if !failed { m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}) } @@ -1640,7 +1738,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} } -func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel string, execModels []string, pooled bool) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult) (*cliproxyexecutor.StreamResult, error) { if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } @@ -1728,7 +1826,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi close(closedCh) remaining = closedCh } - return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining), nil + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult), nil } if lastErr == nil { lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} @@ -2331,7 +2429,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) - models, pooled := m.preparedExecutionModels(auth, routeModel) + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) if len(models) == 0 { continue } @@ -2375,6 +2473,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req continue } m.MarkResult(execCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) return resp, nil } if authErr != nil { @@ -2432,7 +2531,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) - models, pooled := m.preparedExecutionModels(auth, routeModel) + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) if len(models) == 0 { continue } @@ -2476,6 +2575,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, continue } m.MarkResult(execCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) return resp, nil } if authErr != nil { @@ -2531,7 +2631,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } - models, pooled := m.preparedExecutionModels(auth, routeModel) + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) if len(models) == 0 { continue } @@ -2548,7 +2648,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string continue } execReq := sanitizeDownstreamWebsocketFallbackRequest(execCtx, auth, req) - streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, opts, routeModel, models, pooled) + streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, opts, routeModel, models, pooled, aliasResult) if errStream != nil { if errCtx := execCtx.Err(); errCtx != nil { return nil, errCtx @@ -4955,12 +5055,12 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy } c.auth = preparedAuth publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID) - models := m.executionModelCandidates(c.auth, routeModel) + models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) if len(models) == 0 { continue } for _, upstreamModel := range models { - resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, len(models) > 1) + resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts) @@ -4977,6 +5077,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy continue } m.MarkResult(creditsCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) return resp, true, nil } } @@ -5005,11 +5106,11 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl } c.auth = preparedAuth publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID) - models := m.executionModelCandidates(c.auth, routeModel) + models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) if len(models) == 0 { continue } - result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, models, len(models) > 1) + result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, models, pooled, aliasResult) if errStream != nil { continue } diff --git a/sdk/cliproxy/auth/conductor_force_mapping_test.go b/sdk/cliproxy/auth/conductor_force_mapping_test.go new file mode 100644 index 00000000..72738210 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_force_mapping_test.go @@ -0,0 +1,555 @@ +package auth + +import ( + "context" + "net/http" + "strings" + "sync" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type forceMappingExecutor struct { + id string + + mu sync.Mutex + executeModels []string + streamModels []string +} + +func (e *forceMappingExecutor) Identifier() string { return e.id } + +func (e *forceMappingExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.executeModels = append(e.executeModels, req.Model) + e.mu.Unlock() + payload := forceMappingNonStreamUpstreamPayload(e.id, req.Model) + return cliproxyexecutor.Response{Payload: []byte(payload)}, nil +} + +func (e *forceMappingExecutor) ExecuteStream(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.streamModels = append(e.streamModels, req.Model) + e.mu.Unlock() + chunks := forceMappingStreamUpstreamChunks(e.id, req.Model) + ch := make(chan cliproxyexecutor.StreamChunk, len(chunks)) + for _, chunk := range chunks { + ch <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *forceMappingExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *forceMappingExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "CountTokens not implemented"} +} + +func (e *forceMappingExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "HttpRequest not implemented"} +} + +func (e *forceMappingExecutor) ExecuteModels() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.executeModels)) + copy(out, e.executeModels) + return out +} + +func (e *forceMappingExecutor) StreamModels() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.streamModels)) + copy(out, e.streamModels) + return out +} + +type forceMappingCreditsFallbackExecutor struct { + id string + + mu sync.Mutex + executeModels []string + executeCreditsRequested []bool + streamModels []string + streamCreditsRequested []bool +} + +func (e *forceMappingCreditsFallbackExecutor) Identifier() string { return e.id } + +func (e *forceMappingCreditsFallbackExecutor) Execute(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + creditsRequested := AntigravityCreditsRequested(ctx) + e.mu.Lock() + e.executeModels = append(e.executeModels, req.Model) + e.executeCreditsRequested = append(e.executeCreditsRequested, creditsRequested) + e.mu.Unlock() + if !creditsRequested { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "MODEL_CAPACITY_EXHAUSTED"} + } + payload := `{"model":"` + req.Model + `","message":{"model":"` + req.Model + `"}}` + return cliproxyexecutor.Response{Payload: []byte(payload)}, nil +} + +func (e *forceMappingCreditsFallbackExecutor) ExecuteStream(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + creditsRequested := AntigravityCreditsRequested(ctx) + e.mu.Lock() + e.streamModels = append(e.streamModels, req.Model) + e.streamCreditsRequested = append(e.streamCreditsRequested, creditsRequested) + e.mu.Unlock() + ch := make(chan cliproxyexecutor.StreamChunk, 1) + if !creditsRequested { + ch <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "MODEL_CAPACITY_EXHAUSTED"}} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil + } + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"message":{"model":"` + req.Model + `"}}` + "\n\n")} + close(ch) + return &cliproxyexecutor.StreamResult{Chunks: ch}, nil +} + +func (e *forceMappingCreditsFallbackExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *forceMappingCreditsFallbackExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "CountTokens not implemented"} +} + +func (e *forceMappingCreditsFallbackExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "HttpRequest not implemented"} +} + +func (e *forceMappingCreditsFallbackExecutor) ExecuteModels() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.executeModels)) + copy(out, e.executeModels) + return out +} + +func (e *forceMappingCreditsFallbackExecutor) ExecuteCreditsRequested() []bool { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]bool, len(e.executeCreditsRequested)) + copy(out, e.executeCreditsRequested) + return out +} + +func (e *forceMappingCreditsFallbackExecutor) StreamModels() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.streamModels)) + copy(out, e.streamModels) + return out +} + +func (e *forceMappingCreditsFallbackExecutor) StreamCreditsRequested() []bool { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]bool, len(e.streamCreditsRequested)) + copy(out, e.streamCreditsRequested) + return out +} + +func forceMappingPayloadLeaksUpstream(payload, upstreamModel string) bool { + if upstreamModel == "" { + return false + } + return strings.Contains(payload, `"model":"`+upstreamModel+`"`) || + strings.Contains(payload, `"model": "`+upstreamModel+`"`) || + strings.Contains(payload, `"modelVersion":"`+upstreamModel+`"`) +} + +func forceMappingNonStreamUpstreamPayload(provider, upstreamModel string) string { + switch provider { + case "codex": + return strings.Replace(liveCodexResponsesNonStreamUpstream, "gpt-5.4", upstreamModel, 1) + case "kimi": + return strings.Replace(liveKimiMessagesNonStreamUpstream, "kimi-k2.5", upstreamModel, 1) + case "xai": + return `{"type":"message","role":"assistant","model":"` + upstreamModel + `","content":[{"type":"text","text":"hi"}]}` + case "antigravity": + return strings.Replace(liveAntigravityMessagesStartUpstream, "gemini-3-flash", upstreamModel, 1) + default: + return `{"model":"` + upstreamModel + `","message":{"model":"` + upstreamModel + `"}}` + } +} + +func forceMappingStreamUpstreamChunks(provider, upstreamModel string) [][]byte { + switch provider { + case "codex": + created := strings.Replace(liveCodexResponsesCreatedUpstream, "gpt-5.4", upstreamModel, -1) + completed := strings.Replace(liveCodexResponsesCompletedUpstream, "gpt-5.4", upstreamModel, -1) + return [][]byte{ + []byte("event: response.created\n"), + []byte("data: " + created + "\n"), + []byte("\n"), + []byte("event: response.completed\n"), + []byte("data: " + completed + "\n"), + []byte("\n"), + } + case "kimi": + msg := strings.Replace(liveKimiMessagesStartUpstream, "kimi-k2.5", upstreamModel, 1) + chat := strings.Replace(liveKimiChatChunkUpstream, "kimi-k2.5", upstreamModel, 1) + return [][]byte{ + []byte("event:message_start\n"), + []byte("data:" + msg + "\n\n"), + []byte("data: " + chat + "\n\n"), + } + case "xai": + msg := strings.Replace(liveXAIMessagesStartUpstream, "grok-4.3", upstreamModel, 1) + return [][]byte{ + []byte("event: message_start\n"), + []byte("data: " + msg + "\n\n"), + } + case "antigravity": + msg := strings.Replace(liveAntigravityMessagesStartUpstream, "gemini-3-flash", upstreamModel, 1) + return [][]byte{ + []byte("event: message_start\n"), + []byte("data: " + msg + "\n\n"), + } + default: + return [][]byte{ + []byte(`data: {"type":"response.created","response":{"model":"` + upstreamModel + `"}}` + "\n\n"), + } + } +} + +func setupForceMappingManager(t *testing.T, provider, upstreamModel, aliasModel string) (*Manager, *forceMappingExecutor) { + t.Helper() + manager := NewManager(nil, nil, nil) + executor := &forceMappingExecutor{id: provider} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: upstreamModel, + Alias: aliasModel, + Fork: true, + ForceMapping: true, + }}, + }) + + auth := &Auth{ + ID: provider + "-force-mapping-auth", + Provider: provider, + Status: StatusActive, + } + 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 setupForceMappingCreditsFallbackManager(t *testing.T, upstreamModel, aliasModel string) (*Manager, *forceMappingCreditsFallbackExecutor) { + t.Helper() + const provider = "antigravity" + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true}, + }) + manager.SetRetryConfig(0, 0, 1) + executor := &forceMappingCreditsFallbackExecutor{id: provider} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: upstreamModel, + Alias: aliasModel, + Fork: true, + ForceMapping: true, + }}, + }) + + auth := &Auth{ + ID: provider + "-force-mapping-credits-auth", + Provider: provider, + Status: StatusActive, + } + 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_OAuthAliasForceMappingRewritesNonStreamResponse(t *testing.T) { + const ( + provider = "antigravity" + upstreamModel = "gemini-3-flash-preview" + aliasModel = "claude-haiku-4-5-20251001" + ) + + manager, executor := setupForceMappingManager(t, provider, upstreamModel, aliasModel) + resp, errExecute := manager.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute error = %v, want success", errExecute) + } + + gotModels := executor.ExecuteModels() + if len(gotModels) != 1 || gotModels[0] != upstreamModel { + t.Fatalf("execute models = %v, want [%s]", gotModels, upstreamModel) + } + if got := string(resp.Payload); !strings.Contains(got, aliasModel) || forceMappingPayloadLeaksUpstream(got, upstreamModel) { + t.Fatalf("response payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel) + } +} + +func TestManagerExecuteStream_OAuthAliasForceMappingRewritesStreamResponse(t *testing.T) { + const ( + provider = "antigravity" + upstreamModel = "gemini-3-flash-preview" + aliasModel = "claude-haiku-4-5-20251001" + ) + + manager, executor := setupForceMappingManager(t, provider, upstreamModel, aliasModel) + streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute stream error = %v, want success", errExecute) + } + + gotModels := executor.StreamModels() + if len(gotModels) != 1 || gotModels[0] != upstreamModel { + t.Fatalf("stream models = %v, want [%s]", gotModels, 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, aliasModel) || forceMappingPayloadLeaksUpstream(got, upstreamModel) { + t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel) + } +} + +func TestManagerExecuteStream_OAuthAliasForceMappingRewritesCodexStyleLineChunks(t *testing.T) { + const ( + provider = "codex" + upstreamModel = "gpt-5.4" + aliasModel = "gpt-5.4-fast" + ) + + manager, _ := setupForceMappingManager(t, provider, upstreamModel, aliasModel) + + streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute stream error = %v, want success", errExecute) + } + + 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...) + } + got := string(payload) + if !strings.Contains(got, aliasModel) || forceMappingPayloadLeaksUpstream(got, upstreamModel) { + t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel) + } +} + +func TestManagerExecute_OAuthAliasForceMappingRewritesKimiAndXAIResponses(t *testing.T) { + cases := []struct { + provider string + upstreamModel string + aliasModel string + }{ + {provider: "kimi", upstreamModel: "kimi-k2.5", aliasModel: "k2.5"}, + {provider: "xai", upstreamModel: "grok-4.3", aliasModel: "grok-latest"}, + } + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + manager, executor := setupForceMappingManager(t, tc.provider, tc.upstreamModel, tc.aliasModel) + resp, errExecute := manager.Execute(context.Background(), []string{tc.provider}, cliproxyexecutor.Request{Model: tc.aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute error = %v, want success", errExecute) + } + gotModels := executor.ExecuteModels() + if len(gotModels) != 1 || gotModels[0] != tc.upstreamModel { + t.Fatalf("execute models = %v, want [%s]", gotModels, tc.upstreamModel) + } + if got := string(resp.Payload); !strings.Contains(got, tc.aliasModel) || forceMappingPayloadLeaksUpstream(got, tc.upstreamModel) { + t.Fatalf("response payload = %s, want alias %q without upstream %q", got, tc.aliasModel, tc.upstreamModel) + } + }) + } +} + +func TestManagerExecuteStream_OAuthAliasForceMappingRewritesKimiAndXAIResponses(t *testing.T) { + cases := []struct { + provider string + upstreamModel string + aliasModel string + }{ + {provider: "kimi", upstreamModel: "kimi-k2.5", aliasModel: "k2.5"}, + {provider: "xai", upstreamModel: "grok-4.3", aliasModel: "grok-latest"}, + } + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + manager, executor := setupForceMappingManager(t, tc.provider, tc.upstreamModel, tc.aliasModel) + streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{tc.provider}, cliproxyexecutor.Request{Model: tc.aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute stream error = %v, want success", errExecute) + } + gotModels := executor.StreamModels() + if len(gotModels) != 1 || gotModels[0] != tc.upstreamModel { + t.Fatalf("stream models = %v, want [%s]", gotModels, tc.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...) + } + got := string(payload) + if !strings.Contains(got, tc.aliasModel) || forceMappingPayloadLeaksUpstream(got, tc.upstreamModel) { + t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, tc.aliasModel, tc.upstreamModel) + } + }) + } +} + +func TestManagerExecute_LiveDerivedForceMapping_AllProviders(t *testing.T) { + cases := []struct { + provider string + upstreamModel string + aliasModel string + }{ + {provider: "codex", upstreamModel: "gpt-5.4", aliasModel: "gpt-5.4-fast"}, + {provider: "antigravity", upstreamModel: "gemini-3-flash", aliasModel: "claude-haiku-4-5-20251001"}, + {provider: "kimi", upstreamModel: "kimi-k2.5", aliasModel: "k2.5"}, + {provider: "xai", upstreamModel: "grok-4.3", aliasModel: "grok-latest"}, + } + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + manager, executor := setupForceMappingManager(t, tc.provider, tc.upstreamModel, tc.aliasModel) + resp, errExecute := manager.Execute(context.Background(), []string{tc.provider}, cliproxyexecutor.Request{Model: tc.aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute error = %v", errExecute) + } + if got := executor.ExecuteModels(); len(got) != 1 || got[0] != tc.upstreamModel { + t.Fatalf("execute models = %v, want [%s]", got, tc.upstreamModel) + } + gotPayload := string(resp.Payload) + if !strings.Contains(gotPayload, tc.aliasModel) || forceMappingPayloadLeaksUpstream(gotPayload, tc.upstreamModel) { + t.Fatalf("payload = %s, want alias %q without upstream %q", gotPayload, tc.aliasModel, tc.upstreamModel) + } + }) + } +} + +func TestManagerExecuteStream_LiveDerivedForceMapping_AllProviders(t *testing.T) { + cases := []struct { + provider string + upstreamModel string + aliasModel string + }{ + {provider: "codex", upstreamModel: "gpt-5.4", aliasModel: "gpt-5.4-fast"}, + {provider: "antigravity", upstreamModel: "gemini-3-flash", aliasModel: "claude-haiku-4-5-20251001"}, + {provider: "kimi", upstreamModel: "kimi-k2.5", aliasModel: "k2.5"}, + {provider: "xai", upstreamModel: "grok-4.3", aliasModel: "grok-latest"}, + } + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + manager, executor := setupForceMappingManager(t, tc.provider, tc.upstreamModel, tc.aliasModel) + streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{tc.provider}, cliproxyexecutor.Request{Model: tc.aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute stream error = %v", errExecute) + } + if got := executor.StreamModels(); len(got) != 1 || got[0] != tc.upstreamModel { + t.Fatalf("stream models = %v, want [%s]", got, tc.upstreamModel) + } + var payload []byte + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + t.Fatalf("stream error: %v", chunk.Err) + } + payload = append(payload, chunk.Payload...) + } + got := string(payload) + if !strings.Contains(got, tc.aliasModel) { + t.Fatalf("stream payload missing alias %q: %s", tc.aliasModel, got) + } + if forceMappingPayloadLeaksUpstream(got, tc.upstreamModel) { + t.Fatalf("stream payload leaked upstream %q: %s", tc.upstreamModel, got) + } + }) + } +} + +func TestManagerExecute_AntigravityCreditsFallbackForceMappingRewritesResponse(t *testing.T) { + const ( + upstreamModel = "gemini-3-flash-preview" + aliasModel = "claude-haiku-4-5-20251001" + ) + + manager, executor := setupForceMappingCreditsFallbackManager(t, upstreamModel, aliasModel) + resp, errExecute := manager.Execute(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute error = %v, want success", errExecute) + } + + if got := executor.ExecuteModels(); len(got) != 2 || got[0] != upstreamModel || got[1] != upstreamModel { + t.Fatalf("execute models = %v, want [%s %s]", got, upstreamModel, upstreamModel) + } + if got := executor.ExecuteCreditsRequested(); len(got) != 2 || got[0] || !got[1] { + t.Fatalf("credits flags = %v, want [false true]", got) + } + if got := string(resp.Payload); !strings.Contains(got, aliasModel) || forceMappingPayloadLeaksUpstream(got, upstreamModel) { + t.Fatalf("response payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel) + } +} + +func TestManagerExecuteStream_AntigravityCreditsFallbackForceMappingRewritesResponse(t *testing.T) { + const ( + upstreamModel = "gemini-3-flash-preview" + aliasModel = "claude-haiku-4-5-20251001" + ) + + manager, executor := setupForceMappingCreditsFallbackManager(t, upstreamModel, aliasModel) + streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: aliasModel}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("execute stream error = %v, want success", errExecute) + } + + if got := executor.StreamModels(); len(got) != 2 || got[0] != upstreamModel || got[1] != upstreamModel { + t.Fatalf("stream models = %v, want [%s %s]", got, upstreamModel, upstreamModel) + } + if got := executor.StreamCreditsRequested(); len(got) != 2 || got[0] || !got[1] { + t.Fatalf("credits flags = %v, want [false true]", got) + } + 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, aliasModel) || forceMappingPayloadLeaksUpstream(got, upstreamModel) { + t.Fatalf("stream payload = %s, want alias %q without upstream %q", got, aliasModel, upstreamModel) + } +} diff --git a/sdk/cliproxy/auth/force_mapping_live_fixtures_test.go b/sdk/cliproxy/auth/force_mapping_live_fixtures_test.go new file mode 100644 index 00000000..66603b37 --- /dev/null +++ b/sdk/cliproxy/auth/force_mapping_live_fixtures_test.go @@ -0,0 +1,20 @@ +package auth + +// Live CPA-derived upstream response fixtures (2026-06-24, local 8343). +// Executors emit these upstream model names; tests assert client-visible aliases after force-mapping. + +const liveCodexResponsesCreatedUpstream = `{"type":"response.created","response":{"id":"resp_live","object":"response","created_at":1782272843,"status":"in_progress","model":"gpt-5.4","output":[],"parallel_tool_calls":true}}` + +const liveCodexResponsesCompletedUpstream = `{"type":"response.completed","response":{"id":"resp_live","object":"response","created_at":1782272843,"status":"completed","model":"gpt-5.4","output":[{"type":"message","content":[{"type":"output_text","text":"Hi!"}]}]}}` + +const liveAntigravityMessagesStartUpstream = `{"type": "message_start", "message": {"id": "UVM7aqirB-npz7IP8rfZuQQ", "type": "message", "role": "assistant", "content": [], "model": "gemini-3-flash", "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 2, "output_tokens": 1}}}` + +const liveKimiChatChunkUpstream = `{"id":"chatcmpl-McAG6QS2WmxRKmMxSjWvbgWB","object":"chat.completion.chunk","created":1782272842,"model":"kimi-k2.5","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}],"system_fingerprint":"fpv0_b30801d4"}` + +const liveKimiMessagesStartUpstream = `{"type":"message_start","message":{"id":"msg_iFEkPDty2KtvlbdThqOBsN25","type":"message","role":"assistant","content":[],"model":"kimi-k2.5","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1263,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":0,"service_tier":"standard","inference_geo":"not_available","prompt_tokens":1263,"cached_tokens":0}}}` + +const liveXAIMessagesStartUpstream = `{"type":"message_start","message":{"id":"4aeb964a-1190-98f6-9978-a8a7548848d8","type":"message","role":"assistant","model":"grok-4.3","stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0},"content":[],"stop_reason":null}}` + +const liveCodexResponsesNonStreamUpstream = `{"model":"gpt-5.4","output":[{"type":"message","content":[{"type":"output_text","text":"Hi!"}]}]}` + +const liveKimiMessagesNonStreamUpstream = `{"type":"message","role":"assistant","model":"kimi-k2.5","content":[{"type":"text","text":"hi"}]}` diff --git a/sdk/cliproxy/auth/oauth_model_alias.go b/sdk/cliproxy/auth/oauth_model_alias.go index cd0d6179..37b57f05 100644 --- a/sdk/cliproxy/auth/oauth_model_alias.go +++ b/sdk/cliproxy/auth/oauth_model_alias.go @@ -15,9 +15,22 @@ type modelAliasEntry interface { GetAlias() string } +// oauthModelAliasEntry stores the upstream model name and mapping flags for an alias. +type oauthModelAliasEntry struct { + upstreamModel string + forceMapping bool +} + type oauthModelAliasTable struct { - // reverse maps channel -> alias (lower) -> original upstream model name. - reverse map[string]map[string]string + // reverse maps channel -> alias (lower) -> entry with upstream model and flags. + reverse map[string]map[string]oauthModelAliasEntry +} + +// OAuthModelAliasResult contains the resolved upstream model and mapping metadata. +type OAuthModelAliasResult struct { + UpstreamModel string // resolved upstream model name (empty if no mapping found) + ForceMapping bool // whether to rewrite model name in responses + OriginalAlias string // the original requested alias (for response rewriting) } func compileOAuthModelAliasTable(aliases map[string][]internalconfig.OAuthModelAlias) *oauthModelAliasTable { @@ -25,14 +38,14 @@ func compileOAuthModelAliasTable(aliases map[string][]internalconfig.OAuthModelA return &oauthModelAliasTable{} } out := &oauthModelAliasTable{ - reverse: make(map[string]map[string]string, len(aliases)), + reverse: make(map[string]map[string]oauthModelAliasEntry, len(aliases)), } for rawChannel, entries := range aliases { channel := strings.ToLower(strings.TrimSpace(rawChannel)) if channel == "" || len(entries) == 0 { continue } - rev := make(map[string]string, len(entries)) + rev := make(map[string]oauthModelAliasEntry, len(entries)) for _, entry := range entries { name := strings.TrimSpace(entry.Name) alias := strings.TrimSpace(entry.Alias) @@ -46,7 +59,10 @@ func compileOAuthModelAliasTable(aliases map[string][]internalconfig.OAuthModelA if _, exists := rev[aliasKey]; exists { continue } - rev[aliasKey] = name + rev[aliasKey] = oauthModelAliasEntry{ + upstreamModel: name, + forceMapping: entry.ForceMapping, + } } if len(rev) > 0 { out.reverse[channel] = rev @@ -186,12 +202,17 @@ func resolveModelAliasFromConfigModels(requestedModel string, models []modelAlia // the suffix is preserved in the returned model name. However, if the alias's // original name already contains a suffix, the config suffix takes priority. func (m *Manager) resolveOAuthUpstreamModel(auth *Auth, requestedModel string) string { + result := m.resolveOAuthModelAliasWithResult(auth, requestedModel) + return result.UpstreamModel +} + +func (m *Manager) resolveOAuthModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult { channel := modelAliasChannel(auth) if channel == "" { - return "" + return OAuthModelAliasResult{} } - if resolved := resolveUpstreamModelFromAliases(OAuthModelAliasesFromAttributes(authAttributes(auth)), requestedModel); resolved != "" { - return resolved + if result := resolveUpstreamModelFromAliases(OAuthModelAliasesFromAttributes(authAttributes(auth)), requestedModel); result.UpstreamModel != "" { + return result } return resolveUpstreamModelFromAliasTable(m, auth, requestedModel, channel) } @@ -255,13 +276,13 @@ func sanitizeOAuthModelAliases(aliases []internalconfig.OAuthModelAlias) []inter return append([]internalconfig.OAuthModelAlias(nil), clean...) } -func resolveUpstreamModelFromAliases(aliases []internalconfig.OAuthModelAlias, requestedModel string) string { +func resolveUpstreamModelFromAliases(aliases []internalconfig.OAuthModelAlias, requestedModel string) OAuthModelAliasResult { if len(aliases) == 0 { - return "" + return OAuthModelAliasResult{} } requestResult, candidates := modelAliasLookupCandidates(requestedModel) if len(candidates) == 0 { - return "" + return OAuthModelAliasResult{} } baseModel := requestResult.ModelName if baseModel == "" { @@ -279,27 +300,44 @@ func resolveUpstreamModelFromAliases(aliases []internalconfig.OAuthModelAlias, r continue } if strings.EqualFold(original, baseModel) { - return "" + if !entry.ForceMapping { + return OAuthModelAliasResult{} + } + return OAuthModelAliasResult{ + UpstreamModel: preserveResolvedModelSuffix(original, requestResult), + ForceMapping: entry.ForceMapping, + OriginalAlias: requestedModel, + } + } + return OAuthModelAliasResult{ + UpstreamModel: preserveResolvedModelSuffix(original, requestResult), + ForceMapping: entry.ForceMapping, + OriginalAlias: requestedModel, } - return preserveResolvedModelSuffix(original, requestResult) } } - return "" + return OAuthModelAliasResult{} } -func resolveUpstreamModelFromAliasTable(m *Manager, auth *Auth, requestedModel, channel string) string { +func (m *Manager) applyOAuthModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult { + result := m.resolveOAuthModelAliasWithResult(auth, requestedModel) + if result.UpstreamModel == "" { + return OAuthModelAliasResult{UpstreamModel: requestedModel} + } + return result +} + +func resolveUpstreamModelFromAliasTable(m *Manager, auth *Auth, requestedModel, channel string) OAuthModelAliasResult { if m == nil || auth == nil { - return "" + return OAuthModelAliasResult{} } if channel == "" { - return "" + return OAuthModelAliasResult{} } - // Extract thinking suffix from requested model using ParseSuffix requestResult := thinking.ParseSuffix(requestedModel) baseModel := requestResult.ModelName - // Candidate keys to match: base model and raw input (handles suffix-parsing edge cases). candidates := []string{baseModel} if baseModel != requestedModel { candidates = append(candidates, requestedModel) @@ -308,11 +346,11 @@ func resolveUpstreamModelFromAliasTable(m *Manager, auth *Auth, requestedModel, raw := m.oauthModelAlias.Load() table, _ := raw.(*oauthModelAliasTable) if table == nil || table.reverse == nil { - return "" + return OAuthModelAliasResult{} } rev := table.reverse[channel] if rev == nil { - return "" + return OAuthModelAliasResult{} } for _, candidate := range candidates { @@ -320,26 +358,44 @@ func resolveUpstreamModelFromAliasTable(m *Manager, auth *Auth, requestedModel, if key == "" { continue } - original := strings.TrimSpace(rev[key]) - if original == "" { + entry, exists := rev[key] + if !exists { continue } - if strings.EqualFold(original, baseModel) { - return "" + + targetModel := entry.upstreamModel + if targetModel == "" { + continue } - // If config already has suffix, it takes priority. - if thinking.ParseSuffix(original).HasSuffix { - return original + if strings.EqualFold(targetModel, baseModel) { + if !entry.forceMapping { + return OAuthModelAliasResult{} + } + return OAuthModelAliasResult{ + UpstreamModel: preserveResolvedModelSuffix(targetModel, requestResult), + ForceMapping: entry.forceMapping, + OriginalAlias: requestedModel, + } } - // Preserve user's thinking suffix on the resolved model. - if requestResult.HasSuffix && requestResult.RawSuffix != "" { - return original + "(" + requestResult.RawSuffix + ")" + + var upstreamModel string + if thinking.ParseSuffix(targetModel).HasSuffix { + upstreamModel = targetModel + } else if requestResult.HasSuffix && requestResult.RawSuffix != "" { + upstreamModel = targetModel + "(" + requestResult.RawSuffix + ")" + } else { + upstreamModel = targetModel + } + + return OAuthModelAliasResult{ + UpstreamModel: upstreamModel, + ForceMapping: entry.forceMapping, + OriginalAlias: requestedModel, } - return original } - return "" + return OAuthModelAliasResult{} } // modelAliasChannel extracts the OAuth model alias channel from an Auth object. diff --git a/sdk/cliproxy/auth/oauth_model_alias_test.go b/sdk/cliproxy/auth/oauth_model_alias_test.go index 3504a622..f1cdd76a 100644 --- a/sdk/cliproxy/auth/oauth_model_alias_test.go +++ b/sdk/cliproxy/auth/oauth_model_alias_test.go @@ -208,6 +208,49 @@ func TestApplyOAuthModelAlias_SuffixPreservation(t *testing.T) { } } +func TestApplyOAuthModelAlias_ForceMappingSameBasePreservesSuffix(t *testing.T) { + t.Parallel() + + aliases := map[string][]internalconfig.OAuthModelAlias{ + "antigravity": {{ + Name: "gemini-2.5-pro", + Alias: "gemini-2.5-pro(8192)", + ForceMapping: true, + }}, + } + + mgr := NewManager(nil, nil, nil) + mgr.SetConfig(&internalconfig.Config{}) + mgr.SetOAuthModelAlias(aliases) + + auth := &Auth{ID: "test-auth-id", Provider: "antigravity"} + + resolvedModel := mgr.applyOAuthModelAlias(auth, "gemini-2.5-pro(8192)") + if resolvedModel != "gemini-2.5-pro(8192)" { + t.Errorf("applyOAuthModelAlias() model = %q, want %q", resolvedModel, "gemini-2.5-pro(8192)") + } +} + +func TestApplyOAuthModelAlias_PerAuthForceMappingSameBasePreservesSuffix(t *testing.T) { + t.Parallel() + + mgr := NewManager(nil, nil, nil) + mgr.SetConfig(&internalconfig.Config{}) + + auth := &Auth{ + ID: "test-auth-id", + Provider: "antigravity", + Attributes: map[string]string{ + "model_aliases": `[{"name":"gemini-2.5-pro","alias":"gemini-2.5-pro(8192)","force-mapping":true}]`, + }, + } + + resolvedModel := mgr.applyOAuthModelAlias(auth, "gemini-2.5-pro(8192)") + if resolvedModel != "gemini-2.5-pro(8192)" { + t.Errorf("applyOAuthModelAlias() model = %q, want %q", resolvedModel, "gemini-2.5-pro(8192)") + } +} + func TestApplyOAuthModelAlias_PerAuthOverridesGlobalAlias(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/auth/response_model_rewriter.go b/sdk/cliproxy/auth/response_model_rewriter.go new file mode 100644 index 00000000..6ebee5bb --- /dev/null +++ b/sdk/cliproxy/auth/response_model_rewriter.go @@ -0,0 +1,211 @@ +package auth + +import ( + "bytes" + + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +var modelFieldPaths = []string{"model", "modelVersion", "response.model", "response.modelVersion", "message.model"} + +const maxPendingBufSize = 1 << 20 // 1MB limit for pending buffer + +func rewriteSSEPayloadLines(payload []byte, targetModel string) []byte { + if targetModel == "" || len(payload) == 0 { + return payload + } + lines := bytes.Split(payload, []byte("\n")) + out := make([][]byte, 0, len(lines)) + for _, line := range lines { + prefix, jsonData, ok := extractSSEDataLine(line) + if ok && len(jsonData) > 0 && jsonData[0] == '{' && gjson.ValidBytes(jsonData) { + rewritten := rewriteModelInResponse(jsonData, targetModel) + line = append(append([]byte{}, prefix...), rewritten...) + } + out = append(out, line) + } + joined := bytes.Join(out, []byte("\n")) + if len(payload) > 0 && payload[len(payload)-1] == '\n' && (len(joined) == 0 || joined[len(joined)-1] != '\n') { + joined = append(joined, '\n') + } + return joined +} + +func rewriteModelInResponse(data []byte, targetModel string) []byte { + if targetModel == "" || len(data) == 0 { + return data + } + for _, path := range modelFieldPaths { + if gjson.GetBytes(data, path).Exists() { + data, _ = sjson.SetBytes(data, path, targetModel) + log.Debugf("response rewriter: rewrote model at path %s to %s", path, targetModel) + } + } + return data +} + +// StreamRewriteOptions configures the stream rewriter. +type StreamRewriteOptions struct { + RewriteModel string +} + +// StreamRewriter rewrites model names in streaming SSE responses. +type StreamRewriter struct { + options StreamRewriteOptions + pendingBuf []byte +} + +// NewStreamRewriter creates a new stream rewriter. +func NewStreamRewriter(options StreamRewriteOptions) *StreamRewriter { + return &StreamRewriter{ + options: options, + pendingBuf: nil, + } +} + +// RewriteChunk rewrites model names in a single SSE chunk. +func (r *StreamRewriter) RewriteChunk(chunk []byte) []byte { + if r.options.RewriteModel == "" { + return chunk + } + + if len(r.pendingBuf) > 0 { + chunk = append(r.pendingBuf, chunk...) + r.pendingBuf = nil + } + + if len(chunk) > maxPendingBufSize { + return chunk + } + + // Handle raw JSON chunks (Gemini/OpenAI format without SSE "data:" prefix) + trimmed := bytes.TrimSpace(chunk) + if len(trimmed) > 0 && trimmed[0] == '{' && gjson.ValidBytes(trimmed) { + rewritten := trimmed + if r.options.RewriteModel != "" { + rewritten = rewriteModelInResponse(rewritten, r.options.RewriteModel) + } + return rewritten + } + + lastDoubleNewline := bytes.LastIndex(chunk, []byte("\n\n")) + lastNewline := -1 + if len(chunk) > 0 && chunk[len(chunk)-1] == '\n' { + lastNewline = len(chunk) - 1 + } + + var processChunk []byte + if lastDoubleNewline >= 0 { + afterComplete := chunk[lastDoubleNewline+2:] + if len(afterComplete) > 0 && !bytes.Equal(afterComplete, []byte("\n")) { + processChunk = chunk[:lastDoubleNewline+2] + r.pendingBuf = make([]byte, len(afterComplete)) + copy(r.pendingBuf, afterComplete) + } else { + processChunk = chunk + } + } else if lastNewline >= 0 && gjson.ValidBytes(extractLastDataPayload(chunk)) { + processChunk = chunk + } else if len(bytes.TrimSpace(chunk)) == 0 { + return chunk + } else if len(chunk) > 0 { + r.pendingBuf = make([]byte, len(chunk)) + copy(r.pendingBuf, chunk) + return nil + } else { + return chunk + } + + lines := bytes.Split(processChunk, []byte("\n")) + var result [][]byte + var pendingEvent []byte + skipBlanks := false + + for _, line := range lines { + if len(line) == 0 && skipBlanks { + continue + } + if len(line) != 0 && skipBlanks { + skipBlanks = false + } + + if bytes.HasPrefix(line, []byte("event:")) { + pendingEvent = line + continue + } + + dataPrefix, jsonData, found := extractSSEDataLine(line) + if found && len(jsonData) > 0 && jsonData[0] == '{' { + if !gjson.ValidBytes(jsonData) { + if pendingEvent != nil { + r.pendingBuf = append(pendingEvent, '\n') + r.pendingBuf = append(r.pendingBuf, line...) + pendingEvent = nil + } else { + r.pendingBuf = append(r.pendingBuf, line...) + } + continue + } + + if pendingEvent != nil { + result = append(result, pendingEvent) + pendingEvent = nil + } + + rewritten := jsonData + if r.options.RewriteModel != "" { + rewritten = rewriteModelInResponse(jsonData, r.options.RewriteModel) + } + result = append(result, append(dataPrefix, rewritten...)) + continue + } + + if pendingEvent != nil { + result = append(result, pendingEvent) + pendingEvent = nil + } + result = append(result, line) + } + + if pendingEvent != nil { + result = append(result, pendingEvent) + } + + joined := bytes.Join(result, []byte("\n")) + if len(joined) == 0 && len(chunk) > 0 { + return rewriteSSEPayloadLines(chunk, r.options.RewriteModel) + } + return joined +} + +func extractLastDataPayload(chunk []byte) []byte { + lines := bytes.Split(chunk, []byte("\n")) + for i := len(lines) - 1; i >= 0; i-- { + if _, jsonData, found := extractSSEDataLine(lines[i]); found && len(jsonData) > 0 { + return jsonData + } + } + return nil +} + +func extractSSEDataLine(line []byte) (prefix []byte, jsonData []byte, ok bool) { + if jsonData, found := bytes.CutPrefix(line, []byte("data: ")); found { + return []byte("data: "), jsonData, true + } + if jsonData, found := bytes.CutPrefix(line, []byte("data:")); found { + return []byte("data:"), jsonData, true + } + return nil, nil, false +} + +// Finish flushes any buffered partial SSE data at the end of a stream. +func (r *StreamRewriter) Finish() []byte { + if len(r.pendingBuf) == 0 { + return nil + } + chunk := r.RewriteChunk(r.pendingBuf) + r.pendingBuf = nil + return chunk +} diff --git a/sdk/cliproxy/auth/response_model_rewriter_test.go b/sdk/cliproxy/auth/response_model_rewriter_test.go new file mode 100644 index 00000000..d5384782 --- /dev/null +++ b/sdk/cliproxy/auth/response_model_rewriter_test.go @@ -0,0 +1,183 @@ +package auth + +import ( + "strings" + "testing" +) + +func TestStreamRewriter_RewriteChunk_KimiMessagesDataPrefixWithoutSpace(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "k2.5"}) + chunk := []byte("event:message_start\n" + + `data:{"type":"message_start","message":{"model":"kimi-k2.5"}}` + "\n\n") + + got := string(rewriter.RewriteChunk(chunk)) + if !strings.Contains(got, `"model":"k2.5"`) { + t.Fatalf("rewritten chunk = %q, want alias model k2.5", got) + } + if strings.Contains(got, "kimi-k2.5") { + t.Fatalf("rewritten chunk still contains upstream model: %q", got) + } + if !strings.Contains(got, "data:{") { + t.Fatalf("rewritten chunk should preserve data: prefix without space: %q", got) + } +} + +func TestStreamRewriter_RewriteChunk_AnthropicMessagesDataPrefixWithSpace(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "grok-latest"}) + chunk := []byte(`data: {"type":"message_start","message":{"model":"grok-4.3"}}` + "\n\n") + + got := string(rewriter.RewriteChunk(chunk)) + if !strings.Contains(got, `"model":"grok-latest"`) { + t.Fatalf("rewritten chunk = %q, want alias model grok-latest", got) + } + if strings.Contains(got, "grok-4.3") { + t.Fatalf("rewritten chunk still contains upstream model: %q", got) + } + if !strings.Contains(got, "data: {") { + t.Fatalf("rewritten chunk should preserve spaced data: prefix: %q", got) + } +} + +func TestStreamRewriter_Finish_FlushesCodexResponsesEventChunk(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + part1 := []byte("event: response.created\n") + part2 := []byte(`data: {"type":"response.created","response":{"model":"gpt-5.4"}}` + "\n\n") + + got1 := rewriter.RewriteChunk(part1) + if got1 != nil { + t.Fatalf("first partial chunk should buffer, got %q", string(got1)) + } + got2 := string(rewriter.RewriteChunk(part2)) + gotTail := string(rewriter.Finish()) + combined := got2 + gotTail + if !strings.Contains(combined, "gpt-5.4-fast") { + t.Fatalf("combined output = %q, want rewritten alias", combined) + } + if strings.Contains(combined, `"model":"gpt-5.4"`) { + t.Fatalf("combined output still has upstream model: %q", combined) + } +} + +func TestStreamRewriter_RewriteChunk_CodexResponsesLineChunks(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + lines := [][]byte{ + []byte("event: response.created\n"), + []byte(`data: {"type":"response.created","response":{"model":"gpt-5.4"}}` + "\n"), + []byte("\n"), + []byte("event: response.completed\n"), + []byte(`data: {"type":"response.completed","response":{"model":"gpt-5.4"}}` + "\n"), + []byte("\n"), + } + var out []byte + for _, line := range lines { + if rewritten := rewriter.RewriteChunk(line); len(rewritten) > 0 { + out = append(out, rewritten...) + } + } + if tail := rewriter.Finish(); len(tail) > 0 { + out = append(out, tail...) + } + got := string(out) + if !strings.Contains(got, "gpt-5.4-fast") { + t.Fatalf("rewritten output = %q, want alias gpt-5.4-fast", got) + } + if strings.Contains(got, `"model":"gpt-5.4"`) { + t.Fatalf("rewritten output still contains upstream model: %q", got) + } +} + +func TestRewriteForceMappedStreamChunk_CodexLineChunksDoNotDuplicateBufferedEvent(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + chunks := [][]byte{ + []byte("event: response.created\n"), + []byte(`data: {"type":"response.created","response":{"model":"gpt-5.4"}}` + "\n\n"), + } + + var out []byte + for _, chunk := range chunks { + if rewritten := rewriteForceMappedStreamChunk(rewriter, chunk); len(rewritten) > 0 { + out = append(out, rewritten...) + } + } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + out = append(out, tail...) + } + + got := string(out) + if count := strings.Count(got, "event: response.created"); count != 1 { + t.Fatalf("event count = %d, want 1; output=%q", count, got) + } + if !strings.HasSuffix(got, "\n\n") { + t.Fatalf("rewritten output = %q, want complete SSE frame terminator", got) + } + if !strings.Contains(got, `"model":"gpt-5.4-fast"`) { + t.Fatalf("rewritten output = %q, want alias model", got) + } + if strings.Contains(got, `"model":"gpt-5.4"`) { + t.Fatalf("rewritten output still contains upstream model: %q", got) + } +} + +func TestRewriteModelInResponse_AntigravityModelVersion(t *testing.T) { + payload := []byte(`{"response":{"modelVersion":"gemini-3-flash","candidates":[{"content":{"role":"model","parts":[{"text":"AGYMSG"}]}}]}}`) + got := string(rewriteModelInResponse(payload, "claude-haiku-4-5-20251001")) + if !strings.Contains(got, `"modelVersion":"claude-haiku-4-5-20251001"`) { + t.Fatalf("rewritten payload = %q, want alias modelVersion", got) + } + if strings.Contains(got, "gemini-3-flash") { + t.Fatalf("rewritten payload still contains upstream modelVersion: %q", got) + } +} + +func TestStreamRewriter_RewriteChunk_LiveDerivedProviderChunks(t *testing.T) { + cases := []struct { + name string + rewriteModel string + upstream string + chunk string + }{ + { + name: "kimi_chat_stream", + rewriteModel: "k2.5", + upstream: "kimi-k2.5", + chunk: `data:{"id":"chatcmpl-live","object":"chat.completion.chunk","created":1782272323,"model":"kimi-k2.5","choices":[{"index":0,"delta":{"content":"KCHATS"},"finish_reason":null}]}` + "\n\n", + }, + { + name: "kimi_messages_stream", + rewriteModel: "k2.5", + upstream: "kimi-k2.5", + chunk: "event:message_start\n" + `data:{"type":"message_start","message":{"model":"kimi-k2.5"}}` + "\n\n", + }, + { + name: "xai_messages_stream", + rewriteModel: "grok-latest", + upstream: "grok-4.3", + chunk: `data: {"type":"message_start","message":{"model":"grok-4.3"}}` + "\n\n", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: tc.rewriteModel}) + got := string(rewriter.RewriteChunk([]byte(tc.chunk))) + if !strings.Contains(got, tc.rewriteModel) { + t.Fatalf("rewritten chunk = %q, want alias %q", got, tc.rewriteModel) + } + if strings.Contains(got, tc.upstream) { + t.Fatalf("rewritten chunk still contains upstream %q: %q", tc.upstream, got) + } + }) + } +} +func TestRewriteSSEPayloadLines_CodexResponsesLiveFrame(t *testing.T) { + chunk := []byte("event: response.created\n" + + `data: {"type":"response.created","response":{"model":"gpt-5.4"}}` + "\n\n" + + "event: response.completed\n" + + `data: {"type":"response.completed","response":{"model":"gpt-5.4"}}` + "\n\n") + got := string(rewriteSSEPayloadLines(chunk, "gpt-5.4-fast")) + if !strings.Contains(got, "gpt-5.4-fast") { + t.Fatalf("rewritten chunk = %q, want alias gpt-5.4-fast", got) + } + if strings.Contains(got, `"model":"gpt-5.4"`) { + t.Fatalf("rewritten chunk still contains upstream model: %q", got) + } +}