test/a2a/_helpers.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
import asyncio
from collections.abc import Awaitable, Callable, Iterable
from dataclasses import dataclass
from typing import Any
from uuid import uuid4
import grpc.aio
import httpx
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 (
PushNotificationConfigStore,
TaskStore,
TaskUpdater,
)
from a2a.types import Part, Task, TaskState, TaskStatus
from ag2 import Agent
from ag2.a2a import A2AConfig, A2AServer, build_card
from ag2.a2a.testing import (
ToolRoundsConfig,
make_test_client_factory,
make_test_rest_client_factory,
pick_free_port,
)
from ag2.events import (
BaseEvent,
ModelMessageChunk,
)
from ag2.testing import TestConfig, TrackingConfig, Turn
def a2a_config(agent: Agent) -> A2AConfig:
"""The agent's ``A2AConfig``.
``Agent.config`` is declared ``ModelConfig | None``, but the task and
push-notification admin helpers take an ``A2AConfig`` — narrow once here
instead of handing them an unchecked union at every call site.
"""
config = agent.config
assert isinstance(config, A2AConfig), f"{agent.name} is not backed by an A2AConfig"
return config
@dataclass(slots=True)
class A2APair:
server: A2AServer
server_agent: Agent
client: Agent
tracking: TrackingConfig
@dataclass(slots=True)
class ExecutorPair:
server: A2AServer
executor: A2AAgentExecutorBase
client: Agent
@dataclass(slots=True)
class GrpcPair:
server: A2AServer
server_agent: Agent
client: Agent
tracking: TrackingConfig
grpc_url: str
grpc_server: grpc.aio.Server
class PromptThenAckExecutor(A2AAgentExecutorBase):
def __init__(self, prompt: str) -> None:
self._prompt = prompt
self.received_user_text: str | None = None
async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
msg = request_context.message
if msg is None:
return
task_id = msg.task_id or uuid4().hex
context_id = msg.context_id or uuid4().hex
updater = TaskUpdater(event_queue, task_id, context_id)
if request_context.current_task is None:
await event_queue.enqueue_event(
Task(
id=task_id,
context_id=context_id,
status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
),
)
await updater.start_work()
await updater.requires_input(
message=updater.new_agent_message(parts=[Part(text=self._prompt)]),
)
return
text = "".join(p.text for p in msg.parts if p.text)
self.received_user_text = text
await updater.complete(
message=updater.new_agent_message(parts=[Part(text=f"echo: {text}")]),
)
async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
task = request_context.current_task
if task is None:
return
updater = TaskUpdater(event_queue, task.id, task.context_id)
await updater.cancel()
class Switchboard(httpx.AsyncBaseTransport):
"""Send each request to whichever replica is `active` — a load balancer that changes its mind."""
def __init__(self, *servers: A2AServer, url: str) -> None:
self._transports = [httpx.ASGITransport(app=server.build_jsonrpc(url=url)) for server in servers]
self.active = 0
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
return await self._transports[self.active].handle_async_request(request)
class GatedExecutor(A2AAgentExecutorBase):
"""Starts a task, then holds it `working` until the test opens the gate."""
def __init__(self) -> None:
self.gate = asyncio.Event()
self.working = asyncio.Event()
self.task_id = ""
async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
task_id, context_id = request_context.task_id or uuid4().hex, request_context.context_id or uuid4().hex
self.task_id = task_id
updater = TaskUpdater(event_queue, task_id, context_id)
await event_queue.enqueue_event(
Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED))
)
await updater.start_work()
self.working.set()
await self.gate.wait()
await updater.complete(message=updater.new_agent_message(parts=[Part(text="finished")]))
async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
raise NotImplementedError
def _tool_rounds(initial: Turn, after_tool: "Turn | None") -> ToolRoundsConfig:
"""The server answers with ``initial``, then with ``after_tool`` once the tool results come back."""
return ToolRoundsConfig(initial, after_tool if after_tool is not None else "")
def make_pair(
initial: Turn,
after_tool: "Turn | None" = None,
*,
server_tools: Iterable[Callable[..., object]] = (),
client_tools: Iterable[Callable[..., object]] = (),
server_url: str = "http://test",
streaming: bool = True,
task_store: TaskStore | None = None,
push_config_store: PushNotificationConfigStore | None = None,
) -> A2APair:
tracking = TrackingConfig(_tool_rounds(initial, after_tool))
server_agent = Agent("server-agent", config=tracking)
for tool in server_tools:
server_agent.tool(tool)
server_kwargs: dict[str, Any] = {}
if task_store is not None:
server_kwargs["task_store"] = task_store
if push_config_store is not None:
server_kwargs["push_config_store"] = push_config_store
server = A2AServer(server_agent, **server_kwargs)
factory = make_test_client_factory(server, url=server_url)
client_config = A2AConfig(
card_url=server_url,
httpx_client_factory=factory,
streaming=streaming,
)
client_agent = Agent("client-agent", config=client_config)
for tool in client_tools:
client_agent.tool(tool)
return A2APair(server=server, server_agent=server_agent, client=client_agent, tracking=tracking)
def make_executor_pair(
executor: A2AAgentExecutorBase,
*,
server_url: str = "http://test",
streaming: bool = False,
task_store: TaskStore | None = None,
push_config_store: PushNotificationConfigStore | None = None,
hitl_hook: Callable[..., Any] | None = None,
) -> ExecutorPair:
server_agent = Agent("server-stub", config=TestConfig("unused"))
server_kwargs: dict[str, Any] = {}
if task_store is not None:
server_kwargs["task_store"] = task_store
if push_config_store is not None:
server_kwargs["push_config_store"] = push_config_store
server = A2AServer(server_agent, executor=executor, **server_kwargs)
factory = make_test_client_factory(server, url=server_url)
client_kwargs: dict[str, Any] = {}
if hitl_hook is not None:
client_kwargs["hitl_hook"] = hitl_hook
client = Agent(
"client",
config=A2AConfig(card_url=server_url, httpx_client_factory=factory, streaming=streaming),
**client_kwargs,
)
return ExecutorPair(server=server, executor=executor, client=client)
def make_rest_pair(
initial: Turn,
after_tool: "Turn | None" = None,
*,
server_url: str = "http://test",
streaming: bool = False,
) -> A2APair:
tracking = TrackingConfig(_tool_rounds(initial, after_tool))
server_agent = Agent("server-agent", config=tracking)
server = A2AServer(server_agent)
factory = make_test_rest_client_factory(server, url=server_url)
client = Agent(
"client-agent",
config=A2AConfig(
card_url=server_url,
httpx_client_factory=factory,
prefer="rest",
streaming=streaming,
),
)
return A2APair(server=server, server_agent=server_agent, client=client, tracking=tracking)
async def start_grpc_pair(
initial: Turn,
after_tool: "Turn | None" = None,
*,
host: str = "127.0.0.1",
streaming: bool = False,
) -> GrpcPair:
tracking = TrackingConfig(_tool_rounds(initial, after_tool))
server_agent = Agent("server-agent", config=tracking)
server = A2AServer(server_agent)
grpc_url = f"{host}:{pick_free_port(host)}"
card = build_card(server_agent, url=grpc_url, transports=("grpc",), grpc_url=grpc_url)
grpc_server = server.build_grpc(bind=grpc_url, grpc_url=grpc_url, card=card)
await grpc_server.start()
client = Agent(
"client-agent",
config=A2AConfig(
card_url=grpc_url,
preset_card=card,
prefer="grpc",
streaming=streaming,
),
)
return GrpcPair(
server=server,
server_agent=server_agent,
client=client,
tracking=tracking,
grpc_url=grpc_url,
grpc_server=grpc_server,
)
_Script = Callable[[TaskUpdater, asyncio.Event], Awaitable[None]]
class ScriptedExecutor(A2AAgentExecutorBase):
"""Runs a fresh task through ``script``; ``release`` is opened by the test, never by a timer."""
def __init__(self, script: _Script) -> None:
self._script = script
self.release = asyncio.Event()
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 self._script(updater, self.release)
async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
return None
class EventCollector:
"""Subscriber that records the events it is given; ``on_event`` is handed to ``stream.subscribe``."""
def __init__(self) -> None:
self.events: list[BaseEvent] = []
async def on_event(self, event: BaseEvent) -> None:
self.events.append(event)
@property
def chunks(self) -> list[str]:
return [e.content for e in self.events if isinstance(e, ModelMessageChunk)]
_Turn = Callable[[TaskUpdater], Awaitable[None]]
async def _complete_with_ok(updater: TaskUpdater) -> None:
await updater.complete(message=updater.new_agent_message([Part(text="ok")]))
class ToolRoundTripExecutor(A2AAgentExecutorBase):
"""Runs ``first_turn`` for a new task and ``continuation`` once the client answers its request for input."""
def __init__(self, first_turn: _Turn, continuation: _Turn = _complete_with_ok) -> None:
self._first_turn = first_turn
self._continuation = continuation
async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
msg = request_context.message
assert msg is not None
resumed = request_context.current_task is not None
task_id = msg.task_id or uuid4().hex
context_id = msg.context_id or uuid4().hex
if not resumed:
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 (self._continuation if resumed else self._first_turn)(updater)
async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
return None