Python: defer turn-scoped after_run providers to the agent loop boundary - #7289
Python: defer turn-scoped after_run providers to the agent loop boundary#7289Yufeng He (he-yufeng) wants to merge 5 commits into
Conversation
Each AgentLoopMiddleware iteration is a full agent run, so CompactionProvider.after_run fired per iteration and rewrote persisted history mid-task (microsoft#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.
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
Pull request overview
This PR fixes a harness-loop correctness issue where ContextProvider.after_run was firing once per AgentLoopMiddleware iteration (i.e., mid-turn), by allowing providers to opt into turn-scoped after_run semantics and deferring those providers until the loop boundary. This is intended to prevent providers like CompactionProvider from rewriting persisted history while the loop is still actively using it.
Changes:
- Adds
ContextProvider.after_run_once_per_turn(defaultFalse) to mark providers whoseafter_runshould be deferred to the end of anAgentLoopMiddlewareloop. - Introduces loop-iteration tracking via a contextvar and updates
_run_after_providersto skip turn-scoped providers during loop iterations, then fire them once at loop exit. - Marks
CompactionProvideras turn-scoped and adds tests validating once-per-turn vs once-per-run behavior for both streaming and non-streaming runs.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/core/tests/core/test_harness_loop.py | Adds tests asserting turn-scoped providers fire once per loop (streaming/non-streaming) and CompactionProvider opts in. |
| python/packages/core/agent_framework/_sessions.py | Extends ContextProvider with the after_run_once_per_turn opt-in flag and documentation. |
| python/packages/core/agent_framework/_harness/_loop.py | Sets a loop-iteration contextvar around each iteration and fires deferred turn-scoped after_run providers once at loop exit. |
| python/packages/core/agent_framework/_compaction.py | Opts CompactionProvider into once-per-turn after_run semantics. |
| python/packages/core/agent_framework/_agents.py | Adds the loop-iteration contextvar and updates _run_after_providers to implement skip/run logic for turn-scoped providers. |
| 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) |
There was a problem hiding this comment.
Fixed in 213829c. The boundary SessionContext now forwards the original run options, so turn-scoped providers see the same context.options they would on a normal per-run after_run. Covered by a new test that records the options a turn-scoped provider sees at the loop boundary.
| 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 |
There was a problem hiding this comment.
Fixed in 213829c. The contextvar now carries the looping agent instance and _run_after_providers only suppresses when it matches self, so a nested agent.run() inside an iteration keeps firing its own turn-scoped providers. A new test drives an inner agent from a tool inside a loop and asserts its provider still fires once.
…s 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.
GitHub Actions
Inactive
GitHub Actions
Inactive
| f"got {type(inner).__name__}." | ||
| ) | ||
|
|
||
| async for update in inner: |
There was a problem hiding this comment.
Could we avoid retaining _LOOP_ITERATION_ACTIVE across each streamed yield? While the stream is paused, the marker leaks into caller work and suppresses turn-scoped providers on another run of the same agent; pulling the stream from sequential tasks instead raises ValueError during token reset and can invoke the provider twice. This can skip or duplicate CompactionProvider mutations and crash otherwise valid ResponseStream consumption.
There was a problem hiding this comment.
Both points addressed in 1ca036d by replacing the contextvar with a stamp on the run options. The loop writes _agent_loop_iteration into the options of the runs it drives; nothing is set or reset around stream yields anymore, so a paused stream leaves the caller's context untouched and there is no token whose reset could land in a different task. A run the caller starts while the stream is paused gets fresh options and keeps its turn. Covered by a test that runs the same agent from the caller mid-stream and asserts its turn-scoped provider still fires.
| 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) |
There was a problem hiding this comment.
Could the suppression state distinguish a loop invocation rather than only the agent instance? If an iteration recursively runs the same Agent with a separate session, the nested loop's boundary still sees the outer marker and skips its turn-scoped provider; the outer boundary later flushes only the outer session. That leaves the nested session without compaction or any other opted-in persistence hook.
There was a problem hiding this comment.
Same fix covers this: the stamp now travels in run options, which are fresh per agent.run(), so a nested loop on the same agent with its own session is governed by its own stamp, not the outer one. Its boundary flush fires normally and only the outer session's flush is driven by the outer stamp. test_nested_same_agent_run_with_separate_session_is_not_suppressed drives exactly that shape and asserts both boundary flushes fire.
|
Yufeng He (@he-yufeng) I don't see commit 1ca036d showing up here. |
…xtvar 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.
GitHub Actions
Inactive
|
Evan Mattson (@moonbox3) my mistake, the commit landed on the wrong local branch and I pushed that one, so the PR head never had it. It's on the PR now as 6dc2e17 (same change, cherry-picked cleanly, the harness-loop suite passes 96/96 locally). |
GitHub Actions
Inactive
Motivation & Context
With
create_harness_agent(or anyAgentLoopMiddlewareusage) every loop iteration is a full agent run, soCompactionProvider.after_runfires per iteration instead of once per user turn. Fired mid-task, compaction rewrites the persisted history the task still works from (#7236).Description & Review Guide
after_run_once_per_turn(default False onContextProvider). While a loop iteration is active,_run_after_providersskips those providers (tracked through a contextvar the middleware sets around each iteration); when the loop exits, the middleware fires the deferred providers once with the turn-level response viaonly_per_turn=True.CompactionProvideropts in;HistoryProviderkeeps firing per iteration because its persistence is incremental per run._harness/_loop.py(contextvar reset placement and thefinallythat fires the deferred providers) and the skip logic in_run_after_providers.Tests: 4 new tests in test_harness_loop.py (turn-scoped fires once vs run-scoped per iteration, streaming and non-streaming, no-loop baseline, CompactionProvider opt-in). test_harness_loop 92 pass; test_compaction + test_sessions + test_agents 270 pass.
Related Issue
Fixes #7236
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.