chore: merge litellm_internal_staging

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-07 19:21:41 +00:00
commit f97d9b8707
13 changed files with 455 additions and 50 deletions

View file

@ -13,6 +13,33 @@ How it solves it:
- <blah>
- ...
## User Flow
<!-- Two ordered lists, Before and After, walking the same end user through the same task, written strictly from that user's seat
Read the linked issue, ticket, or customer thread first so the flow reflects the real application and the routes its users actually hit; don't invent a generic scenario
Lead each list with one plain sentence saying where the flow fails (Before) or succeeds (After), then number the steps
Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
Keep the two lists step-for-step identical until they diverge, so the changed step is obvious
If the bug had a security or authorization consequence, end each list with what another user could or could no longer do
Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision
Example:
Before: a developer whose app streams chat completions gets no token counts back, so their cost dashboard reads zero
1. They send POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
2. The last SSE chunk arrives with `"usage": null`, so their app records 0 prompt and 0 completion tokens
3. They open https://litellm-domain/ui/?page=logs and see the request logged at $0 spend
After: the same request comes back with real token counts, so the dashboard shows real spend
1. The proxy admin sets `always_include_stream_usage: true` and restarts the proxy
2. The developer sends the same POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
3. The last SSE chunk now carries a `usage` object with real prompt and completion token counts
4. https://litellm-domain/ui/?page=logs shows that request at non-zero spend
-->
## Relevant issues
<!-- e.g., "Fixes #000" -->

View file

@ -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}")

View file

@ -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(

View file

@ -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",

View file

@ -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.

View file

@ -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)
]

View file

@ -46,6 +46,10 @@ class AgentsRepository(PrismaTableRepository):
table_name = "litellm_agentstable"
class ObjectPermissionRepository(PrismaTableRepository):
table_name = "litellm_objectpermissiontable"
class GuardrailsRepository(PrismaTableRepository):
table_name = "litellm_guardrailstable"

View file

@ -135,7 +135,7 @@
"limit": 27
},
"PERF401": {
"limit": 23
"limit": 13
},
"PERF402": {
"limit": 0

View file

@ -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)

View file

@ -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()

View file

@ -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"])

View file

@ -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"]),

View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23232
"limit": 23250
},
"LIT002": {
"limit": 27156
"limit": 27195
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1067
"limit": 1091
},
"LIT007": {
"limit": 0
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16765
"limit": 16777
},
"LIT011": {
"limit": 5602