diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 560d3d51177..319735a8c08 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -246,6 +246,17 @@ _ListedToolsByCaller: TypeAlias = Mapping[str | None, Mapping[str, MCPTool]] _NO_LISTED_TOOLS: Final[_ListedToolsByCaller] = MappingProxyType({}) _LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 + +@dataclass(frozen=True, slots=True) +class ListedToolsCaller: + """Request inputs that select which upstream catalog a caller was shown by tools/list.""" + + user_api_key_auth: UserAPIKeyAuth | None = None + mcp_auth_header: str | dict[str, str] | None = None + raw_headers: Mapping[str, str] | None = None + oauth2_headers: Mapping[str, str] | None = None + + # Auth types whose upstream OAuth endpoints (protected-resource + authorization-server metadata) the # gateway discovers from the upstream itself: interactive oauth2 and the two client-forwarded modes. # OBO/M2M endpoint discovery is decided separately via _obo_needs_endpoint_discovery. Shared by the @@ -1123,6 +1134,25 @@ def _authorization_is_litellm_admission_credential( return bool(user_api_key_auth and user_api_key_auth.api_key and not admission_header) +def _server_auth_header_for( + server: MCPServer, + mcp_server_auth_headers: Mapping[str, str | dict[str, str]] | None, + mcp_auth_header: str | dict[str, str] | None, +) -> str | dict[str, str] | None: + """Server-specific ``x-mcp--authorization`` header, else the deprecated global one.""" + server_specific: Final = ( + lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + access_groups=server.access_groups, + ) + if mcp_server_auth_headers + else None + ) + return mcp_auth_header if server_specific is None else server_specific + + def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str: """Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection. @@ -3720,19 +3750,7 @@ class MCPServerManager: verbose_logger.warning("MCP Server %s not found", server_id) return [] - # Get server-specific auth header if available - server_auth_header: str | dict[str, str] | None = None - if mcp_server_auth_headers: - server_auth_header = lookup_mcp_server_auth_in_headers( - mcp_server_auth_headers, - alias=server.alias, - server_name=server.server_name, - access_groups=server.access_groups, - ) - - # Fall back to deprecated mcp_auth_header if no server-specific header found - if server_auth_header is None: - server_auth_header = mcp_auth_header + server_auth_header: Final = _server_auth_header_for(server, mcp_server_auth_headers, mcp_auth_header) try: tools: Final = await self._get_tools_from_server( @@ -3818,7 +3836,7 @@ class MCPServerManager: def _build_stdio_env( self, server: MCPServer, - raw_headers: dict[str, str] | None = None, + raw_headers: Mapping[str, str] | None = None, ) -> dict[str, str] | None: """Resolve stdio env values, supporting header-driven placeholders.""" @@ -4347,6 +4365,12 @@ class MCPServerManager: verbose_logger.info("_get_tools_from_server for %s...", server.name) client = None + listed_caller: Final = ListedToolsCaller( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + raw_headers=raw_headers, + oauth2_headers=oauth2_headers, + ) try: # Tool *listing* must not be blocked by missing per-user env vars — @@ -4433,14 +4457,14 @@ class MCPServerManager: unprefixed_tools: Final = [ # mutable-ok: returned through the list[MCPTool] listing contract t.model_copy(update=MappingProxyType({"name": t.name[len(registry_prefix) :]})) for t in tools ] - self._record_listed_tools(server, unprefixed_tools, user_api_key_auth) + self._record_listed_tools(server, unprefixed_tools, listed_caller) return tools if add_prefix else unprefixed_tools else: tools = await self._fetch_tools_with_timeout(client, server.name) self._remember_upstream_initialize_instructions(server, client) prefixed_or_original_tools: Final = self._create_prefixed_tools( - tools, server, add_prefix=add_prefix, user_api_key_auth=user_api_key_auth + tools, server, add_prefix=add_prefix, caller=listed_caller ) return prefixed_or_original_tools @@ -4497,16 +4521,49 @@ class MCPServerManager: or server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) ) - def _listed_tools_identity(self, server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None) -> str | None: - if server.spec_path or user_api_key_auth is None or not self._discovers_per_caller(server): + def _listed_tools_identity(self, server: MCPServer, caller: ListedToolsCaller | None) -> str | None: + """Key the listed-tool cache by every request input that can change the upstream catalog. + + Forwarded headers, header-driven stdio env, a relayed caller bearer, and the + server-specific auth header all reach upstream, so two callers differing in any of + them may be shown different tools. Shared servers with none of those stay on the + shared (``None``) slot. OpenAPI servers list from the process-wide registry. + """ + if server.spec_path or caller is None: return None - material: Final = json.dumps((user_api_key_auth.user_id, user_api_key_auth.api_key), separators=(",", ":")) + auth: Final = caller.user_api_key_auth + identity: Final = ( + (auth.user_id, auth.api_key) if auth is not None and self._discovers_per_caller(server) else None + ) + forwarded: Final = self._forwarded_header_values(server, caller.raw_headers) + header_env: Final = self._build_stdio_env(server, caller.raw_headers) + stdio_env: Final = None if header_env == self._build_stdio_env(server) else header_env + relayed_bearer: Final = ( + self._extract_subject_token(caller.oauth2_headers, caller.raw_headers, auth) + if server.is_client_forwarded_token + else None + ) + inputs: Final = (identity, caller.mcp_auth_header, forwarded, stdio_env, relayed_bearer) + if not any(inputs): + return None + material: Final = json.dumps(inputs, sort_keys=True, separators=(",", ":")) return hashlib.sha256(material.encode()).hexdigest() + @staticmethod + def _forwarded_header_values( + server: MCPServer, raw_headers: Mapping[str, str] | None + ) -> tuple[tuple[str, str], ...]: + if not raw_headers or not server.extra_headers: + return () + forwarded_names: Final = frozenset(name.lower() for name in server.extra_headers) + return tuple( + sorted((name.lower(), value) for name, value in raw_headers.items() if name.lower() in forwarded_names) + ) + def _record_listed_tools( - self, server: MCPServer, tools: Sequence[MCPTool], user_api_key_auth: UserAPIKeyAuth | None + self, server: MCPServer, tools: Sequence[MCPTool], caller: ListedToolsCaller | None ) -> None: - identity: Final = self._listed_tools_identity(server, user_api_key_auth) + identity: Final = self._listed_tools_identity(server, caller) listing: Final = MappingProxyType({tool.name: tool for tool in tools}) existing: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS) shared: Final = existing.get(None) @@ -5347,7 +5404,7 @@ class MCPServerManager: tools: list[MCPTool], server: MCPServer, add_prefix: bool = True, - user_api_key_auth: UserAPIKeyAuth | None = None, + caller: ListedToolsCaller | None = None, ) -> list[MCPTool]: """ Create prefixed tools and update tool mapping. @@ -5380,14 +5437,12 @@ class MCPServerManager: for spelling in iter_known_tool_name_spellings(original_name, server): self.tool_name_to_mcp_server_name_mapping[spelling] = prefix - self._record_listed_tools(server, tools, user_api_key_auth) + self._record_listed_tools(server, tools, caller) 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, user_api_key_auth: UserAPIKeyAuth | None = None - ) -> MCPTool | None: - identity: Final = self._listed_tools_identity(server, user_api_key_auth) + def get_listed_tool(self, server: MCPServer, name: str, caller: ListedToolsCaller | None = None) -> MCPTool | None: + identity: Final = self._listed_tools_identity(server, caller) listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity) if not listed: return None @@ -5909,21 +5964,7 @@ class MCPServerManager: GuardrailRaisedException: If guardrails block the call HTTPException: If an HTTP error occurs """ - # Get server-specific auth header if available (case-insensitive) - # FIX: Added case-insensitive matching to handle auth header keys that may not match - # the exact case of server alias/name (e.g., '1litellmagcgateway' vs '1LiteLLMAGCGateway') - server_auth_header: dict[str, str] | str | None = None - if mcp_server_auth_headers: - server_auth_header = lookup_mcp_server_auth_in_headers( - mcp_server_auth_headers, - alias=mcp_server.alias, - server_name=mcp_server.server_name, - access_groups=mcp_server.access_groups, - ) - - # Fall back to deprecated mcp_auth_header if no server-specific header found - if server_auth_header is None: - server_auth_header = mcp_auth_header + server_auth_header: Final = _server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header) # Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows subject_token: str | None = None @@ -6359,6 +6400,12 @@ class MCPServerManager: user_api_key_auth, mcp_auth_header, ) + listed_caller: Final = ListedToolsCaller( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=_server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header), + raw_headers=raw_headers, + oauth2_headers=oauth2_headers, + ) ######################################################### # Pre MCP Tool Call Hook @@ -6374,7 +6421,7 @@ class MCPServerManager: server=mcp_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, - tool=self.get_listed_tool(mcp_server, name, user_api_key_auth), + tool=self.get_listed_tool(mcp_server, name, listed_caller), ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -6390,7 +6437,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, user_api_key_auth), + tool=self.get_listed_tool(mcp_server, name, listed_caller), ) tasks.append(during_hook_task) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 922ad9ac9b2..3b471033235 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2783,8 +2783,10 @@ if MCP_AVAILABLE: return managed_resource_templates - def _registered_tool_metadata(name: str, registered: RegisteredTool) -> MCPTool: - return MCPTool(name=name, description=registered.description, inputSchema=registered.input_schema) + def _registered_tool_metadata(name: str, registered: RegisteredTool, server: MCPServer) -> MCPTool: + overrides: Final = server.tool_name_to_description + description: Final = overrides.get(name, registered.description) if overrides else registered.description + return MCPTool(name=name, description=description, inputSchema=registered.input_schema) def _resolve_display_name_to_original( name: str, @@ -3124,7 +3126,7 @@ if MCP_AVAILABLE: server=mcp_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, - tool=_registered_tool_metadata(original_tool_name, local_tool), + tool=_registered_tool_metadata(original_tool_name, local_tool, mcp_server), ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -3232,7 +3234,7 @@ if MCP_AVAILABLE: server=prefix_server, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, - tool=_registered_tool_metadata(original_tool_name, registered_local_tool), + tool=_registered_tool_metadata(original_tool_name, registered_local_tool, prefix_server), ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args @@ -4019,6 +4021,21 @@ if MCP_AVAILABLE: ) return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + async def _key_granted_single_server( + server: MCPServer, + mcp_servers: Sequence[str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None, + ) -> bool: + """Sign-in challenges are issued only on a single-server connect the key's grant admits, so a key + without access gets the grant's 403 instead of a sign-in it could not use.""" + if len(mcp_servers or []) != 1: + return False + allowed: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip + ) + return any(granted.server_id == server.server_id for granted in allowed) + async def _raise_preemptive_401_for_unauthenticated_servers( scope: Scope, mcp_servers: list[str] | None, @@ -4146,10 +4163,14 @@ if MCP_AVAILABLE: # bearer is the LiteLLM key itself, which admits the caller but is not an exchangeable subject. # Only on the server's own route: the per-server metadata ``resource`` must equal the URL the # client connected to (RFC 9728 3.3), which aggregate ``/mcp`` and multi-server connects never do. + granted_single_server = server is not None and await _key_granted_single_server( + server, mcp_servers, user_api_key_auth, client_ip + ) if server and ( (server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers) or ( - tuple(_get_mcp_servers_in_path(get_route_relative_request_path(scope)) or ()) == (server_name,) + granted_single_server + and tuple(_get_mcp_servers_in_path(get_route_relative_request_path(scope)) or ()) == (server_name,) and not agent_365_subject_token_present(oauth2_headers) and agent_365_authorization_servers(server, user_api_key_auth) ) @@ -4174,17 +4195,7 @@ if MCP_AVAILABLE: # and what each mints from. Gated to single-server routes the key may reach; the # multi-server aggregate keeps absorbing per-server auth failures so one bad server # cannot 401 the whole connect. - if ( - server - and len(mcp_servers or []) == 1 - and server.server_id - in frozenset( - allowed.server_id - for allowed in await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip - ) - ) - ): + if server and granted_single_server: await global_mcp_server_manager.preflight_token_exchange( server=server, oauth2_headers=oauth2_headers, 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 d81fa2f095b..b134e1e7ff6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -136,7 +136,10 @@ class Agent365ThrottledError(Exception): class Agent365Guardrail(CustomGuardrail): - """Pre-MCP-call guardrail enforcing Microsoft Agent 365 tool-evaluation verdicts.""" + """Pre-MCP-call guardrail enforcing Microsoft Agent 365 tool-evaluation verdicts. + + Block-only: it never rewrites the call, so it runs in the post-sequential phase and judges the + arguments the sequential guardrails hand upstream, whatever order the guardrails list uses.""" records_own_guardrail_information: ClassVar[bool] = True @@ -157,6 +160,7 @@ class Agent365Guardrail(CustomGuardrail): super().__init__( guardrail_name=guardrail_name, supported_event_hooks=self.get_supported_event_hooks(), + run_in_parallel=True, **kwargs, ) self.guardrail_provider = "agent_365" @@ -408,11 +412,10 @@ class Agent365Guardrail(CustomGuardrail): @staticmethod def _resolve_conversation_id(data: Mapping[str, object]) -> str: + """The MCP session groups every tool call of one client conversation, so it is the conversation id + when the transport carries one; stateless calls fall back to the per-call id.""" 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 = ( @@ -432,6 +435,9 @@ class Agent365Guardrail(CustomGuardrail): ) if isinstance(session_id, str) and session_id: return session_id + 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: @@ -643,7 +649,7 @@ def _applies_to_caller(guardrail: Agent365Guardrail, user_api_key_auth: "UserAPI def _applicable_guardrails( server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None" ) -> tuple[Agent365Guardrail, ...]: - """Agent 365 guardrails that gate ``server`` for this caller: every registered one for the anonymous + """Agent 365 guardrails that gate ``server`` for this caller: the ``default_on`` ones for the anonymous discovery fetch, otherwise those the caller's key, team, or policies select. Empty when the gateway does not own sign-in for the server.""" if server.auth_type == MCPAuth.oauth2 or not server.advertises_gateway_authorization_server: @@ -654,7 +660,7 @@ def _applicable_guardrails( if isinstance(callback, Agent365Guardrail) ) if user_api_key_auth is None: - return registered + return tuple(g for g in registered if g.default_on) return tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index ebb5775b14e..25164d4b215 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7200,6 +7200,19 @@ async def test_agent_365_prm_defaults_scopeless_server_to_the_gateway_app_scope( } +@pytest.mark.asyncio +async def test_key_selected_agent_365_guardrail_leaves_anonymous_prm_on_the_gateway_issuer(agent_365_guardrail): + """A default-off guardrail gates only the keys that select it, so the anonymous discovery fetch must keep + pointing every other client at the gateway's own authorization server and scopes.""" + agent_365_guardrail.default_on = False + response = await _agent_365_gated_prm(scopes=["mcp:read"]) + assert jsonable_encoder(response) == { + "authorization_servers": ["https://litellm.example.com/mcp"], + "resource": "https://litellm.example.com/mcp/tools", + "scopes_supported": ["mcp:read"], + } + + def _token_request(headers): """A real Starlette request with case-insensitive headers (matches production).""" from starlette.requests import Request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 9a49f411fde..1817f101758 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7052,6 +7052,7 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): fake_server.server_name = "openapi-petstore" fake_server.alias = None fake_server.short_prefix = None + fake_server.tool_name_to_description = None fake_tool = MagicMock() fake_tool.name = "list_pets" @@ -7152,6 +7153,48 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca ) +@pytest.mark.asyncio +async def test_execute_mcp_tool_hands_openapi_hooks_the_admin_description_clients_saw(): + """tools/list shows the admin's tool_name_to_description wording, so the local-registry call path + must hand the pre-call hooks that same wording rather than the generated one.""" + from litellm.proxy._experimental.mcp_server import server as mcp_module + + petstore = MCPServer( + server_id="petstore-id", + name="petstore", + server_name="petstore", + transport=MCPTransport.http, + url=None, + spec_path="https://example.com/petstore.yaml", + tool_name_to_description={"getpetbyid": "ADMIN DESC"}, + ) + schema = {"type": "object", "properties": {"petId": {"type": "integer"}}} + mcp_module.global_mcp_tool_registry.register_tool( + name="petstore-getpetbyid", description="Find pet by ID", input_schema=schema, handler=lambda petId: "ok" + ) + manager = mcp_module.global_mcp_server_manager + manager._listed_tools_by_server_id.pop(petstore.server_id, None) + pre_call_tool_check = AsyncMock(return_value={}) + + try: + with ( + patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore), + patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check), + ): + await mcp_module.execute_mcp_tool( + name="petstore-getpetbyid", + arguments={"petId": 1}, + allowed_mcp_servers=[petstore], + start_time=datetime.now(), + user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"), + ) + finally: + mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") + + handed_tool = pre_call_tool_check.call_args.kwargs["tool"] + assert (handed_tool.description, handed_tool.inputSchema) == ("ADMIN DESC", schema) + + @pytest.mark.asyncio async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_runs_when_names_collide(): """An OpenAPI operation whose name starts with its own server prefix must not be reported to the @@ -8914,6 +8957,7 @@ class TestAgent365ChallengeAtConnect: oauth2_headers: dict[str, str] | None, path: str = "/mcp/tools", mount_scope: dict[str, str] | None = None, + granted: bool = True, ) -> HTTPException | None: from litellm.proxy._experimental.mcp_server import server as server_module @@ -8922,7 +8966,7 @@ class TestAgent365ChallengeAtConnect: server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server ), patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer - server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[]) + server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server] if granted else []) ), ): try: @@ -9034,6 +9078,12 @@ class TestAgent365ChallengeAtConnect: async def test_no_registered_guardrail_means_no_challenge(self): assert await self._connect(self._server([self.GATEWAY_SCOPE]), None) is None + @pytest.mark.asyncio + async def test_key_without_the_server_grant_is_not_sent_to_sign_in(self, agent_365_guardrail): + """Signing in cannot earn a key a server it was never granted, so the connect must fall through to + the ordinary 403 grant denial instead of leading with an Entra challenge the caller cannot use.""" + assert await self._connect(self._server([self.GATEWAY_SCOPE]), None, granted=False) is None + def _make_obo_server(alias: str) -> MCPServer: return MCPServer( 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 56e693e3523..f94d1dfa53b 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 @@ -34,6 +34,7 @@ from pydantic import AnyUrl, TypeAdapter from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + ListedToolsCaller, MCPServerManager, _deserialize_json_dict, _flow_endpoints_missing, @@ -6736,25 +6737,153 @@ class TestMCPServerManager: alice_schema = {"type": "object", "properties": {"path": {"type": "string"}}} bob_schema = {"type": "object", "properties": {"path": {"type": "string"}, "site": {"type": "string"}}} manager._create_prefixed_tools( - [MCPTool(name="read", description="alice view", inputSchema=alice_schema)], server, user_api_key_auth=alice + [MCPTool(name="read", description="alice view", inputSchema=alice_schema)], + server, + caller=ListedToolsCaller(user_api_key_auth=alice), ) manager._create_prefixed_tools( - [MCPTool(name="read", description="bob view", inputSchema=bob_schema)], server, user_api_key_auth=bob + [MCPTool(name="read", description="bob view", inputSchema=bob_schema)], + server, + caller=ListedToolsCaller(user_api_key_auth=bob), ) - alice_tool = manager.get_listed_tool(server, "srv-read", alice) - bob_tool = manager.get_listed_tool(server, "srv-read", bob) - assert alice_tool is not None and (alice_tool.description, alice_tool.inputSchema) == ("alice view", alice_schema) + alice_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=alice)) + bob_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=bob)) + assert alice_tool is not None and (alice_tool.description, alice_tool.inputSchema) == ( + "alice view", + alice_schema, + ) assert bob_tool is not None and (bob_tool.description, bob_tool.inputSchema) == ("bob view", bob_schema) - assert manager.get_listed_tool(server, "srv-read", UserAPIKeyAuth(user_id="carol", api_key="k")) is None + carol = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="carol", api_key="k")) + assert manager.get_listed_tool(server, "srv-read", carol) is None shared = MCPServer(server_id="shared", name="shared", transport=MCPTransport.http, url="http://shared") manager._create_prefixed_tools( - [MCPTool(name="echo", description="everyone", inputSchema={})], shared, user_api_key_auth=alice + [MCPTool(name="echo", description="everyone", inputSchema={})], + shared, + caller=ListedToolsCaller(user_api_key_auth=alice), ) - for_bob = manager.get_listed_tool(shared, "echo", bob) + for_bob = manager.get_listed_tool(shared, "echo", ListedToolsCaller(user_api_key_auth=bob)) assert for_bob is not None and for_bob.description == "everyone" + @pytest.mark.parametrize( + ("server_kwargs", "caller_a", "caller_b"), + [ + pytest.param( + {"extra_headers": ["X-Workspace"]}, + ListedToolsCaller(raw_headers={"x-workspace": "A"}), + ListedToolsCaller(raw_headers={"X-Workspace": "B"}), + id="forwarded-header", + ), + pytest.param( + {"auth_type": MCPAuth.true_passthrough}, + ListedToolsCaller(raw_headers={"authorization": "Bearer upstream-a"}), + ListedToolsCaller(raw_headers={"authorization": "Bearer upstream-b"}), + id="anonymous-passthrough-bearer", + ), + pytest.param( + {"auth_type": MCPAuth.bearer_token}, + ListedToolsCaller(mcp_auth_header="byok-a"), + ListedToolsCaller(mcp_auth_header="byok-b"), + id="per-server-auth-header", + ), + pytest.param( + {"transport": MCPTransport.stdio, "command": "srv", "env": {"WS": "${X-WS}"}}, + ListedToolsCaller(raw_headers={"X-WS": "A"}), + ListedToolsCaller(raw_headers={"X-WS": "B"}), + id="header-driven-stdio-env", + ), + ], + ) + def test_upstream_identity_inputs_keep_listed_tools_apart(self, server_kwargs, caller_a, caller_b): + """Whatever reaches upstream and can change its catalog must also split the listed-tool cache.""" + manager = MCPServerManager() + server = MCPServer( + **{"server_id": "srv", "name": "srv", "transport": MCPTransport.http, "url": "http://srv", **server_kwargs} + ) + manager._create_prefixed_tools( + [MCPTool(name="turn", description="Catalog A", inputSchema={})], server, caller=caller_a + ) + manager._create_prefixed_tools( + [MCPTool(name="turn", description="Catalog B", inputSchema={})], server, caller=caller_b + ) + + for_a = manager.get_listed_tool(server, "srv-turn", caller_a) + for_b = manager.get_listed_tool(server, "srv-turn", caller_b) + assert for_a is not None and for_a.description == "Catalog A" + assert for_b is not None and for_b.description == "Catalog B" + assert manager.get_listed_tool(server, "srv-turn", ListedToolsCaller()) is None + + def test_shared_server_ignores_headers_it_never_forwards(self): + manager = MCPServerManager() + server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") + manager._create_prefixed_tools( + [MCPTool(name="turn", description="everyone", inputSchema={})], + server, + caller=ListedToolsCaller(raw_headers={"authorization": "Bearer sk-litellm", "x-workspace": "A"}), + ) + + other = ListedToolsCaller(raw_headers={"authorization": "Bearer sk-other", "x-workspace": "B"}) + listed = manager.get_listed_tool(server, "turn", other) + assert listed is not None and listed.description == "everyone" + + @pytest.mark.asyncio + async def test_call_tool_hands_hooks_the_catalog_the_same_forwarded_headers_listed(self): + """Interleaved callers on a forwarded-header server: the hook must see the caller's own catalog.""" + manager = MCPServerManager() + server = MCPServer( + server_id="catalog", + name="catalog", + transport=MCPTransport.http, + url="http://catalog", + extra_headers=["X-Workspace"], + ) + manager.registry = {"catalog": server} + catalogs = { + "A": [ + MCPTool( + name="turn", description="Catalog A", inputSchema={"properties": {"turn": {"description": "A"}}} + ) + ], + "B": [ + MCPTool( + name="turn", description="Catalog B", inputSchema={"properties": {"turn": {"description": "B"}}} + ) + ], + } + mock_client = AsyncMock() + mock_client.call_tool.return_value = MagicMock(spec=CallToolResult, content=[], isError=False) + manager._create_mcp_client = AsyncMock(return_value=mock_client) + manager._fetch_tools_with_timeout = AsyncMock(side_effect=lambda client, name: catalogs[client.workspace]) + for workspace in ("A", "B"): + manager._create_mcp_client.return_value.workspace = workspace + await manager._get_tools_from_server( + server=server, + extra_headers={"X-Workspace": workspace}, + raw_headers={"x-workspace": workspace, "authorization": "Bearer sk-litellm"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm", user_id="shared-key"), + ) + + 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) + await manager.call_tool( + server_name="catalog", + name="catalog-turn", + arguments={"turn": "A-1"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm", user_id="shared-key"), + proxy_logging_obj=proxy_logging_obj, + raw_headers={"x-workspace": "A", "authorization": "Bearer sk-litellm"}, + ) + + 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"]) == ( + "Catalog A", + {"properties": {"turn": {"description": "A"}}}, + ) + def test_per_caller_listed_tools_evict_oldest_caller_and_keep_shared(self): from litellm.proxy._experimental.mcp_server.mcp_server_manager import _LISTED_TOOLS_CALLERS_PER_SERVER @@ -6767,20 +6896,25 @@ class TestMCPServerManager: auth_type=MCPAuth.oauth2_token_exchange, ) manager._create_prefixed_tools([MCPTool(name="read", description="shared", inputSchema={})], server) - callers = [UserAPIKeyAuth(user_id=f"u{i}", api_key=f"k{i}") for i in range(_LISTED_TOOLS_CALLERS_PER_SERVER + 1)] + callers = [ + ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id=f"u{i}", api_key=f"k{i}")) + for i in range(_LISTED_TOOLS_CALLERS_PER_SERVER + 1) + ] for caller in callers: manager._create_prefixed_tools( - [MCPTool(name="read", description=caller.user_id, inputSchema={})], server, user_api_key_auth=caller + [MCPTool(name="read", description=caller.user_api_key_auth.user_id, inputSchema={})], + server, + caller=caller, ) manager._create_prefixed_tools( - [MCPTool(name="read", description="u1 again", inputSchema={})], server, user_api_key_auth=callers[1] + [MCPTool(name="read", description="u1 again", inputSchema={})], server, caller=callers[1] ) assert manager.get_listed_tool(server, "srv-read", callers[0]) is None second = manager.get_listed_tool(server, "srv-read", callers[1]) assert second is not None and second.description == "u1 again" newest = manager.get_listed_tool(server, "srv-read", callers[-1]) - assert newest is not None and newest.description == callers[-1].user_id + assert newest is not None and newest.description == callers[-1].user_api_key_auth.user_id assert len(manager._listed_tools_by_server_id[server.server_id]) == _LISTED_TOOLS_CALLERS_PER_SERVER + 1 shared = manager.get_listed_tool(server, "srv-read") assert shared is not None and shared.description == "shared" 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 c9373de8f82..a679c56f561 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 @@ -10,7 +10,9 @@ import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.exceptions import Timeout as LitellmTimeout +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.proxy._types import UserAPIKeyAuth @@ -24,6 +26,7 @@ from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import ( agent_365_authorization_servers, agent_365_scopes_supported, ) +from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, @@ -346,32 +349,47 @@ class TestAllowFlow: class TestConversationId: + """One MCP session is one client conversation, so every tool call it carries must share the + conversationId Agent 365 sees; the per-call id is only for stateless calls without a session.""" + @pytest.mark.asyncio - async def test_request_call_id_beats_logging_obj_and_client_header(self): + async def test_two_calls_in_one_session_share_the_conversation_id(self): + handler: Final = FakeHandler([_token_response(), _allow_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + for call_id in ("call-1", "call-2"): + await _run( + guardrail, + _mcp_data(litellm_call_id=call_id, litellm_logging_obj=_logging_obj(call_id, mcp_session_id="sess-A")), + ) + assert [call.json["conversationId"] for call in handler.calls[1:]] == ["sess-A", "sess-A"] + + @pytest.mark.asyncio + async def test_server_recorded_session_beats_the_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"] == "sess-from-logging" + + @pytest.mark.asyncio + async def test_sessionless_call_falls_back_to_the_request_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-from-data", - litellm_logging_obj=_logging_obj("call-id-from-logging", mcp_session_id="sess-from-logging"), + litellm_logging_obj=_logging_obj("call-id-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): + async def test_sessionless_call_without_request_call_id_uses_the_logging_call_id(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")) + data: Final = _mcp_data(metadata={"headers": {}}, litellm_logging_obj=_logging_obj("call-id-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" + assert handler.calls[1].json["conversationId"] == "call-id-from-logging" @pytest.mark.asyncio async def test_session_id_header_case_insensitive(self): @@ -949,6 +967,56 @@ class TestVeriaHardening: assert records[0]["guardrail_status"] == "guardrail_failed_to_respond" +class _ArgumentMasker(CustomGuardrail): + """Sequential pre_mcp_call guardrail that redacts a marker in the tool arguments the way a content + filter configured with a MASK action does.""" + + def __init__(self, guardrail_name: str) -> None: + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=[GuardrailEventHooks.pre_mcp_call], + event_hook=GuardrailEventHooks.pre_mcp_call, + default_on=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + masked: Final = { + key: value.replace("REWRITE_ME", "[REWRITE_ME_REDACTED]") if isinstance(value, str) else value + for key, value in data["mcp_arguments"].items() + } + data["mcp_arguments"] = masked + data["modified_arguments"] = masked + return data + + +class TestFinalArgumentsEvaluated: + """Agent 365 must judge the arguments that reach the upstream tool. A sibling guardrail that rewrites + them must not be able to slip a different argument state past the verdict, whichever way the two + are ordered in the guardrails list.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("agent_365_first", [True, False], ids=["agent_365_then_masker", "masker_then_agent_365"]) + async def test_agent_365_receives_the_arguments_sent_upstream(self, agent_365_first: bool): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + masker: Final = _ArgumentMasker("arg-rewrite") + registered: Final = (guardrail, masker) if agent_365_first else (masker, guardrail) + for callback in registered: + litellm.logging_callback_manager.add_litellm_callback(callback) + data: Final = _mcp_data(mcp_arguments={"turn": "please REWRITE_ME now"}) + try: + result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=_user(), data=data, call_type="call_mcp_tool" + ) + finally: + for callback in registered: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, callback, require_self=False + ) + assert result["modified_arguments"] == {"turn": "please [REWRITE_ME_REDACTED] now"} + assert handler.calls[1].json["arguments"] == {"turn": "please [REWRITE_ME_REDACTED] now"} + + ENTRA_ISSUER: Final = "https://login.microsoftonline.com/tenant-abc/v2.0" GATEWAY_SCOPE: Final = "api://gateway-app/access_as_user" @@ -1041,7 +1109,8 @@ class TestAgent365AuthorizationServers: ): 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,) + assert agent_365_authorization_servers(server, None) == () + assert agent_365_scopes_supported(_mcp_server(scopes=None), None) == () finally: litellm.logging_callback_manager.remove_callback_from_list_by_object( litellm.callbacks, guardrail, require_self=False