Refactor token management, client initialization, and project handling

- Consolidated `TokenStorage` struct into `internal/auth/models.go` for better organization. - Updated `Client` to use `TokenStorage` for managing email and project ID. - Simplified `SetupUser` method to ensure proper token and project assignment. - Refactored API handlers to leverage new `GetEmail` and `GetProjectID` methods in `Client`. - Cleanup: Removed unused structures and redundant code from `client.go` and `auth.go`. - Adjusted CLI flow in `login.go` and `run.go` for streamlined user onboarding.
2026-02-03 04:50:52 +08:00 · 2025-07-04 17:08:58 +08:00
parent 57ead9a4bc
commit 582280f4c5
7 changed files with 116 additions and 130 deletions
--- a/internal/api/handlers.go
+++ b/internal/api/handlers.go
@@ -407,7 +407,7 @@ func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte)
 		cliClient.RequestMutex.Lock()
 	}

-	log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID)
+	log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
 	jsonTemplate := `{"id":"","object":"chat.completion","created":123456,"model":"model","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}]}`
 	respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools)
 	for {
@@ -501,7 +501,7 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) {
 		cliClient.RequestMutex.Lock()
 	}

-	log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID)
+	log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
 	respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools)
 	for {
 		select {
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"github.com/luispater/CLIProxyAPI/internal/config"
 	log "github.com/sirupsen/logrus"
+	"github.com/skratchdot/open-golang/open"
 	"github.com/tidwall/gjson"
 	"golang.org/x/net/proxy"
 	"io"
@@ -15,7 +16,6 @@ import (
 	"net/url"
 	"time"

-	"github.com/skratchdot/open-golang/open"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
 )
@@ -33,14 +33,6 @@ var (
 	}
 )

-type TokenStorage struct {
-	Token     any    `json:"token"`
-	ProjectID string `json:"project_id"`
-	Email     string `json:"email"`
-	Auto      bool   `json:"auto"`
-	Checked   bool   `json:"checked"`
-}
-
 // GetAuthenticatedClient configures and returns an HTTP client with OAuth2 tokens.
 // It handles the entire flow: loading, refreshing, and fetching new tokens.
 func GetAuthenticatedClient(ctx context.Context, ts *TokenStorage, cfg *config.Config) (*http.Client, error) {
@@ -199,7 +191,8 @@ func getTokenFromWeb(ctx context.Context, config *oauth2.Config) (*oauth2.Token,
 	authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline, oauth2.SetAuthURLParam("prompt", "consent"))
 	log.Debugf("CLI login required.\nAttempting to open authentication page in your browser.\nIf it does not open, please navigate to this URL:\n\n%s\n", authURL)

-	err := open.Run(authURL)
+	var err error
+	err = open.Run(authURL)
 	if err != nil {
 		log.Errorf("Failed to open browser: %v. Please open the URL manually.", err)
 	}
--- a/internal/auth/models.go
+++ b/internal/auth/models.go
@@ -0,0 +1,9 @@
+package auth
+
+type TokenStorage struct {
+	Token     any    `json:"token"`
+	ProjectID string `json:"project_id"`
+	Email     string `json:"email"`
+	Auto      bool   `json:"auto"`
+	Checked   bool   `json:"checked"`
+}
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -22,96 +22,16 @@ import (
 	"time"
 )

-// --- Constants ---
 const (
 	codeAssistEndpoint = "https://cloudcode-pa.googleapis.com"
 	apiVersion         = "v1internal"
-	pluginVersion      = "1.0.0"
+	pluginVersion      = "0.1.9"
 )

-type ErrorMessage struct {
-	StatusCode int
-	Error      error
-}
-
-type GCPProject struct {
-	Projects []GCPProjectProjects `json:"projects"`
-}
-type GCPProjectLabels struct {
-	GenerativeLanguage string `json:"generative-language"`
-}
-type GCPProjectProjects struct {
-	ProjectNumber  string           `json:"projectNumber"`
-	ProjectID      string           `json:"projectId"`
-	LifecycleState string           `json:"lifecycleState"`
-	Name           string           `json:"name"`
-	Labels         GCPProjectLabels `json:"labels"`
-	CreateTime     time.Time        `json:"createTime"`
-}
-
-type Content struct {
-	Role  string `json:"role"`
-	Parts []Part `json:"parts"`
-}
-
-// Part represents a single part of a message's content.
-type Part struct {
-	Text             string            `json:"text,omitempty"`
-	InlineData       *InlineData       `json:"inlineData,omitempty"`
-	FunctionCall     *FunctionCall     `json:"functionCall,omitempty"`
-	FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"`
-}
-
-type InlineData struct {
-	MimeType string `json:"mime_type,omitempty"`
-	Data     string `json:"data,omitempty"`
-}
-
-// FunctionCall represents a tool call requested by the model.
-type FunctionCall struct {
-	Name string                 `json:"name"`
-	Args map[string]interface{} `json:"args"`
-}
-
-// FunctionResponse represents the result of a tool execution.
-type FunctionResponse struct {
-	Name     string                 `json:"name"`
-	Response map[string]interface{} `json:"response"`
-}
-
-// GenerateContentRequest is the request payload for the streamGenerateContent endpoint.
-type GenerateContentRequest struct {
-	Contents         []Content         `json:"contents"`
-	Tools            []ToolDeclaration `json:"tools,omitempty"`
-	GenerationConfig `json:"generationConfig"`
-}
-
-// GenerationConfig defines model generation parameters.
-type GenerationConfig struct {
-	ThinkingConfig GenerationConfigThinkingConfig `json:"thinkingConfig,omitempty"`
-	Temperature    float64                        `json:"temperature,omitempty"`
-	TopP           float64                        `json:"topP,omitempty"`
-	TopK           float64                        `json:"topK,omitempty"`
-	// Temperature, TopP, TopK, etc. can be added here.
-}
-
-type GenerationConfigThinkingConfig struct {
-	IncludeThoughts bool `json:"include_thoughts,omitempty"`
-}
-
-// ToolDeclaration is the structure for declaring tools to the API.
-// For now, we'll assume a simple structure. A more complete implementation
-// would mirror the OpenAPI schema definition.
-type ToolDeclaration struct {
-	FunctionDeclarations []interface{} `json:"functionDeclarations"`
-}
-
 // Client is the main client for interacting with the CLI API.
 type Client struct {
 	httpClient   *http.Client
-	ProjectID    string
 	RequestMutex sync.Mutex
-	Email        string
 	tokenStorage *auth.TokenStorage
 	cfg          *config.Config
 }
@@ -145,9 +65,17 @@ func (c *Client) IsAuto() bool {
 	return c.tokenStorage.Auto
 }

+func (c *Client) GetEmail() string {
+	return c.tokenStorage.Email
+}
+
+func (c *Client) GetProjectID() string {
+	return c.tokenStorage.ProjectID
+}
+
 // SetupUser performs the initial user onboarding and setup.
-func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string, error) {
-	c.Email = email
+func (c *Client) SetupUser(ctx context.Context, email, projectID string) error {
+	c.tokenStorage.Email = email
 	log.Info("Performing user onboarding...")

 	// 1. LoadCodeAssist
@@ -161,7 +89,7 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string
 	var loadAssistResp map[string]interface{}
 	err := c.makeAPIRequest(ctx, "loadCodeAssist", "POST", loadAssistReqBody, &loadAssistResp)
 	if err != nil {
-		return projectID, fmt.Errorf("failed to load code assist: %w", err)
+		return fmt.Errorf("failed to load code assist: %w", err)
 	}

 	// a, _ := json.Marshal(&loadAssistResp)
@@ -197,14 +125,14 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string
 	if onboardProjectID != "" {
 		onboardReqBody["cloudaicompanionProject"] = onboardProjectID
 	} else {
-		return projectID, fmt.Errorf("failed to start user onboarding, need define a project id")
+		return fmt.Errorf("failed to start user onboarding, need define a project id")
 	}

 	for {
 		var lroResp map[string]interface{}
 		err = c.makeAPIRequest(ctx, "onboardUser", "POST", onboardReqBody, &lroResp)
 		if err != nil {
-			return projectID, fmt.Errorf("failed to start user onboarding: %w", err)
+			return fmt.Errorf("failed to start user onboarding: %w", err)
 		}
 		// a, _ := json.Marshal(&lroResp)
 		// log.Debug(string(a))
@@ -214,12 +142,12 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string
 		if doneOk && done {
 			if project, projectOk := lroResp["response"].(map[string]interface{})["cloudaicompanionProject"].(map[string]interface{}); projectOk {
 				if projectID != "" {
-					c.ProjectID = projectID
+					c.tokenStorage.ProjectID = projectID
 				} else {
-					c.ProjectID = project["id"].(string)
+					c.tokenStorage.ProjectID = project["id"].(string)
 				}
-				log.Infof("Onboarding complete. Using Project ID: %s", c.ProjectID)
-				return c.ProjectID, nil
+				log.Infof("Onboarding complete. Using Project ID: %s", c.tokenStorage.ProjectID)
+				return nil
 			}
 		} else {
 			log.Println("Onboarding in progress, waiting 5 seconds...")
@@ -356,7 +284,7 @@ func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model st
 		request.Tools = tools

 		requestBody := map[string]interface{}{
-			"project": c.ProjectID, // Assuming ProjectID is available
+			"project": c.tokenStorage.ProjectID, // Assuming ProjectID is available
 			"request": request,
 			"model":   model,
 		}
--- a/internal/client/models.go
+++ b/internal/client/models.go
@@ -0,0 +1,80 @@
+package client
+
+import "time"
+
+type ErrorMessage struct {
+	StatusCode int
+	Error      error
+}
+
+type GCPProject struct {
+	Projects []GCPProjectProjects `json:"projects"`
+}
+type GCPProjectLabels struct {
+	GenerativeLanguage string `json:"generative-language"`
+}
+type GCPProjectProjects struct {
+	ProjectNumber  string           `json:"projectNumber"`
+	ProjectID      string           `json:"projectId"`
+	LifecycleState string           `json:"lifecycleState"`
+	Name           string           `json:"name"`
+	Labels         GCPProjectLabels `json:"labels"`
+	CreateTime     time.Time        `json:"createTime"`
+}
+
+type Content struct {
+	Role  string `json:"role"`
+	Parts []Part `json:"parts"`
+}
+
+// Part represents a single part of a message's content.
+type Part struct {
+	Text             string            `json:"text,omitempty"`
+	InlineData       *InlineData       `json:"inlineData,omitempty"`
+	FunctionCall     *FunctionCall     `json:"functionCall,omitempty"`
+	FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"`
+}
+
+type InlineData struct {
+	MimeType string `json:"mime_type,omitempty"`
+	Data     string `json:"data,omitempty"`
+}
+
+// FunctionCall represents a tool call requested by the model.
+type FunctionCall struct {
+	Name string                 `json:"name"`
+	Args map[string]interface{} `json:"args"`
+}
+
+// FunctionResponse represents the result of a tool execution.
+type FunctionResponse struct {
+	Name     string                 `json:"name"`
+	Response map[string]interface{} `json:"response"`
+}
+
+// GenerateContentRequest is the request payload for the streamGenerateContent endpoint.
+type GenerateContentRequest struct {
+	Contents         []Content         `json:"contents"`
+	Tools            []ToolDeclaration `json:"tools,omitempty"`
+	GenerationConfig `json:"generationConfig"`
+}
+
+// GenerationConfig defines model generation parameters.
+type GenerationConfig struct {
+	ThinkingConfig GenerationConfigThinkingConfig `json:"thinkingConfig,omitempty"`
+	Temperature    float64                        `json:"temperature,omitempty"`
+	TopP           float64                        `json:"topP,omitempty"`
+	TopK           float64                        `json:"topK,omitempty"`
+	// Temperature, TopP, TopK, etc. can be added here.
+}
+
+type GenerationConfigThinkingConfig struct {
+	IncludeThoughts bool `json:"include_thoughts,omitempty"`
+}
+
+// ToolDeclaration is the structure for declaring tools to the API.
+// For now, we'll assume a simple structure. A more complete implementation
+// would mirror the OpenAPI schema definition.
+type ToolDeclaration struct {
+	FunctionDeclarations []interface{} `json:"functionDeclarations"`
+}
--- a/internal/cmd/login.go
+++ b/internal/cmd/login.go
@@ -29,7 +29,7 @@ func DoLogin(cfg *config.Config, projectID string) {

 	// 3. Initialize CLI Client
 	cliClient := client.NewClient(httpClient, &ts, cfg)
-	projectID, err = cliClient.SetupUser(clientCtx, ts.Email, projectID)
+	err = cliClient.SetupUser(clientCtx, ts.Email, projectID)
 	if err != nil {
 		if err.Error() == "failed to start user onboarding, need define a project id" {
 			log.Error("failed to start user onboarding")
@@ -52,8 +52,7 @@ func DoLogin(cfg *config.Config, projectID string) {
 			log.Fatalf("failed to complete user setup: %v", err)
 		}
 	} else {
-		auto := ts.ProjectID == ""
-		cliClient.SetProjectID(projectID)
+		auto := projectID == ""
 		cliClient.SetIsAuto(auto)

 		if !cliClient.IsChecked() && !cliClient.IsAuto() {
--- a/internal/cmd/run.go
+++ b/internal/cmd/run.go
@@ -57,30 +57,7 @@ func StartService(cfg *config.Config) {

 				// 3. Initialize CLI Client
 				cliClient := client.NewClient(httpClient, &ts, cfg)
-				if _, err = cliClient.SetupUser(clientCtx, ts.Email, ts.ProjectID); err != nil {
-					if err.Error() == "failed to start user onboarding, need define a project id" {
-						log.Error("failed to start user onboarding")
-						project, errGetProjectList := cliClient.GetProjectList(clientCtx)
-						if errGetProjectList != nil {
-							log.Fatalf("failed to complete user setup: %v", err)
-						} else {
-							log.Infof("Your account %s needs specify a project id.", ts.Email)
-							log.Info("========================================================================")
-							for i := 0; i < len(project.Projects); i++ {
-								log.Infof("Project ID: %s", project.Projects[i].ProjectID)
-								log.Infof("Project Name: %s", project.Projects[i].Name)
-								log.Info("========================================================================")
-							}
-							log.Infof("Please run this command to login again:\n\n%s --login --project_id <project_id>\n", os.Args[0])
-						}
-					} else {
-						// Log as a warning because in some cases, the CLI might still be usable
-						// or the user might want to retry setup later.
-						log.Fatalf("failed to complete user setup: %v", err)
-					}
-				} else {
-					cliClients = append(cliClients, cliClient)
-				}
+				cliClients = append(cliClients, cliClient)
 			}
 		}
 		return nil