Merge remote-tracking branch 'origin/main' into contrib/gemini-4-argon-support

Resolve Gemini 4 Argon pricing conflicts in favor of main (#8078), which
already provides the base and Cyber Permissive price cards, thinking-tier
pricing aliases, and the local pricing snapshot.
This commit is contained in:
shaw
2026-10-15 15:37:14 +08:00
31 changed files with 1280 additions and 152 deletions
-5
View File
@@ -63,11 +63,6 @@ Please read the following carefully before using this project:
<td>Real GPT-5.6 series at 3% of OpenAI pricing — <a href="https://codex-everywhere.com">CodexEverywhere</a> is democratizing access to frontier models for developers worldwide. We believe in transparency and honesty, with model quality verified by active community oversight for months. USD and crypto friendly. Start with a free $20 trial at <a href="https://codex-everywhere.com">codex-everywhere.com</a>.</td>
</tr>
<tr>
<td width="180"><a href="https://shop.bmoplus.com/?utm_source=github"><img src="assets/partners/logos/bmoplus.jpg" alt="bmoplus" width="150"></a></td>
<td>Huge thanks to BmoPlus for sponsoring this project! BmoPlus is a highly reliable AI account provider built strictly for heavy AI users and developers. They offer rock-solid, ready-to-use accounts and official top-up services for ChatGPT Plus / ChatGPT Pro (Full Warranty) / Claude Pro / Super Grok / Gemini Pro. By registering and ordering through <a href="https://shop.bmoplus.com/?utm_source=github">BmoPlus - Premium AI Accounts & Top-ups</a>, users can unlock the mind-blowing rate of 10% of the official GPT subscription price (90% OFF)</td>
</tr>
<tr>
<td width="180"><a href="https://bestproxy.com/?keyword=a2e8iuol"><img src="assets/partners/logos/bestproxy.png" alt="bestproxy" width="150"></a></td>
<td>Thanks to Bestproxy for sponsoring this project! <a href="https://bestproxy.com/?keyword=a2e8iuol">Bestproxy</a> provides high-purity residential IPs with dedicated one-IP-per-account support. By combining real home networks with fingerprint isolation, it enables link environment isolation and reduces the probability of association-based risk control.</td>
-5
View File
@@ -64,11 +64,6 @@
<td>Real GPT-5.6 series at 3% of OpenAI pricing — <a href="https://codex-everywhere.com">CodexEverywhere</a> is democratizing access to frontier models for developers worldwide. We believe in transparency and honesty, with model quality verified by active community oversight for months. USD and crypto friendly. Start with a free $20 trial at <a href="https://codex-everywhere.com">codex-everywhere.com</a>.</td>
</tr>
<tr>
<td width="180"><a href="https://shop.bmoplus.com/?utm_source=github"><img src="assets/partners/logos/bmoplus.jpg" alt="bmoplus" width="150"></a></td>
<td>感谢 BmoPlus 赞助了本项目!BmoPlus 是一家专为AI订阅重度用户打造的可靠 AI 账号代充服务商,提供稳定的 ChatGPT Plus / ChatGPT Pro(全程质保) / Claude Pro / Super Grok / Gemini Pro 的官方代充&成品账号。 通过<a href="https://shop.bmoplus.com/?utm_source=github">BmoPlus AI成品号专卖/代充</a>注册下单的用户,可享GPT 官网订阅一折 的震撼价格!</td>
</tr>
<tr>
<td width="180"><a href="https://bestproxy.com/?keyword=a2e8iuol"><img src="assets/partners/logos/bestproxy.png" alt="bestproxy" width="150"></a></td>
<td>感谢 Bestproxy 赞助了本项目!<a href="https://bestproxy.com/?keyword=a2e8iuol">Bestproxy</a> 是一家提供高纯度住宅IP,支持一号一IP独享,结合真实家庭网络与指纹隔离,可实现链路环境隔离,降低关联风控概率。</td>
-5
View File
@@ -63,11 +63,6 @@
<td>OpenAI 公式価格のわずか 3% で本物の GPT-5.6 シリーズを提供 — <a href="https://codex-everywhere.com">CodexEverywhere</a> は世界中の開発者にフロンティアモデルへのアクセスを民主化しています。私たちは透明性と誠実さを信条とし、モデル品質は数か月にわたるアクティブなコミュニティの監視によって検証されています。USD および暗号通貨に対応。<a href="https://codex-everywhere.com">codex-everywhere.com</a> で $20 の無料トライアルから始めましょう。</td>
</tr>
<tr>
<td width="180"><a href="https://shop.bmoplus.com/?utm_source=github"><img src="assets/partners/logos/bmoplus.jpg" alt="bmoplus" width="150"></a></td>
<td>本プロジェクトにご支援いただいた BmoPlus に感謝いたします!BmoPlusは、AIサブスクリプションのヘビーユーザー向けに特化した信頼性の高いAIアカウントサービスプロバイダーであり、安定した ChatGPT Plus / ChatGPT Pro (完全保証) / Claude Pro / Super Grok / Gemini Pro の公式代行チャージおよび即納アカウントを提供しています。こちらの<a href="https://shop.bmoplus.com/?utm_source=github">BmoPlus AIアカウント専門店/代行チャージ</a>経由でご登録・ご注文いただいたユーザー様は、GPTを 公式サイト価格の約1割(90% OFF) という驚異的な価格でご利用いただけます!</td>
</tr>
<tr>
<td width="180"><a href="https://bestproxy.com/?keyword=a2e8iuol"><img src="assets/partners/logos/bestproxy.png" alt="bestproxy" width="150"></a></td>
<td>Bestproxy のご支援に感謝します!<a href="https://bestproxy.com/?keyword=a2e8iuol">Bestproxy</a> は高純度の住宅IPを提供し、1アカウント1IP専有をサポートしています。実際の家庭ネットワークとフィンガープリント分離を組み合わせることで、リンク環境の分離を実現し、関連付けによるリスク管理の確率を低減します。</td>
@@ -2540,6 +2540,11 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "Failed to resolve composite model route")
return
}
if decision.OAuthAmbiguous() {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, decision.AmbiguousMessage())
return
}
if decision.Matched {
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
wsRouteModel = decision.UpstreamModel
@@ -185,6 +185,53 @@ func TestCompositeTargetPlatformMiddlewareRewritesNestedLiveModel(t *testing.T)
require.Equal(t, http.StatusNoContent, w.Code)
}
// Scenario: 回归 —— openai OAuth 与 deepseek API Key 同时映射 gpt-6-astra 时,
// 请求必须在路由层被拒绝,不能落到默认 Anthropic handler 后再按模型名猜平台。
// 仅 API Key 的跨平台歧义保持原行为(放行且不解析平台)。
func TestCompositeTargetPlatformMiddlewareRejectsOnlyOAuthAmbiguousOwnership(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
name string
oauthClaimed bool
wantStatus int
}{
{name: "oauth claimant", oauthClaimed: true, wantStatus: http.StatusBadRequest},
{name: "apikey only", oauthClaimed: false, wantStatus: http.StatusNoContent},
} {
t.Run(tc.name, func(t *testing.T) {
resolver := service.NewCompositeRouteResolver(nil)
resolver.SetModelOwnershipResolver(func(context.Context, int64, string) (service.CompositeModelOwnership, error) {
return service.CompositeModelOwnership{Ambiguous: true, OpenAIGatewayOAuthClaimed: tc.oauthClaimed}, nil
})
router := gin.New()
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(8)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(resolver))
router.POST("/v1/responses", func(c *gin.Context) {
_, resolved := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.False(t, resolved)
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-6-astra"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, tc.wantStatus, w.Code)
if tc.oauthClaimed {
require.Contains(t, w.Body.String(), "explicit composite route")
}
})
}
}
func TestCompositeRequestModelFromMultipartLiveSession(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
@@ -586,6 +586,14 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
c.Abort()
return
}
// 含 OAuth 账号的跨平台歧义是分组配置错误:拒绝而不是放行。否则未解析的
// 请求落到默认 Anthropic handler,handler 再按模型名前缀另猜平台,
// OAuth 凭据被送进不匹配的协议管线(如 ChatGPT token 发往 Anthropic)。
if decision.OAuthAmbiguous() {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"type": "invalid_request_error", "message": decision.AmbiguousMessage()}})
return
}
if decision.Matched {
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
if upstreamModel := strings.TrimSpace(decision.UpstreamModel); upstreamModel != "" && upstreamModel != model && gjson.ValidBytes(body) {
+20 -4
View File
@@ -524,15 +524,27 @@ func (s *BillingService) initFallbackPricing() {
SupportsCacheBreakdown: false,
}
// Gemini 4 Argon (Google AI pricing: $2.00 input / $10.00 output /
// $0.10 cached input per MTok introductory rate through 2026-11-30).
// Gemini 4 Argon (Google AI pricing: $2 input / $10 output / $0.10 cached
// input per MTok, promotional through 2026-11-30; official rates double to
// $4/$20/$0.20 from 2026-12-01). One rate covers every prompt length.
// Antigravity's -high/-low/-medium aliases are matched below so unavailable
// remote pricing never records token-bearing requests at $0.
s.fallbackPrices["gemini-4-argon"] = &ModelPricing{
InputPricePerToken: 2.0e-6,
OutputPricePerToken: 10.0e-6,
InputPricePerToken: 2e-6,
OutputPricePerToken: 10e-6,
CacheReadPricePerToken: 0.1e-6,
SupportsCacheBreakdown: false,
}
// Gemini 4 Argon Cyber Permissive (Vertex AI pricing: $4 input / $20 output /
// $0.20 cached input per MTok). The Gemini 4 Argon promotion does not apply.
s.fallbackPrices["gemini-4-argon-cyber-permissive"] = &ModelPricing{
InputPricePerToken: 4e-6,
OutputPricePerToken: 20e-6,
CacheReadPricePerToken: 0.2e-6,
SupportsCacheBreakdown: false,
}
// OpenAI GPT-5.4(业务指定价格)
s.fallbackPrices["gpt-5.4"] = &ModelPricing{
InputPricePerToken: 2.5e-6, // $2.5 per MTok
@@ -1091,6 +1103,10 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
if strings.Contains(modelLower, "gemini-3.8-flash") || strings.Contains(modelLower, "gemini-3-8-flash") {
return s.fallbackPrices["gemini-3.8-flash"]
}
// Cyber 变体不享受 Gemini 4 Argon 的优惠价,必须先于 "gemini-4-argon" 判。
if strings.Contains(modelLower, "gemini-4-argon-cyber") {
return s.fallbackPrices["gemini-4-argon-cyber-permissive"]
}
if strings.Contains(modelLower, "gemini-4-argon") {
return s.fallbackPrices["gemini-4-argon"]
}
@@ -497,7 +497,7 @@ func minimaxResetTime(v gjson.Result) string {
return time.UnixMilli(ms).UTC().Format(time.RFC3339)
}
// cnZhipuWindow 标识智谱 TOKENS_LIMIT 条目所属窗口。
// cnZhipuWindow 标识智谱额度条目(TOKENS_LIMIT / CREDIT_LIMIT)所属窗口。
type cnZhipuWindow int
const (
@@ -529,7 +529,8 @@ func classifyZhipuWindowUnit(unit int64) cnZhipuWindow {
//
// CREDIT_LIMIT(信用额度)与 TOKENS_LIMIT(token 窗口)度量不同:两者同时返回时
// 只让 TOKENS_LIMIT 参与 5h/weekly 槽位竞争,避免信用额度百分比污染阈值停调
// 快照;仅当无任何 TOKENS_LIMIT 条目时才降级用 CREDIT_LIMIT 展示。
// 快照;仅当无任何 TOKENS_LIMIT 条目时才降级用 CREDIT_LIMIT 展示 —— 此时
// CREDIT_LIMIT 同样按 unit 分类(它也携带 unit),不依赖 reset 时间排序。
// 老套餐只回 1 条 TOKENS_LIMIT,自然降级为仅 5h;新套餐回 2 条。
func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
type entry struct {
@@ -537,6 +538,7 @@ func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
hasReset bool
percentage float64
resetISO string
unit int64
}
var (
fiveHour entry
@@ -546,8 +548,8 @@ func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
unclassified []entry
)
classify := func(item gjson.Result, e entry) {
switch classifyZhipuWindowUnit(item.Get("unit").Int()) {
classify := func(unit int64, e entry) {
switch classifyZhipuWindowUnit(unit) {
case cnZhipuWindow5h:
if !fiveHourSet {
fiveHour, fiveHourSet = e, true
@@ -592,10 +594,10 @@ func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
hasReset = resetISO != ""
}
}
e := entry{resetMs: resetMs, hasReset: hasReset, percentage: percentage, resetISO: resetISO}
e := entry{resetMs: resetMs, hasReset: hasReset, percentage: percentage, resetISO: resetISO, unit: item.Get("unit").Int()}
if limitType == "TOKENS_LIMIT" {
hasTokensLimit = true
classify(item, e)
classify(e.unit, e)
} else {
creditFallback = append(creditFallback, e)
}
@@ -603,8 +605,17 @@ func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
})
// 无任何 TOKENS_LIMIT 条目(部分套餐只报信用额度):降级用 CREDIT_LIMIT 展示。
// CREDIT_LIMIT 条目同样带 unit(3=5h / 6=weekly;实测 GLM Coding Plan 两条窗口
// 都是 CREDIT_LIMIT),必须和 TOKENS_LIMIT 一样先按 unit 显式分类 —— 否则会退化
// 成 reset 时间启发式,而周期末尾周窗口比 5h 更早重置,时间排序必然把两档标反。
if !hasTokensLimit {
unclassified = append(unclassified, creditFallback...)
for _, e := range creditFallback {
if classifyZhipuWindowUnit(e.unit) != cnZhipuWindowUnknown {
classify(e.unit, e)
continue
}
unclassified = append(unclassified, e)
}
}
// 无 reset 的条目排前,再按 reset 升序,依次填入仍空缺的槽位。
@@ -126,6 +126,60 @@ func TestParseZhipuTokenTiers_UnitClassification(t *testing.T) {
require.InDelta(t, 70.0, tiers[1].UsedPercent, 1e-9)
}
// TestParseZhipuTokenTiers_CreditLimitUnitClassification GLM Coding Plan 当前把 5h / weekly
// 两条窗口都标成 CREDIT_LIMIT(实测 level=lite)。CREDIT_LIMIT 同样携带 unit,
// 必须和 TOKENS_LIMIT 一样按 unit 显式分类,否则会退化成 reset 时间排序 ——
// 此处刻意让 weekly 的 nextResetTime 早于 5h(周期末尾),两档会被标反。
func TestParseZhipuTokenTiers_CreditLimitUnitClassification(t *testing.T) {
t.Parallel()
data := gjson.Parse(`{
"limits": [
{"type":"CREDIT_LIMIT","unit":3,"number":5,"percentage":14,"nextResetTime":1700000099999},
{"type":"CREDIT_LIMIT","unit":6,"number":1,"percentage":90,"nextResetTime":1700000000000}
]
}`)
tiers := parseZhipuTokenTiers(data)
require.Len(t, tiers, 2)
require.Equal(t, "5h", tiers[0].Window)
require.InDelta(t, 14.0, tiers[0].UsedPercent, 1e-9)
require.Equal(t, "weekly", tiers[1].Window)
require.InDelta(t, 90.0, tiers[1].UsedPercent, 1e-9)
}
// TestParseZhipuTokenTiers_CreditLimitWithoutUnitFallback CREDIT_LIMIT 缺 unit 时,
// 保留原有 fallback:无 nextResetTime 的优先归 5h,其余按 reset 升序填入剩余槽位。
func TestParseZhipuTokenTiers_CreditLimitWithoutUnitFallback(t *testing.T) {
t.Parallel()
data := gjson.Parse(`{
"limits": [
{"type":"CREDIT_LIMIT","percentage":30,"nextResetTime":1700000000000},
{"type":"CREDIT_LIMIT","percentage":80}
]
}`)
tiers := parseZhipuTokenTiers(data)
require.Len(t, tiers, 2)
require.Equal(t, "5h", tiers[0].Window)
require.InDelta(t, 80.0, tiers[0].UsedPercent, 1e-9)
require.Equal(t, "weekly", tiers[1].Window)
require.InDelta(t, 30.0, tiers[1].UsedPercent, 1e-9)
}
// TestParseZhipuTokenTiers_TokensLimitWinsOverCreditLimit 两种 type 同时返回时,
// TOKENS_LIMIT 仍然优先占槽、CREDIT_LIMIT 不参与(原有语义保持不变)。
func TestParseZhipuTokenTiers_TokensLimitWinsOverCreditLimit(t *testing.T) {
t.Parallel()
data := gjson.Parse(`{
"limits": [
{"type":"TOKENS_LIMIT","unit":3,"percentage":20,"nextResetTime":1700000000000},
{"type":"CREDIT_LIMIT","unit":6,"percentage":99,"nextResetTime":1700000099999}
]
}`)
tiers := parseZhipuTokenTiers(data)
require.Len(t, tiers, 1)
require.Equal(t, "5h", tiers[0].Window)
require.InDelta(t, 20.0, tiers[0].UsedPercent, 1e-9)
}
// TestParseZhipuTokenTiers_SingleTierOldPlan 老套餐仅回 1 条 → 降级为仅 5h。
func TestParseZhipuTokenTiers_SingleTierOldPlan(t *testing.T) {
t.Parallel()
@@ -32,6 +32,11 @@ type CompositeModelOwnership struct {
TargetPlatform string
Matched bool
Ambiguous bool
// OpenAIGatewayOAuthClaimed 表示声明方中含 OpenAI 网关族平台的 OAuth /
// setup-token 凭据。歧义请求无法解析平台,会落到默认的 Anthropic 网关
// handler:Anthropic / Gemini / Antigravity 的 OAuth 在那里有原生或兼容
// 转发,API Key 上游可能兼容多种格式,唯独这类凭据只能走 OpenAI 网关。
OpenAIGatewayOAuthClaimed bool
}
type CompositeModelOwnershipResolver func(context.Context, int64, string) (CompositeModelOwnership, error)
@@ -13,7 +13,7 @@ type compositeOwnershipAccountRepo struct {
accounts []Account
}
func (r *compositeOwnershipAccountRepo) ListSchedulableByGroupID(context.Context, int64) ([]Account, error) {
func (r *compositeOwnershipAccountRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) {
return r.accounts, nil
}
@@ -103,6 +103,35 @@ func TestResolveCompositeModelOwnershipAllowsSamePlatformAndRejectsCrossPlatform
require.Equal(t, CompositeModelOwnership{Ambiguous: true}, ambiguous)
}
// Scenario: 只有 OpenAI 网关族的 OAuth 声明方会被标记(API Key 上游可能兼容多协议,
// Anthropic / Antigravity 的 OAuth 在默认 Anthropic handler 中可正常转发)
func TestResolveCompositeModelOwnershipFlagsOpenAIGatewayOAuthClaimants(t *testing.T) {
groupID := int64(7)
repo := &compositeOwnershipAccountRepo{
accounts: []Account{
{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-6-astra": "gpt-6-astra"}}},
{ID: 2, Platform: PlatformDeepseek, Type: AccountTypeAPIKey, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-6-astra": "deepseek-v4-pro", "apikey-only": "deepseek-v4-pro"}}},
{ID: 3, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"model_mapping": map[string]any{"apikey-only": "gpt-5"}}},
// Antigravity 账号创建时会把默认映射(含 claude-*)写入 model_mapping。
{ID: 4, Platform: PlatformAntigravity, Type: AccountTypeOAuth, Credentials: map[string]any{"model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}}},
{ID: 5, Platform: PlatformAnthropic, Type: AccountTypeSetupToken, Credentials: map[string]any{"model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}}},
},
}
svc := &GatewayService{accountRepo: repo}
withOAuth, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "gpt-6-astra")
require.NoError(t, err)
require.Equal(t, CompositeModelOwnership{Ambiguous: true, OpenAIGatewayOAuthClaimed: true}, withOAuth)
apiKeyOnly, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "apikey-only")
require.NoError(t, err)
require.Equal(t, CompositeModelOwnership{Ambiguous: true}, apiKeyOnly)
anthropicFamily, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "claude-opus-4-8")
require.NoError(t, err)
require.Equal(t, CompositeModelOwnership{Ambiguous: true}, anthropicFamily)
}
func TestNewGatewayServiceWiresCompositeModelOwnershipResolver(t *testing.T) {
groupID := int64(7)
repo := &compositeOwnershipAccountRepo{
@@ -5,8 +5,34 @@ import (
"fmt"
"sort"
"strings"
"github.com/Wei-Shaw/sub2api/internal/domain"
)
const (
// CompositeRouteReasonAmbiguous marks a model whose explicit account
// mappings span several provider platforms.
CompositeRouteReasonAmbiguous = "model is exposed by multiple provider platforms"
// CompositeRouteReasonOAuthAmbiguous is the ambiguous case where a claimant
// holds an OAuth / setup-token credential that only the OpenAI gateway can
// carry, and the model-name fallback would also pick an OpenAI-gateway
// platform. The unresolved request would land in the Anthropic gateway
// handler and select those credentials, so this group misconfiguration must
// be rejected instead of guessing a platform from the model name.
CompositeRouteReasonOAuthAmbiguous = "model is exposed by multiple provider platforms including OAuth accounts"
)
// OAuthAmbiguous reports whether the request must be rejected because the
// cross-platform ownership involves OpenAI-gateway OAuth credentials.
func (d CompositeRouteDecision) OAuthAmbiguous() bool {
return !d.Matched && d.Reason == CompositeRouteReasonOAuthAmbiguous
}
// AmbiguousMessage is the client-facing rejection for an OAuthAmbiguous decision.
func (d CompositeRouteDecision) AmbiguousMessage() string {
return fmt.Sprintf("Model %q is mapped on multiple provider platforms in this group; configure an explicit composite route", d.PublicModel)
}
type CompositeRouteResolver struct {
repo CompositeModelRouteRepository
modelOwnershipResolver CompositeModelOwnershipResolver
@@ -113,7 +139,12 @@ func (r *CompositeRouteResolver) Resolve(ctx context.Context, groupID int64, mod
return decision, fmt.Errorf("resolve account model ownership: %w", err)
}
} else if ownership.Ambiguous {
decision.Reason = "model is exposed by multiple provider platforms"
decision.Reason = CompositeRouteReasonAmbiguous
// 未解析的请求由下游 handler 按模型名前缀兜底选平台;只有兜底落在
// OpenAI 网关族时,OpenAI 网关 OAuth 声明方才会被误送进 Anthropic 管线。
if guessed, ok := DetectModelPlatform(model); ok && ownership.OpenAIGatewayOAuthClaimed && domain.UsesOpenAIGateway(guessed) {
decision.Reason = CompositeRouteReasonOAuthAmbiguous
}
return decision, nil
} else if ownership.Matched {
platform := strings.TrimSpace(ownership.TargetPlatform)
@@ -221,6 +221,33 @@ func TestCompositeRouteResolverDoesNotGuessAmbiguousAccountOwnership(t *testing.
require.False(t, decision.Matched)
require.Empty(t, decision.TargetPlatform)
require.Equal(t, "model is exposed by multiple provider platforms", decision.Reason)
require.False(t, decision.OAuthAmbiguous())
}
// Scenario: 含 OAuth 账号的跨平台歧义必须被识别为需拒绝的配置错误
func TestCompositeRouteResolverFlagsOAuthAmbiguousOwnership(t *testing.T) {
resolver := NewCompositeRouteResolver(nil)
resolver.SetModelOwnershipResolver(func(context.Context, int64, string) (CompositeModelOwnership, error) {
return CompositeModelOwnership{Ambiguous: true, OpenAIGatewayOAuthClaimed: true}, nil
})
// gpt-* 可被内置探测器识别,验证歧义不会退化为按模型名猜平台。
decision, err := resolver.Resolve(context.Background(), 7, "gpt-6-astra", CompositeRouteEndpointResponses)
require.NoError(t, err)
require.False(t, decision.Matched)
require.Empty(t, decision.TargetPlatform)
require.True(t, decision.OAuthAmbiguous())
require.Contains(t, decision.AmbiguousMessage(), `"gpt-6-astra"`)
// 兜底按模型名落到 Anthropic(如 OpenAI OAuth 映射 claude-* 给 Claude Code 用)
// 或无法兜底的自定义别名:OpenAI 网关 OAuth 不会被选中,保持原有行为。
for _, model := range []string{"claude-sonnet-4-6", "team-alias"} {
decision, err := resolver.Resolve(context.Background(), 7, model, CompositeRouteEndpointMessages)
require.NoError(t, err)
require.False(t, decision.OAuthAmbiguous(), "model=%s", model)
require.Equal(t, CompositeRouteReasonAmbiguous, decision.Reason, "model=%s", model)
}
}
func TestCompositeRouteResolverOwnershipLookupErrorFallsBackOnlyForDetectableModels(t *testing.T) {
@@ -22,6 +22,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
s.countTokensError(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
return fmt.Errorf("parse request: empty request")
}
// Antigravity 账户不支持 count_tokens,返回 404 让客户端 fallback 到本地估算。
// 返回 nil 避免 handler 层记录为错误,也不设置 ops 上游错误上下文。
// 凭据只能走其他协议的 OAuth 账号同理,且必须在取 token、建请求之前拦下。
if account != nil && (account.Platform == PlatformAntigravity || !anthropicMessagesCredentialAllowed(account)) {
s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for this platform")
return nil
}
validationModel := parsed.Model
if account != nil {
@@ -105,13 +112,6 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
}
// Antigravity 账户不支持 count_tokens,返回 404 让客户端 fallback 到本地估算。
// 返回 nil 避免 handler 层记录为错误,也不设置 ops 上游错误上下文。
if account.Platform == PlatformAntigravity {
s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for this platform")
return nil
}
// 应用模型映射:
// - APIKey 账号:使用账号级别的显式映射(如果配置),否则透传原始模型名
// - OAuth/SetupToken 账号:使用 Anthropic 标准映射(短ID → 长ID)
@@ -180,6 +180,10 @@ func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context
return out, nil
}
func (s *modelsListAccountRepoStub) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, _ []string, _ bool) ([]Account, error) {
return s.ListSchedulableByGroupID(ctx, *groupID)
}
func (s *modelsListAccountRepoStub) ListSchedulable(ctx context.Context) ([]Account, error) {
s.listAllCalls.Add(1)
if s.err != nil {
+9 -6
View File
@@ -20,6 +20,7 @@ import (
"unsafe"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
@@ -1501,21 +1502,23 @@ func (s *GatewayService) resolveCompositeModelOwnership(ctx context.Context, gro
}
}
accounts, err := s.accountRepo.ListSchedulableByGroupID(ctx, groupID)
// 归属取决于配置而非运行时健康:限流/临时不可调度不应让同一模型在
// 「唯一归属」与「歧义」之间来回翻转,否则配置错误只会间歇暴露。
accounts, err := s.accountRepo.ListModelAvailabilityCandidates(ctx, &groupID, domain.ConcretePlatformIDs(), false)
if err != nil {
return CompositeModelOwnership{}, err
}
platforms := make(map[string]struct{})
ownership := CompositeModelOwnership{}
for _, account := range accounts {
platform := strings.TrimSpace(account.Platform)
if !isConcreteRequestPlatform(platform) || !explicitModelMappingClaims(account, model) {
continue
if explicitModelMappingClaims(account, model) {
platforms[account.Platform] = struct{}{}
ownership.OpenAIGatewayOAuthClaimed = ownership.OpenAIGatewayOAuthClaimed ||
(account.IsOAuth() && domain.UsesOpenAIGateway(account.Platform))
}
platforms[platform] = struct{}{}
}
ownership := CompositeModelOwnership{}
if len(platforms) == 1 {
for platform := range platforms {
ownership.TargetPlatform = platform
@@ -9,6 +9,7 @@ import (
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
@@ -18,7 +19,37 @@ import (
"github.com/gin-gonic/gin"
)
// anthropicMessagesForeignOAuthReason 标记账号凭据与 Anthropic Messages 协议不匹配。
const anthropicMessagesForeignOAuthReason = GatewayFailureReason("anthropic_messages_foreign_oauth")
// anthropicMessagesCredentialAllowed 判断账号凭据能否经 Anthropic Messages 协议
// 出站。OAuth / setup-token 只对签发它的供应商有效(ChatGPT token 只认
// chatgpt.com),送错协议既泄露凭据,上游 401 还会被误记到账号头上;
// API Key 不受限:其上游可能同时兼容多种协议格式。
func anthropicMessagesCredentialAllowed(account *Account) bool {
if !account.IsOAuth() {
return true
}
spec, ok := domain.LookupPlatform(account.Platform)
return ok && spec.Gateway == domain.PlatformGatewayAnthropic
}
// checkAnthropicMessagesCredential 在发出任何上游请求前拦截不匹配的凭据。返回的
// failover 错误只让 handler 换下一个账号:没有上游响应,也就不触发任何账号处罚。
func checkAnthropicMessagesCredential(account *Account) error {
if anthropicMessagesCredentialAllowed(account) {
return nil
}
return &UpstreamFailoverError{
StatusCode: http.StatusBadGateway,
Reason: anthropicMessagesForeignOAuthReason,
}
}
func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) {
if err := checkAnthropicMessagesCredential(account); err != nil {
return nil, nil, err
}
body = stripDeferredToolCacheControl(body)
if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount {
req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream)
@@ -0,0 +1,73 @@
//go:build unit
package service
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// Scenario: OAuth / setup-token 凭据只能走签发方协议,API Key 与其他凭据类型不受限
func TestAnthropicMessagesCredentialAllowed(t *testing.T) {
cases := []struct {
platform string
typ string
allowed bool
}{
{PlatformAnthropic, AccountTypeOAuth, true},
{PlatformAnthropic, AccountTypeSetupToken, true},
{PlatformAnthropic, AccountTypeServiceAccount, true},
{PlatformAnthropic, AccountTypeBedrock, true},
{PlatformTypeSafe, AccountTypeOAuth, true},
{PlatformOpenAI, AccountTypeAPIKey, true},
{PlatformDeepseek, AccountTypeAPIKey, true},
{PlatformOpenAI, AccountTypeOAuth, false},
{PlatformOpenAI, AccountTypeSetupToken, false},
{PlatformGrok, AccountTypeOAuth, false},
{PlatformGemini, AccountTypeOAuth, false},
{PlatformAntigravity, AccountTypeOAuth, false},
}
for _, tc := range cases {
allowed := anthropicMessagesCredentialAllowed(&Account{ID: 1, Platform: tc.platform, Type: tc.typ})
require.Equal(t, tc.allowed, allowed, "%s/%s", tc.platform, tc.typ)
}
}
// Scenario: 回归 —— OpenAI OAuth 账号误入 Anthropic 管线时,在构建请求前即被拒绝
// (ChatGPT token 不会发往 api.anthropic.com),且只让 handler 换号、不处罚账号
func TestGatewayBuildUpstreamRequestRejectsForeignOAuthWithoutPenalty(t *testing.T) {
account := &Account{ID: 25, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
body := []byte(`{"model":"gpt-6-astra","messages":[]}`)
req, _, err := (&GatewayService{}).buildUpstreamRequest(context.Background(), nil, account, body, "chatgpt-token", "oauth", "gpt-6-astra", true, false)
require.Nil(t, req)
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr))
require.True(t, failoverErr.ShouldRetryNextAccount())
require.False(t, failoverErr.RetryableOnSameAccount, "same-account retry would end in temp-unschedulable")
require.False(t, failoverErr.IsCredentialFailure())
require.Equal(t, anthropicMessagesForeignOAuthReason, failoverErr.Reason)
}
// Scenario: count_tokens 对不匹配凭据按「平台不支持」返回 404,不取 token、不出站
func TestForwardCountTokensRejectsForeignOAuthAsUnsupported(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
account := &Account{ID: 25, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
parsed := &ParsedRequest{Model: "gpt-6-astra", Body: NewRequestBodyRef([]byte(`{"model":"gpt-6-astra","messages":[]}`))}
err := (&GatewayService{}).ForwardCountTokens(context.Background(), c, account, parsed)
require.NoError(t, err)
require.Equal(t, http.StatusNotFound, w.Code)
}
@@ -43,7 +43,10 @@ func TestBuiltinGeminiCacheStoragePrice(t *testing.T) {
{"gemini-flash-latest", 1e-6},
{"gemini-flash-lite-latest", 1e-6},
{"gemini-pro-latest", 4.5e-6},
// 无存储价
// 无存储价(Gemini 4 Argon 仅支持隐式缓存)
{"gemini-4-argon", 0},
{"gemini-4-argon-high", 0},
{"gemini-4-argon-cyber-permissive", 0},
{"gemini-embedding-001", 0},
{"gemini-unknown-model", 0},
{"claude-sonnet-5-5", 0},
@@ -112,12 +112,17 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
)
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, err
}
if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil {
return nil, failoverErr
}
return nil, s.writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", err.Error())
}
requestIDHeader = idHeader
@@ -127,11 +132,11 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
return nil, s.handleUpstreamTransportError(ctx, c, account, err)
}
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -196,7 +201,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
respBody := s.readUpstreamErrorBody(resp)
policy := ErrorPolicyNone
if s.rateLimitService != nil {
policy = s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy = s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
}
// 与 messages 兼容层一致:只有 None / Matched 才走账号状态处理。
// Skipped(池模式、或自定义错误码未命中)与 TempUnscheduled 已由策略层裁决完毕。
@@ -0,0 +1,201 @@
//go:build unit
package service
import (
"bytes"
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// Gemini 三条转发路径对每个上游错误响应只检查一次错误策略:
// - 重试循环内已检查的最终响应,循环外的错误处理复用该结果;
// - 账号级 401 规则首次命中只临时不可调度,账号上已有 401 暂停记录时才升级为 error;
// - 循环内未检查的最终响应(如 Claude 兼容层的非签名 400)在循环外检查。
// ---------------------------------------------------------------------------
type geminiSingleCheckRepoStub struct {
errorPolicyRepoStub
tempReason string
}
func (r *geminiSingleCheckRepoStub) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
r.tempCalls++
r.tempReason = reason
return nil
}
func (r *geminiSingleCheckRepoStub) GetByID(_ context.Context, id int64) (*Account, error) {
return &Account{ID: id, TempUnschedulableReason: r.tempReason}, nil
}
type geminiSequenceUpstreamResponse struct {
status int
body string
}
type geminiSequenceUpstreamStub struct {
responses []geminiSequenceUpstreamResponse
calls int
}
func (s *geminiSequenceUpstreamStub) Do(*http.Request, string, int64, int) (*http.Response, error) {
r := s.responses[min(s.calls, len(s.responses)-1)]
s.calls++
return &http.Response{
StatusCode: r.status,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(r.body)),
}, nil
}
func (s *geminiSequenceUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return s.Do(req, proxyURL, accountID, concurrency)
}
func newGeminiSingleCheckService(responses ...geminiSequenceUpstreamResponse) (*GeminiMessagesCompatService, *geminiSequenceUpstreamStub, *geminiSingleCheckRepoStub) {
repo := &geminiSingleCheckRepoStub{}
rateLimit := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
upstream := &geminiSequenceUpstreamStub{responses: responses}
svc := &GeminiMessagesCompatService{
accountRepo: repo,
httpUpstream: upstream,
rateLimitService: rateLimit,
cfg: &config.Config{},
}
return svc, upstream, repo
}
func geminiAPIKeyAccountWithTempRule(rule map[string]any) *Account {
return &Account{
ID: 830,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"api_key": "test-key",
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{rule},
},
}
}
const geminiSingleCheck401Body = `{"error":{"code":401,"message":"API key expired. Please renew the API key.","status":"UNAUTHENTICATED"}}`
func geminiAccountScoped401Rule() map[string]any {
return map[string]any{
"error_code": float64(401),
"keywords": []any{"expired"},
"duration_minutes": float64(30),
"scope": TempUnschedScopeAccount,
}
}
// geminiSingleCheckForwarders 以三条转发路径各发一次请求,返回转发错误与 gin.Context。
var geminiSingleCheckForwarders = []struct {
name string
forward func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context)
}{
{
name: "native",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := newGeminiNativeTestContext(t)
_, err := svc.ForwardNative(context.Background(), c, account, "gemini-2.5-flash", "generateContent", false,
[]byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
return err, c
},
},
{
name: "claude_compat",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
_, err := svc.Forward(context.Background(), c, account, body)
return err, c
},
},
{
name: "chat_completions_compat",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
_, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
return err, c
},
},
}
func TestGeminiForward_First401RuleHitOnlyTempUnschedules(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, fw := range geminiSingleCheckForwarders {
t.Run(fw.name, func(t *testing.T) {
svc, upstream, repo := newGeminiSingleCheckService(geminiSequenceUpstreamResponse{status: http.StatusUnauthorized, body: geminiSingleCheck401Body})
err, _ := fw.forward(t, svc, geminiAPIKeyAccountWithTempRule(geminiAccountScoped401Rule()))
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr), "规则命中应换号,got %T: %v", err, err)
require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
require.Equal(t, 1, upstream.calls)
require.Equal(t, 1, repo.tempCalls, "首次 401 命中规则应临时不可调度")
require.Zero(t, repo.setErrCalls, "首次 401 命中规则不应升级为 error")
})
}
}
func TestGeminiForward_Repeated401RuleHitEscalatesToError(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, fw := range geminiSingleCheckForwarders {
t.Run(fw.name, func(t *testing.T) {
svc, _, repo := newGeminiSingleCheckService(geminiSequenceUpstreamResponse{status: http.StatusUnauthorized, body: geminiSingleCheck401Body})
repo.tempReason = `{"until_unix":1,"triggered_at_unix":1,"status_code":401,"matched_keyword":"expired","rule_index":0,"scope":"account"}`
_, _ = fw.forward(t, svc, geminiAPIKeyAccountWithTempRule(geminiAccountScoped401Rule()))
require.Zero(t, repo.tempCalls, "已有 401 暂停记录时直接升级,不临时不可调度")
require.Equal(t, 1, repo.setErrCalls, "已有 401 暂停记录时应升级为 error")
require.Contains(t, repo.lastErrorMsg, "Authentication failed (401)")
})
}
}
func TestGeminiForward_UncheckedFinalResponseIsCheckedAfterLoop(t *testing.T) {
gin.SetMode(gin.TestMode)
if testing.Short() {
t.Skip("包含一次 1s 的重试退避")
}
// 第一次 503 在循环内检查后重试;第二次非签名 400 走 Claude 兼容层的签名分支,循环内不检查。
svc, upstream, repo := newGeminiSingleCheckService(
geminiSequenceUpstreamResponse{status: http.StatusServiceUnavailable, body: `{"error":{"code":503,"message":"backend busy","status":"UNAVAILABLE"}}`},
geminiSequenceUpstreamResponse{status: http.StatusBadRequest, body: `{"error":{"code":400,"message":"quota project misconfigured","status":"FAILED_PRECONDITION"}}`},
)
account := geminiAPIKeyAccountWithTempRule(map[string]any{
"error_code": float64(400),
"keywords": []any{"misconfigured"},
"duration_minutes": float64(30),
"scope": TempUnschedScopeAccount,
})
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
_, err := svc.Forward(context.Background(), c, account, body)
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr), "400 规则命中应换号,got %T: %v", err, err)
require.Equal(t, 2, upstream.calls)
require.Equal(t, 1, repo.tempCalls, "循环内未检查的最终响应应在循环外检查")
}
@@ -631,30 +631,159 @@ func TestHandleGeminiUpstreamError_VertexRetryInfoUsesParsedDelay(t *testing.T)
require.WithinDuration(t, before.Add(39*time.Second), repo.lastRateLimitReset, 2*time.Second)
}
// Vertex(service_account)429 响应不带任何可解析的重置时间时,走短冷却兜底,
// 不再冷却到 PST 午夜(按量付费没有每日配额,长冷却会导致整组不可用)。
func TestHandleGeminiUpstreamError_VertexFallbackUsesShortCooldown(t *testing.T) {
// Vertex(service_account)共享容量不足的 429 不写账号级限流:账号保持可调度,
// 由请求内的同号退避重试与换号消化。
func TestHandleGeminiUpstreamError_VertexSharedCapacityKeepsAccountSchedulable(t *testing.T) {
tests := []struct {
name string
body string
}{
{
name: "object",
body: `{"error":{"code":429,"message":"Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.","status":"RESOURCE_EXHAUSTED"}}`,
},
{
name: "array_wrapped",
body: `[{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED"}}]`,
},
{
name: "empty_details",
body: `{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED","details":[]}}`,
},
{
name: "documented_comma_wording",
body: `{"error":{"code":429,"message":"Resource exhausted, please try again later.","status":"RESOURCE_EXHAUSTED"}}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 602,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"project_id": "my-vertex-project",
},
}
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, []byte(tt.body))
require.Zero(t, repo.rateLimitCalls)
})
}
}
// Vertex(service_account)项目配额、预留吞吐等账号自身的 429 不带可解析的重置时间时,
// 走短冷却兜底,不冷却到 PST 午夜(按量付费没有每日配额,长冷却会导致整组不可用)。
func TestHandleGeminiUpstreamError_VertexAccountScoped429UsesShortCooldown(t *testing.T) {
tests := []struct {
name string
body string
}{
{
name: "project_quota",
body: `{"error":{"code":429,"message":"Quota exceeded for aiplatform.googleapis.com/generate_content_requests_per_minute_per_project_per_base_model with base model: gemini-2.5-pro. Please submit a quota increase request.","status":"RESOURCE_EXHAUSTED"}}`,
},
{
name: "provisioned_throughput",
body: `{"error":{"code":429,"message":"Too many requests. Exceeded the Provisioned Throughput.","status":"RESOURCE_EXHAUSTED"}}`,
},
{
name: "quota_failure_details",
body: `{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"subject":"project"}]}]}}`,
},
{
name: "other_status",
body: `{"error":{"code":429,"message":"Too many requests.","status":"UNAVAILABLE"}}`,
},
{
name: "unknown_wording",
body: `{"error":{"code":429,"message":"Rate limit exceeded for requests per minute per project.","status":"RESOURCE_EXHAUSTED"}}`,
},
{
name: "missing_message",
body: `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED"}}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 604,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"project_id": "my-vertex-project",
},
}
before := time.Now()
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, []byte(tt.body))
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(604), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(geminiVertexFallbackCooldown), repo.lastRateLimitReset, 2*time.Second)
// 显式守护:不应被冷却超过 15 分钟(更不应到次日 PST 午夜)
require.True(t, repo.lastRateLimitReset.Before(before.Add(15*time.Minute)))
})
}
}
// Vertex 共享容量形态的 429 只对 Vertex 账号豁免冷却:AI Studio 官方 API Key 与 Code Assist OAuth
// 收到同样的响应体仍按各自语义冷却。
func TestHandleGeminiUpstreamError_SharedCapacityBodyStillCoolsNonVertexAccounts(t *testing.T) {
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.","status":"RESOURCE_EXHAUSTED"}}`)
tests := []struct {
name string
account *Account
}{
{
name: "ai_studio_api_key",
account: &Account{ID: 605, Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "k"}},
},
{
name: "code_assist_oauth",
account: &Account{ID: 606, Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{
"oauth_type": "code_assist", "project_id": "my-code-assist-project",
}},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
svc.handleGeminiUpstreamError(context.Background(), tt.account, http.StatusTooManyRequests, http.Header{}, body)
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, tt.account.ID, repo.lastRateLimitID)
})
}
}
// Vertex 429 文案里给出了重试时长时按该时长冷却,不归为共享容量不足。
func TestHandleGeminiUpstreamError_VertexMessageRetryHintUsesParsedDelay(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 602,
ID: 607,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"project_id": "my-vertex-project",
},
}
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED"}}`)
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please retry in 5s","status":"RESOURCE_EXHAUSTED"}}`)
before := time.Now()
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(602), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(geminiVertexFallbackCooldown), repo.lastRateLimitReset, 2*time.Second)
// 显式守护:不应被冷却超过 15 分钟(更不应到次日 PST 午夜)
require.True(t, repo.lastRateLimitReset.Before(before.Add(15*time.Minute)))
require.Equal(t, int64(607), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(5*time.Second), repo.lastRateLimitReset, 2*time.Second)
}
// API Key(AI Studio)无法解析重置时间时仍冷却到 PST 午夜,行为保持不变。
@@ -45,8 +45,8 @@ const (
geminiRetryInfoTypeURL = "type.googleapis.com/google.rpc.RetryInfo"
// RetryInfo.retryDelay 的上限,防止异常值导致账号长时间不可用
geminiRetryInfoMaxDelay = 15 * time.Minute
// Vertex(service_account,按量付费)429 无法解析重置时间时的兜底冷却:
// 按量付费限速是短窗口 RPM/TPM,没有每日配额,不适用 PST 午夜重置
// Vertex(service_account)账号自身的 429(按量付费的项目配额、预留吞吐超额等)无法解析重置时间时的兜底冷却:
// 这类限速是短窗口 RPM/TPM,没有每日配额,不适用 PST 午夜重置;共享容量不足的 429 不冷却
geminiVertexFallbackCooldown = time.Minute
)
@@ -709,13 +709,19 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
}
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
signatureRetryStage := 0
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
checkedPolicy = nil
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, err
}
if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil {
return nil, failoverErr
}
// Local build error: don't retry.
if strings.Contains(err.Error(), "missing project_id") {
return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
@@ -800,11 +806,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
}
// 错误策略优先:匹配则跳过重试直接处理。
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -820,7 +826,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
break
}
if resp.StatusCode == 429 {
// Mark as rate-limited early so concurrent requests avoid this account.
// Apply account-scoped 429 handling early so concurrent requests avoid a rate-limited account.
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
}
if attempt < geminiMaxRetries {
@@ -872,7 +878,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
respBody := s.readUpstreamErrorBody(resp)
// 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil {
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy := s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
switch policy {
case ErrorPolicySkipped:
upstreamReqID := resp.Header.Get(requestIDHeader)
@@ -1239,12 +1245,18 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
checkedPolicy = nil
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, err
}
if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil {
return nil, failoverErr
}
// Local build error: don't retry.
if strings.Contains(err.Error(), "missing project_id") {
return nil, s.writeGoogleError(c, http.StatusBadRequest, err.Error())
@@ -1284,11 +1296,11 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
// 错误策略优先:匹配则跳过重试直接处理。
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -1404,7 +1416,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
// 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil {
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy := s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
switch policy {
case ErrorPolicySkipped:
if failoverErr := s.skippedErrorPolicyFailoverError(c, account, resp.StatusCode, respBody, requestID); failoverErr != nil {
@@ -1573,14 +1585,14 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}, nil
}
// checkErrorPolicyInLoop 在重试循环内预检查错误策略。
// 返回 true 表示策略已匹配(调用者应 break),resp 已重建可直接使用。
// 返回 false 表示 ErrorPolicyNone,resp 已重建,调用者继续走重试逻辑。
// checkErrorPolicyInLoop 在重试循环内对上游错误响应检查错误策略,resp 已重建可直接使用。
// checked 为 nil 表示未检查(非错误响应或未配置限流服务);否则指向检查结果,
// 非 ErrorPolicyNone 时调用者应 break 并由循环外的错误处理复用该结果。
func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
ctx context.Context, account *Account, resp *http.Response, mappedModel string,
) (matched bool, rebuilt *http.Response) {
) (checked *ErrorPolicyResult, rebuilt *http.Response) {
if resp.StatusCode < 400 || s.rateLimitService == nil {
return false, resp
return nil, resp
}
body := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
@@ -1590,7 +1602,20 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
Body: io.NopCloser(bytes.NewReader(body)),
}
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, body, mappedModel)
return policy != ErrorPolicyNone, rebuilt
return &policy, rebuilt
}
// finalErrorPolicy 返回循环外处理最终上游错误响应时使用的错误策略。每个上游响应只检查一次:
// 错误策略检查会触发临时不可调度规则,并据账号上已记录的 401 暂停把 401 升级为 error。
// checked 非 nil 时它就是重试循环内对同一响应的检查结果,直接复用;否则在此检查。
// 调用方保证 s.rateLimitService 非 nil。
func (s *GeminiMessagesCompatService) finalErrorPolicy(
ctx context.Context, account *Account, statusCode int, body []byte, mappedModel string, checked *ErrorPolicyResult,
) ErrorPolicyResult {
if checked != nil {
return *checked
}
return s.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, body, mappedModel)
}
// cachedContentPassthroughInLoop 判断上游响应是否跳过同号重试与错误策略直接处理:引用缓存不存在,
@@ -3026,7 +3051,13 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Google One OAuth, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
}
} else if account.IsVertexServiceAccount() {
// Vertex 按量付费:短窗限速,无每日配额,短冷却后即可恢复
// 共享容量不足描述的是请求当时的上游容量,不是账号额度:账号保持可调度,
// 由请求内的同号退避重试与换号消化。
if isVertexSharedCapacity429(body) {
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Vertex service account, project=%s) shared capacity exhausted, account stays schedulable", account.ID, projectID)
return
}
// 项目配额、预留吞吐等账号自身的限速是短窗口 RPM/TPM,无每日配额,短冷却后即可恢复
ra = time.Now().Add(geminiVertexFallbackCooldown)
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Vertex service account, project=%s) rate limited, cooldown=%v", account.ID, projectID, time.Until(ra).Truncate(time.Second))
} else {
@@ -3051,6 +3082,30 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
account.ID, resetTime, oauthType, tierID)
}
// isVertexSharedCapacity429 判断 Vertex 429 是否为按量付费共享容量不足:RESOURCE_EXHAUSTED、
// 不带任何错误明细(RetryInfo / QuotaFailure / ErrorInfo),文案为 "Resource exhausted ... try again later"
// 且不涉及项目配额或预留吞吐。其他形态的 429 按账号自身限速处理。
func isVertexSharedCapacity429(body []byte) bool {
errObj := gjson.GetBytes(body, "error")
if !errObj.Exists() {
errObj = gjson.GetBytes(body, "0.error")
}
if !errObj.IsObject() {
return false
}
if !strings.EqualFold(strings.TrimSpace(errObj.Get("status").String()), "RESOURCE_EXHAUSTED") {
return false
}
if len(errObj.Get("details").Array()) > 0 {
return false
}
msg := strings.ToLower(errObj.Get("message").String())
if !strings.Contains(msg, "resource exhausted") || !strings.Contains(msg, "try again later") {
return false
}
return !strings.Contains(msg, "quota") && !strings.Contains(msg, "provisioned throughput")
}
// handleGeminiThirdParty429 处理第三方上游的 429:不套用 AI Studio 每日配额语义,
// 只认响应里显式给出的重试时间,否则走全局 429 兜底冷却。
func (s *GeminiMessagesCompatService) handleGeminiThirdParty429(ctx context.Context, account *Account, body []byte) {
@@ -12,6 +12,7 @@ import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
@@ -99,3 +100,66 @@ func TestGeminiForwardAsChatCompletions_VertexAPIKeyUsesAPIKeyHeader(t *testing.
require.Equal(t, "https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-3-flash-preview:generateContent", httpStub.lastReq.URL.String())
requireVertexAPIKeyAuth(t, httpStub.lastReq)
}
type geminiVertex429SequenceUpstream struct {
responses []*http.Response
calls int
}
func (s *geminiVertex429SequenceUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
resp := s.responses[min(s.calls, len(s.responses)-1)]
s.calls++
return resp, nil
}
func (s *geminiVertex429SequenceUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return s.Do(req, proxyURL, accountID, accountConcurrency)
}
// Vertex 共享容量不足的 429 在同一账号上退避重试,成功后账号状态不变;绑定显式缓存的请求同样如此。
func TestGeminiForwardNative_VertexSharedCapacity429RetriesSameAccountWithoutCooldown(t *testing.T) {
gin.SetMode(gin.TestMode)
sharedCapacity := `{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED"}}`
tests := []struct {
name string
ctx context.Context
}{
{name: "unbound", ctx: context.Background()},
{name: "bound_cached_content", ctx: WithGeminiCachedContentBound(context.Background())},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
upstream := &geminiVertex429SequenceUpstream{responses: []*http.Response{
{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(sharedCapacity)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(geminiVertexAPIKeyTestResponse)),
},
}}
repo := &geminiBoundCacheRateLimitRepo{}
svc := &GeminiMessagesCompatService{
httpUpstream: upstream,
cfg: &config.Config{},
accountRepo: repo,
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
}
c, rec := newGeminiNativeTestContext(t)
result, err := svc.ForwardNative(tt.ctx, c, geminiVertexAPIKeyTestAccount("my-project"),
"gemini-3.7-flash", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 2, upstream.calls, "429 后在同一账号上重试")
require.Zero(t, repo.rateLimitedCalls, "不写账号级限流")
require.Zero(t, repo.tempCalls)
require.Zero(t, repo.setErrCalls)
require.Empty(t, repo.modelRateLimitCalls)
})
}
}
@@ -0,0 +1,60 @@
package service
import (
"context"
"log/slog"
"net/http"
"github.com/gin-gonic/gin"
)
// vertexServiceAccountNotFoundReason 标识换取 SA access token 时 SA 已不存在。
const vertexServiceAccountNotFoundReason GatewayFailureReason = "vertex_service_account_not_found"
// geminiVertexCredentialFailoverBody 是 SA 凭据失效时挂在 failover 错误上的 Google 格式错误体,
// 供 handler 在换号耗尽后做透传规则匹配与错误信息提取。
var geminiVertexCredentialFailoverBody = []byte(`{"error":{"code":401,"message":"Upstream authentication failed","status":"UNAUTHENTICATED"}}`)
// handleVertexServiceAccountNotFound 处理 Gemini 三条转发路径(原生 / Claude 兼容 /
// Chat Completions 兼容)构造上游请求时换取 SA access token 失败、且 SA 已不存在的情况:
//
// - 账号置为 error 并停止调度,错误信息带上 token 端点的 error_description;
// - 记录凭据获取阶段(stage=account_auth)的 Ops 事件,token 端点的状态码只出现在 Message 中;
// - 返回 *UpstreamFailoverError(401),由 handler 的 failover 循环换号。
//
// 不是该类失败时返回 nil,由调用方继续处理。本函数不写响应。
func (s *GeminiMessagesCompatService) handleVertexServiceAccountNotFound(ctx context.Context, c *gin.Context, account *Account, err error) *UpstreamFailoverError {
tokenErr, ok := vertexServiceAccountNotFoundError(err)
if !ok {
return nil
}
msg := "Service account token rejected (400): " + tokenErr.Description
stateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIAccountStateUpdateTimeout)
defer cancel()
if s.rateLimitService != nil {
s.rateLimitService.handleAuthError(stateCtx, account, msg)
} else if s.accountRepo != nil {
if setErr := s.accountRepo.SetError(stateCtx, account.ID, msg); setErr != nil {
slog.Warn("account_set_error_failed", "account_id", account.ID, "error", setErr)
}
}
event := OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
Kind: "failover",
Stage: string(GatewayFailureStageAccountAuth),
Scope: string(GatewayFailureScopeAccount),
Reason: string(vertexServiceAccountNotFoundReason),
Message: sanitizeUpstreamErrorMessage(err.Error()),
}
event.ProxyID, event.ProxyName = opsUpstreamProxyAttribution(account)
appendOpsUpstreamError(c, event)
return &UpstreamFailoverError{
StatusCode: http.StatusUnauthorized,
ResponseBody: geminiVertexCredentialFailoverBody,
}
}
@@ -0,0 +1,207 @@
//go:build unit
package service
import (
"bytes"
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/json"
"encoding/pem"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// 换取 SA access token 时 SA 已不存在(token 端点 400 invalid_grant
// "Invalid grant: account not found")的契约,三条转发路径一致:
// - 账号置为 error 并停止调度,错误信息带 token 端点的 error_description;
// - 不发出上游请求,返回 *UpstreamFailoverError(401) 交给 handler 换号,service 不写响应;
// - Ops 事件记为凭据获取阶段(stage=account_auth);
// - 其他换 token 失败不改账号状态、不换号,直接回写 502。
// ---------------------------------------------------------------------------
const vertexAccountNotFoundTokenBody = `{"error":"invalid_grant","error_description":"Invalid grant: account not found"}`
func newVertexTokenTestAccount(t *testing.T) *Account {
t.Helper()
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
pemBytes := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)})
saJSON, err := json.Marshal(map[string]string{
"type": "service_account",
"project_id": "proj-1",
"private_key_id": "kid-1",
"private_key": string(pemBytes),
"client_email": "svc@proj-1.iam.gserviceaccount.com",
})
require.NoError(t, err)
return &Account{
ID: 820,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"service_account_json": string(saJSON),
"project_id": "proj-1",
"location": "global",
},
}
}
func useVertexTokenEndpoint(t *testing.T, status int, body string) {
t.Helper()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_, _ = w.Write([]byte(body))
}))
t.Cleanup(server.Close)
original := vertexServiceAccountTokenURL
vertexServiceAccountTokenURL = server.URL
t.Cleanup(func() { vertexServiceAccountTokenURL = original })
}
func newVertexTokenErrorService() (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub, *errorPolicyRepoStub) {
repo := &errorPolicyRepoStub{}
httpStub := &geminiCompatHTTPUpstreamStub{}
svc := &GeminiMessagesCompatService{
accountRepo: repo,
httpUpstream: httpStub,
tokenProvider: NewGeminiTokenProvider(repo, nil, nil),
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
cfg: &config.Config{},
}
return svc, httpStub, repo
}
func requireVertexAccountNotFoundHandled(t *testing.T, c *gin.Context, err error, httpStub *geminiCompatHTTPUpstreamStub, repo *errorPolicyRepoStub) {
t.Helper()
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr), "SA 不存在应换号,got %T: %v", err, err)
require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
require.Equal(t, string(geminiVertexCredentialFailoverBody), string(failoverErr.ResponseBody))
require.True(t, failoverErr.ShouldRetryNextAccount())
require.False(t, failoverErr.RetryableOnSameAccount)
require.Equal(t, 1, repo.setErrCalls, "SA 不存在应把账号置为 error")
require.Equal(t, "Service account token rejected (400): Invalid grant: account not found", repo.lastErrorMsg)
require.Zero(t, httpStub.calls, "换不到 token 时不应发出上游请求")
require.False(t, c.Writer.Written(), "换号场景不应写客户端响应")
raw, ok := c.Get(OpsUpstreamErrorsKey)
require.True(t, ok, "应记录 ops 上游错误事件")
events := raw.([]*OpsUpstreamErrorEvent)
require.Len(t, events, 1)
require.Equal(t, "failover", events[0].Kind)
require.Equal(t, int64(820), events[0].AccountID)
require.Equal(t, string(GatewayFailureStageAccountAuth), events[0].Stage)
require.Equal(t, string(GatewayFailureScopeAccount), events[0].Scope)
require.Equal(t, string(vertexServiceAccountNotFoundReason), events[0].Reason)
require.Zero(t, events[0].UpstreamStatusCode, "换 token 失败不占用上游状态码")
require.Contains(t, events[0].Message, "service account token request returned 400: Invalid grant: account not found")
}
func TestExchangeVertexServiceAccountToken_AccountNotFoundIsRecognized(t *testing.T) {
useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody)
key, err := parseVertexServiceAccountKey(newVertexTokenTestAccount(t))
require.NoError(t, err)
_, _, err = exchangeVertexServiceAccountToken(contextWithTestTimeout(t), key, "")
require.EqualError(t, err, "service account token request returned 400: Invalid grant: account not found")
tokenErr, ok := vertexServiceAccountNotFoundError(fmt.Errorf("build request: %w", err))
require.True(t, ok)
require.Equal(t, "invalid_grant", tokenErr.Code)
require.Equal(t, "Invalid grant: account not found", tokenErr.Description)
}
func TestVertexServiceAccountNotFoundError_OnlyMatchesAccountNotFound(t *testing.T) {
tests := []struct {
name string
err error
}{
{"other_invalid_grant", &vertexTokenRequestError{StatusCode: 400, Code: "invalid_grant", Description: "Invalid JWT Signature."}},
{"not_400", &vertexTokenRequestError{StatusCode: 401, Code: "invalid_grant", Description: "Invalid grant: account not found"}},
{"other_code", &vertexTokenRequestError{StatusCode: 400, Code: "invalid_request", Description: "Invalid grant: account not found"}},
{"server_error", &vertexTokenRequestError{StatusCode: 500, Message: "backend error"}},
{"plain_error", errors.New("service account token request returned 400: Invalid grant: account not found")},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, ok := vertexServiceAccountNotFoundError(tt.err)
require.False(t, ok)
})
}
}
func TestGeminiForwardNative_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) {
gin.SetMode(gin.TestMode)
useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody)
svc, httpStub, repo := newVertexTokenErrorService()
c, _ := newGeminiNativeTestContext(t)
result, err := svc.ForwardNative(context.Background(), c, newVertexTokenTestAccount(t),
"gemini-3.8-flash", "streamGenerateContent", true, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
require.Nil(t, result)
requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo)
}
func TestGeminiForward_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) {
gin.SetMode(gin.TestMode)
useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody)
svc, httpStub, repo := newVertexTokenErrorService()
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
body := []byte(`{"model":"gemini-3.8-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
result, err := svc.Forward(context.Background(), c, newVertexTokenTestAccount(t), body)
require.Nil(t, result)
requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo)
}
func TestGeminiForwardAsChatCompletions_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) {
gin.SetMode(gin.TestMode)
useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody)
svc, httpStub, repo := newVertexTokenErrorService()
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
body := []byte(`{"model":"gemini-3.8-flash","messages":[{"role":"user","content":"hi"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
result, err := svc.ForwardAsChatCompletions(context.Background(), c, newVertexTokenTestAccount(t), body)
require.Nil(t, result)
requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo)
}
func TestGeminiForwardNative_OtherVertexTokenErrorKeepsAccountState(t *testing.T) {
gin.SetMode(gin.TestMode)
useVertexTokenEndpoint(t, http.StatusBadRequest, `{"error":"invalid_grant","error_description":"Invalid JWT Signature."}`)
svc, httpStub, repo := newVertexTokenErrorService()
c, rec := newGeminiNativeTestContext(t)
result, err := svc.ForwardNative(context.Background(), c, newVertexTokenTestAccount(t),
"gemini-3.8-flash", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
require.Nil(t, result)
require.Error(t, err)
var failoverErr *UpstreamFailoverError
require.False(t, errors.As(err, &failoverErr), "未识别的换 token 失败不换号")
require.Zero(t, repo.setErrCalls, "未识别的换 token 失败不改账号状态")
require.Zero(t, httpStub.calls)
require.Equal(t, http.StatusBadGateway, rec.Code)
require.True(t, strings.Contains(rec.Body.String(), "Invalid JWT Signature."))
}
@@ -48,6 +48,10 @@ func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupID(ctx context.Con
return append([]Account(nil), r.accounts...), nil
}
func (r schedulerTestOpenAIAccountRepo) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, _ []string, _ bool) ([]Account, error) {
return r.ListSchedulableByGroupID(ctx, *groupID)
}
func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
+5 -5
View File
@@ -1306,7 +1306,7 @@ func (s *PricingService) buildModelLookupCandidates(modelLower string) []string
normalized := normalizeModelNameForPricing(modelLower)
// A tier-specific entry should take precedence when the pricing catalog gains
// one later. Antigravity's Gemini Flash thinking tiers share the base rate,
// one later. Antigravity's Gemini thinking tiers share the base rate,
// so the normalized base remains the fallback after the exact aliases.
candidates := rawCandidates
if normalizeGeminiThinkingTierAlias(lastSegment(modelLower)) != lastSegment(modelLower) {
@@ -1368,10 +1368,10 @@ func normalizeModelNameForPricing(model string) string {
return normalizeGeminiThinkingTierAlias(model)
}
// normalizeGeminiThinkingTierAlias maps Antigravity's Gemini
// thinking-tier model IDs to the public base model. The tier controls reasoning
// behavior, not the published token rate, so this keeps -high/-low/-medium and
// -tiered requests on the corresponding base model's price card.
// normalizeGeminiThinkingTierAlias maps Antigravity's Gemini thinking-tier
// model IDs to the public base model. The tier controls reasoning behavior, not
// the published token rate, so this keeps -high/-low/-medium and -tiered
// requests on the corresponding base model's price card.
func normalizeGeminiThinkingTierAlias(model string) string {
for _, baseModel := range []string{"gemini-3.6-flash", "gemini-3.7-flash", "gemini-3.8-flash", "gemini-4-argon"} {
for _, tier := range []string{"-high", "-low", "-medium", "-tiered"} {
@@ -711,74 +711,6 @@ func TestBillingService_Gemini38FlashThinkingTierFallbacksAreBillable(t *testing
}
}
func TestPricingService_Gemini4ArgonThinkingTiersUseBaseModelPricing(t *testing.T) {
basePricing := &LiteLLMModelPricing{
InputCostPerToken: 2.0e-6,
OutputCostPerToken: 10.0e-6,
CacheReadInputTokenCost: 0.1e-6,
}
svc := &PricingService{pricingData: map[string]*LiteLLMModelPricing{
"gemini-4-argon": basePricing,
}}
for _, model := range []string{
"gemini-4-argon",
"gemini-4-argon-high",
"gemini-4-argon-low",
"gemini-4-argon-medium",
"gemini-4-argon-tiered",
} {
t.Run(model, func(t *testing.T) {
require.Same(t, basePricing, svc.GetModelPricing(model))
})
}
}
func TestBillingService_Gemini4ArgonThinkingTierFallbacksAreBillable(t *testing.T) {
svc := NewBillingService(&config.Config{}, nil)
tokens := UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}
for _, model := range []string{
"gemini-4-argon",
"gemini-4-argon-high",
"gemini-4-argon-low",
"gemini-4-argon-medium",
"gemini-4-argon-tiered",
} {
t.Run(model, func(t *testing.T) {
cost, err := svc.CalculateCost(model, tokens, 1)
require.NoError(t, err)
require.InDelta(t, 2.0, cost.InputCost, 1e-12)
require.InDelta(t, 10.0, cost.OutputCost, 1e-12)
require.InDelta(t, 0.1, cost.CacheReadCost, 1e-12)
require.InDelta(t, 12.1, cost.TotalCost, 1e-12)
})
}
}
func TestDefaultPricingIncludesGemini4ArgonRates(t *testing.T) {
data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json"))
require.NoError(t, err)
pricingSvc := &PricingService{}
pricingData, err := pricingSvc.parsePricingData(data)
require.NoError(t, err)
pricingSvc.pricingData = pricingData
for _, model := range []string{"gemini-4-argon", "gemini-4-argon-high", "gemini-4-argon-low", "gemini-4-argon-medium"} {
t.Run(model, func(t *testing.T) {
pricing := pricingSvc.GetModelPricing(model)
require.NotNil(t, pricing)
require.Equal(t, "gemini", pricing.LiteLLMProvider)
require.InDelta(t, 2e-6, pricing.InputCostPerToken, 1e-15)
require.InDelta(t, 10e-6, pricing.OutputCostPerToken, 1e-15)
require.InDelta(t, 0.1e-6, pricing.CacheReadInputTokenCost, 1e-15)
require.Zero(t, pricing.LongContextInputTokenThreshold)
})
}
require.Contains(t, pricingSvc.ListModelNamesByProvider("gemini"), "gemini-4-argon")
}
func TestDefaultPricingIncludesGemini36FlashRates(t *testing.T) {
data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json"))
require.NoError(t, err)
@@ -800,6 +732,93 @@ func TestDefaultPricingIncludesGemini36FlashRates(t *testing.T) {
}
}
func TestPricingService_Gemini4ArgonThinkingTiersUseBasePricing(t *testing.T) {
basePricing := &LiteLLMModelPricing{InputCostPerToken: 2e-6, OutputCostPerToken: 10e-6, CacheReadInputTokenCost: 0.1e-6}
cyberPricing := &LiteLLMModelPricing{InputCostPerToken: 4e-6, OutputCostPerToken: 20e-6, CacheReadInputTokenCost: 0.2e-6}
svc := &PricingService{pricingData: map[string]*LiteLLMModelPricing{
"gemini-4-argon": basePricing,
"gemini-4-argon-cyber-permissive": cyberPricing,
}}
for _, model := range []string{
"gemini-4-argon",
"gemini-4-argon-high",
"gemini-4-argon-low",
"gemini-4-argon-medium",
"models/gemini-4-argon-high",
"projects/p/locations/global/publishers/google/models/gemini-4-argon",
} {
t.Run(model, func(t *testing.T) {
require.Same(t, basePricing, svc.GetModelPricing(model))
})
}
require.Same(t, cyberPricing, svc.GetModelPricing("gemini-4-argon-cyber-permissive"))
}
func TestBillingService_Gemini4ArgonFallbacksAreBillable(t *testing.T) {
svc := NewBillingService(&config.Config{}, nil)
tokens := UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}
for _, tc := range []struct {
model string
input, output, read float64
}{
{"gemini-4-argon", 2, 10, 0.1},
{"gemini-4-argon-high", 2, 10, 0.1},
{"gemini-4-argon-low", 2, 10, 0.1},
{"gemini-4-argon-medium", 2, 10, 0.1},
{"models/gemini-4-argon", 2, 10, 0.1},
{"gemini-4-argon-cyber-permissive", 4, 20, 0.2},
{"gemini-4-argon-cyber", 4, 20, 0.2},
{"publishers/google/models/gemini-4-argon-cyber-permissive", 4, 20, 0.2},
} {
t.Run(tc.model, func(t *testing.T) {
cost, err := svc.CalculateCost(tc.model, tokens, 1)
require.NoError(t, err)
require.InDelta(t, tc.input, cost.InputCost, 1e-12)
require.InDelta(t, tc.output, cost.OutputCost, 1e-12)
require.InDelta(t, tc.read, cost.CacheReadCost, 1e-12)
require.InDelta(t, tc.input+tc.output+tc.read, cost.TotalCost, 1e-12)
})
}
}
func TestDefaultPricingIncludesGemini4ArgonRates(t *testing.T) {
data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json"))
require.NoError(t, err)
pricingSvc := &PricingService{}
pricingData, err := pricingSvc.parsePricingData(data)
require.NoError(t, err)
pricingSvc.pricingData = pricingData
for _, tc := range []struct {
model string
provider string
input, output, read float64
}{
{"gemini-4-argon", "gemini", 2e-6, 10e-6, 0.1e-6},
{"gemini-4-argon-high", "gemini", 2e-6, 10e-6, 0.1e-6},
{"gemini-4-argon-cyber-permissive", "vertex_ai-language-models", 4e-6, 20e-6, 0.2e-6},
} {
t.Run(tc.model, func(t *testing.T) {
pricing := pricingSvc.GetModelPricing(tc.model)
require.NotNil(t, pricing)
require.Equal(t, tc.provider, pricing.LiteLLMProvider)
require.InDelta(t, tc.input, pricing.InputCostPerToken, 1e-15)
require.InDelta(t, tc.output, pricing.OutputCostPerToken, 1e-15)
require.InDelta(t, tc.read, pricing.CacheReadInputTokenCost, 1e-15)
// Priority 档为标准价的 1.8 倍;单一价档不分提示词长度,显式缓存不可用故无存储价。
require.InDelta(t, tc.input*1.8, pricing.InputCostPerTokenPriority, 1e-15)
require.InDelta(t, tc.output*1.8, pricing.OutputCostPerTokenPriority, 1e-15)
require.InDelta(t, tc.read*1.8, pricing.CacheReadInputTokenCostPriority, 1e-15)
require.Zero(t, pricing.LongContextInputTokenThreshold)
require.Zero(t, pricing.CacheStorageCostPerTokenPerHour)
})
}
require.Contains(t, pricingSvc.ListModelNamesByProvider("gemini"), "gemini-4-argon")
}
func TestDefaultPricingUsesCurrentCodexAutoReviewBaseRates(t *testing.T) {
data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json"))
require.NoError(t, err)
@@ -39,6 +39,9 @@ var (
vertexAnthropicAlreadyDatedIDPattern = regexp.MustCompile(`^.+@[0-9]{8}$`)
)
// vertexServiceAccountTokenURL 是换取 SA access token 的端点,测试中替换为本地服务。
var vertexServiceAccountTokenURL = vertexDefaultTokenURL
type vertexServiceAccountKey struct {
Type string `json:"type"`
ProjectID string `json:"project_id"`
@@ -138,7 +141,7 @@ func parseVertexServiceAccountJSON(raw []byte) (*vertexServiceAccountKey, error)
return nil, errors.New("service account json missing project_id")
}
// Always use the well-known Google token endpoint to prevent SSRF via crafted token_uri.
key.TokenURI = vertexDefaultTokenURL
key.TokenURI = vertexServiceAccountTokenURL
return &key, nil
}
@@ -277,7 +280,12 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc
if msg == "" {
msg = string(bytes.TrimSpace(body))
}
return "", 0, fmt.Errorf("service account token request returned %d: %s", resp.StatusCode, msg)
return "", 0, &vertexTokenRequestError{
StatusCode: resp.StatusCode,
Code: strings.TrimSpace(parsed.Error),
Description: strings.TrimSpace(parsed.ErrorDesc),
Message: msg,
}
}
if strings.TrimSpace(parsed.AccessToken) == "" {
return "", 0, errors.New("service account token response missing access_token")
@@ -292,6 +300,33 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc
return parsed.AccessToken, ttl, nil
}
// vertexTokenRequestError 是 token 端点对换取 SA access token 请求返回的非 2xx 响应。
type vertexTokenRequestError struct {
StatusCode int
Code string // 响应的 error 字段,如 invalid_grant
Description string // 响应的 error_description 字段
Message string
}
func (e *vertexTokenRequestError) Error() string {
return fmt.Sprintf("service account token request returned %d: %s", e.StatusCode, e.Message)
}
// vertexServiceAccountNotFoundError 识别 SA 已不存在:token 端点返回
// 400 {"error":"invalid_grant","error_description":"Invalid grant: account not found"}。
// 该凭据再也换不到 token,账号需要人工更换凭据。
func vertexServiceAccountNotFoundError(err error) (*vertexTokenRequestError, bool) {
var tokenErr *vertexTokenRequestError
if !errors.As(err, &tokenErr) {
return nil, false
}
if tokenErr.StatusCode != http.StatusBadRequest || tokenErr.Code != "invalid_grant" ||
!strings.Contains(strings.ToLower(tokenErr.Description), "account not found") {
return nil, false
}
return tokenErr, true
}
// buildVertexGeminiAccountURL 按账号构造 Vertex Gemini 请求地址:未配置 project_id 的 API Key 账号
// 走 Express mode 全局端点,其余账号走项目级端点。
func buildVertexGeminiAccountURL(account *Account, model, action string, stream bool) (string, error) {
@@ -2683,6 +2683,23 @@
"supports_reasoning": true,
"supports_service_tier": true
},
"gemini-4-argon-cyber-permissive": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_priority": 3.6e-07,
"input_cost_per_token": 4e-06,
"input_cost_per_token_priority": 7.2e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 2e-05,
"output_cost_per_token_priority": 3.6e-05,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_service_tier": true
},
"gemini-embedding-001": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "vertex_ai-embedding-models",