diff --git a/src/mcp/__init__.py b/src/mcp/__init__.py index 28bc4703ed..7e80516fb1 100644 --- a/src/mcp/__init__.py +++ b/src/mcp/__init__.py @@ -68,7 +68,12 @@ from .client.stdio import StdioServerParameters, stdio_client from .server.session import ServerSession from .server.stdio import stdio_server -from .shared.exceptions import MCPDeprecationWarning, MCPError, UrlElicitationRequiredError +from .shared.exceptions import ( + MCPDeprecationWarning, + MCPError, + SessionNotFoundError, + UrlElicitationRequiredError, +) from .shared.uri_template import InvalidUriTemplate, UriTemplate __all__ = [ @@ -105,6 +110,7 @@ "LoggingMessageNotification", "MCPDeprecationWarning", "MCPError", + "SessionNotFoundError", "Notification", "PingRequest", "ProgressNotification", diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index 053f1f10a9..cb476f4a78 100644 --- a/src/mcp/client/streamable_http.py +++ b/src/mcp/client/streamable_http.py @@ -40,6 +40,7 @@ sse_within_origin, stream_within_origin, ) +from mcp.shared.exceptions import _SESSION_NOT_FOUND_MARKER from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER from mcp.shared.jsonrpc_dispatcher import cancelled_request_id_from_params from mcp.shared.message import ClientMessageMetadata, SessionMessage @@ -406,9 +407,16 @@ async def _handle_post_request(self, ctx: RequestContext) -> None: body = await response.aread() parsed = jsonrpc_message_adapter.validate_json(body, by_name=False) if isinstance(parsed, JSONRPCError): + error = parsed.error + if response.status_code == 404 and self.session_id is not None: + error = ErrorData( + code=error.code, + message=error.message, + data={_SESSION_NOT_FOUND_MARKER: True, "data": error.data}, + ) # The server may have set `id: null` (request rejected before its # id was parsed); use this request's id so correlation works. - reply = JSONRPCError(jsonrpc="2.0", id=message.id, error=parsed.error) + reply = JSONRPCError(jsonrpc="2.0", id=message.id, error=error) await ctx.read_stream_writer.send(SessionMessage(reply)) return except (httpx2.StreamError, ValidationError): @@ -421,7 +429,11 @@ async def _handle_post_request(self, ctx: RequestContext) -> None: # this method); "Session terminated" 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") + error_data = ErrorData( + code=INVALID_REQUEST, + message="Session terminated", + data={_SESSION_NOT_FOUND_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)) diff --git a/src/mcp/shared/exceptions.py b/src/mcp/shared/exceptions.py index c2a7fd44e7..ac1956be66 100644 --- a/src/mcp/shared/exceptions.py +++ b/src/mcp/shared/exceptions.py @@ -5,6 +5,11 @@ from mcp_types import INVALID_REQUEST, URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError +# Internal transport marker used to preserve the distinction between a lost +# Streamable HTTP session and an ordinary INVALID_REQUEST response. +_SESSION_NOT_FOUND_MARKER = "__mcp_session_not_found__" + + class MCPDeprecationWarning(UserWarning): """A custom deprecation warning for the MCP SDK. @@ -52,6 +57,10 @@ def __str__(self) -> str: return self.message +class SessionNotFoundError(MCPError): + """Raised when a Streamable HTTP server no longer recognizes the session.""" + + class NoBackChannelError(MCPError): """Raised when a server-initiated request has no channel that can deliver it. diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 223373da0f..7e82f6e301 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -52,7 +52,12 @@ coerce_request_id, run_notify_intercept, ) -from mcp.shared.exceptions import MCPError, NoBackChannelError +from mcp.shared.exceptions import ( + MCPError, + NoBackChannelError, + SessionNotFoundError, + _SESSION_NOT_FOUND_MARKER, +) from mcp.shared.message import ( ClientMessageMetadata, MessageMetadata, @@ -458,6 +463,12 @@ async def send_raw_request( receive.close() if isinstance(outcome, ErrorData): + if isinstance(outcome.data, dict) and outcome.data.get(_SESSION_NOT_FOUND_MARKER): + raise SessionNotFoundError( + code=outcome.code, + message=outcome.message, + data=outcome.data.get("data"), + ) raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data) return outcome diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index 655d9941dc..f5fd02afd7 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -44,7 +44,7 @@ from starlette.routing import Mount from starlette.types import Message, Scope -from mcp import MCPError +from mcp import MCPError, SessionNotFoundError from mcp.client import ClientRequestContext, IncomingMessage from mcp.client.session import ClientSession from mcp.client.streamable_http import StreamableHTTPTransport, streamable_http_client @@ -1140,7 +1140,7 @@ async def test_streamable_http_client_session_termination(basic_app: Starlette) ): async with ClientSession(read_stream, write_stream) as session: # pragma: no branch # Attempt to make a request after termination - with pytest.raises(MCPError) as exc_info: # pragma: no branch + with pytest.raises(SessionNotFoundError) as exc_info: # pragma: no branch await session.list_tools() assert exc_info.value.error.code == INVALID_REQUEST assert exc_info.value.error.message == "Session not found" @@ -1201,7 +1201,7 @@ async def __aexit__(self, *args: Any) -> None: ): async with ClientSession(read_stream, write_stream) as session: # pragma: no branch # Attempt to make a request after termination - with pytest.raises(MCPError) as exc_info: # pragma: no branch + with pytest.raises(SessionNotFoundError) as exc_info: # pragma: no branch await session.list_tools() assert exc_info.value.error.code == INVALID_REQUEST assert exc_info.value.error.message == "Session not found"