fast-agent-mcp 0.3.7__py3-none-any.whl → 0.3.9__py3-none-any.whl

fast_agent/mcp/mcp_connection_manager.py

@@ -21,10 +21,9 @@ from mcp.client.sse import sse_client
 from mcp.client.stdio import (
     StdioServerParameters,
     get_default_environment,
-    stdio_client,
 )
-from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
-from mcp.types import JSONRPCMessage, ServerCapabilities
+from mcp.client.streamable_http import GetSessionIdCallback
+from mcp.types import Implementation, JSONRPCMessage, ServerCapabilities
 
 from fast_agent.config import MCPServerSettings
 from fast_agent.context_dependent import ContextDependent
@@ -34,6 +33,9 @@ from fast_agent.event_progress import ProgressAction
 from fast_agent.mcp.logger_textio import get_stderr_handler
 from fast_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
 from fast_agent.mcp.oauth_client import build_oauth_provider
+from fast_agent.mcp.stdio_tracking_simple import tracking_stdio_client
+from fast_agent.mcp.streamable_http_tracking import tracking_streamablehttp_client
+from fast_agent.mcp.transport_tracking import TransportChannelMetrics
 
 if TYPE_CHECKING:
     from fast_agent.context import Context
@@ -107,6 +109,14 @@ class ServerConnection:
 
         # Server instructions from initialization
         self.server_instructions: str | None = None
+        self.server_capabilities: ServerCapabilities | None = None
+        self.server_implementation: Implementation | None = None
+        self.client_capabilities: dict | None = None
+        self.server_instructions_available: bool = False
+        self.server_instructions_enabled: bool = server_config.include_instructions if server_config else True
+        self.session_id: str | None = None
+        self._get_session_id_cb: GetSessionIdCallback | None = None
+        self.transport_metrics: TransportChannelMetrics | None = None
 
     def is_healthy(self) -> bool:
         """Check if the server connection is healthy and ready to use."""
@@ -138,15 +148,32 @@ class ServerConnection:
         result = await self.session.initialize()
 
         self.server_capabilities = result.capabilities
+        # InitializeResult exposes server info via `serverInfo`; keep fallback for older fields
+        implementation = getattr(result, "serverInfo", None)
+        if implementation is None:
+            implementation = getattr(result, "implementation", None)
+        self.server_implementation = implementation
+
+        raw_instructions = getattr(result, "instructions", None)
+        self.server_instructions_available = bool(raw_instructions)
 
         # Store instructions if provided by the server and enabled in config
         if self.server_config.include_instructions:
-            self.server_instructions = getattr(result, 'instructions', None)
+            self.server_instructions = raw_instructions
             if self.server_instructions:
-                logger.debug(f"{self.server_name}: Received server instructions", data={"instructions": self.server_instructions})
+                logger.debug(
+                    f"{self.server_name}: Received server instructions",
+                    data={"instructions": self.server_instructions},
+                )
         else:
             self.server_instructions = None
-            logger.debug(f"{self.server_name}: Server instructions disabled by configuration")
+            if self.server_instructions_available:
+                logger.debug(
+                    f"{self.server_name}: Server instructions disabled by configuration",
+                    data={"instructions": raw_instructions},
+                )
+            else:
+                logger.debug(f"{self.server_name}: No server instructions provided")
 
         # If there's an init hook, run it
 
@@ -175,10 +202,15 @@ class ServerConnection:
         )
 
         session = self._client_session_factory(
-            read_stream, send_stream, read_timeout, server_config=self.server_config
+            read_stream,
+            send_stream,
+            read_timeout,
+            server_config=self.server_config,
+            transport_metrics=self.transport_metrics,
         )
 
         self.session = session
+        self.client_capabilities = getattr(session, "client_capabilities", None)
 
         return session
 
@@ -192,11 +224,30 @@ async def _server_lifecycle_task(server_conn: ServerConnection) -> None:
     try:
         transport_context = server_conn._transport_context_factory()
 
-        async with transport_context as (read_stream, write_stream, _):
+        async with transport_context as (read_stream, write_stream, get_session_id_cb):
+            server_conn._get_session_id_cb = get_session_id_cb
+
+            if get_session_id_cb is not None:
+                try:
+                    server_conn.session_id = get_session_id_cb()
+                except Exception:
+                    logger.debug(f"{server_name}: Unable to retrieve session id from transport")
+            elif server_conn.server_config.transport == "stdio":
+                server_conn.session_id = "local"
+
             server_conn.create_session(read_stream, write_stream)
 
             async with server_conn.session:
                 await server_conn.initialize_session()
