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): authorize targets against current database policy

Browse source
This commit is contained in:
Joshua Valluru 2026-09-30 12:06:11 -07:00
parent 812f6c7024
commit 7e60757407
3 changed files with 55 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

3
tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
View file

@ -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"})