feat(mcp): add Streamable HTTP client to extra/mcp

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

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