github-featured/microsoft--agent-framework
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Scope runtime checkpoint storage to its owning run

Browse Source 
Close the stream-drop race where a dropped run's deferred async-generator finalizer could leave a runtime checkpoint storage override set (inherited by a new run) or clear a successor run's storage. run() now defensively clears any stale override before starting, and _run_core only clears the override if this run still owns it (mirroring the _active_run ownership guard). Adds regression tests for both the inheritance and clobber cases.
This commit is contained in:
Tao Chen
2026-06-24 09:56:04 -07:00
parent 714016a035
commit fcda092fa4
2 changed files with 105 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/_workflows/_workflow.py
+25 -5
View File
@@ -371,6 +371,12 @@ class Workflow(DictConvertible):
# so a subsequent ``run()`` is allowed.
self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None
# Identifies which run currently owns the runtime checkpoint storage override
# on the RunnerContext (keyed by that run's ``_active_run`` weakref). Used to
# scope clearing so a deferred async-generator finalizer from a dropped run
# cannot clobber a successor run's storage.
self._runtime_storage_owner: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None
@property
def status(self) -> WorkflowRunState:
"""Return the current run-level status of this workflow instance.
@@ -756,6 +762,13 @@ class Workflow(DictConvertible):
"Workflow is already running; concurrent runs are not allowed on the same instance."
)
# No run is active, so any runtime checkpoint storage override still set on the
# context is stale - left over from a prior run whose stream was dropped before
# its async-generator finalizer ran. Clear it so this run starts clean and does
# not silently inherit the prior run's runtime checkpoint storage.
self._runner.context.clear_runtime_checkpoint_storage()
self._runtime_storage_owner = None
response_stream = ResponseStream[WorkflowEvent, WorkflowRunResult](
self._run_core(
message=message,
@@ -790,10 +803,6 @@ class Workflow(DictConvertible):
Yields:
WorkflowEvent: The events generated during the workflow execution.
"""
# Enable runtime checkpointing if storage provided
if checkpoint_storage is not None:
self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
# Capture the weakref instance ``run()`` installed for *this* run. We
# compare by object identity in the finally so a stale finalizer (e.g.
# the caller dropped this stream after partial iteration, then started
@@ -804,6 +813,13 @@ class Workflow(DictConvertible):
# here it already points at our own ``ResponseStream``.
my_active_run = self._active_run
# Enable runtime checkpointing if storage provided, and record that this run
# owns the override so the ownership-guarded clear in the finally only fires
# for this run (and a deferred finalizer can't clear a successor's storage).
if checkpoint_storage is not None:
self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
self._runtime_storage_owner = my_active_run
try:
# Async validation: a fresh-message run is only allowed when the
# runner context has fully drained from any prior run. If it still
@@ -853,8 +869,12 @@ class Workflow(DictConvertible):
# it would silently break the successor's concurrency guard.
if self._active_run is my_active_run:
self._active_run = None
if checkpoint_storage is not None:
# Only clear the runtime checkpoint storage if this run still owns it. A
# dropped run's deferred finalizer must not clear a successor run's storage
# (mirrors the ``_active_run`` ownership guard above).
if checkpoint_storage is not None and self._runtime_storage_owner is my_active_run:
self._runner.context.clear_runtime_checkpoint_storage()
self._runtime_storage_owner = None
@staticmethod
def _finalize_events(
python/packages/core/tests/workflow/test_workflow.py
+80
View File
@@ -994,6 +994,86 @@ async def test_workflow_partial_stream_does_not_clobber_successor_active_run() -
await asyncio.sleep(0)
async def test_workflow_stale_runtime_checkpoint_storage_not_inherited() -> None:
"""A new run must not inherit a prior run's leftover runtime checkpoint storage.
If a run that set a runtime ``checkpoint_storage`` override is dropped before
its async-generator finalizer clears it, the override can linger on the
``RunnerContext`` while ``_is_run_active()`` already reports False. ``run()``
defensively clears that stale override so a subsequent run that does not pass
its own ``checkpoint_storage`` does not silently checkpoint into it.
"""
with tempfile.TemporaryDirectory() as temp_dir:
leftover_storage = FileCheckpointStorage(temp_dir)
executor = IncrementExecutor(id="stale_storage_exec", limit=3, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
# Simulate a leftover runtime override from a dropped prior run.
workflow._runner.context.set_runtime_checkpoint_storage(leftover_storage) # type: ignore[attr-defined]
# A fresh run without its own checkpoint_storage must not use the leftover.
result = await workflow.run(NumberMessage(data=0))
assert result.get_final_state() == WorkflowRunState.IDLE
checkpoints = await leftover_storage.list_checkpoints(workflow_name=workflow.name)
assert checkpoints == [], "Stale runtime checkpoint storage must not be inherited by a new run"
assert workflow._runner.context._runtime_checkpoint_storage is None # type: ignore[attr-defined]
async def test_workflow_partial_stream_does_not_clobber_successor_runtime_storage() -> None:
"""A stale ``_run_core`` finalizer must not clear a successor's runtime storage.
Same GC-finalizer race as
``test_workflow_partial_stream_does_not_clobber_successor_active_run`` but for the
runtime checkpoint storage override: the dropped run's deferred ``finally`` must
only clear the override if it still owns it, otherwise it wipes the successor
run's storage.
"""
with (
tempfile.TemporaryDirectory() as temp_dir_a,
tempfile.TemporaryDirectory() as temp_dir_b,
):
storage_a = FileCheckpointStorage(temp_dir_a)
storage_b = FileCheckpointStorage(temp_dir_b)
executor = IncrementExecutor(id="storage_finalizer_exec", limit=100, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
context = workflow._runner.context # type: ignore[attr-defined]
# Step 1: drive stream A's body to its first yield so it set storage_a.
stream_a = workflow.run(NumberMessage(data=0), checkpoint_storage=storage_a, stream=True)
aiter_a = stream_a.__aiter__()
await aiter_a.__anext__()
assert context._runtime_checkpoint_storage is storage_a
# Step 2: drop stream A; the weakref dies and async-gen close is scheduled
# but not run inline.
del stream_a
del aiter_a
gc.collect()
# Step 3: synchronously start stream B with its own storage and drive it to
# its first yield so it set storage_b and took ownership of the override.
stream_b = workflow.run(NumberMessage(data=0), checkpoint_storage=storage_b, stream=True)
aiter_b = stream_b.__aiter__()
await aiter_b.__anext__()
assert context._runtime_checkpoint_storage is storage_b
# Step 4: yield enough for stream A's scheduled aclose to drive its body
# through ``GeneratorExit`` and into its ``finally``.
for _ in range(5):
await asyncio.sleep(0)
# With the ownership guard, stream B's override survives. Without it, A's
# stale finalizer would have cleared it.
assert context._runtime_checkpoint_storage is storage_b
# Tear down stream B.
del stream_b
del aiter_b
gc.collect()
await asyncio.sleep(0)
class _StreamingTestAgent(BaseAgent):
"""Test agent that supports both streaming and non-streaming modes."""