scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(agents): gate the managed path on an admitted policy object

Browse source 
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.
This commit is contained in:
Joshua Valluru 2026-09-29 15:34:48 -07:00
parent 4c4675744c
commit f8dd979d9a
5 changed files with 33 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -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

3
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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]

9
litellm/proxy/agent_endpoints/auth/agent_permission_handler.py
View file

@ -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({}))

10
litellm/proxy/agent_endpoints/auth/managed_authorization.py
View file

@ -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

15
litellm/proxy/auth/auth_checks.py
View file

@ -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 ()