mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
2792887e47
* fix(proxy): give proxy_admin_viewer read parity with proxy_admin Route-level checks already default-allow management GETs for the viewer role, but ~15 handlers compared user_role to PROXY_ADMIN only, dropping viewers into regular-user scoping (/key/list, /user/info, /model/info, guardrails, prompts, agents, memory, workflows, MCP catalog, coordination redis settings, credential migration check, enterprise projects). Swap those read paths to user_api_key_has_admin_view; write gates unchanged. The dashboard now presents the viewer session as Admin for all gating (effectiveSessionRole) so every page fetches with admin visibility, with userRoleLabel/isViewOnly preserving the account-menu label and the playground cost guard. The server remains the write authority. * refactor(agents): remove side-effectful health_check param from GET /v1/agents Addresses a security review finding on the admin viewer read parity change: listing agents with health_check=true made the proxy issue a server-side GET to every agent URL, so a read-scoped caller could trigger request fan-out beyond their object permissions. The list endpoint is now a pure read for every role. Removes the query param, the URL probing helper and its timeouts, the AgentHealthCheck httpx provider tag, and the dashboard's Health Check toggle. Requests still passing health_check=true get the full list back with the param ignored. * fix(proxy): keep credential encryption check proxy_admin only The residual scan behind GET /credentials/migrate-encryption/check loads every model, credential, MCP, team, and verification-token row and runs a decryption attempt on each stored value. Extending it to proxy_admin_viewer let a read-only account repeatedly trigger deployment-wide scans, so the route keeps its original full-admin gate. * fix(agents): restore health_check, keep list fast path proxy_admin only Restores the agent health_check feature exactly as before this PR: the query param, the URL probing helper, the httpx provider tag, and the dashboard toggle all return, so existing callers keep the filtering contract. The viewer expansion is instead reverted at its source: the GET /v1/agents admin fast path stays PROXY_ADMIN only, so a proxy_admin_viewer goes through the object-permission scoped branch as before and cannot fan out health checks beyond their allowlist. The viewer read of a single agent stays viewer-inclusive since it has no side effects.
951 lines
36 KiB
Python
951 lines
36 KiB
Python
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
|
|
from litellm.proxy.agent_endpoints.endpoints import (
|
|
_attach_keys_to_agents,
|
|
_check_agent_management_permission,
|
|
get_agent_daily_activity,
|
|
router,
|
|
user_api_key_auth,
|
|
)
|
|
from litellm.types.agents import AgentResponse
|
|
|
|
|
|
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": [],
|
|
}
|
|
|
|
|
|
def _sample_agent_config() -> dict:
|
|
return {
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"make_public": False},
|
|
}
|
|
|
|
|
|
def _sample_agent_response(
|
|
agent_id: str = "agent-123", agent_name: str = "Test Agent"
|
|
) -> AgentResponse:
|
|
return AgentResponse(
|
|
agent_id=agent_id,
|
|
agent_name=agent_name,
|
|
agent_card_params=_sample_agent_card_params(),
|
|
litellm_params={"make_public": False},
|
|
)
|
|
|
|
|
|
def _make_app_with_role(role: LitellmUserRoles) -> TestClient:
|
|
"""Create a TestClient where the auth dependency returns the given role."""
|
|
test_app = FastAPI()
|
|
test_app.include_router(router)
|
|
test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
|
user_id="test-user", user_role=role
|
|
)
|
|
return TestClient(test_app)
|
|
|
|
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
|
user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
|
|
)
|
|
client = TestClient(app)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_prisma_client():
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock:
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_user_api_key_auth():
|
|
with patch("litellm.proxy.agent_endpoints.endpoints.user_api_key_auth") as mock:
|
|
mock.return_value = UserAPIKeyAuth(
|
|
user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
|
|
)
|
|
yield mock
|
|
|
|
|
|
def test_update_agent_success(mock_prisma_client, mock_user_api_key_auth, monkeypatch):
|
|
existing_agent = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Existing Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=existing_agent
|
|
)
|
|
|
|
mock_registry = MagicMock()
|
|
mock_registry.update_agent_in_db = AsyncMock(
|
|
return_value=_sample_agent_response(agent_id="agent-123")
|
|
)
|
|
mock_registry.deregister_agent = MagicMock()
|
|
mock_registry.register_agent = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
|
|
|
|
response = client.put(
|
|
"/v1/agents/agent-123",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer test-key"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["agent_id"] == "agent-123"
|
|
assert response.json()["agent_name"] == "Test Agent"
|
|
|
|
|
|
def test_update_agent_not_found(
|
|
mock_prisma_client, mock_user_api_key_auth, monkeypatch
|
|
):
|
|
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
|
|
mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
|
|
|
|
response = client.put(
|
|
"/v1/agents/missing-agent",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer test-key"},
|
|
)
|
|
|
|
assert response.status_code == 404
|
|
assert "Agent with ID missing-agent not found" in response.json()["detail"]
|
|
|
|
|
|
def test_get_agent_by_id_not_found(
|
|
mock_prisma_client, mock_user_api_key_auth, monkeypatch
|
|
):
|
|
mock_registry = MagicMock()
|
|
mock_registry.get_agent_by_id = MagicMock(return_value=None)
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
|
|
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
|
|
response = client.get(
|
|
"/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
|
|
)
|
|
|
|
assert response.status_code == 404
|
|
assert "Agent with ID missing-agent not found" in response.json()["detail"]
|
|
|
|
|
|
def test_delete_agent_not_found(
|
|
mock_prisma_client, mock_user_api_key_auth, monkeypatch
|
|
):
|
|
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
|
|
|
|
response = client.delete(
|
|
"/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
|
|
)
|
|
|
|
assert response.status_code == 404
|
|
assert "Agent with ID missing-agent not found in DB." in response.json()["detail"]
|
|
|
|
|
|
def test_agent_error_schema_consistency(
|
|
mock_prisma_client, mock_user_api_key_auth, monkeypatch
|
|
):
|
|
mock_registry = MagicMock()
|
|
mock_registry.get_agent_by_id = MagicMock(return_value=None)
|
|
mock_registry.update_agent_in_db = AsyncMock(
|
|
side_effect=Exception("should not run")
|
|
)
|
|
mock_registry.delete_agent_from_db = AsyncMock(
|
|
side_effect=Exception("should not run")
|
|
)
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
|
|
|
|
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
|
|
|
missing_agent_id = "missing-agent"
|
|
responses = [
|
|
client.get(
|
|
f"/v1/agents/{missing_agent_id}",
|
|
headers={"Authorization": "Bearer test-key"},
|
|
),
|
|
client.put(
|
|
f"/v1/agents/{missing_agent_id}",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer test-key"},
|
|
),
|
|
client.delete(
|
|
f"/v1/agents/{missing_agent_id}",
|
|
headers={"Authorization": "Bearer test-key"},
|
|
),
|
|
]
|
|
|
|
for resp in responses:
|
|
assert resp.status_code == 404
|
|
detail = resp.json()["detail"]
|
|
assert isinstance(detail, str)
|
|
assert missing_agent_id in detail
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agent_daily_activity_admin_param_passing(monkeypatch):
|
|
mock_prisma = AsyncMock()
|
|
mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
|
|
|
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
|
|
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
|
|
monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)
|
|
|
|
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
|
|
result = await get_agent_daily_activity(
|
|
agent_ids="agent-1,agent-2",
|
|
start_date="2024-01-01",
|
|
end_date="2024-01-31",
|
|
model="gpt-4",
|
|
api_key="test-key",
|
|
page=2,
|
|
page_size=5,
|
|
exclude_agent_ids="agent-3",
|
|
user_api_key_dict=auth,
|
|
)
|
|
|
|
get_daily_activity_mock.assert_awaited_once()
|
|
kwargs = get_daily_activity_mock.call_args.kwargs
|
|
assert kwargs["table_name"] == "litellm_dailyagentspend"
|
|
assert kwargs["entity_id_field"] == "agent_id"
|
|
assert kwargs["entity_id"] == ["agent-1", "agent-2"]
|
|
assert kwargs["exclude_entity_ids"] == ["agent-3"]
|
|
assert kwargs["start_date"] == "2024-01-01"
|
|
assert kwargs["end_date"] == "2024-01-31"
|
|
assert kwargs["model"] == "gpt-4"
|
|
assert kwargs["api_key"] == "test-key"
|
|
assert kwargs["page"] == 2
|
|
assert kwargs["page_size"] == 5
|
|
assert result is mocked_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agent_daily_activity_with_agent_names(monkeypatch):
|
|
mock_prisma = AsyncMock()
|
|
mock_agent1 = MagicMock()
|
|
mock_agent1.agent_id = "agent-1"
|
|
mock_agent1.agent_name = "First Agent"
|
|
mock_agent2 = MagicMock()
|
|
mock_agent2.agent_id = "agent-2"
|
|
mock_agent2.agent_name = "Second Agent"
|
|
|
|
mock_prisma.db.litellm_agentstable.find_many = AsyncMock(
|
|
return_value=[mock_agent1, mock_agent2]
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
|
|
|
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
|
|
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
|
|
monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)
|
|
|
|
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
|
|
await get_agent_daily_activity(
|
|
agent_ids="agent-1,agent-2",
|
|
start_date="2024-01-01",
|
|
end_date="2024-01-31",
|
|
model=None,
|
|
api_key=None,
|
|
page=1,
|
|
page_size=10,
|
|
exclude_agent_ids=None,
|
|
user_api_key_dict=auth,
|
|
)
|
|
|
|
kwargs = get_daily_activity_mock.call_args.kwargs
|
|
assert kwargs["entity_metadata_field"] == {
|
|
"agent-1": {"agent_name": "First Agent"},
|
|
"agent-2": {"agent_name": "Second Agent"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attach_keys_to_agents_groups_by_agent_and_omits_secret():
|
|
"""
|
|
The agents response must carry each agent's attached virtual keys (derived
|
|
from the key table's agent_id FK), grouped per agent, exposing only
|
|
non-secret summary fields. Agents with no key get None so the UI renders
|
|
"Needs Setup" rather than a stale badge.
|
|
"""
|
|
|
|
class _Row:
|
|
def __init__(self, token, agent_id, key_alias, key_name):
|
|
self.token = token
|
|
self.agent_id = agent_id
|
|
self.key_alias = key_alias
|
|
self.key_name = key_name
|
|
self.user_id = "secret-owner" # extra field that must NOT leak
|
|
|
|
agent_with_keys = _sample_agent_response(agent_id="agent-1")
|
|
agent_without_keys = _sample_agent_response(agent_id="agent-2")
|
|
|
|
mock_prisma = MagicMock()
|
|
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
|
return_value=[
|
|
_Row("hash-aaa", "agent-1", "primary", "sk-...aaa"),
|
|
_Row("hash-bbb", "agent-1", "backup", "sk-...bbb"),
|
|
]
|
|
)
|
|
|
|
await _attach_keys_to_agents([agent_with_keys, agent_without_keys], mock_prisma)
|
|
|
|
# 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"]}}
|
|
|
|
# agent-1 gets both of its keys; agent-2 gets None.
|
|
assert agent_without_keys.keys is None
|
|
assert agent_with_keys.keys is not None
|
|
assert {k.token for k in agent_with_keys.keys} == {"hash-aaa", "hash-bbb"}
|
|
assert {k.key_alias for k in agent_with_keys.keys} == {"primary", "backup"}
|
|
|
|
# Only summary fields are exposed; the row's user_id must not be carried.
|
|
summary = agent_with_keys.keys[0]
|
|
assert set(summary.model_dump().keys()) == {"token", "key_alias", "key_name"}
|
|
|
|
|
|
class TestAgentByIdKeyRedaction:
|
|
"""GET /v1/agents/{id} surfaces attached keys to admins but never to
|
|
non-admins, even when the agent has keys attached."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
self.mock_registry = MagicMock()
|
|
self.mock_registry.get_agent_by_id = MagicMock(
|
|
return_value=_sample_agent_response()
|
|
)
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
|
|
|
|
def _get_as(self, role: LitellmUserRoles):
|
|
key_row = MagicMock()
|
|
key_row.token = "hash-aaa"
|
|
key_row.agent_id = "agent-123"
|
|
key_row.key_alias = "primary"
|
|
key_row.key_name = "sk-...aaa"
|
|
|
|
test_client = _make_app_with_role(role)
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=None
|
|
)
|
|
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
|
return_value=[key_row]
|
|
)
|
|
return test_client.get(
|
|
"/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
|
|
)
|
|
|
|
def test_admin_sees_attached_keys(self):
|
|
resp = self._get_as(LitellmUserRoles.PROXY_ADMIN)
|
|
assert resp.status_code == 200
|
|
keys = resp.json()["keys"]
|
|
assert keys is not None
|
|
assert keys[0] == {
|
|
"token": "hash-aaa",
|
|
"key_alias": "primary",
|
|
"key_name": "sk-...aaa",
|
|
}
|
|
|
|
def test_non_admin_never_sees_keys(self):
|
|
resp = self._get_as(LitellmUserRoles.INTERNAL_USER)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["keys"] is None
|
|
|
|
def test_view_only_admin_reads_a_denied_agent_but_still_without_keys(self):
|
|
"""proxy_admin_viewer skips the per-agent object_permission gate (denied
|
|
here) yet stays on the redacted response path."""
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
|
|
AsyncMock(return_value=False),
|
|
):
|
|
resp = self._get_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["keys"] is None
|
|
|
|
|
|
# ---------- RBAC enforcement tests ----------
|
|
|
|
|
|
class TestAgentRBACInternalUser:
|
|
"""Internal users should be able to read agents but not create/update/delete."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
self.internal_client = _make_app_with_role(LitellmUserRoles.INTERNAL_USER)
|
|
self.mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
|
|
|
|
def test_should_allow_internal_user_to_list_agents(self, monkeypatch):
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=[])
|
|
resp = self.internal_client.get(
|
|
"/v1/agents", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_should_allow_internal_user_to_get_agent_by_id(self, monkeypatch):
|
|
self.mock_registry.get_agent_by_id = MagicMock(
|
|
return_value=_sample_agent_response()
|
|
)
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=None
|
|
)
|
|
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
|
return_value=[]
|
|
)
|
|
resp = self.internal_client.get(
|
|
"/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_should_block_internal_user_from_creating_agent(self):
|
|
resp = self.internal_client.post(
|
|
"/v1/agents",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 403
|
|
assert "Only proxy admins" in resp.json()["detail"]["error"]
|
|
|
|
def test_should_block_internal_user_from_updating_agent(self):
|
|
resp = self.internal_client.put(
|
|
"/v1/agents/agent-123",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
def test_should_block_internal_user_from_patching_agent(self):
|
|
resp = self.internal_client.patch(
|
|
"/v1/agents/agent-123",
|
|
json={"agent_name": "new-name"},
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
def test_should_block_internal_user_from_deleting_agent(self):
|
|
resp = self.internal_client.delete(
|
|
"/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
|
|
class TestAgentRBACInternalUserViewOnly:
|
|
"""View-only internal users should only be able to read agents."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
self.viewer_client = _make_app_with_role(
|
|
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
|
|
)
|
|
self.mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
|
|
|
|
def test_should_allow_view_only_user_to_list_agents(self):
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=[])
|
|
resp = self.viewer_client.get(
|
|
"/v1/agents", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_should_block_view_only_user_from_creating_agent(self):
|
|
resp = self.viewer_client.post(
|
|
"/v1/agents",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
def test_should_block_view_only_user_from_deleting_agent(self):
|
|
resp = self.viewer_client.delete(
|
|
"/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
|
|
class TestAgentRBACProxyAdminViewOnly:
|
|
"""Read-only proxy admins go through the object-permission scoped branch on
|
|
GET /v1/agents (the admin fast path stays full PROXY_ADMIN only, so viewers
|
|
cannot fan out health checks beyond their allowlist), and secret unredaction
|
|
also stays gated on full PROXY_ADMIN."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
from litellm.proxy.agent_endpoints import agent_registry as ar_mod
|
|
|
|
self.viewer_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
|
self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
|
|
self.agents = [
|
|
AgentResponse(
|
|
agent_id=f"agent-{index}",
|
|
agent_name=f"Agent {index}",
|
|
agent_card_params=_sample_agent_card_params(),
|
|
litellm_params={"api_key": "sk-super-secret-agent-key"},
|
|
)
|
|
for index in (1, 2)
|
|
]
|
|
self.mock_registry = MagicMock()
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=self.agents)
|
|
monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry)
|
|
|
|
self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"])
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
|
|
self.allowed_agents_spy,
|
|
)
|
|
|
|
def _list_agents(self, test_client: TestClient):
|
|
key_row = MagicMock()
|
|
key_row.token = "hash-aaa"
|
|
key_row.agent_id = "agent-1"
|
|
key_row.key_alias = "primary"
|
|
key_row.key_name = "sk-...aaa"
|
|
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
|
mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
|
|
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
|
return_value=[key_row]
|
|
)
|
|
return test_client.get("/v1/agents", headers={"Authorization": "Bearer k"})
|
|
|
|
def test_should_scope_view_only_admin_to_allowed_agents(self):
|
|
"""The key/team allowlist here excludes every registered agent; a viewer
|
|
on the admin fast path would see everything, so an empty response pins
|
|
that viewers stay in the scoped branch."""
|
|
resp = self._list_agents(self.viewer_client)
|
|
|
|
assert resp.status_code == 200
|
|
assert resp.json() == []
|
|
self.allowed_agents_spy.assert_awaited_once()
|
|
|
|
def test_should_still_redact_secrets_for_view_only_admin(self):
|
|
"""An unrestricted viewer (empty allowlist means no restrictions) sees the
|
|
same agents as an admin but with keys stripped and litellm_params masked."""
|
|
self.allowed_agents_spy.return_value = []
|
|
viewer_resp = self._list_agents(self.viewer_client)
|
|
admin_resp = self._list_agents(self.admin_client)
|
|
|
|
assert viewer_resp.status_code == 200
|
|
viewer_by_id = {agent["agent_id"]: agent for agent in viewer_resp.json()}
|
|
assert set(viewer_by_id) == {"agent-1", "agent-2"}
|
|
assert viewer_by_id["agent-1"]["keys"] is None
|
|
assert "sk-super-secret-agent-key" not in viewer_resp.text
|
|
|
|
admin_by_id = {agent["agent_id"]: agent for agent in admin_resp.json()}
|
|
assert admin_by_id["agent-1"]["keys"][0]["token"] == "hash-aaa"
|
|
assert (
|
|
admin_by_id["agent-1"]["litellm_params"]["api_key"]
|
|
== "sk-super-secret-agent-key"
|
|
)
|
|
|
|
|
|
class TestAgentRBACProxyAdmin:
|
|
"""Proxy admins should have full CRUD access to agents."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
|
|
self.mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
|
|
|
|
def test_should_allow_admin_to_create_agent(self, monkeypatch):
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
|
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
|
|
self.mock_registry.add_agent_to_db = AsyncMock(
|
|
return_value=_sample_agent_response()
|
|
)
|
|
self.mock_registry.register_agent = MagicMock()
|
|
resp = self.admin_client.post(
|
|
"/v1/agents",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_create_agent_applies_litellm_merge_to_stored_card(self):
|
|
"""The card stored in the DB must reflect the LiteLLM-fronting merge."""
|
|
with patch("litellm.proxy.proxy_server.prisma_client"):
|
|
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
|
|
self.mock_registry.add_agent_to_db = AsyncMock(
|
|
return_value=_sample_agent_response()
|
|
)
|
|
self.mock_registry.register_agent = MagicMock()
|
|
|
|
self.admin_client.post(
|
|
"/v1/agents",
|
|
json=_sample_agent_config(),
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
|
|
call_kwargs = self.mock_registry.add_agent_to_db.await_args.kwargs
|
|
stored_card = call_kwargs["agent"]["agent_card_params"]
|
|
new_agent_id = call_kwargs["agent_id"]
|
|
|
|
# Top-level url is retained for runtime A2A invocation (the public
|
|
# well-known endpoint rewrites it before exposing to clients);
|
|
# supportedInterfaces points at the proxy.
|
|
assert stored_card["url"] == "http://localhost"
|
|
assert stored_card["supportedInterfaces"][0]["protocolBinding"] == "JSONRPC"
|
|
assert stored_card["supportedInterfaces"][0]["url"].endswith(
|
|
f"/a2a/{new_agent_id}"
|
|
)
|
|
# Security scheme is the LiteLLM scheme.
|
|
assert "LiteLLMKey" in stored_card["securitySchemes"]
|
|
|
|
def test_should_allow_admin_to_delete_agent(self):
|
|
existing = {
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Existing Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value=existing
|
|
)
|
|
self.mock_registry.delete_agent_from_db = AsyncMock()
|
|
self.mock_registry.deregister_agent = MagicMock()
|
|
resp = self.admin_client.delete(
|
|
"/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
|
|
class TestAgentProtocolVersionValidation:
|
|
"""Registration accepts spec-default semver protocolVersion values and still
|
|
rejects genuinely unsupported versions."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
|
|
self.mock_registry = MagicMock()
|
|
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)
|
|
|
|
def _create_agent_with_protocol_version(self, protocol_version: str):
|
|
config = _sample_agent_config()
|
|
config["agent_card_params"]["protocolVersion"] = protocol_version
|
|
with patch("litellm.proxy.proxy_server.prisma_client"):
|
|
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
|
|
self.mock_registry.add_agent_to_db = AsyncMock(
|
|
return_value=_sample_agent_response()
|
|
)
|
|
self.mock_registry.register_agent = MagicMock()
|
|
return self.admin_client.post(
|
|
"/v1/agents",
|
|
json=config,
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
|
|
def test_semver_protocol_version_registers_and_stores_major_minor(self):
|
|
resp = self._create_agent_with_protocol_version("0.3.0")
|
|
assert resp.status_code == 200
|
|
stored_card = self.mock_registry.add_agent_to_db.await_args.kwargs["agent"][
|
|
"agent_card_params"
|
|
]
|
|
assert stored_card["protocolVersion"] == "0.3"
|
|
assert stored_card["supportedInterfaces"][0]["protocolVersion"] == "0.3"
|
|
|
|
def test_unsupported_protocol_version_is_rejected(self):
|
|
resp = self._create_agent_with_protocol_version("0.2.6")
|
|
assert resp.status_code == 400
|
|
assert "Unsupported protocolVersion '0.2.6'" in resp.json()["detail"]
|
|
self.mock_registry.add_agent_to_db.assert_not_awaited()
|
|
|
|
def test_malformed_protocol_version_is_rejected(self):
|
|
resp = self._create_agent_with_protocol_version("0.3.garbage")
|
|
assert resp.status_code == 400
|
|
assert "Unsupported protocolVersion '0.3.garbage'" in resp.json()["detail"]
|
|
self.mock_registry.add_agent_to_db.assert_not_awaited()
|
|
|
|
|
|
class TestCheckAgentManagementPermission:
|
|
"""Unit tests for the _check_agent_management_permission helper."""
|
|
|
|
def test_should_allow_proxy_admin(self):
|
|
auth = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
|
_check_agent_management_permission(auth)
|
|
|
|
@pytest.mark.parametrize(
|
|
"role",
|
|
[
|
|
LitellmUserRoles.INTERNAL_USER,
|
|
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
|
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
|
],
|
|
)
|
|
def test_should_block_non_admin_roles(self, role):
|
|
from fastapi import HTTPException
|
|
|
|
auth = UserAPIKeyAuth(user_id="user", user_role=role)
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
_check_agent_management_permission(auth)
|
|
assert exc_info.value.status_code == 403
|
|
|
|
|
|
class TestAgentRoutesIncludesAgentIdPattern:
|
|
"""Verify that agent_routes includes the {agent_id} pattern for route access."""
|
|
|
|
def test_should_include_agent_id_pattern(self):
|
|
from litellm.proxy._types import LiteLLMRoutes
|
|
|
|
assert "/v1/agents/{agent_id}" in LiteLLMRoutes.agent_routes.value
|
|
|
|
|
|
class TestAgentHealthCheck:
|
|
"""Tests for the health_check query parameter on GET /v1/agents."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup(self, monkeypatch):
|
|
from litellm.proxy.agent_endpoints import agent_registry as ar_mod
|
|
|
|
self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
|
|
self.mock_registry = MagicMock()
|
|
monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry)
|
|
# Ensure prisma_client is None so the endpoint skips DB queries.
|
|
# In CI with parallel workers, a MagicMock can leak from other test
|
|
# scopes, causing "object MagicMock can't be used in 'await'" errors.
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
|
|
|
def _make_agent(self, agent_id: str, url: str | None = None) -> AgentResponse:
|
|
card = _sample_agent_card_params()
|
|
if url is not None:
|
|
card["url"] = url
|
|
else:
|
|
card.pop("url", None)
|
|
return AgentResponse(
|
|
agent_id=agent_id,
|
|
agent_name=f"Agent {agent_id}",
|
|
agent_card_params=card,
|
|
litellm_params={},
|
|
)
|
|
|
|
def test_should_return_all_agents_when_health_check_disabled(self):
|
|
agents = [
|
|
self._make_agent("a1", "http://reachable"),
|
|
self._make_agent("a2", "http://unreachable"),
|
|
]
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=agents)
|
|
|
|
resp = self.admin_client.get(
|
|
"/v1/agents", headers={"Authorization": "Bearer k"}
|
|
)
|
|
assert resp.status_code == 200
|
|
assert len(resp.json()) == 2
|
|
|
|
def test_should_filter_unhealthy_agents_when_health_check_enabled(
|
|
self, monkeypatch
|
|
):
|
|
agents = [
|
|
self._make_agent("a1", "http://reachable"),
|
|
self._make_agent("a2", "http://unreachable"),
|
|
]
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=agents)
|
|
|
|
results = iter(
|
|
[
|
|
{"agent_id": "a1", "healthy": True},
|
|
{"agent_id": "a2", "healthy": False, "error": "Connection refused"},
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
agent_endpoints,
|
|
"_check_agent_url_health",
|
|
AsyncMock(side_effect=lambda agent: next(results)),
|
|
)
|
|
|
|
resp = self.admin_client.get(
|
|
"/v1/agents?health_check=true",
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert len(data) == 1
|
|
assert data[0]["agent_id"] == "a1"
|
|
|
|
def test_should_return_empty_list_when_all_agents_unhealthy(self, monkeypatch):
|
|
agents = [self._make_agent("a1", "http://down")]
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=agents)
|
|
monkeypatch.setattr(
|
|
agent_endpoints,
|
|
"_check_agent_url_health",
|
|
AsyncMock(
|
|
return_value={"agent_id": "a1", "healthy": False, "error": "timeout"}
|
|
),
|
|
)
|
|
|
|
resp = self.admin_client.get(
|
|
"/v1/agents?health_check=true",
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert len(resp.json()) == 0
|
|
|
|
def test_should_return_all_agents_when_all_healthy(self, monkeypatch):
|
|
agents = [
|
|
self._make_agent("a1", "http://ok1"),
|
|
self._make_agent("a2", "http://ok2"),
|
|
]
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=agents)
|
|
|
|
results = iter(
|
|
[
|
|
{"agent_id": "a1", "healthy": True},
|
|
{"agent_id": "a2", "healthy": True},
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
agent_endpoints,
|
|
"_check_agent_url_health",
|
|
AsyncMock(side_effect=lambda agent: next(results)),
|
|
)
|
|
|
|
resp = self.admin_client.get(
|
|
"/v1/agents?health_check=true",
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert len(resp.json()) == 2
|
|
|
|
|
|
class TestCheckAgentUrlHealth:
|
|
"""Unit tests for the _check_agent_url_health helper."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_return_healthy_when_no_url(self):
|
|
from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health
|
|
|
|
agent = AgentResponse(
|
|
agent_id="no-url",
|
|
agent_name="No URL Agent",
|
|
agent_card_params={"name": "test"},
|
|
litellm_params={},
|
|
)
|
|
result = await _check_agent_url_health(agent)
|
|
assert result["healthy"] is True
|
|
assert "error" not in result
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
|
|
async def test_should_return_healthy_for_200(self, mock_get_client):
|
|
from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_client = AsyncMock()
|
|
mock_client.get = AsyncMock(return_value=mock_response)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
agent = AgentResponse(
|
|
agent_id="ok",
|
|
agent_name="OK Agent",
|
|
agent_card_params={"url": "http://example.com"},
|
|
litellm_params={},
|
|
)
|
|
result = await _check_agent_url_health(agent)
|
|
assert result["healthy"] is True
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
|
|
async def test_should_return_unhealthy_for_500(self, mock_get_client):
|
|
from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 500
|
|
mock_client = AsyncMock()
|
|
mock_client.get = AsyncMock(return_value=mock_response)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
agent = AgentResponse(
|
|
agent_id="err",
|
|
agent_name="Error Agent",
|
|
agent_card_params={"url": "http://failing.com"},
|
|
litellm_params={},
|
|
)
|
|
result = await _check_agent_url_health(agent)
|
|
assert result["healthy"] is False
|
|
assert "HTTP 500" in result["error"]
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
|
|
async def test_should_return_unhealthy_on_connection_error(self, mock_get_client):
|
|
from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.get = AsyncMock(side_effect=Exception("Connection refused"))
|
|
mock_get_client.return_value = mock_client
|
|
|
|
agent = AgentResponse(
|
|
agent_id="down",
|
|
agent_name="Down Agent",
|
|
agent_card_params={"url": "http://down.com"},
|
|
litellm_params={},
|
|
)
|
|
result = await _check_agent_url_health(agent)
|
|
assert result["healthy"] is False
|
|
assert "Connection refused" in result["error"]
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
|
|
async def test_should_treat_404_as_healthy(self, mock_get_client):
|
|
"""A 404 means the server is reachable, just not the specific path."""
|
|
from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 404
|
|
mock_client = AsyncMock()
|
|
mock_client.get = AsyncMock(return_value=mock_response)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
agent = AgentResponse(
|
|
agent_id="notfound",
|
|
agent_name="NotFound Agent",
|
|
agent_card_params={"url": "http://example.com/missing"},
|
|
litellm_params={},
|
|
)
|
|
result = await _check_agent_url_health(agent)
|
|
assert result["healthy"] is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"base_url",
|
|
["http://0.0.0.0:4000/", "http://localhost:4000/", "https://api.example.com/"],
|
|
)
|
|
def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url(
|
|
monkeypatch, base_url
|
|
):
|
|
"""Without PROXY_BASE_URL, request.base_url carries a trailing slash; the merged
|
|
card's supportedInterfaces URL must still join cleanly (no `//a2a`)."""
|
|
from litellm.proxy.agent_endpoints.endpoints import _build_merged_agent_card
|
|
|
|
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
|
|
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
|
|
|
|
http_request = MagicMock()
|
|
http_request.base_url = base_url
|
|
|
|
merged = _build_merged_agent_card(
|
|
_sample_agent_card_params(),
|
|
agent_id="agent-xyz",
|
|
http_request=http_request,
|
|
agent_name="Test Agent",
|
|
)
|
|
|
|
interface_url = merged["supportedInterfaces"][0]["url"]
|
|
assert interface_url == f"{base_url.rstrip('/')}/a2a/agent-xyz"
|
|
assert "//a2a" not in interface_url
|