mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
306 lines
16 KiB
Python
306 lines
16 KiB
Python
import json
|
|
import uuid
|
|
from contextlib import ExitStack
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import yaml
|
|
from hypothesis import strategies as st
|
|
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
|
|
|
|
from integration._support.client import Gateway
|
|
from integration._support.database import read_rows
|
|
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names
|
|
|
|
|
|
@pytest.mark.covers("mcp.call_tool.saved_headers.reach_actual_transport")
|
|
def test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "integration" + uuid.uuid4().hex
|
|
identity: Final = register_mcp(
|
|
scenario, peer, alias, static_headers={"X-Integration-Saved": "synthetic-header-value"}
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
for generation in range(2):
|
|
names: Final = tool_names(gateway, key, identity)
|
|
assert set(names) == {"add", "multiply", "fail"}
|
|
peer.drain()
|
|
response: Final = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["isError"] is False
|
|
assert len(response.json()["content"]) == 1
|
|
assert response.json()["content"][0]["type"] == "text"
|
|
assert response.json()["content"][0]["text"] == "8"
|
|
calls: Final = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
|
|
assert len(calls) == 1
|
|
assert calls[0]["headers"][b"x-integration-saved"] == b"synthetic-header-value"
|
|
assert calls[0]["body"]["params"]["name"] == "add"
|
|
assert calls[0]["body"]["params"]["arguments"] == {"a": 3, "b": 5}
|
|
if generation == 0:
|
|
updated: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "server_name": alias + "renamed"}
|
|
)
|
|
assert updated.status_code == 202, updated.text
|
|
rows: Final = read_rows('SELECT server_name FROM "LiteLLM_MCPServerTable" WHERE server_id = %s', (identity,))
|
|
assert rows == [{"server_name": alias + "renamed"}]
|
|
|
|
|
|
@pytest.mark.covers("mcp.call_tool.errors.tool_failure_is_not_success")
|
|
def test_tool_error_remains_error_and_healthy_sibling_returns_value(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
identity: Final = register_mcp(scenario, peer, "integration" + uuid.uuid4().hex)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
names: Final = tool_names(gateway, key, identity)
|
|
failure: Final = call_tool(gateway, key, identity, names["fail"], {})
|
|
assert failure.status_code == 200, failure.text
|
|
assert failure.json()["isError"] is True
|
|
assert failure.json()["content"][0]["text"] == "Error executing tool fail"
|
|
healthy: Final = call_tool(gateway, key, identity, names["multiply"], {"a": 3, "b": 5})
|
|
assert healthy.status_code == 200, healthy.text
|
|
assert healthy.json()["isError"] is False
|
|
assert healthy.json()["content"][0]["text"] == "15"
|
|
|
|
|
|
@pytest.mark.timeout(180)
|
|
@pytest.mark.covers("other.mcp.lifecycle.generated_save_reload_preserves_effective_headers")
|
|
def test_generated_mcp_edits_preserve_actual_headers_and_tool_results(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, bounded_http_requests((gateway,), limit=1500) as budget:
|
|
|
|
class Servers(RuleBasedStateMachine):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.resources = ExitStack()
|
|
self.marker = "first"
|
|
self.name = "integration" + uuid.uuid4().hex
|
|
try:
|
|
scenario = self.resources.enter_context(gateway.scenario())
|
|
self.identity = register_mcp(
|
|
scenario, peer, self.name, static_headers={"X-Integration-Saved": self.marker}
|
|
)
|
|
self.key = scenario.key(object_permission={"mcp_servers": [self.identity]})
|
|
except BaseException:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
raise
|
|
|
|
@rule(value=st.sampled_from(("first", "second", "third")))
|
|
def header(self, value: str) -> None:
|
|
response: Final = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{"server_id": self.identity, "static_headers": {"X-Integration-Saved": value}},
|
|
)
|
|
assert response.status_code == 202, response.text
|
|
self.marker = value
|
|
|
|
@rule(value=st.sampled_from(("original", "renamed")))
|
|
def rename(self, value: str) -> None:
|
|
response: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": self.identity, "server_name": self.name + value}
|
|
)
|
|
assert response.status_code == 202, response.text
|
|
|
|
@invariant()
|
|
def persisted_configuration_controls_actual_tools(self) -> None:
|
|
names: Final = tool_names(gateway, self.key, self.identity)
|
|
assert set(names) == {"add", "multiply", "fail"}
|
|
peer.drain()
|
|
result: Final = call_tool(gateway, self.key, self.identity, names["add"], {"a": 3, "b": 5})
|
|
assert result.status_code == 200 and result.json()["isError"] is False, result.text
|
|
assert result.json()["content"][0]["text"] == "8"
|
|
calls: Final = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
|
|
assert len(calls) == 1 and calls[0]["headers"][b"x-integration-saved"] == self.marker.encode()
|
|
assert (
|
|
len(
|
|
read_rows('SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id=%s', (self.identity,))
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def teardown(self) -> None:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
|
|
run_state_machine_as_test(Servers, settings=LIFECYCLE_SETTINGS)
|
|
|
|
|
|
@pytest.mark.covers("other.mcp.health.restricted_keys_intersect_grants_in_both_modes")
|
|
def test_health_intersects_route_restricted_key_grants_in_both_management_modes(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
for mode in ("restricted", "view_all"):
|
|
config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
config["general_settings"]["user_mcp_management_mode"] = mode
|
|
path = tmp_path / f"health-{mode}.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
with (
|
|
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
|
|
mcp_peer() as peer,
|
|
candidate.scenario() as scenario,
|
|
):
|
|
first = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
|
|
second = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
|
|
control = scenario.key(object_permission={"mcp_servers": [first]})
|
|
names = tool_names(candidate, control, first)
|
|
healthy = call_tool(candidate, control, first, names["add"], {"a": 3, "b": 5})
|
|
assert healthy.status_code == 200 and healthy.json()["content"][0]["text"] == "8", healthy.text
|
|
for grants in ([first], [second], []):
|
|
key = scenario.key(
|
|
allowed_routes=["/v1/mcp/server", "/v1/mcp/server/health"],
|
|
object_permission={"mcp_servers": grants},
|
|
)
|
|
listed = candidate.request("GET", "/v1/mcp/server", key=key)
|
|
assert listed.status_code == 200, listed.text
|
|
assert {row["server_id"] for row in listed.json()} == set(grants), listed.text
|
|
for requested in (None, [second], [first, second]):
|
|
response = candidate.client.get(
|
|
"/v1/mcp/server/health",
|
|
headers={"Authorization": f"Bearer {key}"},
|
|
params=[] if requested is None else [("server_ids", identity) for identity in requested],
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
expected = set(grants) if requested is None else set(grants).intersection(requested)
|
|
assert {row["server_id"] for row in response.json()} == expected, response.text
|
|
assert all(row["status"] == "healthy" for row in response.json())
|
|
|
|
|
|
@pytest.mark.covers("other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic")
|
|
def test_warm_credential_removal_rejects_without_upstream_traffic(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
identity = register_mcp(
|
|
scenario,
|
|
peer,
|
|
"credentials" + uuid.uuid4().hex,
|
|
auth_type="bearer_token",
|
|
static_headers={"Authorization": "Bearer synthetic-upstream-credential"},
|
|
)
|
|
key = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
names = tool_names(gateway, key, identity)
|
|
warm = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
|
assert warm.status_code == 200 and warm.json()["content"][0]["text"] == "8", warm.text
|
|
calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
|
|
assert len(calls) == 1
|
|
assert calls[0]["headers"][b"authorization"] == b"Bearer synthetic-upstream-credential"
|
|
removed = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "static_headers": {}})
|
|
assert removed.status_code == 202, removed.text
|
|
stored = gateway.request("GET", f"/v1/mcp/server/{identity}")
|
|
assert stored.status_code == 200, stored.text
|
|
assert stored.json()["auth_type"] == "bearer_token"
|
|
assert not stored.json().get("static_headers"), stored.text
|
|
peer.drain()
|
|
for operation in ("list", "call"):
|
|
rejected = (
|
|
gateway.client.get(
|
|
"/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": key}
|
|
)
|
|
if operation == "list"
|
|
else call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
|
)
|
|
assert rejected.status_code == 500, rejected.text
|
|
if operation == "list":
|
|
assert rejected.json()["detail"]["error"] == "internal", rejected.text
|
|
assert "Failed to list tools from server" in rejected.json()["detail"]["message"], rejected.text
|
|
else:
|
|
assert "requires a usable upstream credential" in rejected.text, rejected.text
|
|
assert peer.drain() == (), "missing static credential escaped to upstream"
|
|
changed = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"auth_type": "oauth2_token_exchange",
|
|
"token_exchange_endpoint": peer.url + "/token",
|
|
"credentials": {"client_id": "synthetic-client"},
|
|
},
|
|
)
|
|
assert changed.status_code == 202, changed.text
|
|
peer.drain()
|
|
rejected_subject = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
|
assert rejected_subject.status_code == 401, rejected_subject.text
|
|
assert peer.drain() == (), "virtual key cannot supply an OBO subject token"
|
|
control_id = register_mcp(scenario, peer, "control" + uuid.uuid4().hex, auth_type="none")
|
|
control_key = scenario.key(object_permission={"mcp_servers": [control_id]})
|
|
control_names = tool_names(gateway, control_key, control_id)
|
|
control = call_tool(gateway, control_key, control_id, control_names["multiply"], {"a": 3, "b": 5})
|
|
assert control.status_code == 200 and control.json()["content"][0]["text"] == "15", control.text
|
|
|
|
|
|
@pytest.mark.parametrize("authenticated", (False, True), ids=("anonymous", "bearer"))
|
|
@pytest.mark.covers("other.mcp.permissions.same_url_servers_enforce_discovery_and_execution")
|
|
def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
|
|
gateway: Gateway, authenticated: bool
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
aliases: Final = tuple("scope" + uuid.uuid4().hex for _ in range(2))
|
|
servers: Final = tuple(
|
|
register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
auth_type="bearer_token" if authenticated else "none",
|
|
static_headers={
|
|
"X-Integration-Server": alias,
|
|
**({"Authorization": f"Bearer synthetic-{alias}"} if authenticated else {}),
|
|
},
|
|
)
|
|
for alias in aliases
|
|
)
|
|
for virtual in (False, True):
|
|
keys: Final = tuple(
|
|
scenario.key(object_permission={"mcp_servers": [server], "mcp_tool_search_enabled": virtual})
|
|
for server in servers
|
|
)
|
|
for server, alias, key in zip(servers, aliases, keys):
|
|
catalog: Final = gateway.request("GET", "/mcp-rest/tools/list", key=key)
|
|
assert catalog.status_code == 200, catalog.text
|
|
if virtual:
|
|
assert {tool["name"] for tool in catalog.json()["tools"]} == {
|
|
"mcp_tool_search",
|
|
"mcp_tool_call",
|
|
"agent_search",
|
|
"skill_search",
|
|
}, catalog.text
|
|
search: Final = gateway.request(
|
|
"POST",
|
|
"/mcp-rest/tools/call",
|
|
{"name": "mcp_tool_search", "arguments": {"query": "add", "top_k": 10}},
|
|
key=key,
|
|
)
|
|
assert search.status_code == 200 and search.json()["isError"] is False, search.text
|
|
assert [tool["name"] for tool in json.loads(search.json()["content"][0]["text"])] == [
|
|
f"{alias}-add"
|
|
], search.text
|
|
else:
|
|
assert {tool["mcp_info"]["server_id"] for tool in catalog.json()["tools"]} == {server}
|
|
assert {tool["name"] for tool in catalog.json()["tools"]} == {"add", "multiply", "fail"}
|
|
for server_index, caller_index in ((0, 0), (1, 0), (1, 1)):
|
|
peer.drain()
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/mcp-rest/tools/call",
|
|
{
|
|
"name": "mcp_tool_call" if virtual else "add",
|
|
**({} if virtual else {"server_id": servers[server_index]}),
|
|
"arguments": (
|
|
{"tool_name": f"{aliases[server_index]}-add", "arguments": {"a": 3, "b": 5}}
|
|
if virtual
|
|
else {"a": 3, "b": 5}
|
|
),
|
|
},
|
|
key=keys[caller_index],
|
|
)
|
|
observed: Final = peer.drain()
|
|
if server_index != caller_index:
|
|
assert response.status_code == 403 and "not allowed" in response.text, response.text
|
|
assert observed == (), "forbidden server reached the upstream"
|
|
continue
|
|
assert response.status_code == 200 and response.json()["isError"] is False, response.text
|
|
assert response.json()["content"][0]["text"] == "8", response.text
|
|
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
|
|
assert len(calls) == 1
|
|
assert calls[0]["headers"][b"x-integration-server"] == aliases[server_index].encode()
|
|
expected_auth: Final = f"Bearer synthetic-{aliases[server_index]}".encode() if authenticated else None
|
|
assert all(item["headers"].get(b"authorization") == expected_auth for item in observed)
|