diff --git a/internal/api/handlers/management/config_basic.go b/internal/api/handlers/management/config_basic.go index f9069198..9983b4d7 100644 --- a/internal/api/handlers/management/config_basic.go +++ b/internal/api/handlers/management/config_basic.go @@ -241,3 +241,11 @@ func (h *Handler) DeleteProxyURL(c *gin.Context) { h.cfg.ProxyURL = "" h.persist(c) } + +// Prioritize Model Mappings (for Amp CLI) +func (h *Handler) GetPrioritizeModelMappings(c *gin.Context) { + c.JSON(200, gin.H{"prioritize-model-mappings": h.cfg.AmpCode.PrioritizeModelMappings}) +} +func (h *Handler) PutPrioritizeModelMappings(c *gin.Context) { + h.updateBoolField(c, func(v bool) { h.cfg.AmpCode.PrioritizeModelMappings = v }) +} diff --git a/internal/api/modules/amp/amp.go b/internal/api/modules/amp/amp.go index dabb7404..5c7c2708 100644 --- a/internal/api/modules/amp/amp.go +++ b/internal/api/modules/amp/amp.go @@ -100,6 +100,16 @@ func (m *AmpModule) Name() string { return "amp-routing" } +// getPrioritizeModelMappings returns whether model mappings should take precedence over local API keys +func (m *AmpModule) getPrioritizeModelMappings() bool { + m.configMu.RLock() + defer m.configMu.RUnlock() + if m.lastConfig == nil { + return false + } + return m.lastConfig.PrioritizeModelMappings +} + // Register sets up Amp routes if configured. // This implements the RouteModuleV2 interface with Context. // Routes are registered only once via sync.Once for idempotent behavior. diff --git a/internal/api/modules/amp/fallback_handlers.go b/internal/api/modules/amp/fallback_handlers.go index 0cbe0e1a..771e2713 100644 --- a/internal/api/modules/amp/fallback_handlers.go +++ b/internal/api/modules/amp/fallback_handlers.go @@ -77,23 +77,29 @@ func logAmpRouting(routeType AmpRouteType, requestedModel, resolvedModel, provid // FallbackHandler wraps a standard handler with fallback logic to ampcode.com // when the model's provider is not available in CLIProxyAPI type FallbackHandler struct { - getProxy func() *httputil.ReverseProxy - modelMapper ModelMapper + getProxy func() *httputil.ReverseProxy + modelMapper ModelMapper + getPrioritizeModelMappings func() bool } // NewFallbackHandler creates a new fallback handler wrapper // The getProxy function allows lazy evaluation of the proxy (useful when proxy is created after routes) func NewFallbackHandler(getProxy func() *httputil.ReverseProxy) *FallbackHandler { return &FallbackHandler{ - getProxy: getProxy, + getProxy: getProxy, + getPrioritizeModelMappings: func() bool { return false }, } } // NewFallbackHandlerWithMapper creates a new fallback handler with model mapping support -func NewFallbackHandlerWithMapper(getProxy func() *httputil.ReverseProxy, mapper ModelMapper) *FallbackHandler { +func NewFallbackHandlerWithMapper(getProxy func() *httputil.ReverseProxy, mapper ModelMapper, getPrioritize func() bool) *FallbackHandler { + if getPrioritize == nil { + getPrioritize = func() bool { return false } + } return &FallbackHandler{ - getProxy: getProxy, - modelMapper: mapper, + getProxy: getProxy, + modelMapper: mapper, + getPrioritizeModelMappings: getPrioritize, } } @@ -130,34 +136,65 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc // Normalize model (handles Gemini thinking suffixes) normalizedModel, _ := util.NormalizeGeminiThinkingModel(modelName) - // Check if we have providers for this model - providers := util.GetProviderName(normalizedModel) - // Track resolved model for logging (may change if mapping is applied) resolvedModel := normalizedModel usedMapping := false + var providers []string - if len(providers) == 0 { - // No providers configured - check if we have a model mapping + // Check if model mappings should take priority over local API keys + prioritizeMappings := fh.getPrioritizeModelMappings != nil && fh.getPrioritizeModelMappings() + + if prioritizeMappings { + // PRIORITY MODE: Check model mappings FIRST (takes precedence over local API keys) + // This allows users to route Amp requests to their preferred OAuth providers if fh.modelMapper != nil { if mappedModel := fh.modelMapper.MapModel(normalizedModel); mappedModel != "" { - // Mapping found - rewrite the model in request body - bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel) - c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes)) - // Store mapped model in context for handlers that check it (like gemini bridge) - c.Set(MappedModelContextKey, mappedModel) - resolvedModel = mappedModel - usedMapping = true - - // Get providers for the mapped model - providers = util.GetProviderName(mappedModel) - - // Continue to handler with remapped model - goto handleRequest + // Mapping found - check if we have a provider for the mapped model + mappedProviders := util.GetProviderName(mappedModel) + if len(mappedProviders) > 0 { + // Mapping found and provider available - rewrite the model in request body + bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel) + c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + // Store mapped model in context for handlers that check it (like gemini bridge) + c.Set(MappedModelContextKey, mappedModel) + resolvedModel = mappedModel + usedMapping = true + providers = mappedProviders + } } } - // No mapping found - check if we have a proxy for fallback + // If no mapping applied, check for local providers + if !usedMapping { + providers = util.GetProviderName(normalizedModel) + } + } else { + // DEFAULT MODE: Check local providers first, then mappings as fallback + providers = util.GetProviderName(normalizedModel) + + if len(providers) == 0 { + // No providers configured - check if we have a model mapping + if fh.modelMapper != nil { + if mappedModel := fh.modelMapper.MapModel(normalizedModel); mappedModel != "" { + // Mapping found - check if we have a provider for the mapped model + mappedProviders := util.GetProviderName(mappedModel) + if len(mappedProviders) > 0 { + // Mapping found and provider available - rewrite the model in request body + bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel) + c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + // Store mapped model in context for handlers that check it (like gemini bridge) + c.Set(MappedModelContextKey, mappedModel) + resolvedModel = mappedModel + usedMapping = true + providers = mappedProviders + } + } + } + } + } + + // If no providers available, fallback to ampcode.com + if len(providers) == 0 { proxy := fh.getProxy() if proxy != nil { // Log: Forwarding to ampcode.com (uses Amp credits) @@ -175,8 +212,6 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc logAmpRouting(RouteTypeNoProvider, modelName, "", "", requestPath) } - handleRequest: - // Log the routing decision providerName := "" if len(providers) > 0 { diff --git a/internal/api/modules/amp/routes.go b/internal/api/modules/amp/routes.go index 6826dbbe..dedbd444 100644 --- a/internal/api/modules/amp/routes.go +++ b/internal/api/modules/amp/routes.go @@ -171,7 +171,7 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha geminiBridge := createGeminiBridgeHandler(geminiHandlers.GeminiHandler) geminiV1Beta1Fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy { return m.getProxy() - }, m.modelMapper) + }, m.modelMapper, m.getPrioritizeModelMappings) geminiV1Beta1Handler := geminiV1Beta1Fallback.WrapHandler(geminiBridge) // Route POST model calls through Gemini bridge with FallbackHandler. @@ -209,7 +209,7 @@ func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *han // Also includes model mapping support for routing unavailable models to alternatives fallbackHandler := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy { return m.getProxy() - }, m.modelMapper) + }, m.modelMapper, m.getPrioritizeModelMappings) // Provider-specific routes under /api/provider/:provider ampProviders := engine.Group("/api/provider") diff --git a/internal/api/server.go b/internal/api/server.go index 72cb0313..6a2b6400 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -520,6 +520,10 @@ func (s *Server) registerManagementRoutes() { mgmt.PUT("/ws-auth", s.mgmt.PutWebsocketAuth) mgmt.PATCH("/ws-auth", s.mgmt.PutWebsocketAuth) + mgmt.GET("/prioritize-model-mappings", s.mgmt.GetPrioritizeModelMappings) + mgmt.PUT("/prioritize-model-mappings", s.mgmt.PutPrioritizeModelMappings) + mgmt.PATCH("/prioritize-model-mappings", s.mgmt.PutPrioritizeModelMappings) + mgmt.GET("/request-retry", s.mgmt.GetRequestRetry) mgmt.PUT("/request-retry", s.mgmt.PutRequestRetry) mgmt.PATCH("/request-retry", s.mgmt.PutRequestRetry) diff --git a/internal/auth/iflow/iflow_auth.go b/internal/auth/iflow/iflow_auth.go index b3431f84..1d69a821 100644 --- a/internal/auth/iflow/iflow_auth.go +++ b/internal/auth/iflow/iflow_auth.go @@ -321,17 +321,23 @@ func (ia *IFlowAuth) AuthenticateWithCookie(ctx context.Context, cookie string) return nil, fmt.Errorf("iflow cookie authentication: cookie is empty") } - // First, get initial API key information using GET request + // First, get initial API key information using GET request to obtain the name keyInfo, err := ia.fetchAPIKeyInfo(ctx, cookie) if err != nil { return nil, fmt.Errorf("iflow cookie authentication: fetch initial API key info failed: %w", err) } - // Convert to token data format + // Refresh the API key using POST request + refreshedKeyInfo, err := ia.RefreshAPIKey(ctx, cookie, keyInfo.Name) + if err != nil { + return nil, fmt.Errorf("iflow cookie authentication: refresh API key failed: %w", err) + } + + // Convert to token data format using refreshed key data := &IFlowTokenData{ - APIKey: keyInfo.APIKey, - Expire: keyInfo.ExpireTime, - Email: keyInfo.Name, + APIKey: refreshedKeyInfo.APIKey, + Expire: refreshedKeyInfo.ExpireTime, + Email: refreshedKeyInfo.Name, Cookie: cookie, } diff --git a/internal/config/config.go b/internal/config/config.go index 1c72ece4..4746a58e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -151,6 +151,10 @@ type AmpCode struct { // When Amp requests a model that isn't available locally, these mappings // allow routing to an alternative model that IS available. ModelMappings []AmpModelMapping `yaml:"model-mappings" json:"model-mappings"` + + // PrioritizeModelMappings when true, model mappings take precedence over local API keys. + // When false (default), local API keys are used first if available. + PrioritizeModelMappings bool `yaml:"prioritize-model-mappings" json:"prioritize-model-mappings"` } // PayloadConfig defines default and override parameter rules applied to provider payloads. diff --git a/internal/registry/model_definitions.go b/internal/registry/model_definitions.go index aa09f688..b25d91c2 100644 --- a/internal/registry/model_definitions.go +++ b/internal/registry/model_definitions.go @@ -943,8 +943,19 @@ func GetQwenModels() []*ModelInfo { } } -// GetIFlowModels returns supported models for iFlow OAuth accounts. +// GetAntigravityThinkingConfig returns the Thinking configuration for antigravity models. +// Keys use the ALIASED model names (after modelName2Alias conversion) for direct lookup. +func GetAntigravityThinkingConfig() map[string]*ThinkingSupport { + return map[string]*ThinkingSupport{ + "gemini-2.5-flash": {Min: 0, Max: 24576, ZeroAllowed: true, DynamicAllowed: true}, + "gemini-2.5-flash-lite": {Min: 0, Max: 24576, ZeroAllowed: true, DynamicAllowed: true}, + "gemini-3-pro-preview": {Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true}, + "gemini-claude-sonnet-4-5-thinking": {Min: 1024, Max: 200000, ZeroAllowed: false, DynamicAllowed: true}, + "gemini-claude-opus-4-5-thinking": {Min: 1024, Max: 200000, ZeroAllowed: false, DynamicAllowed: true}, + } +} +// GetIFlowModels returns supported models for iFlow OAuth accounts. func GetIFlowModels() []*ModelInfo { entries := []struct { ID string diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index f2f0fdc5..ce836a77 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -370,29 +370,25 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c } now := time.Now().Unix() + thinkingConfig := registry.GetAntigravityThinkingConfig() models := make([]*registry.ModelInfo, 0, len(result.Map())) - for id := range result.Map() { - id = modelName2Alias(id) - if id != "" { + for originalName := range result.Map() { + aliasName := modelName2Alias(originalName) + if aliasName != "" { modelInfo := ®istry.ModelInfo{ - ID: id, - Name: id, - Description: id, - DisplayName: id, - Version: id, + ID: aliasName, + Name: aliasName, + Description: aliasName, + DisplayName: aliasName, + Version: aliasName, Object: "model", Created: now, OwnedBy: antigravityAuthType, Type: antigravityAuthType, } - // Add Thinking support for thinking models - if strings.HasSuffix(id, "-thinking") || strings.Contains(id, "-thinking-") { - modelInfo.Thinking = ®istry.ThinkingSupport{ - Min: 1024, - Max: 100000, - ZeroAllowed: false, - DynamicAllowed: true, - } + // Look up Thinking support from static config using alias name + if thinking, ok := thinkingConfig[aliasName]; ok { + modelInfo.Thinking = thinking } models = append(models, modelInfo) } diff --git a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_request.go b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_request.go index d1914ec8..82e71758 100644 --- a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_request.go +++ b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_request.go @@ -88,6 +88,20 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _ } } + // Claude/Anthropic API format: thinking.type == "enabled" with budget_tokens + // This allows Claude Code and other Claude API clients to pass thinking configuration + if !gjson.GetBytes(out, "request.generationConfig.thinkingConfig").Exists() && util.ModelSupportsThinking(modelName) { + if t := gjson.GetBytes(rawJSON, "thinking"); t.Exists() && t.IsObject() { + if t.Get("type").String() == "enabled" { + if b := t.Get("budget_tokens"); b.Exists() && b.Type == gjson.Number { + budget := util.NormalizeThinkingBudget(modelName, int(b.Int())) + out, _ = sjson.SetBytes(out, "request.generationConfig.thinkingConfig.thinkingBudget", budget) + out, _ = sjson.SetBytes(out, "request.generationConfig.thinkingConfig.include_thoughts", true) + } + } + } + } + // For gemini-3-pro-preview, always send default thinkingConfig when none specified. // This matches the official Gemini CLI behavior which always sends: // { thinkingBudget: -1, includeThoughts: true } diff --git a/internal/translator/openai/claude/openai_claude_request.go b/internal/translator/openai/claude/openai_claude_request.go index bff306cc..3521b2e5 100644 --- a/internal/translator/openai/claude/openai_claude_request.go +++ b/internal/translator/openai/claude/openai_claude_request.go @@ -8,6 +8,7 @@ package claude import ( "bytes" "encoding/json" + "strings" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -242,11 +243,12 @@ func convertClaudeContentPart(part gjson.Result) (string, bool) { switch partType { case "text": - if !part.Get("text").Exists() { + text := part.Get("text").String() + if strings.TrimSpace(text) == "" { return "", false } textContent := `{"type":"text","text":""}` - textContent, _ = sjson.Set(textContent, "text", part.Get("text").String()) + textContent, _ = sjson.Set(textContent, "text", text) return textContent, true case "image":