Compare commits

..

2 Commits

Author SHA1 Message Date
Luis Pater
6d98a71796 feat(login): add interactive project selection and improve onboarding flow
- Introduced `promptForProjectSelection` to enable interactive project selection for better user onboarding.
- Improved project validation and handling when no preset project ID is provided.
- Added a default project prompt mechanism to guide users through project selection seamlessly.
- Refined error handling for onboarding and project selection failures.
2025-09-30 01:52:03 +08:00
Luis Pater
1c91823308 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.
2025-09-30 00:04:58 +08:00

View File

@@ -4,18 +4,45 @@
package cmd
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"os"
"strconv"
"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 +53,90 @@ 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.")
projects, errProjects := fetchGCPProjects(ctx, httpClient)
if errProjects != nil {
log.Fatalf("Failed to get project list: %v", errProjects)
return
}
promptFn := options.Prompt
if promptFn == nil {
promptFn = defaultProjectPrompt()
}
selectedProjectID := promptForProjectSelection(projects, strings.TrimSpace(projectID), promptFn)
if strings.TrimSpace(selectedProjectID) == "" {
log.Fatal("No project selected; aborting login.")
return
}
if errSetup := performGeminiCLISetup(ctx, httpClient, storage, selectedProjectID); 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.Fatalf("Gemini authentication failed: %v", err)
log.Fatalf("Failed to complete user setup: %v", errSetup)
return
}
storage.Auto = false
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 +146,342 @@ 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
}
// promptForProjectSelection prints available projects and returns the chosen project ID.
func promptForProjectSelection(projects []interfaces.GCPProjectProjects, presetID string, promptFn func(string) (string, error)) string {
trimmedPreset := strings.TrimSpace(presetID)
if len(projects) == 0 {
if trimmedPreset != "" {
return trimmedPreset
}
fmt.Println("No Google Cloud projects are available for selection.")
return ""
}
fmt.Println("Available Google Cloud projects:")
defaultIndex := 0
for idx, project := range projects {
fmt.Printf("[%d] %s (%s)\n", idx+1, project.ProjectID, project.Name)
if trimmedPreset != "" && project.ProjectID == trimmedPreset {
defaultIndex = idx
}
}
defaultID := projects[defaultIndex].ProjectID
if trimmedPreset != "" {
for _, project := range projects {
if project.ProjectID == trimmedPreset {
return trimmedPreset
}
}
log.Warnf("Provided project ID %s not found in available projects; please choose from the list.", trimmedPreset)
}
for {
promptMsg := fmt.Sprintf("Enter project ID [%s]: ", defaultID)
answer, errPrompt := promptFn(promptMsg)
if errPrompt != nil {
log.Errorf("Project selection prompt failed: %v", errPrompt)
return defaultID
}
answer = strings.TrimSpace(answer)
if answer == "" {
return defaultID
}
for _, project := range projects {
if project.ProjectID == answer {
return project.ProjectID
}
}
if idx, errAtoi := strconv.Atoi(answer); errAtoi == nil {
if idx >= 1 && idx <= len(projects) {
return projects[idx-1].ProjectID
}
}
fmt.Println("Invalid selection, enter a project ID or a number from the list.")
}
}
func defaultProjectPrompt() func(string) (string, error) {
reader := bufio.NewReader(os.Stdin)
return func(prompt string) (string, error) {
fmt.Print(prompt)
line, errRead := reader.ReadString('\n')
if errRead != nil {
if errors.Is(errRead, io.EOF) {
return strings.TrimSpace(line), nil
}
return "", errRead
}
return strings.TrimSpace(line), 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) {
serviceUsageURL := "https://serviceusage.googleapis.com"
requiredServices := []string{
"geminicloudassist.googleapis.com", // Gemini Cloud Assist API
"cloudaicompanion.googleapis.com", // Gemini for Google Cloud API
}
for _, service := range requiredServices {
checkUrl := fmt.Sprintf("%s/v1/projects/%s/services/%s", serviceUsageURL, projectID, service)
req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, checkUrl, nil)
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)
resp, errDo := httpClient.Do(req)
if errDo != nil {
return false, fmt.Errorf("failed to execute request: %w", errDo)
}
if resp.StatusCode == http.StatusOK {
bodyBytes, _ := io.ReadAll(resp.Body)
if gjson.GetBytes(bodyBytes, "state").String() == "ENABLED" {
_ = resp.Body.Close()
continue
}
}
_ = resp.Body.Close()
enableUrl := fmt.Sprintf("%s/v1/projects/%s/services/%s:enable", serviceUsageURL, projectID, service)
req, errRequest = http.NewRequestWithContext(ctx, http.MethodPost, enableUrl, strings.NewReader("{}"))
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)
resp, errDo = httpClient.Do(req)
if errDo != nil {
return false, fmt.Errorf("failed to execute request: %w", errDo)
}
bodyBytes, _ := io.ReadAll(resp.Body)
errMessage := string(bodyBytes)
errMessageResult := gjson.GetBytes(bodyBytes, "error.message")
if errMessageResult.Exists() {
errMessage = errMessageResult.String()
}
if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
_ = resp.Body.Close()
continue
} else if resp.StatusCode == http.StatusBadRequest {
_ = resp.Body.Close()
if strings.Contains(strings.ToLower(errMessage), "already enabled") {
continue
}
}
return false, fmt.Errorf("project activation required: %s", errMessage)
}
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
}