mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
262 lines
9.3 KiB
Python
262 lines
9.3 KiB
Python
"""Unit tests for the pure merge logic in litellm/proxy/a2a/agent_card.py."""
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.a2a.agent_card import (
|
|
LITELLM_A2A_PROTOCOL_VERSION,
|
|
LITELLM_SECURITY_REQUIREMENTS,
|
|
LITELLM_SECURITY_SCHEMES,
|
|
merge_agent_card,
|
|
normalize_protocol_version,
|
|
resolve_served_protocol_version,
|
|
)
|
|
|
|
PROXY_URL = "https://proxy.example/a2a/agent-xyz"
|
|
PROXY_BASE = "https://proxy.example"
|
|
|
|
|
|
def _full_upstream_card() -> dict:
|
|
return {
|
|
"protocolVersion": "0.9",
|
|
"name": "Upstream Name",
|
|
"description": "Upstream description",
|
|
"url": "http://internal:9999/",
|
|
"version": "1.2.3",
|
|
"capabilities": {
|
|
"streaming": True,
|
|
"pushNotifications": True,
|
|
"stateTransitionHistory": True,
|
|
"extensions": [{"uri": "x"}],
|
|
},
|
|
"skills": [
|
|
{"id": "s1", "name": "skill one", "description": "d", "tags": ["t"]}
|
|
],
|
|
"defaultInputModes": ["text", "audio"],
|
|
"defaultOutputModes": ["text"],
|
|
"securitySchemes": {"upstreamKey": {"type": "apiKey"}},
|
|
"security": [{"upstreamKey": []}],
|
|
"provider": {"organization": "UpstreamCo", "url": "https://upstream.example"},
|
|
"iconUrl": "https://upstream.example/icon.png",
|
|
"documentationUrl": "https://upstream.example/docs",
|
|
"somethingNotInSchema": "should be stripped",
|
|
}
|
|
|
|
|
|
def test_preserves_top_level_url_for_runtime_invocation():
|
|
# The runtime A2A invocation path reads ``agent_card_params['url']`` to
|
|
# know where to proxy requests, so the merge must keep the upstream URL
|
|
# on the stored card. The public well-known endpoint rewrites this field
|
|
# to the proxy URL before exposing it to clients.
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["url"] == "http://internal:9999/"
|
|
|
|
|
|
def test_unsupported_protocol_version_defaults_to_1_0():
|
|
# The fixture card pins "0.9", which LiteLLM does not serve; it falls back to
|
|
# the default rather than advertising a version the proxy can't honor.
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["protocolVersion"] == LITELLM_A2A_PROTOCOL_VERSION
|
|
|
|
|
|
def test_serves_pinned_protocol_version():
|
|
for version in ("0.3", "1.0"):
|
|
card = _full_upstream_card()
|
|
card["protocolVersion"] = version
|
|
merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["protocolVersion"] == version
|
|
assert merged["supportedInterfaces"][0]["protocolVersion"] == version
|
|
|
|
|
|
def test_absent_protocol_version_defaults_to_1_0():
|
|
card = _full_upstream_card()
|
|
card.pop("protocolVersion", None)
|
|
merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["protocolVersion"] == "1.0"
|
|
|
|
|
|
def test_overrides_name_and_description_when_provided():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(),
|
|
proxy_url=PROXY_URL,
|
|
proxy_base_url=PROXY_BASE,
|
|
name="UI Name",
|
|
description="UI Description",
|
|
)
|
|
assert merged["name"] == "UI Name"
|
|
assert merged["description"] == "UI Description"
|
|
|
|
|
|
def test_keeps_upstream_name_and_description_when_not_overridden():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["name"] == "Upstream Name"
|
|
assert merged["description"] == "Upstream description"
|
|
|
|
|
|
def test_filters_capabilities_to_allowlist():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
# Only ``streaming`` is allowlisted today.
|
|
assert merged["capabilities"] == {"streaming": True}
|
|
|
|
|
|
def test_drops_streaming_when_upstream_disables_it():
|
|
upstream = _full_upstream_card()
|
|
upstream["capabilities"]["streaming"] = False
|
|
merged = merge_agent_card(upstream, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["capabilities"] == {}
|
|
|
|
|
|
def test_replaces_security_schemes_and_requirements():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["securitySchemes"] == LITELLM_SECURITY_SCHEMES
|
|
assert merged["security"] == LITELLM_SECURITY_REQUIREMENTS
|
|
assert "securityRequirements" not in merged
|
|
|
|
|
|
def test_emits_supported_interfaces_pointing_at_proxy():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["supportedInterfaces"] == [
|
|
{
|
|
"url": PROXY_URL,
|
|
"protocolBinding": "JSONRPC",
|
|
"protocolVersion": LITELLM_A2A_PROTOCOL_VERSION,
|
|
}
|
|
]
|
|
|
|
|
|
def test_passes_through_skills_modes_provider_icon_docs():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["skills"] == _full_upstream_card()["skills"]
|
|
assert merged["defaultInputModes"] == ["text", "audio"]
|
|
assert merged["defaultOutputModes"] == ["text"]
|
|
assert merged["provider"] == {
|
|
"organization": "UpstreamCo",
|
|
"url": "https://upstream.example",
|
|
}
|
|
assert merged["iconUrl"] == "https://upstream.example/icon.png"
|
|
assert merged["documentationUrl"] == "https://upstream.example/docs"
|
|
|
|
|
|
def test_strips_fields_not_in_v1_schema():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert "somethingNotInSchema" not in merged
|
|
|
|
|
|
def test_defaults_for_missing_skills_and_modes():
|
|
sparse = {"name": "x", "description": "y", "version": "1"}
|
|
merged = merge_agent_card(sparse, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["skills"] and merged["skills"][0]["id"] == "chat"
|
|
assert merged["defaultInputModes"] == ["text"]
|
|
assert merged["defaultOutputModes"] == ["text"]
|
|
|
|
|
|
def test_defaults_version_when_upstream_omits_it():
|
|
sparse = {"name": "x", "description": "y"}
|
|
merged = merge_agent_card(sparse, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["version"] == "1.0.0"
|
|
|
|
|
|
def test_preserves_upstream_version_when_present():
|
|
merged = merge_agent_card(
|
|
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
|
|
)
|
|
assert merged["version"] == "1.2.3"
|
|
|
|
|
|
def test_falls_back_to_litellm_provider_when_upstream_lacks_one():
|
|
sparse = {"name": "x", "description": "y", "version": "1"}
|
|
merged = merge_agent_card(sparse, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["provider"] == {
|
|
"organization": "LiteLLM Proxy",
|
|
"url": PROXY_BASE,
|
|
}
|
|
|
|
|
|
def test_handles_none_upstream_card():
|
|
merged = merge_agent_card(None, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["protocolVersion"] == LITELLM_A2A_PROTOCOL_VERSION
|
|
assert merged["supportedInterfaces"][0]["url"] == PROXY_URL
|
|
assert merged["securitySchemes"] == LITELLM_SECURITY_SCHEMES
|
|
|
|
|
|
def test_does_not_mutate_input():
|
|
upstream = _full_upstream_card()
|
|
snapshot = dict(upstream)
|
|
merge_agent_card(upstream, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert upstream == snapshot
|
|
|
|
|
|
def test_strips_additional_interfaces_to_prevent_backend_url_leak():
|
|
upstream = _full_upstream_card()
|
|
upstream["additionalInterfaces"] = [
|
|
{"url": "http://internal-backend:8080/", "transport": "JSONRPC"},
|
|
{"url": "grpc://internal-backend:50051", "transport": "GRPC"},
|
|
]
|
|
merged = merge_agent_card(upstream, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert "additionalInterfaces" not in merged
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("raw", "expected"),
|
|
[
|
|
("0.3", "0.3"),
|
|
("0.3.0", "0.3"),
|
|
("1.0", "1.0"),
|
|
("1.0.0", "1.0"),
|
|
("1.0.1", "1.0"),
|
|
("0.3.0-rc1", "0.3"),
|
|
("1.0.0-rc.1+build.5", "1.0"),
|
|
("0.2.6", None),
|
|
("2.0", None),
|
|
("0.30", None),
|
|
("0.3.garbage", None),
|
|
("0.3.", None),
|
|
("1.0.not-semver", None),
|
|
("0.3.0.0", None),
|
|
("0.3-rc1", None),
|
|
("garbage", None),
|
|
("", None),
|
|
(None, None),
|
|
(1.0, None),
|
|
],
|
|
)
|
|
def test_normalize_protocol_version(raw, expected):
|
|
assert normalize_protocol_version(raw) == expected
|
|
|
|
|
|
def test_resolve_served_protocol_version_canonicalizes_semver_pins():
|
|
assert resolve_served_protocol_version({"protocolVersion": "0.3.0"}) == "0.3"
|
|
assert resolve_served_protocol_version({"protocolVersion": "1.0.0"}) == "1.0"
|
|
assert resolve_served_protocol_version({"protocolVersion": "0.3"}) == "0.3"
|
|
assert resolve_served_protocol_version({"protocolVersion": "1.0"}) == "1.0"
|
|
|
|
|
|
def test_resolve_served_protocol_version_falls_back_for_unsupported():
|
|
assert (
|
|
resolve_served_protocol_version({"protocolVersion": "0.2.6"})
|
|
== LITELLM_A2A_PROTOCOL_VERSION
|
|
)
|
|
assert resolve_served_protocol_version(None) == LITELLM_A2A_PROTOCOL_VERSION
|
|
|
|
|
|
def test_serves_semver_pinned_protocol_version_as_major_minor():
|
|
card = _full_upstream_card()
|
|
card["protocolVersion"] = "0.3.0"
|
|
merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE)
|
|
assert merged["protocolVersion"] == "0.3"
|
|
assert merged["supportedInterfaces"][0]["protocolVersion"] == "0.3"
|