diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index e2ff23f1..0f855a03 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -1188,6 +1188,30 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) { } ts.ProjectID = strings.Join(projects, ",") ts.Checked = true + } else if strings.EqualFold(requestedProjectID, "GOOGLE_ONE") { + ts.Auto = false + if errSetup := performGeminiCLISetup(ctx, gemClient, &ts, ""); errSetup != nil { + log.Errorf("Google One auto-discovery failed: %v", errSetup) + SetOAuthSessionError(state, "Google One auto-discovery failed") + return + } + if strings.TrimSpace(ts.ProjectID) == "" { + log.Error("Google One auto-discovery returned empty project ID") + SetOAuthSessionError(state, "Google One auto-discovery returned empty project ID") + return + } + isChecked, errCheck := checkCloudAPIIsEnabled(ctx, gemClient, ts.ProjectID) + if errCheck != nil { + log.Errorf("Failed to verify Cloud AI API status: %v", errCheck) + SetOAuthSessionError(state, "Failed to verify Cloud AI API status") + return + } + ts.Checked = isChecked + if !isChecked { + log.Error("Cloud AI API is not enabled for the auto-discovered project") + SetOAuthSessionError(state, "Cloud AI API not enabled") + return + } } else { if errEnsure := ensureGeminiProjectAndOnboard(ctx, gemClient, &ts, requestedProjectID); errEnsure != nil { log.Errorf("Failed to complete Gemini CLI onboarding: %v", errEnsure) @@ -2036,7 +2060,48 @@ func performGeminiCLISetup(ctx context.Context, httpClient *http.Client, storage } } if projectID == "" { - return &projectSelectionRequiredError{} + // Auto-discovery: try onboardUser without specifying a project + // to let Google auto-provision one (matches Gemini CLI headless behavior + // and Antigravity's FetchProjectID pattern). + autoOnboardReq := map[string]any{ + "tierId": tierID, + "metadata": metadata, + } + + autoCtx, autoCancel := context.WithTimeout(ctx, 30*time.Second) + defer autoCancel() + for attempt := 1; ; attempt++ { + var onboardResp map[string]any + if errOnboard := callGeminiCLI(autoCtx, httpClient, "onboardUser", autoOnboardReq, &onboardResp); errOnboard != nil { + return fmt.Errorf("auto-discovery onboardUser: %w", errOnboard) + } + + if done, okDone := onboardResp["done"].(bool); okDone && done { + if resp, okResp := onboardResp["response"].(map[string]any); okResp { + switch v := resp["cloudaicompanionProject"].(type) { + case string: + projectID = strings.TrimSpace(v) + case map[string]any: + if id, okID := v["id"].(string); okID { + projectID = strings.TrimSpace(id) + } + } + } + break + } + + log.Debugf("Auto-discovery: onboarding in progress, attempt %d...", attempt) + select { + case <-autoCtx.Done(): + return &projectSelectionRequiredError{} + case <-time.After(2 * time.Second): + } + } + + if projectID == "" { + return &projectSelectionRequiredError{} + } + log.Infof("Auto-discovered project ID via onboarding: %s", projectID) } onboardReqBody := map[string]any{ diff --git a/internal/auth/gemini/gemini_token.go b/internal/auth/gemini/gemini_token.go index 0ec7da17..f7fca810 100644 --- a/internal/auth/gemini/gemini_token.go +++ b/internal/auth/gemini/gemini_token.go @@ -71,17 +71,17 @@ func (ts *GeminiTokenStorage) SaveTokenToFile(authFilePath string) error { // CredentialFileName returns the filename used to persist Gemini CLI credentials. // When projectID represents multiple projects (comma-separated or literal ALL), -// the suffix is normalized to "all" and a "gemini-" prefix is enforced to keep +// the suffix is normalized to "all" and a "geminicli-" prefix is enforced to keep // web and CLI generated files consistent. func CredentialFileName(email, projectID string, includeProviderPrefix bool) string { email = strings.TrimSpace(email) project := strings.TrimSpace(projectID) if strings.EqualFold(project, "all") || strings.Contains(project, ",") { - return fmt.Sprintf("gemini-%s-all.json", email) + return fmt.Sprintf("geminicli-%s-all.json", email) } prefix := "" if includeProviderPrefix { - prefix = "gemini-" + prefix = "geminicli-" } return fmt.Sprintf("%s%s-%s.json", prefix, email, project) } diff --git a/internal/cmd/login.go b/internal/cmd/login.go index b5129cfd..3286e7a7 100644 --- a/internal/cmd/login.go +++ b/internal/cmd/login.go @@ -100,49 +100,75 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) { log.Info("Authentication successful.") - projects, errProjects := fetchGCPProjects(ctx, httpClient) - if errProjects != nil { - log.Errorf("Failed to get project list: %v", errProjects) - return + var activatedProjects []string + + useGoogleOne := false + if trimmedProjectID == "" && promptFn != nil { + fmt.Println("\nSelect login mode:") + fmt.Println(" 1. Code Assist (GCP project, manual selection)") + fmt.Println(" 2. Google One (personal account, auto-discover project)") + choice, errPrompt := promptFn("Enter choice [1/2] (default: 1): ") + if errPrompt == nil && strings.TrimSpace(choice) == "2" { + useGoogleOne = true + } } - selectedProjectID := promptForProjectSelection(projects, trimmedProjectID, promptFn) - projectSelections, errSelection := resolveProjectSelections(selectedProjectID, projects) - if errSelection != nil { - log.Errorf("Invalid project selection: %v", errSelection) - return - } - if len(projectSelections) == 0 { - log.Error("No project selected; aborting login.") - return - } - - activatedProjects := make([]string, 0, len(projectSelections)) - seenProjects := make(map[string]bool) - for _, candidateID := range projectSelections { - log.Infof("Activating project %s", candidateID) - if errSetup := performGeminiCLISetup(ctx, httpClient, storage, candidateID); errSetup != nil { - var projectErr *projectSelectionRequiredError - if errors.As(errSetup, &projectErr) { - log.Error("Failed to start user onboarding: A project ID is required.") - showProjectSelectionHelp(storage.Email, projects) - return - } - log.Errorf("Failed to complete user setup: %v", errSetup) + if useGoogleOne { + log.Info("Google One mode: auto-discovering project...") + if errSetup := performGeminiCLISetup(ctx, httpClient, storage, ""); errSetup != nil { + log.Errorf("Google One auto-discovery failed: %v", errSetup) return } - finalID := strings.TrimSpace(storage.ProjectID) - if finalID == "" { - finalID = candidateID + autoProject := strings.TrimSpace(storage.ProjectID) + if autoProject == "" { + log.Error("Google One auto-discovery returned empty project ID") + return + } + log.Infof("Auto-discovered project: %s", autoProject) + activatedProjects = []string{autoProject} + } else { + projects, errProjects := fetchGCPProjects(ctx, httpClient) + if errProjects != nil { + log.Errorf("Failed to get project list: %v", errProjects) + return } - // Skip duplicates - if seenProjects[finalID] { - log.Infof("Project %s already activated, skipping", finalID) - continue + selectedProjectID := promptForProjectSelection(projects, trimmedProjectID, promptFn) + projectSelections, errSelection := resolveProjectSelections(selectedProjectID, projects) + if errSelection != nil { + log.Errorf("Invalid project selection: %v", errSelection) + return + } + if len(projectSelections) == 0 { + log.Error("No project selected; aborting login.") + return + } + + seenProjects := make(map[string]bool) + for _, candidateID := range projectSelections { + log.Infof("Activating project %s", candidateID) + if errSetup := performGeminiCLISetup(ctx, httpClient, storage, candidateID); errSetup != nil { + var projectErr *projectSelectionRequiredError + if errors.As(errSetup, &projectErr) { + log.Error("Failed to start user onboarding: A project ID is required.") + showProjectSelectionHelp(storage.Email, projects) + return + } + log.Errorf("Failed to complete user setup: %v", errSetup) + return + } + finalID := strings.TrimSpace(storage.ProjectID) + if finalID == "" { + finalID = candidateID + } + + if seenProjects[finalID] { + log.Infof("Project %s already activated, skipping", finalID) + continue + } + seenProjects[finalID] = true + activatedProjects = append(activatedProjects, finalID) } - seenProjects[finalID] = true - activatedProjects = append(activatedProjects, finalID) } storage.Auto = false @@ -235,7 +261,48 @@ func performGeminiCLISetup(ctx context.Context, httpClient *http.Client, storage } } if projectID == "" { - return &projectSelectionRequiredError{} + // Auto-discovery: try onboardUser without specifying a project + // to let Google auto-provision one (matches Gemini CLI headless behavior + // and Antigravity's FetchProjectID pattern). + autoOnboardReq := map[string]any{ + "tierId": tierID, + "metadata": metadata, + } + + autoCtx, autoCancel := context.WithTimeout(ctx, 30*time.Second) + defer autoCancel() + for attempt := 1; ; attempt++ { + var onboardResp map[string]any + if errOnboard := callGeminiCLI(autoCtx, httpClient, "onboardUser", autoOnboardReq, &onboardResp); errOnboard != nil { + return fmt.Errorf("auto-discovery onboardUser: %w", errOnboard) + } + + if done, okDone := onboardResp["done"].(bool); okDone && done { + if resp, okResp := onboardResp["response"].(map[string]any); okResp { + switch v := resp["cloudaicompanionProject"].(type) { + case string: + projectID = strings.TrimSpace(v) + case map[string]any: + if id, okID := v["id"].(string); okID { + projectID = strings.TrimSpace(id) + } + } + } + break + } + + log.Debugf("Auto-discovery: onboarding in progress, attempt %d...", attempt) + select { + case <-autoCtx.Done(): + return &projectSelectionRequiredError{} + case <-time.After(2 * time.Second): + } + } + + if projectID == "" { + return &projectSelectionRequiredError{} + } + log.Infof("Auto-discovered project ID via onboarding: %s", projectID) } onboardReqBody := map[string]any{ @@ -617,7 +684,7 @@ func updateAuthRecord(record *cliproxyauth.Auth, storage *gemini.GeminiTokenStor return } - finalName := gemini.CredentialFileName(storage.Email, storage.ProjectID, false) + finalName := gemini.CredentialFileName(storage.Email, storage.ProjectID, true) if record.Metadata == nil { record.Metadata = make(map[string]any) diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go index 0bb7ff7d..795bba0d 100644 --- a/sdk/auth/filestore.go +++ b/sdk/auth/filestore.go @@ -4,8 +4,10 @@ import ( "context" "encoding/json" "fmt" + "io" "io/fs" "net/http" + "net/url" "os" "path/filepath" "strings" @@ -186,15 +188,21 @@ func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, if provider == "" { provider = "unknown" } - if provider == "antigravity" { + if provider == "antigravity" || provider == "gemini" { projectID := "" if pid, ok := metadata["project_id"].(string); ok { projectID = strings.TrimSpace(pid) } if projectID == "" { - accessToken := "" - if token, ok := metadata["access_token"].(string); ok { - accessToken = strings.TrimSpace(token) + accessToken := extractAccessToken(metadata) + // For gemini type, the stored access_token is likely expired (~1h lifetime). + // Refresh it using the long-lived refresh_token before querying. + if provider == "gemini" { + if tokenMap, ok := metadata["token"].(map[string]any); ok { + if refreshed, errRefresh := refreshGeminiAccessToken(tokenMap, http.DefaultClient); errRefresh == nil { + accessToken = refreshed + } + } } if accessToken != "" { fetchedProjectID, errFetch := FetchAntigravityProjectID(context.Background(), accessToken, http.DefaultClient) @@ -304,6 +312,67 @@ func (s *FileTokenStore) baseDirSnapshot() string { return s.baseDir } +func extractAccessToken(metadata map[string]any) string { + if at, ok := metadata["access_token"].(string); ok { + if v := strings.TrimSpace(at); v != "" { + return v + } + } + if tokenMap, ok := metadata["token"].(map[string]any); ok { + if at, ok := tokenMap["access_token"].(string); ok { + if v := strings.TrimSpace(at); v != "" { + return v + } + } + } + return "" +} + +func refreshGeminiAccessToken(tokenMap map[string]any, httpClient *http.Client) (string, error) { + refreshToken, _ := tokenMap["refresh_token"].(string) + clientID, _ := tokenMap["client_id"].(string) + clientSecret, _ := tokenMap["client_secret"].(string) + tokenURI, _ := tokenMap["token_uri"].(string) + + if refreshToken == "" || clientID == "" || clientSecret == "" { + return "", fmt.Errorf("missing refresh credentials") + } + if tokenURI == "" { + tokenURI = "https://oauth2.googleapis.com/token" + } + + data := url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {refreshToken}, + "client_id": {clientID}, + "client_secret": {clientSecret}, + } + + resp, err := httpClient.PostForm(tokenURI, data) + if err != nil { + return "", fmt.Errorf("refresh request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("refresh failed: status %d", resp.StatusCode) + } + + var result map[string]any + if errUnmarshal := json.Unmarshal(body, &result); errUnmarshal != nil { + return "", fmt.Errorf("decode refresh response: %w", errUnmarshal) + } + + newAccessToken, _ := result["access_token"].(string) + if newAccessToken == "" { + return "", fmt.Errorf("no access_token in refresh response") + } + + tokenMap["access_token"] = newAccessToken + return newAccessToken, nil +} + // jsonEqual compares two JSON blobs by parsing them into Go objects and deep comparing. func jsonEqual(a, b []byte) bool { var objA any diff --git a/sdk/auth/filestore_test.go b/sdk/auth/filestore_test.go new file mode 100644 index 00000000..9e135ad4 --- /dev/null +++ b/sdk/auth/filestore_test.go @@ -0,0 +1,80 @@ +package auth + +import "testing" + +func TestExtractAccessToken(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + metadata map[string]any + expected string + }{ + { + "antigravity top-level access_token", + map[string]any{"access_token": "tok-abc"}, + "tok-abc", + }, + { + "gemini nested token.access_token", + map[string]any{ + "token": map[string]any{"access_token": "tok-nested"}, + }, + "tok-nested", + }, + { + "top-level takes precedence over nested", + map[string]any{ + "access_token": "tok-top", + "token": map[string]any{"access_token": "tok-nested"}, + }, + "tok-top", + }, + { + "empty metadata", + map[string]any{}, + "", + }, + { + "whitespace-only access_token", + map[string]any{"access_token": " "}, + "", + }, + { + "wrong type access_token", + map[string]any{"access_token": 12345}, + "", + }, + { + "token is not a map", + map[string]any{"token": "not-a-map"}, + "", + }, + { + "nested whitespace-only", + map[string]any{ + "token": map[string]any{"access_token": " "}, + }, + "", + }, + { + "fallback to nested when top-level empty", + map[string]any{ + "access_token": "", + "token": map[string]any{"access_token": "tok-fallback"}, + }, + "tok-fallback", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := extractAccessToken(tt.metadata) + if got != tt.expected { + t.Errorf("extractAccessToken() = %q, want %q", got, tt.expected) + } + }) + } +}