+
+                if get_session_id_cb is not None:
+                    try:
+                        server_conn.session_id = get_session_id_cb() or server_conn.session_id
+                    except Exception:
+                        logger.debug(f"{server_name}: Unable to refresh session id after init")
+                elif server_conn.server_config.transport == "stdio":
+                    server_conn.session_id = "local"
+
                 await server_conn.wait_for_shutdown_request()
 
     except HTTPStatusError as http_exc:
@@ -353,6 +404,8 @@ class MCPConnectionManager(ContextDependent):
 
         logger.debug(f"{server_name}: Found server configuration=", data=config.model_dump())
 
+        transport_metrics = TransportChannelMetrics() if config.transport in ("http", "stdio") else None
+
         def transport_context_factory():
             if config.transport == "stdio":
                 if not config.command:
@@ -369,7 +422,11 @@ class MCPConnectionManager(ContextDependent):
                 error_handler = get_stderr_handler(server_name)
                 # Explicitly ensure we're using our custom logger for stderr
                 logger.debug(f"{server_name}: Creating stdio client with custom error handler")
-                return _add_none_to_context(stdio_client(server_params, errlog=error_handler))
+
+                channel_hook = transport_metrics.record_event if transport_metrics else None
+                return _add_none_to_context(
+                    tracking_stdio_client(server_params, channel_hook=channel_hook, errlog=error_handler)
+                )
             elif config.transport == "sse":
                 if not config.url:
                     raise ValueError(
@@ -401,7 +458,23 @@ class MCPConnectionManager(ContextDependent):
                 if oauth_auth is not None:
                     headers.pop("Authorization", None)
                     headers.pop("X-HF-Authorization", None)
-                return streamablehttp_client(config.url, headers, auth=oauth_auth)
+                channel_hook = None
+                if transport_metrics is not None:
+                    def channel_hook(event):
+                        try:
+                            transport_metrics.record_event(event)
+                        except Exception:  # pragma: no cover - defensive guard
+                            logger.debug(
+                                "%s: transport metrics hook failed", server_name,
+                                exc_info=True,
+                            )
+
+                return tracking_streamablehttp_client(
+                    config.url,
+                    headers,
+                    auth=oauth_auth,
+                    channel_hook=channel_hook,
+                )
             else:
                 raise ValueError(f"Unsupported transport: {config.transport}")
 
@@ -412,6 +485,9 @@ class MCPConnectionManager(ContextDependent):
             client_session_factory=client_session_factory,
         )
 
+        if transport_metrics is not None:
+            server_conn.transport_metrics = transport_metrics
+
         async with self._lock:
             # Check if already running
             if server_name in self.running_servers:

fast_agent/mcp/stdio_tracking_simple.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, AsyncGenerator, Callable
+
+from mcp.client.stdio import StdioServerParameters, stdio_client
+
+from fast_agent.mcp.transport_tracking import ChannelEvent
+
+if TYPE_CHECKING:
+    from anyio.abc import ObjectReceiveStream, ObjectSendStream
+    from mcp.shared.message import SessionMessage
+
+logger = logging.getLogger(__name__)
+
+ChannelHook = Callable[[ChannelEvent], None]
+
+
+@asynccontextmanager
+async def tracking_stdio_client(
+    server_params: StdioServerParameters,
+    *,
+    channel_hook: ChannelHook | None = None,
+    errlog: Callable[[str], None] | None = None,
+) -> AsyncGenerator[
+    tuple[ObjectReceiveStream[SessionMessage | Exception], ObjectSendStream[SessionMessage]], None
+]:
+    """Context manager for stdio client with basic connection tracking."""
+
+    def emit_channel_event(event_type: str, detail: str | None = None) -> None:
+        if channel_hook is None:
+            return
+        try:
+            channel_hook(
+                ChannelEvent(
+                    channel="stdio",
+                    event_type=event_type,  # type: ignore[arg-type]
+                    detail=detail,
+                )
+            )
+        except Exception:  # pragma: no cover - hook errors must not break transport
+            logger.exception("Channel hook raised an exception")
+
+    try:
+        # Emit connection event
+        emit_channel_event("connect")
+
+        # Use the original stdio_client without stream interception
+        async with stdio_client(server_params, errlog=errlog) as (read_stream, write_stream):
+            yield read_stream, write_stream
+
+    except Exception as exc:
+        # Emit error event
+        emit_channel_event("error", detail=str(exc))
+        raise
+    finally:
+        # Emit disconnection event
+        emit_channel_event("disconnect")

fast_agent/mcp/streamable_http_tracking.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, AsyncGenerator, Awaitable, Callable
+
+import anyio
+import httpx
+from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
+from mcp.client.streamable_http import (
+    RequestContext,
+    RequestId,
+    StreamableHTTPTransport,
+    StreamWriter,
+)
+from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
+from mcp.shared.message import SessionMessage
+from mcp.types import JSONRPCError, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
+
+from fast_agent.mcp.transport_tracking import ChannelEvent, ChannelName
+
+if TYPE_CHECKING:
+    from datetime import timedelta
+
+    from anyio.abc import ObjectReceiveStream, ObjectSendStream
+
+logger = logging.getLogger(__name__)
+
+ChannelHook = Callable[[ChannelEvent], None]
+
+
+class ChannelTrackingStreamableHTTPTransport(StreamableHTTPTransport):
+    """Streamable HTTP transport that emits channel events before dispatching."""
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        headers: dict[str, str] | None = None,
+        timeout: float | timedelta = 30,
+        sse_read_timeout: float | timedelta = 60 * 5,
+        auth: httpx.Auth | None = None,
+        channel_hook: ChannelHook | None = None,
+    ) -> None:
+        super().__init__(
+            url,
+            headers=headers,
+            timeout=timeout,
+            sse_read_timeout=sse_read_timeout,
+            auth=auth,
+        )
+        self._channel_hook = channel_hook
+
+    def _emit_channel_event(
+        self,
+        channel: ChannelName,
+        event_type: str,
+        *,
+        message: JSONRPCMessage | None = None,
+        raw_event: str | None = None,
+        detail: str | None = None,
+        status_code: int | None = None,
+    ) -> None:
+        if self._channel_hook is None:
+            return
+        try:
+            self._channel_hook(
+                ChannelEvent(
+                    channel=channel,
+                    event_type=event_type,  # type: ignore[arg-type]
+                    message=message,
+                    raw_event=raw_event,
+                    detail=detail,
+                    status_code=status_code,
+                )
+            )
+        except Exception:  # pragma: no cover - hook errors must not break transport
+            logger.exception("Channel hook raised an exception")
+
+    async def _handle_json_response(  # type: ignore[override]
+        self,
+        response: httpx.Response,
+        read_stream_writer: StreamWriter,
+        is_initialization: bool = False,
+    ) -> None:
+        try:
+            content = await response.aread()
+            message = JSONRPCMessage.model_validate_json(content)
+
+            if is_initialization:
+                self._maybe_extract_protocol_version_from_message(message)
+
+            self._emit_channel_event("post-json", "message", message=message)
+            await read_stream_writer.send(SessionMessage(message))
+        except Exception as exc:  # pragma: no cover - propagate to session
+            logger.exception("Error parsing JSON response")
+            await read_stream_writer.send(exc)
+
+    async def _handle_sse_event_with_channel(
+        self,
+        channel: ChannelName,
+        sse: ServerSentEvent,
+        read_stream_writer: StreamWriter,
+        original_request_id: RequestId | None = None,
+        resumption_callback: Callable[[str], Awaitable[None]] | None = None,
+        is_initialization: bool = False,
+    ) -> bool:
+        if sse.event != "message":
+            # Treat non-message events (e.g. ping) as keepalive notifications
+            self._emit_channel_event(channel, "keepalive", raw_event=sse.event or "keepalive")
+            return False
+
+        try:
+            message = JSONRPCMessage.model_validate_json(sse.data)
+            if is_initialization:
+                self._maybe_extract_protocol_version_from_message(message)
+
+            if original_request_id is not None and isinstance(
+                message.root, (JSONRPCResponse, JSONRPCError)
+            ):
+                message.root.id = original_request_id
+
+            self._emit_channel_event(channel, "message", message=message)
+            await read_stream_writer.send(SessionMessage(message))
+
+            if sse.id and resumption_callback:
+                await resumption_callback(sse.id)
+
+            return isinstance(message.root, (JSONRPCResponse, JSONRPCError))
+        except Exception as exc:  # pragma: no cover - propagate to session
+            logger.exception("Error parsing SSE message")
+            await read_stream_writer.send(exc)
+            return False
+
+    async def handle_get_stream(  # type: ignore[override]
+        self,
+        client: httpx.AsyncClient,
+        read_stream_writer: StreamWriter,
+    ) -> None:
+        if not self.session_id:
+            return
+
+        headers = self._prepare_request_headers(self.request_headers)
+        connected = False
+        try:
+            async with aconnect_sse(
+                client,
+                "GET",
+                self.url,
+                headers=headers,
+                timeout=httpx.Timeout(self.timeout, read=self.sse_read_timeout),
+            ) as event_source:
+                event_source.response.raise_for_status()
+                self._emit_channel_event("get", "connect")
+                connected = True
+                async for sse in event_source.aiter_sse():
+                    await self._handle_sse_event_with_channel(
+                        "get",
+                        sse,
+                        read_stream_writer,
+                    )
+        except Exception as exc:  # pragma: no cover - non fatal stream errors
+            logger.debug("GET stream error (non-fatal): %s", exc)
+            status_code = None
+            detail = str(exc)
+            if isinstance(exc, httpx.HTTPStatusError):
+                if exc.response is not None:
+                    status_code = exc.response.status_code
+                    reason = exc.response.reason_phrase or ""
+                    if not reason:
+                        try:
+                            reason = (exc.response.text or "").strip()
+                        except Exception:
+                            reason = ""
+                    detail = f"HTTP {status_code}: {reason or 'response'}"
+                else:
+                    status_code = exc.response.status_code if hasattr(exc, "response") else None
+            self._emit_channel_event("get", "error", detail=detail, status_code=status_code)
+        finally:
+            if connected:
+                self._emit_channel_event("get", "disconnect")
+
+    async def _handle_resumption_request(  # type: ignore[override]
+        self,
+        ctx: RequestContext,
+    ) -> None:
+        headers = self._prepare_request_headers(ctx.headers)
+        if ctx.metadata and ctx.metadata.resumption_token:
+            headers["last-event-id"] = ctx.metadata.resumption_token
+        else:  # pragma: no cover - defensive
+            raise ValueError("Resumption request requires a resumption token")
+
+        original_request_id: RequestId | None = None
+        if isinstance(ctx.session_message.message.root, JSONRPCRequest):
+            original_request_id = ctx.session_message.message.root.id
+
+        async with aconnect_sse(
+            ctx.client,
+            "GET",
+            self.url,
+            headers=headers,
+            timeout=httpx.Timeout(self.timeout, read=self.sse_read_timeout),
+        ) as event_source:
+            event_source.response.raise_for_status()
+            async for sse in event_source.aiter_sse():
+                is_complete = await self._handle_sse_event_with_channel(
+                    "resumption",
+                    sse,
+                    ctx.read_stream_writer,
+                    original_request_id,
+                    ctx.metadata.on_resumption_token_update if ctx.metadata else None,
+                )
+                if is_complete:
+                    await event_source.response.aclose()
+                    break
+
+    async def _handle_sse_response(  # type: ignore[override]
+        self,
+        response: httpx.Response,
+        ctx: RequestContext,
+        is_initialization: bool = False,
+    ) -> None:
+        try:
+            event_source = EventSource(response)
+            async for sse in event_source.aiter_sse():
+                is_complete = await self._handle_sse_event_with_channel(
+                    "post-sse",
+                    sse,
+                    ctx.read_stream_writer,
+                    resumption_callback=(
+                        ctx.metadata.on_resumption_token_update if ctx.metadata else None
+                    ),
+                    is_initialization=is_initialization,
+                )
+                if is_complete:
+                    await response.aclose()
+                    break
+        except Exception as exc:  # pragma: no cover - propagate to session
+            logger.exception("Error reading SSE stream")
+            await ctx.read_stream_writer.send(exc)
+
+
+@asynccontextmanager
+async def tracking_streamablehttp_client(
+    url: str,
+    headers: dict[str, str] | None = None,
+    *,
+    timeout: float | timedelta = 30,
+    sse_read_timeout: float | timedelta = 60 * 5,
+    terminate_on_close: bool = True,
+    httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
+    auth: httpx.Auth | None = None,
+    channel_hook: ChannelHook | None = None,
+) -> AsyncGenerator[
+    tuple[
+        ObjectReceiveStream[SessionMessage | Exception],
+        ObjectSendStream[SessionMessage],
+        Callable[[], str | None],
+    ],
+    None,
+]:
+    """Context manager mirroring streamablehttp_client with channel tracking."""
+
+    transport = ChannelTrackingStreamableHTTPTransport(
+        url,
+        headers=headers,
+        timeout=timeout,
+        sse_read_timeout=sse_read_timeout,
+        auth=auth,
+        channel_hook=channel_hook,
+    )
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](
+        0
+    )
+    write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
+
+    async with anyio.create_task_group() as tg:
+        try:
+            async with httpx_client_factory(
+                headers=transport.request_headers,
+                timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
+                auth=transport.auth,
+            ) as client:
+
+                def start_get_stream() -> None:
+                    tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
+
+                tg.start_soon(
+                    transport.post_writer,
+                    client,
+                    write_stream_reader,
+                    read_stream_writer,
+                    write_stream,
+                    start_get_stream,
+                    tg,
+                )
+
+                try:
+                    yield read_stream, write_stream, transport.get_session_id
+                finally:
+                    if transport.session_id and terminate_on_close:
+                        await transport.terminate_session(client)
+                    tg.cancel_scope.cancel()
+        finally:
+            await read_stream_writer.aclose()
+            await read_stream.aclose()
+            await write_stream_reader.aclose()
+            await write_stream.aclose()