diff --git a/src/band/integrations/a2a/adapter.py b/src/band/integrations/a2a/adapter.py index bb17a73a0..15042ac84 100644 --- a/src/band/integrations/a2a/adapter.py +++ b/src/band/integrations/a2a/adapter.py @@ -37,6 +37,13 @@ logger = logging.getLogger(__name__) +# httpx's read timeout resets on every chunk received, so this bounds the gap +# between SSE events, not the turn as a whole. Generous enough for the +# multi-second silences of a live LLM call or tool loop; still finite, so a +# peer that accepts the connection and then hangs eventually fails the turn +# instead of blocking the room forever. +_SSE_READ_TIMEOUT_S = 300.0 + class A2AAdapter(SimpleAdapter[A2ASessionState]): """Adapter that forwards messages to a remote A2A agent. @@ -107,7 +114,14 @@ async def on_started(self, agent_name: str, agent_description: str) -> None: headers = self.auth.to_headers() if self.auth else {} - self._http_client = httpx.AsyncClient(headers=headers) + # httpx's default 5s read timeout fires on the normal, multi-second + # gap between SSE events during a real remote turn (a live LLM call, + # a tool loop) -- not a hang. Use a generous bound instead of the + # default so a genuinely dead peer still fails promptly. + self._http_client = httpx.AsyncClient( + headers=headers, + timeout=httpx.Timeout(10.0, read=_SSE_READ_TIMEOUT_S), + ) factory = ClientFactory( ClientConfig(streaming=self.streaming, httpx_client=self._http_client) ) @@ -319,12 +333,14 @@ async def on_cleanup(self, room_id: str) -> None: async def cleanup_all(self) -> None: """Close the owned A2A client and its HTTP transport.""" - if self._client is not None: - await self._client.close() - self._client = None - if self._http_client is not None: - await self._http_client.aclose() - self._http_client = None + client, self._client = self._client, None + http_client, self._http_client = self._http_client, None + try: + if client is not None: + await client.close() + finally: + if http_client is not None: + await http_client.aclose() async def _emit_task_event( self, tools: AgentToolsProtocol, task: Task, state: TaskState diff --git a/src/band/integrations/a2a/gateway/adapter.py b/src/band/integrations/a2a/gateway/adapter.py index b7709bc00..29a9b6f4f 100644 --- a/src/band/integrations/a2a/gateway/adapter.py +++ b/src/band/integrations/a2a/gateway/adapter.py @@ -189,7 +189,7 @@ async def on_started(self, agent_name: str, agent_description: str) -> None: ) await self._server.start() - logger.info("Gateway HTTP server started on port %d", self.port) + logger.info("Gateway HTTP server started on port %d", self._server.bound_port) @property def rest(self) -> AsyncRestClient: @@ -344,6 +344,7 @@ async def _execute_a2a( request.context_id, request.pending.task.id, ) + await request.pending.fail("A2A request failed") raise else: if completed: diff --git a/src/band/integrations/a2a/gateway/server.py b/src/band/integrations/a2a/gateway/server.py index 6d0fa28ed..0ba7a5e31 100644 --- a/src/band/integrations/a2a/gateway/server.py +++ b/src/band/integrations/a2a/gateway/server.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import logging from collections.abc import Awaitable, Callable from typing import Any @@ -21,32 +20,29 @@ from starlette.responses import JSONResponse from starlette.routing import BaseRoute, Route +from band.integrations.uvicorn_server import ( + SERVER_START_TIMEOUT_S, + SERVER_STOP_TIMEOUT_S, + ManagedUvicornServer, +) from band_rest import Peer logger = logging.getLogger(__name__) ExecutorFactory = Callable[[str], AgentExecutor] -# uvicorn's own default (None) waits forever for existing connections to close -# on stop() -- and a live message:stream SSE response has no other way to end -# on its own. sse_starlette normally closes it cooperatively on shutdown, but -# that mechanism is a process-global switch any co-located -# band.integrations.mcp.local_server permanently disables (see that module's -# AppStatus.disable_automatic_graceful_drain() call) -- so this bound is the -# only thing that keeps stop() from hanging once that happens. -SERVER_STOP_TIMEOUT_S = 5 +# sse_starlette's shutdown-drain footgun (see uvicorn_server's docstring) +# is disabled by importing that module, not here. # The REST endpoints the gateway serves per peer: the messaging binding and # the compat card. The upstream factory also returns task read/cancel/list # and push-config routes — an unauthenticated window into past conversations. MESSAGING_REST_SUFFIXES = ("/message:send", "/message:stream", "/card") -# The JSON-RPC methods the gateway serves (1.0 names and their v0.3-compat -# spellings). Sends create work; the per-task operations are gated by the -# unguessable task UUID the server minted for the caller. Everything else -# stays closed: with no auth layer every caller shares one identity, so -# enumeration (ListTasks) and the push-config/extended-card methods would -# disclose or disrupt other callers' conversations. +# JSON-RPC methods the gateway serves (1.0 + v0.3-compat spellings). Sends +# create work; per-task ops are gated by the unguessable task UUID. Everything +# else stays closed -- with no auth layer, enumeration/push-config methods +# would disclose or disrupt other callers' conversations. ALLOWED_JSONRPC_METHODS = frozenset( { "SendMessage", @@ -78,8 +74,7 @@ def __init__( self.port = port self.executor_factory = executor_factory self._app: Starlette | None = None - self._uvicorn: Any | None = None - self._server_task: asyncio.Task[Any] | None = None + self._runtime: ManagedUvicornServer | None = None def _agent_card(self, slug: str, peer: Peer) -> AgentCard: rpc_url = f"{self.gateway_url}/agents/{slug}" @@ -193,10 +188,9 @@ def _messaging_rest_routes( ) -> list[BaseRoute]: """The REST binding, reduced to the endpoints this gateway serves. - Beyond the unauthenticated task routes, the upstream factory ends with - a multi-tenant catch-all ``Mount("/{tenant}")``; peers here are - namespaced by path, and the first alias's mount would shadow every - later alias's flat routes. + The upstream factory also returns a catch-all ``Mount("/{tenant}")``; + since peers are namespaced by path here, the first alias's mount + would shadow every later alias's routes. """ return [ route @@ -237,38 +231,33 @@ async def _handle_list_peers(self, _request: Request) -> JSONResponse: ] return JSONResponse({"peers": peers, "count": len(peers)}) - async def start(self) -> None: - import uvicorn + @property + def bound_port(self) -> int: + """The actual listening port -- resolves ``port=0`` to whatever the + OS assigned.""" + if self._runtime is None: + raise RuntimeError("A2A Gateway server has not started") + return self._runtime.bound_port + async def start(self) -> None: self._app = self._build_app() - self._uvicorn = uvicorn.Server( - uvicorn.Config( - self._app, - host="0.0.0.0", - port=self.port, - log_level="warning", - timeout_graceful_shutdown=SERVER_STOP_TIMEOUT_S, - ) + self._runtime = ManagedUvicornServer( + self._app, + host="0.0.0.0", + port=self.port, + start_timeout_s=SERVER_START_TIMEOUT_S, + stop_timeout_s=SERVER_STOP_TIMEOUT_S, ) - self._server_task = asyncio.create_task(self._uvicorn.serve()) + await self._runtime.start() logger.info( "Starting A2A Gateway server on port %d with %d peers", - self.port, + self.bound_port, len(self.peers), ) async def stop(self) -> None: - if self._uvicorn is None or self._server_task is None: + if self._runtime is None: return - # Ask uvicorn to exit rather than cancelling serve(): cancellation - # skips its shutdown phase and leaks the listening socket. - self._uvicorn.should_exit = True - try: - await self._server_task - except asyncio.CancelledError: - raise - except BaseException: # uvicorn raises SystemExit on startup failure - logger.exception("A2A Gateway server exited with error") - self._uvicorn = None - self._server_task = None + await self._runtime.stop() + self._runtime = None logger.info("A2A Gateway server stopped") diff --git a/src/band/integrations/mcp/local_server.py b/src/band/integrations/mcp/local_server.py index 1e039e034..b4ec4bc4e 100644 --- a/src/band/integrations/mcp/local_server.py +++ b/src/band/integrations/mcp/local_server.py @@ -1,15 +1,12 @@ """The embedded MCP front door: run one ``LocalMCPServer`` per adapter. Ephemeral-port scanning starts from a random offset (dodges a just-freed- -port wedge). Mounts ``engine.py``'s FastMCP app rather than hand-rolling a -lowlevel ``Server``; building the tool-registration list itself is -``engine.py``'s job too (``build_band_mcp_tool_registrations`` / -``build_resolved_band_mcp_tool_registrations``) -- this module only runs -the server once it has that list. - -Every lifecycle transition (``start()``/``stop()``) routes through one lock, -with cleanup in ``finally`` -- so a serve-task crash always closes the -socket and resets state, and concurrent start/stop calls can't race. +port wedge). Mounts ``engine.py``'s FastMCP app; ``engine.py`` also builds +the tool-registration list, this module only runs the server once it has one. + +Every ``start()``/``stop()`` routes through one lock with cleanup in +``finally``, so a serve-task crash always closes the socket and resets +state, and concurrent start/stop calls can't race. """ from __future__ import annotations @@ -23,7 +20,6 @@ import uvicorn from mcp.server.fastmcp import FastMCP -from sse_starlette.sse import AppStatus from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import PlainTextResponse @@ -35,6 +31,11 @@ build_engine, validate_unique_tool_names, ) +from band.integrations.uvicorn_server import ( + SERVER_START_TIMEOUT_S, + SERVER_STOP_TIMEOUT_S, + wait_until_started, +) logger = logging.getLogger(__name__) @@ -45,50 +46,22 @@ LOCAL_MCP_HTTP_PATH = "/mcp" LOCAL_MCP_MESSAGE_PATH = "/messages/" LOCAL_MCP_HEALTH_PATH = "/healthz" -SERVER_START_TIMEOUT_S = 5.0 -# uvicorn's own default (None) waits forever for existing connections to close -# on `stop()` -- fatal here, since an MCP client (e.g. OpenCode) holds its `/sse` -# GET open for the life of its session and may never close it on its own after -# we deregister. Bound it so `stop()` force-cancels that connection instead of -# hanging the adapter's cleanup indefinitely. -SERVER_STOP_TIMEOUT_S = 5 - -# sse_starlette's EventSourceResponse watches a process-global AppStatus for a -# shutdown signal, closing every open SSE stream right after its headers once -# latched -- from either of two sources: our own signal handler (neutralized -# by EmbeddedUvicornServer.capture_signals below), or *any other* -# uvicorn.Server anywhere in this process whose handle_exit() ever fires, -# since AppStatus.should_exit is a bare class attribute with no notion of -# "which server." The second case is real: a Windows CI hang traced to -# exactly this, a live SSE connection closing right after its headers with no -# code of ours involved. LocalMCPServer.stop() already forces its own socket -# closed and cancels its serve task directly, so it never needed -# sse_starlette's automatic drain-on-shutdown; disabling it here removes the -# dependency on that global entirely. Process-wide and one-time by nature -# (AppStatus has no per-instance scope), hence a module-level call rather than -# something threaded through LocalMCPServer's API. -# -# Cost of that global scope: any *other* sse_starlette consumer in the same -# process -- e.g. the A2A gateway's own message:stream responses -# (src/band/integrations/a2a/gateway/server.py) -- loses this same -# cooperative-drain signal too, permanently, the moment this module is -# imported anywhere in the process. That server's own uvicorn.Config sets -# timeout_graceful_shutdown precisely so its stop() still bounds how long it -# waits on a live stream, rather than relying on the now-disabled signal. -AppStatus.disable_automatic_graceful_drain() + +# The process-global sse_starlette shutdown-drain footgun (see +# band.integrations.uvicorn_server's docstring) is disabled by importing +# that module above, not here -- a Windows CI hang was traced to exactly +# this before that fix existed. class EmbeddedUvicornServer(uvicorn.Server): """A uvicorn server that leaves process signal handling to its host. - uvicorn's ``serve()`` captures SIGINT/SIGTERM for itself -- fine for a - standalone process, but this server is embedded in a host that may run - several servers over its lifetime and already owns its own signal - handling. It's also the other half of the sse_starlette bug documented at - the ``AppStatus.disable_automatic_graceful_drain()`` call above: capturing - signals here would let sse_starlette latch its process-global shutdown - flag through *this* server's handler too. Shutdown is driven - programmatically instead, via ``should_exit`` (see ``LocalMCPServer.stop``). + uvicorn's ``serve()`` captures SIGINT/SIGTERM by default -- wrong for a + server embedded in a host that already owns signal handling, and it's + the other half of the sse_starlette footgun (see uvicorn_server's + docstring): capturing signals here would let sse_starlette latch its + shutdown flag through this server's handler too. Shutdown goes through + ``should_exit`` instead (see ``LocalMCPServer.stop``). """ @contextmanager @@ -200,9 +173,8 @@ async def start(self) -> None: return reserved_socket, port = self._reserve_socket() - # Tracked immediately, before anything below can raise: `stop()`'s - # cleanup closes `self._socket` unconditionally, so a failure in - # engine/app/uvicorn construction still gets the socket closed + # Tracked immediately: stop()'s cleanup closes self._socket + # unconditionally, so a failure below still gets it closed # instead of leaking a bound-and-listening fd. self._socket = reserved_socket self._port = port @@ -238,7 +210,9 @@ async def start(self) -> None: self._uvicorn_server = uvicorn_server self._serve_task = serve_task - await self._wait_until_started() + await wait_until_started( + uvicorn_server, serve_task, timeout_s=SERVER_START_TIMEOUT_S + ) except Exception: await self._stop_locked() raise @@ -259,9 +233,9 @@ async def stop(self) -> None: async def _stop_locked(self) -> None: """The actual teardown, run only while ``_lifecycle_lock`` is held. - Cleanup lives in ``finally``: the previous version's bare ``await - self._serve_task`` re-raised past the socket-close/state-reset code - below it whenever the serve task crashed with anything but + Cleanup lives in ``finally`` -- a bare ``await self._serve_task`` + outside one would re-raise past the socket-close/state-reset code + below it if the serve task crashed with anything but ``CancelledError``, leaking the socket and leaving stale state for the next ``start()``. """ @@ -288,11 +262,10 @@ async def _stop_locked(self) -> None: def _build_app(self, mcp: FastMCP) -> Starlette: """Mount the engine's SSE + streamable-HTTP routes onto one host app. - ``streamable_http_app()`` lazily creates ``mcp.session_manager`` and - returns its own Starlette app whose lifespan runs it -- but a mounted - sub-app's lifespan is never invoked by the ASGI server, only the - top-level app's is. So the host lifespan below enters - ``session_manager.run()`` itself (verified by the step-1 spike). + ``streamable_http_app()`` lazily creates ``mcp.session_manager``, but + a mounted sub-app's lifespan is never invoked by the ASGI server -- + only the top-level app's is. So the host lifespan below enters + ``session_manager.run()`` itself. """ sse_routes = list(mcp.sse_app().routes) http_routes = list(mcp.streamable_http_app().routes) @@ -322,11 +295,10 @@ def _reserve_socket(self) -> tuple[socket.socket, int]: port = reserved_socket.getsockname()[1] return _listen(reserved_socket), port - # Scan the range from a random starting offset (wrapping around), not - # first-fit from port_min: first-fit hands a new server the port a - # just-stopped sibling freed moments ago, and that port's previous - # consumers (e.g. an MCP client subprocess still winding down) keep - # sending stale session traffic that wedges the new server's transport. + # Random starting offset, not first-fit from port_min: first-fit + # reuses the port a just-stopped sibling freed, and that port's old + # consumers (an MCP client subprocess still winding down) keep + # sending stale traffic that wedges the new server's transport. last_error: OSError | None = None span = self._port_max - self._port_min + 1 start = random.randrange(span) @@ -345,15 +317,3 @@ def _reserve_socket(self) -> tuple[socket.socket, int]: "Could not find a free localhost MCP port in range " f"{self._port_min}-{self._port_max}" ) from last_error - - async def _wait_until_started(self) -> None: - if self._serve_task is None or self._uvicorn_server is None: - raise RuntimeError("Local MCP server task not initialized") - - deadline = asyncio.get_running_loop().time() + SERVER_START_TIMEOUT_S - while not self._uvicorn_server.started: - if self._serve_task.done(): - await self._serve_task - if asyncio.get_running_loop().time() >= deadline: - raise TimeoutError("Timed out waiting for local MCP server startup") - await asyncio.sleep(0.05) diff --git a/src/band/integrations/uvicorn_server.py b/src/band/integrations/uvicorn_server.py new file mode 100644 index 000000000..c8e10e593 --- /dev/null +++ b/src/band/integrations/uvicorn_server.py @@ -0,0 +1,129 @@ +"""Shared lifecycle for integrations that embed their own uvicorn server. + +Used by ``mcp.local_server``, ``a2a.gateway.server``, and the A2A baseline +test fixture. Importing this module also disables sse_starlette's +automatic graceful-drain watcher (see the ``AppStatus`` call below): its +shutdown signal is a bare process-global with no notion of "which server," +so every embedder needs it disabled once, not per caller. +""" + +from __future__ import annotations + +import asyncio +import logging + +import uvicorn +from sse_starlette.sse import AppStatus +from starlette.types import ASGIApp + +logger = logging.getLogger(__name__) + +POLL_INTERVAL_S = 0.05 + +# How long start() waits for uvicorn to report ready -- without it, a caller +# dialing in right after start() returns could race a socket that isn't +# listening yet. +SERVER_START_TIMEOUT_S = 5.0 + +# uvicorn's own default (None) waits forever for an open connection to close +# on stop() -- fatal for a caller holding one open on purpose (a live +# message:stream response, an MCP client's long-lived /sse GET). Bound it so +# stop() force-closes the connection instead of hanging indefinitely. +SERVER_STOP_TIMEOUT_S = 5 + +# Process-global footgun -- see module docstring. Disabled once, on import. +AppStatus.disable_automatic_graceful_drain() + + +async def wait_until_started( + server: uvicorn.Server, + serve_task: asyncio.Task[object], + *, + timeout_s: float, +) -> None: + """Block until ``server`` reports ready. + + Polls ``server.started`` since ``serve_task`` only returns once the + server stops. A task that ends first -- raising or not -- means the + server will never start, so that's fatal immediately rather than + busy-waited to the timeout. + """ + deadline = asyncio.get_running_loop().time() + timeout_s + while not server.started: + if serve_task.done(): + await serve_task # re-raises if the task itself failed + raise RuntimeError("uvicorn server task ended before ever starting") + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError( + f"uvicorn server did not report ready within {timeout_s}s" + ) + await asyncio.sleep(POLL_INTERVAL_S) + + +class ManagedUvicornServer: + """Runs one ASGI app on a background uvicorn server. + + Starts it, waits for readiness, tears it down -- no knowledge of what + the app is or does. + """ + + def __init__( + self, + app: ASGIApp, + *, + host: str, + port: int, + start_timeout_s: float = SERVER_START_TIMEOUT_S, + stop_timeout_s: int = SERVER_STOP_TIMEOUT_S, + ) -> None: + self._app = app + self._host = host + self._port = port + self._start_timeout_s = start_timeout_s + self._stop_timeout_s = stop_timeout_s + self._server: uvicorn.Server | None = None + self._task: asyncio.Task[None] | None = None + + @property + def bound_port(self) -> int: + """The actual listening port -- resolves ``port=0`` to whatever the + OS assigned.""" + if self._server is None: + raise RuntimeError("server has not started") + return self._server.servers[0].sockets[0].getsockname()[1] + + async def start(self) -> None: + server = uvicorn.Server( + uvicorn.Config( + self._app, + host=self._host, + port=self._port, + log_level="warning", + timeout_graceful_shutdown=self._stop_timeout_s, + ) + ) + task = asyncio.create_task(server.serve()) + self._server = server + self._task = task + try: + await wait_until_started(server, task, timeout_s=self._start_timeout_s) + except BaseException: + # A failed/timed-out startup still leaves the task running and + # the socket bound; stop() unwinds both. + await self.stop() + raise + + async def stop(self) -> None: + if self._server is None or self._task is None: + return + # Ask uvicorn to exit rather than cancelling serve(): cancellation + # skips its shutdown phase and leaks the listening socket. + self._server.should_exit = True + try: + await self._task + except asyncio.CancelledError: + raise + except BaseException: # uvicorn raises SystemExit on startup failure + logger.exception("Embedded uvicorn server exited with error") + self._server = None + self._task = None diff --git a/tests/e2e/baseline/README.md b/tests/e2e/baseline/README.md index 150dbb968..ee239d5cc 100644 --- a/tests/e2e/baseline/README.md +++ b/tests/e2e/baseline/README.md @@ -239,7 +239,7 @@ These three are why the toolkit is shaped the way it is — keep them when exten | `smoke/matrix/` | runs across the adapter matrix: `test_adapter_matrix.py`, `test_capability_matrix.py` (memory store + recall), `test_context_recall.py` (in-session + rejoin), `test_rehydration_offline.py` / `test_rehydration_partial.py` (cold-boot / partial-reboot `/context` recall), `test_rehydration_cross_framework.py` (a different-framework `peer=` authors, A rehydrates), `test_room_isolation.py`, `test_noisy_room.py`, `test_tool_round_trip.py` (custom-tool subgroup) | | `smoke/behavior/` | platform/transport + scenario behavior: `test_delivery_status.py`, `test_processing_barrier.py`, `test_isolation.py`, `test_agent_scenarios.py` | | `smoke/inspection/` | `capture.*` observation worked-examples: `test_tool_calls.py`, `test_events.py`, `test_memory.py`, `test_usage.py` (the `Emit.USAGE` fan), plus `test_next_actionable_semantics.py` (a platform `/next` invariant, `@lane`-pinned since it runs no adapter) | -| `smoke/adapters/` | adapter-specific showcases: `test_agno.py`, `test_copilot_acp.py`, `test_copilot_sdk.py`, `test_crewai.py`, `test_letta.py`, `test_opencode.py`, `test_parlant.py` | +| `smoke/adapters/` | adapter-specific showcases: `test_a2a.py`, `test_a2a_gateway.py`, `test_a2a_roundtrip.py` (+ their shared `a2aServer.py` fixture -- a scripted, non-Band A2A counterparty), `test_agno.py`, `test_copilot_acp.py`, `test_copilot_sdk.py`, `test_crewai.py`, `test_letta.py`, `test_opencode.py`, `test_parlant.py` | The `toolkit/` modules are pytest-free and reusable anywhere. The package root (`settings`, `requires`, `agents`, `conftest`) is the pytest wiring. diff --git a/tests/e2e/baseline/smoke/adapters/a2aServer.py b/tests/e2e/baseline/smoke/adapters/a2aServer.py new file mode 100644 index 000000000..d5414870d --- /dev/null +++ b/tests/e2e/baseline/smoke/adapters/a2aServer.py @@ -0,0 +1,138 @@ +"""Standalone A2A counterparty server for baseline E2E smokes. + +Built directly on a2a-sdk's own server primitives (not Band's gateway), so +``test_a2a.py`` can point a live ``A2AAdapter`` at a real, independent A2A +implementation -- proving the SDK's outbound client against something other +than our own gateway. Scripted, not LLM-backed, so it needs no LLM key: a +request whose text carries ``ERROR_MARKER`` fails the task; every other +request completes with ``CANNED_REPLY``. +""" + +from __future__ import annotations + +from a2a.helpers import get_message_text, new_task_from_user_message +from a2a.server.agent_execution import AgentExecutor, RequestContext +from a2a.server.events import EventQueue +from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.routes.agent_card_routes import create_agent_card_routes +from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes +from a2a.server.tasks import InMemoryTaskStore, TaskUpdater +from a2a.types import ( + AgentCapabilities, + AgentCard, + AgentInterface, + AgentSkill, + Part, + UnsupportedOperationError, +) +from a2a.utils.constants import PROTOCOL_VERSION_CURRENT +from starlette.applications import Starlette + +from band.integrations.uvicorn_server import ( + SERVER_START_TIMEOUT_S, + SERVER_STOP_TIMEOUT_S, + ManagedUvicornServer, +) + +from tests.ports import reserve_port + +CANNED_REPLY = "a2a-fixture-canned-reply" +ERROR_MARKER = "a2a-fixture-trigger-error" + + +class ScriptedExecutor(AgentExecutor): + """A deterministic counterparty: a canned reply, or a scripted failure. + + No LLM involved -- ``ERROR_MARKER`` in the request text is the only + branch, so a smoke can trigger either path deterministically. + """ + + async def execute(self, context: RequestContext, event_queue: EventQueue) -> None: + if context.message is None: + raise ValueError("A2A request is missing its message") + task = context.current_task or new_task_from_user_message(context.message) + if context.current_task is None: + await event_queue.enqueue_event(task) + + updater = TaskUpdater(event_queue, task.id, task.context_id) + if ERROR_MARKER in get_message_text(context.message): + await updater.failed( + updater.new_agent_message([Part(text="fixture: scripted failure")]) + ) + return + await updater.complete(updater.new_agent_message([Part(text=CANNED_REPLY)])) + + async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: + raise UnsupportedOperationError() + + +class A2ACounterparty: + """A live, standalone A2A JSON-RPC server for a smoke to point an + ``A2AAdapter`` at. + + Binds an OS-assigned free port (via ``reserve_port``) so parallel runs + never collide. The port must be known before the agent card is built (the + card advertises it), so it is reserved up front in ``__init__`` rather + than left to uvicorn to pick at ``start()`` time. + """ + + def __init__(self) -> None: + self.port = reserve_port() + self._runtime: ManagedUvicornServer | None = None + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + def _agent_card(self) -> AgentCard: + return AgentCard( + name="A2A Smoke Fixture", + description="Deterministic A2A counterparty for baseline E2E smokes.", + supported_interfaces=[ + AgentInterface( + protocol_binding="JSONRPC", + protocol_version=PROTOCOL_VERSION_CURRENT, + url=self.url, + ), + ], + version="1.0.0", + capabilities=AgentCapabilities(streaming=True), + skills=[ + AgentSkill( + id="default", + name="Scripted reply", + description="Replies with a deterministic canned marker.", + tags=["smoke"], + ) + ], + default_input_modes=["text/plain"], + default_output_modes=["text/plain"], + ) + + def _build_app(self) -> Starlette: + card = self._agent_card() + handler = DefaultRequestHandler( + agent_executor=ScriptedExecutor(), + task_store=InMemoryTaskStore(), + agent_card=card, + ) + routes = create_agent_card_routes(card) + create_jsonrpc_routes( + handler, rpc_url="/", enable_v0_3_compat=True + ) + return Starlette(routes=routes) + + async def start(self) -> None: + self._runtime = ManagedUvicornServer( + self._build_app(), + host="127.0.0.1", + port=self.port, + start_timeout_s=SERVER_START_TIMEOUT_S, + stop_timeout_s=SERVER_STOP_TIMEOUT_S, + ) + await self._runtime.start() + + async def stop(self) -> None: + if self._runtime is None: + return + await self._runtime.stop() + self._runtime = None diff --git a/tests/e2e/baseline/smoke/adapters/test_a2a.py b/tests/e2e/baseline/smoke/adapters/test_a2a.py new file mode 100644 index 000000000..472a48f12 --- /dev/null +++ b/tests/e2e/baseline/smoke/adapters/test_a2a.py @@ -0,0 +1,107 @@ +"""A2AAdapter showcase smoke -- a live A2AAdapter driven against +``a2aServer.A2ACounterparty``, a minimal scripted A2A server (not Band's +own gateway), proving the outbound adapter against an independent +implementation. Deterministic, not LLM-backed, so neither side needs an +LLM key. + +A2A is a protocol bridge, not an LLM-agent adapter (``NON_AGENT_ADAPTERS``), +so this is a bespoke, non-matrix smoke like ``test_parlant.py``: the +adapter is built directly and handed to ``running_provisioned_agent`` so +provisioning, capture, and reaping share the same plumbing as every other +baseline test. + +Run with: + E2E_TESTS_ENABLED=true uv run pytest \\ + tests/e2e/baseline/smoke/adapters/test_a2a.py -v -s --no-cov +""" + +from __future__ import annotations + +import pytest + +from band.integrations.a2a import A2AAdapter + +from tests.e2e.baseline.agents import Lane, lane +from tests.e2e.baseline.settings import BaselineSettings +from tests.e2e.baseline.smoke.adapters.a2aServer import ( + CANNED_REPLY, + ERROR_MARKER, + A2ACounterparty, +) +from tests.e2e.baseline.toolkit.capture import CaptureFactory +from tests.e2e.baseline.toolkit.provisioning import ( + ResourceManager, + running_provisioned_agent, +) +from tests.e2e.baseline.toolkit.user_ops import UserOps +from tests.lifecycle import running + + +# Not in the adapter registry, so the lane selector can't derive a home +# lane and would run it in every lane. Pin to core -- needs no provider key +# (the counterparty is scripted), only the always-on Band-platform gate. +@lane(Lane.CORE) +@pytest.mark.timeout(extra=60) +@pytest.mark.asyncio(loop_scope="session") +async def test_a2a_adapter_relays_a_real_counterparty_reply( + resource_manager: ResourceManager, + user_ops: UserOps, + reply_capture: CaptureFactory, + baseline_settings: BaselineSettings, +) -> None: + """A live ``A2AAdapter`` forwards a Band room message to a real, + independent A2A server and relays its reply back into the room.""" + async with running(A2ACounterparty()) as counterparty: + adapter = A2AAdapter(remote_url=counterparty.url, streaming=True) + async with running_provisioned_agent( + adapter, resource_manager, label="a2a" + ) as agent: + room_id = await resource_manager.provision_room( + title="e2e-a2a-reply", participants=[agent.id] + ) + async with reply_capture(room_id) as capture: + mid = await user_ops.send_message( + room_id, + "Please say hello.", + mention_id=agent.id, + mention_name=agent.name, + ) + replies = await capture.wait_for_reply( + mid, agent.id, deadline_s=baseline_settings.e2e_timeout + ) + + replies.assert_contains_any([CANNED_REPLY]) + + +@lane(Lane.CORE) +@pytest.mark.timeout(extra=60) +@pytest.mark.asyncio(loop_scope="session") +async def test_a2a_adapter_surfaces_a_remote_task_failure( + resource_manager: ResourceManager, + user_ops: UserOps, + reply_capture: CaptureFactory, + baseline_settings: BaselineSettings, +) -> None: + """A terminal FAILED task from the remote A2A server surfaces as a room + error event, not a silently dropped turn.""" + async with running(A2ACounterparty()) as counterparty: + adapter = A2AAdapter(remote_url=counterparty.url, streaming=True) + async with running_provisioned_agent( + adapter, resource_manager, label="a2a" + ) as agent: + room_id = await resource_manager.provision_room( + title="e2e-a2a-failure", participants=[agent.id] + ) + async with reply_capture(room_id) as capture: + mid = await user_ops.send_message( + room_id, + f"trigger a scripted failure: {ERROR_MARKER}", + mention_id=agent.id, + mention_name=agent.name, + ) + await capture.wait_for_processed( + mid, agent.id, deadline_s=baseline_settings.e2e_timeout + ) + errors = await capture.errors(sender_id=agent.id) + + errors.assert_present() diff --git a/tests/e2e/baseline/smoke/adapters/test_a2a_gateway.py b/tests/e2e/baseline/smoke/adapters/test_a2a_gateway.py new file mode 100644 index 000000000..f49a86a45 --- /dev/null +++ b/tests/e2e/baseline/smoke/adapters/test_a2a_gateway.py @@ -0,0 +1,84 @@ +"""A2AGatewayAdapter showcase smoke -- driving a live gateway with the +official a2a-sdk reference client, independent of ``A2AAdapter``. + +A2A is a protocol bridge, not an LLM-agent adapter (listed in +``NON_AGENT_ADAPTERS``), so the gateway itself is built bespoke and handed to +``running_provisioned_agent``, like ``test_parlant.py``. The *target* peer it +exposes, though, is an ordinary Anthropic-backed Band agent provisioned +through ``@with_adapters`` -- which also derives this smoke's home lane +(``core``), so no ``@lane`` pin is needed here. + +Drives the gateway with a real ``a2a.client.Client`` (the a2a-sdk reference +client), not our own ``A2AAdapter`` -- this validates the gateway's JSON-RPC +server against the actual upstream implementation, independent of any bug +the two could otherwise share. + +Run with: + E2E_TESTS_ENABLED=true uv run pytest \\ + tests/e2e/baseline/smoke/adapters/test_a2a_gateway.py -v -s --no-cov +""" + +from __future__ import annotations + +import httpx +import pytest +from a2a.client import ClientConfig, ClientFactory +from a2a.helpers import get_message_text, new_text_message +from a2a.types import Role, SendMessageRequest + +from band.integrations.a2a.adapter import _SSE_READ_TIMEOUT_S +from band.integrations.a2a.gateway import A2AGatewayAdapter + +from tests.e2e.baseline.agents import Adapter, with_adapters +from tests.e2e.baseline.flaky import flaky_infra +from tests.e2e.baseline.toolkit.provisioning import ( + ProvisionedAgent, + ResourceManager, + running_provisioned_agent, +) +from tests.ports import reserve_port + +_SHORT = "You are a friendly assistant in a chat room. Reply in one short sentence." + + +@with_adapters(Adapter.ANTHROPIC, prompt=_SHORT) +@flaky_infra("retry a transient live-turn timeout; assertion failures fail loud") +@pytest.mark.timeout(extra=120) +@pytest.mark.asyncio(loop_scope="session") +async def test_gateway_serves_a_real_a2a_client( + agent: ProvisionedAgent, + resource_manager: ResourceManager, +) -> None: + """A raw a2a-sdk client drives the gateway's JSON-RPC endpoint for a live + Band peer and receives its real reply back over A2A.""" + port = reserve_port() + gateway = A2AGatewayAdapter(gateway_url=f"http://127.0.0.1:{port}", port=port) + + async with running_provisioned_agent(gateway, resource_manager, label="a2a-gw"): + # The Anthropic peer's own id is a stable alias the gateway always + # serves (alongside its slug), so the client needs no slug lookup. + # Same generous-but-bounded read timeout as A2AAdapter itself (see + # its module docstring): httpx's default 5s fires on the normal, + # multi-second gap between SSE events during a real LLM turn. + http_client = httpx.AsyncClient( + timeout=httpx.Timeout(10.0, read=_SSE_READ_TIMEOUT_S) + ) + factory = ClientFactory(ClientConfig(streaming=True, httpx_client=http_client)) + client = await factory.create_from_url( + f"http://127.0.0.1:{port}/agents/{agent.id}" + ) + reply_text = "" + try: + message = new_text_message("Please say hello.", role=Role.ROLE_USER) + async for event in client.send_message(SendMessageRequest(message=message)): + if event.HasField( + "status_update" + ) and event.status_update.status.HasField("message"): + text = get_message_text(event.status_update.status.message) + if text: + reply_text = text + finally: + await client.close() + await http_client.aclose() + + assert reply_text, "expected a reply relayed from the live Band peer over A2A" diff --git a/tests/e2e/baseline/smoke/adapters/test_a2a_roundtrip.py b/tests/e2e/baseline/smoke/adapters/test_a2a_roundtrip.py new file mode 100644 index 000000000..56f3f0c6e --- /dev/null +++ b/tests/e2e/baseline/smoke/adapters/test_a2a_roundtrip.py @@ -0,0 +1,86 @@ +"""Full A2A round-trip smoke: Band Agent A -> ``A2AAdapter`` -> a live +gateway -> Band Agent B (Anthropic), and the reply flowing all the way back. + +Reuses ``test_a2a_gateway.py``'s target-peer + gateway setup, but swaps the +raw a2a-sdk reference client for a live ``A2AAdapter`` as the caller -- +proving the realistic Band-to-Band usage pattern end to end: a Band room +message forwarded over real A2A JSON-RPC to a gateway-exposed Band peer, with +that peer's real LLM reply relayed all the way back into Agent A's own room. + +``@with_adapters(Adapter.ANTHROPIC)`` (for the target peer B) also derives +this smoke's home lane (``core``), so no ``@lane`` pin is needed -- the +gateway and the caller ``A2AAdapter`` are themselves protocol bridges +(``NON_AGENT_ADAPTERS``) built bespoke, same as ``test_a2a_gateway.py``. + +Run with: + E2E_TESTS_ENABLED=true uv run pytest \\ + tests/e2e/baseline/smoke/adapters/test_a2a_roundtrip.py -v -s --no-cov +""" + +from __future__ import annotations + +import pytest + +from band.integrations.a2a import A2AAdapter +from band.integrations.a2a.gateway import A2AGatewayAdapter + +from tests.e2e.baseline.agents import Adapter, with_adapters +from tests.e2e.baseline.flaky import flaky_infra +from tests.e2e.baseline.settings import BaselineSettings +from tests.e2e.baseline.toolkit.capture import CaptureFactory +from tests.e2e.baseline.toolkit.provisioning import ( + ProvisionedAgent, + ResourceManager, + running_provisioned_agent, +) +from tests.e2e.baseline.toolkit.user_ops import UserOps +from tests.ports import reserve_port + +_SHORT = "You are a friendly assistant in a chat room. Reply in one short sentence." + + +@with_adapters(Adapter.ANTHROPIC, prompt=_SHORT) +@flaky_infra("retry a transient live-turn timeout; assertion failures fail loud") +# Two relayed hops (Agent A -> gateway -> Agent B) plus one real LLM turn on +# top of a gateway + two adapter cold starts: 240s outer, leaving 120s +# overhead beyond the 2x e2e_timeout barrier deadline below. +@pytest.mark.timeout(extra=240) +@pytest.mark.asyncio(loop_scope="session") +async def test_band_to_band_round_trip_over_real_a2a( + agent: ProvisionedAgent, + resource_manager: ResourceManager, + user_ops: UserOps, + reply_capture: CaptureFactory, + baseline_settings: BaselineSettings, +) -> None: + """Band Agent A relays a room message through a live A2A gateway to Band + Agent B (a real Anthropic turn) and posts B's reply back into A's room.""" + port = reserve_port() + gateway = A2AGatewayAdapter(gateway_url=f"http://127.0.0.1:{port}", port=port) + + async with running_provisioned_agent(gateway, resource_manager, label="a2a-gw"): + caller = A2AAdapter( + remote_url=f"http://127.0.0.1:{port}/agents/{agent.id}", streaming=True + ) + async with running_provisioned_agent( + caller, resource_manager, label="a2a-caller" + ) as caller_agent: + room_id = await resource_manager.provision_room( + title="e2e-a2a-roundtrip", participants=[caller_agent.id] + ) + async with reply_capture(room_id) as capture: + mid = await user_ops.send_message( + room_id, + "Please say hello.", + mention_id=caller_agent.id, + mention_name=caller_agent.name, + ) + replies = await capture.wait_for_reply( + mid, + caller_agent.id, + deadline_s=baseline_settings.e2e_timeout * 2, + ) + + replies.assert_present( + what="a reply relayed over a live Band-to-Band A2A round trip" + ) diff --git a/tests/integrations/a2a/gateway/helpers.py b/tests/integrations/a2a/gateway/helpers.py index 1ff13dc3f..9e1149b27 100644 --- a/tests/integrations/a2a/gateway/helpers.py +++ b/tests/integrations/a2a/gateway/helpers.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + from band_rest import Peer @@ -16,3 +18,9 @@ def make_peer(peer_id: str, name: str, description: str = "") -> Peer: is_contact=False, source="registry", ) + + +def peers_page(peers: list[Peer]) -> SimpleNamespace: + """A fake ``list_agent_peers`` response page -- only ``.data`` matters + to ``_fetch_all_peers``, which pages until a page comes back short.""" + return SimpleNamespace(data=peers) diff --git a/tests/integrations/a2a/gateway/test_adapter.py b/tests/integrations/a2a/gateway/test_adapter.py index 2cb1169d7..b108cc5e9 100644 --- a/tests/integrations/a2a/gateway/test_adapter.py +++ b/tests/integrations/a2a/gateway/test_adapter.py @@ -26,7 +26,7 @@ from band.integrations.a2a.gateway.adapter import BandAgentExecutor from band.integrations.a2a.gateway.types import GatewaySessionState, PendingA2ATask from band.testing import FakeAgentTools -from tests.integrations.a2a.gateway.helpers import make_peer +from tests.integrations.a2a.gateway.helpers import make_peer, peers_page def make_platform_message( @@ -56,10 +56,18 @@ def make_request(content: str = "What is the weather?") -> RequestContext: return RequestContext(None, request=SendMessageRequest(message=message)) -def configure_room_creation(adapter: A2AGatewayAdapter) -> None: +def room_creation_response(room_id: str) -> MagicMock: response = MagicMock() - response.data.id = "room-123" - adapter._rest.agent_api_chats.create_agent_chat = AsyncMock(return_value=response) + response.data.id = room_id + return response + + +def configure_room_creation( + adapter: A2AGatewayAdapter, *, room_id: str = "room-123" +) -> None: + adapter._rest.agent_api_chats.create_agent_chat = AsyncMock( + return_value=room_creation_response(room_id) + ) adapter._rest.agent_api_participants.add_agent_chat_participant = AsyncMock() adapter._rest.agent_api_messages.create_agent_chat_message = AsyncMock() adapter._rest.agent_api_events.create_agent_chat_event = AsyncMock() @@ -97,10 +105,8 @@ class TestGatewayStartup: @pytest.mark.asyncio async def test_discovers_peers_and_starts_server(self) -> None: adapter = A2AGatewayAdapter(rest_client=MagicMock()) - response = MagicMock() - response.data = [make_peer("weather", "Weather Agent")] adapter._rest.agent_api_peers.list_agent_peers = AsyncMock( - return_value=response + return_value=peers_page([make_peer("weather", "Weather Agent")]) ) with patch( @@ -268,6 +274,57 @@ async def test_timeout_returns_terminal_failure( "A2A request completed" in record.message for record in caplog.records ), "a timed-out request must not be logged as completed" + @pytest.mark.asyncio + async def test_send_failure_publishes_terminal_failure(self) -> None: + """A REST failure while posting to Band must not leave the remote + A2A caller waiting on a stuck WORKING task.""" + adapter = A2AGatewayAdapter(rest_client=MagicMock()) + adapter._peers = {"weather": make_peer("weather", "Weather Agent")} + configure_room_creation(adapter) + adapter._rest.agent_api_messages.create_agent_chat_message = AsyncMock( + side_effect=RuntimeError("Band unavailable") + ) + queue = EventQueueLegacy() + + with pytest.raises(RuntimeError, match="Band unavailable"): + await BandAgentExecutor(adapter, "weather").execute(make_request(), queue) + + initial = await queue.dequeue_event() + terminal = await queue.dequeue_event() + assert initial.status.state == TaskState.TASK_STATE_WORKING + assert terminal.status.state == TaskState.TASK_STATE_FAILED + assert terminal.status.message.parts[0].text == "A2A request failed" + assert "Band unavailable" not in terminal.status.message.parts[0].text + assert adapter._pending_tasks == {} + + @pytest.mark.asyncio + async def test_establish_request_raises_when_peer_missing(self) -> None: + adapter = A2AGatewayAdapter(rest_client=MagicMock()) + + with pytest.raises(ValueError, match="Peer not found"): + await adapter._establish_request( + "missing", make_request(), EventQueueLegacy() + ) + + @pytest.mark.asyncio + async def test_fetch_all_peers_accumulates_across_pages(self) -> None: + adapter = A2AGatewayAdapter(rest_client=MagicMock()) + full_page = [make_peer(f"peer-{i}", f"Peer {i}") for i in range(100)] + partial_page = [make_peer("peer-100", "Peer 100")] + adapter._rest.agent_api_peers.list_agent_peers = AsyncMock( + side_effect=[peers_page(full_page), peers_page(partial_page)] + ) + + peers = await adapter._fetch_all_peers() + + assert len(peers) == 101 + assert adapter._rest.agent_api_peers.list_agent_peers.await_count == 2 + first_call, second_call = ( + adapter._rest.agent_api_peers.list_agent_peers.call_args_list + ) + assert first_call.kwargs["page"] == 1 + assert second_call.kwargs["page"] == 2 + @pytest.mark.asyncio async def test_cleanup_all_stops_the_hosted_server(self) -> None: """Agent.stop() reaches the adapter only via cleanup_all, so the @@ -338,12 +395,7 @@ def adapter(self) -> A2AGatewayAdapter: "weather": make_peer("weather", "Weather Agent"), "data": make_peer("data", "Data Agent"), } - response = MagicMock() - response.data.id = "new-room" - adapter._rest.agent_api_chats.create_agent_chat = AsyncMock( - return_value=response - ) - adapter._rest.agent_api_participants.add_agent_chat_participant = AsyncMock() + configure_room_creation(adapter, room_id="new-room") return adapter @pytest.mark.asyncio @@ -366,13 +418,11 @@ async def test_context_reuses_room_and_adds_new_peer( async def test_different_contexts_get_different_rooms( self, adapter: A2AGatewayAdapter ) -> None: - responses = [] - for room_id in ("room-a", "room-b"): - response = MagicMock() - response.data.id = room_id - responses.append(response) adapter._rest.agent_api_chats.create_agent_chat = AsyncMock( - side_effect=responses + side_effect=[ + room_creation_response("room-a"), + room_creation_response("room-b"), + ] ) room_a, _ = await adapter._get_or_create_room("ctx-a", "weather") @@ -382,6 +432,23 @@ async def test_different_contexts_get_different_rooms( "distinct A2A contexts must not share a Band room" ) + @pytest.mark.asyncio + async def test_participant_add_failure_leaves_no_partial_room_state( + self, adapter: A2AGatewayAdapter + ) -> None: + """Regression coverage: a REST failure after room creation currently + leaves no context/room mapping behind, so a retry creates a brand + new room rather than reusing the one that was just orphaned.""" + adapter._rest.agent_api_participants.add_agent_chat_participant = AsyncMock( + side_effect=RuntimeError("participant add failed") + ) + + with pytest.raises(RuntimeError, match="participant add failed"): + await adapter._get_or_create_room("ctx", "weather") + + assert adapter._context_to_room == {} + assert adapter._room_participants == {} + def test_rehydrate_merges_without_overwriting_live_context(self) -> None: adapter = A2AGatewayAdapter(rest_client=MagicMock()) adapter._context_to_room["ctx"] = "live-room" diff --git a/tests/integrations/a2a/gateway/test_server.py b/tests/integrations/a2a/gateway/test_server.py index 4b019e0f5..c6084aca4 100644 --- a/tests/integrations/a2a/gateway/test_server.py +++ b/tests/integrations/a2a/gateway/test_server.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from contextlib import suppress from uuid import uuid4 @@ -17,16 +17,9 @@ from a2a.utils.constants import PROTOCOL_VERSION_0_3 from httpx import ASGITransport -# Side effect, not used directly: importing this module disables -# sse_starlette's automatic graceful drain process-wide (see its own -# AppStatus.disable_automatic_graceful_drain() call) -- the exact real-world -# coexistence (an ACP/opencode backend in the same process as this gateway) -# that test_stop_returns_promptly_with_a_still_open_message_stream guards -# against. Imported explicitly so the test is deterministic regardless of -# whether some other test file happened to import it first. -import band.integrations.mcp.local_server # noqa: F401 from band.integrations.a2a.gateway.server import SERVER_STOP_TIMEOUT_S, GatewayServer from tests.integrations.a2a.gateway.helpers import make_peer +from tests.lifecycle import elapsed, held_open, running class FakeExecutor(AgentExecutor): @@ -50,13 +43,17 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None raise NotImplementedError -def build_server() -> GatewayServer: +def build_server( + *, + port: int = 10000, + executor_factory: Callable[[str], AgentExecutor] | None = None, +) -> GatewayServer: peer = make_peer("uuid-weather", "Weather Agent", "Gets weather info") return GatewayServer( peers={"weather-agent": peer}, - gateway_url="http://localhost:10000", - port=10000, - executor_factory=lambda _slug: FakeExecutor(), + gateway_url=f"http://localhost:{port}", + port=port, + executor_factory=executor_factory or (lambda _slug: FakeExecutor()), ) @@ -182,6 +179,79 @@ async def test_jsonrpc_method_errors_are_upstream_owned( assert response.json()["error"]["code"] == -32601 +async def test_malformed_json_body_falls_through_to_upstream_dispatch( + gateway_client: httpx.AsyncClient, +) -> None: + """A body ``request.json()`` can't parse degrades ``method`` to ``None``, + which skips the closed-method guard entirely -- the upstream dispatcher + owns reporting the parse error, same as any other non-blocked method.""" + response = await gateway_client.post( + "/agents/weather-agent", + headers={"A2A-Version": "1.0", "Content-Type": "application/json"}, + content=b"{not valid json", + ) + + assert response.status_code == 200 + assert response.json()["error"]["code"] == -32700 + + +async def test_non_scalar_request_id_is_normalized_to_null( + gateway_client: httpx.AsyncClient, +) -> None: + response = await gateway_client.post( + "/agents/weather-agent", + headers={"A2A-Version": "1.0"}, + json={ + "jsonrpc": "2.0", + "id": [1, 2, 3], + "method": "ListTasks", + "params": {}, + }, + ) + + body = response.json() + assert body["id"] is None, "a non-str/int id must not be echoed back verbatim" + assert body["error"]["code"] == -32601 + + +async def test_send_streaming_message_runs_through_official_handler( + gateway_client: httpx.AsyncClient, +) -> None: + response = await gateway_client.post( + "/agents/weather-agent", + headers={"A2A-Version": "1.0"}, + json={ + "jsonrpc": "2.0", + "id": "request-1", + "method": "SendStreamingMessage", + "params": { + "message": { + "role": "ROLE_USER", + "messageId": "message-1", + "parts": [{"text": "Hello"}], + } + }, + }, + ) + + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + + +@pytest.mark.parametrize("method", ["GetTask", "CancelTask"]) +async def test_task_methods_without_id_are_rejected_by_upstream_handler( + gateway_client: httpx.AsyncClient, method: str +) -> None: + response = await gateway_client.post( + "/agents/weather-agent", + headers={"A2A-Version": "1.0"}, + json={"jsonrpc": "2.0", "id": "request-1", "method": method, "params": {}}, + ) + + assert response.status_code == 200 + assert response.json()["error"]["code"] == -32602 + + async def test_jsonrpc_send_runs_through_official_handler_and_executor( gateway_client: httpx.AsyncClient, ) -> None: @@ -337,6 +407,88 @@ async def test_v03_jsonrpc_stream_accepts_legacy_payload( assert "text/event-stream" in response.headers["content-type"] +async def test_start_returns_only_once_the_server_is_listening() -> None: + """A caller dialing in right after ``on_started()`` returns (e.g. a real + A2A client, or one of the E2E smokes) must not race a socket that isn't + accepting connections yet.""" + async with running(build_server(port=0)) as server: + async with httpx.AsyncClient() as client: + response = await client.get(f"http://127.0.0.1:{server.bound_port}/peers") + assert response.status_code == 200 + + +class DelayedTwoStepExecutor(AgentExecutor): + """Task, then a working update, then (after a pause) completion -- three + distinct SSE events, so a stream cut short is distinguishable from a + healthy one.""" + + async def execute(self, context: RequestContext, event_queue: EventQueue) -> None: + if context.message is None: + raise ValueError("A2A request is missing its message") + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + await event_queue.enqueue_event( + TaskStatusUpdateEvent( + task_id=task.id, + context_id=task.context_id, + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ) + await asyncio.sleep(0.3) + await event_queue.enqueue_event( + TaskStatusUpdateEvent( + task_id=task.id, + context_id=task.context_id, + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + ) + ) + + async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: + raise NotImplementedError + + +async def test_a_second_server_is_not_poisoned_by_a_prior_servers_shutdown() -> None: + """Regression: a second GatewayServer's live stream, opened only after a + first one has stopped in the same process, must still deliver every + event.""" + async with running(build_server(port=0)) as first: + port1 = first.bound_port + async with httpx.AsyncClient(timeout=None) as client: + async with client.stream( + "POST", + f"http://127.0.0.1:{port1}/agents/weather-agent/message:stream", + headers={"A2A-Version": "1.0"}, + json=hello_message_body(), + ) as response: + async for _ in response.aiter_bytes(): + pass + + second = GatewayServer( + peers={"other-agent": make_peer("uuid-other", "Other Agent", "")}, + gateway_url="http://localhost:0", + port=0, + executor_factory=lambda _slug: DelayedTwoStepExecutor(), + ) + async with running(second): + port2 = second.bound_port + events: list[str] = [] + async with httpx.AsyncClient(timeout=None) as client: + async with client.stream( + "POST", + f"http://127.0.0.1:{port2}/agents/other-agent/message:stream", + headers={"A2A-Version": "1.0"}, + json=hello_message_body(), + ) as response: + async for line in response.aiter_lines(): + if line.startswith("data:"): + events.append(line) + + assert len(events) >= 3, ( + f"got {len(events)} events, expected 3 (task, working, completed) -- " + "the second server's stream was cut short by the first server's shutdown" + ) + + class NeverFinishingExecutor(AgentExecutor): """Enqueues one event, then never returns -- holding the SSE response open indefinitely, the way a real long-running agent task would.""" @@ -357,71 +509,43 @@ async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None @pytest.mark.timeout(SERVER_STOP_TIMEOUT_S + 15.0) async def test_stop_returns_promptly_with_a_still_open_message_stream() -> None: - """Regression: sse_starlette's cooperative shutdown drain is a process- - global switch that band.integrations.mcp.local_server permanently - disables the moment it's imported anywhere in the process -- a real - coexistence scenario (an ACP/opencode backend sharing the process with - this gateway). A live message:stream connection then has no other way - to end on its own, so stop() must bound its wait via - timeout_graceful_shutdown instead of hanging forever. - - Measures wall-clock time around a bare ``await server.stop()`` (no - wrapping ``asyncio.wait_for``, which would cancel ``stop()`` from the - outside and mask a real hang as a false pass) -- same rationale as - LocalMCPServer's own equivalent regression test. + """Regression: disabling sse_starlette's drain means a live + message:stream connection has no other way to end -- stop() must bound + its wait via timeout_graceful_shutdown instead of hanging forever. + + Measures wall-clock time around a bare ``server.stop()`` (wrapping it + in ``asyncio.wait_for`` would cancel it externally and mask a real hang). """ - peer = make_peer("uuid-weather", "Weather Agent", "Gets weather info") - server = GatewayServer( - peers={"weather-agent": peer}, - gateway_url="http://localhost:0", - port=0, - executor_factory=lambda _slug: NeverFinishingExecutor(), + server = build_server( + port=0, executor_factory=lambda _slug: NeverFinishingExecutor() ) - await server.start() - # start() doesn't wait for uvicorn's own startup phase to finish -- it - # only schedules serve() as a background task. Poll for it directly - # since GatewayServer exposes no readiness signal of its own. - for _ in range(50): - if server._uvicorn.started: - break - await asyncio.sleep(0.05) - port = server._uvicorn.servers[0].sockets[0].getsockname()[1] - - connection_ready = asyncio.Event() - - async def hold_connection_open() -> None: - with suppress(Exception): - # timeout=None: httpx's default 5s read timeout would otherwise - # give up waiting for the next chunk and disconnect on its own - # around the same mark as SERVER_STOP_TIMEOUT_S -- masking a real - # server-side hang as a false pass, since the connection would - # end for the wrong reason (a bored client) rather than proving - # stop() itself is bounded. - async with ( - httpx.AsyncClient(timeout=None) as client, - client.stream( - "POST", - f"http://127.0.0.1:{port}/agents/weather-agent/message:stream", - headers={"A2A-Version": "1.0"}, - json=hello_message_body(), - ) as response, - ): - async for _ in response.aiter_bytes(): - connection_ready.set() - - holder = asyncio.create_task(hold_connection_open()) - try: - await asyncio.wait_for(connection_ready.wait(), timeout=5.0) - - started_at = asyncio.get_running_loop().time() - await server.stop() - elapsed = asyncio.get_running_loop().time() - started_at - - assert elapsed < SERVER_STOP_TIMEOUT_S + 5.0, ( - f"stop() took {elapsed:.1f}s -- graceful shutdown is not " + async with running(server): + port = server.bound_port + + async def connect(ready: asyncio.Event) -> None: + with suppress(Exception): + # timeout=None: httpx's default 5s read timeout would otherwise + # give up waiting for the next chunk and disconnect on its own + # around the same mark as SERVER_STOP_TIMEOUT_S -- masking a + # real server-side hang as a false pass, since the connection + # would end for the wrong reason (a bored client) rather than + # proving stop() itself is bounded. + async with ( + httpx.AsyncClient(timeout=None) as client, + client.stream( + "POST", + f"http://127.0.0.1:{port}/agents/weather-agent/message:stream", + headers={"A2A-Version": "1.0"}, + json=hello_message_body(), + ) as response, + ): + async for _ in response.aiter_bytes(): + ready.set() + + async with held_open(connect): + stop_elapsed = await elapsed(server.stop()) + + assert stop_elapsed < SERVER_STOP_TIMEOUT_S + 5.0, ( + f"stop() took {stop_elapsed:.1f}s -- graceful shutdown is not " "bounded by SERVER_STOP_TIMEOUT_S" ) - finally: - holder.cancel() - with suppress(asyncio.CancelledError): - await holder diff --git a/tests/integrations/a2a/test_adapter.py b/tests/integrations/a2a/test_adapter.py index 95ff32005..1d7ed508e 100644 --- a/tests/integrations/a2a/test_adapter.py +++ b/tests/integrations/a2a/test_adapter.py @@ -2,6 +2,8 @@ from __future__ import annotations +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 @@ -22,6 +24,7 @@ from band.core.types import PlatformMessage from band.integrations.a2a import A2AAdapter, A2AAuth, A2ASessionState +from band.integrations.a2a.adapter import _SSE_READ_TIMEOUT_S from band.testing import FakeAgentTools @@ -99,6 +102,25 @@ async def stream(*events: StreamResponse): yield event +@asynccontextmanager +async def started_adapter( + adapter: A2AAdapter, +) -> AsyncIterator[tuple[MagicMock, MagicMock]]: + """Start ``adapter`` against a patched ``ClientFactory`` and clean it up + afterward -- yields ``(client, factory_type)`` so a test states only its + own setup and assertions, not the patch/cleanup dance.""" + client = MagicMock() + with patch("band.integrations.a2a.adapter.ClientFactory") as factory_type: + factory = factory_type.return_value + factory.create_from_url = AsyncMock(return_value=client) + await adapter.on_started("Agent", "Description") + try: + yield client, factory_type + finally: + client.close = AsyncMock() + await adapter.cleanup_all() + + class TestA2AAuth: def test_to_headers_combines_authentication_methods(self) -> None: auth = A2AAuth( @@ -121,26 +143,31 @@ async def test_creates_client_with_auth_headers(self) -> None: remote_url="http://localhost:10000", auth=A2AAuth(api_key="key"), ) - client = MagicMock() - - with patch("band.integrations.a2a.adapter.ClientFactory") as factory_type: - factory = factory_type.return_value - factory.create_from_url = AsyncMock(return_value=client) - await adapter.on_started("Agent", "Description") + async with started_adapter(adapter) as (client, factory_type): + assert adapter._client is client + config = factory_type.call_args.args[0] + assert config.streaming is True + assert adapter._http_client is not None + assert adapter._http_client.headers["X-API-Key"] == "key" + assert config.httpx_client is adapter._http_client, ( + "the factory must receive the adapter's own client — this " + "identity is what carries auth to card resolution and every " + "A2A request" + ) - assert adapter._client is client - config = factory_type.call_args.args[0] - assert config.streaming is True - assert adapter._http_client is not None - assert adapter._http_client.headers["X-API-Key"] == "key" - assert config.httpx_client is adapter._http_client, ( - "the factory must receive the adapter's own client — this identity " - "is what carries auth to card resolution and every A2A request" - ) + @pytest.mark.asyncio + async def test_owned_http_client_has_a_generous_bounded_read_timeout(self) -> None: + """A real remote turn (a live LLM call, a tool loop) routinely leaves + several seconds of silence between SSE events -- httpx's 5s default + read timeout would misreport that as a dead connection. The bound + must still be finite, though, so a peer that hangs after accepting + the connection fails the turn instead of blocking the room forever.""" + adapter = A2AAdapter(remote_url="http://localhost:10000") - client.close = AsyncMock() - await adapter.cleanup_all() + async with started_adapter(adapter): + assert adapter._http_client is not None + assert adapter._http_client.timeout.read == _SSE_READ_TIMEOUT_S class TestA2AAdapterMessageFlow: @@ -283,6 +310,32 @@ async def test_terminal_task_is_finalized_even_when_band_delivery_fails( assert adapter._task_cache == {} assert adapter._task_senders == {} + @pytest.mark.asyncio + async def test_auth_required_task_is_posted_as_error_event( + self, adapter: A2AAdapter + ) -> None: + tools = FakeAgentTools() + + await adapter._handle_event( + task_event( + make_task( + TaskState.TASK_STATE_AUTH_REQUIRED, + status_message="Please authenticate", + ) + ), + tools, + "room-123", + "user-456", + "Test User", + ) + + error_events = [ + event for event in tools.events_sent if event["message_type"] == "error" + ] + assert error_events, "an auth-required task must produce an error event" + assert error_events[-1]["content"] == "Please authenticate" + assert error_events[-1]["metadata"]["a2a_state"] == "TASK_STATE_AUTH_REQUIRED" + @pytest.mark.asyncio async def test_input_required_is_forwarded_and_persisted( self, adapter: A2AAdapter @@ -441,6 +494,28 @@ async def test_cleanup_all_closes_owned_clients(self) -> None: assert adapter._client is None assert adapter._http_client is None + @pytest.mark.asyncio + async def test_cleanup_all_closes_http_transport_even_if_client_close_fails( + self, + ) -> None: + """A broken remote client must not leak the owned httpx transport.""" + adapter = A2AAdapter(remote_url="http://localhost:10000") + adapter._client = MagicMock() + adapter._client.close = AsyncMock( + side_effect=RuntimeError("client close failed") + ) + adapter._http_client = httpx.AsyncClient() + http_client = adapter._http_client + + with pytest.raises(RuntimeError, match="client close failed"): + await adapter.cleanup_all() + + assert http_client.is_closed, ( + "http transport must close even if client.close() raises" + ) + assert adapter._client is None + assert adapter._http_client is None + class TestA2AAdapterSession: @pytest.mark.asyncio diff --git a/tests/integrations/a2a/test_protocol.py b/tests/integrations/a2a/test_protocol.py index 8ca0e3a38..b6a3ce422 100644 --- a/tests/integrations/a2a/test_protocol.py +++ b/tests/integrations/a2a/test_protocol.py @@ -78,6 +78,38 @@ def test_appended_artifact_chunks_are_combined_before_response_extraction() -> N assert task_response_text(task) == "Part one. \nPart two." +def test_artifact_overwrite_replaces_existing_artifact_content() -> None: + first = StreamResponse( + artifact_update={ + "task_id": "task-123", + "context_id": "context-123", + "artifact": Artifact( + artifact_id="artifact-123", + parts=[Part(text="stale content")], + ), + "append": False, + } + ) + overwrite = StreamResponse( + artifact_update={ + "task_id": "task-123", + "context_id": "context-123", + "artifact": Artifact( + artifact_id="artifact-123", + parts=[Part(text="fresh content")], + ), + "append": False, + } + ) + + task = apply_task_stream_event(None, first) + task = apply_task_stream_event(task, overwrite) + + assert task is not None + assert len(task.artifacts) == 1 + assert task_response_text(task) == "fresh content" + + def test_task_stream_snapshot_does_not_alias_the_event() -> None: event = StreamResponse( task=Task( diff --git a/tests/integrations/mcp/test_local_server.py b/tests/integrations/mcp/test_local_server.py index 88daf8159..85377c1e5 100644 --- a/tests/integrations/mcp/test_local_server.py +++ b/tests/integrations/mcp/test_local_server.py @@ -15,7 +15,6 @@ from mcp.server.transport_security import TransportSecuritySettings from mcp.types import CallToolResult, TextContent from pydantic import BaseModel -from sse_starlette.sse import AppStatus from band.core.exceptions import BandToolError from band.integrations.mcp.engine import ( @@ -32,6 +31,8 @@ from band.runtime.custom_tools import get_custom_tool_name from band.runtime.tools import AgentTools +from tests.lifecycle import elapsed, held_open, running + class EchoInput(BaseModel): """Echo text back to the caller.""" @@ -195,31 +196,6 @@ def test_accepts_explicit_non_loopback_bind_host(self) -> None: ) assert server._host == "0.0.0.0" - def test_disables_sse_starlette_automatic_graceful_drain(self) -> None: - """Regression, traced live on Windows CI: sse_starlette's - AppStatus.should_exit is a bare process-global class attribute with - no notion of "which server" -- ANY OTHER uvicorn.Server's signal - handler firing handle_exit() anywhere in the process (not just - ours) used to latch it, closing every subsequent SSE response -- - including a fresh, healthy LocalMCPServer's that never touched that - other server -- right after its headers. Importing local_server - must disable the automatic drain so handle_exit() (the real - 2-argument signal-handler call, not our own) becomes a no-op for - this flag; original_handler is swapped out for the duration since - it expects a bound Server instance, not this direct call. - """ - assert AppStatus.enable_automatic_graceful_drain is False - - original_should_exit = AppStatus.should_exit - original_handler = AppStatus.original_handler - AppStatus.original_handler = None - try: - AppStatus.handle_exit(0, None) - assert AppStatus.should_exit is False - finally: - AppStatus.should_exit = original_should_exit - AppStatus.original_handler = original_handler - @pytest.mark.asyncio async def test_serves_sse_tools_on_localhost(self) -> None: server = LocalMCPServer( @@ -229,8 +205,7 @@ async def test_serves_sse_tools_on_localhost(self) -> None: port_max=0, ) - await server.start() - try: + async with running(server): assert server.url.startswith(f"http://{LOCAL_MCP_HOST}:") async with sse_client(server.url) as (read_stream, write_stream): @@ -238,27 +213,22 @@ async def test_serves_sse_tools_on_localhost(self) -> None: await session.initialize() await _session_lists_only_echo(session) await _call_echo(session, "hello") - finally: - await server.stop() @pytest.mark.timeout(SERVER_STOP_TIMEOUT_S + 15.0) @pytest.mark.asyncio async def test_stop_returns_promptly_with_a_still_open_sse_connection( self, ) -> None: - """Regression: an MCP client (e.g. OpenCode) holds its `/sse` GET open - for the life of its own session and may never close it after we ask - it to deregister. uvicorn's own graceful-shutdown wait is unbounded by - default, so ``stop()`` used to hang forever waiting for a connection - that never closes on its own; it must now force it closed instead. - - Measures wall-clock time around a bare ``await server.stop()`` (no - wrapping ``asyncio.wait_for``, which would cancel ``stop()`` from the - outside and let its own ``except CancelledError`` swallow that - cancellation -- masking a real hang as a false pass). The - ``pytest.mark.timeout`` above is the sole backstop, matching how the - live baseline run actually surfaced this hang (pytest-timeout, not an - internal asyncio timeout). + """Regression: an MCP client (e.g. OpenCode) holds its `/sse` GET + open for the life of its session and may never close it -- stop() + must force it closed rather than hang on uvicorn's unbounded default + graceful-shutdown wait. + + Measures wall-clock time around a bare ``server.stop()`` (wrapping + it in ``asyncio.wait_for`` would cancel it externally and mask a + real hang). ``pytest.mark.timeout`` above is the backstop, matching + how this hang is actually caught (pytest-timeout, not an asyncio + timeout). """ server = LocalMCPServer( name="test-local-mcp-stop", @@ -266,34 +236,23 @@ async def test_stop_returns_promptly_with_a_still_open_sse_connection( port_min=0, port_max=0, ) - await server.start() + async with running(server): - connection_ready = asyncio.Event() - - async def hold_connection_open() -> None: - with suppress(Exception): - async with sse_client(server.url) as (read_stream, write_stream): - async with ClientSession(read_stream, write_stream) as session: - await session.initialize() - connection_ready.set() - await asyncio.sleep(60) # never closes on its own - - holder = asyncio.create_task(hold_connection_open()) - try: - await asyncio.wait_for(connection_ready.wait(), timeout=5.0) + async def connect(ready: asyncio.Event) -> None: + with suppress(Exception): + async with sse_client(server.url) as (read_stream, write_stream): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + ready.set() + await asyncio.sleep(60) # never closes on its own - started_at = asyncio.get_running_loop().time() - await server.stop() - elapsed = asyncio.get_running_loop().time() - started_at + async with held_open(connect): + stop_elapsed = await elapsed(server.stop()) - assert elapsed < SERVER_STOP_TIMEOUT_S + 5.0, ( - f"stop() took {elapsed:.1f}s -- graceful shutdown is not " + assert stop_elapsed < SERVER_STOP_TIMEOUT_S + 5.0, ( + f"stop() took {stop_elapsed:.1f}s -- graceful shutdown is not " "bounded by SERVER_STOP_TIMEOUT_S" ) - finally: - holder.cancel() - with suppress(asyncio.CancelledError): - await holder # 30s default barely fits on GitHub Actions Python 3.12 runners — the # streamable-HTTP loopback initialization spends most of that on uvicorn @@ -309,8 +268,7 @@ async def test_serves_streamable_http_tools_on_localhost(self) -> None: port_max=0, ) - await server.start() - try: + async with running(server): assert server.http_url.startswith(f"http://{LOCAL_MCP_HOST}:") async with streamablehttp_client(server.http_url) as ( @@ -322,8 +280,6 @@ async def test_serves_streamable_http_tools_on_localhost(self) -> None: await session.initialize() await _session_lists_only_echo(session) await _call_echo(session, "hello") - finally: - await server.stop() @pytest.mark.asyncio async def test_stop_cleans_up_state_even_if_serve_task_crashed(self) -> None: @@ -375,11 +331,10 @@ async def _raise() -> None: async def test_start_forwards_real_host_to_build_engine( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """Regression (found live via the Letta lane): build_engine must be - told the real bind host, or FastMCP wrongly assumes loopback and - locks DNS-rebinding protection to 127.0.0.1/localhost only -- even - for a server explicitly bound to a non-loopback host for a Docker - callback (see LocalMCPServer's own class docstring).""" + """Regression: build_engine must be told the real bind host, or + FastMCP wrongly assumes loopback and locks DNS-rebinding protection + to 127.0.0.1/localhost only -- even for a non-loopback Docker- + callback bind (see LocalMCPServer's class docstring).""" import band.integrations.mcp.local_server as local_server_mod seen_hosts: list[str] = [] @@ -413,10 +368,8 @@ def spy_build_engine( port_min=0, port_max=0, ) - try: - await server.start() - finally: - await server.stop() + async with running(server): + pass assert seen_hosts == ["0.0.0.0"] @@ -497,8 +450,7 @@ async def test_start_stop_start_cycle_rebuilds_engine(self) -> None: await server.start() await server.stop() - await server.start() - try: + async with running(server): async with streamablehttp_client(server.http_url) as ( read_stream, write_stream, @@ -507,5 +459,3 @@ async def test_start_stop_start_cycle_rebuilds_engine(self) -> None: async with ClientSession(read_stream, write_stream) as session: await session.initialize() await _call_echo(session, "hi") - finally: - await server.stop() diff --git a/tests/integrations/test_uvicorn_server.py b/tests/integrations/test_uvicorn_server.py new file mode 100644 index 000000000..5ecb7a7a6 --- /dev/null +++ b/tests/integrations/test_uvicorn_server.py @@ -0,0 +1,172 @@ +"""Behavior tests for the shared embedded-uvicorn-server lifecycle. + +``mcp.local_server``, ``a2a.gateway.server``, and the A2A baseline test +fixture each embed their own uvicorn server -- ``wait_until_started`` and +``ManagedUvicornServer`` live here, once, instead of every caller +re-deriving (and re-testing) the same startup/shutdown correctness. +""" + +from __future__ import annotations + +import asyncio + +import httpx +import pytest +from sse_starlette.sse import AppStatus + +from band.integrations.uvicorn_server import ManagedUvicornServer, wait_until_started + +from tests.lifecycle import backgrounded, running + + +class FakeUvicornServer: + def __init__(self, *, started: bool = False) -> None: + self.started = started + + +async def _minimal_asgi_app(scope: dict, receive: object, send: object) -> None: + if scope["type"] != "http": + return + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + +@pytest.mark.asyncio +async def test_returns_once_the_server_flips_ready() -> None: + server = FakeUvicornServer() + + async def flip_ready_soon() -> None: + await asyncio.sleep(0.1) + server.started = True + + async with ( + backgrounded(asyncio.sleep(10)) as serve_task, + backgrounded(flip_ready_soon()), + ): + await asyncio.wait_for( + wait_until_started(server, serve_task, timeout_s=5.0), timeout=2.0 + ) + + +@pytest.mark.asyncio +async def test_surfaces_a_serve_task_failure_immediately() -> None: + """A serve task that dies before the server ever reports ready (e.g. a + port already in use) must surface its real exception right away -- + busy-waiting the full timeout and raising a generic one instead would + hide the actual cause.""" + + async def fail_immediately() -> None: + raise OSError("address already in use") + + server = FakeUvicornServer() + serve_task = asyncio.create_task(fail_immediately()) + + start = asyncio.get_running_loop().time() + with pytest.raises(OSError, match="address already in use"): + await wait_until_started(server, serve_task, timeout_s=30.0) + + assert asyncio.get_running_loop().time() - start < 1.0 + + +@pytest.mark.asyncio +async def test_surfaces_a_clean_task_completion_that_never_started() -> None: + """A serve task that ends without ever setting ``started`` -- e.g. an + early shutdown signal, not a raised exception -- must be treated as + fatal immediately too; only distinguishing "done and raised" from "done" + would still busy-wait the full timeout on this path.""" + server = FakeUvicornServer() + serve_task = asyncio.create_task(asyncio.sleep(0)) + + start = asyncio.get_running_loop().time() + with pytest.raises(RuntimeError, match="ended before ever starting"): + await wait_until_started(server, serve_task, timeout_s=30.0) + + assert asyncio.get_running_loop().time() - start < 1.0 + + +@pytest.mark.asyncio +async def test_times_out_if_the_server_never_reports_ready() -> None: + server = FakeUvicornServer() + async with backgrounded(asyncio.sleep(10)) as serve_task: + with pytest.raises(TimeoutError): + await wait_until_started(server, serve_task, timeout_s=0.2) + + +@pytest.mark.asyncio +async def test_managed_server_starts_and_bound_port_resolves_real_port() -> None: + server = ManagedUvicornServer( + _minimal_asgi_app, + host="127.0.0.1", + port=0, + start_timeout_s=5.0, + stop_timeout_s=5, + ) + async with running(server): + async with httpx.AsyncClient() as client: + response = await client.get(f"http://127.0.0.1:{server.bound_port}/") + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_stop_before_start_is_a_no_op() -> None: + server = ManagedUvicornServer( + _minimal_asgi_app, + host="127.0.0.1", + port=0, + start_timeout_s=5.0, + stop_timeout_s=5, + ) + await server.stop() # must not raise + with pytest.raises(RuntimeError, match="has not started"): + _ = server.bound_port + + +@pytest.mark.asyncio +async def test_start_cleans_up_when_the_startup_wait_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed/timed-out startup wait must still tell uvicorn to exit and + clear server state, not leave a listening socket and stray task behind.""" + import band.integrations.uvicorn_server as uvicorn_server_module + + async def failing_wait_until_started(*args: object, **kwargs: object) -> None: + raise TimeoutError("simulated startup failure") + + monkeypatch.setattr( + uvicorn_server_module, "wait_until_started", failing_wait_until_started + ) + + server = ManagedUvicornServer( + _minimal_asgi_app, + host="127.0.0.1", + port=0, + start_timeout_s=5.0, + stop_timeout_s=5, + ) + + with pytest.raises(TimeoutError, match="simulated startup failure"): + await server.start() + + with pytest.raises(RuntimeError, match="has not started"): + _ = server.bound_port + + +def test_disables_sse_starlette_automatic_graceful_drain() -> None: + """Regression: AppStatus.should_exit is a process-global with no notion + of "which server" -- any uvicorn.Server's handle_exit() latches it, + cutting off every other embedded server's SSE responses too. Importing + this module must disable the automatic drain so handle_exit() (the real + signal-handler call) is a no-op for this flag; original_handler is + swapped since it expects a bound Server, not this direct call. + """ + assert AppStatus.enable_automatic_graceful_drain is False + + original_should_exit = AppStatus.should_exit + original_handler = AppStatus.original_handler + AppStatus.original_handler = None + try: + AppStatus.handle_exit(0, None) + assert AppStatus.should_exit is False + finally: + AppStatus.should_exit = original_should_exit + AppStatus.original_handler = original_handler diff --git a/tests/lifecycle.py b/tests/lifecycle.py new file mode 100644 index 000000000..9ad039c4c --- /dev/null +++ b/tests/lifecycle.py @@ -0,0 +1,64 @@ +"""Generic async lifecycle helpers for tests driving a background +server or task end to end.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine +from contextlib import asynccontextmanager, suppress +from typing import Any, Protocol + + +class Startable(Protocol): + """``stop()`` must be a safe no-op if ``start()`` was never called or + failed partway -- ``running()`` below relies on that to clean up + unconditionally after a failed start.""" + + async def start(self) -> None: ... + async def stop(self) -> None: ... + + +@asynccontextmanager +async def running(server: Startable) -> AsyncIterator[Startable]: + """Start ``server``, yield it, and always stop it -- even on failure.""" + try: + await server.start() + yield server + finally: + await server.stop() + + +@asynccontextmanager +async def backgrounded( + coro: Coroutine[Any, Any, object], +) -> AsyncIterator[asyncio.Task[object]]: + """Run ``coro`` as a background task for the block; always cancelled + and awaited afterward, even on failure.""" + task = asyncio.create_task(coro) + try: + yield task + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + +@asynccontextmanager +async def held_open( + connect: Callable[[asyncio.Event], Awaitable[None]], +) -> AsyncIterator[None]: + """Run ``connect`` in the background until it signals its ready event, + keeping whatever connection it opens alive for the block.""" + ready = asyncio.Event() + async with backgrounded(connect(ready)): + await asyncio.wait_for(ready.wait(), timeout=5.0) + yield + + +async def elapsed(coro: Awaitable[None]) -> float: + """Wall-clock seconds ``coro`` took -- for asserting a bounded-shutdown + guarantee without an external ``wait_for`` that would mask a real hang + by cancelling the call from outside.""" + start = asyncio.get_running_loop().time() + await coro + return asyncio.get_running_loop().time() - start