Skip to content

Commit 26fcd43

Browse files
adrienducheminadrienduchemin
authored
feat(mcp): add Streamable HTTP client to extra/mcp (#592)
Adds MCPClientStreamableHTTP + StreamableHTTPServerParams alongside the existing SSE and stdio MCP clients, for servers that speak the Streamable HTTP transport (--transport http). Credentials are set on the httpx client so they ride every request including initialize; the transport's own headers arg is deprecated.
1 parent b0613c7 commit 26fcd43

2 files changed

Lines changed: 164 additions & 0 deletions

File tree

‎src/mistralai/extra/mcp/streamable_http.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import logging
2+
from contextlib import AsyncExitStack
3+
from typing import Any
4+
5+
import httpx
6+
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
7+
from mcp.client.streamable_http import streamable_http_client # pyright: ignore[reportMissingImports]
8+
from mcp.shared.message import SessionMessage # pyright: ignore[reportMissingImports]
9+
10+
from mistralai.extra.mcp.base import (
11+
MCPClientBase,
12+
)
13+
14+
from mistralai.client.types import BaseModel
15+
16+
logger = logging.getLogger(__name__)
17+
18+
19+
class StreamableHTTPServerParams(BaseModel):
20+
"""Parameters required for a MCPClient with Streamable HTTP transport"""
21+
22+
url: str
23+
headers: dict[str, Any] | None = None
24+
timeout: float = 30
25+
# Whether the httpx client trusts the ambient environment (HTTP(S)_PROXY,
26+
# SSL_CERT_FILE/DIR, .netrc). Defaults to httpx's default (True). Set False to
27+
# reach the endpoint directly without an ambient egress proxy, e.g. for an
28+
# in-cluster URL on a host whose external egress is proxied.
29+
trust_env: bool = True
30+
# Whether httpx follows HTTP redirects (3xx). Defaults to False (also httpx's
31+
# default). On a redirect httpx only strips Authorization on cross-origin hops,
32+
# so the ``headers`` above (e.g. a per-request integration token) would be resent
33+
# verbatim to the redirect target, letting a compromised or open-redirecting
34+
# server exfiltrate them. Set True only if the server relies on redirects and
35+
# every host it can redirect to is trusted.
36+
follow_redirects: bool = False
37+
38+
39+
class MCPClientStreamableHTTP(MCPClientBase):
40+
"""MCP client that uses the Streamable HTTP transport for communication.
41+
42+
Credentials (for example a bearer token, or a per-request integration token)
43+
are provided as ``headers`` and set as the default headers of the underlying
44+
``httpx.AsyncClient``, so they are sent on every request including the
45+
initialize call. Recent ``mcp`` releases deprecate and ignore the transport's
46+
own ``headers`` argument, so configuring them on the client is required.
47+
"""
48+
49+
_params: StreamableHTTPServerParams
50+
51+
def __init__(
52+
self,
53+
params: StreamableHTTPServerParams,
54+
name: str | None = None,
55+
):
56+
super().__init__(name=name)
57+
self._params = params
58+
59+
async def _get_transport(
60+
self, exit_stack: AsyncExitStack
61+
) -> tuple[
62+
MemoryObjectReceiveStream[SessionMessage | Exception],
63+
MemoryObjectSendStream[SessionMessage],
64+
]:
65+
# trust_env controls whether the client inherits the ambient
66+
# HTTP(S)_PROXY / cert / netrc env. Set it False (see params) to reach an
67+
# in-cluster endpoint directly, bypassing a proxy meant for external
68+
# traffic that would make the connection hang.
69+
#
70+
# follow_redirects defaults False (see params): the secret default headers
71+
# would otherwise be resent to a redirect target on same-origin hops.
72+
http_client = await exit_stack.enter_async_context(
73+
httpx.AsyncClient(
74+
headers=self._params.headers,
75+
timeout=self._params.timeout,
76+
follow_redirects=self._params.follow_redirects,
77+
trust_env=self._params.trust_env,
78+
)
79+
)
80+
read_stream, write_stream, _ = await exit_stack.enter_async_context(
81+
streamable_http_client(url=self._params.url, http_client=http_client)
82+
)
83+
return read_stream, write_stream

‎src/mistralai/extra/tests/test_mcp_streamable_http.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
"""Tests for the Streamable HTTP MCP client."""
2+
3+
import unittest
4+
from collections.abc import AsyncIterator
5+
from contextlib import AsyncExitStack, asynccontextmanager
6+
from typing import Any
7+
from unittest import mock
8+
9+
from mistralai.extra.mcp.streamable_http import (
10+
MCPClientStreamableHTTP,
11+
StreamableHTTPServerParams,
12+
)
13+
14+
15+
class TestMCPClientStreamableHTTP(unittest.IsolatedAsyncioTestCase):
16+
async def test_headers_are_set_on_the_http_client(self) -> None:
17+
"""Credentials must live on the httpx client (the transport's own headers
18+
argument is deprecated/ignored), so they are sent on every request."""
19+
captured: dict[str, Any] = {}
20+
21+
@asynccontextmanager
22+
async def fake_streamable_http_client(url: str, http_client: Any) -> AsyncIterator[Any]:
23+
captured["url"] = url
24+
captured["headers"] = dict(http_client.headers)
25+
yield object(), object(), lambda: None
26+
27+
client = MCPClientStreamableHTTP(
28+
params=StreamableHTTPServerParams(
29+
url="http://mcp.example/mcp",
30+
headers={"Authorization": "Bearer gate", "Notion-Token": "ntn_x"},
31+
),
32+
name="test",
33+
)
34+
35+
with mock.patch(
36+
"mistralai.extra.mcp.streamable_http.streamable_http_client",
37+
fake_streamable_http_client,
38+
):
39+
async with AsyncExitStack() as stack:
40+
read_stream, write_stream = await client._get_transport(stack)
41+
42+
self.assertIsNotNone(read_stream)
43+
self.assertIsNotNone(write_stream)
44+
self.assertEqual(captured["url"], "http://mcp.example/mcp")
45+
# httpx normalizes header names to lower case
46+
self.assertEqual(captured["headers"]["authorization"], "Bearer gate")
47+
self.assertEqual(captured["headers"]["notion-token"], "ntn_x")
48+
49+
async def test_follow_redirects_defaults_false_and_is_configurable(self) -> None:
50+
"""Redirects are not followed by default (secret headers would be resent to
51+
the redirect target); the flag is forwarded to the httpx client when set."""
52+
captured: dict[str, Any] = {}
53+
54+
@asynccontextmanager
55+
async def fake_streamable_http_client(url: str, http_client: Any) -> AsyncIterator[Any]:
56+
captured["follow_redirects"] = http_client.follow_redirects
57+
yield object(), object(), lambda: None
58+
59+
with mock.patch(
60+
"mistralai.extra.mcp.streamable_http.streamable_http_client",
61+
fake_streamable_http_client,
62+
):
63+
default_client = MCPClientStreamableHTTP(
64+
params=StreamableHTTPServerParams(url="http://mcp.example/mcp"),
65+
name="test",
66+
)
67+
async with AsyncExitStack() as stack:
68+
await default_client._get_transport(stack)
69+
self.assertFalse(captured["follow_redirects"])
70+
71+
opted_in_client = MCPClientStreamableHTTP(
72+
params=StreamableHTTPServerParams(url="http://mcp.example/mcp", follow_redirects=True),
73+
name="test",
74+
)
75+
async with AsyncExitStack() as stack:
76+
await opted_in_client._get_transport(stack)
77+
self.assertTrue(captured["follow_redirects"])
78+
79+
80+
if __name__ == "__main__":
81+
unittest.main()

0 commit comments

Comments
 (0)