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

chuk_tool_processor/mcp/transport/sse_transport.py

@@ -19,6 +19,7 @@ from typing import Any
 import httpx
 
 from .base_transport import MCPBaseTransport
+from .models import TimeoutConfig, TransportMetrics
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +39,8 @@ class SSETransport(MCPBaseTransport):
         connection_timeout: float = 30.0,
         default_timeout: float = 60.0,
         enable_metrics: bool = True,
+        oauth_refresh_callback: Any | None = None,
+        timeout_config: TimeoutConfig | None = None,
     ):
         """
         Initialize SSE transport.
@@ -45,9 +48,16 @@ class SSETransport(MCPBaseTransport):
         self.url = url.rstrip("/")
         self.api_key = api_key
         self.configured_headers = headers or {}
-        self.connection_timeout = connection_timeout
-        self.default_timeout = default_timeout
         self.enable_metrics = enable_metrics
+        self.oauth_refresh_callback = oauth_refresh_callback
+
+        # Use timeout config or create from individual parameters
+        if timeout_config is None:
+            timeout_config = TimeoutConfig(connect=connection_timeout, operation=default_timeout)
+
+        self.timeout_config = timeout_config
+        self.connection_timeout = timeout_config.connect
+        self.default_timeout = timeout_config.operation
 
         logger.debug("SSE Transport initialized with URL: %s", self.url)
 
@@ -73,18 +83,8 @@ class SSETransport(MCPBaseTransport):
         self._connection_grace_period = 30.0  # NEW: Grace period after initialization
         self._initialization_time = None  # NEW: Track when we initialized
 
-        # Performance metrics
-        self._metrics = {
-            "total_calls": 0,
-            "successful_calls": 0,
-            "failed_calls": 0,
-            "total_time": 0.0,
-            "avg_response_time": 0.0,
-            "last_ping_time": None,
-            "initialization_time": None,
-            "session_discoveries": 0,
-            "stream_errors": 0,
-        }
+        # Performance metrics - use Pydantic model
+        self._metrics = TransportMetrics() if enable_metrics else None
 
     def _construct_sse_url(self, base_url: str) -> str:
         """Construct the SSE endpoint URL from the base URL."""
@@ -110,8 +110,9 @@ class SSETransport(MCPBaseTransport):
         if self.configured_headers:
             headers.update(self.configured_headers)
 
-        # Add API key as Bearer token if provided
-        if self.api_key:
+        # Add API key as Bearer token if provided and no Authorization header exists
+        # This prevents clobbering OAuth tokens from configured_headers
+        if self.api_key and "Authorization" not in headers:
             headers["Authorization"] = f"Bearer {self.api_key}"
 
         return headers
@@ -173,7 +174,7 @@ class SSETransport(MCPBaseTransport):
 
             # Wait for session discovery
             logger.debug("Waiting for session discovery...")
-            session_timeout = 10.0
+            session_timeout = self.timeout_config.connect
             session_start = time.time()
 
             while not self.message_url and (time.time() - session_start) < session_timeout:
@@ -183,17 +184,17 @@ class SSETransport(MCPBaseTransport):
                 if self.sse_task.done():
                     exception = self.sse_task.exception()
                     if exception:
-                        logger.error(f"SSE task died during session discovery: {exception}")
+                        logger.debug(f"SSE task died during session discovery: {exception}")
                         await self._cleanup()
                         return False
 
             if not self.message_url:
-                logger.error("Failed to discover session endpoint within %.1fs", session_timeout)
+                logger.warning("Failed to discover session endpoint within %.1fs", session_timeout)
                 await self._cleanup()
                 return False
 
-            if self.enable_metrics:
-                self._metrics["session_discoveries"] += 1
+            if self.enable_metrics and self._metrics:
+                self._metrics.session_discoveries += 1
 
             logger.debug("Session endpoint discovered: %s", self.message_url)
 
@@ -210,7 +211,7 @@ class SSETransport(MCPBaseTransport):
                 )
 
                 if "error" in init_response:
-                    logger.error("MCP initialize failed: %s", init_response["error"])
+                    logger.warning("MCP initialize failed: %s", init_response["error"])
                     await self._cleanup()
                     return False
 
@@ -223,9 +224,9 @@ class SSETransport(MCPBaseTransport):
                 self._last_successful_ping = time.time()
                 self._consecutive_failures = 0  # Reset failure count
 
-                if self.enable_metrics:
+                if self.enable_metrics and self._metrics:
                     init_time = time.time() - start_time
-                    self._metrics["initialization_time"] = init_time
+                    self._metrics.initialization_time = init_time
 
                 logger.debug("SSE transport initialized successfully in %.3fs", time.time() - start_time)
                 return True
@@ -269,12 +270,30 @@ class SSETransport(MCPBaseTransport):
                         # Extract session ID from URL if present
                         if "session_id=" in data_part:
                             self.session_id = data_part.split("session_id=")[1].split("&")[0]
+                        elif "sessionId=" in data_part:
+                            self.session_id = data_part.split("sessionId=")[1].split("&")[0]
                         else:
                             self.session_id = str(uuid.uuid4())
 
                         logger.debug("Session endpoint discovered via event format: %s", self.message_url)
                         continue
 
+                    # RELATIVE PATH FORMAT: event: endpoint + data: /sse/message?sessionId=...
+                    elif current_event == "endpoint" and data_part.startswith("/"):
+                        endpoint_path = data_part
+                        self.message_url = f"{self.url}{endpoint_path}"
+
+                        # Extract session ID if present
+                        if "session_id=" in endpoint_path:
+                            self.session_id = endpoint_path.split("session_id=")[1].split("&")[0]
+                        elif "sessionId=" in endpoint_path:
+                            self.session_id = endpoint_path.split("sessionId=")[1].split("&")[0]
+                        else:
+                            self.session_id = str(uuid.uuid4())
+
+                        logger.debug("Session endpoint discovered via relative path: %s", self.message_url)
+                        continue
+
                     # OLD FORMAT: data: /messages/... (backwards compatibility)
                     elif "/messages/" in data_part:
                         endpoint_path = data_part
@@ -315,8 +334,8 @@ class SSETransport(MCPBaseTransport):
                         logger.debug("Non-JSON data in SSE stream (ignoring): %s", e)
 
         except Exception as e:
-            if self.enable_metrics:
-                self._metrics["stream_errors"] += 1
+            if self.enable_metrics and self._metrics:
+                self._metrics.stream_errors += 1
             logger.error("SSE stream processing error: %s", e)
             # FIXED: Don't increment consecutive failures for stream processing errors
             # These are often temporary and don't indicate connection health
@@ -400,7 +419,7 @@ class SSETransport(MCPBaseTransport):
         start_time = time.time()
         try:
             # Use tools/list as a lightweight ping since not all servers support ping
-            response = await self._send_request("tools/list", {}, timeout=10.0)
+            response = await self._send_request("tools/list", {}, timeout=self.timeout_config.quick)
 
             success = "error" not in response
 
@@ -408,9 +427,9 @@ class SSETransport(MCPBaseTransport):
                 self._last_successful_ping = time.time()
                 # FIXED: Don't reset consecutive failures here - let tool calls do that
 
-            if self.enable_metrics:
+            if self.enable_metrics and self._metrics:
                 ping_time = time.time() - start_time
-                self._metrics["last_ping_time"] = ping_time
+                self._metrics.last_ping_time = ping_time
                 logger.debug("SSE ping completed in %.3fs: %s", ping_time, success)
 
             return success
@@ -457,7 +476,7 @@ class SSETransport(MCPBaseTransport):
     async def get_tools(self) -> list[dict[str, Any]]:
         """Get list of available tools from the server."""
         if not self._initialized:
-            logger.error("Cannot get tools: transport not initialized")
+            logger.debug("Cannot get tools: transport not initialized")
             return []
 
         start_time = time.time()
@@ -465,7 +484,7 @@ class SSETransport(MCPBaseTransport):
             response = await self._send_request("tools/list", {})
 
             if "error" in response:
-                logger.error("Error getting tools: %s", response["error"])
+                logger.warning("Error getting tools: %s", response["error"])
                 return []
 
             tools = response.get("result", {}).get("tools", [])
@@ -488,8 +507,8 @@ class SSETransport(MCPBaseTransport):
             return {"isError": True, "error": "Transport not initialized"}
 
         start_time = time.time()
-        if self.enable_metrics:
-            self._metrics["total_calls"] += 1
+        if self.enable_metrics and self._metrics:
+            self._metrics.total_calls += 1
 
         try:
             logger.debug("Calling tool '%s' with arguments: %s", tool_name, arguments)
@@ -498,11 +517,55 @@ class SSETransport(MCPBaseTransport):
                 "tools/call", {"name": tool_name, "arguments": arguments}, timeout=timeout
             )
 
+            # Check for errors
             if "error" in response:
+                error_msg = response["error"].get("message", "Unknown error")
+
+                # NEW: Check for OAuth errors and attempt refresh if callback is available
+                if self._is_oauth_error(error_msg):
+                    logger.warning("OAuth error detected: %s", error_msg)
+
+                    if self.oauth_refresh_callback:
+                        logger.debug("Attempting OAuth token refresh...")
+                        try:
+                            # Call the refresh callback
+                            new_headers = await self.oauth_refresh_callback()
+
+                            if new_headers and "Authorization" in new_headers:
+                                # Update configured headers with new token
+                                self.configured_headers.update(new_headers)
+                                logger.debug("OAuth token refreshed, retrying tool call...")
+
+                                # Retry the tool call once with new token
+                                response = await self._send_request(
+                                    "tools/call", {"name": tool_name, "arguments": arguments}, timeout=timeout
+                                )
+
+                                # Check if retry succeeded
+                                if "error" not in response:
+                                    logger.debug("Tool call succeeded after token refresh")
+                                    result = response.get("result", {})
+                                    normalized_result = self._normalize_mcp_response({"result": result})
+
+                                    if self.enable_metrics:
+                                        self._update_metrics(time.time() - start_time, True)
+
+                                    return normalized_result
+                                else:
+                                    error_msg = response["error"].get("message", "Unknown error")
+                                    logger.error("Tool call failed after token refresh: %s", error_msg)
+                            else:
+                                logger.warning("Token refresh did not return valid Authorization header")
+                        except Exception as refresh_error:
+                            logger.error("OAuth token refresh failed: %s", refresh_error)
+                    else:
+                        logger.warning("OAuth error detected but no refresh callback configured")
+
+                # Return error (original or from failed retry)
                 if self.enable_metrics:
                     self._update_metrics(time.time() - start_time, False)
 
-                return {"isError": True, "error": response["error"].get("message", "Unknown error")}
+                return {"isError": True, "error": error_msg}
 
             # Extract and normalize result using base class method
             result = response.get("result", {})
@@ -527,14 +590,41 @@ class SSETransport(MCPBaseTransport):
 
     def _update_metrics(self, response_time: float, success: bool) -> None:
         """Update performance metrics."""
-        if success:
-            self._metrics["successful_calls"] += 1
-        else:
-            self._metrics["failed_calls"] += 1
+        if not self._metrics:
+            return
+
+        self._metrics.update_call_metrics(response_time, success)
+
+    def _is_oauth_error(self, error_msg: str) -> bool:
+        """
+        Detect if error is OAuth-related per RFC 6750 and MCP OAuth spec.
+
+        Checks for:
+        - RFC 6750 Section 3.1 Bearer token errors (invalid_token, insufficient_scope)
+        - OAuth 2.1 token refresh errors (invalid_grant)
+        - MCP spec OAuth validation failures (401/403 responses)
+        """
+        if not error_msg:
+            return False
 
-        self._metrics["total_time"] += response_time
-        if self._metrics["total_calls"] > 0:
-            self._metrics["avg_response_time"] = self._metrics["total_time"] / self._metrics["total_calls"]
+        error_lower = error_msg.lower()
+        oauth_indicators = [
+            # RFC 6750 Section 3.1 - Standard Bearer token errors
+            "invalid_token",  # Token expired, revoked, malformed, or invalid
+            "insufficient_scope",  # Request requires higher privileges (403 Forbidden)
+            # OAuth 2.1 token refresh errors
+            "invalid_grant",  # Refresh token errors
+            # MCP spec - OAuth validation failures (401 Unauthorized)
+            "oauth validation",
+            "unauthorized",
+            # Common OAuth error descriptions
+            "expired token",
+            "token expired",
+            "authentication failed",
+            "invalid access token",
+        ]
+
+        return any(indicator in error_lower for indicator in oauth_indicators)
 
     async def list_resources(self) -> dict[str, Any]:
         """List available resources from the server."""
@@ -542,7 +632,7 @@ class SSETransport(MCPBaseTransport):
             return {}
 
         try:
-            response = await self._send_request("resources/list", {}, timeout=10.0)
+            response = await self._send_request("resources/list", {}, timeout=self.timeout_config.operation)
             if "error" in response:
                 logger.debug("Resources not supported: %s", response["error"])
                 return {}
@@ -557,7 +647,7 @@ class SSETransport(MCPBaseTransport):
             return {}
 
         try:
-            response = await self._send_request("prompts/list", {}, timeout=10.0)
+            response = await self._send_request("prompts/list", {}, timeout=self.timeout_config.operation)
             if "error" in response:
                 logger.debug("Prompts not supported: %s", response["error"])
                 return {}
@@ -572,12 +662,12 @@ class SSETransport(MCPBaseTransport):
             return
 
         # Log final metrics
-        if self.enable_metrics and self._metrics["total_calls"] > 0:
+        if self.enable_metrics and self._metrics and self._metrics.total_calls > 0:
             logger.debug(
                 "SSE transport closing - Total calls: %d, Success rate: %.1f%%, Avg response time: %.3fs",
-                self._metrics["total_calls"],
-                (self._metrics["successful_calls"] / self._metrics["total_calls"] * 100),
-                self._metrics["avg_response_time"],
+                self._metrics.total_calls,
+                (self._metrics.successful_calls / self._metrics.total_calls * 100),
+                self._metrics.avg_response_time,
             )
 
         await self._cleanup()
@@ -626,7 +716,10 @@ class SSETransport(MCPBaseTransport):
 
     def get_metrics(self) -> dict[str, Any]:
         """Get performance and connection metrics with health info."""
-        metrics = self._metrics.copy()
+        if not self._metrics:
+            return {}
+
+        metrics = self._metrics.to_dict()
         metrics.update(
             {
                 "is_connected": self.is_connected(),
@@ -646,17 +739,20 @@ class SSETransport(MCPBaseTransport):
 
     def reset_metrics(self) -> None:
         """Reset performance metrics."""
-        self._metrics = {
-            "total_calls": 0,
-            "successful_calls": 0,
-            "failed_calls": 0,
-            "total_time": 0.0,
-            "avg_response_time": 0.0,
-            "last_ping_time": self._metrics.get("last_ping_time"),
-            "initialization_time": self._metrics.get("initialization_time"),
-            "session_discoveries": self._metrics.get("session_discoveries", 0),
-            "stream_errors": 0,
-        }
+        if not self._metrics:
+            return
+
+        # Preserve important historical values
+        preserved_last_ping = self._metrics.last_ping_time
+        preserved_init_time = self._metrics.initialization_time
+        preserved_discoveries = self._metrics.session_discoveries
+
+        # Create new metrics instance with preserved values
+        self._metrics = TransportMetrics(
+            last_ping_time=preserved_last_ping,
+            initialization_time=preserved_init_time,
+            session_discoveries=preserved_discoveries,
+        )
 
     def get_streams(self) -> list[tuple]:
         """SSE transport doesn't expose raw streams."""

chuk_tool_processor/mcp/transport/stdio_transport.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
+import os
 import time
 from typing import Any
 
@@ -54,13 +55,28 @@ class StdioTransport(MCPBaseTransport):
         """
         # Convert dict to StdioParameters if needed
         if isinstance(server_params, dict):
+            # Merge provided env with system environment to ensure PATH is available
+            merged_env = os.environ.copy()
+            if server_params.get("env"):
+                merged_env.update(server_params["env"])
+
             self.server_params = StdioParameters(
                 command=server_params.get("command", "python"),
                 args=server_params.get("args", []),
-                env=server_params.get("env"),
+                env=merged_env,
             )
         else:
-            self.server_params = server_params
+            # Also handle StdioParameters object - merge env if provided
+            # Create a new StdioParameters with merged env (Pydantic models are immutable)
+            merged_env = os.environ.copy()
+            if hasattr(server_params, "env") and server_params.env:
+                merged_env.update(server_params.env)
+
+            self.server_params = StdioParameters(
+                command=server_params.command,
+                args=server_params.args,
+                env=merged_env,
+            )
 
         self.connection_timeout = connection_timeout
         self.default_timeout = default_timeout
@@ -184,7 +200,8 @@ class StdioTransport(MCPBaseTransport):
                 # Enhanced health verification (like SSE)
                 logger.debug("Verifying connection with ping...")
                 ping_start = time.time()
-                ping_success = await asyncio.wait_for(send_ping(*self._streams), timeout=10.0)
+                # Use default timeout for initial ping verification
+                ping_success = await asyncio.wait_for(send_ping(*self._streams), timeout=self.default_timeout)
                 ping_time = time.time() - ping_start
 
                 if ping_success:
@@ -204,7 +221,7 @@ class StdioTransport(MCPBaseTransport):
                     )
                     return True
                 else:
-                    logger.warning("STDIO connection established but ping failed")
+                    logger.debug("STDIO connection established but ping failed")
                     # Still consider it initialized
                     self._initialized = True
                     self._consecutive_failures = 1
@@ -212,7 +229,7 @@ class StdioTransport(MCPBaseTransport):
                         self._metrics["initialization_time"] = time.time() - start_time
                     return True
             else:
-                logger.error("STDIO initialization failed")
+                logger.warning("STDIO initialization failed")
                 await self._cleanup()
                 return False
 
@@ -365,16 +382,47 @@ class StdioTransport(MCPBaseTransport):
     async def get_tools(self) -> list[dict[str, Any]]:
         """Enhanced tools retrieval with recovery."""
         if not self._initialized:
-            logger.error("Cannot get tools: transport not initialized")
+            logger.debug("Cannot get tools: transport not initialized")
             return []
 
         start_time = time.time()
         try:
             response = await asyncio.wait_for(send_tools_list(*self._streams), timeout=self.default_timeout)
 
-            # Normalize response
-            if isinstance(response, dict):
-                tools = response.get("tools", [])
+            # Normalize response - handle multiple formats including Pydantic models
+            # 1. Check if it's a Pydantic model with tools attribute (e.g., ListToolsResult from chuk_mcp)
+            if hasattr(response, "tools"):
+                tools = response.tools
+                # Convert Pydantic Tool models to dicts if needed
+                if tools and len(tools) > 0 and hasattr(tools[0], "model_dump"):
+                    tools = [tool.model_dump() if hasattr(tool, "model_dump") else tool for tool in tools]
+                elif tools and len(tools) > 0 and hasattr(tools[0], "dict"):
+                    # Older Pydantic versions use dict() instead of model_dump()
+                    tools = [tool.dict() if hasattr(tool, "dict") else tool for tool in tools]
+            # 2. Check if it's a Pydantic model that can be dumped
+            elif hasattr(response, "model_dump"):
+                dumped = response.model_dump()
+                tools = dumped.get("tools", [])
+            # 3. Handle dict responses
+            elif isinstance(response, dict):
+                # Check for tools at top level
+                if "tools" in response:
+                    tools = response["tools"]
+                # Check for nested result.tools (common in some MCP implementations)
+                elif "result" in response and isinstance(response["result"], dict):
+                    tools = response["result"].get("tools", [])
+                # Check if response itself is the result with MCP structure
+                elif "jsonrpc" in response and "result" in response:
+                    result = response["result"]
+                    if isinstance(result, dict):
+                        tools = result.get("tools", [])
+                    elif isinstance(result, list):
+                        tools = result
+                    else:
+                        tools = []
+                else:
+                    tools = []
+            # 4. Handle list responses
             elif isinstance(response, list):
                 tools = response
             else:
@@ -426,7 +474,7 @@ class StdioTransport(MCPBaseTransport):
                     return {"isError": True, "error": "Failed to recover connection"}
 
             response = await asyncio.wait_for(
-                send_tools_call(*self._streams, tool_name, arguments), timeout=tool_timeout
+                send_tools_call(*self._streams, tool_name, arguments, timeout=tool_timeout), timeout=tool_timeout
             )
 
             response_time = time.time() - start_time

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."""