mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge 6ff17267b1 into e16d79d635
This commit is contained in:
mrinal-berri
GitHub
commit
d94c107774
2 changed files with 218 additions and 81 deletions
|
|
@ -219,6 +219,8 @@ def register(
|
|||
name: str,
|
||||
url: str,
|
||||
litellm_params: Mapping[str, JsonValue] | None,
|
||||
*,
|
||||
database_url: str | None = None,
|
||||
) -> str:
|
||||
body: dict[str, JsonValue] = {
|
||||
"agent_name": name,
|
||||
|
|
@ -243,7 +245,14 @@ def register(
|
|||
def delete() -> 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,)) == []
|
||||
assert (
|
||||
read_rows(
|
||||
'SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s',
|
||||
(identity,),
|
||||
database_url=database_url,
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
cleanups(delete)
|
||||
return identity
|
||||
|
|
@ -277,7 +286,9 @@ def sse_frames(text: str) -> tuple[dict[str, JsonValue], ...]:
|
|||
)
|
||||
|
||||
|
||||
def spend_rows(request_id: str) -> list[dict[str, JsonValue]]:
|
||||
def spend_rows(request_id: str, *, database_url: str | None = None) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT request_id, model, call_type, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)
|
||||
'SELECT request_id, model, call_type, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
||||
(request_id,),
|
||||
database_url=database_url,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,12 +11,15 @@ from __future__ import annotations
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.database import read_rows, scratch_database
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.process import owned_proxy
|
||||
from integration.compatibility._a2a_peers import (
|
||||
HINT,
|
||||
KIND_ERROR,
|
||||
|
|
@ -38,6 +41,26 @@ PINNED: Final = {"a2a_protocol_version": "0.3"}
|
|||
ROUTES: Final = ("/a2a/{id}", "/a2a/{id}/message/send", "/v1/a2a/{id}/message/send")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Isolated:
|
||||
gateway: Gateway
|
||||
database_url: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def isolated(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Isolated]:
|
||||
with (
|
||||
scratch_database() as database_url,
|
||||
owned_proxy(
|
||||
session_gateway,
|
||||
tmp_path_factory.mktemp("a2a-override"),
|
||||
{"DATABASE_URL": database_url},
|
||||
workers=PROXY_WORKERS,
|
||||
) as proxy,
|
||||
):
|
||||
yield Isolated(proxy, database_url)
|
||||
|
||||
|
||||
def _artifact_text(result: dict[str, JsonValue]) -> str:
|
||||
task: Final = result.get("task", result)
|
||||
assert isinstance(task, dict), result
|
||||
|
|
@ -65,8 +88,10 @@ def _assert_v1_wire(sent: tuple[dict[str, JsonValue], ...], marker: str, method:
|
|||
assert message["messageId"] == marker + "-in", sent
|
||||
|
||||
|
||||
def _one_spend_row(request_id: str, call_type: str, status: str) -> None:
|
||||
rows: Final = eventually(lambda: spend_rows(request_id), lambda values: len(values) == 1, seconds=70)
|
||||
def _one_spend_row(database_url: str, request_id: str, call_type: str, status: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: spend_rows(request_id, database_url=database_url), lambda values: len(values) == 1, seconds=70
|
||||
)
|
||||
assert rows[0]["call_type"] == call_type and rows[0]["status"] == status, rows
|
||||
|
||||
|
||||
|
|
@ -88,18 +113,32 @@ def _send_kind_error(gateway: Gateway, identity: str, marker: str) -> str:
|
|||
|
||||
|
||||
@pytest.mark.parametrize("route", ROUTES)
|
||||
def test_pinned_langgraph_agent_answers_message_send_in_the_legacy_dialect(gateway: Gateway, route: str) -> None:
|
||||
with peer("lgsend", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
_send_ok(gateway, identity, agent.marker, route)
|
||||
def test_pinned_langgraph_agent_answers_message_send_in_the_legacy_dialect(isolated: Isolated, route: str) -> None:
|
||||
with peer("lgsend", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(isolated.gateway, identity, agent.marker, route)
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_pinned_langgraph_agent_streams_message_stream_to_the_final_frame(gateway: Gateway) -> None:
|
||||
with peer("lgstream", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
response: Final = gateway.request(
|
||||
def test_pinned_langgraph_agent_streams_message_stream_to_the_final_frame(isolated: Isolated) -> None:
|
||||
with peer("lgstream", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
response: Final = isolated.gateway.request(
|
||||
"POST", f"/a2a/{identity}", legacy_send(agent.marker, method="message/stream")
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
|
@ -113,15 +152,22 @@ def test_pinned_langgraph_agent_streams_message_stream_to_the_final_frame(gatewa
|
|||
assert results[1]["artifactUpdate"]["artifact"]["parts"] == [{"text": "echo: synthetic ping"}] # type: ignore[index]
|
||||
assert results[2]["statusUpdate"]["status"]["state"] == "TASK_STATE_COMPLETED", response.text # type: ignore[index]
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker, "message/stream") # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message_streaming", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message_streaming", "success")
|
||||
|
||||
|
||||
def test_v1_caller_reaches_pinned_legacy_agent_in_its_own_envelope(gateway: Gateway) -> None:
|
||||
with peer("lgv1caller", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
response: Final = gateway.client.post(
|
||||
def test_v1_caller_reaches_pinned_legacy_agent_in_its_own_envelope(isolated: Isolated) -> None:
|
||||
with peer("lgv1caller", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
response: Final = isolated.gateway.client.post(
|
||||
f"/a2a/{identity}",
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "a2a-version": "1.0"},
|
||||
headers={"Authorization": f"Bearer {isolated.gateway.key}", "a2a-version": "1.0"},
|
||||
json={
|
||||
"jsonrpc": "2.0",
|
||||
"id": agent.marker,
|
||||
|
|
@ -142,101 +188,156 @@ def test_v1_caller_reaches_pinned_legacy_agent_in_its_own_envelope(gateway: Gate
|
|||
assert "kind" not in task and task["status"]["state"] == "TASK_STATE_COMPLETED", response.text
|
||||
assert task["artifacts"][0]["parts"] == [{"text": "echo: synthetic ping"}], response.text
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ("0.3.0", 0.3, " 0.3 "), ids=("dotted", "number", "spaced"))
|
||||
def test_protocol_version_aliases_pin_the_legacy_transport(gateway: Gateway, value: JsonValue) -> None:
|
||||
with peer("lgalias", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
def test_protocol_version_aliases_pin_the_legacy_transport(isolated: Isolated, value: JsonValue) -> None:
|
||||
with peer("lgalias", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
gateway, scenario.cleanups.callback, agent.marker, agent.url, {"a2a_protocol_version": value}
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
{"a2a_protocol_version": value},
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_genuine_v1_agent_pinned_one_point_zero_is_unchanged(gateway: Gateway) -> None:
|
||||
with peer("v1pinned", "v1") as agent, gateway.scenario() as scenario:
|
||||
def test_genuine_v1_agent_pinned_one_point_zero_is_unchanged(isolated: Isolated) -> None:
|
||||
with peer("v1pinned", "v1") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
gateway, scenario.cleanups.callback, agent.marker, agent.url, {"a2a_protocol_version": "1.0"}
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
{"a2a_protocol_version": "1.0"},
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_genuine_v1_agent_unpinned_is_unchanged(gateway: Gateway) -> None:
|
||||
with peer("v1plain", "v1") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, None)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
def test_genuine_v1_agent_unpinned_is_unchanged(isolated: Isolated) -> None:
|
||||
with peer("v1plain", "v1") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
None,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pin", (None, "0.3"), ids=("unpinned", "pinned-legacy"))
|
||||
def test_lowercase_binding_still_downgrades_to_the_legacy_transport(gateway: Gateway, pin: str | None) -> None:
|
||||
def test_lowercase_binding_still_downgrades_to_the_legacy_transport(isolated: Isolated, pin: str | None) -> None:
|
||||
with (
|
||||
peer("lglower", "langgraph", interfaces=[{"protocolBinding": "jsonrpc"}]) as agent,
|
||||
gateway.scenario() as scenario,
|
||||
isolated.gateway.scenario() as scenario,
|
||||
):
|
||||
params: Final = None if pin is None else {"a2a_protocol_version": pin}
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, params)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
params,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_explicit_one_point_zero_opts_out_of_the_lowercase_downgrade(gateway: Gateway) -> None:
|
||||
with peer("v1lower", "v1", interfaces=[{"protocolBinding": "jsonrpc"}]) as agent, gateway.scenario() as scenario:
|
||||
def test_explicit_one_point_zero_opts_out_of_the_lowercase_downgrade(isolated: Isolated) -> None:
|
||||
with (
|
||||
peer("v1lower", "v1", interfaces=[{"protocolBinding": "jsonrpc"}]) as agent,
|
||||
isolated.gateway.scenario() as scenario,
|
||||
):
|
||||
identity: Final = register(
|
||||
gateway, scenario.cleanups.callback, agent.marker, agent.url, {"a2a_protocol_version": "1.0"}
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
{"a2a_protocol_version": "1.0"},
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_one_point_zero_pin_on_a_legacy_card_sends_v1(gateway: Gateway) -> None:
|
||||
def test_one_point_zero_pin_on_a_legacy_card_sends_v1(isolated: Isolated) -> None:
|
||||
legacy_card: Final = [{"protocolBinding": "JSONRPC", "protocolVersion": "0.3"}]
|
||||
with peer("v1oncard03", "v1", interfaces=legacy_card) as agent, gateway.scenario() as scenario:
|
||||
with peer("v1oncard03", "v1", interfaces=legacy_card) as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
gateway, scenario.cleanups.callback, agent.marker, agent.url, {"a2a_protocol_version": "1.0"}
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
{"a2a_protocol_version": "1.0"},
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_override_rewrites_only_canonical_interfaces(gateway: Gateway) -> None:
|
||||
with peer("lgmulti", "langgraph") as agent, peer("lgdecoy", "langgraph") as decoy, gateway.scenario() as scenario:
|
||||
def test_override_rewrites_only_canonical_interfaces(isolated: Isolated) -> None:
|
||||
with (
|
||||
peer("lgmulti", "langgraph") as agent,
|
||||
peer("lgdecoy", "langgraph") as decoy,
|
||||
isolated.gateway.scenario() as scenario,
|
||||
):
|
||||
agent.interfaces = [
|
||||
{"protocolBinding": "JSONRPC", "url": agent.url},
|
||||
{"protocolBinding": "custom-ws", "url": decoy.url},
|
||||
{"protocolBinding": "HTTP+JSON", "url": decoy.url},
|
||||
]
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
assert decoy.wire.drain() == (), "the HTTP+JSON and custom bindings must not be selected" # type: ignore[union-attr]
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_override_added_and_removed_by_patch_takes_effect_on_both_workers(gateway: Gateway) -> None:
|
||||
with peer("lgpatch", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, None)
|
||||
assert HINT in _send_kind_error(gateway, identity, agent.marker + "-before")
|
||||
def test_override_added_and_removed_by_patch_takes_effect_on_both_workers(isolated: Isolated) -> None:
|
||||
with peer("lgpatch", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
None,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
assert HINT in _send_kind_error(isolated.gateway, identity, agent.marker + "-before")
|
||||
before: Final = posts(agent.wire) # type: ignore[arg-type]
|
||||
assert [item["method"] for item in before] == ["SendMessage"], before
|
||||
|
||||
def send(suffix: str) -> int:
|
||||
response: Final = gateway.request(
|
||||
response: Final = isolated.gateway.request(
|
||||
"POST", f"/a2a/{identity}", legacy_send(agent.marker + suffix), headers=FRESH_CONNECTION
|
||||
)
|
||||
return response.status_code
|
||||
|
||||
def settle(params: dict[str, JsonValue], expected: int) -> None:
|
||||
patched: Final = gateway.request("PATCH", f"/v1/agents/{identity}", {"litellm_params": params})
|
||||
patched: Final = isolated.gateway.request("PATCH", f"/v1/agents/{identity}", {"litellm_params": params})
|
||||
assert patched.status_code == 200, patched.text
|
||||
assert patched.json()["litellm_params"].get("a2a_protocol_version") == params.get("a2a_protocol_version")
|
||||
written_at: Final = time.monotonic()
|
||||
|
|
@ -249,50 +350,75 @@ def test_override_added_and_removed_by_patch_takes_effect_on_both_workers(gatewa
|
|||
agent.wire.drain() # type: ignore[union-attr]
|
||||
|
||||
settle(dict(PINNED), 200)
|
||||
response: Final = gateway.request("POST", f"/a2a/{identity}", legacy_send(agent.marker))
|
||||
response: Final = isolated.gateway.request("POST", f"/a2a/{identity}", legacy_send(agent.marker))
|
||||
assert response.status_code == 200, response.text
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
settle({}, 500)
|
||||
assert HINT in _send_kind_error(gateway, identity, agent.marker + "-after")
|
||||
assert HINT in _send_kind_error(isolated.gateway, identity, agent.marker + "-after")
|
||||
|
||||
|
||||
def test_repeated_pinned_sends_log_one_row_each(gateway: Gateway) -> None:
|
||||
with peer("lgrepeat", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
def test_repeated_pinned_sends_log_one_row_each(isolated: Isolated) -> None:
|
||||
with peer("lgrepeat", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
markers: Final = tuple(f"{agent.marker}-{index}" for index in range(10))
|
||||
for marker in markers:
|
||||
response = gateway.request("POST", f"/a2a/{identity}", legacy_send(marker))
|
||||
response = isolated.gateway.request("POST", f"/a2a/{identity}", legacy_send(marker))
|
||||
assert response.status_code == 200 and response.json()["id"] == marker, response.text
|
||||
sent: Final = posts(agent.wire) # type: ignore[arg-type]
|
||||
assert [item["params"]["message"]["messageId"] for item in sent] == [m + "-in" for m in markers], sent # type: ignore[index]
|
||||
assert len({item["id"] for item in sent}) == len(markers), sent
|
||||
assert {item["method"] for item in sent} == {"message/send"}, sent
|
||||
for marker in markers:
|
||||
_one_spend_row(marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, marker, "asend_message", "success")
|
||||
|
||||
|
||||
def test_registered_agent_persists_protocol_version(gateway: Gateway) -> None:
|
||||
with peer("lgstored", "langgraph") as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
assert gateway.get(f"/v1/agents/{identity}")["litellm_params"]["a2a_protocol_version"] == "0.3" # type: ignore[index]
|
||||
rows: Final = read_rows('SELECT litellm_params FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (identity,))
|
||||
def test_registered_agent_persists_protocol_version(isolated: Isolated) -> None:
|
||||
with peer("lgstored", "langgraph") as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
assert isolated.gateway.get(f"/v1/agents/{identity}")["litellm_params"]["a2a_protocol_version"] == "0.3" # type: ignore[index]
|
||||
rows: Final = read_rows(
|
||||
'SELECT litellm_params FROM "LiteLLM_AgentsTable" WHERE agent_id=%s',
|
||||
(identity,),
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
stored: Final = rows[0]["litellm_params"]
|
||||
stored_params: Final = json.loads(stored) if isinstance(stored, str) else stored
|
||||
assert stored_params["a2a_protocol_version"] == "0.3", rows # type: ignore[index]
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
_assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_pinned_agent_survives_localhost_retry(gateway: Gateway) -> None:
|
||||
def test_pinned_agent_survives_localhost_retry(isolated: Isolated) -> None:
|
||||
dead: Final = f"http://localhost:{allocate_port()}/a2a/unreachable"
|
||||
with peer("lgretry", "langgraph", card_url=dead) as agent, gateway.scenario() as scenario:
|
||||
identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED)
|
||||
_send_ok(gateway, identity, agent.marker)
|
||||
with peer("lgretry", "langgraph", card_url=dead) as agent, isolated.gateway.scenario() as scenario:
|
||||
identity: Final = register(
|
||||
isolated.gateway,
|
||||
scenario.cleanups.callback,
|
||||
agent.marker,
|
||||
agent.url,
|
||||
PINNED,
|
||||
database_url=isolated.database_url,
|
||||
)
|
||||
_send_ok(isolated.gateway, identity, agent.marker)
|
||||
traffic: Final = agent.wire.drain() # type: ignore[union-attr]
|
||||
sent: Final = tuple(json.loads(item.body) for item in traffic if item.method == "POST")
|
||||
assert tuple(item.target for item in traffic if item.method == "POST") == ("/a2a/" + agent.marker + "/",), (
|
||||
traffic
|
||||
)
|
||||
_assert_legacy_wire(sent, agent.marker)
|
||||
_one_spend_row(agent.marker, "asend_message", "success")
|
||||
_one_spend_row(isolated.database_url, agent.marker, "asend_message", "success")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue