chuk-tool-processor 0.6.4__py3-none-any.whl → 0.9.7__py3-none-any.whl

chuk_tool_processor/execution/wrappers/circuit_breaker.py

@@ -0,0 +1,370 @@
+# chuk_tool_processor/execution/wrappers/circuit_breaker.py
+"""
+Circuit breaker pattern for tool execution.
+
+Prevents cascading failures by tracking failure rates and temporarily
+blocking calls to failing tools. Implements a state machine:
+
+CLOSED → OPEN → HALF_OPEN → CLOSED (or back to OPEN)
+
+States:
+- CLOSED: Normal operation, requests pass through
+- OPEN: Too many failures, requests blocked immediately
+- HALF_OPEN: Testing if service recovered, limited requests allowed
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from chuk_tool_processor.core.exceptions import ToolCircuitOpenError
+from chuk_tool_processor.logging import get_logger
+from chuk_tool_processor.models.tool_call import ToolCall
+from chuk_tool_processor.models.tool_result import ToolResult
+
+logger = get_logger("chuk_tool_processor.execution.wrappers.circuit_breaker")
+
+# Optional observability imports
+try:
+    from chuk_tool_processor.observability.metrics import get_metrics
+    from chuk_tool_processor.observability.tracing import trace_circuit_breaker
+
+    _observability_available = True
+except ImportError:
+    _observability_available = False
+
+    # No-op functions when observability not available
+    def get_metrics():
+        return None
+
+    def trace_circuit_breaker(*_args, **_kwargs):
+        from contextlib import nullcontext
+
+        return nullcontext()
+
+
+# --------------------------------------------------------------------------- #
+# Circuit breaker state
+# --------------------------------------------------------------------------- #
+class CircuitState(str, Enum):
+    """Circuit breaker states."""
+
+    CLOSED = "closed"  # Normal operation
+    OPEN = "open"  # Blocking requests due to failures
+    HALF_OPEN = "half_open"  # Testing recovery with limited requests
+
+
+class CircuitBreakerConfig:
+    """Configuration for circuit breaker behavior."""
+
+    def __init__(
+        self,
+        failure_threshold: int = 5,
+        success_threshold: int = 2,
+        reset_timeout: float = 60.0,
+        half_open_max_calls: int = 1,
+        timeout_threshold: float | None = None,
+    ):
+        """
+        Initialize circuit breaker configuration.
+
+        Args:
+            failure_threshold: Number of failures before opening circuit
+            success_threshold: Number of successes in HALF_OPEN to close circuit
+            reset_timeout: Seconds to wait before trying HALF_OPEN
+            half_open_max_calls: Max concurrent calls in HALF_OPEN state
+            timeout_threshold: Optional timeout (s) to consider as failure
+        """
+        self.failure_threshold = failure_threshold
+        self.success_threshold = success_threshold
+        self.reset_timeout = reset_timeout
+        self.half_open_max_calls = half_open_max_calls
+        self.timeout_threshold = timeout_threshold
+
+
+class CircuitBreakerState:
+    """Per-tool circuit breaker state tracking."""
+
+    def __init__(self, config: CircuitBreakerConfig):
+        self.config = config
+        self.state = CircuitState.CLOSED
+        self.failure_count = 0
+        self.success_count = 0
+        self.last_failure_time: float | None = None
+        self.opened_at: float | None = None
+        self.half_open_calls = 0
+        self._lock = asyncio.Lock()
+
+    async def record_success(self) -> None:
+        """Record a successful call."""
+        async with self._lock:
+            if self.state == CircuitState.HALF_OPEN:
+                self.success_count += 1
+                logger.debug(f"Circuit HALF_OPEN: success {self.success_count}/{self.config.success_threshold}")
+
+                # Enough successes? Close the circuit
+                if self.success_count >= self.config.success_threshold:
+                    logger.info("Circuit breaker: Transitioning to CLOSED (service recovered)")
+                    self.state = CircuitState.CLOSED
+                    self.failure_count = 0
+                    self.success_count = 0
+                    self.opened_at = None
+                    self.half_open_calls = 0
+            else:
+                # In CLOSED state, just reset failure count
+                self.failure_count = 0
+
+    async def record_failure(self) -> None:
+        """Record a failed call."""
+        async with self._lock:
+            self.failure_count += 1
+            self.last_failure_time = time.monotonic()
+            logger.debug(f"Circuit: failure {self.failure_count}/{self.config.failure_threshold}")
+
+            if self.state == CircuitState.CLOSED:
+                # Check if we should open
+                if self.failure_count >= self.config.failure_threshold:
+                    logger.warning(f"Circuit breaker: OPENING after {self.failure_count} failures")
+                    self.state = CircuitState.OPEN
+                    self.opened_at = time.monotonic()
+            elif self.state == CircuitState.HALF_OPEN:
+                # Failed during test → back to OPEN
+                logger.warning("Circuit breaker: Back to OPEN (test failed)")
+                self.state = CircuitState.OPEN
+                self.success_count = 0
+                self.opened_at = time.monotonic()
+                self.half_open_calls = 0
+
+    async def can_execute(self) -> bool:
+        """Check if a call should be allowed through."""
+        async with self._lock:
+            if self.state == CircuitState.CLOSED:
+                return True
+
+            if self.state == CircuitState.HALF_OPEN:
+                # Limit concurrent calls in HALF_OPEN
+                if self.half_open_calls < self.config.half_open_max_calls:
+                    self.half_open_calls += 1
+                    return True
+                return False
+
+            # OPEN state: check if we should try HALF_OPEN
+            if self.opened_at is not None:
+                elapsed = time.monotonic() - self.opened_at
+                if elapsed >= self.config.reset_timeout:
+                    logger.info("Circuit breaker: Transitioning to HALF_OPEN (testing recovery)")
+                    self.state = CircuitState.HALF_OPEN
+                    self.half_open_calls = 1
+                    self.success_count = 0
+                    return True
+
+            return False
+
+    async def release_half_open_slot(self) -> None:
+        """Release a HALF_OPEN slot after call completes."""
+        async with self._lock:
+            if self.state == CircuitState.HALF_OPEN:
+                self.half_open_calls = max(0, self.half_open_calls - 1)
+
+    def get_state(self) -> dict[str, Any]:
+        """Get current state as dict."""
+        return {
+            "state": self.state.value,
+            "failure_count": self.failure_count,
+            "success_count": self.success_count,
+            "opened_at": self.opened_at,
+            "time_until_half_open": (
+                max(0, self.config.reset_timeout - (time.monotonic() - self.opened_at))
+                if self.opened_at and self.state == CircuitState.OPEN
+                else None
+            ),
+        }
+
+
+# --------------------------------------------------------------------------- #
+# Circuit breaker executor wrapper
+# --------------------------------------------------------------------------- #
+class CircuitBreakerExecutor:
+    """
+    Executor wrapper that implements circuit breaker pattern.
+
+    Tracks failures per tool and opens circuit breakers to prevent
+    cascading failures when tools are consistently failing.
+    """
+
+    def __init__(
+        self,
+        executor: Any,
+        *,
+        default_config: CircuitBreakerConfig | None = None,
+        tool_configs: dict[str, CircuitBreakerConfig] | None = None,
+    ):
+        """
+        Initialize circuit breaker executor.
+
+        Args:
+            executor: Underlying executor to wrap
+            default_config: Default circuit breaker configuration
+            tool_configs: Per-tool circuit breaker configurations
+        """
+        self.executor = executor
+        self.default_config = default_config or CircuitBreakerConfig()
+        self.tool_configs = tool_configs or {}
+        self._states: dict[str, CircuitBreakerState] = {}
+        self._states_lock = asyncio.Lock()
+
+    async def _get_state(self, tool: str) -> CircuitBreakerState:
+        """Get or create circuit breaker state for a tool."""
+        if tool not in self._states:
+            async with self._states_lock:
+                if tool not in self._states:
+                    config = self.tool_configs.get(tool, self.default_config)
+                    self._states[tool] = CircuitBreakerState(config)
+        return self._states[tool]
+
+    async def execute(
+        self,
+        calls: list[ToolCall],
+        *,
+        timeout: float | None = None,
+        use_cache: bool = True,
+    ) -> list[ToolResult]:
+        """
+        Execute tool calls with circuit breaker protection.
+
+        Args:
+            calls: List of tool calls to execute
+            timeout: Optional timeout for execution
+            use_cache: Whether to use cached results
+
+        Returns:
+            List of tool results
+        """
+        if not calls:
+            return []
+
+        results: list[ToolResult] = []
+
+        for call in calls:
+            state = await self._get_state(call.tool)
+
+            # Record circuit breaker state
+            metrics = get_metrics()
+            if metrics:
+                metrics.record_circuit_breaker_state(call.tool, state.state.value)
+
+            # Check if circuit allows execution with tracing
+            with trace_circuit_breaker(call.tool, state.state.value):
+                can_execute = await state.can_execute()
+
+            if not can_execute:
+                # Circuit is OPEN - reject immediately
+                state_info = state.get_state()
+                logger.warning(f"Circuit breaker OPEN for {call.tool} (failures: {state.failure_count})")
+
+                reset_time = state_info.get("time_until_half_open")
+                error = ToolCircuitOpenError(
+                    tool_name=call.tool,
+                    failure_count=state.failure_count,
+                    reset_timeout=reset_time,
+                )
+
+                now = datetime.now(UTC)
+                results.append(
+                    ToolResult(
+                        tool=call.tool,
+                        result=None,
+                        error=str(error),
+                        start_time=now,
+                        end_time=now,
+                        machine="circuit_breaker",
+                        pid=0,
+                    )
+                )
+                continue
+
+            # Execute the call
+            start_time = time.monotonic()
+            try:
+                # Execute single call
+                executor_kwargs = {"timeout": timeout}
+                if hasattr(self.executor, "use_cache"):
+                    executor_kwargs["use_cache"] = use_cache
+
+                result_list = await self.executor.execute([call], **executor_kwargs)
+                result = result_list[0]
+
+                # Check if successful
+                duration = time.monotonic() - start_time
+
+                # Determine success/failure
+                is_timeout = state.config.timeout_threshold is not None and duration > state.config.timeout_threshold
+                is_error = result.error is not None
+
+                if is_error or is_timeout:
+                    await state.record_failure()
+                    # Record circuit breaker failure metric
+                    if metrics:
+                        metrics.record_circuit_breaker_failure(call.tool)
+                else:
+                    await state.record_success()
+
+                results.append(result)
+
+            except Exception as e:
+                # Exception during execution
+                await state.record_failure()
+
+                now = datetime.now(UTC)
+                results.append(
+                    ToolResult(
+                        tool=call.tool,
+                        result=None,
+                        error=f"Circuit breaker caught exception: {str(e)}",
+                        start_time=now,
+                        end_time=now,
+                        machine="circuit_breaker",
+                        pid=0,
+                    )
+                )
+
+            finally:
+                # Release HALF_OPEN slot if applicable
+                if state.state == CircuitState.HALF_OPEN:
+                    await state.release_half_open_slot()
+
+        return results
+
+    async def get_circuit_states(self) -> dict[str, dict[str, Any]]:
+        """
+        Get current state of all circuit breakers.
+
+        Returns:
+            Dict mapping tool name to state info
+        """
+        states = {}
+        async with self._states_lock:
+            for tool, state in self._states.items():
+                states[tool] = state.get_state()
+        return states
+
+    async def reset_circuit(self, tool: str) -> None:
+        """
+        Manually reset a circuit breaker.
+
+        Args:
+            tool: Tool name to reset
+        """
+        if tool in self._states:
+            state = self._states[tool]
+            async with state._lock:
+                state.state = CircuitState.CLOSED
+                state.failure_count = 0
+                state.success_count = 0
+                state.opened_at = None
+                state.half_open_calls = 0
+            logger.info(f"Manually reset circuit breaker for {tool}")

chuk_tool_processor/execution/wrappers/rate_limiting.py

@@ -7,44 +7,64 @@ Two layers of limits are enforced:
 * **Global** - ``<N requests> / <period>`` over *all* tools.
 * **Per-tool** - independent ``<N requests> / <period>`` windows.
 
-A simple sliding-window algorithm with timestamp queues is used.  
+A simple sliding-window algorithm with timestamp queues is used.
 `asyncio.Lock` guards shared state so the wrapper can be used safely from
 multiple coroutines.
 """
