feat(gemini): 优化 OAuth 和配额展示
主要改进: - 修复 google_one OAuth scopes 配置问题 - 添加 Gemini 账号配额展示组件 - 优化 Code Assist 类型检测逻辑 - 添加 OAuth 测试用例
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user