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

chuk_tool_processor/mcp/transport/http_streamable_transport.py

@@ -16,12 +16,15 @@ from chuk_mcp.protocol.messages import (  # type: ignore[import-untyped]
     send_tools_call,
     send_tools_list,
 )
+from chuk_mcp.transports.http.parameters import StreamableHTTPParameters  # type: ignore[import-untyped]
 
 # Import chuk-mcp HTTP transport components
-from chuk_mcp.transports.http import http_client  # type: ignore[import-untyped]
-from chuk_mcp.transports.http.parameters import StreamableHTTPParameters  # type: ignore[import-untyped]
+from chuk_mcp.transports.http.transport import (
+    StreamableHTTPTransport as ChukHTTPTransport,  # type: ignore[import-untyped]
+)
 
 from .base_transport import MCPBaseTransport
+from .models import TimeoutConfig, TransportMetrics
 
 logger = logging.getLogger(__name__)
 
@@ -38,11 +41,13 @@ class HTTPStreamableTransport(MCPBaseTransport):
         self,
         url: str,
         api_key: str | None = None,
-        headers: dict[str, str] | None = None,  # NEW: Headers support
+        headers: dict[str, str] | None = None,
         connection_timeout: float = 30.0,
         default_timeout: float = 30.0,
         session_id: str | None = None,
         enable_metrics: bool = True,
+        oauth_refresh_callback: Any | None = None,
+        timeout_config: TimeoutConfig | None = None,
     ):
         """
         Initialize HTTP Streamable transport with enhanced configuration.
@@ -50,11 +55,13 @@ class HTTPStreamableTransport(MCPBaseTransport):
         Args:
             url: HTTP server URL (should end with /mcp)
             api_key: Optional API key for authentication
-            headers: Optional custom headers (NEW)
-            connection_timeout: Timeout for initial connection
-            default_timeout: Default timeout for operations
+            headers: Optional custom headers
+            connection_timeout: Timeout for initial connection (overrides timeout_config.connect)
+            default_timeout: Default timeout for operations (overrides timeout_config.operation)
             session_id: Optional session ID for stateful connections
             enable_metrics: Whether to track performance metrics
+            oauth_refresh_callback: Optional async callback to refresh OAuth tokens
+            timeout_config: Optional timeout configuration model with connect/operation/quick/shutdown
         """
         # Ensure URL points to the /mcp endpoint
         if not url.endswith("/mcp"):
@@ -63,11 +70,18 @@ class HTTPStreamableTransport(MCPBaseTransport):
             self.url = url
 
         self.api_key = api_key
-        self.configured_headers = headers or {}  # NEW: Store configured headers
-        self.connection_timeout = connection_timeout
-        self.default_timeout = default_timeout
+        self.configured_headers = headers or {}
         self.session_id = session_id
         self.enable_metrics = enable_metrics
+        self.oauth_refresh_callback = oauth_refresh_callback
+
+        # Use timeout config or create from individual parameters
+        if timeout_config is None:
+            timeout_config = TimeoutConfig(connect=connection_timeout, operation=default_timeout)
+
+        self.timeout_config = timeout_config
+        self.connection_timeout = timeout_config.connect
+        self.default_timeout = timeout_config.operation
 
         logger.debug("HTTP Streamable transport initialized with URL: %s", self.url)
         if self.api_key:
@@ -78,7 +92,7 @@ class HTTPStreamableTransport(MCPBaseTransport):
             logger.debug("Session ID configured: %s", self.session_id)
 
         # State tracking (enhanced like SSE)
-        self._http_context = None
+        self._http_transport = None
         self._read_stream = None
         self._write_stream = None
         self._initialized = False
@@ -88,20 +102,8 @@ class HTTPStreamableTransport(MCPBaseTransport):
         self._consecutive_failures = 0
         self._max_consecutive_failures = 3
 
-        # Performance metrics (enhanced like SSE)
-        self._metrics = {
-            "total_calls": 0,
-            "successful_calls": 0,
-            "failed_calls": 0,
-            "total_time": 0.0,
-            "avg_response_time": 0.0,
-            "last_ping_time": None,
-            "initialization_time": None,
-            "connection_resets": 0,
-            "stream_errors": 0,
-            "connection_errors": 0,  # NEW
-            "recovery_attempts": 0,  # NEW
-        }
+        # Performance metrics (enhanced like SSE) - use Pydantic model
+        self._metrics = TransportMetrics() if enable_metrics else None
 
     def _get_headers(self) -> dict[str, str]:
         """Get headers with authentication and custom headers (like SSE)."""