+
 from __future__ import annotations
 
 import asyncio
 import inspect
 import time
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
+from chuk_tool_processor.logging import get_logger
 from chuk_tool_processor.models.tool_call import ToolCall
 from chuk_tool_processor.models.tool_result import ToolResult
-from chuk_tool_processor.logging import get_logger
 
 logger = get_logger("chuk_tool_processor.execution.wrappers.rate_limiting")
 
+# Optional observability imports
+try:
+    from chuk_tool_processor.observability.metrics import get_metrics
+    from chuk_tool_processor.observability.tracing import trace_rate_limit
+
+    _observability_available = True
+except ImportError:
+    _observability_available = False
+
+    # No-op functions when observability not available
+    def get_metrics():
+        return None
+
+    def trace_rate_limit(*_args, **_kwargs):
+        from contextlib import nullcontext
+
+        return nullcontext()
+
+
 # --------------------------------------------------------------------------- #
 # Core limiter
 # --------------------------------------------------------------------------- #
 class RateLimiter:
     """
     Async-native rate limiter for controlling execution frequency.
-    
+
     Implements a sliding window algorithm to enforce rate limits both globally
     and per-tool. All operations are thread-safe using asyncio locks.
     """
-    
+
     def __init__(
         self,
         *,
-        global_limit: Optional[int] = None,
+        global_limit: int | None = None,
         global_period: float = 60.0,
-        tool_limits: Optional[Dict[str, Tuple[int, float]]] = None,
+        tool_limits: dict[str, tuple[int, float]] | None = None,
     ) -> None:
         """
         Initialize the rate limiter.
-        
+
         Args:
             global_limit: Maximum global requests per period (None = no limit)
             global_period: Time period in seconds for the global limit
@@ -55,13 +75,13 @@ class RateLimiter:
         self.tool_limits = tool_limits or {}
 
         # Timestamp queues
-        self._global_ts: List[float] = []
-        self._tool_ts: Dict[str, List[float]] = {}
+        self._global_ts: list[float] = []
+        self._tool_ts: dict[str, list[float]] = {}
 
         # Locks for thread safety
         self._global_lock = asyncio.Lock()
-        self._tool_locks: Dict[str, asyncio.Lock] = {}
-        
+        self._tool_locks: dict[str, asyncio.Lock] = {}
+
         logger.debug(
             f"Initialized rate limiter: global={global_limit}/{global_period}s, "
             f"tool-specific={len(self.tool_limits)} tools"
@@ -77,7 +97,7 @@ class RateLimiter:
             async with self._global_lock:
                 now = time.monotonic()
                 cutoff = now - self.global_period
-                
+
                 # Prune expired timestamps
                 self._global_ts = [t for t in self._global_ts if t > cutoff]
 
@@ -88,7 +108,7 @@ class RateLimiter:
 
                 # Calculate wait time until a slot becomes available
                 wait = (self._global_ts[0] + self.global_period) - now
-                
+
             logger.debug(f"Global rate limit reached, waiting {wait:.2f}s")
             await asyncio.sleep(wait)
 
@@ -105,7 +125,7 @@ class RateLimiter:
             async with lock:
                 now = time.monotonic()
                 cutoff = now - period
-                
+
                 # Prune expired timestamps in-place
                 buf[:] = [t for t in buf if t > cutoff]
 
@@ -116,7 +136,7 @@ class RateLimiter:
 
                 # Calculate wait time until a slot becomes available
                 wait = (buf[0] + period) - now
-                
+
             logger.debug(f"Tool '{tool}' rate limit reached, waiting {wait:.2f}s")
             await asyncio.sleep(wait)
 
@@ -124,32 +144,32 @@ class RateLimiter:
     async def wait(self, tool: str) -> None:
         """
         Block until rate limits allow execution.
