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

chuk_tool_processor/mcp/setup_mcp_http_streamable.py

@@ -34,15 +34,17 @@ async def setup_mcp_http_streamable(
     server_names: dict[int, str] | None = None,
     connection_timeout: float = 30.0,
     default_timeout: float = 30.0,
+    initialization_timeout: float = 60.0,
     max_concurrency: int | None = None,
     enable_caching: bool = True,
     cache_ttl: int = 300,
     enable_rate_limiting: bool = False,
     global_rate_limit: int | None = None,
     tool_rate_limits: dict[str, tuple] | None = None,
-    enable_retries: bool = True,
-    max_retries: int = 3,
+    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]:
     """
     Initialize HTTP Streamable transport MCP + a :class:`ToolProcessor`.
@@ -58,6 +60,7 @@ async def setup_mcp_http_streamable(
         server_names: Optional mapping of server indices to names
         connection_timeout: Timeout for initial HTTP connection setup
         default_timeout: Default timeout for tool execution
+        initialization_timeout: Timeout for complete initialization (default 60s, increase to 120s+ for slow servers like Notion)
         max_concurrency: Maximum concurrent operations
         enable_caching: Whether to enable response caching
         cache_ttl: Cache time-to-live in seconds
@@ -67,6 +70,7 @@ async def setup_mcp_http_streamable(
         enable_retries: Whether to enable automatic retries
         max_retries: Maximum retry attempts
         namespace: Namespace for registered tools
+        oauth_refresh_callback: Optional async callback to refresh OAuth tokens (NEW)
 
     Returns:
         Tuple of (ToolProcessor, StreamManager)
@@ -90,12 +94,49 @@ async def setup_mcp_http_streamable(
         server_names=server_names,
         connection_timeout=connection_timeout,
         default_timeout=default_timeout,
+        initialization_timeout=initialization_timeout,
+        oauth_refresh_callback=oauth_refresh_callback,  # NEW: Pass OAuth callback
     )
 
     # 2️⃣  pull the remote tool list and register each one locally
     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
+    # Based on RFC 6750 (Bearer Token Usage) and MCP OAuth spec
+    oauth_error_patterns = [
+        # 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",
+    ]
+
+    # 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,
@@ -106,6 +147,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

@@ -30,15 +30,17 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
     server_names: dict[int, str] | None = None,
     connection_timeout: float = 30.0,  # 🔧 INCREASED DEFAULT: was 10.0
     default_timeout: float = 30.0,  # 🔧 INCREASED DEFAULT: was 10.0
+    initialization_timeout: float = 60.0,
     max_concurrency: int | None = None,
     enable_caching: bool = True,
     cache_ttl: int = 300,
     enable_rate_limiting: bool = False,
     global_rate_limit: int | None = None,
     tool_rate_limits: dict[str, tuple] | None = None,
-    enable_retries: bool = True,
-    max_retries: int = 3,
+    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]:
     """
     Initialise SSE-transport MCP + a :class:`ToolProcessor`.
@@ -50,6 +52,7 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
         server_names: Optional mapping of server indices to names
         connection_timeout: Timeout for initial SSE connection setup
         default_timeout: Default timeout for tool execution
+        initialization_timeout: Timeout for complete initialization (default 60s, increase for slow servers)
         max_concurrency: Maximum concurrent operations
         enable_caching: Whether to enable response caching
         cache_ttl: Cache time-to-live in seconds
@@ -59,6 +62,7 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
         enable_retries: Whether to enable automatic retries
         max_retries: Maximum retry attempts
         namespace: Namespace for registered tools
+        oauth_refresh_callback: Optional async callback to refresh OAuth tokens (NEW)
 
     Returns:
         Tuple of (ToolProcessor, StreamManager)
