mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
8e65265bb4
* fix(agents): redact secret litellm_params fields from all /v1/agents responses Secret-bearing litellm_params fields (aws_secret_access_key, api_key, and similar) are now write-only: list, get, create, update, and patch responses always replace them with a fixed marker, regardless of caller role. Editing an agent no longer requires resending a real credential -- an update that omits a sensitive field, or echoes the marker back, preserves the stored value; a real value still rotates it. * fix(agents): redact secrets nested inside dicts/lists in litellm_params too Greptile found that a secret nested one level down under a non-sensitively-named key, or inside a list of per-provider configs, was neither redacted on read nor restored symmetrically on write (the marker string could get persisted as the real value). Recurse into lists on the read side, and mirror that recursion on the write side so restoration isn't limited to top-level keys. Also fixes a regression the redact rewrite introduced (a plain string leaf like a model name was being misinterpreted as a JSON blob and redacted), and suppresses 3 new test-quality-gate findings on an established repo-wide mocking pattern this PR's new tests also use. * fix(agents): guard list-position credential restore against misassignment Two more real gaps Greptile/veria found in the recursive redact/restore mechanism, verified directly against the exact reported shape (litellm_params.model_list, each entry carrying its own nested litellm_params.api_key/aws_secret_access_key) before fixing: - Positional restoration inside a list could attach one entry's stored credential to a different entry if the list were reordered or resized between GET and PUT/PATCH. Restoration by index now only fires when the incoming and existing entries match on every non-secret field; otherwise the caller's own value is used (never a guessed cross-entry secret). - A subtree collapsed to the flat REDACTED_BY_LITELM marker by the read-side recursion depth cap couldn't be recovered on write (the marker string itself would get persisted). Restore now recognizes that shape and recovers the whole existing subtree. Both covered by regression tests mirroring the exact model_list shape reported, mutation-verified. * fix(agents): simplify list-entry credential restore to positional matching The content-match guard from the previous commit fixed one Greptile finding (cross-entry misassignment on reorder) but introduced a worse one: it also rejected restoration whenever an entry's own non-secret fields changed, which is the common case (rename a model_list entry while leaving its own secret masked) -- silently dropping the stored credential on an ordinary edit. There is no stable per-element identity in a plain dict[str, object] schema, so no rule can satisfy both 'restore whenever the entry itself only had its secret masked' and 'never restore across a reorder' at once. Positional correspondence is what every other part of this restore (and the endpoints' full-replace-on-PUT semantics) already assumes, so drop the content-match gate and rely on it here too: this fixes the common case correctly and accepts cross-entry misassignment on a simultaneous reorder-plus-masked-echo as a known, narrow, documented limitation (not a leak between different agents or tenants, since it only reshuffles one agent's own stored values). Tests updated to pin the accepted trade-off explicitly rather than asserting it away, and to cover the previously broken ordinary-edit case.
1064 lines
41 KiB
Python
1064 lines
41 KiB
Python
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.constants import REDACTED_BY_LITELM_STRING
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
|
|
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
|
RestrictedAgentAccess,
|
|
UnrestrictedAgentAccess,
|
|
)
|
|
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
|
|
|
|
|
|
SENTINEL_AGENT_API_KEY = "sk-test-sentinel-do-not-use"
|
|
|
|
|
|
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). litellm_params
|
|
secrets are redacted for every caller, admin included (LIT-6736); only the
|
|
virtual-key/header visibility 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": SENTINEL_AGENT_API_KEY},
|
|
)
|
|
for index in (1, 2)
|
|
]
|
|
self.mock_registry = MagicMock()
|
|
self.mock_registry.get_agent_list = MagicMock(return_value=self.agents)
|
|
self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id}))
|
|
monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry)
|
|
|
|
self.allowed_agents_spy = AsyncMock(
|
|
return_value=RestrictedAgentAccess(frozenset({"someone-elses-agent"}))
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
|
|
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 sees the same agents as an admin but with keys
|
|
stripped; litellm_params secrets never appear in either response."""
|
|
self.allowed_agents_spy.return_value = UnrestrictedAgentAccess()
|
|
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 SENTINEL_AGENT_API_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 SENTINEL_AGENT_API_KEY not in admin_resp.text
|
|
assert admin_by_id["agent-1"]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING
|
|
|
|
|
|
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_create_agent_response_never_echoes_secret(self):
|
|
"""LIT-6736: POST /v1/agents must not echo the stored secret back, even
|
|
though it's the caller's own value and even for a proxy admin."""
|
|
with patch("litellm.proxy.proxy_server.prisma_client"): # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
|
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
|
|
self.mock_registry.add_agent_to_db = AsyncMock(
|
|
return_value=AgentResponse(
|
|
agent_id="agent-123",
|
|
agent_name="Test Agent",
|
|
agent_card_params=_sample_agent_card_params(),
|
|
litellm_params={
|
|
"aws_secret_access_key": SENTINEL_AGENT_API_KEY,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
},
|
|
)
|
|
)
|
|
self.mock_registry.register_agent = MagicMock()
|
|
|
|
resp = self.admin_client.post(
|
|
"/v1/agents",
|
|
json={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {
|
|
"aws_secret_access_key": SENTINEL_AGENT_API_KEY,
|
|
"model": "bedrock/agentcore/my-agent",
|
|
},
|
|
},
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
|
|
assert resp.status_code == 200
|
|
assert SENTINEL_AGENT_API_KEY not in resp.text
|
|
body = resp.json()
|
|
assert body["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
assert body["litellm_params"]["model"] == "bedrock/agentcore/my-agent"
|
|
|
|
def test_update_agent_response_never_echoes_secret(self):
|
|
"""LIT-6736: PUT /v1/agents/{id} must not echo the stored secret back."""
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value={
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Existing Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
)
|
|
self.mock_registry.update_agent_in_db = AsyncMock(
|
|
return_value=AgentResponse(
|
|
agent_id="agent-123",
|
|
agent_name="Test Agent",
|
|
agent_card_params=_sample_agent_card_params(),
|
|
litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY},
|
|
)
|
|
)
|
|
self.mock_registry.deregister_agent = MagicMock()
|
|
self.mock_registry.register_agent = MagicMock()
|
|
|
|
resp = self.admin_client.put(
|
|
"/v1/agents/agent-123",
|
|
json={
|
|
"agent_name": "Test Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
"litellm_params": {"aws_secret_access_key": REDACTED_BY_LITELM_STRING},
|
|
},
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
|
|
assert resp.status_code == 200
|
|
assert SENTINEL_AGENT_API_KEY not in resp.text
|
|
assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
|
|
def test_patch_agent_response_never_echoes_secret(self):
|
|
"""LIT-6736: PATCH /v1/agents/{id} must not echo the stored secret back."""
|
|
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
|
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
|
return_value={
|
|
"agent_id": "agent-123",
|
|
"agent_name": "Existing Agent",
|
|
"agent_card_params": _sample_agent_card_params(),
|
|
}
|
|
)
|
|
self.mock_registry.patch_agent_in_db = AsyncMock(
|
|
return_value=AgentResponse(
|
|
agent_id="agent-123",
|
|
agent_name="Renamed Agent",
|
|
agent_card_params=_sample_agent_card_params(),
|
|
litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY},
|
|
)
|
|
)
|
|
self.mock_registry.deregister_agent = MagicMock()
|
|
self.mock_registry.register_agent = MagicMock()
|
|
|
|
resp = self.admin_client.patch(
|
|
"/v1/agents/agent-123",
|
|
json={"agent_name": "Renamed Agent"},
|
|
headers={"Authorization": "Bearer k"},
|
|
)
|
|
|
|
assert resp.status_code == 200
|
|
assert SENTINEL_AGENT_API_KEY not in resp.text
|
|
assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
|
|
|
|
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
|