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,32 @@ 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, DefaultCircuitBreakerConfig())
+}
+
+// NewRequestBridgeWithCircuitBreaker creates a new *[RequestBridge] with custom circuit breaker configuration.
+// See [NewRequestBridge] for more details.
+func NewRequestBridgeWithCircuitBreaker(ctx context.Context, providers []Provider, recorder Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, metrics *Metrics, tracer trace.Tracer, cbConfig 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(float64(to))
+			if to == gobreaker.StateOpen {
+				metrics.CircuitBreakerTrips.WithLabelValues(provider, endpoint).Inc()
+			}
+		}
+	}
+	cbs := NewCircuitBreakers(cbConfig, 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))
+			mux.HandleFunc(path, newInterceptionProcessor(provider, recorder, mcpProxy, logger, metrics, tracer, cbs))
 		}
 
 		// Any requests which passthrough to this will be reverse-proxied to the upstream.
@@ -77,11 +102,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 +179,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,131 @@
+package aibridge
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/sony/gobreaker/v2"
+)
+
+// CircuitBreakerConfig holds configuration for circuit breakers.
+// Fields match gobreaker.Settings for clarity.
+type CircuitBreakerConfig struct {
+	// Enabled controls whether circuit breakers are active.
+	Enabled bool
+	// 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
+}
+
+// DefaultCircuitBreakerConfig returns sensible defaults for circuit breaker configuration.
+func DefaultCircuitBreakerConfig() CircuitBreakerConfig {
+	return CircuitBreakerConfig{
+		Enabled:          false, // Disabled by default for backward compatibility
+		FailureThreshold: 5,
+		Interval:         10 * time.Second,
+		Timeout:          30 * time.Second,
+		MaxRequests:      3,
+	}
+}
+
+// isCircuitBreakerFailure returns true if the given HTTP status code
+// should count as a failure for circuit breaker purposes.
+func isCircuitBreakerFailure(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]
+	config   CircuitBreakerConfig
+	onChange func(name string, from, to gobreaker.State)
+}
+
+// NewCircuitBreakers creates a new circuit breaker manager.
+func NewCircuitBreakers(config CircuitBreakerConfig, onChange func(name string, from, to gobreaker.State)) *CircuitBreakers {
+	return &CircuitBreakers{
+		config:   config,
+		onChange: onChange,
+	}
+}
+
+// Allow checks if a request to provider/endpoint should be allowed.
+func (c *CircuitBreakers) Allow(provider, endpoint string) bool {
+	if !c.config.Enabled {
+		return true
+	}
+	cb := c.getOrCreate(provider, endpoint)
+	return cb.State() != gobreaker.StateOpen
+}
+
+// RecordSuccess records a successful request.
+func (c *CircuitBreakers) RecordSuccess(provider, endpoint string) {
+	if !c.config.Enabled {
+		return
+	}
+	cb := c.getOrCreate(provider, endpoint)
+	_, _ = cb.Execute(func() (any, error) { return nil, nil })
+}
+
+// RecordFailure records a failed request. Returns true if this caused the circuit to open.
+func (c *CircuitBreakers) RecordFailure(provider, endpoint string, statusCode int) bool {
+	if !c.config.Enabled || !isCircuitBreakerFailure(statusCode) {
+		return false
+	}
+	cb := c.getOrCreate(provider, endpoint)
+	before := cb.State()
+	_, _ = cb.Execute(func() (any, error) {
+		return nil, fmt.Errorf("upstream error: %d", statusCode)
+	})
+	return before != gobreaker.StateOpen && cb.State() == gobreaker.StateOpen
+}
+
+// State returns the current state for a provider/endpoint.
+func (c *CircuitBreakers) State(provider, endpoint string) gobreaker.State {
+	if !c.config.Enabled {
+		return gobreaker.StateClosed
+	}
+	cb := c.getOrCreate(provider, endpoint)
+	return cb.State()
+}
+
+func (c *CircuitBreakers) getOrCreate(provider, endpoint string) *gobreaker.CircuitBreaker[any] {
+	key := provider + ":" + endpoint
+	if v, ok := c.breakers.Load(key); ok {
+		return v.(*gobreaker.CircuitBreaker[any])
+	}
+
+	settings := gobreaker.Settings{
+		Name:        key,
+		MaxRequests: c.config.MaxRequests,
+		Interval:    c.config.Interval,
+		Timeout:     c.config.Timeout,
+		ReadyToTrip: func(counts gobreaker.Counts) bool {
+			return counts.ConsecutiveFailures >= c.config.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])
+}