test/a2a/test_lifecycle.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from uuid import uuid4
import httpx
import pytest
from a2a.server.agent_execution import AgentExecutor as A2AAgentExecutorBase
from a2a.server.agent_execution import RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import InMemoryPushNotificationConfigStore, TaskUpdater
from a2a.types import Task, TaskState, TaskStatus
from fast_depends.pydantic import PydanticSerializer
from ag2 import Agent, Context, MemoryStream
from ag2.a2a import A2AConfig, A2AServer, build_card
from ag2.a2a.client import A2AClient
from ag2.a2a.errors import A2ATaskFailedError
from ag2.a2a.testing import make_test_client_factory
from ag2.events import ModelMessage, ModelMessageChunk, ModelRequest, ModelResponse, TextInput
from ag2.testing import TestConfig
# What an ``Agent`` hands its ``LLMClient``, built the same way.
_SERIALIZER = PydanticSerializer(pydantic_config={"arbitrary_types_allowed": True}, use_fastdepends_errors=False)
class _SpyAsyncClient(httpx.AsyncClient):
aclose_count = 0
async def aclose(self) -> None:
type(self).aclose_count += 1
await super().aclose()
def _make_spy_factory(server: A2AServer, url: str) -> Callable[[], httpx.AsyncClient]:
transport = httpx.ASGITransport(app=server.build_jsonrpc(url=url))
def factory() -> httpx.AsyncClient:
return _SpyAsyncClient(transport=transport, base_url=url, timeout=30.0)
return factory
class _AlwaysFailingExecutor(A2AAgentExecutorBase):
async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
msg = request_context.message
assert msg is not None
task_id = msg.task_id or uuid4().hex
context_id = msg.context_id or uuid4().hex
await event_queue.enqueue_event(
Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED)),
)
updater = TaskUpdater(event_queue, task_id, context_id)
await updater.start_work()
await updater.failed()
async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
return None
@pytest.mark.asyncio
class TestHttpxLifecycle:
async def test_client_closes_httpx_after_successful_ask(self) -> None:
_SpyAsyncClient.aclose_count = 0
server = A2AServer(Agent("server", config=TestConfig("hi")))
url = "http://test"
client = Agent("client", config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url)))
await client.ask("ping")
assert _SpyAsyncClient.aclose_count == 1
async def test_client_closes_httpx_when_task_fails(self) -> None:
_SpyAsyncClient.aclose_count = 0
server = A2AServer(
Agent("stub", config=TestConfig("unused")),
executor=_AlwaysFailingExecutor(),
)
url = "http://test"
client = Agent(
"client",
config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url), streaming=False),
)
with pytest.raises(A2ATaskFailedError):
await client.ask("ping")
assert _SpyAsyncClient.aclose_count == 1
async def test_closing_leaves_the_client_reusable_for_the_next_turn(self) -> None:
# ``A2AClient`` is reused across ``reply.ask(...)`` follow-ups and
# closes its httpx client at the end of every turn, so each turn has
# to reconnect from scratch rather than reach for a closed client.
_SpyAsyncClient.aclose_count = 0
server = A2AServer(Agent("server", config=TestConfig("hi")))
url = "http://test"
client = Agent("client", config=A2AConfig(card_url=url, httpx_client_factory=_make_spy_factory(server, url)))
first = await client.ask("ping")
second = await first.ask("ping again")
assert second.body == "hi"
assert _SpyAsyncClient.aclose_count == 2
async def test_aclose_on_a_client_that_never_connected_is_a_no_op(self) -> None:
_SpyAsyncClient.aclose_count = 0
client = A2AClient(card_url="http://test", httpx_client_factory=_SpyAsyncClient)
await client.aclose()
await client.aclose()
assert _SpyAsyncClient.aclose_count == 0
async def test_aclose_after_a_turn_does_not_close_the_same_connection_twice(self) -> None:
# The turn closes in its own ``finally``; an explicit ``aclose()``
# on top must not reach the httpx client again. Driven directly
# because ``Agent`` never hands the client back to the caller.
_SpyAsyncClient.aclose_count = 0
server = A2AServer(Agent("server", config=TestConfig("hi")))
url = "http://test"
client = A2AClient(card_url=url, httpx_client_factory=_make_spy_factory(server, url))
await client(
[ModelRequest([TextInput("ping")])],
Context(stream=MemoryStream()),
tools=[],
response_schema=None,
serializer=_SERIALIZER,
)
await client.aclose()
assert _SpyAsyncClient.aclose_count == 1
@pytest.mark.asyncio
async def test_streamed_chunks_not_duplicated_in_final_message() -> None:
turn = TestConfig(ModelMessageChunk("he"), ModelMessageChunk("llo"), ModelResponse(message=ModelMessage("hello")))
server = A2AServer(Agent("server", config=turn))
url = "http://test"
factory = make_test_client_factory(server, url=url)
client = Agent(
"client",
config=A2AConfig(card_url=url, httpx_client_factory=factory, streaming=True),
)
reply = await client.ask("ping")
assert reply.response.content == "hello"
@pytest.mark.asyncio
class TestCardImmutability:
async def test_build_jsonrpc_does_not_mutate_input_card(self) -> None:
agent = Agent("server", config=TestConfig("hi"))
card = build_card(agent, url="http://test", transports=("jsonrpc",))
before_extended = card.capabilities.extended_agent_card
before_push = card.capabilities.push_notifications
server = A2AServer(
agent,
extended_card=card,
push_config_store=InMemoryPushNotificationConfigStore(),
)
server.build_jsonrpc(url="http://test", card=card)
assert card.capabilities.extended_agent_card == before_extended
assert card.capabilities.push_notifications == before_push
async def test_build_rest_does_not_mutate_input_card(self) -> None:
agent = Agent("server", config=TestConfig("hi"))
card = build_card(agent, url="http://test", transports=("rest",))
before_extended = card.capabilities.extended_agent_card
before_push = card.capabilities.push_notifications
server = A2AServer(
agent,
extended_card=card,
push_config_store=InMemoryPushNotificationConfigStore(),
)
server.build_rest(url="http://test", card=card)
assert card.capabilities.extended_agent_card == before_extended
assert card.capabilities.push_notifications == before_push
async def test_build_grpc_does_not_mutate_input_card(self) -> None:
agent = Agent("server", config=TestConfig("hi"))
card = build_card(agent, url="http://test", transports=("grpc",), grpc_url="localhost:50051")
before_extended = card.capabilities.extended_agent_card
before_push = card.capabilities.push_notifications
server = A2AServer(
agent,
extended_card=card,
push_config_store=InMemoryPushNotificationConfigStore(),
)
server.build_grpc(bind="localhost:0", grpc_url="localhost:50051", card=card)
assert card.capabilities.extended_agent_card == before_extended
assert card.capabilities.push_notifications == before_push