Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions src/fastmcp/client/transports.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import shutil
import sys
import warnings
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Callable
from pathlib import Path
from typing import Any, Literal, TextIO, TypeVar, cast, overload

Expand Down Expand Up @@ -254,6 +254,8 @@ def __init__(
sse_read_timeout = datetime.timedelta(seconds=float(sse_read_timeout))
self.sse_read_timeout = sse_read_timeout

self._get_session_id_cb: Callable[[], str | None] | None = None

def _set_auth(self, auth: httpx.Auth | Literal["oauth"] | str | None):
if auth == "oauth":
auth = OAuth(self.url, httpx_client_factory=self.httpx_client_factory)
Expand Down Expand Up @@ -287,12 +289,25 @@ async def connect_session(
auth=self.auth,
**client_kwargs,
) as transport:
read_stream, write_stream, _ = transport
read_stream, write_stream, get_session_id = transport
self._get_session_id_cb = get_session_id
async with ClientSession(
read_stream, write_stream, **session_kwargs
) as session:
yield session

def get_session_id(self) -> str | None:
if self._get_session_id_cb:
try:
return self._get_session_id_cb()
except Exception:
return None
return None

async def close(self):
# Reset the session id callback
self._get_session_id_cb = None

def __repr__(self) -> str:
return f"<StreamableHttpTransport(url='{self.url}')>"

Expand Down
9 changes: 9 additions & 0 deletions tests/client/test_streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,15 @@ async def test_http_headers(streamable_http_server: str):
assert json_result["x-demo-header"] == "ABC"


async def test_session_id_callback(streamable_http_server: str):
"""Test getting mcp-session-id from the transport."""
transport = StreamableHttpTransport(streamable_http_server)
assert transport.get_session_id() is None
async with Client(transport=transport):
session_id = transport.get_session_id()
assert session_id is not None


@pytest.mark.parametrize("streamable_http_server", [True, False], indirect=True)
async def test_greet_with_progress_tool(streamable_http_server: str):
"""Test calling the greet tool."""
Expand Down
Loading