mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
d76db470f9
* test: update integration tests left stale by recent product changes * test: expect proxy Messages calls to stay on Python with Rust enabled
389 lines
18 KiB
Python
389 lines
18 KiB
Python
import itertools
|
|
import re
|
|
import uuid
|
|
from collections.abc import Mapping
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.mcp import (
|
|
ENTRY_POINTS,
|
|
EntryPoint,
|
|
JsonRpc,
|
|
McpCaller,
|
|
Outcome,
|
|
ScriptedTool,
|
|
disconnecting_tool,
|
|
echo_tool,
|
|
listed_tools,
|
|
mcp_peer,
|
|
register_mcp,
|
|
scripted_peer,
|
|
slow_tool,
|
|
text_result,
|
|
tool_calls,
|
|
)
|
|
from integration._support.process import owned_proxy_process
|
|
from integration._support.wire import Reply
|
|
from pydantic import BaseModel, TypeAdapter
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
_ECHO: Final = "catalog-echo:"
|
|
_PROBE: Final = "catalog-probe"
|
|
_DESCRIPTION: Final = "Look up one record"
|
|
_FETCHED: Final = re.compile(rf"{_DESCRIPTION} [0-9a-f]{{32}}")
|
|
_GUARDRAIL_CODE: Final = (
|
|
"def apply_guardrail(inputs, request_data, input_type):\n"
|
|
' texts = list(inputs.get("texts") or [])\n'
|
|
' function = inputs.get("tools", [{}])[0].get("function", {})\n'
|
|
f' if "{_PROBE}" in texts:\n'
|
|
f' return block("{_ECHO}" + json_stringify({{"description": function.get("description")}}))\n'
|
|
" return allow()\n"
|
|
)
|
|
_BURST: Final = 20
|
|
_OUTAGE: Final = 6
|
|
_SPEND_NONCES: Final = (
|
|
"SELECT status, metadata->'mcp_tool_call_metadata'->'arguments'->>'nonce' AS nonce"
|
|
' FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s'
|
|
)
|
|
_OBJECTS: Final = TypeAdapter(Mapping[str, object])
|
|
_STRINGS: Final = TypeAdapter(Mapping[str, str])
|
|
_ECHOED: Final = TypeAdapter(Mapping[str, str | None])
|
|
|
|
|
|
class _Content(BaseModel):
|
|
text: str
|
|
|
|
|
|
class _Result(BaseModel):
|
|
content: tuple[_Content, ...]
|
|
|
|
|
|
class _RpcReply(BaseModel):
|
|
id: int
|
|
result: _Result
|
|
|
|
|
|
class _ListedTool(BaseModel):
|
|
name: str
|
|
description: str | None = None
|
|
|
|
|
|
class _ListedTools(BaseModel):
|
|
tools: tuple[_ListedTool, ...]
|
|
|
|
|
|
class _ListReply(BaseModel):
|
|
id: int
|
|
result: _ListedTools
|
|
|
|
|
|
class _SessionsReport(BaseModel):
|
|
worker_pid: int
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def echo_config(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
|
base: Final = _OBJECTS.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
|
guardrail: Final = {
|
|
"guardrail_name": "catalog-echo-" + uuid.uuid4().hex[:8],
|
|
"litellm_params": {
|
|
"guardrail": "custom_code",
|
|
"mode": "pre_mcp_call",
|
|
"default_on": True,
|
|
"custom_code": _GUARDRAIL_CODE,
|
|
},
|
|
}
|
|
path: Final = tmp_path_factory.mktemp("failure-recovery") / "config.yaml"
|
|
path.write_text(yaml.safe_dump({**base, "guardrails": [guardrail]}))
|
|
return path
|
|
|
|
|
|
def _call(caller: McpCaller, name: str, arguments: dict[str, object], entry: EntryPoint, identity: str) -> Outcome:
|
|
return caller.call(name, arguments, identity if entry == "rest" else None)
|
|
|
|
|
|
def _health(gateway: Gateway, key: str, identity: str) -> str:
|
|
response: Final = gateway.client.get(
|
|
"/v1/mcp/server/health", headers={"x-litellm-api-key": key}, params={"server_ids": [identity]}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
statuses: Final = {entry["server_id"]: entry["status"] for entry in response.json()}
|
|
assert identity in statuses, response.text
|
|
return str(statuses[identity])
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ENTRY_POINTS)
|
|
def test_tool_error_surfaces_as_error_with_the_peer_message_and_never_as_success(
|
|
gateway: Gateway, entry: EntryPoint
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "toolerr" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(gateway, key, entry, alias)
|
|
peer.drain()
|
|
outcome: Final = _call(caller, f"{alias}-fail", {}, entry, identity)
|
|
assert outcome.error is not None, f"failing tool reported success: {outcome.raw}"
|
|
assert "Error executing tool fail" in str(outcome.raw), outcome.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
recovered: Final = _call(caller, f"{alias}-add", {"a": 2, "b": 3}, entry, identity)
|
|
assert recovered.text == "5", recovered.raw
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ENTRY_POINTS)
|
|
def test_unreachable_peer_errors_while_a_healthy_sibling_keeps_serving(gateway: Gateway, entry: EntryPoint) -> None:
|
|
with mcp_peer() as healthy, gateway.scenario() as scenario:
|
|
good: Final = "good" + uuid.uuid4().hex[:8]
|
|
bad: Final = "bad" + uuid.uuid4().hex[:8]
|
|
good_id: Final = register_mcp(scenario, healthy, good)
|
|
bad_id: Final = register_mcp(scenario, healthy, bad, url="http://127.0.0.1:9/mcp")
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [good_id, bad_id]})
|
|
caller: Final = McpCaller(gateway, key, entry, good)
|
|
listing: Final = caller.list_tools(good_id if entry == "rest" else None)
|
|
assert listing.error is None, listing.raw
|
|
assert {f"{good}-add", "add"} & set(listing.tools), listing.raw
|
|
assert not {f"{bad}-add"} & set(listing.tools) or entry != "rest", listing.raw
|
|
healthy.drain()
|
|
served: Final = _call(caller, f"{good}-add", {"a": 2, "b": 3}, entry, good_id)
|
|
assert served.text == "5", served.raw
|
|
assert len(tool_calls(healthy.drain())) == 1
|
|
if entry == "server_mcp":
|
|
return
|
|
failed: Final = _call(McpCaller(gateway, key, entry, bad), f"{bad}-add", {"a": 2, "b": 3}, entry, bad_id)
|
|
assert failed.error is not None, f"call to unreachable peer succeeded: {failed.raw}"
|
|
assert failed.text != "5"
|
|
|
|
|
|
def test_unreachable_peer_is_reported_unhealthy_and_healthy_peer_healthy(gateway: Gateway) -> None:
|
|
with mcp_peer() as healthy, gateway.scenario() as scenario:
|
|
good: Final = "hgood" + uuid.uuid4().hex[:8]
|
|
bad: Final = "hbad" + uuid.uuid4().hex[:8]
|
|
good_id: Final = register_mcp(scenario, healthy, good)
|
|
bad_id: Final = register_mcp(scenario, healthy, bad, url="http://127.0.0.1:9/mcp")
|
|
assert _health(gateway, gateway.key, good_id) == "healthy"
|
|
assert _health(gateway, gateway.key, bad_id) == "unhealthy"
|
|
|
|
|
|
def test_slow_peer_beyond_configured_timeout_errors_and_does_not_hang_the_gateway(gateway: Gateway) -> None:
|
|
with scripted_peer(slow_tool("nap", 4), echo_tool("echo")) as peer, gateway.scenario() as scenario:
|
|
alias: Final = "slow" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, timeout=1)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(gateway, key, "mcp", alias)
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-nap", {})
|
|
assert outcome.error is not None, f"call past the timeout succeeded: {outcome.raw}"
|
|
assert outcome.text != "slept"
|
|
quick: Final = caller.call(f"{alias}-echo", {"k": "v"})
|
|
assert quick.text == '{"k": "v"}', quick.raw
|
|
|
|
|
|
def test_peer_disconnecting_mid_response_errors_and_the_next_call_succeeds(gateway: Gateway) -> None:
|
|
with scripted_peer(disconnecting_tool("drop"), echo_tool("echo")) as peer, gateway.scenario() as scenario:
|
|
alias: Final = "drop" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
for entry in ("mcp", "rest"):
|
|
caller = McpCaller(gateway, key, entry, alias)
|
|
dropped = caller.call(f"{alias}-drop", {}, identity if entry == "rest" else None)
|
|
assert dropped.error is not None, f"half-written reply became success on {entry}: {dropped.raw}"
|
|
recovered = caller.call(f"{alias}-echo", {"n": 1}, identity if entry == "rest" else None)
|
|
assert recovered.text == '{"n": 1}', recovered.raw
|
|
|
|
|
|
def test_peer_restart_on_the_same_url_is_picked_up_without_gateway_restart(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
alias: Final = "restart" + uuid.uuid4().hex[:8]
|
|
with mcp_peer() as first:
|
|
identity: Final = register_mcp(scenario, first, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
assert set(listed_tools(gateway, key, identity)) == {"add", "multiply", "fail"}
|
|
caller: Final = McpCaller(gateway, key, "mcp", alias)
|
|
down: Final = caller.call(f"{alias}-add", {"a": 1, "b": 1})
|
|
assert down.error is not None, down.raw
|
|
with scripted_peer(echo_tool("add")) as replacement:
|
|
edited: Final = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{"server_id": identity, "server_name": alias, "alias": alias, **replacement.registration()},
|
|
)
|
|
assert edited.status_code in (200, 202), edited.text
|
|
back: Final = eventually(
|
|
lambda: caller.call(f"{alias}-add", {"a": 1, "b": 1}), lambda outcome: outcome.error is None, seconds=40
|
|
)
|
|
assert back.text == '{"a": 1, "b": 1}', back.raw
|
|
assert len(tool_calls(replacement.drain())) >= 1
|
|
|
|
|
|
def _rpc_payload(raw: str) -> str:
|
|
data: Final = tuple(line[5:].strip() for line in raw.splitlines() if line.startswith("data:"))
|
|
return data[-1] if data else raw
|
|
|
|
|
|
def _rpc_reply(raw: str) -> _RpcReply:
|
|
return _RpcReply.model_validate_json(_rpc_payload(raw))
|
|
|
|
|
|
def _call_params(call: Mapping[str, object]) -> Mapping[str, object]:
|
|
return _OBJECTS.validate_python(_OBJECTS.validate_python(call["body"])["params"])
|
|
|
|
|
|
def _call_nonce(call: Mapping[str, object]) -> str:
|
|
return _STRINGS.validate_python(_call_params(call)["arguments"])["nonce"]
|
|
|
|
|
|
def _listed(caller: McpCaller, name: str) -> Outcome:
|
|
listing: Final = eventually(caller.list_tools, lambda outcome: name in outcome.tools, seconds=45)
|
|
assert listing.error is None, (caller.gateway.client.base_url, listing.raw)
|
|
return listing
|
|
|
|
|
|
def _listed_description(caller: McpCaller, name: str) -> str:
|
|
listing: Final = _listed(caller, name)
|
|
tools: Final = _ListReply.model_validate_json(_rpc_payload(listing.raw)).result.tools
|
|
described: Final = next(tool.description for tool in tools if tool.name == name)
|
|
assert described is not None and _FETCHED.fullmatch(described), listing.raw
|
|
return described
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Worker:
|
|
caller: McpCaller
|
|
pid: int
|
|
|
|
|
|
def _worker(proxy: Gateway, key: str, alias: str) -> _Worker:
|
|
sessions: Final = proxy.client.get("/v1/mcp/sessions", headers={"x-litellm-api-key": proxy.key})
|
|
assert sessions.status_code == 200, sessions.text
|
|
return _Worker(McpCaller(proxy, key, "mcp", alias), _SessionsReport.model_validate_json(sessions.text).worker_pid)
|
|
|
|
|
|
def _served(worker: _Worker, name: str, nonce: str) -> None:
|
|
served: Final = worker.caller.call(name, {"nonce": nonce})
|
|
assert served.text == "found", (worker.pid, served.raw)
|
|
|
|
|
|
def _probed_description(worker: _Worker, name: str) -> str | None:
|
|
"""The description the pre_mcp_call guardrail on that worker was handed, recovered from its block reason."""
|
|
blocked: Final = worker.caller.call(name, {"nonce": _PROBE})
|
|
assert blocked.error is not None, (worker.pid, blocked.raw)
|
|
carrier: Final = next((item.text for item in _rpc_reply(blocked.raw).result.content if _ECHO in item.text), None)
|
|
assert carrier is not None, (worker.pid, blocked.raw)
|
|
return _ECHOED.validate_json(carrier.split(_ECHO, 1)[1])["description"]
|
|
|
|
|
|
def _catalog_is_cold(worker: _Worker, name: str, listings: tuple[str, ...]) -> bool:
|
|
probed: Final = _probed_description(worker, name)
|
|
return probed is not None and _FETCHED.fullmatch(probed) is not None and probed not in listings
|
|
|
|
|
|
@pytest.mark.timeout(600)
|
|
def test_worker_restart_cools_its_listed_catalog_while_the_sibling_worker_keeps_serving(
|
|
gateway: Gateway, echo_config: Path, tmp_path: Path
|
|
) -> None:
|
|
tool: Final = ScriptedTool(
|
|
"lookup", lambda _: text_result("found"), description=lambda _: f"{_DESCRIPTION} {uuid.uuid4().hex}"
|
|
)
|
|
with scripted_peer(tool) as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cold" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
name: Final = f"{alias}-lookup"
|
|
with owned_proxy_process(gateway, tmp_path / "sibling", {}, config=echo_config) as sibling_proxy:
|
|
sibling: Final = _worker(sibling_proxy.gateway, key, alias)
|
|
with owned_proxy_process(gateway, tmp_path / "first", {}, config=echo_config) as first_proxy:
|
|
first: Final = _worker(first_proxy.gateway, key, alias)
|
|
assert first.pid != sibling.pid
|
|
_served(first, name, "first-unlisted")
|
|
_served(sibling, name, "sibling-unlisted")
|
|
assert _catalog_is_cold(first, name, ()) and _catalog_is_cold(sibling, name, ())
|
|
first_listing: Final = _listed_description(first.caller, name)
|
|
assert _probed_description(first, name) == first_listing
|
|
assert _catalog_is_cold(sibling, name, (first_listing,)), "a listing on one worker warmed its sibling"
|
|
sibling_listing: Final = _listed_description(sibling.caller, name)
|
|
assert _probed_description(sibling, name) == sibling_listing
|
|
_served(sibling, name, "sibling-alone")
|
|
with owned_proxy_process(gateway, tmp_path / "restarted", {}, config=echo_config) as restarted_proxy:
|
|
restarted: Final = _worker(restarted_proxy.gateway, key, alias)
|
|
assert restarted.pid not in (first.pid, sibling.pid)
|
|
assert _catalog_is_cold(restarted, name, (first_listing, sibling_listing)), (
|
|
"a restarted worker kept the old process's catalog"
|
|
)
|
|
assert _probed_description(sibling, name) == sibling_listing
|
|
_served(restarted, name, "restarted-unlisted")
|
|
restarted_listing: Final = _listed_description(restarted.caller, name)
|
|
assert _probed_description(restarted, name) == restarted_listing
|
|
calls: Final = tool_calls(peer.drain())
|
|
assert [_call_nonce(call) for call in calls] == [
|
|
"first-unlisted",
|
|
"sibling-unlisted",
|
|
"sibling-alone",
|
|
"restarted-unlisted",
|
|
], calls
|
|
assert all(set(_call_params(call)) - {"_meta"} == {"name", "arguments"} for call in calls), calls
|
|
|
|
|
|
def _outage_echo(name: str, failures: int) -> ScriptedTool:
|
|
attempts: Final = itertools.count(1)
|
|
|
|
def respond(params: JsonRpc) -> Reply | JsonRpc:
|
|
if next(attempts) <= failures:
|
|
return Reply(status=503, body=b'{"error": "scripted outage"}')
|
|
return text_result(_STRINGS.validate_python(params["arguments"])["nonce"])
|
|
|
|
return ScriptedTool(name, respond)
|
|
|
|
|
|
def _echo_call(caller: McpCaller, name: str, nonce: str) -> Outcome:
|
|
return caller.call(name, {"nonce": nonce})
|
|
|
|
|
|
def _logged_nonces(key: str, count: int) -> tuple[tuple[str, str], ...]:
|
|
rows: Final = eventually(
|
|
lambda: read_rows(_SPEND_NONCES, (sha256(key.encode()).hexdigest(), "call_mcp_tool")),
|
|
lambda found: len(found) >= count,
|
|
seconds=70,
|
|
)
|
|
return tuple(sorted((str(row["status"]), str(row["nonce"])) for row in rows))
|
|
|
|
|
|
def test_peer_outage_during_a_bounded_burst_fails_exactly_the_outage_calls_and_lands_each_call_once(
|
|
gateway: Gateway, peer: Gateway
|
|
) -> None:
|
|
with scripted_peer(_outage_echo("echo", _OUTAGE)) as upstream, gateway.scenario() as scenario:
|
|
alias: Final = "burst" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, upstream, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
name: Final = f"{alias}-echo"
|
|
callers: Final = (McpCaller(gateway, key, "mcp", alias), McpCaller(peer, key, "mcp", alias))
|
|
for caller in callers:
|
|
_listed(caller, name)
|
|
upstream.drain()
|
|
nonces: Final = tuple(uuid.uuid4().hex for _ in range(_BURST))
|
|
with ThreadPoolExecutor(max_workers=_BURST) as pool:
|
|
outcomes: Final = tuple(pool.map(_echo_call, itertools.cycle(callers), itertools.repeat(name), nonces))
|
|
raws: Final = [outcome.raw for outcome in outcomes]
|
|
failed: Final = tuple(nonce for nonce, outcome in zip(nonces, outcomes) if outcome.error is not None)
|
|
assert len(failed) == _OUTAGE, raws
|
|
assert all(outcome.error is not None or outcome.text == nonce for nonce, outcome in zip(nonces, outcomes)), raws
|
|
assert all(_rpc_reply(outcome.raw).id == 1 for outcome in outcomes), raws
|
|
burst_calls: Final = tool_calls(upstream.drain())
|
|
assert sorted(_call_nonce(call) for call in burst_calls) == sorted(nonces), burst_calls
|
|
assert all(_call_params(call)["arguments"] == {"nonce": _call_nonce(call)} for call in burst_calls), burst_calls
|
|
assert all(set(_call_params(call)) - {"_meta"} == {"name", "arguments"} for call in burst_calls), burst_calls
|
|
recovered: Final = tuple(
|
|
_echo_call(caller, name, nonce) for caller, nonce in zip(itertools.cycle(callers), failed)
|
|
)
|
|
assert [outcome.text for outcome in recovered] == list(failed), [outcome.raw for outcome in recovered]
|
|
assert sorted(_call_nonce(call) for call in tool_calls(upstream.drain())) == sorted(failed)
|
|
assert _logged_nonces(key, _BURST + len(failed)) == tuple(
|
|
sorted([("success", nonce) for nonce in nonces] + [("failure", nonce) for nonce in failed])
|
|
)
|