mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #42039 from BerriAI/revert-41731-litellm_mcp_integration_regressions_4506
Revert "test(mcp): verify scoped execution and OAuth credential isolation"
This commit is contained in:
joshua-berri
GitHub
commit
93ebc1e281
9 changed files with 1268 additions and 407 deletions
|
|
@ -9,7 +9,7 @@ import httpx
|
|||
from integration._support.asgi import asgi_server
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.database import read_rows
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp_tests.mcp_e2e_upstream_server import add, multiply
|
||||
from starlette.requests import Request
|
||||
|
|
@ -27,7 +27,12 @@ class McpPeer:
|
|||
|
||||
@contextmanager
|
||||
def mcp_peer() -> Iterator[McpPeer]:
|
||||
service: Final = MCPServer("integration-math")
|
||||
service: Final = FastMCP(
|
||||
"integration-math",
|
||||
stateless_http=True,
|
||||
json_response=True,
|
||||
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
|
||||
)
|
||||
service.add_tool(add)
|
||||
service.add_tool(multiply)
|
||||
|
||||
|
|
@ -35,11 +40,7 @@ def mcp_peer() -> Iterator[McpPeer]:
|
|||
def fail() -> str:
|
||||
raise ValueError("synthetic tool failure")
|
||||
|
||||
app: Final = service.streamable_http_app(
|
||||
stateless_http=True,
|
||||
json_response=True,
|
||||
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
|
||||
)
|
||||
app: Final = service.streamable_http_app()
|
||||
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
|
||||
|
||||
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
|
|
@ -93,7 +94,9 @@ def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]:
|
|||
}
|
||||
|
||||
|
||||
def call_tool(gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]) -> httpx.Response:
|
||||
def call_tool(
|
||||
gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]
|
||||
) -> httpx.Response:
|
||||
return gateway.client.post(
|
||||
"/mcp-rest/tools/call",
|
||||
headers={"x-litellm-api-key": key},
|
||||
|
|
|
|||
|
|
@ -1311,27 +1311,6 @@
|
|||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [
|
||||
"quota_management.spend_tracking.scripted_wire.logs_cost"
|
||||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [
|
||||
"other.mcp.health.restricted_keys_intersect_grants_in_both_modes"
|
||||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_warm_credential_removal_rejects_without_upstream_traffic": [
|
||||
"other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic"
|
||||
],
|
||||
"tests/integration/observability/test_guardrail_effects.py::test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls": [
|
||||
"other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution"
|
||||
],
|
||||
"tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[revoke]": [
|
||||
"other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server"
|
||||
],
|
||||
"tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[expire]": [
|
||||
"other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server"
|
||||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[anonymous]": [
|
||||
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
|
||||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
|
||||
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
|
||||
]
|
||||
},
|
||||
"browser": {
|
||||
|
|
|
|||
|
|
@ -1,28 +0,0 @@
|
|||
# MCP security regression coverage
|
||||
|
||||
[LIT-4506](https://linear.app/litellm-ai/issue/LIT-4506) tracks ten gateway guards and the later JWT/OAuth acceptance. This inventory distinguishes executable assertions from unresolved coverage. A listed test counts as verified only when its exact commit has an executed, passing result
|
||||
|
||||
Run the controlled gateway cases through `python tests/integration/run.py extensions`. They use real HTTP, PostgreSQL, scoped non-master keys and an SDK upstream. The existing runner supplies test entitlement; these tests do not validate licenses or external-provider consent. Canonical nodes and contract IDs live in `../contracts.json`
|
||||
|
||||
| Requested guard | Existing or added coverage | Remaining limitation and owner |
|
||||
| --- | --- | --- |
|
||||
| 1. Discovery scoped by org/team/user/key | `test_mcp_lifecycle.py` checks the exact key-granted catalog and health visibility in both management modes. [PR #38680](https://github.com/BerriAI/litellm/pull/38680) adds team/org/user toolset E2E assertions | Per-principal native MCP coverage is not established by REST results; reuse #38680 rather than duplicate it |
|
||||
| 2. Users cannot attach unauthorized servers to their own keys | Existing live probes are recorded on LIT-4506; they are not durable endpoint regression tests | Own-key create/update escalation and its permission-validator boundary remain on existing management security tickets, including [LIT-4502](https://linear.app/litellm-ai/issue/LIT-4502). A generic route denial does not prove that validator ran |
|
||||
| 3. UI/API permission parity | Existing dashboard tests cover admin operations | The same non-admin actor must be tested through browser and API; admin UI tests do not establish parity. Retained with [LIT-3644](https://linear.app/litellm-ai/issue/LIT-3644) |
|
||||
| 4. Server ID determines identity | `test_mcp_lifecycle.py` grants one of two servers sharing a URL and denies calls to the other, using explicit server IDs for direct REST calls and server-qualified search results for virtual calls, with and without bearer credentials | Virtual calls identify the target by the searched tool name, not the REST `server_id` field. Bare names such as `add` are ambiguous across servers; duplicate aliases/names and unprefixed protocol routing remain with [LIT-4500](https://linear.app/litellm-ai/issue/LIT-4500) |
|
||||
| 5. Same-URL servers do not share credentials | `test_oauth_configuration.py` crosses two gateway users with two server IDs and four distinct stored OAuth tokens. It checks actual upstream headers and successful results, then invalidates only one tuple | Controlled stored-token tests do not prove separate external-provider accounts or consent flows |
|
||||
| 6. OAuth never falls back to anonymous | OAuth isolation variants remove a stored token or expire it without refresh, require separate list/call auth failures and no upstream requests, and preserve all other valid tuples. `test_mcp_lifecycle.py` also covers warm static-header removal and OBO without a caller JWT | External upstream revocation, refresh/reauthorization and aggregate challenges remain with [LIT-4501](https://linear.app/litellm-ai/issue/LIT-4501), [LIT-3433](https://linear.app/litellm-ai/issue/LIT-3433), [LIT-4422](https://linear.app/litellm-ai/issue/LIT-4422) and [LIT-4436](https://linear.app/litellm-ai/issue/LIT-4436) |
|
||||
| 7. Stateful HTTP/session continuity | Legacy public-client tests exercise initialized sessions | No claim here proves upstream session state continuity; retained under [LIT-3143](https://linear.app/litellm-ai/issue/LIT-3143) |
|
||||
| 8. Production guardrails/hooks run | `../observability/test_guardrail_effects.py` checks selected pre-call guards on direct and virtual execution, key/team/request selection, allowed results and zero denied executions. The two legacy test-owned dispatcher files are removed | This does not establish every post-call/output-scanning or concurrent hook contract |
|
||||
| 9. Permissions enforced at discovery and execution | Exact key catalog and virtual search results plus forbidden direct/virtual calls in `test_mcp_lifecycle.py`; existing `../compatibility/test_persisted_toolsets.py` checks tool-level ceiling, denied sibling and allowed control | All principal/transport combinations are not established; link #38680's evidence for its additional principal cases |
|
||||
| 10. Stateless/stateful matrix | These controlled peers use stateless HTTP upstreams | Stateful combinations depend on LIT-3143 and shared conformance runs. Modern-agent/legacy-upstream interaction remains deferred; legacy passes do not establish modern conformance |
|
||||
|
||||
## Additional JWT/OAuth acceptance
|
||||
|
||||
[LIT-3467 / PR #41909](https://github.com/BerriAI/litellm/pull/41909) owns one shared real login/consent, immediate list/call and cold-restart implementation, with aggregate SSO and explicitly configured per-server JWT variants. Reuse that implementation and its protected login secret; do not create another browser bootstrap here. Credit its exact-commit evidence separately from these controlled credential tests
|
||||
|
||||
The two-user/two-server cases here create non-admin users and scoped API keys through management APIs. They store synthetic upstream OAuth credentials through the real credential endpoint and assert the actual bearer at the owned upstream. This deliberately isolates credential lookup, expiry and revocation from consent. No gateway API key may replace the expected upstream token
|
||||
|
||||
Gateway JWT precedence, invalid/expired gateway JWTs, inactive-user denial, and their MCP-specific interaction with isolated credential lookup remain unverified by these API-key cases. General JWT unit/API tests are useful existing coverage but do not substitute for those MCP outcomes. Real-provider auth failures should extend LIT-3467's settled helpers; its explicit-header case must not be described as an uninterrupted Authorization-only OAuth flow
|
||||
|
||||
[PR #41718 / LIT-7737](https://github.com/BerriAI/litellm/pull/41718) owns dependency and public-client compatibility checks. This suite consumes the merged SDK2 API and keeps the existing dependency constraints. Its result must be reported independently of an installation-matrix pass
|
||||
|
|
@ -1,18 +1,14 @@
|
|||
import json
|
||||
import uuid
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from hypothesis import strategies as st
|
||||
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
|
||||
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names
|
||||
|
||||
|
||||
|
|
@ -57,7 +53,7 @@ def test_tool_error_remains_error_and_healthy_sibling_returns_value(gateway: Gat
|
|||
failure: Final = call_tool(gateway, key, identity, names["fail"], {})
|
||||
assert failure.status_code == 200, failure.text
|
||||
assert failure.json()["isError"] is True
|
||||
assert failure.json()["content"][0]["text"] == "Error executing tool fail"
|
||||
assert "synthetic tool failure" in failure.json()["content"][0]["text"]
|
||||
healthy: Final = call_tool(gateway, key, identity, names["multiply"], {"a": 3, "b": 5})
|
||||
assert healthy.status_code == 200, healthy.text
|
||||
assert healthy.json()["isError"] is False
|
||||
|
|
@ -125,182 +121,3 @@ def test_generated_mcp_edits_preserve_actual_headers_and_tool_results(gateway: G
|
|||
self.resources.close()
|
||||
|
||||
run_state_machine_as_test(Servers, settings=LIFECYCLE_SETTINGS)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.mcp.health.restricted_keys_intersect_grants_in_both_modes")
|
||||
def test_health_intersects_route_restricted_key_grants_in_both_management_modes(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
for mode in ("restricted", "view_all"):
|
||||
config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"]["user_mcp_management_mode"] = mode
|
||||
path = tmp_path / f"health-{mode}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
|
||||
mcp_peer() as peer,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
first = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
|
||||
second = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
|
||||
control = scenario.key(object_permission={"mcp_servers": [first]})
|
||||
names = tool_names(candidate, control, first)
|
||||
healthy = call_tool(candidate, control, first, names["add"], {"a": 3, "b": 5})
|
||||
assert healthy.status_code == 200 and healthy.json()["content"][0]["text"] == "8", healthy.text
|
||||
for grants in ([first], [second], []):
|
||||
key = scenario.key(
|
||||
allowed_routes=["/v1/mcp/server", "/v1/mcp/server/health"],
|
||||
object_permission={"mcp_servers": grants},
|
||||
)
|
||||
listed = candidate.request("GET", "/v1/mcp/server", key=key)
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert {row["server_id"] for row in listed.json()} == set(grants), listed.text
|
||||
for requested in (None, [second], [first, second]):
|
||||
response = candidate.client.get(
|
||||
"/v1/mcp/server/health",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
params=[] if requested is None else [("server_ids", identity) for identity in requested],
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
expected = set(grants) if requested is None else set(grants).intersection(requested)
|
||||
assert {row["server_id"] for row in response.json()} == expected, response.text
|
||||
assert all(row["status"] == "healthy" for row in response.json())
|
||||
|
||||
|
||||
@pytest.mark.covers("other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic")
|
||||
def test_warm_credential_removal_rejects_without_upstream_traffic(gateway: Gateway) -> None:
|
||||
with mcp_peer() as peer, gateway.scenario() as scenario:
|
||||
identity = register_mcp(
|
||||
scenario,
|
||||
peer,
|
||||
"credentials" + uuid.uuid4().hex,
|
||||
auth_type="bearer_token",
|
||||
static_headers={"Authorization": "Bearer synthetic-upstream-credential"},
|
||||
)
|
||||
key = scenario.key(object_permission={"mcp_servers": [identity]})
|
||||
names = tool_names(gateway, key, identity)
|
||||
warm = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
||||
assert warm.status_code == 200 and warm.json()["content"][0]["text"] == "8", warm.text
|
||||
calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["headers"][b"authorization"] == b"Bearer synthetic-upstream-credential"
|
||||
removed = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "static_headers": {}})
|
||||
assert removed.status_code == 202, removed.text
|
||||
stored = gateway.request("GET", f"/v1/mcp/server/{identity}")
|
||||
assert stored.status_code == 200, stored.text
|
||||
assert stored.json()["auth_type"] == "bearer_token"
|
||||
assert not stored.json().get("static_headers"), stored.text
|
||||
peer.drain()
|
||||
for operation in ("list", "call"):
|
||||
rejected = (
|
||||
gateway.client.get(
|
||||
"/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": key}
|
||||
)
|
||||
if operation == "list"
|
||||
else call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
||||
)
|
||||
assert rejected.status_code == 500, rejected.text
|
||||
if operation == "list":
|
||||
assert rejected.json()["detail"]["error"] == "internal", rejected.text
|
||||
assert "Failed to list tools from server" in rejected.json()["detail"]["message"], rejected.text
|
||||
else:
|
||||
assert "requires a usable upstream credential" in rejected.text, rejected.text
|
||||
assert peer.drain() == (), "missing static credential escaped to upstream"
|
||||
changed = gateway.request(
|
||||
"PUT",
|
||||
"/v1/mcp/server",
|
||||
{
|
||||
"server_id": identity,
|
||||
"auth_type": "oauth2_token_exchange",
|
||||
"token_exchange_endpoint": peer.url + "/token",
|
||||
"credentials": {"client_id": "synthetic-client"},
|
||||
},
|
||||
)
|
||||
assert changed.status_code == 202, changed.text
|
||||
peer.drain()
|
||||
rejected_subject = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
|
||||
assert rejected_subject.status_code == 401, rejected_subject.text
|
||||
assert peer.drain() == (), "virtual key cannot supply an OBO subject token"
|
||||
control_id = register_mcp(scenario, peer, "control" + uuid.uuid4().hex, auth_type="none")
|
||||
control_key = scenario.key(object_permission={"mcp_servers": [control_id]})
|
||||
control_names = tool_names(gateway, control_key, control_id)
|
||||
control = call_tool(gateway, control_key, control_id, control_names["multiply"], {"a": 3, "b": 5})
|
||||
assert control.status_code == 200 and control.json()["content"][0]["text"] == "15", control.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("authenticated", (False, True), ids=("anonymous", "bearer"))
|
||||
@pytest.mark.covers("other.mcp.permissions.same_url_servers_enforce_discovery_and_execution")
|
||||
def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
|
||||
gateway: Gateway, authenticated: bool
|
||||
) -> None:
|
||||
with mcp_peer() as peer, gateway.scenario() as scenario:
|
||||
aliases: Final = tuple("scope" + uuid.uuid4().hex for _ in range(2))
|
||||
servers: Final = tuple(
|
||||
register_mcp(
|
||||
scenario,
|
||||
peer,
|
||||
alias,
|
||||
auth_type="bearer_token" if authenticated else "none",
|
||||
static_headers={
|
||||
"X-Integration-Server": alias,
|
||||
**({"Authorization": f"Bearer synthetic-{alias}"} if authenticated else {}),
|
||||
},
|
||||
)
|
||||
for alias in aliases
|
||||
)
|
||||
for virtual in (False, True):
|
||||
keys: Final = tuple(
|
||||
scenario.key(object_permission={"mcp_servers": [server], "mcp_tool_search_enabled": virtual})
|
||||
for server in servers
|
||||
)
|
||||
for server, alias, key in zip(servers, aliases, keys):
|
||||
catalog: Final = gateway.request("GET", "/mcp-rest/tools/list", key=key)
|
||||
assert catalog.status_code == 200, catalog.text
|
||||
if virtual:
|
||||
assert {tool["name"] for tool in catalog.json()["tools"]} == {
|
||||
"mcp_tool_search",
|
||||
"mcp_tool_call",
|
||||
"agent_search",
|
||||
"skill_search",
|
||||
}, catalog.text
|
||||
search: Final = gateway.request(
|
||||
"POST",
|
||||
"/mcp-rest/tools/call",
|
||||
{"name": "mcp_tool_search", "arguments": {"query": "add", "top_k": 10}},
|
||||
key=key,
|
||||
)
|
||||
assert search.status_code == 200 and search.json()["isError"] is False, search.text
|
||||
assert [tool["name"] for tool in json.loads(search.json()["content"][0]["text"])] == [
|
||||
f"{alias}-add"
|
||||
], search.text
|
||||
else:
|
||||
assert {tool["mcp_info"]["server_id"] for tool in catalog.json()["tools"]} == {server}
|
||||
assert {tool["name"] for tool in catalog.json()["tools"]} == {"add", "multiply", "fail"}
|
||||
for server_index, caller_index in ((0, 0), (1, 0), (1, 1)):
|
||||
peer.drain()
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/mcp-rest/tools/call",
|
||||
{
|
||||
"name": "mcp_tool_call" if virtual else "add",
|
||||
**({} if virtual else {"server_id": servers[server_index]}),
|
||||
"arguments": (
|
||||
{"tool_name": f"{aliases[server_index]}-add", "arguments": {"a": 3, "b": 5}}
|
||||
if virtual
|
||||
else {"a": 3, "b": 5}
|
||||
),
|
||||
},
|
||||
key=keys[caller_index],
|
||||
)
|
||||
observed: Final = peer.drain()
|
||||
if server_index != caller_index:
|
||||
assert response.status_code == 403 and "not allowed" in response.text, response.text
|
||||
assert observed == (), "forbidden server reached the upstream"
|
||||
continue
|
||||
assert response.status_code == 200 and response.json()["isError"] is False, response.text
|
||||
assert response.json()["content"][0]["text"] == "8", response.text
|
||||
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["headers"][b"x-integration-server"] == aliases[server_index].encode()
|
||||
expected_auth: Final = f"Bearer synthetic-{aliases[server_index]}".encode() if authenticated else None
|
||||
assert all(item["headers"].get(b"authorization") == expected_auth for item in observed)
|
||||
|
|
|
|||
|
|
@ -2,14 +2,14 @@ import json
|
|||
import queue
|
||||
import uuid
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
from typing import Final, Literal
|
||||
from typing import Final
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, tool_names
|
||||
from integration._support.mcp import McpPeer, register_mcp
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
|
|
@ -102,87 +102,3 @@ def test_partial_discovery_and_unrelated_edit_keep_actual_authorization_destinat
|
|||
"PUT", "/v1/mcp/server", {"server_id": identity, "server_name": alias + "renamed"}
|
||||
)
|
||||
assert updated.status_code == 202, updated.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server")
|
||||
@pytest.mark.parametrize("transition", ("revoke", "expire"))
|
||||
def test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server(
|
||||
gateway: Gateway,
|
||||
transition: Literal["revoke", "expire"],
|
||||
) -> None:
|
||||
with mcp_peer() as peer, gateway.scenario() as scenario:
|
||||
servers: Final = tuple(
|
||||
register_mcp(
|
||||
scenario,
|
||||
peer,
|
||||
"oauth" + uuid.uuid4().hex,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
authorization_url=peer.url + "/authorize",
|
||||
token_url=peer.url + "/token",
|
||||
credentials={"client_id": "synthetic-oauth-client"},
|
||||
)
|
||||
for _ in range(2)
|
||||
)
|
||||
users: Final = tuple(scenario.user(user_role="internal_user") for _ in range(2))
|
||||
keys: Final = tuple(
|
||||
scenario.key(user_id=user, object_permission={"mcp_servers": list(servers)}) for user in users
|
||||
)
|
||||
for user_index, key in enumerate(keys):
|
||||
for server_index, server_id in enumerate(servers):
|
||||
stored: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1/mcp/server/{server_id}/oauth-user-credential",
|
||||
{"access_token": f"synthetic-user-{user_index}-server-{server_index}", "expires_in": 3600},
|
||||
key=key,
|
||||
)
|
||||
assert stored.status_code == 200 and stored.json()["has_credential"] is True, stored.text
|
||||
scenario.cleanups.callback(
|
||||
gateway.request,
|
||||
"DELETE",
|
||||
f"/v1/mcp/server/{server_id}/oauth-user-credential",
|
||||
key=key,
|
||||
)
|
||||
names: Final = tuple(tool_names(gateway, keys[0], server) for server in servers)
|
||||
for generation in range(2):
|
||||
for user_index, key in enumerate(keys):
|
||||
for server_index, server_id in enumerate(servers):
|
||||
peer.drain()
|
||||
discovery: Final = gateway.request(
|
||||
"GET",
|
||||
"/mcp-rest/tools/list",
|
||||
key=key,
|
||||
params={"server_id": server_id},
|
||||
)
|
||||
call: Final = call_tool(gateway, key, server_id, names[server_index]["add"], {"a": 3, "b": 5})
|
||||
observed: Final = peer.drain()
|
||||
if generation == 1 and user_index == 0 and server_index == 0:
|
||||
for rejected in (discovery, call):
|
||||
assert rejected.status_code == 401, rejected.text
|
||||
assert rejected.json() == {"detail": "Unauthorized"}, rejected.text
|
||||
assert "resource_metadata=" in rejected.headers["www-authenticate"]
|
||||
assert observed == (), "unusable credentials must not fall back to another user or server"
|
||||
else:
|
||||
assert discovery.status_code == 200, discovery.text
|
||||
assert {tool["name"] for tool in discovery.json()["tools"]} == set(names[server_index].values())
|
||||
assert call.status_code == 200 and call.json()["isError"] is False, call.text
|
||||
assert call.json()["content"][0]["text"] == "8", call.text
|
||||
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
|
||||
assert len(calls) == 1
|
||||
expected: Final = f"Bearer synthetic-user-{user_index}-server-{server_index}".encode()
|
||||
assert calls[0]["headers"][b"authorization"] == expected
|
||||
assert all(item["headers"].get(b"authorization") == expected for item in observed)
|
||||
if generation == 0:
|
||||
changed: Final = gateway.request(
|
||||
"DELETE" if transition == "revoke" else "POST",
|
||||
f"/v1/mcp/server/{servers[0]}/oauth-user-credential",
|
||||
None
|
||||
if transition == "revoke"
|
||||
else {
|
||||
"access_token": "synthetic-expired-user-0-server-0",
|
||||
"expires_in": -60,
|
||||
},
|
||||
key=keys[0],
|
||||
)
|
||||
assert changed.status_code == 200, changed.text
|
||||
assert changed.json()["has_credential"] is (transition == "expire"), changed.text
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ import yaml
|
|||
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.mcp import mcp_peer, register_mcp, tool_names
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
|
|
@ -144,73 +143,3 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa
|
|||
)
|
||||
assert len(observed.get("/__observations").json()["requests"]) == 1
|
||||
assert len(policy.drain()) == 2
|
||||
|
||||
|
||||
@pytest.mark.covers("other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution")
|
||||
def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: Gateway, tmp_path: Path) -> None:
|
||||
guardrail = "mcp-policy-" + uuid.uuid4().hex
|
||||
config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": guardrail,
|
||||
"litellm_params": {
|
||||
"guardrail": "custom_code",
|
||||
"mode": "pre_mcp_call",
|
||||
"default_on": False,
|
||||
"custom_code": (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "add":\n'
|
||||
' return block("integration resolved add denied")\n'
|
||||
" return allow()\n"
|
||||
),
|
||||
},
|
||||
}
|
||||
]
|
||||
path = tmp_path / "mcp-guardrail.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
|
||||
mcp_peer() as peer,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
identity = register_mcp(scenario, peer, "guardrail" + uuid.uuid4().hex)
|
||||
permission = {"mcp_servers": [identity], "mcp_tool_search_enabled": True}
|
||||
key = scenario.key(object_permission=permission)
|
||||
key_selected = scenario.key(object_permission=permission, guardrails=[guardrail])
|
||||
team = scenario.team(guardrails=[guardrail], object_permission={"mcp_servers": [identity]})
|
||||
team_selected = scenario.key(team_id=team, object_permission=permission)
|
||||
catalog_key = scenario.key(object_permission={"mcp_servers": [identity]})
|
||||
names = tool_names(candidate, catalog_key, identity)
|
||||
assert set(names) == {"add", "multiply", "fail"}
|
||||
for virtual in (False, True):
|
||||
for caller, selected, tool, expected in (
|
||||
(key, [], "add", 8),
|
||||
(key, [guardrail], "add", None),
|
||||
(key_selected, [], "add", None),
|
||||
(team_selected, [], "add", None),
|
||||
(key, [guardrail], "multiply", 15),
|
||||
):
|
||||
arguments = {"a": 3, "b": 5}
|
||||
peer.drain()
|
||||
response = candidate.client.post(
|
||||
"/mcp-rest/tools/call",
|
||||
headers={"x-litellm-api-key": caller},
|
||||
json={
|
||||
"server_id": identity,
|
||||
"name": "mcp_tool_call" if virtual else names[tool],
|
||||
"arguments": {"tool_name": names[tool], "arguments": arguments} if virtual else arguments,
|
||||
"guardrails": selected,
|
||||
},
|
||||
)
|
||||
calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
|
||||
if expected is None:
|
||||
assert response.status_code == 400, response.text
|
||||
assert "integration resolved add denied" in response.text, response.text
|
||||
assert calls == (), "pre-call denial must prevent upstream execution"
|
||||
else:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["isError"] is False
|
||||
assert response.json()["content"][0]["text"] == str(expected), response.text
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["body"]["params"]["name"] == tool
|
||||
assert calls[0]["body"]["params"]["arguments"] == arguments
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""Deterministic upstream MCP server for the mcp e2e suite.
|
||||
|
||||
A tiny MCP server exposing `add` and `multiply` over streamable-http so the
|
||||
A tiny FastMCP server exposing `add` and `multiply` over streamable-http so the
|
||||
suite has a self-hosted, offline upstream to register and exercise. DNS-rebinding
|
||||
protection is turned off because the litellm container reaches this over the
|
||||
compose network by service name (`mcp-upstream:8090`), not localhost, and the
|
||||
|
|
@ -9,10 +9,15 @@ stack is an isolated throwaway. Bind host/port come from MCP_HOST/MCP_PORT.
|
|||
|
||||
import os
|
||||
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
|
||||
mcp: MCPServer = MCPServer("e2e-math")
|
||||
mcp: FastMCP = FastMCP(
|
||||
"e2e-math",
|
||||
host=os.getenv("MCP_HOST", "0.0.0.0"),
|
||||
port=int(os.getenv("MCP_PORT", "8090")),
|
||||
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
|
|
@ -28,12 +33,7 @@ def multiply(a: int, b: int) -> int:
|
|||
|
||||
|
||||
def main() -> None:
|
||||
mcp.run(
|
||||
transport="streamable-http",
|
||||
host=os.getenv("MCP_HOST", "0.0.0.0"),
|
||||
port=int(os.getenv("MCP_PORT", "8090")),
|
||||
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
|
||||
)
|
||||
mcp.run(transport="streamable-http")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
770
tests/mcp_tests/test_mcp_guardrails.py
Normal file
770
tests/mcp_tests/test_mcp_guardrails.py
Normal file
|
|
@ -0,0 +1,770 @@
|
|||
"""
|
||||
Test file for MCP Guardrails Feature
|
||||
|
||||
This file tests the MCP guardrails functionality for both pre and during MCP call hooks,
|
||||
including various guardrail types and proper exception handling.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from typing import Optional, Dict, Any
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
# Add the project root to the path
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.types.mcp import (
|
||||
MCPPreCallRequestObject,
|
||||
MCPPreCallResponseObject,
|
||||
MCPDuringCallRequestObject,
|
||||
MCPDuringCallResponseObject,
|
||||
)
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
class MockPiiGuardrail(CustomGuardrail):
|
||||
"""Mock PII guardrail that raises BlockedPiiEntityError"""
|
||||
|
||||
def __init__(self, should_block: bool = True, entity_type: str = "EMAIL_ADDRESS"):
|
||||
super().__init__()
|
||||
self.should_block = should_block
|
||||
self.entity_type = entity_type
|
||||
self.guardrail_name = "mock-pii-guardrail"
|
||||
self.call_count = 0
|
||||
|
||||
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
|
||||
"""Always run for testing"""
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
"""Mock pre-call hook that raises BlockedPiiEntityError"""
|
||||
self.call_count += 1
|
||||
|
||||
if self.should_block:
|
||||
raise BlockedPiiEntityError(
|
||||
entity_type=self.entity_type,
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MockContentGuardrail(CustomGuardrail):
|
||||
"""Mock content guardrail that raises GuardrailRaisedException"""
|
||||
|
||||
def __init__(self, should_block: bool = True):
|
||||
super().__init__()
|
||||
self.should_block = should_block
|
||||
self.guardrail_name = "mock-content-guardrail"
|
||||
self.call_count = 0
|
||||
|
||||
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
|
||||
"""Always run for testing"""
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
"""Mock pre-call hook that raises GuardrailRaisedException"""
|
||||
self.call_count += 1
|
||||
|
||||
if self.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name, message="Content violates policy"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MockHttpGuardrail(CustomGuardrail):
|
||||
"""Mock HTTP guardrail that raises HTTPException"""
|
||||
|
||||
def __init__(self, should_block: bool = True):
|
||||
super().__init__()
|
||||
self.should_block = should_block
|
||||
self.guardrail_name = "mock-http-guardrail"
|
||||
self.call_count = 0
|
||||
|
||||
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
|
||||
"""Always run for testing"""
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
"""Mock pre-call hook that raises HTTPException"""
|
||||
self.call_count += 1
|
||||
|
||||
if self.should_block:
|
||||
raise HTTPException(
|
||||
status_code=400, detail={"error": "Violated guardrail policy"}
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MockDuringCallGuardrail(CustomGuardrail):
|
||||
"""Mock guardrail for during-call testing"""
|
||||
|
||||
def __init__(self, should_block: bool = True):
|
||||
super().__init__()
|
||||
self.should_block = should_block
|
||||
self.guardrail_name = "mock-during-guardrail"
|
||||
self.call_count = 0
|
||||
|
||||
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
|
||||
"""Always run for testing"""
|
||||
return True
|
||||
|
||||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: str,
|
||||
):
|
||||
"""Mock during-call hook that raises exceptions"""
|
||||
self.call_count += 1
|
||||
|
||||
if self.should_block:
|
||||
raise BlockedPiiEntityError(
|
||||
entity_type="PHONE_NUMBER",
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MockProxyLogging:
|
||||
"""Mock proxy logging object for testing MCP guardrails"""
|
||||
|
||||
def __init__(self, guardrails: Optional[list] = None):
|
||||
self.guardrails = guardrails if guardrails is not None else []
|
||||
self.call_details = {"user_api_key_cache": DualCache()}
|
||||
self.dynamic_success_callbacks = []
|
||||
self.call_count = 0
|
||||
|
||||
def get_combined_callback_list(self, dynamic_success_callbacks, global_callbacks):
|
||||
"""Return the guardrails for testing"""
|
||||
return self.guardrails
|
||||
|
||||
def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict:
|
||||
"""Convert MCP tool call to LLM message format"""
|
||||
tool_call_content = (
|
||||
f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
||||
)
|
||||
|
||||
return {
|
||||
"messages": [{"role": "user", "content": tool_call_content}],
|
||||
"model": kwargs.get("model", "mcp-tool-call"),
|
||||
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
|
||||
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
|
||||
}
|
||||
|
||||
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj):
|
||||
"""Convert LLM result back to MCP response format"""
|
||||
return None # For testing, we don't need to convert back
|
||||
|
||||
def _parse_pre_mcp_call_hook_response(self, response, original_request):
|
||||
"""Parse pre MCP call hook response"""
|
||||
return response
|
||||
|
||||
async def async_pre_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict,
|
||||
request_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> Optional[Any]:
|
||||
"""Mock pre MCP tool call hook"""
|
||||
self.call_count += 1
|
||||
|
||||
# Simulate the actual hook logic
|
||||
for guardrail in self.guardrails:
|
||||
if isinstance(guardrail, CustomGuardrail):
|
||||
try:
|
||||
synthetic_data = self._convert_mcp_to_llm_format(
|
||||
request_obj, kwargs
|
||||
)
|
||||
|
||||
# Check if guardrail should run
|
||||
if not guardrail.should_run_guardrail(
|
||||
synthetic_data, GuardrailEventHooks.pre_mcp_call
|
||||
):
|
||||
continue
|
||||
|
||||
result = await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=kwargs.get("user_api_key_auth"),
|
||||
cache=self.call_details["user_api_key_cache"],
|
||||
data=synthetic_data,
|
||||
call_type="mcp_call",
|
||||
)
|
||||
if result is not None:
|
||||
return self._parse_pre_mcp_call_hook_response(
|
||||
result, request_obj
|
||||
)
|
||||
except (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
HTTPException,
|
||||
) as e:
|
||||
# Re-raise guardrail exceptions
|
||||
raise e
|
||||
except Exception as e:
|
||||
# Log non-guardrail exceptions as non-blocking
|
||||
print(
|
||||
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
async def async_during_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict,
|
||||
request_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> Optional[Any]:
|
||||
"""Mock during MCP tool call hook"""
|
||||
self.call_count += 1
|
||||
|
||||
# Simulate the actual hook logic
|
||||
for guardrail in self.guardrails:
|
||||
if isinstance(guardrail, CustomGuardrail):
|
||||
try:
|
||||
synthetic_data = self._convert_mcp_to_llm_format(
|
||||
request_obj, kwargs
|
||||
)
|
||||
result = await guardrail.async_moderation_hook(
|
||||
data=synthetic_data,
|
||||
user_api_key_dict=kwargs.get("user_api_key_auth"),
|
||||
call_type="mcp_call",
|
||||
)
|
||||
if result is not None:
|
||||
return result
|
||||
except (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
HTTPException,
|
||||
) as e:
|
||||
# Re-raise guardrail exceptions
|
||||
raise e
|
||||
except Exception as e:
|
||||
# Log non-guardrail exceptions as non-blocking
|
||||
print(
|
||||
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user_api_key():
|
||||
"""Mock user API key for testing"""
|
||||
return UserAPIKeyAuth(api_key="test_key", user_id="test_user")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_cache():
|
||||
"""Mock cache for testing"""
|
||||
return DualCache()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_pii_guardrail():
|
||||
"""Mock PII guardrail that blocks"""
|
||||
return MockPiiGuardrail(should_block=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_pii_guardrail_allow():
|
||||
"""Mock PII guardrail that allows"""
|
||||
return MockPiiGuardrail(should_block=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_content_guardrail():
|
||||
"""Mock content guardrail that blocks"""
|
||||
return MockContentGuardrail(should_block=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_http_guardrail():
|
||||
"""Mock HTTP guardrail that blocks"""
|
||||
return MockHttpGuardrail(should_block=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_during_guardrail():
|
||||
"""Mock during-call guardrail that blocks"""
|
||||
return MockDuringCallGuardrail(should_block=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_proxy_logging():
|
||||
"""Mock proxy logging object"""
|
||||
return MockProxyLogging()
|
||||
|
||||
|
||||
class TestMCPGuardrailsPreCall:
|
||||
"""Test MCP guardrails for pre-call hooks"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pii_guardrail_blocks_pre_call(
|
||||
self, mock_pii_guardrail, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test that PII guardrail properly blocks pre-call"""
|
||||
proxy_logging = MockProxyLogging([mock_pii_guardrail])
|
||||
|
||||
# Create MCP request
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="email_tool",
|
||||
arguments={"email": "test@example.com"},
|
||||
server_name="email_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "email_tool",
|
||||
"arguments": {"email": "test@example.com"},
|
||||
"server_name": "email_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that BlockedPiiEntityError is raised
|
||||
with pytest.raises(BlockedPiiEntityError) as excinfo:
|
||||
await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Verify the error details
|
||||
assert excinfo.value.entity_type == "EMAIL_ADDRESS"
|
||||
assert excinfo.value.guardrail_name == "mock-pii-guardrail"
|
||||
assert mock_pii_guardrail.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pii_guardrail_allows_pre_call(
|
||||
self, mock_pii_guardrail_allow, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test that PII guardrail allows pre-call when configured to allow"""
|
||||
proxy_logging = MockProxyLogging([mock_pii_guardrail_allow])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="email_tool",
|
||||
arguments={"email": "test@example.com"},
|
||||
server_name="email_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "email_tool",
|
||||
"arguments": {"email": "test@example.com"},
|
||||
"server_name": "email_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that no exception is raised
|
||||
result = await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert mock_pii_guardrail_allow.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_content_guardrail_blocks_pre_call(
|
||||
self, mock_content_guardrail, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test that content guardrail properly blocks pre-call"""
|
||||
proxy_logging = MockProxyLogging([mock_content_guardrail])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="content_tool",
|
||||
arguments={"content": "sensitive content"},
|
||||
server_name="content_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "content_tool",
|
||||
"arguments": {"content": "sensitive content"},
|
||||
"server_name": "content_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that GuardrailRaisedException is raised
|
||||
with pytest.raises(GuardrailRaisedException) as excinfo:
|
||||
await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Verify the error details
|
||||
assert "Content violates policy" in str(excinfo.value)
|
||||
assert excinfo.value.guardrail_name == "mock-content-guardrail"
|
||||
assert mock_content_guardrail.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_guardrail_blocks_pre_call(
|
||||
self, mock_http_guardrail, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test that HTTP guardrail properly blocks pre-call"""
|
||||
proxy_logging = MockProxyLogging([mock_http_guardrail])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="http_tool",
|
||||
arguments={"url": "http://example.com"},
|
||||
server_name="http_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "http_tool",
|
||||
"arguments": {"url": "http://example.com"},
|
||||
"server_name": "http_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that HTTPException is raised
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Verify the error details
|
||||
assert excinfo.value.status_code == 400
|
||||
assert "Violated guardrail policy" in str(excinfo.value.detail)
|
||||
assert mock_http_guardrail.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_guardrails_pre_call(
|
||||
self, mock_pii_guardrail, mock_content_guardrail, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test multiple guardrails - first one should block"""
|
||||
proxy_logging = MockProxyLogging([mock_pii_guardrail, mock_content_guardrail])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="test_tool",
|
||||
arguments={"email": "test@example.com"},
|
||||
server_name="test_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "test_tool",
|
||||
"arguments": {"email": "test@example.com"},
|
||||
"server_name": "test_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that first guardrail blocks
|
||||
with pytest.raises(BlockedPiiEntityError):
|
||||
await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Verify only first guardrail was called
|
||||
assert mock_pii_guardrail.call_count == 1
|
||||
assert mock_content_guardrail.call_count == 0
|
||||
|
||||
|
||||
class TestMCPGuardrailsDuringCall:
|
||||
"""Test MCP guardrails for during-call hooks"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_guardrail_blocks(
|
||||
self, mock_during_guardrail, mock_user_api_key, mock_cache
|
||||
):
|
||||
"""Test that during-call guardrail properly blocks execution"""
|
||||
proxy_logging = MockProxyLogging([mock_during_guardrail])
|
||||
|
||||
request_obj = MCPDuringCallRequestObject(
|
||||
tool_name="phone_tool",
|
||||
arguments={"phone": "555-123-4567"},
|
||||
server_name="phone_server",
|
||||
start_time=datetime.now().timestamp(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "phone_tool",
|
||||
"arguments": {"phone": "555-123-4567"},
|
||||
"server_name": "phone_server",
|
||||
}
|
||||
|
||||
# Test that BlockedPiiEntityError is raised
|
||||
with pytest.raises(BlockedPiiEntityError) as excinfo:
|
||||
await proxy_logging.async_during_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Verify the error details
|
||||
assert excinfo.value.entity_type == "PHONE_NUMBER"
|
||||
assert excinfo.value.guardrail_name == "mock-during-guardrail"
|
||||
assert mock_during_guardrail.call_count == 1
|
||||
|
||||
|
||||
class TestMCPGuardrailsIntegration:
|
||||
"""Test MCP guardrails integration with MCP server manager"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_server_manager_with_guardrails(self):
|
||||
"""Test MCP server manager with guardrail integration"""
|
||||
|
||||
mock_proxy_logging = MockProxyLogging([MockPiiGuardrail(should_block=True)])
|
||||
|
||||
# Test that guardrail exception is properly raised in the hook
|
||||
with pytest.raises(BlockedPiiEntityError):
|
||||
await mock_proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs={
|
||||
"name": "email_tool",
|
||||
"arguments": {"email": "test@example.com"},
|
||||
},
|
||||
request_obj=MagicMock(),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_exception_propagation(self):
|
||||
"""Test that guardrail exceptions properly propagate through the system"""
|
||||
# Test BlockedPiiEntityError
|
||||
with pytest.raises(BlockedPiiEntityError):
|
||||
raise BlockedPiiEntityError(
|
||||
entity_type="EMAIL_ADDRESS", guardrail_name="test-guardrail"
|
||||
)
|
||||
|
||||
# Test GuardrailRaisedException
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name="test-guardrail", message="Test message"
|
||||
)
|
||||
|
||||
# Test HTTPException
|
||||
with pytest.raises(HTTPException):
|
||||
raise HTTPException(status_code=400, detail={"error": "Test error"})
|
||||
|
||||
|
||||
class TestMCPGuardrailsErrorHandling:
|
||||
"""Test MCP guardrails error handling scenarios"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_guardrail_exception_logging(self, mock_user_api_key, mock_cache):
|
||||
"""Test that non-guardrail exceptions are logged as non-blocking"""
|
||||
|
||||
class MockFailingGuardrail(CustomGuardrail):
|
||||
def should_run_guardrail(
|
||||
self, data: dict, event_type: GuardrailEventHooks
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
raise Exception("Non-guardrail error")
|
||||
|
||||
proxy_logging = MockProxyLogging([MockFailingGuardrail()])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="test_tool",
|
||||
arguments={"test": "data"},
|
||||
server_name="test_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "test_tool",
|
||||
"arguments": {"test": "data"},
|
||||
"server_name": "test_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that non-guardrail exceptions are handled gracefully
|
||||
result = await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Should return None (not raise exception)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_should_not_run(self, mock_user_api_key, mock_cache):
|
||||
"""Test that guardrails don't run when should_run_guardrail returns False"""
|
||||
|
||||
class MockConditionalGuardrail(CustomGuardrail):
|
||||
def should_run_guardrail(
|
||||
self, data: dict, event_type: GuardrailEventHooks
|
||||
) -> bool:
|
||||
return False # Don't run
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail")
|
||||
|
||||
proxy_logging = MockProxyLogging([MockConditionalGuardrail()])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="test_tool",
|
||||
arguments={"test": "data"},
|
||||
server_name="test_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "test_tool",
|
||||
"arguments": {"test": "data"},
|
||||
"server_name": "test_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Test that guardrail doesn't run and no exception is raised
|
||||
result = await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Should return None (guardrail didn't run)
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestMCPGuardrailsEdgeCases:
|
||||
"""Test MCP guardrails edge cases and error conditions"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_guardrails_list(self, mock_user_api_key, mock_cache):
|
||||
"""Test behavior with empty guardrails list"""
|
||||
proxy_logging = MockProxyLogging([]) # No guardrails
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="test_tool",
|
||||
arguments={"test": "data"},
|
||||
server_name="test_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "test_tool",
|
||||
"arguments": {"test": "data"},
|
||||
"server_name": "test_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Should return None without any issues
|
||||
result = await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_with_invalid_data(self, mock_user_api_key, mock_cache):
|
||||
"""Test guardrail behavior with invalid data"""
|
||||
|
||||
class MockInvalidDataGuardrail(CustomGuardrail):
|
||||
def should_run_guardrail(
|
||||
self, data: dict, event_type: GuardrailEventHooks
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
# Try to access invalid data
|
||||
invalid_data = data.get("invalid_key", {})
|
||||
if invalid_data.get("should_fail"):
|
||||
raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail")
|
||||
return None
|
||||
|
||||
proxy_logging = MockProxyLogging([MockInvalidDataGuardrail()])
|
||||
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="test_tool",
|
||||
arguments={"test": "data"},
|
||||
server_name="test_server",
|
||||
user_api_key_auth=mock_user_api_key.model_dump(),
|
||||
hidden_params=HiddenParams(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"name": "test_tool",
|
||||
"arguments": {"test": "data"},
|
||||
"server_name": "test_server",
|
||||
"user_api_key_auth": mock_user_api_key,
|
||||
}
|
||||
|
||||
# Should handle invalid data gracefully
|
||||
result = await proxy_logging.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
475
tests/mcp_tests/test_mcp_hooks.py
Normal file
475
tests/mcp_tests/test_mcp_hooks.py
Normal file
|
|
@ -0,0 +1,475 @@
|
|||
"""
|
||||
Test file for MCP Hook Architecture
|
||||
|
||||
This file demonstrates the new MCP hook system with comprehensive examples
|
||||
and validation tests.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.mcp import (
|
||||
MCPPreCallRequestObject,
|
||||
MCPPreCallResponseObject,
|
||||
MCPDuringCallRequestObject,
|
||||
MCPDuringCallResponseObject,
|
||||
MCPPostCallResponseObject,
|
||||
)
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
|
||||
|
||||
class TestMCPAccessControlHook(CustomLogger):
|
||||
"""Test hook for access control functionality"""
|
||||
|
||||
def __init__(self):
|
||||
self.allowed_tools = {"github/create_issue", "zapier/send_email"}
|
||||
self.blocked_users = {"user123", "user456"}
|
||||
self.call_count = 0
|
||||
|
||||
async def async_pre_mcp_tool_call_hook(
|
||||
self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time
|
||||
) -> Optional[MCPPreCallResponseObject]:
|
||||
"""Test access control validation"""
|
||||
self.call_count += 1
|
||||
|
||||
tool_name = request_obj.tool_name
|
||||
user_id = kwargs.get("user_api_key_auth", {}).get("user_id")
|
||||
|
||||
# Check if user is blocked
|
||||
if user_id in self.blocked_users:
|
||||
return MCPPreCallResponseObject(
|
||||
should_proceed=False,
|
||||
error_message=f"User {user_id} is not authorized to use MCP tools",
|
||||
)
|
||||
|
||||
# Check if tool is allowed
|
||||
if tool_name not in self.allowed_tools:
|
||||
return MCPPreCallResponseObject(
|
||||
should_proceed=False,
|
||||
error_message=f"Tool {tool_name} is not authorized",
|
||||
)
|
||||
|
||||
return None # Allow execution to proceed
|
||||
|
||||
|
||||
class TestMCPCostTrackingHook(CustomLogger):
|
||||
"""Test hook for cost tracking functionality"""
|
||||
|
||||
def __init__(self):
|
||||
self.cost_map = {
|
||||
"github/create_issue": 0.10,
|
||||
"zapier/send_email": 0.05,
|
||||
"default": 0.01,
|
||||
}
|
||||
self.call_count = 0
|
||||
|
||||
async def async_post_mcp_tool_call_hook(
|
||||
self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time
|
||||
) -> Optional[MCPPostCallResponseObject]:
|
||||
"""Test cost calculation after tool execution"""
|
||||
self.call_count += 1
|
||||
|
||||
tool_name = kwargs.get("name", "")
|
||||
cost = self.cost_map.get(tool_name, self.cost_map["default"])
|
||||
|
||||
# Set the response cost
|
||||
response_obj.hidden_params.response_cost = cost
|
||||
|
||||
return response_obj
|
||||
|
||||
|
||||
class TestMCPMonitoringHook(CustomLogger):
|
||||
"""Test hook for real-time monitoring functionality"""
|
||||
|
||||
def __init__(self):
|
||||
self.max_execution_time = 30.0 # seconds
|
||||
self.call_count = 0
|
||||
|
||||
async def async_during_mcp_tool_call_hook(
|
||||
self, kwargs, request_obj: MCPDuringCallRequestObject, start_time, end_time
|
||||
) -> Optional[MCPDuringCallResponseObject]:
|
||||
"""Test execution time monitoring"""
|
||||
self.call_count += 1
|
||||
|
||||
tool_name = request_obj.tool_name
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Check if execution is taking too long
|
||||
if execution_time > self.max_execution_time:
|
||||
return MCPDuringCallResponseObject(
|
||||
should_continue=False,
|
||||
error_message=f"Tool {tool_name} execution timeout after {execution_time}s",
|
||||
)
|
||||
|
||||
return None # Allow execution to continue
|
||||
|
||||
|
||||
class TestMCPArgumentValidationHook(CustomLogger):
|
||||
"""Test hook for argument validation functionality"""
|
||||
|
||||
def __init__(self):
|
||||
self.call_count = 0
|
||||
|
||||
async def async_pre_mcp_tool_call_hook(
|
||||
self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time
|
||||
) -> Optional[MCPPreCallResponseObject]:
|
||||
"""Test argument validation and sanitization"""
|
||||
self.call_count += 1
|
||||
|
||||
tool_name = request_obj.tool_name
|
||||
arguments = request_obj.arguments.copy() # Create a copy to modify
|
||||
|
||||
# Example: Validate GitHub issue creation
|
||||
if tool_name == "github/create_issue":
|
||||
if not arguments.get("title"):
|
||||
return MCPPreCallResponseObject(
|
||||
should_proceed=False, error_message="GitHub issue title is required"
|
||||
)
|
||||
|
||||
# Sanitize the title
|
||||
title = arguments["title"]
|
||||
if len(title) > 100:
|
||||
title = title[:97] + "..."
|
||||
arguments["title"] = title
|
||||
|
||||
# Example: Validate email sending
|
||||
elif tool_name == "zapier/send_email":
|
||||
if not arguments.get("to"):
|
||||
return MCPPreCallResponseObject(
|
||||
should_proceed=False, error_message="Email recipient is required"
|
||||
)
|
||||
|
||||
return MCPPreCallResponseObject(
|
||||
should_proceed=True, modified_arguments=arguments
|
||||
)
|
||||
|
||||
|
||||
# Test fixtures
|
||||
@pytest.fixture
|
||||
def access_control_hook():
|
||||
return TestMCPAccessControlHook()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cost_tracking_hook():
|
||||
return TestMCPCostTrackingHook()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def monitoring_hook():
|
||||
return TestMCPMonitoringHook()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def argument_validation_hook():
|
||||
return TestMCPArgumentValidationHook()
|
||||
|
||||
|
||||
# Test cases
|
||||
class TestMCPHooks:
|
||||
"""Test cases for MCP hook functionality"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_access_control_hook_allowed_tool(self, access_control_hook):
|
||||
"""Test that allowed tools pass validation"""
|
||||
kwargs = {
|
||||
"user_api_key_auth": {"user_id": "user789"},
|
||||
"name": "github/create_issue",
|
||||
}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue",
|
||||
arguments={"title": "Test issue"},
|
||||
user_api_key_auth={"user_id": "user789"},
|
||||
)
|
||||
|
||||
result = await access_control_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is None # Should allow execution
|
||||
assert access_control_hook.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_access_control_hook_blocked_user(self, access_control_hook):
|
||||
"""Test that blocked users are rejected"""
|
||||
kwargs = {
|
||||
"user_api_key_auth": {"user_id": "user123"},
|
||||
"name": "github/create_issue",
|
||||
}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue",
|
||||
arguments={"title": "Test issue"},
|
||||
user_api_key_auth={"user_id": "user123"},
|
||||
)
|
||||
|
||||
result = await access_control_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is False
|
||||
assert "not authorized" in result.error_message
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_access_control_hook_unauthorized_tool(self, access_control_hook):
|
||||
"""Test that unauthorized tools are rejected"""
|
||||
kwargs = {
|
||||
"user_api_key_auth": {"user_id": "user789"},
|
||||
"name": "unauthorized_tool",
|
||||
}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="unauthorized_tool",
|
||||
arguments={"param": "value"},
|
||||
user_api_key_auth={"user_id": "user789"},
|
||||
)
|
||||
|
||||
result = await access_control_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is False
|
||||
assert "not authorized" in result.error_message
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cost_tracking_hook(self, cost_tracking_hook):
|
||||
"""Test cost tracking functionality"""
|
||||
kwargs = {"name": "github/create_issue"}
|
||||
response_obj = MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[], hidden_params=HiddenParams()
|
||||
)
|
||||
|
||||
result = await cost_tracking_hook.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.hidden_params.response_cost == 0.10
|
||||
assert cost_tracking_hook.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cost_tracking_hook_default_cost(self, cost_tracking_hook):
|
||||
"""Test default cost assignment"""
|
||||
kwargs = {"name": "unknown_tool"}
|
||||
response_obj = MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[], hidden_params=HiddenParams()
|
||||
)
|
||||
|
||||
result = await cost_tracking_hook.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.hidden_params.response_cost == 0.01 # Default cost
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitoring_hook_normal_execution(self, monitoring_hook):
|
||||
"""Test monitoring hook with normal execution time"""
|
||||
kwargs = {"name": "test_tool"}
|
||||
request_obj = MCPDuringCallRequestObject(
|
||||
tool_name="test_tool", arguments={}, start_time=datetime.now().timestamp()
|
||||
)
|
||||
|
||||
result = await monitoring_hook.async_during_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is None # Should allow execution to continue
|
||||
assert monitoring_hook.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_argument_validation_hook_valid_github_issue(
|
||||
self, argument_validation_hook
|
||||
):
|
||||
"""Test argument validation for valid GitHub issue"""
|
||||
kwargs = {"name": "github/create_issue"}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue", arguments={"title": "Valid issue title"}
|
||||
)
|
||||
|
||||
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is True
|
||||
assert result.modified_arguments == {"title": "Valid issue title"}
|
||||
assert argument_validation_hook.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_argument_validation_hook_missing_title(
|
||||
self, argument_validation_hook
|
||||
):
|
||||
"""Test argument validation for missing GitHub issue title"""
|
||||
kwargs = {"name": "github/create_issue"}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue", arguments={} # Missing title
|
||||
)
|
||||
|
||||
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is False
|
||||
assert "title is required" in result.error_message
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_argument_validation_hook_long_title_sanitization(
|
||||
self, argument_validation_hook
|
||||
):
|
||||
"""Test argument validation with title sanitization"""
|
||||
kwargs = {"name": "github/create_issue"}
|
||||
long_title = "A" * 150 # Very long title
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue", arguments={"title": long_title}
|
||||
)
|
||||
|
||||
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is True
|
||||
assert len(result.modified_arguments["title"]) == 100 # Truncated
|
||||
assert result.modified_arguments["title"].endswith("...")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_argument_validation_hook_email_validation(
|
||||
self, argument_validation_hook
|
||||
):
|
||||
"""Test argument validation for email sending"""
|
||||
kwargs = {"name": "zapier/send_email"}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="zapier/send_email",
|
||||
arguments={"to": "test@example.com", "subject": "Test"},
|
||||
)
|
||||
|
||||
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is True
|
||||
assert result.modified_arguments == {
|
||||
"to": "test@example.com",
|
||||
"subject": "Test",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_argument_validation_hook_missing_email_recipient(
|
||||
self, argument_validation_hook
|
||||
):
|
||||
"""Test argument validation for missing email recipient"""
|
||||
kwargs = {"name": "zapier/send_email"}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="zapier/send_email",
|
||||
arguments={"subject": "Test"}, # Missing 'to' field
|
||||
)
|
||||
|
||||
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.should_proceed is False
|
||||
assert "recipient is required" in result.error_message
|
||||
|
||||
|
||||
# Integration test
|
||||
class TestMCPHookIntegration:
|
||||
"""Integration tests for MCP hook system"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_chain_execution(self):
|
||||
"""Test that multiple hooks can work together"""
|
||||
access_hook = TestMCPAccessControlHook()
|
||||
cost_hook = TestMCPCostTrackingHook()
|
||||
validation_hook = TestMCPArgumentValidationHook()
|
||||
|
||||
# Test data
|
||||
kwargs = {
|
||||
"user_api_key_auth": {"user_id": "user789"},
|
||||
"name": "github/create_issue",
|
||||
}
|
||||
request_obj = MCPPreCallRequestObject(
|
||||
tool_name="github/create_issue",
|
||||
arguments={"title": "Integration test issue"},
|
||||
user_api_key_auth={"user_id": "user789"},
|
||||
)
|
||||
|
||||
# Execute pre-hooks
|
||||
access_result = await access_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
validation_result = await validation_hook.async_pre_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
request_obj=request_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Both hooks should allow execution
|
||||
assert access_result is None
|
||||
assert validation_result is not None
|
||||
assert validation_result.should_proceed is True
|
||||
|
||||
# Simulate post-hook execution
|
||||
response_obj = MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[], hidden_params=HiddenParams()
|
||||
)
|
||||
|
||||
cost_result = await cost_hook.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert cost_result is not None
|
||||
assert cost_result.hidden_params.response_cost == 0.10
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
Loading…
Add table
Reference in a new issue