diff --git a/src/mcp/server/streamable_http_manager.py b/src/mcp/server/streamable_http_manager.py index 31f587ee66..3f2ae8ef3e 100644 --- a/src/mcp/server/streamable_http_manager.py +++ b/src/mcp/server/streamable_http_manager.py @@ -281,6 +281,9 @@ async def _handle_stateful_request(self, scope: Scope, receive: Receive, send: S if transport.idle_scope is not None and self.session_idle_timeout is not None: transport.idle_scope.deadline = anyio.current_time() + self.session_idle_timeout # pragma: no cover await transport.handle_request(scope, receive, send) + if transport.is_terminated: + self._server_instances.pop(request_mcp_session_id, None) + self._session_owners.pop(request_mcp_session_id, None) return if request_mcp_session_id is None: diff --git a/tests/server/test_streamable_http_manager.py b/tests/server/test_streamable_http_manager.py index 70440d9d03..7cb90e6b2b 100644 --- a/tests/server/test_streamable_http_manager.py +++ b/tests/server/test_streamable_http_manager.py @@ -304,6 +304,36 @@ async def mock_receive(): assert not manager._server_instances, "No sessions should be tracked after the only session exits gracefully" +@pytest.mark.anyio +async def test_terminated_existing_session_is_removed_from_registry( + running_manager: tuple[StreamableHTTPSessionManager, Server], +): + manager, _app = running_manager + session_id = "terminated-session" + transport = AsyncMock() + transport.is_terminated = True + transport.idle_scope = None + manager._server_instances[session_id] = transport + + scope = { + "type": "http", + "method": "DELETE", + "path": "/mcp", + "headers": [(MCP_SESSION_ID_HEADER.encode(), session_id.encode())], + } + + async def mock_receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def mock_send(_message: Message): + pass + + await manager.handle_request(scope, mock_receive, mock_send) + + assert session_id not in manager._server_instances + assert session_id not in manager._session_owners + + @pytest.mark.anyio async def test_stateful_session_cleanup_on_exception(running_manager: tuple[StreamableHTTPSessionManager, Server]): manager, _app = running_manager