Coverage for agentos/tools/circuit_breaker.py: 21%
175 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +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 with granular hooks ----------
164 def allow_request(self) -> bool:
165 """Check whether a request is allowed (used by ToolExecutor)."""
166 with self._lock:
167 if self._state == CircuitState.CLOSED:
168 return True
169 if self._state == CircuitState.OPEN:
170 elapsed = time.time() - self._opened_at
171 if elapsed >= self.recovery_timeout:
172 self._transition(CircuitState.HALF_OPEN)
173 self._half_open_calls += 1
174 return True
175 return False
176 if self._state == CircuitState.HALF_OPEN:
177 if self._half_open_calls >= self.half_open_max_calls:
178 return False
179 self._half_open_calls += 1
180 return True
181 return False
183 def record_success(self) -> None:
184 """Record a successful call (used by ToolExecutor)."""
185 with self._lock:
186 self._total_calls += 1
187 self._total_successes += 1
188 self._last_success_time = time.time()
189 if self._state == CircuitState.HALF_OPEN:
190 self._success_count += 1
191 if self._success_count >= self.success_threshold:
192 self._transition(CircuitState.CLOSED)
193 elif self._state == CircuitState.CLOSED:
194 self._failure_count = 0
196 def record_failure(self) -> None:
197 """Record a failed call (used by ToolExecutor)."""
198 with self._lock:
199 self._total_calls += 1
200 self._total_failures += 1
201 self._last_failure_time = time.time()
202 self._failure_count += 1
203 if self._state == CircuitState.HALF_OPEN:
204 self._transition(CircuitState.OPEN)
205 elif (
206 self._state == CircuitState.CLOSED
207 and self._failure_count >= self.failure_threshold
208 ):
209 self._transition(CircuitState.OPEN)
211 # ---------- manual control ----------
213 def reset(self) -> None:
214 """Force circuit back to CLOSED."""
215 with self._lock:
216 self._failure_count = 0
217 self._success_count = 0
218 self._half_open_calls = 0
219 self._transition(CircuitState.CLOSED)
221 def trip(self) -> None:
222 """Force circuit OPEN."""
223 with self._lock:
224 self._failure_count = self.failure_threshold
225 self._transition(CircuitState.OPEN)
227 # ---------- stats ----------
229 @property
230 def stats(self) -> dict[str, Any]:
231 with self._lock:
232 return {
233 "name": self.name,
234 "state": self._state.name,
235 "failure_count": self._failure_count,
236 "half_open_calls": self._half_open_calls,
237 "total_calls": self._total_calls,
238 "total_successes": self._total_successes,
239 "total_failures": self._total_failures,
240 "last_failure": self._last_failure_time,
241 "last_success": self._last_success_time,
242 "opened_at": self._opened_at,
243 }
246# ============================================================================
247# Errors
248# ============================================================================
251class CircuitOpenError(Exception):
252 """Raised when a call is attempted on an OPEN circuit."""
256# ============================================================================
257# CircuitRegistry — manage multiple breakers by name
258# ============================================================================
261class CircuitRegistry:
262 """Global registry for named circuit breakers."""
264 def __init__(self):
265 self._breakers: dict[str, CircuitBreaker] = {}
266 self._lock = threading.Lock()
268 def get(self, name: str, **kwargs) -> CircuitBreaker:
269 with self._lock:
270 if name not in self._breakers:
271 self._breakers[name] = CircuitBreaker(name=name, **kwargs)
272 return self._breakers[name]
274 def remove(self, name: str) -> bool:
275 with self._lock:
276 return self._breakers.pop(name, None) is not None
278 def list_breakers(self) -> dict[str, str]:
279 with self._lock:
280 return {n: b.state.name for n, b in self._breakers.items()}
282 def reset_all(self) -> None:
283 with self._lock:
284 for b in self._breakers.values():
285 b.reset()
288_default_registry: CircuitRegistry | None = None
289_registry_lock = threading.Lock()
292def get_circuit_registry() -> CircuitRegistry:
293 global _default_registry
294 if _default_registry is None:
295 with _registry_lock:
296 if _default_registry is None:
297 _default_registry = CircuitRegistry()
298 return _default_registry