diff --git a/src/mistralai/extra/mcp/streamable_http.py b/src/mistralai/extra/mcp/streamable_http.py new file mode 100644 index 00000000..e69953ba --- /dev/null +++ b/src/mistralai/extra/mcp/streamable_http.py @@ -0,0 +1,63 @@ +import logging +from contextlib import AsyncExitStack +from typing import Any + +import httpx +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream +from mcp.client.streamable_http import streamable_http_client # pyright: ignore[reportMissingImports] +from mcp.shared.message import SessionMessage # pyright: ignore[reportMissingImports] + +from mistralai.extra.mcp.base import ( + MCPClientBase, +) + +from mistralai.client.types import BaseModel + +logger = logging.getLogger(__name__) + + +class StreamableHTTPServerParams(BaseModel): + """Parameters required for a MCPClient with Streamable HTTP transport""" + + url: str + headers: dict[str, Any] | None = None + timeout: float = 30 + + +class MCPClientStreamableHTTP(MCPClientBase): + """MCP client that uses the Streamable HTTP transport for communication. + + Credentials (for example a bearer token, or a per-request integration token) + are provided as ``headers`` and set as the default headers of the underlying + ``httpx.AsyncClient``, so they are sent on every request including the + initialize call. Recent ``mcp`` releases deprecate and ignore the transport's + own ``headers`` argument, so configuring them on the client is required. + """ + + _params: StreamableHTTPServerParams + + def __init__( + self, + params: StreamableHTTPServerParams, + name: str | None = None, + ): + super().__init__(name=name) + self._params = params + + async def _get_transport( + self, exit_stack: AsyncExitStack + ) -> tuple[ + MemoryObjectReceiveStream[SessionMessage | Exception], + MemoryObjectSendStream[SessionMessage], + ]: + http_client = await exit_stack.enter_async_context( + httpx.AsyncClient( + headers=self._params.headers, + timeout=self._params.timeout, + follow_redirects=True, + ) + ) + read_stream, write_stream, _ = await exit_stack.enter_async_context( + streamable_http_client(url=self._params.url, http_client=http_client) + ) + return read_stream, write_stream diff --git a/src/mistralai/extra/tests/test_mcp_streamable_http.py b/src/mistralai/extra/tests/test_mcp_streamable_http.py new file mode 100644 index 00000000..c58ac5d8 --- /dev/null +++ b/src/mistralai/extra/tests/test_mcp_streamable_http.py @@ -0,0 +1,51 @@ +"""Tests for the Streamable HTTP MCP client.""" + +import unittest +from collections.abc import AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager +from typing import Any +from unittest import mock + +from mistralai.extra.mcp.streamable_http import ( + MCPClientStreamableHTTP, + StreamableHTTPServerParams, +) + + +class TestMCPClientStreamableHTTP(unittest.IsolatedAsyncioTestCase): + async def test_headers_are_set_on_the_http_client(self) -> None: + """Credentials must live on the httpx client (the transport's own headers + argument is deprecated/ignored), so they are sent on every request.""" + captured: dict[str, Any] = {} + + @asynccontextmanager + async def fake_streamable_http_client(url: str, http_client: Any) -> AsyncIterator[Any]: + captured["url"] = url + captured["headers"] = dict(http_client.headers) + yield object(), object(), lambda: None + + client = MCPClientStreamableHTTP( + params=StreamableHTTPServerParams( + url="http://mcp.example/mcp", + headers={"Authorization": "Bearer gate", "Notion-Token": "ntn_x"}, + ), + name="test", + ) + + with mock.patch( + "mistralai.extra.mcp.streamable_http.streamable_http_client", + fake_streamable_http_client, + ): + async with AsyncExitStack() as stack: + read_stream, write_stream = await client._get_transport(stack) + + self.assertIsNotNone(read_stream) + self.assertIsNotNone(write_stream) + self.assertEqual(captured["url"], "http://mcp.example/mcp") + # httpx normalizes header names to lower case + self.assertEqual(captured["headers"]["authorization"], "Bearer gate") + self.assertEqual(captured["headers"]["notion-token"], "ntn_x") + + +if __name__ == "__main__": + unittest.main()