From a185ad114460382f9ca8cfe48182ff42fff1878c Mon Sep 17 00:00:00 2001 From: IanShaw027 <131567472+IanShaw027@users.noreply.github.com> Date: Sun, 4 Jan 2026 15:36:00 +0800 Subject: [PATCH] =?UTF-8?q?feat(gemini):=20=E5=AE=8C=E5=96=84=20Gemini=20O?= =?UTF-8?q?Auth=20=E9=85=8D=E9=A2=9D=E7=B3=BB=E7=BB=9F=E5=92=8C=E7=94=A8?= =?UTF-8?q?=E9=87=8F=E6=98=BE=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 主要改动： - 后端：重构 Gemini 配额服务，支持多层级配额策略（GCP Standard/Free, Google One, AI Studio, Code Assist） - 后端：优化 OAuth 服务，增强 tier_id 识别和存储逻辑 - 后端：改进用量统计服务，支持不同平台的配额查询 - 后端：优化限流服务，增加临时解除调度状态管理 - 前端：统一四种授权方式的用量显示格式和徽标样式 - 前端：增强账户配额信息展示，支持多种配额类型 - 前端：改进创建和重新授权模态框的用户体验 - 国际化：完善中英文配额相关文案 - 移除 CHANGELOG.md 文件 测试：所有单元测试通过 --- CHANGELOG.md | 17 - .../handler/admin/gemini_oauth_handler.go | 8 +- backend/internal/pkg/geminicli/oauth.go | 18 +- .../repository/gemini_oauth_client.go | 6 +- backend/internal/service/account.go | 5 +- .../internal/service/account_usage_service.go | 60 +++- backend/internal/service/gateway_service.go | 7 + .../service/gemini_messages_compat_service.go | 18 + .../internal/service/gemini_oauth_service.go | 325 ++++++++++++++---- .../service/gemini_oauth_service_test.go | 153 +++++++-- backend/internal/service/gemini_quota.go | 244 +++++++++++-- backend/internal/service/ratelimit_service.go | 128 +++++-- frontend/src/api/admin/gemini.ts | 2 + .../components/account/AccountQuotaInfo.vue | 83 +++-- .../components/account/AccountUsageCell.vue | 280 ++++++++++----- .../components/account/CreateAccountModal.vue | 79 ++++- .../components/account/ReAuthAccountModal.vue | 48 ++- frontend/src/composables/useGeminiOAuth.ts | 8 +- frontend/src/i18n/locales/en.ts | 32 +- frontend/src/i18n/locales/zh.ts | 31 +- frontend/src/types/index.ts | 21 +- 21 files changed, 1205 insertions(+), 368 deletions(-) delete mode 100644 CHANGELOG.md diff --git a/CHANGELOG.md b/CHANGELOG.md deleted file mode 100644 index 1cab8802..00000000 --- a/CHANGELOG.md +++ /dev/null @@ -1,17 +0,0 @@ -# Changelog - -All notable changes to this project are documented in this file. - -The format is based on Keep a Changelog, and this project aims to follow Semantic Versioning. - -## [Unreleased] - -### Breaking Changes - -- Admin ops error logs: `GET /api/v1/admin/ops/error-logs` now enforces `limit <= 500` (previously `<= 5000`). Requests with `limit > 500` return `400 Bad Request` (`Invalid limit (must be 1-500)`). - -### Migration - -- Prefer the paginated endpoint `GET /api/v1/admin/ops/errors` using `page` / `page_size`. -- If you must keep using `.../error-logs`, reduce `limit` to `<= 500` and fetch multiple pages by splitting queries (e.g., by time window) instead of requesting a single large result set. - diff --git a/backend/internal/handler/admin/gemini_oauth_handler.go b/backend/internal/handler/admin/gemini_oauth_handler.go index 037800e2..8b4b13d5 100644 --- a/backend/internal/handler/admin/gemini_oauth_handler.go +++ b/backend/internal/handler/admin/gemini_oauth_handler.go @@ -30,6 +30,8 @@ type GeminiGenerateAuthURLRequest struct { // OAuth 类型: "code_assist" (需要 project_id) 或 "ai_studio" (不需要 project_id) // 默认为 "code_assist" 以保持向后兼容 OAuthType string `json:"oauth_type"` + // TierID is a user-selected tier to be used when auto detection is unavailable or fails. + TierID string `json:"tier_id"` } // GenerateAuthURL generates Google OAuth authorization URL for Gemini. @@ -54,7 +56,7 @@ func (h *GeminiOAuthHandler) GenerateAuthURL(c *gin.Context) { // Always pass the "hosted" callback URI; the OAuth service may override it depending on // oauth_type and whether the built-in Gemini CLI OAuth client is used. redirectURI := deriveGeminiRedirectURI(c) - result, err := h.geminiOAuthService.GenerateAuthURL(c.Request.Context(), req.ProxyID, redirectURI, req.ProjectID, oauthType) + result, err := h.geminiOAuthService.GenerateAuthURL(c.Request.Context(), req.ProxyID, redirectURI, req.ProjectID, oauthType, req.TierID) if err != nil { msg := err.Error() // Treat missing/invalid OAuth client configuration as a user/config error. @@ -76,6 +78,9 @@ type GeminiExchangeCodeRequest struct { ProxyID *int64 `json:"proxy_id"` // OAuth 类型: "code_assist" 或 "ai_studio"，需要与 GenerateAuthURL 时的类型一致 OAuthType string `json:"oauth_type"` + // TierID is a user-selected tier to be used when auto detection is unavailable or fails. + // This field is optional; when omitted, the server uses the tier stored in the OAuth session. + TierID string `json:"tier_id"` } // ExchangeCode exchanges authorization code for tokens. @@ -103,6 +108,7 @@ func (h *GeminiOAuthHandler) ExchangeCode(c *gin.Context) { Code: req.Code, ProxyID: req.ProxyID, OAuthType: oauthType, + TierID: req.TierID, }) if err != nil { response.BadRequest(c, "Failed to exchange code: "+err.Error()) diff --git a/backend/internal/pkg/geminicli/oauth.go b/backend/internal/pkg/geminicli/oauth.go index 83b3d491..473017a2 100644 --- a/backend/internal/pkg/geminicli/oauth.go +++ b/backend/internal/pkg/geminicli/oauth.go @@ -19,13 +19,17 @@ type OAuthConfig struct { } type OAuthSession struct { - State string `json:"state"` - CodeVerifier string `json:"code_verifier"` - ProxyURL string `json:"proxy_url,omitempty"` - RedirectURI string `json:"redirect_uri"` - ProjectID string `json:"project_id,omitempty"` - OAuthType string `json:"oauth_type"` // "code_assist" 或 "ai_studio" - CreatedAt time.Time `json:"created_at"` + State string `json:"state"` + CodeVerifier string `json:"code_verifier"` + ProxyURL string `json:"proxy_url,omitempty"` + RedirectURI string `json:"redirect_uri"` + ProjectID string `json:"project_id,omitempty"` + // TierID is a user-selected fallback tier. + // For oauth types that support auto detection (google_one/code_assist), the server will prefer + // the detected tier and fall back to TierID when detection fails. + TierID string `json:"tier_id,omitempty"` + OAuthType string `json:"oauth_type"` // "code_assist" 或 "ai_studio" + CreatedAt time.Time `json:"created_at"` } type SessionStore struct { diff --git a/backend/internal/repository/gemini_oauth_client.go b/backend/internal/repository/gemini_oauth_client.go index b1c86853..14ecfc89 100644 --- a/backend/internal/repository/gemini_oauth_client.go +++ b/backend/internal/repository/gemini_oauth_client.go @@ -30,14 +30,14 @@ func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, c // Use different OAuth clients based on oauthType: // - code_assist: always use built-in Gemini CLI OAuth client (public) - // - google_one: same as code_assist, uses built-in client for personal Google accounts + // - google_one: uses configured OAuth client when provided; otherwise falls back to built-in client // - ai_studio: requires a user-provided OAuth client oauthCfgInput := geminicli.OAuthConfig{ ClientID: c.cfg.Gemini.OAuth.ClientID, ClientSecret: c.cfg.Gemini.OAuth.ClientSecret, Scopes: c.cfg.Gemini.OAuth.Scopes, } - if oauthType == "code_assist" || oauthType == "google_one" { + if oauthType == "code_assist" { oauthCfgInput.ClientID = "" oauthCfgInput.ClientSecret = "" } @@ -78,7 +78,7 @@ func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refresh ClientSecret: c.cfg.Gemini.OAuth.ClientSecret, Scopes: c.cfg.Gemini.OAuth.Scopes, } - if oauthType == "code_assist" || oauthType == "google_one" { + if oauthType == "code_assist" { oauthCfgInput.ClientID = "" oauthCfgInput.ClientSecret = "" } diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 953090c2..40b765ad 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -105,10 +105,7 @@ func (a *Account) GeminiOAuthType() string { func (a *Account) GeminiTierID() string { tierID := strings.TrimSpace(a.GetCredential("tier_id")) - if tierID == "" { - return "" - } - return strings.ToUpper(tierID) + return tierID } func (a *Account) IsGeminiCodeAssist() bool { diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index eb992cd4..0a1227dd 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -107,6 +107,8 @@ type UsageProgress struct { ResetsAt *time.Time `json:"resets_at"` // 重置时间 RemainingSeconds int `json:"remaining_seconds"` // 距重置剩余秒数 WindowStats *WindowStats `json:"window_stats,omitempty"` // 窗口期统计（从窗口开始到当前的使用量） + UsedRequests int64 `json:"used_requests,omitempty"` + LimitRequests int64 `json:"limit_requests,omitempty"` } // AntigravityModelQuota Antigravity 单个模型的配额信息 @@ -117,12 +119,16 @@ type AntigravityModelQuota struct { // UsageInfo 账号使用量信息 type UsageInfo struct { - UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间 - FiveHour *UsageProgress `json:"five_hour"` // 5小时窗口 - SevenDay *UsageProgress `json:"seven_day,omitempty"` // 7天窗口 - SevenDaySonnet *UsageProgress `json:"seven_day_sonnet,omitempty"` // 7天Sonnet窗口 - GeminiProDaily *UsageProgress `json:"gemini_pro_daily,omitempty"` // Gemini Pro 日配额 - GeminiFlashDaily *UsageProgress `json:"gemini_flash_daily,omitempty"` // Gemini Flash 日配额 + UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间 + FiveHour *UsageProgress `json:"five_hour"` // 5小时窗口 + SevenDay *UsageProgress `json:"seven_day,omitempty"` // 7天窗口 + SevenDaySonnet *UsageProgress `json:"seven_day_sonnet,omitempty"` // 7天Sonnet窗口 + GeminiSharedDaily *UsageProgress `json:"gemini_shared_daily,omitempty"` // Gemini shared pool RPD (Google One / Code Assist) + GeminiProDaily *UsageProgress `json:"gemini_pro_daily,omitempty"` // Gemini Pro 日配额 + GeminiFlashDaily *UsageProgress `json:"gemini_flash_daily,omitempty"` // Gemini Flash 日配额 + GeminiSharedMinute *UsageProgress `json:"gemini_shared_minute,omitempty"` // Gemini shared pool RPM (Google One / Code Assist) + GeminiProMinute *UsageProgress `json:"gemini_pro_minute,omitempty"` // Gemini Pro RPM + GeminiFlashMinute *UsageProgress `json:"gemini_flash_minute,omitempty"` // Gemini Flash RPM // Antigravity 多模型配额 AntigravityQuota map[string]*AntigravityModelQuota `json:"antigravity_quota,omitempty"` @@ -258,17 +264,44 @@ func (s *AccountUsageService) getGeminiUsage(ctx context.Context, account *Accou return usage, nil } - start := geminiDailyWindowStart(now) - stats, err := s.usageLogRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID) + dayStart := geminiDailyWindowStart(now) + stats, err := s.usageLogRepo.GetModelStatsWithFilters(ctx, dayStart, now, 0, 0, account.ID) if err != nil { return nil, fmt.Errorf("get gemini usage stats failed: %w", err) } - totals := geminiAggregateUsage(stats) - resetAt := geminiDailyResetTime(now) + dayTotals := geminiAggregateUsage(stats) + dailyResetAt := geminiDailyResetTime(now) - usage.GeminiProDaily = buildGeminiUsageProgress(totals.ProRequests, quota.ProRPD, resetAt, totals.ProTokens, totals.ProCost, now) - usage.GeminiFlashDaily = buildGeminiUsageProgress(totals.FlashRequests, quota.FlashRPD, resetAt, totals.FlashTokens, totals.FlashCost, now) + // Daily window (RPD) + if quota.SharedRPD > 0 { + totalReq := dayTotals.ProRequests + dayTotals.FlashRequests + totalTokens := dayTotals.ProTokens + dayTotals.FlashTokens + totalCost := dayTotals.ProCost + dayTotals.FlashCost + usage.GeminiSharedDaily = buildGeminiUsageProgress(totalReq, quota.SharedRPD, dailyResetAt, totalTokens, totalCost, now) + } else { + usage.GeminiProDaily = buildGeminiUsageProgress(dayTotals.ProRequests, quota.ProRPD, dailyResetAt, dayTotals.ProTokens, dayTotals.ProCost, now) + usage.GeminiFlashDaily = buildGeminiUsageProgress(dayTotals.FlashRequests, quota.FlashRPD, dailyResetAt, dayTotals.FlashTokens, dayTotals.FlashCost, now) + } + + // Minute window (RPM) - fixed-window approximation: current minute [truncate(now), truncate(now)+1m) + minuteStart := now.Truncate(time.Minute) + minuteResetAt := minuteStart.Add(time.Minute) + minuteStats, err := s.usageLogRepo.GetModelStatsWithFilters(ctx, minuteStart, now, 0, 0, account.ID) + if err != nil { + return nil, fmt.Errorf("get gemini minute usage stats failed: %w", err) + } + minuteTotals := geminiAggregateUsage(minuteStats) + + if quota.SharedRPM > 0 { + totalReq := minuteTotals.ProRequests + minuteTotals.FlashRequests + totalTokens := minuteTotals.ProTokens + minuteTotals.FlashTokens + totalCost := minuteTotals.ProCost + minuteTotals.FlashCost + usage.GeminiSharedMinute = buildGeminiUsageProgress(totalReq, quota.SharedRPM, minuteResetAt, totalTokens, totalCost, now) + } else { + usage.GeminiProMinute = buildGeminiUsageProgress(minuteTotals.ProRequests, quota.ProRPM, minuteResetAt, minuteTotals.ProTokens, minuteTotals.ProCost, now) + usage.GeminiFlashMinute = buildGeminiUsageProgress(minuteTotals.FlashRequests, quota.FlashRPM, minuteResetAt, minuteTotals.FlashTokens, minuteTotals.FlashCost, now) + } return usage, nil } @@ -508,6 +541,7 @@ func (s *AccountUsageService) estimateSetupTokenUsage(account *Account) *UsageIn } func buildGeminiUsageProgress(used, limit int64, resetAt time.Time, tokens int64, cost float64, now time.Time) *UsageProgress { + // limit <= 0 means "no local quota window" (unknown or unlimited). if limit <= 0 { return nil } @@ -521,6 +555,8 @@ func buildGeminiUsageProgress(used, limit int64, resetAt time.Time, tokens int64 Utilization: utilization, ResetsAt: &resetCopy, RemainingSeconds: remainingSeconds, + UsedRequests: used, + LimitRequests: limit, WindowStats: &WindowStats{ Requests: used, Tokens: tokens, diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 7cb976b9..4e4b180a 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1064,6 +1064,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } // 不需要重试（成功或不可重试的错误），跳出循环 + // DEBUG: 输出响应 headers（用于检测 rate limit 信息） + if account.Platform == PlatformGemini && resp.StatusCode < 400 { + log.Printf("[DEBUG] Gemini API Response Headers for account %d:", account.ID) + for k, v := range resp.Header { + log.Printf("[DEBUG] %s: %v", k, v) + } + } break } defer func() { _ = resp.Body.Close() }() diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 3466c734..69f8a018 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -1628,6 +1628,15 @@ type UpstreamHTTPResult struct { } func (s *GeminiMessagesCompatService) handleNativeNonStreamingResponse(c *gin.Context, resp *http.Response, isOAuth bool) (*ClaudeUsage, error) { + // Log response headers for debugging + log.Printf("[GeminiAPI] ========== Response Headers ==========") + for key, values := range resp.Header { + if strings.HasPrefix(strings.ToLower(key), "x-ratelimit") { + log.Printf("[GeminiAPI] %s: %v", key, values) + } + } + log.Printf("[GeminiAPI] ========================================") + respBody, err := io.ReadAll(resp.Body) if err != nil { return nil, err @@ -1658,6 +1667,15 @@ func (s *GeminiMessagesCompatService) handleNativeNonStreamingResponse(c *gin.Co } func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, isOAuth bool) (*geminiNativeStreamResult, error) { + // Log response headers for debugging + log.Printf("[GeminiAPI] ========== Streaming Response Headers ==========") + for key, values := range resp.Header { + if strings.HasPrefix(strings.ToLower(key), "x-ratelimit") { + log.Printf("[GeminiAPI] %s: %v", key, values) + } + } + log.Printf("[GeminiAPI] ====================================================") + c.Status(resp.StatusCode) c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") diff --git a/backend/internal/service/gemini_oauth_service.go b/backend/internal/service/gemini_oauth_service.go index fdf7bee6..32526e42 100644 --- a/backend/internal/service/gemini_oauth_service.go +++ b/backend/internal/service/gemini_oauth_service.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "log" "net/http" "regexp" "strconv" @@ -18,12 +19,23 @@ import ( ) const ( - TierAIPremium = "AI_PREMIUM" - TierGoogleOneStandard = "GOOGLE_ONE_STANDARD" - TierGoogleOneBasic = "GOOGLE_ONE_BASIC" - TierFree = "FREE" - TierGoogleOneUnknown = "GOOGLE_ONE_UNKNOWN" - TierGoogleOneUnlimited = "GOOGLE_ONE_UNLIMITED" + // Canonical tier IDs used by sub2api (2026-aligned). + GeminiTierGoogleOneFree = "google_one_free" + GeminiTierGoogleAIPro = "google_ai_pro" + GeminiTierGoogleAIUltra = "google_ai_ultra" + GeminiTierGCPStandard = "gcp_standard" + GeminiTierGCPEnterprise = "gcp_enterprise" + GeminiTierAIStudioFree = "aistudio_free" + GeminiTierAIStudioPaid = "aistudio_paid" + GeminiTierGoogleOneUnknown = "google_one_unknown" + + // Legacy/compat tier IDs that may exist in historical data or upstream responses. + legacyTierAIPremium = "AI_PREMIUM" + legacyTierGoogleOneStandard = "GOOGLE_ONE_STANDARD" + legacyTierGoogleOneBasic = "GOOGLE_ONE_BASIC" + legacyTierFree = "FREE" + legacyTierGoogleOneUnknown = "GOOGLE_ONE_UNKNOWN" + legacyTierGoogleOneUnlimited = "GOOGLE_ONE_UNLIMITED" ) const ( @@ -84,7 +96,7 @@ type GeminiAuthURLResult struct { State string `json:"state"` } -func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64, redirectURI, projectID, oauthType string) (*GeminiAuthURLResult, error) { +func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64, redirectURI, projectID, oauthType, tierID string) (*GeminiAuthURLResult, error) { state, err := geminicli.GenerateState() if err != nil { return nil, fmt.Errorf("failed to generate state: %w", err) @@ -109,14 +121,14 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64 // OAuth client selection: // - code_assist: always use built-in Gemini CLI OAuth client (public), regardless of configured client_id/secret. - // - google_one: same as code_assist, uses built-in client for personal Google accounts. + // - google_one: uses configured OAuth client when provided; otherwise falls back to built-in client. // - ai_studio: requires a user-provided OAuth client. oauthCfg := geminicli.OAuthConfig{ ClientID: s.cfg.Gemini.OAuth.ClientID, ClientSecret: s.cfg.Gemini.OAuth.ClientSecret, Scopes: s.cfg.Gemini.OAuth.Scopes, } - if oauthType == "code_assist" || oauthType == "google_one" { + if oauthType == "code_assist" { oauthCfg.ClientID = "" oauthCfg.ClientSecret = "" } @@ -127,6 +139,7 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64 ProxyURL: proxyURL, RedirectURI: redirectURI, ProjectID: strings.TrimSpace(projectID), + TierID: canonicalGeminiTierIDForOAuthType(oauthType, tierID), OAuthType: oauthType, CreatedAt: time.Now(), } @@ -146,9 +159,9 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64 } // Redirect URI strategy: - // - code_assist: use Gemini CLI redirect URI (codeassist.google.com/authcode) - // - ai_studio: use localhost callback for manual copy/paste flow - if oauthType == "code_assist" { + // - built-in Gemini CLI OAuth client: use upstream redirect URI (codeassist.google.com/authcode) + // - custom OAuth client: use localhost callback for manual copy/paste flow + if isBuiltinClient { redirectURI = geminicli.GeminiCLIRedirectURI } else { redirectURI = geminicli.AIStudioOAuthRedirectURI @@ -174,6 +187,9 @@ type GeminiExchangeCodeInput struct { Code string ProxyID *int64 OAuthType string // "code_assist" 或 "ai_studio" + // TierID is a user-selected tier to be used when auto detection is unavailable or fails. + // If empty, the service will fall back to the tier stored in the OAuth session (if any). + TierID string } type GeminiTokenInfo struct { @@ -185,7 +201,7 @@ type GeminiTokenInfo struct { Scope string `json:"scope,omitempty"` ProjectID string `json:"project_id,omitempty"` OAuthType string `json:"oauth_type,omitempty"` // "code_assist" 或 "ai_studio" - TierID string `json:"tier_id,omitempty"` // Gemini Code Assist tier: LEGACY/PRO/ULTRA + TierID string `json:"tier_id,omitempty"` // Canonical tier id (e.g. google_one_free, gcp_standard, aistudio_free) Extra map[string]any `json:"extra,omitempty"` // Drive metadata } @@ -204,6 +220,90 @@ func validateTierID(tierID string) error { return nil } +func canonicalGeminiTierID(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return "" + } + + lower := strings.ToLower(raw) + switch lower { + case GeminiTierGoogleOneFree, + GeminiTierGoogleAIPro, + GeminiTierGoogleAIUltra, + GeminiTierGCPStandard, + GeminiTierGCPEnterprise, + GeminiTierAIStudioFree, + GeminiTierAIStudioPaid, + GeminiTierGoogleOneUnknown: + return lower + } + + upper := strings.ToUpper(raw) + switch upper { + // Google One legacy tiers + case legacyTierAIPremium: + return GeminiTierGoogleAIPro + case legacyTierGoogleOneUnlimited: + return GeminiTierGoogleAIUltra + case legacyTierFree, legacyTierGoogleOneBasic, legacyTierGoogleOneStandard: + return GeminiTierGoogleOneFree + case legacyTierGoogleOneUnknown: + return GeminiTierGoogleOneUnknown + + // Code Assist legacy tiers + case "STANDARD", "PRO", "LEGACY": + return GeminiTierGCPStandard + case "ENTERPRISE", "ULTRA": + return GeminiTierGCPEnterprise + } + + // Some Code Assist responses use kebab-case tier identifiers. + switch lower { + case "standard-tier", "pro-tier": + return GeminiTierGCPStandard + case "ultra-tier": + return GeminiTierGCPEnterprise + } + + return "" +} + +func canonicalGeminiTierIDForOAuthType(oauthType, tierID string) string { + oauthType = strings.ToLower(strings.TrimSpace(oauthType)) + canonical := canonicalGeminiTierID(tierID) + if canonical == "" { + return "" + } + + switch oauthType { + case "google_one": + switch canonical { + case GeminiTierGoogleOneFree, GeminiTierGoogleAIPro, GeminiTierGoogleAIUltra: + return canonical + default: + return "" + } + case "code_assist": + switch canonical { + case GeminiTierGCPStandard, GeminiTierGCPEnterprise: + return canonical + default: + return "" + } + case "ai_studio": + switch canonical { + case GeminiTierAIStudioFree, GeminiTierAIStudioPaid: + return canonical + default: + return "" + } + default: + // Unknown oauth type: accept canonical tier. + return canonical + } +} + // extractTierIDFromAllowedTiers extracts tierID from LoadCodeAssist response // Prioritizes IsDefault tier, falls back to first non-empty tier func extractTierIDFromAllowedTiers(allowedTiers []geminicli.AllowedTier) string { @@ -233,55 +333,35 @@ func inferGoogleOneTier(storageBytes int64) string { if storageBytes <= 0 { log.Printf("[GeminiOAuth] inferGoogleOneTier - storageBytes <= 0, returning UNKNOWN") - return TierGoogleOneUnknown + return GeminiTierGoogleOneUnknown } if storageBytes > StorageTierUnlimited { log.Printf("[GeminiOAuth] inferGoogleOneTier - > %d bytes (100TB), returning UNLIMITED", StorageTierUnlimited) - return TierGoogleOneUnlimited + return GeminiTierGoogleAIUltra } if storageBytes >= StorageTierAIPremium { - log.Printf("[GeminiOAuth] inferGoogleOneTier - >= %d bytes (2TB), returning AI_PREMIUM", StorageTierAIPremium) - return TierAIPremium - } - if storageBytes >= StorageTierStandard { - log.Printf("[GeminiOAuth] inferGoogleOneTier - >= %d bytes (200GB), returning STANDARD", StorageTierStandard) - return TierGoogleOneStandard - } - if storageBytes >= StorageTierBasic { - log.Printf("[GeminiOAuth] inferGoogleOneTier - >= %d bytes (100GB), returning BASIC", StorageTierBasic) - return TierGoogleOneBasic + log.Printf("[GeminiOAuth] inferGoogleOneTier - >= %d bytes (2TB), returning google_ai_pro", StorageTierAIPremium) + return GeminiTierGoogleAIPro } if storageBytes >= StorageTierFree { log.Printf("[GeminiOAuth] inferGoogleOneTier - >= %d bytes (15GB), returning FREE", StorageTierFree) - return TierFree + return GeminiTierGoogleOneFree } log.Printf("[GeminiOAuth] inferGoogleOneTier - < %d bytes (15GB), returning UNKNOWN", StorageTierFree) - return TierGoogleOneUnknown + return GeminiTierGoogleOneUnknown } -// fetchGoogleOneTier fetches Google One tier from Drive API or LoadCodeAssist API +// fetchGoogleOneTier fetches Google One tier from Drive API +// Note: LoadCodeAssist API is NOT called for Google One accounts because: +// 1. It's designed for GCP IAM (enterprise), not personal Google accounts +// 2. Personal accounts will get 403/404 from cloudaicompanion.googleapis.com +// 3. Google consumer (Google One) and enterprise (GCP) systems are physically isolated func (s *GeminiOAuthService) FetchGoogleOneTier(ctx context.Context, accessToken, proxyURL string) (string, *geminicli.DriveStorageInfo, error) { - log.Printf("[GeminiOAuth] Starting FetchGoogleOneTier") + log.Printf("[GeminiOAuth] Starting FetchGoogleOneTier (Google One personal account)") - // First try LoadCodeAssist API (works for accounts with GCP projects) - if s.codeAssist != nil { - log.Printf("[GeminiOAuth] Trying LoadCodeAssist API...") - loadResp, err := s.codeAssist.LoadCodeAssist(ctx, accessToken, proxyURL, nil) - if err != nil { - log.Printf("[GeminiOAuth] LoadCodeAssist failed: %v", err) - } else if loadResp != nil { - if tier := loadResp.GetTier(); tier != "" { - log.Printf("[GeminiOAuth] Got tier from LoadCodeAssist: %s (skipping Drive API)", tier) - return tier, nil, nil - } else { - log.Printf("[GeminiOAuth] LoadCodeAssist returned no tier, falling back to Drive API") - } - } - } - - // Fallback to Drive API (requires drive.readonly scope) + // Use Drive API to infer tier from storage quota (requires drive.readonly scope) log.Printf("[GeminiOAuth] Calling Drive API for storage quota...") driveClient := geminicli.NewDriveClient() @@ -290,11 +370,11 @@ func (s *GeminiOAuthService) FetchGoogleOneTier(ctx context.Context, accessToken // Check if it's a 403 (scope not granted) if strings.Contains(err.Error(), "status 403") { log.Printf("[GeminiOAuth] Drive API scope not available (403): %v", err) - return TierGoogleOneUnknown, nil, err + return GeminiTierGoogleOneUnknown, nil, err } // Other errors log.Printf("[GeminiOAuth] Failed to fetch Drive storage: %v", err) - return TierGoogleOneUnknown, nil, err + return GeminiTierGoogleOneUnknown, nil, err } log.Printf("[GeminiOAuth] Drive API response - Limit: %d bytes (%.2f TB), Usage: %d bytes (%.2f GB)", @@ -362,11 +442,16 @@ func (s *GeminiOAuthService) RefreshAccountGoogleOneTier( } func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExchangeCodeInput) (*GeminiTokenInfo, error) { + log.Printf("[GeminiOAuth] ========== ExchangeCode START ==========") + log.Printf("[GeminiOAuth] SessionID: %s", input.SessionID) + session, ok := s.sessionStore.Get(input.SessionID) if !ok { + log.Printf("[GeminiOAuth] ERROR: Session not found or expired") return nil, fmt.Errorf("session not found or expired") } if strings.TrimSpace(input.State) == "" || input.State != session.State { + log.Printf("[GeminiOAuth] ERROR: Invalid state") return nil, fmt.Errorf("invalid state") } @@ -377,6 +462,7 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch proxyURL = proxy.URL() } } + log.Printf("[GeminiOAuth] ProxyURL: %s", proxyURL) redirectURI := session.RedirectURI @@ -385,6 +471,8 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch if oauthType == "" { oauthType = "code_assist" } + log.Printf("[GeminiOAuth] OAuth Type: %s", oauthType) + log.Printf("[GeminiOAuth] Project ID from session: %s", session.ProjectID) // If the session was created for AI Studio OAuth, ensure a custom OAuth client is configured. if oauthType == "ai_studio" { @@ -410,8 +498,13 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch tokenResp, err := s.oauthClient.ExchangeCode(ctx, oauthType, input.Code, session.CodeVerifier, redirectURI, proxyURL) if err != nil { + log.Printf("[GeminiOAuth] ERROR: Failed to exchange code: %v", err) return nil, fmt.Errorf("failed to exchange code: %w", err) } + log.Printf("[GeminiOAuth] Token exchange successful") + log.Printf("[GeminiOAuth] Token scope: %s", tokenResp.Scope) + log.Printf("[GeminiOAuth] Token expires_in: %d seconds", tokenResp.ExpiresIn) + sessionProjectID := strings.TrimSpace(session.ProjectID) s.sessionStore.Delete(input.SessionID) @@ -427,36 +520,63 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch projectID := sessionProjectID var tierID string + fallbackTierID := canonicalGeminiTierIDForOAuthType(oauthType, input.TierID) + if fallbackTierID == "" { + fallbackTierID = canonicalGeminiTierIDForOAuthType(oauthType, session.TierID) + } + + log.Printf("[GeminiOAuth] ========== Account Type Detection START ==========") + log.Printf("[GeminiOAuth] OAuth Type: %s", oauthType) // 对于 code_assist 模式，project_id 是必需的，需要调用 Code Assist API // 对于 google_one 模式，使用个人 Google 账号，不需要 project_id，配额由 Google 网关自动识别 // 对于 ai_studio 模式，project_id 是可选的（不影响使用 AI Studio API） switch oauthType { case "code_assist": + log.Printf("[GeminiOAuth] Processing code_assist OAuth type") if projectID == "" { + log.Printf("[GeminiOAuth] No project_id provided, attempting to fetch from LoadCodeAssist API...") var err error projectID, tierID, err = s.fetchProjectID(ctx, tokenResp.AccessToken, proxyURL) if err != nil { // 记录警告但不阻断流程，允许后续补充 project_id fmt.Printf("[GeminiOAuth] Warning: Failed to fetch project_id during token exchange: %v\n", err) + log.Printf("[GeminiOAuth] WARNING: Failed to fetch project_id: %v", err) + } else { + log.Printf("[GeminiOAuth] Successfully fetched project_id: %s, tier_id: %s", projectID, tierID) } } else { + log.Printf("[GeminiOAuth] User provided project_id: %s, fetching tier_id...", projectID) // 用户手动填了 project_id，仍需调用 LoadCodeAssist 获取 tierID _, fetchedTierID, err := s.fetchProjectID(ctx, tokenResp.AccessToken, proxyURL) if err != nil { fmt.Printf("[GeminiOAuth] Warning: Failed to fetch tierID: %v\n", err) + log.Printf("[GeminiOAuth] WARNING: Failed to fetch tier_id: %v", err) } else { tierID = fetchedTierID + log.Printf("[GeminiOAuth] Successfully fetched tier_id: %s", tierID) } } if strings.TrimSpace(projectID) == "" { + log.Printf("[GeminiOAuth] ERROR: Missing project_id for Code Assist OAuth") return nil, fmt.Errorf("missing project_id for Code Assist OAuth: please fill Project ID (optional field) and regenerate the auth URL, or ensure your Google account has an ACTIVE GCP project") } - // tierID 缺失时使用默认值 + // Prefer auto-detected tier; fall back to user-selected tier. + tierID = canonicalGeminiTierIDForOAuthType(oauthType, tierID) if tierID == "" { - tierID = "LEGACY" + if fallbackTierID != "" { + tierID = fallbackTierID + log.Printf("[GeminiOAuth] Using fallback tier_id from user/session: %s", tierID) + } else { + tierID = GeminiTierGCPStandard + log.Printf("[GeminiOAuth] Using default tier_id: %s", tierID) + } } + log.Printf("[GeminiOAuth] Final code_assist result - project_id: %s, tier_id: %s", projectID, tierID) + case "google_one": + log.Printf("[GeminiOAuth] Processing google_one OAuth type") + log.Printf("[GeminiOAuth] Attempting to fetch Google One tier from Drive API...") // Attempt to fetch Drive storage tier var storageInfo *geminicli.DriveStorageInfo var err error @@ -464,9 +584,27 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch if err != nil { // Log warning but don't block - use fallback fmt.Printf("[GeminiOAuth] Warning: Failed to fetch Drive tier: %v\n", err) - tierID = TierGoogleOneUnknown + log.Printf("[GeminiOAuth] WARNING: Failed to fetch Drive tier: %v", err) + tierID = "" + } else { + log.Printf("[GeminiOAuth] Successfully fetched Drive tier: %s", tierID) + if storageInfo != nil { + log.Printf("[GeminiOAuth] Drive storage - Limit: %d bytes (%.2f TB), Usage: %d bytes (%.2f GB)", + storageInfo.Limit, float64(storageInfo.Limit)/float64(TB), + storageInfo.Usage, float64(storageInfo.Usage)/float64(GB)) + } } - fmt.Printf("[GeminiOAuth] Google One tierID after fetch: %s\n", tierID) + tierID = canonicalGeminiTierIDForOAuthType(oauthType, tierID) + if tierID == "" || tierID == GeminiTierGoogleOneUnknown { + if fallbackTierID != "" { + tierID = fallbackTierID + log.Printf("[GeminiOAuth] Using fallback tier_id from user/session: %s", tierID) + } else { + tierID = GeminiTierGoogleOneFree + log.Printf("[GeminiOAuth] Using default tier_id: %s", tierID) + } + } + fmt.Printf("[GeminiOAuth] Google One tierID after normalization: %s\n", tierID) // Store Drive info in extra field for caching if storageInfo != nil { @@ -486,10 +624,23 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch "drive_tier_updated_at": time.Now().Format(time.RFC3339), }, } + log.Printf("[GeminiOAuth] ========== ExchangeCode END (google_one with storage info) ==========") return tokenInfo, nil } + + case "ai_studio": + // No automatic tier detection for AI Studio OAuth; rely on user selection. + if fallbackTierID != "" { + tierID = fallbackTierID + } else { + tierID = GeminiTierAIStudioFree + } + + default: + log.Printf("[GeminiOAuth] Processing %s OAuth type (no tier detection)", oauthType) } - // ai_studio 模式不设置 tierID，保持为空 + + log.Printf("[GeminiOAuth] ========== Account Type Detection END ==========") result := &GeminiTokenInfo{ AccessToken: tokenResp.AccessToken, @@ -502,7 +653,8 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch TierID: tierID, OAuthType: oauthType, } - fmt.Printf("[GeminiOAuth] ExchangeCode returning tierID: %s\n", result.TierID) + log.Printf("[GeminiOAuth] Final result - OAuth Type: %s, Project ID: %s, Tier ID: %s", result.OAuthType, result.ProjectID, result.TierID) + log.Printf("[GeminiOAuth] ========== ExchangeCode END ==========") return result, nil } @@ -599,6 +751,17 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A err = nil } } + // Backward compatibility for google_one: + // - New behavior: when a custom OAuth client is configured, google_one will use it. + // - Old behavior: google_one always used the built-in Gemini CLI OAuth client. + // If an existing account was authorized with the built-in client, refreshing with the custom client + // will fail with "unauthorized_client". Retry with the built-in client (code_assist path forces it). + if err != nil && oauthType == "google_one" && strings.Contains(err.Error(), "unauthorized_client") && s.GetOAuthConfig().AIStudioOAuthEnabled { + if alt, altErr := s.RefreshToken(ctx, "code_assist", refreshToken, proxyURL); altErr == nil { + tokenInfo = alt + err = nil + } + } if err != nil { // Provide a more actionable error for common OAuth client mismatch issues. if strings.Contains(err.Error(), "unauthorized_client") { @@ -624,13 +787,14 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A case "code_assist": // 先设置默认值或保留旧值，确保 tier_id 始终有值 if existingTierID != "" { - tokenInfo.TierID = existingTierID - } else { - tokenInfo.TierID = "LEGACY" // 默认值 + tokenInfo.TierID = canonicalGeminiTierIDForOAuthType(oauthType, existingTierID) + } + if tokenInfo.TierID == "" { + tokenInfo.TierID = GeminiTierGCPStandard } // 尝试自动探测 project_id 和 tier_id - needDetect := strings.TrimSpace(tokenInfo.ProjectID) == "" || existingTierID == "" + needDetect := strings.TrimSpace(tokenInfo.ProjectID) == "" || tokenInfo.TierID == "" if needDetect { projectID, tierID, err := s.fetchProjectID(ctx, tokenInfo.AccessToken, proxyURL) if err != nil { @@ -639,9 +803,10 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A if strings.TrimSpace(tokenInfo.ProjectID) == "" && projectID != "" { tokenInfo.ProjectID = projectID } - // 只有当原来没有 tier_id 且探测成功时才更新 - if existingTierID == "" && tierID != "" { - tokenInfo.TierID = tierID + if tierID != "" { + if canonical := canonicalGeminiTierIDForOAuthType(oauthType, tierID); canonical != "" { + tokenInfo.TierID = canonical + } } } } @@ -650,6 +815,7 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A return nil, fmt.Errorf("failed to auto-detect project_id: empty result") } case "google_one": + canonicalExistingTier := canonicalGeminiTierIDForOAuthType(oauthType, existingTierID) // Check if tier cache is stale (> 24 hours) needsRefresh := true if account.Extra != nil { @@ -658,30 +824,37 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A if time.Since(updatedAt) <= 24*time.Hour { needsRefresh = false // Use cached tier - if existingTierID != "" { - tokenInfo.TierID = existingTierID - } + tokenInfo.TierID = canonicalExistingTier } } } } + if tokenInfo.TierID == "" { + tokenInfo.TierID = canonicalExistingTier + } + if needsRefresh { tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, tokenInfo.AccessToken, proxyURL) - if err == nil && storageInfo != nil { - tokenInfo.TierID = tierID - tokenInfo.Extra = map[string]any{ - "drive_storage_limit": storageInfo.Limit, - "drive_storage_usage": storageInfo.Usage, - "drive_tier_updated_at": time.Now().Format(time.RFC3339), + if err == nil { + if canonical := canonicalGeminiTierIDForOAuthType(oauthType, tierID); canonical != "" && canonical != GeminiTierGoogleOneUnknown { + tokenInfo.TierID = canonical } + if storageInfo != nil { + tokenInfo.Extra = map[string]any{ + "drive_storage_limit": storageInfo.Limit, + "drive_storage_usage": storageInfo.Usage, + "drive_tier_updated_at": time.Now().Format(time.RFC3339), + } + } + } + } + + if tokenInfo.TierID == "" || tokenInfo.TierID == GeminiTierGoogleOneUnknown { + if canonicalExistingTier != "" { + tokenInfo.TierID = canonicalExistingTier } else { - // Fallback to cached or unknown - if existingTierID != "" { - tokenInfo.TierID = existingTierID - } else { - tokenInfo.TierID = TierGoogleOneUnknown - } + tokenInfo.TierID = GeminiTierGoogleOneFree } } } diff --git a/backend/internal/service/gemini_oauth_service_test.go b/backend/internal/service/gemini_oauth_service_test.go index 026e6dc2..eb3d86e6 100644 --- a/backend/internal/service/gemini_oauth_service_test.go +++ b/backend/internal/service/gemini_oauth_service_test.go @@ -1,50 +1,129 @@ package service -import "testing" +import ( + "context" + "net/url" + "strings" + "testing" -func TestInferGoogleOneTier(t *testing.T) { - tests := []struct { - name string - storageBytes int64 - expectedTier string - }{ - {"Negative storage", -1, TierGoogleOneUnknown}, - {"Zero storage", 0, TierGoogleOneUnknown}, + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" +) - // Free tier boundary (15GB) - {"Below free tier", 10 * GB, TierGoogleOneUnknown}, - {"Just below free tier", StorageTierFree - 1, TierGoogleOneUnknown}, - {"Free tier (15GB)", StorageTierFree, TierFree}, +func TestGeminiOAuthService_GenerateAuthURL_RedirectURIStrategy(t *testing.T) { + t.Parallel() - // Basic tier boundary (100GB) - {"Between free and basic", 50 * GB, TierFree}, - {"Just below basic tier", StorageTierBasic - 1, TierFree}, - {"Basic tier (100GB)", StorageTierBasic, TierGoogleOneBasic}, + type testCase struct { + name string + cfg *config.Config + oauthType string + projectID string + wantClientID string + wantRedirect string + wantScope string + wantProjectID string + wantErrSubstr string + } - // Standard tier boundary (200GB) - {"Between basic and standard", 150 * GB, TierGoogleOneBasic}, - {"Just below standard tier", StorageTierStandard - 1, TierGoogleOneBasic}, - {"Standard tier (200GB)", StorageTierStandard, TierGoogleOneStandard}, - - // AI Premium tier boundary (2TB) - {"Between standard and premium", 1 * TB, TierGoogleOneStandard}, - {"Just below AI Premium tier", StorageTierAIPremium - 1, TierGoogleOneStandard}, - {"AI Premium tier (2TB)", StorageTierAIPremium, TierAIPremium}, - - // Unlimited tier boundary (> 100TB) - {"Between premium and unlimited", 50 * TB, TierAIPremium}, - {"At unlimited threshold (100TB)", StorageTierUnlimited, TierAIPremium}, - {"Unlimited tier (100TB+)", StorageTierUnlimited + 1, TierGoogleOneUnlimited}, - {"Unlimited tier (101TB+)", 101 * TB, TierGoogleOneUnlimited}, - {"Very large storage", 1000 * TB, TierGoogleOneUnlimited}, + tests := []testCase{ + { + name: "google_one uses built-in client when not configured and redirects to upstream", + cfg: &config.Config{ + Gemini: config.GeminiConfig{ + OAuth: config.GeminiOAuthConfig{}, + }, + }, + oauthType: "google_one", + wantClientID: geminicli.GeminiCLIOAuthClientID, + wantRedirect: geminicli.GeminiCLIRedirectURI, + wantScope: geminicli.DefaultCodeAssistScopes, + wantProjectID: "", + }, + { + name: "google_one uses custom client when configured and redirects to localhost", + cfg: &config.Config{ + Gemini: config.GeminiConfig{ + OAuth: config.GeminiOAuthConfig{ + ClientID: "custom-client-id", + ClientSecret: "custom-client-secret", + }, + }, + }, + oauthType: "google_one", + wantClientID: "custom-client-id", + wantRedirect: geminicli.AIStudioOAuthRedirectURI, + wantScope: geminicli.DefaultGoogleOneScopes, + wantProjectID: "", + }, + { + name: "code_assist always forces built-in client even when custom client configured", + cfg: &config.Config{ + Gemini: config.GeminiConfig{ + OAuth: config.GeminiOAuthConfig{ + ClientID: "custom-client-id", + ClientSecret: "custom-client-secret", + }, + }, + }, + oauthType: "code_assist", + projectID: "my-gcp-project", + wantClientID: geminicli.GeminiCLIOAuthClientID, + wantRedirect: geminicli.GeminiCLIRedirectURI, + wantScope: geminicli.DefaultCodeAssistScopes, + wantProjectID: "my-gcp-project", + }, + { + name: "ai_studio requires custom client", + cfg: &config.Config{ + Gemini: config.GeminiConfig{ + OAuth: config.GeminiOAuthConfig{}, + }, + }, + oauthType: "ai_studio", + wantErrSubstr: "AI Studio OAuth requires a custom OAuth Client", + }, } for _, tt := range tests { + tt := tt t.Run(tt.name, func(t *testing.T) { - result := inferGoogleOneTier(tt.storageBytes) - if result != tt.expectedTier { - t.Errorf("inferGoogleOneTier(%d) = %s, want %s", - tt.storageBytes, result, tt.expectedTier) + t.Parallel() + + svc := NewGeminiOAuthService(nil, nil, nil, tt.cfg) + got, err := svc.GenerateAuthURL(context.Background(), nil, "https://example.com/auth/callback", tt.projectID, tt.oauthType, "") + if tt.wantErrSubstr != "" { + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantErrSubstr) + } + if !strings.Contains(err.Error(), tt.wantErrSubstr) { + t.Fatalf("expected error containing %q, got: %v", tt.wantErrSubstr, err) + } + return + } + if err != nil { + t.Fatalf("GenerateAuthURL returned error: %v", err) + } + + parsed, err := url.Parse(got.AuthURL) + if err != nil { + t.Fatalf("failed to parse auth_url: %v", err) + } + q := parsed.Query() + + if gotState := q.Get("state"); gotState != got.State { + t.Fatalf("state mismatch: query=%q result=%q", gotState, got.State) + } + if gotClientID := q.Get("client_id"); gotClientID != tt.wantClientID { + t.Fatalf("client_id mismatch: got=%q want=%q", gotClientID, tt.wantClientID) + } + if gotRedirect := q.Get("redirect_uri"); gotRedirect != tt.wantRedirect { + t.Fatalf("redirect_uri mismatch: got=%q want=%q", gotRedirect, tt.wantRedirect) + } + if gotScope := q.Get("scope"); gotScope != tt.wantScope { + t.Fatalf("scope mismatch: got=%q want=%q", gotScope, tt.wantScope) + } + if gotProjectID := q.Get("project_id"); gotProjectID != tt.wantProjectID { + t.Fatalf("project_id mismatch: got=%q want=%q", gotProjectID, tt.wantProjectID) } }) } diff --git a/backend/internal/service/gemini_quota.go b/backend/internal/service/gemini_quota.go index 47ffbfe8..3a70232c 100644 --- a/backend/internal/service/gemini_quota.go +++ b/backend/internal/service/gemini_quota.go @@ -20,13 +20,24 @@ const ( geminiModelFlash geminiModelClass = "flash" ) -type GeminiDailyQuota struct { - ProRPD int64 - FlashRPD int64 +type GeminiQuota struct { + // SharedRPD is a shared requests-per-day pool across models. + // When SharedRPD > 0, callers should treat ProRPD/FlashRPD as not applicable for daily quota checks. + SharedRPD int64 `json:"shared_rpd,omitempty"` + // SharedRPM is a shared requests-per-minute pool across models. + // When SharedRPM > 0, callers should treat ProRPM/FlashRPM as not applicable for minute quota checks. + SharedRPM int64 `json:"shared_rpm,omitempty"` + + // Per-model quotas (AI Studio / API key). + // A value of -1 means "unlimited" (pay-as-you-go). + ProRPD int64 `json:"pro_rpd,omitempty"` + ProRPM int64 `json:"pro_rpm,omitempty"` + FlashRPD int64 `json:"flash_rpd,omitempty"` + FlashRPM int64 `json:"flash_rpm,omitempty"` } type GeminiTierPolicy struct { - Quota GeminiDailyQuota + Quota GeminiQuota Cooldown time.Duration } @@ -45,10 +56,27 @@ type GeminiUsageTotals struct { const geminiQuotaCacheTTL = time.Minute -type geminiQuotaOverrides struct { +type geminiQuotaOverridesV1 struct { Tiers map[string]config.GeminiTierQuotaConfig `json:"tiers"` } +type geminiQuotaOverridesV2 struct { + QuotaRules map[string]geminiQuotaRuleOverride `json:"quota_rules"` +} + +type geminiQuotaRuleOverride struct { + SharedRPD *int64 `json:"shared_rpd,omitempty"` + SharedRPM *int64 `json:"rpm,omitempty"` + GeminiPro *geminiModelQuotaOverride `json:"gemini_pro,omitempty"` + GeminiFlash *geminiModelQuotaOverride `json:"gemini_flash,omitempty"` + Desc *string `json:"desc,omitempty"` +} + +type geminiModelQuotaOverride struct { + RPD *int64 `json:"rpd,omitempty"` + RPM *int64 `json:"rpm,omitempty"` +} + type GeminiQuotaService struct { cfg *config.Config settingRepo SettingRepository @@ -82,11 +110,17 @@ func (s *GeminiQuotaService) Policy(ctx context.Context) *GeminiQuotaPolicy { if s.cfg != nil { policy.ApplyOverrides(s.cfg.Gemini.Quota.Tiers) if strings.TrimSpace(s.cfg.Gemini.Quota.Policy) != "" { - var overrides geminiQuotaOverrides - if err := json.Unmarshal([]byte(s.cfg.Gemini.Quota.Policy), &overrides); err != nil { - log.Printf("gemini quota: parse config policy failed: %v", err) + raw := []byte(s.cfg.Gemini.Quota.Policy) + var overridesV2 geminiQuotaOverridesV2 + if err := json.Unmarshal(raw, &overridesV2); err == nil && len(overridesV2.QuotaRules) > 0 { + policy.ApplyQuotaRulesOverrides(overridesV2.QuotaRules) } else { - policy.ApplyOverrides(overrides.Tiers) + var overridesV1 geminiQuotaOverridesV1 + if err := json.Unmarshal(raw, &overridesV1); err != nil { + log.Printf("gemini quota: parse config policy failed: %v", err) + } else { + policy.ApplyOverrides(overridesV1.Tiers) + } } } } @@ -96,11 +130,17 @@ func (s *GeminiQuotaService) Policy(ctx context.Context) *GeminiQuotaPolicy { if err != nil && !errors.Is(err, ErrSettingNotFound) { log.Printf("gemini quota: load setting failed: %v", err) } else if strings.TrimSpace(value) != "" { - var overrides geminiQuotaOverrides - if err := json.Unmarshal([]byte(value), &overrides); err != nil { - log.Printf("gemini quota: parse setting failed: %v", err) + raw := []byte(value) + var overridesV2 geminiQuotaOverridesV2 + if err := json.Unmarshal(raw, &overridesV2); err == nil && len(overridesV2.QuotaRules) > 0 { + policy.ApplyQuotaRulesOverrides(overridesV2.QuotaRules) } else { - policy.ApplyOverrides(overrides.Tiers) + var overridesV1 geminiQuotaOverridesV1 + if err := json.Unmarshal(raw, &overridesV1); err != nil { + log.Printf("gemini quota: parse setting failed: %v", err) + } else { + policy.ApplyOverrides(overridesV1.Tiers) + } } } } @@ -113,12 +153,20 @@ func (s *GeminiQuotaService) Policy(ctx context.Context) *GeminiQuotaPolicy { return policy } -func (s *GeminiQuotaService) QuotaForAccount(ctx context.Context, account *Account) (GeminiDailyQuota, bool) { - if account == nil || !account.IsGeminiCodeAssist() { - return GeminiDailyQuota{}, false +func (s *GeminiQuotaService) QuotaForAccount(ctx context.Context, account *Account) (GeminiQuota, bool) { + if account == nil || account.Platform != PlatformGemini { + return GeminiQuota{}, false } + + // Map (oauth_type + tier_id) to a canonical policy tier key. + // This keeps the policy table stable even if upstream tier_id strings vary. + tierKey := geminiQuotaTierKeyForAccount(account) + if tierKey == "" { + return GeminiQuota{}, false + } + policy := s.Policy(ctx) - return policy.QuotaForTier(account.GeminiTierID()) + return policy.QuotaForTier(tierKey) } func (s *GeminiQuotaService) CooldownForTier(ctx context.Context, tierID string) time.Duration { @@ -126,12 +174,36 @@ func (s *GeminiQuotaService) CooldownForTier(ctx context.Context, tierID string) return policy.CooldownForTier(tierID) } +func (s *GeminiQuotaService) CooldownForAccount(ctx context.Context, account *Account) time.Duration { + if s == nil || account == nil || account.Platform != PlatformGemini { + return 5 * time.Minute + } + tierKey := geminiQuotaTierKeyForAccount(account) + if strings.TrimSpace(tierKey) == "" { + return 5 * time.Minute + } + return s.CooldownForTier(ctx, tierKey) +} + func newGeminiQuotaPolicy() *GeminiQuotaPolicy { return &GeminiQuotaPolicy{ tiers: map[string]GeminiTierPolicy{ - "LEGACY": {Quota: GeminiDailyQuota{ProRPD: 50, FlashRPD: 1500}, Cooldown: 30 * time.Minute}, - "PRO": {Quota: GeminiDailyQuota{ProRPD: 1500, FlashRPD: 4000}, Cooldown: 5 * time.Minute}, - "ULTRA": {Quota: GeminiDailyQuota{ProRPD: 2000, FlashRPD: 0}, Cooldown: 5 * time.Minute}, + // --- AI Studio / API Key (per-model) --- + // aistudio_free: + // - gemini_pro: 50 RPD / 2 RPM + // - gemini_flash: 1500 RPD / 15 RPM + GeminiTierAIStudioFree: {Quota: GeminiQuota{ProRPD: 50, ProRPM: 2, FlashRPD: 1500, FlashRPM: 15}, Cooldown: 30 * time.Minute}, + // aistudio_paid: -1 means "unlimited/pay-as-you-go" for RPD. + GeminiTierAIStudioPaid: {Quota: GeminiQuota{ProRPD: -1, ProRPM: 1000, FlashRPD: -1, FlashRPM: 2000}, Cooldown: 5 * time.Minute}, + + // --- Google One (shared pool) --- + GeminiTierGoogleOneFree: {Quota: GeminiQuota{SharedRPD: 1000, SharedRPM: 60}, Cooldown: 30 * time.Minute}, + GeminiTierGoogleAIPro: {Quota: GeminiQuota{SharedRPD: 1500, SharedRPM: 120}, Cooldown: 5 * time.Minute}, + GeminiTierGoogleAIUltra: {Quota: GeminiQuota{SharedRPD: 2000, SharedRPM: 120}, Cooldown: 5 * time.Minute}, + + // --- GCP Code Assist (shared pool) --- + GeminiTierGCPStandard: {Quota: GeminiQuota{SharedRPD: 1500, SharedRPM: 120}, Cooldown: 5 * time.Minute}, + GeminiTierGCPEnterprise: {Quota: GeminiQuota{SharedRPD: 2000, SharedRPM: 120}, Cooldown: 5 * time.Minute}, }, } } @@ -149,11 +221,22 @@ func (p *GeminiQuotaPolicy) ApplyOverrides(tiers map[string]config.GeminiTierQuo if !ok { policy = GeminiTierPolicy{Cooldown: 5 * time.Minute} } + // Backward-compatible overrides: + // - If the tier uses shared quota, interpret pro_rpd as shared_rpd. + // - Otherwise apply per-model overrides. if override.ProRPD != nil { - policy.Quota.ProRPD = clampGeminiQuotaInt64(*override.ProRPD) + if policy.Quota.SharedRPD > 0 { + policy.Quota.SharedRPD = clampGeminiQuotaInt64WithUnlimited(*override.ProRPD) + } else { + policy.Quota.ProRPD = clampGeminiQuotaInt64WithUnlimited(*override.ProRPD) + } } if override.FlashRPD != nil { - policy.Quota.FlashRPD = clampGeminiQuotaInt64(*override.FlashRPD) + if policy.Quota.SharedRPD > 0 { + // No separate flash RPD for shared tiers. + } else { + policy.Quota.FlashRPD = clampGeminiQuotaInt64WithUnlimited(*override.FlashRPD) + } } if override.CooldownMinutes != nil { minutes := clampGeminiQuotaInt(*override.CooldownMinutes) @@ -163,10 +246,51 @@ func (p *GeminiQuotaPolicy) ApplyOverrides(tiers map[string]config.GeminiTierQuo } } -func (p *GeminiQuotaPolicy) QuotaForTier(tierID string) (GeminiDailyQuota, bool) { +func (p *GeminiQuotaPolicy) ApplyQuotaRulesOverrides(rules map[string]geminiQuotaRuleOverride) { + if p == nil || len(rules) == 0 { + return + } + for rawID, override := range rules { + tierID := normalizeGeminiTierID(rawID) + if tierID == "" { + continue + } + policy, ok := p.tiers[tierID] + if !ok { + policy = GeminiTierPolicy{Cooldown: 5 * time.Minute} + } + + if override.SharedRPD != nil { + policy.Quota.SharedRPD = clampGeminiQuotaInt64WithUnlimited(*override.SharedRPD) + } + if override.SharedRPM != nil { + policy.Quota.SharedRPM = clampGeminiQuotaRPM(*override.SharedRPM) + } + if override.GeminiPro != nil { + if override.GeminiPro.RPD != nil { + policy.Quota.ProRPD = clampGeminiQuotaInt64WithUnlimited(*override.GeminiPro.RPD) + } + if override.GeminiPro.RPM != nil { + policy.Quota.ProRPM = clampGeminiQuotaRPM(*override.GeminiPro.RPM) + } + } + if override.GeminiFlash != nil { + if override.GeminiFlash.RPD != nil { + policy.Quota.FlashRPD = clampGeminiQuotaInt64WithUnlimited(*override.GeminiFlash.RPD) + } + if override.GeminiFlash.RPM != nil { + policy.Quota.FlashRPM = clampGeminiQuotaRPM(*override.GeminiFlash.RPM) + } + } + + p.tiers[tierID] = policy + } +} + +func (p *GeminiQuotaPolicy) QuotaForTier(tierID string) (GeminiQuota, bool) { policy, ok := p.policyForTier(tierID) if !ok { - return GeminiDailyQuota{}, false + return GeminiQuota{}, false } return policy.Quota, true } @@ -184,22 +308,43 @@ func (p *GeminiQuotaPolicy) policyForTier(tierID string) (GeminiTierPolicy, bool return GeminiTierPolicy{}, false } normalized := normalizeGeminiTierID(tierID) - if normalized == "" { - normalized = "LEGACY" - } if policy, ok := p.tiers[normalized]; ok { return policy, true } - policy, ok := p.tiers["LEGACY"] - return policy, ok + return GeminiTierPolicy{}, false } func normalizeGeminiTierID(tierID string) string { - return strings.ToUpper(strings.TrimSpace(tierID)) + tierID = strings.TrimSpace(tierID) + if tierID == "" { + return "" + } + // Prefer canonical mapping (handles legacy tier strings). + if canonical := canonicalGeminiTierID(tierID); canonical != "" { + return canonical + } + // Accept older policy keys that used uppercase names. + switch strings.ToUpper(tierID) { + case "AISTUDIO_FREE": + return GeminiTierAIStudioFree + case "AISTUDIO_PAID": + return GeminiTierAIStudioPaid + case "GOOGLE_ONE_FREE": + return GeminiTierGoogleOneFree + case "GOOGLE_AI_PRO": + return GeminiTierGoogleAIPro + case "GOOGLE_AI_ULTRA": + return GeminiTierGoogleAIUltra + case "GCP_STANDARD": + return GeminiTierGCPStandard + case "GCP_ENTERPRISE": + return GeminiTierGCPEnterprise + } + return strings.ToLower(tierID) } -func clampGeminiQuotaInt64(value int64) int64 { - if value < 0 { +func clampGeminiQuotaInt64WithUnlimited(value int64) int64 { + if value < -1 { return 0 } return value @@ -212,11 +357,46 @@ func clampGeminiQuotaInt(value int) int { return value } +func clampGeminiQuotaRPM(value int64) int64 { + if value < 0 { + return 0 + } + return value +} + func geminiCooldownForTier(tierID string) time.Duration { policy := newGeminiQuotaPolicy() return policy.CooldownForTier(tierID) } +func geminiQuotaTierKeyForAccount(account *Account) string { + if account == nil || account.Platform != PlatformGemini { + return "" + } + + // Note: GeminiOAuthType() already defaults legacy (project_id present) to code_assist. + oauthType := strings.ToLower(strings.TrimSpace(account.GeminiOAuthType())) + rawTier := strings.TrimSpace(account.GeminiTierID()) + + // Prefer the canonical tier stored in credentials. + if tierID := canonicalGeminiTierIDForOAuthType(oauthType, rawTier); tierID != "" && tierID != GeminiTierGoogleOneUnknown { + return tierID + } + + // Fallback defaults when tier_id is missing or unknown. + switch oauthType { + case "google_one": + return GeminiTierGoogleOneFree + case "code_assist": + return GeminiTierGCPStandard + case "ai_studio": + return GeminiTierAIStudioFree + default: + // API Key accounts (type=apikey) have empty oauth_type and are treated as AI Studio. + return GeminiTierAIStudioFree + } +} + func geminiModelClassFromName(model string) geminiModelClass { name := strings.ToLower(strings.TrimSpace(model)) if strings.Contains(name, "flash") || strings.Contains(name, "lite") { diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 2282bcfa..196f1643 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -92,7 +92,7 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc // PreCheckUsage proactively checks local quota before dispatching a request. // Returns false when the account should be skipped. func (s *RateLimitService) PreCheckUsage(ctx context.Context, account *Account, requestedModel string) (bool, error) { - if account == nil || !account.IsGeminiCodeAssist() || strings.TrimSpace(requestedModel) == "" { + if account == nil || account.Platform != PlatformGemini { return true, nil } if s.usageRepo == nil || s.geminiQuotaService == nil { @@ -104,44 +104,99 @@ func (s *RateLimitService) PreCheckUsage(ctx context.Context, account *Account, return true, nil } - var limit int64 - switch geminiModelClassFromName(requestedModel) { - case geminiModelFlash: - limit = quota.FlashRPD - default: - limit = quota.ProRPD - } - if limit <= 0 { - return true, nil - } - now := time.Now() - start := geminiDailyWindowStart(now) - totals, ok := s.getGeminiUsageTotals(account.ID, start, now) - if !ok { - stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID) - if err != nil { - return true, err + modelClass := geminiModelClassFromName(requestedModel) + + // 1) Daily quota precheck (RPD; resets at PST midnight) + { + var limit int64 + if quota.SharedRPD > 0 { + limit = quota.SharedRPD + } else { + switch modelClass { + case geminiModelFlash: + limit = quota.FlashRPD + default: + limit = quota.ProRPD + } + } + + if limit > 0 { + start := geminiDailyWindowStart(now) + totals, ok := s.getGeminiUsageTotals(account.ID, start, now) + if !ok { + stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID) + if err != nil { + return true, err + } + totals = geminiAggregateUsage(stats) + s.setGeminiUsageTotals(account.ID, start, now, totals) + } + + var used int64 + if quota.SharedRPD > 0 { + used = totals.ProRequests + totals.FlashRequests + } else { + switch modelClass { + case geminiModelFlash: + used = totals.FlashRequests + default: + used = totals.ProRequests + } + } + + if used >= limit { + resetAt := geminiDailyResetTime(now) + // NOTE: + // - This is a local precheck to reduce upstream 429s. + // - Do NOT mark the account as rate-limited here; rate_limit_reset_at should reflect real upstream 429s. + log.Printf("[Gemini PreCheck] Account %d reached daily quota (%d/%d), skip until %v", account.ID, used, limit, resetAt) + return false, nil + } } - totals = geminiAggregateUsage(stats) - s.setGeminiUsageTotals(account.ID, start, now, totals) } - var used int64 - switch geminiModelClassFromName(requestedModel) { - case geminiModelFlash: - used = totals.FlashRequests - default: - used = totals.ProRequests - } - - if used >= limit { - resetAt := geminiDailyResetTime(now) - if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil { - log.Printf("SetRateLimited failed for account %d: %v", account.ID, err) + // 2) Minute quota precheck (RPM; fixed window current minute) + { + var limit int64 + if quota.SharedRPM > 0 { + limit = quota.SharedRPM + } else { + switch modelClass { + case geminiModelFlash: + limit = quota.FlashRPM + default: + limit = quota.ProRPM + } + } + + if limit > 0 { + start := now.Truncate(time.Minute) + stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID) + if err != nil { + return true, err + } + totals := geminiAggregateUsage(stats) + + var used int64 + if quota.SharedRPM > 0 { + used = totals.ProRequests + totals.FlashRequests + } else { + switch modelClass { + case geminiModelFlash: + used = totals.FlashRequests + default: + used = totals.ProRequests + } + } + + if used >= limit { + resetAt := start.Add(time.Minute) + // Do not persist "rate limited" status from local precheck. See note above. + log.Printf("[Gemini PreCheck] Account %d reached minute quota (%d/%d), skip until %v", account.ID, used, limit, resetAt) + return false, nil + } } - log.Printf("[Gemini PreCheck] Account %d reached daily quota (%d/%d), rate limited until %v", account.ID, used, limit, resetAt) - return false, nil } return true, nil @@ -186,7 +241,10 @@ func (s *RateLimitService) GeminiCooldown(ctx context.Context, account *Account) if account == nil { return 5 * time.Minute } - return s.geminiQuotaService.CooldownForTier(ctx, account.GeminiTierID()) + if s.geminiQuotaService == nil { + return 5 * time.Minute + } + return s.geminiQuotaService.CooldownForAccount(ctx, account) } // handleAuthError 处理认证类错误(401/403)，停止账号调度 diff --git a/frontend/src/api/admin/gemini.ts b/frontend/src/api/admin/gemini.ts index 4b40dc17..6113f468 100644 --- a/frontend/src/api/admin/gemini.ts +++ b/frontend/src/api/admin/gemini.ts @@ -20,6 +20,7 @@ export interface GeminiAuthUrlRequest { proxy_id?: number project_id?: string oauth_type?: 'code_assist' | 'google_one' | 'ai_studio' + tier_id?: string } export interface GeminiExchangeCodeRequest { @@ -28,6 +29,7 @@ export interface GeminiExchangeCodeRequest { code: string proxy_id?: number oauth_type?: 'code_assist' | 'google_one' | 'ai_studio' + tier_id?: string } export type GeminiTokenInfo = { diff --git a/frontend/src/components/account/AccountQuotaInfo.vue b/frontend/src/components/account/AccountQuotaInfo.vue index 9f919386..2f7f80de 100644 --- a/frontend/src/components/account/AccountQuotaInfo.vue +++ b/frontend/src/components/account/AccountQuotaInfo.vue @@ -65,32 +65,33 @@ const tierLabel = computed(() => { const creds = props.account.credentials as GeminiCredentials | undefined if (isCodeAssist.value) { - // GCP Code Assist: 显示 GCP tier - const tierMap: Record = { - LEGACY: 'Free', - PRO: 'Pro', - ULTRA: 'Ultra', - 'standard-tier': 'Standard', - 'pro-tier': 'Pro', - 'ultra-tier': 'Ultra' - } - return tierMap[creds?.tier_id || ''] || (creds?.tier_id ? 'GCP' : 'Unknown') + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'gcp_enterprise') return 'GCP Enterprise' + if (tier === 'gcp_standard') return 'GCP Standard' + // Backward compatibility + const upper = (creds?.tier_id || '').toString().trim().toUpperCase() + if (upper.includes('ULTRA') || upper.includes('ENTERPRISE')) return 'GCP Enterprise' + if (upper) return `GCP ${upper}` + return 'GCP' } if (isGoogleOne.value) { - // Google One: tier 映射 - const tierMap: Record = { - AI_PREMIUM: 'AI Premium', - GOOGLE_ONE_STANDARD: 'Standard', - GOOGLE_ONE_BASIC: 'Basic', - FREE: 'Free', - GOOGLE_ONE_UNKNOWN: 'Personal', - GOOGLE_ONE_UNLIMITED: 'Unlimited' - } - return tierMap[creds?.tier_id || ''] || 'Personal' + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'google_ai_ultra') return 'Google AI Ultra' + if (tier === 'google_ai_pro') return 'Google AI Pro' + if (tier === 'google_one_free') return 'Google One Free' + // Backward compatibility + const upper = (creds?.tier_id || '').toString().trim().toUpperCase() + if (upper === 'AI_PREMIUM') return 'Google AI Pro' + if (upper === 'GOOGLE_ONE_UNLIMITED') return 'Google AI Ultra' + if (upper) return `Google One ${upper}` + return 'Google One' } // API Key: 显示 AI Studio + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'aistudio_paid') return 'AI Studio Pay-as-you-go' + if (tier === 'aistudio_free') return 'AI Studio Free Tier' return 'AI Studio' }) @@ -99,35 +100,31 @@ const tierBadgeClass = computed(() => { const creds = props.account.credentials as GeminiCredentials | undefined if (isCodeAssist.value) { - // GCP Code Assist 样式 - const tierColorMap: Record = { - LEGACY: 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300', - PRO: 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300', - ULTRA: 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300', - 'standard-tier': 'bg-green-100 text-green-600 dark:bg-green-900/40 dark:text-green-300', - 'pro-tier': 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300', - 'ultra-tier': 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' - } - return ( - tierColorMap[creds?.tier_id || ''] || - 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' - ) + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'gcp_enterprise') return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + if (tier === 'gcp_standard') return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' + // Backward compatibility + const upper = (creds?.tier_id || '').toString().trim().toUpperCase() + if (upper.includes('ULTRA') || upper.includes('ENTERPRISE')) return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' } if (isGoogleOne.value) { - // Google One tier 样式 - const tierColorMap: Record = { - AI_PREMIUM: 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300', - GOOGLE_ONE_STANDARD: 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300', - GOOGLE_ONE_BASIC: 'bg-green-100 text-green-600 dark:bg-green-900/40 dark:text-green-300', - FREE: 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300', - GOOGLE_ONE_UNKNOWN: 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300', - GOOGLE_ONE_UNLIMITED: 'bg-amber-100 text-amber-600 dark:bg-amber-900/40 dark:text-amber-300' - } - return tierColorMap[creds?.tier_id || ''] || 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'google_ai_ultra') return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + if (tier === 'google_ai_pro') return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' + if (tier === 'google_one_free') return 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' + // Backward compatibility + const upper = (creds?.tier_id || '').toString().trim().toUpperCase() + if (upper === 'GOOGLE_ONE_UNLIMITED') return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + if (upper === 'AI_PREMIUM') return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' + return 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' } // AI Studio 默认样式：蓝色 + const tier = (creds?.tier_id || '').toString().trim().toLowerCase() + if (tier === 'aistudio_paid') return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' + if (tier === 'aistudio_free') return 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' }) diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index ddae4e3e..c0212c5a 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -241,23 +241,16 @@
{{ error }}
- +
-

