Coverage for agentos/tools/circuit_breaker.py: 25%
137 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +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 enum import Enum, auto
15from typing import Any, Callable, Dict, Optional, TypeVar
17T = TypeVar("T")
20# ============================================================================
21# Enums & Types
22# ============================================================================
24class CircuitState(Enum):
25 CLOSED = auto() # Normal operation
26 OPEN = auto() # Fast-fail, no calls allowed
27 HALF_OPEN = auto() # Probe mode, limited calls allowed
30CircuitCallback = Callable[["CircuitBreaker", CircuitState, CircuitState], None]
33# ============================================================================
34# CircuitBreaker
35# ============================================================================
37class CircuitBreaker:
38 """Thread-safe circuit breaker.
40 Parameters:
41 failure_threshold: consecutive/max failures before tripping
42 recovery_timeout: seconds before transitioning OPEN → HALF_OPEN
43 half_open_max_calls: max probe calls in HALF_OPEN before deciding
44 success_threshold: successes needed in HALF_OPEN to close circuit
45 """
47 def __init__(
48 self,
49 name: str = "default",
50 failure_threshold: int = 5,
51 recovery_timeout: float = 30.0,
52 half_open_max_calls: int = 3,
53 success_threshold: int = 2,
54 on_state_change: Optional[CircuitCallback] = None,
55 ):
56 self.name = name
57 self.failure_threshold = failure_threshold
58 self.recovery_timeout = recovery_timeout
59 self.half_open_max_calls = half_open_max_calls
60 self.success_threshold = success_threshold
61 self.on_state_change = on_state_change
63 self._lock = threading.RLock()
64 self._state: CircuitState = CircuitState.CLOSED
65 self._failure_count: int = 0
66 self._success_count: int = 0
67 self._half_open_calls: int = 0
68 self._last_failure_time: float = 0.0
69 self._last_success_time: float = 0.0
70 self._total_calls: int = 0
71 self._total_failures: int = 0
72 self._total_successes: int = 0
73 self._opened_at: float = 0.0
75 # ---------- state management ----------
77 def _transition(self, new_state: CircuitState) -> None:
78 old = self._state
79 if old == new_state:
80 return
81 self._state = new_state
82 if new_state == CircuitState.OPEN:
83 self._opened_at = time.time()
84 elif new_state == CircuitState.HALF_OPEN:
85 self._success_count = 0
86 self._half_open_calls = 0
87 elif new_state == CircuitState.CLOSED:
88 self._failure_count = 0
89 if self.on_state_change:
90 try:
91 self.on_state_change(self, old, new_state)
92 except Exception:
93 pass
95 @property
96 def state(self) -> CircuitState:
97 with self._lock:
98 return self._state
100 # ---------- call execution ----------
102 def call(self, fn: Callable[..., T], *args, **kwargs) -> T:
103 """Execute fn through the circuit breaker. Raises CircuitOpenError if open."""
104 self._check_state()
105 self._total_calls += 1
106 try:
107 result = fn(*args, **kwargs)
108 self._on_success()
109 return result
110 except Exception:
111 self._on_failure()
112 raise
114 def _check_state(self) -> None:
115 with self._lock:
116 if self._state == CircuitState.CLOSED:
117 return
118 if self._state == CircuitState.OPEN:
119 elapsed = time.time() - self._opened_at
120 if elapsed >= self.recovery_timeout:
121 self._transition(CircuitState.HALF_OPEN)
122 self._half_open_calls += 1 # count this probe
123 return
124 raise CircuitOpenError(
125 f"Circuit '{self.name}' is OPEN "
126 f"(recovery in {self.recovery_timeout - elapsed:.1f}s)"
127 )
128 if self._state == CircuitState.HALF_OPEN:
129 if self._half_open_calls >= self.half_open_max_calls:
130 raise CircuitOpenError(
131 f"Circuit '{self.name}' HALF_OPEN limit reached "
132 f"({self._half_open_calls}/{self.half_open_max_calls})"
133 )
134 self._half_open_calls += 1
136 def _on_success(self) -> None:
137 with self._lock:
138 self._total_successes += 1
139 self._last_success_time = time.time()
140 if self._state == CircuitState.HALF_OPEN:
141 self._success_count += 1
142 if self._success_count >= self.success_threshold:
143 self._transition(CircuitState.CLOSED)
144 elif self._state == CircuitState.CLOSED:
145 self._failure_count = 0
147 def _on_failure(self) -> None:
148 with self._lock:
149 self._total_failures += 1
150 self._last_failure_time = time.time()
151 self._failure_count += 1
152 if self._state == CircuitState.HALF_OPEN:
153 self._transition(CircuitState.OPEN)
154 elif self._state == CircuitState.CLOSED and self._failure_count >= self.failure_threshold:
155 self._transition(CircuitState.OPEN)
157 # ---------- manual control ----------
159 def reset(self) -> None:
160 """Force circuit back to CLOSED."""
161 with self._lock:
162 self._failure_count = 0
163 self._success_count = 0
164 self._half_open_calls = 0
165 self._transition(CircuitState.CLOSED)
167 def trip(self) -> None:
168 """Force circuit OPEN."""
169 with self._lock:
170 self._failure_count = self.failure_threshold
171 self._transition(CircuitState.OPEN)
173 # ---------- stats ----------
175 @property
176 def stats(self) -> Dict[str, Any]:
177 with self._lock:
178 return {
179 "name": self.name,
180 "state": self._state.name,
181 "failure_count": self._failure_count,
182 "half_open_calls": self._half_open_calls,
183 "total_calls": self._total_calls,
184 "total_successes": self._total_successes,
185 "total_failures": self._total_failures,
186 "last_failure": self._last_failure_time,
187 "last_success": self._last_success_time,
188 "opened_at": self._opened_at,
189 }
192# ============================================================================
193# Errors
194# ============================================================================
196class CircuitOpenError(Exception):
197 """Raised when a call is attempted on an OPEN circuit."""
198 pass
201# ============================================================================
202# CircuitRegistry — manage multiple breakers by name
203# ============================================================================
205class CircuitRegistry:
206 """Global registry for named circuit breakers."""
208 def __init__(self):
209 self._breakers: Dict[str, CircuitBreaker] = {}
210 self._lock = threading.Lock()
212 def get(self, name: str, **kwargs) -> CircuitBreaker:
213 with self._lock:
214 if name not in self._breakers:
215 self._breakers[name] = CircuitBreaker(name=name, **kwargs)
216 return self._breakers[name]
218 def remove(self, name: str) -> bool:
219 with self._lock:
220 return self._breakers.pop(name, None) is not None
222 def list_breakers(self) -> Dict[str, str]:
223 with self._lock:
224 return {n: b.state.name for n, b in self._breakers.items()}
226 def reset_all(self) -> None:
227 with self._lock:
228 for b in self._breakers.values():
229 b.reset()
232_default_registry: Optional[CircuitRegistry] = None
233_registry_lock = threading.Lock()
236def get_circuit_registry() -> CircuitRegistry:
237 global _default_registry
238 if _default_registry is None:
239 with _registry_lock:
240 if _default_registry is None:
241 _default_registry = CircuitRegistry()
242 return _default_registry