-        
+
         This method blocks until both global and tool-specific rate limits
         allow one more execution of the specified tool.
-        
+
         Args:
             tool: Name of the tool being executed
         """
         await self._acquire_global()
         await self._acquire_tool(tool)
-        
-    async def check_limits(self, tool: str) -> Tuple[bool, bool]:
+
+    async def check_limits(self, tool: str) -> tuple[bool, bool]:
         """
         Check if the tool would be rate limited without consuming a slot.
-        
+
         This is a non-blocking method useful for checking limits without
         affecting the rate limiting state.
-        
+
         Args:
             tool: Name of the tool to check
-            
+
         Returns:
             Tuple of (global_limit_reached, tool_limit_reached)
         """
         global_limited = False
         tool_limited = False
-        
+
         # Check global limit
         if self.global_limit is not None:
             async with self._global_lock:
@@ -157,7 +177,7 @@ class RateLimiter:
                 cutoff = now - self.global_period
                 active_ts = [t for t in self._global_ts if t > cutoff]
                 global_limited = len(active_ts) >= self.global_limit
-        
+
         # Check tool limit
         if tool in self.tool_limits:
             limit, period = self.tool_limits[tool]
@@ -167,7 +187,7 @@ class RateLimiter:
                 buf = self._tool_ts.get(tool, [])
                 active_ts = [t for t in buf if t > cutoff]
                 tool_limited = len(active_ts) >= limit
-                
+
         return global_limited, tool_limited
 
 
@@ -177,7 +197,7 @@ class RateLimiter:
 class RateLimitedToolExecutor:
     """
     Executor wrapper that applies rate limiting to tool executions.
-    
+
     This wrapper delegates to another executor but ensures that all
     tool calls respect the configured rate limits.
     """
@@ -185,48 +205,60 @@ class RateLimitedToolExecutor:
     def __init__(self, executor: Any, limiter: RateLimiter) -> None:
         """
         Initialize the rate-limited executor.
-        
+
         Args:
             executor: The underlying executor to wrap
             limiter: The RateLimiter that controls execution frequency
         """
         self.executor = executor
         self.limiter = limiter
-        logger.debug(f"Initialized rate-limited executor")
+        logger.debug("Initialized rate-limited executor")
 
     async def execute(
         self,
-        calls: List[ToolCall],
-        timeout: Optional[float] = None,
+        calls: list[ToolCall],
+        timeout: float | None = None,
         use_cache: bool = True,
-    ) -> List[ToolResult]:
+    ) -> list[ToolResult]:
         """
         Execute tool calls while respecting rate limits.
-        
+
         This method blocks until rate limits allow execution, then delegates
         to the underlying executor.
-        
+
         Args:
             calls: List of tool calls to execute
             timeout: Optional timeout for execution
             use_cache: Whether to use cached results (forwarded to underlying executor)
-            
+
         Returns:
             List of tool results
         """
         if not calls:
             return []
-            
+
         # Block for each call *before* dispatching to the wrapped executor
+        metrics = get_metrics()
+
         for c in calls:
-            await self.limiter.wait(c.tool)
-            
+            # Check limits first for metrics
+            global_limited, tool_limited = await self.limiter.check_limits(c.tool)
+            allowed = not (global_limited or tool_limited)
+
+            # Trace rate limit check
+            with trace_rate_limit(c.tool, allowed):
+                await self.limiter.wait(c.tool)
+
+            # Record metrics
+            if metrics:
+                metrics.record_rate_limit_check(c.tool, allowed)
+
         # Check if the executor has a use_cache parameter
         if hasattr(self.executor, "execute"):
             sig = inspect.signature(self.executor.execute)
             if "use_cache" in sig.parameters:
                 return await self.executor.execute(calls, timeout=timeout, use_cache=use_cache)
-                
+
         # Fall back to standard execute method
         return await self.executor.execute(calls, timeout=timeout)
 
@@ -237,7 +269,7 @@ class RateLimitedToolExecutor:
 def rate_limited(limit: int, period: float = 60.0):
     """
     Class decorator that marks a Tool with default rate-limit metadata.
-    
+
     This allows higher-level code to detect and configure rate limiting
     for the tool class.
 
@@ -246,17 +278,18 @@ def rate_limited(limit: int, period: float = 60.0):
         class WeatherTool:
             async def execute(self, location: str) -> Dict[str, Any]:
                 # Implementation
-                
+
     Args:
         limit: Maximum number of calls allowed in the period
         period: Time period in seconds
-        
+
     Returns:
         Decorated class with rate limit metadata
     """
+
     def decorator(cls):
         cls._rate_limit = limit
         cls._rate_period = period
         return cls
 
-    return decorator
+    return decorator