@@ -115,8 +117,9 @@ class HTTPStreamableTransport(MCPBaseTransport):
         if self.configured_headers:
             headers.update(self.configured_headers)
 
-        # Add API key as Bearer token if provided
-        if self.api_key:
+        # Add API key as Bearer token if provided and no Authorization header exists
+        # This prevents clobbering OAuth tokens from configured_headers
+        if self.api_key and "Authorization" not in headers:
             headers["Authorization"] = f"Bearer {self.api_key}"
 
         # Add session ID if provided
@@ -130,7 +133,7 @@ class HTTPStreamableTransport(MCPBaseTransport):
         try:
             import httpx
 
-            async with httpx.AsyncClient(timeout=5.0) as client:
+            async with httpx.AsyncClient(timeout=self.timeout_config.quick) as client:
                 # Test basic connectivity to base URL
                 base_url = self.url.replace("/mcp", "")
                 response = await client.get(f"{base_url}/health", headers=self._get_headers())
@@ -159,33 +162,38 @@ class HTTPStreamableTransport(MCPBaseTransport):
             headers = self._get_headers()
             logger.debug("Using headers: %s", list(headers.keys()))
 
-            # Create StreamableHTTPParameters with proper configuration
+            # Create StreamableHTTPParameters with minimal configuration
+            # NOTE: Keep params minimal - extra params can break message routing
             http_params = StreamableHTTPParameters(
                 url=self.url,
                 timeout=self.default_timeout,
                 headers=headers,
-                bearer_token=None,  # Don't duplicate auth - it's in headers
-                session_id=self.session_id,
                 enable_streaming=True,
-                max_concurrent_requests=5,
-                max_retries=2,
-                retry_delay=1.0,
-                user_agent="chuk-tool-processor/1.0.0",
             )
 
-            # Create and enter the HTTP context
-            self._http_context = http_client(http_params)
+            # Create and store transport (will be managed via async with in parent scope)
+            self._http_transport = ChukHTTPTransport(http_params)
 
+            # IMPORTANT: Must use async with for proper stream setup
             logger.debug("Establishing HTTP connection...")
