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

refactor(agents): expose shared permission resolution helpers

Browse source
This commit is contained in:
Joshua Valluru 2026-09-26 14:31:34 -07:00
parent 7cbce50513
commit 97ddb52e0d
4 changed files with 30 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1841,7 +1841,7 @@ class MCPRequestHandler:
return scoped
@staticmethod
async def _admitted_subject_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
async def admitted_subject_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
"""The independent sources a keyless admitted subject reaches MCP servers through: their own
direct grants, plus every team they are a live roster member of.
@ -1951,7 +1951,7 @@ class MCPRequestHandler:
roster instead of by grant charged unrelated teams' buckets)."""
return [
(source, set(await MCPRequestHandler.get_allowed_mcp_servers(source, keyless_source=True)))
for source in await MCPRequestHandler._admitted_subject_sources(auth)
for source in await MCPRequestHandler.admitted_subject_sources(auth)
]
@staticmethod

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

@ -91,7 +91,7 @@ class AgentRequestHandler:
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:
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)
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)
own_access: Final = _intersect_agent_access(key_team_access, caller_access)
agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling)
@ -109,13 +109,13 @@ class AgentRequestHandler:
return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth)
@staticmethod
async def _resolve_key_team_agent_access(
async def resolve_key_team_agent_access(
user_api_key_auth: UserAPIKeyAuth | None,
*,
strict: bool = False,
) -> AgentAccess:
try:
key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth, strict=strict)
key_access: Final = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth, strict=strict)
team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(
user_api_key_auth, strict=strict
)
@ -240,7 +240,7 @@ class AgentRequestHandler:
return team_obj.object_permission
@staticmethod
async def _get_allowed_agents_for_key(
async def get_allowed_agents_for_key(
user_api_key_auth: UserAPIKeyAuth | None = None,
*,
strict: bool = False,
@ -315,7 +315,7 @@ class AgentRequestHandler:
2. Also includes agents from team's access_group_ids (unified access groups)
Fetches the team object once and reuses it for both permission sources.
Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`.
Declared-but-empty grants stay restricted; see `get_allowed_agents_for_key`.
"""
if user_api_key_auth is None:
return UnrestrictedAgentAccess()
@ -623,7 +623,7 @@ async def accessible_agents(
async def _strict_agent_access(auth: UserAPIKeyAuth) -> AgentAccess:
if auth.managed_agent_policy is not None:
return await _managed_actor_agent_access(auth)
return await AgentRequestHandler._resolve_key_team_agent_access(auth, strict=True)
return await AgentRequestHandler.resolve_key_team_agent_access(auth, strict=True)
async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess:
@ -632,7 +632,7 @@ async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess:
return RestrictedAgentAccess(frozenset())
permission: Final = LiteLLM_ObjectPermissionTable.model_validate(agent.object_permission or MappingProxyType({}))
own_auth: Final = UserAPIKeyAuth(object_permission=permission)
own: Final = _granted_ids(await AgentRequestHandler._get_allowed_agents_for_key(own_auth, strict=True))
own: Final = _granted_ids(await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True))
from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings
@ -653,6 +653,6 @@ async def verified_human_agent_grants(user_id: str | None) -> frozenset[str]:
if user_id is None:
return frozenset()
human: Final = await MCPRequestHandler.reload_admitted_user(user_id, requires_fresh_policy=True)
sources: Final = await MCPRequestHandler._admitted_subject_sources(human)
sources: Final = await MCPRequestHandler.admitted_subject_sources(human)
human_access: Final = await asyncio.gather(*(_strict_agent_access(source) for source in sources))
return frozenset().union(*(_granted_ids(access) for access in human_access))

4
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -8308,7 +8308,7 @@ class TestUserSubjectTeamUnion:
) == ["t1"]
# An admitted subject never fans out HERE: it resolves one source per team first, and each of
# those pins a team_id, so this helper only ever answers the single-team question. The fan-out
# itself is _admitted_subject_sources' job, asserted below.
# itself is admitted_subject_sources' job, asserted below.
with self._patch(teams_by_id={}, user_teams=["t2", "t3"]):
assert await MCPRequestHandler._team_ids_for_mcp_grant(_make_admitted_subject("u")) == []
# keyless, no user_id -> nothing
@ -8871,7 +8871,7 @@ class TestUserSubjectTeamUnion:
teams["t-member"].organization_id = "org-a"
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["t-member", "t-stale"]):
sources = await MCPRequestHandler._admitted_subject_sources(auth)
sources = await MCPRequestHandler.admitted_subject_sources(auth)
assert [(s.team_id, s.org_id) for s in sources] == [(None, None), ("t-member", "org-a")]
# The user's own source carries their grants; a team source must NOT, or the team would be

36
tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py
View file

@ -67,7 +67,7 @@ class TestAgentRequestHandler:
# Case 1: Both key and team have agents - intersection
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_key"
AgentRequestHandler, "get_allowed_agents_for_key"
) as mock_key:
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_team"
@ -86,7 +86,7 @@ class TestAgentRequestHandler:
# Case 2: Team has agents, key has none - inherit from team
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_key"
AgentRequestHandler, "get_allowed_agents_for_key"
) as mock_key:
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_team"
@ -105,7 +105,7 @@ class TestAgentRequestHandler:
# Case 3: Key has agents, team has none - key restrictions stand
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_key"
AgentRequestHandler, "get_allowed_agents_for_key"
) as mock_key:
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_team"
@ -120,7 +120,7 @@ class TestAgentRequestHandler:
# Case 4: No grant anywhere - unrestricted (documented open-by-default)
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_key"
AgentRequestHandler, "get_allowed_agents_for_key"
) as mock_key:
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_team"
@ -141,7 +141,7 @@ class TestAgentRequestHandler:
api_key="test-key", user_id="test-user", team_id="test-team"
)
with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key:
with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key:
with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team:
mock_key.return_value = RestrictedAgentAccess(frozenset({"agent-alpha"}))
mock_team.return_value = RestrictedAgentAccess(frozenset({"agent-beta"}))
@ -298,7 +298,7 @@ class TestAgentRequestHandler:
) as mock_groups:
mock_groups.return_value = []
assert await AgentRequestHandler._get_allowed_agents_for_key(
assert await AgentRequestHandler.get_allowed_agents_for_key(
user_api_key_auth=mock_user_auth
) == RestrictedAgentAccess(frozenset())
@ -314,7 +314,7 @@ class TestAgentRequestHandler:
) as mock_groups:
mock_groups.side_effect = Exception("DB Error")
assert await AgentRequestHandler._get_allowed_agents_for_key(
assert await AgentRequestHandler.get_allowed_agents_for_key(
user_api_key_auth=mock_user_auth
) == UnrestrictedAgentAccess()
@ -403,7 +403,7 @@ class TestAgentRequestHandler:
)
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_key"
AgentRequestHandler, "get_allowed_agents_for_key"
) as mock_key:
with patch.object(
AgentRequestHandler, "_get_allowed_agents_for_team"
@ -488,9 +488,9 @@ class TestAgentRequestHandler:
listed: Final = await accessible_agents(session, registry.get_agent_list(), resolve_access, effective_contexts)
assert {agent.agent_name for agent in listed} == {"alpha", "beta"}
async def test_get_allowed_agents_for_key_via_access_group_ids(self):
async def testget_allowed_agents_for_key_via_access_group_ids(self):
"""
Test that _get_allowed_agents_for_key includes agents from key's access_group_ids
Test that get_allowed_agents_for_key includes agents from key's access_group_ids
(unified access groups) when key has no native object_permission.
"""
mock_user_auth = UserAPIKeyAuth(
@ -507,16 +507,16 @@ class TestAgentRequestHandler:
new_callable=AsyncMock,
return_value=["agent-from-ag-1", "agent-from-ag-2"],
):
result = await AgentRequestHandler._get_allowed_agents_for_key(
result = await AgentRequestHandler.get_allowed_agents_for_key(
user_api_key_auth=mock_user_auth
)
assert result == RestrictedAgentAccess(
frozenset({"agent-from-ag-1", "agent-from-ag-2"})
)
async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self):
async def testget_allowed_agents_for_key_combines_native_and_access_groups(self):
"""
Test that _get_allowed_agents_for_key combines agents from native object_permission
Test that get_allowed_agents_for_key combines agents from native object_permission
and key's access_group_ids (unified access groups).
"""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
@ -539,7 +539,7 @@ class TestAgentRequestHandler:
new_callable=AsyncMock,
return_value=["agent-from-ag"],
):
result = await AgentRequestHandler._get_allowed_agents_for_key(
result = await AgentRequestHandler.get_allowed_agents_for_key(
user_api_key_auth=mock_user_auth
)
assert result == RestrictedAgentAccess(
@ -610,7 +610,7 @@ class TestAgentRequestHandler:
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
registry,
):
with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key:
with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key:
with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team:
for key_grant, team_grant in (
(
@ -787,11 +787,11 @@ async def test_strict_legacy_group_grants_ignore_stale_replica(monkeypatch: pyte
object_permission_id="permission", agent_access_groups=["group"]
)
)
assert await AgentRequestHandler._get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess(
assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess(
frozenset({"revoked"})
)
database.writer_db.litellm_agentstable.find_many.return_value = []
assert await AgentRequestHandler._get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess(
assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess(
frozenset()
)
database.db.litellm_agentstable.find_many.assert_not_awaited()
@ -823,7 +823,7 @@ async def test_strict_invocation_policy_outage_denies_instead_of_allowing_all(
),
)
with pytest.raises(HTTPException, match="policy is unavailable") as denied:
await AgentRequestHandler._resolve_key_team_agent_access(auth, strict=True)
await AgentRequestHandler.resolve_key_team_agent_access(auth, strict=True)
assert denied.value.status_code == 503