scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

fix(mcp): send debug headers immediately for GET streams

Browse source
This commit is contained in:
Joshua Valluru 2026-09-09 16:22:41 -07:00
parent feeda1a36c
commit c972bbe80f
3 changed files with 19 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/_experimental/mcp_server/mcp_debug.py
View file

@ -308,14 +308,18 @@ class MCPDebug:
@staticmethod
def wrap_send_with_debug_headers(
send: Send, debug_headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]] | None = None
send: Send,
debug_headers: Mapping[str, str],
resolution: Callable[[], Mapping[str, str]] | None = None,
*,
request_method: str | None = None,
) -> Send:
"""
Return a new ASGI ``send`` callable that injects *debug_headers*
into the ``http.response.start`` message.
"""
if resolution is not None:
if resolution is not None and request_method == "POST":
return _DiagnosticSend(send, debug_headers, resolution)
async def _send_with_debug(message: Message) -> None:

4
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -4482,7 +4482,9 @@ if MCP_AVAILABLE:
diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None
if diagnostics is not None:
scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics
send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers, diagnostics.headers)
send = MCPDebug.wrap_send_with_debug_headers(
send, _debug_headers, diagnostics.headers, request_method=scope.get("method")
)
# Ensure session managers are initialized
if not _SESSION_MANAGERS_INITIALIZED:

15
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py
View file

@ -208,20 +208,25 @@ class TestWrapSendWithDebugHeaders:
@pytest.mark.asyncio
@pytest.mark.parametrize("source", tuple(AuthResolution))
async def test_debug_uses_resolution_recorded_after_response_start(source: AuthResolution) -> None:
@pytest.mark.parametrize("method", ("GET", "DELETE", "POST"))
async def test_debug_defers_resolution_until_first_frame_only_for_post(source: AuthResolution, method: str) -> None:
captured: Final[list[Message]] = []
diagnostics: Final = MCPAuthDiagnostics()
async def send(message: Message) -> None:
captured.append(message)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(
send, diagnostics.headers(), diagnostics.headers, request_method=method
)
await wrapped({"type": "http.response.start", "status": 200, "headers": []})
assert captured == []
assert len(captured) == (0 if method == "POST" else 1)
diagnostics.record("s1", source)
body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True}
await wrapped(body)
assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == source.value.encode()
assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == (
source.value.encode() if method == "POST" else b"unresolved"
)
assert captured[1] == body
@ -233,7 +238,7 @@ async def test_early_stream_frame_reports_unresolved_without_waiting() -> None:
async def send(message: Message) -> None:
captured.append(message)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers, request_method="POST")
await wrapped({"type": "http.response.start", "status": 200, "headers": []})
await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True})
diagnostics.record("s1", AuthResolution.stored_user_token)