From aff30763cda30e8d9413865d36c7e6340add3b8c Mon Sep 17 00:00:00 2001 From: Leandro Barbosa Date: Thu, 19 Feb 2026 13:32:04 -0300 Subject: [PATCH] feat: wire auth rotation into config and provider factory Add api_keys array to ProviderConfig for multi-key support. ResolveAPIKeys() returns api_keys if set, otherwise wraps api_key. Factory detects multiple keys and creates AuthRotatingProvider. Backward compatible: single api_key works unchanged. --- pkg/config/config.go | 24 +++++++++++++++++++----- pkg/providers/factory.go | 23 ++++++++++++++++++++--- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index d814edeb8..25a221200 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -278,11 +278,25 @@ type ProvidersConfig struct { } type ProviderConfig struct { - APIKey string `json:"api_key" env:"PICOCLAW_PROVIDERS_{{.Name}}_API_KEY"` - APIBase string `json:"api_base" env:"PICOCLAW_PROVIDERS_{{.Name}}_API_BASE"` - Proxy string `json:"proxy,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_PROXY"` - AuthMethod string `json:"auth_method,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_AUTH_METHOD"` - ConnectMode string `json:"connect_mode,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_CONNECT_MODE"` //only for Github Copilot, `stdio` or `grpc` + APIKey string `json:"api_key" env:"PICOCLAW_PROVIDERS_{{.Name}}_API_KEY"` + APIKeys []string `json:"api_keys,omitempty"` // multiple keys for auth rotation (takes precedence over api_key) + APIBase string `json:"api_base" env:"PICOCLAW_PROVIDERS_{{.Name}}_API_BASE"` + Proxy string `json:"proxy,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_PROXY"` + AuthMethod string `json:"auth_method,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_AUTH_METHOD"` + ConnectMode string `json:"connect_mode,omitempty" env:"PICOCLAW_PROVIDERS_{{.Name}}_CONNECT_MODE"` //only for Github Copilot, `stdio` or `grpc` +} + +// ResolveAPIKeys returns the effective list of API keys for this provider. +// If APIKeys is set, returns it. Otherwise wraps APIKey as a single-element slice. +// Returns nil if no keys are configured. +func (pc *ProviderConfig) ResolveAPIKeys() []string { + if len(pc.APIKeys) > 0 { + return pc.APIKeys + } + if pc.APIKey != "" { + return []string{pc.APIKey} + } + return nil } type OpenAIProviderConfig struct { diff --git a/pkg/providers/factory.go b/pkg/providers/factory.go index e39cfe32b..c4a9f61f5 100644 --- a/pkg/providers/factory.go +++ b/pkg/providers/factory.go @@ -26,7 +26,9 @@ const ( type providerSelection struct { providerType providerType + providerName string // resolved provider name (e.g. "openrouter", "anthropic") apiKey string + apiKeys []string // multiple keys for auth rotation (nil = single key) apiBase string proxy string model string @@ -120,7 +122,9 @@ func resolveProviderSelection(cfg *config.Config) (providerSelection, error) { } } case "openrouter": - if cfg.Providers.OpenRouter.APIKey != "" { + if cfg.Providers.OpenRouter.APIKey != "" || len(cfg.Providers.OpenRouter.APIKeys) > 0 { + sel.providerName = "openrouter" + sel.apiKeys = cfg.Providers.OpenRouter.ResolveAPIKeys() sel.apiKey = cfg.Providers.OpenRouter.APIKey sel.proxy = cfg.Providers.OpenRouter.Proxy if cfg.Providers.OpenRouter.APIBase != "" { @@ -227,6 +231,8 @@ func resolveProviderSelection(cfg *config.Config) (providerSelection, error) { strings.HasPrefix(model, "meta-llama/") || strings.HasPrefix(model, "deepseek/") || strings.HasPrefix(model, "google/"): + sel.providerName = "openrouter" + sel.apiKeys = cfg.Providers.OpenRouter.ResolveAPIKeys() sel.apiKey = cfg.Providers.OpenRouter.APIKey sel.proxy = cfg.Providers.OpenRouter.Proxy if cfg.Providers.OpenRouter.APIBase != "" { @@ -307,7 +313,9 @@ func resolveProviderSelection(cfg *config.Config) (providerSelection, error) { sel.apiBase = cfg.Providers.VLLM.APIBase sel.proxy = cfg.Providers.VLLM.Proxy default: - if cfg.Providers.OpenRouter.APIKey != "" { + if cfg.Providers.OpenRouter.APIKey != "" || len(cfg.Providers.OpenRouter.APIKeys) > 0 { + sel.providerName = "openrouter" + sel.apiKeys = cfg.Providers.OpenRouter.ResolveAPIKeys() sel.apiKey = cfg.Providers.OpenRouter.APIKey sel.proxy = cfg.Providers.OpenRouter.Proxy if cfg.Providers.OpenRouter.APIBase != "" { @@ -322,7 +330,7 @@ func resolveProviderSelection(cfg *config.Config) (providerSelection, error) { } if sel.providerType == providerTypeHTTPCompat { - if sel.apiKey == "" && !strings.HasPrefix(model, "bedrock/") { + if sel.apiKey == "" && len(sel.apiKeys) == 0 && !strings.HasPrefix(model, "bedrock/") { return providerSelection{}, fmt.Errorf("no API key configured for provider (model: %s)", model) } if sel.apiBase == "" { @@ -355,6 +363,15 @@ func CreateProvider(cfg *config.Config) (LLMProvider, error) { case providerTypeGitHubCopilot: return NewGitHubCopilotProvider(sel.apiBase, sel.connectMode, sel.model) default: + // Auth rotation: wrap with AuthRotatingProvider if multiple keys configured. + if len(sel.apiKeys) > 1 { + profiles := BuildAuthProfiles(sel.providerName, sel.apiKeys) + cooldown := NewCooldownTracker() + factory := func(apiKey string) LLMProvider { + return NewHTTPProvider(apiKey, sel.apiBase, sel.proxy) + } + return NewAuthRotatingProvider(profiles, cooldown, factory), nil + } return NewHTTPProvider(sel.apiKey, sel.apiBase, sel.proxy), nil } }