[andrewpark-scaleai]

Why

A live trajectory follower can outlive the step that started it. On the agent-env workers, 26 followers kept polling agents that had already gone away, on 11 pods. Each one holds a Temporal activity slot and logs Live trajectory read failed about 5 times a second. Two bugs combine:

1. Python 3.11's asyncio.wait_for drops the cancel. A step's teardown sets stop and cancels the follower in the same tick. If the follower is waiting in wait_for(stop.wait()) at that moment, 3.11 returns normally and swallows the cancel. The workers run 3.11.
2. After stop is set, a failed read retries every 0.2 s with no limit. An agent that is gone never answers, so the follower loops until it is cancelled. When the cancel was dropped, it loops forever, and following_live_trajectory's await asyncio.wait(turn.followers) hangs the step with it.

What

• The follower waits with asyncio.timeout, which does not drop a cancel.
• Once stop is set, a failed read ends following. The turn has ended by then, so an agent that stops answering is gone, not busy. following_live_trajectory already marks a turn's end when its follower stops early, and the final trajectory still holds every event.

Validation

I ran a step against an agent that closes every connection, on the deployed worker image (Python 3.11.15, agentenv-framework 0.9.1490). The step either was cancelled, raised, or returned.

Step	Before	After
cancelled	never ends (224 failed reads in 90 s)	ends at once (1 failed read)
raises	never ends (216 failed reads in 90 s)	ends at once (1 failed read)
returns	ends after the 60 s drain (150 failed reads)	ends at once (2 failed reads)

Two new unit tests:

• test_a_follower_cancelled_as_it_is_stopped_ends fails on the old follower under 3.11 and passes under 3.12, where wait_for was fixed. CI runs 3.12, so it guards only a 3.11 run.
• test_an_agent_that_fails_after_the_turn_ends_ends_following fails on the old follower under both versions.

With the change, the follower and live-trajectory tests pass on 3.11 and 3.12 (31 each).

In the full unit suite, 14,522 tests pass. The 1 failure, test_mcp_cli_builder.py::test_codegen_loads_sibling_env_file, fails on main too, because its subprocess can't import httpx here.

Confidence Score: 5/5

The PR appears safe to merge; both previous test issues are fixed.

Summary

The PR keeps trajectory followers from outliving their steps:

• A2A trajectory followers now stop with their step.

Both previous findings are addressed. No new actionable issues were found in the changes since the previous review. Tests were not run during this review.

_{Reviews (2) · Last reviewed commit: "test: the cancel test fails when only th..." · Reviewed by Greptile}