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 df429f66268..ec92bbba677 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 @@ -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 diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index f0d85f4e09b..4b9238f1868 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -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)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index ede304742e0..5aaa3d3c446 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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 diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index a5c13d8bf0e..e5dfeb369e2 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -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