feat(auth): enhance DoLogin to include Gemini CLI user onboarding flow

- Integrated Gemini CLI user setup into the `DoLogin` flow for streamlined authentication. - Added project selection handling, automatic project detection, and validation of Cloud AI API enablement. - Implemented new helper functions for Gemini CLI operations, project fetching, and onboarding logic. - Enhanced token storage and metadata updates for better user and project management.
2026-02-03 04:50:52 +08:00 · 2025-09-30 00:04:58 +08:00
parent 352a67857b
commit 1c91823308
1 changed files with 346 additions and 26 deletions
--- a/internal/cmd/login.go
+++ b/internal/cmd/login.go
@@ -4,18 +4,44 @@
 package cmd

 import (
+	"bufio"
+	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"time"

+	"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
 	sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
+	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
 	log "github.com/sirupsen/logrus"
+	"github.com/tidwall/gjson"
 )

+const (
+	geminiCLIEndpoint       = "https://cloudcode-pa.googleapis.com"
+	geminiCLIVersion        = "v1internal"
+	geminiCLIUserAgent      = "google-api-nodejs-client/9.15.1"
+	geminiCLIApiClient      = "gl-node/22.17.0"
+	geminiCLIClientMetadata = "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI"
+)
+
+type projectSelectionRequiredError struct{}
+
+func (e *projectSelectionRequiredError) Error() string {
+	return "gemini cli: project selection required"
+}
+
 // DoLogin handles Google Gemini authentication using the shared authentication manager.
-// It initiates the OAuth flow for Google Gemini services and saves the authentication
-// tokens to the configured auth directory.
+// It initiates the OAuth flow for Google Gemini services, performs the legacy CLI user setup,
+// and saves the authentication tokens to the configured auth directory.
 //
 // Parameters:
 //   - cfg: The application configuration
@@ -26,38 +52,78 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
 		options = &LoginOptions{}
 	}

-	manager := newAuthManager()
+	ctx := context.Background()

