From f8dd979d9a3229a424af35ded3c64a4b38ca4098 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:34:48 -0700 Subject: [PATCH] fix(agents): gate the managed path on an admitted policy object Ten call sites branched on `managed_agent_policy is not None`, which any MagicMock attribute satisfies, so the managed path fired on unmanaged subjects and died in Pydantic validation as a 503. Route every check through a shared helper that requires a real AgentResponse. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 14 ++++++++------ .../mcp_server/mcp_server_manager.py | 3 ++- .../auth/agent_permission_handler.py | 9 +++++---- .../agent_endpoints/auth/managed_authorization.py | 10 ++++++++++ litellm/proxy/auth/auth_checks.py | 15 ++++++++------- 5 files changed, 33 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 60241d0fc40..2da7fc72bba 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -68,6 +68,7 @@ from litellm.repositories.table_repositories import ( MCPServerRepository, ) from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -1600,7 +1601,7 @@ class MCPRequestHandler: """ from litellm.proxy.proxy_server import general_settings - if user_api_key_auth is not None and user_api_key_auth.managed_agent_policy is not None: + if managed_agent_policy(user_api_key_auth) is not None: from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers return MCPServerAccess(server_ids=await managed_agent_servers(user_api_key_auth), scope="scoped") @@ -2250,7 +2251,7 @@ class MCPRequestHandler: if not user_api_key_auth: return None - if user_api_key_auth.managed_agent_policy is not None: + if managed_agent_policy(user_api_key_auth) is not None: from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_tools return await managed_agent_tools(server_id, user_api_key_auth) @@ -3409,8 +3410,9 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None - if user_api_key_auth.managed_agent_policy is not None: - permission: Final = user_api_key_auth.managed_agent_policy.object_permission + managed: Final = managed_agent_policy(user_api_key_auth) + if managed is not None: + permission: Final = managed.object_permission return LiteLLM_ObjectPermissionTable.model_validate(permission) if permission is not None else None if prisma_client is None: @@ -3478,7 +3480,7 @@ class MCPRequestHandler: inline_tools: Final = global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions) return list({*expanded_direct_servers, *access_group_servers, *toolset_grants, *inline_tools}) except Exception as e: - if user_api_key_auth.managed_agent_policy is not None or isinstance(e, UnloadableEntitlementError): + if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError): raise verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] @@ -3551,7 +3553,7 @@ class MCPRequestHandler: agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) return list(agent_tools) if agent_tools is not None else None except Exception as e: - if user_api_key_auth.managed_agent_policy is not None or isinstance(e, UnloadableEntitlementError): + if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError): raise verbose_logger.warning("Failed to get agent tool permissions for server: %s", e) return None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e939962069c..1db525bc760 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -210,6 +210,7 @@ from litellm.types.mcp_server.mcp_server_manager import ( parse_pinned_tools, ) from litellm.types.utils import CallTypes +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -3479,7 +3480,7 @@ class MCPServerManager: 2. If admin and no object_permission, return all servers 3. Otherwise, use standard permission checks """ - if user_api_key_auth is not None and user_api_key_auth.managed_agent_policy is not None: + if managed_agent_policy(user_api_key_auth) is not None: managed: Final = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) return managed if access is None else [server for server in managed if server in access.server_ids] diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 4686796480e..49171c313f4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -29,6 +29,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy @dataclass(frozen=True, slots=True) @@ -89,7 +90,7 @@ class AgentRequestHandler: ) -> AgentAccess: """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" - if user_api_key_auth is not None and user_api_key_auth.managed_agent_policy is not None: + if managed_agent_policy(user_api_key_auth) is not None: return await _managed_actor_agent_access(user_api_key_auth) key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access(user_api_key_auth) caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth) @@ -183,7 +184,7 @@ class AgentRequestHandler: authority: Final = ( await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) if key_hash - and user_api_key_auth.managed_agent_policy is None + and managed_agent_policy(user_api_key_auth) is None and not user_api_key_auth.is_session_token else user_api_key_auth ) @@ -633,13 +634,13 @@ async def accessible_agents( async def _strict_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: - if auth.managed_agent_policy is not None: + if managed_agent_policy(auth) is not None: return await _managed_actor_agent_access(auth) return await AgentRequestHandler.resolve_key_team_agent_access(auth, strict=True) async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: - agent: Final = auth.managed_agent_policy + agent: Final = managed_agent_policy(auth) if agent is None or not agent.object_permission: return RestrictedAgentAccess(frozenset()) permission: Final = LiteLLM_ObjectPermissionTable.model_validate(agent.object_permission or MappingProxyType({})) diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index 29acd77d78b..67d0a497b09 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -7,6 +7,16 @@ from litellm.types.agents import AgentResponse from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext +def managed_agent_policy(auth: "UserAPIKeyAuth | None") -> AgentResponse | None: + """The admitted managed policy, or ``None`` when the subject was never admitted as a managed agent. + + ``admit_managed_actor`` only assigns ``managed_agent_policy`` after ``actor_admission_failure`` + has verified the bound context, so an ``AgentResponse`` here means admission succeeded. + """ + policy: Final = auth.managed_agent_policy if auth is not None else None + return policy if isinstance(policy, AgentResponse) else None + + async def admit_managed_actor(auth: UserAPIKeyAuth, store: AgentIdentityStore | None) -> None: delegation_verified: Final = auth._managed_delegation_verified auth._managed_delegation_verified = False diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 74993110964..68101ee73b2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -142,6 +142,7 @@ from litellm.router import Router from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from .auth_checks_organization import ( add_team_org_context_to_request_body, @@ -1058,10 +1059,9 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) - if _model and valid_token is not None and valid_token.managed_agent_policy is not None: - managed_models: Final = (valid_token.managed_agent_policy.object_permission or MappingProxyType({})).get( - "models", () - ) + managed_policy: Final = managed_agent_policy(valid_token) + if _model and managed_policy is not None: + managed_models: Final = (managed_policy.object_permission or MappingProxyType({})).get("models", ()) if not isinstance(managed_models, (list, tuple)) or not managed_models: raise HTTPException(403, "This agent has no model grants") _can_object_call_model( @@ -4548,10 +4548,11 @@ async def _check_agent_access_group_model_access( from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if valid_token.managed_agent_policy is None else None + managed: Final = managed_agent_policy(valid_token) + unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if managed is None else None ceilings: Final = ( - await resolve_managed_agent_ceilings(valid_token.managed_agent_policy) - if valid_token.managed_agent_policy is not None + await resolve_managed_agent_ceilings(managed) + if managed is not None else (unmanaged,) if unmanaged is not None else ()