refactor(agents): keep the model listing cap within the type-discipline budget

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

litellm/proxy/agent_endpoints/auth/agent_access_groups.py

@@ -1,33 +1,41 @@
-"""
-Ceiling that an agent's attached access groups place on requests made with that agent's key.
-
-Keys and teams use access groups as grants. An agent uses them the way it already uses its
-``object_permission``: the union of the attached groups caps what the agent's key can reach,
-on top of whatever the key and team allow. A group that cannot be loaded contributes nothing,
-so a missing or unreadable group can only narrow the agent, never widen it.
-"""
-
 import asyncio
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass
-from typing import Final, TypeAlias
+from typing import Final, Protocol, 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.types.agents import AgentResponse
+from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl
 
-AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]]  # mutable-ok: Callable parameter syntax
+
+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:
     """Everything the agent's attached access groups allow. An empty set denies that resource kind."""
 
-    access_group_ids: tuple[str, ...]
+    access_group_ids: AccessGroupIds
     models: frozenset[str]
     mcp_server_ids: frozenset[str]
     agent_ids: frozenset[str]
@@ -36,10 +44,69 @@ class AgentAccessGroupCeiling:
 CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]]  # mutable-ok: Callable params
 
 
-async def _load_agent(agent_id: str) -> AgentResponse | None:
+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
 
-    return await get_agent_with_read_through(agent_id)
+    agent: Final = await get_agent_with_read_through(agent_id)
+    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:
@@ -65,12 +132,11 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup:
 
 async def resolve_agent_access_group_ceiling(
     agent_id: str,
-    load_agent: AgentLoader = _load_agent,
+    load_access_group_ids: AccessGroupIdsLoader = _load_agent_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."""
-    agent: Final = await load_agent(agent_id)
-    access_group_ids: Final = tuple(agent.access_group_ids or ()) if agent is not None else ()
+    access_group_ids: Final = await load_access_group_ids(agent_id)
     if not access_group_ids:
         return None
 

litellm/proxy/utils.py

@@ -122,6 +122,7 @@ from litellm.proxy._types import (
     Member,
     UserAPIKeyAuth,
 )
+from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -7980,6 +7981,51 @@ async def _get_access_group_models(
     return tuple(dict.fromkeys((*team_group_models, *key_group_models)))
 
 
+async def _agent_access_group_visible_models(
+    user_api_key_dict: "UserAPIKeyAuth",
+    llm_router: Optional["Router"],
+    include_model_access_groups: bool,
+    return_wildcard_routes: bool,
+    team_id: str | None,
+    resolve_agent_ceiling: CeilingResolver,
+) -> frozenset[str] | None:
+    """Models an agent key may still list once its attached access groups cap it, ``None`` when
+    nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on."""
+    from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models
+
+    if not user_api_key_dict.agent_id:
+        return None
+    ceiling: Final = await resolve_agent_ceiling(user_api_key_dict.agent_id)
+    if ceiling is None:
+        return None
+    if llm_router is None:
+        return ceiling.models
+    proxy_model_list: Final = llm_router.get_model_names()
+    model_access_groups: Final = llm_router.get_model_access_groups()
+    granted: Final = get_team_models(
+        team_models=sorted(ceiling.models),
+        proxy_model_list=proxy_model_list,
+        model_access_groups=model_access_groups,
+        include_model_access_groups=include_model_access_groups,
+    )
+    if not granted:
+        return frozenset()
+    return frozenset(
+        get_complete_model_list(
+            key_models=granted,
+            team_models=(),
+            proxy_model_list=proxy_model_list,
+            user_model=None,
+            infer_model_from_keys=False,
+            return_wildcard_routes=return_wildcard_routes,
+            llm_router=llm_router,
+            model_access_groups=model_access_groups,
+            include_model_access_groups=include_model_access_groups,
+            team_id=team_id,
+        )
+    )
+
+
 async def get_available_models_for_user(
     user_api_key_dict: "UserAPIKeyAuth",
     llm_router: Optional["Router"],
@@ -7992,6 +8038,7 @@ async def get_available_models_for_user(
     only_model_access_groups: bool = False,
     return_wildcard_routes: bool = False,
     user_api_key_cache: Optional["UserApiKeyCache"] = None,
+    resolve_agent_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
 ) -> list[str]:
     """
     Get the list of models available to a user based on their API key and team permissions.
@@ -8095,7 +8142,18 @@ async def get_available_models_for_user(
         team_id=effective_team_id,
     )
 
-    return all_models
+    agent_visible: Final = await _agent_access_group_visible_models(
+        user_api_key_dict=user_api_key_dict,
+        llm_router=llm_router,
+        include_model_access_groups=include_model_access_groups,
+        return_wildcard_routes=return_wildcard_routes,
+        team_id=effective_team_id,
+        resolve_agent_ceiling=resolve_agent_ceiling,
+    )
+    if agent_visible is None:
+        return all_models
+    capped: Final = [m for m in all_models if m in agent_visible]  # mutable-ok: callers expect the list all_models is
+    return capped
 
 
 def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: