diff --git a/litellm/proxy/_experimental/mcp_server/scope_challenge.py b/litellm/proxy/_experimental/mcp_server/scope_challenge.py index c326418bee7..121234e305d 100644 --- a/litellm/proxy/_experimental/mcp_server/scope_challenge.py +++ b/litellm/proxy/_experimental/mcp_server/scope_challenge.py @@ -33,9 +33,8 @@ class OAuthScopeResponse: self._request_path = request_path self._ready = asyncio.Event() self._lock = asyncio.Lock() - self._start: Message | None = None self._challenge: MCPUpstreamAuthError | None = None - self._refused = False + self._closed = False @property def pending(self) -> bool: @@ -50,31 +49,33 @@ class OAuthScopeResponse: self._ready.set() async def __call__(self, message: Message) -> None: - if message["type"] == "http.response.start": - self._start = message - if message["status"] != 200: - self.allow() - return - await self._ready.wait() async with self._lock: - if self._refused: + if self._closed: return - if self._start is not None: - start: Final = self._start - self._start = None - if self._challenge is not None: - error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) - response: Final = JSONResponse( - status_code=error.status_code, content={"detail": error.detail}, headers=error.headers - ) - self._refused = True - await self._send( - {"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers} - ) - await self._send({"type": "http.response.body", "body": bytes(response.body)}) - return - await self._send(start) - await self._send(message) + try: + if message["type"] == "http.response.start": + if message["status"] != 200: + self.allow() + await self._ready.wait() + if self._challenge is not None: + error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) + response: Final = JSONResponse( + status_code=error.status_code, content={"detail": error.detail}, headers=error.headers + ) + self._closed = True + await self._send( + { + "type": "http.response.start", + "status": response.status_code, + "headers": response.raw_headers, + } + ) + await self._send({"type": "http.response.body", "body": bytes(response.body)}) + return + await self._send(message) + except BaseException: + self._closed = True + raise def scope_response_from_request(request: object) -> OAuthScopeResponse | None: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index b0d53f9c93b..34b021df2b9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -5,6 +5,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import logging import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -531,12 +532,12 @@ async def test_scope_challenge_replaces_pending_stream_headers_once(): send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"})) await asyncio.sleep(0) send.assert_not_awaited() response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) - await asyncio.wait_for(body, 1) + await asyncio.wait_for(asyncio.gather(headers, body), 1) await response({"type": "http.response.body", "body": b"private tool error"}) assert send.await_count == 2 start, result = [call.args[0] for call in send.await_args_list] @@ -567,7 +568,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) await asyncio.sleep(0) for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]: @@ -582,7 +583,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re upstream = MCPClient(server_url="https://upstream.example/mcp") async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client: async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}): - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) assert send.await_args_list[-1].args[0]["body"] == b"progress" gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) await gate({"type": "http.response.body", "body": b"result"}) @@ -594,21 +595,105 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re @pytest.mark.asyncio @pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) -async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure): +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation( + failure: type[BaseException], fail_send: bool, +) -> None: from mcp.server.context import ServerRequestContext from starlette.requests import Request from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response - send = AsyncMock() - gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") - context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + send: Final = AsyncMock(side_effect=BrokenPipeError("peer disconnected") if fail_send else None) + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context: Final = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) with pytest.raises(failure): await finish_scope_response(context, AsyncMock(side_effect=failure)) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + await headers await gate({"type": "http.response.body", "body": b"error"}) assert not gate.pending - assert send.await_args_list[-1].args[0]["body"] == b"error" + assert send.await_args_list[-1].args[0]["type"] == ("http.response.start" if fail_send else "http.response.body") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_transport_owns_header_send_cancellation_and_errors(fail_send: bool) -> None: + import httpx2 + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send_started: Final = asyncio.Event() + send_stopped: Final = asyncio.Event() + writable: Final = asyncio.Event() + sent: Final[list[Message]] = [] + + async def send(message: Message) -> None: + sent.append(message) + assert message["type"] == "http.response.start" + send_started.set() + try: + if fail_send: + raise BrokenPipeError("peer disconnected") + await writable.wait() + finally: + send_stopped.set() + + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + body: Final = asyncio.create_task(gate({"type": "http.response.body", "body": b"result"})) + await asyncio.sleep(0) + assert not send_started.is_set() + response: Final = httpx2.Response(200, request=httpx2.Request( + "POST", "https://upstream.example/mcp", json={"method": "tools/call"}, + )) + try: + await record_upstream_tool_authorization(response, pending=gate) + await send_started.wait() + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + headers.cancel() + with pytest.raises(asyncio.CancelledError): + await headers + await body + assert len(sent) == 1 + assert send_stopped.is_set() + finally: + writable.set() + await send_stopped.wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_before_start", [False, True]) +async def test_scope_authorization_sends_stream_headers_before_the_first_body_chunk(allowed_before_start: bool) -> None: + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + start: Final[Message] = {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]} + if allowed_before_start: + gate.allow() + headers: Final = asyncio.create_task(gate(start)) + await asyncio.sleep(0) + if not allowed_before_start: + send.assert_not_awaited() + gate.allow() + await headers + send.assert_awaited_once_with(start) + await gate({"type": "http.response.body", "body": b"event: message\n", "more_body": True}) + assert send.await_args_list[-1].args[0]["body"] == b"event: message\n" + assert send.await_count == 2 @pytest.mark.asyncio @@ -697,13 +782,14 @@ async def test_scope_challenges_are_isolated_between_concurrent_requests(): context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) await asyncio.sleep(0) assert get_scope_response() is gate if denied: gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) else: await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response()) + await headers await gate({"type": "http.response.body", "body": b"done"}) return send.await_args_list[0].args[0]["status"] finally: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7315315a619..ecd08e17cda 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7681,7 +7681,7 @@ class TestMCPServerManager: request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}), ) token = active_mcp_request_ctx_var.set(context) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) async def upstream_call(*args, **kwargs): @@ -7701,7 +7701,7 @@ class TestMCPServerManager: assert not result.is_error finally: gate.allow() - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) active_mcp_request_ctx_var.reset(token) @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index b8e28c3f9b8..a873e2ab1ae 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -280,7 +280,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response}) - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( @@ -328,6 +328,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args] assert any("Upstream auth failure" in m for m in info_calls) response.allow() + await headers await response({"type": "http.response.body", "body": b"tool error"}) start = send.await_args_list[0].args[0] assert start["status"] == (403 if required_scope is not None else 200)