mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
test(mcp): enforce security regression contracts through live gateway
This commit is contained in:
Joshua Valluru
parent
02a20fe264
commit
15b45839e1
5 changed files with 179 additions and 1245 deletions
|
|
@ -213,6 +213,15 @@
|
|||
],
|
||||
"tests/integration/sdk/test_http2_wire.py::test_sync_handler_negotiates_http2_only_when_enabled": [
|
||||
"other.sdk_wire.http2.sync_handler_negotiates_h2_only_when_enabled"
|
||||
],
|
||||
"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"
|
||||
]
|
||||
},
|
||||
"browser": {
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
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
|
||||
|
||||
|
||||
|
|
@ -121,3 +124,100 @@ 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)
|
||||
owned = {first, second}
|
||||
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()}.intersection(owned) == set(grants)
|
||||
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()}.intersection(owned) == expected, response.text
|
||||
assert all(row["status"] == "healthy" for row in response.json() if row["server_id"] in owned)
|
||||
|
||||
|
||||
@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
|
||||
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
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ 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
|
||||
|
||||
|
|
@ -143,3 +144,72 @@ 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])
|
||||
team_selected = scenario.key(team_id=team, object_permission=permission)
|
||||
names = tool_names(candidate, 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,770 +0,0 @@
|
|||
"""
|
||||
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__])
|
||||
|
|
@ -1,475 +0,0 @@
|
|||
"""
|
||||
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