mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(agents): show non-admin callers the proxy URL in agent cards (#49483)
* test(agents): cover backend card URL leak to non-admin callers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): assert backend URLs before card parity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): scope unit card regressions to one agent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): isolate card redaction from access policy loading Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): remove unused proxy fixture binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): normalize unit test imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): pass endpoint base URL to redaction helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): cover proxy card deep-copy behavior Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(agents): rewrite non-admin agent card URLs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(agents): simplify proxy card URL helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(agents): drop card mapping adapter import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(agents): use typed card mapping helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): cover alternate card interfaces and real access checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(agents): drop alternate backend interfaces from non-admin cards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(agents): probe the real health check in card URL test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
moe-berri
GitHub
parent
9312a25f18
commit
87e32e7c6d
6 changed files with 356 additions and 29 deletions
|
|
@ -12,6 +12,8 @@ from collections.abc import Mapping
|
|||
from copy import deepcopy
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
|
||||
SupportedA2AVersion = Literal["0.3", "1.0"]
|
||||
|
||||
# Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent;
|
||||
|
|
@ -21,7 +23,6 @@ SUPPORTED_A2A_PROTOCOL_VERSIONS: Final[tuple[SupportedA2AVersion, ...]] = ("0.3"
|
|||
# Default served version when the agent card does not pin one.
|
||||
LITELLM_A2A_PROTOCOL_VERSION: Final = "1.0"
|
||||
|
||||
|
||||
_PROTOCOL_VERSION_PATTERN: Final = re.compile(
|
||||
r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$"
|
||||
)
|
||||
|
|
@ -50,6 +51,24 @@ def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str:
|
|||
return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION
|
||||
|
||||
|
||||
def point_agent_card_at_proxy(
|
||||
agent_card: Mapping[str, object],
|
||||
proxy_url: str,
|
||||
) -> dict[str, object]: # mutable-ok: the public helper contract requires a dictionary
|
||||
interfaces: Final = agent_card.get("supportedInterfaces")
|
||||
interface_items: Final = (
|
||||
tuple(iter_items(agent_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
|
||||
)
|
||||
first: Final[object] = interface_items[0] if interface_items else None
|
||||
first_mapping: Final = as_str_mapping(first)
|
||||
pointed_interfaces: Final = (
|
||||
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}, *interface_items[1:]]}
|
||||
if isinstance(first, dict) and first_mapping is not None
|
||||
else {}
|
||||
)
|
||||
return deepcopy({**agent_card, "url": proxy_url, **pointed_interfaces})
|
||||
|
||||
|
||||
# Security scheme exposed by the LiteLLM-fronted agent card. Always replaces
|
||||
# whatever upstream advertised — the client must authenticate to the proxy,
|
||||
# not the upstream agent.
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
|
|||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
|
||||
from copy import deepcopy
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -29,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
|||
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
|
||||
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy
|
||||
from litellm.proxy.a2a.version_convert import (
|
||||
A2AVersion,
|
||||
normalize_agent_card,
|
||||
|
|
@ -750,16 +750,12 @@ async def get_agent_card(
|
|||
)
|
||||
|
||||
proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}")
|
||||
agent_card = deepcopy(agent.agent_card_params)
|
||||
agent_card["url"] = proxy_url
|
||||
interfaces: Final = agent_card.get("supportedInterfaces")
|
||||
if isinstance(interfaces, list) and interfaces:
|
||||
interfaces[0]["url"] = proxy_url
|
||||
agent_card: Final = point_agent_card_at_proxy(agent.agent_card_params, proxy_url)
|
||||
served_version: Final = _served_version(agent, request)
|
||||
agent_card = normalize_agent_card(agent_card, served_version)
|
||||
normalized_agent_card: Final = normalize_agent_card(agent_card, served_version)
|
||||
|
||||
verbose_proxy_logger.debug("Returning agent card for '%s' with proxy URL: %s", agent_id, proxy_url)
|
||||
return JSONResponse(content=agent_card)
|
||||
return JSONResponse(content=normalized_agent_card)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from typing_extensions import ReadOnly, Required, assert_never
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, AGENT_HEALTH_CHECK_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
|
|
@ -33,6 +34,7 @@ from litellm.proxy.a2a.agent_card import (
|
|||
SUPPORTED_A2A_PROTOCOL_VERSIONS,
|
||||
merge_agent_card,
|
||||
normalize_protocol_version,
|
||||
point_agent_card_at_proxy,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
AgentIdWhere,
|
||||
|
|
@ -180,10 +182,36 @@ def _redact_agent_litellm_params_dict(
|
|||
return dict(parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params)))
|
||||
|
||||
|
||||
def _non_admin_agent_card(
|
||||
card: Mapping[str, object], proxy_url: str
|
||||
) -> dict[str, object]: # mutable-ok: AgentResponse.agent_card_params requires a mutable dictionary
|
||||
pointed_card: Final = point_agent_card_at_proxy(card, proxy_url)
|
||||
interfaces: Final = pointed_card.get("supportedInterfaces")
|
||||
interface_items: Final = (
|
||||
tuple(iter_items(pointed_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
|
||||
)
|
||||
first_interface: Final[object] = interface_items[0] if interface_items else None
|
||||
first_mapping: Final = as_str_mapping(first_interface)
|
||||
proxy_interface: Final = (
|
||||
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}]}
|
||||
if isinstance(first_interface, dict) and first_mapping is not None and first_mapping.get("url") == proxy_url
|
||||
else {}
|
||||
)
|
||||
return {
|
||||
**{
|
||||
key: value
|
||||
for key, value in pointed_card.items()
|
||||
if key not in ("additionalInterfaces", "supportedInterfaces")
|
||||
},
|
||||
**proxy_interface,
|
||||
}
|
||||
|
||||
|
||||
def _redact_sensitive_agent_fields(
|
||||
agents: Sequence[AgentResponse],
|
||||
*,
|
||||
is_admin: bool,
|
||||
request_base_url: str,
|
||||
) -> list[AgentResponse]:
|
||||
"""
|
||||
Return copies of the given agents with credential-bearing litellm_params
|
||||
|
|
@ -207,6 +235,11 @@ def _redact_sensitive_agent_fields(
|
|||
copy.keys = None
|
||||
copy.kill_switch = None
|
||||
copy.identity = None
|
||||
if copy.agent_card_params:
|
||||
copy.agent_card_params = _non_admin_agent_card(
|
||||
copy.agent_card_params,
|
||||
get_custom_url(request_base_url, route=f"a2a/{copy.agent_id}"),
|
||||
)
|
||||
if copy.litellm_params:
|
||||
copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
|
||||
copy.kill_switch = redact_kill_switch(copy.kill_switch)
|
||||
|
|
@ -461,14 +494,6 @@ async def get_agents(
|
|||
global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
|
||||
)
|
||||
|
||||
# litellm_params secrets are always redacted; keys/headers stay
|
||||
# admin-only.
|
||||
is_admin: Final = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin)
|
||||
|
||||
if health_check:
|
||||
agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")]
|
||||
agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")]
|
||||
|
|
@ -493,6 +518,18 @@ async def get_agents(
|
|||
healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
|
||||
returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url
|
||||
|
||||
# litellm_params secrets are always redacted; keys/headers stay
|
||||
# admin-only.
|
||||
is_admin: Final = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
returned_agents = _redact_sensitive_agent_fields(
|
||||
returned_agents,
|
||||
is_admin=is_admin,
|
||||
request_base_url=str(request.base_url),
|
||||
)
|
||||
|
||||
if query is None:
|
||||
return returned_agents
|
||||
return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict)
|
||||
|
|
@ -691,7 +728,11 @@ async def create_agent(
|
|||
|
||||
# The caller is a proxy admin (enforced above); litellm_params
|
||||
# secrets are still never echoed back in the response.
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
|
|
@ -708,6 +749,7 @@ async def create_agent(
|
|||
)
|
||||
async def get_agent_by_id(
|
||||
agent_id: str,
|
||||
request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
|
|
@ -770,7 +812,11 @@ async def get_agent_by_id(
|
|||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0]
|
||||
agent = _redact_sensitive_agent_fields(
|
||||
(agent,),
|
||||
is_admin=is_admin,
|
||||
request_base_url=str(request.base_url),
|
||||
)[0]
|
||||
|
||||
return agent
|
||||
except HTTPException:
|
||||
|
|
@ -879,7 +925,11 @@ async def update_agent(
|
|||
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
|
||||
)
|
||||
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
@ -987,7 +1037,11 @@ async def patch_agent(
|
|||
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
|
||||
)
|
||||
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -66,6 +66,13 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
|||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"supportedInterfaces": [
|
||||
{
|
||||
"url": url + "/",
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "0.3",
|
||||
}
|
||||
],
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
|
|
@ -73,10 +80,15 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
|||
}
|
||||
|
||||
|
||||
def _register_agent(gateway: Gateway, scenario: Scenario, name: str, url: str) -> str:
|
||||
created: Final = gateway.request(
|
||||
"POST", "/v1/agents", {"agent_name": name, "agent_card_params": _peer_card(url, name)}
|
||||
)
|
||||
def _register_agent(
|
||||
gateway: Gateway,
|
||||
scenario: Scenario,
|
||||
name: str,
|
||||
url: str,
|
||||
agent_card: dict[str, JsonValue] | None = None,
|
||||
) -> str:
|
||||
card: Final = _peer_card(url, name) if agent_card is None else agent_card
|
||||
created: Final = gateway.request("POST", "/v1/agents", {"agent_name": name, "agent_card_params": card})
|
||||
assert created.status_code == 200, created.text
|
||||
identity: Final = created.json()["agent_id"]
|
||||
|
||||
|
|
@ -416,3 +428,68 @@ def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys
|
|||
)
|
||||
assert (denied_legacy.status_code, denied_legacy.text) == (403, refusal(denied)), denied_legacy.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_agent_list_and_detail_cards_use_proxy_urls_for_scoped_keys(gateway: Gateway) -> None:
|
||||
marker: Final = "a2acardurl" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
assert request.method == "GET", request.method
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
backend_url: Final = wire.url + "/"
|
||||
peer_card: Final = _peer_card(wire.url, marker)
|
||||
peer_interfaces: Final = peer_card["supportedInterfaces"]
|
||||
assert isinstance(peer_interfaces, list)
|
||||
agent_card: Final = {
|
||||
**peer_card,
|
||||
"additionalInterfaces": [{"url": backend_url, "transport": "GRPC"}],
|
||||
"supportedInterfaces": [
|
||||
*peer_interfaces,
|
||||
{
|
||||
"url": backend_url,
|
||||
"protocolBinding": "GRPC",
|
||||
"protocolVersion": "0.3",
|
||||
},
|
||||
],
|
||||
}
|
||||
agent_id: Final = _register_agent(gateway, scenario, marker, wire.url, agent_card)
|
||||
user_id: Final = scenario.user(user_role="internal_user")
|
||||
scoped_key: Final = scenario.key(user_id=user_id, object_permission={"agents": [agent_id]})
|
||||
|
||||
public_card: Final = gateway.request(
|
||||
"GET",
|
||||
f"/a2a/{agent_id}/.well-known/agent-card.json",
|
||||
key=scoped_key,
|
||||
)
|
||||
listed: Final = gateway.request("GET", "/v1/agents", key=scoped_key)
|
||||
detailed: Final = gateway.request("GET", f"/v1/agents/{agent_id}", key=scoped_key)
|
||||
|
||||
assert public_card.status_code == 200, public_card.text
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert detailed.status_code == 200, detailed.text
|
||||
assert wire.url not in listed.text
|
||||
assert wire.url not in detailed.text
|
||||
assert "additionalInterfaces" not in listed.text
|
||||
assert "additionalInterfaces" not in detailed.text
|
||||
public_urls: Final = (
|
||||
public_card.json()["url"],
|
||||
public_card.json()["supportedInterfaces"][0]["url"],
|
||||
)
|
||||
listed_agent: Final = next(item for item in listed.json() if item["agent_id"] == agent_id)
|
||||
for body in (listed_agent, detailed.json()):
|
||||
card: Final = body["agent_card_params"]
|
||||
assert (card["url"], card["supportedInterfaces"][0]["url"]) == public_urls
|
||||
assert card["supportedInterfaces"] == [
|
||||
{
|
||||
"url": public_urls[1],
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "0.3",
|
||||
}
|
||||
]
|
||||
|
||||
admin_listed: Final = gateway.request("GET", "/v1/agents")
|
||||
assert admin_listed.status_code == 200, admin_listed.text
|
||||
admin_agent: Final = next(item for item in admin_listed.json() if item["agent_id"] == agent_id)
|
||||
assert admin_agent["agent_card_params"]["url"] == backend_url
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.proxy.a2a.agent_card import (
|
|||
LITELLM_SECURITY_SCHEMES,
|
||||
merge_agent_card,
|
||||
normalize_protocol_version,
|
||||
point_agent_card_at_proxy,
|
||||
resolve_served_protocol_version,
|
||||
)
|
||||
|
||||
|
|
@ -15,6 +16,32 @@ PROXY_URL = "https://proxy.example/a2a/agent-xyz"
|
|||
PROXY_BASE = "https://proxy.example"
|
||||
|
||||
|
||||
def test_point_agent_card_at_proxy_returns_a_deep_copy() -> None:
|
||||
card: dict[str, object] = {
|
||||
"url": "http://internal/",
|
||||
"metadata": {"nested": ["original"]},
|
||||
"supportedInterfaces": [
|
||||
{"url": "http://internal/", "nested": {"value": "original"}},
|
||||
{"url": "http://secondary/"},
|
||||
],
|
||||
}
|
||||
|
||||
pointed: dict[str, object] = point_agent_card_at_proxy(card, PROXY_URL)
|
||||
|
||||
assert pointed["url"] == PROXY_URL
|
||||
assert pointed["metadata"] == card["metadata"]
|
||||
assert pointed["metadata"] is not card["metadata"]
|
||||
assert pointed["supportedInterfaces"] == [
|
||||
{"url": PROXY_URL, "nested": {"value": "original"}},
|
||||
{"url": "http://secondary/"},
|
||||
]
|
||||
assert card["url"] == "http://internal/"
|
||||
assert card["supportedInterfaces"] == [
|
||||
{"url": "http://internal/", "nested": {"value": "original"}},
|
||||
{"url": "http://secondary/"},
|
||||
]
|
||||
|
||||
|
||||
def _full_upstream_card() -> dict:
|
||||
return {
|
||||
"protocolVersion": "0.9",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
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, patch
|
||||
|
|
@ -13,7 +12,12 @@ from fastapi.testclient import TestClient
|
|||
from prisma.models import LiteLLM_AgentsTable
|
||||
|
||||
from litellm.constants import REDACTED_BY_LITELM_STRING
|
||||
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_AuditLogs,
|
||||
LitellmTableNames,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
RestrictedAgentAccess,
|
||||
|
|
@ -43,6 +47,29 @@ def _sample_agent_card_params() -> dict:
|
|||
}
|
||||
|
||||
|
||||
_AGENT_CARD_BACKEND_URL: Final = "http://agent.internal:9000/plain/"
|
||||
|
||||
|
||||
def _sample_agent_card_params_with_backend_url() -> dict[str, object]:
|
||||
return {
|
||||
**_sample_agent_card_params(),
|
||||
"url": _AGENT_CARD_BACKEND_URL,
|
||||
"supportedInterfaces": [
|
||||
{
|
||||
"url": _AGENT_CARD_BACKEND_URL,
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "1.0",
|
||||
},
|
||||
{
|
||||
"url": _AGENT_CARD_BACKEND_URL,
|
||||
"protocolBinding": "GRPC",
|
||||
"protocolVersion": "1.0",
|
||||
},
|
||||
],
|
||||
"additionalInterfaces": [{"url": _AGENT_CARD_BACKEND_URL, "transport": "GRPC"}],
|
||||
}
|
||||
|
||||
|
||||
def _sample_agent_config() -> dict:
|
||||
return {
|
||||
"agent_name": "Test Agent",
|
||||
|
|
@ -641,7 +668,7 @@ class TestAgentRBACProxyAdmin:
|
|||
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:
|
||||
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()
|
||||
|
|
@ -1015,6 +1042,129 @@ class TestAgentHealthCheck:
|
|||
assert len(resp.json()) == 2
|
||||
|
||||
|
||||
def _client_with_agent_card(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
role: LitellmUserRoles,
|
||||
agent: AgentResponse,
|
||||
) -> TestClient:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.agent_endpoints import agent_registry
|
||||
|
||||
registry: Final = MagicMock()
|
||||
registry.get_agent_by_id = MagicMock(return_value=agent)
|
||||
registry.get_agent_list = MagicMock(return_value=(agent,))
|
||||
registry.ids_for_agent = MagicMock(return_value=frozenset({agent.agent_id}))
|
||||
registry.stable_agent_id = MagicMock(side_effect=lambda agent_id: agent_id)
|
||||
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
|
||||
monkeypatch.setattr(agent_registry, "global_agent_registry", registry)
|
||||
prisma: Final = MagicMock()
|
||||
prisma.writer_db = prisma.db
|
||||
prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
|
||||
prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
test_app: Final = FastAPI()
|
||||
test_app.include_router(router)
|
||||
test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="test-user",
|
||||
user_role=role,
|
||||
object_permission={"object_permission_id": "agent-card-test", "agents": [agent.agent_id]},
|
||||
)
|
||||
return TestClient(test_app)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/v1/agents", "/v1/agents/agent-123"])
|
||||
def test_non_admin_agent_card_urls_point_at_proxy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
path: str,
|
||||
) -> None:
|
||||
agent: Final = AgentResponse(
|
||||
agent_id="agent-123",
|
||||
agent_name="Test Agent",
|
||||
agent_card_params=_sample_agent_card_params_with_backend_url(),
|
||||
)
|
||||
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
|
||||
|
||||
response: Final = test_client.get(path, headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
agents: Final = response.json()
|
||||
result: Final = agents[0] if path == "/v1/agents" else agents
|
||||
card: Final = result["agent_card_params"]
|
||||
proxy_url: Final = "http://testserver/a2a/agent-123"
|
||||
assert card["url"] == proxy_url
|
||||
assert "agent.internal" not in response.text
|
||||
assert "additionalInterfaces" not in card
|
||||
assert card["supportedInterfaces"] == [
|
||||
{
|
||||
"url": proxy_url,
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "1.0",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_proxy_admin_agent_listing_keeps_backend_card_urls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
agent: Final = AgentResponse(
|
||||
agent_id="agent-123",
|
||||
agent_name="Test Agent",
|
||||
agent_card_params=_sample_agent_card_params_with_backend_url(),
|
||||
)
|
||||
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.PROXY_ADMIN, agent)
|
||||
|
||||
response: Final = test_client.get("/v1/agents", headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
card: Final = response.json()[0]["agent_card_params"]
|
||||
assert card["url"] == _AGENT_CARD_BACKEND_URL
|
||||
assert card["supportedInterfaces"] == [
|
||||
{
|
||||
"url": _AGENT_CARD_BACKEND_URL,
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "1.0",
|
||||
},
|
||||
{
|
||||
"url": _AGENT_CARD_BACKEND_URL,
|
||||
"protocolBinding": "GRPC",
|
||||
"protocolVersion": "1.0",
|
||||
},
|
||||
]
|
||||
assert card["additionalInterfaces"] == [{"url": _AGENT_CARD_BACKEND_URL, "transport": "GRPC"}]
|
||||
|
||||
|
||||
def test_non_admin_health_check_probes_backend_and_returns_proxy_card(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
agent: Final = AgentResponse(
|
||||
agent_id="agent-123",
|
||||
agent_name="Test Agent",
|
||||
agent_card_params=_sample_agent_card_params_with_backend_url(),
|
||||
)
|
||||
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
|
||||
http_get: Final = AsyncMock(return_value=httpx.Response(200))
|
||||
http_client: Final = MagicMock(get=http_get)
|
||||
monkeypatch.setattr(
|
||||
agent_endpoints,
|
||||
"get_async_httpx_client",
|
||||
MagicMock(return_value=http_client),
|
||||
)
|
||||
|
||||
response: Final = test_client.get("/v1/agents?health_check=true", headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert http_get.await_args.args[0] == _AGENT_CARD_BACKEND_URL
|
||||
assert response.json()[0]["agent_card_params"]["url"] == "http://testserver/a2a/agent-123"
|
||||
|
||||
|
||||
def test_non_admin_agent_card_without_interfaces_does_not_add_them(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
agent: Final = _sample_agent_response()
|
||||
test_client: Final = _client_with_agent_card(monkeypatch, LitellmUserRoles.INTERNAL_USER, agent)
|
||||
|
||||
response: Final = test_client.get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["agent_card_params"]["url"] == "http://testserver/a2a/agent-123"
|
||||
assert "supportedInterfaces" not in response.json()["agent_card_params"]
|
||||
|
||||
|
||||
class TestCheckAgentUrlHealth:
|
||||
"""Unit tests for the _check_agent_url_health helper."""
|
||||
|
||||
|
|
@ -1230,7 +1380,11 @@ def test_jwt_authentication_status_does_not_require_virtual_keys(
|
|||
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
|
||||
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
|
||||
agent: Final = _sample_agent_response()
|
||||
response: Final = agent_endpoints._redact_sensitive_agent_fields((agent,), is_admin=True)[0]
|
||||
response: Final = agent_endpoints._redact_sensitive_agent_fields(
|
||||
(agent,),
|
||||
is_admin=True,
|
||||
request_base_url="http://testserver",
|
||||
)[0]
|
||||
assert response.jwt_auth_configured is expected
|
||||
assert agent.jwt_auth_configured is False
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue