diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 05606d0a5f..dddadfdfe7 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -84,6 +84,15 @@ logger = logging.getLogger("agent_framework") +# 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. 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) else: @@ -545,6 +554,7 @@ 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. @@ -557,6 +567,10 @@ async def _run_after_providers( 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. """ if _defer_run_persistence(partial(self._run_after_providers, session=session, context=context)): return @@ -570,9 +584,17 @@ 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 ) + # 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 + 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 96e0d1a126..e8b46a008a 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -1529,6 +1529,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 7640d624cf..7713555767 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_TOKEN_KEY from .._feature_stage import ExperimentalFeature, experimental from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination +from .._sessions import SessionContext from .._types import ( AgentResponse, AgentResponseUpdate, @@ -472,6 +474,35 @@ 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), + options=dict(context.options or {}), + ) + 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, @@ -487,53 +518,45 @@ async def _process_non_streaming( aggregated: list[Message] = [] aggregated_usage: UsageDetails | None = None final_result: AgentResponse | None = None - while True: - await call_next() - 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__}." - ) + 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() + 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, @@ -543,18 +566,37 @@ 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) + 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 +613,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, - ) + 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: + 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, + ) - 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 +668,33 @@ 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: + 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: 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 893bb96e52..6b637e5d9f 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -753,8 +753,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 b3e98f1724..781e54155f 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, @@ -31,6 +32,7 @@ background_tasks_running, background_tasks_running_message, set_agent_mode, + tool, todos_remaining, todos_remaining_message, ) @@ -41,6 +43,8 @@ AgentLoopMiddleware, ) +from .conftest import MockBaseChatClient + class RecordingChatClient(BaseChatClient[ChatOptions[None]]): """A minimal chat client that records inputs and returns scripted responses. @@ -1395,3 +1399,191 @@ 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 + + +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 + + +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