From eb3c8c168f533640bf16d70fdeea148e18f65de0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 12:19:49 -0700 Subject: [PATCH] fix(proxy): derive config agent ids from agent_name so grants survive secret rotation (#36020) * fix(proxy): derive config agent ids from agent_name so grants survive secret rotation Config-defined A2A agents were identified by a sha256 of the whole resolved config entry, secrets included, so rotating an os.environ secret re-minted the agent_id on restart and orphaned every object_permission.agents grant while grant-less keys kept access (LIT-5144). The id now hashes only agent_name, and the old full-entry hash is kept as a legacy alias: permission checks, GET /v1/agents filtering, spend and key attachment, and public_agent_groups all normalize legacy ids so pre-upgrade grants keep working * fix(proxy): persist stable agent ids into stored grants at startup The runtime alias only translates a legacy grant while the current config still hashes to it, so a secret rotation after upgrading would orphan the grant, and an orphaned grant intersecting a stable team grant collapses to an empty list that downstream reads as allow-all. Rewriting the stored ids once at boot removes both. This cannot be a SQL migration because only the running proxy can recompute the legacy hash from resolved config secrets * fix(proxy): make the grant id migration a compare-and-swap A grant edited between the migration's read and write kept the stale snapshot. The update now predicates on the agents array read at scan time via update_many, so a concurrently modified row is skipped and the runtime alias covers it until the next boot retries * fix(proxy): retry the grant id migration and stay within the LIT002 ceiling The one-shot startup task now retries up to three times with a short delay so a transient DB error at boot cannot leave a legacy grant unmigrated until an operator's next restart is the rotation itself. The new list constructions in the migration and the alias-expanded agent id lookups are tuples now, keeping the branch under the mutable-collection budget * fix(proxy): count compare-and-swap misses in the grant id migration migrate_legacy_grant_ids now returns rewritten and missed counts from the update_many results instead of reporting scanned rows as migrated, and the startup task retries while any rows remain unmigrated, not just on errors * fix(lint): clear basedpyright budget breaches in agent id aliasing --- .../proxy/agent_endpoints/agent_registry.py | 119 ++++++++++++-- .../auth/agent_permission_handler.py | 38 +++-- litellm/proxy/agent_endpoints/endpoints.py | 34 +++- litellm/proxy/proxy_server.py | 31 ++++ .../public_endpoints/public_endpoints.py | 2 +- litellm/repositories/table_repositories.py | 4 + ruff-strict-budget.json | 2 +- .../auth/test_agent_permission_handler.py | 80 +++++++++ .../agent_endpoints/test_agent_registry.py | 155 +++++++++++++++++- .../proxy/agent_endpoints/test_endpoints.py | 3 +- .../public_endpoints/test_public_endpoints.py | 2 + type-discipline-budget.json | 4 +- 12 files changed, 426 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 476bd725c73..742fdf35b1e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,8 +1,10 @@ +import asyncio import hashlib import json from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Protocol, TypedDict +from types import MappingProxyType +from typing import Any, Final, NamedTuple, Protocol, TypedDict import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -10,7 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import AgentsRepository +from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -86,10 +88,32 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient: return table +class ObjectPermissionGrantRecord(Protocol): + object_permission_id: str + agents: list[str] | None + + +class ObjectPermissionTableClient(Protocol): + async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ... + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: + table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table + return table + + +class GrantMigrationResult(NamedTuple): + rewritten: int + missed: int + + class AgentRegistry: def __init__(self): self.agent_list: list[AgentResponse] = [] self.config_agents: tuple[AgentConfig, ...] = () + self.config_agent_legacy_ids: Mapping[str, str] = MappingProxyType({}) def reset_agent_list(self): self.agent_list = [] @@ -100,23 +124,33 @@ class AgentRegistry: def deregister_agent(self, agent_name: str): self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] - def get_agent_list(self, agent_names: Sequence[str] | None = None): + def get_agent_list(self, agent_names: Sequence[str] | None = None) -> tuple[AgentResponse, ...]: if agent_names is not None: - return [agent for agent in self.agent_list if agent.agent_name in agent_names] - return self.agent_list + return tuple(agent for agent in self.agent_list if agent.agent_name in agent_names) + return tuple(self.agent_list) - def get_public_agent_list(self) -> list[AgentResponse]: - public_agent_list: Final[list[AgentResponse]] = [] - if litellm.public_agent_groups is None: - return public_agent_list - for agent in self.agent_list: - if agent.agent_id in litellm.public_agent_groups: - public_agent_list.append(agent) - return public_agent_list + def get_public_agent_list(self) -> tuple[AgentResponse, ...]: + public_agent_groups: Final = litellm.public_agent_groups + if public_agent_groups is None: + return () + return tuple( + agent for agent in self.agent_list if not self.ids_for_agent(agent.agent_id).isdisjoint(public_agent_groups) + ) def _create_agent_id(self, agent_config: AgentConfig) -> str: + return hashlib.sha256(agent_config["agent_name"].encode()).hexdigest() + + def _create_legacy_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() + def ids_for_agent(self, agent_id: str) -> frozenset[str]: + return frozenset( + {agent_id, *(legacy for legacy, stable in self.config_agent_legacy_ids.items() if stable == agent_id)} + ) + + def stable_agent_id(self, agent_id: str) -> str: + return self.config_agent_legacy_ids.get(agent_id, agent_id) + def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): """ Register the agents declared in config.yaml and remember them for later rebuilds. @@ -131,12 +165,20 @@ class AgentRegistry: if agent_config is None: return - self.config_agents = tuple(agent_config) - for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") + self.config_agents = tuple(agent_config) + self.config_agent_legacy_ids = MappingProxyType( + { + self._create_legacy_agent_id(agent_config_item): self._create_agent_id(agent_config_item) + for agent_config_item in agent_config + if agent_config_item.get("agent_name") and agent_config_item.get("agent_card_params") + } + ) + + for agent_config_item in agent_config: agent_name = agent_config_item.get("agent_name") agent_card_params = agent_config_item.get("agent_card_params") if not all([agent_name, agent_card_params]): @@ -180,6 +222,45 @@ class AgentRegistry: self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) return self.agent_list + async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult: + """ + Rewrite object_permission.agents rows holding a legacy full-entry hash to the + stable name-derived id. + + Only the running proxy can do this: the legacy hash is computed from the + resolved config entry (secrets included), so no SQL migration can know it. + Persisting the stable id here is what keeps a grant alive across a later + secret rotation, which re-mints the legacy hash and would otherwise orphan + the stored value. Idempotent; runs of it after the first find no rows. + + Each write is a compare-and-swap against the agents array read above, so a + grant edited concurrently is left untouched; the runtime alias keeps covering + it and the next boot retries the rewrite. + """ + legacy_ids: Final = tuple(legacy for legacy, stable in self.config_agent_legacy_ids.items() if legacy != stable) + if not legacy_ids: + return GrantMigrationResult(rewritten=0, missed=0) + rows: Final = await table.find_many(where={"agents": {"has_some": legacy_ids}}) + updates: Final = tuple( + ( + row.object_permission_id, + tuple(row.agents or ()), + tuple(dict.fromkeys(self.stable_agent_id(agent_id) for agent_id in row.agents or ())), + ) + for row in rows + ) + counts: Final = await asyncio.gather( + *( + table.update_many( + where={"object_permission_id": object_permission_id, "agents": {"equals": snapshot_agents}}, + data={"agents": translated_agents}, + ) + for object_permission_id, snapshot_agents, translated_agents in updates + ) + ) + rewritten: Final = sum(counts) + return GrantMigrationResult(rewritten=rewritten, missed=len(updates) - rewritten) + ########################################################### ########### DB management helpers for agents ########### ############################################################ @@ -492,6 +573,14 @@ class AgentRegistry: if agent.agent_id == agent_id: return agent + translated_id: Final = self.config_agent_legacy_ids.get(agent_id) + if translated_id is None: + return None + + for agent in self.agent_list: + if agent.agent_id == translated_id: + return agent + return None except Exception as e: raise Exception(f"Error getting agent from DB: {e}") diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index af7e731b385..6baeab3dfff 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -42,24 +42,25 @@ class AgentRequestHandler: List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). """ try: - allowed_agents: list[str] = [] - allowed_agents_for_key: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - allowed_agents_for_team: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + raw_key_grants: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + raw_team_grants: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + allowed_agents_for_key: Final = frozenset( + global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_key_grants + ) + allowed_agents_for_team: Final = frozenset( + global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_team_grants + ) # If team has agent restrictions, handle inheritance and intersection logic - if len(allowed_agents_for_team) > 0: - if len(allowed_agents_for_key) > 0: - # Key has its own agent permissions - use intersection with team permissions - for agent_id in allowed_agents_for_key: - if agent_id in allowed_agents_for_team: - allowed_agents.append(agent_id) - else: - # Key has no agent permissions - inherit from team - allowed_agents = allowed_agents_for_team - else: - allowed_agents = allowed_agents_for_key - - return list(set(allowed_agents)) + if allowed_agents_for_team and allowed_agents_for_key: + # Key has its own agent permissions - use intersection with team permissions + return sorted(allowed_agents_for_key & allowed_agents_for_team) + if allowed_agents_for_team: + # Key has no agent permissions - inherit from team + return sorted(allowed_agents_for_team) + return sorted(allowed_agents_for_key) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) return [] @@ -79,13 +80,16 @@ class AgentRequestHandler: Returns: bool: True if agent is allowed, False otherwise """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) # Empty list means no restrictions - allow all if len(allowed_agents) == 0: return True - return agent_id in allowed_agents + stable_id: Final = global_agent_registry.stable_agent_id(agent_id) + return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agents) @staticmethod def _get_key_object_permission( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 1f9c6e1cc05..028885aa065 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -101,7 +101,11 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) foreign key. Mirrors how spend is joined into the agent response so the UI never has to cross-reference a full key dump client-side. Only non-secret fields are exposed (alias, masked key_name, hashed token).""" - agent_ids: Final = [agent.agent_id for agent in agents] + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent_ids: Final = tuple( + alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if not agent_ids: return key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( @@ -117,7 +121,12 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) ) ) for agent in agents: - agent.keys = keys_by_agent.get(agent.agent_id) + matched_keys = [ + key_summary + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + for key_summary in keys_by_agent.get(alias_id) or () + ] + agent.keys = matched_keys or None def _redact_sensitive_agent_fields( @@ -266,23 +275,32 @@ async def get_agents( from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - agent_ids: Final = [agent.agent_id for agent in returned_agents] + agent_ids: Final = tuple( + alias_id + for agent in returned_agents + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if agent_ids: db_agents: Final = await agents_table(prisma_client).find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map: Final = {a.agent_id: a.spend for a in db_agents} for agent in returned_agents: - if agent.agent_id in spend_map: - agent.spend = spend_map[agent.agent_id] + matched_spends = tuple( + spend_map[alias_id] + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + if alias_id in spend_map + ) + if matched_spends: + agent.spend = sum(matched_spends) await _attach_keys_to_agents(returned_agents, prisma_client) # add is_public field to each agent - we do it this way, to allow setting config agents as public for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and ( - agent.agent_id in litellm.public_agent_groups + agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not ( + global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ) # Redact sensitive fields for non-admin users @@ -863,7 +881,7 @@ async def make_agent_public( if litellm.public_agent_groups is None: litellm.public_agent_groups = [] # handle duplicates - if agent.agent_id in litellm.public_agent_groups: + if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups): raise HTTPException( status_code=400, detail=f"Agent with name {agent.agent_name} already in public agent groups", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d016eb57dd4..90daaaeae6b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1016,6 +1016,37 @@ async def proxy_startup_event(app: FastAPI): asyncio.create_task(_run_pw_migration()) + async def _run_agent_grant_id_migration() -> None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + object_permission_table, + ) + + for attempt in range(3): + try: + result = await global_agent_registry.migrate_legacy_grant_ids( + table=object_permission_table(prisma_client) + ) + if result.rewritten: + verbose_proxy_logger.info( + "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten + ) + if result.missed == 0: + return + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated", + attempt + 1, + result.missed, + ) + except Exception as e: # noqa: BLE001 # startup task must survive any DB error and retry + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 failed: %s", attempt + 1, e + ) + if attempt < 2: + await asyncio.sleep(5) + + asyncio.create_task(_run_agent_grant_id_migration()) + ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before ## the coordination Redis is published to its consumers below. diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 78c3e9fd31b..79791607b2e 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -220,7 +220,7 @@ async def get_agents(request: Request): "url": get_custom_url(str(request.base_url), route=f"a2a/{agent.agent_id}"), } for agent in agents - if agent.agent_id in litellm.public_agent_groups + if not global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ] diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 66e0b6d59e7..be19f290ba6 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -42,6 +42,10 @@ class AgentsRepository(PrismaTableRepository): table_name = "litellm_agentstable" +class ObjectPermissionRepository(PrismaTableRepository): + table_name = "litellm_objectpermissiontable" + + class GuardrailsRepository(PrismaTableRepository): table_name = "litellm_guardrailstable" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 60356eda05b..afd8c0107ea 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -135,7 +135,7 @@ "limit": 27 }, "PERF401": { - "limit": 23 + "limit": 13 }, "PERF402": { "limit": 0 diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 6bdea9c2615..c57296dfc3f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -2,8 +2,11 @@ Unit tests for AgentRequestHandler - Agent permission management for keys and teams. """ +import hashlib +import json import os import sys +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -11,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../..")) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) @@ -212,3 +216,79 @@ class TestAgentRequestHandler: user_api_key_auth=mock_user_auth ) assert sorted(result) == ["agent-from-ag", "native-agent-1"] + + async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): + """LIT-5144: object_permission grants stored under the pre-fix full-entry hash + must keep authorizing the agent after its id became name-based.""" + entry: Final = { + "agent_name": "granted-agent", + "agent_card_params": { + "name": "Granted Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("granted-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "get_allowed_agents") as mock_get_allowed: + for grant, expected in ( + ([legacy_id], True), + ([agent.agent_id], True), + (["unrelated-agent-id"], False), + ([], True), + ): + mock_get_allowed.return_value = grant + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=mock_user_auth, + ) + is expected + ), grant + + async def test_get_allowed_agents_intersects_legacy_team_grant_with_stable_key_grant(self): + """LIT-5144: a team grant stored under the pre-fix full-entry hash and a key grant + stored under the name-based id name the same agent; the intersection must resolve + to that agent instead of collapsing to the allow-all empty list.""" + entry: Final = { + "agent_name": "shared-agent", + "agent_card_params": { + "name": "Shared Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("shared-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + for key_grant, team_grant in ( + ([agent.agent_id], [legacy_id]), + ([legacy_id], [agent.agent_id]), + ([legacy_id], []), + ): + mock_key.return_value = key_grant + mock_team.return_value = team_grant + assert await AgentRequestHandler.get_allowed_agents(user_api_key_auth=mock_user_auth) == [ + agent.agent_id + ], (key_grant, team_grant) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index f506535b5e1..a87f5384c6f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -1,10 +1,14 @@ """Unit tests for AgentRegistry DB operations.""" +import hashlib +import json +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, GrantMigrationResult def _sample_agent_card_params() -> dict: @@ -153,7 +157,7 @@ def test_load_agents_from_db_and_config_skips_incomplete_config_entries(): registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [] + assert registry.get_agent_list() == () @pytest.mark.parametrize( @@ -278,4 +282,149 @@ def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents assert registry.config_agents == () registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [], "a removed config agent must not come back on the next rebuild" + assert registry.get_agent_list() == (), "a removed config agent must not come back on the next rebuild" + + +def test_config_agent_id_survives_static_header_secret_rotation(): + """LIT-5144: the id was a hash of the whole entry, so rotating a static_headers secret silently + re-identified the agent and orphaned every grant pointing at it.""" + base_entry: Final = { + "agent_name": "rotating-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry_v1: Final = AgentRegistry() + registry_v1.load_agents_from_config([base_entry]) + agent_v1: Final = registry_v1.get_agent_by_name("rotating-agent") + assert agent_v1 is not None + + registry_v2: Final = AgentRegistry() + registry_v2.load_agents_from_config([{**base_entry, "static_headers": {"x-upstream-token": "token-v2"}}]) + agent_v2: Final = registry_v2.get_agent_by_name("rotating-agent") + assert agent_v2 is not None + + assert agent_v1.agent_id == agent_v2.agent_id + + +def test_config_agent_ids_differ_when_only_the_agent_name_differs(): + """Two entries identical except for agent_name must not collapse onto one id.""" + registry: Final = AgentRegistry() + registry.load_agents_from_config( + [ + {"agent_name": "agent-a", "agent_card_params": _sample_agent_card_params()}, + {"agent_name": "agent-b", "agent_card_params": _sample_agent_card_params()}, + ] + ) + + ids: Final = {agent.agent_id for agent in registry.get_agent_list()} + assert len(ids) == 2 + + +def test_legacy_full_entry_hash_still_resolves_the_config_agent(): + """Grants and clients created before LIT-5144 hold the old full-entry hash; it must keep resolving.""" + entry: Final = { + "agent_name": "legacy-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("legacy-agent") + assert agent is not None + + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + assert registry.config_agent_legacy_ids[legacy_id] == agent.agent_id + assert legacy_id in registry.ids_for_agent(agent.agent_id) + assert agent.agent_id in registry.ids_for_agent(agent.agent_id) + + resolved: Final = registry.get_agent_by_id(legacy_id) + assert resolved is not None + assert resolved.agent_id == agent.agent_id + assert registry.get_agent_by_id("nonexistent-id") is None + + +def test_public_agent_groups_holding_the_legacy_id_still_mark_the_config_agent_public(monkeypatch): + """LIT-5144: config.yaml written before the fix stores the full-entry hash in + public_agent_groups; the agent must stay public after its id became name-based.""" + import litellm + + entry: Final = { + "agent_name": "public-agent", + "agent_card_params": _sample_agent_card_params(), + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("public-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + + monkeypatch.setattr(litellm, "public_agent_groups", [legacy_id]) + assert [a.agent_id for a in registry.get_public_agent_list()] == [agent.agent_id] + + monkeypatch.setattr(litellm, "public_agent_groups", ["unrelated-id"]) + assert registry.get_public_agent_list() == () + + monkeypatch.setattr(litellm, "public_agent_groups", None) + assert registry.get_public_agent_list() == () + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_persists_stable_ids_into_grant_rows(): + """LIT-5144: the startup migration rewrites stored legacy full-entry hashes to the stable + name id, so a later secret rotation (which re-mints the legacy hash) cannot orphan grants.""" + entry: Final = { + "agent_name": "migrated-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("migrated-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id, "unrelated-id", agent.agent_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=1) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=1, missed=0) + table.find_many.assert_awaited_once_with(where={"agents": {"has_some": (legacy_id,)}}) + table.update_many.assert_awaited_once_with( + where={"object_permission_id": "op-1", "agents": {"equals": (legacy_id, "unrelated-id", agent.agent_id)}}, + data={"agents": (agent.agent_id, "unrelated-id")}, + ) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_reports_compare_and_swap_misses(): + """A concurrently edited row makes the CAS update affect zero rows; the result must + surface that as missed so the startup task knows to retry instead of reporting success.""" + entry: Final = { + "agent_name": "contended-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=0) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=1) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_no_ops_without_config_agents(): + """Without config agents there are no legacy hashes to translate, so the DB is never queried.""" + registry: Final = AgentRegistry() + table: Final = MagicMock() + table.find_many = AsyncMock() + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=0) + table.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 3e097711ad7..6dbd7475be5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -308,7 +308,7 @@ async def test_attach_keys_to_agents_groups_by_agent_and_omits_secret(): # Query is scoped to the agents being returned, not the whole key table. where = mock_prisma.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] - assert where == {"agent_id": {"in": ["agent-1", "agent-2"]}} + assert where == {"agent_id": {"in": ("agent-1", "agent-2")}} # agent-1 gets both of its keys; agent-2 gets None. assert agent_without_keys.keys is None @@ -503,6 +503,7 @@ class TestAgentRBACProxyAdminViewOnly: ] self.mock_registry = MagicMock() self.mock_registry.get_agent_list = MagicMock(return_value=self.agents) + self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"]) diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 8b8b7871cce..88dc07e741b 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -582,6 +582,7 @@ def test_public_agent_hub_rewrites_upstream_url_to_proxy(): mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), @@ -631,6 +632,7 @@ def test_public_agent_hub_serializes_http_security_scheme_without_bearer_format( mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2921592acbd..b824f5bb550 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23256 + "limit": 23250 }, "LIT002": { "limit": 27195 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16783 + "limit": 16777 }, "LIT011": { "limit": 5602