@@ -69,12 +73,49 @@ async def setup_mcp_sse(  # noqa: C901 - long but just a config facade
         server_names=server_names,
         connection_timeout=connection_timeout,  # 🔧 ADD THIS LINE
         default_timeout=default_timeout,  # 🔧 ADD THIS LINE
+        initialization_timeout=initialization_timeout,
+        oauth_refresh_callback=oauth_refresh_callback,  # NEW: Pass OAuth callback
     )
 
     # 2️⃣  pull the remote tool list and register each one locally
     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
+    # Based on RFC 6750 (Bearer Token Usage) and MCP OAuth spec
+    oauth_error_patterns = [
+        # 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",
+    ]
+
+    # 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,
@@ -85,6 +126,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/mcp/setup_mcp_stdio.py

@@ -30,6 +30,7 @@ async def setup_mcp_stdio(  # noqa: C901 - long but just a config facade
     servers: list[str],
     server_names: dict[int, str] | None = None,
     default_timeout: float = 10.0,
+    initialization_timeout: float = 60.0,
     max_concurrency: int | None = None,
     enable_caching: bool = True,
     cache_ttl: int = 300,
@@ -53,6 +54,7 @@ async def setup_mcp_stdio(  # noqa: C901 - long but just a config facade
         server_names=server_names,
         transport_type="stdio",
         default_timeout=default_timeout,  # 🔧 ADD THIS LINE
+        initialization_timeout=initialization_timeout,
     )
 
     # 2️⃣  pull the remote tool list and register each one locally

chuk_tool_processor/mcp/stream_manager.py

@@ -21,6 +21,7 @@ from chuk_tool_processor.mcp.transport import (
     MCPBaseTransport,
     SSETransport,
     StdioTransport,
+    TimeoutConfig,
 )
 
 logger = get_logger("chuk_tool_processor.mcp.stream_manager")
@@ -38,7 +39,7 @@ class StreamManager:
     - HTTP Streamable (modern replacement for SSE, spec 2025-03-26) with graceful headers handling
     """
 
-    def __init__(self) -> None:
+    def __init__(self, timeout_config: TimeoutConfig | None = None) -> None:
         self.transports: dict[str, MCPBaseTransport] = {}
         self.server_info: list[dict[str, Any]] = []
         self.tool_to_server_map: dict[str, str] = {}
@@ -46,7 +47,7 @@ class StreamManager:
         self.all_tools: list[dict[str, Any]] = []
         self._lock = asyncio.Lock()
         self._closed = False  # Track if we've been closed
-        self._shutdown_timeout = 2.0  # Maximum time to spend on shutdown
+        self.timeout_config = timeout_config or TimeoutConfig()
 
     # ------------------------------------------------------------------ #
     #  factory helpers with enhanced error handling                      #
@@ -62,16 +63,16 @@ class StreamManager:
         initialization_timeout: float = 60.0,  # NEW: Timeout for entire initialization
     ) -> StreamManager:
         """Create StreamManager with timeout protection."""
-        try:
-            inst = cls()
-            await asyncio.wait_for(
-                inst.initialize(config_file, servers, server_names, transport_type, default_timeout=default_timeout),
-                timeout=initialization_timeout,
-            )
-            return inst
-        except TimeoutError:
-            logger.error("StreamManager initialization timed out after %ss", initialization_timeout)
-            raise RuntimeError(f"StreamManager initialization timed out after {initialization_timeout}s")
+        inst = cls()
+        await inst.initialize(
+            config_file,
+            servers,
+            server_names,
+            transport_type,
+            default_timeout=default_timeout,
+            initialization_timeout=initialization_timeout,
+        )
+        return inst
 
     @classmethod
     async def create_with_sse(
@@ -81,20 +82,19 @@ class StreamManager:
         connection_timeout: float = 10.0,
         default_timeout: float = 30.0,
         initialization_timeout: float = 60.0,  # NEW
+        oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
     ) -> StreamManager:
         """Create StreamManager with SSE transport and timeout protection."""
-        try:
-            inst = cls()
-            await asyncio.wait_for(
-                inst.initialize_with_sse(
-                    servers, server_names, connection_timeout=connection_timeout, default_timeout=default_timeout
-                ),
-                timeout=initialization_timeout,
-            )
-            return inst
-        except TimeoutError:
-            logger.error("SSE StreamManager initialization timed out after %ss", initialization_timeout)
-            raise RuntimeError(f"SSE StreamManager initialization timed out after {initialization_timeout}s")
+        inst = cls()
+        await inst.initialize_with_sse(
+            servers,
+            server_names,
+            connection_timeout=connection_timeout,
+            default_timeout=default_timeout,
+            initialization_timeout=initialization_timeout,
+            oauth_refresh_callback=oauth_refresh_callback,  # NEW: Pass OAuth callback
+        )
+        return inst
 
     @classmethod
     async def create_with_http_streamable(
@@ -104,22 +104,19 @@ class StreamManager:
         connection_timeout: float = 30.0,
         default_timeout: float = 30.0,
         initialization_timeout: float = 60.0,  # NEW
+        oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
     ) -> StreamManager:
         """Create StreamManager with HTTP Streamable transport and timeout protection."""
-        try:
-            inst = cls()
-            await asyncio.wait_for(
-                inst.initialize_with_http_streamable(
-                    servers, server_names, connection_timeout=connection_timeout, default_timeout=default_timeout
-                ),
-                timeout=initialization_timeout,
-            )
-            return inst
-        except TimeoutError:
-            logger.error("HTTP Streamable StreamManager initialization timed out after %ss", initialization_timeout)
-            raise RuntimeError(
-                f"HTTP Streamable StreamManager initialization timed out after {initialization_timeout}s"
-            )
+        inst = cls()
+        await inst.initialize_with_http_streamable(
+            servers,
+            server_names,
+            connection_timeout=connection_timeout,
+            default_timeout=default_timeout,
+            initialization_timeout=initialization_timeout,
+            oauth_refresh_callback=oauth_refresh_callback,  # NEW: Pass OAuth callback
+        )
+        return inst
 
     # ------------------------------------------------------------------ #
     #  NEW: Context manager support for automatic cleanup               #
@@ -167,6 +164,7 @@ class StreamManager:
         server_names: dict[int, str] | None = None,
         transport_type: str = "stdio",
         default_timeout: float = 30.0,
+        initialization_timeout: float = 60.0,
     ) -> None:
         """Initialize with graceful headers handling for all transport types."""
         if self._closed:
@@ -178,13 +176,24 @@ class StreamManager:
             for idx, server_name in enumerate(servers):
                 try:
                     if transport_type == "stdio":
-                        params = await load_config(config_file, server_name)
-                        transport: MCPBaseTransport = StdioTransport(params)
+                        params, server_timeout = await load_config(config_file, server_name)
+                        # Use per-server timeout if specified, otherwise use global default
+                        effective_timeout = server_timeout if server_timeout is not None else default_timeout
+                        logger.info(
+                            f"Server '{server_name}' using timeout: {effective_timeout}s (per-server: {server_timeout}, default: {default_timeout})"
+                        )
+                        # Use initialization_timeout for connection_timeout since subprocess
+                        # launch can take time (e.g., uvx downloading packages)
+                        transport: MCPBaseTransport = StdioTransport(
+                            params, connection_timeout=initialization_timeout, default_timeout=effective_timeout
+                        )
                     elif transport_type == "sse":
-                        logger.warning(
+                        logger.debug(
                             "Using SSE transport in initialize() - consider using initialize_with_sse() instead"
                         )
-                        params = await load_config(config_file, server_name)
+                        params, server_timeout = await load_config(config_file, server_name)
+                        # Use per-server timeout if specified, otherwise use global default
+                        effective_timeout = server_timeout if server_timeout is not None else default_timeout
 
                         if isinstance(params, dict) and "url" in params:
                             sse_url = params["url"]
@@ -194,20 +203,22 @@ class StreamManager:
                             sse_url = "http://localhost:8000"
                             api_key = None
                             headers = {}
-                            logger.warning("No URL configured for SSE transport, using default: %s", sse_url)
+                            logger.debug("No URL configured for SSE transport, using default: %s", sse_url)
 
                         # Build SSE transport with optional headers
-                        transport_params = {"url": sse_url, "api_key": api_key, "default_timeout": default_timeout}
+                        transport_params = {"url": sse_url, "api_key": api_key, "default_timeout": effective_timeout}
                         if headers:
                             transport_params["headers"] = headers
 
                         transport = SSETransport(**transport_params)
 
                     elif transport_type == "http_streamable":
-                        logger.warning(
+                        logger.debug(
                             "Using HTTP Streamable transport in initialize() - consider using initialize_with_http_streamable() instead"
                         )
-                        params = await load_config(config_file, server_name)
+                        params, server_timeout = await load_config(config_file, server_name)
+                        # Use per-server timeout if specified, otherwise use global default
+                        effective_timeout = server_timeout if server_timeout is not None else default_timeout
 
                         if isinstance(params, dict) and "url" in params:
                             http_url = params["url"]
@@ -219,15 +230,13 @@ class StreamManager:
                             api_key = None
                             headers = {}
                             session_id = None
-                            logger.warning(
-                                "No URL configured for HTTP Streamable transport, using default: %s", http_url
-                            )
+                            logger.debug("No URL configured for HTTP Streamable transport, using default: %s", http_url)
 
                         # Build HTTP transport (headers not supported yet)
                         transport_params = {
                             "url": http_url,
                             "api_key": api_key,
-                            "default_timeout": default_timeout,
+                            "default_timeout": effective_timeout,
                             "session_id": session_id,
                         }
                         # Note: headers not added until HTTPStreamableTransport supports them
@@ -241,15 +250,23 @@ class StreamManager:
                         continue
 
                     # Initialize with timeout protection
-                    if not await asyncio.wait_for(transport.initialize(), timeout=default_timeout):
-                        logger.error("Failed to init %s", server_name)
+                    try:
+                        if not await asyncio.wait_for(transport.initialize(), timeout=initialization_timeout):
+                            logger.warning("Failed to init %s", server_name)
+                            continue
+                    except TimeoutError:
+                        logger.error("Timeout initialising %s (timeout=%ss)", server_name, initialization_timeout)
                         continue
 
                     self.transports[server_name] = transport
 
-                    # Ping and get tools with timeout protection
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=5.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=10.0)
+                    # Ping and get tools with timeout protection (use longer timeouts for slow servers)
+                    status = (
+                        "Up"
+                        if await asyncio.wait_for(transport.send_ping(), timeout=self.timeout_config.operation)
+                        else "Down"
+                    )
+                    tools = await asyncio.wait_for(transport.get_tools(), timeout=self.timeout_config.operation)
 
                     for t in tools:
                         name = t.get("name")
@@ -283,6 +300,8 @@ class StreamManager:
         server_names: dict[int, str] | None = None,
         connection_timeout: float = 10.0,
         default_timeout: float = 30.0,
+        initialization_timeout: float = 60.0,
+        oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
     ) -> None:
         """Initialize with SSE transport with optional headers support."""
         if self._closed:
@@ -311,15 +330,29 @@ class StreamManager:
                         logger.debug("SSE %s: Using configured headers: %s", name, list(headers.keys()))
                         transport_params["headers"] = headers
 
+                    # Add OAuth refresh callback if provided (NEW)
+                    if oauth_refresh_callback:
+                        transport_params["oauth_refresh_callback"] = oauth_refresh_callback
+                        logger.debug("SSE %s: OAuth refresh callback configured", name)
+
                     transport = SSETransport(**transport_params)
 
-                    if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
-                        logger.error("Failed to init SSE %s", name)
+                    try:
+                        if not await asyncio.wait_for(transport.initialize(), timeout=initialization_timeout):
+                            logger.warning("Failed to init SSE %s", name)
+                            continue
+                    except TimeoutError:
+                        logger.error("Timeout initialising SSE %s (timeout=%ss)", name, initialization_timeout)
                         continue
 
                     self.transports[name] = transport
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=5.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=10.0)
+                    # Use longer timeouts for slow servers (ping can take time after initialization)
+                    status = (
+                        "Up"
+                        if await asyncio.wait_for(transport.send_ping(), timeout=self.timeout_config.operation)
+                        else "Down"
+                    )
+                    tools = await asyncio.wait_for(transport.get_tools(), timeout=self.timeout_config.operation)
 
                     for t in tools:
                         tname = t.get("name")
@@ -346,11 +379,15 @@ class StreamManager:
         server_names: dict[int, str] | None = None,
         connection_timeout: float = 30.0,
         default_timeout: float = 30.0,
+        initialization_timeout: float = 60.0,
+        oauth_refresh_callback: any | None = None,  # NEW: OAuth token refresh callback
     ) -> None:
         """Initialize with HTTP Streamable transport with graceful headers handling."""
         if self._closed:
             raise RuntimeError("Cannot initialize a closed StreamManager")
 
+        logger.debug(f"initialize_with_http_streamable: initialization_timeout={initialization_timeout}")
+
         async with self._lock:
             self.server_names = server_names or {}
 
@@ -369,22 +406,39 @@ class StreamManager:
                         "session_id": cfg.get("session_id"),
                     }
 
-                    # Handle headers if provided (for future HTTPStreamableTransport support)
+                    # Handle headers if provided
                     headers = cfg.get("headers", {})
                     if headers:
-                        logger.debug("HTTP Streamable %s: Headers provided but not yet supported in transport", name)
-                        # TODO: Add headers support when HTTPStreamableTransport is updated
-                        # transport_params['headers'] = headers
+                        transport_params["headers"] = headers
+                        logger.debug("HTTP Streamable %s: Custom headers configured: %s", name, list(headers.keys()))
+
+                    # Add OAuth refresh callback if provided (NEW)
+                    if oauth_refresh_callback:
+                        transport_params["oauth_refresh_callback"] = oauth_refresh_callback
+                        logger.debug("HTTP Streamable %s: OAuth refresh callback configured", name)
 
                     transport = HTTPStreamableTransport(**transport_params)
 
-                    if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
-                        logger.error("Failed to init HTTP Streamable %s", name)
+                    logger.debug(f"Calling transport.initialize() for {name} with timeout={initialization_timeout}s")
+                    try:
+                        if not await asyncio.wait_for(transport.initialize(), timeout=initialization_timeout):
+                            logger.warning("Failed to init HTTP Streamable %s", name)
+                            continue
+                    except TimeoutError:
+                        logger.error(
+                            "Timeout initialising HTTP Streamable %s (timeout=%ss)", name, initialization_timeout
+                        )
                         continue
+                    logger.debug(f"Successfully initialized {name}")
 
                     self.transports[name] = transport
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=5.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=10.0)
+                    # Use longer timeouts for slow servers (ping can take time after initialization)
+                    status = (
+                        "Up"
+                        if await asyncio.wait_for(transport.send_ping(), timeout=self.timeout_config.operation)
+                        else "Down"
+                    )
+                    tools = await asyncio.wait_for(transport.get_tools(), timeout=self.timeout_config.operation)
 
                     for t in tools:
                         tname = t.get("name")
@@ -430,7 +484,7 @@ class StreamManager:
         transport = self.transports[server_name]
 
         try:
-            tools = await asyncio.wait_for(transport.get_tools(), timeout=10.0)
+            tools = await asyncio.wait_for(transport.get_tools(), timeout=self.timeout_config.operation)
             logger.debug("Found %d tools for server %s", len(tools), server_name)
             return tools
         except TimeoutError:
@@ -449,7 +503,7 @@ class StreamManager:
 
         async def _ping_one(name: str, tr: MCPBaseTransport):
             try:
-                ok = await asyncio.wait_for(tr.send_ping(), timeout=5.0)
+                ok = await asyncio.wait_for(tr.send_ping(), timeout=self.timeout_config.quick)
             except Exception:
                 ok = False
             return {"server": name, "ok": ok}
@@ -464,7 +518,7 @@ class StreamManager:
 
         async def _one(name: str, tr: MCPBaseTransport):
             try:
-                res = await asyncio.wait_for(tr.list_resources(), timeout=10.0)
+                res = await asyncio.wait_for(tr.list_resources(), timeout=self.timeout_config.operation)
                 resources = res.get("resources", []) if isinstance(res, dict) else res
                 for item in resources:
                     item = dict(item)
@@ -484,7 +538,7 @@ class StreamManager:
 
         async def _one(name: str, tr: MCPBaseTransport):
             try:
-                res = await asyncio.wait_for(tr.list_prompts(), timeout=10.0)
+                res = await asyncio.wait_for(tr.list_prompts(), timeout=self.timeout_config.operation)
                 prompts = res.get("prompts", []) if isinstance(res, dict) else res
                 for item in prompts:
                     item = dict(item)
@@ -611,7 +665,7 @@ class StreamManager:
             try:
                 results = await asyncio.wait_for(
                     asyncio.gather(*[task for _, task in close_tasks], return_exceptions=True),
-                    timeout=self._shutdown_timeout,
+                    timeout=self.timeout_config.shutdown,
                 )
 
                 # Process results
@@ -634,7 +688,8 @@ class StreamManager:
                 # Brief wait for cancellations to complete
                 with contextlib.suppress(TimeoutError):
                     await asyncio.wait_for(
-                        asyncio.gather(*[task for _, task in close_tasks], return_exceptions=True), timeout=0.5
+                        asyncio.gather(*[task for _, task in close_tasks], return_exceptions=True),
+                        timeout=self.timeout_config.shutdown,
                     )
 
     async def _sequential_close(self, transport_items: list[tuple[str, MCPBaseTransport]], close_results: list) -> None:
@@ -643,7 +698,7 @@ class StreamManager:
             try:
                 await asyncio.wait_for(
                     self._close_single_transport(name, transport),
-                    timeout=0.5,  # Short timeout per transport
+                    timeout=self.timeout_config.shutdown,
                 )
                 logger.debug("Closed transport: %s", name)
                 close_results.append((name, True, None))
@@ -735,7 +790,7 @@ class StreamManager:
 
         for name, transport in self.transports.items():
             try:
-                ping_ok = await asyncio.wait_for(transport.send_ping(), timeout=5.0)
+                ping_ok = await asyncio.wait_for(transport.send_ping(), timeout=self.timeout_config.quick)
                 health_info["transports"][name] = {
                     "status": "healthy" if ping_ok else "unhealthy",
                     "ping_success": ping_ok,

chuk_tool_processor/mcp/transport/__init__.py

@@ -11,6 +11,12 @@ All transports now follow the same interface and provide consistent behavior:
 
 from .base_transport import MCPBaseTransport
 from .http_streamable_transport import HTTPStreamableTransport
+from .models import (
+    HeadersConfig,
+    ServerInfo,
+    TimeoutConfig,
+    TransportMetrics,
+)
 from .sse_transport import SSETransport
 from .stdio_transport import StdioTransport
 
@@ -19,4 +25,8 @@ __all__ = [
     "StdioTransport",
     "SSETransport",
     "HTTPStreamableTransport",
+    "TimeoutConfig",
+    "TransportMetrics",
+    "ServerInfo",
+    "HeadersConfig",
 ]