979114db45
问题：Google One Ultra 等已注册用户在 OAuth 授权时，如果 LoadCodeAssist 返回了 currentTier/paidTier 但没有返回 cloudaicompanionProject，之前的 逻辑会继续调用 onboardUser，导致 INVALID_ARGUMENT 错误。 修复：对齐 Gemini CLI 的处理逻辑： - 当检测到用户已注册（有 currentTier/paidTier）时，不再调用 onboardUser - 先尝试从 Cloud Resource Manager 获取可用项目 - 如果仍无法获取，返回友好的错误提示，引导用户手动填写 Project ID 这个修复解决了 Google One 订阅用户无法正常授权的问题。
1101 lines
39 KiB
Go
1101 lines
39 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"regexp"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||
)
|
||
|
||
const (
|
||
// 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 (
|
||
GB = 1024 * 1024 * 1024
|
||
TB = 1024 * GB
|
||
|
||
StorageTierUnlimited = 100 * TB // 100TB
|
||
StorageTierAIPremium = 2 * TB // 2TB
|
||
StorageTierStandard = 200 * GB // 200GB
|
||
StorageTierBasic = 100 * GB // 100GB
|
||
StorageTierFree = 15 * GB // 15GB
|
||
)
|
||
|
||
type GeminiOAuthService struct {
|
||
sessionStore *geminicli.SessionStore
|
||
proxyRepo ProxyRepository
|
||
oauthClient GeminiOAuthClient
|
||
codeAssist GeminiCliCodeAssistClient
|
||
cfg *config.Config
|
||
}
|
||
|
||
type GeminiOAuthCapabilities struct {
|
||
AIStudioOAuthEnabled bool `json:"ai_studio_oauth_enabled"`
|
||
RequiredRedirectURIs []string `json:"required_redirect_uris"`
|
||
}
|
||
|
||
func NewGeminiOAuthService(
|
||
proxyRepo ProxyRepository,
|
||
oauthClient GeminiOAuthClient,
|
||
codeAssist GeminiCliCodeAssistClient,
|
||
cfg *config.Config,
|
||
) *GeminiOAuthService {
|
||
return &GeminiOAuthService{
|
||
sessionStore: geminicli.NewSessionStore(),
|
||
proxyRepo: proxyRepo,
|
||
oauthClient: oauthClient,
|
||
codeAssist: codeAssist,
|
||
cfg: cfg,
|
||
}
|
||
}
|
||
|
||
func (s *GeminiOAuthService) GetOAuthConfig() *GeminiOAuthCapabilities {
|
||
// AI Studio OAuth is only enabled when the operator configures a custom OAuth client.
|
||
clientID := strings.TrimSpace(s.cfg.Gemini.OAuth.ClientID)
|
||
clientSecret := strings.TrimSpace(s.cfg.Gemini.OAuth.ClientSecret)
|
||
enabled := clientID != "" && clientSecret != "" &&
|
||
(clientID != geminicli.GeminiCLIOAuthClientID || clientSecret != geminicli.GeminiCLIOAuthClientSecret)
|
||
|
||
return &GeminiOAuthCapabilities{
|
||
AIStudioOAuthEnabled: enabled,
|
||
RequiredRedirectURIs: []string{geminicli.AIStudioOAuthRedirectURI},
|
||
}
|
||
}
|
||
|
||
type GeminiAuthURLResult struct {
|
||
AuthURL string `json:"auth_url"`
|
||
SessionID string `json:"session_id"`
|
||
State string `json:"state"`
|
||
}
|
||
|
||
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)
|
||
}
|
||
codeVerifier, err := geminicli.GenerateCodeVerifier()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to generate code verifier: %w", err)
|
||
}
|
||
codeChallenge := geminicli.GenerateCodeChallenge(codeVerifier)
|
||
sessionID, err := geminicli.GenerateSessionID()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to generate session ID: %w", err)
|
||
}
|
||
|
||
var proxyURL string
|
||
if proxyID != nil {
|
||
proxy, err := s.proxyRepo.GetByID(ctx, *proxyID)
|
||
if err == nil && proxy != nil {
|
||
proxyURL = proxy.URL()
|
||
}
|
||
}
|
||
|
||
// OAuth client selection:
|
||
// - code_assist: always use built-in Gemini CLI OAuth client (public)
|
||
// - google_one: always use built-in Gemini CLI OAuth client (public)
|
||
// - 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" {
|
||
// Force use of built-in Gemini CLI OAuth client
|
||
oauthCfg.ClientID = ""
|
||
oauthCfg.ClientSecret = ""
|
||
}
|
||
|
||
session := &geminicli.OAuthSession{
|
||
State: state,
|
||
CodeVerifier: codeVerifier,
|
||
ProxyURL: proxyURL,
|
||
RedirectURI: redirectURI,
|
||
ProjectID: strings.TrimSpace(projectID),
|
||
TierID: canonicalGeminiTierIDForOAuthType(oauthType, tierID),
|
||
OAuthType: oauthType,
|
||
CreatedAt: time.Now(),
|
||
}
|
||
s.sessionStore.Set(sessionID, session)
|
||
|
||
effectiveCfg, err := geminicli.EffectiveOAuthConfig(oauthCfg, oauthType)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
isBuiltinClient := effectiveCfg.ClientID == geminicli.GeminiCLIOAuthClientID &&
|
||
effectiveCfg.ClientSecret == geminicli.GeminiCLIOAuthClientSecret
|
||
|
||
// AI Studio OAuth requires a user-provided OAuth client (built-in Gemini CLI client is scope-restricted).
|
||
if oauthType == "ai_studio" && isBuiltinClient {
|
||
return nil, fmt.Errorf("AI Studio OAuth requires a custom OAuth Client (GEMINI_OAUTH_CLIENT_ID / GEMINI_OAUTH_CLIENT_SECRET). If you don't want to configure an OAuth client, please use an AI Studio API Key account instead")
|
||
}
|
||
|
||
// Redirect URI strategy:
|
||
// - 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
|
||
}
|
||
session.RedirectURI = redirectURI
|
||
s.sessionStore.Set(sessionID, session)
|
||
|
||
authURL, err := geminicli.BuildAuthorizationURL(effectiveCfg, state, codeChallenge, redirectURI, session.ProjectID, oauthType)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return &GeminiAuthURLResult{
|
||
AuthURL: authURL,
|
||
SessionID: sessionID,
|
||
State: state,
|
||
}, nil
|
||
}
|
||
|
||
type GeminiExchangeCodeInput struct {
|
||
SessionID string
|
||
State string
|
||
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 {
|
||
AccessToken string `json:"access_token"`
|
||
RefreshToken string `json:"refresh_token"`
|
||
ExpiresIn int64 `json:"expires_in"`
|
||
ExpiresAt int64 `json:"expires_at"`
|
||
TokenType string `json:"token_type"`
|
||
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"` // Canonical tier id (e.g. google_one_free, gcp_standard, aistudio_free)
|
||
Extra map[string]any `json:"extra,omitempty"` // Drive metadata
|
||
}
|
||
|
||
// validateTierID validates tier_id format and length
|
||
func validateTierID(tierID string) error {
|
||
if tierID == "" {
|
||
return nil // Empty is allowed
|
||
}
|
||
if len(tierID) > 64 {
|
||
return fmt.Errorf("tier_id exceeds maximum length of 64 characters")
|
||
}
|
||
// Allow alphanumeric, underscore, hyphen, and slash (for tier paths)
|
||
if !regexp.MustCompile(`^[a-zA-Z0-9_/-]+$`).MatchString(tierID) {
|
||
return fmt.Errorf("tier_id contains invalid characters")
|
||
}
|
||
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 {
|
||
tierID := "LEGACY"
|
||
// First pass: look for default tier
|
||
for _, tier := range allowedTiers {
|
||
if tier.IsDefault && strings.TrimSpace(tier.ID) != "" {
|
||
tierID = strings.TrimSpace(tier.ID)
|
||
break
|
||
}
|
||
}
|
||
// Second pass: if still LEGACY, take first non-empty tier
|
||
if tierID == "LEGACY" {
|
||
for _, tier := range allowedTiers {
|
||
if strings.TrimSpace(tier.ID) != "" {
|
||
tierID = strings.TrimSpace(tier.ID)
|
||
break
|
||
}
|
||
}
|
||
}
|
||
return tierID
|
||
}
|
||
|
||
// inferGoogleOneTier infers Google One tier from Drive storage limit
|
||
func inferGoogleOneTier(storageBytes int64) string {
|
||
log.Printf("[GeminiOAuth] inferGoogleOneTier - input: %d bytes (%.2f TB)", storageBytes, float64(storageBytes)/float64(TB))
|
||
|
||
if storageBytes <= 0 {
|
||
log.Printf("[GeminiOAuth] inferGoogleOneTier - storageBytes <= 0, returning UNKNOWN")
|
||
return GeminiTierGoogleOneUnknown
|
||
}
|
||
|
||
if storageBytes > StorageTierUnlimited {
|
||
log.Printf("[GeminiOAuth] inferGoogleOneTier - > %d bytes (100TB), returning UNLIMITED", StorageTierUnlimited)
|
||
return GeminiTierGoogleAIUltra
|
||
}
|
||
if storageBytes >= StorageTierAIPremium {
|
||
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 GeminiTierGoogleOneFree
|
||
}
|
||
|
||
log.Printf("[GeminiOAuth] inferGoogleOneTier - < %d bytes (15GB), returning UNKNOWN", StorageTierFree)
|
||
return GeminiTierGoogleOneUnknown
|
||
}
|
||
|
||
// 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 (Google One personal account)")
|
||
|
||
// 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()
|
||
|
||
storageInfo, err := driveClient.GetStorageQuota(ctx, accessToken, proxyURL)
|
||
if err != nil {
|
||
// 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 GeminiTierGoogleOneUnknown, nil, err
|
||
}
|
||
// Other errors
|
||
log.Printf("[GeminiOAuth] Failed to fetch Drive storage: %v", err)
|
||
return GeminiTierGoogleOneUnknown, nil, err
|
||
}
|
||
|
||
log.Printf("[GeminiOAuth] Drive API response - Limit: %d bytes (%.2f TB), Usage: %d bytes (%.2f GB)",
|
||
storageInfo.Limit, float64(storageInfo.Limit)/float64(TB),
|
||
storageInfo.Usage, float64(storageInfo.Usage)/float64(GB))
|
||
|
||
tierID := inferGoogleOneTier(storageInfo.Limit)
|
||
log.Printf("[GeminiOAuth] Inferred tier from storage: %s", tierID)
|
||
|
||
return tierID, storageInfo, nil
|
||
}
|
||
|
||
// RefreshAccountGoogleOneTier 刷新单个账号的 Google One Tier
|
||
func (s *GeminiOAuthService) RefreshAccountGoogleOneTier(
|
||
ctx context.Context,
|
||
account *Account,
|
||
) (tierID string, extra map[string]any, credentials map[string]any, err error) {
|
||
if account == nil {
|
||
return "", nil, nil, fmt.Errorf("account is nil")
|
||
}
|
||
|
||
// 验证账号类型
|
||
oauthType, ok := account.Credentials["oauth_type"].(string)
|
||
if !ok || oauthType != "google_one" {
|
||
return "", nil, nil, fmt.Errorf("not a google_one OAuth account")
|
||
}
|
||
|
||
// 获取 access_token
|
||
accessToken, ok := account.Credentials["access_token"].(string)
|
||
if !ok || accessToken == "" {
|
||
return "", nil, nil, fmt.Errorf("missing access_token")
|
||
}
|
||
|
||
// 获取 proxy URL
|
||
var proxyURL string
|
||
if account.ProxyID != nil && account.Proxy != nil {
|
||
proxyURL = account.Proxy.URL()
|
||
}
|
||
|
||
// 调用 Drive API
|
||
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, accessToken, proxyURL)
|
||
if err != nil {
|
||
return "", nil, nil, err
|
||
}
|
||
|
||
// 构建 extra 数据（保留原有 extra 字段）
|
||
extra = make(map[string]any)
|
||
for k, v := range account.Extra {
|
||
extra[k] = v
|
||
}
|
||
if storageInfo != nil {
|
||
extra["drive_storage_limit"] = storageInfo.Limit
|
||
extra["drive_storage_usage"] = storageInfo.Usage
|
||
extra["drive_tier_updated_at"] = time.Now().Format(time.RFC3339)
|
||
}
|
||
|
||
// 构建 credentials 数据
|
||
credentials = make(map[string]any)
|
||
for k, v := range account.Credentials {
|
||
credentials[k] = v
|
||
}
|
||
credentials["tier_id"] = tierID
|
||
|
||
return tierID, extra, credentials, nil
|
||
}
|
||
|
||
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")
|
||
}
|
||
|
||
proxyURL := session.ProxyURL
|
||
if input.ProxyID != nil {
|
||
proxy, err := s.proxyRepo.GetByID(ctx, *input.ProxyID)
|
||
if err == nil && proxy != nil {
|
||
proxyURL = proxy.URL()
|
||
}
|
||
}
|
||
log.Printf("[GeminiOAuth] ProxyURL: %s", proxyURL)
|
||
|
||
redirectURI := session.RedirectURI
|
||
|
||
// Resolve oauth_type early (defaults to code_assist for backward compatibility).
|
||
oauthType := session.OAuthType
|
||
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" {
|
||
effectiveCfg, err := geminicli.EffectiveOAuthConfig(geminicli.OAuthConfig{
|
||
ClientID: s.cfg.Gemini.OAuth.ClientID,
|
||
ClientSecret: s.cfg.Gemini.OAuth.ClientSecret,
|
||
Scopes: s.cfg.Gemini.OAuth.Scopes,
|
||
}, "ai_studio")
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
isBuiltinClient := effectiveCfg.ClientID == geminicli.GeminiCLIOAuthClientID &&
|
||
effectiveCfg.ClientSecret == geminicli.GeminiCLIOAuthClientSecret
|
||
if isBuiltinClient {
|
||
return nil, fmt.Errorf("AI Studio OAuth requires a custom OAuth Client. Please use an AI Studio API Key account, or configure GEMINI_OAUTH_CLIENT_ID / GEMINI_OAUTH_CLIENT_SECRET and re-authorize")
|
||
}
|
||
}
|
||
|
||
// code_assist always uses the built-in client and its fixed redirect URI.
|
||
if oauthType == "code_assist" {
|
||
redirectURI = geminicli.GeminiCLIRedirectURI
|
||
}
|
||
|
||
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)
|
||
|
||
// 计算过期时间：减去 5 分钟安全时间窗口（考虑网络延迟和时钟偏差）
|
||
// 同时设置下界保护，防止 expires_in 过小导致过去时间（引发刷新风暴）
|
||
const safetyWindow = 300 // 5 minutes
|
||
const minTTL = 30 // minimum 30 seconds
|
||
expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - safetyWindow
|
||
minExpiresAt := time.Now().Unix() + minTTL
|
||
if expiresAt < minExpiresAt {
|
||
expiresAt = minExpiresAt
|
||
}
|
||
|
||
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")
|
||
}
|
||
// Prefer auto-detected tier; fall back to user-selected tier.
|
||
tierID = canonicalGeminiTierIDForOAuthType(oauthType, tierID)
|
||
if tierID == "" {
|
||
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")
|
||
|
||
// Google One accounts use cloudaicompanion API, which requires a project_id.
|
||
// For personal accounts, Google auto-assigns a project_id via the LoadCodeAssist API.
|
||
if projectID == "" {
|
||
log.Printf("[GeminiOAuth] No project_id provided, attempting to fetch from LoadCodeAssist API...")
|
||
var err error
|
||
projectID, _, err = s.fetchProjectID(ctx, tokenResp.AccessToken, proxyURL)
|
||
if err != nil {
|
||
log.Printf("[GeminiOAuth] ERROR: Failed to fetch project_id: %v", err)
|
||
return nil, fmt.Errorf("google One accounts require a project_id, failed to auto-detect: %w", err)
|
||
}
|
||
log.Printf("[GeminiOAuth] Successfully fetched project_id: %s", projectID)
|
||
}
|
||
|
||
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
|
||
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)
|
||
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))
|
||
}
|
||
}
|
||
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 {
|
||
tokenInfo := &GeminiTokenInfo{
|
||
AccessToken: tokenResp.AccessToken,
|
||
RefreshToken: tokenResp.RefreshToken,
|
||
TokenType: tokenResp.TokenType,
|
||
ExpiresIn: tokenResp.ExpiresIn,
|
||
ExpiresAt: expiresAt,
|
||
Scope: tokenResp.Scope,
|
||
ProjectID: projectID,
|
||
TierID: tierID,
|
||
OAuthType: oauthType,
|
||
Extra: map[string]any{
|
||
"drive_storage_limit": storageInfo.Limit,
|
||
"drive_storage_usage": storageInfo.Usage,
|
||
"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)
|
||
}
|
||
|
||
log.Printf("[GeminiOAuth] ========== Account Type Detection END ==========")
|
||
|
||
result := &GeminiTokenInfo{
|
||
AccessToken: tokenResp.AccessToken,
|
||
RefreshToken: tokenResp.RefreshToken,
|
||
TokenType: tokenResp.TokenType,
|
||
ExpiresIn: tokenResp.ExpiresIn,
|
||
ExpiresAt: expiresAt,
|
||
Scope: tokenResp.Scope,
|
||
ProjectID: projectID,
|
||
TierID: tierID,
|
||
OAuthType: oauthType,
|
||
}
|
||
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
|
||
}
|
||
|
||
func (s *GeminiOAuthService) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*GeminiTokenInfo, error) {
|
||
var lastErr error
|
||
|
||
for attempt := 0; attempt <= 3; attempt++ {
|
||
if attempt > 0 {
|
||
backoff := time.Duration(1<<uint(attempt-1)) * time.Second
|
||
if backoff > 30*time.Second {
|
||
backoff = 30 * time.Second
|
||
}
|
||
time.Sleep(backoff)
|
||
}
|
||
|
||
tokenResp, err := s.oauthClient.RefreshToken(ctx, oauthType, refreshToken, proxyURL)
|
||
if err == nil {
|
||
// 计算过期时间：减去 5 分钟安全时间窗口（考虑网络延迟和时钟偏差）
|
||
// 同时设置下界保护，防止 expires_in 过小导致过去时间（引发刷新风暴）
|
||
const safetyWindow = 300 // 5 minutes
|
||
const minTTL = 30 // minimum 30 seconds
|
||
expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - safetyWindow
|
||
minExpiresAt := time.Now().Unix() + minTTL
|
||
if expiresAt < minExpiresAt {
|
||
expiresAt = minExpiresAt
|
||
}
|
||
return &GeminiTokenInfo{
|
||
AccessToken: tokenResp.AccessToken,
|
||
RefreshToken: tokenResp.RefreshToken,
|
||
TokenType: tokenResp.TokenType,
|
||
ExpiresIn: tokenResp.ExpiresIn,
|
||
ExpiresAt: expiresAt,
|
||
Scope: tokenResp.Scope,
|
||
}, nil
|
||
}
|
||
|
||
if isNonRetryableGeminiOAuthError(err) {
|
||
return nil, err
|
||
}
|
||
lastErr = err
|
||
}
|
||
|
||
return nil, fmt.Errorf("token refresh failed after retries: %w", lastErr)
|
||
}
|
||
|
||
func isNonRetryableGeminiOAuthError(err error) bool {
|
||
msg := err.Error()
|
||
nonRetryable := []string{
|
||
"invalid_grant",
|
||
"invalid_client",
|
||
"unauthorized_client",
|
||
"access_denied",
|
||
}
|
||
for _, needle := range nonRetryable {
|
||
if strings.Contains(msg, needle) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*GeminiTokenInfo, error) {
|
||
if account.Platform != PlatformGemini || account.Type != AccountTypeOAuth {
|
||
return nil, fmt.Errorf("account is not a Gemini OAuth account")
|
||
}
|
||
|
||
refreshToken := account.GetCredential("refresh_token")
|
||
if strings.TrimSpace(refreshToken) == "" {
|
||
return nil, fmt.Errorf("no refresh token available")
|
||
}
|
||
|
||
// Preserve oauth_type from the account (defaults to code_assist for backward compatibility).
|
||
oauthType := strings.TrimSpace(account.GetCredential("oauth_type"))
|
||
if oauthType == "" {
|
||
oauthType = "code_assist"
|
||
}
|
||
|
||
var proxyURL string
|
||
if account.ProxyID != nil {
|
||
proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID)
|
||
if err == nil && proxy != nil {
|
||
proxyURL = proxy.URL()
|
||
}
|
||
}
|
||
|
||
tokenInfo, err := s.RefreshToken(ctx, oauthType, refreshToken, proxyURL)
|
||
// Backward compatibility:
|
||
// Older versions could refresh Code Assist tokens using a user-provided OAuth client when configured.
|
||
// If the refresh token was originally issued to that custom client, forcing the built-in client will
|
||
// fail with "unauthorized_client". In that case, retry with the custom client (ai_studio path) when available.
|
||
if err != nil && oauthType == "code_assist" && strings.Contains(err.Error(), "unauthorized_client") && s.GetOAuthConfig().AIStudioOAuthEnabled {
|
||
if alt, altErr := s.RefreshToken(ctx, "ai_studio", refreshToken, proxyURL); altErr == nil {
|
||
tokenInfo = alt
|
||
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") {
|
||
return nil, fmt.Errorf("%w (OAuth client mismatch: the refresh_token is bound to the OAuth client used during authorization; please re-authorize this account or restore the original GEMINI_OAUTH_CLIENT_ID/SECRET)", err)
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
tokenInfo.OAuthType = oauthType
|
||
|
||
// Preserve account's project_id when present.
|
||
existingProjectID := strings.TrimSpace(account.GetCredential("project_id"))
|
||
if existingProjectID != "" {
|
||
tokenInfo.ProjectID = existingProjectID
|
||
}
|
||
|
||
// 尝试从账号凭证获取 tierID（向后兼容）
|
||
existingTierID := strings.TrimSpace(account.GetCredential("tier_id"))
|
||
|
||
// For Code Assist, project_id is required. Auto-detect if missing.
|
||
// For AI Studio OAuth, project_id is optional and should not block refresh.
|
||
switch oauthType {
|
||
case "code_assist":
|
||
// 先设置默认值或保留旧值，确保 tier_id 始终有值
|
||
if existingTierID != "" {
|
||
tokenInfo.TierID = canonicalGeminiTierIDForOAuthType(oauthType, existingTierID)
|
||
}
|
||
if tokenInfo.TierID == "" {
|
||
tokenInfo.TierID = GeminiTierGCPStandard
|
||
}
|
||
|
||
// 尝试自动探测 project_id 和 tier_id
|
||
needDetect := strings.TrimSpace(tokenInfo.ProjectID) == "" || tokenInfo.TierID == ""
|
||
if needDetect {
|
||
projectID, tierID, err := s.fetchProjectID(ctx, tokenInfo.AccessToken, proxyURL)
|
||
if err != nil {
|
||
fmt.Printf("[GeminiOAuth] Warning: failed to auto-detect project/tier: %v\n", err)
|
||
} else {
|
||
if strings.TrimSpace(tokenInfo.ProjectID) == "" && projectID != "" {
|
||
tokenInfo.ProjectID = projectID
|
||
}
|
||
if tierID != "" {
|
||
if canonical := canonicalGeminiTierIDForOAuthType(oauthType, tierID); canonical != "" {
|
||
tokenInfo.TierID = canonical
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if strings.TrimSpace(tokenInfo.ProjectID) == "" {
|
||
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 {
|
||
if updatedAtStr, ok := account.Extra["drive_tier_updated_at"].(string); ok {
|
||
if updatedAt, err := time.Parse(time.RFC3339, updatedAtStr); err == nil {
|
||
if time.Since(updatedAt) <= 24*time.Hour {
|
||
needsRefresh = false
|
||
// Use cached tier
|
||
tokenInfo.TierID = canonicalExistingTier
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if tokenInfo.TierID == "" {
|
||
tokenInfo.TierID = canonicalExistingTier
|
||
}
|
||
|
||
if needsRefresh {
|
||
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, tokenInfo.AccessToken, proxyURL)
|
||
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 {
|
||
tokenInfo.TierID = GeminiTierGoogleOneFree
|
||
}
|
||
}
|
||
}
|
||
|
||
return tokenInfo, nil
|
||
}
|
||
|
||
func (s *GeminiOAuthService) BuildAccountCredentials(tokenInfo *GeminiTokenInfo) map[string]any {
|
||
creds := map[string]any{
|
||
"access_token": tokenInfo.AccessToken,
|
||
"expires_at": strconv.FormatInt(tokenInfo.ExpiresAt, 10),
|
||
}
|
||
if tokenInfo.RefreshToken != "" {
|
||
creds["refresh_token"] = tokenInfo.RefreshToken
|
||
}
|
||
if tokenInfo.TokenType != "" {
|
||
creds["token_type"] = tokenInfo.TokenType
|
||
}
|
||
if tokenInfo.Scope != "" {
|
||
creds["scope"] = tokenInfo.Scope
|
||
}
|
||
if tokenInfo.ProjectID != "" {
|
||
creds["project_id"] = tokenInfo.ProjectID
|
||
}
|
||
if tokenInfo.TierID != "" {
|
||
// 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)
|
||
}
|
||
if tokenInfo.OAuthType != "" {
|
||
creds["oauth_type"] = tokenInfo.OAuthType
|
||
}
|
||
// Store extra metadata (Drive info) if present
|
||
if len(tokenInfo.Extra) > 0 {
|
||
for k, v := range tokenInfo.Extra {
|
||
creds[k] = v
|
||
}
|
||
}
|
||
return creds
|
||
}
|
||
|
||
func (s *GeminiOAuthService) Stop() {
|
||
s.sessionStore.Stop()
|
||
}
|
||
|
||
func (s *GeminiOAuthService) fetchProjectID(ctx context.Context, accessToken, proxyURL string) (string, string, error) {
|
||
if s.codeAssist == nil {
|
||
return "", "", errors.New("code assist client not configured")
|
||
}
|
||
|
||
loadResp, loadErr := s.codeAssist.LoadCodeAssist(ctx, accessToken, proxyURL, nil)
|
||
|
||
// Extract tierID from response (works whether CloudAICompanionProject is set or not)
|
||
tierID := "LEGACY"
|
||
if loadResp != nil {
|
||
// 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
|
||
if loadErr == nil && loadResp != nil && strings.TrimSpace(loadResp.CloudAICompanionProject) != "" {
|
||
return strings.TrimSpace(loadResp.CloudAICompanionProject), tierID, nil
|
||
}
|
||
|
||
// 关键逻辑：对齐 Gemini CLI 对“已注册用户”的处理方式。
|
||
// 当 LoadCodeAssist 返回了 currentTier / paidTier（表示账号已注册）但没有返回 cloudaicompanionProject 时：
|
||
// - 不要再调用 onboardUser（通常不会再分配 project_id，且可能触发 INVALID_ARGUMENT）
|
||
// - 先尝试从 Cloud Resource Manager 获取可用项目；仍失败则提示用户手动填写 project_id
|
||
if loadResp != nil {
|
||
registeredTierID := strings.TrimSpace(loadResp.GetTier())
|
||
if registeredTierID != "" {
|
||
// 已注册但未返回 cloudaicompanionProject，这在 Google One 用户中较常见：需要用户自行提供 project_id。
|
||
log.Printf("[GeminiOAuth] User has tier (%s) but no cloudaicompanionProject, trying Cloud Resource Manager...", registeredTierID)
|
||
|
||
// Try to get project from Cloud Resource Manager
|
||
fallback, fbErr := fetchProjectIDFromResourceManager(ctx, accessToken, proxyURL)
|
||
if fbErr == nil && strings.TrimSpace(fallback) != "" {
|
||
log.Printf("[GeminiOAuth] Found project from Cloud Resource Manager: %s", fallback)
|
||
return strings.TrimSpace(fallback), tierID, nil
|
||
}
|
||
|
||
// No project found - user must provide project_id manually
|
||
log.Printf("[GeminiOAuth] No project found from Cloud Resource Manager, user must provide project_id manually")
|
||
return "", tierID, fmt.Errorf("user is registered (tier: %s) but no project_id available. Please provide Project ID manually in the authorization form, or create a project at https://console.cloud.google.com", registeredTierID)
|
||
}
|
||
}
|
||
|
||
// 未检测到 currentTier/paidTier，视为新用户，继续调用 onboardUser
|
||
log.Printf("[GeminiOAuth] No currentTier/paidTier found, proceeding with onboardUser (tierID: %s)", tierID)
|
||
|
||
req := &geminicli.OnboardUserRequest{
|
||
TierID: tierID,
|
||
Metadata: geminicli.LoadCodeAssistMetadata{
|
||
IDEType: "ANTIGRAVITY",
|
||
Platform: "PLATFORM_UNSPECIFIED",
|
||
PluginType: "GEMINI",
|
||
},
|
||
}
|
||
|
||
maxAttempts := 5
|
||
for attempt := 1; attempt <= maxAttempts; attempt++ {
|
||
resp, err := s.codeAssist.OnboardUser(ctx, accessToken, proxyURL, req)
|
||
if err != nil {
|
||
// If Code Assist onboarding fails (e.g. INVALID_ARGUMENT), fallback to Cloud Resource Manager projects.
|
||
fallback, fbErr := fetchProjectIDFromResourceManager(ctx, accessToken, proxyURL)
|
||
if fbErr == nil && strings.TrimSpace(fallback) != "" {
|
||
return strings.TrimSpace(fallback), tierID, nil
|
||
}
|
||
return "", tierID, err
|
||
}
|
||
if resp.Done {
|
||
if resp.Response != nil && resp.Response.CloudAICompanionProject != nil {
|
||
switch v := resp.Response.CloudAICompanionProject.(type) {
|
||
case string:
|
||
return strings.TrimSpace(v), tierID, nil
|
||
case map[string]any:
|
||
if id, ok := v["id"].(string); ok {
|
||
return strings.TrimSpace(id), tierID, nil
|
||
}
|
||
}
|
||
}
|
||
|
||
fallback, fbErr := fetchProjectIDFromResourceManager(ctx, accessToken, proxyURL)
|
||
if fbErr == nil && strings.TrimSpace(fallback) != "" {
|
||
return strings.TrimSpace(fallback), tierID, nil
|
||
}
|
||
return "", tierID, errors.New("onboardUser completed but no project_id returned")
|
||
}
|
||
time.Sleep(2 * time.Second)
|
||
}
|
||
|
||
fallback, fbErr := fetchProjectIDFromResourceManager(ctx, accessToken, proxyURL)
|
||
if fbErr == nil && strings.TrimSpace(fallback) != "" {
|
||
return strings.TrimSpace(fallback), tierID, nil
|
||
}
|
||
if loadErr != nil {
|
||
return "", tierID, fmt.Errorf("loadCodeAssist failed (%v) and onboardUser timeout after %d attempts", loadErr, maxAttempts)
|
||
}
|
||
return "", tierID, fmt.Errorf("onboardUser timeout after %d attempts", maxAttempts)
|
||
}
|
||
|
||
type googleCloudProject struct {
|
||
ProjectID string `json:"projectId"`
|
||
DisplayName string `json:"name"`
|
||
LifecycleState string `json:"lifecycleState"`
|
||
}
|
||
|
||
type googleCloudProjectsResponse struct {
|
||
Projects []googleCloudProject `json:"projects"`
|
||
}
|
||
|
||
func fetchProjectIDFromResourceManager(ctx context.Context, accessToken, proxyURL string) (string, error) {
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://cloudresourcemanager.googleapis.com/v1/projects", nil)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to create resource manager request: %w", err)
|
||
}
|
||
|
||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||
req.Header.Set("User-Agent", geminicli.GeminiCLIUserAgent)
|
||
|
||
client, err := httpclient.GetClient(httpclient.Options{
|
||
ProxyURL: strings.TrimSpace(proxyURL),
|
||
Timeout: 30 * time.Second,
|
||
ValidateResolvedIP: true,
|
||
})
|
||
if err != nil {
|
||
client = &http.Client{Timeout: 30 * time.Second}
|
||
}
|
||
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("resource manager request failed: %w", err)
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
|
||
bodyBytes, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to read resource manager response: %w", err)
|
||
}
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
return "", fmt.Errorf("resource manager HTTP %d: %s", resp.StatusCode, string(bodyBytes))
|
||
}
|
||
|
||
var projectsResp googleCloudProjectsResponse
|
||
if err := json.Unmarshal(bodyBytes, &projectsResp); err != nil {
|
||
return "", fmt.Errorf("failed to parse resource manager response: %w", err)
|
||
}
|
||
|
||
active := make([]googleCloudProject, 0, len(projectsResp.Projects))
|
||
for _, p := range projectsResp.Projects {
|
||
if p.LifecycleState == "ACTIVE" && strings.TrimSpace(p.ProjectID) != "" {
|
||
active = append(active, p)
|
||
}
|
||
}
|
||
if len(active) == 0 {
|
||
return "", errors.New("no ACTIVE projects found from resource manager")
|
||
}
|
||
|
||
// Prefer likely companion projects first.
|
||
for _, p := range active {
|
||
id := strings.ToLower(strings.TrimSpace(p.ProjectID))
|
||
name := strings.ToLower(strings.TrimSpace(p.DisplayName))
|
||
if strings.Contains(id, "cloud-ai-companion") || strings.Contains(name, "cloud ai companion") || strings.Contains(name, "code assist") {
|
||
return strings.TrimSpace(p.ProjectID), nil
|
||
}
|
||
}
|
||
// Then prefer "default".
|
||
for _, p := range active {
|
||
id := strings.ToLower(strings.TrimSpace(p.ProjectID))
|
||
name := strings.ToLower(strings.TrimSpace(p.DisplayName))
|
||
if strings.Contains(id, "default") || strings.Contains(name, "default") {
|
||
return strings.TrimSpace(p.ProjectID), nil
|
||
}
|
||
}
|
||
|
||
return strings.TrimSpace(active[0].ProjectID), nil
|
||
}
|