Skip to content

fix(labs): persist Antigravity resume state when a turn ends early - #6765

Open
LHMQ878 wants to merge 1 commit into
google:mainfrom
LHMQ878:fix/antigravity-persist-resume-state
Open

fix(labs): persist Antigravity resume state when a turn ends early#6765
LHMQ878 wants to merge 1 commit into
google:mainfrom
LHMQ878:fix/antigravity-persist-resume-state

Conversation

@LHMQ878

@LHMQ878 LHMQ878 commented Aug 17, 2026

Copy link
Copy Markdown

What's wrong

AntigravityAgent._run_async_impl renamed the harness trajectory and wrote the resume step index after the async with block, so neither happened unless the turn ran to the end of the step stream. Two ordinary paths skip it.

A caller that stops reading. The runner wraps agent generators in contextlib.aclosing (runners.py:1415), so a disconnected SSE client or a cancelled run closes this generator at a yield. __aexit__ still runs and the harness still flushes traj-<random>, but the statements after the with do not.

A harness error mid-stream. Same skip, and here the events emitted before the failure are already in the session.

Reproduced

Driving _run_async_impl with a stubbed SDK agent, on 735402d:

--- A: caller breaks after the first event ---
  files:                      ['traj-harness-random']
  renamed to traj-<cid>?      False
  resume index written?       False
  next turn would resume?     False        <-- starts a fresh conversation

--- B: harness raises after emitting step 2 (resumed turn, index was 1) ---
  events emitted to session:  ['new']
  resume index on disk:       1            <-- step 2 replays next turn

So in A the conversation is silently abandoned: the ADK session holds the turn's events, the harness is told nothing about it next time, and traj-harness-random is orphaned in save_dir permanently. In B step 2 gets recorded twice.

The fix

Move the rename and the index write into a finally, and read active_agent.conversation_id before the step loop rather than after it, so it is available when the loop does not finish.

Both writes are plain synchronous file operations, which is what makes them safe in a finally that may be running under GeneratorExit — nothing is awaited and nothing is yielded there.

The condition is resumed or harness_conversation_id, so a turn that never reached the harness at all (say Agent(config) raises) still writes nothing: there is no trajectory to rename and no step worth remembering, and a stray traj-<cid>.resume would describe a conversation that never began.

Tests

Three added to tests/unittests/labs/antigravity/test_antigravity_agent.py:

test without the fix
test_trajectory_is_claimed_when_the_caller_stops_reading fails — no rename, no index
test_resume_index_survives_a_harness_error_mid_turn fails — index stays 1
test_a_turn_that_never_reaches_the_harness_records_nothing passes

The third passes on main by construction — it is a guard against this change writing state for a conversation that does not exist, not a regression test for the bug. Calling that out rather than presenting three red-to-green tests.

Verification

  • pytest tests/unittests/labs/antigravity/50 passed (was 47 passed)
  • without the src change, the same run is 2 failed, 48 passed
  • pytest tests/unittests/labs/ tests/unittests/agents/986 passed, 2 xfailed, no failures
  • pyink --check2 files would be left unchanged; isort --check-only → exit 0

Deliberately not in this PR

_trajectory_files.save_resume_step_index truncates and then writes, so an interruption during the write leaves an empty file. load_resume_step_index treats that as -1, and with the trajectory still present the whole conversation then replays. Putting the write in a finally makes that window slightly more reachable, so it is worth a temp-file-plus-os.replace follow-up — but it needs a crash to trigger and I have no test that demonstrates it, so I left it out rather than mixing a shown bug with a reasoned one. Happy to add it here if you would rather have both together.

Follows on from #6040.

_run_async_impl renamed the trajectory and wrote the resume step index after
the `async with`, so neither ran unless the turn reached the end of the step
stream. Two ordinary paths skip it.

A caller that stops reading closes the generator at a yield: the runner wraps
agent generators in `aclosing` (runners.py:1415), so a disconnected client or
a cancelled run does exactly that. The harness has already flushed
traj-<random>, but the rename never happens, so the next turn finds no
trajectory under the derived id, starts a brand new conversation, and leaves
the old file behind for good.

A harness error mid-stream loses the resume index the same way. Steps emitted
before the failure are already recorded in the session, so the next turn
replays them and records them a second time.

Move both writes into a finally, and read conversation_id before the step loop
rather than after it, so it is available on those paths. A turn that never
reached the harness still writes nothing: there is no trajectory to rename and
no step to remember.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants