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..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,6 +125,60 @@ async def test_a_server_error_is_retried(monkeypatch): assert batches[-1].events == [{"i": 0}] +@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 + + 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()) + 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. 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() + async with _following(agent, stop, poll_interval_seconds=30) as follower: + stop.set() + follower.cancel() + + done, _ = await asyncio.wait({follower}, timeout=2) + + 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() + async with _following(agent, stop, poll_interval_seconds=30) as follower: + 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"}))