From 17dd6e55cac1b60dfb91d7ba37f8afb6e5896658 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 23 Jul 2026 19:30:24 +0800 Subject: [PATCH 1/3] Python: defer turn-scoped after_run providers to the agent loop boundary Each AgentLoopMiddleware iteration is a full agent run, so CompactionProvider.after_run fired per iteration and rewrote persisted history mid-task (#7236). Providers can now opt into turn scope with after_run_once_per_turn; iterations defer them via a contextvar, and the loop fires them once at the boundary. CompactionProvider opts in; HistoryProvider keeps its incremental per-run persistence. --- .../packages/core/agent_framework/_agents.py | 17 ++ .../core/agent_framework/_compaction.py | 4 + .../core/agent_framework/_harness/_loop.py | 177 +++++++++++------- .../core/agent_framework/_sessions.py | 7 + .../core/tests/core/test_harness_loop.py | 76 ++++++++ 5 files changed, 215 insertions(+), 66 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c6d20f85abb..9f453464a02 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -7,6 +7,7 @@ import sys from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack +from contextvars import ContextVar from copy import deepcopy from functools import partial from itertools import chain @@ -73,6 +74,11 @@ logger = logging.getLogger("agent_framework") +# Set by AgentLoopMiddleware while a loop iteration is running so that +# providers scoped to the whole user turn (``after_run_once_per_turn``) skip +# their per-iteration ``after_run`` and only fire once at the loop boundary. +_LOOP_ITERATION_ACTIVE: ContextVar[bool] = ContextVar("agent_loop_iteration_active", default=False) + if TYPE_CHECKING: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) else: @@ -528,12 +534,17 @@ async def _run_after_providers( *, session: AgentSession | None, context: SessionContext, + only_per_turn: bool = False, ) -> None: """Run after_run on all context providers in reverse order. Keyword Args: session: The conversation session. context: The invocation context with response populated. + only_per_turn: When True, run only providers that opted into + once-per-turn semantics (``after_run_once_per_turn``); used by + AgentLoopMiddleware when a loop ends. When False, those + providers are skipped while a loop iteration is in progress. """ provider_session = session if provider_session is None and self.context_providers: @@ -545,9 +556,15 @@ async def _run_after_providers( per_service_call_history_required = self.require_per_service_call_history_persistence and any( isinstance(provider, HistoryProvider) for provider in self.context_providers ) + in_loop_iteration = _LOOP_ITERATION_ACTIVE.get() for provider in reversed(self.context_providers): if per_service_call_history_required and isinstance(provider, HistoryProvider): continue + once_per_turn = getattr(provider, "after_run_once_per_turn", False) + if only_per_turn and not once_per_turn: + continue + if in_loop_iteration and once_per_turn: + continue if provider_session is None: raise RuntimeError("Provider session must be available when context providers are configured.") await provider.after_run( diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 59abb10a468..904b8981e39 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -1285,6 +1285,10 @@ class CompactionProvider(ContextProvider): await agent.run("Hello", session=session) """ + # Compacting persisted history mid-task rewrites the transcript the loop + # still works from, so defer it to the end of the user turn. + after_run_once_per_turn = True + def __init__( self, *, diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index 7640d624cfd..dc221569c48 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -34,8 +34,10 @@ from pydantic import BaseModel, Field from typing_extensions import Self +from .._agents import _LOOP_ITERATION_ACTIVE from .._feature_stage import ExperimentalFeature, experimental from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination +from .._sessions import SessionContext from .._types import ( AgentResponse, AgentResponseUpdate, @@ -472,6 +474,34 @@ def _restore_session(session: Any, snapshot: dict[str, Any]) -> None: session.service_session_id = restored.service_session_id session.state = restored.state + async def _fire_turn_scoped_after_providers( + self, + context: AgentContext, + response: AgentResponse | None, + input_messages: list[Message], + ) -> None: + """Fire deferred ``after_run`` for turn-scoped providers at the loop boundary. + + Iterations suppress providers that set ``after_run_once_per_turn``; the + loop runs them here once, with the turn-level response. + """ + agent = context.agent + run_after = getattr(agent, "_run_after_providers", None) + if response is None or run_after is None: + return + if not any( + getattr(provider, "after_run_once_per_turn", False) + for provider in getattr(agent, "context_providers", []) + ): + return + session_context = SessionContext( + session_id=context.session.session_id if context.session else None, + service_session_id=context.session.service_session_id if context.session else None, + input_messages=list(input_messages), + ) + session_context._response = response + await run_after(session=context.session, context=session_context, only_per_turn=True) + async def _process_non_streaming( self, context: AgentContext, @@ -488,7 +518,11 @@ async def _process_non_streaming( aggregated_usage: UsageDetails | None = None final_result: AgentResponse | None = None while True: - await call_next() + loop_token = _LOOP_ITERATION_ACTIVE.set(True) + try: + await call_next() + finally: + _LOOP_ITERATION_ACTIVE.reset(loop_token) iteration += 1 result = context.result @@ -555,6 +589,11 @@ async def _process_non_streaming( if not self.return_final_only: context.result = self._aggregate_response(final_result, aggregated, aggregated_usage) + await self._fire_turn_scoped_after_providers( + context, + context.result if isinstance(context.result, AgentResponse) else final_result, + original_messages, + ) def _process_streaming( self, @@ -571,58 +610,52 @@ async def _generator() -> Any: iteration = 0 work_iterations = 0 progress: list[str] = [] - while True: - try: - await call_next() - inner = context.result - if not isinstance(inner, ResponseStream): - raise TypeError( - "AgentLoopMiddleware expected a ResponseStream from a streaming run, " - f"got {type(inner).__name__}." - ) - - async for update in inner: - yield update - - holder["final"] = await inner.get_final_response() - except MiddlewareTermination: - # The pipeline's MiddlewareTermination suppression is no longer active once - # 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 - - iteration += 1 - - messages_used = context.messages - final = holder["final"] - # 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 - loop_kwargs = self._build_loop_kwargs( - context=context, - iteration=iteration, - last_result=final, - messages_used=messages_used, - original_messages=original_messages, - progress=progress, - ) + try: + while True: + loop_token = _LOOP_ITERATION_ACTIVE.set(True) + try: + await call_next() + inner = context.result + if not isinstance(inner, ResponseStream): + raise TypeError( + "AgentLoopMiddleware expected a ResponseStream from a streaming run, " + f"got {type(inner).__name__}." + ) + + async for update in inner: + yield update + + holder["final"] = await inner.get_final_response() + except MiddlewareTermination: + # The pipeline's MiddlewareTermination suppression is no longer active once + # 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 + finally: + _LOOP_ITERATION_ACTIVE.reset(loop_token) + + iteration += 1 + + messages_used = context.messages + final = holder["final"] + # 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 + loop_kwargs = self._build_loop_kwargs( + context=context, + iteration=iteration, + last_result=final, + messages_used=messages_used, + original_messages=original_messages, + progress=progress, + ) - work_iterations += 1 - # Decide whether to stop and capture any feedback from should_continue first, so the - # feedback is available to both the progress and next-message callables this iteration. - stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) - loop_kwargs = self._build_loop_kwargs( - context=context, - iteration=iteration, - last_result=final, - messages_used=messages_used, - original_messages=original_messages, - progress=progress, - feedback=feedback, - ) - if await self._record_progress(final, loop_kwargs, progress): + work_iterations += 1 + # Decide whether to stop and capture any feedback from should_continue first, so the + # feedback is available to both the progress and next-message callables this iteration. + stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -632,20 +665,32 @@ async def _generator() -> Any: progress=progress, feedback=feedback, ) - if stop: - return - 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 - # been propagated and is safe to discard here. - self._restore_session(context.session, snapshot) - next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages) - context.messages = next_messages - # Surface the injected "nudge" messages in the stream so consumers see the user - # turns that drive each subsequent iteration (the equivalent of the aggregated - # transcript that non-streaming runs return). - for message in next_messages: - yield self._message_to_update(message) + if await self._record_progress(final, loop_kwargs, progress): + loop_kwargs = self._build_loop_kwargs( + context=context, + iteration=iteration, + last_result=final, + messages_used=messages_used, + original_messages=original_messages, + progress=progress, + feedback=feedback, + ) + if stop: + return + 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 + # been propagated and is safe to discard here. + self._restore_session(context.session, snapshot) + next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages) + context.messages = next_messages + # Surface the injected "nudge" messages in the stream so consumers see the user + # turns that drive each subsequent iteration (the equivalent of the aggregated + # transcript that non-streaming runs return). + for message in next_messages: + yield self._message_to_update(message) + finally: + 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/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 9a6dab9a7c1..810961dc77a 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -424,8 +424,15 @@ class ContextProvider: Attributes: source_id: Unique identifier for this provider instance (required). Used for message/tool attribution so other providers can filter. + after_run_once_per_turn: When True, ``after_run`` is scoped to the user + turn instead of the individual agent run: inside an + ``AgentLoopMiddleware`` loop it is deferred until the loop ends. + Providers that mutate persisted history (e.g. compaction) opt in, + since firing mid-task would rewrite history the task still needs. """ + after_run_once_per_turn: bool = False + def __init__(self, source_id: str): """Initialize the provider. diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index b3e98f1724b..61b9f3a2cb0 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -22,6 +22,7 @@ ChatResponse, ChatResponseUpdate, Content, + ContextProvider, JudgeVerdict, Message, MiddlewareTermination, @@ -1395,3 +1396,78 @@ def should_continue(*, iteration: int, **kwargs: Any) -> bool: assert client.call_count == 1 assert calls == [] assert AgentLoopMiddleware._has_pending_approval_request(final) is True + + +# region turn-scoped after_run providers (#7236) + + +class _TurnScopedRecordingProvider(ContextProvider): + after_run_once_per_turn = True + + def __init__(self) -> None: + super().__init__("turn-scoped-test") + self.after_calls = 0 + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + self.after_calls += 1 + + +class _RunScopedRecordingProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("run-scoped-test") + self.after_calls = 0 + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + self.after_calls += 1 + + +async def test_turn_scoped_after_run_fires_once_per_loop() -> None: + client = RecordingChatClient() + turn_scoped = _TurnScopedRecordingProvider() + run_scoped = _RunScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=3)], + context_providers=[turn_scoped, run_scoped], + ) + + await agent.run("start") + + assert client.call_count == 3 + assert turn_scoped.after_calls == 1 + assert run_scoped.after_calls == 3 + + +async def test_turn_scoped_after_run_fires_once_per_loop_streaming() -> None: + client = RecordingChatClient() + turn_scoped = _TurnScopedRecordingProvider() + run_scoped = _RunScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=3)], + context_providers=[turn_scoped, run_scoped], + ) + + stream = agent.run("start", stream=True) + _ = [update async for update in stream] + await stream.get_final_response() + + assert client.call_count == 3 + assert turn_scoped.after_calls == 1 + assert run_scoped.after_calls == 3 + + +async def test_turn_scoped_after_run_fires_normally_without_loop() -> None: + client = RecordingChatClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent(client=client, context_providers=[turn_scoped]) + + await agent.run("start") + + assert turn_scoped.after_calls == 1 + + +def test_compaction_provider_defers_to_turn_boundary() -> None: + from agent_framework._compaction import CompactionProvider + + assert CompactionProvider(after_strategy=None).after_run_once_per_turn is True From 213829c0c2677c05d06e1bb040d853b04ef2a1b1 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 24 Jul 2026 03:22:13 +0800 Subject: [PATCH 2/3] Python: key loop suppression to the looping agent and pass run options through Two review follow-ups: the contextvar now carries the agent instance so a nested agent.run() inside a loop iteration is not suppressed as if it were an iteration, and the boundary SessionContext forwards the original run options to turn-scoped providers. --- .../packages/core/agent_framework/_agents.py | 12 ++-- .../core/agent_framework/_harness/_loop.py | 5 +- .../core/tests/core/test_harness_loop.py | 59 +++++++++++++++++++ 3 files changed, 69 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 9f453464a02..b044e07f23b 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -74,10 +74,12 @@ logger = logging.getLogger("agent_framework") -# Set by AgentLoopMiddleware while a loop iteration is running so that -# providers scoped to the whole user turn (``after_run_once_per_turn``) skip -# their per-iteration ``after_run`` and only fire once at the loop boundary. -_LOOP_ITERATION_ACTIVE: ContextVar[bool] = ContextVar("agent_loop_iteration_active", default=False) +# Set by AgentLoopMiddleware to the looping agent while a loop iteration is +# running, so providers scoped to the whole user turn +# (``after_run_once_per_turn``) skip their per-iteration ``after_run`` and only +# fire once at the loop boundary. Keyed to the agent instance: a nested +# ``agent.run()`` inside the iteration is a separate run, not a loop iteration. +_LOOP_ITERATION_ACTIVE: ContextVar[Any] = ContextVar("agent_loop_iteration_active", default=None) if TYPE_CHECKING: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) @@ -556,7 +558,7 @@ async def _run_after_providers( per_service_call_history_required = self.require_per_service_call_history_persistence and any( isinstance(provider, HistoryProvider) for provider in self.context_providers ) - in_loop_iteration = _LOOP_ITERATION_ACTIVE.get() + in_loop_iteration = _LOOP_ITERATION_ACTIVE.get() is self for provider in reversed(self.context_providers): if per_service_call_history_required and isinstance(provider, HistoryProvider): continue diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index dc221569c48..41b889fb5ef 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -498,6 +498,7 @@ async def _fire_turn_scoped_after_providers( session_id=context.session.session_id if context.session else None, service_session_id=context.session.service_session_id if context.session else None, input_messages=list(input_messages), + options=dict(context.options or {}), ) session_context._response = response await run_after(session=context.session, context=session_context, only_per_turn=True) @@ -518,7 +519,7 @@ async def _process_non_streaming( aggregated_usage: UsageDetails | None = None final_result: AgentResponse | None = None while True: - loop_token = _LOOP_ITERATION_ACTIVE.set(True) + loop_token = _LOOP_ITERATION_ACTIVE.set(context.agent) try: await call_next() finally: @@ -612,7 +613,7 @@ async def _generator() -> Any: progress: list[str] = [] try: while True: - loop_token = _LOOP_ITERATION_ACTIVE.set(True) + loop_token = _LOOP_ITERATION_ACTIVE.set(context.agent) try: await call_next() inner = context.result diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index 61b9f3a2cb0..4f3eb29890e 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -32,6 +32,7 @@ background_tasks_running, background_tasks_running_message, set_agent_mode, + tool, todos_remaining, todos_remaining_message, ) @@ -42,6 +43,8 @@ AgentLoopMiddleware, ) +from .conftest import MockBaseChatClient + class RecordingChatClient(BaseChatClient[ChatOptions[None]]): """A minimal chat client that records inputs and returns scripted responses. @@ -1471,3 +1474,59 @@ def test_compaction_provider_defers_to_turn_boundary() -> None: from agent_framework._compaction import CompactionProvider assert CompactionProvider(after_strategy=None).after_run_once_per_turn is True + + +class _OptionsRecordingProvider(ContextProvider): + after_run_once_per_turn = True + + def __init__(self) -> None: + super().__init__("options-recorder") + self.seen_options: list[dict[str, Any]] = [] + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + self.seen_options.append(dict(context.options)) + + +async def test_turn_scoped_after_run_sees_run_options() -> None: + client = RecordingChatClient() + provider = _OptionsRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=2)], + context_providers=[provider], + ) + + await agent.run("start", options={"custom_flag": "yes"}) + + assert provider.seen_options == [{"custom_flag": "yes"}] + + +async def test_nested_agent_run_inside_loop_iteration_is_not_suppressed() -> None: + inner_provider = _TurnScopedRecordingProvider() + inner_agent = Agent(client=RecordingChatClient(), context_providers=[inner_provider]) + + @tool(name="run_inner", approval_mode="never_require") + async def run_inner() -> str: + await inner_agent.run("nested task") + return "done" + + outer_client = MockBaseChatClient() + outer_client.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="run_inner", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["finished"])), + ] + outer = Agent( + client=outer_client, + tools=[run_inner], + middleware=[AgentLoopMiddleware(always_continue, max_iterations=1)], + context_providers=[_TurnScopedRecordingProvider()], + ) + + await outer.run("start") + + assert inner_provider.after_calls == 1 From 6dc2e1763da2577f50ce22e846cdcb899b6492ad Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 4 Aug 2026 02:13:39 +0800 Subject: [PATCH 3/3] fix(core): carry the loop-iteration stamp in run options, not a contextvar The contextvar marker leaked in two ways. Held across a streamed yield it bled into the caller's context, suppressing turn-scoped providers on an unrelated same-agent run while the stream was paused, and a reset from a different consuming task raised on the token. Keyed to the agent instance, it also swallowed the boundary flush of a nested loop on the same agent with its own session. Stamp the runs the loop drives through their options instead. Run options reach only the inner runs (they never enter the model request), a nested or concurrent run starts with fresh options and keeps its own turn, and there is no token to reset, so stream consumption is safe from any task. --- .../packages/core/agent_framework/_agents.py | 17 ++- .../core/agent_framework/_harness/_loop.py | 125 +++++++++--------- .../core/tests/core/test_harness_loop.py | 57 ++++++++ 3 files changed, 131 insertions(+), 68 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index b044e07f23b..fccff1befb3 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -7,7 +7,6 @@ import sys from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack -from contextvars import ContextVar from copy import deepcopy from functools import partial from itertools import chain @@ -74,12 +73,14 @@ logger = logging.getLogger("agent_framework") -# Set by AgentLoopMiddleware to the looping agent while a loop iteration is -# running, so providers scoped to the whole user turn +# AgentLoopMiddleware stamps this key into the run options while a loop +# iteration is running, so providers scoped to the whole user turn # (``after_run_once_per_turn``) skip their per-iteration ``after_run`` and only -# fire once at the loop boundary. Keyed to the agent instance: a nested -# ``agent.run()`` inside the iteration is a separate run, not a loop iteration. -_LOOP_ITERATION_ACTIVE: ContextVar[Any] = ContextVar("agent_loop_iteration_active", default=None) +# fire once at the loop boundary. It rides the run's options rather than a +# context variable: options reach only the runs the loop itself drives, so a +# nested ``agent.run()`` (fresh options, its own session) keeps its own turn, +# and nothing leaks into the caller's context while a stream is paused. +_LOOP_ITERATION_TOKEN_KEY = "_agent_loop_iteration" if TYPE_CHECKING: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) @@ -558,7 +559,9 @@ async def _run_after_providers( per_service_call_history_required = self.require_per_service_call_history_persistence and any( isinstance(provider, HistoryProvider) for provider in self.context_providers ) - in_loop_iteration = _LOOP_ITERATION_ACTIVE.get() is self + # The loop stamps the runs it drives via their options; anything else + # (nested run, caller-side run while a stream is paused) is its own turn. + in_loop_iteration = context.options.get(_LOOP_ITERATION_TOKEN_KEY) is not None for provider in reversed(self.context_providers): if per_service_call_history_required and isinstance(provider, HistoryProvider): continue diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index 41b889fb5ef..77135557672 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -34,7 +34,7 @@ from pydantic import BaseModel, Field from typing_extensions import Self -from .._agents import _LOOP_ITERATION_ACTIVE +from .._agents import _LOOP_ITERATION_TOKEN_KEY from .._feature_stage import ExperimentalFeature, experimental from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination from .._sessions import SessionContext @@ -518,57 +518,45 @@ async def _process_non_streaming( aggregated: list[Message] = [] aggregated_usage: UsageDetails | None = None final_result: AgentResponse | None = None - while True: - loop_token = _LOOP_ITERATION_ACTIVE.set(context.agent) - try: + stamped_options = dict(context.options) if context.options is not None else {} + stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object() + context.options = stamped_options + try: + while True: await call_next() - finally: - _LOOP_ITERATION_ACTIVE.reset(loop_token) - iteration += 1 - - result = context.result - if not isinstance(result, AgentResponse): - raise TypeError( - "AgentLoopMiddleware expected an AgentResponse from a non-streaming run, " - f"got {type(result).__name__}." - ) + iteration += 1 - final_result = result - aggregated.extend(result.messages) - if result.usage_details is not None: - aggregated_usage = add_usage_details(aggregated_usage, result.usage_details) - - # Escape hatch: if this iteration is asking for tool approval, stop and return the - # response so the caller can approve, instead of continuing or injecting next_message. - if self._has_pending_approval_request(result): - break - - messages_used = context.messages - loop_kwargs = self._build_loop_kwargs( - context=context, - iteration=iteration, - last_result=result, - messages_used=messages_used, - original_messages=original_messages, - progress=progress, - ) + result = context.result + if not isinstance(result, AgentResponse): + raise TypeError( + "AgentLoopMiddleware expected an AgentResponse from a non-streaming run, " + f"got {type(result).__name__}." + ) - work_iterations += 1 - # Decide whether to stop and capture any feedback from should_continue first, so the - # feedback is available to both the progress and next-message callables this iteration. - stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) - loop_kwargs = self._build_loop_kwargs( - context=context, - iteration=iteration, - last_result=result, - messages_used=messages_used, - original_messages=original_messages, - progress=progress, - feedback=feedback, - ) - # Capture this iteration's progress entry, then refresh loop_kwargs so the next-message - # resolution sees the latest entry. - if await self._record_progress(result, loop_kwargs, progress): + final_result = result + aggregated.extend(result.messages) + if result.usage_details is not None: + aggregated_usage = add_usage_details(aggregated_usage, result.usage_details) + + # Escape hatch: if this iteration is asking for tool approval, stop and return the + # response so the caller can approve, instead of continuing or injecting next_message. + if self._has_pending_approval_request(result): + break + + messages_used = context.messages + loop_kwargs = self._build_loop_kwargs( + context=context, + iteration=iteration, + last_result=result, + messages_used=messages_used, + original_messages=original_messages, + progress=progress, + ) + + work_iterations += 1 + # Decide whether to stop and capture any feedback from should_continue first, so the + # feedback is available to both the progress and next-message callables this iteration. + stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -578,15 +566,29 @@ async def _process_non_streaming( progress=progress, feedback=feedback, ) - if stop: - break - if snapshot is not None and context.session is not None: - # Reset the session to the pre-loop baseline so the next run starts fresh; only the - # progress log (injected by _resolve_next_message) carries continuity forward. - self._restore_session(context.session, snapshot) - next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages) - context.messages = next_messages - aggregated.extend(next_messages) + # Capture this iteration's progress entry, then refresh loop_kwargs so the next-message + # resolution sees the latest entry. + if await self._record_progress(result, loop_kwargs, progress): + loop_kwargs = self._build_loop_kwargs( + context=context, + iteration=iteration, + last_result=result, + messages_used=messages_used, + original_messages=original_messages, + progress=progress, + feedback=feedback, + ) + if stop: + break + if snapshot is not None and context.session is not None: + # Reset the session to the pre-loop baseline so the next run starts fresh; only the + # progress log (injected by _resolve_next_message) carries continuity forward. + self._restore_session(context.session, snapshot) + next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages) + context.messages = next_messages + aggregated.extend(next_messages) + finally: + context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) if not self.return_final_only: context.result = self._aggregate_response(final_result, aggregated, aggregated_usage) @@ -611,9 +613,11 @@ async def _generator() -> Any: iteration = 0 work_iterations = 0 progress: list[str] = [] + stamped_options = dict(context.options) if context.options is not None else {} + stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object() + context.options = stamped_options try: while True: - loop_token = _LOOP_ITERATION_ACTIVE.set(context.agent) try: await call_next() inner = context.result @@ -633,8 +637,6 @@ async def _generator() -> Any: # 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 - finally: - _LOOP_ITERATION_ACTIVE.reset(loop_token) iteration += 1 @@ -691,6 +693,7 @@ async def _generator() -> Any: for message in next_messages: yield self._message_to_update(message) finally: + context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages) def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index 4f3eb29890e..781e54155fd 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -1530,3 +1530,60 @@ async def run_inner() -> str: await outer.run("start") assert inner_provider.after_calls == 1 + + +async def test_nested_same_agent_run_with_separate_session_is_not_suppressed() -> None: + turn_scoped = _TurnScopedRecordingProvider() + holder: dict[str, Any] = {} + + @tool(name="run_nested", approval_mode="never_require") + async def run_nested() -> str: + # Recursively running the same agent on its own session is a separate + # turn boundary, so its turn-scoped providers must still fire. + await holder["agent"].run("nested task", session=AgentSession()) + return "done" + + outer_client = MockBaseChatClient() + outer_client.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="run_nested", arguments="{}")], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["finished"])), + ] + outer = Agent( + client=outer_client, + tools=[run_nested], + middleware=[AgentLoopMiddleware(always_continue, max_iterations=1)], + context_providers=[turn_scoped], + ) + holder["agent"] = outer + + await outer.run("start", session=AgentSession()) + + # once for the loop's turn boundary, once for the nested run's own turn + assert turn_scoped.after_calls == 2 + + +async def test_concurrent_same_agent_run_during_stream_pause_is_not_suppressed() -> None: + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=RecordingChatClient(), + middleware=[AgentLoopMiddleware(always_continue, max_iterations=1)], + context_providers=[turn_scoped], + ) + + stream = agent.run("start", stream=True) + first = True + async for _update in stream: + if first: + first = False + # a run started by the caller while the outer stream is paused is + # its own turn, not a loop iteration + await agent.run("concurrent turn", session=AgentSession()) + await stream.get_final_response() + + # once for the loop boundary, once for the concurrent run's own boundary + assert turn_scoped.after_calls == 2