fix(client): reinitialize expired Streamable HTTP sessions

src/mcp/client/_transport.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from contextlib import AbstractAsyncContextManager
-from typing import Protocol
+from typing import Final, Protocol
 
 from mcp.shared._stream_protocols import ReadStream, WriteStream
 from mcp.shared.message import SessionMessage
@@ -12,6 +12,11 @@
 
 TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]]
 
+# SDK-private signal emitted by StreamableHTTPTransport after a request with an
+# established MCP session receives HTTP 404. It never crosses the wire.
+SESSION_EXPIRED: Final = -32003
+SESSION_EXPIRED_MARKER: Final = "mcp.client.session_expired"
+
 
 class Transport(AbstractAsyncContextManager[TransportStreams], Protocol):
     """Protocol for MCP transports.

src/mcp/client/session.py

@@ -37,7 +37,7 @@
 from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, ValidationError
 from typing_extensions import Self, TypeVar, deprecated
 
-from mcp.client._transport import ReadStream, WriteStream
+from mcp.client._transport import SESSION_EXPIRED, SESSION_EXPIRED_MARKER, ReadStream, WriteStream
 from mcp.client.extension import NotificationBinding, ResultClaim, UnexpectedClaimedResult
 from mcp.client.subscriptions import ListenRoute
 from mcp.shared._compat import resync_tracer
@@ -415,6 +415,8 @@ def __init__(
         self._negotiated_version: str | None = None
         self._stamp: Callable[[dict[str, Any], CallOptions], None] = _preconnect_stamp
         self._task_group: anyio.abc.TaskGroup | None = None
+        self._session_recovery_lock = anyio.Lock()
+        self._session_generation = 0
         # subscriptions/listen demux routes; membership decides ack consumption (raw listens are never registered)
         self._listen_routes: dict[RequestId, ListenRoute] = {}
         if dispatcher is not None:
@@ -504,26 +506,15 @@ async def _deliver_bound_notifications(
                 # A raising handler costs only that delivery, as in _on_notify.
                 logger.exception("notification binding handler for %r raised", binding.method)
 
-    async def send_request(
+    async def _send_request_once(
         self,
         request: types.ClientRequest | types.Request[Any, Any],
         result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT],
-        request_read_timeout_seconds: float | None = None,
-        metadata: ClientMessageMetadata | None = None,
-        progress_callback: ProgressFnT | None = None,
+        request_read_timeout_seconds: float | None,
+        metadata: ClientMessageMetadata | None,
+        progress_callback: ProgressFnT | None,
     ) -> ReceiveResultT:
-        """Send a request and wait for its typed result.
-
-        Args:
-            metadata: Streamable HTTP resumption hints.
-
-        Raises:
-            MCPError: Error response, read timeout, or connection closed.
-            RuntimeError: Called before entering the context manager.
-            ValueError: The request declares `name_param` but its params carry no string name.
-            pydantic.ValidationError: The server returned a result that does not
-                conform to the negotiated protocol version.
-        """
+        """Send one typed request without session-expiry recovery."""
         data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
         method: str = data["method"]
         opts: CallOptions = {}
@@ -562,6 +553,76 @@ async def send_request(
             return result_type.validate_python(raw, by_name=False)
         return result_type.model_validate(raw, by_name=False)
 
+    async def _recover_expired_session(self, generation: int) -> None:
+        """Reinitialize once for all requests that observed one expired legacy session."""
+        async with self._session_recovery_lock:
+            if generation != self._session_generation:
+                return
+            self._initialize_result = None
+            self._discover_result = None
+            self._discover_server_info = None
+            self._negotiated_version = None
+            self._stamp = _preconnect_stamp
+            self._active_claims = {}
+            self._call_tool_adapter = _CallToolResultAdapter
+            self._x_mcp_header_maps.clear()
+            self._tool_output_schemas.clear()
+            self._tool_output_validators.clear()
+            await self.initialize()
+            self._session_generation += 1
+
+    async def send_request(
+        self,
+        request: types.ClientRequest | types.Request[Any, Any],
+        result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT],
+        request_read_timeout_seconds: float | None = None,
+        metadata: ClientMessageMetadata | None = None,
+        progress_callback: ProgressFnT | None = None,
+    ) -> ReceiveResultT:
+        """Send a request and wait for its typed result.
+
+        An established legacy Streamable HTTP session that receives a 404 is
+        reinitialized once and the original request is retried once, as required
+        by the transport specification.
+
+        Args:
+            metadata: Streamable HTTP resumption hints.
+
+        Raises:
+            MCPError: Error response, read timeout, connection closed, or a
+                repeated session-expiry response.
+            RuntimeError: Called before entering the context manager.
+            ValueError: The request declares `name_param` but its params carry no string name.
+            pydantic.ValidationError: The server returned a result that does not
+                conform to the negotiated protocol version.
+        """
+        generation = self._session_generation
+        sent_on_legacy_session = self._initialize_result is not None and self._discover_result is None
+        try:
+            return await self._send_request_once(
+                request,
+                result_type,
+                request_read_timeout_seconds,
+                metadata,
+                progress_callback,
+            )
+        except MCPError as exc:
+            error_data = exc.data
+            is_transport_expiry = (
+                isinstance(error_data, Mapping)
+                and cast(Mapping[str, object], error_data).get(SESSION_EXPIRED_MARKER) is True
+            )
+            if exc.code != SESSION_EXPIRED or not is_transport_expiry or not sent_on_legacy_session:
+                raise
+        await self._recover_expired_session(generation)
+        return await self._send_request_once(
+            request,
+            result_type,
+            request_read_timeout_seconds,
+            metadata,
+            progress_callback,
+        )
+
     async def send_notification(self, notification: types.ClientNotification) -> None:
         """Send a one-way notification. Usable before entering the context manager.
 

