Python: Next workflow run loses its state writes when a dropped run is finalized late #9473

Observed Behavior

If I stop iterating a workflow stream in the middle of a superstep and start a new run on the same Workflow before the dropped stream's async generator has been finalized, the late finalizer discards the new run's pending state writes. The new run still finishes IDLE with its outputs, so the lost write is silent: nothing is raised or logged.

Expected Behavior

Cleanup from the abandoned run should only touch its own superstep. The second run's ctx.set_state("b_key", ...) should be committed and visible in the workflow state after it finishes.

Steps to Reproduce

1. Build a workflow whose executor stages a state write and then blocks (run A).
2. Iterate run A's stream until executor_invoked, then drop the stream and call gc.collect(). The generator's aclose() is now scheduled on the loop but has not run.
3. Start run B on the same workflow. Its handler stages b_key and waits.
4. Let the loop run A's finalizer, then let B finish.
5. b_key is missing from the state, although B completed normally.

Minimal Reproduction

"""A dropped workflow stream, finalized late, discards the next run's state writes.

Run A stages a state write and blocks inside its superstep. The caller stops
iterating A's stream and starts run B on the same Workflow before A's async
generator finalizer has run. B stages its own write. When A's finalizer runs,
it discards the pending writes, which now belong to B. B finishes IDLE with its
output and its write is gone.
"""

import asyncio
import gc
from dataclasses import dataclass

from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler


@dataclass
class Msg:
    run: str


a_staged = asyncio.Event()
b_staged = asyncio.Event()
release_b = asyncio.Event()


class Writer(Executor):
    @handler
    async def handle(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None:
        if message.run == "A":
            ctx.set_state("a_key", "from-run-A")
            a_staged.set()
            await asyncio.Event().wait()  # never finishes; A is abandoned
        else:
            ctx.set_state("b_key", "from-run-B")
            b_staged.set()
            await release_b.wait()
            await ctx.yield_output("b-done")


async def main() -> None:
    workflow = WorkflowBuilder(start_executor=Writer(id="writer")).build()

    # Run A: iterate until the executor is invoked, then drop the stream.
    stream_a = workflow.run(Msg("A"), stream=True)
    it_a = stream_a.__aiter__()
    while True:
        event = await it_a.__anext__()
        if event.type == "executor_invoked":
            break
    await a_staged.wait()
    del stream_a, it_a, event
    gc.collect()  # A's finalizer is now scheduled on the loop, not yet run

    # Run B on the same workflow.
    async def consume_b():
        outputs, final_state = [], None
        async for ev in workflow.run(Msg("B"), stream=True):
            if ev.type == "output":
                outputs.append(ev.data)
            elif ev.type == "status":
                final_state = ev.state
        return outputs, final_state

    task_b = asyncio.create_task(consume_b())
    await asyncio.wait_for(b_staged.wait(), 5)

    # Let A's deferred finalizer run, then let B finish.
    for _ in range(20):
        await asyncio.sleep(0.01)
    release_b.set()
    outputs, final_state = await asyncio.wait_for(task_b, 5)

    b_key = workflow._runner.state.export_state().get("b_key")
    print(f"run B outputs: {outputs}")
    print(f"run B final state: {final_state}")
    print(f"observed b_key after run B: {b_key!r}")
    print("expected b_key after run B: 'from-run-B'")
    print("BUG REPRODUCED" if b_key != "from-run-B" else "not reproduced")


asyncio.run(main())

Output:

run B outputs: ['b-done']
run B final state: WorkflowRunState.IDLE
observed b_key after run B: None
expected b_key after run B: 'from-run-B'
BUG REPRODUCED

With State.discard patched to a no-op the same script prints b_key as 'from-run-B', so the write is lost in the discard.

Error Messages and Stack Traces

None. No exception and no log output.

Package Versions

agent-framework-core: 1.22.0 (python-1.22.0)

Python Version

Python 3.12.3

Operating System

Linux

Regression

Yes

Additional Context

The cause is the finally: if not committed: self._state.discard() around each superstep in RunnerImpl.run_until_convergence (python/packages/core/agent_framework/_workflows/_runner.py lines 127-171 at 44fe223, the discard at 169-171), added in #8819 for #7859. The State object is created once in Workflow.__init__ (_workflow.py line 395) and shared by every run on that instance, so when run A's generator is closed late its discard() clears whatever is pending at that moment, which is run B's write. python-1.21.0 has no discard in the runner, and with the discard disabled the same script commits run B's write.

Workflow._run_core already handles this ordering for its own cleanup: the comment at _workflow.py 921-928 describes exactly this case (a stream dropped after partial iteration, then a new run started before the finalizer throws GeneratorExit), and owns_run = self._active_run is my_active_run at line 1020 keeps the stale finalizer from clearing the successor's run lock. The runner's discard has no such check.

Proposed fix: scope the discard to the run that owns the superstep, for example record a per-run token on the runner when a superstep starts and skip discard() in the finally when the token no longer matches (task cancellation still runs), the same way owns_run guards the cleanup in _run_core. Another option is for a new run to close and await any still-live previous runner generator before it touches State. A regression test modelled on test_workflow_partial_stream_does_not_clobber_successor_active_run that asserts the successor's write survives would pin it. I'm happy to open a PR for this.

Acknowledgements

• I searched existing issues and did not find a duplicate.
• I personally verified this behavior and the reproduction details are authentic.
• I will wait for explicit maintainer agreement before starting implementation of a non-trivial change.