fast-agent-mcp 0.4.7__py3-none-any.whl

fast_agent/mcp/sse_tracking.py

@@ -0,0 +1,287 @@
+"""SSE transport wrapper that emits channel events for UI display."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable
+from urllib.parse import parse_qs, urljoin, urlparse
+
+import anyio
+import httpx
+import mcp.types as types
+from httpx_sse import aconnect_sse
+from httpx_sse._exceptions import SSEError
+from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
+from mcp.shared.message import SessionMessage
+
+from fast_agent.mcp.transport_tracking import ChannelEvent, ChannelName
+
+if TYPE_CHECKING:
+    from anyio.abc import TaskStatus
+    from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+logger = logging.getLogger(__name__)
+
+ChannelHook = Callable[[ChannelEvent], None]
+
+
+def _extract_session_id(endpoint_url: str) -> str | None:
+    parsed = urlparse(endpoint_url)
+    query_params = parse_qs(parsed.query)
+    for key in ("sessionId", "session_id", "session"):
+        values = query_params.get(key)
+        if values:
+            return values[0]
+    return None
+
+
+def _emit_channel_event(
+    channel_hook: ChannelHook | None,
+    channel: ChannelName,
+    event_type: str,
+    *,
+    message: types.JSONRPCMessage | None = None,
+    raw_event: str | None = None,
+    detail: str | None = None,
+    status_code: int | None = None,
+) -> None:
+    if channel_hook is None:
+        return
+    try:
+        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:
+        logger.debug("Channel hook raised an exception", exc_info=True)
+
+
+def _format_http_error(exc: httpx.HTTPStatusError) -> tuple[int | None, str]:
+    status_code: int | None = None
+    detail = str(exc)
+    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'}"
+    return status_code, detail
+
+
+@asynccontextmanager
+async def tracking_sse_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5,
+    sse_read_timeout: float = 60 * 5,
+    httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
+    auth: httpx.Auth | None = None,
+    channel_hook: ChannelHook | None = None,
+) -> AsyncGenerator[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        Callable[[], str | None],
+    ],
+    None,
+]:
+    """
+    Client transport for SSE with channel activity tracking.
+    """
+
+    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)
+
+    session_id: str | None = None
+
+    def get_session_id() -> str | None:
+        return session_id
+
+    async with anyio.create_task_group() as tg:
+        try:
+            logger.debug("Connecting to SSE endpoint: %s", url)
+            async with httpx_client_factory(
+                headers=headers,
+                auth=auth,
+                timeout=httpx.Timeout(timeout, read=sse_read_timeout),
+            ) as client:
+                connected = False
+                post_connected = False
+
+                async def sse_reader(
+                    task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED,
+                ):
+                    try:
+                        async for sse in event_source.aiter_sse():
+                            if sse.event == "endpoint":
+                                endpoint_url = urljoin(url, sse.data)
+                                logger.debug("Received SSE endpoint URL: %s", endpoint_url)
+
+                                url_parsed = urlparse(url)
+                                endpoint_parsed = urlparse(endpoint_url)
+                                if (
+                                    url_parsed.scheme != endpoint_parsed.scheme
+                                    or url_parsed.netloc != endpoint_parsed.netloc
+                                ):
+                                    error_msg = (
+                                        "Endpoint origin does not match connection origin: "
+                                        f"{endpoint_url}"
+                                    )
+                                    logger.error(error_msg)
+                                    _emit_channel_event(
+                                        channel_hook,
+                                        "get",
+                                        "error",
+                                        detail=error_msg,
+                                    )
+                                    raise ValueError(error_msg)
+
+                                nonlocal session_id
+                                session_id = _extract_session_id(endpoint_url)
+                                task_status.started(endpoint_url)
+                            elif sse.event == "message":
+                                try:
+                                    message = types.JSONRPCMessage.model_validate_json(sse.data)
+                                except Exception as exc:
+                                    logger.exception("Error parsing server message")
+                                    _emit_channel_event(
+                                        channel_hook,
+                                        "get",
+                                        "error",
+                                        detail="Error parsing server message",
+                                    )
+                                    await read_stream_writer.send(exc)
+                                    continue
+
+                                _emit_channel_event(channel_hook, "get", "message", message=message)
+                                await read_stream_writer.send(SessionMessage(message))
+                            else:
+                                _emit_channel_event(
+                                    channel_hook,
+                                    "get",
+                                    "keepalive",
+                                    raw_event=sse.event or "keepalive",
+                                )
+                    except SSEError as sse_exc:
+                        logger.exception("Encountered SSE exception")
+                        _emit_channel_event(
+                            channel_hook,
+                            "get",
+                            "error",
+                            detail=str(sse_exc),
+                        )
+                        raise
+                    except Exception as exc:
+                        logger.exception("Error in sse_reader")
+                        _emit_channel_event(
+                            channel_hook,
+                            "get",
+                            "error",
+                            detail=str(exc),
+                        )
+                        await read_stream_writer.send(exc)
+                    finally:
+                        await read_stream_writer.aclose()
+
+                async def post_writer(endpoint_url: str):
+                    try:
+                        async with write_stream_reader:
+                            async for session_message in write_stream_reader:
+                                try:
+                                    payload = session_message.message.model_dump(
+                                        by_alias=True,
+                                        mode="json",
+                                        exclude_none=True,
+                                    )
+                                except Exception:
+                                    logger.exception("Invalid session message payload")
+                                    continue
+
+                                _emit_channel_event(
+                                    channel_hook,
+                                    "post-sse",
+                                    "message",
+                                    message=session_message.message,
+                                )
+
+                                try:
+                                    response = await client.post(endpoint_url, json=payload)
+                                    response.raise_for_status()
+                                except httpx.HTTPStatusError as exc:
+                                    status_code, detail = _format_http_error(exc)
+                                    _emit_channel_event(
+                                        channel_hook,
+                                        "post-sse",
+                                        "error",
+                                        detail=detail,
+                                        status_code=status_code,
+                                    )
+                                    raise
+                    except httpx.HTTPStatusError:
+                        logger.exception("HTTP error in post_writer")
+                    except Exception:
+                        logger.exception("Error in post_writer")
+                        _emit_channel_event(
+                            channel_hook,
+                            "post-sse",
+                            "error",
+                            detail="Error sending client message",
+                        )
+                    finally:
+                        await write_stream.aclose()
+
+                try:
+                    async with aconnect_sse(
+                        client,
+                        "GET",
+                        url,
+                    ) as event_source:
+                        try:
+                            event_source.response.raise_for_status()
+                        except httpx.HTTPStatusError as exc:
+                            status_code, detail = _format_http_error(exc)
+                            _emit_channel_event(
+                                channel_hook,
+                                "get",
+                                "error",
+                                detail=detail,
+                                status_code=status_code,
+                            )
+                            raise
+
+                        _emit_channel_event(channel_hook, "get", "connect")
+                        connected = True
+
+                        endpoint_url = await tg.start(sse_reader)
+                        _emit_channel_event(channel_hook, "post-sse", "connect")
+                        post_connected = True
+                        tg.start_soon(post_writer, endpoint_url)
+
+                        try:
+                            yield read_stream, write_stream, get_session_id
+                        finally:
+                            tg.cancel_scope.cancel()
+                except Exception:
+                    raise
+                finally:
+                    if connected:
+                        _emit_channel_event(channel_hook, "get", "disconnect")
+                    if post_connected:
+                        _emit_channel_event(channel_hook, "post-sse", "disconnect")
+        finally:
+            await read_stream_writer.aclose()
+            await read_stream.aclose()
+            await write_stream_reader.aclose()
+            await write_stream.aclose()

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()