From 63242e7a64728915e9742aca5dbc0e3a6fa6ebd9 Mon Sep 17 00:00:00 2001 From: LHMQ878 Date: Tue, 18 Aug 2026 01:03:12 +0800 Subject: [PATCH] fix(labs): persist Antigravity resume state when a turn ends early _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-, 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. --- .../labs/antigravity/_antigravity_agent.py | 69 +++++++------ .../antigravity/test_antigravity_agent.py | 96 +++++++++++++++++++ 2 files changed, 137 insertions(+), 28 deletions(-) diff --git a/src/google/adk/labs/antigravity/_antigravity_agent.py b/src/google/adk/labs/antigravity/_antigravity_agent.py index 8f06cf2335..3f7e215601 100644 --- a/src/google/adk/labs/antigravity/_antigravity_agent.py +++ b/src/google/adk/labs/antigravity/_antigravity_agent.py @@ -209,35 +209,48 @@ async def _run_async_impl( ctx.run_config and ctx.run_config.streaming_mode == StreamingMode.SSE ) - async with self._sdk_agent_cls(config) as active_agent: - await active_agent.conversation.send(prompt) - - async for step in active_agent.conversation.receive_steps(): - if step.step_index <= resume_step_index: - continue - max_step_index = max(max_step_index, step.step_index) - for event in _event_converter.convert_step_to_events( - step, - ctx=ctx, - author=self.name, - seen_tool_calls=seen_tool_calls, - seen_tool_results=seen_tool_results, - streaming=streaming, - ): - yield event - - harness_conversation_id = active_agent.conversation_id - - # On a fresh turn the harness wrote traj- (flushed when the session - # exits above); rename it. No id under single-turn, so that case skips. - if save_dir and conversation_id: - if not resumed and harness_conversation_id: - _trajectory_files.rename_trajectory( - save_dir, conversation_id, harness_conversation_id + harness_conversation_id: str | None = None + + try: + async with self._sdk_agent_cls(config) as active_agent: + await active_agent.conversation.send(prompt) + # Read before the step loop rather than after it: a turn that ends early + # still has to rename the trajectory the harness has by then created. + harness_conversation_id = active_agent.conversation_id + + async for step in active_agent.conversation.receive_steps(): + if step.step_index <= resume_step_index: + continue + max_step_index = max(max_step_index, step.step_index) + for event in _event_converter.convert_step_to_events( + step, + ctx=ctx, + author=self.name, + seen_tool_calls=seen_tool_calls, + seen_tool_results=seen_tool_results, + streaming=streaming, + ): + yield event + finally: + # In a finally because a turn does not have to run to completion. The + # runner wraps this generator in `aclosing`, so a caller that stops + # reading - a disconnected client, a cancelled run - closes it at a + # `yield`, and the harness itself can fail mid-stream. Every event + # yielded above is already recorded in the session either way, so + # leaving this out replays them on the next turn, and leaving the rename + # out orphans the trajectory and silently starts a fresh conversation. + # + # On a fresh turn the harness wrote traj-, flushed when the + # session exits above; rename it. No id under single-turn, so that case + # skips, as does a turn that never reached the harness at all. + if save_dir and conversation_id and (resumed or harness_conversation_id): + if not resumed and harness_conversation_id: + _trajectory_files.rename_trajectory( + save_dir, conversation_id, harness_conversation_id + ) + _trajectory_files.save_resume_step_index( + save_dir, conversation_id, max_step_index ) - _trajectory_files.save_resume_step_index( - save_dir, conversation_id, max_step_index - ) @override async def _run_impl( diff --git a/tests/unittests/labs/antigravity/test_antigravity_agent.py b/tests/unittests/labs/antigravity/test_antigravity_agent.py index de778d4a71..6611a3ffcf 100644 --- a/tests/unittests/labs/antigravity/test_antigravity_agent.py +++ b/tests/unittests/labs/antigravity/test_antigravity_agent.py @@ -448,6 +448,102 @@ async def _receive_steps(): assert (save_dir / f'traj-{conversation_id}.resume').read_text() == '2' +@pytest.mark.asyncio +async def test_trajectory_is_claimed_when_the_caller_stops_reading(tmp_path): + """A turn closed early still renames the trajectory and records its index. + + The runner wraps agent generators in `aclosing`, so a caller that stops + reading - a disconnected client, a cancelled run - closes this one at a + yield. Without the rename the next turn finds no trajectory under the + derived id, so it starts a fresh conversation and abandons this one's file. + """ + + async def _receive_steps(): + yield _text_step(0, 'one') + yield _text_step(1, 'two') + + active_agent = _fake_active_agent( + _receive_steps, conversation_id='harness-random' + ) + # The harness has already created its trajectory under its own random id. + (tmp_path / 'traj-harness-random').write_bytes(b'data') + agent = AntigravityAgent( + name='agy', config=_make_config(save_dir=str(tmp_path)) + ) + conversation_id = _antigravity_agent._derive_conversation_id( + 'sess_456', 'agy' + ) + + with patch.object(_antigravity_agent, 'Agent', return_value=active_agent): + agen = agent._run_async_impl(_mock_run_ctx()) + async for _ in agen: + break + await agen.aclose() + + assert (tmp_path / f'traj-{conversation_id}').exists() + assert not (tmp_path / 'traj-harness-random').exists() + # Step 0 was emitted and recorded, so the next turn must not replay it. + assert (tmp_path / f'traj-{conversation_id}.resume').read_text() == '0' + + +@pytest.mark.asyncio +async def test_resume_index_survives_a_harness_error_mid_turn(tmp_path): + """Steps emitted before a mid-turn failure are not replayed on the next turn. + + The events are in the session as soon as they are yielded, so a step whose + turn later failed still has to count as emitted. + """ + + async def _receive_steps(): + yield _text_step(0, 'old-1') + yield _text_step(1, 'old-2') + yield _text_step(2, 'new') + raise RuntimeError('harness died mid-turn') + + conversation_id = _antigravity_agent._derive_conversation_id( + 'sess_456', 'agy' + ) + active_agent = _fake_active_agent( + _receive_steps, conversation_id=conversation_id + ) + (tmp_path / f'traj-{conversation_id}').write_bytes(b'data') + (tmp_path / f'traj-{conversation_id}.resume').write_text('1') + agent = AntigravityAgent( + name='agy', config=_make_config(save_dir=str(tmp_path)) + ) + + emitted = [] + with patch.object(_antigravity_agent, 'Agent', return_value=active_agent): + with pytest.raises(RuntimeError, match='harness died mid-turn'): + async for event in agent._run_async_impl(_mock_run_ctx()): + emitted.append(event.content.parts[0].text) + + assert emitted == ['new'] + assert (tmp_path / f'traj-{conversation_id}.resume').read_text() == '2' + + +@pytest.mark.asyncio +async def test_a_turn_that_never_reaches_the_harness_records_nothing(tmp_path): + """A failure before the conversation exists leaves save_dir untouched. + + There is no trajectory to rename and no step to remember, so writing a + resume index would only leave a file for a conversation that never began. + """ + agent = AntigravityAgent( + name='agy', config=_make_config(save_dir=str(tmp_path)) + ) + + def _refuse(_config): + raise RuntimeError('harness unavailable') + + with patch.object(_antigravity_agent, 'Agent', _refuse): + with pytest.raises(RuntimeError, match='harness unavailable'): + async for _ in agent._run_async_impl(_mock_run_ctx()): + pass + + assert list(tmp_path.iterdir()) == [] + + @pytest.mark.asyncio async def test_node_input_becomes_the_prompt(tmp_path): """The parent's composed request wins over the original user message.