Coverage for agentos/tools/circuit_breaker.py: 25%
137 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 01:44 +0800
1"""
2Circuit Breaker for AgentOS.
4Protects against cascading failures with three states:
5- CLOSED: normal operation, track failures
6- OPEN: circuit tripped, fast-fail all calls
7- HALF_OPEN: probe with limited calls to test recovery
9Supports failure/success thresholds, recovery timeout, and callbacks.
10"""
12import threading
13import time
14from collections.abc import Callable
15from enum import Enum, auto
16from typing import Any, TypeVar
18T = TypeVar("T")
21# ============================================================================
22# Enums & Types
23# ============================================================================
26class CircuitState(Enum):
27 CLOSED = auto() # Normal operation
28 OPEN = auto() # Fast-fail, no calls allowed
29 HALF_OPEN = auto() # Probe mode, limited calls allowed
32CircuitCallback = Callable[["CircuitBreaker", CircuitState, CircuitState], None]
35# ============================================================================
36# CircuitBreaker
37# ============================================================================
40class CircuitBreaker:
41 """Thread-safe circuit breaker.
43 Parameters:
44 failure_threshold: consecutive/max failures before tripping
45 recovery_timeout: seconds before transitioning OPEN → HALF_OPEN
46 half_open_max_calls: max probe calls in HALF_OPEN before deciding
47 success_threshold: successes needed in HALF_OPEN to close circuit
48 """
50 def __init__(
51 self,
52 name: str = "default",
53 failure_threshold: int = 5,
54 recovery_timeout: float = 30.0,
55 half_open_max_calls: int = 3,
56 success_threshold: int = 2,
57 on_state_change: CircuitCallback | None = None,
58 ):
59 self.name = name
60 self.failure_threshold = failure_threshold
61 self.recovery_timeout = recovery_timeout
62 self.half_open_max_calls = half_open_max_calls
63 self.success_threshold = success_threshold
64 self.on_state_change = on_state_change
66 self._lock = threading.RLock()
67 self._state: CircuitState = CircuitState.CLOSED
68 self._failure_count: int = 0
69 self._success_count: int = 0
70 self._half_open_calls: int = 0
71 self._last_failure_time: float = 0.0
72 self._last_success_time: float = 0.0
73 self._total_calls: int = 0
74 self._total_failures: int = 0
75 self._total_successes: int = 0
76 self._opened_at: float = 0.0
78 # ---------- state management ----------
80 def _transition(self, new_state: CircuitState) -> None:
81 old = self._state
82 if old == new_state:
83 return
84 self._state = new_state
85 if new_state == CircuitState.OPEN:
86 self._opened_at = time.time()
87 elif new_state == CircuitState.HALF_OPEN:
88 self._success_count = 0
89 self._half_open_calls = 0
90 elif new_state == CircuitState.CLOSED:
91 self._failure_count = 0
92 if self.on_state_change:
93 try:
94 self.on_state_change(self, old, new_state)
95 except Exception:
96 pass
98 @property
99 def state(self) -> CircuitState:
100 with self._lock:
101 return self._state
103 # ---------- call execution ----------
105 def call(self, fn: Callable[..., T], *args, **kwargs) -> T:
106 """Execute fn through the circuit breaker. Raises CircuitOpenError if open."""
107 self._check_state()
108 self._total_calls += 1
109 try:
110 result = fn(*args, **kwargs)
111 self._on_success()
112 return result
113 except Exception:
114 self._on_failure()
115 raise
117 def _check_state(self) -> None:
118 with self._lock:
119 if self._state == CircuitState.CLOSED:
120 return
121 if self._state == CircuitState.OPEN:
122 elapsed = time.time() - self._opened_at
123 if elapsed >= self.recovery_timeout:
124 self._transition(CircuitState.HALF_OPEN)
125 self._half_open_calls += 1 # count this probe
126 return
127 raise CircuitOpenError(
128 f"Circuit '{self.name}' is OPEN "
129 f"(recovery in {self.recovery_timeout - elapsed:.1f}s)"
130 )
131 if self._state == CircuitState.HALF_OPEN:
132 if self._half_open_calls >= self.half_open_max_calls:
133 raise CircuitOpenError(
134 f"Circuit '{self.name}' HALF_OPEN limit reached "
135 f"({self._half_open_calls}/{self.half_open_max_calls})"
136 )
137 self._half_open_calls += 1
139 def _on_success(self) -> None:
140 with self._lock:
141 self._total_successes += 1
142 self._last_success_time = time.time()
143 if self._state == CircuitState.HALF_OPEN:
144 self._success_count += 1
145 if self._success_count >= self.success_threshold:
146 self._transition(CircuitState.CLOSED)
147 elif self._state == CircuitState.CLOSED:
148 self._failure_count = 0
150 def _on_failure(self) -> None:
151 with self._lock:
152 self._total_failures += 1
153 self._last_failure_time = time.time()
154 self._failure_count += 1
155 if self._state == CircuitState.HALF_OPEN:
156 self._transition(CircuitState.OPEN)
157 elif (
158 self._state == CircuitState.CLOSED and self._failure_count >= self.failure_threshold
159 ):
160 self._transition(CircuitState.OPEN)
162 # ---------- manual control ----------
164 def reset(self) -> None:
165 """Force circuit back to CLOSED."""
166 with self._lock:
167 self._failure_count = 0
168 self._success_count = 0
169 self._half_open_calls = 0
170 self._transition(CircuitState.CLOSED)
172 def trip(self) -> None:
173 """Force circuit OPEN."""
174 with self._lock:
175 self._failure_count = self.failure_threshold
176 self._transition(CircuitState.OPEN)
178 # ---------- stats ----------
180 @property
181 def stats(self) -> dict[str, Any]:
182 with self._lock:
183 return {
184 "name": self.name,
185 "state": self._state.name,
186 "failure_count": self._failure_count,
187 "half_open_calls": self._half_open_calls,
188 "total_calls": self._total_calls,
189 "total_successes": self._total_successes,
190 "total_failures": self._total_failures,
191 "last_failure": self._last_failure_time,
192 "last_success": self._last_success_time,
193 "opened_at": self._opened_at,
194 }
197# ============================================================================
198# Errors
199# ============================================================================
202class CircuitOpenError(Exception):
203 """Raised when a call is attempted on an OPEN circuit."""
207# ============================================================================
208# CircuitRegistry — manage multiple breakers by name
209# ============================================================================
212class CircuitRegistry:
213 """Global registry for named circuit breakers."""
215 def __init__(self):
216 self._breakers: dict[str, CircuitBreaker] = {}
217 self._lock = threading.Lock()
219 def get(self, name: str, **kwargs) -> CircuitBreaker:
220 with self._lock:
221 if name not in self._breakers:
222 self._breakers[name] = CircuitBreaker(name=name, **kwargs)
223 return self._breakers[name]
225 def remove(self, name: str) -> bool:
226 with self._lock:
227 return self._breakers.pop(name, None) is not None
229 def list_breakers(self) -> dict[str, str]:
230 with self._lock:
231 return {n: b.state.name for n, b in self._breakers.items()}
233 def reset_all(self) -> None:
234 with self._lock:
235 for b in self._breakers.values():
236 b.reset()
239_default_registry: CircuitRegistry | None = None
240_registry_lock = threading.Lock()
243def get_circuit_registry() -> CircuitRegistry:
244 global _default_registry
245 if _default_registry is None:
246 with _registry_lock:
247 if _default_registry is None:
248 _default_registry = CircuitRegistry()
249 return _default_registry