diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 54cc26885ab..4e8880d37c5 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -177,11 +177,10 @@ class AgentRequestHandler: registered: Final = global_agent_registry.get_agent_by_id(agent_id) registry_managed: Final = isinstance(registered, AgentResponse) and registered.identity_managed - if registry_managed or (registered is None and prisma_client is not None): + if registry_managed or prisma_client is not None: target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) if isinstance(target, AgentIdentityFailure): - if registry_managed: - raise_identity_failure(target) + raise_identity_failure(target) elif target is None and registry_managed: return False elif isinstance(target, AgentResponse) and target.identity_managed: 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 e3834d7d499..a1a022fdd35 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 @@ -1109,3 +1109,53 @@ async def test_delegated_target_grants_do_not_borrow_another_teams_authority( object_permission={"object_permission_id": "own", "agents": ["direct", "a-only", "b-only"]}, ) assert await AgentRequestHandler.resolve_agent_access(auth) == RestrictedAgentAccess(frozenset(expected)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "managed,enabled,grant,outage,allowed", + [ + (True, True, False, False, False), + (True, True, True, False, True), + (True, False, True, False, False), + (False, True, False, False, True), + (True, True, False, True, False), + ], +) +async def test_target_authorization_uses_live_policy_despite_stale_unmanaged_registry( + monkeypatch: pytest.MonkeyPatch, managed: bool, enabled: bool, grant: bool, outage: bool, allowed: bool +) -> None: + from unittest.mock import MagicMock + + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints import agent_registry + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding + + stale: Final = AgentResponse(agent_id="target", agent_name="Target", agent_card_params={}) + registry: Final = AgentRegistry() + registry.register_agent(stale) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + binding: Final = AgentIdentityBinding( + agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", + issuer="issuer", revision="current", + ) + current: Final = stale.model_copy(update={ + "identity_managed": managed, "identity": binding if managed else None, "enabled": enabled, + }) + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock( + return_value=current, side_effect=ConnectionError("writer unavailable") if outage else None, + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["target"]) + auth: Final = UserAPIKeyAuth(object_permission=permission if grant else None) + + if outage: + with pytest.raises(HTTPException) as denied: + await AgentRequestHandler.is_agent_allowed("target", auth) + assert denied.value.status_code == 503 + return + assert await AgentRequestHandler.is_agent_allowed("target", auth) is allowed diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 526f24c5221..bd43cb7ce13 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -350,6 +350,7 @@ class TestAgentByIdKeyRedaction: test_client = _make_app_with_role(role) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value=None ) @@ -412,6 +413,7 @@ class TestAgentRBACInternalUser: return_value=_sample_agent_response() ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value=None ) @@ -1342,6 +1344,7 @@ def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_other def _get_as(role: LitellmUserRoles): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"})