mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +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
718 lines
34 KiB
Python
718 lines
34 KiB
Python
import base64
|
|
import re
|
|
import textwrap
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
|
|
from integration._support.database import read_rows
|
|
from integration._support.mcp import (
|
|
ENTRY_POINTS,
|
|
EntryPoint,
|
|
McpCaller,
|
|
McpPeer,
|
|
Outcome,
|
|
ScriptedTool,
|
|
call_tool,
|
|
forget_mcp,
|
|
mcp_peer,
|
|
openapi_peer,
|
|
register_mcp,
|
|
scripted_peer,
|
|
text_result,
|
|
tool_calls,
|
|
tool_names,
|
|
)
|
|
from integration._support.oauth_server import oauth_server
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import TypeAdapter
|
|
|
|
ADD: Final = {"a": 2, "b": 3}
|
|
STATIC_MODES: Final = (
|
|
("api_key", b"x-api-key", "{secret}"),
|
|
("bearer_token", b"authorization", "Bearer {secret}"),
|
|
("basic", b"authorization", "Basic {basic}"),
|
|
("authorization", b"authorization", "{secret}"),
|
|
)
|
|
|
|
|
|
def _header(call: dict[str, object], name: bytes) -> bytes | None:
|
|
headers: Final = call["headers"]
|
|
assert isinstance(headers, dict)
|
|
value: Final = headers.get(name)
|
|
return value if isinstance(value, bytes) else None
|
|
|
|
|
|
def _one_call(peer: McpPeer) -> dict[str, object]:
|
|
sent: Final = tool_calls(peer.drain())
|
|
assert len(sent) == 1, sent
|
|
return sent[0]
|
|
|
|
|
|
def _plaintext_rows(identity: str, secret: str) -> list[dict[str, object]]:
|
|
return read_rows(
|
|
'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id = %s '
|
|
"AND (credentials::text LIKE %s OR static_headers::text LIKE %s)",
|
|
(identity, f"%{secret}%", f"%{secret}%"),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(("auth_type", "header", "shape"), STATIC_MODES)
|
|
def test_static_credential_reaches_the_peer_in_its_mode_shape_and_is_encrypted_at_rest(
|
|
gateway: Gateway, auth_type: str, header: bytes, shape: str
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
secret: Final = "user:" + uuid.uuid4().hex
|
|
basic: Final = base64.b64encode(secret.encode()).decode()
|
|
identity: Final = register_mcp(
|
|
scenario, peer, "cred" + uuid.uuid4().hex[:8], auth_type=auth_type, credentials={"auth_value": secret}
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
peer.drain()
|
|
response: Final = call_tool(gateway, key, identity, tool_names(gateway, key, identity)["add"], ADD)
|
|
assert response.status_code == 200, response.text
|
|
assert _header(_one_call(peer), header) == shape.format(secret=secret, basic=basic).encode()
|
|
assert _plaintext_rows(identity, secret) == [], "credential stored in plaintext"
|
|
|
|
|
|
def test_create_response_does_not_echo_the_stored_credential(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
secret: Final = "create-secret-" + uuid.uuid4().hex
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/mcp/server",
|
|
{
|
|
"server_name": alias,
|
|
"alias": alias,
|
|
**peer.registration(),
|
|
"auth_type": "bearer_token",
|
|
"credentials": {"auth_value": secret},
|
|
},
|
|
)
|
|
scenario.cleanups.callback(forget_mcp, gateway, response.json()["server_id"])
|
|
assert response.status_code == 201, response.text
|
|
assert response.json()["alias"] == alias
|
|
assert response.json().get("credentials") is None
|
|
assert secret not in response.text
|
|
|
|
|
|
def test_editing_the_credential_rotates_what_the_peer_receives(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
first: Final = "cred-" + uuid.uuid4().hex
|
|
second: Final = "cred-" + uuid.uuid4().hex
|
|
identity: Final = register_mcp(
|
|
scenario, peer, "cred" + uuid.uuid4().hex[:8], auth_type="bearer_token", credentials={"auth_value": first}
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
name: Final = tool_names(gateway, key, identity)["add"]
|
|
peer.drain()
|
|
assert call_tool(gateway, key, identity, name, ADD).status_code == 200
|
|
assert _header(_one_call(peer), b"authorization") == f"Bearer {first}".encode()
|
|
rotated: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "credentials": {"auth_value": second}}
|
|
)
|
|
assert rotated.status_code == 202, rotated.text
|
|
observed: Final = eventually(
|
|
lambda: (call_tool(gateway, key, identity, name, ADD).status_code, tool_calls(peer.drain())),
|
|
lambda value: any(_header(call, b"authorization") == f"Bearer {second}".encode() for call in value[1]),
|
|
)
|
|
assert all(_header(call, b"authorization") != f"Bearer {first}".encode() for call in observed[1][-1:])
|
|
assert _plaintext_rows(identity, second) == [] and _plaintext_rows(identity, first) == []
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ENTRY_POINTS)
|
|
def test_caller_headers_for_other_servers_and_unknown_headers_never_reach_the_peer(
|
|
gateway: Gateway, entry: EntryPoint
|
|
) -> None:
|
|
with mcp_peer() as peer, mcp_peer() as other, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
other_alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
other_id: Final = register_mcp(scenario, other, other_alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity, other_id]})
|
|
leak: Final = "leak-" + uuid.uuid4().hex
|
|
caller: Final = McpCaller(
|
|
gateway,
|
|
key,
|
|
entry,
|
|
alias,
|
|
headers={
|
|
f"x-mcp-{other_alias}-authorization": f"Bearer {leak}",
|
|
"x-integration-unknown": leak,
|
|
"cookie": f"session={leak}",
|
|
},
|
|
)
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry in ("mcp", "root", "sse", "rest") else None)
|
|
assert outcome.ok, outcome.raw
|
|
call: Final = _one_call(peer)
|
|
assert leak.encode() not in b"".join(_header(call, name) or b"" for name in call["headers"]), call["headers"]
|
|
assert tool_calls(other.drain()) == ()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("auth_type", "per_server_header"),
|
|
(("none", True), ("true_passthrough", False)),
|
|
)
|
|
def test_callers_own_litellm_key_never_reaches_the_peer_over_rest(
|
|
gateway: Gateway, auth_type: str, per_server_header: bool
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, auth_type=auth_type)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller_headers: Final = {
|
|
"x-litellm-api-key": f"Bearer {key}",
|
|
(f"x-mcp-{alias}-authorization" if per_server_header else "Authorization"): f"Bearer {key}",
|
|
}
|
|
peer.drain()
|
|
response: Final = call_tool(
|
|
gateway, key, identity, tool_names(gateway, key, identity)["add"], ADD, headers=caller_headers
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["isError"] is False
|
|
observed: Final = peer.drain()
|
|
calls: Final = tool_calls(observed)
|
|
assert len(calls) == 1, calls
|
|
header_sets: Final = tuple(
|
|
TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed
|
|
)
|
|
assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), (
|
|
"caller key appeared in the recorded MCP header set"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ("server_mcp", "rest"))
|
|
def test_empty_primary_header_does_not_expose_authorization_key(gateway: Gateway, entry: EntryPoint) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, extra_headers=["x-upstream-token", "x-tenant"])
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(
|
|
gateway,
|
|
key,
|
|
entry,
|
|
alias,
|
|
headers={
|
|
"x-litellm-api-key": "",
|
|
"Authorization": f"Bearer {key}",
|
|
f"x-mcp-{alias}-authorization": f"Bearer {key}",
|
|
"x-upstream-token": key,
|
|
"x-tenant": "tenant-control",
|
|
},
|
|
)
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry == "rest" else None)
|
|
assert outcome.ok, outcome.raw
|
|
observed: Final = peer.drain()
|
|
calls: Final = tool_calls(observed)
|
|
assert len(calls) == 1, calls
|
|
assert _header(calls[0], b"x-tenant") == b"tenant-control"
|
|
header_sets: Final = tuple(
|
|
TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed
|
|
)
|
|
assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), (
|
|
"empty primary header prevented scrubbing the admitted Authorization key"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ENTRY_POINTS)
|
|
def test_extra_headers_cannot_forward_gateway_admission_key(gateway: Gateway, entry: EntryPoint) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario, peer, alias, extra_headers=["X-LiteLLM-API-Key", "x-tenant"],
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(gateway, key, entry, alias, headers={"x-tenant": "tenant-control"})
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry != "server_mcp" else None)
|
|
assert outcome.ok, outcome.raw
|
|
observed: Final = peer.drain()
|
|
calls: Final = tool_calls(observed)
|
|
assert len(calls) == 1, calls
|
|
assert _header(calls[0], b"x-tenant") == b"tenant-control"
|
|
header_sets: Final = tuple(
|
|
TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed
|
|
)
|
|
assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), (
|
|
"gateway admission key reached the upstream through Extra Headers"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ("server_mcp", "rest"))
|
|
def test_openapi_extra_headers_cannot_forward_gateway_admission_key(gateway: Gateway, entry: EntryPoint) -> None:
|
|
with openapi_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario, peer, alias, extra_headers=["X-LiteLLM-API-Key", "x-tenant"],
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(gateway, key, entry, alias, headers={"x-tenant": "tenant-control"})
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-getpet", {"petId": "7"}, identity if entry == "rest" else None)
|
|
assert outcome.ok, outcome.raw
|
|
observed: Final = peer.drain()
|
|
assert [(request["method"], request["path"]) for request in observed] == [("GET", "/pets/7")]
|
|
headers: Final = TypeAdapter(dict[bytes, bytes]).validate_python(observed[0]["headers"])
|
|
assert headers.get(b"x-tenant") == b"tenant-control"
|
|
assert all(key.encode() not in value for value in headers.values()), "gateway admission key reached OpenAPI"
|
|
|
|
|
|
def test_server_scoped_caller_header_reaches_only_its_server(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, mcp_peer() as other, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
other_alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
other_id: Final = register_mcp(scenario, other, other_alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity, other_id]})
|
|
token: Final = "user-" + uuid.uuid4().hex
|
|
caller: Final = McpCaller(
|
|
gateway, key, "mcp", None, headers={f"x-mcp-{alias}-authorization": f"Bearer {token}"}
|
|
)
|
|
peer.drain()
|
|
other.drain()
|
|
assert caller.call(f"{alias}-add", ADD).ok
|
|
assert caller.call(f"{other_alias}-add", ADD).ok
|
|
assert _header(_one_call(peer), b"authorization") == f"Bearer {token}".encode()
|
|
assert _header(_one_call(other), b"authorization") is None
|
|
|
|
|
|
def test_extra_headers_allowlist_forwards_only_named_headers(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "cred" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, extra_headers=["x-tenant"])
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(gateway, key, "server_mcp", alias, headers={"x-tenant": "acme", "x-other": "no"})
|
|
peer.drain()
|
|
assert caller.call(f"{alias}-add", ADD).ok
|
|
call: Final = _one_call(peer)
|
|
assert _header(call, b"x-tenant") == b"acme"
|
|
assert _header(call, b"x-other") is None
|
|
|
|
|
|
def test_byok_server_uses_the_calling_users_stored_credential_and_fails_closed_without_one(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "byok" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, auth_type="api_key", is_byok=True)
|
|
owner: Final = scenario.user()
|
|
stranger: Final = scenario.user()
|
|
owner_key: Final = scenario.key(user_id=owner, object_permission={"mcp_servers": [identity]})
|
|
stranger_key: Final = scenario.key(user_id=stranger, object_permission={"mcp_servers": [identity]})
|
|
secret: Final = "byok-" + uuid.uuid4().hex
|
|
stored: Final = gateway.client.post(
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
json={"credential": secret},
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
assert stored.status_code in (200, 201), stored.text
|
|
scenario.cleanups.callback(
|
|
gateway.client.delete,
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
assert (
|
|
read_rows(
|
|
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE server_id = %s AND credential_b64 LIKE %s',
|
|
(identity, f"%{secret}%"),
|
|
)
|
|
== []
|
|
)
|
|
name: Final = f"{alias}-add"
|
|
peer.drain()
|
|
granted: Final = call_tool(gateway, owner_key, identity, name, ADD)
|
|
assert granted.status_code == 200, granted.text
|
|
assert _header(_one_call(peer), b"x-api-key") == secret.encode()
|
|
denied: Final = call_tool(gateway, stranger_key, identity, name, ADD)
|
|
assert denied.status_code == 401, denied.text
|
|
assert tool_calls(peer.drain()) == ()
|
|
removed: Final = gateway.client.delete(
|
|
f"/v1/mcp/server/{identity}/user-credential", headers={"x-litellm-api-key": owner_key}
|
|
)
|
|
assert removed.status_code in (200, 204), removed.text
|
|
eventually(lambda: call_tool(gateway, owner_key, identity, name, ADD), lambda value: value.status_code == 401)
|
|
assert tool_calls(peer.drain()) == ()
|
|
|
|
|
|
def _listings(peer: McpPeer) -> tuple[dict[str, object], ...]:
|
|
return tuple(
|
|
item
|
|
for item in peer.drain()
|
|
if isinstance(item.get("body"), dict) and item["body"].get("method") == "tools/list"
|
|
)
|
|
|
|
|
|
def test_oauth2_byok_listing_sends_the_minted_token_not_the_users_stored_secret(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
|
|
alias: Final = "cc" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
auth_type="oauth2",
|
|
oauth2_flow="client_credentials",
|
|
is_byok=True,
|
|
token_url=auth.issuer + "/token",
|
|
credentials={"client_id": "cc-client", "client_secret": "cc-secret-" + uuid.uuid4().hex},
|
|
)
|
|
owner: Final = scenario.user()
|
|
owner_key: Final = scenario.key(user_id=owner, object_permission={"mcp_servers": [identity]})
|
|
secret: Final = "byok-" + uuid.uuid4().hex
|
|
stored: Final = gateway.client.post(
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
json={"credential": secret},
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
assert stored.status_code in (200, 201), stored.text
|
|
scenario.cleanups.callback(
|
|
gateway.client.delete,
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
peer.drain()
|
|
auth.drain()
|
|
response: Final = gateway.client.get(
|
|
"/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": owner_key}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert "add" in {tool["name"] for tool in response.json()["tools"]}, response.text
|
|
assert [request["grant_type"] for request in auth.token_requests()] == ["client_credentials"]
|
|
listings: Final = _listings(peer)
|
|
assert len(listings) == 1, listings
|
|
sent: Final = _header(listings[0], b"authorization")
|
|
assert sent is not None and auth.is_live(sent.decode().removeprefix("Bearer ")), sent
|
|
assert secret.encode() not in sent, "stored BYOK secret replaced the minted token on tools/list"
|
|
|
|
|
|
@pytest.mark.parametrize(("auth_type", "header", "shape"), STATIC_MODES[:2])
|
|
def test_byok_rest_listing_sends_the_servers_static_credential_not_the_users_stored_secret(
|
|
gateway: Gateway, auth_type: str, header: bytes, shape: str
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "byok" + uuid.uuid4().hex[:8]
|
|
static: Final = "static-" + uuid.uuid4().hex
|
|
identity: Final = register_mcp(
|
|
scenario, peer, alias, auth_type=auth_type, is_byok=True, credentials={"auth_value": static}
|
|
)
|
|
owner: Final = scenario.user()
|
|
owner_key: Final = scenario.key(user_id=owner, object_permission={"mcp_servers": [identity]})
|
|
secret: Final = "byok-" + uuid.uuid4().hex
|
|
stored: Final = gateway.client.post(
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
json={"credential": secret},
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
assert stored.status_code in (200, 201), stored.text
|
|
scenario.cleanups.callback(
|
|
gateway.client.delete,
|
|
f"/v1/mcp/server/{identity}/user-credential",
|
|
headers={"x-litellm-api-key": owner_key},
|
|
)
|
|
peer.drain()
|
|
response: Final = gateway.client.get(
|
|
"/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": owner_key}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert "add" in {tool["name"] for tool in response.json()["tools"]}, response.text
|
|
listings: Final = _listings(peer)
|
|
assert len(listings) == 1, listings
|
|
assert _header(listings[0], header) == shape.format(secret=static, basic="").encode(), listings[0]["headers"]
|
|
peer.drain()
|
|
called: Final = call_tool(gateway, owner_key, identity, f"{alias}-add", ADD)
|
|
assert called.status_code == 200, called.text
|
|
assert _header(_one_call(peer), header) == shape.format(secret=secret, basic="").encode()
|
|
|
|
|
|
def test_deprecated_string_x_mcp_auth_lists_a_byok_server_for_a_key_without_a_user(gateway: Gateway) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "byok" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, auth_type="bearer_token", is_byok=True)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
peer.drain()
|
|
response: Final = gateway.client.get(
|
|
"/mcp-rest/tools/list",
|
|
params={"server_id": identity},
|
|
headers={"x-litellm-api-key": key, "x-mcp-auth": "Bearer hdr"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
names: Final = {tool["name"] for tool in response.json()["tools"]}
|
|
assert "add" in names, names
|
|
listings: Final = _listings(peer)
|
|
assert len(listings) == 1, listings
|
|
assert listings[0]["headers"].get(b"authorization") == b"Bearer hdr"
|
|
|
|
|
|
_PROBE: Final = "catalog-probe"
|
|
_ECHO: Final = "catalog-echo"
|
|
_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)
|
|
|
|
|
|
_PROBE_ARGUMENTS: Final = {"probe": _PROBE}
|
|
_HEADERS: Final = TypeAdapter(dict[str, str])
|
|
|
|
|
|
def _add_described_anew_per_fetch_for(header: str) -> ScriptedTool:
|
|
return ScriptedTool(
|
|
"add",
|
|
lambda _: text_result("3"),
|
|
description=lambda headers: f"Adds for {headers.get(header, 'nobody')} at {uuid.uuid4().hex}",
|
|
)
|
|
|
|
|
|
def _fetched_for(credential: str) -> re.Pattern[str]:
|
|
return re.compile(rf"Adds for {re.escape(credential)} at [0-9a-f]{{32}}")
|
|
|
|
|
|
def _listed_description(listing: Outcome, credential: str) -> str:
|
|
found: Final = _fetched_for(credential).search(listing.raw)
|
|
assert found is not None, listing.raw
|
|
return found.group(0)
|
|
|
|
|
|
def _sent_header(request: dict[str, object], name: str) -> str | None:
|
|
return _HEADERS.validate_python(request["headers"]).get(name)
|
|
|
|
|
|
def _store_byok_credential(scenario: Scenario, identity: str, key: str, secret: str) -> None:
|
|
stored: Final = scenario.gateway.client.post(
|
|
f"/v1/mcp/server/{identity}/user-credential", json={"credential": secret}, headers={"x-litellm-api-key": key}
|
|
)
|
|
assert stored.status_code in (200, 201), stored.text
|
|
scenario.cleanups.callback(
|
|
scenario.gateway.client.delete, f"/v1/mcp/server/{identity}/user-credential", headers={"x-litellm-api-key": key}
|
|
)
|
|
|
|
|
|
def test_rotating_the_credential_drops_the_callers_listing_until_it_lists_again(echo_rig: Gateway) -> None:
|
|
with scripted_peer(_add_described_anew_per_fetch_for("authorization")) as peer, echo_rig.scenario() as scenario:
|
|
first: Final = "cred-" + uuid.uuid4().hex
|
|
second: Final = "cred-" + uuid.uuid4().hex
|
|
alias: Final = "rot" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario, peer, alias, auth_type="bearer_token", credentials={"auth_value": first}
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
caller: Final = McpCaller(echo_rig, key, "mcp", alias)
|
|
listed: Final = caller.list_tools()
|
|
assert listed.ok, listed.raw
|
|
served: Final = _listed_description(listed, f"Bearer {first}")
|
|
assert _echoed_description(caller.call(f"{alias}-add", _PROBE_ARGUMENTS)) == served
|
|
rotated: Final = echo_rig.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "credentials": {"auth_value": second}}
|
|
)
|
|
assert rotated.status_code == 202, rotated.text
|
|
eventually(
|
|
lambda: _echoed_description(caller.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
lambda seen: _fetched_for(f"Bearer {second}").fullmatch(seen) is not None,
|
|
)
|
|
peer.drain()
|
|
relisting: Final = caller.list_tools()
|
|
assert relisting.ok, relisting.raw
|
|
served_again: Final = _listed_description(relisting, f"Bearer {second}")
|
|
assert _echoed_description(caller.call(f"{alias}-add", _PROBE_ARGUMENTS)) == served_again
|
|
relisted: Final = _listings(peer)
|
|
assert len(relisted) == 1, relisted
|
|
assert _sent_header(relisted[0], "authorization") == f"Bearer {second}", relisted[0]["headers"]
|
|
|
|
|
|
def test_byok_callers_are_evaluated_against_their_own_listing_and_the_stored_secret_never_keys_the_slot(
|
|
echo_rig: Gateway,
|
|
) -> None:
|
|
with scripted_peer(_add_described_anew_per_fetch_for("x-api-key")) as peer, echo_rig.scenario() as scenario:
|
|
alias: Final = "byok" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, auth_type="api_key", is_byok=True)
|
|
owner_key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
|
|
stranger_key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
|
|
owner_secret: Final = "byok-" + uuid.uuid4().hex
|
|
stranger_secret: Final = "byok-" + uuid.uuid4().hex
|
|
replacement: Final = "byok-" + uuid.uuid4().hex
|
|
_store_byok_credential(scenario, identity, owner_key, owner_secret)
|
|
_store_byok_credential(scenario, identity, stranger_key, stranger_secret)
|
|
owner: Final = McpCaller(echo_rig, owner_key, "mcp", alias)
|
|
stranger: Final = McpCaller(echo_rig, stranger_key, "mcp", alias)
|
|
peer.drain()
|
|
listed: Final = owner.list_tools()
|
|
assert listed.ok, listed.raw
|
|
listings: Final = _listings(peer)
|
|
assert [_sent_header(item, "x-api-key") for item in listings] == [owner_secret], listings
|
|
served: Final = _listed_description(listed, owner_secret)
|
|
own: Final = _echoed_description(owner.call(f"{alias}-add", _PROBE_ARGUMENTS))
|
|
other: Final = _echoed_description(stranger.call(f"{alias}-add", _PROBE_ARGUMENTS))
|
|
assert own == served, (own, served)
|
|
assert _fetched_for(stranger_secret).fullmatch(other), (other, served)
|
|
_store_byok_credential(scenario, identity, owner_key, replacement)
|
|
sent: Final = eventually(
|
|
lambda: (owner.call(f"{alias}-add", ADD).ok, tool_calls(peer.drain())),
|
|
lambda value: any(_sent_header(call, "x-api-key") == replacement for call in value[1]),
|
|
)
|
|
assert sent[0], sent
|
|
assert _echoed_description(owner.call(f"{alias}-add", _PROBE_ARGUMENTS)) == served, (
|
|
"the slot is keyed by the client-supplied header, never by the stored credential"
|
|
)
|
|
|
|
|
|
def test_callers_with_different_server_scoped_auth_headers_are_evaluated_against_their_own_listings(
|
|
echo_rig: Gateway,
|
|
) -> None:
|
|
tool: Final = ScriptedTool(
|
|
"add",
|
|
lambda _: text_result("3"),
|
|
description=lambda headers: "Adds for " + headers.get("authorization", "nobody"),
|
|
)
|
|
with scripted_peer(tool) as peer, echo_rig.scenario() as scenario:
|
|
alias: Final = "scoped" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
acme_token: Final = "acme-" + uuid.uuid4().hex
|
|
globex_token: Final = "globex-" + uuid.uuid4().hex
|
|
acme: Final = McpCaller(echo_rig, key, "mcp", alias, {f"x-mcp-{alias}-authorization": f"Bearer {acme_token}"})
|
|
globex: Final = McpCaller(
|
|
echo_rig, key, "mcp", alias, {f"x-mcp-{alias}-authorization": f"Bearer {globex_token}"}
|
|
)
|
|
assert acme.list_tools().ok and globex.list_tools().ok
|
|
seen: Final = (
|
|
_echoed_description(acme.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
_echoed_description(globex.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
)
|
|
assert seen == (f"Adds for Bearer {acme_token}", f"Adds for Bearer {globex_token}"), seen
|
|
assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer"
|
|
|
|
|
|
def test_deprecated_string_x_mcp_auth_callers_on_a_user_less_key_own_separate_listings(echo_rig: Gateway) -> None:
|
|
with scripted_peer(_add_described_anew_per_fetch_for("authorization")) as peer, echo_rig.scenario() as scenario:
|
|
alias: Final = "legacy" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(scenario, peer, alias, auth_type="bearer_token")
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
first_token: Final = "first-" + uuid.uuid4().hex
|
|
second_token: Final = "second-" + uuid.uuid4().hex
|
|
first: Final = McpCaller(echo_rig, key, "mcp", alias, {"x-mcp-auth": f"Bearer {first_token}"})
|
|
second: Final = McpCaller(echo_rig, key, "mcp", alias, {"x-mcp-auth": f"Bearer {second_token}"})
|
|
first_cold: Final = _echoed_description(first.call(f"{alias}-add", _PROBE_ARGUMENTS))
|
|
assert _fetched_for(f"Bearer {first_token}").fullmatch(first_cold), first_cold
|
|
peer.drain()
|
|
first_listing: Final = first.list_tools()
|
|
assert first_listing.ok, first_listing.raw
|
|
listings: Final = _listings(peer)
|
|
assert len(listings) == 1, listings
|
|
assert _sent_header(listings[0], "authorization") == f"Bearer {first_token}", listings[0]["headers"]
|
|
first_served: Final = _listed_description(first_listing, f"Bearer {first_token}")
|
|
warm: Final = eventually(
|
|
lambda: _echoed_description(first.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
lambda seen: seen == first_served,
|
|
)
|
|
assert warm == first_served != first_cold, (first_cold, first_served, warm)
|
|
second_cold: Final = _echoed_description(second.call(f"{alias}-add", _PROBE_ARGUMENTS))
|
|
assert _fetched_for(f"Bearer {second_token}").fullmatch(second_cold), (second_cold, first_served)
|
|
second_listing: Final = second.list_tools()
|
|
assert second_listing.ok, second_listing.raw
|
|
second_served: Final = _listed_description(second_listing, f"Bearer {second_token}")
|
|
seen: Final = eventually(
|
|
lambda: (
|
|
_echoed_description(first.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
_echoed_description(second.call(f"{alias}-add", _PROBE_ARGUMENTS)),
|
|
),
|
|
lambda pair: pair == (first_served, second_served),
|
|
)
|
|
assert seen == (first_served, second_served), seen
|
|
assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer"
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ("mcp", "server_mcp", "sse"))
|
|
@pytest.mark.parametrize("forwarded_token", (False, True))
|
|
def test_oauth_passthrough_probe_uses_server_token_without_admission_key(
|
|
gateway: Gateway, entry: EntryPoint, forwarded_token: bool
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
alias: Final = "probe" + uuid.uuid4().hex[:8]
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
auth_type="none",
|
|
oauth_passthrough=True,
|
|
extra_headers=["Authorization", "x-integration-caller"],
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
upstream_token: Final = "upstream-" + uuid.uuid4().hex
|
|
caller_trace: Final = "caller-" + uuid.uuid4().hex
|
|
caller: Final = McpCaller(
|
|
gateway,
|
|
key,
|
|
entry,
|
|
alias,
|
|
headers={
|
|
"x-integration-caller": caller_trace,
|
|
"Authorization": f"Bearer {upstream_token if forwarded_token else key}",
|
|
f"x-mcp-{alias}-authorization": "Bearer expired-token"
|
|
if forwarded_token
|
|
else f"Bearer {upstream_token}",
|
|
},
|
|
)
|
|
peer.drain()
|
|
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry != "server_mcp" else None)
|
|
assert outcome.ok, outcome.raw
|
|
observed: Final = peer.drain()
|
|
forwarded: Final = f"Bearer {upstream_token}".encode()
|
|
probes: Final = tuple(
|
|
request
|
|
for request in observed
|
|
if TypeAdapter(dict[str, object]).validate_python(request["body"]).get("id") == "litellm-mcp-auth-probe"
|
|
)
|
|
assert probes, "expected an upstream initialize authentication probe"
|
|
assert all(_header(probe, b"authorization") == forwarded for probe in probes)
|
|
callers_own: Final = tuple(
|
|
request for request in observed if _header(request, b"x-integration-caller") == caller_trace.encode()
|
|
)
|
|
assert len(tool_calls(observed)) == 1
|
|
assert len(tool_calls(callers_own)) == 1, "the tool call lost the caller's forwarded headers"
|
|
attributed: Final = probes + callers_own
|
|
foreign: Final = tuple(request for request in attributed if _header(request, b"authorization") != forwarded)
|
|
assert foreign == (), foreign
|
|
credentialed_background: Final = tuple(
|
|
request
|
|
for request in observed
|
|
if request not in attributed and _header(request, b"authorization") is not None
|
|
)
|
|
assert credentialed_background == (), credentialed_background
|
|
header_sets: Final = tuple(
|
|
TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed
|
|
)
|
|
assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), (
|
|
"upstream authentication probe exposed the gateway admission key"
|
|
)
|