tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 04:28:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Return 200 response for stale dlete requests from client

Browse Source
This commit is contained in:
Sameer Kankute
2026-02-12 19:21:14 +05:30
parent af3acdda18
commit e8f97bbfca
2 changed files with 217 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/server.py
+47 -11
View File
@@ -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:
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py
+170 -19
View File
@@ -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():
"""