mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
bf2def7eb5
commit
eb3c8c168f
12 changed files with 426 additions and 48 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -135,7 +135,7 @@
|
|||
"limit": 27
|
||||
},
|
||||
"PERF401": {
|
||||
"limit": 23
|
||||
"limit": 13
|
||||
},
|
||||
"PERF402": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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"]),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue