opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Revert "revert(providers): revert cooldown tracker per-model changes"

Browse source 
This reverts commit 673280ce2f.
This commit is contained in:
muava12 2026-02-24 10:38:48 +08:00
parent 673280ce2f
commit 92c00a3a81
3 changed files with 32 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
pkg/providers/cooldown.go
View file

@ -37,14 +37,14 @@ func NewCooldownTracker() *CooldownTracker {
} }
} }
// MarkFailure records a failure for a provider and sets appropriate cooldown. // MarkFailure records a failure for a specific model key and sets appropriate cooldown.
// Resets error counts if last failure was more than failureWindow ago. // Resets error counts if last failure was more than failureWindow ago.
func (ct *CooldownTracker) MarkFailure(provider string, reason FailoverReason) { func (ct *CooldownTracker) MarkFailure(key string, reason FailoverReason) {
ct.mu.Lock() ct.mu.Lock()
defer ct.mu.Unlock() defer ct.mu.Unlock()
now := ct.nowFunc() now := ct.nowFunc()
entry := ct.getOrCreate(provider) entry := ct.getOrCreate(key)
// 24h failure window reset: if no failure in failureWindow, reset counters. // 24h failure window reset: if no failure in failureWindow, reset counters.
if !entry.LastFailure.IsZero() && now.Sub(entry.LastFailure) > ct.failureWindow { if !entry.LastFailure.IsZero() && now.Sub(entry.LastFailure) > ct.failureWindow {
@ -65,12 +65,12 @@ func (ct *CooldownTracker) MarkFailure(provider string, reason FailoverReason) {
} }
} }
// MarkSuccess resets all counters and cooldowns for a provider. // MarkSuccess resets all counters and cooldowns for a specific model key.
func (ct *CooldownTracker) MarkSuccess(provider string) { func (ct *CooldownTracker) MarkSuccess(key string) {
ct.mu.Lock() ct.mu.Lock()
defer ct.mu.Unlock() defer ct.mu.Unlock()
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
return return
} }
@ -82,12 +82,12 @@ func (ct *CooldownTracker) MarkSuccess(provider string) {
entry.DisabledReason = "" entry.DisabledReason = ""
} }
// IsAvailable returns true if the provider is not in cooldown or disabled. // IsAvailable returns true if the specific model key is not in cooldown or disabled.
func (ct *CooldownTracker) IsAvailable(provider string) bool { func (ct *CooldownTracker) IsAvailable(key string) bool {
ct.mu.RLock() ct.mu.RLock()
defer ct.mu.RUnlock() defer ct.mu.RUnlock()
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
return true return true
} }
@ -107,13 +107,13 @@ func (ct *CooldownTracker) IsAvailable(provider string) bool {
return true return true
} }
// CooldownRemaining returns how long until the provider becomes available. // CooldownRemaining returns how long until the specific model key becomes available.
// Returns 0 if already available. // Returns 0 if already available.
func (ct *CooldownTracker) CooldownRemaining(provider string) time.Duration { func (ct *CooldownTracker) CooldownRemaining(key string) time.Duration {
ct.mu.RLock() ct.mu.RLock()
defer ct.mu.RUnlock() defer ct.mu.RUnlock()
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
return 0 return 0
} }
@ -138,37 +138,37 @@ func (ct *CooldownTracker) CooldownRemaining(provider string) time.Duration {
return remaining return remaining
} }
// ErrorCount returns the current error count for a provider. // ErrorCount returns the current error count for a specific model key.
func (ct *CooldownTracker) ErrorCount(provider string) int { func (ct *CooldownTracker) ErrorCount(key string) int {
ct.mu.RLock() ct.mu.RLock()
defer ct.mu.RUnlock() defer ct.mu.RUnlock()
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
return 0 return 0
} }
return entry.ErrorCount return entry.ErrorCount
} }
// FailureCount returns the failure count for a specific reason. // FailureCount returns the failure count for a specific reason and model key.
func (ct *CooldownTracker) FailureCount(provider string, reason FailoverReason) int { func (ct *CooldownTracker) FailureCount(key string, reason FailoverReason) int {
ct.mu.RLock() ct.mu.RLock()
defer ct.mu.RUnlock() defer ct.mu.RUnlock()
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
return 0 return 0
} }
return entry.FailureCounts[reason] return entry.FailureCounts[reason]
} }
func (ct *CooldownTracker) getOrCreate(provider string) *cooldownEntry { func (ct *CooldownTracker) getOrCreate(key string) *cooldownEntry {
entry := ct.entries[provider] entry := ct.entries[key]
if entry == nil { if entry == nil {
entry = &cooldownEntry{ entry = &cooldownEntry{
FailureCounts: make(map[FailoverReason]int), FailureCounts: make(map[FailoverReason]int),
} }
ct.entries[provider] = entry ct.entries[key] = entry
} }
return entry return entry
} }

