feat(gemini): 优化 OAuth 和配额展示

主要改进:
- 修复 google_one OAuth scopes 配置问题
- 添加 Gemini 账号配额展示组件
- 优化 Code Assist 类型检测逻辑
- 添加 OAuth 测试用例
This commit is contained in:
ianshaw
2026-01-03 06:32:04 -08:00
parent 26438f7232
commit 26106eb0ac
11 changed files with 363 additions and 109 deletions

View File

@@ -251,8 +251,20 @@ func inferGoogleOneTier(storageBytes int64) string {
return TierGoogleOneUnknown
}
// FetchGoogleOneTier fetches Google One tier from Drive API
// fetchGoogleOneTier fetches Google One tier from Drive API or LoadCodeAssist API
func (s *GeminiOAuthService) FetchGoogleOneTier(ctx context.Context, accessToken, proxyURL string) (string, *geminicli.DriveStorageInfo, error) {
// First try LoadCodeAssist API (works for accounts with GCP projects)
if s.codeAssist != nil {
loadResp, err := s.codeAssist.LoadCodeAssist(ctx, accessToken, proxyURL, nil)
if err == nil && loadResp != nil {
if tier := loadResp.GetTier(); tier != "" {
fmt.Printf("[GeminiOAuth] Got tier from LoadCodeAssist: %s\n", tier)
return tier, nil, nil
}
}
}
// Fallback to Drive API (requires drive.readonly scope)
driveClient := geminicli.NewDriveClient()
storageInfo, err := driveClient.GetStorageQuota(ctx, accessToken, proxyURL)
@@ -422,12 +434,15 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
}
case "google_one":
// Attempt to fetch Drive storage tier
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, tokenResp.AccessToken, proxyURL)
var storageInfo *geminicli.DriveStorageInfo
var err error
tierID, storageInfo, err = s.FetchGoogleOneTier(ctx, tokenResp.AccessToken, proxyURL)
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
}
fmt.Printf("[GeminiOAuth] Google One tierID after fetch: %s\n", tierID)
// Store Drive info in extra field for caching
if storageInfo != nil {
@@ -452,7 +467,7 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
}
// ai_studio 模式不设置 tierID保持为空
return &GeminiTokenInfo{
result := &GeminiTokenInfo{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
TokenType: tokenResp.TokenType,
@@ -462,7 +477,9 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
ProjectID: projectID,
TierID: tierID,
OAuthType: oauthType,
}, nil
}
fmt.Printf("[GeminiOAuth] ExchangeCode returning tierID: %s\n", result.TierID)
return result, nil
}
func (s *GeminiOAuthService) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*GeminiTokenInfo, error) {
@@ -669,6 +686,9 @@ func (s *GeminiOAuthService) BuildAccountCredentials(tokenInfo *GeminiTokenInfo)
// Validate tier_id before storing
if err := validateTierID(tokenInfo.TierID); err == nil {
creds["tier_id"] = tokenInfo.TierID
fmt.Printf("[GeminiOAuth] Storing tier_id: %s\n", tokenInfo.TierID)
} else {
fmt.Printf("[GeminiOAuth] Invalid tier_id %s: %v\n", tokenInfo.TierID, err)
}
// Silently skip invalid tier_id (don't block account creation)
}
@@ -698,7 +718,13 @@ func (s *GeminiOAuthService) fetchProjectID(ctx context.Context, accessToken, pr
// Extract tierID from response (works whether CloudAICompanionProject is set or not)
tierID := "LEGACY"
if loadResp != nil {
tierID = extractTierIDFromAllowedTiers(loadResp.AllowedTiers)
// First try to get tier from currentTier/paidTier fields
if tier := loadResp.GetTier(); tier != "" {
tierID = tier
} else {
// Fallback to extracting from allowedTiers
tierID = extractTierIDFromAllowedTiers(loadResp.AllowedTiers)
}
}
// If LoadCodeAssist returned a project, use it