mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
901b7f762f
* test(mcp): cancel disconnected HTTP tool handlers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): verify owned peer disconnect cleanup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
318 lines
13 KiB
Python
318 lines
13 KiB
Python
import asyncio
|
|
from collections import deque
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from starlette.types import Message, Scope
|
|
from starlette.websockets import WebSocket
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", (OSError("client disconnected"), RuntimeError("unexpected upstream failure")))
|
|
async def test_realtime_disconnect_after_session_creation_keeps_unexpected_errors_visible(
|
|
failure: Exception, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.upstream import Provider
|
|
from integration.cost_calculation.cost_tracking_case import RealtimeResponse
|
|
|
|
provider: Final = Provider()
|
|
response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},))
|
|
provider.scenario_store.put("owned-scenario", response)
|
|
incoming: Final[deque[Message]] = deque(
|
|
(
|
|
{"type": "websocket.connect"},
|
|
{"type": "websocket.receive", "text": '{"type":"response.create"}'},
|
|
{"type": "websocket.disconnect", "code": 1006},
|
|
)
|
|
)
|
|
outgoing: Final[deque[Message]] = deque()
|
|
|
|
async def receive() -> Message:
|
|
return incoming.popleft()
|
|
|
|
async def send(message: Message) -> None:
|
|
if message["type"] == "websocket.send" and any(item["type"] == "websocket.send" for item in outgoing):
|
|
raise failure
|
|
outgoing.append(message)
|
|
|
|
scope: Final[Scope] = {
|
|
"type": "websocket",
|
|
"path": "/v1/realtime",
|
|
"headers": [(b"authorization", b"Bearer owned-scenario")],
|
|
"query_string": b"model=synthetic",
|
|
}
|
|
socket: Final = WebSocket(scope, receive, send)
|
|
if isinstance(failure, RuntimeError):
|
|
with pytest.raises(RuntimeError, match="unexpected upstream failure"):
|
|
await provider.realtime(socket)
|
|
else:
|
|
await provider.realtime(socket)
|
|
assert [item["type"] for item in outgoing] == ["websocket.accept", "websocket.send"]
|
|
observed: Final = provider.observations.snapshot()
|
|
assert [(item.method, item.body) for item in observed] == [
|
|
("WEBSOCKET", {"query": [["model", "synthetic"]]}),
|
|
("WEBSOCKET_FRAME", {"type": "response.create"}),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_observations_read_without_draining_until_explicitly_reset(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.upstream import Provider
|
|
|
|
transport: Final = httpx.ASGITransport(app=Provider().app())
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://upstream") as client:
|
|
sent: Final = await client.post(
|
|
"/v1/chat/completions",
|
|
json={"model": "synthetic", "messages": [{"role": "user", "content": "observed-marker"}]},
|
|
)
|
|
assert sent.status_code == 200, sent.text
|
|
first: Final = (await client.get("/__observations")).json()["requests"]
|
|
second: Final = (await client.get("/__observations")).json()["requests"]
|
|
assert [(item["path"], item["body"]["messages"][0]["content"]) for item in first] == [
|
|
("/v1/chat/completions", "observed-marker")
|
|
]
|
|
assert second == first
|
|
reset: Final = await client.delete("/__observations")
|
|
assert reset.status_code == 204
|
|
assert (await client.get("/__observations")).json()["requests"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("method,status", (("POST", 200), ("POST", 400), ("DELETE", 200)))
|
|
async def test_mcp_notification_ordering_preserves_other_responses(
|
|
method: str, status: int, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _acknowledge_notifications_after_dispatch
|
|
from starlette.types import Receive, Send
|
|
|
|
outgoing: Final[list[Message]] = []
|
|
messages: Final[tuple[Message, ...]] = (
|
|
{"type": "http.response.start", "status": status, "headers": []},
|
|
{"type": "http.response.body", "body": b"result", "more_body": True},
|
|
{"type": "http.response.body", "body": b"", "more_body": False},
|
|
)
|
|
|
|
async def receive() -> Message:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
outgoing.append(message)
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
for index, message in enumerate(messages):
|
|
await send(message)
|
|
assert outgoing == list(messages[: index + 1]), "Ordinary response streaming was buffered"
|
|
|
|
await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": method}, receive, send)
|
|
assert outgoing == list(messages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_http_disconnect_cancels_an_in_flight_handler(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _cancel_on_disconnect
|
|
from starlette.types import Receive, Send
|
|
|
|
started: Final = asyncio.Event()
|
|
stopped: Final = asyncio.Event()
|
|
body: Final[Message] = {"type": "http.request", "body": b"tool-call", "more_body": False}
|
|
pending: Final = iter((body,))
|
|
|
|
async def receive() -> Message:
|
|
message: Final = next(pending, None)
|
|
if message is not None:
|
|
return message
|
|
await started.wait()
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
pytest.fail(f"Disconnected handler sent a response: {message}")
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
assert await receive() == body
|
|
started.set()
|
|
try:
|
|
await asyncio.Future[None]()
|
|
finally:
|
|
stopped.set()
|
|
|
|
await asyncio.wait_for(_cancel_on_disconnect(app)({"type": "http"}, receive, send), 2)
|
|
assert stopped.is_set(), "The disconnected HTTP handler was left running"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", (RuntimeError("dispatch failed"), asyncio.CancelledError()))
|
|
async def test_mcp_notification_dispatch_failure_is_not_acknowledged(
|
|
failure: BaseException, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _acknowledge_notifications_after_dispatch
|
|
from starlette.types import Receive, Send
|
|
|
|
outgoing: Final[list[Message]] = []
|
|
|
|
async def receive() -> Message:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
outgoing.append(message)
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["method"] == "DELETE":
|
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
|
await send({"type": "http.response.body", "body": b""})
|
|
return
|
|
await send({"type": "http.response.start", "status": 202, "headers": []})
|
|
await send({"type": "http.response.body", "body": b""})
|
|
raise failure
|
|
|
|
ordered: Final = _acknowledge_notifications_after_dispatch(app)
|
|
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"failed-session")]}
|
|
with pytest.raises(type(failure)) as caught:
|
|
await ordered({**scope, "method": "POST"}, receive, send)
|
|
assert caught.value is failure
|
|
assert outgoing == [], "Failed dispatch must not send an acknowledgement"
|
|
await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2)
|
|
assert outgoing[0]["status"] == 200
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"responses",
|
|
(
|
|
(),
|
|
({"type": "http.response.body", "body": b"unexpected"},),
|
|
({"type": "http.response.body", "body": b"", "more_body": True},),
|
|
({"type": "http.response.start", "status": 500, "headers": []},),
|
|
({"type": "http.response.body", "body": b""}, {"type": "http.response.body", "body": b""}),
|
|
),
|
|
)
|
|
async def test_mcp_notification_ordering_rejects_invalid_acknowledgements(
|
|
responses: tuple[Message, ...], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _acknowledge_notifications_after_dispatch
|
|
from starlette.types import Receive, Send
|
|
|
|
outgoing: Final[list[Message]] = []
|
|
|
|
async def receive() -> Message:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
outgoing.append(message)
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
await send({"type": "http.response.start", "status": 202, "headers": []})
|
|
for message in responses:
|
|
await send(message)
|
|
|
|
with pytest.raises(RuntimeError, match="MCP notification acknowledgement"):
|
|
await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": "POST"}, receive, send)
|
|
assert outgoing == [], "Invalid acknowledgement must not be reported as accepted"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_session_deletion_waits_for_concurrent_notification_dispatch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _acknowledge_notifications_after_dispatch
|
|
from starlette.types import Receive, Send
|
|
|
|
started: Final = [asyncio.Event(), asyncio.Event()]
|
|
release: Final = asyncio.Event()
|
|
deletion_started: Final = asyncio.Event()
|
|
deleted: Final = asyncio.Event()
|
|
other_deleted: Final = asyncio.Event()
|
|
completed: Final[list[str]] = []
|
|
|
|
async def receive() -> Message:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
return None
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["method"] == "POST":
|
|
started[int(scope["path"])].set()
|
|
await release.wait()
|
|
completed.append(scope["path"])
|
|
await send({"type": "http.response.start", "status": 202, "headers": []})
|
|
await send({"type": "http.response.body", "body": b""})
|
|
else:
|
|
if scope["headers"][0][1] == b"other-session":
|
|
other_deleted.set()
|
|
else:
|
|
deleted.set()
|
|
completed.append("DELETE")
|
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
|
await send({"type": "http.response.body", "body": b""})
|
|
|
|
ordered: Final = _acknowledge_notifications_after_dispatch(app)
|
|
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"owned-session")]}
|
|
posts: Final = [
|
|
asyncio.create_task(ordered({**scope, "method": "POST", "path": str(index)}, receive, send))
|
|
for index in range(2)
|
|
]
|
|
|
|
async def delete() -> None:
|
|
deletion_started.set()
|
|
await ordered({**scope, "method": "DELETE"}, receive, send)
|
|
|
|
deletion = None
|
|
try:
|
|
await asyncio.wait_for(asyncio.gather(*(event.wait() for event in started)), 2)
|
|
deletion = asyncio.create_task(delete())
|
|
await deletion_started.wait()
|
|
await asyncio.wait_for(
|
|
ordered({**scope, "method": "DELETE", "headers": [(b"mcp-session-id", b"other-session")]}, receive, send), 2
|
|
)
|
|
assert other_deleted.is_set(), "Unrelated session deletion was blocked"
|
|
assert not deleted.is_set(), "Session deletion overtook in-flight notification dispatch"
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(asyncio.gather(*posts, *([deletion] if deletion is not None else [])), 2)
|
|
assert set(completed[:2]) == {"0", "1"}
|
|
assert completed[-1] == "DELETE"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_session_deletion_can_terminate_streaming_request(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
|
|
from integration._support.mcp import _acknowledge_notifications_after_dispatch
|
|
from starlette.types import Receive, Send
|
|
|
|
streaming: Final = asyncio.Event()
|
|
deleted: Final = asyncio.Event()
|
|
|
|
async def receive() -> Message:
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(message: Message) -> None:
|
|
return None
|
|
|
|
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
|
if scope["method"] == "POST":
|
|
streaming.set()
|
|
await deleted.wait()
|
|
else:
|
|
deleted.set()
|
|
await send({"type": "http.response.body", "body": b""})
|
|
|
|
ordered: Final = _acknowledge_notifications_after_dispatch(app)
|
|
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"streaming-session")]}
|
|
post: Final = asyncio.create_task(ordered({**scope, "method": "POST"}, receive, send))
|
|
try:
|
|
await asyncio.wait_for(streaming.wait(), 2)
|
|
await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2)
|
|
assert deleted.is_set(), "Session deletion must terminate an ordinary streaming request"
|
|
finally:
|
|
deleted.set()
|
|
await asyncio.wait_for(post, 2)
|