feat: add circuit breaker for upstream provider overload protection

bridge.go

@@ -4,14 +4,15 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strings"
 	"sync"
 	"sync/atomic"
 
 	"cdr.dev/slog"
 	"github.com/coder/aibridge/mcp"
-	"go.opentelemetry.io/otel/trace"
-
 	"github.com/hashicorp/go-multierror"
+	"github.com/sony/gobreaker/v2"
+	"go.opentelemetry.io/otel/trace"
 )
 
 // RequestBridge is an [http.Handler] which is capable of masquerading as AI providers' APIs;
@@ -30,6 +31,10 @@ type RequestBridge struct {
 
 	mcpProxy mcp.ServerProxier
 
+	// circuitBreakers manages circuit breakers for upstream providers.
+	// When enabled, it protects against cascading failures from upstream rate limits.
+	circuitBreakers *CircuitBreakers
+
 	inflightReqs atomic.Int32
 	inflightWG   sync.WaitGroup // For graceful shutdown.
 
@@ -49,12 +54,35 @@ var _ http.Handler = &RequestBridge{}
 //
 // mcpProxy will be closed when the [RequestBridge] is closed.
 func NewRequestBridge(ctx context.Context, providers []Provider, recorder Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, metrics *Metrics, tracer trace.Tracer) (*RequestBridge, error) {
+	return NewRequestBridgeWithCircuitBreaker(ctx, providers, recorder, mcpProxy, logger, metrics, tracer, nil)
+}
+
+// NewRequestBridgeWithCircuitBreaker creates a new *[RequestBridge] with per-provider circuit breaker configuration.
+// The cbConfigs map is keyed by provider name. Providers not in the map will not have circuit breaker protection.
+// Pass nil to disable circuit breakers entirely.
+func NewRequestBridgeWithCircuitBreaker(ctx context.Context, providers []Provider, recorder Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, metrics *Metrics, tracer trace.Tracer, cbConfigs map[string]CircuitBreakerConfig) (*RequestBridge, error) {
 	mux := http.NewServeMux()
 
+	// Create circuit breakers with metrics callback
+	var onChange func(name string, from, to gobreaker.State)
+	if metrics != nil {
+		onChange = func(name string, from, to gobreaker.State) {
+			provider, endpoint, _ := strings.Cut(name, ":")
+			metrics.CircuitBreakerState.WithLabelValues(provider, endpoint).Set(stateToGaugeValue(to))
+			if to == gobreaker.StateOpen {
+				metrics.CircuitBreakerTrips.WithLabelValues(provider, endpoint).Inc()
+			}
+		}
+	}
+	cbs := NewCircuitBreakers(cbConfigs, onChange)
+
 	for _, provider := range providers {
 		// Add the known provider-specific routes which are bridged (i.e. intercepted and augmented).
 		for _, path := range provider.BridgedRoutes() {
-			mux.HandleFunc(path, newInterceptionProcessor(provider, recorder, mcpProxy, logger, metrics, tracer))
+			handler := newInterceptionProcessor(provider, recorder, mcpProxy, logger, metrics, tracer)
+			// Wrap with circuit breaker middleware if configured for this provider
+			handler = CircuitBreakerMiddleware(cbs, metrics, provider.Name())(handler).ServeHTTP
+			mux.HandleFunc(path, handler)
 		}
 
 		// Any requests which passthrough to this will be reverse-proxied to the upstream.
@@ -77,11 +105,12 @@ func NewRequestBridge(ctx context.Context, providers []Provider, recorder Record
 
 	inflightCtx, cancel := context.WithCancel(context.Background())
 	return &RequestBridge{
-		mux:            mux,
-		logger:         logger,
-		mcpProxy:       mcpProxy,
-		inflightCtx:    inflightCtx,
-		inflightCancel: cancel,
+		mux:             mux,
+		logger:          logger,
+		mcpProxy:        mcpProxy,
+		circuitBreakers: cbs,
+		inflightCtx:     inflightCtx,
+		inflightCancel:  cancel,
 
 		closed: make(chan struct{}, 1),
 	}, nil
@@ -153,6 +182,11 @@ func (b *RequestBridge) InflightRequests() int32 {
 	return b.inflightReqs.Load()
 }
 
+// CircuitBreakers returns the circuit breakers for this bridge.
+func (b *RequestBridge) CircuitBreakers() *CircuitBreakers {
+	return b.circuitBreakers
+}
+
 // mergeContexts merges two contexts together, so that if either is cancelled
 // the returned context is cancelled. The context values will only be used from
 // the first context.

circuit_breaker.go

@@ -0,0 +1,200 @@
+package aibridge
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/sony/gobreaker/v2"
+)
+
+// CircuitBreakerConfig holds configuration for circuit breakers.
+// Fields match gobreaker.Settings for clarity.
+type CircuitBreakerConfig struct {
+	// MaxRequests is the maximum number of requests allowed in half-open state.
+	MaxRequests uint32
+	// Interval is the cyclic period of the closed state for clearing internal counts.
+	Interval time.Duration
+	// Timeout is how long the circuit stays open before transitioning to half-open.
+	Timeout time.Duration
+	// FailureThreshold is the number of consecutive failures that triggers the circuit to open.
+	FailureThreshold uint32
+	// IsFailure determines if a status code should count as a failure.
+	// If nil, defaults to 429, 503, and 529 (Anthropic overloaded).
+	IsFailure func(statusCode int) bool
+}
+
+// DefaultCircuitBreakerConfig returns sensible defaults for circuit breaker configuration.
+func DefaultCircuitBreakerConfig() CircuitBreakerConfig {
+	return CircuitBreakerConfig{
+		FailureThreshold: 5,
+		Interval:         10 * time.Second,
+		Timeout:          30 * time.Second,
+		MaxRequests:      3,
+		IsFailure:        DefaultIsFailure,
+	}
+}
+
+// DefaultIsFailure returns true for status codes that typically indicate
+// upstream overload: 429 (Too Many Requests), 503 (Service Unavailable),
+// and 529 (Anthropic Overloaded).
+func DefaultIsFailure(statusCode int) bool {
+	switch statusCode {
+	case http.StatusTooManyRequests, // 429
+		http.StatusServiceUnavailable, // 503
+		529:                           // Anthropic "Overloaded"
+		return true
+	default:
+		return false
+	}
+}
+
+// CircuitBreakers manages per-endpoint circuit breakers using sony/gobreaker.
+// Circuit breakers are keyed by "provider:endpoint" for per-endpoint isolation.
+type CircuitBreakers struct {
+	breakers sync.Map // map[string]*gobreaker.CircuitBreaker[any]
+	configs  map[string]CircuitBreakerConfig
+	onChange func(name string, from, to gobreaker.State)
+}
+
+// NewCircuitBreakers creates a new circuit breaker manager with per-provider configs.
+// The configs map is keyed by provider name. Providers not in the map will not have
+// circuit breaker protection.
+func NewCircuitBreakers(configs map[string]CircuitBreakerConfig, onChange func(name string, from, to gobreaker.State)) *CircuitBreakers {
+	return &CircuitBreakers{
+		configs:  configs,
+		onChange: onChange,
+	}
+}
+
+// getConfig returns the config for a provider, or nil if not configured.
+func (c *CircuitBreakers) getConfig(provider string) *CircuitBreakerConfig {
+	if c.configs == nil {
+		return nil
+	}
+	cfg, ok := c.configs[provider]
+	if !ok {
+		return nil
+	}
+	return &cfg
+}
+
+// getOrCreate returns the circuit breaker for a provider/endpoint, creating if needed.
+// Returns nil if the provider is not configured.
+func (c *CircuitBreakers) getOrCreate(provider, endpoint string) *gobreaker.CircuitBreaker[any] {
+	cfg := c.getConfig(provider)
+	if cfg == nil {
+		return nil
+	}
+
+	key := provider + ":" + endpoint
+	if v, ok := c.breakers.Load(key); ok {
+		return v.(*gobreaker.CircuitBreaker[any])
+	}
+
+	settings := gobreaker.Settings{
+		Name:        key,
+		MaxRequests: cfg.MaxRequests,
+		Interval:    cfg.Interval,
+		Timeout:     cfg.Timeout,
+		ReadyToTrip: func(counts gobreaker.Counts) bool {
+			return counts.ConsecutiveFailures >= cfg.FailureThreshold
+		},
+		OnStateChange: func(name string, from, to gobreaker.State) {
+			if c.onChange != nil {
+				c.onChange(name, from, to)
+			}
+		},
+	}
+
+	cb := gobreaker.NewCircuitBreaker[any](settings)
+	actual, _ := c.breakers.LoadOrStore(key, cb)
+	return actual.(*gobreaker.CircuitBreaker[any])
+}
+
+// statusCapturingWriter wraps http.ResponseWriter to capture the status code.
+type statusCapturingWriter struct {
+	http.ResponseWriter
+	statusCode    int
+	headerWritten bool
+}
+
+func (w *statusCapturingWriter) WriteHeader(code int) {
+	if !w.headerWritten {
+		w.statusCode = code
+		w.headerWritten = true
+	}
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *statusCapturingWriter) Write(b []byte) (int, error) {
+	if !w.headerWritten {
+		w.statusCode = http.StatusOK
+		w.headerWritten = true
+	}
+	return w.ResponseWriter.Write(b)
+}
+
+// CircuitBreakerMiddleware returns middleware that wraps handlers with circuit breaker protection.
+// It captures the response status code to determine success/failure without provider-specific logic.
+// If the provider is not configured, returns a noop middleware (passes through without any circuit breaker).
+func CircuitBreakerMiddleware(cbs *CircuitBreakers, metrics *Metrics, provider string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		cfg := cbs.getConfig(provider)
+		if cfg == nil {
+			// Noop: no config for this provider, pass through without circuit breaker
+			return next
+		}
+
+		isFailure := cfg.IsFailure
+		if isFailure == nil {
+			isFailure = DefaultIsFailure
+		}
+
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			endpoint := strings.TrimPrefix(r.URL.Path, fmt.Sprintf("/%s", provider))
+
+			// Check if circuit is open
+			cb := cbs.getOrCreate(provider, endpoint)
+			if cb != nil && cb.State() == gobreaker.StateOpen {
+				if metrics != nil {
+					metrics.CircuitBreakerRejects.WithLabelValues(provider, endpoint).Inc()
+				}
+				http.Error(w, "circuit breaker is open", http.StatusServiceUnavailable)
+				return
+			}
+
+			// Wrap response writer to capture status code
+			sw := &statusCapturingWriter{ResponseWriter: w, statusCode: http.StatusOK}
+			next.ServeHTTP(sw, r)
+
+			// Record result
+			if cb != nil {
+				if isFailure(sw.statusCode) {
+					_, _ = cb.Execute(func() (any, error) {
+						return nil, fmt.Errorf("upstream error: %d", sw.statusCode)
+					})
+				} else {
+					_, _ = cb.Execute(func() (any, error) { return nil, nil })
+				}
+			}
+		})
+	}
+}
+
+// stateToGaugeValue converts gobreaker.State to a gauge value.
+// closed=0, half-open=0.5, open=1
+func stateToGaugeValue(s gobreaker.State) float64 {
+	switch s {
+	case gobreaker.StateClosed:
+		return 0
+	case gobreaker.StateHalfOpen:
+		return 0.5
+	case gobreaker.StateOpen:
+		return 1
+	default:
+		return 0
+	}
+}