src/mcp/client/streamable_http.py

@@ -30,7 +30,7 @@
 from mcp_types.version import MODERN_PROTOCOL_VERSIONS
 from pydantic import ValidationError
 
-from mcp.client._transport import TransportStreams
+from mcp.client._transport import SESSION_EXPIRED, SESSION_EXPIRED_MARKER, TransportStreams
 from mcp.shared._compat import resync_tracer
 from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
 from mcp.shared._httpx_utils import create_mcp_http_client
@@ -359,13 +359,24 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                             pass
                         logger.debug("Non-2xx body was not a JSON-RPC error; using fallback")
                     if response.status_code == 404:
-                        if self.session_id is None:
+                        request_session_id = headers.get(MCP_SESSION_ID)
+                        if request_session_id is None:
                             # No session yet → 404 is the HTTP-level spelling of
                             # METHOD_NOT_FOUND (gateway / legacy server doesn't know
-                            # this method); "Session terminated" would be a lie here.
+                            # this method); session recovery would be a lie here.
                             error_data = ErrorData(code=METHOD_NOT_FOUND, message="Not Found")
                         else:
-                            error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated")
+                            # A post-session 404 means the server discarded this
+                            # request's session. Clear only if this request still
+                            # owns the current generation: a delayed old 404 must
+                            # not erase a session another request just recovered.
+                            if self.session_id == request_session_id:
+                                self.session_id = None
+                            error_data = ErrorData(
+                                code=SESSION_EXPIRED,
+                                message="Session expired",
+                                data={SESSION_EXPIRED_MARKER: True},
+                            )
                     else:
                         error_data = ErrorData(code=INTERNAL_ERROR, message="Server returned an error response")
                     session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))

tests/client/test_notification_response.py

@@ -6,6 +6,7 @@
 
 import json
 
+import anyio
 import httpx2
 import mcp_types as types
 import pytest
@@ -17,6 +18,7 @@
 
 from mcp import ClientSession, MCPError
 from mcp.client import IncomingMessage
+from mcp.client._transport import SESSION_EXPIRED, SESSION_EXPIRED_MARKER
 from mcp.client.streamable_http import streamable_http_client
 
 pytestmark = pytest.mark.anyio