-            self._read_stream, self._write_stream = await asyncio.wait_for(
-                self._http_context.__aenter__(), timeout=self.connection_timeout
+            self._http_context_entered = await asyncio.wait_for(
+                self._http_transport.__aenter__(), timeout=self.connection_timeout
             )
 
+            # Get streams after context entered
+            self._read_stream, self._write_stream = await self._http_transport.get_streams()
+
+            # Give the transport's message handler task time to start
+            await asyncio.sleep(0.1)
+
             # Enhanced MCP initialize sequence
             logger.debug("Sending MCP initialize request...")
             init_start = time.time()
 
-            await asyncio.wait_for(send_initialize(self._read_stream, self._write_stream), timeout=self.default_timeout)
+            await asyncio.wait_for(
+                send_initialize(self._read_stream, self._write_stream, timeout=self.default_timeout),
+                timeout=self.default_timeout,
+            )
 
             init_time = time.time() - init_start
             logger.debug("MCP initialize completed in %.3fs", init_time)
@@ -193,9 +201,11 @@ class HTTPStreamableTransport(MCPBaseTransport):
             # Verify connection with ping (enhanced like SSE)
             logger.debug("Verifying connection with ping...")
             ping_start = time.time()
+            # Use connect timeout for initial ping - some servers (like Notion) are slow
+            ping_timeout = self.timeout_config.connect
             ping_success = await asyncio.wait_for(
-                send_ping(self._read_stream, self._write_stream),
-                timeout=10.0,  # Longer timeout for initial ping
+                send_ping(self._read_stream, self._write_stream, timeout=ping_timeout),
+                timeout=ping_timeout,
             )
             ping_time = time.time() - ping_start
 
@@ -205,9 +215,9 @@ class HTTPStreamableTransport(MCPBaseTransport):
                 self._consecutive_failures = 0
 
                 total_init_time = time.time() - start_time
-                if self.enable_metrics:
-                    self._metrics["initialization_time"] = total_init_time
-                    self._metrics["last_ping_time"] = ping_time
+                if self.enable_metrics and self._metrics:
+                    self._metrics.initialization_time = total_init_time
+                    self._metrics.last_ping_time = ping_time
 
                 logger.debug(
                     "HTTP Streamable transport initialized successfully in %.3fs (ping: %.3fs)",
@@ -216,31 +226,31 @@ class HTTPStreamableTransport(MCPBaseTransport):
                 )
                 return True
             else:
-                logger.warning("HTTP connection established but ping failed")
+                logger.debug("HTTP connection established but ping failed")
                 # Still consider it initialized since connection was established
                 self._initialized = True
                 self._consecutive_failures = 1  # Mark one failure
-                if self.enable_metrics:
-                    self._metrics["initialization_time"] = time.time() - start_time
+                if self.enable_metrics and self._metrics:
+                    self._metrics.initialization_time = time.time() - start_time
                 return True
 
         except TimeoutError:
             logger.error("HTTP Streamable initialization timed out after %ss", self.connection_timeout)
             await self._cleanup()
-            if self.enable_metrics:
-                self._metrics["connection_errors"] += 1
-            return False
+            if self.enable_metrics and self._metrics:
+                self._metrics.connection_errors += 1
+            raise  # Re-raise for OAuth error detection in mcp-cli
         except Exception as e:
             logger.error("Error initializing HTTP Streamable transport: %s", e, exc_info=True)
             await self._cleanup()
-            if self.enable_metrics:
-                self._metrics["connection_errors"] += 1
-            return False
+            if self.enable_metrics and self._metrics:
+                self._metrics.connection_errors += 1
+            raise  # Re-raise for OAuth error detection in mcp-cli
 
     async def _attempt_recovery(self) -> bool:
         """Attempt to recover from connection issues (NEW - like SSE resilience)."""
-        if self.enable_metrics:
-            self._metrics["recovery_attempts"] += 1
+        if self.enable_metrics and self._metrics:
+            self._metrics.recovery_attempts += 1
 
         logger.debug("Attempting HTTP connection recovery...")
 
@@ -260,21 +270,21 @@ class HTTPStreamableTransport(MCPBaseTransport):
             return
 
         # Enhanced metrics logging (like SSE)
-        if self.enable_metrics and self._metrics["total_calls"] > 0:
-            success_rate = self._metrics["successful_calls"] / self._metrics["total_calls"] * 100
+        if self.enable_metrics and self._metrics and self._metrics.total_calls > 0:
+            success_rate = self._metrics.successful_calls / self._metrics.total_calls * 100
             logger.debug(
                 "HTTP Streamable transport closing - Calls: %d, Success: %.1f%%, "
                 "Avg time: %.3fs, Recoveries: %d, Errors: %d",
-                self._metrics["total_calls"],
+                self._metrics.total_calls,
                 success_rate,
-                self._metrics["avg_response_time"],
-                self._metrics["recovery_attempts"],
-                self._metrics["connection_errors"],
+                self._metrics.avg_response_time,
+                self._metrics.recovery_attempts,
+                self._metrics.connection_errors,
             )
 
         try:
-            if self._http_context is not None:
-                await self._http_context.__aexit__(None, None, None)
+            if self._http_transport is not None:
+                await self._http_transport.__aexit__(None, None, None)
                 logger.debug("HTTP Streamable context closed")
 
         except Exception as e:
@@ -284,7 +294,7 @@ class HTTPStreamableTransport(MCPBaseTransport):
 
     async def _cleanup(self) -> None:
         """Enhanced cleanup with state reset."""
-        self._http_context = None
+        self._http_transport = None
         self._read_stream = None
         self._write_stream = None
         self._initialized = False
@@ -292,13 +302,14 @@ class HTTPStreamableTransport(MCPBaseTransport):
     async def send_ping(self) -> bool:
         """Enhanced ping with health monitoring (like SSE)."""
         if not self._initialized or not self._read_stream:
-            logger.error("Cannot send ping: transport not initialized")
+            logger.debug("Cannot send ping: transport not initialized")
             return False
 
         start_time = time.time()
         try:
             result = await asyncio.wait_for(
-                send_ping(self._read_stream, self._write_stream), timeout=self.default_timeout
+                send_ping(self._read_stream, self._write_stream, timeout=self.default_timeout),
+                timeout=self.default_timeout,
             )
 
             success = bool(result)
@@ -309,9 +320,9 @@ class HTTPStreamableTransport(MCPBaseTransport):
             else:
                 self._consecutive_failures += 1
 
-            if self.enable_metrics:
+            if self.enable_metrics and self._metrics:
                 ping_time = time.time() - start_time
-                self._metrics["last_ping_time"] = ping_time
+                self._metrics.last_ping_time = ping_time
                 logger.debug("HTTP Streamable ping completed in %.3fs: %s", ping_time, success)
 
             return success
@@ -322,8 +333,8 @@ class HTTPStreamableTransport(MCPBaseTransport):
         except Exception as e:
             logger.error("HTTP Streamable ping failed: %s", e)
             self._consecutive_failures += 1
-            if self.enable_metrics:
-                self._metrics["stream_errors"] += 1
+            if self.enable_metrics and self._metrics:
+                self._metrics.stream_errors += 1
             return False
 
     def is_connected(self) -> bool:
@@ -341,18 +352,29 @@ class HTTPStreamableTransport(MCPBaseTransport):
     async def get_tools(self) -> list[dict[str, Any]]:
         """Enhanced tools retrieval with error handling."""
         if not self._initialized:
-            logger.error("Cannot get tools: transport not initialized")
+            logger.debug("Cannot get tools: transport not initialized")
             return []
 
         start_time = time.time()
         try:
             tools_response = await asyncio.wait_for(
-                send_tools_list(self._read_stream, self._write_stream), timeout=self.default_timeout
+                send_tools_list(self._read_stream, self._write_stream, timeout=self.default_timeout),
+                timeout=self.default_timeout,
             )
 
-            # Normalize response
-            if isinstance(tools_response, dict):
+            # Normalize response - handle multiple formats including Pydantic models
+            # 1. Check if it's a Pydantic model with tools attribute (e.g., ListToolsResult from chuk_mcp)
+            if hasattr(tools_response, "tools"):
+                tools = tools_response.tools
+                # Convert Pydantic Tool models to dicts if needed
+                if tools and len(tools) > 0 and hasattr(tools[0], "model_dump"):
+                    tools = [t.model_dump() for t in tools]
+                elif tools and len(tools) > 0 and hasattr(tools[0], "dict"):
+                    tools = [t.dict() for t in tools]
+            # 2. Check if it's a dict with "tools" key
+            elif isinstance(tools_response, dict):
                 tools = tools_response.get("tools", [])
+            # 3. Check if it's already a list
             elif isinstance(tools_response, list):
                 tools = tools_response
             else:
@@ -375,8 +397,8 @@ class HTTPStreamableTransport(MCPBaseTransport):
         except Exception as e:
             logger.error("Error getting tools: %s", e)
             self._consecutive_failures += 1
-            if self.enable_metrics:
-                self._metrics["stream_errors"] += 1
+            if self.enable_metrics and self._metrics:
+                self._metrics.stream_errors += 1
             return []
 
     async def call_tool(
@@ -389,8 +411,8 @@ class HTTPStreamableTransport(MCPBaseTransport):
         tool_timeout = timeout or self.default_timeout
         start_time = time.time()
 
-        if self.enable_metrics:
-            self._metrics["total_calls"] += 1
+        if self.enable_metrics and self._metrics:
+            self._metrics.total_calls += 1
 
         try:
             logger.debug("Calling tool '%s' with timeout %ss", tool_name, tool_timeout)
@@ -410,9 +432,44 @@ class HTTPStreamableTransport(MCPBaseTransport):
             response_time = time.time() - start_time
             result = self._normalize_mcp_response(raw_response)
 
+            # NEW: Check for OAuth errors and attempt refresh if callback is available
+            if result.get("isError", False) and self._is_oauth_error(result.get("error", "")):
+                logger.warning("OAuth error detected: %s", result.get("error"))
+
+                if self.oauth_refresh_callback:
+                    logger.debug("Attempting OAuth token refresh...")
+                    try:
+                        # Call the refresh callback
+                        new_headers = await self.oauth_refresh_callback()
+
+                        if new_headers and "Authorization" in new_headers:
+                            # Update configured headers with new token
+                            self.configured_headers.update(new_headers)
+                            logger.debug("OAuth token refreshed, reconnecting...")
+
+                            # Reconnect with new token
+                            if await self._attempt_recovery():
+                                logger.debug("Retrying tool call after token refresh...")
+                                # Retry the tool call once with new token
+                                raw_response = await asyncio.wait_for(
+                                    send_tools_call(self._read_stream, self._write_stream, tool_name, arguments),
+                                    timeout=tool_timeout,
+                                )
+                                result = self._normalize_mcp_response(raw_response)
+                                logger.debug("Tool call retry completed")
+                            else:
+                                logger.error("Failed to reconnect after token refresh")
+                        else:
+                            logger.warning("Token refresh did not return valid Authorization header")
+                    except Exception as refresh_error:
+                        logger.error("OAuth token refresh failed: %s", refresh_error)
+                else:
+                    logger.warning("OAuth error detected but no refresh callback configured")
+
             # Reset failure count on success
-            self._consecutive_failures = 0
-            self._last_successful_ping = time.time()  # Update health timestamp
+            if not result.get("isError", False):
+                self._consecutive_failures = 0
+                self._last_successful_ping = time.time()  # Update health timestamp
 
             if self.enable_metrics:
                 self._update_metrics(response_time, not result.get("isError", False))
@@ -438,17 +495,17 @@ class HTTPStreamableTransport(MCPBaseTransport):
         except Exception as e:
             response_time = time.time() - start_time
             self._consecutive_failures += 1
-            if self.enable_metrics:
+            if self.enable_metrics and self._metrics:
                 self._update_metrics(response_time, False)
-                self._metrics["stream_errors"] += 1
+                self._metrics.stream_errors += 1
 
             # Enhanced connection error detection
             error_str = str(e).lower()
             if any(indicator in error_str for indicator in ["connection", "disconnected", "broken pipe", "eof"]):
                 logger.warning("Connection error detected: %s", e)
                 self._initialized = False
-                if self.enable_metrics:
-                    self._metrics["connection_errors"] += 1
+                if self.enable_metrics and self._metrics:
+                    self._metrics.connection_errors += 1
 
             error_msg = f"Tool execution failed: {str(e)}"
             logger.error("Tool '%s' error: %s", tool_name, error_msg)
@@ -456,14 +513,41 @@ class HTTPStreamableTransport(MCPBaseTransport):
 
     def _update_metrics(self, response_time: float, success: bool) -> None:
         """Enhanced metrics tracking (like SSE)."""
-        if success:
-            self._metrics["successful_calls"] += 1
-        else:
-            self._metrics["failed_calls"] += 1
+        if not self._metrics:
+            return
+
+        self._metrics.update_call_metrics(response_time, success)
+
+    def _is_oauth_error(self, error_msg: str) -> bool:
+        """
+        Detect if error is OAuth-related per RFC 6750 and MCP OAuth spec.
 
-        self._metrics["total_time"] += response_time
-        if self._metrics["total_calls"] > 0:
-            self._metrics["avg_response_time"] = self._metrics["total_time"] / self._metrics["total_calls"]
+        Checks for:
+        - RFC 6750 Section 3.1 Bearer token errors (invalid_token, insufficient_scope)
+        - OAuth 2.1 token refresh errors (invalid_grant)
+        - MCP spec OAuth validation failures (401/403 responses)
+        """
+        if not error_msg:
+            return False
+
+        error_lower = error_msg.lower()
+        oauth_indicators = [
+            # RFC 6750 Section 3.1 - Standard Bearer token errors
+            "invalid_token",  # Token expired, revoked, malformed, or invalid
+            "insufficient_scope",  # Request requires higher privileges (403 Forbidden)
+            # OAuth 2.1 token refresh errors
+            "invalid_grant",  # Refresh token errors
+            # MCP spec - OAuth validation failures (401 Unauthorized)
+            "oauth validation",
+            "unauthorized",
+            # Common OAuth error descriptions
+            "expired token",
+            "token expired",
+            "authentication failed",
+            "invalid access token",
+        ]
+
+        return any(indicator in error_lower for indicator in oauth_indicators)
 
     async def list_resources(self) -> dict[str, Any]:
         """Enhanced resource listing with error handling."""
@@ -544,7 +628,10 @@ class HTTPStreamableTransport(MCPBaseTransport):
 
     def get_metrics(self) -> dict[str, Any]:
         """Enhanced metrics with health information."""
-        metrics = self._metrics.copy()
+        if not self._metrics:
+            return {}
+
+        metrics = self._metrics.to_dict()
         metrics.update(
             {
                 "is_connected": self.is_connected(),
@@ -557,22 +644,20 @@ class HTTPStreamableTransport(MCPBaseTransport):
 
     def reset_metrics(self) -> None:
         """Enhanced metrics reset preserving health state."""
-        preserved_init_time = self._metrics.get("initialization_time")
-        preserved_last_ping = self._metrics.get("last_ping_time")
-
-        self._metrics = {
-            "total_calls": 0,
-            "successful_calls": 0,
-            "failed_calls": 0,
-            "total_time": 0.0,
-            "avg_response_time": 0.0,
-            "last_ping_time": preserved_last_ping,
-            "initialization_time": preserved_init_time,
-            "connection_resets": self._metrics.get("connection_resets", 0),
-            "stream_errors": 0,
-            "connection_errors": 0,
-            "recovery_attempts": 0,
-        }
+        if not self._metrics:
+            return
+
+        # Preserve important historical values
+        preserved_init_time = self._metrics.initialization_time
+        preserved_last_ping = self._metrics.last_ping_time
+        preserved_resets = self._metrics.connection_resets
+
+        # Create new metrics instance with preserved values
+        self._metrics = TransportMetrics(
+            initialization_time=preserved_init_time,
+            last_ping_time=preserved_last_ping,
+            connection_resets=preserved_resets,
+        )
 
     def get_streams(self) -> list[tuple]:
         """Enhanced streams access with connection check."""

chuk_tool_processor/mcp/transport/models.py

@@ -0,0 +1,100 @@
+# chuk_tool_processor/mcp/transport/models.py
+"""
+Pydantic models for MCP transport configuration and metrics.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TimeoutConfig(BaseModel):
+    """
+    Unified timeout configuration for all MCP operations.
+
+    Just 4 simple, logical timeout categories:
+    - connect: Connection establishment, initialization, session discovery (30s default)
+    - operation: Normal operations like tool calls, listing resources (30s default)
+    - quick: Fast health checks and pings (5s default)
+    - shutdown: Cleanup and shutdown operations (2s default)
+    """
+
+    connect: float = Field(
+        default=30.0, description="Timeout for connection establishment, initialization, and session discovery"
+    )
+    operation: float = Field(
+        default=30.0, description="Timeout for normal operations (tool calls, listing tools/resources/prompts)"
+    )
+    quick: float = Field(default=5.0, description="Timeout for quick health checks and pings")
+    shutdown: float = Field(default=2.0, description="Timeout for shutdown and cleanup operations")
+
+
+class TransportMetrics(BaseModel):
+    """Performance and connection metrics for transports."""
+
+    model_config = {"validate_assignment": True}
+
+    total_calls: int = Field(default=0, description="Total number of calls made")
+    successful_calls: int = Field(default=0, description="Number of successful calls")
+    failed_calls: int = Field(default=0, description="Number of failed calls")
+    total_time: float = Field(default=0.0, description="Total time spent on calls")
+    avg_response_time: float = Field(default=0.0, description="Average response time")
+    last_ping_time: float | None = Field(default=None, description="Time taken for last ping")
+    initialization_time: float | None = Field(default=None, description="Time taken for initialization")
+    connection_resets: int = Field(default=0, description="Number of connection resets")
+    stream_errors: int = Field(default=0, description="Number of stream errors")
+    connection_errors: int = Field(default=0, description="Number of connection errors")
+    recovery_attempts: int = Field(default=0, description="Number of recovery attempts")
+    session_discoveries: int = Field(default=0, description="Number of session discoveries (SSE)")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format."""
+        return self.model_dump()
+
+    def update_call_metrics(self, response_time: float, success: bool) -> None:
+        """Update metrics after a call."""
+        if success:
+            self.successful_calls += 1
+        else:
+            self.failed_calls += 1
+
+        self.total_time += response_time
+        if self.total_calls > 0:
+            self.avg_response_time = self.total_time / self.total_calls
+
+
+class ServerInfo(BaseModel):
+    """Information about a server in StreamManager."""
+
+    id: int = Field(description="Server ID")
+    name: str = Field(description="Server name")
+    tools: int = Field(description="Number of tools available")
+    status: str = Field(description="Server status (Up/Down)")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format."""
+        return self.model_dump()
+
+
+class HeadersConfig(BaseModel):
+    """Configuration for HTTP headers."""
+
+    headers: dict[str, str] = Field(default_factory=dict, description="Custom HTTP headers")
+
+    def get_headers(self) -> dict[str, str]:
+        """Get headers as dict."""
+        return self.headers.copy()
+
+    def update_headers(self, new_headers: dict[str, str]) -> None:
+        """Update headers with new values."""
+        self.headers.update(new_headers)
+
+    def has_authorization(self) -> bool:
+        """Check if Authorization header is present."""
+        return "Authorization" in self.headers
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format."""
+        return self.model_dump()