diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index b27440f3..c36f771c 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -600,11 +600,12 @@ func mapAntigravityModel(account *Account, requestedModel string) string { return mapped } - // 如果 mapped == requestedModel，检查是否在映射表 key 中显式配置 + // 如果 mapped == requestedModel，检查是否在映射表中配置（精确或通配符） // 这区分两种情况： // 1. 映射表中有 "model-a": "model-a"（显式透传）→ 返回 model-a - // 2. 映射表中没有 model-a 的配置 → 返回空（不支持） - if _, exists := mapping[requestedModel]; exists { + // 2. 通配符匹配 "claude-*": "claude-sonnet-4-5" 恰好目标等于请求名 → 返回 model-a + // 3. 映射表中没有 model-a 的配置 → 返回空（不支持） + if account.IsModelSupported(requestedModel) { return requestedModel } @@ -2042,7 +2043,7 @@ func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { // - waitDuration: 等待时间（智能重试时使用，shouldRateLimitModel=true 时为 0） // - modelName: 限流的模型名称 func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string) { - if !account.IsOAuth() { + if account.Platform != PlatformAntigravity { return false, false, 0, "" } diff --git a/backend/internal/service/antigravity_rate_limit_test.go b/backend/internal/service/antigravity_rate_limit_test.go index f70a30de..4e5ad24d 100644 --- a/backend/internal/service/antigravity_rate_limit_test.go +++ b/backend/internal/service/antigravity_rate_limit_test.go @@ -524,8 +524,8 @@ func TestParseAntigravitySmartRetryInfo(t *testing.T) { } func TestShouldTriggerAntigravitySmartRetry(t *testing.T) { - oauthAccount := &Account{Type: AccountTypeOAuth} - setupTokenAccount := &Account{Type: AccountTypeSetupToken} + oauthAccount := &Account{Type: AccountTypeOAuth, Platform: PlatformAntigravity} + setupTokenAccount := &Account{Type: AccountTypeSetupToken, Platform: PlatformAntigravity} apiKeyAccount := &Account{Type: AccountTypeAPIKey} tests := []struct { diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index 95ef2489..f6f43aad 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -343,13 +343,13 @@ func TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError(t *testi require.Equal(t, "gemini-3-pro-high", repo.modelRateLimitCalls[0].modelKey) } -// TestHandleSmartRetry_NonOAuthAccount_ContinuesDefaultLogic 测试非 OAuth 账号走默认逻辑 -func TestHandleSmartRetry_NonOAuthAccount_ContinuesDefaultLogic(t *testing.T) { +// TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic 测试非 Antigravity 平台账号走默认逻辑 +func TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic(t *testing.T) { account := &Account{ ID: 4, Name: "acc-4", - Type: AccountTypeAPIKey, // 非 OAuth 账号 - Platform: PlatformAntigravity, + Type: AccountTypeAPIKey, // 非 Antigravity 平台账号 + Platform: PlatformAnthropic, } // 即使是模型限流响应，非 OAuth 账号也应该走默认逻辑 @@ -385,7 +385,7 @@ func TestHandleSmartRetry_NonOAuthAccount_ContinuesDefaultLogic(t *testing.T) { result := handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs) require.NotNil(t, result) - require.Equal(t, smartRetryActionContinue, result.action, "non-OAuth account should continue default logic") + require.Equal(t, smartRetryActionContinue, result.action, "non-Antigravity platform account should continue default logic") require.Nil(t, result.resp) require.Nil(t, result.err) require.Nil(t, result.switchError) diff --git a/backend/internal/service/antigravity_token_provider.go b/backend/internal/service/antigravity_token_provider.go index 94eca94d..1eb740f9 100644 --- a/backend/internal/service/antigravity_token_provider.go +++ b/backend/internal/service/antigravity_token_provider.go @@ -42,7 +42,18 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account * if account == nil { return "", errors.New("account is nil") } - if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { + if account.Platform != PlatformAntigravity { + return "", errors.New("not an antigravity account") + } + // upstream 类型：直接从 credentials 读取 api_key，不走 OAuth 刷新流程 + if account.Type == AccountTypeUpstream { + apiKey := account.GetCredential("api_key") + if apiKey == "" { + return "", errors.New("upstream account missing api_key in credentials") + } + return apiKey, nil + } + if account.Type != AccountTypeOAuth { return "", errors.New("not an antigravity oauth account") } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 250f7bff..dd2aad34 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -2587,6 +2587,9 @@ func (s *GatewayService) isModelSupportedByAccountWithContext(ctx context.Contex // 应用 thinking 后缀后检查最终模型是否在账号映射中 if enabled, ok := ctx.Value(ctxkey.ThinkingEnabled).(bool); ok { finalModel := applyThinkingModelSuffix(mapped, enabled) + if finalModel == mapped { + return true // thinking 后缀未改变模型名，映射已通过 + } return account.IsModelSupported(finalModel) } return true diff --git a/backend/internal/service/gateway_service_antigravity_whitelist_test.go b/backend/internal/service/gateway_service_antigravity_whitelist_test.go index 553dc55b..c078be32 100644 --- a/backend/internal/service/gateway_service_antigravity_whitelist_test.go +++ b/backend/internal/service/gateway_service_antigravity_whitelist_test.go @@ -80,21 +80,21 @@ func TestGatewayService_isModelSupportedByAccountWithContext_ThinkingMode(t *tes thinkingEnabled bool expected bool }{ - // 场景 1: 配置 claude-sonnet-4-5-thinking，请求 claude-sonnet-4-5 + thinking=true - // 最终模型名 = claude-sonnet-4-5-thinking，应该匹配 + // 场景 1: 只配置 claude-sonnet-4-5-thinking，请求 claude-sonnet-4-5 + thinking=true + // mapAntigravityModel 找不到 claude-sonnet-4-5 的映射 → 返回 false { - name: "thinking_enabled_matches_thinking_model", + name: "thinking_enabled_no_base_mapping_returns_false", modelMapping: map[string]any{ "claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking", }, requestedModel: "claude-sonnet-4-5", thinkingEnabled: true, - expected: true, + expected: false, }, // 场景 2: 只配置 claude-sonnet-4-5-thinking，请求 claude-sonnet-4-5 + thinking=false - // 最终模型名 = claude-sonnet-4-5，不在 mapping 中，应该不匹配 + // mapAntigravityModel 找不到 claude-sonnet-4-5 的映射 → 返回 false { - name: "thinking_disabled_no_match_thinking_only_mapping", + name: "thinking_disabled_no_base_mapping_returns_false", modelMapping: map[string]any{ "claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking", }, @@ -145,15 +145,16 @@ func TestGatewayService_isModelSupportedByAccountWithContext_ThinkingMode(t *tes thinkingEnabled: true, expected: true, // claude-sonnet-4-5-thinking 匹配 claude-* }, - // 场景 7: 其他模型（非 sonnet-4-5）的 thinking 不受影响 + // 场景 7: 只配置 thinking 变体但没有基础模型映射 → 返回 false + // mapAntigravityModel 找不到 claude-opus-4-6 的映射 { - name: "opus_thinking_unchanged", + name: "opus_thinking_no_base_mapping_returns_false", modelMapping: map[string]any{ "claude-opus-4-6-thinking": "claude-opus-4-6-thinking", }, requestedModel: "claude-opus-4-6", thinkingEnabled: true, - expected: true, // claude-opus-4-6 映射到 claude-opus-4-6-thinking，匹配 + expected: false, }, } @@ -175,3 +176,65 @@ func TestGatewayService_isModelSupportedByAccountWithContext_ThinkingMode(t *tes }) } } + +// TestGatewayService_isModelSupportedByAccount_CustomMappingNotInDefault 测试自定义模型映射中 +// 不在 DefaultAntigravityModelMapping 中的模型能通过调度 +func TestGatewayService_isModelSupportedByAccount_CustomMappingNotInDefault(t *testing.T) { + svc := &GatewayService{} + + // 自定义映射中包含不在默认映射中的模型 + account := &Account{ + Platform: PlatformAntigravity, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "my-custom-model": "actual-upstream-model", + "gpt-4o": "some-upstream-model", + "llama-3-70b": "llama-3-70b-upstream", + "claude-sonnet-4-5": "claude-sonnet-4-5", + }, + }, + } + + // 自定义模型应该通过（不在 DefaultAntigravityModelMapping 中也可以） + require.True(t, svc.isModelSupportedByAccount(account, "my-custom-model")) + require.True(t, svc.isModelSupportedByAccount(account, "gpt-4o")) + require.True(t, svc.isModelSupportedByAccount(account, "llama-3-70b")) + require.True(t, svc.isModelSupportedByAccount(account, "claude-sonnet-4-5")) + + // 不在自定义映射中的模型不通过 + require.False(t, svc.isModelSupportedByAccount(account, "gpt-3.5-turbo")) + require.False(t, svc.isModelSupportedByAccount(account, "unknown-model")) + + // 空模型允许 + require.True(t, svc.isModelSupportedByAccount(account, "")) +} + +// TestGatewayService_isModelSupportedByAccountWithContext_CustomMappingThinking +// 测试自定义映射 + thinking 模式的交互 +func TestGatewayService_isModelSupportedByAccountWithContext_CustomMappingThinking(t *testing.T) { + svc := &GatewayService{} + + // 自定义映射同时配置基础模型和 thinking 变体 + account := &Account{ + Platform: PlatformAntigravity, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "claude-sonnet-4-5": "claude-sonnet-4-5", + "claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking", + "my-custom-model": "upstream-model", + }, + }, + } + + // thinking=true: claude-sonnet-4-5 → mapped=claude-sonnet-4-5 → +thinking → check IsModelSupported(claude-sonnet-4-5-thinking)=true + ctx := context.WithValue(context.Background(), ctxkey.ThinkingEnabled, true) + require.True(t, svc.isModelSupportedByAccountWithContext(ctx, account, "claude-sonnet-4-5")) + + // thinking=false: claude-sonnet-4-5 → mapped=claude-sonnet-4-5 → check IsModelSupported(claude-sonnet-4-5)=true + ctx = context.WithValue(context.Background(), ctxkey.ThinkingEnabled, false) + require.True(t, svc.isModelSupportedByAccountWithContext(ctx, account, "claude-sonnet-4-5")) + + // 自定义模型（非 claude）不受 thinking 后缀影响，mapped 成功即通过 + ctx = context.WithValue(context.Background(), ctxkey.ThinkingEnabled, true) + require.True(t, svc.isModelSupportedByAccountWithContext(ctx, account, "my-custom-model")) +}