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

chuk_tool_processor/mcp/stream_manager.py

@@ -1,24 +1,28 @@
 # chuk_tool_processor/mcp/stream_manager.py
 """
-StreamManager for CHUK Tool Processor - Enhanced with robust shutdown handling
+StreamManager for CHUK Tool Processor - Enhanced with robust shutdown handling and headers support
 """
+
 from __future__ import annotations
 
 import asyncio
-from typing import Any, Dict, List, Optional, Tuple
+import contextlib
 from contextlib import asynccontextmanager
+from typing import Any
 
 # --------------------------------------------------------------------------- #
 #  CHUK imports                                                               #
 # --------------------------------------------------------------------------- #
-from chuk_mcp.config import load_config
+from chuk_mcp.config import load_config  # type: ignore[import-untyped]
+
+from chuk_tool_processor.logging import get_logger
 from chuk_tool_processor.mcp.transport import (
+    HTTPStreamableTransport,
     MCPBaseTransport,
-    StdioTransport,
     SSETransport,
-    HTTPStreamableTransport,
+    StdioTransport,
+    TimeoutConfig,
 )
-from chuk_tool_processor.logging import get_logger
 
 logger = get_logger("chuk_tool_processor.mcp.stream_manager")
 
@@ -26,24 +30,24 @@ logger = get_logger("chuk_tool_processor.mcp.stream_manager")
 class StreamManager:
     """
     Manager for MCP server streams with support for multiple transport types.
-    
-    Enhanced with robust shutdown handling to prevent event loop closure issues.
-    
+
+    Enhanced with robust shutdown handling and proper headers support.
+
     Updated to support the latest transports:
     - STDIO (process-based)
-    - SSE (Server-Sent Events) 
-    - HTTP Streamable (modern replacement for SSE, spec 2025-03-26)
+    - SSE (Server-Sent Events) with headers support
+    - HTTP Streamable (modern replacement for SSE, spec 2025-03-26) with graceful headers handling
     """
 
-    def __init__(self) -> None:
-        self.transports: Dict[str, MCPBaseTransport] = {}
-        self.server_info: List[Dict[str, Any]] = []
-        self.tool_to_server_map: Dict[str, str] = {}
-        self.server_names: Dict[int, str] = {}
-        self.all_tools: List[Dict[str, Any]] = []
+    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] = {}
+        self.server_names: dict[int, str] = {}
+        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                      #
@@ -52,81 +56,67 @@ class StreamManager:
     async def create(
         cls,
         config_file: str,
-        servers: List[str],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[str],
+        server_names: dict[int, str] | None = None,
         transport_type: str = "stdio",
         default_timeout: float = 30.0,
         initialization_timeout: float = 60.0,  # NEW: Timeout for entire initialization
-    ) -> "StreamManager":
+    ) -> 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 asyncio.TimeoutError:
-            logger.error(f"StreamManager initialization timed out after {initialization_timeout}s")
-            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(
         cls,
-        servers: List[Dict[str, str]],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[dict[str, str]],
+        server_names: dict[int, str] | None = None,
         connection_timeout: float = 10.0,
         default_timeout: float = 30.0,
         initialization_timeout: float = 60.0,  # NEW
-    ) -> "StreamManager":
+        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 asyncio.TimeoutError:
-            logger.error(f"SSE StreamManager initialization timed out after {initialization_timeout}s")
-            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(
         cls,
-        servers: List[Dict[str, str]],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[dict[str, str]],
+        server_names: dict[int, str] | None = None,
         connection_timeout: float = 30.0,
         default_timeout: float = 30.0,
         initialization_timeout: float = 60.0,  # NEW
-    ) -> "StreamManager":
+        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 asyncio.TimeoutError:
-            logger.error(f"HTTP Streamable StreamManager initialization timed out after {initialization_timeout}s")
-            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               #
@@ -144,8 +134,8 @@ class StreamManager:
     async def create_managed(
         cls,
         config_file: str,
-        servers: List[str],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[str],
+        server_names: dict[int, str] | None = None,
         transport_type: str = "stdio",
         default_timeout: float = 30.0,
     ):
@@ -170,73 +160,113 @@ class StreamManager:
     async def initialize(
         self,
         config_file: str,
-        servers: List[str],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[str],
+        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:
             raise RuntimeError("Cannot initialize a closed StreamManager")
-            
+
         async with self._lock:
             self.server_names = server_names or {}
 
             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("Using SSE transport in initialize() - consider using initialize_with_sse() instead")
-                        params = await load_config(config_file, server_name)
-                        
-                        if isinstance(params, dict) and 'url' in params:
-                            sse_url = params['url']
-                            api_key = params.get('api_key')
+                        logger.debug(
+                            "Using SSE transport in initialize() - consider using initialize_with_sse() instead"
+                        )
+                        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"]
+                            api_key = params.get("api_key")
+                            headers = params.get("headers", {})
                         else:
                             sse_url = "http://localhost:8000"
                             api_key = None
-                            logger.warning(f"No URL configured for SSE transport, using default: {sse_url}")
-                        
-                        transport = SSETransport(
-                            sse_url,
-                            api_key,
-                            default_timeout=default_timeout
-                        )
+                            headers = {}
+                            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": effective_timeout}
+                        if headers:
+                            transport_params["headers"] = headers
+
+                        transport = SSETransport(**transport_params)
+
                     elif transport_type == "http_streamable":
-                        logger.warning("Using HTTP Streamable transport in initialize() - consider using initialize_with_http_streamable() instead")
-                        params = await load_config(config_file, server_name)
-                        
-                        if isinstance(params, dict) and 'url' in params:
-                            http_url = params['url']
-                            api_key = params.get('api_key')
-                            session_id = params.get('session_id')
+                        logger.debug(
+                            "Using HTTP Streamable transport in initialize() - consider using initialize_with_http_streamable() instead"
+                        )
+                        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"]
+                            api_key = params.get("api_key")
+                            headers = params.get("headers", {})
+                            session_id = params.get("session_id")
                         else:
                             http_url = "http://localhost:8000"
                             api_key = None
+                            headers = {}
                             session_id = None
-                            logger.warning(f"No URL configured for HTTP Streamable transport, using default: {http_url}")
-                        
-                        transport = HTTPStreamableTransport(
-                            http_url,
-                            api_key,
-                            default_timeout=default_timeout,
-                            session_id=session_id
-                        )
+                            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": effective_timeout,
+                            "session_id": session_id,
+                        }
+                        # Note: headers not added until HTTPStreamableTransport supports them
+                        if headers:
+                            logger.debug("Headers provided but not supported in HTTPStreamableTransport yet")
+
+                        transport = HTTPStreamableTransport(**transport_params)
+
                     else:
                         logger.error("Unsupported transport type: %s", transport_type)
                         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")
@@ -252,13 +282,13 @@ class StreamManager:
                             "status": status,
                         }
                     )
-                    logger.info("Initialised %s - %d tool(s)", server_name, len(tools))
-                except asyncio.TimeoutError:
+                    logger.debug("Initialised %s - %d tool(s)", server_name, len(tools))
+                except TimeoutError:
                     logger.error("Timeout initialising %s", server_name)
                 except Exception as exc:
                     logger.error("Error initialising %s: %s", server_name, exc)
 
