From 3a86567c9d3ffcd12507407215ea14d9d89d2184 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:31:03 +0000 Subject: [PATCH] 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> --- .../proxy/agent_endpoints/agent_registry.py | 3 + litellm/proxy/agent_endpoints/endpoints.py | 4 + .../access_group_endpoints.py | 2 + .../auth/test_agent_access_groups.py | 91 ++++++++++++++++++- .../proxy/utils/helpers/test_model_access.py | 87 ++++++++++++++++-- 5 files changed, 177 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c8948d8d70e..d6b12e830e1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/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: ... diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index aa8979a73c6..62783fb412a 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/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")) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index b4923b0a2dc..2694d00b17f 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/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 ) 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 8c31c9428e7..266ce52e3a9 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,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 diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 5fb4392eec6..7f77f938323 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/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"])