10
pkg/providers/fallback.go
View file

@ -122,9 +122,11 @@ func (fc *FallbackChain) Execute(
return nil, context.Canceled return nil, context.Canceled
} }
key := ModelKey(candidate.Provider, candidate.Model)
// Check cooldown. // Check cooldown.
if !fc.cooldown.IsAvailable(candidate.Provider) { if !fc.cooldown.IsAvailable(key) {
remaining := fc.cooldown.CooldownRemaining(candidate.Provider) remaining := fc.cooldown.CooldownRemaining(key)
result.Attempts = append(result.Attempts, FallbackAttempt{ result.Attempts = append(result.Attempts, FallbackAttempt{
Provider: candidate.Provider, Provider: candidate.Provider,
Model: candidate.Model, Model: candidate.Model,
@ -145,7 +147,7 @@ func (fc *FallbackChain) Execute(
if err == nil { if err == nil {
// Success. // Success.
fc.cooldown.MarkSuccess(candidate.Provider) fc.cooldown.MarkSuccess(key)
result.Response = resp result.Response = resp
result.Provider = candidate.Provider result.Provider = candidate.Provider
result.Model = candidate.Model result.Model = candidate.Model
@ -191,7 +193,7 @@ func (fc *FallbackChain) Execute(
} }
// Retriable error: mark failure and continue to next candidate. // Retriable error: mark failure and continue to next candidate.
fc.cooldown.MarkFailure(candidate.Provider, failErr.Reason) fc.cooldown.MarkFailure(key, failErr.Reason)
result.Attempts = append(result.Attempts, FallbackAttempt{ result.Attempts = append(result.Attempts, FallbackAttempt{
Provider: candidate.Provider, Provider: candidate.Provider,
Model: candidate.Model, Model: candidate.Model,

10
pkg/providers/fallback_test.go
View file

@ -164,7 +164,7 @@ func TestFallback_CooldownSkip(t *testing.T) {
fc := NewFallbackChain(ct) fc := NewFallbackChain(ct)
// Put openai in cooldown // Put openai in cooldown
ct.MarkFailure("openai", FailoverRateLimit) ct.MarkFailure("openai/gpt-4", FailoverRateLimit)
candidates := []FallbackCandidate{ candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"), makeCandidate("openai", "gpt-4"),
@ -202,8 +202,8 @@ func TestFallback_AllInCooldown(t *testing.T) {
fc := NewFallbackChain(ct) fc := NewFallbackChain(ct)
// Put all providers in cooldown // Put all providers in cooldown
ct.MarkFailure("openai", FailoverRateLimit) ct.MarkFailure("openai/gpt-4", FailoverRateLimit)
ct.MarkFailure("anthropic", FailoverBilling) ct.MarkFailure("anthropic/claude", FailoverBilling)
candidates := []FallbackCandidate{ candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"), makeCandidate("openai", "gpt-4"),
@ -284,7 +284,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
attempt++ attempt++
if attempt == 1 { if attempt == 1 {
ct.MarkFailure("openai", FailoverRateLimit) // simulate failure tracked elsewhere ct.MarkFailure("openai/gpt-4", FailoverRateLimit) // simulate failure tracked elsewhere
} }
return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
} }
@ -293,7 +293,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
if !ct.IsAvailable("openai") { if !ct.IsAvailable("openai/gpt-4") {
t.Error("success should reset cooldown") t.Error("success should reset cooldown")
} }
} }