-            logger.info(
+            logger.debug(
                 "StreamManager ready - %d server(s), %d tool(s)",
                 len(self.transports),
                 len(self.all_tools),
@@ -266,14 +296,17 @@ class StreamManager:
 
     async def initialize_with_sse(
         self,
-        servers: List[Dict[str, str]],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[dict[str, str]],
+        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:
             raise RuntimeError("Cannot initialize a closed StreamManager")
-            
+
         async with self._lock:
             self.server_names = server_names or {}
 
@@ -283,20 +316,43 @@ class StreamManager:
                     logger.error("Bad server config: %s", cfg)
                     continue
                 try:
-                    transport = SSETransport(
-                        url, 
-                        cfg.get("api_key"),
-                        connection_timeout=connection_timeout,
-                        default_timeout=default_timeout
-                    )
-                    
-                    if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
-                        logger.error("Failed to init SSE %s", name)
+                    # Build SSE transport parameters with optional headers
+                    transport_params = {
+                        "url": url,
+                        "api_key": cfg.get("api_key"),
+                        "connection_timeout": connection_timeout,
+                        "default_timeout": default_timeout,
+                    }
+
+                    # Add headers if provided
+                    headers = cfg.get("headers", {})
+                    if headers:
+                        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)
+
+                    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")
@@ -304,16 +360,14 @@ class StreamManager:
                             self.tool_to_server_map[tname] = name
                     self.all_tools.extend(tools)
 
-                    self.server_info.append(
-                        {"id": idx, "name": name, "tools": len(tools), "status": status}
-                    )
-                    logger.info("Initialised SSE %s - %d tool(s)", name, len(tools))
-                except asyncio.TimeoutError:
+                    self.server_info.append({"id": idx, "name": name, "tools": len(tools), "status": status})
+                    logger.debug("Initialised SSE %s - %d tool(s)", name, len(tools))
+                except TimeoutError:
                     logger.error("Timeout initialising SSE %s", name)
                 except Exception as exc:
                     logger.error("Error initialising SSE %s: %s", name, exc)
 
-            logger.info(
+            logger.debug(
                 "StreamManager ready - %d SSE server(s), %d tool(s)",
                 len(self.transports),
                 len(self.all_tools),
@@ -321,15 +375,19 @@ class StreamManager:
 
     async def initialize_with_http_streamable(
         self,
-        servers: List[Dict[str, str]],
-        server_names: Optional[Dict[int, str]] = None,
+        servers: list[dict[str, str]],
+        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 (modern MCP spec 2025-03-26)."""
+        """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 {}
 
@@ -339,21 +397,48 @@ class StreamManager:
                     logger.error("Bad server config: %s", cfg)
                     continue
                 try:
-                    transport = HTTPStreamableTransport(
-                        url, 
-                        cfg.get("api_key"),
-                        connection_timeout=connection_timeout,
-                        default_timeout=default_timeout,
-                        session_id=cfg.get("session_id")
-                    )
-                    
-                    if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
-                        logger.error("Failed to init HTTP Streamable %s", name)
+                    # Build HTTP Streamable transport parameters
+                    transport_params = {
+                        "url": url,
+                        "api_key": cfg.get("api_key"),
+                        "connection_timeout": connection_timeout,
+                        "default_timeout": default_timeout,
+                        "session_id": cfg.get("session_id"),
+                    }
+
+                    # Handle headers if provided
+                    headers = cfg.get("headers", {})
+                    if 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)
+
+                    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")
@@ -361,16 +446,14 @@ class StreamManager:
                             self.tool_to_server_map[tname] = name
                     self.all_tools.extend(tools)
 
-                    self.server_info.append(
-                        {"id": idx, "name": name, "tools": len(tools), "status": status}
-                    )
-                    logger.info("Initialised HTTP Streamable %s - %d tool(s)", name, len(tools))
-                except asyncio.TimeoutError:
+                    self.server_info.append({"id": idx, "name": name, "tools": len(tools), "status": status})
+                    logger.debug("Initialised HTTP Streamable %s - %d tool(s)", name, len(tools))
+                except TimeoutError:
                     logger.error("Timeout initialising HTTP Streamable %s", name)
                 except Exception as exc:
                     logger.error("Error initialising HTTP Streamable %s: %s", name, exc)
 
-            logger.info(
+            logger.debug(
                 "StreamManager ready - %d HTTP Streamable server(s), %d tool(s)",
                 len(self.transports),
                 len(self.all_tools),
@@ -379,66 +462,64 @@ class StreamManager:
     # ------------------------------------------------------------------ #
     #  queries                                                           #
     # ------------------------------------------------------------------ #
-    def get_all_tools(self) -> List[Dict[str, Any]]:
+    def get_all_tools(self) -> list[dict[str, Any]]:
         return self.all_tools
 
-    def get_server_for_tool(self, tool_name: str) -> Optional[str]:
+    def get_server_for_tool(self, tool_name: str) -> str | None:
         return self.tool_to_server_map.get(tool_name)
 
-    def get_server_info(self) -> List[Dict[str, Any]]:
+    def get_server_info(self) -> list[dict[str, Any]]:
         return self.server_info
-    
-    async def list_tools(self, server_name: str) -> List[Dict[str, Any]]:
+
+    async def list_tools(self, server_name: str) -> list[dict[str, Any]]:
         """List all tools available from a specific server."""
         if self._closed:
             logger.warning("Cannot list tools: StreamManager is closed")
             return []
-            
+
         if server_name not in self.transports:
-            logger.error(f"Server '{server_name}' not found in transports")
+            logger.error("Server '%s' not found in transports", server_name)
             return []
-        
+
         transport = self.transports[server_name]
-        
+
         try:
-            tools = await asyncio.wait_for(transport.get_tools(), timeout=10.0)
-            logger.debug(f"Found {len(tools)} tools for server {server_name}")
+            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 asyncio.TimeoutError:
-            logger.error(f"Timeout listing tools for server {server_name}")
+        except TimeoutError:
+            logger.error("Timeout listing tools for server %s", server_name)
             return []
         except Exception as e:
-            logger.error(f"Error listing tools for server {server_name}: {e}")
+            logger.error("Error listing tools for server %s: %s", server_name, e)
             return []
 
     # ------------------------------------------------------------------ #
     #  EXTRA HELPERS - ping / resources / prompts                        #
     # ------------------------------------------------------------------ #
-    async def ping_servers(self) -> List[Dict[str, Any]]:
+    async def ping_servers(self) -> list[dict[str, Any]]:
         if self._closed:
             return []
-            
+
         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}
 
         return await asyncio.gather(*(_ping_one(n, t) for n, t in self.transports.items()), return_exceptions=True)
 
-    async def list_resources(self) -> List[Dict[str, Any]]:
+    async def list_resources(self) -> list[dict[str, Any]]:
         if self._closed:
             return []
-            
-        out: List[Dict[str, Any]] = []
+
+        out: list[dict[str, Any]] = []
 
         async def _one(name: str, tr: MCPBaseTransport):
             try:
-                res = await asyncio.wait_for(tr.list_resources(), timeout=10.0)
-                resources = (
-                    res.get("resources", []) if isinstance(res, dict) else res
-                )
+                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)
                     item["server"] = name
@@ -449,15 +530,15 @@ class StreamManager:
         await asyncio.gather(*(_one(n, t) for n, t in self.transports.items()), return_exceptions=True)
         return out
 
-    async def list_prompts(self) -> List[Dict[str, Any]]:
+    async def list_prompts(self) -> list[dict[str, Any]]:
         if self._closed:
             return []
-            
-        out: List[Dict[str, Any]] = []
+
+        out: list[dict[str, Any]] = []
 
         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)
@@ -475,45 +556,40 @@ class StreamManager:
     async def call_tool(
         self,
         tool_name: str,
-        arguments: Dict[str, Any],
-        server_name: Optional[str] = None,
-        timeout: Optional[float] = None,
-    ) -> Dict[str, Any]:
+        arguments: dict[str, Any],
+        server_name: str | None = None,
+        timeout: float | None = None,
+    ) -> dict[str, Any]:
         """Call a tool on the appropriate server with timeout support."""
         if self._closed:
             return {
                 "isError": True,
                 "error": "StreamManager is closed",
             }
-            
+
         server_name = server_name or self.get_server_for_tool(tool_name)
         if not server_name or server_name not in self.transports:
             return {
                 "isError": True,
                 "error": f"No server found for tool: {tool_name}",
             }
-        
+
         transport = self.transports[server_name]
-        
+
         if timeout is not None:
             logger.debug("Calling tool '%s' with %ss timeout", tool_name, timeout)
             try:
-                if hasattr(transport, 'call_tool'):
+                if hasattr(transport, "call_tool"):
                     import inspect
+
                     sig = inspect.signature(transport.call_tool)
-                    if 'timeout' in sig.parameters:
+                    if "timeout" in sig.parameters:
                         return await transport.call_tool(tool_name, arguments, timeout=timeout)
                     else:
-                        return await asyncio.wait_for(
-                            transport.call_tool(tool_name, arguments),
-                            timeout=timeout
-                        )
+                        return await asyncio.wait_for(transport.call_tool(tool_name, arguments), timeout=timeout)
                 else:
-                    return await asyncio.wait_for(
-                        transport.call_tool(tool_name, arguments),
-                        timeout=timeout
-                    )
-            except asyncio.TimeoutError:
+                    return await asyncio.wait_for(transport.call_tool(tool_name, arguments), timeout=timeout)
+            except TimeoutError:
                 logger.warning("Tool '%s' timed out after %ss", tool_name, timeout)
                 return {
                     "isError": True,
@@ -521,28 +597,28 @@ class StreamManager:
                 }
         else:
             return await transport.call_tool(tool_name, arguments)
-        
+
     # ------------------------------------------------------------------ #
     #  ENHANCED shutdown with robust error handling                      #
     # ------------------------------------------------------------------ #
     async def close(self) -> None:
         """
         Close all transports safely with enhanced error handling.
-        
+
         ENHANCED: Uses asyncio.shield() to protect critical cleanup and
         provides multiple fallback strategies for different failure modes.
         """
         if self._closed:
             logger.debug("StreamManager already closed")
             return
-            
+
         if not self.transports:
             logger.debug("No transports to close")
             self._closed = True
             return
-        
-        logger.debug(f"Closing {len(self.transports)} transports...")
-        
+
+        logger.debug("Closing %d transports...", len(self.transports))
+
         try:
             # Use shield to protect the cleanup operation from cancellation
             await asyncio.shield(self._do_close_all_transports())
@@ -551,7 +627,7 @@ class StreamManager:
             logger.debug("Close operation cancelled, performing synchronous cleanup")
             self._sync_cleanup()
         except Exception as e:
-            logger.debug(f"Error during close: {e}")
+            logger.debug("Error during close: %s", e)
             self._sync_cleanup()
         finally:
             self._closed = True
@@ -560,99 +636,91 @@ class StreamManager:
         """Protected cleanup implementation with multiple strategies."""
         close_results = []
         transport_items = list(self.transports.items())
-        
+
         # Strategy 1: Try concurrent close with timeout
         try:
             await self._concurrent_close(transport_items, close_results)
         except Exception as e:
-            logger.debug(f"Concurrent close failed: {e}, falling back to sequential close")
+            logger.debug("Concurrent close failed: %s, falling back to sequential close", e)
             # Strategy 2: Fall back to sequential close
             await self._sequential_close(transport_items, close_results)
-        
+
         # Always clean up state
         self._cleanup_state()
-        
+
         # Log summary
         if close_results:
             successful_closes = sum(1 for _, success, _ in close_results if success)
-            logger.debug(f"Transport cleanup: {successful_closes}/{len(close_results)} closed successfully")
+            logger.debug("Transport cleanup: %d/%d closed successfully", successful_closes, len(close_results))
 
-    async def _concurrent_close(self, transport_items: List[Tuple[str, MCPBaseTransport]], close_results: List) -> None:
+    async def _concurrent_close(self, transport_items: list[tuple[str, MCPBaseTransport]], close_results: list) -> None:
         """Try to close all transports concurrently."""
         close_tasks = []
         for name, transport in transport_items:
-            task = asyncio.create_task(
-                self._close_single_transport(name, transport),
-                name=f"close_{name}"
-            )
+            task = asyncio.create_task(self._close_single_transport(name, transport), name=f"close_{name}")
             close_tasks.append((name, task))
-        
+
         # Wait for all tasks with a reasonable timeout
         if close_tasks:
             try:
                 results = await asyncio.wait_for(
-                    asyncio.gather(
-                        *[task for _, task in close_tasks],
-                        return_exceptions=True
-                    ),
-                    timeout=self._shutdown_timeout
+                    asyncio.gather(*[task for _, task in close_tasks], return_exceptions=True),
+                    timeout=self.timeout_config.shutdown,
                 )
-                
+
                 # Process results
                 for i, (name, _) in enumerate(close_tasks):
                     result = results[i] if i < len(results) else None
                     if isinstance(result, Exception):
-                        logger.debug(f"Transport {name} close failed: {result}")
+                        logger.debug("Transport %s close failed: %s", name, result)
                         close_results.append((name, False, str(result)))
                     else:
-                        logger.debug(f"Transport {name} closed successfully")
+                        logger.debug("Transport %s closed successfully", name)
                         close_results.append((name, True, None))
-                        
-            except asyncio.TimeoutError:
+
+            except TimeoutError:
                 # Cancel any remaining tasks
                 for name, task in close_tasks:
                     if not task.done():
                         task.cancel()
                         close_results.append((name, False, "timeout"))
-                
+
                 # Brief wait for cancellations to complete
-                try:
+                with contextlib.suppress(TimeoutError):
                     await asyncio.wait_for(
                         asyncio.gather(*[task for _, task in close_tasks], return_exceptions=True),
-                        timeout=0.5
+                        timeout=self.timeout_config.shutdown,
                     )
-                except asyncio.TimeoutError:
-                    pass  # Some tasks may not cancel cleanly
 
-    async def _sequential_close(self, transport_items: List[Tuple[str, MCPBaseTransport]], close_results: List) -> None:
+    async def _sequential_close(self, transport_items: list[tuple[str, MCPBaseTransport]], close_results: list) -> None:
         """Close transports one by one as fallback."""
         for name, transport in transport_items:
             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(f"Closed transport: {name}")
+                logger.debug("Closed transport: %s", name)
                 close_results.append((name, True, None))
-            except asyncio.TimeoutError:
-                logger.debug(f"Transport {name} close timed out (normal during shutdown)")
+            except TimeoutError:
+                logger.debug("Transport %s close timed out (normal during shutdown)", name)
                 close_results.append((name, False, "timeout"))
             except asyncio.CancelledError:
-                logger.debug(f"Transport {name} close cancelled during event loop shutdown")
+                logger.debug("Transport %s close cancelled during event loop shutdown", name)
                 close_results.append((name, False, "cancelled"))
             except Exception as e:
-                logger.debug(f"Error closing transport {name}: {e}")
+                logger.debug("Error closing transport %s: %s", name, e)
                 close_results.append((name, False, str(e)))
 
     async def _close_single_transport(self, name: str, transport: MCPBaseTransport) -> None:
         """Close a single transport with error handling."""
         try:
-            if hasattr(transport, 'close') and callable(transport.close):
+            if hasattr(transport, "close") and callable(transport.close):
                 await transport.close()
             else:
-                logger.debug(f"Transport {name} has no close method")
+                logger.debug("Transport %s has no close method", name)
         except Exception as e:
-            logger.debug(f"Error closing transport {name}: {e}")
+            logger.debug("Error closing transport %s: %s", name, e)
             raise
 
     def _sync_cleanup(self) -> None:
@@ -660,9 +728,9 @@ class StreamManager:
         try:
             transport_count = len(self.transports)
             self._cleanup_state()
-            logger.debug(f"Synchronous cleanup completed for {transport_count} transports")
+            logger.debug("Synchronous cleanup completed for %d transports", transport_count)
         except Exception as e:
-            logger.debug(f"Error during synchronous cleanup: {e}")
+            logger.debug("Error during synchronous cleanup: %s", e)
 
     def _cleanup_state(self) -> None:
         """Clean up internal state synchronously."""
@@ -673,17 +741,17 @@ class StreamManager:
             self.all_tools.clear()
             self.server_names.clear()
         except Exception as e:
-            logger.debug(f"Error during state cleanup: {e}")
+            logger.debug("Error during state cleanup: %s", e)
 
     # ------------------------------------------------------------------ #
     #  backwards-compat: streams helper                                  #
     # ------------------------------------------------------------------ #
-    def get_streams(self) -> List[Tuple[Any, Any]]:
+    def get_streams(self) -> list[tuple[Any, Any]]:
         """Return a list of (read_stream, write_stream) tuples for all transports."""
         if self._closed:
             return []
-            
-        pairs: List[Tuple[Any, Any]] = []
+
+        pairs: list[tuple[Any, Any]] = []
 
         for tr in self.transports.values():
             if hasattr(tr, "get_streams") and callable(tr.get_streams):
@@ -698,7 +766,7 @@ class StreamManager:
         return pairs
 
     @property
-    def streams(self) -> List[Tuple[Any, Any]]:
+    def streams(self) -> list[tuple[Any, Any]]:
         """Convenience alias for get_streams()."""
         return self.get_streams()
 
@@ -713,34 +781,23 @@ class StreamManager:
         """Get the number of active transports."""
         return len(self.transports)
 
-    async def health_check(self) -> Dict[str, Any]:
+    async def health_check(self) -> dict[str, Any]:
         """Perform a health check on all transports."""
         if self._closed:
             return {"status": "closed", "transports": {}}
-        
-        health_info = {
-            "status": "active",
-            "transport_count": len(self.transports),
-            "transports": {}
-        }
-        
+
+        health_info = {"status": "active", "transport_count": len(self.transports), "transports": {}}
+
         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
-                }
-            except asyncio.TimeoutError:
-                health_info["transports"][name] = {
-                    "status": "timeout",
-                    "ping_success": False
+                    "ping_success": ping_ok,
                 }
+            except TimeoutError:
+                health_info["transports"][name] = {"status": "timeout", "ping_success": False}
             except Exception as e:
-                health_info["transports"][name] = {
-                    "status": "error",
-                    "ping_success": False,
-                    "error": str(e)
-                }
-        
-        return health_info
+                health_info["transports"][name] = {"status": "error", "ping_success": False, "error": str(e)}
+
+        return health_info