improved timeouts, pydantic models etc

Author: chrishayuk

README.md

@@ -913,6 +913,79 @@ async def test_calculator():
 
 ## Configuration
 
+### Timeout Configuration
+
+CHUK Tool Processor uses a unified timeout configuration system that applies to all MCP transports (HTTP Streamable, SSE, STDIO) and the StreamManager. Instead of managing dozens of individual timeout values, there are just **4 logical timeout categories**:
+
+```python
+from chuk_tool_processor.mcp.transport import TimeoutConfig
+
+# Create custom timeout configuration
+timeout_config = TimeoutConfig(
+    connect=30.0,     # Connection establishment, initialization, session discovery
+    operation=30.0,   # Normal operations (tool calls, listing tools/resources/prompts)
+    quick=5.0,        # Fast health checks and pings
+    shutdown=2.0      # Cleanup and shutdown operations
+)
+```
+
+**Using timeout configuration with StreamManager:**
+
+```python
+from chuk_tool_processor.mcp.stream_manager import StreamManager
+from chuk_tool_processor.mcp.transport import TimeoutConfig
+
+# Create StreamManager with custom timeouts
+timeout_config = TimeoutConfig(
+    connect=60.0,     # Longer for slow initialization
+    operation=45.0,   # Longer for heavy operations
+    quick=3.0,        # Faster health checks
+    shutdown=5.0      # More time for cleanup
+)
+
+manager = StreamManager(timeout_config=timeout_config)
+```
+
+**Timeout categories explained:**
+
+| Category | Default | Used For | Examples |
+|----------|---------|----------|----------|
+| `connect` | 30.0s | Connection setup, initialization, discovery | HTTP connection, SSE session discovery, STDIO subprocess launch |
+| `operation` | 30.0s | Normal tool operations | Tool calls, listing tools/resources/prompts, get_tools() |
+| `quick` | 5.0s | Fast health/status checks | Ping operations, health checks |
+| `shutdown` | 2.0s | Cleanup and teardown | Transport close, connection cleanup |
+
+**Why this matters:**
+- ✅ **Simple**: 4 timeout values instead of 20+
+- ✅ **Consistent**: Same timeout behavior across all transports
+- ✅ **Configurable**: Adjust timeouts based on your environment (slow networks, large datasets, etc.)
+- ✅ **Type-safe**: Pydantic validation ensures correct values
+
+**Example: Adjusting for slow environments**
+
+```python
+from chuk_tool_processor.mcp import setup_mcp_stdio
+from chuk_tool_processor.mcp.transport import TimeoutConfig
+
+# For slow network or resource-constrained environments
+slow_timeouts = TimeoutConfig(
+    connect=120.0,    # Allow more time for package downloads
+    operation=60.0,   # Allow more time for heavy operations
+    quick=10.0,       # Be patient with health checks
+    shutdown=10.0     # Allow thorough cleanup
+)
+
+processor, manager = await setup_mcp_stdio(
+    config_file="mcp_config.json",
+    servers=["sqlite"],
+    namespace="db",
+    initialization_timeout=120.0
+)
+
+# Set custom timeouts on the manager
+manager.timeout_config = slow_timeouts
+```
+
 ### Environment Variables
 
 | Variable | Default | Description |

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chuk-tool-processor"
-version = "0.6.28"
+version = "0.6.29"
 description = "Async-native framework for registering, discovering, and executing tools referenced in LLM responses"
 readme = "README.md"
 requires-python = ">=3.11"

src/chuk_tool_processor/mcp/mcp_tool.py

@@ -370,7 +370,7 @@ async def _record_success(self) -> None:
             self._circuit_open = False
             self._circuit_open_time = None
             self.connection_state = ConnectionState.HEALTHY
-            logger.info(f"Circuit breaker closed for tool '{self.tool_name}' after successful execution")
+            logger.debug(f"Circuit breaker closed for tool '{self.tool_name}' after successful execution")
 
     async def _record_failure(self, is_connection_error: bool = False) -> None:
         """Record a failed execution."""
