From b00a8a54fd94c2ac4b6fb648847fdfef9ebee98e Mon Sep 17 00:00:00 2001 From: Leandro Barbosa Date: Thu, 19 Feb 2026 13:31:57 -0300 Subject: [PATCH] feat: add auth profile rotation with per-key cooldown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AuthRotator manages round-robin selection across multiple API keys per provider. Each key tracks its own cooldown state via CooldownTracker (2-track: transient 1min→1hr exponential, billing 5h→24h). AuthRotatingProvider wraps multiple LLM providers and delegates to the best available key on each request. On retriable failures, the failing key is put in cooldown and subsequent requests use the next available key. --- pkg/providers/auth_rotation.go | 185 +++++++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 pkg/providers/auth_rotation.go diff --git a/pkg/providers/auth_rotation.go b/pkg/providers/auth_rotation.go new file mode 100644 index 000000000..eaef1d631 --- /dev/null +++ b/pkg/providers/auth_rotation.go @@ -0,0 +1,185 @@ +package providers + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// AuthProfile represents a single API key with rotation metadata. +type AuthProfile struct { + ID string // unique identifier (e.g. "openrouter:0") + APIKey string +} + +// AuthRotator manages round-robin selection across multiple API keys, +// with per-key cooldown tracking via CooldownTracker. +type AuthRotator struct { + profiles []AuthProfile + cooldown *CooldownTracker + mu sync.Mutex + lastUsed map[string]time.Time +} + +// NewAuthRotator creates a rotator for the given profiles. +// Uses the provided CooldownTracker for per-key cooldown state. +func NewAuthRotator(profiles []AuthProfile, cooldown *CooldownTracker) *AuthRotator { + lastUsed := make(map[string]time.Time, len(profiles)) + for _, p := range profiles { + lastUsed[p.ID] = time.Time{} // never used + } + return &AuthRotator{ + profiles: profiles, + cooldown: cooldown, + lastUsed: lastUsed, + } +} + +// NextAvailable returns the best available profile using round-robin +// (oldest lastUsed first), skipping profiles in cooldown. +// Returns nil if all profiles are in cooldown. +func (r *AuthRotator) NextAvailable() *AuthProfile { + r.mu.Lock() + defer r.mu.Unlock() + + var best *AuthProfile + var bestTime time.Time + first := true + + for i := range r.profiles { + p := &r.profiles[i] + if !r.cooldown.IsAvailable(p.ID) { + continue + } + lu := r.lastUsed[p.ID] + if first || lu.Before(bestTime) { + best = p + bestTime = lu + first = false + } + } + + if best != nil { + r.lastUsed[best.ID] = time.Now() + } + return best +} + +// MarkFailure records a failure for a specific profile. +func (r *AuthRotator) MarkFailure(profileID string, reason FailoverReason) { + r.cooldown.MarkFailure(profileID, reason) + logger.WarnCF("auth_rotation", "Profile marked as failed", map[string]interface{}{ + "profile_id": profileID, + "reason": string(reason), + "remaining": r.cooldown.CooldownRemaining(profileID).Round(time.Second).String(), + }) +} + +// MarkSuccess resets counters for a specific profile. +func (r *AuthRotator) MarkSuccess(profileID string) { + r.cooldown.MarkSuccess(profileID) +} + +// AvailableCount returns the number of profiles not in cooldown. +func (r *AuthRotator) AvailableCount() int { + count := 0 + for _, p := range r.profiles { + if r.cooldown.IsAvailable(p.ID) { + count++ + } + } + return count +} + +// ProfileCount returns the total number of profiles. +func (r *AuthRotator) ProfileCount() int { + return len(r.profiles) +} + +// AuthRotatingProvider wraps multiple LLM providers (one per API key) +// and rotates between them using AuthRotator. +type AuthRotatingProvider struct { + providers map[string]LLMProvider // profileID -> provider + rotator *AuthRotator + model string // default model from first provider +} + +// NewAuthRotatingProvider creates a rotating provider. +// factory is called once per profile to create the underlying provider. +func NewAuthRotatingProvider( + profiles []AuthProfile, + cooldown *CooldownTracker, + factory func(apiKey string) LLMProvider, +) *AuthRotatingProvider { + providerMap := make(map[string]LLMProvider, len(profiles)) + var defaultModel string + for _, p := range profiles { + prov := factory(p.APIKey) + providerMap[p.ID] = prov + if defaultModel == "" { + defaultModel = prov.GetDefaultModel() + } + } + + rotator := NewAuthRotator(profiles, cooldown) + + logger.InfoCF("auth_rotation", "Auth rotation initialized", map[string]interface{}{ + "profiles": len(profiles), + }) + + return &AuthRotatingProvider{ + providers: providerMap, + rotator: rotator, + model: defaultModel, + } +} + +// Chat selects the best available profile and delegates to its provider. +// On failure, marks the profile and returns the error (FallbackChain handles retry). +func (p *AuthRotatingProvider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + opts map[string]interface{}, +) (*LLMResponse, error) { + profile := p.rotator.NextAvailable() + if profile == nil { + return nil, fmt.Errorf("all auth profiles in cooldown (%d total)", p.rotator.ProfileCount()) + } + + provider := p.providers[profile.ID] + resp, err := provider.Chat(ctx, messages, tools, model, opts) + + if err != nil { + // Classify and record failure against this specific profile. + if failErr := ClassifyError(err, profile.ID, model); failErr != nil && failErr.IsRetriable() { + p.rotator.MarkFailure(profile.ID, failErr.Reason) + } + return nil, err + } + + p.rotator.MarkSuccess(profile.ID) + return resp, nil +} + +// GetDefaultModel returns the default model from the underlying providers. +func (p *AuthRotatingProvider) GetDefaultModel() string { + return p.model +} + +// BuildAuthProfiles creates AuthProfile entries from a list of API keys. +// Profile IDs follow the pattern "provider:N" (e.g. "openrouter:0"). +func BuildAuthProfiles(providerName string, apiKeys []string) []AuthProfile { + profiles := make([]AuthProfile, len(apiKeys)) + for i, key := range apiKeys { + profiles[i] = AuthProfile{ + ID: fmt.Sprintf("%s:%d", providerName, i), + APIKey: key, + } + } + return profiles +}