From 4564623e8c0717ad4df12f16312d84987896ae61 Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:05:28 +0300 Subject: [PATCH 1/4] Python: Keep a dropped run's finalizer from discarding state writes RunnerImpl discarded pending state in the finally of every uncommitted superstep, but State is shared by all runs on a Workflow. When a dropped stream was finalized after a new run had started, that discard removed the new run's staged writes. The runner now records a token for the current superstep and only discards when the finishing superstep still owns it. --- .../agent_framework/_workflows/_runner.py | 7 ++- .../core/tests/workflow/test_workflow.py | 55 +++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index f41183cb1c2..b7eb044e58f 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -80,6 +80,9 @@ def __init__( self._iteration = 0 self._max_iterations = max_iterations self._state = state + # Identifies the superstep that owns pending state writes, so a stale run's finalizer + # does not discard writes staged by a newer run on the same shared State. + self._superstep_token: object | None = None # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None @@ -125,6 +128,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: yield WorkflowEvent.superstep_started(iteration=self._iteration + 1) committed = False + superstep_token = object() + self._superstep_token = superstep_token try: # Wake on either a live event or iteration completion, including silent supersteps. iteration_task = asyncio.create_task(self._run_iteration()) @@ -167,7 +172,7 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: self._state.commit() committed = True finally: - if not committed: + if not committed and self._superstep_token is superstep_token: self._state.discard() # Create checkpoint after each superstep iteration diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 76a285e6d93..864d7c22ac4 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1483,6 +1483,61 @@ async def test_workflow_partial_stream_does_not_clobber_successor_active_run() - await asyncio.sleep(0) +async def test_workflow_partial_stream_does_not_discard_successor_state_writes() -> None: + """A stale runner finalizer must not discard a successor run's pending state writes. + + Same GC-finalizer race as + ``test_workflow_partial_stream_does_not_clobber_successor_active_run`` but inside a + superstep: stream A is dropped while its executor has staged a write, stream B stages + its own write, then A's deferred ``finally`` runs. It must only discard A's superstep, + so B's write is committed when B completes. + """ + a_staged = asyncio.Event() + b_staged = asyncio.Event() + release_b = asyncio.Event() + + class _Writer(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None: + if message == "A": + ctx.set_state("a_key", "from-run-A") + a_staged.set() + await asyncio.Event().wait() + else: + ctx.set_state("b_key", "from-run-B") + b_staged.set() + await release_b.wait() + await ctx.yield_output("b-done") + + workflow = WorkflowBuilder(start_executor=_Writer(id="writer")).build() + + # Drive stream A until its executor has staged a write and is blocked, then drop it. + stream_a = workflow.run("A", stream=True) + aiter_a = stream_a.__aiter__() + while (await aiter_a.__anext__()).type != "executor_invoked": + pass + await a_staged.wait() + del stream_a + del aiter_a + gc.collect() + + async def consume_b() -> WorkflowRunResult: + return await workflow.run("B", stream=True).get_final_response() + + task_b = asyncio.create_task(consume_b()) + await asyncio.wait_for(b_staged.wait(), 5) + + # Let stream A's scheduled aclose run its finally while B's write is pending. + for _ in range(20): + await asyncio.sleep(0.01) + release_b.set() + result_b = await asyncio.wait_for(task_b, 5) + + assert result_b.get_outputs() == ["b-done"] + assert result_b.get_final_state() == WorkflowRunState.IDLE + assert workflow._runner.state.get("b_key") == "from-run-B" # pyright: ignore[reportPrivateUsage] + + 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. From cc54a1496975869ea760c0eb7f7ad1576fe0d0c2 Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:44:01 +0300 Subject: [PATCH 2/4] Python: Claim pending state writes before checkpoint preparation A dropped run's deferred finalizer could still own the pending-writes token while a successor run built its entry checkpoint, so it discarded executor snapshots staged between awaited on_checkpoint_save hooks and the checkpoint stored a stale snapshot. Checkpoint preparation now takes ownership of pending writes before calling the hooks. A checkpoint-enabled regression test suspends a save hook while the stale finalizer runs. --- .../agent_framework/_workflows/_runner.py | 5 +- .../core/tests/workflow/test_workflow.py | 75 +++++++++++++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index b7eb044e58f..c7478679673 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -80,8 +80,8 @@ def __init__( self._iteration = 0 self._max_iterations = max_iterations self._state = state - # Identifies the superstep that owns pending state writes, so a stale run's finalizer - # does not discard writes staged by a newer run on the same shared State. + # Identifies the superstep or checkpoint preparation that owns pending state writes, so a + # stale run's finalizer does not discard writes staged by a newer run on the same shared State. self._superstep_token: object | None = None # Checkpointing related attributes @@ -248,6 +248,7 @@ async def _prepare_checkpoint_state(self) -> None: This is used by checkpoint capture paths that need a complete, restorable state payload without necessarily writing to a checkpoint storage backend. """ + self._superstep_token = object() await self._save_executor_states() self._save_edge_runner_states() self._state.commit() diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 864d7c22ac4..78b1c3a3df3 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -39,6 +39,7 @@ handler, response_handler, ) +from agent_framework._workflows._const import EXECUTOR_STATE_KEY @dataclass @@ -1538,6 +1539,80 @@ async def consume_b() -> WorkflowRunResult: assert workflow._runner.state.get("b_key") == "from-run-B" # pyright: ignore[reportPrivateUsage] +async def test_workflow_partial_stream_does_not_discard_successor_entry_checkpoint_snapshots() -> None: + """A stale runner finalizer must not discard snapshots staged by a successor's entry checkpoint. + + Stream A is dropped mid-superstep. Stream B's entry checkpoint stages the writer's snapshot, + then suspends in the next executor's ``on_checkpoint_save`` while A's deferred ``finally`` + runs. B's entry checkpoint must still hold B's writer snapshot. + """ + a_staged = asyncio.Event() + hook_suspended = asyncio.Event() + release_hook = asyncio.Event() + gate_armed = False + + class _Writer(Executor): + def __init__(self, id: str) -> None: + super().__init__(id=id) + self.saves = 0 + + @handler + async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None: + if message == "A": + a_staged.set() + await asyncio.Event().wait() + else: + await ctx.yield_output("b-done") + + async def on_checkpoint_save(self) -> dict[str, Any]: + self.saves += 1 + return {"saves": self.saves} + + class _SuspendingSave(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext) -> None: + pass + + async def on_checkpoint_save(self) -> dict[str, Any]: + if gate_armed: + hook_suspended.set() + await release_hook.wait() + return {} + + storage = InMemoryCheckpointStorage() + writer = _Writer(id="writer") + workflow = ( + WorkflowBuilder(start_executor=writer, checkpoint_storage=storage) + .add_edge(writer, _SuspendingSave(id="gate")) + .build() + ) + + stream_a = workflow.run("A", stream=True) + aiter_a = stream_a.__aiter__() + while (await aiter_a.__anext__()).type != "executor_invoked": + pass + await a_staged.wait() + a_checkpoint_ids = set(await storage.list_checkpoint_ids(workflow_name=workflow.name)) + del stream_a + del aiter_a + gc.collect() + + gate_armed = True + task_b = asyncio.create_task(workflow.run("B", stream=True).get_final_response()) + await asyncio.wait_for(hook_suspended.wait(), 5) + + # Let stream A's scheduled aclose run its finally while B's writer snapshot is pending. + for _ in range(20): + await asyncio.sleep(0.01) + release_hook.set() + result_b = await asyncio.wait_for(task_b, 5) + assert result_b.get_outputs() == ["b-done"] + + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + [b_entry] = [c for c in checkpoints if c.checkpoint_id not in a_checkpoint_ids and c.iteration_count == 0] + assert b_entry.state[EXECUTOR_STATE_KEY]["writer"] == {"saves": 2} + + 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. From f10c1bf49138bf1e85311ac547f49576bf304882 Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:03:37 +0300 Subject: [PATCH 3/4] Python: Discard a dropped run's pending state writes at run start A new run discards any pending state writes before it commits its run kwargs. Those writes can only come from a superstep or checkpoint save that never committed, such as a dropped stream whose finalizer has not run yet, so committing them put a half-finished superstep into the successor run's state. --- python/packages/core/agent_framework/_workflows/_workflow.py | 3 +++ python/packages/core/tests/workflow/test_workflow.py | 1 + 2 files changed, 4 insertions(+) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c60617dabd5..bf52108f1da 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -589,6 +589,9 @@ async def _run_workflow_with_tracing( if not is_continuation: self._runner.reset_iteration_count() + # Pending writes here come from a superstep or checkpoint save that never committed; do not commit them. + self._runner.state.discard() + # Store run kwargs in State so executors can access them. # Per-run kwargs semantics: # - On a fresh message run, prior kwargs go away (set to {} diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 78b1c3a3df3..a86513ea04a 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1537,6 +1537,7 @@ async def consume_b() -> WorkflowRunResult: assert result_b.get_outputs() == ["b-done"] assert result_b.get_final_state() == WorkflowRunState.IDLE assert workflow._runner.state.get("b_key") == "from-run-B" # pyright: ignore[reportPrivateUsage] + assert workflow._runner.state.get("a_key") is None # pyright: ignore[reportPrivateUsage] async def test_workflow_partial_stream_does_not_discard_successor_entry_checkpoint_snapshots() -> None: From e92c3f9f8493fcb866fe6cb6dd1d32984308ec59 Mon Sep 17 00:00:00 2001 From: digit50 <93092234+digit50@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:27:51 +0300 Subject: [PATCH 4/4] Python: Scope pending state writes to the run that made them Pending state writes are now tied to the superstep that owns the shared State, and writes from executor tasks of a superstep that no longer owns it are ignored. A dropped run's executor that writes state during cancellation after a newer run has started could otherwise have its write committed by the newer run. --- .../agent_framework/_workflows/_runner.py | 13 ++-- .../core/agent_framework/_workflows/_state.py | 40 +++++++++++ .../agent_framework/_workflows/_workflow.py | 4 +- .../core/tests/workflow/test_workflow.py | 67 +++++++++++++++++++ 4 files changed, 115 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index c7478679673..7b86fc90b2b 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -80,9 +80,6 @@ def __init__( self._iteration = 0 self._max_iterations = max_iterations self._state = state - # Identifies the superstep or checkpoint preparation that owns pending state writes, so a - # stale run's finalizer does not discard writes staged by a newer run on the same shared State. - self._superstep_token: object | None = None # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None @@ -128,11 +125,10 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: yield WorkflowEvent.superstep_started(iteration=self._iteration + 1) committed = False - superstep_token = object() - self._superstep_token = superstep_token + superstep_epoch = self._state.begin_epoch() try: # Wake on either a live event or iteration completion, including silent supersteps. - iteration_task = asyncio.create_task(self._run_iteration()) + iteration_task = self._state.bound_context().run(asyncio.create_task, self._run_iteration()) event_task: asyncio.Task[WorkflowEvent] | None = None try: while not iteration_task.done(): @@ -172,7 +168,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: self._state.commit() committed = True finally: - if not committed and self._superstep_token is superstep_token: + # A stale run's finalizer must not discard pending writes owned by a newer superstep. + if not committed and self._state.epoch is superstep_epoch: self._state.discard() # Create checkpoint after each superstep iteration @@ -248,7 +245,7 @@ async def _prepare_checkpoint_state(self) -> None: This is used by checkpoint capture paths that need a complete, restorable state payload without necessarily writing to a checkpoint storage backend. """ - self._superstep_token = object() + self._state.begin_epoch() await self._save_executor_states() self._save_edge_runner_states() self._state.commit() diff --git a/python/packages/core/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 136b70de14e..157066399ce 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/python/packages/core/agent_framework/_workflows/_state.py @@ -1,9 +1,17 @@ # Copyright (c) Microsoft. All rights reserved. import copy +import logging from collections.abc import Callable +from contextvars import Context, ContextVar, copy_context from typing import Any +logger = logging.getLogger(__name__) + +# The State and epoch that writes from the current task belong to. Set for each superstep's iteration +# task so executor tasks inherit it; writes made outside any superstep are always accepted. +_bound_epoch: ContextVar[tuple["State", object] | None] = ContextVar("_bound_epoch", default=None) + class State: """Manages shared state across executors within a workflow. @@ -18,6 +26,8 @@ class State: - `get()` checks pending buffer first, then committed state - `commit()` moves all pending changes to committed state (called by Runner at superstep boundary) - `discard()` clears pending changes without committing + - `begin_epoch()` starts a new owner of pending writes; writes, deletes and clears from tasks + bound to an earlier epoch (for example a dropped run's cancelled executor) are ignored Reserved Keys: Keys starting with underscore (_) are reserved for internal framework use. @@ -28,6 +38,30 @@ def __init__(self) -> None: """Initialize the state.""" self._committed: dict[str, Any] = {} self._pending: dict[str, Any] = {} + self._epoch = object() + + @property + def epoch(self) -> object: + """The current write epoch.""" + return self._epoch + + def begin_epoch(self) -> object: + """Start a new write epoch and return it.""" + self._epoch = object() + return self._epoch + + def bound_context(self) -> Context: + """Return a copy of the current context whose writes, and its tasks' writes, belong to the current epoch.""" + context = copy_context() + context.run(_bound_epoch.set, (self, self._epoch)) + return context + + def _is_stale_write(self) -> bool: + bound = _bound_epoch.get() + if bound is None or bound[0] is not self or bound[1] is self._epoch: + return False + logger.debug("Ignoring a state write from a superstep that no longer owns this state.") + return True def set(self, key: str, value: Any) -> None: """Set a value in the pending state buffer. @@ -43,6 +77,8 @@ def set(self, key: str, value: Any) -> None: the .NET behavior and the superstep execution model where all executors in a superstep see the same committed state at the start. """ + if self._is_stale_write(): + return self._pending[key] = copy.deepcopy(value) def get(self, key: str, default: Any = None) -> Any: @@ -88,6 +124,8 @@ def delete(self, key: str) -> None: to indicate deletion at commit time. If it only exists in pending, it is removed from pending. """ + if self._is_stale_write(): + return if key not in self._pending and key not in self._committed: raise KeyError(f"Key '{key}' not found in state.") @@ -100,6 +138,8 @@ def delete(self, key: str) -> None: def clear(self) -> None: """Clear both committed and pending state.""" + if self._is_stale_write(): + return self._committed.clear() self._pending.clear() diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index bf52108f1da..bc60178385e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -589,7 +589,9 @@ async def _run_workflow_with_tracing( if not is_continuation: self._runner.reset_iteration_count() - # Pending writes here come from a superstep or checkpoint save that never committed; do not commit them. + # Pending writes here come from a superstep or checkpoint save that never committed; do not + # commit them, and ignore later writes from that superstep's executors if it is still finalizing. + self._runner.state.begin_epoch() self._runner.state.discard() # Store run kwargs in State so executors can access them. diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index a86513ea04a..9a43e650ad4 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1614,6 +1614,73 @@ async def on_checkpoint_save(self) -> dict[str, Any]: assert b_entry.state[EXECUTOR_STATE_KEY]["writer"] == {"saves": 2} +async def test_workflow_partial_stream_cancellation_write_not_committed_by_successor() -> None: + """A dropped run's cancellation cleanup must not write into a successor run's state. + + Stream A is dropped while its executor is blocked. Stream B passes its run start and stages its + own write, then A is finalized: its executor's ``finally`` writes state during cancellation. + B's superstep commit must keep B's write and must not persist A's cleanup write. + """ + a_blocked = asyncio.Event() + b_staged = asyncio.Event() + release_b = asyncio.Event() + + class _Router(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message, target_id=f"worker_{message.lower()}") + + class _WorkerA(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext) -> None: + try: + a_blocked.set() + await asyncio.Event().wait() + finally: + ctx.set_state("a_cleanup", "from-run-A") + + class _WorkerB(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None: + ctx.set_state("b_key", "from-run-B") + b_staged.set() + await release_b.wait() + await ctx.yield_output("b-done") + + # Separate workers so B's superstep is not serialized behind A's blocked handler. + router = _Router(id="router") + workflow = ( + WorkflowBuilder(start_executor=router) + .add_edge(router, _WorkerA(id="worker_a")) + .add_edge(router, _WorkerB(id="worker_b")) + .build() + ) + + stream_a = workflow.run("A", stream=True) + aiter_a = stream_a.__aiter__() + while (event := await aiter_a.__anext__()).type != "executor_invoked" or event.executor_id != "worker_a": + pass + await a_blocked.wait() + # Keep A's generator so its finalization can be ordered after B's run start. + generator_a = cast(Any, stream_a)._iterator + del stream_a + del aiter_a + gc.collect() + + task_b = asyncio.create_task(workflow.run("B", stream=True).get_final_response()) + await asyncio.wait_for(b_staged.wait(), 5) + + # Finalize A the way the event loop does for a dropped stream. + await generator_a.aclose() + release_b.set() + result_b = await asyncio.wait_for(task_b, 5) + + assert result_b.get_outputs() == ["b-done"] + committed = workflow._runner.state.export_state() # pyright: ignore[reportPrivateUsage] + assert committed["b_key"] == "from-run-B" + assert "a_cleanup" not in committed + + 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.