From 7dac7826616f92c9afb0cff4192a3503d0dc8068 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 11 Nov 2026 17:07:50 +0000 Subject: [PATCH 1/2] fix(a2a_agent): end a trajectory follower with its step --- .../a2a_agent/trajectory_follower.py | 20 ++++++--- .../a2a_agent/test_trajectory_follower.py | 45 +++++++++++++++++++ 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/src/agent_env/a2a_agent/trajectory_follower.py b/src/agent_env/a2a_agent/trajectory_follower.py index 9bcaed794..01ffbdf02 100644 --- a/src/agent_env/a2a_agent/trajectory_follower.py +++ b/src/agent_env/a2a_agent/trajectory_follower.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio +import contextlib import logging from collections.abc import Awaitable, Callable from contextlib import nullcontext @@ -71,7 +72,8 @@ async def follow_trajectory( events and the last one, until the task has ended and every event is read. Between reads it waits ``poll_interval_seconds``, or less once ``stop`` is set. Transport errors - and 5xx answers are retried with backoff; a task the agent no longer knows (404) ends following.""" + and 5xx answers are retried with backoff until ``stop`` is set, and end following after it; a + task the agent no longer knows (404) ends following.""" after = 0 failures = 0 async with nullcontext(client) if client is not None else httpx.AsyncClient() as http: @@ -113,11 +115,15 @@ async def follow_trajectory( if page.has_more: continue if stop.is_set(): + if failures: + # The turn has ended, so an agent that stops answering now is gone, not busy: retrying + # would poll a dead address until the step's drain gives up. The step marks the turn's end. + logger.warning("Agent stopped answering after task %s ended; stopped following its trajectory", task_id) + return await asyncio.sleep(_STOPPED_INTERVAL_SECONDS) continue - try: - await asyncio.wait_for( - stop.wait(), timeout=min(poll_interval_seconds + failures * 5, _MAX_BACKOFF_SECONDS) - ) - except TimeoutError: - pass + # Not `wait_for`: on 3.11 it drops a cancel that lands as `stop` is set, the step's own + # teardown, and the follower then outlives its step. + with contextlib.suppress(TimeoutError): + async with asyncio.timeout(min(poll_interval_seconds + failures * 5, _MAX_BACKOFF_SECONDS)): + await stop.wait() diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index a3cdc71c9..febc7de0e 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -124,6 +124,51 @@ async def test_a_server_error_is_retried(monkeypatch): assert batches[-1].events == [{"i": 0}] +async def _start_following(agent: _Agent, stop: asyncio.Event, *, poll_interval_seconds: float) -> asyncio.Task: + async def sink(batch: TrajectoryBatch) -> None: + pass + + async def follow() -> None: + async with agent.client() as client: + await follow_trajectory( + _ENDPOINT, "task-1", sink, poll_interval_seconds=poll_interval_seconds, stop=stop, client=client + ) + + follower = asyncio.create_task(follow()) + async with asyncio.timeout(5): + while not agent.afters: + await asyncio.sleep(0.01) + return follower + + +@pytest.mark.asyncio +async def test_a_follower_cancelled_as_it_is_stopped_ends(): + """A step's teardown sets ``stop`` and cancels its follower in the same tick, against an agent + that is gone; on Python 3.11 ``asyncio.wait_for`` dropped that cancel and the follower ran on.""" + agent = _Agent(httpx.Response(200, json=_page(0, [], "running")), *[httpx.Response(503)] * 100) + stop = asyncio.Event() + follower = await _start_following(agent, stop, poll_interval_seconds=30) + + stop.set() + follower.cancel() + + done, _ = await asyncio.wait({follower}, timeout=2) + assert done == {follower} + + +@pytest.mark.asyncio +async def test_an_agent_that_fails_after_the_turn_ends_ends_following(): + agent = _Agent(httpx.Response(200, json=_page(0, [], "running")), *[httpx.Response(503)] * 100) + stop = asyncio.Event() + follower = await _start_following(agent, stop, poll_interval_seconds=30) + + stop.set() + + async with asyncio.timeout(2): + await follower + assert agent.afters == [0, 0] + + @pytest.mark.asyncio async def test_a_task_the_agent_no_longer_knows_ends_following(): agent = _Agent(httpx.Response(404, json={"detail": "Unknown task"})) From f9bee8a5208385c3604f32759ee11516e1be214f Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 11 Nov 2026 17:36:30 +0000 Subject: [PATCH 2/2] test: the cancel test fails when only the cancel is lost, and no test leaves its follower running --- .../a2a_agent/test_trajectory_follower.py | 44 ++++++++++++------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/tst/unit/a2a_agent/test_trajectory_follower.py b/tst/unit/a2a_agent/test_trajectory_follower.py index febc7de0e..8aac86070 100644 --- a/tst/unit/a2a_agent/test_trajectory_follower.py +++ b/tst/unit/a2a_agent/test_trajectory_follower.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import json import time @@ -124,7 +125,11 @@ async def test_a_server_error_is_retried(monkeypatch): assert batches[-1].events == [{"i": 0}] -async def _start_following(agent: _Agent, stop: asyncio.Event, *, poll_interval_seconds: float) -> asyncio.Task: +@contextlib.asynccontextmanager +async def _following(agent: _Agent, stop: asyncio.Event, *, poll_interval_seconds: float): + """A follower of ``agent``'s task, yielded once it has made its first read; cancelled on the way out + if it is still running, so a failing test leaves no follower behind.""" + async def sink(batch: TrajectoryBatch) -> None: pass @@ -135,37 +140,42 @@ async def follow() -> None: ) follower = asyncio.create_task(follow()) - async with asyncio.timeout(5): - while not agent.afters: - await asyncio.sleep(0.01) - return follower + try: + async with asyncio.timeout(5): + while not agent.afters: + await asyncio.sleep(0.01) + yield follower + finally: + follower.cancel() + await asyncio.wait({follower}) @pytest.mark.asyncio async def test_a_follower_cancelled_as_it_is_stopped_ends(): - """A step's teardown sets ``stop`` and cancels its follower in the same tick, against an agent - that is gone; on Python 3.11 ``asyncio.wait_for`` dropped that cancel and the follower ran on.""" - agent = _Agent(httpx.Response(200, json=_page(0, [], "running")), *[httpx.Response(503)] * 100) + """A step's teardown sets ``stop`` and cancels its follower in the same tick. The agent keeps + answering, so only the cancel can end the follower; on Python 3.11 ``asyncio.wait_for`` dropped it.""" + agent = _Agent(*[httpx.Response(200, json=_page(0, [], "running"))] * 100) stop = asyncio.Event() - follower = await _start_following(agent, stop, poll_interval_seconds=30) + async with _following(agent, stop, poll_interval_seconds=30) as follower: + stop.set() + follower.cancel() - stop.set() - follower.cancel() + done, _ = await asyncio.wait({follower}, timeout=2) - done, _ = await asyncio.wait({follower}, timeout=2) - assert done == {follower} + assert done == {follower} + assert follower.cancelled() @pytest.mark.asyncio async def test_an_agent_that_fails_after_the_turn_ends_ends_following(): agent = _Agent(httpx.Response(200, json=_page(0, [], "running")), *[httpx.Response(503)] * 100) stop = asyncio.Event() - follower = await _start_following(agent, stop, poll_interval_seconds=30) + async with _following(agent, stop, poll_interval_seconds=30) as follower: + stop.set() - stop.set() + async with asyncio.timeout(2): + await follower - async with asyncio.timeout(2): - await follower assert agent.afters == [0, 0]