persona-community-5/pkg/routing/cooldown.go

package routing

import "time"

// Default cooldown periods.
const (
	// DefaultCooldownPeriod is the cooldown for rate limits and quota errors.
	// 1 hour gives provider APIs time to reset their quotas.
	DefaultCooldownPeriod = 1 * time.Hour

	// TransientCooldownPeriod is the cooldown for transient server errors (5xx).
	// These errors typically resolve quickly, so we use a shorter cooldown to
	// trigger fallback but allow quick recovery.
	TransientCooldownPeriod = 30 * time.Second
)

// CooldownTracker defines the interface for tracking provider cooldowns.
// CircuitBreaker (in-memory) implements this interface. Additional implementations
// (e.g., file-based persistence) can be added via CombinedCooldown.
//
// IMPORTANT: All cooldown tracking in the codebase MUST use this interface.
// Do NOT implement custom cooldown tracking elsewhere.
type CooldownTracker interface {
	// IsAvailable returns true if the provider is not in cooldown.
	IsAvailable(providerName string) bool

	// CooldownRemaining returns time until cooldown expires (0 if available).
	CooldownRemaining(providerName string) time.Duration

	// RecordFailure records a failure and potentially enters cooldown.
	// Returns true if the provider entered cooldown.
	// Note: Exempt providers (ExemptProviders) never enter cooldown.
	RecordFailure(providerName string, err error) bool

	// Reset removes a provider from cooldown immediately.
	Reset(providerName string)

	// ResetAll clears all cooldowns.
	ResetAll()
}

// CombinedCooldown wraps multiple CooldownTrackers, checking all of them.
// A provider is only available if ALL trackers report it as available.
//
// Typical usage: combine multiple CooldownTracker implementations
// (e.g., in-memory CircuitBreaker with custom persistence) for layered tracking.
type CombinedCooldown struct {
	trackers []CooldownTracker
}

// NewCombinedCooldown creates a tracker that combines multiple sources.
// Pass nil trackers to skip them (they will be filtered out).
func NewCombinedCooldown(trackers ...CooldownTracker) *CombinedCooldown {
	// Filter out nil trackers
	nonNil := make([]CooldownTracker, 0, len(trackers))
	for _, t := range trackers {
		if t != nil {
			nonNil = append(nonNil, t)
		}
	}
	return &CombinedCooldown{trackers: nonNil}
}

// IsAvailable returns true only if ALL trackers report the provider as available.
func (c *CombinedCooldown) IsAvailable(providerName string) bool {
	for _, t := range c.trackers {
		if !t.IsAvailable(providerName) {
			return false
		}
	}
	return true
}

// CooldownRemaining returns the maximum remaining cooldown from all trackers.
func (c *CombinedCooldown) CooldownRemaining(providerName string) time.Duration {
	var max time.Duration
	for _, t := range c.trackers {
		if remaining := t.CooldownRemaining(providerName); remaining > max {
			max = remaining
		}
	}
	return max
}

// RecordFailure records the failure to all trackers.
// Returns true if any tracker opened a cooldown.
func (c *CombinedCooldown) RecordFailure(providerName string, err error) bool {
	var anyOpened bool
	for _, t := range c.trackers {
		if t.RecordFailure(providerName, err) {
			anyOpened = true
		}
	}
	return anyOpened
}

// Reset removes the provider from cooldown in all trackers.
func (c *CombinedCooldown) Reset(providerName string) {
	for _, t := range c.trackers {
		t.Reset(providerName)
	}
}

// ResetAll clears all cooldowns in all trackers.
func (c *CombinedCooldown) ResetAll() {
	for _, t := range c.trackers {
		t.ResetAll()
	}
}

// Compile-time interface check
var _ CooldownTracker = (*CombinedCooldown)(nil)

// CooldownConfig holds the configuration for building a CooldownTracker.
// This struct is used by mediagen.Manager and textgen.Manager to configure
// their cooldown behavior.
//
// Usage: Provide a CircuitBreaker for in-memory tracking, or leave nil
// to auto-create one with the specified CooldownPeriod.
type CooldownConfig struct {
	// CircuitBreaker provides in-memory cooldown tracking.
	// Good for long-running services where state is maintained in memory.
	// If nil, a default CircuitBreaker is created with CooldownPeriod.
	CircuitBreaker *CircuitBreaker

	// CooldownPeriod is the default cooldown duration for rate-limited providers.
	// Only used when creating a default CircuitBreaker (when CircuitBreaker is nil).
	// Defaults to DefaultCooldownPeriod (1 hour) if zero.
	CooldownPeriod time.Duration
}

// BuildCooldownTracker creates a CooldownTracker from the configuration.
// This is the standard way to construct cooldown trackers in the codebase.
//
// Logic:
//  1. If CircuitBreaker is provided, returns it directly
//  2. If CircuitBreaker is nil, creates a default one with the specified
//     cooldown period (or DefaultCooldownPeriod if zero)
func BuildCooldownTracker(config CooldownConfig) CooldownTracker {
	if config.CircuitBreaker != nil {
		return config.CircuitBreaker
	}

	// Create default circuit breaker
	cooldown := config.CooldownPeriod
	if cooldown == 0 {
		cooldown = DefaultCooldownPeriod
	}
	return NewCircuitBreaker(cooldown)
}