* {{ t('admin.accounts.gemini.quotaPolicy.simulatedNote') || 'Simulated quota' }} @@ -288,7 +281,7 @@ import { ref, computed, onMounted } from 'vue' import { useI18n } from 'vue-i18n' import { adminAPI } from '@/api/admin' -import type { Account, AccountUsageInfo, GeminiCredentials } from '@/types' +import type { Account, AccountUsageInfo, GeminiCredentials, WindowStats } from '@/types' import UsageProgressBar from './UsageProgressBar.vue' import AccountQuotaInfo from './AccountQuotaInfo.vue' @@ -303,16 +296,18 @@ const error = ref(null) const usageInfo = ref(null) // Show usage windows for OAuth and Setup Token accounts -const showUsageWindows = computed( - () => props.account.type === 'oauth' || props.account.type === 'setup-token' -) +const showUsageWindows = computed(() => { + // Gemini: we can always compute local usage windows from DB logs (simulated quotas). + if (props.account.platform === 'gemini') return true + return props.account.type === 'oauth' || props.account.type === 'setup-token' +}) const shouldFetchUsage = computed(() => { if (props.account.platform === 'anthropic') { return props.account.type === 'oauth' || props.account.type === 'setup-token' } if (props.account.platform === 'gemini') { - return props.account.type === 'oauth' + return true } if (props.account.platform === 'antigravity') { return props.account.type === 'oauth' @@ -322,8 +317,12 @@ const shouldFetchUsage = computed(() => { const geminiUsageAvailable = computed(() => { return ( + !!usageInfo.value?.gemini_shared_daily || !!usageInfo.value?.gemini_pro_daily || - !!usageInfo.value?.gemini_flash_daily + !!usageInfo.value?.gemini_flash_daily || + !!usageInfo.value?.gemini_shared_minute || + !!usageInfo.value?.gemini_pro_minute || + !!usageInfo.value?.gemini_flash_minute ) }) @@ -569,6 +568,12 @@ const geminiTier = computed(() => { return creds?.tier_id || null }) +const geminiOAuthType = computed(() => { + if (props.account.platform !== 'gemini') return null + const creds = props.account.credentials as GeminiCredentials | undefined + return (creds?.oauth_type || '').trim() || null +}) + // Gemini 是否为 Code Assist OAuth const isGeminiCodeAssist = computed(() => { if (props.account.platform !== 'gemini') return false @@ -576,109 +581,208 @@ const isGeminiCodeAssist = computed(() => { return creds?.oauth_type === 'code_assist' || (!creds?.oauth_type && !!creds?.project_id) }) -// Gemini 认证类型 + Tier 组合标签（简洁版） -const geminiAuthTypeLabel = computed(() => { - const creds = props.account.credentials as GeminiCredentials | undefined - const oauthType = creds?.oauth_type +const geminiChannelShort = computed((): 'ai studio' | 'gcp' | 'google one' | 'client' | null => { + if (props.account.platform !== 'gemini') return null - // For API Key accounts, don't show auth type label - if (props.account.type !== 'oauth') return null + // API Key accounts are AI Studio. + if (props.account.type === 'apikey') return 'ai studio' - if (oauthType === 'google_one') { - // Google One: show "Google One" + tier - const tierMap: Record = { - AI_PREMIUM: 'AI Premium', - GOOGLE_ONE_STANDARD: 'Standard', - GOOGLE_ONE_BASIC: 'Basic', - FREE: 'Free', - GOOGLE_ONE_UNKNOWN: 'Personal', - GOOGLE_ONE_UNLIMITED: 'Unlimited' - } - const tierLabel = geminiTier.value ? tierMap[geminiTier.value] || 'Personal' : 'Personal' - return `Google One ${tierLabel}` - } else if (oauthType === 'code_assist' || (!oauthType && isGeminiCodeAssist.value)) { - // Code Assist: show "GCP" + tier - const tierMap: Record = { - LEGACY: 'Free', - PRO: 'Pro', - ULTRA: 'Ultra' - } - const tierLabel = geminiTier.value ? tierMap[geminiTier.value] || 'Free' : 'Free' - return `GCP ${tierLabel}` - } else if (oauthType === 'ai_studio') { - // 自定义 OAuth Client: show "Client" (no tier) - return 'Client' + if (geminiOAuthType.value === 'google_one') return 'google one' + if (isGeminiCodeAssist.value) return 'gcp' + if (geminiOAuthType.value === 'ai_studio') return 'client' + + // Fallback (unknown legacy data): treat as AI Studio. + return 'ai studio' +}) + +const geminiUserLevel = computed((): string | null => { + if (props.account.platform !== 'gemini') return null + + const tier = (geminiTier.value || '').toString().trim() + const tierLower = tier.toLowerCase() + const tierUpper = tier.toUpperCase() + + // Google One: free / pro / ultra + if (geminiOAuthType.value === 'google_one') { + if (tierLower === 'google_one_free') return 'free' + if (tierLower === 'google_ai_pro') return 'pro' + if (tierLower === 'google_ai_ultra') return 'ultra' + + // Backward compatibility (legacy tier markers) + if (tierUpper === 'AI_PREMIUM' || tierUpper === 'GOOGLE_ONE_STANDARD') return 'pro' + if (tierUpper === 'GOOGLE_ONE_UNLIMITED') return 'ultra' + if (tierUpper === 'FREE' || tierUpper === 'GOOGLE_ONE_BASIC' || tierUpper === 'GOOGLE_ONE_UNKNOWN' || tierUpper === '') return 'free' + + return null + } + + // GCP Code Assist: standard / enterprise + if (isGeminiCodeAssist.value) { + if (tierLower === 'gcp_enterprise') return 'enterprise' + if (tierLower === 'gcp_standard') return 'standard' + + // Backward compatibility + if (tierUpper.includes('ULTRA') || tierUpper.includes('ENTERPRISE')) return 'enterprise' + return 'standard' + } + + // AI Studio (API Key) and Client OAuth: free / paid + if (props.account.type === 'apikey' || geminiOAuthType.value === 'ai_studio') { + if (tierLower === 'aistudio_paid') return 'paid' + if (tierLower === 'aistudio_free') return 'free' + + // Backward compatibility + if (tierUpper.includes('PAID') || tierUpper.includes('PAYG') || tierUpper.includes('PAY')) return 'paid' + if (tierUpper.includes('FREE')) return 'free' + if (props.account.type === 'apikey') return 'free' + return null } return null }) +// Gemini 认证类型（按要求：授权方式简称 + 用户等级） +const geminiAuthTypeLabel = computed(() => { + if (props.account.platform !== 'gemini') return null + if (!geminiChannelShort.value) return null + return geminiUserLevel.value ? `${geminiChannelShort.value} ${geminiUserLevel.value}` : geminiChannelShort.value +}) + // Gemini 账户类型徽章样式（统一样式） const geminiTierClass = computed(() => { - const creds = props.account.credentials as GeminiCredentials | undefined - const oauthType = creds?.oauth_type + // Use channel+level to choose a stable color without depending on raw tier_id variants. + const channel = geminiChannelShort.value + const level = geminiUserLevel.value - // Client (自定义 OAuth): 使用蓝色（与 AI Studio 一致） - if (oauthType === 'ai_studio') { + if (channel === 'client' || channel === 'ai studio') { return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' } - if (!geminiTier.value) return '' - - const isGoogleOne = creds?.oauth_type === 'google_one' - - if (isGoogleOne) { - // Google One tier 颜色 - const colorMap: Record = { - AI_PREMIUM: 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300', - GOOGLE_ONE_STANDARD: 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300', - GOOGLE_ONE_BASIC: 'bg-green-100 text-green-600 dark:bg-green-900/40 dark:text-green-300', - FREE: 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300', - GOOGLE_ONE_UNKNOWN: 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300', - GOOGLE_ONE_UNLIMITED: 'bg-amber-100 text-amber-600 dark:bg-amber-900/40 dark:text-amber-300' - } - return colorMap[geminiTier.value] || 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' + if (channel === 'google one') { + if (level === 'ultra') return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + if (level === 'pro') return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' + return 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' } - // Code Assist tier 颜色 - switch (geminiTier.value) { - case 'LEGACY': - return 'bg-gray-100 text-gray-600 dark:bg-gray-700 dark:text-gray-300' - case 'PRO': - return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' - case 'ULTRA': - return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' - default: - return '' + if (channel === 'gcp') { + if (level === 'enterprise') return 'bg-purple-100 text-purple-600 dark:bg-purple-900/40 dark:text-purple-300' + return 'bg-blue-100 text-blue-600 dark:bg-blue-900/40 dark:text-blue-300' } + + return '' }) // Gemini 配额政策信息 const geminiQuotaPolicyChannel = computed(() => { + if (geminiOAuthType.value === 'google_one') { + return t('admin.accounts.gemini.quotaPolicy.rows.googleOne.channel') + } if (isGeminiCodeAssist.value) { - return t('admin.accounts.gemini.quotaPolicy.rows.cli.channel') + return t('admin.accounts.gemini.quotaPolicy.rows.gcp.channel') } return t('admin.accounts.gemini.quotaPolicy.rows.aiStudio.channel') }) const geminiQuotaPolicyLimits = computed(() => { - if (isGeminiCodeAssist.value) { - if (geminiTier.value === 'PRO' || geminiTier.value === 'ULTRA') { - return t('admin.accounts.gemini.quotaPolicy.rows.cli.limitsPremium') + const tierLower = (geminiTier.value || '').toString().trim().toLowerCase() + + if (geminiOAuthType.value === 'google_one') { + if (tierLower === 'google_ai_ultra' || geminiUserLevel.value === 'ultra') { + return t('admin.accounts.gemini.quotaPolicy.rows.googleOne.limitsUltra') } - return t('admin.accounts.gemini.quotaPolicy.rows.cli.limitsFree') + if (tierLower === 'google_ai_pro' || geminiUserLevel.value === 'pro') { + return t('admin.accounts.gemini.quotaPolicy.rows.googleOne.limitsPro') + } + return t('admin.accounts.gemini.quotaPolicy.rows.googleOne.limitsFree') + } + + if (isGeminiCodeAssist.value) { + if (tierLower === 'gcp_enterprise' || geminiUserLevel.value === 'enterprise') { + return t('admin.accounts.gemini.quotaPolicy.rows.gcp.limitsEnterprise') + } + return t('admin.accounts.gemini.quotaPolicy.rows.gcp.limitsStandard') + } + + // AI Studio (API Key / custom OAuth) + if (tierLower === 'aistudio_paid' || geminiUserLevel.value === 'paid') { + return t('admin.accounts.gemini.quotaPolicy.rows.aiStudio.limitsPaid') } - // AI Studio - 默认显示免费层限制 return t('admin.accounts.gemini.quotaPolicy.rows.aiStudio.limitsFree') }) const geminiQuotaPolicyDocsUrl = computed(() => { - if (isGeminiCodeAssist.value) { - return 'https://cloud.google.com/products/gemini/code-assist#pricing' + if (geminiOAuthType.value === 'google_one' || isGeminiCodeAssist.value) { + return 'https://developers.google.com/gemini-code-assist/resources/quotas' } return 'https://ai.google.dev/pricing' }) +const geminiUsesSharedDaily = computed(() => { + if (props.account.platform !== 'gemini') return false + // Per requirement: Google One & GCP are shared RPD pools (no per-model breakdown). + return ( + !!usageInfo.value?.gemini_shared_daily || + !!usageInfo.value?.gemini_shared_minute || + geminiOAuthType.value === 'google_one' || + isGeminiCodeAssist.value + ) +}) + +const geminiUsageBars = computed(() => { + if (props.account.platform !== 'gemini') return [] + if (!usageInfo.value) return [] + + const bars: Array<{ + key: string + label: string + utilization: number + resetsAt: string | null + windowStats?: WindowStats | null + color: 'indigo' | 'emerald' + }> = [] + + if (geminiUsesSharedDaily.value) { + const sharedDaily = usageInfo.value.gemini_shared_daily + if (sharedDaily) { + bars.push({ + key: 'shared_daily', + label: '1d', + utilization: sharedDaily.utilization, + resetsAt: sharedDaily.resets_at, + windowStats: sharedDaily.window_stats, + color: 'indigo' + }) + } + return bars + } + + const pro = usageInfo.value.gemini_pro_daily + if (pro) { + bars.push({ + key: 'pro_daily', + label: 'pro', + utilization: pro.utilization, + resetsAt: pro.resets_at, + windowStats: pro.window_stats, + color: 'indigo' + }) + } + + const flash = usageInfo.value.gemini_flash_daily + if (flash) { + bars.push({ + key: 'flash_daily', + label: 'flash', + utilization: flash.utilization, + resetsAt: flash.resets_at, + windowStats: flash.window_stats, + color: 'emerald' + }) + } + + return bars +}) + // 账户类型显示标签 const antigravityTierLabel = computed(() => { switch (antigravityTier.value) { diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 50bcf690..33342e7a 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -653,6 +653,41 @@

