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

src/mistralai/extra/mcp/streamable_http.py

@@ -0,0 +1,83 @@
+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
+    # Whether the httpx client trusts the ambient environment (HTTP(S)_PROXY,
+    # SSL_CERT_FILE/DIR, .netrc). Defaults to httpx's default (True). Set False to
+    # reach the endpoint directly without an ambient egress proxy, e.g. for an
+    # in-cluster URL on a host whose external egress is proxied.
+    trust_env: bool = True
+    # Whether httpx follows HTTP redirects (3xx). Defaults to False (also httpx's
+    # default). On a redirect httpx only strips Authorization on cross-origin hops,
+    # so the ``headers`` above (e.g. a per-request integration token) would be resent
+    # verbatim to the redirect target, letting a compromised or open-redirecting
+    # server exfiltrate them. Set True only if the server relies on redirects and
+    # every host it can redirect to is trusted.
+    follow_redirects: bool = False
+
+
+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],
+    ]:
+        # trust_env controls whether the client inherits the ambient
+        # HTTP(S)_PROXY / cert / netrc env. Set it False (see params) to reach an
+        # in-cluster endpoint directly, bypassing a proxy meant for external
+        # traffic that would make the connection hang.
+        #
+        # follow_redirects defaults False (see params): the secret default headers
+        # would otherwise be resent to a redirect target on same-origin hops.
+        http_client = await exit_stack.enter_async_context(
+            httpx.AsyncClient(
+                headers=self._params.headers,
+                timeout=self._params.timeout,
+                follow_redirects=self._params.follow_redirects,
+                trust_env=self._params.trust_env,
+            )
+        )
+        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,81 @@
+"""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")
+
+    async def test_follow_redirects_defaults_false_and_is_configurable(self) -> None:
+        """Redirects are not followed by default (secret headers would be resent to
+        the redirect target); the flag is forwarded to the httpx client when set."""
+        captured: dict[str, Any] = {}
+
+        @asynccontextmanager
+        async def fake_streamable_http_client(url: str, http_client: Any) -> AsyncIterator[Any]:
+            captured["follow_redirects"] = http_client.follow_redirects
+            yield object(), object(), lambda: None
+
+        with mock.patch(
+            "mistralai.extra.mcp.streamable_http.streamable_http_client",
+            fake_streamable_http_client,
+        ):
+            default_client = MCPClientStreamableHTTP(
+                params=StreamableHTTPServerParams(url="http://mcp.example/mcp"),
+                name="test",
+            )
+            async with AsyncExitStack() as stack:
+                await default_client._get_transport(stack)
+            self.assertFalse(captured["follow_redirects"])
+
+            opted_in_client = MCPClientStreamableHTTP(
+                params=StreamableHTTPServerParams(url="http://mcp.example/mcp", follow_redirects=True),
+                name="test",
+            )
+            async with AsyncExitStack() as stack:
+                await opted_in_client._get_transport(stack)
+            self.assertTrue(captured["follow_redirects"])
+
+
+if __name__ == "__main__":
+    unittest.main()