mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
25109a523b
* test(proxy): move utils, agent_endpoints and endpoint tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): package moved unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): exclude proxy-db-owned files from the misc target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop the redundant fixture docstrings in the proxy conftest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1633 lines
64 KiB
Python
1633 lines
64 KiB
Python
"""Unit tests for AgentRegistry DB operations."""
|
|
|
|
import hashlib
|
|
import json
|
|
from collections.abc import Mapping
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from prisma.models import LiteLLM_AgentsTable
|
|
|
|
from litellm.constants import REDACTED_BY_LITELM_STRING
|
|
from litellm.proxy.agent_endpoints.agent_registry import (
|
|
AgentRegistry,
|
|
GrantMigrationResult,
|
|
_restore_redacted_litellm_params,
|
|
redact_sensitive_agent_litellm_params,
|
|
)
|
|
from litellm.types.agents import PatchAgentRequest
|
|
|
|
# Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression
|
|
# fixtures) -- never a real key shape, and must never appear in any response.
|
|
SENTINEL_AWS_ACCESS_KEY_ID: Final = "AKIATESTSENTINEL0000"
|
|
SENTINEL_AWS_SECRET_ACCESS_KEY: Final = "test-sentinel-do-not-use-secret-value"
|
|
|
|
|
|
def _sample_agent_card_params() -> dict:
|
|
return {
|
|
"protocolVersion": "1.0",
|
|
"name": "Test Agent",
|
|
"description": "desc",
|
|
"url": "http://localhost",
|
|
"version": "1.0.0",
|
|
"capabilities": {"streaming": True},
|
|
"defaultInputModes": ["text"],
|
|
"defaultOutputModes": ["text"],
|
|
"skills": [],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_clears_static_headers_and_extra_headers_when_omitted():
|
|
"""
|
|
PUT (full-replace) should clear static_headers and extra_headers when omitted.
|
|
Previously, omitting these fields left stale DB values intact.
|
|
"""
|
|
registry = AgentRegistry()
|
|
mock_prisma = MagicMock()
|
|
|
|
# Simulate existing agent that had headers set
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Updated Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"static_headers": {},
|
|
"extra_headers": [],
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
|
|
# Agent config WITHOUT static_headers or extra_headers (omitted)
|
|
agent_config = {
|
|
"agent_name": "Updated Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
}
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent=agent_config,
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
mock_update.assert_awaited_once()
|
|
call_kwargs = mock_update.call_args.kwargs
|
|
update_data = call_kwargs["data"]
|
|
|
|
# Should include static_headers and extra_headers with empty defaults
|
|
assert "static_headers" in update_data
|
|
assert update_data["static_headers"] == "{}"
|
|
assert "extra_headers" in update_data
|
|
assert update_data["extra_headers"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_preserves_explicit_static_headers_and_extra_headers():
|
|
"""PUT with explicit values should still work correctly."""
|
|
registry = AgentRegistry()
|
|
mock_prisma = MagicMock()
|
|
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Updated Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"static_headers": {"Authorization": "Bearer xyz"},
|
|
"extra_headers": ["X-Custom-Header"],
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
|
|
agent_config = {
|
|
"agent_name": "Updated Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"static_headers": {"Authorization": "Bearer xyz"},
|
|
"extra_headers": ["X-Custom-Header"],
|
|
}
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent=agent_config,
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
call_kwargs = mock_update.call_args.kwargs
|
|
update_data = call_kwargs["data"]
|
|
|
|
assert update_data["static_headers"] == '{"Authorization": "Bearer xyz"}'
|
|
assert update_data["extra_headers"] == ["X-Custom-Header"]
|
|
|
|
|
|
def test_load_agents_from_db_and_config_retains_previously_loaded_config_agents():
|
|
"""
|
|
A DB reload calls this without an explicit agent_config. It resets the
|
|
registry, so it must fall back to the agents remembered from config.yaml
|
|
instead of dropping them.
|
|
"""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_config(
|
|
[
|
|
{
|
|
"agent_name": "config-agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
]
|
|
)
|
|
|
|
registry.load_agents_from_db_and_config(
|
|
db_agents=[
|
|
{
|
|
"agent_id": "db-id",
|
|
"agent_name": "db-agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
]
|
|
)
|
|
|
|
assert sorted(agent.agent_name for agent in registry.get_agent_list()) == [
|
|
"config-agent",
|
|
"db-agent",
|
|
]
|
|
|
|
|
|
def test_load_agents_from_db_and_config_skips_incomplete_config_entries():
|
|
"""Config entries missing agent_card_params are skipped, not registered half-built."""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_config([{"agent_name": "no-card"}])
|
|
|
|
registry.load_agents_from_db_and_config(db_agents=None)
|
|
|
|
assert registry.get_agent_list() == ()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"db_agent_names",
|
|
[
|
|
("shared-name", "other-db-agent"),
|
|
("other-db-agent", "shared-name"),
|
|
],
|
|
ids=["colliding-db-row-first", "colliding-db-row-last"],
|
|
)
|
|
def test_load_agents_from_db_and_config_lets_db_rows_win_a_name_collision(db_agent_names):
|
|
"""
|
|
A config entry that reuses a DB agent's name must not be registered next to it.
|
|
|
|
Name lookups and deregistration both address a single agent, so two entries
|
|
under one name would shadow the DB row for routing and let a single delete
|
|
drop both. Parametrized over both DB row orders because the registry is a
|
|
list and a first-match lookup would otherwise pass on ordering luck.
|
|
"""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_config(
|
|
[
|
|
{
|
|
"agent_name": "shared-name",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
},
|
|
{
|
|
"agent_name": "config-only",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
},
|
|
]
|
|
)
|
|
|
|
registry.load_agents_from_db_and_config(
|
|
db_agents=[
|
|
{
|
|
"agent_id": f"db-id-{name}",
|
|
"agent_name": name,
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
for name in db_agent_names
|
|
]
|
|
)
|
|
|
|
registered = registry.get_agent_list()
|
|
assert sorted(agent.agent_name for agent in registered) == [
|
|
"config-only",
|
|
"other-db-agent",
|
|
"shared-name",
|
|
]
|
|
|
|
shared = registry.get_agent_by_name("shared-name")
|
|
assert shared is not None
|
|
assert shared.agent_id == "db-id-shared-name", "the DB row must win the shared name, not the config entry"
|
|
|
|
registry.deregister_agent("shared-name")
|
|
assert [agent.agent_name for agent in registry.get_agent_list() if agent.agent_name == "shared-name"] == [], (
|
|
"deleting the DB agent must not leave a shadowed config duplicate behind"
|
|
)
|
|
|
|
|
|
def test_load_agents_from_config_skips_a_name_already_held_by_a_db_agent():
|
|
"""
|
|
The one-agent-per-name rule is enforced by the loader, not by the call order.
|
|
|
|
A config load that lands on a registry already holding DB rows must not append a
|
|
colliding entry, otherwise the duplicate is reachable any time the two sources are
|
|
loaded in the other order.
|
|
"""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_db_and_config(
|
|
db_agents=[
|
|
{
|
|
"agent_id": "db-id",
|
|
"agent_name": "shared-name",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
]
|
|
)
|
|
|
|
registry.load_agents_from_config(
|
|
[
|
|
{
|
|
"agent_name": "shared-name",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
]
|
|
)
|
|
|
|
assert [agent.agent_id for agent in registry.get_agent_list()] == ["db-id"]
|
|
|
|
|
|
def test_load_agents_from_config_registers_one_agent_per_name_within_the_config():
|
|
"""Two config entries sharing a name collapse to the first, keeping name lookups unambiguous."""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_config(
|
|
[
|
|
{"agent_name": "dupe", "agent_card_params": _sample_agent_card_params()},
|
|
{"agent_name": "dupe", "agent_card_params": {**_sample_agent_card_params(), "url": "http://second"}},
|
|
]
|
|
)
|
|
|
|
registered = registry.get_agent_list()
|
|
assert len(registered) == 1
|
|
assert registered[0].agent_card_params["url"] == "http://localhost"
|
|
|
|
|
|
def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents():
|
|
"""
|
|
An explicitly empty config must forget the previously loaded agents.
|
|
|
|
Otherwise the next DB rebuild replays agents the operator removed from config.yaml.
|
|
``None`` still means "no opinion" and leaves them alone.
|
|
"""
|
|
registry = AgentRegistry()
|
|
registry.load_agents_from_config(
|
|
[{"agent_name": "removed-agent", "agent_card_params": _sample_agent_card_params()}]
|
|
)
|
|
assert registry.config_agents != ()
|
|
|
|
registry.load_agents_from_config([])
|
|
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"
|
|
|
|
|
|
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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_raises_when_row_deleted_mid_update():
|
|
"""Prisma's update returns None when the row vanished between read and write. Without a
|
|
guard the code dereferences None and reports an opaque AttributeError instead of the id."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(SimpleNamespace(litellm_params={}, object_permission_id=None))
|
|
)
|
|
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
|
|
|
|
with pytest.raises(Exception, match="Agent not found") as exc_info:
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"agent_name": "Updated Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
assert str(exc_info.value) == "Agent not found, passed agent_id=agent-123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_raises_when_row_deleted_mid_update():
|
|
"""Same race on PATCH: the existing row is read, then deleted before the update lands."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{"agent_id": "agent-123", "agent_name": "Old Agent", "object_permission_id": None}
|
|
)
|
|
)
|
|
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
|
|
|
|
with pytest.raises(Exception, match="Agent not found") as exc_info:
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={"agent_name": "Patched Agent"},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
assert str(exc_info.value) == "Agent not found, passed agent_id=agent-123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_agent_from_db_raises_when_row_already_gone() -> None:
|
|
registry: Final = AgentRegistry()
|
|
database: Final = MagicMock()
|
|
tx: Final = database.tx.return_value.__aenter__.return_value
|
|
tx.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
with pytest.raises(ValueError, match="Agent not found, passed agent_id=agent-123"):
|
|
await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=database)
|
|
tx.litellm_verificationtoken.delete_many.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("managed", [True, False])
|
|
async def test_agent_deletion_revokes_managed_keys_and_keeps_identity_history(managed: bool) -> None:
|
|
registry: Final = AgentRegistry()
|
|
database: Final = MagicMock()
|
|
tx: Final = database.tx.return_value.__aenter__.return_value
|
|
row: Final = _stored_agent_row({"agent_id": "agent-123", "identity_managed": managed})
|
|
tx.litellm_agentstable.find_unique = AsyncMock(return_value=row)
|
|
tx.litellm_agentstable.delete = AsyncMock(return_value=row)
|
|
tx.litellm_verificationtoken.delete_many = AsyncMock(return_value=2)
|
|
tx.litellm_retiredagent.upsert = AsyncMock()
|
|
result: Final = await registry.delete_agent_from_db("agent-123", database)
|
|
assert result["agent_id"] == "agent-123"
|
|
tx.litellm_agentstable.delete.assert_awaited_once_with(where={"agent_id": "agent-123"})
|
|
if managed:
|
|
tx.litellm_retiredagent.upsert.assert_awaited_once_with(
|
|
where={"original_agent_id": "agent-123"},
|
|
data={"create": {"original_agent_id": "agent-123"}, "update": {}},
|
|
)
|
|
tx.litellm_verificationtoken.delete_many.assert_awaited_once_with(where={"agent_id": "agent-123"})
|
|
else:
|
|
tx.litellm_retiredagent.upsert.assert_not_awaited()
|
|
tx.litellm_verificationtoken.delete_many.assert_not_awaited()
|
|
|
|
|
|
# ---------- LIT-6736: agent litellm_params secret redaction ----------
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_masks_secrets_keeps_the_rest():
|
|
"""The sentinel secret must never appear in the redacted output; non-secret
|
|
keys (model reference, is_public) must survive untouched."""
|
|
redacted = redact_sensitive_agent_litellm_params(
|
|
{
|
|
"aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID,
|
|
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
"is_public": True,
|
|
}
|
|
)
|
|
|
|
assert SENTINEL_AWS_ACCESS_KEY_ID not in json.dumps(redacted)
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted)
|
|
assert redacted["aws_access_key_id"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["model"] == "bedrock/agentcore/my-agent"
|
|
assert redacted["is_public"] is True
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_recurses_into_nested_dicts():
|
|
"""A secret nested one level down (e.g. a per-provider sub-config) must
|
|
also be redacted, not just top-level keys."""
|
|
redacted = redact_sensitive_agent_litellm_params(
|
|
{"provider_config": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}}
|
|
)
|
|
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted)
|
|
assert redacted["provider_config"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["provider_config"]["region"] == "us-east-1"
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_handles_none_and_json_string():
|
|
assert redact_sensitive_agent_litellm_params(None) is None
|
|
|
|
serialized = json.dumps({"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"})
|
|
redacted = redact_sensitive_agent_litellm_params(serialized)
|
|
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY not in redacted
|
|
assert json.loads(redacted)["api_key"] == REDACTED_BY_LITELM_STRING
|
|
assert json.loads(redacted)["model"] == "gpt-4"
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_recurses_into_lists_of_dicts():
|
|
"""A secret nested inside a list of provider sub-configs (a shape a
|
|
non-sensitively-named key can legitimately hold) must also be redacted,
|
|
not silently returned as-is."""
|
|
redacted = redact_sensitive_agent_litellm_params(
|
|
{
|
|
"provider_configs": [
|
|
{"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"},
|
|
{"aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-west-2"},
|
|
]
|
|
}
|
|
)
|
|
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted)
|
|
assert redacted["provider_configs"][0]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["provider_configs"][0]["region"] == "us-east-1"
|
|
assert redacted["provider_configs"][1]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["provider_configs"][1]["region"] == "us-west-2"
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_redacts_secrets_inside_model_list():
|
|
"""The exact shape flagged in review: litellm_params.model_list, where each
|
|
entry carries its own nested litellm_params with a provider credential."""
|
|
redacted = redact_sensitive_agent_litellm_params(
|
|
{
|
|
"model_list": [
|
|
{
|
|
"model_name": "gpt-4",
|
|
"litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"},
|
|
},
|
|
{
|
|
"model_name": "claude",
|
|
"litellm_params": {
|
|
"aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY,
|
|
"model": "bedrock/claude",
|
|
},
|
|
},
|
|
]
|
|
}
|
|
)
|
|
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted)
|
|
assert redacted["model_list"][0]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["model_list"][0]["litellm_params"]["model"] == "gpt-4"
|
|
assert redacted["model_list"][1]["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert redacted["model_list"][1]["litellm_params"]["model"] == "bedrock/claude"
|
|
|
|
|
|
def test_restore_redacted_litellm_params_preserves_secret_inside_model_list():
|
|
"""The write-side counterpart: a caller editing a model_list entry's own
|
|
non-secret field (renaming it) while leaving that same entry's nested
|
|
secret masked must not corrupt the stored per-deployment credential.
|
|
List entries correspond by position (see the module docstring on
|
|
``_restore_redacted_nested_value``), so this -- the common "edit this
|
|
entry, keep its secret" pattern -- must keep working."""
|
|
existing = {
|
|
"agent_name": "my-agent",
|
|
"model_list": [
|
|
{
|
|
"model_name": "gpt-4",
|
|
"litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"},
|
|
},
|
|
],
|
|
}
|
|
incoming = {
|
|
"agent_name": "my-agent-renamed",
|
|
"model_list": [
|
|
{
|
|
"model_name": "gpt-4-renamed",
|
|
"litellm_params": {"api_key": REDACTED_BY_LITELM_STRING, "model": "gpt-4"},
|
|
},
|
|
],
|
|
}
|
|
|
|
restored = _restore_redacted_litellm_params(incoming, existing)
|
|
|
|
assert SENTINEL_AWS_SECRET_ACCESS_KEY == restored["model_list"][0]["litellm_params"]["api_key"]
|
|
assert restored["model_list"][0]["model_name"] == "gpt-4-renamed"
|
|
assert restored["agent_name"] == "my-agent-renamed"
|
|
|
|
|
|
def test_restore_redacted_litellm_params_matches_list_entries_by_position():
|
|
"""Documents the accepted trade-off: a list has no stable per-element
|
|
identity in a plain ``dict[str, object]`` schema, so restoration matches
|
|
entries by index, the same correspondence every other part of this merge
|
|
(and the endpoints' full-replace-on-PUT semantics) already assumes. If a
|
|
caller both reorders the list AND echoes back a masked marker in the same
|
|
request, a credential can end up attached to a different logical entry.
|
|
That is a known, narrow limitation -- not a leak between different
|
|
agents or tenants, since it only reshuffles one agent's own stored
|
|
values -- and this test pins the current, deliberate behavior rather
|
|
than asserting it away."""
|
|
existing = {
|
|
"model_list": [
|
|
{"model_name": "gpt-4", "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY}},
|
|
{"model_name": "claude", "litellm_params": {"api_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY}},
|
|
],
|
|
}
|
|
incoming = {
|
|
"model_list": [
|
|
# Same index (0) now holds what used to be at index 1's entry.
|
|
{"model_name": "claude", "litellm_params": {"api_key": REDACTED_BY_LITELM_STRING}},
|
|
],
|
|
}
|
|
|
|
restored = _restore_redacted_litellm_params(incoming, existing)
|
|
|
|
assert restored["model_list"][0]["litellm_params"]["api_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
|
|
|
|
|
def test_restore_redacted_litellm_params_recovers_a_whole_subtree_collapsed_by_the_depth_cap():
|
|
"""Past the read-side recursion depth cap, a whole nested subtree is
|
|
collapsed to the flat REDACTED_BY_LITELM marker rather than a dict/list.
|
|
If the caller echoes that flat marker back unchanged, the whole
|
|
subtree -- not just the literal marker string -- must be restored."""
|
|
existing_subtree = {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}
|
|
incoming = {"provider_config": REDACTED_BY_LITELM_STRING}
|
|
existing = {"provider_config": existing_subtree}
|
|
|
|
restored = _restore_redacted_litellm_params(incoming, existing)
|
|
|
|
assert restored["provider_config"] == existing_subtree
|
|
|
|
|
|
def test_redact_sensitive_agent_litellm_params_does_not_reinterpret_plain_string_values_as_json():
|
|
"""A plain non-JSON string value (most string leaves) must pass through
|
|
unchanged rather than failing to parse and getting redacted."""
|
|
redacted = redact_sensitive_agent_litellm_params({"model": "bedrock/agentcore/my-agent", "is_public": True})
|
|
|
|
assert redacted["model"] == "bedrock/agentcore/my-agent"
|
|
assert redacted["is_public"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_agent_to_db_drops_a_sentinel_value_instead_of_storing_the_placeholder():
|
|
"""A create has nothing stored to restore behind a redaction marker, so a
|
|
sensitive key submitted as the literal marker is dropped rather than
|
|
persisted as the placeholder string itself."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
created_agent = MagicMock()
|
|
created_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
created_agent.object_permission = None
|
|
mock_create = AsyncMock(return_value=created_agent)
|
|
mock_prisma.db.litellm_agentstable.create = mock_create
|
|
|
|
await registry.add_agent_to_db(
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {
|
|
"aws_secret_access_key": REDACTED_BY_LITELM_STRING,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_create.call_args.kwargs["data"]["litellm_params"])
|
|
assert "aws_secret_access_key" not in stored_params
|
|
assert stored_params["model"] == "bedrock/agentcore/my-agent"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|
"""PUT round-trips the GET response, which shows the secret redacted. Saving
|
|
an unrelated field change must not overwrite the real stored credential
|
|
with the redaction marker."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(
|
|
litellm_params={
|
|
"aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID,
|
|
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
},
|
|
object_permission_id=None,
|
|
)
|
|
)
|
|
)
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Renamed Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"agent_name": "Renamed Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
# The UI round-tripped the redacted secret and the untouched
|
|
# access key id verbatim; only agent_name actually changed.
|
|
"litellm_params": {
|
|
"aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID,
|
|
"aws_secret_access_key": REDACTED_BY_LITELM_STRING,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
|
assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
|
assert stored_params["aws_access_key_id"] == SENTINEL_AWS_ACCESS_KEY_ID
|
|
assert stored_params["model"] == "bedrock/agentcore/my-agent"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely():
|
|
"""Omitting the sensitive key altogether must fall back to the stored
|
|
value too, not just an explicit redaction-marker round-trip."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(
|
|
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
|
object_permission_id=None,
|
|
)
|
|
)
|
|
)
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"model": "bedrock/agentcore/my-agent"},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
|
assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_key():
|
|
"""A secret nested inside a dict held by a non-sensitively-named key
|
|
(e.g. a per-provider sub-config) must also survive an echoed-back
|
|
redaction marker, not just top-level secret keys."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(
|
|
litellm_params={
|
|
"provider_config": {
|
|
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
|
"region": "us-east-1",
|
|
}
|
|
},
|
|
object_permission_id=None,
|
|
)
|
|
)
|
|
)
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {
|
|
# The GET response redacted the nested secret; the caller
|
|
# round-trips it verbatim while changing nothing.
|
|
"provider_config": {
|
|
"aws_secret_access_key": REDACTED_BY_LITELM_STRING,
|
|
"region": "us-west-2",
|
|
}
|
|
},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
|
assert stored_params["provider_config"]["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
|
assert stored_params["provider_config"]["region"] == "us-west-2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_clears_secret_on_explicit_empty_value():
|
|
"""An explicit empty string is a deliberate clear, distinct from an omitted
|
|
key or the redaction marker, and must actually clear the stored secret."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(
|
|
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
|
object_permission_id=None,
|
|
)
|
|
)
|
|
)
|
|
updated_agent = MagicMock()
|
|
updated_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
updated_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=updated_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"aws_secret_access_key": ""},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
|
assert stored_params["aws_secret_access_key"] == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_preserves_secret_when_litellm_params_omitted():
|
|
"""A PATCH that only renames the agent must not touch (let alone drop) the
|
|
stored litellm_params secret."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Old Name",
|
|
"litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
|
"object_permission_id": None,
|
|
}
|
|
)
|
|
)
|
|
patched_agent = MagicMock()
|
|
patched_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "New Name",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
|
"object_permission": None,
|
|
}
|
|
patched_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=patched_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={"agent_name": "New Name"},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
update_data: Final = mock_update.call_args.kwargs["data"]
|
|
assert "litellm_params" not in update_data
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|
"""A PATCH that includes litellm_params (e.g. to flip an unrelated flag)
|
|
with the secret round-tripped as the redaction marker must not clobber
|
|
the stored credential."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"litellm_params": {
|
|
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
|
"is_public": False,
|
|
},
|
|
"object_permission_id": None,
|
|
}
|
|
)
|
|
)
|
|
patched_agent = MagicMock()
|
|
patched_agent.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
}
|
|
patched_agent.object_permission = None
|
|
mock_update = AsyncMock(return_value=patched_agent)
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"litellm_params": {
|
|
"aws_secret_access_key": REDACTED_BY_LITELM_STRING,
|
|
"is_public": True,
|
|
}
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="test-user",
|
|
)
|
|
|
|
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
|
assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
|
assert stored_params["is_public"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["patch", "put"])
|
|
async def test_runtime_update_drops_legacy_identity_and_keeps_agent_id(operation: str) -> None:
|
|
registry: Final = AgentRegistry()
|
|
prisma: Final = MagicMock()
|
|
identity: Final = {
|
|
"provider": "microsoft_entra",
|
|
"tenant_id": "11111111-1111-4111-8111-111111111111",
|
|
"client_id": "22222222-2222-4222-8222-222222222222",
|
|
}
|
|
existing_params: Final = {"identity": identity, "model": "old"}
|
|
existing: Final = (
|
|
SimpleNamespace(litellm_params=existing_params, object_permission_id=None)
|
|
if operation == "put"
|
|
else {"agent_name": "Readable agent", "litellm_params": existing_params}
|
|
)
|
|
prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row(existing))
|
|
saved: Final = MagicMock()
|
|
saved.object_permission = None
|
|
saved.model_dump.return_value = {
|
|
"agent_id": "unchanged-id",
|
|
"agent_name": "Renamed agent",
|
|
"agent_card_params": {},
|
|
"litellm_params": {"model": "new"},
|
|
}
|
|
prisma.db.litellm_agentstable.update = AsyncMock(return_value=saved)
|
|
update: Final = registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db
|
|
result: Final = await update(
|
|
agent_id="unchanged-id",
|
|
agent={"agent_name": "Renamed agent", "agent_card_params": {}, "litellm_params": {"model": "new"}},
|
|
prisma_client=prisma,
|
|
updated_by="admin",
|
|
)
|
|
stored: Final = prisma.db.litellm_agentstable.update.call_args.kwargs
|
|
assert stored["where"] == {"agent_id": "unchanged-id"}
|
|
assert json.loads(stored["data"]["litellm_params"]) == {"model": "new"}, (
|
|
"a stored litellm_params.identity must not be resurrected once the JWT path no longer honours it"
|
|
)
|
|
assert result.agent_id == "unchanged-id"
|
|
assert "object_permission_id" not in stored["data"]
|
|
|
|
|
|
def _agent_row_mock(access_group_ids: list[str]) -> MagicMock:
|
|
row: Final = MagicMock()
|
|
row.model_dump.return_value = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {},
|
|
"object_permission": None,
|
|
"access_group_ids": access_group_ids,
|
|
}
|
|
row.object_permission = None
|
|
return row
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_agent_to_db_persists_deduplicated_access_group_ids():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_create = AsyncMock(return_value=_agent_row_mock(["ag-1", "ag-2"]))
|
|
mock_prisma.db.litellm_agentstable.create = mock_create
|
|
|
|
result: Final = await registry.add_agent_to_db(
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"access_group_ids": ["ag-1", "ag-2", "ag-1"],
|
|
},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert tuple(mock_create.call_args.kwargs["data"]["access_group_ids"]) == ("ag-1", "ag-2")
|
|
assert result.access_group_ids == ["ag-1", "ag-2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_default():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_create = AsyncMock(return_value=_agent_row_mock([]))
|
|
mock_prisma.db.litellm_agentstable.create = mock_create
|
|
|
|
await registry.add_agent_to_db(
|
|
agent={"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params()},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert "access_group_ids" not in mock_create.call_args.kwargs["data"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("patch_body", "expected"),
|
|
[
|
|
({"access_group_ids": ["ag-2", "ag-3"]}, ["ag-2", "ag-3"]),
|
|
({"access_group_ids": []}, []),
|
|
({"access_group_ids": None}, []),
|
|
],
|
|
)
|
|
async def test_patch_agent_in_db_replaces_access_group_ids_when_provided(
|
|
patch_body: PatchAgentRequest, expected: list[str]
|
|
):
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"litellm_params": {},
|
|
"object_permission_id": None,
|
|
"access_group_ids": ["ag-1"],
|
|
}
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123", agent=patch_body, prisma_client=mock_prisma, updated_by="test-user"
|
|
)
|
|
|
|
assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Old Name",
|
|
"litellm_params": {},
|
|
"object_permission_id": None,
|
|
"access_group_ids": ["ag-1"],
|
|
}
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"]))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123", agent={"agent_name": "New Name"}, prisma_client=mock_prisma, updated_by="test-user"
|
|
)
|
|
|
|
assert "access_group_ids" not in mock_update.call_args.kwargs["data"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("body_access_group_ids", "expected"),
|
|
[(["ag-9", "ag-9"], ["ag-9"]), (None, []), ("omitted", [])],
|
|
)
|
|
async def test_update_agent_in_db_always_writes_access_group_ids(body_access_group_ids, expected: list[str]):
|
|
"""PUT is a full replacement: omitting the field clears any previously attached groups."""
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"])
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
body: Final = {
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"model": "bedrock/agentcore/my-agent"},
|
|
**({} if body_access_group_ids == "omitted" else {"access_group_ids": body_access_group_ids}),
|
|
}
|
|
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123", agent=body, prisma_client=mock_prisma, updated_by="test-user"
|
|
)
|
|
|
|
assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected)
|
|
|
|
|
|
def _stored_agent_row(values: Mapping[str, object] | SimpleNamespace) -> LiteLLM_AgentsTable:
|
|
fields: Final = vars(values) if isinstance(values, SimpleNamespace) else values
|
|
return LiteLLM_AgentsTable.model_validate(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": "{}",
|
|
"extra_headers": [],
|
|
"agent_access_groups": [],
|
|
"access_group_ids": [],
|
|
"created_at": datetime.now(timezone.utc),
|
|
"updated_at": datetime.now(timezone.utc),
|
|
"created_by": "admin",
|
|
"updated_by": "admin",
|
|
"spend": 0,
|
|
"identity_managed": False,
|
|
"enabled": True,
|
|
"execution_mode": "autonomous",
|
|
**{
|
|
key: json.dumps(value)
|
|
if key in ("litellm_params", "agent_card_params", "kill_switch", "static_headers") and not isinstance(value, str)
|
|
else value
|
|
for key, value in fields.items()
|
|
},
|
|
}
|
|
)
|
|
|
|
|
|
_KILL_SWITCH: Final = {
|
|
"url": "https://ops.example.com/kill",
|
|
"method": "POST",
|
|
"headers": {"X-Env": "prod"},
|
|
"query_params": {"reason": "manual"},
|
|
"body": {"action": "stop"},
|
|
"auth": {"type": "bearer", "token": "tok-real"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unset():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_create = AsyncMock(return_value=_agent_row_mock([]))
|
|
mock_prisma.db.litellm_agentstable.create = mock_create
|
|
|
|
await registry.add_agent_to_db(
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"kill_switch": _KILL_SWITCH,
|
|
},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
assert json.loads(mock_create.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH
|
|
|
|
await registry.add_agent_to_db(
|
|
agent={"agent_name": "Plain Agent", "agent_card_params": _sample_agent_card_params()},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
assert mock_create.call_args.kwargs["data"]["kill_switch"] == json.dumps(None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.create = AsyncMock(return_value=_agent_row_mock([]))
|
|
|
|
with pytest.raises(Exception, match="absolute http"):
|
|
await registry.add_agent_to_db(
|
|
agent={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"kill_switch": {**_KILL_SWITCH, "url": "ops.example.com/kill"},
|
|
},
|
|
prisma_client=mock_prisma,
|
|
created_by="test-user",
|
|
)
|
|
mock_prisma.db.litellm_agentstable.create.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on_null():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Old",
|
|
"litellm_params": {},
|
|
"object_permission_id": None,
|
|
"kill_switch": _KILL_SWITCH,
|
|
}
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123", agent={"agent_name": "New"}, prisma_client=mock_prisma, updated_by="u"
|
|
)
|
|
assert "kill_switch" not in mock_update.call_args.kwargs["data"]
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123", agent={"kill_switch": None}, prisma_client=mock_prisma, updated_by="u"
|
|
)
|
|
assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None), (
|
|
"prisma-client-py silently drops None, so the clear must be written as the JSON literal null"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_the_marker():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
{
|
|
"agent_id": "agent-123",
|
|
"agent_name": "A",
|
|
"litellm_params": {},
|
|
"object_permission_id": None,
|
|
"kill_switch": _KILL_SWITCH,
|
|
}
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
|
|
await registry.patch_agent_in_db(
|
|
agent_id="agent-123",
|
|
agent={
|
|
"kill_switch": {
|
|
**_KILL_SWITCH,
|
|
"url": "https://ops.example.com/v2/kill",
|
|
"auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING},
|
|
}
|
|
},
|
|
prisma_client=mock_prisma,
|
|
updated_by="u",
|
|
)
|
|
|
|
assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == {
|
|
**_KILL_SWITCH,
|
|
"url": "https://ops.example.com/v2/kill",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_secret_when_echoed():
|
|
registry: Final = AgentRegistry()
|
|
mock_prisma: Final = MagicMock()
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=_stored_agent_row(
|
|
SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH))
|
|
)
|
|
)
|
|
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
|
mock_prisma.db.litellm_agentstable.update = mock_update
|
|
base: Final = {"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params(), "litellm_params": {}}
|
|
|
|
await registry.update_agent_in_db(agent_id="agent-123", agent=base, prisma_client=mock_prisma, updated_by="u")
|
|
assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None)
|
|
|
|
echoed: Final = {**_KILL_SWITCH, "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}}
|
|
await registry.update_agent_in_db(
|
|
agent_id="agent-123", agent={**base, "kill_switch": echoed}, prisma_client=mock_prisma, updated_by="u"
|
|
)
|
|
assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH
|
|
|
|
|
|
def test_load_agents_from_config_exposes_a_typed_kill_switch():
|
|
registry: Final = AgentRegistry()
|
|
|
|
registry.load_agents_from_config(
|
|
[{"agent_name": "cfg-agent", "agent_card_params": _sample_agent_card_params(), "kill_switch": _KILL_SWITCH}]
|
|
)
|
|
|
|
(agent,) = registry.get_agent_list()
|
|
assert agent.kill_switch is not None
|
|
assert agent.kill_switch.model_dump() == _KILL_SWITCH
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("bound", [False, True])
|
|
async def test_agent_listing_preserves_stored_identity_bindings(bound: bool) -> None:
|
|
from datetime import datetime, timezone
|
|
|
|
from prisma.models import LiteLLM_AgentIdentity, LiteLLM_AgentsTable
|
|
|
|
from litellm.types.agents import AgentResponse
|
|
|
|
binding: Final = LiteLLM_AgentIdentity(
|
|
agent_id="agent",
|
|
provider="microsoft_entra",
|
|
issuer="issuer",
|
|
tenant_id="tenant",
|
|
client_id="client",
|
|
active=True,
|
|
required_roles=[],
|
|
required_scopes=["user_impersonation"],
|
|
revision="revision",
|
|
)
|
|
row: Final = LiteLLM_AgentsTable(
|
|
agent_id="agent",
|
|
agent_name="Bound agent",
|
|
agent_card_params="{}",
|
|
identity_managed=bound,
|
|
identity=binding if bound else None,
|
|
enabled=True,
|
|
execution_mode="autonomous",
|
|
spend=0.0,
|
|
agent_access_groups=[],
|
|
access_group_ids=[],
|
|
extra_headers=[],
|
|
created_by="admin",
|
|
updated_by="admin",
|
|
created_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
|
updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
|
)
|
|
client: Final = MagicMock()
|
|
client.db.litellm_agentstable.find_many = AsyncMock(return_value=[row])
|
|
listed: Final = await AgentRegistry.get_all_agents_from_db(client)
|
|
response: Final = AgentResponse.model_validate(listed[0])
|
|
if bound:
|
|
assert response.identity is not None
|
|
assert response.identity.client_id == binding.client_id
|
|
assert response.identity.revision == binding.revision
|
|
else:
|
|
assert response.identity is None
|
|
client.db.litellm_agentstable.find_many.assert_awaited_once_with(
|
|
order={"created_at": "desc"},
|
|
include={"object_permission": True, "identity": True},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
|
async def test_agent_permissions_are_written_atomically_with_the_registration(operation: str) -> None:
|
|
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
|
|
|
|
registry: Final = AgentRegistry()
|
|
client: Final = MagicMock()
|
|
existing: Final = _stored_agent_row({"agent_id": "agent-123", "object_permission_id": "permissions"})
|
|
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=existing)
|
|
client.db.litellm_agentstable.create = AsyncMock(return_value=existing)
|
|
client.db.litellm_agentstable.update = AsyncMock(return_value=existing)
|
|
client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
|
return_value=(
|
|
LiteLLM_ObjectPermissionTable(object_permission_id="permissions", models=["prior"], mcp_servers=["slack"])
|
|
if operation != "create"
|
|
else None
|
|
)
|
|
)
|
|
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}, "object_permission": {"models": ["new"]}}
|
|
if operation == "create":
|
|
await registry.add_agent_to_db(incoming, client, created_by="admin")
|
|
else:
|
|
update: Final = registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db
|
|
await update("agent-123", incoming, client, updated_by="admin")
|
|
write: Final = (
|
|
client.db.litellm_agentstable.create if operation == "create" else client.db.litellm_agentstable.update
|
|
)
|
|
permission: Final = write.call_args.kwargs["data"]["object_permission"][
|
|
"create" if operation == "create" else "update"
|
|
]
|
|
assert permission["models"] == ["new"]
|
|
if operation != "create":
|
|
assert permission["mcp_servers"] == ["slack"]
|
|
assert permission["object_permission_id"] == "permissions"
|
|
client.db.litellm_objectpermissiontable.update.assert_not_called()
|
|
client.db.litellm_objectpermissiontable.create.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
|
async def test_invalid_identity_fails_before_registration_is_written(operation: str) -> None:
|
|
from fastapi import HTTPException
|
|
|
|
registry: Final = AgentRegistry()
|
|
client: Final = MagicMock()
|
|
client.db.litellm_agentstable.create = AsyncMock()
|
|
client.db.litellm_agentstable.update = AsyncMock()
|
|
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row({"agent_id": "agent-123"}))
|
|
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}, "identity": {"provider": "unknown"}}
|
|
write: Final = (
|
|
registry.add_agent_to_db(incoming, client, created_by="admin")
|
|
if operation == "create"
|
|
else (registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db)(
|
|
"agent-123", incoming, client, updated_by="admin"
|
|
)
|
|
)
|
|
with pytest.raises(HTTPException) as failure:
|
|
await write
|
|
assert failure.value.status_code == 400
|
|
client.db.litellm_agentstable.create.assert_not_awaited()
|
|
client.db.litellm_agentstable.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
|
async def test_duplicate_agent_binding_returns_conflict_for_every_write(operation: str) -> None:
|
|
from fastapi import HTTPException
|
|
from prisma.errors import UniqueViolationError
|
|
|
|
registry: Final = AgentRegistry()
|
|
client: Final = MagicMock()
|
|
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row({"agent_id": "agent-123"}))
|
|
failure: Final = UniqueViolationError(
|
|
{
|
|
"user_facing_error": {
|
|
"message": "Unique constraint failed",
|
|
"meta": {"target": ["client_id"]},
|
|
"error_code": "P2002",
|
|
}
|
|
}
|
|
)
|
|
client.db.litellm_agentstable.create = AsyncMock(side_effect=failure)
|
|
client.db.litellm_agentstable.update = AsyncMock(side_effect=failure)
|
|
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}}
|
|
write: Final = (
|
|
registry.add_agent_to_db(incoming, client, created_by="admin")
|
|
if operation == "create"
|
|
else (registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db)(
|
|
"agent-123", incoming, client, updated_by="admin"
|
|
)
|
|
)
|
|
with pytest.raises(HTTPException) as denied:
|
|
await write
|
|
assert denied.value.status_code == 409
|
|
assert denied.value.detail == "Agent name or Entra application is already registered"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
|
@pytest.mark.parametrize("owner", ["previous-agent", None])
|
|
async def test_retired_application_cannot_transfer_to_another_agent(operation: str, owner: str | None) -> None:
|
|
from fastapi import HTTPException
|
|
|
|
registry: Final = AgentRegistry()
|
|
client: Final = MagicMock()
|
|
row: Final = _stored_agent_row({"agent_id": "agent-123"})
|
|
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=row)
|
|
client.db.litellm_agentstable.create = AsyncMock(return_value=row)
|
|
client.db.litellm_agentstable.update = AsyncMock(return_value=row)
|
|
client.writer_db.litellm_retiredagentidentity.find_unique = AsyncMock(return_value=SimpleNamespace(agent_id=owner))
|
|
incoming: Final = {
|
|
"agent_name": "Agent",
|
|
"agent_card_params": {},
|
|
"identity": {
|
|
"provider": "microsoft_entra",
|
|
"tenant_id": "11111111-1111-4111-8111-111111111111",
|
|
"client_id": "22222222-2222-4222-8222-222222222222",
|
|
"service_principal_id": "33333333-3333-4333-8333-333333333333",
|
|
},
|
|
}
|
|
write: Final = (
|
|
registry.add_agent_to_db(incoming, client, created_by="admin")
|
|
if operation == "create"
|
|
else (registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db)(
|
|
"agent-123", incoming, client, updated_by="admin"
|
|
)
|
|
)
|
|
with pytest.raises(HTTPException) as denied:
|
|
await write
|
|
assert denied.value.status_code == 409
|
|
client.db.litellm_agentstable.create.assert_not_awaited()
|
|
client.db.litellm_agentstable.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
|
@pytest.mark.parametrize("prior_owner", [False, True])
|
|
async def test_application_registration_preserves_its_existing_owner(operation: str, prior_owner: bool) -> None:
|
|
registry: Final = AgentRegistry()
|
|
client: Final = MagicMock()
|
|
row: Final = _stored_agent_row({"agent_id": "agent-123"})
|
|
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=row)
|
|
client.db.litellm_agentstable.create = AsyncMock(return_value=row)
|
|
client.db.litellm_agentstable.update = AsyncMock(return_value=row)
|
|
client.writer_db.litellm_retiredagentidentity.find_unique = AsyncMock(
|
|
return_value=SimpleNamespace(agent_id="agent-123") if prior_owner and operation != "create" else None
|
|
)
|
|
incoming: Final = {
|
|
"agent_name": "Agent",
|
|
"agent_card_params": {},
|
|
"identity": {
|
|
"provider": "microsoft_entra",
|
|
"tenant_id": "11111111-1111-4111-8111-111111111111",
|
|
"client_id": "22222222-2222-4222-8222-222222222222",
|
|
"service_principal_id": "33333333-3333-4333-8333-333333333333",
|
|
},
|
|
}
|
|
if operation == "create":
|
|
result: Final = await registry.add_agent_to_db(incoming, client, created_by="admin")
|
|
else:
|
|
update: Final = registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db
|
|
result = await update("agent-123", incoming, client, updated_by="admin")
|
|
assert result.agent_id == "agent-123"
|
|
write: Final = (
|
|
client.db.litellm_agentstable.create if operation == "create" else client.db.litellm_agentstable.update
|
|
)
|
|
data: Final = write.call_args.kwargs["data"]
|
|
if prior_owner and operation != "create":
|
|
assert "retired_identities" not in data
|
|
else:
|
|
assert data["retired_identities"] == {
|
|
"create": {
|
|
**{key: value for key, value in incoming["identity"].items() if key != "service_principal_id"},
|
|
"issuer": "https://login.microsoftonline.com/11111111-1111-4111-8111-111111111111/v2.0",
|
|
}
|
|
}
|