From df66046b1466d63c29876cd7250fb108ad56bef3 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 19 Sep 2025 01:53:38 +0800 Subject: [PATCH] feat: add client availability tracking and error handling improvements - Introduced `IsAvailable` and `SetUnavailable` methods to clients for availability tracking. - Integrated availability checks in client selection logic to skip unavailable clients. - Enhanced error handling by marking clients unavailable on specific error codes (e.g., 401, 402). - Removed redundant quota verification logs in client reordering logic. --- internal/api/handlers/claude/code_handlers.go | 4 ++ .../handlers/gemini/gemini-cli_handlers.go | 16 ++++++++ .../api/handlers/gemini/gemini_handlers.go | 16 ++++++++ internal/api/handlers/handlers.go | 22 +--------- .../api/handlers/openai/openai_handlers.go | 40 +++++++++++++++++++ .../openai/openai_responses_handlers.go | 16 ++++++++ internal/client/claude_client.go | 12 ++++++ internal/client/client.go | 3 ++ internal/client/codex_client.go | 12 ++++++ internal/client/gemini-cli_client.go | 12 ++++++ internal/client/gemini-web_client.go | 11 +++++ internal/client/gemini_client.go | 11 +++++ .../client/openai-compatibility_client.go | 11 +++++ internal/client/qwen_client.go | 11 +++++ internal/interfaces/client.go | 7 ++++ 15 files changed, 183 insertions(+), 21 deletions(-) diff --git a/internal/api/handlers/claude/code_handlers.go b/internal/api/handlers/claude/code_handlers.go index 0a6d93d0..6d3f7928 100644 --- a/internal/api/handlers/claude/code_handlers.go +++ b/internal/api/handlers/claude/code_handlers.go @@ -205,9 +205,13 @@ outLoop: err := cliClient.RefreshTokens(cliCtx) if err != nil { log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() } retryCount++ continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(errInfo.StatusCode) diff --git a/internal/api/handlers/gemini/gemini-cli_handlers.go b/internal/api/handlers/gemini/gemini-cli_handlers.go index 82d52ddc..f319dfcc 100644 --- a/internal/api/handlers/gemini/gemini-cli_handlers.go +++ b/internal/api/handlers/gemini/gemini-cli_handlers.go @@ -221,6 +221,18 @@ outLoop: log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue outLoop + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -293,9 +305,13 @@ func (h *GeminiCLIAPIHandler) handleInternalGenerateContent(c *gin.Context, rawJ errRefreshTokens := cliClient.RefreshTokens(cliCtx) if errRefreshTokens != nil { log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() } retryCount++ continue + case 402: + cliClient.SetUnavailable() + continue default: // Forward other errors directly to the client c.Status(err.StatusCode) diff --git a/internal/api/handlers/gemini/gemini_handlers.go b/internal/api/handlers/gemini/gemini_handlers.go index d162cc6d..71b6e2a3 100644 --- a/internal/api/handlers/gemini/gemini_handlers.go +++ b/internal/api/handlers/gemini/gemini_handlers.go @@ -276,6 +276,18 @@ outLoop: log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue outLoop + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -406,9 +418,13 @@ func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName strin errRefreshTokens := cliClient.RefreshTokens(cliCtx) if errRefreshTokens != nil { log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() } retryCount++ continue + case 402: + cliClient.SetUnavailable() + continue default: // Forward other errors directly to the client c.Status(err.StatusCode) diff --git a/internal/api/handlers/handlers.go b/internal/api/handlers/handlers.go index fbe52027..ffc181e2 100644 --- a/internal/api/handlers/handlers.go +++ b/internal/api/handlers/handlers.go @@ -8,11 +8,9 @@ import ( "sync" "github.com/gin-gonic/gin" - "github.com/luispater/CLIProxyAPI/v5/internal/client" "github.com/luispater/CLIProxyAPI/v5/internal/config" "github.com/luispater/CLIProxyAPI/v5/internal/interfaces" "github.com/luispater/CLIProxyAPI/v5/internal/util" - log "github.com/sirupsen/logrus" "golang.org/x/net/context" ) @@ -97,7 +95,7 @@ func (h *BaseAPIHandler) UpdateClients(clients []interfaces.Client, cfg *config. func (h *BaseAPIHandler) GetClient(modelName string, isGenerateContent ...bool) (interfaces.Client, *interfaces.ErrorMessage) { clients := make([]interfaces.Client, 0) for i := 0; i < len(h.CliClients); i++ { - if h.CliClients[i].CanProvideModel(modelName) { + if h.CliClients[i].CanProvideModel(modelName) && h.CliClients[i].IsAvailable() && !h.CliClients[i].IsModelQuotaExceeded(modelName) { clients = append(clients, h.CliClients[i]) } } @@ -126,24 +124,6 @@ func (h *BaseAPIHandler) GetClient(modelName string, isGenerateContent ...bool) reorderedClients := make([]interfaces.Client, 0) for i := 0; i < len(clients); i++ { cliClient = clients[(startIndex+1+i)%len(clients)] - if cliClient.IsModelQuotaExceeded(modelName) { - if cliClient.Provider() == "gemini-cli" { - log.Debugf("Gemini Model %s is quota exceeded for account %s, project id: %s", modelName, cliClient.GetEmail(), cliClient.(*client.GeminiCLIClient).GetProjectID()) - } else if cliClient.Provider() == "gemini" { - log.Debugf("Gemini Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) - } else if cliClient.Provider() == "codex" { - log.Debugf("Codex Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) - } else if cliClient.Provider() == "claude" { - log.Debugf("Claude Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) - } else if cliClient.Provider() == "qwen" { - log.Debugf("Qwen Model %s is quota exceeded for account %s", modelName, cliClient.GetEmail()) - } else if cliClient.Type() == "openai-compatibility" { - log.Debugf("OpenAI Compatibility Model %s is quota exceeded for provider %s", modelName, cliClient.Provider()) - } - cliClient = nil - continue - - } reorderedClients = append(reorderedClients, cliClient) } diff --git a/internal/api/handlers/openai/openai_handlers.go b/internal/api/handlers/openai/openai_handlers.go index f8b43279..90ec76e8 100644 --- a/internal/api/handlers/openai/openai_handlers.go +++ b/internal/api/handlers/openai/openai_handlers.go @@ -442,9 +442,13 @@ func (h *OpenAIAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON [] errRefreshTokens := cliClient.RefreshTokens(cliCtx) if errRefreshTokens != nil { log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() } retryCount++ continue + case 402: + cliClient.SetUnavailable() + continue default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -557,6 +561,18 @@ outLoop: log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue outLoop + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -632,6 +648,18 @@ func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context, log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue + case 402: + cliClient.SetUnavailable() + continue default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -755,6 +783,18 @@ outLoop: log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue outLoop + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(err.StatusCode) diff --git a/internal/api/handlers/openai/openai_responses_handlers.go b/internal/api/handlers/openai/openai_responses_handlers.go index 224941ed..f49b971d 100644 --- a/internal/api/handlers/openai/openai_responses_handlers.go +++ b/internal/api/handlers/openai/openai_responses_handlers.go @@ -146,9 +146,13 @@ func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, r errRefreshTokens := cliClient.RefreshTokens(cliCtx) if errRefreshTokens != nil { log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() } retryCount++ continue + case 402: + cliClient.SetUnavailable() + continue default: // Forward other errors directly to the client c.Status(err.StatusCode) @@ -260,6 +264,18 @@ outLoop: log.Debugf("http status code %d, switch client", err.StatusCode) retryCount++ continue outLoop + case 401: + log.Debugf("unauthorized request, try to refresh token, %s", util.HideAPIKey(cliClient.GetEmail())) + errRefreshTokens := cliClient.RefreshTokens(cliCtx) + if errRefreshTokens != nil { + log.Debugf("refresh token failed, switch client, %s", util.HideAPIKey(cliClient.GetEmail())) + cliClient.SetUnavailable() + } + retryCount++ + continue outLoop + case 402: + cliClient.SetUnavailable() + continue outLoop default: // Forward other errors directly to the client c.Status(err.StatusCode) diff --git a/internal/client/claude_client.go b/internal/client/claude_client.go index ba990feb..540af78b 100644 --- a/internal/client/claude_client.go +++ b/internal/client/claude_client.go @@ -67,6 +67,7 @@ func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeC cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), tokenStorage: ts, + isAvailable: true, }, claudeAuth: claude.NewClaudeAuth(cfg), apiKeyIndex: -1, @@ -102,6 +103,7 @@ func NewClaudeClientWithKey(cfg *config.Config, apiKeyIndex int) *ClaudeClient { cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), tokenStorage: &empty.EmptyStorage{}, + isAvailable: true, }, claudeAuth: claude.NewClaudeAuth(cfg), apiKeyIndex: apiKeyIndex, @@ -581,3 +583,13 @@ func (c *ClaudeClient) IsModelQuotaExceeded(model string) bool { func (c *ClaudeClient) GetRequestMutex() *sync.Mutex { return nil } + +// IsAvailable returns true if the client is available for use. +func (c *ClaudeClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *ClaudeClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/client.go b/internal/client/client.go index 0e54b6a5..dd4a69d8 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -41,6 +41,9 @@ type ClientBase struct { // modelRegistry is the global model registry for tracking model availability. modelRegistry *registry.ModelRegistry + + // unavailable tracks whether the client is unavailable + isAvailable bool } // GetRequestMutex returns the mutex used to synchronize requests for this client. diff --git a/internal/client/codex_client.go b/internal/client/codex_client.go index 59327f77..5892a57d 100644 --- a/internal/client/codex_client.go +++ b/internal/client/codex_client.go @@ -65,6 +65,7 @@ func NewCodexClient(cfg *config.Config, ts *codex.CodexTokenStorage) (*CodexClie cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), tokenStorage: ts, + isAvailable: true, }, codexAuth: codex.NewCodexAuth(cfg), apiKeyIndex: -1, @@ -100,6 +101,7 @@ func NewCodexClientWithKey(cfg *config.Config, apiKeyIndex int) *CodexClient { cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), tokenStorage: &empty.EmptyStorage{}, + isAvailable: true, }, codexAuth: codex.NewCodexAuth(cfg), apiKeyIndex: apiKeyIndex, @@ -557,3 +559,13 @@ func (c *CodexClient) IsModelQuotaExceeded(model string) bool { func (c *CodexClient) GetRequestMutex() *sync.Mutex { return nil } + +// IsAvailable returns true if the client is available for use. +func (c *CodexClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *CodexClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/gemini-cli_client.go b/internal/client/gemini-cli_client.go index 474870fd..e058a00e 100644 --- a/internal/client/gemini-cli_client.go +++ b/internal/client/gemini-cli_client.go @@ -69,6 +69,7 @@ func NewGeminiCLIClient(httpClient *http.Client, ts *geminiAuth.GeminiTokenStora cfg: cfg, tokenStorage: ts, modelQuotaExceeded: make(map[string]*time.Time), + isAvailable: true, }, } @@ -871,7 +872,18 @@ func (c *GeminiCLIClient) GetRequestMutex() *sync.Mutex { return nil } +// RefreshTokens is not applicable for Gemini CLI clients as they use API keys. func (c *GeminiCLIClient) RefreshTokens(ctx context.Context) error { // API keys don't need refreshing return nil } + +// IsAvailable returns true if the client is available for use. +func (c *GeminiCLIClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *GeminiCLIClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/gemini-web_client.go b/internal/client/gemini-web_client.go index ae8796b7..308473a9 100644 --- a/internal/client/gemini-web_client.go +++ b/internal/client/gemini-web_client.go @@ -96,6 +96,7 @@ func NewGeminiWebClient(cfg *config.Config, ts *gemini.GeminiWebTokenStorage, to cfg: cfg, tokenStorage: ts, modelQuotaExceeded: make(map[string]*time.Time), + isAvailable: true, }, tokenFilePath: tokenFilePath, convStore: make(map[string][]string), @@ -1072,3 +1073,13 @@ func (c *GeminiWebClient) storeConversationJSON(model string, history []geminiWe c.convMutex.Unlock() _ = geminiWeb.SaveConvData(geminiWeb.ConvDataPath(c.tokenFilePath), items, index) } + +// IsAvailable returns true if the client is available for use. +func (c *GeminiWebClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *GeminiWebClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/gemini_client.go b/internal/client/gemini_client.go index 23d49de0..10e43d2a 100644 --- a/internal/client/gemini_client.go +++ b/internal/client/gemini_client.go @@ -54,6 +54,7 @@ func NewGeminiClient(httpClient *http.Client, cfg *config.Config, glAPIKey strin httpClient: httpClient, cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), + isAvailable: true, }, glAPIKey: glAPIKey, } @@ -445,3 +446,13 @@ func (c *GeminiClient) RefreshTokens(ctx context.Context) error { // API keys don't need refreshing return nil } + +// IsAvailable returns true if the client is available for use. +func (c *GeminiClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *GeminiClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/openai-compatibility_client.go b/internal/client/openai-compatibility_client.go index d4c86477..990bc610 100644 --- a/internal/client/openai-compatibility_client.go +++ b/internal/client/openai-compatibility_client.go @@ -68,6 +68,7 @@ func NewOpenAICompatibilityClient(cfg *config.Config, compatConfig *config.OpenA httpClient: httpClient, cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), + isAvailable: true, }, compatConfig: compatConfig, currentAPIKeyIndex: apiKeyIndex, @@ -425,3 +426,13 @@ func (c *OpenAICompatibilityClient) RefreshTokens(ctx context.Context) error { func (c *OpenAICompatibilityClient) GetRequestMutex() *sync.Mutex { return nil } + +// IsAvailable returns true if the client is available for use. +func (c *OpenAICompatibilityClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *OpenAICompatibilityClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/client/qwen_client.go b/internal/client/qwen_client.go index ae3473f6..59b2670b 100644 --- a/internal/client/qwen_client.go +++ b/internal/client/qwen_client.go @@ -61,6 +61,7 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient { cfg: cfg, modelQuotaExceeded: make(map[string]*time.Time), tokenStorage: ts, + isAvailable: true, }, qwenAuth: qwen.NewQwenAuth(cfg), } @@ -447,3 +448,13 @@ func (c *QwenClient) IsModelQuotaExceeded(model string) bool { func (c *QwenClient) GetRequestMutex() *sync.Mutex { return nil } + +// IsAvailable returns true if the client is available for use. +func (c *QwenClient) IsAvailable() bool { + return c.isAvailable +} + +// SetUnavailable sets the client to unavailable. +func (c *QwenClient) SetUnavailable() { + c.isAvailable = false +} diff --git a/internal/interfaces/client.go b/internal/interfaces/client.go index a9bd5987..ae7a641f 100644 --- a/internal/interfaces/client.go +++ b/internal/interfaces/client.go @@ -52,5 +52,12 @@ type Client interface { // Provider returns the name of the AI service provider (e.g., "gemini", "claude"). Provider() string + // RefreshTokens refreshes the access tokens if needed RefreshTokens(ctx context.Context) error + + // IsAvailable returns true if the client is available for use. + IsAvailable() bool + + // SetUnavailable sets the client to unavailable. + SetUnavailable() }