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.
This commit is contained in:
Luis Pater
2025-09-19 01:53:38 +08:00
parent 9ec8478b41
commit df66046b14
15 changed files with 183 additions and 21 deletions

View File

@@ -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
}

View File

@@ -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.

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}