-	metadata := map[string]string{}
-	if projectID != "" {
-		metadata["project_id"] = projectID
-	}
-
-	authOpts := &sdkAuth.LoginOptions{
+	loginOpts := &sdkAuth.LoginOptions{
 		NoBrowser: options.NoBrowser,
-		ProjectID: projectID,
-		Metadata:  metadata,
+		ProjectID: strings.TrimSpace(projectID),
+		Metadata:  map[string]string{},
 		Prompt:    options.Prompt,
 	}

-	_, savedPath, err := manager.Login(context.Background(), "gemini", cfg, authOpts)
-	if err != nil {
-		var selectionErr *sdkAuth.ProjectSelectionError
-		if errors.As(err, &selectionErr) {
-			fmt.Println(selectionErr.Error())
-			projects := selectionErr.ProjectsDisplay()
-			if len(projects) > 0 {
-				fmt.Println("========================================================================")
-				for _, p := range projects {
-					fmt.Printf("Project ID: %s\n", p.ProjectID)
-					fmt.Printf("Project Name: %s\n", p.Name)
-					fmt.Println("------------------------------------------------------------------------")
-				}
-				fmt.Println("Please rerun the login command with --project_id <project_id>.")
+	authenticator := sdkAuth.NewGeminiAuthenticator()
+	record, errLogin := authenticator.Login(ctx, cfg, loginOpts)
+	if errLogin != nil {
+		log.Fatalf("Gemini authentication failed: %v", errLogin)
+		return
+	}
+
+	storage, okStorage := record.Storage.(*gemini.GeminiTokenStorage)
+	if !okStorage || storage == nil {
+		log.Fatal("Gemini authentication failed: unsupported token storage")
+		return
+	}
+
+	geminiAuth := gemini.NewGeminiAuth()
+	httpClient, errClient := geminiAuth.GetAuthenticatedClient(ctx, storage, cfg, options.NoBrowser)
+	if errClient != nil {
+		log.Fatalf("Gemini authentication failed: %v", errClient)
+		return
+	}
+
+	log.Info("Authentication successful.")
+
+	if errSetup := performGeminiCLISetup(ctx, httpClient, storage, strings.TrimSpace(projectID)); errSetup != nil {
+		var projectErr *projectSelectionRequiredError
+		if errors.As(errSetup, &projectErr) {
+			log.Error("Failed to start user onboarding: A project ID is required.")
+			projects, errProjects := fetchGCPProjects(ctx, httpClient)
+			if errProjects != nil {
+				log.Fatalf("Failed to get project list: %v", errProjects)
+				return
 			}
+			showProjectSelectionHelp(storage.Email, projects)
 			return
 		}
-		log.Fatalf("Gemini authentication failed: %v", err)
+		log.Fatalf("Failed to complete user setup: %v", errSetup)
+		return
+	}
+
+	storage.Auto = strings.TrimSpace(projectID) == ""
+
+	if !storage.Auto && !storage.Checked {
+		isChecked, errCheck := checkCloudAPIIsEnabled(ctx, httpClient, storage.ProjectID)
+		if errCheck != nil {
+			log.Fatalf("Failed to check if Cloud AI API is enabled: %v", errCheck)
+			return
+		}
+		storage.Checked = isChecked
+		if !isChecked {
+			log.Fatal("Failed to check if Cloud AI API is enabled. If you encounter an error message, please create an issue.")
+			return
+		}
+	}
+
+	updateAuthRecord(record, storage)
+
+	store := sdkAuth.GetTokenStore()
+	if setter, okSetter := store.(interface{ SetBaseDir(string) }); okSetter && cfg != nil {
+		setter.SetBaseDir(cfg.AuthDir)
+	}
+
+	savedPath, errSave := store.Save(ctx, record)
+	if errSave != nil {
+		log.Fatalf("Failed to save token to file: %v", errSave)
 		return
 	}

@@ -67,3 +133,257 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {

 	fmt.Println("Gemini authentication successful!")
 }
+
+func performGeminiCLISetup(ctx context.Context, httpClient *http.Client, storage *gemini.GeminiTokenStorage, requestedProject string) error {
+	metadata := map[string]string{
+		"ideType":    "IDE_UNSPECIFIED",
+		"platform":   "PLATFORM_UNSPECIFIED",
+		"pluginType": "GEMINI",
+	}
+
+	loadReqBody := map[string]any{
+		"metadata": metadata,
+	}
+	if requestedProject != "" {
+		loadReqBody["cloudaicompanionProject"] = requestedProject
+	}
+
+	var loadResp map[string]any
+	if errLoad := callGeminiCLI(ctx, httpClient, "loadCodeAssist", loadReqBody, &loadResp); errLoad != nil {
+		return fmt.Errorf("load code assist: %w", errLoad)
+	}
+
+	tierID := "legacy-tier"
+	if tiers, okTiers := loadResp["allowedTiers"].([]any); okTiers {
+		for _, rawTier := range tiers {
+			tier, okTier := rawTier.(map[string]any)
+			if !okTier {
+				continue
+			}
+			if isDefault, okDefault := tier["isDefault"].(bool); okDefault && isDefault {
+				if id, okID := tier["id"].(string); okID && strings.TrimSpace(id) != "" {
+					tierID = strings.TrimSpace(id)
+					break
+				}
+			}
+		}
+	}
+
+	projectID := strings.TrimSpace(requestedProject)
+	if projectID == "" {
+		if id, okProject := loadResp["cloudaicompanionProject"].(string); okProject {
+			projectID = strings.TrimSpace(id)
+		}
+	}
+	if projectID == "" {
+		return &projectSelectionRequiredError{}
+	}
+
+	onboardReqBody := map[string]any{
+		"tierId":                  tierID,
+		"metadata":                metadata,
+		"cloudaicompanionProject": projectID,
+	}
+
+	// Store the requested project as a fallback in case the response omits it.
+	storage.ProjectID = projectID
+
+	for {
+		var onboardResp map[string]any
+		if errOnboard := callGeminiCLI(ctx, httpClient, "onboardUser", onboardReqBody, &onboardResp); errOnboard != nil {
+			return fmt.Errorf("onboard user: %w", errOnboard)
+		}
+
+		if done, okDone := onboardResp["done"].(bool); okDone && done {
+			if resp, okResp := onboardResp["response"].(map[string]any); okResp {
+				if project, okProject := resp["cloudaicompanionProject"].(map[string]any); okProject {
+					if id, okID := project["id"].(string); okID && strings.TrimSpace(id) != "" {
+						storage.ProjectID = strings.TrimSpace(id)
+					}
+				}
+			}
+			storage.ProjectID = strings.TrimSpace(storage.ProjectID)
+			if storage.ProjectID == "" {
+				storage.ProjectID = projectID
+			}
+			if storage.ProjectID == "" {
+				return fmt.Errorf("onboard user completed without project id")
+			}
+			log.Infof("Onboarding complete. Using Project ID: %s", storage.ProjectID)
+			return nil
+		}
+
+		log.Println("Onboarding in progress, waiting 5 seconds...")
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func callGeminiCLI(ctx context.Context, httpClient *http.Client, endpoint string, body any, result any) error {
+	url := fmt.Sprintf("%s/%s:%s", geminiCLIEndpoint, geminiCLIVersion, endpoint)
+	if strings.HasPrefix(endpoint, "operations/") {
+		url = fmt.Sprintf("%s/%s", geminiCLIEndpoint, endpoint)
+	}
+
+	var reader io.Reader
+	if body != nil {
+		rawBody, errMarshal := json.Marshal(body)
+		if errMarshal != nil {
+			return fmt.Errorf("marshal request body: %w", errMarshal)
+		}
+		reader = bytes.NewReader(rawBody)
+	}
+
+	req, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, url, reader)
+	if errRequest != nil {
+		return fmt.Errorf("create request: %w", errRequest)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("User-Agent", geminiCLIUserAgent)
+	req.Header.Set("X-Goog-Api-Client", geminiCLIApiClient)
+	req.Header.Set("Client-Metadata", geminiCLIClientMetadata)
+
+	resp, errDo := httpClient.Do(req)
+	if errDo != nil {
+		return fmt.Errorf("execute request: %w", errDo)
+	}
+	defer func() {
+		if errClose := resp.Body.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+	}()
+
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("api request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
+	}
+
+	if result == nil {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		return nil
+	}
+
+	if errDecode := json.NewDecoder(resp.Body).Decode(result); errDecode != nil {
+		return fmt.Errorf("decode response body: %w", errDecode)
+	}
+
+	return nil
+}
+
+func fetchGCPProjects(ctx context.Context, httpClient *http.Client) ([]interfaces.GCPProjectProjects, error) {
+	req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, "https://cloudresourcemanager.googleapis.com/v1/projects", nil)
+	if errRequest != nil {
+		return nil, fmt.Errorf("could not create project list request: %w", errRequest)
+	}
+
+	resp, errDo := httpClient.Do(req)
+	if errDo != nil {
+		return nil, fmt.Errorf("failed to execute project list request: %w", errDo)
+	}
+	defer func() {
+		if errClose := resp.Body.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+	}()
+
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("project list request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
+	}
+
+	var projects interfaces.GCPProject
+	if errDecode := json.NewDecoder(resp.Body).Decode(&projects); errDecode != nil {
+		return nil, fmt.Errorf("failed to unmarshal project list: %w", errDecode)
+	}
+
+	return projects.Projects, nil
+}
+
+func showProjectSelectionHelp(email string, projects []interfaces.GCPProjectProjects) {
+	if email != "" {
+		log.Infof("Your account %s needs to specify a project ID.", email)
+	} else {
+		log.Info("You need to specify a project ID.")
+	}
+
+	if len(projects) > 0 {
+		fmt.Println("========================================================================")
+		for _, p := range projects {
+			fmt.Printf("Project ID: %s\n", p.ProjectID)
+			fmt.Printf("Project Name: %s\n", p.Name)
+			fmt.Println("------------------------------------------------------------------------")
+		}
+	} else {
+		fmt.Println("No active projects were returned for this account.")
+	}
+
+	fmt.Printf("Please run this command to login again with a specific project:\n\n%s --login --project_id <project_id>\n", os.Args[0])
+}
+
+func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projectID string) (bool, error) {
+	payload := fmt.Sprintf(`{"project":"%s","request":{"contents":[{"role":"user","parts":[{"text":"Be concise. What is the capital of France?"}]}],"generationConfig":{"thinkingConfig":{"include_thoughts":false,"thinkingBudget":0}}},"model":"gemini-2.5-flash"}`, projectID)
+
+	url := fmt.Sprintf("%s/%s:%s?alt=sse", geminiCLIEndpoint, geminiCLIVersion, "streamGenerateContent")
+	req, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(payload))
+	if errRequest != nil {
+		return false, fmt.Errorf("failed to create request: %w", errRequest)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("User-Agent", geminiCLIUserAgent)
+	req.Header.Set("X-Goog-Api-Client", geminiCLIApiClient)
+	req.Header.Set("Client-Metadata", geminiCLIClientMetadata)
+	req.Header.Set("Accept", "text/event-stream")
+
+	resp, errDo := httpClient.Do(req)
+	if errDo != nil {
+		return false, fmt.Errorf("failed to execute request: %w", errDo)
+	}
+	defer func() {
+		if errClose := resp.Body.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+	}()
+
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		if resp.StatusCode == http.StatusForbidden {
+			activationURL := gjson.GetBytes(bodyBytes, "0.error.details.0.metadata.activationUrl").String()
+			if activationURL != "" {
+				log.Warnf("\n\nPlease activate your account with this url:\n\n%s\n\n And execute this command again:\n%s --login --project_id %s", activationURL, os.Args[0], projectID)
+				return false, nil
+			}
+			log.Warnf("\n\nPlease copy this message and create an issue.\n\n%s\n\n", strings.TrimSpace(string(bodyBytes)))
+			return false, nil
+		}
+		return false, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
+	}
+
+	scanner := bufio.NewScanner(resp.Body)
+	for scanner.Scan() {
+		// Consume the stream to ensure the request succeeds.
+	}
+	if errScan := scanner.Err(); errScan != nil {
+		return false, fmt.Errorf("stream read failed: %w", errScan)
+	}
+
+	return true, nil
+}
+
+func updateAuthRecord(record *cliproxyauth.Auth, storage *gemini.GeminiTokenStorage) {
+	if record == nil || storage == nil {
+		return
+	}
+
+	finalName := fmt.Sprintf("%s-%s.json", storage.Email, storage.ProjectID)
+
+	if record.Metadata == nil {
+		record.Metadata = make(map[string]any)
+	}
+	record.Metadata["email"] = storage.Email
+	record.Metadata["project_id"] = storage.ProjectID
+	record.Metadata["auto"] = storage.Auto
+	record.Metadata["checked"] = storage.Checked
+
+	record.ID = finalName
+	record.FileName = finalName
+	record.Storage = storage
+}