fix(a2a_agent): a trajectory follower ends with its step

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()

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"}))