// Package handlers provides core API handler functionality for the CLI Proxy API server. // It includes common types, client management, load balancing, and error handling // shared across all API endpoint handlers (OpenAI, Claude, Gemini). package handlers import ( "fmt" "sync" "github.com/gin-gonic/gin" "github.com/luispater/CLIProxyAPI/internal/client" "github.com/luispater/CLIProxyAPI/internal/config" "github.com/luispater/CLIProxyAPI/internal/util" log "github.com/sirupsen/logrus" "golang.org/x/net/context" ) // ErrorResponse represents a standard error response format for the API. // It contains a single ErrorDetail field. type ErrorResponse struct { // Error contains detailed information about the error that occurred. Error ErrorDetail `json:"error"` } // ErrorDetail provides specific information about an error that occurred. // It includes a human-readable message, an error type, and an optional error code. type ErrorDetail struct { // Message is a human-readable message providing more details about the error. Message string `json:"message"` // Type is the category of error that occurred (e.g., "invalid_request_error"). Type string `json:"type"` // Code is a short code identifying the error, if applicable. Code string `json:"code,omitempty"` } // APIHandlers contains the handlers for API endpoints. // It holds a pool of clients to interact with the backend service and manages // load balancing, client selection, and configuration. type APIHandlers struct { // CliClients is the pool of available AI service clients. CliClients []client.Client // Cfg holds the current application configuration. Cfg *config.Config // Mutex ensures thread-safe access to shared resources. Mutex *sync.Mutex // LastUsedClientIndex tracks the last used client index for each provider // to implement round-robin load balancing. LastUsedClientIndex map[string]int // apiResponseData recording provider api response data apiResponseData map[*gin.Context][]byte } // NewAPIHandlers creates a new API handlers instance. // It takes a slice of clients and configuration as input. // // Parameters: // - cliClients: A slice of AI service clients // - cfg: The application configuration // // Returns: // - *APIHandlers: A new API handlers instance func NewAPIHandlers(cliClients []client.Client, cfg *config.Config) *APIHandlers { return &APIHandlers{ CliClients: cliClients, Cfg: cfg, Mutex: &sync.Mutex{}, LastUsedClientIndex: make(map[string]int), apiResponseData: make(map[*gin.Context][]byte), } } // UpdateClients updates the handlers' client list and configuration. // This method is called when the configuration or authentication tokens change. // // Parameters: // - clients: The new slice of AI service clients // - cfg: The new application configuration func (h *APIHandlers) UpdateClients(clients []client.Client, cfg *config.Config) { h.CliClients = clients h.Cfg = cfg } // GetClient returns an available client from the pool using round-robin load balancing. // It checks for quota limits and tries to find an unlocked client for immediate use. // The modelName parameter is used to check quota status for specific models. // // Parameters: // - modelName: The name of the model to be used // - isGenerateContent: Optional parameter to indicate if this is for content generation // // Returns: // - client.Client: An available client for the requested model // - *client.ErrorMessage: An error message if no client is available func (h *APIHandlers) GetClient(modelName string, isGenerateContent ...bool) (client.Client, *client.ErrorMessage) { provider := util.GetProviderName(modelName) clients := make([]client.Client, 0) if provider == "gemini" { for i := 0; i < len(h.CliClients); i++ { if cli, ok := h.CliClients[i].(*client.GeminiClient); ok { clients = append(clients, cli) } } } else if provider == "gpt" { for i := 0; i < len(h.CliClients); i++ { if cli, ok := h.CliClients[i].(*client.CodexClient); ok { clients = append(clients, cli) } } } else if provider == "claude" { for i := 0; i < len(h.CliClients); i++ { if cli, ok := h.CliClients[i].(*client.ClaudeClient); ok { clients = append(clients, cli) } } } else if provider == "qwen" { for i := 0; i < len(h.CliClients); i++ { if cli, ok := h.CliClients[i].(*client.QwenClient); ok { clients = append(clients, cli) } } } if _, hasKey := h.LastUsedClientIndex[provider]; !hasKey { h.LastUsedClientIndex[provider] = 0 } if len(clients) == 0 { return nil, &client.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("no clients available")} } var cliClient client.Client // Lock the mutex to update the last used client index h.Mutex.Lock() startIndex := h.LastUsedClientIndex[provider] if (len(isGenerateContent) > 0 && isGenerateContent[0]) || len(isGenerateContent) == 0 { currentIndex := (startIndex + 1) % len(clients) h.LastUsedClientIndex[provider] = currentIndex } h.Mutex.Unlock() // Reorder the client to start from the last used index reorderedClients := make([]client.Client, 0) for i := 0; i < len(clients); i++ { cliClient = clients[(startIndex+1+i)%len(clients)] if cliClient.IsModelQuotaExceeded(modelName) { if provider == "gemini" { log.Debugf("Gemini Model %s is quota exceeded for account %s, project id: %s", modelName, cliClient.GetEmail(), cliClient.(*client.GeminiClient).GetProjectID()) } else if provider == "gpt" { log.Debugf("Codex Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) } else if provider == "claude" { log.Debugf("Claude Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) } else if provider == "qwen" { log.Debugf("Qwen Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) } cliClient = nil continue } reorderedClients = append(reorderedClients, cliClient) } if len(reorderedClients) == 0 { if provider == "claude" { // log.Debugf("Claude Model %s is quota exceeded for all accounts", modelName) return nil, &client.ErrorMessage{StatusCode: 429, Error: fmt.Errorf(`{"type":"error","error":{"type":"rate_limit_error","message":"This request would exceed your account's rate limit. Please try again later."}}`)} } return nil, &client.ErrorMessage{StatusCode: 429, Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName)} } locked := false for i := 0; i < len(reorderedClients); i++ { cliClient = reorderedClients[i] if cliClient.GetRequestMutex().TryLock() { locked = true break } } if !locked { cliClient = clients[0] cliClient.GetRequestMutex().Lock() } return cliClient, nil } // GetAlt extracts the 'alt' parameter from the request query string. // It checks both 'alt' and '$alt' parameters and returns the appropriate value. // // Parameters: // - c: The Gin context containing the HTTP request // // Returns: // - string: The alt parameter value, or empty string if it's "sse" func (h *APIHandlers) GetAlt(c *gin.Context) string { var alt string var hasAlt bool alt, hasAlt = c.GetQuery("alt") if !hasAlt { alt, _ = c.GetQuery("$alt") } if alt == "sse" { return "" } return alt } func (h *APIHandlers) GetContextWithCancel(c *gin.Context, ctx context.Context) (context.Context, APIHandlerCancelFunc) { newCtx, cancel := context.WithCancel(ctx) newCtx = context.WithValue(newCtx, "gin", c) return newCtx, func(params ...interface{}) { if h.Cfg.RequestLog { if len(params) == 1 { data := params[0] switch data.(type) { case []byte: c.Set("API_RESPONSE", data.([]byte)) case error: c.Set("API_RESPONSE", []byte(data.(error).Error())) case string: c.Set("API_RESPONSE", []byte(data.(string))) case bool: case nil: } } else { if _, hasKey := h.apiResponseData[c]; hasKey { c.Set("API_RESPONSE", h.apiResponseData[c]) delete(h.apiResponseData, c) } } } cancel() } } func (h *APIHandlers) AddAPIResponseData(c *gin.Context, data []byte) { if h.Cfg.RequestLog { if _, hasKey := h.apiResponseData[c]; !hasKey { h.apiResponseData[c] = make([]byte, 0) } h.apiResponseData[c] = append(h.apiResponseData[c], data...) } } type APIHandlerCancelFunc func(params ...interface{})