mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
739 lines
36 KiB
Python
739 lines
36 KiB
Python
import itertools
|
|
import uuid
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import JSON_OBJECT, Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.mcp import (
|
|
McpCaller,
|
|
McpPeer,
|
|
call_tool,
|
|
delete_mcp,
|
|
forget_mcp,
|
|
listed_tools,
|
|
mcp_peer,
|
|
openapi_peer,
|
|
register_mcp,
|
|
tool_calls,
|
|
tool_names,
|
|
)
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import BaseModel, JsonValue
|
|
|
|
from litellm.models.user import LiteLLM_UserTable
|
|
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
ADD: Final = {"a": 4, "b": 5}
|
|
|
|
|
|
def _dashboard_ui_session_token(user_id: str) -> str:
|
|
user: Final = LiteLLM_UserTable(user_id=user_id, user_role="internal_user", models=[])
|
|
return ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user)
|
|
|
|
|
|
def _servers(gateway: Gateway, key: str | None = None) -> dict[str, dict[str, object]]:
|
|
response: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": key or gateway.key})
|
|
assert response.status_code == 200, response.text
|
|
return {server["server_id"]: server for server in response.json()}
|
|
|
|
|
|
def test_non_admin_key_cannot_create_edit_or_delete_servers(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
plain: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
headers: Final = {"x-litellm-api-key": plain}
|
|
created: Final = gateway.client.post(
|
|
"/v1/mcp/server",
|
|
json={"server_name": alias + "x", "alias": alias + "x", **peer.registration()},
|
|
headers=headers,
|
|
)
|
|
assert created.status_code == 403, created.text
|
|
edited: Final = gateway.client.put(
|
|
"/v1/mcp/server", json={"server_id": identity, "server_name": "hijacked"}, headers=headers
|
|
)
|
|
assert edited.status_code == 403, edited.text
|
|
deleted: Final = gateway.client.delete(f"/v1/mcp/server/{identity}", headers=headers)
|
|
assert deleted.status_code == 403, deleted.text
|
|
assert _servers(gateway)[identity]["server_name"] == alias
|
|
assert call_tool(gateway, plain, identity, tool_names(gateway, plain, identity)["add"], ADD).status_code == 200
|
|
|
|
|
|
def test_secrets_never_appear_in_server_listing_or_detail(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
secret: Final = "shh-" + uuid.uuid4().hex
|
|
header_secret: Final = "hdr-" + uuid.uuid4().hex
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
auth_type="bearer_token",
|
|
credentials={"auth_value": secret},
|
|
static_headers={"X-Integration-Secret": header_secret},
|
|
)
|
|
viewer: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
for key in (gateway.key, viewer):
|
|
listing: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": key})
|
|
detail: Final = gateway.client.get(f"/v1/mcp/server/{identity}", headers={"x-litellm-api-key": key})
|
|
assert listing.status_code == 200 and detail.status_code == 200, (listing.text, detail.text)
|
|
assert secret not in listing.text + detail.text, key == gateway.key
|
|
viewed: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": viewer})
|
|
assert header_secret not in viewed.text, viewed.text
|
|
peer.drain()
|
|
assert (
|
|
call_tool(gateway, viewer, identity, tool_names(gateway, viewer, identity)["add"], ADD).status_code == 200
|
|
)
|
|
sent: Final = tool_calls(peer.drain())
|
|
assert [call["headers"][b"authorization"] for call in sent] == [f"Bearer {secret}".encode()]
|
|
assert [call["headers"][b"x-integration-secret"] for call in sent] == [header_secret.encode()]
|
|
|
|
|
|
def test_edit_url_moves_calls_to_the_new_peer_without_touching_grants(gateway: Gateway) -> None:
|
|
with mcp_peer() as first, mcp_peer() as second, gateway.scenario() as scenario:
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, first, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
name: Final = tool_names(gateway, key, identity)["add"]
|
|
assert call_tool(gateway, key, identity, name, ADD).status_code == 200
|
|
assert len(tool_calls(first.drain())) == 1
|
|
moved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "url": second.url})
|
|
assert moved.status_code == 202, moved.text
|
|
second.drain()
|
|
response: Final = eventually(
|
|
lambda: call_tool(gateway, key, identity, name, ADD),
|
|
lambda value: value.status_code == 200 and len(tool_calls(second.drain())) == 1,
|
|
)
|
|
assert response.json()["content"][0]["text"] == "9", response.text
|
|
assert tool_calls(first.drain()) == ()
|
|
|
|
|
|
def test_delete_removes_listing_calls_and_database_row(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
name: Final = tool_names(gateway, key, identity)["add"]
|
|
delete_mcp(gateway, identity)
|
|
assert identity not in _servers(gateway)
|
|
listing: Final = gateway.client.get(
|
|
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
|
|
)
|
|
assert listing.status_code >= 400 or listing.json() == [], listing.text
|
|
peer.drain()
|
|
response: Final = call_tool(gateway, key, identity, name, ADD)
|
|
assert response.status_code >= 400, response.text
|
|
assert tool_calls(peer.drain()) == ()
|
|
caller: Final = McpCaller(gateway, key, "server_mcp", alias)
|
|
assert caller.list_tools().tools == (), caller.list_tools().raw
|
|
|
|
|
|
def test_duplicate_alias_is_rejected_so_tool_prefixes_cannot_collide(gateway: Gateway) -> None:
|
|
import concurrent.futures
|
|
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
|
|
duplicate: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, **peer.registration()}
|
|
)
|
|
assert duplicate.status_code == 400, duplicate.text
|
|
assert alias in duplicate.json()["detail"]["error"], duplicate.text
|
|
|
|
same_alias: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": alias + "other", "alias": alias, **peer.registration()}
|
|
)
|
|
assert same_alias.status_code == 400, same_alias.text
|
|
assert alias in same_alias.json()["detail"]["error"], same_alias.text
|
|
|
|
case_variant: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": alias.upper(), "alias": alias.upper(), **peer.registration()}
|
|
)
|
|
assert case_variant.status_code == 400, case_variant.text
|
|
|
|
same_name_no_alias: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": alias, **peer.registration()}
|
|
)
|
|
assert same_name_no_alias.status_code == 400, same_name_no_alias.text
|
|
|
|
second_alias: Final = alias + "2"
|
|
second_identity: Final = register_mcp(scenario, peer, second_alias)
|
|
colliding_rename: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": second_identity, "alias": alias}
|
|
)
|
|
assert colliding_rename.status_code == 400, colliding_rename.text
|
|
|
|
cleared_alias: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": second_identity, "alias": None})
|
|
assert cleared_alias.status_code == 202, cleared_alias.text
|
|
|
|
name: Final = tool_names(gateway, key, identity)["add"]
|
|
response: Final = call_tool(gateway, key, identity, name, ADD)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["content"][0]["text"] == "9", response.text
|
|
|
|
racing_alias: Final = "race" + uuid.uuid4().hex[:8]
|
|
|
|
def try_register() -> int:
|
|
response: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": racing_alias, "alias": racing_alias, **peer.registration()}
|
|
)
|
|
return response.status_code
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
|
|
statuses: Final = tuple(pool.map(lambda _i: try_register(), range(8)))
|
|
|
|
assert statuses.count(201) == 1, statuses
|
|
assert statuses.count(400) == 7, statuses
|
|
winner: Final = next(
|
|
server["server_id"] for server in _servers(gateway).values() if server["alias"] == racing_alias
|
|
)
|
|
scenario.cleanups.callback(forget_mcp, gateway, winner)
|
|
|
|
|
|
def _openapi_server_lists_and_calls_only_its_own_tools(
|
|
gateway: Gateway, key: str, peer: McpPeer, identity: str
|
|
) -> None:
|
|
listed: Final = set(listed_tools(gateway, key, identity))
|
|
assert listed == {"getpet", "createpet"}, (identity, listed)
|
|
peer.drain()
|
|
called: Final = call_tool(gateway, key, identity, "getpet", {"petId": "7"})
|
|
assert called.status_code == 200, called.text
|
|
assert [(item["method"], item["path"]) for item in peer.drain()] == [("GET", "/pets/7")], identity
|
|
|
|
|
|
def test_openapi_listing_is_scoped_to_the_exact_alias_when_aliases_overlap(gateway: Gateway) -> None:
|
|
with openapi_peer() as short, openapi_peer() as long, gateway.scenario() as scenario:
|
|
stem: Final = "pet" + uuid.uuid4().hex[:8]
|
|
servers: Final = tuple(
|
|
(peer, alias, register_mcp(scenario, peer, alias))
|
|
for peer, alias in ((short, stem), (long, stem + "store"))
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity for _, _, identity in servers]})
|
|
for peer, _, identity in servers:
|
|
_openapi_server_lists_and_calls_only_its_own_tools(gateway, key, peer, identity)
|
|
aggregate: Final = McpCaller(gateway, key, "mcp").list_tools()
|
|
assert aggregate.ok, aggregate.raw
|
|
assert sorted(aggregate.tools) == sorted(
|
|
f"{prefix}-{tool}" for prefix, tool in itertools.product((stem, stem + "store"), ("getpet", "createpet"))
|
|
), aggregate.tools
|
|
assert all(peer.drain() == () for peer, _, _ in servers), "listing must not reach any OpenAPI upstream"
|
|
|
|
|
|
def test_config_declared_openapi_server_with_a_space_in_its_name_lists_its_tools(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
with openapi_peer() as peer:
|
|
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
|
name: Final = "pet store " + uuid.uuid4().hex[:8]
|
|
config["mcp_servers"] = {name: peer.registration()}
|
|
path: Final = tmp_path / "openapi-space.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
|
|
identity: Final = next(i for i, s in _servers(candidate).items() if s["server_name"] == name)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
_openapi_server_lists_and_calls_only_its_own_tools(candidate, key, peer, identity)
|
|
aggregate: Final = McpCaller(candidate, key, "mcp").list_tools()
|
|
assert aggregate.ok, aggregate.raw
|
|
prefix: Final = name.replace(" ", "_")
|
|
assert sorted(aggregate.tools) == [f"{prefix}-createpet", f"{prefix}-getpet"], aggregate.tools
|
|
|
|
|
|
def test_invalid_registrations_are_rejected(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
register_mcp(scenario, peer, alias)
|
|
no_url: Final = gateway.request("POST", "/v1/mcp/server", {"server_name": alias + "b", "transport": "http"})
|
|
assert no_url.status_code in (400, 422), no_url.text
|
|
bad_command: Final = gateway.request(
|
|
"POST",
|
|
"/v1/mcp/server",
|
|
{"server_name": alias + "c", "transport": "stdio", "command": "/bin/sh", "args": ["-c", "true"]},
|
|
)
|
|
assert bad_command.status_code in (400, 422), bad_command.text
|
|
hyphenless: Final = gateway.request(
|
|
"POST", "/v1/mcp/server", {"server_name": "bad name!", **peer.registration()}
|
|
)
|
|
assert hyphenless.status_code in (400, 422), hyphenless.text
|
|
assert len([s for s in _servers(gateway).values() if str(s["server_name"]).startswith(alias)]) == 1
|
|
|
|
|
|
def test_access_group_membership_follows_edits(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
group: Final = "grp" + uuid.uuid4().hex[:8]
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, mcp_access_groups=[group])
|
|
key: Final = scenario.key(object_permission={"mcp_access_groups": [group]})
|
|
groups: Final = gateway.client.get("/v1/mcp/access_groups", headers={"x-litellm-api-key": gateway.key})
|
|
assert groups.status_code == 200 and group in groups.text, groups.text
|
|
assert "add" in tool_names(gateway, key, identity)
|
|
removed: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "mcp_access_groups": []})
|
|
assert removed.status_code == 202, removed.text
|
|
eventually(
|
|
lambda: gateway.client.get(
|
|
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
|
|
),
|
|
lambda value: value.status_code >= 400 or value.json() == [],
|
|
)
|
|
peer.drain()
|
|
denied: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
|
|
assert denied.status_code >= 400, denied.text
|
|
assert tool_calls(peer.drain()) == ()
|
|
|
|
|
|
def test_peer_worker_observes_create_edit_and_delete_without_restart(gateway: Gateway, tmp_path: Path) -> None:
|
|
config: Final = tmp_path / "peer.yaml"
|
|
config.write_text(yaml.safe_dump({
|
|
"model_list": [],
|
|
"general_settings": {"master_key": gateway.key, "store_model_in_db": True, "proxy_config_reload_interval_seconds": 3600},
|
|
}))
|
|
with (
|
|
owned_proxy(gateway, tmp_path, {}, config=config, database_setup=()) as peer,
|
|
mcp_peer() as first,
|
|
mcp_peer() as second,
|
|
gateway.scenario() as scenario,
|
|
):
|
|
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, first, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
listing: Final = peer.client.get(
|
|
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
|
|
)
|
|
assert listing.status_code == 200 and listing.json() != [], listing.text
|
|
names: Final = tool_names(peer, key, identity)
|
|
assert call_tool(peer, key, identity, names["add"], ADD).status_code == 200
|
|
assert len(tool_calls(first.drain())) == 1
|
|
moved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "url": second.url})
|
|
assert moved.status_code == 202, moved.text
|
|
called: Final = call_tool(peer, key, identity, names["add"], ADD)
|
|
assert called.status_code == 200, called.text
|
|
assert tool_calls(first.drain()) == ()
|
|
assert len(tool_calls(second.drain())) == 1
|
|
delete_mcp(gateway, identity)
|
|
deleted: Final = peer.client.get(
|
|
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
|
|
)
|
|
assert deleted.status_code in (403, 404) or (deleted.status_code == 200 and deleted.json() == []), deleted.text
|
|
assert call_tool(peer, key, identity, names["add"], ADD).status_code in (403, 404)
|
|
assert tool_calls(second.drain()) == ()
|
|
|
|
|
|
def test_config_declared_server_behaves_like_database_server_but_is_read_only(gateway: Gateway, tmp_path: Path) -> None:
|
|
with mcp_peer() as declared_peer, mcp_peer() as database_peer:
|
|
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
|
declared: Final = "declared" + uuid.uuid4().hex[:8]
|
|
config["mcp_servers"] = {declared: {**declared_peer.registration(), "static_headers": {"X-From": "config"}}}
|
|
path: Final = tmp_path / "mcp.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
|
|
servers: Final = _servers(candidate)
|
|
declared_id: Final = next(identity for identity, s in servers.items() if s["server_name"] == declared)
|
|
created: Final = register_mcp(scenario, database_peer, "database" + uuid.uuid4().hex[:8])
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [declared_id, created]})
|
|
declared_names: Final = tool_names(candidate, key, declared_id)
|
|
assert set(declared_names) == set(tool_names(candidate, key, created)) == {"add", "multiply", "fail"}
|
|
declared_peer.drain()
|
|
response: Final = call_tool(candidate, key, declared_id, declared_names["add"], ADD)
|
|
assert response.status_code == 200 and response.json()["content"][0]["text"] == "9", response.text
|
|
sent: Final = tool_calls(declared_peer.drain())
|
|
assert [call["headers"][b"x-from"] for call in sent] == [b"config"]
|
|
edited: Final = candidate.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": declared_id, "url": database_peer.url}
|
|
)
|
|
assert edited.status_code >= 400, edited.text
|
|
deleted: Final = candidate.request("DELETE", f"/v1/mcp/server/{declared_id}")
|
|
assert deleted.status_code >= 400, deleted.text
|
|
assert declared_id in _servers(candidate)
|
|
assert call_tool(candidate, key, declared_id, declared_names["add"], ADD).status_code == 200
|
|
assert len(tool_calls(declared_peer.drain())) == 1 and tool_calls(database_peer.drain()) == ()
|
|
|
|
|
|
def test_team_granted_database_server_detail_is_available_to_team_key(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit3974_team_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
team_id: Final = scenario.team(object_permission={"mcp_servers": [server_id]})
|
|
key: Final = scenario.key(team_id=team_id)
|
|
|
|
response: Final = gateway.request("GET", f"/v1/mcp/server/{server_id}", key=key)
|
|
|
|
assert response.status_code == 200, f"Team-granted server detail access should succeed: {response.text}"
|
|
assert response.json()["server_id"] == server_id, response.text
|
|
assert response.json()["alias"] == alias, response.text
|
|
|
|
|
|
def test_ui_session_lists_and_fetches_team_granted_config_server(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit3974_config_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = "lit3974-" + uuid.uuid4().hex[:12]
|
|
team_id: Final = scenario.team(object_permission={"mcp_servers": [server_id]})
|
|
user_id: Final = scenario.user(user_role="internal_user", teams=[team_id])
|
|
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
|
config["mcp_servers"] = {alias: {**peer.registration(), "alias": alias, "server_id": server_id}}
|
|
config_path: Final = tmp_path / "lit3974-mcp.yaml"
|
|
config_path.write_text(yaml.safe_dump(config))
|
|
|
|
with owned_proxy(gateway, tmp_path, {}, config=config_path) as candidate:
|
|
token: Final = _dashboard_ui_session_token(user_id)
|
|
headers: Final = {"Authorization": f"Bearer {token}"}
|
|
listed: Final = candidate.client.get("/v1/mcp/server", headers=headers)
|
|
assert listed.status_code == 200, listed.text
|
|
assert [server["server_id"] for server in listed.json()] == [server_id], listed.text
|
|
|
|
detail: Final = candidate.client.get(f"/v1/mcp/server/{server_id}", headers=headers)
|
|
|
|
assert detail.status_code == 200, f"Team-granted server detail access should succeed: {detail.text}"
|
|
assert detail.json()["server_id"] == server_id, detail.text
|
|
assert detail.json()["alias"] == alias, detail.text
|
|
|
|
|
|
@pytest.mark.parametrize("explicit_transport", [False, True])
|
|
def test_modern_sse_registration_rejected_without_saving(gateway: Gateway, explicit_transport: bool) -> None:
|
|
identity: Final = str(uuid.uuid4())
|
|
with mcp_peer() as peer:
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"server_name": "invalid" + uuid.uuid4().hex[:8],
|
|
"url": peer.url,
|
|
"mcp_info": {"protocol_version": "2026-07-28"},
|
|
**({"transport": "sse"} if explicit_transport else {}),
|
|
},
|
|
)
|
|
try:
|
|
assert response.status_code == 422, response.text
|
|
assert "Modern MCP requires HTTP or stdio" in response.text
|
|
assert identity not in _servers(gateway)
|
|
assert peer.drain() == (), "Rejected configuration reached upstream"
|
|
finally:
|
|
if response.status_code == 201:
|
|
delete_mcp(gateway, identity)
|
|
|
|
|
|
def test_protocol_transport_updates_validate_effective_configuration(gateway: Gateway) -> None:
|
|
from integration._support.database import read_rows
|
|
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
identity: Final = register_mcp(scenario, peer, "protocol" + uuid.uuid4().hex[:8])
|
|
modern: Final = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"mcp_info": {"protocol_version": "2026-07-28"},
|
|
},
|
|
)
|
|
assert modern.status_code == 202, modern.text
|
|
assert modern.json()["transport"] == "http"
|
|
changed: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "renamed"})
|
|
assert changed.status_code == 202, changed.text
|
|
assert changed.json()["mcp_info"]["protocol_version"] == "2026-07-28"
|
|
snapshot: Final = read_rows(
|
|
'SELECT transport, mcp_info, updated_at::text FROM "LiteLLM_MCPServerTable" WHERE server_id=%s', (identity,)
|
|
)
|
|
peer.drain()
|
|
rejected: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "transport": "sse", "url": peer.url}
|
|
)
|
|
assert rejected.status_code == 400, rejected.text
|
|
assert "Modern MCP requires HTTP or stdio" in rejected.text
|
|
assert (
|
|
read_rows(
|
|
'SELECT transport, mcp_info, updated_at::text FROM "LiteLLM_MCPServerTable" WHERE server_id=%s',
|
|
(identity,),
|
|
)
|
|
== snapshot
|
|
)
|
|
assert tool_calls(peer.drain()) == ()
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
assert call_tool(gateway, key, identity, "add", ADD).status_code == 200
|
|
|
|
legacy: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "transport": "sse", "url": peer.url, "mcp_info": {}}
|
|
)
|
|
assert legacy.status_code == 202, legacy.text
|
|
legacy_snapshot: Final = read_rows(
|
|
'SELECT transport, mcp_info, updated_at::text FROM "LiteLLM_MCPServerTable" WHERE server_id=%s', (identity,)
|
|
)
|
|
rejected_protocol: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "mcp_info": {"protocol_version": "2026-07-28"}}
|
|
)
|
|
assert rejected_protocol.status_code == 400, rejected_protocol.text
|
|
assert (
|
|
read_rows(
|
|
'SELECT transport, mcp_info, updated_at::text FROM "LiteLLM_MCPServerTable" WHERE server_id=%s',
|
|
(identity,),
|
|
)
|
|
== legacy_snapshot
|
|
)
|
|
repaired: Final = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"transport": "http",
|
|
"url": peer.url,
|
|
"mcp_info": {"protocol_version": "2026-07-28"},
|
|
},
|
|
)
|
|
assert repaired.status_code == 202, repaired.text
|
|
assert call_tool(gateway, key, identity, "add", ADD).status_code == 200
|
|
|
|
|
|
@pytest.mark.parametrize("rename", [False, True])
|
|
def test_concurrent_protocol_transport_edits_cannot_save_incompatible_configuration(
|
|
gateway: Gateway, peer: Gateway, rename: bool
|
|
) -> None:
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import psycopg
|
|
from integration._support.database import read_rows
|
|
|
|
with mcp_peer() as upstream, gateway.scenario() as scenario:
|
|
identity: Final = register_mcp(scenario, upstream, "race" + uuid.uuid4().hex[:8])
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
# Hold the row so both workers read the old configuration before either can write.
|
|
with psycopg.connect(os.environ["DATABASE_URL"]) as blocker:
|
|
blocker.execute(
|
|
'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id=%s FOR UPDATE', (identity,)
|
|
)
|
|
protocol = pool.submit(
|
|
gateway.request,
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"mcp_info": {"protocol_version": "2026-07-28"},
|
|
**({"alias": "renamed" + uuid.uuid4().hex[:8]} if rename else {}),
|
|
},
|
|
)
|
|
transport = pool.submit(
|
|
peer.request,
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"transport": "sse",
|
|
"url": upstream.url,
|
|
},
|
|
)
|
|
eventually(
|
|
lambda: read_rows(
|
|
"SELECT pid FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock' AND query LIKE %s",
|
|
("%LiteLLM_MCPServerTable%",),
|
|
),
|
|
lambda rows: len(rows) >= 2,
|
|
seconds=3,
|
|
)
|
|
responses: Final = [protocol.result(), transport.result()]
|
|
assert sorted(response.status_code for response in responses) == [202, 400], [r.text for r in responses]
|
|
saved: Final = read_rows(
|
|
'SELECT transport, mcp_info FROM "LiteLLM_MCPServerTable" WHERE server_id=%s', (identity,)
|
|
)[0]
|
|
assert saved["transport"] == "http" or saved["mcp_info"] != {"protocol_version": "2026-07-28"}
|
|
repaired: Final = gateway.request(
|
|
"PUT",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_id": identity,
|
|
"transport": "http",
|
|
"url": upstream.url,
|
|
"mcp_info": {"protocol_version": "2026-07-28"},
|
|
},
|
|
)
|
|
assert repaired.status_code == 202, repaired.text
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
assert call_tool(gateway, key, identity, "add", ADD).status_code == 200
|
|
|
|
|
|
class Submission(BaseModel):
|
|
server_id: str
|
|
approval_status: str
|
|
submitted_by: str | None = None
|
|
review_notes: str | None = None
|
|
|
|
|
|
class Submissions(BaseModel):
|
|
items: list[Submission]
|
|
|
|
|
|
class ListedServerInfo(BaseModel):
|
|
server_id: str
|
|
|
|
|
|
class ListedTool(BaseModel):
|
|
name: str
|
|
mcp_info: ListedServerInfo
|
|
|
|
|
|
class ToolListing(BaseModel):
|
|
tools: list[ListedTool]
|
|
|
|
|
|
def _submission_row(identity: str) -> list[dict[str, JsonValue]]:
|
|
return read_rows(
|
|
'SELECT approval_status, submitted_by, review_notes FROM "LiteLLM_MCPServerTable" WHERE server_id = %s',
|
|
(identity,),
|
|
)
|
|
|
|
|
|
def _submission(gateway: Gateway, identity: str) -> Submission:
|
|
response: Final = gateway.client.get("/v1/mcp/server/submissions", headers={"x-litellm-api-key": gateway.key})
|
|
assert response.status_code == 200, response.text
|
|
matches: Final = [
|
|
item for item in Submissions.model_validate_json(response.content).items if item.server_id == identity
|
|
]
|
|
assert len(matches) == 1, response.text
|
|
return matches[0]
|
|
|
|
|
|
def _served(gateway: Gateway, key: str, identity: str) -> bool:
|
|
response: Final = gateway.client.post(
|
|
"/mcp-rest/tools/call",
|
|
headers={"x-litellm-api-key": key},
|
|
json={"server_id": identity, "name": "add", "arguments": ADD},
|
|
)
|
|
return response.status_code == 200 and JSON_OBJECT.validate_json(response.content)["content"] == [
|
|
{"type": "text", "text": "9", "annotations": None, "_meta": None}
|
|
]
|
|
|
|
|
|
def _not_found(gateway: Gateway, key: str, identity: str) -> bool:
|
|
response: Final = gateway.client.post(
|
|
"/mcp-rest/tools/call",
|
|
headers={"x-litellm-api-key": key},
|
|
json={"server_id": identity, "name": "add", "arguments": ADD},
|
|
)
|
|
return response.status_code == 404 and JSON_OBJECT.validate_json(response.content) == {
|
|
"detail": {"error": "server_not_found", "message": f"MCP server '{identity}' was not found"}
|
|
}
|
|
|
|
|
|
def _refusal(gateway: Gateway, key: str, identity: str) -> tuple[int, dict[str, JsonValue]]:
|
|
response: Final = gateway.client.post(
|
|
"/mcp-rest/tools/call",
|
|
headers={"x-litellm-api-key": key},
|
|
json={"server_id": identity, "name": "add", "arguments": ADD},
|
|
)
|
|
return response.status_code, JSON_OBJECT.validate_json(response.content)
|
|
|
|
|
|
def _called_tools(upstream: McpPeer) -> list[tuple[object, object]]:
|
|
return [
|
|
(call["body"]["params"]["name"], call["body"]["params"]["arguments"])
|
|
for call in tool_calls(upstream.drain())
|
|
if isinstance(call["body"], dict) and isinstance(call["body"]["params"], dict)
|
|
]
|
|
|
|
|
|
def test_team_submission_is_inert_until_approved_serves_both_workers_and_stops_after_reject(
|
|
gateway: Gateway, peer: Gateway
|
|
) -> None:
|
|
with mcp_peer() as upstream, mcp_peer() as throwaway_upstream, gateway.scenario() as scenario:
|
|
alias: Final = "submitted" + uuid.uuid4().hex[:8]
|
|
user: Final = scenario.user()
|
|
team: Final = scenario.team()
|
|
member: Final = scenario.key(team_id=team, user_id=user)
|
|
outsider: Final = scenario.key(team_id=scenario.team())
|
|
registered: Final = gateway.client.post(
|
|
"/v1/mcp/server/register",
|
|
headers={"x-litellm-api-key": member},
|
|
json={"server_name": alias, "alias": alias, **upstream.registration()},
|
|
)
|
|
assert registered.status_code == 201, registered.text
|
|
submitted: Final = Submission.model_validate_json(registered.content)
|
|
identity: Final = submitted.server_id
|
|
scenario.cleanups.callback(forget_mcp, gateway, identity)
|
|
assert (submitted.approval_status, submitted.submitted_by) == ("pending_review", user), registered.text
|
|
assert _submission_row(identity) == [
|
|
{"approval_status": "pending_review", "submitted_by": user, "review_notes": None}
|
|
]
|
|
assert _submission(gateway, identity).approval_status == "pending_review"
|
|
throwaway_alias: Final = "reload" + uuid.uuid4().hex[:8]
|
|
throwaway_registered: Final = gateway.client.post(
|
|
"/v1/mcp/server/register",
|
|
headers={"x-litellm-api-key": member},
|
|
json={"server_name": throwaway_alias, "alias": throwaway_alias, **throwaway_upstream.registration()},
|
|
)
|
|
assert throwaway_registered.status_code == 201, throwaway_registered.text
|
|
throwaway_submission: Final = Submission.model_validate_json(throwaway_registered.content)
|
|
throwaway_identity: Final = throwaway_submission.server_id
|
|
scenario.cleanups.callback(forget_mcp, gateway, throwaway_identity)
|
|
assert (throwaway_submission.approval_status, throwaway_submission.submitted_by) == ("pending_review", user), (
|
|
throwaway_registered.text
|
|
)
|
|
granted: Final = gateway.request(
|
|
"POST",
|
|
"/team/update",
|
|
{"team_id": team, "object_permission": {"mcp_servers": [identity, throwaway_identity]}},
|
|
)
|
|
assert granted.status_code == 200, granted.text
|
|
throwaway_approved: Final = gateway.request("PUT", f"/v1/mcp/server/{throwaway_identity}/approve")
|
|
assert throwaway_approved.status_code == 200, throwaway_approved.text
|
|
assert Submission.model_validate_json(throwaway_approved.content).approval_status == "active", (
|
|
throwaway_approved.text
|
|
)
|
|
assert _served(gateway, member, throwaway_identity), "the approving worker did not serve the throwaway"
|
|
assert eventually(partial(_served, peer, member, throwaway_identity), bool, seconds=70)
|
|
throwaway_upstream.drain()
|
|
upstream.drain()
|
|
not_found: Final = (
|
|
404,
|
|
{"detail": {"error": "server_not_found", "message": f"MCP server '{identity}' was not found"}},
|
|
)
|
|
for worker, key in itertools.product((gateway, peer), (member, outsider)):
|
|
pending_list = worker.client.get("/mcp-rest/tools/list", headers={"x-litellm-api-key": key})
|
|
assert pending_list.status_code == 200, pending_list.text
|
|
listed = ToolListing.model_validate_json(pending_list.content).tools
|
|
assert [tool.name for tool in listed if tool.mcp_info.server_id == identity] == [], pending_list.text
|
|
assert _refusal(worker, key, identity) == not_found, worker.client.base_url
|
|
assert _called_tools(upstream) == [], "a pending submission reached the peer"
|
|
|
|
approved: Final = gateway.request("PUT", f"/v1/mcp/server/{identity}/approve")
|
|
assert approved.status_code == 200, approved.text
|
|
assert Submission.model_validate_json(approved.content).approval_status == "active", approved.text
|
|
assert _submission_row(identity)[0]["approval_status"] == "active"
|
|
assert _submission(gateway, identity).approval_status == "active"
|
|
assert _served(gateway, member, identity), "the approving worker did not serve the approved submission"
|
|
assert eventually(partial(_served, peer, member, identity), bool, seconds=70)
|
|
denied: Final = (
|
|
403,
|
|
{"detail": {"error": "access_denied", "message": f"The key is not allowed to access server {identity}"}},
|
|
)
|
|
for worker in (gateway, peer):
|
|
assert _refusal(worker, outsider, identity) == denied, worker.client.base_url
|
|
assert _called_tools(upstream) == [("add", ADD), ("add", ADD)], "expected one served call per worker"
|
|
|
|
notes: Final = "rejected by integration review " + uuid.uuid4().hex[:8]
|
|
rejected: Final = gateway.request("PUT", f"/v1/mcp/server/{identity}/reject", {"review_notes": notes})
|
|
assert rejected.status_code == 200, rejected.text
|
|
rejection: Final = Submission.model_validate_json(rejected.content)
|
|
assert (rejection.approval_status, rejection.review_notes) == ("rejected", notes), rejected.text
|
|
assert _submission_row(identity) == [
|
|
{"approval_status": "rejected", "submitted_by": user, "review_notes": notes}
|
|
]
|
|
submission: Final = _submission(gateway, identity)
|
|
assert (submission.approval_status, submission.review_notes) == ("rejected", notes)
|
|
assert _refusal(gateway, member, identity) == not_found, "the rejecting worker still served the submission"
|
|
assert _called_tools(upstream) == [], "the rejecting worker forwarded a call for the rejected submission"
|
|
assert eventually(partial(_not_found, peer, member, identity), bool, seconds=70)
|
|
upstream.drain()
|
|
for worker in (gateway, peer):
|
|
assert _refusal(worker, member, identity) == not_found, worker.client.base_url
|
|
assert _refusal(worker, outsider, identity) == not_found, worker.client.base_url
|
|
assert _called_tools(upstream) == [], "a rejected submission reached the peer"
|