mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
test(integration): cover agents_a2a endpoint contracts (#44650)
* test(integration): cover agents_a2a endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin agent daily spend to the logged request day and assert cleanups Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expect agent daily activity newest day first Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover A2A 1.0 task aliases, agent id headers, declared access groups and per-token spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert full A2A stream, header, alias, task and spend contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert forwarded a2a-version header on A2A sends Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): release A2A stream gate only after a full SSE frame Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin A2A trace id, stream headers, denied alias and task methods Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin A2A card protocol version Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the whole accepted push config response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1c0f722aae
commit
117684cce5
4 changed files with 1583 additions and 1 deletions
334
tests/integration/authorization/test_agent_access.py
Normal file
334
tests/integration/authorization/test_agent_access.py
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
|
||||
class _JsonRpcResult(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
jsonrpc: str
|
||||
id: str
|
||||
result: dict[str, JsonValue]
|
||||
|
||||
|
||||
_STREAM_EVENTS: Final = (
|
||||
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "working"}, "final": False},
|
||||
{
|
||||
"kind": "artifact-update",
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "chunk"}]},
|
||||
},
|
||||
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "completed"}, "final": True},
|
||||
)
|
||||
_STREAM_RESULTS: Final = (
|
||||
{"contextId": "c1", "final": False, "kind": "status-update", "status": {"state": "working"}, "taskId": "t1"},
|
||||
{
|
||||
"append": False,
|
||||
"artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "chunk"}]},
|
||||
"contextId": "c1",
|
||||
"kind": "artifact-update",
|
||||
"lastChunk": False,
|
||||
"taskId": "t1",
|
||||
},
|
||||
{"contextId": "c1", "final": True, "kind": "status-update", "status": {"state": "completed"}, "taskId": "t1"},
|
||||
)
|
||||
|
||||
|
||||
def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"protocolVersion": "0.3",
|
||||
"name": name,
|
||||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
|
||||
|
||||
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)}
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
identity: Final = created.json()["agent_id"]
|
||||
|
||||
def cleanup() -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/agents/{identity}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (identity,)) == []
|
||||
|
||||
scenario.cleanups.callback(cleanup)
|
||||
return identity
|
||||
|
||||
|
||||
def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_card_before_the_upstream(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = "a2aacl" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
if body["method"] == "message/stream":
|
||||
frames: Final = tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': body['id'], 'result': event})}\n\n".encode()
|
||||
for event in _STREAM_EVENTS
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": "granted"}],
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
allowed_name: Final = marker + "-allowed"
|
||||
denied_name: Final = marker + "-denied"
|
||||
allowed: Final = _register_agent(gateway, scenario, allowed_name, wire.url)
|
||||
denied: Final = _register_agent(gateway, scenario, denied_name, wire.url)
|
||||
group: Final = gateway.request(
|
||||
"POST", "/v1/access_group", {"access_group_name": marker + "-group", "access_agent_ids": [allowed]}
|
||||
)
|
||||
assert group.status_code == 201, group.text
|
||||
group_id: Final = group.json()["access_group_id"]
|
||||
|
||||
def delete_group() -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group_id}")
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
assert (
|
||||
read_rows(
|
||||
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group_id,)
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
scenario.cleanups.callback(delete_group)
|
||||
tag: Final = marker + "-tag"
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_AgentsTable" SET agent_access_groups = ARRAY[%s] WHERE agent_id = %s', (tag, allowed)
|
||||
)
|
||||
callers: Final = {
|
||||
"key grant": scenario.key(object_permission={"agents": [allowed]}),
|
||||
"team grant": scenario.key(team_id=scenario.team(object_permission={"agents": [allowed]})),
|
||||
"team grant narrowed by key": scenario.key(
|
||||
team_id=scenario.team(object_permission={"agents": [allowed, denied]}),
|
||||
object_permission={"agents": [allowed]},
|
||||
),
|
||||
"access group": scenario.key(access_group_ids=[group_id]),
|
||||
"declared agent access group": scenario.key(object_permission={"agent_access_groups": [tag]}),
|
||||
}
|
||||
message: Final = {
|
||||
"kind": "message",
|
||||
"role": "user",
|
||||
"messageId": marker + "-in",
|
||||
"parts": [{"kind": "text", "text": "ping"}],
|
||||
}
|
||||
send_payload: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "message/send",
|
||||
"params": {"message": message},
|
||||
}
|
||||
stream_payload: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "message/stream",
|
||||
"params": {"message": message},
|
||||
}
|
||||
|
||||
def refusal(identity: str) -> str:
|
||||
return json.dumps(
|
||||
{"detail": f"Agent '{identity}' is not allowed for your key/team. Contact proxy admin for access."},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
spellings: Final = (
|
||||
("send by id", lambda agent, name: f"/a2a/{agent}", send_payload, "message/send"),
|
||||
("send by name", lambda agent, name: f"/a2a/{name}", send_payload, "message/send"),
|
||||
("send v1 alias", lambda agent, name: f"/v1/a2a/{agent}/message/send", send_payload, "message/send"),
|
||||
("send alias", lambda agent, name: f"/a2a/{agent}/message/send", send_payload, "message/send"),
|
||||
("stream", lambda agent, name: f"/a2a/{agent}", stream_payload, "message/stream"),
|
||||
)
|
||||
denied_tasks: Final = (
|
||||
(
|
||||
"tasks/get",
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "tasks/get",
|
||||
"params": {"id": "t1"},
|
||||
},
|
||||
),
|
||||
(
|
||||
"GetTask",
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "GetTask",
|
||||
"params": {"id": "t1"},
|
||||
},
|
||||
),
|
||||
(
|
||||
"tasks/pushNotificationConfig/set",
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "tasks/pushNotificationConfig/set",
|
||||
"params": {
|
||||
"taskId": "t1",
|
||||
"pushNotificationConfig": {"url": "https://callback.example.com/hook"},
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
def attempt(key: str, agent: str, name: str) -> tuple[tuple[str, int, str], ...]:
|
||||
def call(
|
||||
spelling: str, path_of: Callable[[str, str], str], payload: dict[str, JsonValue]
|
||||
) -> tuple[str, int, str]:
|
||||
headers: Final = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
**({"Accept": "text/event-stream"} if payload is stream_payload else {}),
|
||||
}
|
||||
response = gateway.client.post(path_of(agent, name), headers=headers, json=payload)
|
||||
return spelling, response.status_code, response.read().decode()
|
||||
|
||||
return tuple(call(label, path_of, payload) for label, path_of, payload, _ in spellings)
|
||||
|
||||
def card(key: str, identity: str) -> tuple[int, str]:
|
||||
response = gateway.client.get(
|
||||
f"/a2a/{identity}/.well-known/agent-card.json", headers={"Authorization": f"Bearer {key}"}
|
||||
)
|
||||
return response.status_code, response.text
|
||||
|
||||
for label, key in callers.items():
|
||||
for spelling, status, text in attempt(key, denied, denied_name):
|
||||
refused: Final = refusal(denied_name if "by name" in spelling else denied)
|
||||
assert (status, text) == (403, refused), f"{label} {spelling}: {status} {text}"
|
||||
for spelling, payload in denied_tasks:
|
||||
response: Final = gateway.client.post(
|
||||
f"/a2a/{denied}", headers={"Authorization": f"Bearer {key}"}, json=payload
|
||||
)
|
||||
assert (response.status_code, response.text) == (403, refusal(denied)), (
|
||||
f"{label} {spelling}: {response.status_code} {response.text}"
|
||||
)
|
||||
assert card(key, denied) == (403, refusal(denied)), label
|
||||
assert wire.drain() == (), label
|
||||
for spelling, status, text in attempt(key, allowed, allowed_name):
|
||||
assert status == 200, f"{label} {spelling}: {status} {text}"
|
||||
if "stream" in spelling:
|
||||
frames: Final = text.split("\n\n")
|
||||
assert frames[-1] == "" and all(frame.startswith("data: ") for frame in frames[:-1]), (
|
||||
f"{label} {spelling}: {text}"
|
||||
)
|
||||
parsed: Final = tuple(
|
||||
_JsonRpcResult.model_validate_json(frame.removeprefix("data: ")) for frame in frames[:-1]
|
||||
)
|
||||
assert tuple(frame.result for frame in parsed) == _STREAM_RESULTS, f"{label} {spelling}: {text}"
|
||||
assert {(frame.jsonrpc, frame.id) for frame in parsed} == {("2.0", marker)}, (
|
||||
f"{label} {spelling}: {text}"
|
||||
)
|
||||
continue
|
||||
assert _JsonRpcResult.model_validate_json(text) == _JsonRpcResult(
|
||||
jsonrpc="2.0",
|
||||
id=marker,
|
||||
result={
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": "granted"}],
|
||||
},
|
||||
), f"{label} {spelling}: {text}"
|
||||
status, text = card(key, allowed)
|
||||
assert status == 200 and json.loads(text)["url"].endswith(f"/a2a/{allowed}"), f"{label}: {text}"
|
||||
forwarded = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
wire_methods: Final = tuple(name for _, _, _, name in spellings)
|
||||
assert forwarded == tuple(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[index]["id"] if index < len(forwarded) else "<none>",
|
||||
"method": wire_method,
|
||||
"params": {"configuration": {"blocking": True}, "message": message},
|
||||
}
|
||||
for index, wire_method in enumerate(wire_methods)
|
||||
), f"{label}: {forwarded}"
|
||||
assert all(item["id"] != marker and str(uuid.UUID(item["id"])) == item["id"] for item in forwarded), label
|
||||
|
||||
emptied: Final = gateway.request("PUT", f"/v1/access_group/{group_id}", {"access_agent_ids": []})
|
||||
assert emptied.status_code == 200, emptied.text
|
||||
emptied_outcomes: Final = attempt(callers["access group"], allowed, allowed_name)
|
||||
assert emptied_outcomes == tuple(
|
||||
(spelling, 403, refusal(allowed_name if "by name" in spelling else allowed))
|
||||
for spelling, _, _ in emptied_outcomes
|
||||
), emptied_outcomes
|
||||
assert wire.drain() == (), emptied_outcomes
|
||||
|
||||
|
||||
def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: GET /a2a/{agent_id}/.well-known/agent.json returns 401 'Only proxy admin can be used' for non-admin keys because agent_inference_routes lists only agent-card.json"
|
||||
)
|
||||
marker: Final = "a2alegacy" + 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())
|
||||
|
||||
def refusal(identity: str) -> str:
|
||||
return json.dumps(
|
||||
{"detail": f"Agent '{identity}' is not allowed for your key/team. Contact proxy admin for access."},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
allowed_name: Final = marker + "-allowed"
|
||||
denied_name: Final = marker + "-denied"
|
||||
allowed: Final = _register_agent(gateway, scenario, allowed_name, wire.url)
|
||||
denied: Final = _register_agent(gateway, scenario, denied_name, wire.url)
|
||||
granted_key: Final = scenario.key(object_permission={"agents": [allowed]})
|
||||
unrestricted_key: Final = scenario.key()
|
||||
proxy_base: Final = str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
granted_legacy: Final = gateway.client.get(
|
||||
f"/a2a/{allowed}/.well-known/agent.json", headers={"Authorization": f"Bearer {granted_key}"}
|
||||
)
|
||||
granted_card: Final = gateway.client.get(
|
||||
f"/a2a/{allowed}/.well-known/agent-card.json", headers={"Authorization": f"Bearer {granted_key}"}
|
||||
)
|
||||
assert granted_legacy.status_code == 200, granted_legacy.text
|
||||
assert granted_legacy.json()["url"] == f"{proxy_base}/a2a/{allowed}", granted_legacy.text
|
||||
assert granted_legacy.json() == granted_card.json(), granted_legacy.text
|
||||
|
||||
unrestricted_legacy: Final = gateway.client.get(
|
||||
f"/a2a/{allowed}/.well-known/agent.json", headers={"Authorization": f"Bearer {unrestricted_key}"}
|
||||
)
|
||||
assert unrestricted_legacy.status_code == 200, unrestricted_legacy.text
|
||||
assert unrestricted_legacy.json()["url"] == f"{proxy_base}/a2a/{allowed}", unrestricted_legacy.text
|
||||
wire.drain()
|
||||
|
||||
denied_legacy: Final = gateway.client.get(
|
||||
f"/a2a/{denied}/.well-known/agent.json", headers={"Authorization": f"Bearer {granted_key}"}
|
||||
)
|
||||
assert (denied_legacy.status_code, denied_legacy.text) == (403, refusal(denied)), denied_legacy.text
|
||||
assert wire.drain() == ()
|
||||
|
|
@ -1,11 +1,92 @@
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from a2a.client import A2ACardResolver, ClientConfig, create_client
|
||||
from a2a.types import a2a_pb2
|
||||
from google.protobuf.json_format import MessageToDict
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
|
||||
class _JsonRpcResult(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
jsonrpc: str
|
||||
id: str
|
||||
result: dict[str, JsonValue]
|
||||
|
||||
|
||||
_MESSAGE_V03: Final = {
|
||||
"kind": "message",
|
||||
"messageId": "client-message",
|
||||
"parts": [{"kind": "text", "text": "synthetic ping"}],
|
||||
"role": "user",
|
||||
}
|
||||
_MESSAGE_V10: Final = {"messageId": "client-message", "parts": [{"text": "synthetic ping"}], "role": "ROLE_USER"}
|
||||
|
||||
|
||||
def _peer_card(url: str, name: str, version: str = "0.3") -> dict[str, JsonValue]:
|
||||
return {
|
||||
"protocolVersion": version,
|
||||
"name": name,
|
||||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
|
||||
|
||||
def _is_uuid(value: str) -> bool:
|
||||
return (
|
||||
re.fullmatch(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", value) is not None
|
||||
and str(uuid.UUID(value)) == value
|
||||
)
|
||||
|
||||
|
||||
def _register_agent(gateway: Gateway, scenario: Scenario, body: dict[str, JsonValue]) -> str:
|
||||
created: Final = gateway.request("POST", "/v1/agents", body)
|
||||
assert created.status_code == 200, created.text
|
||||
identity: Final = created.json()["agent_id"]
|
||||
|
||||
def cleanup() -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/agents/{identity}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (identity,)) == []
|
||||
|
||||
scenario.cleanups.callback(cleanup)
|
||||
return identity
|
||||
|
||||
|
||||
def _proxy_base(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
def _message_reply(body: dict[str, JsonValue], text: str) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": text}],
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.compatibility.a2a.supported_versions_preserve_literal_envelopes")
|
||||
|
|
@ -214,3 +295,748 @@ def test_agent_serving_its_card_only_at_versioned_path_is_reached_with_bearer_an
|
|||
actual: Final = wire.drain()
|
||||
assert tuple(item.target for item in actual if item.method == "GET")[-1] == "/agentCard/v1.0", actual
|
||||
assert tuple(item.target for item in actual if item.method == "POST") == ("/",), actual
|
||||
|
||||
|
||||
_STREAM_EVENTS: Final = (
|
||||
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "working"}, "final": False},
|
||||
{
|
||||
"kind": "artifact-update",
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "chunk"}]},
|
||||
},
|
||||
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "completed"}, "final": True},
|
||||
)
|
||||
_STREAM_RESULTS: Final = {
|
||||
"0.3": (
|
||||
{"contextId": "c1", "final": False, "kind": "status-update", "status": {"state": "working"}, "taskId": "t1"},
|
||||
{
|
||||
"append": False,
|
||||
"artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "chunk"}]},
|
||||
"contextId": "c1",
|
||||
"kind": "artifact-update",
|
||||
"lastChunk": False,
|
||||
"taskId": "t1",
|
||||
},
|
||||
{"contextId": "c1", "final": True, "kind": "status-update", "status": {"state": "completed"}, "taskId": "t1"},
|
||||
),
|
||||
"1.0": (
|
||||
{"statusUpdate": {"taskId": "t1", "contextId": "c1", "status": {"state": "TASK_STATE_WORKING"}}},
|
||||
{
|
||||
"artifactUpdate": {
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"artifact": {"artifactId": "a1", "parts": [{"text": "chunk"}]},
|
||||
}
|
||||
},
|
||||
{"statusUpdate": {"taskId": "t1", "contextId": "c1", "status": {"state": "TASK_STATE_COMPLETED"}}},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
_TRANSPORT_HEADERS: Final = frozenset(
|
||||
{
|
||||
"host",
|
||||
"accept",
|
||||
"accept-encoding",
|
||||
"connection",
|
||||
"content-length",
|
||||
"content-type",
|
||||
"user-agent",
|
||||
"cache-control",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("version", "method", "message"),
|
||||
(("0.3", "message/stream", _MESSAGE_V03), ("1.0", "SendStreamingMessage", _MESSAGE_V10)),
|
||||
ids=("v0.3", "v1.0"),
|
||||
)
|
||||
def test_a2a_stream_relays_each_upstream_event_as_its_own_sse_frame_in_the_pinned_dialect(
|
||||
gateway: Gateway, version: str, method: str, message: dict[str, JsonValue]
|
||||
) -> None:
|
||||
marker: Final = "a2astream" + uuid.uuid4().hex
|
||||
first_frame_seen: Final = threading.Event()
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
assert request.target == "/.well-known/agent-card.json", request.target
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
frames: Final = tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': 'upstream-' + marker, 'result': event})}\n\n".encode()
|
||||
for event in _STREAM_EVENTS
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=frames, gate_after_first=first_frame_seen)
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
user: Final = scenario.user()
|
||||
key: Final = scenario.key(user_id=user)
|
||||
identity: Final = _register_agent(
|
||||
gateway,
|
||||
scenario,
|
||||
{
|
||||
"agent_name": marker,
|
||||
"agent_card_params": _peer_card(wire.url, marker, version),
|
||||
"static_headers": {"Authorization": "Bearer stream-server-token"},
|
||||
},
|
||||
)
|
||||
with gateway.client.stream(
|
||||
"POST",
|
||||
f"/a2a/{identity}",
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"a2a-version": version,
|
||||
"x-litellm-trace-id": marker + "-trace",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
json={"jsonrpc": "2.0", "id": marker, "method": method, "params": {"message": message}},
|
||||
) as response:
|
||||
assert response.status_code == 200, response.read().decode()
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8", response.headers
|
||||
pieces = response.iter_text()
|
||||
received = ""
|
||||
for piece in pieces:
|
||||
received += piece
|
||||
if "\n\n" in received:
|
||||
break
|
||||
first_frame_seen.set()
|
||||
received += "".join(pieces)
|
||||
frames: Final = received.split("\n\n")
|
||||
assert frames[-1] == "" and all(frame.startswith("data: ") for frame in frames[:-1]), received
|
||||
parsed: Final = tuple(_JsonRpcResult.model_validate_json(frame.removeprefix("data: ")) for frame in frames[:-1])
|
||||
assert tuple(frame.result for frame in parsed) == _STREAM_RESULTS[version], received
|
||||
assert {(frame.jsonrpc, frame.id) for frame in parsed} == {("2.0", marker)}, received
|
||||
posts: Final = tuple(item for item in wire.drain() if item.method == "POST")
|
||||
assert len(posts) == 1, posts
|
||||
assert posts[0].target == "/", posts[0].target
|
||||
assert posts[0].headers["accept"] == "text/event-stream", posts[0].headers
|
||||
forwarded: Final = {name: value for name, value in posts[0].headers.items() if name not in _TRANSPORT_HEADERS}
|
||||
assert forwarded == {
|
||||
"authorization": "Bearer stream-server-token",
|
||||
"a2a-version": "0.3",
|
||||
"x-litellm-trace-id": marker + "-trace",
|
||||
"x-litellm-user-id": user,
|
||||
"x-litellm-agent-id": identity,
|
||||
}, posts[0].headers
|
||||
body: Final = json.loads(posts[0].body)
|
||||
assert str(uuid.UUID(body["id"])) == body["id"] and body["id"] != marker, posts[0].body
|
||||
assert body == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"method": "message/stream",
|
||||
"params": {"configuration": {"blocking": True}, "message": _MESSAGE_V03},
|
||||
}, posts[0].body
|
||||
|
||||
|
||||
def test_a2a_send_forwards_configured_headers_and_minted_identity_but_never_spoofed_or_proxy_credentials(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = "a2ahdr" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
return _message_reply(json.loads(request.body), "headers seen")
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = _register_agent(
|
||||
gateway,
|
||||
scenario,
|
||||
{
|
||||
"agent_name": marker,
|
||||
"agent_card_params": _peer_card(wire.url, marker),
|
||||
"static_headers": {"Authorization": "Bearer server-token"},
|
||||
"extra_headers": ["x-api-key"],
|
||||
},
|
||||
)
|
||||
team: Final = scenario.team()
|
||||
member: Final = scenario.user()
|
||||
team_key: Final = scenario.key(team_id=team, user_id=member)
|
||||
solo_user: Final = scenario.user()
|
||||
solo_key: Final = scenario.key(user_id=solo_user)
|
||||
payload: Final = {"jsonrpc": "2.0", "id": marker, "method": "message/send", "params": {"message": _MESSAGE_V03}}
|
||||
spoofs: Final = {
|
||||
"X-API-Key": "client-secret",
|
||||
f"x-a2a-{marker}-authorization": "Bearer agent-token",
|
||||
f"x-a2a-{marker}-x-tenant": "tenant-1",
|
||||
f"x-a2a-{marker}-x-litellm-team-id": "spoofed-team",
|
||||
f"x-a2a-{identity}-x-region": "eu-1",
|
||||
"x-a2a-other-agent-authorization": "Bearer leak",
|
||||
"X-LiteLLM-User-Id": "spoofed-user",
|
||||
"X-LiteLLM-Team-Id": "spoofed-team",
|
||||
}
|
||||
for headers in (
|
||||
{"Authorization": f"Bearer {team_key}", "x-litellm-trace-id": marker + "-trace", **spoofs},
|
||||
{"x-litellm-api-key": solo_key, "Authorization": "Bearer backend-token", **spoofs},
|
||||
):
|
||||
response = gateway.client.post(f"/a2a/{identity}", headers=headers, json=payload)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JsonRpcResult.model_validate_json(response.content).result["parts"] == [
|
||||
{"kind": "text", "text": "headers seen"}
|
||||
], response.text
|
||||
posts: Final = tuple(item for item in wire.drain() if item.method == "POST")
|
||||
assert len(posts) == 2, posts
|
||||
for post, (user, team_header, key) in zip(
|
||||
posts, ((member, team, team_key), (solo_user, None, solo_key)), strict=True
|
||||
):
|
||||
forwarded = {name: value for name, value in post.headers.items() if name not in _TRANSPORT_HEADERS}
|
||||
minted: Final = forwarded.get("x-litellm-trace-id", "")
|
||||
if team_header is None:
|
||||
assert _is_uuid(minted), post.headers
|
||||
expected = {
|
||||
"authorization": "Bearer server-token",
|
||||
"x-api-key": "client-secret",
|
||||
"x-tenant": "tenant-1",
|
||||
"x-region": "eu-1",
|
||||
"a2a-version": "0.3",
|
||||
"x-litellm-agent-id": identity,
|
||||
"x-litellm-trace-id": marker + "-trace" if team_header is not None else minted,
|
||||
"x-litellm-user-id": user,
|
||||
**({} if team_header is None else {"x-litellm-team-id": team_header}),
|
||||
}
|
||||
assert forwarded == expected, post.headers
|
||||
body = json.loads(post.body)
|
||||
assert str(uuid.UUID(body["id"])) == body["id"], post.body
|
||||
assert body == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"method": "message/send",
|
||||
"params": {"configuration": {"blocking": True}, "message": _MESSAGE_V03},
|
||||
}, post.body
|
||||
assert all(
|
||||
key not in value and gateway.key not in value and "leak" not in value for value in post.headers.values()
|
||||
), post.headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", ("0.3", "1.0"))
|
||||
def test_a2a_card_routes_front_the_agent_with_the_proxy_url_and_sdk_clients_call_back_through_the_proxy(
|
||||
gateway: Gateway, version: str
|
||||
) -> None:
|
||||
marker: Final = "a2acard" + uuid.uuid4().hex
|
||||
|
||||
def upstream_card() -> dict[str, JsonValue]:
|
||||
return {
|
||||
**_peer_card(wire.url, marker, version),
|
||||
"supportedInterfaces": [{"url": wire.url + "/", "protocolBinding": "JSONRPC", "protocolVersion": version}],
|
||||
}
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(upstream_card()).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
if body["method"] in ("message/stream", "SendStreamingMessage"):
|
||||
events: Final = (
|
||||
_STREAM_EVENTS
|
||||
if version == "0.3"
|
||||
else tuple(_STREAM_RESULTS["1.0"][index] for index in range(len(_STREAM_EVENTS)))
|
||||
)
|
||||
frames: Final = tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': body['id'], 'result': event})}\n\n".encode()
|
||||
for event in events
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
if version == "0.3":
|
||||
return _message_reply(body, "via proxy")
|
||||
message: Final = {"messageId": "peer-message", "role": "ROLE_AGENT", "parts": [{"text": "via proxy"}]}
|
||||
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": {"message": message}}).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
user: Final = scenario.user()
|
||||
key: Final = scenario.key(user_id=user)
|
||||
identity: Final = _register_agent(
|
||||
gateway,
|
||||
scenario,
|
||||
{
|
||||
"agent_name": marker,
|
||||
"agent_card_params": upstream_card(),
|
||||
"static_headers": {"Authorization": "Bearer card-server-token"},
|
||||
},
|
||||
)
|
||||
proxy_url: Final = f"{_proxy_base(gateway)}/a2a/{identity}"
|
||||
cards: Final = tuple(
|
||||
gateway.request("GET", f"/a2a/{identity}/.well-known/{name}") for name in ("agent-card.json", "agent.json")
|
||||
)
|
||||
for card in cards:
|
||||
assert card.status_code == 200, card.text
|
||||
assert wire.url not in card.text, card.text
|
||||
assert card.json() == cards[0].json(), card.text
|
||||
served: Final = cards[0].json()
|
||||
assert served["url"] == proxy_url, cards[0].text
|
||||
assert served["protocolVersion"] == version, cards[0].text
|
||||
assert served["supportedInterfaces"] == [
|
||||
{"url": proxy_url, "protocolBinding": "JSONRPC", "protocolVersion": version}
|
||||
], cards[0].text
|
||||
assert served["name"] == marker, cards[0].text
|
||||
|
||||
async def call_through_sdk() -> tuple[a2a_pb2.AgentCard, tuple[a2a_pb2.StreamResponse, ...]]:
|
||||
async with httpx.AsyncClient(headers={"Authorization": f"Bearer {key}"}, timeout=30) as http:
|
||||
resolved = await A2ACardResolver(httpx_client=http, base_url=proxy_url).get_agent_card()
|
||||
client = await create_client(resolved, ClientConfig(httpx_client=http, streaming=False))
|
||||
request = a2a_pb2.SendMessageRequest(
|
||||
message=a2a_pb2.Message(
|
||||
message_id=marker + "-in", role=a2a_pb2.Role.ROLE_USER, parts=[a2a_pb2.Part(text="ping")]
|
||||
)
|
||||
)
|
||||
return resolved, tuple([event async for event in client.send_message(request)])
|
||||
|
||||
resolved, events = asyncio.run(call_through_sdk())
|
||||
assert [interface.url for interface in resolved.supported_interfaces] == [proxy_url], resolved
|
||||
assert [event.message.parts[0].text for event in events] == ["via proxy"], events
|
||||
|
||||
async def stream_through_sdk() -> tuple[a2a_pb2.StreamResponse, ...]:
|
||||
async with httpx.AsyncClient(headers={"Authorization": f"Bearer {key}"}, timeout=30) as http:
|
||||
client = await create_client(resolved, ClientConfig(httpx_client=http, streaming=True))
|
||||
request = a2a_pb2.SendMessageRequest(
|
||||
message=a2a_pb2.Message(
|
||||
message_id=marker + "-stream", role=a2a_pb2.Role.ROLE_USER, parts=[a2a_pb2.Part(text="ping")]
|
||||
)
|
||||
)
|
||||
return tuple([event async for event in client.send_message(request)])
|
||||
|
||||
stream_events: Final = asyncio.run(stream_through_sdk())
|
||||
assert [MessageToDict(event, preserving_proto_field_name=False) for event in stream_events] == [
|
||||
{"statusUpdate": {"taskId": "t1", "contextId": "c1", "status": {"state": "TASK_STATE_WORKING"}}},
|
||||
{
|
||||
"artifactUpdate": {
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"artifact": {"artifactId": "a1", "parts": [{"text": "chunk"}]},
|
||||
}
|
||||
},
|
||||
{"statusUpdate": {"taskId": "t1", "contextId": "c1", "status": {"state": "TASK_STATE_COMPLETED"}}},
|
||||
], stream_events
|
||||
posts: Final = tuple(item for item in wire.drain() if item.method == "POST")
|
||||
assert len(posts) == 2, posts
|
||||
forwarded_headers: Final = tuple(
|
||||
{name: value for name, value in post.headers.items() if name not in _TRANSPORT_HEADERS} for post in posts
|
||||
)
|
||||
assert all(_is_uuid(headers.get("x-litellm-trace-id", "")) for headers in forwarded_headers), forwarded_headers
|
||||
assert forwarded_headers == tuple(
|
||||
{
|
||||
"authorization": "Bearer card-server-token",
|
||||
"x-litellm-user-id": user,
|
||||
"x-litellm-agent-id": identity,
|
||||
"x-litellm-trace-id": headers.get("x-litellm-trace-id", ""),
|
||||
"a2a-version": version,
|
||||
}
|
||||
for headers in forwarded_headers
|
||||
), [post.headers for post in posts]
|
||||
forwarded: Final = tuple(json.loads(post.body) for post in posts)
|
||||
assert all(str(uuid.UUID(post["id"])) == post["id"] for post in forwarded), [post.body for post in posts]
|
||||
assert forwarded == (
|
||||
{
|
||||
"0.3": {
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[0]["id"],
|
||||
"method": "message/send",
|
||||
"params": {
|
||||
"configuration": {"blocking": True},
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"messageId": marker + "-in",
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "ping"}],
|
||||
},
|
||||
},
|
||||
},
|
||||
"1.0": {
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[0]["id"],
|
||||
"method": "SendMessage",
|
||||
"params": {
|
||||
"configuration": {},
|
||||
"message": {"messageId": marker + "-in", "role": "ROLE_USER", "parts": [{"text": "ping"}]},
|
||||
},
|
||||
},
|
||||
}[version],
|
||||
{
|
||||
"0.3": {
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[1]["id"],
|
||||
"method": "message/stream",
|
||||
"params": {
|
||||
"configuration": {"blocking": True},
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"messageId": marker + "-stream",
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "ping"}],
|
||||
},
|
||||
},
|
||||
},
|
||||
"1.0": {
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[1]["id"],
|
||||
"method": "SendStreamingMessage",
|
||||
"params": {
|
||||
"configuration": {},
|
||||
"message": {
|
||||
"messageId": marker + "-stream",
|
||||
"role": "ROLE_USER",
|
||||
"parts": [{"text": "ping"}],
|
||||
},
|
||||
},
|
||||
},
|
||||
}[version],
|
||||
), [post.body for post in posts]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streaming", (False, True), ids=("send", "stream"))
|
||||
def test_a2a_message_send_aliases_and_agent_name_reach_the_same_agent_with_the_same_conversion(
|
||||
gateway: Gateway, streaming: bool
|
||||
) -> None:
|
||||
marker: Final = "a2aalias" + uuid.uuid4().hex
|
||||
method: Final = "SendStreamingMessage" if streaming else "SendMessage"
|
||||
wire_method: Final = "message/stream" if streaming else "message/send"
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
assert (request.method, request.target) == ("POST", "/"), request.target
|
||||
body: Final = json.loads(request.body)
|
||||
if body["method"] == "message/stream":
|
||||
frames: Final = tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': 'upstream-' + str(body['id']), 'result': event})}\n\n".encode()
|
||||
for event in _STREAM_EVENTS
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
return _message_reply(body, "alias pong")
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = _register_agent(
|
||||
gateway, scenario, {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker, "1.0")}
|
||||
)
|
||||
paths: Final = (f"/a2a/{identity}/message/send", f"/v1/a2a/{identity}/message/send", f"/a2a/{marker}")
|
||||
for path in paths:
|
||||
headers: Final = {"Authorization": f"Bearer {gateway.key}", "a2a-version": "1.0"}
|
||||
payload: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": method,
|
||||
"params": {"message": _MESSAGE_V10},
|
||||
}
|
||||
if streaming:
|
||||
with gateway.client.stream(
|
||||
"POST", path, headers={**headers, "Accept": "text/event-stream"}, json=payload
|
||||
) as response:
|
||||
received: Final = response.read().decode()
|
||||
assert response.status_code == 200, f"{path}: {received}"
|
||||
frames: Final = received.split("\n\n")
|
||||
assert frames[-1] == "" and all(frame.startswith("data: ") for frame in frames[:-1]), (
|
||||
f"{path}: {received}"
|
||||
)
|
||||
parsed: Final = tuple(
|
||||
_JsonRpcResult.model_validate_json(frame.removeprefix("data: ")) for frame in frames[:-1]
|
||||
)
|
||||
assert tuple(frame.result for frame in parsed) == _STREAM_RESULTS["1.0"], f"{path}: {received}"
|
||||
assert {(frame.jsonrpc, frame.id) for frame in parsed} == {("2.0", marker)}, f"{path}: {received}"
|
||||
continue
|
||||
response = gateway.client.post(path, headers=headers, json=payload)
|
||||
assert response.status_code == 200, f"{path}: {response.text}"
|
||||
assert _JsonRpcResult.model_validate_json(response.content) == _JsonRpcResult(
|
||||
jsonrpc="2.0",
|
||||
id=marker,
|
||||
result={
|
||||
"message": {"messageId": "peer-message", "role": "ROLE_AGENT", "parts": [{"text": "alias pong"}]}
|
||||
},
|
||||
), f"{path}: {response.text}"
|
||||
posts: Final = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
assert len(posts) == len(paths), posts
|
||||
assert all(str(uuid.UUID(post["id"])) == post["id"] and post["id"] != marker for post in posts), posts
|
||||
assert posts == tuple(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": post["id"],
|
||||
"method": wire_method,
|
||||
"params": {"configuration": {"blocking": True}, "message": _MESSAGE_V03},
|
||||
}
|
||||
for post in posts
|
||||
), posts
|
||||
|
||||
|
||||
def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated_before_forwarding(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = "a2atask" + uuid.uuid4().hex
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
if body["method"] == "tasks/resubscribe":
|
||||
frames: Final = tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': body['id'], 'result': event})}\n\n".encode()
|
||||
for event in _STREAM_EVENTS
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
result: JsonValue = (
|
||||
_peer_card(wire.url, marker)
|
||||
if body["method"] == "agent/getAuthenticatedExtendedCard"
|
||||
else {"kind": "task", "id": "t1", "contextId": "c1", "status": {"state": body["method"]}}
|
||||
)
|
||||
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": result}).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = _register_agent(
|
||||
gateway, scenario, {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker)}
|
||||
)
|
||||
|
||||
def call(method: str, params: dict[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.client.post(
|
||||
f"/a2a/{identity}",
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
json={"jsonrpc": "2.0", "id": marker + method, "method": method, "params": params},
|
||||
)
|
||||
|
||||
task_params: Final = {"id": "t1", "historyLength": 2}
|
||||
push_params: Final = {"id": "t1", "pushNotificationConfigId": "c1"}
|
||||
for method, canonical, params in (
|
||||
("tasks/get", "tasks/get", task_params),
|
||||
("GetTask", "tasks/get", task_params),
|
||||
("tasks/list", "tasks/list", {"contextId": "c1"}),
|
||||
("ListTasks", "tasks/list", {"contextId": "c1"}),
|
||||
("tasks/cancel", "tasks/cancel", {"id": "t1"}),
|
||||
("CancelTask", "tasks/cancel", {"id": "t1"}),
|
||||
("tasks/pushNotificationConfig/get", "tasks/pushNotificationConfig/get", push_params),
|
||||
("GetTaskPushNotificationConfig", "tasks/pushNotificationConfig/get", push_params),
|
||||
("tasks/pushNotificationConfig/list", "tasks/pushNotificationConfig/list", {"id": "t1"}),
|
||||
("ListTaskPushNotificationConfigs", "tasks/pushNotificationConfig/list", {"id": "t1"}),
|
||||
("tasks/pushNotificationConfig/delete", "tasks/pushNotificationConfig/delete", push_params),
|
||||
("DeleteTaskPushNotificationConfig", "tasks/pushNotificationConfig/delete", push_params),
|
||||
):
|
||||
response = call(method, params)
|
||||
assert response.status_code == 200, f"{method}: {response.text}"
|
||||
assert response.json() == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + method,
|
||||
"result": {"kind": "task", "id": "t1", "contextId": "c1", "status": {"state": canonical}},
|
||||
}, f"{method}: {response.text}"
|
||||
forwarded = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
assert forwarded == (
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + method,
|
||||
"method": canonical,
|
||||
"params": params,
|
||||
},
|
||||
), f"{method}: {forwarded}"
|
||||
|
||||
for method in ("tasks/resubscribe", "SubscribeToTask"):
|
||||
with gateway.client.stream(
|
||||
"POST",
|
||||
f"/a2a/{identity}",
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "Accept": "text/event-stream"},
|
||||
json={"jsonrpc": "2.0", "id": marker + method, "method": method, "params": {"id": "t1"}},
|
||||
) as response:
|
||||
received: Final = response.read().decode()
|
||||
assert response.status_code == 200, f"{method}: {received}"
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8", response.headers
|
||||
frames: Final = received.split("\n\n")
|
||||
assert frames[-1] == "" and all(frame.startswith("data: ") for frame in frames[:-1]), (
|
||||
f"{method}: {received}"
|
||||
)
|
||||
parsed: Final = tuple(
|
||||
_JsonRpcResult.model_validate_json(frame.removeprefix("data: ")) for frame in frames[:-1]
|
||||
)
|
||||
assert tuple(frame.result for frame in parsed) == _STREAM_EVENTS, f"{method}: {received}"
|
||||
assert {(frame.jsonrpc, frame.id) for frame in parsed} == {("2.0", marker + method)}, (
|
||||
f"{method}: {received}"
|
||||
)
|
||||
forwarded = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
assert forwarded == (
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + method,
|
||||
"method": "tasks/resubscribe",
|
||||
"params": {"id": "t1"},
|
||||
},
|
||||
), f"{method}: {forwarded}"
|
||||
|
||||
for callback, detail in (
|
||||
("http://169.254.169.254/latest/meta-data/", "Push notification URL must use HTTPS"),
|
||||
(
|
||||
"https://10.0.0.1/callback",
|
||||
"URL targets a blocked address (10.0.0.1). If this is a legitimate internal service, "
|
||||
"add the host to `user_url_allowed_hosts` in litellm_settings.",
|
||||
),
|
||||
):
|
||||
for params in (
|
||||
{"taskId": "t1", "pushNotificationConfig": {"url": callback}},
|
||||
{"taskId": "t1", "url": callback},
|
||||
):
|
||||
rejected = call("tasks/pushNotificationConfig/set", params)
|
||||
assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), (
|
||||
f"{params}: {rejected.text}"
|
||||
)
|
||||
assert tuple(item for item in wire.drain() if item.method == "POST") == ()
|
||||
|
||||
public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}}
|
||||
accepted: Final = call("tasks/pushNotificationConfig/set", public)
|
||||
assert (accepted.status_code, accepted.json()) == (
|
||||
200,
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + "tasks/pushNotificationConfig/set",
|
||||
"result": {
|
||||
"kind": "task",
|
||||
"id": "t1",
|
||||
"contextId": "c1",
|
||||
"status": {"state": "tasks/pushNotificationConfig/set"},
|
||||
},
|
||||
},
|
||||
), accepted.text
|
||||
assert tuple(json.loads(item.body)["params"] for item in wire.drain() if item.method == "POST") == (public,)
|
||||
|
||||
extended: Final = call("agent/getAuthenticatedExtendedCard", {})
|
||||
assert extended.status_code == 200, extended.text
|
||||
assert extended.json()["result"] == {
|
||||
**_peer_card(wire.url, marker),
|
||||
"url": f"{_proxy_base(gateway)}/a2a/{identity}",
|
||||
}, extended.text
|
||||
assert wire.url not in extended.text, extended.text
|
||||
assert tuple(json.loads(item.body)["method"] for item in wire.drain() if item.method == "POST") == (
|
||||
"agent/getAuthenticatedExtendedCard",
|
||||
)
|
||||
|
||||
|
||||
def _v1_card(url: str, name: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
**_peer_card(url, name, "1.0"),
|
||||
"supportedInterfaces": [{"url": url + "/", "protocolBinding": "JSONRPC", "protocolVersion": "1.0"}],
|
||||
}
|
||||
|
||||
|
||||
def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_in_1_0(gateway: Gateway) -> None:
|
||||
marker: Final = "a2atask10" + uuid.uuid4().hex
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_v1_card(wire.url, marker)).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
results: Final[dict[str, JsonValue]] = {
|
||||
"tasks/get": {"kind": "task", "id": "t1", "contextId": "c1", "status": {"state": "working"}},
|
||||
"tasks/cancel": {"kind": "task", "id": "t1", "contextId": "c1", "status": {"state": "canceled"}},
|
||||
"tasks/pushNotificationConfig/set": {
|
||||
"taskId": "t1",
|
||||
"pushNotificationConfig": {"id": "c1", "url": "https://1.1.1.1/callback"},
|
||||
},
|
||||
"agent/getAuthenticatedExtendedCard": _v1_card(wire.url, marker),
|
||||
}
|
||||
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": results[body["method"]]}).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = _register_agent(
|
||||
gateway, scenario, {"agent_name": marker, "agent_card_params": _v1_card(wire.url, marker)}
|
||||
)
|
||||
proxy_url: Final = f"{_proxy_base(gateway)}/a2a/{identity}"
|
||||
|
||||
def call(method: str, params: dict[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.client.post(
|
||||
f"/a2a/{identity}",
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "a2a-version": "1.0"},
|
||||
json={"jsonrpc": "2.0", "id": marker + method, "method": method, "params": params},
|
||||
)
|
||||
|
||||
def forwarded() -> tuple[JsonValue, ...]:
|
||||
return tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
|
||||
for method, canonical, state, params in (
|
||||
("GetTask", "tasks/get", "TASK_STATE_WORKING", {"id": "t1", "historyLength": 2}),
|
||||
("CancelTask", "tasks/cancel", "TASK_STATE_CANCELED", {"id": "t1"}),
|
||||
):
|
||||
response = call(method, params)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + method,
|
||||
"result": {"id": "t1", "contextId": "c1", "status": {"state": state}},
|
||||
}, response.text
|
||||
assert forwarded() == ({"jsonrpc": "2.0", "id": marker + method, "method": canonical, "params": params},)
|
||||
|
||||
for params in (
|
||||
{"parent": "tasks/t1", "configId": "c1", "config": {"url": "http://169.254.169.254/latest/meta-data/"}},
|
||||
{"taskId": "t1", "url": "http://169.254.169.254/latest/meta-data/"},
|
||||
):
|
||||
rejected = call("CreateTaskPushNotificationConfig", params)
|
||||
assert (rejected.status_code, rejected.json()) == (
|
||||
400,
|
||||
{"detail": "Push notification URL must use HTTPS"},
|
||||
), f"{params}: {rejected.text}"
|
||||
assert forwarded() == ()
|
||||
|
||||
accepted: Final = call(
|
||||
"CreateTaskPushNotificationConfig",
|
||||
{"parent": "tasks/t1", "configId": "c1", "config": {"url": "https://1.1.1.1/callback"}},
|
||||
)
|
||||
assert accepted.status_code == 200, accepted.text
|
||||
assert forwarded() == (
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + "CreateTaskPushNotificationConfig",
|
||||
"method": "tasks/pushNotificationConfig/set",
|
||||
"params": {"taskId": "t1", "pushNotificationConfig": {"id": "c1", "url": "https://1.1.1.1/callback"}},
|
||||
},
|
||||
)
|
||||
|
||||
extended: Final = call("GetExtendedAgentCard", {})
|
||||
assert extended.status_code == 200, extended.text
|
||||
assert extended.json() == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + "GetExtendedAgentCard",
|
||||
"result": {
|
||||
**_v1_card(wire.url, marker),
|
||||
"url": proxy_url,
|
||||
"supportedInterfaces": [{"url": proxy_url, "protocolBinding": "JSONRPC", "protocolVersion": "1.0"}],
|
||||
},
|
||||
}, extended.text
|
||||
assert wire.url not in extended.text, extended.text
|
||||
assert forwarded() == (
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker + "GetExtendedAgentCard",
|
||||
"method": "agent/getAuthenticatedExtendedCard",
|
||||
"params": {},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_a2a_v1_push_config_envelope_sent_to_a_0_3_agent_has_its_callback_validated_before_forwarding(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
pytest.skip(
|
||||
"BUG: CreateTaskPushNotificationConfig {parent, configId, config: {url: http://169.254.169.254/...}} with "
|
||||
"a2a-version 1.0 to an agent pinned to 0.3 returns 200 and forwards the callback unvalidated, because "
|
||||
"validation reads only params.url and pushNotificationConfig.url and 0.3-pinned params are never lowered"
|
||||
)
|
||||
marker: Final = "a2apush03" + uuid.uuid4().hex
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": body["params"]}).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = _register_agent(
|
||||
gateway, scenario, {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker)}
|
||||
)
|
||||
rejected: Final = gateway.client.post(
|
||||
f"/a2a/{identity}",
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "a2a-version": "1.0"},
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "CreateTaskPushNotificationConfig",
|
||||
"params": {
|
||||
"parent": "tasks/t1",
|
||||
"configId": "c1",
|
||||
"config": {"url": "http://169.254.169.254/latest/meta-data/"},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert (rejected.status_code, rejected.json()) == (
|
||||
400,
|
||||
{"detail": "Push notification URL must use HTTPS"},
|
||||
), rejected.text
|
||||
assert tuple(item for item in wire.drain() if item.method == "POST") == ()
|
||||
|
|
|
|||
145
tests/integration/management/test_agent_management.py
Normal file
145
tests/integration/management/test_agent_management.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
import json
|
||||
import uuid
|
||||
from typing import Final
|
||||
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
_REDACTED: Final = "REDACTED_BY_LITELM"
|
||||
|
||||
|
||||
class _AgentRead(BaseModel):
|
||||
agent_id: str
|
||||
agent_name: str
|
||||
litellm_params: dict[str, JsonValue] | None
|
||||
static_headers: dict[str, str] | None
|
||||
|
||||
|
||||
def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"protocolVersion": "0.3",
|
||||
"name": name,
|
||||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"capabilities": {"streaming": False},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
|
||||
|
||||
def _register_agent(gateway: Gateway, scenario: Scenario, name: str, url: str, secret: str) -> str:
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/agents",
|
||||
{
|
||||
"agent_name": name,
|
||||
"agent_card_params": _peer_card(url, name),
|
||||
"litellm_params": {"api_key": secret, "cost_per_query": 0.1},
|
||||
"static_headers": {"Authorization": f"Bearer {secret}"},
|
||||
},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
assert secret not in json.dumps(created.json()["litellm_params"]), created.text
|
||||
identity: Final = created.json()["agent_id"]
|
||||
|
||||
def cleanup() -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/agents/{identity}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (identity,)) == []
|
||||
|
||||
scenario.cleanups.callback(cleanup)
|
||||
return identity
|
||||
|
||||
|
||||
def test_agent_reads_redact_credentials_filter_by_key_permission_and_keep_the_stored_credential_for_calls(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = "agentread" + uuid.uuid4().hex[:12]
|
||||
secret: Final = "sk-agent-" + uuid.uuid4().hex
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
body: Final = json.loads(request.body)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": "authorized"}],
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
visible: Final = _register_agent(gateway, scenario, marker + "-visible", wire.url, secret)
|
||||
hidden: Final = _register_agent(gateway, scenario, marker + "-hidden", wire.url, secret)
|
||||
restricted: Final = scenario.key(object_permission={"agents": [visible]})
|
||||
redacted_params: Final = {"api_key": _REDACTED, "cost_per_query": 0.1, "is_public": False}
|
||||
|
||||
listed: Final = gateway.client.get("/v1/agents", headers={"Authorization": f"Bearer {restricted}"})
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert secret not in listed.text, listed.text
|
||||
assert [
|
||||
(agent.agent_id, agent.litellm_params, agent.static_headers)
|
||||
for agent in map(_AgentRead.model_validate, listed.json())
|
||||
] == [(visible, redacted_params, None)], listed.text
|
||||
|
||||
fetched: Final = gateway.client.get(f"/v1/agents/{visible}", headers={"Authorization": f"Bearer {restricted}"})
|
||||
assert fetched.status_code == 200, fetched.text
|
||||
assert secret not in fetched.text, fetched.text
|
||||
read: Final = _AgentRead.model_validate_json(fetched.content)
|
||||
assert (read.agent_id, read.litellm_params, read.static_headers) == (visible, redacted_params, None), (
|
||||
fetched.text
|
||||
)
|
||||
|
||||
refused: Final = gateway.client.get(f"/v1/agents/{hidden}", headers={"Authorization": f"Bearer {restricted}"})
|
||||
assert (refused.status_code, refused.json()) == (
|
||||
403,
|
||||
{"detail": f"Agent '{hidden}' is not allowed for your key/team. Contact proxy admin for access."},
|
||||
), refused.text
|
||||
|
||||
admin_list: Final = gateway.request("GET", "/v1/agents")
|
||||
assert admin_list.status_code == 200, admin_list.text
|
||||
admin_rows: Final = {
|
||||
agent.agent_id: agent.litellm_params
|
||||
for agent in map(_AgentRead.model_validate, admin_list.json())
|
||||
if agent.agent_id in (visible, hidden)
|
||||
}
|
||||
assert admin_rows == {visible: redacted_params, hidden: redacted_params}, admin_list.text
|
||||
|
||||
stored: Final = read_rows(
|
||||
"SELECT litellm_params->>'api_key' AS api_key FROM \"LiteLLM_AgentsTable\" WHERE agent_id=%s", (visible,)
|
||||
)
|
||||
assert stored == [{"api_key": secret}], stored
|
||||
|
||||
invoked: Final = gateway.client.post(
|
||||
f"/a2a/{visible}",
|
||||
headers={"Authorization": f"Bearer {restricted}"},
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": marker,
|
||||
"method": "message/send",
|
||||
"params": {
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"role": "user",
|
||||
"messageId": marker + "-in",
|
||||
"parts": [{"kind": "text", "text": "ping"}],
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
assert invoked.status_code == 200, invoked.text
|
||||
assert invoked.json()["result"]["parts"] == [{"kind": "text", "text": "authorized"}], invoked.text
|
||||
posts: Final = tuple(item for item in wire.drain() if item.method == "POST")
|
||||
assert [post.headers["authorization"] for post in posts] == [f"Bearer {secret}"], posts
|
||||
277
tests/integration/spend/test_agent_spend.py
Normal file
277
tests/integration/spend/test_agent_spend.py
Normal file
|
|
@ -0,0 +1,277 @@
|
|||
import json
|
||||
import uuid
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
|
||||
class _ActivityMetrics(BaseModel):
|
||||
spend: float
|
||||
api_requests: int
|
||||
successful_requests: int
|
||||
|
||||
|
||||
class _ActivityDay(BaseModel):
|
||||
date: str
|
||||
metrics: _ActivityMetrics
|
||||
|
||||
|
||||
class _Activity(BaseModel):
|
||||
results: tuple[_ActivityDay, ...]
|
||||
|
||||
|
||||
_CALL_TYPES: Final = {
|
||||
"message/send": "asend_message",
|
||||
"message/stream": "asend_message_streaming",
|
||||
"SendMessage": "asend_message",
|
||||
"SendStreamingMessage": "asend_message_streaming",
|
||||
}
|
||||
_PER_TOKEN: Final[dict[str, JsonValue]] = {"input_cost_per_token": 0.125, "output_cost_per_token": 0.5}
|
||||
_PER_TOKEN_COST: Final = 1 * 0.125 + 2 * 0.5
|
||||
|
||||
_MESSAGE_V03: Final = {"kind": "message", "role": "user", "messageId": "", "parts": [{"kind": "text", "text": "ping"}]}
|
||||
|
||||
|
||||
def _peer_card(url: str, name: str, version: str = "0.3") -> dict[str, JsonValue]:
|
||||
return {
|
||||
"protocolVersion": version,
|
||||
"name": name,
|
||||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
|
||||
|
||||
def _upstream_reply(body: dict[str, JsonValue]) -> Reply:
|
||||
if body["method"] == "message/stream":
|
||||
events: Final = (
|
||||
{
|
||||
"kind": "artifact-update",
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "billed"}]},
|
||||
},
|
||||
{
|
||||
"kind": "status-update",
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"status": {"state": "completed"},
|
||||
"final": True,
|
||||
},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(
|
||||
f"data: {json.dumps({'jsonrpc': '2.0', 'id': body['id'], 'result': event})}\n\n".encode()
|
||||
for event in events
|
||||
),
|
||||
)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": "billed"}],
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _assert_a2a_calls_are_billed(
|
||||
gateway: Gateway,
|
||||
pricing: dict[str, JsonValue],
|
||||
cost: float,
|
||||
calls: tuple[tuple[str, str], ...],
|
||||
version: str = "0.3",
|
||||
) -> None:
|
||||
marker: Final = "a2aspend" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
return _upstream_reply(json.loads(request.body))
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/agents",
|
||||
{
|
||||
"agent_name": marker,
|
||||
"agent_card_params": _peer_card(wire.url, marker, version),
|
||||
"litellm_params": pricing,
|
||||
},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
agent: Final = created.json()["agent_id"]
|
||||
|
||||
def cleanup() -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/agents/{agent}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent,)) == []
|
||||
|
||||
scenario.cleanups.callback(cleanup)
|
||||
key: Final = scenario.key()
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
message: Final = (
|
||||
{"kind": "message", "role": "user", "messageId": marker + "-in", "parts": [{"kind": "text", "text": "ping"}]}
|
||||
if version == "0.3"
|
||||
else {"role": "ROLE_USER", "messageId": marker + "-in", "parts": [{"text": "ping"}]}
|
||||
)
|
||||
wire_message: Final = {
|
||||
**_MESSAGE_V03,
|
||||
"messageId": marker + "-in",
|
||||
}
|
||||
methods: Final = tuple(method for method, _ in calls)
|
||||
wire_methods: Final = tuple(
|
||||
"message/stream" if "stream" in method.lower() else "message/send" for method in methods
|
||||
)
|
||||
for method, path in calls:
|
||||
payload = {"jsonrpc": "2.0", "id": f"{marker}-{method}", "method": method, "params": {"message": message}}
|
||||
headers: Final = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
**({"a2a-version": version} if version == "1.0" else {}),
|
||||
}
|
||||
if "stream" not in method.lower():
|
||||
sent = gateway.client.post(
|
||||
path.format(agent=agent, name=marker), headers=headers, json=payload
|
||||
)
|
||||
assert sent.status_code == 200, sent.text
|
||||
result: Final = sent.json()["result"]
|
||||
parts: Final = result["parts"] if version == "0.3" else result["message"]["parts"]
|
||||
expected_parts: Final = (
|
||||
[{"kind": "text", "text": "billed"}] if version == "0.3" else [{"text": "billed"}]
|
||||
)
|
||||
assert parts == expected_parts, sent.text
|
||||
continue
|
||||
with gateway.client.stream(
|
||||
"POST",
|
||||
path.format(agent=agent, name=marker),
|
||||
headers={**headers, "Accept": "text/event-stream"},
|
||||
json=payload,
|
||||
) as streamed:
|
||||
text = streamed.read().decode()
|
||||
assert streamed.status_code == 200, text
|
||||
assert text.startswith("data: ") and (
|
||||
'"state": "completed"' in text if version == "0.3" else '"state": "TASK_STATE_COMPLETED"' in text
|
||||
), text
|
||||
forwarded: Final = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
|
||||
assert forwarded == tuple(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": forwarded[index]["id"] if index < len(forwarded) else "<none>",
|
||||
"method": wire_method,
|
||||
"params": {"configuration": {"blocking": True}, "message": wire_message},
|
||||
}
|
||||
for index, wire_method in enumerate(wire_methods)
|
||||
), forwarded
|
||||
assert all(str(uuid.UUID(item["id"])) == item["id"] for item in forwarded), forwarded
|
||||
|
||||
logs: Final = eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT request_id, agent_id, spend, api_key, call_type, "
|
||||
"""to_char("startTime", 'YYYY-MM-DD') AS day FROM "LiteLLM_SpendLogs" """
|
||||
"WHERE api_key=%s ORDER BY request_id",
|
||||
(digest,),
|
||||
),
|
||||
lambda rows: len(rows) == len(methods),
|
||||
seconds=30,
|
||||
)
|
||||
assert [{name: value for name, value in row.items() if name != "day"} for row in logs] == [
|
||||
{
|
||||
"request_id": f"{marker}-{method}",
|
||||
"agent_id": agent,
|
||||
"spend": cost,
|
||||
"api_key": digest,
|
||||
"call_type": _CALL_TYPES[method],
|
||||
}
|
||||
for method in sorted(methods)
|
||||
], logs
|
||||
total: Final = cost * len(methods)
|
||||
assert eventually(
|
||||
lambda: read_rows('SELECT spend FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent,)),
|
||||
lambda rows: rows == [{"spend": total}],
|
||||
seconds=20,
|
||||
) == [{"spend": total}]
|
||||
assert eventually(
|
||||
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
||||
lambda rows: rows == [{"spend": total}],
|
||||
seconds=20,
|
||||
) == [{"spend": total}]
|
||||
requests_per_day: Final = {
|
||||
day: sum(1 for row in logs if row["day"] == day) for day in sorted({str(row["day"]) for row in logs})
|
||||
}
|
||||
daily: Final = eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT agent_id, date, api_key, spend, api_requests::text AS api_requests, "
|
||||
'successful_requests::text AS successful_requests FROM "LiteLLM_DailyAgentSpend" WHERE agent_id=%s '
|
||||
"ORDER BY date",
|
||||
(agent,),
|
||||
),
|
||||
lambda rows: sum(int(str(row["api_requests"])) for row in rows) == len(methods),
|
||||
seconds=45,
|
||||
)
|
||||
assert daily == [
|
||||
{
|
||||
"agent_id": agent,
|
||||
"date": day,
|
||||
"api_key": digest,
|
||||
"spend": cost * count,
|
||||
"api_requests": str(count),
|
||||
"successful_requests": str(count),
|
||||
}
|
||||
for day, count in requests_per_day.items()
|
||||
], daily
|
||||
activity: Final = gateway.request(
|
||||
"GET",
|
||||
"/agent/daily/activity",
|
||||
params={"agent_ids": agent, "start_date": min(requests_per_day), "end_date": max(requests_per_day)},
|
||||
)
|
||||
assert activity.status_code == 200, activity.text
|
||||
days: Final = _Activity.model_validate_json(activity.content).results
|
||||
assert [(day.date, day.metrics) for day in days] == [
|
||||
(day, _ActivityMetrics(spend=cost * count, api_requests=count, successful_requests=count))
|
||||
for day, count in sorted(requests_per_day.items(), reverse=True)
|
||||
], activity.text
|
||||
assert [
|
||||
result["breakdown"]["models"][f"a2a_agent/{marker}"]["api_key_breakdown"][digest]["metrics"]["spend"]
|
||||
for result in activity.json()["results"]
|
||||
] == [cost * count for _, count in sorted(requests_per_day.items(), reverse=True)], activity.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", ("0.3", "1.0"))
|
||||
def test_a2a_send_and_stream_bill_cost_per_query_to_the_agent_the_key_and_daily_agent_activity(
|
||||
gateway: Gateway, version: str
|
||||
) -> None:
|
||||
calls: Final = (
|
||||
(("message/send", "/a2a/{agent}"), ("message/stream", "/a2a/{agent}"))
|
||||
if version == "0.3"
|
||||
else (("SendMessage", "/v1/a2a/{agent}/message/send"), ("SendStreamingMessage", "/a2a/{name}"))
|
||||
)
|
||||
_assert_a2a_calls_are_billed(gateway, {"cost_per_query": 0.25}, 0.25, calls, version)
|
||||
|
||||
|
||||
def test_a2a_send_bills_input_and_output_tokens_at_the_agent_per_token_prices(gateway: Gateway) -> None:
|
||||
_assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/send", "/a2a/{agent}"),))
|
||||
|
||||
|
||||
def test_a2a_stream_bills_input_and_output_tokens_at_the_agent_per_token_prices(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: message/stream to an agent priced with input_cost_per_token/output_cost_per_token logs spend 0.0 "
|
||||
"with prompt_tokens 0 and completion_tokens 0, because A2AStreamingIterator reads text from the request "
|
||||
"Part RootModels without dumping them and ignores artifact-update text, while message/send bills 1.125"
|
||||
)
|
||||
_assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/stream", "/a2a/{agent}"),))
|
||||
Loading…
Add table
Reference in a new issue