diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index bafe33d0a6b..5b0bcc254b2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -5,13 +5,14 @@ import secrets import time from collections.abc import Callable, Mapping from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -81,6 +82,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import ( + agent_365_authorization_servers, + agent_365_scopes_supported, +) from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod @@ -2398,7 +2403,7 @@ async def _build_oauth_protected_resource_response( request: Request, mcp_server_name: str | None, use_standard_pattern: bool, -) -> dict: +) -> Mapping[str, object]: """ Build OAuth protected resource response with the appropriate URL pattern. @@ -2497,6 +2502,15 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response + agent_365_issuers: Final = agent_365_authorization_servers(mcp_server, None) if mcp_server else () + if mcp_server is not None and agent_365_issuers: + agent_365_metadata: Final[_ProtectedResourceMetadata] = { + "authorization_servers": agent_365_issuers, + "resource": resource_url, + "scopes_supported": agent_365_scopes_supported(mcp_server, None), + } + return agent_365_metadata + if mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], @@ -2516,6 +2530,12 @@ async def _build_oauth_protected_resource_response( } +class _ProtectedResourceMetadata(TypedDict): + authorization_servers: ReadOnly[tuple[str, ...]] + resource: ReadOnly[str] + scopes_supported: ReadOnly[tuple[str, ...]] + + def _obo_protected_resource_response(mcp_server: MCPServer | None, resource_url: str) -> dict | None: """The OBO (token_exchange) PRM, or None when this server is not OBO / no issuer is configured. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3291d5effc2..d439579a3ab 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -242,6 +242,9 @@ _user_env_vars_cache: Final[dict[tuple[str, str], tuple[dict[str, str], float]]] _USER_ENV_VARS_CACHE_TTL: Final = 60 # seconds _USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth +_NO_LISTED_TOOLS: Final[Mapping[str | None, Mapping[str, MCPTool]]] = MappingProxyType({}) +_LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 + # 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 @@ -1947,6 +1950,9 @@ class MCPServerManager: "gmail_send_email": "zapier_mcp_server", } """ + self._listed_tools_by_server_id: dict[ + str, Mapping[str | None, Mapping[str, MCPTool]] + ] = {} # mutable-ok: refreshed per tools/list 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 @@ -2642,7 +2648,7 @@ class MCPServerManager: self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) - self._invalidate_discovery_lists(server_id) + self._invalidate_server_definition_caches(server_id) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, @@ -2844,7 +2850,7 @@ class MCPServerManager: global_mcp_tool_registry, ) - self._invalidate_discovery_lists(server.server_id) + self._invalidate_server_definition_caches(server.server_id) prefix_root: Final = normalize_server_name(get_server_prefix(server)) if server.spec_path and prefix_root: openapi_key_prefix: Final = prefix_root + MCP_TOOL_PREFIX_SEPARATOR @@ -3221,7 +3227,7 @@ class MCPServerManager: # env_vars_are_encrypted=False. new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) self._assign_unique_short_prefix(new_server) - self._invalidate_discovery_lists(mcp_server.server_id) + self._invalidate_server_definition_caches(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -3258,7 +3264,7 @@ class MCPServerManager: previous_server=self.registry[mcp_server.server_id], ) self._assign_unique_short_prefix(new_server) - self._invalidate_discovery_lists(mcp_server.server_id) + self._invalidate_server_definition_caches(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -4092,6 +4098,7 @@ class MCPServerManager: oauth2_headers: dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, raw_headers: Mapping[str, str] | None = None, + resource_metadata_url: str | None = None, ) -> None: """Mint an exchange-backed server's upstream credential at the transport edge. @@ -4129,7 +4136,9 @@ class MCPServerManager: if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): return if subject_token is None and isinstance(spec.config, TokenExchangeConfig): - raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path()) + raise_token_exchange_challenge( + resolved_server, root_path=get_request_root_path(), resource_metadata_url=resource_metadata_url + ) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return @@ -4139,6 +4148,7 @@ class MCPServerManager: resolved_server, root_path=get_request_root_path(), claims=err.unauthorized.claims, + resource_metadata_url=resource_metadata_url, ) raise_public(err) @@ -4431,29 +4441,25 @@ class MCPServerManager: if server.spec_path: # OpenAPI tools were stored in the registry under the prefix # active at registration time — fetch by that same prefix. - _tools: Final = global_mcp_tool_registry.list_tools(tool_prefix=get_server_prefix(server)) + registry_prefix: Final = normalize_server_name(get_server_prefix(server)) + MCP_TOOL_PREFIX_SEPARATOR + _tools: Final = global_mcp_tool_registry.list_tools(tool_prefix=registry_prefix) tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(_tools) # OpenAPI tools are stored in the registry with their prefix already # 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 + 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) + 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) + prefixed_or_original_tools: Final = self._create_prefixed_tools( + tools, server, add_prefix=add_prefix, user_api_key_auth=user_api_key_auth + ) return prefixed_or_original_tools @@ -4497,6 +4503,40 @@ class MCPServerManager: self._resource_discovery_cache.invalidate(server_id) self._template_discovery_cache.invalidate(server_id) + def _invalidate_server_definition_caches(self, server_id: str) -> None: + self._invalidate_discovery_lists(server_id) + self._listed_tools_by_server_id.pop(server_id, None) + + def _discovers_per_caller(self, server: MCPServer) -> bool: + return ( + server.requires_per_user_auth + or self._references_per_user_env_var(server) + or server.delegate_auth_to_upstream + 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): + return None + material: Final = json.dumps((user_api_key_auth.user_id, user_api_key_auth.api_key), separators=(",", ":")) + return hashlib.sha256(material.encode()).hexdigest() + + def _record_listed_tools( + self, server: MCPServer, tools: Sequence[MCPTool], user_api_key_auth: UserAPIKeyAuth | None + ) -> None: + identity: Final = self._listed_tools_identity(server, user_api_key_auth) + 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) + callers: Final = tuple((key, value) for key, value in existing.items() if key not in (None, identity)) + evicted: Final = 0 if identity is None else max(len(callers) + 1 - _LISTED_TOOLS_CALLERS_PER_SERVER, 0) + entries: Final = ( + *(() if shared is None else ((None, shared),)), + *callers[evicted:], + (identity, listing), + ) + self._listed_tools_by_server_id[server.server_id] = MappingProxyType(dict(entries)) + def _discovery_key( self, server: MCPServer, @@ -4507,12 +4547,7 @@ class MCPServerManager: subject_token: str | None, credential_fingerprint: str | None = None, ) -> _DiscoveryKey: - per_user: Final = ( - server.requires_per_user_auth - or self._references_per_user_env_var(server) - or server.delegate_auth_to_upstream - or server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) - ) + per_user: Final = self._discovers_per_caller(server) if not (per_user or mcp_auth_header or extra_headers or stdio_env or subject_token): return server.server_id, None identity: Final = ( @@ -5325,7 +5360,13 @@ class MCPServerManager: "attempts; the 3-character prefix space is too crowded." ) - def _create_prefixed_tools(self, tools: list[MCPTool], server: MCPServer, add_prefix: bool = True) -> list[MCPTool]: + def _create_prefixed_tools( + self, + tools: list[MCPTool], + server: MCPServer, + add_prefix: bool = True, + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> list[MCPTool]: """ Create prefixed tools and update tool mapping. @@ -5357,9 +5398,19 @@ 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) 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) + listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity) + if not listed: + return None + return listed.get(name) or listed.get(strip_known_server_prefix(name, server)) + def _create_prefixed_prompts( self, prompts: Sequence[Prompt], server: MCPServer, add_prefix: bool = True ) -> list[Prompt]: @@ -5598,6 +5649,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. @@ -5611,6 +5663,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 @@ -5664,6 +5719,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 @@ -5718,6 +5775,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. @@ -5732,6 +5790,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(), ) @@ -6328,6 +6388,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), ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -6343,6 +6404,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), ) tasks.append(during_hook_task) @@ -6610,7 +6672,7 @@ class MCPServerManager: for server_id in previous_registry.keys() | registered_registry.keys(): if previous_registry.get(server_id) != registered_registry.get(server_id): - self._invalidate_discovery_lists(server_id) + self._invalidate_server_definition_caches(server_id) self.registry = registered_registry # A discovery task may have published into ``previous_registry`` while # this replacement was being staged. Reconcile every published entry diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 39865a35ec6..50eca9de2b5 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -199,7 +199,7 @@ def get_route_relative_request_path(scope: Scope) -> str: :func:`litellm.proxy.auth.auth_utils.get_request_route`, which the rest of the MCP auth path already routes through, so ``/litellmfoo`` is not truncated under ``root_path=/litellm``.""" raw_path = str(scope.get("_original_path") or scope.get("path", "") or "") - root_path = str(scope.get("app_root_path") or scope.get("root_path") or "").rstrip("/") + root_path = str(scope.get("app_root_path", scope.get("root_path")) or "").rstrip("/") if root_path and (raw_path == root_path or raw_path.startswith(f"{root_path}/")): return raw_path[len(root_path) :] return raw_path diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 77979a15199..2401ad4a776 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -358,6 +358,7 @@ def raise_token_exchange_challenge( *, root_path: str, claims: str | None = None, + resource_metadata_url: str | None = None, ) -> NoReturn: """Raise the RFC 9728 / RFC 6750 challenge an OBO (``token_exchange``) server returns when the caller's subject token is missing or the IdP rejected it. @@ -375,8 +376,13 @@ def raise_token_exchange_challenge( ``error="invalid_token"`` and is byte-identical to the static one. Both the error value (one of two literals) and the base64 claims draw from a fixed alphabet, so nothing from the IdP body reaches the header unescaped. + + ``resource_metadata_url`` overrides the derived path with the absolute metadata URL matching the + route spelling the request arrived on: RFC 9728 §3.3 clients reject a ``resource`` that differs + from the URL they connected to, and a ``/{server}/mcp`` connect must not be sent to the + ``/mcp/{server}`` document. """ - resource_metadata: Final = oauth_protected_resource_path(root_path, server) + resource_metadata: Final = resource_metadata_url or oauth_protected_resource_path(root_path, server) encoded_claims: Final = base64.b64encode(claims.encode()).decode() if claims else None error: Final = "insufficient_claims" if encoded_claims else "invalid_token" error_description: Final = ( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fc87db69e16..62e13675a8b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -57,6 +57,7 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_byok_www_authenticate, + get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, get_route_relative_request_path, well_known_root_suffix, @@ -80,12 +81,17 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import ( + agent_365_authorization_servers, + agent_365_subject_token_present, +) from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, get_chain_id_from_headers, ) from litellm.types.mcp import MCPAuth, MCPSpecVersion from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer +from litellm.types.mcp_server.tool_registry import MCPTool as RegisteredTool from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import Rules, client, function_setup @@ -2777,6 +2783,9 @@ 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 _resolve_display_name_to_original( name: str, allowed_mcp_servers: list[MCPServer], @@ -3115,6 +3124,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), ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -3180,7 +3190,8 @@ if MCP_AVAILABLE: # not in the registry either, `_handle_local_mcp_tool` below reports # 404 and nothing runs, so demanding a server here would turn every # unknown tool name into a misleading 503. - if global_mcp_tool_registry.get_tool(original_tool_name) is not None: + registered_local_tool: Final = global_mcp_tool_registry.get_tool(original_tool_name) + if registered_local_tool is not None: # `mcp_server` is None here because the tool name is not in the # tool -> server mapping, but the name still carries a prefix # that the server-level check above compared against the @@ -3221,6 +3232,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), ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args @@ -4129,8 +4141,19 @@ if MCP_AVAILABLE: # (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata # so the client discovers the IdP, SSOs, and retries with a subject token, which LiteLLM # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the - # header lost, so the discovery flow needs this pre-emptive challenge. - if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers: + # header lost, so the discovery flow needs this pre-emptive challenge. Servers gated by an + # Agent 365 guardrail (OBO to the evaluate API) get the same challenge, also when the only + # 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. + 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,) + and not agent_365_subject_token_present(oauth2_headers) + and agent_365_authorization_servers(server, user_api_key_auth) + ) + ): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph raise_token_exchange_challenge, ) @@ -4138,7 +4161,11 @@ if MCP_AVAILABLE: get_request_root_path, ) - raise_token_exchange_challenge(server, root_path=get_request_root_path()) + raise_token_exchange_challenge( + server, + root_path=get_request_root_path(), + resource_metadata_url=get_passthrough_resource_metadata_url(scope=scope, server_name=server_name), + ) # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run # the exchange here at the transport edge, so a rejected subject raises the RFC 9728 @@ -4163,6 +4190,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, user_api_key_auth=user_api_key_auth, raw_headers=raw_headers, + resource_metadata_url=get_passthrough_resource_metadata_url(scope=scope, server_name=server_name), ) # Pass-through OAuth: when the admin has opted a server into diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index cb7a18cd107..58af43f04a3 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -9932,7 +9932,7 @@ }, "unreachable_fallback": { "default": "fail_closed", - "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.", + "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.", "enum": [ "fail_closed", "fail_open" @@ -10894,6 +10894,18 @@ "description": "Custom advisory message template used when on_flagged='inject_system_message'. Must contain a {reason} placeholder. Defaults to a generic advisory message if unset.", "title": "Advisory System Message" }, + "agent_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Agent identity reported to Agent 365 with every tool evaluation. When unset, the caller's key alias is used.", + "title": "Agent Id" + }, "akto_account_id": { "anyOf": [ { @@ -11396,6 +11408,30 @@ "title": "Chunk Budget Chars", "type": "integer" }, + "client_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Client id of the gateway's Entra app registration (a confidential client). Falls back to the AGENT365_CLIENT_ID environment variable.", + "title": "Client Id" + }, + "client_secret": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.", + "title": "Client Secret" + }, "confidence_threshold": { "default": 0.5, "default_value": 0.5, @@ -12436,6 +12472,18 @@ "description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.", "title": "Realtime Violation Message" }, + "resource_app_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Application id of the Agent 365 resource the OBO token is minted for. Defaults to the production resource ea9ffc3e-8a23-4a7d-836d-234d7c7565c1; the Test and PreProd environments use a different id. Falls back to the AGENT365_RESOURCE_APP_ID environment variable.", + "title": "Resource App Id" + }, "rules": { "anyOf": [ { @@ -12673,6 +12721,18 @@ "description": "The ID of your Model Armor template", "title": "Template Id" }, + "tenant_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Entra tenant id used for the On-Behalf-Of token exchange. Falls back to the AGENT365_TENANT_ID environment variable.", + "title": "Tenant Id" + }, "timeout": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py new file mode 100644 index 00000000000..9aacdec0602 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py @@ -0,0 +1,63 @@ +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( + AGENT_365_PROD_API_BASE, + AGENT_365_PROD_RESOURCE_APP_ID, +) + +from .agent_365 import Agent365Guardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> Agent365Guardrail: + import litellm + from litellm.secret_managers.main import get_secret_str + + tenant_id: Final = litellm_params.tenant_id or get_secret_str("AGENT365_TENANT_ID") + client_id: Final = litellm_params.client_id or get_secret_str("AGENT365_CLIENT_ID") + client_secret: Final = ( + litellm_params.client_secret or litellm_params.api_key or get_secret_str("AGENT365_CLIENT_SECRET") + ) + api_base: Final = litellm_params.api_base or get_secret_str("AGENT365_API_BASE") + resource_app_id: Final = litellm_params.resource_app_id or get_secret_str("AGENT365_RESOURCE_APP_ID") + + if not tenant_id: + raise ValueError("Microsoft Agent 365: tenant_id is required") + if not client_id: + raise ValueError("Microsoft Agent 365: client_id is required") + if not client_secret: + raise ValueError( + "Microsoft Agent 365: client secret is required. Set client_secret, api_key, or AGENT365_CLIENT_SECRET" + ) + + guardrail_name: Final = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Microsoft Agent 365: guardrail_name is required") + + agent_365_guardrail: Final = Agent365Guardrail( + guardrail_name=guardrail_name, + tenant_id=tenant_id, + client_id=client_id, + client_secret=client_secret, + api_base=api_base or AGENT_365_PROD_API_BASE, + resource_app_id=resource_app_id or AGENT_365_PROD_RESOURCE_APP_ID, + agent_id=litellm_params.agent_id, + request_timeout=litellm_params.timeout if litellm_params.timeout is not None else 10.0, + unreachable_fallback=litellm_params.unreachable_fallback, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(agent_365_guardrail) + return agent_365_guardrail + + +guardrail_initializer_registry: Final = { # mutable-ok: registry auto-discovery requires a dict instance + SupportedGuardrailIntegrations.AGENT_365.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { # mutable-ok: registry auto-discovery requires a dict instance + SupportedGuardrailIntegrations.AGENT_365.value: Agent365Guardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py new file mode 100644 index 00000000000..58d8ba5f865 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -0,0 +1,695 @@ +"""Microsoft Agent 365 governance guardrail for MCP tool calls. + +Before the gateway executes an MCP tool, the pending call is sent to the +Agent 365 tool-evaluation endpoint, where Microsoft Defender scores it and +Agent 365 records it for observability. The returned allow/block verdict is +enforced here. Authentication is the Entra On-Behalf-Of flow: the caller's +incoming bearer token (audienced to this gateway's app registration) is +exchanged for a delegated Agent 365 token, so Defender evaluates and audits +as the signed-in user. +""" + +import hashlib +import threading +import time +import uuid +from collections import OrderedDict +from collections.abc import Mapping +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 ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + 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, + AGENT_365_SCOPE_NAME, + Agent365GuardrailConfigModel, +) + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + 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" +GATEWAY_SCOPE_TEMPLATE: Final = "api://{client_id}/access_as_user" +_GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset( + {"invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"} +) +_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 + + +def _parse_expires_in(raw: object) -> float: + if not isinstance(raw, (int, float, str)): + return _DEFAULT_TOKEN_TTL_SECONDS + try: + return float(raw) + except ValueError: + 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] + message: ReadOnly[str | None] + + +class _EvaluateResponse(TypedDict, total=False): + allowed: ReadOnly[bool] + defender: ReadOnly[_DefenderResult] + correlationId: 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): + error: ReadOnly[str] + message: ReadOnly[str] + tool: ReadOnly[str] + + +class _BlockedDetail(TypedDict): + error: ReadOnly[str] + message: ReadOnly[str] + tool: ReadOnly[str] + correlation_id: ReadOnly[str | None] + + +class Agent365TokenExchangeError(Exception): + def __init__(self, status_code: int, error_code: str, description: str) -> None: + super().__init__(f"{error_code}: {description}") + self.status_code = status_code + self.error_code = error_code + self.description = description + + +class Agent365MalformedResponseError(Exception): + pass + + +class Agent365ThrottledError(Exception): + def __init__(self, status_code: int) -> None: + super().__init__(f"HTTP {status_code}") + self.status_code = status_code + + +class Agent365Guardrail(CustomGuardrail): + """Pre-MCP-call guardrail enforcing Microsoft Agent 365 tool-evaluation verdicts.""" + + records_own_guardrail_information: ClassVar[bool] = True + + def __init__( + self, + guardrail_name: str, + tenant_id: str, + client_id: str, + client_secret: str, + api_base: str = AGENT_365_PROD_API_BASE, + resource_app_id: str = AGENT_365_PROD_RESOURCE_APP_ID, + agent_id: str | None = None, + request_timeout: float = 10.0, + unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + async_handler: AsyncHTTPHandler | None = None, + **kwargs, # noqa: ANN003 # kwargs-ok: forwarded verbatim to CustomGuardrail (event_hook, default_on) + ) -> None: + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=self.get_supported_event_hooks(), + **kwargs, + ) + self.guardrail_provider = "agent_365" + self.tenant_id = tenant_id + self.client_id = client_id + self.client_secret = client_secret + self.api_base = api_base.rstrip("/") + self.resource_app_id = resource_app_id + self.agent_id = agent_id + self.request_timeout = request_timeout + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( + "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" + ) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + self._obo_token_cache: OrderedDict[str, tuple[str, float]] = OrderedDict() # mutable-ok: lock-guarded LRU + self._obo_cache_lock = threading.Lock() + verbose_proxy_logger.info("Initialized Microsoft Agent 365 guardrail: %s", guardrail_name) + + @staticmethod + def get_config_model() -> "type[GuardrailConfigModel] | None": + return Agent365GuardrailConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail contract + return [GuardrailEventHooks.pre_mcp_call] # mutable-ok: CustomGuardrail contract expects a list + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: "DualCache", + data: dict, # mutable-ok: hook contract; guardrail logging appends into the request metadata in place + call_type: str, + ) -> Exception | str | dict | None: # mutable-ok: CustomGuardrail.async_pre_call_hook contract + if call_type not in _MCP_CALL_TYPES: + return data + if "mcp_tool_name" not in data: + return data + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_mcp_call) is not True: + return data + + tool_name: Final = str(data.get("mcp_tool_name") or "") + assertion: Final = entra_assertion(data.get("incoming_bearer_token")) + if assertion is None: + self._handle_caller_fault( + data=data, + tool_name=tool_name, + status_code=401, + reason=( + "the caller did not present an Entra bearer token; the Agent 365 guardrail " + "authorizes tool calls On-Behalf-Of the signed-in user" + ), + ) + + try: + obo_token: Final = await self._get_obo_token(assertion) + except Agent365TokenExchangeError as exc: + if exc.error_code in _GATEWAY_OWNED_TOKEN_ERRORS: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=( + f"Entra rejected the gateway's own Agent 365 credentials ({exc.error_code}); " + "check the guardrail's client_id, client_secret and resource_app_id" + ), + ) + self._handle_caller_fault( + data=data, + tool_name=tool_name, + status_code=401, + reason=f"the Entra On-Behalf-Of token exchange was rejected ({exc.error_code})", + ) + except Agent365ThrottledError as exc: + self._handle_throttled( + data=data, + tool_name=tool_name, + reason=f"the Entra token endpoint returned HTTP {exc.status_code}", + latency_ms=None, + ) + except (httpx.HTTPError, LitellmTimeout, TimeoutError) as exc: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=f"the Entra token endpoint could not be reached ({type(exc).__name__})", + ) + except Agent365MalformedResponseError as exc: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=str(exc), + ) + + start: Final = time.perf_counter() + try: + response: Final = await self._post_allowing_error_status( + url=f"{self.api_base}{EVALUATE_PATH}", + json=self._build_evaluate_payload(data=data, user_api_key_dict=user_api_key_dict), + headers={"Authorization": f"Bearer {obo_token}"}, # mutable-ok: httpx header dict + ) + except (httpx.HTTPError, LitellmTimeout, TimeoutError) as exc: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=f"the Agent 365 endpoint could not be reached ({type(exc).__name__})", + ) + latency_ms: Final = (time.perf_counter() - start) * 1000.0 + fallback: Final = self._handle_evaluate_error( + data=data, tool_name=tool_name, assertion=assertion, response=response, latency_ms=latency_ms + ) + if fallback is not None: + return fallback + return self._enforce_verdict(data=data, tool_name=tool_name, response=response, latency_ms=latency_ms) + + def _handle_evaluate_error( + self, + data: dict, # mutable-ok: guardrail logging appends into the request metadata in place + tool_name: str, + assertion: str, + response: httpx.Response, + latency_ms: float, + ) -> dict | None: # mutable-ok: returns the request data dict per hook contract on fail_open + if response.status_code in (408, 429): + self._handle_throttled( + data=data, + tool_name=tool_name, + reason=f"the Agent 365 endpoint returned HTTP {response.status_code}", + latency_ms=latency_ms, + ) + if 400 <= response.status_code < 500: + if response.status_code == 401: + self._evict_obo_token(assertion) + self._record_verdict( + data=data, + verdict="Rejected", + guardrail_status="guardrail_intervened", + defender_status=None, + correlation_id=None, + latency_ms=latency_ms, + reason=f"HTTP {response.status_code}: {response.text[:512]}", + ) + rejected_detail: Final[_UnavailableDetail] = { + "error": "Agent 365 rejected the tool evaluation request", + "message": response.text[:512] + if response.status_code == 400 + else f"the Agent 365 evaluation request failed with HTTP {response.status_code}", + "tool": tool_name, + } + raise HTTPException(status_code=400, detail=rejected_detail) + if response.status_code != 200: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=f"the Agent 365 endpoint returned HTTP {response.status_code}", + ) + return None + + def _enforce_verdict( + self, + data: dict, # mutable-ok: guardrail logging appends into the request metadata in place + tool_name: str, + response: httpx.Response, + latency_ms: float, + ) -> dict: # mutable-ok: returns the request data dict per hook contract + try: + parsed_verdict: Final = response.json() + except ValueError: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason="the Agent 365 endpoint returned a non-JSON body", + ) + if not isinstance(parsed_verdict, dict): + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason="the Agent 365 endpoint returned a non-object JSON body", + ) + verdict: Final[_EvaluateResponse] = parsed_verdict + allowed: Final = verdict.get("allowed") + if not isinstance(allowed, bool): + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason="the Agent 365 endpoint returned a verdict without a boolean 'allowed' field", + ) + raw_defender: Final = verdict.get("defender") + 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_status, + correlation_id=correlation_id, + latency_ms=latency_ms, + ) + if not allowed: + blocked_detail: Final[_BlockedDetail] = { + "error": "Blocked by Microsoft Defender", + "message": ( + defender.get("message") + or f"Invocation of '{tool_name}' is blocked by Microsoft Threat Detection policies " + "configured by your administrator." + ), + "tool": tool_name, + "correlation_id": correlation_id, + } + raise HTTPException(status_code=400, detail=blocked_detail) + return data + + def _build_evaluate_payload( + self, + data: Mapping[str, object], + user_api_key_dict: "UserAPIKeyAuth", + ) -> dict[str, object]: # mutable-ok: JSON body for AsyncHTTPHandler.post, which requires dict + tool_name: Final = str(data.get("mcp_tool_name") or "") + 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 + 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.model_dump(by_alias=True, exclude_none=True), + "serverName": server_name, + "conversationId": self._resolve_conversation_id(data), + } + if isinstance(arguments, dict): + payload["arguments"] = arguments + if agent_id: + payload["agentId"] = str(agent_id) + return payload + + @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, + ) + headers: Final = metadata.get("headers") if isinstance(metadata, Mapping) else None + if isinstance(headers, Mapping): + session_id: Final = next( + (value for name, value in headers.items() if str(name).lower() == MCP_SESSION_ID_HEADER), + None, + ) + if isinstance(session_id, str) and session_id: + return session_id + return str(uuid.uuid4()) + + async def _get_obo_token(self, assertion: str) -> str: + cache_key: Final = hashlib.sha256(assertion.encode("utf-8")).hexdigest() + now: Final = time.time() + with self._obo_cache_lock: + cached: Final = self._obo_token_cache.get(cache_key) + if cached and cached[1] > now + _TOKEN_EXPIRY_SLACK_SECONDS: + self._obo_token_cache.move_to_end(cache_key) + return cached[0] + + response: Final = await self._post_allowing_error_status( + url=TOKEN_ENDPOINT_TEMPLATE.format(tenant_id=self.tenant_id), + data={ # mutable-ok: OAuth form body; AsyncHTTPHandler.post requires dict + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "client_id": self.client_id, + "client_secret": self.client_secret, + "assertion": assertion, + "scope": f"{self.resource_app_id}/{AGENT_365_SCOPE_NAME}", + "requested_token_use": "on_behalf_of", + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, # mutable-ok: httpx header dict + ) + if response.status_code in (408, 429): + raise Agent365ThrottledError(status_code=response.status_code) + if response.status_code >= 500: + raise httpx.HTTPStatusError( + f"Entra token endpoint returned {response.status_code}", + request=response.request, + response=response, + ) + try: + parsed_body: Final = response.json() + except ValueError as exc: + raise Agent365MalformedResponseError("the Entra token endpoint returned a non-JSON body") from exc + if not isinstance(parsed_body, dict): + raise Agent365MalformedResponseError("the Entra token endpoint returned a non-object JSON body") + body: Final = parsed_body + if response.status_code >= 400: + raise Agent365TokenExchangeError( + status_code=response.status_code, + error_code=str(body.get("error", "invalid_grant")), + description=str(body.get("error_description", ""))[:512], + ) + if "access_token" not in body: + raise Agent365MalformedResponseError("the Entra token endpoint returned no access_token") + raw_access_token: Final = body.get("access_token") + if not isinstance(raw_access_token, str) or not raw_access_token: + raise Agent365MalformedResponseError("the Entra token endpoint returned a non-string access_token") + access_token: Final = raw_access_token + expires_at: Final = time.time() + _parse_expires_in(body.get("expires_in", 3599)) + with self._obo_cache_lock: + self._obo_token_cache[cache_key] = (access_token, expires_at) + self._obo_token_cache.move_to_end(cache_key) + while len(self._obo_token_cache) > _OBO_CACHE_MAX_ENTRIES: + self._obo_token_cache.popitem(last=False) + return access_token + + async def _post_allowing_error_status( + self, + url: str, + headers: dict[str, str], # mutable-ok: AsyncHTTPHandler.post requires dict + data: dict[str, str] | None = None, # mutable-ok: AsyncHTTPHandler.post requires dict + json: dict[str, object] | None = None, # mutable-ok: AsyncHTTPHandler.post requires dict + ) -> httpx.Response: + try: + return await self.async_handler.post( + url=url, + data=data, + json=json, + headers=headers, + timeout=self.request_timeout, + ) + except httpx.HTTPStatusError as exc: + return exc.response + + def _handle_caller_fault( + self, + data: dict, # mutable-ok: guardrail logging appends into the request metadata in place + tool_name: str, + status_code: int, + reason: str, + ) -> NoReturn: + self._record_verdict( + data=data, + verdict="Rejected", + guardrail_status="guardrail_intervened", + defender_status=None, + correlation_id=None, + latency_ms=None, + reason=reason, + ) + caller_fault_detail: Final[_UnavailableDetail] = { + "error": "Agent 365 guardrail rejected the tool call", + "message": f"Tool call '{tool_name}' was blocked because {reason}.", + "tool": tool_name, + } + raise HTTPException(status_code=status_code, detail=caller_fault_detail) + + def _handle_throttled( + self, + data: dict, # mutable-ok: guardrail logging appends into the request metadata in place + tool_name: str, + reason: str, + latency_ms: float | None, + ) -> NoReturn: + self._record_verdict( + data=data, + verdict="Throttled", + guardrail_status="guardrail_failed_to_respond", + defender_status=None, + correlation_id=None, + latency_ms=latency_ms, + reason=reason, + ) + throttled_detail: Final[_UnavailableDetail] = { + "error": "Agent 365 guardrail could not authorize the tool call", + "message": f"Tool call '{tool_name}' was blocked because {reason}; " + "throttled evaluations block regardless of unreachable_fallback.", + "tool": tool_name, + } + raise HTTPException(status_code=503, detail=throttled_detail) + + def _evict_obo_token(self, assertion: str) -> None: + cache_key: Final = hashlib.sha256(assertion.encode("utf-8")).hexdigest() + with self._obo_cache_lock: + self._obo_token_cache.pop(cache_key, None) + + def _handle_unavailable( + self, + 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( + "Agent 365 guardrail (%s): %s; unreachable_fallback='fail_open', allowing tool call '%s' unscanned", + self.guardrail_name, + reason, + tool_name, + ) + self._record_verdict( + data=data, + verdict="Unscanned", + guardrail_status="guardrail_failed_to_respond", + defender_status=defender_status, + correlation_id=correlation_id, + latency_ms=latency_ms, + reason=reason, + ) + return data + self._record_verdict( + data=data, + verdict="Unavailable", + guardrail_status="guardrail_failed_to_respond", + defender_status=defender_status, + correlation_id=correlation_id, + latency_ms=latency_ms, + reason=reason, + ) + unavailable_detail: Final[_UnavailableDetail] = { + "error": "Agent 365 guardrail could not authorize the tool call", + "message": f"Tool call '{tool_name}' was blocked because {reason} and unreachable_fallback is " + "'fail_closed'.", + "tool": tool_name, + } + raise HTTPException(status_code=503, detail=unavailable_detail) + + def _record_verdict( + self, + data: dict[str, object], # mutable-ok: standard guardrail logging appends into the request metadata in place + verdict: str, + guardrail_status: "GuardrailStatus", + defender_status: str | None, + correlation_id: str | None, + latency_ms: float | None, + reason: str | None = None, + ) -> None: + payload: Final[dict[str, object]] = {"verdict": verdict} # mutable-ok: optional fields added below + if defender_status: + payload["defender_status"] = defender_status + if correlation_id: + payload["correlation_id"] = correlation_id + if latency_ms is not None: + payload["latency_ms"] = round(latency_ms, 1) + if reason: + payload["reason"] = reason + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=payload, + request_data=data, + guardrail_status=guardrail_status, + duration=(latency_ms / 1000.0) if latency_ms is not None else None, + 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 _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 + 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: + return () + registered: Final = tuple( + callback + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(Agent365Guardrail) + if isinstance(callback, Agent365Guardrail) + ) + if user_api_key_auth is None: + return registered + return tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) + + +def entra_assertion(value: object) -> str | None: + """``value`` when it is a compact JWS, the only bearer shape the OBO exchange accepts as its assertion. + A LiteLLM virtual key, session bearer, or opaque upstream token in ``Authorization`` yields ``None``.""" + return value if isinstance(value, str) and value.count(".") == 2 else None + + +def agent_365_subject_token_present(oauth2_headers: Mapping[str, str] | None) -> bool: + """Whether the request's ``Authorization`` carries an Entra assertion the guardrail can exchange.""" + authorization: Final = oauth2_headers.get("Authorization", "") if oauth2_headers else "" + if not authorization.lower().startswith("bearer "): + return False + return entra_assertion(authorization[len("bearer ") :].strip()) is not None + + +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.""" + return tuple( + dict.fromkeys( + ENTRA_ISSUER_TEMPLATE.format(tenant_id=g.tenant_id) + for g in _applicable_guardrails(server, user_api_key_auth) + ) + ) + + +def agent_365_scopes_supported(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]: + """Scopes the client requests from Entra for ``server``: the admin's ``scopes`` when set, otherwise the + ``access_as_user`` scope of each gating guardrail's gateway app registration (``api://``).""" + if server.scopes: + return tuple(server.scopes) + return tuple( + dict.fromkeys( + GATEWAY_SCOPE_TEMPLATE.format(client_id=g.client_id) + for g in _applicable_guardrails(server, user_api_key_auth) + ) + ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 89625021e37..7369ec13356 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1231,6 +1231,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 @@ -1452,6 +1454,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/guardrails.py b/litellm/types/guardrails.py index 69cb88bfa2f..9a6c62808ee 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -8,6 +8,9 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_valida from typing_extensions import Required, TypedDict from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS +from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( + Agent365GuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( AktoConfigModel, ) @@ -137,6 +140,7 @@ class SupportedGuardrailIntegrations(Enum): COMPRESR = "compresr" STRAIKER = "straiker" ALICE = "alice" + AGENT_365 = "agent_365" CONDUCT = "conduct" @@ -945,7 +949,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default="fail_closed", description=( "Behavior when a guardrail endpoint is unreachable due to network errors. " - "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. " + "Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. " "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed." ), ) @@ -1083,6 +1087,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, SingulrGuardrailConfigModel, + Agent365GuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: str | list[str] | Mode = Field( 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/litellm/types/proxy/guardrails/guardrail_hooks/agent_365.py b/litellm/types/proxy/guardrails/guardrail_hooks/agent_365.py new file mode 100644 index 00000000000..dd3d7fe5f74 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/agent_365.py @@ -0,0 +1,66 @@ +from typing import Final + +from pydantic import Field + +from .base import GuardrailConfigModel + +AGENT_365_PROD_API_BASE: Final = "https://agent365.svc.cloud.microsoft" +AGENT_365_PROD_RESOURCE_APP_ID: Final = "ea9ffc3e-8a23-4a7d-836d-234d7c7565c1" +AGENT_365_SCOPE_NAME: Final = "ThreatProtection.Evaluate.All" + + +class Agent365GuardrailConfigModel(GuardrailConfigModel): + tenant_id: str | None = Field( + default=None, + description=( + "Entra tenant id used for the On-Behalf-Of token exchange. " + "Falls back to the AGENT365_TENANT_ID environment variable." + ), + ) + + client_id: str | None = Field( + default=None, + description=( + "Client id of the gateway's Entra app registration (a confidential client). " + "Falls back to the AGENT365_CLIENT_ID environment variable." + ), + ) + + client_secret: str | None = Field( + default=None, + description=( + "Client secret of the gateway's Entra app registration, used to perform the " + "On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable." + ), + ) + + api_base: str | None = Field( + default=None, + description=( + "Base URL of the Microsoft Agent 365 tool-evaluation endpoint. " + f"Defaults to the production endpoint {AGENT_365_PROD_API_BASE}. " + "Falls back to the AGENT365_API_BASE environment variable." + ), + ) + + resource_app_id: str | None = Field( + default=None, + description=( + "Application id of the Agent 365 resource the OBO token is minted for. " + f"Defaults to the production resource {AGENT_365_PROD_RESOURCE_APP_ID}; " + "the Test and PreProd environments use a different id. " + "Falls back to the AGENT365_RESOURCE_APP_ID environment variable." + ), + ) + + agent_id: str | None = Field( + default=None, + description=( + "Agent identity reported to Agent 365 with every tool evaluation. " + "When unset, the caller's key alias is used." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Microsoft Agent 365" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 1b003e11993..d650c5ea6c1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -584,6 +584,22 @@ def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): assert "error_description=" in www +def test_raise_token_exchange_challenge_explicit_resource_metadata_url_wins(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_token_exchange_challenge, + ) + + with pytest.raises(HTTPException) as exc_info: + raise_token_exchange_challenge( + _server(alias="obo-srv"), + root_path="/", + resource_metadata_url="https://gw.example.com/.well-known/oauth-protected-resource/obo-srv/mcp", + ) + www = exc_info.value.headers["WWW-Authenticate"] + assert 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/obo-srv/mcp"' in www + assert "/mcp/obo-srv" not in www + + def test_raise_token_exchange_challenge_includes_server_root_path(monkeypatch): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, 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 9ea870d3210..ebb5775b14e 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 @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from fastapi.encoders import jsonable_encoder from litellm.types.mcp import MCPAuth @@ -7124,6 +7125,81 @@ async def test_build_oauth_protected_resource_response_obo_end_to_end(): global_mcp_server_manager.registry.clear() +@pytest.fixture +def agent_365_guardrail(): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.agent_365 import Agent365Guardrail + + guardrail = Agent365Guardrail( + guardrail_name="agent-365-guard", + tenant_id="tenant-abc", + client_id="client-xyz", + client_secret="secret-123", + async_handler=AsyncMock(), + event_hook="pre_mcp_call", + default_on=True, + ) + 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 + ) + + +async def _agent_365_gated_prm(scopes): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry["tools"] = MCPServer( + server_id="tools", + name="tools", + server_name="tools", + alias="tools", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + scopes=scopes, + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + try: + return await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name="tools", use_standard_pattern=True + ) + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_agent_365_gated_server_prm_names_the_entra_tenant(agent_365_guardrail): + response = await _agent_365_gated_prm(scopes=["api://gateway-app/access_as_user"]) + assert jsonable_encoder(response) == { + "authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"], + "resource": "https://litellm.example.com/mcp/tools", + "scopes_supported": ["api://gateway-app/access_as_user"], + } + + +@pytest.mark.asyncio +async def test_agent_365_prm_defaults_scopeless_server_to_the_gateway_app_scope(agent_365_guardrail): + response = await _agent_365_gated_prm(scopes=None) + assert jsonable_encoder(response) == { + "authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"], + "resource": "https://litellm.example.com/mcp/tools", + "scopes_supported": ["api://client-xyz/access_as_user"], + } + + 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_guardrail_usage_monitor.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py index 24e6d2de10d..956f87f0da1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py @@ -48,6 +48,7 @@ def _bare_manager() -> MOD.MCPServerManager: reaches the guardrail hooks; they have their own coverage elsewhere. """ mgr = MOD.MCPServerManager.__new__(MOD.MCPServerManager) + mgr._listed_tools_by_server_id = {} mgr.check_allowed_or_banned_tools = lambda name, server: True mgr.validate_allowed_params = lambda tool_name, arguments, server: None 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 7eb4396e60d..cac754feb3f 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 @@ -7109,6 +7109,8 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} start_time = datetime.now(timezone.utc) litellm_logging_obj, _ = function_setup( @@ -7159,6 +7161,99 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): assert litellm_logging_obj.model == "MCP: list_pets" +@pytest.mark.asyncio +async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_call_hooks(): + """OpenAPI-generated tools dispatch through the local registry, so the pre-call hooks must get the + registered description and input schema on that path too, even when no tools/list ran first.""" + 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", + ) + schema = {"type": "object", "properties": {"limit": {"type": "integer"}}} + mcp_module.global_mcp_tool_registry.register_tool( + name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda limit: "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-list_pets", + arguments={"limit": 10}, + 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.name, handed_tool.description, handed_tool.inputSchema) == ( + "list_pets", + "List the pets", + 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 + pre-call hooks with the metadata of the shorter operation, since that is not the one that runs.""" + 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", + ) + registry = mcp_module.global_mcp_tool_registry + registry.register_tool(name="petstore-get_pet", description="short", input_schema={}, handler=lambda: "short") + registry.register_tool( + name="petstore-petstore-get_pet", + description="long", + input_schema={"type": "object", "properties": {"petId": {"type": "integer"}}}, + handler=lambda: "long", + ) + manager = mcp_module.global_mcp_server_manager + 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), + ): + result = await mcp_module.execute_mcp_tool( + name="petstore-petstore-get_pet", + arguments={}, + allowed_mcp_servers=[petstore], + start_time=datetime.now(), + user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"), + ) + finally: + registry.unregister_tools_with_prefix("petstore-") + + handed_tool = pre_call_tool_check.call_args.kwargs["tool"] + assert (handed_tool.description, handed_tool.inputSchema) == ( + "long", + {"type": "object", "properties": {"petId": {"type": "integer"}}}, + ) + assert result.content[0].text == "long" + + @pytest.mark.asyncio async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server(): """A prefixed REST name that resolves to no tool must still dispatch to the server_id. @@ -8826,6 +8921,174 @@ class TestSingleServerPreflightReachesIdJag: preflight.assert_not_awaited() +class TestAgent365ChallengeAtConnect: + """A missing Entra bearer on an Agent 365 gated server is challenged at connect (RFC 9728), where the + WWW-Authenticate header survives, instead of only inside the tools/call JSON-RPC error.""" + + GATEWAY_SCOPE = "api://gateway-app/access_as_user" + + def _server(self, scopes: list[str] | None) -> MCPServer: + return MCPServer( + server_id="id-tools", + name="tools", + alias="tools", + server_name="tools", + url="https://tools.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + scopes=scopes, + mcp_info={"server_name": "tools"}, + ) + + @pytest.fixture + def agent_365_guardrail(self): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.agent_365 import Agent365Guardrail + + guardrail = Agent365Guardrail( + guardrail_name="agent-365-guard", + tenant_id="tenant-abc", + client_id="client-xyz", + client_secret="secret-123", + async_handler=AsyncMock(), + event_hook="pre_mcp_call", + default_on=True, + ) + 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 + ) + + async def _connect( + self, + server: MCPServer, + oauth2_headers: dict[str, str] | None, + path: str = "/mcp/tools", + mount_scope: dict[str, str] | None = None, + ) -> HTTPException | None: + from litellm.proxy._experimental.mcp_server import server as server_module + + with ( + patch.object( # test-quality-ok: route wiring must use the manager's configured server + 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=[]) + ), + ): + try: + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={ + "type": "http", + "method": "POST", + "path": path, + "scheme": "https", + "server": ("gw.example.com", 443), + "headers": [], + **(mount_scope or {}), + }, + mcp_servers=["tools"], + oauth2_headers=oauth2_headers, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + except HTTPException as challenge: + return challenge + return None + + @pytest.mark.asyncio + async def test_no_bearer_gets_the_discovery_challenge(self, agent_365_guardrail): + challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), None) + + assert challenge is not None and challenge.status_code == 401 + www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "") + assert 'error="invalid_token"' in www_authenticate + assert ( + 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"' + in www_authenticate + ) + + @pytest.mark.asyncio + async def test_legacy_route_challenge_points_at_its_own_metadata(self, agent_365_guardrail): + """RFC 9728 3.3: the metadata's ``resource`` must equal the URL the client connected to, so a + ``/{server}/mcp`` connect is sent to the ``/{server}/mcp`` document, not the ``/mcp/{server}`` one.""" + challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), None, path="/tools/mcp") + + assert challenge is not None and challenge.status_code == 401 + www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "") + assert ( + 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/tools/mcp"' + in www_authenticate + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "server_root, mount_scope", + [ + ("", {"root_path": "/mcp", "app_root_path": ""}), + ("/litellm", {"root_path": "/litellm/mcp", "app_root_path": "/litellm"}), + ], + ) + async def test_mounted_standard_route_is_challenged(self, agent_365_guardrail, server_root, mount_scope): + """``/mcp/{server}`` is served by the ``/mcp`` Mount, which moves the mount prefix into + ``root_path`` and leaves the app root (empty or SERVER_ROOT_PATH) in ``app_root_path``.""" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": server_root}): + challenge = await self._connect( + self._server([self.GATEWAY_SCOPE]), None, path=f"{server_root}/mcp/tools", mount_scope=mount_scope + ) + + assert challenge is not None and challenge.status_code == 401 + www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "") + assert ( + f'resource_metadata="https://gw.example.com{server_root}' + f'/.well-known/oauth-protected-resource{server_root}/mcp/tools"' in www_authenticate + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("path", ["/mcp", "/mcp/tools,other"]) + async def test_aggregate_route_is_not_challenged_at_connect(self, agent_365_guardrail, path): + """The per-server metadata's ``resource`` can never equal the aggregate ``/mcp`` URL the client + connected to (RFC 9728 3.3), and one guarded server must not 401 a multi-server connect, so the + Agent 365 challenge is left to tools/call there.""" + assert await self._connect(self._server([self.GATEWAY_SCOPE]), None, path=path) is None + + @pytest.mark.asyncio + async def test_entra_assertion_present_connects(self, agent_365_guardrail): + bearer = {"Authorization": "Bearer eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1LTEifQ.c2ln"} + assert await self._connect(self._server([self.GATEWAY_SCOPE]), bearer) is None + + @pytest.mark.asyncio + async def test_litellm_key_in_authorization_is_still_challenged(self, agent_365_guardrail): + """A LiteLLM virtual key admits the caller but is no Entra assertion, so the tools/call would fail + 401 inside JSON-RPC with the WWW-Authenticate header lost. The connect must challenge instead.""" + challenge = await self._connect( + self._server([self.GATEWAY_SCOPE]), {"Authorization": "Bearer sk-litellm-virtual-key"} + ) + + assert challenge is not None and challenge.status_code == 401 + www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "") + assert 'error="invalid_token"' in www_authenticate + assert ( + 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"' + in www_authenticate + ) + + @pytest.mark.asyncio + async def test_scopeless_server_is_still_challenged(self, agent_365_guardrail): + challenge = await self._connect(self._server(None), None) + + assert challenge is not None and challenge.status_code == 401 + assert 'error="invalid_token"' in (challenge.headers or {}).get("WWW-Authenticate", "") + + @pytest.mark.asyncio + async def test_no_registered_guardrail_means_no_challenge(self): + assert await self._connect(self._server([self.GATEWAY_SCOPE]), None) is None + + def _make_obo_server(alias: str) -> MCPServer: return MCPServer( server_id=f"id-{alias}", @@ -8897,7 +9160,11 @@ class TestOboPreflightScopedToAllowedServers: _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key) preflight.assert_awaited_once_with( - server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None + server=requested, + oauth2_headers=self.SUBJECT_HEADERS, + user_api_key_auth=key, + raw_headers=None, + resource_metadata_url="/.well-known/oauth-protected-resource/mcp/obo_tools", ) 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 d2987c5112e..a504946a0d8 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 @@ -6505,6 +6505,276 @@ 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 + + def test_server_definition_change_drops_listed_tools(self): + manager = MCPServerManager() + server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") + other = MCPServer(server_id="other", name="other", transport=MCPTransport.http, url="http://other") + manager._create_prefixed_tools([MCPTool(name="echo", description="old", inputSchema={})], server) + manager._create_prefixed_tools([MCPTool(name="ping", description="kept", inputSchema={})], other) + + manager._invalidate_server_definition_caches(server.server_id) + + assert manager.get_listed_tool(server, "echo") is None + kept = manager.get_listed_tool(other, "ping") + assert kept is not None and kept.description == "kept" + + @pytest.mark.asyncio + async def test_user_oauth_refresh_keeps_listed_tools(self): + """Tool definitions are server-wide, so one user's re-auth must not blank the metadata other + callers' tool calls hand to pre-call guardrails.""" + manager = MCPServerManager() + server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") + manager._create_prefixed_tools([MCPTool(name="echo", description="shared", inputSchema={})], server) + + await manager.invalidate_user_oauth_token_cache("alice", server.server_id) + + listed = manager.get_listed_tool(server, "echo") + assert listed is not None and listed.description == "shared" + + def test_per_caller_server_keeps_listed_tools_per_identity(self): + manager = MCPServerManager() + server = MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + url="http://srv", + auth_type=MCPAuth.oauth2_token_exchange, + ) + alice = UserAPIKeyAuth(user_id="alice", api_key="hashed-alice") + bob = UserAPIKeyAuth(user_id="bob", api_key="hashed-bob") + 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 + ) + manager._create_prefixed_tools( + [MCPTool(name="read", description="bob view", inputSchema=bob_schema)], server, 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) + 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 + + 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 + ) + for_bob = manager.get_listed_tool(shared, "echo", bob) + assert for_bob is not None and for_bob.description == "everyone" + + 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 + + manager = MCPServerManager() + server = MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + url="http://srv", + 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)] + for caller in callers: + manager._create_prefixed_tools( + [MCPTool(name="read", description=caller.user_id, inputSchema={})], server, user_api_key_auth=caller + ) + manager._create_prefixed_tools( + [MCPTool(name="read", description="u1 again", inputSchema={})], server, user_api_key_auth=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 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" + + @pytest.mark.asyncio + @pytest.mark.parametrize("add_prefix", [True, False]) + async def test_openapi_listing_records_listed_tools(self, add_prefix): + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + server = MCPServer( + server_id="petstore-id", + name="petstore", + alias="petstore", + transport=MCPTransport.http, + url=None, + spec_path="/spec.yaml", + ) + manager = MCPServerManager() + manager._create_mcp_client = AsyncMock(return_value=AsyncMock()) + + async def _handler(**kwargs): + return None + + global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") + global_mcp_tool_registry.register_tool( + name="petstore-list_pets", + description="List pets", + input_schema={"type": "object", "properties": {"limit": {"type": "integer"}}}, + handler=_handler, + ) + try: + listed = await manager._get_tools_from_server(server=server, add_prefix=add_prefix) + finally: + global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") + + assert [t.name for t in listed] == ["petstore-list_pets" if add_prefix else "list_pets"] + for name in ("list_pets", "petstore-list_pets"): + tool = manager.get_listed_tool(server, name) + assert tool is not None and tool.description == "List pets" + assert tool.inputSchema["properties"] == {"limit": {"type": "integer"}} + + @pytest.mark.asyncio + async def test_openapi_listing_ignores_overlapping_server_prefix(self): + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + server = MCPServer( + server_id="pet-id", + name="pet", + alias="pet", + transport=MCPTransport.http, + url=None, + spec_path="/spec.yaml", + ) + manager = MCPServerManager() + manager._create_mcp_client = AsyncMock(return_value=AsyncMock()) + + async def _handler(**kwargs): + return None + + for prefix in ("pet-", "petstore-"): + global_mcp_tool_registry.unregister_tools_with_prefix(prefix) + global_mcp_tool_registry.register_tool( + name="pet-petstore-list", + description="Local pet tool", + input_schema={"type": "object", "properties": {"limit": {"type": "integer"}}}, + handler=_handler, + ) + global_mcp_tool_registry.register_tool( + name="petstore-list", + description="Foreign petstore tool", + input_schema={"type": "object", "properties": {"status": {"type": "string"}}}, + handler=_handler, + ) + try: + listed = await manager._get_tools_from_server(server=server, add_prefix=True) + finally: + for prefix in ("pet-", "petstore-"): + global_mcp_tool_registry.unregister_tools_with_prefix(prefix) + + assert [t.name for t in listed] == ["pet-petstore-list"] + tool = manager.get_listed_tool(server, "petstore-list") + assert tool is not None and tool.description == "Local pet tool" + assert tool.inputSchema["properties"] == {"limit": {"type": "integer"}} + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 64614c094ba..c534cf73219 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -43,6 +43,8 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock(return_value={}) handle_local = AsyncMock(return_value=[]) @@ -124,6 +126,8 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): fake_tool = MagicMock() fake_tool.name = "delete_pet" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock( side_effect=HTTPException(status_code=403, detail="not allowed") @@ -186,6 +190,8 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): fake_tool = MagicMock() fake_tool.name = "list_pets" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} pre_call = AsyncMock(return_value={}) handle_local = AsyncMock(return_value=[]) @@ -270,6 +276,8 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): fake_tool = MagicMock() fake_tool.name = "get_values" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} captured: dict = {} async def handle_local(_name, _arguments): @@ -616,6 +624,8 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc if dispatch_arm == "local_registry": fake_tool = MagicMock() fake_tool.name = "list_reports" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} with ( patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server), patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool), @@ -687,6 +697,8 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st fake_tool = MagicMock() fake_tool.name = "list_reports" + fake_tool.description = "test tool" + fake_tool.input_schema = {"type": "object"} fake_tool.handler = raising_handler server = MCPServer( server_id="srv-openapi", 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 new file mode 100644 index 00000000000..c9373de8f82 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -0,0 +1,1048 @@ +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 +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.agent_365 import ( + Agent365Guardrail, + guardrail_class_registry, + guardrail_initializer_registry, + initialize_guardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import ( + agent_365_authorization_servers, + agent_365_scopes_supported, +) +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, + Agent365GuardrailConfigModel, +) + +FAKE_ASSERTION: Final = "eyJhbGciOi.eyJhdWQiOi.c2lnbmF0dXJl" +TOKEN_URL: Final = "https://login.microsoftonline.com/tenant-abc/oauth2/v2.0/token" +EVALUATE_URL: Final = f"{AGENT_365_PROD_API_BASE}/agents/tool-evaluation/evaluate" + + +def _response(status_code: int, payload: Any = None, text: str | None = None) -> httpx.Response: + request: Final = httpx.Request("POST", "https://example.test") + if payload is not None: + return httpx.Response(status_code=status_code, json=payload, request=request) + return httpx.Response(status_code=status_code, text=text or "", request=request) + + +def _token_response(access_token: str = "obo-access-token", expires_in: int = 3599) -> httpx.Response: + return _response(200, {"access_token": access_token, "expires_in": expires_in}) + + +def _allow_response(correlation_id: str = "corr-1") -> httpx.Response: + return _response( + 200, + { + "allowed": True, + "defender": {"status": "Evaluated", "verdict": "Allow", "message": None}, + "observability": {"status": "Recorded"}, + "correlationId": correlation_id, + }, + ) + + +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": 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) + self.calls: list[SimpleNamespace] = [] + + async def post(self, *, url, headers=None, data=None, json=None, timeout=None): + self.calls.append(SimpleNamespace(url=url, headers=headers, data=data, json=json, timeout=timeout)) + if not self._items: + raise AssertionError("FakeHandler ran out of programmed responses") + item = self._items.pop(0) + if isinstance(item, BaseException): + raise item + if item.status_code >= 400: + raise httpx.HTTPStatusError("error status", request=item.request, response=item) + return item + + +def _make_guardrail( + handler: FakeHandler, + *, + unreachable_fallback: str = "fail_closed", + agent_id: str | None = None, + api_base: str = AGENT_365_PROD_API_BASE, +) -> Agent365Guardrail: + return Agent365Guardrail( + guardrail_name="agent-365-guard", + tenant_id="tenant-abc", + client_id="client-xyz", + client_secret="secret-123", + api_base=api_base, + agent_id=agent_id, + unreachable_fallback=unreachable_fallback, + async_handler=handler, + event_hook="pre_mcp_call", + default_on=True, + ) + + +def _mcp_data(**overrides: Any) -> dict: + data: Final[dict] = { + "mcp_tool_name": "send_email", + "mcp_arguments": {"to": "user@example.com", "body": "hello"}, + "mcp_server_name": "outlook_mcp", + "incoming_bearer_token": FAKE_ASSERTION, + "metadata": {"headers": {"mcp-session-id": "sess-123"}}, + } + data.update(overrides) + return data + + +def _user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed-key", key_alias="my-agent-key") + + +async def _run(guardrail: Agent365Guardrail, data: dict, call_type: str = "call_mcp_tool"): + return await guardrail.async_pre_call_hook( + user_api_key_dict=_user(), + cache=None, + data=data, + call_type=call_type, + ) + + +class TestRegistryWiring: + def test_enum_member_exists(self): + assert SupportedGuardrailIntegrations.AGENT_365.value == "agent_365" + + def test_initializer_registry(self): + assert guardrail_initializer_registry["agent_365"] is initialize_guardrail + + def test_class_registry(self): + assert guardrail_class_registry["agent_365"] is Agent365Guardrail + + def test_config_model_wired(self): + assert Agent365Guardrail.get_config_model() is Agent365GuardrailConfigModel + assert Agent365GuardrailConfigModel.ui_friendly_name() == "Microsoft Agent 365" + + def test_supported_event_hooks(self): + assert Agent365Guardrail.get_supported_event_hooks() == [GuardrailEventHooks.pre_mcp_call] + + +class TestInitializeGuardrail: + def test_requires_tenant_id(self, monkeypatch): + monkeypatch.delenv("AGENT365_TENANT_ID", raising=False) + params: Final = LitellmParams( + guardrail="agent_365", + mode="pre_mcp_call", + client_id="client-xyz", + api_key="secret-123", + ) + with pytest.raises(ValueError, match="tenant_id is required"): + initialize_guardrail(params, {"guardrail_name": "a365"}) + + def test_requires_client_secret(self, monkeypatch): + monkeypatch.delenv("AGENT365_CLIENT_SECRET", raising=False) + params: Final = LitellmParams( + guardrail="agent_365", + mode="pre_mcp_call", + tenant_id="tenant-abc", + client_id="client-xyz", + ) + with pytest.raises(ValueError, match="client_secret") as exc_info: + initialize_guardrail(params, {"guardrail_name": "a365"}) + assert redact_string(str(exc_info.value)) == str(exc_info.value) + + def test_env_var_fallbacks(self, monkeypatch): + monkeypatch.delenv("AGENT365_RESOURCE_APP_ID", raising=False) + monkeypatch.setenv("AGENT365_TENANT_ID", "env-tenant") + monkeypatch.setenv("AGENT365_CLIENT_ID", "env-client") + monkeypatch.setenv("AGENT365_CLIENT_SECRET", "env-secret") + monkeypatch.setenv("AGENT365_API_BASE", "https://env.example.test") + params: Final = LitellmParams(guardrail="agent_365", mode="pre_mcp_call") + guardrail: Final = initialize_guardrail(params, {"guardrail_name": "a365-env"}) + assert guardrail.tenant_id == "env-tenant" + assert guardrail.client_id == "env-client" + assert guardrail.client_secret == "env-secret" + assert guardrail.api_base == "https://env.example.test" + assert guardrail.resource_app_id == AGENT_365_PROD_RESOURCE_APP_ID + assert guardrail.unreachable_fallback == "fail_closed" + + def test_explicit_params_win(self, monkeypatch): + monkeypatch.setenv("AGENT365_TENANT_ID", "env-tenant") + params: Final = LitellmParams( + guardrail="agent_365", + mode="pre_mcp_call", + tenant_id="param-tenant", + client_id="client-xyz", + client_secret="param-secret", + agent_id="agent-007", + unreachable_fallback="fail_open", + timeout=5, + ) + guardrail: Final = initialize_guardrail(params, {"guardrail_name": "a365-params"}) + assert guardrail.tenant_id == "param-tenant" + assert guardrail.client_secret == "param-secret" + assert guardrail.agent_id == "agent-007" + assert guardrail.unreachable_fallback == "fail_open" + assert guardrail.request_timeout == 5.0 + + def test_wrong_mode_rejected(self): + params: Final = LitellmParams( + guardrail="agent_365", + mode="post_call", + tenant_id="tenant-abc", + client_id="client-xyz", + api_key="secret-123", + ) + with pytest.raises(Exception, match="post_call"): + initialize_guardrail(params, {"guardrail_name": "a365-badmode"}) + + +def _guardrail_info(data: dict) -> dict: + entries: Final = data["metadata"]["standard_logging_guardrail_information"] + return entries[-1] + + +class TestAllowFlow: + @pytest.mark.asyncio + async def test_allowed_call_passes_through(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "success" + assert info["guardrail_provider"] == "agent_365" + assert info["guardrail_response"]["verdict"] == "Allow" + assert info["guardrail_response"]["defender_status"] == "Evaluated" + assert info["guardrail_response"]["correlation_id"] == "corr-1" + assert info["guardrail_response"]["latency_ms"] >= 0 + + @pytest.mark.asyncio + async def test_obo_exchange_form(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + token_call: Final = handler.calls[0] + assert token_call.url == TOKEN_URL + assert token_call.data["grant_type"] == "urn:ietf:params:oauth:grant-type:jwt-bearer" + assert token_call.data["requested_token_use"] == "on_behalf_of" + assert token_call.data["assertion"] == FAKE_ASSERTION + assert token_call.data["client_id"] == "client-xyz" + assert token_call.data["client_secret"] == "secret-123" + assert token_call.data["scope"] == f"{AGENT_365_PROD_RESOURCE_APP_ID}/ThreatProtection.Evaluate.All" + + @pytest.mark.asyncio + async def test_evaluate_payload(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler, agent_id="agent-007") + await _run(guardrail, _mcp_data()) + evaluate_call: Final = handler.calls[1] + assert evaluate_call.url == EVALUATE_URL + assert evaluate_call.headers["Authorization"] == "Bearer obo-access-token" + assert evaluate_call.json["tool"] == {"name": "send_email"} + assert evaluate_call.json["serverName"] == "outlook_mcp" + assert evaluate_call.json["arguments"] == {"to": "user@example.com", "body": "hello"} + 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()]) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + assert handler.calls[1].json["agentId"] == "my-agent-key" + + @pytest.mark.asyncio + async def test_non_mcp_call_type_skipped(self): + handler: Final = FakeHandler([]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + result: Final = await _run(guardrail, data, call_type="completion") + assert result is data + assert handler.calls == [] + + +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()]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data(metadata={"headers": {"Mcp-Session-Id": "sess-CASED"}}) + await _run(guardrail, data) + assert handler.calls[1].json["conversationId"] == "sess-CASED" + + @pytest.mark.asyncio + async def test_generates_uuid_when_no_identifier_available(self): + handler: Final = FakeHandler([_token_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + 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: + @pytest.mark.asyncio + async def test_blocked_call_raises_400(self): + handler: Final = FakeHandler([_token_response(), _block_response(message="Injection detected")]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Blocked by Microsoft Defender" + assert exc_info.value.detail["message"] == "Injection detected" + assert exc_info.value.detail["tool"] == "send_email" + assert exc_info.value.detail["correlation_id"] == "corr-2" + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_intervened" + assert info["guardrail_response"]["verdict"] == "Block" + + @pytest.mark.asyncio + async def test_blocked_even_with_fail_open(self): + handler: Final = FakeHandler([_token_response(), _block_response()]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + with pytest.raises(HTTPException) as exc_info: + 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")]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 400 + assert "rejected" in exc_info.value.detail["error"] + + +class TestUnreachableFallback: + @pytest.mark.asyncio + async def test_evaluate_litellm_timeout_fail_closed(self): + handler: Final = FakeHandler( + [ + _token_response(), + LitellmTimeout(message="Connection timed out", model="default-model-name", llm_provider="httpx"), + ] + ) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + async def test_evaluate_timeout_fail_closed(self): + handler: Final = FakeHandler([_token_response(), httpx.ReadTimeout("timed out")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "fail_closed" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_evaluate_timeout_fail_open(self): + handler: Final = FakeHandler([_token_response(), httpx.ReadTimeout("timed out")]) + 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" + + @pytest.mark.asyncio + async def test_evaluate_5xx_fail_closed(self): + handler: Final = FakeHandler([_token_response(), _response(502, text="bad gateway")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "502" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_missing_bearer_token_fail_closed(self): + handler: Final = FakeHandler([]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data(incoming_bearer_token=None)) + assert exc_info.value.status_code == 401 + assert handler.calls == [] + + @pytest.mark.asyncio + async def test_non_jwt_bearer_token_fail_closed(self): + handler: Final = FakeHandler([]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data(incoming_bearer_token="sk-litellm-virtual-key")) + assert exc_info.value.status_code == 401 + + @pytest.mark.asyncio + async def test_missing_bearer_token_blocks_even_fail_open(self): + handler: Final = FakeHandler([]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data(incoming_bearer_token=None) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 401 + assert handler.calls == [] + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_intervened" + assert info["guardrail_response"]["verdict"] == "Rejected" + + @pytest.mark.asyncio + async def test_obo_rejected_blocks_even_fail_open(self): + handler: Final = FakeHandler( + [_response(400, {"error": "invalid_grant", "error_description": "AADSTS50013: bad assertion"})] + ) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 401 + assert "invalid_grant" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_evaluate_4xx_blocks_even_fail_open(self): + handler: Final = FakeHandler([_token_response(), _response(403, text="obo token lacks the scope")]) + 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 + assert "403" in exc_info.value.detail["message"] + assert "lacks the scope" not in exc_info.value.detail["message"] + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_intervened" + assert info["guardrail_response"]["reason"] == "HTTP 403: obo token lacks the scope" + + @pytest.mark.asyncio + async def test_obo_rejected_fail_closed(self): + handler: Final = FakeHandler( + [_response(400, {"error": "invalid_grant", "error_description": "AADSTS50013: bad assertion"})] + ) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 401 + assert "invalid_grant" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error_code", ["invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"] + ) + async def test_gateway_credential_rejection_is_unavailable_not_a_caller_401(self, error_code: str): + handler: Final = FakeHandler( + [_response(401, {"error": error_code, "error_description": "AADSTS7000215: invalid client secret"})] + ) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 503 + assert exc_info.value.headers is None or "WWW-Authenticate" not in exc_info.value.headers + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Unavailable" + assert error_code in info["guardrail_response"]["reason"] + assert "client_secret" in info["guardrail_response"]["reason"] + + @pytest.mark.asyncio + async def test_gateway_credential_rejection_follows_fail_open(self): + handler: Final = FakeHandler( + [_response(401, {"error": "invalid_client", "error_description": "AADSTS7000215: invalid client secret"})] + ) + 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 "invalid_client" in info["guardrail_response"]["reason"] + + @pytest.mark.asyncio + async def test_obo_endpoint_5xx_fail_open(self): + handler: Final = FakeHandler([_response(503, text="entra down")]) + 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" + + +class TestOboTokenCache: + @pytest.mark.asyncio + async def test_same_assertion_reuses_token(self): + handler: Final = FakeHandler([_token_response(), _allow_response(), _allow_response()]) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + await _run(guardrail, _mcp_data()) + token_calls: Final = [c for c in handler.calls if c.url == TOKEN_URL] + assert len(token_calls) == 1 + + @pytest.mark.asyncio + async def test_different_assertions_get_distinct_tokens(self): + other_assertion: Final = "eyJhbGciOi.eyJvdGhlciI.b3RoZXJzaWc" + handler: Final = FakeHandler( + [ + _token_response(access_token="token-a"), + _allow_response(), + _token_response(access_token="token-b"), + _allow_response(), + ] + ) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + await _run(guardrail, _mcp_data(incoming_bearer_token=other_assertion)) + token_calls: Final = [c for c in handler.calls if c.url == TOKEN_URL] + assert len(token_calls) == 2 + assert handler.calls[3].headers["Authorization"] == "Bearer token-b" + + @pytest.mark.asyncio + async def test_expired_token_refreshed(self): + handler: Final = FakeHandler( + [ + _token_response(access_token="short-lived", expires_in=1), + _allow_response(), + _token_response(access_token="fresh"), + _allow_response(), + ] + ) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + await _run(guardrail, _mcp_data()) + token_calls: Final = [c for c in handler.calls if c.url == TOKEN_URL] + assert len(token_calls) == 2 + assert handler.calls[3].headers["Authorization"] == "Bearer fresh" + + +class TestEarlyPhasePassthrough: + @pytest.mark.asyncio + async def test_rest_body_shape_without_mcp_fields_skipped(self): + handler: Final = FakeHandler([]) + guardrail: Final = _make_guardrail(handler) + data: Final = { + "server_id": "266024044f9612bf481c78f6cfef1ff0", + "name": "deepwiki-read_wiki_structure", + "arguments": {"repoName": "BerriAI/litellm"}, + "metadata": {"headers": {"mcp-session-id": "sess-123"}}, + } + result: Final = await _run(guardrail, data) + assert result is data + assert handler.calls == [] + assert "standard_logging_guardrail_information" not in data["metadata"] + + +class TestRegistryDiscovery: + def test_auto_discovery_finds_agent_365(self): + from litellm.proxy.guardrails.guardrail_registry import ( + get_guardrail_class_from_hooks, + get_guardrail_initializer_from_hooks, + ) + + assert "agent_365" in get_guardrail_initializer_from_hooks() + assert get_guardrail_class_from_hooks()["agent_365"] is Agent365Guardrail + + +class TestMalformedResponses: + @pytest.mark.asyncio + async def test_obo_html_body_fail_open(self): + handler: Final = FakeHandler([_response(200, text="blocked by egress proxy")]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Unscanned" + + @pytest.mark.asyncio + async def test_obo_html_body_fail_closed(self): + handler: Final = FakeHandler([_response(200, text="outage")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "non-JSON" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_obo_non_object_json_fail_closed(self): + handler: Final = FakeHandler([_response(200, ["not", "a", "dict"])]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + async def test_evaluate_html_body_fail_open(self): + handler: Final = FakeHandler([_token_response(), _response(200, text="waf page")]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Unscanned" + + @pytest.mark.asyncio + async def test_evaluate_html_body_fail_closed(self): + handler: Final = FakeHandler([_token_response(), _response(200, text="waf page")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + async def test_evaluate_non_object_json_fail_closed(self): + handler: Final = FakeHandler([_token_response(), _response(200, "allowed")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "verdict", + [{}, {"allowed": None}, {"allowed": "true"}, {"allowed": 1}, {"allowed": "false"}], + ids=["missing", "null", "string-true", "int-one", "string-false"], + ) + async def test_evaluate_non_boolean_allowed_fail_closed(self, verdict: dict): + handler: Final = FakeHandler([_token_response(), _response(200, verdict)]) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 503 + assert "boolean 'allowed'" in exc_info.value.detail["message"] + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Unavailable" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "verdict", + [{}, {"allowed": None}, {"allowed": "true"}, {"allowed": 1}, {"allowed": "false"}], + ids=["missing", "null", "string-true", "int-one", "string-false"], + ) + async def test_evaluate_non_boolean_allowed_fail_open(self, verdict: dict): + handler: Final = FakeHandler([_token_response(), _response(200, verdict)]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Unscanned" + assert _guardrail_info(data)["guardrail_status"] == "guardrail_failed_to_respond" + + @pytest.mark.asyncio + async def test_bad_expires_in_still_allows(self): + handler: Final = FakeHandler( + [_response(200, {"access_token": "tok-1", "expires_in": "soon"}), _allow_response()] + ) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + + @pytest.mark.asyncio + async def test_obo_litellm_timeout_fail_open(self): + handler: Final = FakeHandler( + [LitellmTimeout(message="Connection timed out", model="default-model-name", llm_provider="httpx")] + ) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + assert _guardrail_info(data)["guardrail_status"] == "guardrail_failed_to_respond" + + +class TestDeltaHardening: + @pytest.mark.asyncio + async def test_non_string_access_token_fail_closed(self): + handler: Final = FakeHandler([_response(200, {"access_token": None, "expires_in": 3599})]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "access_token" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_numeric_string_expires_in_honored(self): + handler: Final = FakeHandler( + [_response(200, {"access_token": "tok-9", "expires_in": "120"}), _allow_response()] + ) + guardrail: Final = _make_guardrail(handler) + await _run(guardrail, _mcp_data()) + entries: Final = list(guardrail._obo_token_cache.values()) + assert len(entries) == 1 + assert entries[0][1] - time.time() < 200 + + @pytest.mark.asyncio + async def test_evaluate_400_records_intervention(self): + handler: Final = FakeHandler([_token_response(), _response(400, text="bad request shape")]) + guardrail: Final = _make_guardrail(handler) + 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"] == "Rejected" + + +class TestVeriaHardening: + @pytest.mark.asyncio + async def test_evaluate_401_evicts_cached_obo_token(self): + handler: Final = FakeHandler( + [ + _token_response(), + _response(401, text="token expired"), + _token_response(access_token="tok-2"), + _allow_response(), + ] + ) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException): + await _run(guardrail, _mcp_data()) + result: Final = await _run(guardrail, _mcp_data()) + assert result is not None + token_calls: Final = [c for c in handler.calls if c.url == TOKEN_URL] + assert len(token_calls) == 2 + + @pytest.mark.asyncio + async def test_evaluate_429_blocks_even_fail_open_as_throttled(self): + handler: Final = FakeHandler([_token_response(), _response(429, text="slow down")]) + 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 == 503 + assert "429" in exc_info.value.detail["message"] + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Throttled" + + @pytest.mark.asyncio + async def test_evaluate_500_is_unavailable(self): + handler: Final = FakeHandler([_token_response(), _response(500, text="oops")]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "500" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_token_endpoint_429_blocks_even_fail_open_as_throttled(self): + handler: Final = FakeHandler( + [_response(429, {"error": "temporarily_throttled", "error_description": "AADSTS90056"})] + ) + 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 == 503 + assert "429" in exc_info.value.detail["message"] + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Throttled" + + @pytest.mark.asyncio + async def test_token_endpoint_408_non_json_blocks_as_throttled(self): + handler: Final = FakeHandler([_response(408, text="Request Timeout")]) + 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 == 503 + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Throttled" + + @pytest.mark.asyncio + async def test_token_endpoint_4xx_html_stays_infra_fail_open(self): + handler: Final = FakeHandler([_response(403, text="waf block page")]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + assert _guardrail_info(data)["guardrail_response"]["verdict"] == "Unscanned" + + @pytest.mark.asyncio + async def test_entra_200_missing_access_token_is_malformed(self): + handler: Final = FakeHandler([_response(200, {"token_type": "Bearer"})]) + guardrail: Final = _make_guardrail(handler) + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, _mcp_data()) + assert exc_info.value.status_code == 503 + assert "access_token" in exc_info.value.detail["message"] + + @pytest.mark.asyncio + async def test_evaluate_5xx_fail_open_allows_unscanned_once(self): + handler: Final = FakeHandler([_token_response(), _response(502, text='{"error": "bad gateway"}')]) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + records: Final = data["metadata"]["standard_logging_guardrail_information"] + 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,) + + @pytest.mark.parametrize("scopes", [None, []], ids=["unset", "empty"]) + def test_scopeless_server_signs_in_with_the_gateway_app_scope(self, registered_guardrail, scopes): + server: Final = _mcp_server(scopes=scopes) + assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,) + assert agent_365_scopes_supported(server, None) == ("api://client-xyz/access_as_user",) + + def test_admin_scopes_override_the_default_gateway_scope(self, registered_guardrail): + assert agent_365_scopes_supported(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (GATEWAY_SCOPE,) + + def test_no_default_scope_when_no_guardrail_gates_the_server(self): + assert agent_365_scopes_supported(_mcp_server(scopes=None), 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( # test-quality-ok: key-selected guardrails read the proxy server premium global, no injection seam + "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 = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index d0afc896260..10ca58294b9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -318,6 +318,13 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + agent_365: { + provider: "Agent365", + guardrailNameSuggestion: "Microsoft Agent 365 Guardrail", + mode: "pre_mcp_call", + // MCP-only: default_on is the only activation path on the MCP hook + defaultOn: true, + }, conduct: { provider: "Conduct", guardrailNameSuggestion: "Conduct Guard", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index 9a9ab3a61d7..eb5d47d7891 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -28,6 +28,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { repelloai: "repelloai.png", straiker: "straiker.svg", alice: "alice.svg", + agent_365: "microsoft_azure.svg", conduct: "conduct.png", }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index 165bd8f9967..d88a333d6f1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -474,6 +474,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Content Moderation", "Prompt Injection", "PII", "Policy"], providerKey: "Alice", }, + { + id: "agent_365", + name: "Microsoft Agent 365", + description: + "Microsoft Agent 365 tool-call governance: Defender threat evaluation and observability for MCP tool calls, acting on behalf of the signed-in user", + category: "partner", + logo: guardrailLogoMap["Microsoft Agent 365"], + tags: ["Agentic", "MCP", "Tool Misuse", "Observability"], + providerKey: "Agent365", + }, { id: "conduct", name: "Conduct Guard", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index fb3cf8f309a..476bcd3a8ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -210,6 +210,7 @@ export const guardrailLogoMap = { "RepelloAI Argus": repelloAiLogo.src, Straiker: straikerLogo.src, Alice: aliceLogo.src, + "Microsoft Agent 365": microsoftAzureLogo.src, "Conduct Guard": conductLogo.src, } satisfies Record; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9e844b992b2..ef06e679e60 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23898,7 +23898,7 @@ export interface components { timeout?: number | null; /** * Unreachable Fallback - * @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. + * @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. * @default fail_closed * @enum {string} */ @@ -30547,6 +30547,11 @@ export interface components { * @description Custom advisory message template used when on_flagged='inject_system_message'. Must contain a {reason} placeholder. Defaults to a generic advisory message if unset. */ advisory_system_message?: string | null; + /** + * Agent Id + * @description Agent identity reported to Agent 365 with every tool evaluation. When unset, the caller's key alias is used. + */ + agent_id?: string | null; /** * Akto Account Id * @description Akto account ID for multi-tenant deployments. Env: AKTO_ACCOUNT_ID. Default: '1000000'. @@ -30748,6 +30753,16 @@ export interface components { * @default 25000 */ chunk_budget_chars: number; + /** + * Client Id + * @description Client id of the gateway's Entra app registration (a confidential client). Falls back to the AGENT365_CLIENT_ID environment variable. + */ + client_id?: string | null; + /** + * Client Secret + * @description Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable. + */ + client_secret?: string | null; /** * Confidence Threshold * @description Only block or mask when detection confidence >= this value; below threshold, allow or log_only. @@ -31181,6 +31196,11 @@ export interface components { * @description The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set. */ realtime_violation_message?: string | null; + /** + * Resource App Id + * @description Application id of the Agent 365 resource the OBO token is minted for. Defaults to the production resource ea9ffc3e-8a23-4a7d-836d-234d7c7565c1; the Test and PreProd environments use a different id. Falls back to the AGENT365_RESOURCE_APP_ID environment variable. + */ + resource_app_id?: string | null; /** * Rules * @description Ordered allow/deny rules. Patterns use regex for tool names/types and optional regex constraints on tool arguments. @@ -31282,6 +31302,11 @@ export interface components { * @description The ID of your Model Armor template */ template_id?: string | null; + /** + * Tenant Id + * @description Entra tenant id used for the On-Behalf-Of token exchange. Falls back to the AGENT365_TENANT_ID environment variable. + */ + tenant_id?: string | null; /** * Timeout * @description Per-request timeout for the guardrail provider API call (seconds). Accepts int, float, or numeric string; coerced to float on load. Each guardrail handler chooses its own default when unset.