mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
chore: merge litellm_internal_staging
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
f97d9b8707
13 changed files with 455 additions and 50 deletions
27
.github/pull_request_template.md
vendored
27
.github/pull_request_template.md
vendored
|
|
@ -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" -->
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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,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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue