|
| 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 |
0 commit comments