@@ -407,7 +407,7 @@ def _is_circuit_open(self) -> bool:
             self._circuit_open = False
             self._circuit_open_time = None
             self.connection_state = ConnectionState.HEALTHY
-            logger.info(f"Circuit breaker reset for tool '{self.tool_name}' after timeout")
+            logger.debug(f"Circuit breaker reset for tool '{self.tool_name}' after timeout")
             return False
 
         return True
@@ -462,12 +462,12 @@ def reset_circuit_breaker(self) -> None:
         self._circuit_open_time = None
         self._consecutive_failures = 0
         self.connection_state = ConnectionState.HEALTHY
-        logger.info(f"Circuit breaker manually reset for tool '{self.tool_name}'")
+        logger.debug(f"Circuit breaker manually reset for tool '{self.tool_name}'")
 
     def disable_resilience(self) -> None:
         """Disable resilience features for this tool instance."""
         self.enable_resilience = False
-        logger.info(f"Resilience features disabled for tool '{self.tool_name}'")
+        logger.debug(f"Resilience features disabled for tool '{self.tool_name}'")
 
     def set_stream_manager(self, stream_manager: StreamManager | None) -> None:
         """
@@ -482,7 +482,7 @@ def set_stream_manager(self, stream_manager: StreamManager | None) -> None:
             if self._circuit_open:
                 self._circuit_open = False
                 self._circuit_open_time = None
-                logger.info(f"Circuit breaker closed for tool '{self.tool_name}' due to new stream manager")
+                logger.debug(f"Circuit breaker closed for tool '{self.tool_name}' due to new stream manager")
         else:
             self.connection_state = ConnectionState.DISCONNECTED
 

src/chuk_tool_processor/mcp/stream_manager.py

@@ -21,6 +21,7 @@
     MCPBaseTransport,
     SSETransport,
     StdioTransport,
+    TimeoutConfig,
 )
 
 logger = get_logger("chuk_tool_processor.mcp.stream_manager")
@@ -38,15 +39,15 @@ 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] = {}
         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                      #
@@ -251,8 +252,12 @@ async def initialize(
                     self.transports[server_name] = transport
 
                     # Ping and get tools with timeout protection (use longer timeouts for slow servers)
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=30.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=30.0)
+                    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")
@@ -333,8 +338,12 @@ async def initialize_with_sse(
 
                     self.transports[name] = transport
                     # Use longer timeouts for slow servers (ping can take time after initialization)
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=30.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=30.0)
+                    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")
@@ -415,8 +424,12 @@ async def initialize_with_http_streamable(
 
                     self.transports[name] = transport
                     # Use longer timeouts for slow servers (ping can take time after initialization)
-                    status = "Up" if await asyncio.wait_for(transport.send_ping(), timeout=30.0) else "Down"
-                    tools = await asyncio.wait_for(transport.get_tools(), timeout=30.0)
+                    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")
@@ -462,7 +475,7 @@ async def list_tools(self, server_name: str) -> list[dict[str, Any]]:
         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:
@@ -481,7 +494,7 @@ async def ping_servers(self) -> list[dict[str, Any]]:
 
         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}
@@ -496,7 +509,7 @@ async def list_resources(self) -> list[dict[str, Any]]:
 
         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)
@@ -516,7 +529,7 @@ async def list_prompts(self) -> 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)
@@ -643,7 +656,7 @@ async def _concurrent_close(self, transport_items: list[tuple[str, MCPBaseTransp
             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
@@ -666,7 +679,8 @@ async def _concurrent_close(self, transport_items: list[tuple[str, MCPBaseTransp
                 # 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:
@@ -675,7 +689,7 @@ async def _sequential_close(self, transport_items: list[tuple[str, MCPBaseTransp
             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))
@@ -767,7 +781,7 @@ async def health_check(self) -> dict[str, Any]:
 
         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,

src/chuk_tool_processor/mcp/transport/__init__.py

@@ -11,6 +11,12 @@
 
 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 @@
     "StdioTransport",
     "SSETransport",
     "HTTPStreamableTransport",
+    "TimeoutConfig",
+    "TransportMetrics",
+    "ServerInfo",
+    "HeadersConfig",
 ]