Fix proxy transport semantics

Author: Kludex

docs/proxying.md

@@ -0,0 +1,141 @@
+# Proxying MCP Transports
+
+The `mcp_proxy()` helper bridges two MCP transports and forwards messages in both directions.
+
+It is useful when you want to put a transport boundary between an MCP client and an upstream MCP server without
+rewriting the forwarding loop yourself.
+
+## What It Does
+
+`mcp_proxy()` takes two transport pairs:
+
+- a transport facing the downstream client
+- a transport facing the upstream server
+
+While the context manager is active, it:
+
+- forwards `SessionMessage` objects from client to server
+- forwards `SessionMessage` objects from server to client
+- sends transport exceptions to an optional `on_error` callback
+- closes the paired write side when the corresponding read side stops
+
+## What It Does Not Do
+
+`mcp_proxy()` is a transport relay, not a full proxy server.
+
+It does not add:
+
+- authentication
+- authorization
+- request or response rewriting
+- routing across multiple upstream servers
+- retries or buffering policies
+- metrics or tracing by default
+
+If you need those behaviors, build them around the helper.
+
+## Weather Service Example
+
+This example proxies a small weather service. The upstream service is defined with `MCPServer` and exposed over
+streamable HTTP. The proxy bridges a downstream transport to that upstream transport.
+
+- `get_weather(city)` for a structured weather snapshot
+- `get_weather_alerts(region)` for active alerts
+
+The client talks only to the downstream side of the proxy.
+
+```python
+import anyio
+import uvicorn
+
+from mcp.client.session import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+from mcp.proxy import mcp_proxy
+from mcp.server.mcpserver import MCPServer
+from mcp.shared.memory import create_client_server_memory_streams
+
+
+app = MCPServer("Weather Service")
+
+
+@app.tool()
+def get_weather(city: str) -> dict[str, str | float]:
+    return {
+        "city": city,
+        "temperature_c": 22.5,
+        "condition": "partly cloudy",
+        "wind_speed_kmh": 12.3,
+    }
+
+
+@app.tool()
+def get_weather_alerts(region: str) -> dict[str, object]:
+    return {
+        "region": region,
+        "alerts": [{"severity": "medium", "title": "Heat advisory"}],
+    }
+
+
+async def main() -> None:
+    starlette_app = app.streamable_http_app(streamable_http_path="/mcp")
+    config = uvicorn.Config(starlette_app, host="127.0.0.1", port=8765, log_level="warning")
+    upstream_server = uvicorn.Server(config)
+
+    async with (
+        create_client_server_memory_streams() as (client_streams, proxy_client_streams),
+        streamable_http_client("http://127.0.0.1:8765/mcp") as proxy_server_streams,
+        anyio.create_task_group() as tg,
+    ):
+        tg.start_soon(upstream_server.serve)
+
+        async with mcp_proxy(
+            proxy_client_streams,
+            proxy_server_streams,
+        ):
+            async with ClientSession(client_streams[0], client_streams[1]) as session:
+                await session.initialize()
+                weather = await session.call_tool("get_weather", {"city": "London"})
+                alerts = await session.call_tool("get_weather_alerts", {"region": "California"})
+
+                print(weather.content[0].text)
+                print(alerts.content[0].text)
+
+        upstream_server.should_exit = True
+        tg.cancel_scope.cancel()
+
+
+anyio.run(main)
+```
+
+## Error Handling
+
+Use `on_error` to observe transport-level exceptions:
+
+```python
+async with mcp_proxy(
+    downstream_transport,
+    upstream_transport,
+    on_error=handle_transport_error,
+):
+    ...
+```
+
+`on_error` is keyword-only. It may be either:
+
+- an async callable
+- a sync callable, which will run in a worker thread
+
+Exceptions raised by `on_error` are swallowed. Transport exceptions still terminate the proxy instead of being silently
+consumed.
+
+## When To Use It
+
+`mcp_proxy()` is a good fit when you are:
+
+- exposing an upstream MCP server through a different transport boundary
+- inserting middleware-like behavior between two MCP transports
+- building a local relay for testing or development
+- experimenting with transport adapters
+
+If all you need is to test a server directly, prefer [`Client`](testing.md), which already provides an in-memory
+transport for that use case.

mkdocs.yml

@@ -17,6 +17,7 @@ nav:
   - Documentation:
       - Concepts: concepts.md
       - Low-Level Server: low-level-server.md
+      - Proxying Transports: proxying.md
       - Authorization: authorization.md
       - Testing: testing.md
   - Experimental:

src/mcp/proxy.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
+import contextvars
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from functools import partial
-from typing import cast
+from typing import Any, Protocol, cast
 
 import anyio
 from anyio import to_thread
@@ -18,6 +19,10 @@
 ErrorHandler = Callable[[Exception], None | Awaitable[None]]
 
 
+class ContextualWriteStream(Protocol):
+    async def send_with_context(self, context: contextvars.Context, item: SessionMessage | Exception) -> None: ...
+
+
 @asynccontextmanager
 async def mcp_proxy(
     transport_to_client: MessageStream,
@@ -49,16 +54,38 @@ async def _forward_messages(
                 async for item in read_stream:
                     if isinstance(item, Exception):
                         await _run_error_handler(item, on_error)
-                        continue
+                        raise item
 
                     try:
-                        await write_stream.send(item)
+                        await _forward_message(item, write_stream, read_stream)
                     except anyio.ClosedResourceError:
                         break
     except anyio.ClosedResourceError:
         return
 
 
+async def _forward_message(
+    item: SessionMessage,
+    write_stream: WriteStream[SessionMessage],
+    read_stream: ReadStream[SessionMessage | Exception],
+) -> None:
+    sender_context: contextvars.Context | None = getattr(read_stream, "last_context", None)
+    context_write_stream = cast(ContextualWriteStream | None, _get_contextual_write_stream(write_stream))
+
+    if sender_context is not None and context_write_stream is not None:
+        await context_write_stream.send_with_context(sender_context, item)
+        return
+
+    await write_stream.send(item)
+
+
+def _get_contextual_write_stream(write_stream: WriteStream[SessionMessage]) -> Any:
+    send_with_context = getattr(write_stream, "send_with_context", None)
+    if callable(send_with_context):
+        return write_stream
+    return None
+
+
 async def _run_error_handler(error: Exception, on_error: ErrorHandler | None) -> None:
     if on_error is None:
         return

src/mcp/shared/_context_streams.py

@@ -36,6 +36,9 @@ def __init__(self, inner: MemoryObjectSendStream[_Envelope[T]]) -> None:
     async def send(self, item: T) -> None:
         await self._inner.send((contextvars.copy_context(), item))
 
+    async def send_with_context(self, context: contextvars.Context, item: T) -> None:
+        await self._inner.send((context, item))
+
     def close(self) -> None:
         self._inner.close()