From 79acea5976504594f0f96fd0e39d595e8a12b0d8 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 4 Jul 2025 00:43:15 +0800 Subject: [PATCH] Refactor authentication and service initialization code - Moved login and service management logic to `internal/cmd` package (`login.go` and `run.go`). - Introduced `DoLogin` and `StartService` functions for modularity. - Enhanced error handling by using structured `ErrorMessage` in `Client`. - Improved token file saving process and added project-specific token identification. - Updated API handlers to handle more detailed error responses, including status codes. --- cmd/server/main.go | 154 +------------------------------------- internal/api/handlers.go | 12 +-- internal/auth/auth.go | 59 ++++++++++----- internal/client/client.go | 53 ++++++++----- internal/cmd/login.go | 62 +++++++++++++++ internal/cmd/run.go | 122 ++++++++++++++++++++++++++++++ 6 files changed, 267 insertions(+), 195 deletions(-) create mode 100644 internal/cmd/login.go create mode 100644 internal/cmd/run.go diff --git a/cmd/server/main.go b/cmd/server/main.go index 204ea2e1..e6e9e686 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -2,23 +2,14 @@ package main import ( "bytes" - "context" - "encoding/json" "flag" "fmt" - "github.com/luispater/CLIProxyAPI/internal/api" - "github.com/luispater/CLIProxyAPI/internal/auth" - "github.com/luispater/CLIProxyAPI/internal/client" + "github.com/luispater/CLIProxyAPI/internal/cmd" "github.com/luispater/CLIProxyAPI/internal/config" log "github.com/sirupsen/logrus" - "io/fs" "os" - "os/signal" "path" - "path/filepath" "strings" - "syscall" - "time" ) type LogFormatter struct { @@ -95,147 +86,8 @@ func main() { } if login { - var ts auth.TokenStorage - if projectID != "" { - ts.ProjectID = projectID - } - - // 2. Initialize authenticated HTTP Client - clientCtx := context.Background() - - log.Info("Initializing authentication...") - httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg) - if errGetClient != nil { - log.Fatalf("failed to get authenticated client: %v", errGetClient) - return - } - log.Info("Authentication successful.") - - // 3. Initialize CLI Client - cliClient := client.NewClient(httpClient) - if err = cliClient.SetupUser(clientCtx, ts.Email, 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 \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) - } - } + cmd.DoLogin(cfg, projectID) } else { - // Create API server configuration - apiConfig := &api.ServerConfig{ - Port: fmt.Sprintf("%d", cfg.Port), - Debug: cfg.Debug, - ApiKeys: cfg.ApiKeys, - } - - cliClients := make([]*client.Client, 0) - err = filepath.Walk(cfg.AuthDir, func(path string, info fs.FileInfo, err error) error { - if err != nil { - return err - } - - if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") { - log.Debugf(path) - f, errOpen := os.Open(path) - if errOpen != nil { - return errOpen - } - defer func() { - _ = f.Close() - }() - - var ts auth.TokenStorage - if err = json.NewDecoder(f).Decode(&ts); err == nil { - // 2. Initialize authenticated HTTP Client - clientCtx := context.Background() - - log.Info("Initializing authentication...") - httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg) - if errGetClient != nil { - log.Fatalf("failed to get authenticated client: %v", errGetClient) - return errGetClient - } - log.Info("Authentication successful.") - - // 3. Initialize CLI Client - cliClient := client.NewClient(httpClient) - 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 \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) - } - } - } - return nil - }) - - // Create API server - apiServer := api.NewServer(apiConfig, cliClients) - log.Infof("Starting API server on port %s", apiConfig.Port) - if err = apiServer.Start(); err != nil { - log.Fatalf("API server failed to start: %v", err) - return - } - - // Set up graceful shutdown - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - - for { - select { - case <-sigChan: - log.Debugf("Received shutdown signal. Cleaning up...") - - // Create shutdown context - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - _ = ctx // Mark ctx as used to avoid error, as apiServer.Stop(ctx) is commented out - - // Stop API server - if err = apiServer.Stop(ctx); err != nil { - log.Debugf("Error stopping API server: %v", err) - } - cancel() - - log.Debugf("Cleanup completed. Exiting...") - os.Exit(0) - case <-time.After(5 * time.Second): - - } - } + cmd.StartService(cfg) } - } diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 8c58c9ee..012ec4c6 100644 --- 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", cliClient.Email) + log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID) 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 { @@ -429,8 +429,8 @@ func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte) } case err, okError := <-errChan: if okError { - c.Status(http.StatusInternalServerError) - _, _ = fmt.Fprint(c.Writer, err.Error()) + c.Status(err.StatusCode) + _, _ = fmt.Fprint(c.Writer, err.Error.Error()) flusher.Flush() // c.JSON(http.StatusInternalServerError, ErrorResponse{ // Error: ErrorDetail{ @@ -501,7 +501,7 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) { cliClient.RequestMutex.Lock() } - log.Debugf("Request use account: %s", cliClient.Email) + log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID) respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools) for { select { @@ -526,8 +526,8 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) { } case err, okError := <-errChan: if okError { - c.Status(http.StatusInternalServerError) - _, _ = fmt.Fprint(c.Writer, err.Error()) + c.Status(err.StatusCode) + _, _ = fmt.Fprint(c.Writer, err.Error.Error()) flusher.Flush() // c.JSON(http.StatusInternalServerError, ErrorResponse{ // Error: ErrorDetail{ diff --git a/internal/auth/auth.go b/internal/auth/auth.go index a0d146ae..44092f28 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -39,6 +39,7 @@ type TokenStorage struct { Token any `json:"token"` ProjectID string `json:"project_id"` Email string `json:"email"` + Auto bool `json:"auto"` } // GetAuthenticatedClient configures and returns an HTTP client with OAuth2 tokens. @@ -95,11 +96,12 @@ func GetAuthenticatedClient(ctx context.Context, ts *TokenStorage, cfg *config.C if err != nil { return nil, fmt.Errorf("failed to get token from web: %w", err) } - ts, err = saveTokenToFile(ctx, conf, token, ts.ProjectID, cfg.AuthDir) - if err != nil { - // Log the error but proceed, as we have a valid token for the session. + newTs, errSaveTokenToFile := saveTokenToFile(ctx, conf, token, ts.ProjectID, cfg.AuthDir) + if errSaveTokenToFile != nil { log.Errorf("Warning: failed to save token to file: %v", err) + return nil, errSaveTokenToFile } + *ts = *newTs } tsToken, _ := json.Marshal(ts.Token) if err = json.Unmarshal(tsToken, &token); err != nil { @@ -139,19 +141,6 @@ func saveTokenToFile(ctx context.Context, config *oauth2.Config, token *oauth2.T log.Info("Failed to get user email from token") } - log.Infof("Saving credentials to %s", filepath.Join(authDir, fmt.Sprintf("%s.json", emailResult.String()))) - if err = os.MkdirAll(authDir, 0700); err != nil { - return nil, fmt.Errorf("failed to create directory: %w", err) - } - - f, err := os.Create(filepath.Join(authDir, fmt.Sprintf("%s.json", emailResult.String()))) - if err != nil { - return nil, fmt.Errorf("failed to create token file: %w", err) - } - defer func() { - _ = f.Close() - }() - var ifToken map[string]any jsonData, _ := json.Marshal(token) err = json.Unmarshal(jsonData, &ifToken) @@ -171,12 +160,46 @@ func saveTokenToFile(ctx context.Context, config *oauth2.Config, token *oauth2.T Email: emailResult.String(), } - if err = json.NewEncoder(f).Encode(ts); err != nil { - return nil, fmt.Errorf("failed to write token to file: %w", err) + if err = os.MkdirAll(authDir, 0700); err != nil { + return nil, fmt.Errorf("failed to create directory: %w", err) + } + + if projectID != "" { + log.Infof("Saving credentials to %s", filepath.Join(authDir, fmt.Sprintf("%s-%s.json", emailResult.String(), projectID))) + + f, errCreate := os.Create(filepath.Join(authDir, fmt.Sprintf("%s-%s.json", emailResult.String(), projectID))) + if errCreate != nil { + return nil, fmt.Errorf("failed to create token file: %w", errCreate) + } + defer func() { + _ = f.Close() + }() + + if err = json.NewEncoder(f).Encode(ts); err != nil { + return nil, fmt.Errorf("failed to write token to file: %w", err) + } } return &ts, nil } +func SaveTokenToFile(ts *TokenStorage, cfg *config.Config, auto bool) error { + ts.Auto = auto + fileName := filepath.Join(cfg.AuthDir, fmt.Sprintf("%s-%s.json", ts.Email, ts.ProjectID)) + log.Infof("Saving credentials to %s", fileName) + f, err := os.Create(fileName) + if err != nil { + return fmt.Errorf("failed to create token file: %w", err) + } + defer func() { + _ = f.Close() + }() + + if err = json.NewEncoder(f).Encode(ts); err != nil { + return fmt.Errorf("failed to write token to file: %w", err) + } + return nil +} + // getTokenFromWeb starts a local server to handle the OAuth2 flow. func getTokenFromWeb(ctx context.Context, config *oauth2.Config) (*oauth2.Token, error) { // Use a channel to pass the authorization code from the HTTP handler to the main function. diff --git a/internal/client/client.go b/internal/client/client.go index 20a96e7c..50194314 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -25,6 +25,11 @@ const ( pluginVersion = "1.0.0" ) +type ErrorMessage struct { + StatusCode int + Error error +} + type GCPProject struct { Projects []GCPProjectProjects `json:"projects"` } @@ -100,7 +105,7 @@ type ToolDeclaration struct { // Client is the main client for interacting with the CLI API. type Client struct { httpClient *http.Client - projectID string + ProjectID string RequestMutex sync.Mutex Email string } @@ -113,7 +118,7 @@ func NewClient(httpClient *http.Client) *Client { } // SetupUser performs the initial user onboarding and setup. -func (c *Client) SetupUser(ctx context.Context, email, projectID string) error { +func (c *Client) SetupUser(ctx context.Context, email, projectID string, auto bool) (string, error) { c.Email = email log.Info("Performing user onboarding...") @@ -128,11 +133,14 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) error { var loadAssistResp map[string]interface{} err := c.makeAPIRequest(ctx, "loadCodeAssist", "POST", loadAssistReqBody, &loadAssistResp) if err != nil { - return fmt.Errorf("failed to load code assist: %w", err) + return projectID, fmt.Errorf("failed to load code assist: %w", err) } // a, _ := json.Marshal(&loadAssistResp) // log.Debug(string(a)) + // + // a, _ = json.Marshal(loadAssistReqBody) + // log.Debug(string(a)) // 2. OnboardUser var onboardTierID = "legacy-tier" @@ -161,13 +169,13 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) error { if onboardProjectID != "" { onboardReqBody["cloudaicompanionProject"] = onboardProjectID } else { - return fmt.Errorf("failed to start user onboarding, need define a project id") + return projectID, fmt.Errorf("failed to start user onboarding, need define a project id") } var lroResp map[string]interface{} err = c.makeAPIRequest(ctx, "onboardUser", "POST", onboardReqBody, &lroResp) if err != nil { - return fmt.Errorf("failed to start user onboarding: %w", err) + return projectID, fmt.Errorf("failed to start user onboarding: %w", err) } // a, _ = json.Marshal(&lroResp) @@ -176,12 +184,17 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) error { // 3. Poll Long-Running Operation (LRO) if done, doneOk := lroResp["done"].(bool); doneOk && done { if project, projectOk := lroResp["response"].(map[string]interface{})["cloudaicompanionProject"].(map[string]interface{}); projectOk { - c.projectID = project["id"].(string) - log.Infof("Onboarding complete. Using Project ID: %s", c.projectID) - return nil + if projectID != "" && !auto { + c.ProjectID = projectID + log.Infof("Onboarding complete. Project ID: %s is being enforced. Maybe you need to enable 'Gemini for Google Cloud' once in Google Cloud Console.", c.ProjectID) + } else { + c.ProjectID = project["id"].(string) + log.Infof("Onboarding complete. Using Project ID: %s", c.ProjectID) + } + return c.ProjectID, nil } } - return fmt.Errorf("failed to get operation name from onboarding response: %v", lroResp) + return projectID, fmt.Errorf("failed to get operation name from onboarding response: %v", lroResp) } // makeAPIRequest handles making requests to the CLI API endpoints. @@ -240,7 +253,7 @@ func (c *Client) makeAPIRequest(ctx context.Context, endpoint, method string, bo } // StreamAPIRequest handles making streaming requests to the CLI API endpoints. -func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body interface{}) (io.ReadCloser, error) { +func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body interface{}) (io.ReadCloser, *ErrorMessage) { var jsonBody []byte var err error if byteBody, ok := body.([]byte); ok { @@ -248,7 +261,7 @@ func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body int } else { jsonBody, err = json.Marshal(body) if err != nil { - return nil, fmt.Errorf("failed to marshal request body: %w", err) + return nil, &ErrorMessage{500, fmt.Errorf("failed to marshal request body: %w", err)} } } // log.Debug(string(jsonBody)) @@ -259,12 +272,12 @@ func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body int req, err := http.NewRequestWithContext(ctx, "POST", url, reqBody) if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + return nil, &ErrorMessage{500, fmt.Errorf("failed to create request: %w", err)} } token, err := c.httpClient.Transport.(*oauth2.Transport).Source.Token() if err != nil { - return nil, fmt.Errorf("failed to get token: %w", err) + return nil, &ErrorMessage{500, fmt.Errorf("failed to get token: %w", err)} } // Set headers @@ -276,7 +289,7 @@ func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body int resp, err := c.httpClient.Do(req) if err != nil { - return nil, fmt.Errorf("failed to execute request: %w", err) + return nil, &ErrorMessage{500, fmt.Errorf("failed to execute request: %w", err)} } if resp.StatusCode < 200 || resp.StatusCode >= 300 { @@ -285,7 +298,7 @@ func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body int }() bodyBytes, _ := io.ReadAll(resp.Body) - return nil, fmt.Errorf(string(bodyBytes)) + return nil, &ErrorMessage{resp.StatusCode, fmt.Errorf(string(bodyBytes))} // return nil, fmt.Errorf("api streaming request failed with status %d: %s", resp.StatusCode, string(bodyBytes)) } @@ -293,9 +306,9 @@ func (c *Client) StreamAPIRequest(ctx context.Context, endpoint string, body int } // SendMessageStream handles a single conversational turn, including tool calls. -func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model string, contents []Content, tools []ToolDeclaration) (<-chan []byte, <-chan error) { +func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model string, contents []Content, tools []ToolDeclaration) (<-chan []byte, <-chan *ErrorMessage) { dataTag := []byte("data: ") - errChan := make(chan error) + errChan := make(chan *ErrorMessage) dataChan := make(chan []byte) go func() { defer close(errChan) @@ -312,7 +325,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.ProjectID, // Assuming ProjectID is available "request": request, "model": model, } @@ -370,9 +383,9 @@ func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model st } } - if err = scanner.Err(); err != nil { + if errScanner := scanner.Err(); errScanner != nil { // log.Println(err) - errChan <- err + errChan <- &ErrorMessage{500, errScanner} _ = stream.Close() return } diff --git a/internal/cmd/login.go b/internal/cmd/login.go new file mode 100644 index 00000000..66a04a20 --- /dev/null +++ b/internal/cmd/login.go @@ -0,0 +1,62 @@ +package cmd + +import ( + "context" + "github.com/luispater/CLIProxyAPI/internal/auth" + "github.com/luispater/CLIProxyAPI/internal/client" + "github.com/luispater/CLIProxyAPI/internal/config" + log "github.com/sirupsen/logrus" + "os" +) + +func DoLogin(cfg *config.Config, projectID string) { + var err error + var ts auth.TokenStorage + if projectID != "" { + ts.ProjectID = projectID + } + + // 2. Initialize authenticated HTTP Client + clientCtx := context.Background() + + log.Info("Initializing authentication...") + httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg) + if errGetClient != nil { + log.Fatalf("failed to get authenticated client: %v", errGetClient) + return + } + log.Info("Authentication successful.") + + // 3. Initialize CLI Client + cliClient := client.NewClient(httpClient) + projectID, err = cliClient.SetupUser(clientCtx, ts.Email, projectID, ts.Auto) + if 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 \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 { + auto := ts.ProjectID == "" + ts.ProjectID = projectID + err = auth.SaveTokenToFile(&ts, cfg, auto) + if err != nil { + log.Fatal(err) + } + } +} diff --git a/internal/cmd/run.go b/internal/cmd/run.go new file mode 100644 index 00000000..52500037 --- /dev/null +++ b/internal/cmd/run.go @@ -0,0 +1,122 @@ +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "github.com/luispater/CLIProxyAPI/internal/api" + "github.com/luispater/CLIProxyAPI/internal/auth" + "github.com/luispater/CLIProxyAPI/internal/client" + "github.com/luispater/CLIProxyAPI/internal/config" + log "github.com/sirupsen/logrus" + "io/fs" + "os" + "os/signal" + "path/filepath" + "strings" + "syscall" + "time" +) + +func StartService(cfg *config.Config) { + // Create API server configuration + apiConfig := &api.ServerConfig{ + Port: fmt.Sprintf("%d", cfg.Port), + Debug: cfg.Debug, + ApiKeys: cfg.ApiKeys, + } + + cliClients := make([]*client.Client, 0) + err := filepath.Walk(cfg.AuthDir, func(path string, info fs.FileInfo, err error) error { + if err != nil { + return err + } + + if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") { + log.Debugf(path) + f, errOpen := os.Open(path) + if errOpen != nil { + return errOpen + } + defer func() { + _ = f.Close() + }() + + var ts auth.TokenStorage + if err = json.NewDecoder(f).Decode(&ts); err == nil { + // 2. Initialize authenticated HTTP Client + clientCtx := context.Background() + + log.Info("Initializing authentication...") + httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg) + if errGetClient != nil { + log.Fatalf("failed to get authenticated client: %v", errGetClient) + return errGetClient + } + log.Info("Authentication successful.") + + // 3. Initialize CLI Client + cliClient := client.NewClient(httpClient) + if _, err = cliClient.SetupUser(clientCtx, ts.Email, ts.ProjectID, ts.Auto); 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 \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) + } + } + } + return nil + }) + + // Create API server + apiServer := api.NewServer(apiConfig, cliClients) + log.Infof("Starting API server on port %s", apiConfig.Port) + if err = apiServer.Start(); err != nil { + log.Fatalf("API server failed to start: %v", err) + return + } + + // Set up graceful shutdown + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + + for { + select { + case <-sigChan: + log.Debugf("Received shutdown signal. Cleaning up...") + + // Create shutdown context + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + _ = ctx // Mark ctx as used to avoid error, as apiServer.Stop(ctx) is commented out + + // Stop API server + if err = apiServer.Stop(ctx); err != nil { + log.Debugf("Error stopping API server: %v", err) + } + cancel() + + log.Debugf("Cleanup completed. Exiting...") + os.Exit(0) + case <-time.After(5 * time.Second): + + } + } +}