Skip to content

Commit ae9e8ec

Browse files
feat(client): support explicit streamable-http session resumption
1 parent 92c693b commit ae9e8ec

File tree

6 files changed

+36
-6
lines changed

6 files changed

+36
-6
lines changed

src/mcp/client/client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,9 @@ async def main():
9595
elicitation_callback: ElicitationFnT | None = None
9696
"""Callback for handling elicitation requests."""
9797

98+
streamable_http_session_id: str | None = None
99+
"""Optional pre-existing MCP session ID used when server is a StreamableHTTP URL."""
100+
98101
_session: ClientSession | None = field(init=False, default=None)
99102
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
100103
_transport: Transport = field(init=False)
@@ -103,7 +106,7 @@ def __post_init__(self) -> None:
103106
if isinstance(self.server, Server | MCPServer):
104107
self._transport = InMemoryTransport(self.server, raise_exceptions=self.raise_exceptions)
105108
elif isinstance(self.server, str):
106-
self._transport = streamable_http_client(self.server)
109+
self._transport = streamable_http_client(self.server, session_id=self.streamable_http_session_id)
107110
else:
108111
self._transport = self.server
109112

src/mcp/client/session_group.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@ class StreamableHttpParameters(BaseModel):
6363
# Close the client session when the transport closes.
6464
terminate_on_close: bool = True
6565

66+
# Optional pre-existing MCP session ID for explicit session resumption.
67+
session_id: str | None = None
68+
6669

6770
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
6871

@@ -296,6 +299,7 @@ async def _establish_session(
296299
url=server_params.url,
297300
http_client=httpx_client,
298301
terminate_on_close=server_params.terminate_on_close,
302+
session_id=server_params.session_id,
299303
)
300304
read, write = await session_stack.enter_async_context(client)
301305

src/mcp/client/streamable_http.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,14 +72,15 @@ class RequestContext:
7272
class StreamableHTTPTransport:
7373
"""StreamableHTTP client transport implementation."""
7474

75-
def __init__(self, url: str) -> None:
75+
def __init__(self, url: str, session_id: str | None = None) -> None:
7676
"""Initialize the StreamableHTTP transport.
7777
7878
Args:
7979
url: The endpoint URL.
80+
session_id: Optional pre-existing MCP session ID to resume.
8081
"""
8182
self.url = url
82-
self.session_id: str | None = None
83+
self.session_id: str | None = session_id
8384
self.protocol_version: str | None = None
8485

8586
def _prepare_headers(self) -> dict[str, str]:
@@ -512,6 +513,7 @@ async def streamable_http_client(
512513
*,
513514
http_client: httpx.AsyncClient | None = None,
514515
terminate_on_close: bool = True,
516+
session_id: str | None = None,
515517
) -> AsyncGenerator[TransportStreams, None]:
516518
"""Client transport for StreamableHTTP.
517519
@@ -521,6 +523,8 @@ async def streamable_http_client(
521523
client with recommended MCP timeouts will be created. To configure headers,
522524
authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here.
523525
terminate_on_close: If True, send a DELETE request to terminate the session when the context exits.
526+
session_id: Optional pre-existing MCP session ID to include in requests,
527+
enabling explicit session resumption.
524528
525529
Yields:
526530
Tuple containing:
@@ -538,7 +542,7 @@ async def streamable_http_client(
538542
# Create default client with recommended MCP timeouts
539543
client = create_mcp_http_client()
540544

541-
transport = StreamableHTTPTransport(url)
545+
transport = StreamableHTTPTransport(url, session_id=session_id)
542546

543547
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
544548

tests/client/test_client.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,13 @@ async def test_complete_with_prompt_reference(simple_server: Server):
307307
def test_client_with_url_initializes_streamable_http_transport():
308308
with patch("mcp.client.client.streamable_http_client") as mock:
309309
_ = Client("http://localhost:8000/mcp")
310-
mock.assert_called_once_with("http://localhost:8000/mcp")
310+
mock.assert_called_once_with("http://localhost:8000/mcp", session_id=None)
311+
312+
313+
def test_client_with_url_and_session_id_initializes_streamable_http_transport():
314+
with patch("mcp.client.client.streamable_http_client") as mock:
315+
_ = Client("http://localhost:8000/mcp", streamable_http_session_id="resume-session-id")
316+
mock.assert_called_once_with("http://localhost:8000/mcp", session_id="resume-session-id")
311317

312318

313319
async def test_client_uses_transport_directly(app: MCPServer):

tests/client/test_session_group.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,11 @@ async def test_client_session_group_disconnect_non_existent_server():
294294
"mcp.client.session_group.sse_client",
295295
), # url, headers, timeout, sse_read_timeout
296296
(
297-
StreamableHttpParameters(url="http://test.com/stream", terminate_on_close=False),
297+
StreamableHttpParameters(
298+
url="http://test.com/stream",
299+
terminate_on_close=False,
300+
session_id="resumed-session-id",
301+
),
298302
"streamablehttp",
299303
"mcp.client.session_group.streamable_http_client",
300304
), # url, headers, timeout, sse_read_timeout, terminate_on_close
@@ -363,6 +367,7 @@ async def test_client_session_group_establish_session_parameterized(
363367
call_args = mock_specific_client_func.call_args
364368
assert call_args.kwargs["url"] == server_params_instance.url
365369
assert call_args.kwargs["terminate_on_close"] == server_params_instance.terminate_on_close
370+
assert call_args.kwargs["session_id"] == server_params_instance.session_id
366371
assert isinstance(call_args.kwargs["http_client"], httpx.AsyncClient)
367372

368373
mock_client_cm_instance.__aenter__.assert_awaited_once()

tests/shared/test_streamable_http.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1801,6 +1801,14 @@ async def test_handle_sse_event_skips_empty_data():
18011801
await read_stream.aclose()
18021802

18031803

1804+
def test_streamable_http_transport_includes_seeded_session_id_header():
1805+
transport = StreamableHTTPTransport(url="http://localhost:8000/mcp", session_id="resume-session-id")
1806+
1807+
headers = transport._prepare_headers()
1808+
1809+
assert headers["mcp-session-id"] == "resume-session-id"
1810+
1811+
18041812
@pytest.mark.anyio
18051813
async def test_priming_event_not_sent_for_old_protocol_version():
18061814
"""Test that _maybe_send_priming_event skips for old protocol versions (backwards compat)."""

0 commit comments

Comments
 (0)