diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 67bb43638e0..e0e5193d0a4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 950ac5e9906..95c221a8f20 100644 --- a/litellm/proxy/utils.py +++ b/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: