chuk-tool-processor 0.7.0__py3-none-any.whl → 0.8__py3-none-any.whl

chuk_tool_processor/execution/wrappers/circuit_breaker.py

@@ -0,0 +1,343 @@
+# 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")
+
+
+# --------------------------------------------------------------------------- #
+# 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)
+
+            # Check if circuit allows execution
+            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()
+                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/retry.py

@@ -36,6 +36,7 @@ class RetryConfig:
         jitter: bool = True,
         retry_on_exceptions: list[type[Exception]] | None = None,
         retry_on_error_substrings: list[str] | None = None,
+        skip_retry_on_error_substrings: list[str] | None = None,
     ):
         if max_retries < 0:
             raise ValueError("max_retries cannot be negative")
@@ -45,6 +46,7 @@ class RetryConfig:
         self.jitter = jitter
         self.retry_on_exceptions = retry_on_exceptions or []
         self.retry_on_error_substrings = retry_on_error_substrings or []
+        self.skip_retry_on_error_substrings = skip_retry_on_error_substrings or []
 
     # --------------------------------------------------------------------- #
     # Decision helpers
@@ -60,6 +62,14 @@ class RetryConfig:
         if attempt >= self.max_retries:
             return False
 
+        # Check skip list first - these errors should never be retried
+        # (e.g., OAuth errors that need to be handled at transport layer)
+        if error_str and self.skip_retry_on_error_substrings:
+            error_lower = error_str.lower()
+            if any(skip_pattern.lower() in error_lower for skip_pattern in self.skip_retry_on_error_substrings):
+                logger.debug(f"Skipping retry for error matching skip pattern: {error_str[:100]}")
+                return False
+
         # Nothing specified → always retry until max_retries reached
         if not self.retry_on_exceptions and not self.retry_on_error_substrings:
             return True
@@ -246,6 +256,7 @@ def retryable(
     jitter: bool = True,
     retry_on_exceptions: list[type[Exception]] | None = None,
     retry_on_error_substrings: list[str] | None = None,
+    skip_retry_on_error_substrings: list[str] | None = None,
 ):
     """
     Class decorator that attaches a :class:`RetryConfig` to a *tool* class.
@@ -267,6 +278,7 @@ def retryable(
             jitter=jitter,
             retry_on_exceptions=retry_on_exceptions,
             retry_on_error_substrings=retry_on_error_substrings,
+            skip_retry_on_error_substrings=skip_retry_on_error_substrings,
         )
         return cls
 

chuk_tool_processor/mcp/setup_mcp_http_streamable.py

@@ -41,8 +41,8 @@ async def setup_mcp_http_streamable(
     enable_rate_limiting: bool = False,
     global_rate_limit: int | None = None,
     tool_rate_limits: dict[str, tuple] | None = None,
-    enable_retries: bool = False,  # CHANGED: Disabled to allow OAuth refresh to work properly
-    max_retries: int = 0,  # CHANGED: 0 retries for HTTP (OAuth refresh happens at transport level)
+    enable_retries: bool = True,  # CHANGED: Enabled with OAuth errors excluded
+    max_retries: int = 2,  # Retry non-OAuth errors (OAuth handled at transport level)
     namespace: str = "http",
     oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
 ) -> tuple[ToolProcessor, StreamManager]:
@@ -102,6 +102,34 @@ async def setup_mcp_http_streamable(
     registered = await register_mcp_tools(stream_manager, namespace=namespace)
 
     # 3️⃣  build a processor instance configured to your taste
+    # IMPORTANT: Retries are enabled but OAuth errors are excluded
+    # OAuth refresh happens at transport level with automatic retry
+
+    # Import RetryConfig to configure OAuth error exclusion
+    from chuk_tool_processor.execution.wrappers.retry import RetryConfig
+
+    # Define OAuth error patterns that should NOT be retried at this level
+    # These will be handled by the transport layer's OAuth refresh mechanism
+    oauth_error_patterns = [
+        "invalid_token",
+        "oauth validation",
+        "unauthorized",
+        "expired token",
+        "token expired",
+        "authentication failed",
+        "invalid access token",
+    ]
+
+    # Create retry config that skips OAuth errors
+    retry_config = (
+        RetryConfig(
+            max_retries=max_retries,
+            skip_retry_on_error_substrings=oauth_error_patterns,
+        )
+        if enable_retries
+        else None
+    )
+
     processor = ToolProcessor(
         default_timeout=default_timeout,
         max_concurrency=max_concurrency,
@@ -112,6 +140,7 @@ async def setup_mcp_http_streamable(
         tool_rate_limits=tool_rate_limits,
         enable_retries=enable_retries,
         max_retries=max_retries,
+        retry_config=retry_config,  # NEW: Pass OAuth-aware retry config
     )
 
     logger.debug(

chuk_tool_processor/mcp/setup_mcp_sse.py

@@ -37,8 +37,8 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
     enable_rate_limiting: bool = False,
     global_rate_limit: int | None = None,
     tool_rate_limits: dict[str, tuple] | None = None,
-    enable_retries: bool = False,  # CHANGED: Disabled to allow OAuth refresh to work properly
-    max_retries: int = 0,  # CHANGED: 0 retries for SSE (OAuth refresh happens at transport level)
+    enable_retries: bool = True,  # CHANGED: Enabled with OAuth errors excluded
+    max_retries: int = 2,  # Retry non-OAuth errors (OAuth handled at transport level)
     namespace: str = "sse",
     oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
 ) -> tuple[ToolProcessor, StreamManager]:
@@ -81,6 +81,34 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
     registered = await register_mcp_tools(stream_manager, namespace=namespace)
 
     # 3️⃣  build a processor instance configured to your taste
+    # IMPORTANT: Retries are enabled but OAuth errors are excluded
+    # OAuth refresh happens at transport level with automatic retry
+
+    # Import RetryConfig to configure OAuth error exclusion
+    from chuk_tool_processor.execution.wrappers.retry import RetryConfig
+
+    # Define OAuth error patterns that should NOT be retried at this level
+    # These will be handled by the transport layer's OAuth refresh mechanism
+    oauth_error_patterns = [
+        "invalid_token",
+        "oauth validation",
+        "unauthorized",
+        "expired token",
+        "token expired",
+        "authentication failed",
+        "invalid access token",
+    ]
+
+    # Create retry config that skips OAuth errors
+    retry_config = (
+        RetryConfig(
+            max_retries=max_retries,
+            skip_retry_on_error_substrings=oauth_error_patterns,
+        )
+        if enable_retries
+        else None
+    )
+
     processor = ToolProcessor(
         default_timeout=default_timeout,
         max_concurrency=max_concurrency,
@@ -91,6 +119,7 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
         tool_rate_limits=tool_rate_limits,
         enable_retries=enable_retries,
         max_retries=max_retries,
+        retry_config=retry_config,  # NEW: Pass OAuth-aware retry config
     )
 
     logger.debug(

chuk_tool_processor/models/__init__.py

@@ -1 +1,21 @@
 # chuk_tool_processor/models/__init__.py
+"""Data models for the tool processor."""
+
+from chuk_tool_processor.models.execution_strategy import ExecutionStrategy
+from chuk_tool_processor.models.streaming_tool import StreamingTool
+from chuk_tool_processor.models.tool_call import ToolCall
+from chuk_tool_processor.models.tool_result import ToolResult
+from chuk_tool_processor.models.tool_spec import ToolCapability, ToolSpec, tool_spec
+from chuk_tool_processor.models.validated_tool import ValidatedTool, with_validation
+
+__all__ = [
+    "ExecutionStrategy",
+    "StreamingTool",
+    "ToolCall",
+    "ToolResult",
+    "ToolSpec",
+    "ToolCapability",
+    "tool_spec",
+    "ValidatedTool",
+    "with_validation",
+]

chuk_tool_processor/models/tool_call.py

@@ -5,10 +5,12 @@ Model representing a tool call with arguments.
 
 from __future__ import annotations
 
+import hashlib
+import json
 import uuid
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 class ToolCall(BaseModel):
@@ -20,6 +22,7 @@ class ToolCall(BaseModel):
         tool: Name of the tool to call
         namespace: Namespace the tool belongs to
         arguments: Arguments to pass to the tool
+        idempotency_key: Optional key for deduplicating duplicate calls (auto-generated)
     """
 
     model_config = ConfigDict(extra="ignore")
@@ -28,6 +31,36 @@ class ToolCall(BaseModel):
     tool: str = Field(..., min_length=1, description="Name of the tool to call; must be non-empty")
     namespace: str = Field(default="default", description="Namespace the tool belongs to")
     arguments: dict[str, Any] = Field(default_factory=dict, description="Arguments to pass to the tool")
+    idempotency_key: str | None = Field(
+        None,
+        description="Idempotency key for deduplication. Auto-generated if not provided.",
+    )
+
+    @model_validator(mode="after")
+    def generate_idempotency_key(self) -> ToolCall:
+        """Generate idempotency key if not provided."""
+        if self.idempotency_key is None:
+            self.idempotency_key = self._compute_idempotency_key()
+        return self
+
+    def _compute_idempotency_key(self) -> str:
+        """
+        Compute a stable idempotency key from tool name, namespace, and arguments.
+
+        Uses SHA256 hash of the sorted JSON representation.
+        Returns first 16 characters of the hex digest for brevity.
+        """
+        # Create a stable representation
+        payload = {
+            "tool": self.tool,
+            "namespace": self.namespace,
+            "arguments": self.arguments,
+        }
+        # Sort keys for stability
+        json_str = json.dumps(payload, sort_keys=True, default=str)
+        # Hash it
+        hash_obj = hashlib.sha256(json_str.encode(), usedforsecurity=False)
+        return hash_obj.hexdigest()[:16]  # Use first 16 chars for brevity
 
     async def to_dict(self) -> dict[str, Any]:
         """Convert to a dictionary for serialization."""