diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 91b1bd86954..1520a9488a5 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -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": { diff --git a/tests/integration/mcp/test_mcp_lifecycle.py b/tests/integration/mcp/test_mcp_lifecycle.py index 7ded23794be..946a3eaae75 100644 --- a/tests/integration/mcp/test_mcp_lifecycle.py +++ b/tests/integration/mcp/test_mcp_lifecycle.py @@ -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 diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 645af77526f..cd44c06cd82 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -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 diff --git a/tests/mcp_tests/test_mcp_guardrails.py b/tests/mcp_tests/test_mcp_guardrails.py deleted file mode 100644 index 04401992449..00000000000 --- a/tests/mcp_tests/test_mcp_guardrails.py +++ /dev/null @@ -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__]) diff --git a/tests/mcp_tests/test_mcp_hooks.py b/tests/mcp_tests/test_mcp_hooks.py deleted file mode 100644 index 6dac7da6d07..00000000000 --- a/tests/mcp_tests/test_mcp_hooks.py +++ /dev/null @@ -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"])