diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 116bcc715d4..d8ac46a1bdf 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1779,6 +1779,7 @@ class MCPServerManager: "gmail_send_email": "zapier_mcp_server", } """ + self._listed_tools_by_server_id: dict[str, dict[str, MCPTool]] = {} self._upstream_initialize_instructions_by_server_id: dict[str, str] = {} # Per-server monotonic timestamp of last upstream prefetch attempt (success, # empty result, or failure). Used to throttle re-probes for servers that do @@ -4239,18 +4240,18 @@ class MCPServerManager: # applied (e.g. "test_petstore-getinventory"). Do NOT pass them # through _create_prefixed_tools — that would add the prefix a second # time producing "test_petstore-test_petstore-getinventory". - if not add_prefix: - prefix: Final = get_server_prefix(server) - sep: Final = MCP_TOOL_PREFIX_SEPARATOR - tools = [ - ( - t.model_copy(update={"name": t.name[len(prefix) + len(sep) :]}) - if t.name.startswith(f"{prefix}{sep}") - else t - ) - for t in tools - ] - return tools + prefix: Final = get_server_prefix(server) + sep: Final = MCP_TOOL_PREFIX_SEPARATOR + bare_tools: Final = [ + ( + t.model_copy(update={"name": t.name[len(prefix) + len(sep) :]}) + if t.name.startswith(f"{prefix}{sep}") + else t + ) + for t in tools + ] + self._listed_tools_by_server_id[server.server_id] = {t.name: t for t in bare_tools} + return tools if add_prefix else bare_tools else: tools = await self._fetch_tools_with_timeout(client, server.name) self._remember_upstream_initialize_instructions(server, client) @@ -5131,9 +5132,14 @@ class MCPServerManager: for spelling in iter_known_tool_name_spellings(original_name, server): self.tool_name_to_mcp_server_name_mapping[spelling] = prefix + self._listed_tools_by_server_id[server.server_id] = {tool.name: tool for tool in tools} verbose_logger.info("Successfully fetched %s tools from server %s", len(prefixed_tools), server.name) return prefixed_tools + def get_listed_tool(self, server: MCPServer, name: str) -> MCPTool | None: + listed: Final = self._listed_tools_by_server_id.get(server.server_id, {}) + return listed.get(name) or listed.get(strip_known_server_prefix(name, server)) + def _create_prefixed_prompts( self, prompts: list[Prompt], server: MCPServer, add_prefix: bool = True ) -> list[Prompt]: @@ -5372,6 +5378,7 @@ class MCPServerManager: server: MCPServer, raw_headers: dict[str, str] | None = None, litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + tool: MCPTool | None = None, ) -> dict[str, Any]: """ Run pre-call checks and guardrail hooks for an MCP tool call. @@ -5385,6 +5392,9 @@ class MCPServerManager: ``pre_mcp_call`` evaluation (or a block) on the spend-log row the Guardrails Monitor counts. It stays optional so callers that do no logging are unchanged. + ``tool`` is the upstream tool definition when one was listed, so guardrails + can see its description and input schema, not just the name and arguments. + Returns a dict that may contain: - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks @@ -5438,6 +5448,8 @@ class MCPServerManager: "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None), "incoming_bearer_token": incoming_bearer_token, "headers": logging_safe_mcp_headers(raw_headers), + "tool_description": tool.description if tool is not None else None, + "tool_input_schema": tool.inputSchema if tool is not None else None, } # Create MCP request object for processing @@ -5492,6 +5504,7 @@ class MCPServerManager: proxy_logging_obj: ProxyLogging, start_time: datetime.datetime, litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + tool: MCPTool | None = None, ): """Create and return a during hook task for MCP tool calls. @@ -5506,6 +5519,8 @@ class MCPServerManager: tool_name=name, arguments=arguments, server_name=server_name_from_prefix, + tool_description=tool.description if tool is not None else None, + tool_input_schema=tool.inputSchema if tool is not None else None, start_time=start_time.timestamp() if start_time else None, hidden_params=HiddenParams(), ) @@ -6101,6 +6116,7 @@ class MCPServerManager: server=mcp_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, + tool=self.get_listed_tool(mcp_server, name), ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -6116,6 +6132,7 @@ class MCPServerManager: proxy_logging_obj=proxy_logging_obj, start_time=start_time, litellm_logging_obj=litellm_logging_obj, + tool=self.get_listed_tool(mcp_server, name), ) tasks.append(during_hook_task) diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index 4830bfcafb0..5813240f01b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -19,8 +19,10 @@ from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn import httpx from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.exceptions import Timeout as LitellmTimeout from litellm.integrations.custom_guardrail import ( @@ -33,7 +35,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( AGENT_365_PROD_API_BASE, AGENT_365_PROD_RESOURCE_APP_ID, @@ -48,9 +53,12 @@ if TYPE_CHECKING: from litellm.types.utils import GuardrailStatus TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token" +ENTRA_ISSUER_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/v2.0" EVALUATE_PATH: Final = "/agents/tool-evaluation/evaluate" MCP_SESSION_ID_HEADER: Final = "mcp-session-id" +DEFENDER_STATUS_EVALUATED: Final = "Evaluated" _MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool") +_TOOL_INPUT_SCHEMA_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBO_CACHE_MAX_ENTRIES: Final = 1000 _DEFAULT_TOKEN_TTL_SECONDS: Final = 3599.0 _TOKEN_EXPIRY_SLACK_SECONDS: Final = 60.0 @@ -65,6 +73,13 @@ def _parse_expires_in(raw: object) -> float: return _DEFAULT_TOKEN_TTL_SECONDS +def _parse_tool_input_schema(raw: object) -> dict[str, object] | None: + try: + return _TOOL_INPUT_SCHEMA_ADAPTER.validate_python(raw) + except ValidationError: + return None + + class _DefenderResult(TypedDict, total=False): status: ReadOnly[str] verdict: ReadOnly[str | None] @@ -77,8 +92,12 @@ class _EvaluateResponse(TypedDict, total=False): correlationId: ReadOnly[str] -class _ToolReference(TypedDict): - name: ReadOnly[str] +class _ToolReference(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str | None = None + input_schema: dict[str, object] | None = Field(default=None, serialization_alias="inputSchema") class _UnavailableDetail(TypedDict): @@ -316,11 +335,21 @@ class Agent365Guardrail(CustomGuardrail): defender: Final = raw_defender if isinstance(raw_defender, dict) else _DefenderResult() raw_correlation_id: Final = verdict.get("correlationId") correlation_id: Final = raw_correlation_id if isinstance(raw_correlation_id, str) else None + defender_status: Final = defender.get("status") + if allowed and defender_status != DEFENDER_STATUS_EVALUATED: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=f"Microsoft Defender did not evaluate the call (defender.status={defender_status or 'missing'})", + defender_status=defender_status, + correlation_id=correlation_id, + latency_ms=latency_ms, + ) self._record_verdict( data=data, verdict="Allow" if allowed else "Block", guardrail_status="success" if allowed else "guardrail_intervened", - defender_status=defender.get("status"), + defender_status=defender_status, correlation_id=correlation_id, latency_ms=latency_ms, ) @@ -347,9 +376,14 @@ class Agent365Guardrail(CustomGuardrail): arguments: Final = data.get("mcp_arguments") server_name: Final = str(data.get("mcp_server_name") or "litellm") agent_id: Final = self.agent_id or user_api_key_dict.key_alias - tool_reference: Final[_ToolReference] = {"name": tool_name} + description: Final = data.get("mcp_tool_description") + tool_reference: Final = _ToolReference( + name=tool_name, + description=description if isinstance(description, str) and description else None, + input_schema=_parse_tool_input_schema(data.get("mcp_tool_input_schema")), + ) payload: Final[dict[str, object]] = { # mutable-ok: JSON body with optional fields added below - "tool": tool_reference, + "tool": tool_reference.model_dump(by_alias=True, exclude_none=True), "serverName": server_name, "conversationId": self._resolve_conversation_id(data), } @@ -361,6 +395,18 @@ class Agent365Guardrail(CustomGuardrail): @staticmethod def _resolve_conversation_id(data: Mapping[str, object]) -> str: + raw_logging_obj: Final = data.get("litellm_logging_obj") + logging_obj: Final = raw_logging_obj if isinstance(raw_logging_obj, LiteLLMLoggingObj) else None + call_id: Final = data.get("litellm_call_id") or (logging_obj.litellm_call_id if logging_obj else None) + if isinstance(call_id, str) and call_id: + return call_id + if logging_obj is not None: + tool_call_metadata: Final = logging_obj.model_call_details.get("mcp_tool_call_metadata") + session_from_logging: Final = ( + tool_call_metadata.get("mcp_session_id") if isinstance(tool_call_metadata, Mapping) else None + ) + if isinstance(session_from_logging, str) and session_from_logging: + return session_from_logging metadata: Final = next( (m for m in (data.get("metadata"), data.get("litellm_metadata")) if isinstance(m, Mapping)), None, @@ -373,18 +419,6 @@ class Agent365Guardrail(CustomGuardrail): ) if isinstance(session_id, str) and session_id: return session_id - raw_logging_obj: Final = data.get("litellm_logging_obj") - logging_obj: Final = raw_logging_obj if isinstance(raw_logging_obj, LiteLLMLoggingObj) else None - if logging_obj is not None: - tool_call_metadata: Final = logging_obj.model_call_details.get("mcp_tool_call_metadata") - session_from_logging: Final = ( - tool_call_metadata.get("mcp_session_id") if isinstance(tool_call_metadata, Mapping) else None - ) - if isinstance(session_from_logging, str) and session_from_logging: - return session_from_logging - call_id: Final = data.get("litellm_call_id") or (logging_obj.litellm_call_id if logging_obj else None) - if isinstance(call_id, str) and call_id: - return call_id return str(uuid.uuid4()) async def _get_obo_token(self, assertion: str) -> str: @@ -518,6 +552,9 @@ class Agent365Guardrail(CustomGuardrail): data: dict, # mutable-ok: guardrail logging appends into the request metadata in place tool_name: str, reason: str, + defender_status: str | None = None, + correlation_id: str | None = None, + latency_ms: float | None = None, ) -> dict: # mutable-ok: returns the request data dict per hook contract if self.unreachable_fallback == "fail_open": verbose_proxy_logger.warning( @@ -530,9 +567,9 @@ class Agent365Guardrail(CustomGuardrail): data=data, verdict="Unscanned", guardrail_status="guardrail_failed_to_respond", - defender_status=None, - correlation_id=None, - latency_ms=None, + defender_status=defender_status, + correlation_id=correlation_id, + latency_ms=latency_ms, reason=reason, ) return data @@ -540,9 +577,9 @@ class Agent365Guardrail(CustomGuardrail): data=data, verdict="Unavailable", guardrail_status="guardrail_failed_to_respond", - defender_status=None, - correlation_id=None, - latency_ms=None, + defender_status=defender_status, + correlation_id=correlation_id, + latency_ms=latency_ms, reason=reason, ) unavailable_detail: Final[_UnavailableDetail] = { @@ -580,3 +617,33 @@ class Agent365Guardrail(CustomGuardrail): guardrail_provider=self.guardrail_provider, event_type=GuardrailEventHooks.pre_mcp_call, ) + + +def _applies_to_caller(guardrail: Agent365Guardrail, user_api_key_auth: "UserAPIKeyAuth") -> bool: + probe: Final[dict[str, object]] = {"metadata": {}} # mutable-ok: filled in place by the key resolver + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_auth, data=probe, metadata_variable_name="metadata" + ) + return guardrail.should_run_guardrail(data=probe, event_type=GuardrailEventHooks.pre_mcp_call) + + +def agent_365_authorization_servers(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]: + """Entra issuers an MCP client signs in with before calling ``server`` through an Agent 365 guardrail. + + Empty unless the admin advertised the server's ``scopes`` (the audience the client requests), the gateway + would otherwise own sign-in for the server, and an Agent 365 guardrail applies: every registered one for the + anonymous discovery fetch, otherwise those the caller's key, team, or policies select. + """ + if not server.scopes or server.auth_type == MCPAuth.oauth2 or not server.advertises_gateway_authorization_server: + return () + registered: Final = tuple( + callback + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(Agent365Guardrail) + if isinstance(callback, Agent365Guardrail) + ) + applicable: Final = ( + registered + if user_api_key_auth is None + else tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) + ) + return tuple(dict.fromkeys(ENTRA_ISSUER_TEMPLATE.format(tenant_id=g.tenant_id) for g in applicable)) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00ccad33b6d..420fb29634f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1216,6 +1216,8 @@ class ProxyLogging: "user_api_key_request_route": kwargs.get("user_api_key_request_route"), "mcp_tool_name": request_obj.tool_name, # Keep original for reference "mcp_arguments": request_obj.arguments, # Keep original for reference + "mcp_tool_description": request_obj.tool_description, + "mcp_tool_input_schema": request_obj.tool_input_schema, # Surface the per-MCP-server rate-limit identity so the # ParallelRequestLimiterV3 hook can apply mcp_rpm_limit on the # synthetic call_mcp_tool payload (otherwise a key with @@ -1437,6 +1439,8 @@ class ProxyLogging: tool_name=kwargs.get("name", ""), arguments=kwargs.get("arguments", {}), server_name=kwargs.get("server_name"), + tool_description=kwargs.get("tool_description"), + tool_input_schema=kwargs.get("tool_input_schema"), user_api_key_auth=user_api_key_auth_dict, hidden_params=HiddenParams(), ) diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index a59fcb1bcb5..13a46d8b84f 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -391,6 +391,8 @@ class MCPPreCallRequestObject(BaseModel): tool_name: str arguments: dict[str, Any] server_name: str | None = None + tool_description: str | None = None + tool_input_schema: dict[str, object] | None = None user_api_key_auth: dict[str, Any] | None = None hidden_params: HiddenParams = HiddenParams() @@ -414,6 +416,8 @@ class MCPDuringCallRequestObject(BaseModel): tool_name: str arguments: dict[str, Any] server_name: str | None = None + tool_description: str | None = None + tool_input_schema: dict[str, object] | None = None start_time: float | None = None hidden_params: HiddenParams = HiddenParams() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index adf985e9a21..ba7b76eae91 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -6403,6 +6403,136 @@ class TestMCPServerManager: # Verify the MCP client call was awaited exactly once assert mock_client.call_tool.await_count == 1 + @staticmethod + def _manager_ready_for_call_tool(listed_tools: list[MCPTool]) -> tuple[MCPServerManager, MagicMock]: + from mcp.types import CallToolResult + + manager = MCPServerManager() + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + url="http://test-server.com", + ) + manager.registry = {"test-server": server} + manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server" + manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server" + manager._create_prefixed_tools(listed_tools, server) + + mock_client = AsyncMock() + mock_client.call_tool.return_value = MagicMock(spec=CallToolResult, content=[], isError=False) + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + proxy_logging_obj = MagicMock() + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + return manager, proxy_logging_obj + + @staticmethod + def _unrestricted_auth() -> MagicMock: + user_api_key_auth = MagicMock() + user_api_key_auth.object_permission = None + user_api_key_auth.object_permission_id = None + return user_api_key_auth + + @pytest.mark.asyncio + async def test_call_tool_hands_listed_tool_description_and_schema_to_pre_call_hooks(self): + schema = {"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]} + listed = [MCPTool(name="test_tool", description="Runs the test tool", inputSchema=schema)] + manager, proxy_logging_obj = self._manager_ready_for_call_tool(listed) + + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={"param": "value"}, + user_api_key_auth=self._unrestricted_auth(), + proxy_logging_obj=proxy_logging_obj, + ) + + hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0] + assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == ("Runs the test tool", schema) + + @pytest.mark.asyncio + async def test_call_tool_hands_listed_tool_metadata_to_during_call_hooks_through_real_conversion(self): + schema = {"type": "object", "properties": {"param": {"type": "string"}}} + listed = [MCPTool(name="test_tool", description="Runs the test tool", inputSchema=schema)] + manager, _ = self._manager_ready_for_call_tool(listed) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={"param": "value"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + during_data = proxy_logging_obj.during_call_hook.call_args.kwargs["data"] + assert (during_data["mcp_tool_description"], during_data["mcp_tool_input_schema"]) == ( + "Runs the test tool", + schema, + ) + + @pytest.mark.asyncio + async def test_call_tool_passes_no_tool_metadata_when_tool_was_never_listed(self): + manager, proxy_logging_obj = self._manager_ready_for_call_tool( + [MCPTool(name="other_tool", description="Unrelated", inputSchema={"type": "object"})] + ) + + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={"param": "value"}, + user_api_key_auth=self._unrestricted_auth(), + proxy_logging_obj=proxy_logging_obj, + ) + + hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0] + assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == (None, None) + + def test_get_listed_tool_resolves_prefixed_name_and_latest_listing(self): + manager = MCPServerManager() + server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") + manager._create_prefixed_tools([MCPTool(name="echo", description="v1", inputSchema={})], server) + manager._create_prefixed_tools([MCPTool(name="echo", description="v2", inputSchema={})], server) + + by_prefixed_name = manager.get_listed_tool(server, "srv-echo") + assert by_prefixed_name is not None and by_prefixed_name.description == "v2" + assert manager.get_listed_tool(server, "missing") is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("add_prefix", [True, False]) + async def test_openapi_listing_records_tool_metadata_for_pre_call_hooks(self, add_prefix): + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.types.mcp_server.tool_registry import MCPTool as RegistryTool + + server = MCPServer( + server_id="petstore-id", + name="petstore", + server_name="petstore", + transport=MCPTransport.http, + url=None, + spec_path="https://example.com/petstore.yaml", + ) + schema = {"type": "object", "properties": {"petId": {"type": "integer"}}} + registered = RegistryTool( + name="petstore-get_pet", description="Fetch a pet", input_schema=schema, handler=lambda: None + ) + manager = MCPServerManager() + manager._create_mcp_client = AsyncMock(return_value=AsyncMock()) + + with patch.dict(global_mcp_tool_registry.tools, {"petstore-get_pet": registered}, clear=True): + listed = await manager._get_tools_from_server(server=server, add_prefix=add_prefix) + + assert [t.name for t in listed] == ["petstore-get_pet" if add_prefix else "get_pet"] + for spelling in ("get_pet", "petstore-get_pet"): + tool = manager.get_listed_tool(server, spelling) + assert tool is not None and (tool.description, tool.inputSchema) == ("Fetch a pet", schema) + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py index 04460b150f5..b928ccad4c6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -1,11 +1,15 @@ import time +import uuid +from collections.abc import Iterator from types import SimpleNamespace from typing import Any, Final +from unittest.mock import patch import httpx import pytest from fastapi import HTTPException +import litellm from litellm.exceptions import Timeout as LitellmTimeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.secret_redaction import redact_string @@ -16,11 +20,14 @@ from litellm.proxy.guardrails.guardrail_hooks.agent_365 import ( guardrail_initializer_registry, initialize_guardrail, ) +from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import agent_365_authorization_servers from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, SupportedGuardrailIntegrations, ) +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( AGENT_365_PROD_API_BASE, AGENT_365_PROD_RESOURCE_APP_ID, @@ -55,17 +62,46 @@ def _allow_response(correlation_id: str = "corr-1") -> httpx.Response: ) -def _block_response(message: str = "Blocked by policy", correlation_id: str = "corr-2") -> httpx.Response: +def _block_response( + message: str = "Blocked by policy", correlation_id: str = "corr-2", status: str = "Evaluated" +) -> httpx.Response: return _response( 200, { "allowed": False, - "defender": {"status": "Evaluated", "verdict": "Block", "message": message}, + "defender": {"status": status, "verdict": "Block", "message": message}, "correlationId": correlation_id, }, ) +def _not_evaluated_response(status: str, correlation_id: str = "corr-3") -> httpx.Response: + return _response( + 200, + { + "allowed": True, + "defender": {"status": status, "verdict": None, "message": None}, + "observability": {"status": "Unavailable"}, + "correlationId": correlation_id, + }, + ) + + +def _logging_obj(litellm_call_id: str, mcp_session_id: str | None = None) -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="mcp", + messages=[], + stream=False, + call_type="call_mcp_tool", + start_time=None, + litellm_call_id=litellm_call_id, + function_id="fn-1", + ) + if mcp_session_id is not None: + logging_obj.model_call_details["mcp_tool_call_metadata"] = {"mcp_session_id": mcp_session_id} + return logging_obj + + class FakeHandler: def __init__(self, items: list[Any]): self._items = list(items) @@ -266,6 +302,29 @@ class TestAllowFlow: assert evaluate_call.json["conversationId"] == "sess-123" assert evaluate_call.json["agentId"] == "agent-007" + @pytest.mark.asyncio + async def test_evaluate_payload_includes_listed_tool_metadata(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + schema: Final = {"type": "object", "properties": {"to": {"type": "string"}}, "required": ["to"]} + await _run(guardrail, _mcp_data(mcp_tool_description="Send an email", mcp_tool_input_schema=schema)) + assert handler.calls[1].json["tool"] == { + "name": "send_email", + "description": "Send an email", + "inputSchema": schema, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("description", "schema"), + [(None, None), ("", None), (None, ["not", "a", "schema"]), (42, "type: object")], + ) + async def test_evaluate_payload_omits_missing_or_malformed_tool_metadata(self, description, schema): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data(mcp_tool_description=description, mcp_tool_input_schema=schema)) + assert handler.calls[1].json["tool"] == {"name": "send_email"} + @pytest.mark.asyncio async def test_agent_id_falls_back_to_key_alias(self): handler: Final = FakeHandler([_token_response(), _allow_response()]) @@ -284,6 +343,33 @@ class TestAllowFlow: class TestConversationId: + @pytest.mark.asyncio + async def test_request_call_id_beats_logging_obj_and_client_header(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data( + litellm_call_id="call-id-from-data", + litellm_logging_obj=_logging_obj("call-id-from-logging", mcp_session_id="sess-from-logging"), + ) + await _run(guardrail, data) + assert handler.calls[1].json["conversationId"] == "call-id-from-data" + + @pytest.mark.asyncio + async def test_logging_obj_call_id_beats_session_metadata_and_client_header(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data(litellm_logging_obj=_logging_obj("call-id-1", mcp_session_id="sess-from-logging")) + await _run(guardrail, data) + assert handler.calls[1].json["conversationId"] == "call-id-1" + + @pytest.mark.asyncio + async def test_logging_obj_session_id_beats_client_header(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data(litellm_logging_obj=_logging_obj("", mcp_session_id="sess-from-logging")) + await _run(guardrail, data) + assert handler.calls[1].json["conversationId"] == "sess-from-logging" + @pytest.mark.asyncio async def test_session_id_header_case_insensitive(self): handler: Final = FakeHandler([_token_response(), _allow_response()]) @@ -293,30 +379,12 @@ class TestConversationId: assert handler.calls[1].json["conversationId"] == "sess-CASED" @pytest.mark.asyncio - async def test_falls_back_to_logging_obj_session_id(self): + async def test_generates_uuid_when_no_identifier_available(self): handler: Final = FakeHandler([_token_response(), _allow_response()]) guardrail: Final = _make_guardrail(handler) - logging_obj: Final = LiteLLMLoggingObj( - model="mcp", - messages=[], - stream=False, - call_type="call_mcp_tool", - start_time=None, - litellm_call_id="call-id-1", - function_id="fn-1", - ) - logging_obj.model_call_details["mcp_tool_call_metadata"] = {"mcp_session_id": "sess-from-logging"} - data: Final = _mcp_data(metadata={"headers": {}}, litellm_logging_obj=logging_obj) - await _run(guardrail, data) - assert handler.calls[1].json["conversationId"] == "sess-from-logging" - - @pytest.mark.asyncio - async def test_falls_back_to_litellm_call_id(self): - handler: Final = FakeHandler([_token_response(), _allow_response()]) - guardrail: Final = _make_guardrail(handler) - data: Final = _mcp_data(metadata={"headers": {}}, litellm_call_id="call-id-2") - await _run(guardrail, data) - assert handler.calls[1].json["conversationId"] == "call-id-2" + await _run(guardrail, _mcp_data(metadata={"headers": {}}, litellm_logging_obj=_logging_obj(""))) + conversation_id: Final = handler.calls[1].json["conversationId"] + assert uuid.UUID(conversation_id).version == 4 class TestBlockFlow: @@ -344,6 +412,65 @@ class TestBlockFlow: await _run(guardrail, _mcp_data()) assert exc_info.value.status_code == 400 + @pytest.mark.asyncio + @pytest.mark.parametrize("status", ["Skipped", "FailedOpen"]) + async def test_explicit_block_wins_over_non_evaluated_status(self, status): + handler: Final = FakeHandler([_token_response(), _block_response(status=status)]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 400 + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_intervened" + assert info["guardrail_response"]["verdict"] == "Block" + assert info["guardrail_response"]["defender_status"] == status + + +class TestDefenderNotEvaluated: + @pytest.mark.asyncio + @pytest.mark.parametrize("status", ["Skipped", "FailedOpen"]) + async def test_fail_closed_blocks_allowed_but_unevaluated_call(self, status): + handler: Final = FakeHandler([_token_response(), _not_evaluated_response(status)]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_closed") + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 503 + assert f"defender.status={status}" in exc_info.value.detail["message"] + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Unavailable" + assert info["guardrail_response"]["defender_status"] == status + assert info["guardrail_response"]["correlation_id"] == "corr-3" + assert info["guardrail_response"]["latency_ms"] >= 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize("status", ["Skipped", "FailedOpen"]) + async def test_fail_open_allows_unevaluated_call_as_unscanned(self, status): + handler: Final = FakeHandler([_token_response(), _not_evaluated_response(status)]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Unscanned" + assert info["guardrail_response"]["defender_status"] == status + assert info["guardrail_response"]["correlation_id"] == "corr-3" + + @pytest.mark.asyncio + @pytest.mark.parametrize("payload", [{"allowed": True}, {"allowed": True, "defender": {"verdict": "Allow"}}]) + async def test_allowed_without_defender_status_is_not_an_evaluated_allow(self, payload): + handler: Final = FakeHandler([_token_response(), _response(200, payload)]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_closed") + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 503 + assert "defender.status=missing" in exc_info.value.detail["message"] + assert "defender_status" not in _guardrail_info(data)["guardrail_response"] + @pytest.mark.asyncio async def test_http_400_always_blocks_even_fail_open(self): handler: Final = FakeHandler([_token_response(), _response(400, text="Bad request: serverName missing")]) @@ -783,3 +910,92 @@ class TestVeriaHardening: assert len(records) == 1 assert records[0]["guardrail_response"]["verdict"] == "Unscanned" assert records[0]["guardrail_status"] == "guardrail_failed_to_respond" + + +ENTRA_ISSUER: Final = "https://login.microsoftonline.com/tenant-abc/v2.0" +GATEWAY_SCOPE: Final = "api://gateway-app/access_as_user" + + +def _mcp_server(auth_type: MCPAuth = MCPAuth.none, scopes: list[str] | None = None, **fields: Any) -> MCPServer: + return MCPServer( + server_id="tools-id", + name="tools", + server_name="tools", + transport=MCPTransport.http, + url="https://tools.test/mcp", + auth_type=auth_type, + scopes=scopes, + **fields, + ) + + +@pytest.fixture +def registered_guardrail() -> Iterator[Agent365Guardrail]: + guardrail: Final = _make_guardrail(FakeHandler([])) + litellm.logging_callback_manager.add_litellm_callback(guardrail) + try: + yield guardrail + finally: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, guardrail, require_self=False + ) + + +class TestAgent365AuthorizationServers: + def test_names_the_guardrail_tenant_for_a_scoped_gateway_signed_in_server(self, registered_guardrail): + assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (ENTRA_ISSUER,) + assert agent_365_authorization_servers( + _mcp_server(MCPAuth.api_key, scopes=[GATEWAY_SCOPE], auth_value="k"), None + ) == (ENTRA_ISSUER,) + + def test_silent_without_advertised_scopes(self, registered_guardrail): + assert agent_365_authorization_servers(_mcp_server(scopes=None), None) == () + assert agent_365_authorization_servers(_mcp_server(scopes=[]), None) == () + + def test_silent_when_no_guardrail_is_registered(self): + assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == () + + @pytest.mark.parametrize( + "server", + [ + _mcp_server(MCPAuth.oauth2, scopes=[GATEWAY_SCOPE]), + _mcp_server(MCPAuth.oauth2_token_exchange, scopes=[GATEWAY_SCOPE], token_exchange_endpoint="https://i/t"), + _mcp_server(MCPAuth.oauth2_id_jag, scopes=[GATEWAY_SCOPE]), + _mcp_server(MCPAuth.true_passthrough, scopes=[GATEWAY_SCOPE]), + _mcp_server(MCPAuth.oauth_delegate, scopes=[GATEWAY_SCOPE]), + _mcp_server(MCPAuth.none, scopes=[GATEWAY_SCOPE], extra_headers=["Authorization"]), + ], + ids=["oauth2", "token_exchange", "id_jag", "true_passthrough", "oauth_delegate", "forwards_authorization"], + ) + def test_leaves_servers_whose_own_auth_mode_owns_sign_in_alone(self, registered_guardrail, server): + assert agent_365_authorization_servers(server, None) == () + + def test_dedupes_guardrails_sharing_a_tenant(self, registered_guardrail): + twin: Final = _make_guardrail(FakeHandler([])) + twin.guardrail_name = "agent-365-twin" + litellm.logging_callback_manager.add_litellm_callback(twin) + try: + assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (ENTRA_ISSUER,) + finally: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, twin, require_self=False + ) + + def test_key_selected_guardrail_challenges_only_that_key(self): + guardrail: Final = _make_guardrail(FakeHandler([])) + guardrail.default_on = False + litellm.logging_callback_manager.add_litellm_callback(guardrail) + server: Final = _mcp_server(scopes=[GATEWAY_SCOPE]) + plain_key: Final = UserAPIKeyAuth(api_key="sk-plain", user_id="u-1") + guarded_key: Final = UserAPIKeyAuth( + api_key="sk-guarded", user_id="u-2", metadata={"guardrails": ["agent-365-guard"]} + ) + try: + with patch("litellm.proxy.proxy_server.premium_user", True): + assert agent_365_authorization_servers(server, plain_key) == () + assert agent_365_authorization_servers(server, guarded_key) == (ENTRA_ISSUER,) + assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,) + finally: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, guardrail, require_self=False + ) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 56057dce7e0..75dd11ea16b 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py @@ -381,6 +381,26 @@ def test_create_mcp_request_object_from_kwargs_full(proxy_logging, make_user_api assert snapshot == {"tool_name": "calc", "arguments": {"x": 1}, "server_name": "math", "auth_user_id": "u-1"} +def test_mcp_tool_metadata_flows_from_kwargs_to_synthetic_data(proxy_logging): + schema = {"type": "object", "properties": {"x": {"type": "integer"}}} + obj = proxy_logging._create_mcp_request_object_from_kwargs( + kwargs={ + "name": "calc", + "arguments": {"x": 1}, + "tool_description": "Adds numbers", + "tool_input_schema": schema, + } + ) + out = proxy_logging._convert_mcp_to_llm_format(request_obj=obj, kwargs={}) + assert (out["mcp_tool_description"], out["mcp_tool_input_schema"]) == ("Adds numbers", schema) + + +def test_mcp_tool_metadata_absent_when_tool_was_never_listed(proxy_logging): + obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs={"name": "calc", "arguments": {}}) + out = proxy_logging._convert_mcp_to_llm_format(request_obj=obj, kwargs={}) + assert (out["mcp_tool_description"], out["mcp_tool_input_schema"]) == (None, None) + + def test_create_mcp_request_object_from_kwargs_empty(proxy_logging): obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs={}) snapshot = {