fix(agents): evict the cached agent access groups on every agent write and cap the model listing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/agent_endpoints/agent_registry.py

@@ -67,6 +67,9 @@ class AgentRecord(Protocol):
     @property
     def object_permission(self) -> AgentObjectPermissionRecord | None: ...
 
+    @property
+    def access_group_ids(self) -> Sequence[str] | None: ...
+
     @property
     def spend(self) -> float: ...
 

litellm/proxy/agent_endpoints/endpoints.py

@@ -44,6 +44,7 @@ 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
@@ -696,6 +697,7 @@ 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"))
@@ -799,6 +801,7 @@ 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"))
@@ -861,6 +864,7 @@ 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"))
 

litellm/proxy/management_endpoints/access_group_endpoints.py

@@ -16,6 +16,7 @@ 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,
@@ -782,6 +783,7 @@ 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
         )

tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py

@@ -1,11 +1,16 @@
+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
@@ -35,8 +40,8 @@ def _group(
 
 
 def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]):
-    async def load_agent(agent_id: str) -> AgentResponse | None:
-        return agent
+    async def load_agent(agent_id: str) -> tuple[str, ...]:
+        return tuple(agent.access_group_ids or ()) if agent is not None else ()
 
     async def load_group(group_id: str) -> LiteLLM_AccessGroupTable | None:
         return groups.get(group_id)
@@ -105,6 +110,88 @@ 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"])})
+
+    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)
+
+    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)
+
+
 @pytest.mark.asyncio
 async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch):
     from litellm.proxy import proxy_server

tests/test_litellm/proxy/utils/helpers/test_model_access.py

@@ -110,8 +110,7 @@ def test_create_model_info_response_happy_path_no_metadata():
         "owned_by": result["owned_by"],
         "created_is_int": isinstance(result["created"], int),
         "metadata_absent": "metadata" not in result,
-        "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int)
-        and result["max_input_tokens"] > 0,
+        "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) and result["max_input_tokens"] > 0,
         "max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int)
         and result["max_output_tokens"] > 0,
     }
@@ -205,9 +204,7 @@ def test_validate_model_access_happy_path_single_model_in_list():
 
 def test_validate_model_access_happy_path_batch_all_accessible():
     summary = {
-        "result": validate_model_access(
-            "gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]
-        ),
+        "result": validate_model_access("gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]),
         "input": "gpt-4o,claude-haiku",
         "available": ["gpt-4o", "claude-haiku", "gemini"],
     }
@@ -389,9 +386,7 @@ async def test_get_available_models_for_user_error_path_complete_list_raises(
     def _boom(**_kwargs):
         raise RuntimeError("downstream failure")
 
-    monkeypatch.setattr(
-        "litellm.proxy.auth.model_checks.get_complete_model_list", _boom
-    )
+    monkeypatch.setattr("litellm.proxy.auth.model_checks.get_complete_model_list", _boom)
     user_api_key_dict = UserAPIKeyAuth(
         api_key="sk-test-key",
         user_id="user-1",
@@ -481,6 +476,7 @@ async def test_get_available_models_for_user_without_access_groups_grants_nothin
     )
     assert result == []
 
+
 @pytest.mark.asyncio
 async def test_get_available_models_for_user_resolves_key_access_group_models(
     monkeypatch,
@@ -521,3 +517,78 @@ async def test_get_available_models_for_user_resolves_key_access_group_models(
         user_api_key_cache=MagicMock(),
     )
     assert result == ["model-b"]
+
+
+def _agent_ceiling(models: frozenset[str] | None):
+    from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
+
+    async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
+        if models is None:
+            return None
+        return AgentAccessGroupCeiling(
+            access_group_ids=("ag-agent",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset()
+        )
+
+    return resolve
+
+
+def _agent_key(models: list[str]) -> UserAPIKeyAuth:
+    return UserAPIKeyAuth(api_key="sk-agent-key", user_id="user-1", agent_id="agent-1", models=models)
+
+
+@pytest.mark.asyncio
+async def test_agent_key_listing_is_capped_to_its_access_groups():
+    result = await get_available_models_for_user(
+        user_api_key_dict=_agent_key(["model-a", "model-b", "model-c"]),
+        llm_router=_router_with_models(["model-a", "model-b", "model-c"]),
+        general_settings={},
+        user_model=None,
+        resolve_agent_ceiling=_agent_ceiling(frozenset({"model-b", "model-d"})),
+    )
+    assert result == ["model-b"]
+
+
+@pytest.mark.asyncio
+async def test_agent_key_listing_is_empty_when_its_groups_grant_no_model():
+    result = await get_available_models_for_user(
+        user_api_key_dict=_agent_key(["model-a"]),
+        llm_router=_router_with_models(["model-a"]),
+        general_settings={},
+        user_model=None,
+        resolve_agent_ceiling=_agent_ceiling(frozenset()),
+    )
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_agent_ceiling_expands_a_model_access_group_name_for_listing():
+    router = _router_with_models(["model-a", "model-b"])
+    router.get_model_access_groups.return_value = {"fast-models": ["model-b"]}
+    result = await get_available_models_for_user(
+        user_api_key_dict=_agent_key(["model-a", "model-b"]),
+        llm_router=router,
+        general_settings={},
+        user_model=None,
+        resolve_agent_ceiling=_agent_ceiling(frozenset({"fast-models"})),
+    )
+    assert result == ["model-b"]
+
+
+@pytest.mark.asyncio
+async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups():
+    router = _router_with_models(["model-a", "model-b"])
+    plain_key = await get_available_models_for_user(
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-plain", user_id="user-1", models=["model-a", "model-b"]),
+        llm_router=router,
+        general_settings={},
+        user_model=None,
+        resolve_agent_ceiling=_agent_ceiling(frozenset({"model-a"})),
+    )
+    agent_without_groups = await get_available_models_for_user(
+        user_api_key_dict=_agent_key(["model-a", "model-b"]),
+        llm_router=router,
+        general_settings={},
+        user_model=None,
+        resolve_agent_ceiling=_agent_ceiling(None),
+    )
+    assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"])