mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
7efce7f8e0
* test(integration): give chaos and timing tests a harness API and gate raw signals Chaos tests killed and suspended proxy workers with raw os.kill and psutil calls and bounded elapsed time with fixed literals, so they raced the uvicorn worker healthcheck and failed on slow CI boxes OwnedProxy now offers kill_worker and replacing_worker, which return only once a replacement answers readiness with its own pid, and suspended, which is refused unless the proxy was launched with a lenient worker healthcheck. timing.py bounds durations against a measured control or the injected delay. run.py measures the proxy import time once and derives readiness deadlines and the per-test timeout from it, writing boot-calibration.json check_test_quality.py adds TQ010 (raw process signals) and TQ011 (elapsed time asserted under a numeric literal) for tests/integration outside _support, and every call site is migrated so both report zero with no suppressions * test(integration): steer the assistants wire burst onto both workers The mixed-status burst assumed the OS would spread 48 connections across both workers and failed when one worker took them all. Each half is now sent while the other worker is suspended, and the test asserts an exact 24 and 24 split * test(integration): restore product-enforced timing bounds lost in the migration Some TQ011 migrations dropped or loosened a bound whose time is the behaviour under test. The websocket session cap only kept its floor, so a cap closing at 300 s would have passed, and several timeouts were bounded by the injected stall instead of the configured product value assert_enforced_within bounds an elapsed time by the configured value plus explicit slack, with an optional floor. The websocket cap, the MCP validation timeout, the guardrail litellm_params timeout, the prisma command timeouts, the baseline projection lock timeout and the key metadata statement timeout use it. The device code refusal keeps its 15 s margin, and the langfuse retried prompt fetch is now bounded by a cold fetch of a steady prompt as well as by the ignored retry-after * fix(tests): wait for the supervisor to reap killed workers and fix real-kill migrations The lenient worker healthcheck was the whole case budget, so uvicorn's supervisor sat in a ping on the killed worker for up to that long before respawning it. It now uses the graceful stop budget, and a worker kill returns only after the victim is reaped, so pooled connections no longer race the dead worker's socket. Several tests on main picked the multiprocessing resource tracker as their victim and never killed a worker. Now that they kill a real one, presidio counts streams cut by the death, delete_model waits for the config reload that removes a deleted model from a replacement worker, and the wedged coordination Redis test kills before the wedge and waits for every parked publish to fail before resuming Redis. The redis cluster test imports the harness through the tests package so it collects in the unit job, and the per-model fallback rig exposes the owned proxy it now needs. * test(observability): only require S3 deliveries the killed worker could not have buffered A real SIGKILL drops the victim's unflushed batch, so requests it already answered can't be expected to land. The test now requires every request served after the replacement to land, some of the surviving worker's burst share to land, and nothing unsent to land --------- Co-authored-by: Yuneng Jiang <yuneng@berri.ai>
596 lines
32 KiB
Python
596 lines
32 KiB
Python
import re
|
|
import secrets
|
|
import signal
|
|
import textwrap
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import yaml
|
|
from pydantic import TypeAdapter
|
|
from integration._support.client import Gateway, Scenario, gateway_from_environment, object_value
|
|
from integration._support.mcp import (
|
|
INITIALIZE,
|
|
Outcome,
|
|
ScriptedTool,
|
|
_outcome_from_rest,
|
|
_outcome_from_rpc,
|
|
mcp_peer,
|
|
register_mcp,
|
|
scripted_peer,
|
|
text_result,
|
|
tool_calls,
|
|
)
|
|
from integration._support.mcp_grants import create_toolset
|
|
from integration._support.process import owned_proxy, owned_proxy_process
|
|
|
|
from litellm.models.user import LiteLLM_UserTable
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX, ExperimentalUIJWTToken
|
|
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token
|
|
|
|
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 _toolset(scenario: Scenario, server_id: str, tool: str) -> tuple[str, str]:
|
|
name: Final = "lit6029_" + uuid.uuid4().hex[:10]
|
|
return create_toolset(scenario, ((server_id, tool),), toolset_name=name), name
|
|
|
|
|
|
def _toolset_rpc(
|
|
gateway: Gateway, headers: dict[str, str], name: str, method: str, params: dict[str, object]
|
|
) -> Outcome:
|
|
def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
|
|
return gateway.client.post(
|
|
f"/toolset/{name}/mcp",
|
|
json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
|
|
headers={**headers, "Accept": "application/json, text/event-stream"},
|
|
)
|
|
|
|
initialized: Final = _outcome_from_rpc(post("initialize", dict(INITIALIZE)))
|
|
if not initialized.ok:
|
|
return initialized
|
|
return _outcome_from_rpc(post(method, params))
|
|
|
|
|
|
def _listed_toolset_ids(gateway: Gateway, headers: dict[str, str]) -> tuple[str, ...]:
|
|
response: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
|
|
assert response.status_code == 200, response.text
|
|
return tuple(toolset["toolset_id"] for toolset in response.json())
|
|
|
|
|
|
def _assert_team_grants_only(gateway: Gateway, team_id: str, key: str, toolset_id: str) -> None:
|
|
team: Final = object_value(gateway.get("/team/info", {"team_id": team_id})["team_info"])
|
|
assert object_value(team["object_permission"])["mcp_toolsets"] == [toolset_id], team
|
|
key_info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
|
assert key_info.get("object_permission") is None, f"key must carry no grant of its own: {key_info}"
|
|
|
|
|
|
def test_team_granted_toolset_is_listed_and_served_to_a_team_key(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
key: Final = scenario.key(team_id=team_id)
|
|
_assert_team_grants_only(gateway, team_id, key, granted_id)
|
|
headers: Final = {"Authorization": f"Bearer {key}"}
|
|
|
|
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
|
|
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
|
|
assert detail.status_code == 200, detail.text
|
|
assert detail.json()["toolset_name"] == granted_name, detail.text
|
|
withheld_detail: Final = gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=headers)
|
|
assert withheld_detail.status_code == 403, withheld_detail.text
|
|
|
|
listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
|
|
assert listed.ok, listed.raw
|
|
assert listed.tools == (f"{alias}-add",), listed.raw
|
|
peer.drain()
|
|
called: Final = _toolset_rpc(
|
|
gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
|
|
)
|
|
assert called.ok and called.text == "9", called.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
denied: Final = _toolset_rpc(gateway, headers, withheld_name, "tools/list", {})
|
|
assert denied.status == 403, denied.raw
|
|
|
|
|
|
def test_dashboard_session_of_a_team_member_lists_the_team_granted_toolset(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
_toolset(scenario, server_id, "multiply")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
user_id: Final = scenario.user(user_role="internal_user", teams=[team_id])
|
|
user: Final = object_value(gateway.get("/user/info", {"user_id": user_id})["user_info"])
|
|
assert user["teams"] == [team_id], user
|
|
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
|
|
|
|
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
|
|
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
|
|
assert detail.status_code == 200, detail.text
|
|
assert detail.json()["toolset_name"] == granted_name, detail.text
|
|
|
|
|
|
def test_direct_grants_no_grants_and_admin_listing_are_unchanged_by_team_resolution(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
|
|
direct: Final = {"Authorization": f"Bearer {scenario.key(object_permission={'mcp_toolsets': [granted_id]})}"}
|
|
ungranted_team: Final = scenario.team()
|
|
no_grant: Final = {"Authorization": f"Bearer {scenario.key(team_id=ungranted_team)}"}
|
|
admin: Final = {"Authorization": f"Bearer {gateway.key}"}
|
|
|
|
assert _listed_toolset_ids(gateway, direct) == (granted_id,)
|
|
assert _toolset_rpc(gateway, direct, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
|
|
assert _toolset_rpc(gateway, direct, withheld_name, "tools/list", {}).status == 403
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=direct).status_code == 403
|
|
|
|
assert _listed_toolset_ids(gateway, no_grant) == ()
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=no_grant).status_code == 403
|
|
assert _toolset_rpc(gateway, no_grant, granted_name, "tools/list", {}).status == 403
|
|
|
|
assert {granted_id, withheld_id} <= set(_listed_toolset_ids(gateway, admin))
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=admin).status_code == 200
|
|
assert _toolset_rpc(gateway, admin, withheld_name, "tools/list", {}).tools == (f"{alias}-multiply",)
|
|
|
|
|
|
def test_a_key_with_its_own_toolset_grant_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
own_id, own_name = _toolset(scenario, server_id, "add")
|
|
team_only_id, team_only_name = _toolset(scenario, server_id, "multiply")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [own_id, team_only_id]})
|
|
key: Final = scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [own_id]})
|
|
headers: Final = {"Authorization": f"Bearer {key}"}
|
|
|
|
assert _listed_toolset_ids(gateway, headers) == (own_id,)
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{team_only_id}", headers=headers).status_code == 403
|
|
assert _toolset_rpc(gateway, headers, team_only_name, "tools/list", {}).status == 403
|
|
assert _toolset_rpc(gateway, headers, own_name, "tools/list", {}).tools == (f"{alias}-add",)
|
|
|
|
|
|
def _team_member_with_own_grant(scenario: Scenario, team_id: str, own_server_id: str) -> str:
|
|
user_id: Final = scenario.user(user_role="internal_user", object_permission={"mcp_servers": [own_server_id]})
|
|
scenario.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}})
|
|
return user_id
|
|
|
|
|
|
def test_dashboard_session_serves_the_team_toolset_despite_a_disjoint_grant_on_the_user_row(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""The member's own row grants a different server outright. That grant must not cap the team's toolset
|
|
to nothing, and the team's sibling toolset must not leak onto the granted toolset's route."""
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
own_server_id: Final = register_mcp(scenario, peer, "lit6029_own_" + uuid.uuid4().hex[:8])
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
sibling_id, sibling_name = _toolset(scenario, server_id, "multiply")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id, sibling_id]})
|
|
user_id: Final = _team_member_with_own_grant(scenario, team_id, own_server_id)
|
|
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
|
|
|
|
assert set(_listed_toolset_ids(gateway, headers)) == {granted_id, sibling_id}
|
|
listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
|
|
assert listed.ok, listed.raw
|
|
assert listed.tools == (f"{alias}-add",), listed.raw
|
|
assert _toolset_rpc(gateway, headers, sibling_name, "tools/list", {}).tools == (f"{alias}-multiply",)
|
|
peer.drain()
|
|
called: Final = _toolset_rpc(
|
|
gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
|
|
)
|
|
assert called.ok and called.text == "9", called.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
stranger: Final = {
|
|
"Authorization": f"Bearer {_dashboard_ui_session_token(scenario.user(user_role='internal_user'))}"
|
|
}
|
|
assert _toolset_rpc(gateway, stranger, granted_name, "tools/list", {}).status == 403
|
|
|
|
|
|
def test_a_member_removed_from_the_team_loses_its_toolset_on_the_dashboard_session(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
user_id: Final = scenario.member(team_id)
|
|
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
|
|
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
|
|
assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
|
|
|
|
gateway.post("/team/member_delete", {"team_id": team_id, "user_id": user_id})
|
|
|
|
assert _listed_toolset_ids(gateway, headers) == ()
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
|
|
assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).status == 403
|
|
|
|
|
|
def _bearer(token: str) -> dict[str, str]:
|
|
return {"Authorization": f"Bearer {token}"}
|
|
|
|
|
|
def _expired_dashboard_token(user_id: str) -> str:
|
|
expired: Final = (datetime.now(timezone.utc) - timedelta(minutes=5)).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
|
|
stale: Final = UserAPIKeyAuth(
|
|
token="ui-token",
|
|
key_name="ui-token",
|
|
key_alias="ui-token",
|
|
expires=expired + "+00:00",
|
|
user_id=user_id,
|
|
team_id="litellm-dashboard",
|
|
models=[],
|
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
|
)
|
|
return encrypt_bearer_token(stale.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)
|
|
|
|
|
|
def _rest_list(gateway: Gateway, headers: dict[str, str], params: object) -> httpx.Response:
|
|
return gateway.client.get("/mcp-rest/tools/list", headers=headers, params=params)
|
|
|
|
|
|
def _rest_call(gateway: Gateway, headers: dict[str, str], name: str, server_id: str) -> Outcome:
|
|
return _outcome_from_rest(
|
|
gateway.client.post(
|
|
"/mcp-rest/tools/call",
|
|
headers=headers,
|
|
json={"name": name, "arguments": dict(ADD), "server_id": server_id},
|
|
)
|
|
)
|
|
|
|
|
|
def _route_tools(gateway: Gateway, headers: dict[str, str], name: str) -> Outcome:
|
|
return _toolset_rpc(gateway, headers, name, "tools/list", {})
|
|
|
|
|
|
def _route_call(gateway: Gateway, headers: dict[str, str], name: str, tool: str) -> Outcome:
|
|
return _toolset_rpc(gateway, headers, name, "tools/call", {"name": tool, "arguments": dict(ADD)})
|
|
|
|
|
|
def _strict_config(directory: Path) -> Path:
|
|
base: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
|
strict: Final = {
|
|
**base,
|
|
"general_settings": {**base.get("general_settings", {}), "require_key_mcp_access_defined": True},
|
|
}
|
|
path: Final = directory / "require_key_mcp_access.yaml"
|
|
path.write_text(yaml.safe_dump(strict))
|
|
return path
|
|
|
|
|
|
def test_mcp_rest_toolset_name_narrows_the_list_and_serves_the_call_for_a_team_key_and_a_dashboard_member(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029rest" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
_, withheld_name = _toolset(scenario, server_id, "multiply")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
key: Final = scenario.key(team_id=team_id)
|
|
member: Final = scenario.member(team_id)
|
|
for headers in (_bearer(key), _bearer(_dashboard_ui_session_token(member))):
|
|
listed: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
|
|
assert listed.status_code == 200, listed.text
|
|
listed_names: Final = tuple(tool["name"] for tool in listed.json()["tools"])
|
|
assert len(listed_names) == 1 and listed_names[0].endswith("add"), listed.text
|
|
denied: Final = _rest_list(gateway, headers, {"toolset_name": withheld_name})
|
|
assert denied.status_code == 200 and denied.json()["tools"] == [], denied.text
|
|
assert "does not have access to toolset" in denied.json()["message"], denied.text
|
|
peer.drain()
|
|
called: Final = _rest_call(gateway, headers, listed_names[0], server_id)
|
|
assert called.ok and called.text == "9", called.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
|
|
|
|
def test_a_key_restricted_to_its_own_servers_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029own" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
other_id: Final = register_mcp(scenario, peer, "lit6029other" + uuid.uuid4().hex[:6])
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id], "mcp_servers": [other_id]})
|
|
key: Final = scenario.key(team_id=team_id, object_permission={"mcp_servers": [other_id]})
|
|
headers: Final = _bearer(key)
|
|
assert _listed_toolset_ids(gateway, headers) == ()
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
|
|
assert _route_tools(gateway, headers, granted_name).status == 403
|
|
assert tool_calls(peer.drain()) == ()
|
|
|
|
|
|
def test_a_team_with_an_empty_or_absent_toolset_grant_gives_its_keys_nothing(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
server_id: Final = register_mcp(scenario, peer, "lit6029empty" + uuid.uuid4().hex[:6])
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
teams: Final = (
|
|
scenario.team(object_permission={"mcp_toolsets": []}),
|
|
scenario.team(object_permission={"mcp_toolsets": None}),
|
|
scenario.team(),
|
|
)
|
|
for team_id in teams:
|
|
headers: Final = _bearer(scenario.key(team_id=team_id))
|
|
assert _listed_toolset_ids(gateway, headers) == (), team_id
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
|
|
assert _route_tools(gateway, headers, granted_name).status == 403, team_id
|
|
assert tool_calls(peer.drain()) == ()
|
|
|
|
|
|
def test_an_unknown_or_malformed_toolset_name_is_refused_without_peer_traffic_and_the_route_keeps_serving(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029bad" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
_, withheld_name = _toolset(scenario, server_id, "multiply")
|
|
headers: Final = _bearer(scenario.key(team_id=scenario.team(object_permission={"mcp_toolsets": [granted_id]})))
|
|
unknown: Final = "missing" + uuid.uuid4().hex[:8]
|
|
assert _route_tools(gateway, headers, unknown).status == 404
|
|
assert _rest_list(gateway, headers, {"toolset_name": unknown}).status_code == 404
|
|
malformed: Final = (
|
|
[("toolset_name", granted_name), ("toolset_name", withheld_name)],
|
|
{"toolset_name": "x" * 5000},
|
|
{"toolset_name": ""},
|
|
{"toolset_name": granted_name + "\x00"},
|
|
)
|
|
statuses: Final = tuple(_rest_list(gateway, headers, params).status_code for params in malformed)
|
|
assert all(status < 500 for status in statuses), statuses
|
|
assert tool_calls(peer.drain()) == ()
|
|
assert gateway.client.get("/health/liveliness").status_code == 200
|
|
served: Final = _route_tools(gateway, headers, granted_name)
|
|
assert served.tools == (f"{alias}-add",), served.raw
|
|
|
|
|
|
def test_garbage_expired_and_tampered_credentials_are_refused_on_every_toolset_surface(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
server_id: Final = register_mcp(scenario, peer, "lit6029cred" + uuid.uuid4().hex[:6])
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
member: Final = scenario.member(team_id)
|
|
forged: Final = (
|
|
"sk-" + secrets.token_urlsafe(24),
|
|
_expired_dashboard_token(member),
|
|
"llm_session_" + secrets.token_urlsafe(32),
|
|
)
|
|
for bearer in forged:
|
|
headers: Final = _bearer(bearer)
|
|
listed: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
|
|
assert listed.status_code == 401, (bearer[:12], listed.text)
|
|
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
|
|
assert detail.status_code == 401, (bearer[:12], detail.text)
|
|
routed: Final = _route_tools(gateway, headers, granted_name)
|
|
assert routed.status == 401, (bearer[:12], routed.raw)
|
|
rest: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
|
|
assert rest.status_code == 401, (bearer[:12], rest.text)
|
|
assert peer.drain() == ()
|
|
|
|
|
|
def test_a_dashboard_member_of_a_deleted_team_loses_the_toolset_while_a_direct_grant_survives(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029gone" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
team_id_, team_name = _toolset(scenario, server_id, "add")
|
|
own_id, own_name = _toolset(scenario, server_id, "multiply")
|
|
doomed: Final = scenario.gateway.post(
|
|
"/team/new",
|
|
{"team_alias": f"integration-{uuid.uuid4().hex}", "object_permission": {"mcp_toolsets": [team_id_]}},
|
|
)
|
|
doomed_team: Final = str(doomed["team_id"])
|
|
member: Final = scenario.user(user_role="internal_user", teams=[doomed_team])
|
|
granted: Final = scenario.user(
|
|
user_role="internal_user", teams=[doomed_team], object_permission={"mcp_toolsets": [own_id]}
|
|
)
|
|
member_headers: Final = _bearer(_dashboard_ui_session_token(member))
|
|
granted_headers: Final = _bearer(_dashboard_ui_session_token(granted))
|
|
assert _listed_toolset_ids(gateway, member_headers) == (team_id_,)
|
|
assert set(_listed_toolset_ids(gateway, granted_headers)) == {team_id_, own_id}
|
|
scenario.delete_team(doomed_team)
|
|
assert _listed_toolset_ids(gateway, member_headers) == ()
|
|
assert gateway.client.get(f"/v1/mcp/toolset/{team_id_}", headers=member_headers).status_code == 403
|
|
assert _route_tools(gateway, member_headers, team_name).status == 403
|
|
assert _listed_toolset_ids(gateway, granted_headers) == (own_id,)
|
|
assert _route_tools(gateway, granted_headers, team_name).status == 403
|
|
assert _route_tools(gateway, granted_headers, own_name).tools == (f"{alias}-multiply",)
|
|
peer.drain()
|
|
kept: Final = _route_call(gateway, granted_headers, own_name, f"{alias}-multiply")
|
|
assert kept.ok and kept.text == "20", kept.raw
|
|
assert [call["body"]["params"]["name"] for call in tool_calls(peer.drain())] == ["multiply"]
|
|
|
|
|
|
def test_require_key_mcp_access_defined_stops_key_inheritance_but_not_the_dashboard_member(
|
|
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with (
|
|
owned_proxy(gateway, tmp_path, {}, config=_strict_config(tmp_path), workers=2) as strict,
|
|
mcp_peer() as peer,
|
|
strict.scenario() as scenario,
|
|
):
|
|
alias: Final = "lit6029strict" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
inheriting: Final = _bearer(scenario.key(team_id=team_id))
|
|
own: Final = _bearer(scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [granted_id]}))
|
|
member: Final = _bearer(_dashboard_ui_session_token(scenario.member(team_id)))
|
|
assert _listed_toolset_ids(strict, inheriting) == ()
|
|
assert _route_tools(strict, inheriting, granted_name).status == 403
|
|
assert _listed_toolset_ids(strict, own) == (granted_id,)
|
|
own_tools: Final = _route_tools(strict, own, granted_name)
|
|
assert own_tools.tools == (f"{alias}-add",), (own_tools.status, own_tools.raw)
|
|
assert _listed_toolset_ids(strict, member) == (granted_id,)
|
|
member_tools: Final = _route_tools(strict, member, granted_name)
|
|
assert member_tools.tools == (f"{alias}-add",), (member_tools.status, member_tools.raw)
|
|
peer.drain()
|
|
called: Final = _route_call(strict, member, granted_name, f"{alias}-add")
|
|
assert called.ok and called.text == "9", called.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
|
|
|
|
def test_a_key_with_only_a_vector_store_grant_still_inherits_the_team_toolset(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029vs" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
headers: Final = _bearer(
|
|
scenario.key(team_id=team_id, object_permission={"vector_stores": ["vs-" + uuid.uuid4().hex[:8]]})
|
|
)
|
|
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
|
|
assert _route_tools(gateway, headers, granted_name).tools == (f"{alias}-add",)
|
|
peer.drain()
|
|
called: Final = _route_call(gateway, headers, granted_name, f"{alias}-add")
|
|
assert called.ok and called.text == "9", called.raw
|
|
assert len(tool_calls(peer.drain())) == 1
|
|
|
|
|
|
def test_a_member_added_after_the_team_was_cached_sees_the_toolset_on_every_following_request(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
server_id: Final = register_mcp(scenario, peer, "lit6029cache" + uuid.uuid4().hex[:6])
|
|
granted_id, _ = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
user_id: Final = scenario.user(user_role="internal_user")
|
|
headers: Final = _bearer(_dashboard_ui_session_token(user_id))
|
|
warm: Final = _bearer(scenario.key(team_id=team_id))
|
|
assert tuple(_listed_toolset_ids(gateway, warm) for _ in range(4)) == ((granted_id,),) * 4
|
|
assert tuple(_listed_toolset_ids(gateway, headers) for _ in range(4)) == ((),) * 4
|
|
added: Final = gateway.request(
|
|
"POST", "/team/member_add", {"team_id": team_id, "member": {"user_id": user_id, "role": "user"}}
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
listings: Final = tuple(_listed_toolset_ids(gateway, headers) for _ in range(8))
|
|
assert listings == ((granted_id,),) * 8, listings
|
|
|
|
|
|
def test_a_member_of_two_teams_sees_the_union_and_each_route_stays_narrowed_to_its_own_toolset(
|
|
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "lit6029two" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
first_id, first_name = _toolset(scenario, server_id, "add")
|
|
second_id, second_name = _toolset(scenario, server_id, "multiply")
|
|
teams: Final = (
|
|
scenario.team(object_permission={"mcp_toolsets": [first_id]}),
|
|
scenario.team(object_permission={"mcp_toolsets": [second_id]}),
|
|
)
|
|
headers: Final = _bearer(
|
|
_dashboard_ui_session_token(scenario.user(user_role="internal_user", teams=list(teams)))
|
|
)
|
|
assert set(_listed_toolset_ids(gateway, headers)) == {first_id, second_id}
|
|
assert _route_tools(gateway, headers, first_name).tools == (f"{alias}-add",)
|
|
assert _route_tools(gateway, headers, second_name).tools == (f"{alias}-multiply",)
|
|
peer.drain()
|
|
crossed: Final = _route_call(gateway, headers, first_name, f"{alias}-multiply")
|
|
assert not crossed.ok, crossed.raw
|
|
assert tool_calls(peer.drain()) == ()
|
|
|
|
|
|
def test_owned_proxy_loss_retains_exit_status_and_first_http_failure(gateway: Gateway, tmp_path: Path) -> None:
|
|
with pytest.raises(httpx.TransportError) as failure: # noqa: PT012 # HTTP failure must propagate through proxy cleanup
|
|
with owned_proxy_process(gateway, tmp_path, {}) as owned:
|
|
response: Final = owned.gateway.client.get("/health/readiness")
|
|
assert response.status_code == 200, response.text
|
|
owned.kill_supervisor()
|
|
owned.gateway.client.get("/health/readiness")
|
|
record: Final = TypeAdapter(dict[str, int | float | str | None]).validate_json(
|
|
owned.log.with_suffix(".exit.json").read_text()
|
|
)
|
|
assert record["root_pid"] == owned.process.pid
|
|
assert record["exit_before_stop"] == -signal.SIGKILL
|
|
assert record["exit_after_stop"] == -signal.SIGKILL
|
|
assert record["exception_before_stop"] == failure.type.__name__
|
|
assert record["exception_at_record"] == failure.type.__name__
|
|
|
|
|
|
_PROBE: Final = "catalog-probe"
|
|
_ECHO: Final = "catalog-echo"
|
|
_UNLISTED: Final = ""
|
|
_GUARDRAIL_CODE: Final = (
|
|
"def apply_guardrail(inputs, request_data, input_type):\n"
|
|
f' if "{_PROBE}" not in list(inputs.get("texts") or []):\n'
|
|
" return allow()\n"
|
|
' function = inputs.get("tools", [{}])[0].get("function", {})\n'
|
|
f' return block("{_ECHO}[" + function.get("description") + "]")\n'
|
|
)
|
|
|
|
|
|
_ECHO_GUARDRAIL_YAML: Final = (
|
|
"guardrails:\n"
|
|
" - guardrail_name: catalog-echo\n"
|
|
" litellm_params:\n"
|
|
" guardrail: custom_code\n"
|
|
" mode: pre_mcp_call\n"
|
|
" default_on: true\n"
|
|
" custom_code: |\n" + textwrap.indent(_GUARDRAIL_CODE, 8 * " ")
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def echo_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
|
directory: Final = tmp_path_factory.mktemp("catalog-echo")
|
|
path: Final = directory / "catalog_echo.yaml"
|
|
path.write_text((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text() + _ECHO_GUARDRAIL_YAML)
|
|
with gateway_from_environment() as gateway, owned_proxy(gateway, directory, {}, config=path, workers=2) as rig:
|
|
yield rig
|
|
|
|
|
|
def _echoed_description(outcome: Outcome) -> str:
|
|
found: Final = re.search(rf"{_ECHO}\[(.*?)\]", outcome.raw)
|
|
assert found is not None, outcome.raw
|
|
return found.group(1)
|
|
|
|
|
|
def test_a_team_keys_toolset_route_listing_feeds_its_own_calls_but_not_a_team_mates(echo_rig: Gateway) -> None:
|
|
described: Final = "Adds for the team " + uuid.uuid4().hex[:8]
|
|
tool: Final = ScriptedTool("add", lambda _: text_result("9"), description=described)
|
|
with scripted_peer(tool) as peer, echo_rig.scenario() as scenario:
|
|
alias: Final = "lit6029echo" + uuid.uuid4().hex[:6]
|
|
server_id: Final = register_mcp(scenario, peer, alias)
|
|
granted_id, granted_name = _toolset(scenario, server_id, "add")
|
|
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
|
|
key: Final = scenario.key(team_id=team_id)
|
|
team_mate: Final = scenario.key(team_id=team_id)
|
|
_assert_team_grants_only(echo_rig, team_id, key, granted_id)
|
|
listed: Final = _toolset_rpc(echo_rig, _bearer(key), granted_name, "tools/list", {})
|
|
assert listed.ok and listed.tools == (f"{alias}-add",), listed.raw
|
|
probe: Final[dict[str, object]] = {"name": f"{alias}-add", "arguments": {"probe": _PROBE}}
|
|
own: Final = _echoed_description(_toolset_rpc(echo_rig, _bearer(key), granted_name, "tools/call", probe))
|
|
mate: Final = _echoed_description(_toolset_rpc(echo_rig, _bearer(team_mate), granted_name, "tools/call", probe))
|
|
assert (own, mate) == (described, _UNLISTED), (
|
|
"the slot is keyed by the hashed key, so a team-mate that never listed is handed nothing"
|
|
)
|
|
assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer"
|