Coverage for agentos/core/circuit_breaker.py: 98%
154 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 23:40 +0800
1"""AgentOS Circuit Breaker — protect against cascading failures.
3Production-grade circuit breaker pattern with:
4- 3 states: CLOSED → OPEN → HALF_OPEN → CLOSED
5- Configurable failure threshold, timeout, and half-open probe limit
6- Per-endpoint isolation with shared registry
7- Optional fallback function support
8- Thread-safe, asyncio-native
10Design: ~280 lines, zero external deps beyond stdlib + asyncio.
11"""
13from __future__ import annotations
15import asyncio
16import logging
17import time
18from collections.abc import Awaitable, Callable
19from dataclasses import dataclass
20from enum import StrEnum
21from functools import wraps
22from typing import Any, TypeVar
24logger = logging.getLogger(__name__)
26T = TypeVar("T")
29# ============================================================================
30# State machine
31# ============================================================================
34class CircuitState(StrEnum):
35 CLOSED = "closed" # Normal operation — requests pass through
36 OPEN = "open" # Failing — requests are rejected immediately
37 HALF_OPEN = "half_open" # Probing — limited requests allowed to test recovery
40@dataclass
41class CircuitConfig:
42 """Circuit breaker configuration."""
44 failure_threshold: int = 5 # Consecutive failures to trip OPEN
45 success_threshold: int = 2 # Consecutive successes in HALF_OPEN to reset
46 timeout_seconds: float = 60.0 # Seconds in OPEN before transitioning to HALF_OPEN
47 half_open_max_requests: int = 1 # Max concurrent requests in HALF_OPEN
48 excluded_exceptions: tuple = () # Exceptions that don't count as failures
51@dataclass
52class CircuitStats:
53 """Per-circuit statistics."""
55 state: CircuitState = CircuitState.CLOSED
56 failure_count: int = 0
57 success_count: int = 0
58 last_failure_time: float = 0.0
59 last_success_time: float = 0.0
60 total_failures: int = 0
61 total_successes: int = 0
62 opened_at: float = 0.0
63 half_open_requests: int = 0
65 def reset(self) -> None:
66 self.failure_count = 0
67 self.success_count = 0
68 self.half_open_requests = 0
71# ============================================================================
72# Circuit Breaker
73# ============================================================================
76class CircuitBreaker:
77 """Thread-safe circuit breaker for a single endpoint."""
79 def __init__(
80 self,
81 name: str,
82 config: CircuitConfig | None = None,
83 ):
84 self.name = name
85 self.config = config or CircuitConfig()
86 self.stats = CircuitStats()
87 self._lock = asyncio.Lock()
89 @property
90 def state(self) -> CircuitState:
91 return self.stats.state
93 async def _transition(self, new_state: CircuitState) -> None:
94 old = self.stats.state
95 if old == new_state:
96 return
97 self.stats.state = new_state
98 if new_state == CircuitState.OPEN:
99 self.stats.opened_at = time.monotonic()
100 self.stats.reset()
101 logger.warning(
102 "Circuit '%s' OPENED after %d consecutive failures",
103 self.name,
104 self.config.failure_threshold,
105 )
106 elif new_state == CircuitState.CLOSED:
107 self.stats.reset()
108 logger.info("Circuit '%s' CLOSED — service recovered", self.name)
109 elif new_state == CircuitState.HALF_OPEN:
110 self.stats.reset()
111 logger.info("Circuit '%s' HALF_OPEN — probing", self.name)
113 def _should_retry_open(self) -> bool:
114 """Check if OPEN state has expired and should move to HALF_OPEN."""
115 elapsed = time.monotonic() - self.stats.opened_at
116 return elapsed >= self.config.timeout_seconds
118 async def _on_success(self) -> None:
119 async with self._lock:
120 self.stats.total_successes += 1
121 self.stats.last_success_time = time.monotonic()
123 if self.stats.state == CircuitState.HALF_OPEN:
124 self.stats.success_count += 1
125 if self.stats.success_count >= self.config.success_threshold:
126 await self._transition(CircuitState.CLOSED)
127 else:
128 self.stats.failure_count = 0 # Reset on success when CLOSED
130 async def _on_failure(self, exc: Exception) -> None:
131 async with self._lock:
132 if isinstance(exc, self.config.excluded_exceptions):
133 return
135 self.stats.total_failures += 1
136 self.stats.last_failure_time = time.monotonic()
138 if self.stats.state == CircuitState.HALF_OPEN:
139 await self._transition(CircuitState.OPEN)
140 else:
141 self.stats.failure_count += 1
142 if self.stats.failure_count >= self.config.failure_threshold:
143 await self._transition(CircuitState.OPEN)
145 async def acquire(self) -> bool:
146 """Try to acquire permission to make a request.
148 Returns True if request should proceed, False if circuit is OPEN.
149 """
150 async with self._lock:
151 state = self.stats.state
153 if state == CircuitState.CLOSED:
154 return True
156 if state == CircuitState.OPEN:
157 if self._should_retry_open():
158 await self._transition(CircuitState.HALF_OPEN)
159 # Re-read state after transition
160 state = self.stats.state
161 else:
162 return False
164 if state == CircuitState.HALF_OPEN:
165 if self.stats.half_open_requests < self.config.half_open_max_requests:
166 self.stats.half_open_requests += 1
167 return True
168 return False
170 return False # unreachable, but defensive
172 async def release(self) -> None:
173 """Release half-open slot after request completes."""
174 async with self._lock:
175 if self.stats.state == CircuitState.HALF_OPEN:
176 self.stats.half_open_requests = max(0, self.stats.half_open_requests - 1)
178 async def call(
179 self,
180 fn: Callable[..., Awaitable[T]],
181 *args: Any,
182 fallback: Callable[..., Awaitable[T]] | None = None,
183 **kwargs: Any,
184 ) -> T:
185 """Execute fn through the circuit breaker.
187 Raises CircuitOpenError if circuit is open and no fallback provided.
188 """
189 if not await self.acquire():
190 if fallback is not None:
191 logger.debug("Circuit '%s' open — using fallback", self.name)
192 return await fallback(*args, **kwargs)
193 raise CircuitOpenError(f"Circuit '{self.name}' is OPEN — request rejected")
195 try:
196 result = await fn(*args, **kwargs)
197 await self._on_success()
198 return result
199 except Exception as exc:
200 await self._on_failure(exc)
201 raise
202 finally:
203 await self.release()
206# ============================================================================
207# Registry
208# ============================================================================
211class CircuitRegistry:
212 """Global registry of circuit breakers, keyed by endpoint name."""
214 def __init__(self):
215 self._circuits: dict[str, CircuitBreaker] = {}
216 self._lock = asyncio.Lock()
218 async def get_or_create(
219 self,
220 name: str,
221 config: CircuitConfig | None = None,
222 ) -> CircuitBreaker:
223 async with self._lock:
224 if name not in self._circuits:
225 self._circuits[name] = CircuitBreaker(name=name, config=config)
226 return self._circuits[name]
228 def get_all_stats(self) -> dict[str, CircuitStats]:
229 """Export all circuit stats for monitoring."""
230 return {name: cb.stats for name, cb in self._circuits.items()}
232 async def reset_all(self) -> None:
233 """Reset all circuits to CLOSED (for testing/admin)."""
234 async with self._lock:
235 for cb in self._circuits.values():
236 cb.stats.reset()
237 cb.stats.state = CircuitState.CLOSED
239 async def force_open(self, name: str) -> None:
240 """Force a circuit OPEN (for manual intervention)."""
241 async with self._lock:
242 if name in self._circuits:
243 await self._circuits[name]._transition(CircuitState.OPEN)
245 async def force_closed(self, name: str) -> None:
246 """Force a circuit CLOSED (for manual intervention)."""
247 async with self._lock:
248 if name in self._circuits:
249 await self._circuits[name]._transition(CircuitState.CLOSED)
252# ============================================================================
253# Decorator
254# ============================================================================
257def circuit_breaker(
258 name: str,
259 failure_threshold: int = 5,
260 timeout_seconds: float = 60.0,
261 success_threshold: int = 2,
262 excluded_exceptions: tuple = (),
263):
264 """Decorator: wrap an async function with a circuit breaker.
266 Usage:
267 @circuit_breaker("llm_api", failure_threshold=3, timeout_seconds=30)
268 async def call_llm(prompt: str) -> str: ...
269 """
270 config = CircuitConfig(
271 failure_threshold=failure_threshold,
272 timeout_seconds=timeout_seconds,
273 success_threshold=success_threshold,
274 excluded_exceptions=excluded_exceptions,
275 )
276 cb = CircuitBreaker(name=name, config=config)
278 def decorator(fn: Callable[..., Awaitable[T]]):
279 @wraps(fn)
280 async def wrapper(*args, **kwargs):
281 return await cb.call(fn, *args, **kwargs)
283 wrapper._circuit_breaker = cb # type: ignore[attr-defined]
284 return wrapper
286 return decorator
289# ============================================================================
290# Exceptions
291# ============================================================================
294class CircuitOpenError(Exception):
295 """Raised when request is rejected because circuit is OPEN."""
299# ============================================================================
300# Module-level defaults
301# ============================================================================
303# Global registry — use this shared instance across the app
304default_registry = CircuitRegistry()