+ +
+ +
+ + + + + +
+

{{ t('admin.accounts.gemini.tier.hint') }}

+
+
{{ apiKeyHint }}

+ +
+ + +

{{ t('admin.accounts.gemini.tier.aiStudioHint') }}

+
+
@@ -1816,6 +1861,24 @@ const geminiOAuthType = ref<'code_assist' | 'google_one' | 'ai_studio'>('google_ const geminiAIStudioOAuthEnabled = ref(false) const showAdvancedOAuth = ref(false) +// Gemini tier selection (used as fallback when auto-detection is unavailable/fails) +const geminiTierGoogleOne = ref<'google_one_free' | 'google_ai_pro' | 'google_ai_ultra'>('google_one_free') +const geminiTierGcp = ref<'gcp_standard' | 'gcp_enterprise'>('gcp_standard') +const geminiTierAIStudio = ref<'aistudio_free' | 'aistudio_paid'>('aistudio_free') + +const geminiSelectedTier = computed(() => { + if (form.platform !== 'gemini') return '' + if (accountCategory.value === 'apikey') return geminiTierAIStudio.value + switch (geminiOAuthType.value) { + case 'google_one': + return geminiTierGoogleOne.value + case 'code_assist': + return geminiTierGcp.value + default: + return geminiTierAIStudio.value + } +}) + const geminiQuotaDocs = { codeAssist: 'https://developers.google.com/gemini-code-assist/resources/quotas', aiStudio: 'https://ai.google.dev/pricing', @@ -2143,6 +2206,9 @@ const resetForm = () => { tempUnschedEnabled.value = false tempUnschedRules.value = [] geminiOAuthType.value = 'code_assist' + geminiTierGoogleOne.value = 'google_one_free' + geminiTierGcp.value = 'gcp_standard' + geminiTierAIStudio.value = 'aistudio_free' oauth.resetState() openaiOAuth.resetState() geminiOAuth.resetState() @@ -2184,6 +2250,9 @@ const handleSubmit = async () => { base_url: apiKeyBaseUrl.value.trim() || defaultBaseUrl, api_key: apiKeyValue.value.trim() } + if (form.platform === 'gemini') { + credentials.tier_id = geminiTierAIStudio.value + } // Add model mapping if configured const modelMapping = buildModelMappingObject(modelRestrictionMode.value, allowedModels.value, modelMappings.value) @@ -2237,7 +2306,12 @@ const handleGenerateUrl = async () => { if (form.platform === 'openai') { await openaiOAuth.generateAuthUrl(form.proxy_id) } else if (form.platform === 'gemini') { - await geminiOAuth.generateAuthUrl(form.proxy_id, oauthFlowRef.value?.projectId, geminiOAuthType.value) + await geminiOAuth.generateAuthUrl( + form.proxy_id, + oauthFlowRef.value?.projectId, + geminiOAuthType.value, + geminiSelectedTier.value + ) } else if (form.platform === 'antigravity') { await antigravityOAuth.generateAuthUrl(form.proxy_id) } else { @@ -2318,7 +2392,8 @@ const handleGeminiExchange = async (authCode: string) => { sessionId: geminiOAuth.sessionId.value, state: stateToUse, proxyId: form.proxy_id, - oauthType: geminiOAuthType.value + oauthType: geminiOAuthType.value, + tierId: geminiSelectedTier.value }) if (!tokenInfo) return diff --git a/frontend/src/components/account/ReAuthAccountModal.vue b/frontend/src/components/account/ReAuthAccountModal.vue index 3ecefbfb..26320451 100644 --- a/frontend/src/components/account/ReAuthAccountModal.vue +++ b/frontend/src/components/account/ReAuthAccountModal.vue @@ -88,7 +88,35 @@
{{ t('admin.accounts.oauth.gemini.oauthTypeLabel') }} -
+
+ +