diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index b25ed27becc..ff30ef99ebd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b53c55ef5fa..72d256e90fc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index 5039209c351..7d46bb0237a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -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)