From 1206fa802b851fb485cf494e234f6a628f86813c Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:51:04 +0000 Subject: [PATCH] refactor(agents): resolve attached access groups from the agent registry instead of the DB on the request path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_access_groups.py | 81 +--------------- litellm/proxy/agent_endpoints/endpoints.py | 4 - .../access_group_endpoints.py | 2 - .../auth/test_agent_access_groups.py | 93 +++---------------- 4 files changed, 14 insertions(+), 166 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index e0e5193d0a4..49e5407ff88 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,35 +1,18 @@ import asyncio -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from fastapi import HTTPException -from pydantic import TypeAdapter, ValidationError -from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_AccessGroupTable -from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl - - -class _AgentAccessGroupsRecord(Protocol): - @property - def access_group_ids(self) -> Sequence[str] | None: ... - - -class _AgentIdWhere(TypedDict): - agent_id: ReadOnly[str] - AccessGroupIds: TypeAlias = tuple[str, ...] AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params -AgentRecordFinder: TypeAlias = Callable[[str], Awaitable[_AgentAccessGroupsRecord | None]] # mutable-ok: Callable LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax -_CACHED_IDS: Final = TypeAdapter(list[str]) - @dataclass(frozen=True, slots=True) class AgentAccessGroupCeiling: @@ -44,19 +27,6 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params -def agent_access_group_ids_cache_key(agent_id: str) -> str: - return f"agent_access_group_ids:{agent_id}" - - -def _cached_access_group_ids(cached: object) -> AccessGroupIds | None: - if cached is None: - return None - try: - return tuple(_CACHED_IDS.validate_python(cached)) - except ValidationError: - return None - - async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -64,51 +34,6 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: return tuple(agent.access_group_ids or ()) if agent is not None else () -async def load_agent_access_group_ids( - agent_id: str, - cache: DualCache, - find_agent: AgentRecordFinder, - fallback: AccessGroupIdsLoader, -) -> AccessGroupIds: - """The agent row's groups, cached for the management-object TTL and evicted on every agent write.""" - cache_key: Final = agent_access_group_ids_cache_key(agent_id) - cached: Final = _cached_access_group_ids(await cache.async_get_cache(key=cache_key)) - if cached is not None: - return cached - try: - record: Final = await find_agent(agent_id) - except Exception as e: # noqa: BLE001 # prisma raises many error types; the registry snapshot answers instead - verbose_proxy_logger.warning("Failed to read access groups for agent %r, using registry: %s", agent_id, e) - return await fallback(agent_id) - access_group_ids: Final = tuple(record.access_group_ids or ()) if record is not None else () - await cache.async_set_cache(key=cache_key, value=access_group_ids, ttl=get_management_object_ttl(cache)) - return access_group_ids - - -async def _load_agent_access_group_ids(agent_id: str) -> AccessGroupIds: - from litellm.proxy.agent_endpoints.agent_registry import agents_table - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client is None: - return await _registry_access_group_ids(agent_id) - db: Final = prisma_client - - async def find_agent(row_agent_id: str) -> _AgentAccessGroupsRecord | None: - return await agents_table(db).find_unique(where=_AgentIdWhere(agent_id=row_agent_id)) - - return await load_agent_access_group_ids(agent_id, user_api_key_cache, find_agent, _registry_access_group_ids) - - -async def evict_agent_access_group_ids(agent_ids: Sequence[str]) -> None: - from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast - from litellm.proxy.proxy_server import user_api_key_cache - - await evict_and_broadcast( - cache_keys=tuple(agent_access_group_ids_cache_key(agent_id) for agent_id in agent_ids), - user_api_key_cache=user_api_key_cache, - ) - - async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: from litellm.proxy.auth.auth_checks import get_access_object from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache @@ -132,7 +57,7 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: async def resolve_agent_access_group_ceiling( agent_id: str, - load_access_group_ids: AccessGroupIdsLoader = _load_agent_access_group_ids, + load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids, load_access_group: AccessGroupLoader = _load_access_group, ) -> AgentAccessGroupCeiling | None: """``None`` when the agent has no access groups attached, so nothing is capped.""" diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 62783fb412a..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -44,7 +44,6 @@ from litellm.proxy.agent_endpoints.agent_search import ( global_agent_search_index, search_agents, ) -from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -697,7 +696,6 @@ async def update_agent( prisma_client=prisma_client, updated_by=updated_by, ) - await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -801,7 +799,6 @@ async def patch_agent( prisma_client=prisma_client, updated_by=updated_by, ) - await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -864,7 +861,6 @@ async def delete_agent( raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.") await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client) - await evict_agent_access_group_ids((agent_id,)) AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 2694d00b17f..b4923b0a2dc 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -16,7 +16,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -783,7 +782,6 @@ async def delete_access_group( await invalidate_access_group_cache(access_group_id) _detach_access_group_from_agent_registry(detached_agent_ids, access_group_id) - await evict_agent_access_group_ids(detached_agent_ids) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py index 266ce52e3a9..e744e84d671 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -1,16 +1,11 @@ -from collections.abc import Sequence -from dataclasses import dataclass from typing import Final import pytest from fastapi import HTTPException -from litellm.caching.dual_cache import DualCache from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( AgentAccessGroupCeiling, - agent_access_group_ids_cache_key, - load_agent_access_group_ids, resolve_agent_access_group_ceiling, ) from litellm.types.agents import AgentResponse @@ -110,86 +105,20 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): assert ceiling.agent_ids == frozenset() -@dataclass(frozen=True, slots=True) -class _AgentRow: - access_group_ids: Sequence[str] | None - - -class _FakeAgentTable: - def __init__(self, rows: dict[str, _AgentRow], failing: bool = False) -> None: - self._rows: Final = rows - self._failing: Final = failing - self.reads = 0 - - async def find_agent(self, agent_id: str) -> _AgentRow | None: - self.reads += 1 - if self._failing: - raise RuntimeError("db down") - return self._rows.get(agent_id) - - -async def _registry_snapshot(agent_id: str) -> tuple[str, ...]: - return ("registry-group",) - - @pytest.mark.asyncio -async def test_agent_row_is_read_once_then_served_from_cache(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1", "g2"])}) +async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - first: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - second: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + _, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))}) + global_agent_registry.register_agent(_agent(["g1"])) + try: + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group) + finally: + global_agent_registry.deregister_agent("agent") - assert (first, second, table.reads) == (("g1", "g2"), ("g1", "g2"), 1) - - -@pytest.mark.asyncio -async def test_agent_with_no_row_or_no_groups_caches_an_empty_answer(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({"bare": _AgentRow(None)}) - - bare: Final = await load_agent_access_group_ids("bare", cache, table.find_agent, _registry_snapshot) - missing: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) - again: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) - - assert (bare, missing, again, table.reads) == ((), (), (), 2) - - -@pytest.mark.asyncio -async def test_evicted_cache_entry_picks_up_the_patched_row(): - cache: Final = DualCache() - rows: Final = {"agent-1": _AgentRow(["g1"])} - table: Final = _FakeAgentTable(rows) - await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - rows["agent-1"] = _AgentRow(["g2"]) - stale: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - await cache.async_delete_cache(key=agent_access_group_ids_cache_key("agent-1")) - fresh: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert (stale, fresh) == (("g1",), ("g2",)) - - -@pytest.mark.asyncio -async def test_unreadable_row_falls_back_to_the_registry_without_caching(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({}, failing=True) - - answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert answer == ("registry-group",) - assert await cache.async_get_cache(key=agent_access_group_ids_cache_key("agent-1")) is None - - -@pytest.mark.asyncio -async def test_garbage_in_the_cache_is_treated_as_a_miss(): - cache: Final = DualCache() - await cache.async_set_cache(key=agent_access_group_ids_cache_key("agent-1"), value={"not": "a list"}) - table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1"])}) - - answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert (answer, table.reads) == (("g1",), 1) + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset() + ) @pytest.mark.asyncio