@@ -254,15 +256,183 @@ async def test_client_falls_back_to_generic_error_when_non_2xx_body_is_a_jsonrpc
                 assert exc.value.error.code == types.INTERNAL_ERROR
 
 
-async def test_client_falls_back_to_session_terminated_when_404_body_is_malformed_json() -> None:
-    """SDK-defined: an unparseable ``application/json`` body on a 404 response is swallowed
-    and the status-derived ``INVALID_REQUEST`` (session-terminated) fallback resolves the
-    pending request — the parse failure never propagates."""
+async def test_client_reports_session_expiry_after_a_404_recovery_retry_has_malformed_json() -> None:
+    """SDK-defined: a malformed 404 body still triggers one session recovery attempt before failing.
+
+    The parse failure is not surfaced because HTTP 404 is the transport's session-expiry signal. A second
+    404 after recovery is bounded and returns the private session-expired error rather than looping.
+    """
     app = _create_non_2xx_json_body_app(404, b"not valid json{{{")
     async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
         async with streamable_http_client("http://localhost/mcp", http_client=client) as (read_stream, write_stream):
             async with ClientSession(read_stream, write_stream) as session:  # pragma: no branch
                 await session.initialize()
                 with pytest.raises(MCPError) as exc:
                     await session.list_tools()
-                assert exc.value.error.code == types.INVALID_REQUEST
+                assert exc.value.error.code == SESSION_EXPIRED
+
+
+def _create_expired_session_recovery_app(requests: list[tuple[str, str | None]]) -> Starlette:
+    """Return a fresh session after rejecting one established-session request."""
+    initialize_count = 0
+    expire_once = True
+
+    async def handle_mcp_request(request: Request) -> Response:
+        nonlocal expire_once, initialize_count
+        data = json.loads(await request.body())
+        method = data.get("method")
+        session_id = request.headers.get("mcp-session-id")
+        requests.append((method, session_id))
+
+        if method == "initialize":
+            initialize_count += 1
+            return JSONResponse(
+                {"jsonrpc": "2.0", "id": data["id"], "result": INIT_RESPONSE},
+                headers={"mcp-session-id": f"session-{initialize_count}"},
+            )
+        if method == "notifications/initialized":
+            return Response(status_code=202)
+        if method == "tools/list" and expire_once:
+            expire_once = False
+            assert session_id == "session-1"
+            return Response(status_code=404)
+        if method == "tools/list":
+            assert session_id == "session-2"
+            return JSONResponse({"jsonrpc": "2.0", "id": data["id"], "result": {"tools": []}})
+        return Response(status_code=500)
+
+    return Starlette(debug=True, routes=[Route("/mcp", handle_mcp_request, methods=["POST"])])
+
+
+async def test_client_reinitializes_once_after_an_established_session_returns_404() -> None:
+    """Spec-mandated: a 404 for an established legacy session creates a fresh session and retries once.
+
+    The recovery initialize must omit the expired session id; its initialized notification and the retried
+    request must carry the new id. This drives the public ``ClientSession`` API through an in-process ASGI app.
+    """
+    requests: list[tuple[str, str | None]] = []
+    app = _create_expired_session_recovery_app(requests)
+
+    async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
+        async with streamable_http_client("http://localhost/mcp", http_client=client) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                result = await session.initialize()
+                assert result.server_info.name == "test-non-sdk-server"
+
+                tools = await session.list_tools()
+
+    assert tools.tools == []
+    assert requests == [
+        ("initialize", None),
+        ("notifications/initialized", "session-1"),
+        ("tools/list", "session-1"),
+        ("initialize", None),
+        ("notifications/initialized", "session-2"),
+        ("tools/list", "session-2"),
+    ]
+
+
+def _create_repeated_expired_session_app(requests: list[tuple[str, str | None]]) -> Starlette:
+    """Always expire requests from an established session."""
+    initialize_count = 0
+
+    async def handle_mcp_request(request: Request) -> Response:
+        nonlocal initialize_count
+        data = json.loads(await request.body())
+        method = data.get("method")
+        session_id = request.headers.get("mcp-session-id")
+        requests.append((method, session_id))
+
+        if method == "initialize":
+            initialize_count += 1
+            return JSONResponse(
+                {"jsonrpc": "2.0", "id": data["id"], "result": INIT_RESPONSE},
+                headers={"mcp-session-id": f"session-{initialize_count}"},
+            )
+        if method == "notifications/initialized":
+            return Response(status_code=202)
+        if method == "tools/list":
+            return Response(status_code=404)
+        return Response(status_code=500)
+
+    return Starlette(debug=True, routes=[Route("/mcp", handle_mcp_request, methods=["POST"])])
+
+
+async def test_client_retries_an_expired_session_request_only_once() -> None:
+    """SDK-defined: a retry that also receives 404 surfaces an error instead of opening a recovery loop."""
+    requests: list[tuple[str, str | None]] = []
+    app = _create_repeated_expired_session_app(requests)
+
+    async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
+        async with streamable_http_client("http://localhost/mcp", http_client=client) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                with pytest.raises(MCPError) as exc_info:
+                    await session.list_tools()
+
+    assert exc_info.value.code == SESSION_EXPIRED
+    assert exc_info.value.data == {SESSION_EXPIRED_MARKER: True}
+    assert requests == [
+        ("initialize", None),
+        ("notifications/initialized", "session-1"),
+        ("tools/list", "session-1"),
+        ("initialize", None),
+        ("notifications/initialized", "session-2"),
+        ("tools/list", "session-2"),
+    ]
+
+
+async def test_concurrent_expired_session_requests_share_one_reinitialization() -> None:
+    """SDK-defined: concurrent 404 responses recover one session generation, not one per caller."""
+    requests: list[tuple[str, str | None]] = []
+    old_requests_started = 0
+    old_requests_ready = anyio.Event()
+    release_old_requests = anyio.Event()
+    initialize_count = 0
+
+    async def handle_mcp_request(request: Request) -> Response:
+        nonlocal initialize_count, old_requests_started
+        data = json.loads(await request.body())
+        method = data.get("method")
+        session_id = request.headers.get("mcp-session-id")
+        requests.append((method, session_id))
+
+        if method == "initialize":
+            initialize_count += 1
+            return JSONResponse(
+                {"jsonrpc": "2.0", "id": data["id"], "result": INIT_RESPONSE},
+                headers={"mcp-session-id": f"session-{initialize_count}"},
+            )
+        if method == "notifications/initialized":
+            return Response(status_code=202)
+        if method == "tools/list" and session_id == "session-1":
+            old_requests_started += 1
+            if old_requests_started == 2:
+                old_requests_ready.set()
+            await release_old_requests.wait()
+            return Response(status_code=404)
+        if method == "tools/list" and session_id == "session-2":
+            return JSONResponse({"jsonrpc": "2.0", "id": data["id"], "result": {"tools": []}})
+        return Response(status_code=500)
+
+    app = Starlette(debug=True, routes=[Route("/mcp", handle_mcp_request, methods=["POST"])])
+    results: list[types.ListToolsResult] = []
+
+    async def append_list_tools_result(session: ClientSession) -> None:
+        results.append(await session.list_tools())
+
+    async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
+        async with streamable_http_client("http://localhost/mcp", http_client=client) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                async with anyio.create_task_group() as task_group:
+                    task_group.start_soon(append_list_tools_result, session)
+                    task_group.start_soon(append_list_tools_result, session)
+                    with anyio.fail_after(5):
+                        await old_requests_ready.wait()
+                    release_old_requests.set()
+
+    assert [result.tools for result in results] == [[], []]
+    assert requests.count(("initialize", None)) == 2
+    assert requests.count(("notifications/initialized", "session-2")) == 1
+    assert requests.count(("tools/list", "session-2")) == 2