Coverage for agentos/core/circuit_breaker.py: 0%
154 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +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 dataclasses import dataclass
19from enum import Enum
20from functools import wraps
21from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar
23logger = logging.getLogger(__name__)
25T = TypeVar("T")
28# ============================================================================
29# State machine
30# ============================================================================
32class CircuitState(str, Enum):
33 CLOSED = "closed" # Normal operation — requests pass through
34 OPEN = "open" # Failing — requests are rejected immediately
35 HALF_OPEN = "half_open" # Probing — limited requests allowed to test recovery
38@dataclass
39class CircuitConfig:
40 """Circuit breaker configuration."""
42 failure_threshold: int = 5 # Consecutive failures to trip OPEN
43 success_threshold: int = 2 # Consecutive successes in HALF_OPEN to reset
44 timeout_seconds: float = 60.0 # Seconds in OPEN before transitioning to HALF_OPEN
45 half_open_max_requests: int = 1 # Max concurrent requests in HALF_OPEN
46 excluded_exceptions: tuple = () # Exceptions that don't count as failures
49@dataclass
50class CircuitStats:
51 """Per-circuit statistics."""
53 state: CircuitState = CircuitState.CLOSED
54 failure_count: int = 0
55 success_count: int = 0
56 last_failure_time: float = 0.0
57 last_success_time: float = 0.0
58 total_failures: int = 0
59 total_successes: int = 0
60 opened_at: float = 0.0
61 half_open_requests: int = 0
63 def reset(self) -> None:
64 self.failure_count = 0
65 self.success_count = 0
66 self.half_open_requests = 0
69# ============================================================================
70# Circuit Breaker
71# ============================================================================
73class CircuitBreaker:
74 """Thread-safe circuit breaker for a single endpoint."""
76 def __init__(
77 self,
78 name: str,
79 config: Optional[CircuitConfig] = None,
80 ):
81 self.name = name
82 self.config = config or CircuitConfig()
83 self.stats = CircuitStats()
84 self._lock = asyncio.Lock()
86 @property
87 def state(self) -> CircuitState:
88 return self.stats.state
90 async def _transition(self, new_state: CircuitState) -> None:
91 old = self.stats.state
92 if old == new_state:
93 return
94 self.stats.state = new_state
95 if new_state == CircuitState.OPEN:
96 self.stats.opened_at = time.monotonic()
97 self.stats.reset()
98 logger.warning(
99 "Circuit '%s' OPENED after %d consecutive failures",
100 self.name, self.config.failure_threshold,
101 )
102 elif new_state == CircuitState.CLOSED:
103 self.stats.reset()
104 logger.info("Circuit '%s' CLOSED — service recovered", self.name)
105 elif new_state == CircuitState.HALF_OPEN:
106 self.stats.reset()
107 logger.info("Circuit '%s' HALF_OPEN — probing", self.name)
109 def _should_retry_open(self) -> bool:
110 """Check if OPEN state has expired and should move to HALF_OPEN."""
111 elapsed = time.monotonic() - self.stats.opened_at
112 return elapsed >= self.config.timeout_seconds
114 async def _on_success(self) -> None:
115 async with self._lock:
116 self.stats.total_successes += 1
117 self.stats.last_success_time = time.monotonic()
119 if self.stats.state == CircuitState.HALF_OPEN:
120 self.stats.success_count += 1
121 if self.stats.success_count >= self.config.success_threshold:
122 await self._transition(CircuitState.CLOSED)
123 else:
124 self.stats.failure_count = 0 # Reset on success when CLOSED
126 async def _on_failure(self, exc: Exception) -> None:
127 async with self._lock:
128 if isinstance(exc, self.config.excluded_exceptions):
129 return
131 self.stats.total_failures += 1
132 self.stats.last_failure_time = time.monotonic()
134 if self.stats.state == CircuitState.HALF_OPEN:
135 await self._transition(CircuitState.OPEN)
136 else:
137 self.stats.failure_count += 1
138 if self.stats.failure_count >= self.config.failure_threshold:
139 await self._transition(CircuitState.OPEN)
141 async def acquire(self) -> bool:
142 """Try to acquire permission to make a request.
144 Returns True if request should proceed, False if circuit is OPEN.
145 """
146 async with self._lock:
147 state = self.stats.state
149 if state == CircuitState.CLOSED:
150 return True
152 if state == CircuitState.OPEN:
153 if self._should_retry_open():
154 await self._transition(CircuitState.HALF_OPEN)
155 # Re-read state after transition
156 state = self.stats.state
157 else:
158 return False
160 if state == CircuitState.HALF_OPEN:
161 if self.stats.half_open_requests < self.config.half_open_max_requests:
162 self.stats.half_open_requests += 1
163 return True
164 return False
166 return False # unreachable, but defensive
168 async def release(self) -> None:
169 """Release half-open slot after request completes."""
170 async with self._lock:
171 if self.stats.state == CircuitState.HALF_OPEN:
172 self.stats.half_open_requests = max(0, self.stats.half_open_requests - 1)
174 async def call(
175 self,
176 fn: Callable[..., Awaitable[T]],
177 *args: Any,
178 fallback: Optional[Callable[..., Awaitable[T]]] = None,
179 **kwargs: Any,
180 ) -> T:
181 """Execute fn through the circuit breaker.
183 Raises CircuitOpenError if circuit is open and no fallback provided.
184 """
185 if not await self.acquire():
186 if fallback is not None:
187 logger.debug("Circuit '%s' open — using fallback", self.name)
188 return await fallback(*args, **kwargs)
189 raise CircuitOpenError(
190 f"Circuit '{self.name}' is OPEN — request rejected"
191 )
193 try:
194 result = await fn(*args, **kwargs)
195 await self._on_success()
196 return result
197 except Exception as exc:
198 await self._on_failure(exc)
199 raise
200 finally:
201 await self.release()
204# ============================================================================
205# Registry
206# ============================================================================
208class CircuitRegistry:
209 """Global registry of circuit breakers, keyed by endpoint name."""
211 def __init__(self):
212 self._circuits: Dict[str, CircuitBreaker] = {}
213 self._lock = asyncio.Lock()
215 async def get_or_create(
216 self,
217 name: str,
218 config: Optional[CircuitConfig] = None,
219 ) -> CircuitBreaker:
220 async with self._lock:
221 if name not in self._circuits:
222 self._circuits[name] = CircuitBreaker(name=name, config=config)
223 return self._circuits[name]
225 def get_all_stats(self) -> Dict[str, CircuitStats]:
226 """Export all circuit stats for monitoring."""
227 return {name: cb.stats for name, cb in self._circuits.items()}
229 async def reset_all(self) -> None:
230 """Reset all circuits to CLOSED (for testing/admin)."""
231 async with self._lock:
232 for cb in self._circuits.values():
233 cb.stats.reset()
234 cb.stats.state = CircuitState.CLOSED
236 async def force_open(self, name: str) -> None:
237 """Force a circuit OPEN (for manual intervention)."""
238 async with self._lock:
239 if name in self._circuits:
240 await self._circuits[name]._transition(CircuitState.OPEN)
242 async def force_closed(self, name: str) -> None:
243 """Force a circuit CLOSED (for manual intervention)."""
244 async with self._lock:
245 if name in self._circuits:
246 await self._circuits[name]._transition(CircuitState.CLOSED)
249# ============================================================================
250# Decorator
251# ============================================================================
253def circuit_breaker(
254 name: str,
255 failure_threshold: int = 5,
256 timeout_seconds: float = 60.0,
257 success_threshold: int = 2,
258 excluded_exceptions: tuple = (),
259):
260 """Decorator: wrap an async function with a circuit breaker.
262 Usage:
263 @circuit_breaker("llm_api", failure_threshold=3, timeout_seconds=30)
264 async def call_llm(prompt: str) -> str: ...
265 """
266 config = CircuitConfig(
267 failure_threshold=failure_threshold,
268 timeout_seconds=timeout_seconds,
269 success_threshold=success_threshold,
270 excluded_exceptions=excluded_exceptions,
271 )
272 cb = CircuitBreaker(name=name, config=config)
274 def decorator(fn: Callable[..., Awaitable[T]]):
275 @wraps(fn)
276 async def wrapper(*args, **kwargs):
277 return await cb.call(fn, *args, **kwargs)
278 wrapper._circuit_breaker = cb # type: ignore[attr-defined]
279 return wrapper
280 return decorator
283# ============================================================================
284# Exceptions
285# ============================================================================
287class CircuitOpenError(Exception):
288 """Raised when request is rejected because circuit is OPEN."""
289 pass
292# ============================================================================
293# Module-level defaults
294# ============================================================================
296# Global registry — use this shared instance across the app
297default_registry = CircuitRegistry()