Merge pull request #1543 from sususu98/feat/gemini-cli-google-one

feat(gemini-cli): add Google One login and improve auto-discovery

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 == "" {
+		// 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{

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)
 }

internal/cmd/login.go

@@ -100,6 +100,33 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
 
 	log.Info("Authentication successful.")
 
+	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
+		}
+	}
+
+	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
+		}
+		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)
@@ -117,7 +144,6 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
 			return
 		}
 
-	activatedProjects := make([]string, 0, len(projectSelections))
 		seenProjects := make(map[string]bool)
 		for _, candidateID := range projectSelections {
 			log.Infof("Activating project %s", candidateID)
@@ -136,7 +162,6 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
 				finalID = candidateID
 			}
 
-		// Skip duplicates
 			if seenProjects[finalID] {
 				log.Infof("Project %s already activated, skipping", finalID)
 				continue
@@ -144,6 +169,7 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
 			seenProjects[finalID] = true
 			activatedProjects = append(activatedProjects, finalID)
 		}
+	}
 
 	storage.Auto = false
 	storage.ProjectID = strings.Join(activatedProjects, ",")
@@ -235,7 +261,48 @@ func performGeminiCLISetup(ctx context.Context, httpClient *http.Client, storage
 		}
 	}
 	if projectID == "" {
+		// 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)

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 := ""
+			accessToken := extractAccessToken(metadata)
-			if token, ok := metadata["access_token"].(string); ok {
+			// For gemini type, the stored access_token is likely expired (~1h lifetime).
-				accessToken = strings.TrimSpace(token)
+			// 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

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)
+			}
+		})
+	}
+}