diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index 6ab54d6dae..803b3c8d17 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -644,6 +644,12 @@ async def _generator() -> Any: iteration = 0 work_iterations = 0 progress: list[str] = [] + # Only a clean loop exit (should_continue stop, max_iterations, approval escape + # hatch, or MiddlewareTermination) should fire the turn-scoped after-providers. An + # exception unwinding the generator, or an early ``aclose()``, must not - otherwise + # they run with a stale prior-iteration final, diverging from the non-streaming path + # and the normal (non-loop) run contract (both of which skip after_run on error). + completed_cleanly = False stamped_options = dict(context.options) if context.options is not None else {} stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object() context.options = stamped_options @@ -668,7 +674,7 @@ async def _generator() -> Any: # process() has returned (the stream is consumed lazily), so a termination # raised by a downstream middleware or during stream consumption surfaces here. # Stop cleanly and keep whatever final response we have from a prior iteration. - return + break iteration += 1 @@ -677,7 +683,7 @@ async def _generator() -> Any: # Escape hatch: if this iteration is asking for tool approval, stop the loop and # let the caller approve, instead of continuing or injecting next_message. if self._has_pending_approval_request(final): - return + break loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -711,7 +717,7 @@ async def _generator() -> Any: feedback=feedback, ) if stop: - return + break if snapshot is not None and context.session is not None: # Reset the session to the pre-loop baseline before the next run. The final # response was already awaited above, so the service-side conversation id has @@ -724,9 +730,13 @@ async def _generator() -> Any: # transcript that non-streaming runs return). for message in next_messages: yield self._message_to_update(message) + # Reached only by breaking out of the loop on a clean stop; an exception + # propagates past this point straight to ``finally`` with the flag unset. + completed_cleanly = True finally: context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) - await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages) + if completed_cleanly: + await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages) def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: if holder["final"] is not None: diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index 1ad4663027..a5f4922de8 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -1652,6 +1652,128 @@ async def test_turn_scoped_after_run_fires_normally_without_loop() -> None: assert turn_scoped.after_calls == 1 +class _RaiseOnSecondIterationClient(RecordingChatClient): + """Succeeds on the first loop iteration and raises on the second. + + By the time the second iteration fails, a prior-iteration final response already + exists, which is the condition under which a loop could erroneously run the + turn-scoped ``after_run`` with a stale response. + """ + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + if stream: + + async def _gen() -> AsyncIterable[ChatResponseUpdate]: + self.call_count += 1 + if self.call_count >= 2: + raise RuntimeError("boom on iteration 2") + yield ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant", finish_reason="stop") + + return ResponseStream(_gen(), finalizer=ChatResponse.from_updates) + + async def _get() -> ChatResponse: + self.call_count += 1 + if self.call_count >= 2: + raise RuntimeError("boom on iteration 2") + return ChatResponse(messages=Message(role="assistant", contents=["first"])) + + return _get() + + +async def test_turn_scoped_after_run_skipped_on_loop_exception_non_streaming() -> None: + """A loop iteration that raises must not fire turn-scoped after_run (non-streaming). + + Matches the normal (non-loop) run contract, where after_run does not fire when the + run raises. + """ + client = _RaiseOnSecondIterationClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=5)], + context_providers=[turn_scoped], + ) + + with pytest.raises(RuntimeError, match="boom"): + await agent.run("start") + + assert turn_scoped.after_calls == 0 + + +async def test_turn_scoped_after_run_skipped_on_loop_exception_streaming() -> None: + """A loop iteration that raises must not fire turn-scoped after_run (streaming). + + Regression: the streaming loop previously fired the turn-scoped after-providers from a + ``finally`` block, so an exception on a later iteration ran them with the prior + iteration's stale final response - diverging from both the non-streaming path and the + normal (non-loop) run contract. + """ + client = _RaiseOnSecondIterationClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=5)], + context_providers=[turn_scoped], + ) + + with pytest.raises(RuntimeError, match="boom"): + stream = agent.run("start", stream=True) + _ = [update async for update in stream] + await stream.get_final_response() + + assert turn_scoped.after_calls == 0 + + +async def test_turn_scoped_after_run_skipped_on_streaming_early_close() -> None: + """Dropping a streaming loop early (GeneratorExit) must not fire turn-scoped after_run.""" + + class _TwoUpdatesPerIterationClient(RecordingChatClient): + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + async def _gen() -> AsyncIterable[ChatResponseUpdate]: + self.call_count += 1 + yield ChatResponseUpdate(contents=[Content.from_text("a")], role="assistant") + yield ChatResponseUpdate(contents=[Content.from_text("b")], role="assistant", finish_reason="stop") + + return ResponseStream(_gen(), finalizer=ChatResponse.from_updates) + + client = _TwoUpdatesPerIterationClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=5)], + context_providers=[turn_scoped], + ) + + stream = agent.run("start", stream=True) + # Pull until iteration 2's client call has started: that guarantees iteration 1 fully + # finalized (so a prior-iteration final response exists) before we drop the stream. + pulled = 0 + async for _update in stream: + pulled += 1 + if client.call_count >= 2: + break + if pulled > 100: # pragma: no cover - safety bound against hanging the test + pytest.fail("second loop iteration never started") + assert client.call_count >= 2 + await stream.close() + + assert turn_scoped.after_calls == 0 + + def test_compaction_provider_defers_to_turn_boundary() -> None: from agent_framework._compaction import CompactionProvider