diff --git a/tests/integration/authorization/test_agent_access.py b/tests/integration/authorization/test_agent_access.py new file mode 100644 index 00000000000..09e2afd289d --- /dev/null +++ b/tests/integration/authorization/test_agent_access.py @@ -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 "", + "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() == () diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 2823911ace3..621dcfe060e 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -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") == () diff --git a/tests/integration/management/test_agent_management.py b/tests/integration/management/test_agent_management.py new file mode 100644 index 00000000000..5c0fdc68e09 --- /dev/null +++ b/tests/integration/management/test_agent_management.py @@ -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 diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py new file mode 100644 index 00000000000..58fcd4e9de2 --- /dev/null +++ b/tests/integration/spend/test_agent_spend.py @@ -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 "", + "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}"),))