From e8f97bbfca722a4b46719477184b036194daba89 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 12 Feb 2026 19:21:14 +0530 Subject: [PATCH] Return 200 response for stale dlete requests from client --- .../proxy/_experimental/mcp_server/server.py | 58 +++++- .../mcp_server/test_mcp_stale_session.py | 189 ++++++++++++++++-- 2 files changed, 217 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 58cd8c99e7..9f7d728896 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -24,6 +24,7 @@ from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict from starlette.requests import Request as StarletteRequest from starlette.types import Receive, Scope, Send +from starlette.responses import JSONResponse from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -1907,18 +1908,27 @@ if MCP_AVAILABLE: raw_headers, ) - def _strip_stale_mcp_session_header( + async def _handle_stale_mcp_session( scope: Scope, + receive: Receive, + send: Send, mgr: "StreamableHTTPSessionManager", - ) -> None: + ) -> bool: """ - Strip stale ``mcp-session-id`` headers so the session manager - creates a fresh session instead of returning 404 "Session not found". + Handle stale MCP session IDs to prevent "Session not found" errors. - When clients like VSCode reconnect after a reload they may resend a - session id that has already been cleaned up. Rather than letting the - SDK return a 404 error loop, we detect the stale id and remove the - header so a brand-new session is created transparently. + When clients reconnect after a server restart or session cleanup, they may + send a session ID that no longer exists. This function handles two scenarios: + + 1. Non-DELETE requests: Strip the stale session ID header so the session + manager creates a fresh session transparently. + + 2. DELETE requests: Return success (200) immediately for idempotent behavior, + since the desired state (session doesn't exist) is already achieved. + + Returns: + True if the request was handled (DELETE on non-existent session) + False if the request should continue to the session manager Fixes https://github.com/BerriAI/litellm/issues/20292 """ @@ -1930,10 +1940,30 @@ if MCP_AVAILABLE: break if _session_id is None: - return + return False known_sessions = getattr(mgr, "_server_instances", None) - if known_sessions is not None and _session_id not in known_sessions: + if known_sessions is None or _session_id in known_sessions: + # Session exists or we can't check - let the session manager handle it + return False + + # Session doesn't exist - handle based on request method + method = scope.get("method", "").upper() + + if method == "DELETE": + # Idempotent DELETE: session doesn't exist, return success + verbose_logger.info( + f"DELETE request for non-existent MCP session '{_session_id}'. " + "Returning success (idempotent DELETE)." + ) + success_response = JSONResponse( + status_code=200, + content={"message": "Session terminated successfully"} + ) + await success_response(scope, receive, send) + return True + else: + # Non-DELETE: strip stale session ID to allow new session creation verbose_logger.warning( "MCP session ID '%s' not found in active sessions. " "Stripping stale header to force new session creation.", @@ -1943,6 +1973,7 @@ if MCP_AVAILABLE: (k, v) for k, v in scope["headers"] if k != _mcp_session_header ] + return False async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send @@ -2005,7 +2036,12 @@ if MCP_AVAILABLE: # Give it a moment to start up await asyncio.sleep(0.1) - _strip_stale_mcp_session_header(scope, session_manager) + # Handle stale session IDs - either strip them for reconnection + # or return success for idempotent DELETE operations + handled = await _handle_stale_mcp_session(scope, receive, send, session_manager) + if handled: + # Request was fully handled (e.g., DELETE on non-existent session) + return await session_manager.handle_request(scope, receive, send) except HTTPException: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index a447ee6af0..5eb8c1e51a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -1,105 +1,193 @@ """ Tests for MCP stale session ID handling (Fixes #20292). -When VSCode reconnects to LiteLLM's MCP endpoint after a reload, it sends a stale -`mcp-session-id` header. The session manager returns a 404 because the old session -was cleaned up. This test verifies that stale session IDs are detected and stripped -so a new session is created automatically. +When clients reconnect to LiteLLM's MCP endpoint after a server restart or reload, +they may send a stale `mcp-session-id` header. This test verifies that: +1. For non-DELETE requests: stale session IDs are stripped so new sessions are created +2. For DELETE requests: idempotent behavior returns success even if session doesn't exist """ import pytest from unittest.mock import AsyncMock, MagicMock, patch -class TestStripStaleMcpSessionHeader: - """Unit tests for the _strip_stale_mcp_session_header helper.""" +class TestHandleStaleMcpSession: + """Unit tests for the _handle_stale_mcp_session helper.""" - def test_strips_stale_session_id(self): + @pytest.mark.asyncio + async def test_strips_stale_session_id_for_non_delete(self): + """Non-DELETE requests should have stale session IDs stripped.""" try: from litellm.proxy._experimental.mcp_server.server import ( - _strip_stale_mcp_session_header, + _handle_stale_mcp_session, ) except ImportError: pytest.skip("MCP server not available") scope = { + "method": "POST", "headers": [ (b"content-type", b"application/json"), (b"mcp-session-id", b"stale-id"), ], } + receive = AsyncMock() + send = AsyncMock() mgr = MagicMock() mgr._server_instances = {} # no active sessions - _strip_stale_mcp_session_header(scope, mgr) + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + # Should not be fully handled (returns False) + assert handled is False + # Header should be stripped header_names = [k for k, _ in scope["headers"]] assert b"mcp-session-id" not in header_names - def test_preserves_valid_session_id(self): + @pytest.mark.asyncio + async def test_delete_stale_session_returns_success(self): + """DELETE requests for non-existent sessions should return success (idempotent).""" try: from litellm.proxy._experimental.mcp_server.server import ( - _strip_stale_mcp_session_header, + _handle_stale_mcp_session, ) except ImportError: pytest.skip("MCP server not available") scope = { + "type": "http", + "method": "DELETE", + "headers": [ + (b"content-type", b"application/json"), + (b"mcp-session-id", b"stale-id"), + ], + } + receive = AsyncMock() + send = AsyncMock() + mgr = MagicMock() + mgr._server_instances = {} # no active sessions + + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + + # Should be fully handled (returns True) + assert handled is True + # Should have sent a success response + assert send.called + # Header should NOT be stripped (DELETE needs the session ID) + header_names = [k for k, _ in scope["headers"]] + assert b"mcp-session-id" in header_names + + @pytest.mark.asyncio + async def test_preserves_valid_session_id(self): + """Valid session IDs should not be modified.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _handle_stale_mcp_session, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "method": "POST", "headers": [ (b"content-type", b"application/json"), (b"mcp-session-id", b"valid-id"), ], } + receive = AsyncMock() + send = AsyncMock() mgr = MagicMock() mgr._server_instances = {"valid-id": MagicMock()} - _strip_stale_mcp_session_header(scope, mgr) + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + # Should not be handled (returns False) + assert handled is False + # Header should be preserved header_names = [k for k, _ in scope["headers"]] assert b"mcp-session-id" in header_names - def test_no_op_when_no_session_header(self): + @pytest.mark.asyncio + async def test_no_op_when_no_session_header(self): + """No session header should result in no-op.""" try: from litellm.proxy._experimental.mcp_server.server import ( - _strip_stale_mcp_session_header, + _handle_stale_mcp_session, ) except ImportError: pytest.skip("MCP server not available") scope = { + "method": "POST", "headers": [ (b"content-type", b"application/json"), ], } + receive = AsyncMock() + send = AsyncMock() mgr = MagicMock() mgr._server_instances = {} - _strip_stale_mcp_session_header(scope, mgr) + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + assert handled is False assert len(scope["headers"]) == 1 - def test_no_op_when_server_instances_missing(self): + @pytest.mark.asyncio + async def test_no_op_when_server_instances_missing(self): """If _server_instances attr doesn't exist, don't crash.""" try: from litellm.proxy._experimental.mcp_server.server import ( - _strip_stale_mcp_session_header, + _handle_stale_mcp_session, ) except ImportError: pytest.skip("MCP server not available") scope = { + "method": "POST", "headers": [ (b"mcp-session-id", b"some-id"), ], } + receive = AsyncMock() + send = AsyncMock() mgr = MagicMock(spec=[]) # no attributes - _strip_stale_mcp_session_header(scope, mgr) + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) - # Should keep the header since we can't verify + # Should not be handled, header should be kept + assert handled is False header_names = [k for k, _ in scope["headers"]] assert b"mcp-session-id" in header_names + @pytest.mark.asyncio + async def test_delete_valid_session_not_handled(self): + """DELETE requests for existing sessions should not be intercepted.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _handle_stale_mcp_session, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "method": "DELETE", + "headers": [ + (b"mcp-session-id", b"valid-id"), + ], + } + receive = AsyncMock() + send = AsyncMock() + mgr = MagicMock() + mgr._server_instances = {"valid-id": MagicMock()} + + handled = await _handle_stale_mcp_session(scope, receive, send, mgr) + + # Should not be handled - let session manager handle it + assert handled is False + # Should not have sent any response + assert not send.called + @pytest.mark.asyncio async def test_stale_mcp_session_id_is_stripped(): @@ -166,6 +254,69 @@ async def test_stale_mcp_session_id_is_stripped(): ) +@pytest.mark.asyncio +async def test_delete_stale_mcp_session_returns_success(): + """ + When a DELETE request is made for a session that no longer exists, + handle_streamable_http_mcp should return success (200) immediately + without forwarding to the session manager (idempotent DELETE). + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + stale_session_id = "stale-session-id-12345" + + scope = { + "type": "http", + "method": "DELETE", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"mcp-session-id", stale_session_id.encode()), + (b"authorization", b"Bearer test-key"), + ], + } + + receive = AsyncMock() + send = AsyncMock() + + # Mock handle_request should NOT be called for stale DELETE + mock_handle_request = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, None, None, None, None), + ), patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", + ), patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), patch.object( + session_manager, + "handle_request", + side_effect=mock_handle_request, + ), patch.object( + session_manager, + "_server_instances", + {}, # Empty dict = no active sessions + ): + await handle_streamable_http_mcp(scope, receive, send) + + # Verify session manager was NOT called (request was handled early) + assert not mock_handle_request.called, ( + "Session manager should not be called for DELETE on non-existent session" + ) + + # Verify a success response was sent + assert send.called, "A response should have been sent" + + @pytest.mark.asyncio async def test_valid_mcp_session_id_is_preserved(): """