From c332197e35d258376a5facf6cf2267cfbed618e7 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:07:24 +0000 Subject: [PATCH 1/3] Python: Add reset() to workflow Introduce Workflow.reset(), which restores a workflow instance to its captured initial state so the same instance can be re-run cleanly (e.g. in a hosted-agent session where the instance is created once). The initial state is captured once per instance as an in-memory baseline checkpoint before the first run; reset() is rejected while a run is active. Also harden Runner.capture_checkpoint_object() to accept optional metadata and to reject capture while in-flight executor messages are present (mid-superstep state is not a clean baseline), raising WorkflowCheckpointException. This is part of breaking down the changes in #6407 into smaller PRs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 4 +- .../agent_framework/_workflows/_runner.py | 35 ++++++-- .../agent_framework/_workflows/_workflow.py | 46 ++++++++++- .../core/tests/workflow/test_workflow.py | 82 +++++++++++++++++++ 4 files changed, 156 insertions(+), 11 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 31124b70dff..ebf8300377b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -154,7 +154,9 @@ agent_framework/ ### Workflows (`_workflows/`) -- **`Workflow`** - Graph-based workflow definition +- **`Workflow`** - Graph-based workflow definition. `reset()` restores a workflow instance to + its captured initial state (an in-memory baseline checkpoint taken before the first run) so + the same instance can be reused for a fresh run; it is only allowed while no run is active. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 360e85099a0..d63b6791163 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -340,22 +340,38 @@ async def restore_from_checkpoint( logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}") raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint_id}") from e - async def capture_checkpoint_object(self) -> WorkflowCheckpoint: + async def capture_checkpoint_object( + self, + metadata: dict[str, Any] | None = None, + ) -> WorkflowCheckpoint: """Capture the current runner state as an in-memory ``WorkflowCheckpoint``. - Builds a checkpoint from committed shared state, executor snapshots, any in-flight - messages, and any pending request_info events, without writing to a storage backend. - The caller owns the returned object - for example, a parent ``WorkflowExecutor`` - embedding a child workflow's checkpoint in its own checkpoint payload. + Builds a checkpoint from committed shared state, executor snapshots, and any + pending request_info events, without writing to a storage backend. The caller + owns the returned object - for example, a parent ``WorkflowExecutor`` embedding + a child workflow's checkpoint in its own checkpoint payload, or a ``Workflow`` + capturing its initial baseline for :meth:`Workflow.reset`. + + Capture is only valid when the runner is quiescent with respect to message + delivery: it rejects capture while in-flight executor messages are present, + since those represent mid-superstep state that would not form a clean baseline. + This mirrors the normal per-superstep checkpoint path and makes no assumption + about which caller is taking the checkpoint. - Like any checkpoint, the snapshot is only internally consistent when captured at a - stable point (e.g. a superstep boundary) while no iteration is concurrently mutating - the runner. This mirrors the normal per-superstep checkpoint path and makes no - assumption about which caller is taking the checkpoint. + Args: + metadata: Optional metadata to associate with the checkpoint. Returns: A ``WorkflowCheckpoint`` snapshot of the current runner state. + + Raises: + WorkflowCheckpointException: If in-flight executor messages are present. """ + if await self._ctx.has_messages(): + raise WorkflowCheckpointException( + "Cannot capture checkpoint while in-flight executor messages are present." + ) + # Persist executor snapshots into committed shared state before exporting it. await self._prepare_checkpoint_state() return await self._ctx.create_checkpoint_object( @@ -364,6 +380,7 @@ async def capture_checkpoint_object(self) -> WorkflowCheckpoint: self._state, None, self._iteration, + metadata=metadata, ) async def restore_from_checkpoint_object(self, checkpoint: WorkflowCheckpoint) -> None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 099e03df2ff..dd11eeb8695 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -20,7 +20,7 @@ from .._types import ResponseStream from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span -from ._checkpoint import CheckpointStorage +from ._checkpoint import CheckpointStorage, WorkflowCheckpoint from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY from ._edge import ( EdgeGroup, @@ -371,6 +371,10 @@ def __init__( # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # In-memory initial checkpoint captured from the just-built workflow state. + # This is internal-only and used by ``reset()``. + self._initial_checkpoint: WorkflowCheckpoint | None = None + @property def status(self) -> WorkflowRunState: """Return the current run-level status of this workflow instance. @@ -474,6 +478,44 @@ def get_executors_list(self) -> list[Executor]: """Get the list of executors in the workflow.""" return list(self.executors.values()) + async def _ensure_initial_checkpoint(self) -> None: + """Capture the in-memory initial checkpoint once for this workflow instance.""" + if self._initial_checkpoint is not None: + return + + self._initial_checkpoint = await self._runner.capture_checkpoint_object( + metadata={"kind": "initial_in_memory"}, + ) + + async def reset(self) -> None: + """Reset the workflow instance to its captured initial checkpoint state. + + The initial checkpoint is captured in memory once per workflow instance and + is not persisted to external checkpoint storage. + + Raises: + WorkflowException: If called while a workflow run is active. + """ + if self._is_run_active(): + raise WorkflowException( + "Cannot reset workflow while a run is active. " + "Reset is only allowed between runs when the workflow is idle." + ) + + # Capture the baseline if it doesn't exist yet. This is idempotent: on a + # normal reset after one or more runs it's a no-op (the snapshot was taken + # before the first run); when reset is the first operation it captures the + # pristine just-built state so the workflow stays runnable. + await self._ensure_initial_checkpoint() + if self._initial_checkpoint is None: + raise WorkflowException("Workflow initial checkpoint is unavailable.") + + # Restore runner state, executor snapshots, and runtime bookkeeping from the + # in-memory initial checkpoint. + await self._runner.restore_from_checkpoint_object(self._initial_checkpoint) + + self._status = WorkflowRunState.IDLE + async def _run_workflow_with_tracing( self, initial_executor_fn: Callable[[], Awaitable[None]] | None = None, @@ -831,6 +873,8 @@ async def _run_core( "checkpointing; there is no in-process recovery path." ) + await self._ensure_initial_checkpoint() + initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage) async for event in self._run_workflow_with_tracing( diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index c77688c03c2..82735bd83d5 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1612,3 +1612,85 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None # endregion + + +# region Workflow.reset + + +class CounterStateExecutor(Executor): + """Executor with local mutable state used to verify checkpoint-based reset.""" + + def __init__(self, id: str) -> None: + super().__init__(id=id) + self.counter = 0 + + @handler + async def handle(self, message: str, ctx: WorkflowContext[str, int]) -> None: + self.counter += 1 + await ctx.yield_output(self.counter) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"counter": self.counter} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.counter = int(state.get("counter", 0)) + + +class TestWorkflowReset: + """Tests for :meth:`Workflow.reset`.""" + + async def test_reset_restores_initial_shared_state(self) -> None: + """Reset clears accumulated workflow state back to the initial baseline.""" + executor = StateTrackingExecutor(id="state_executor") + workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build() + + result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1")) + assert result1.get_outputs()[0] == ["run1:message1"] + + result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2")) + assert result2.get_outputs()[0] == ["run1:message1", "run2:message2"] + + await workflow.reset() + + result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3")) + assert result3.get_outputs()[0] == ["run3:message3"] + + async def test_reset_restores_executor_checkpoint_state(self) -> None: + """Reset restores per-executor local state captured in the initial checkpoint.""" + executor = CounterStateExecutor(id="counter_executor") + workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build() + + result1 = await workflow.run("one") + assert result1.get_outputs() == [1] + + result2 = await workflow.run("two") + assert result2.get_outputs() == [2] + + await workflow.reset() + + result3 = await workflow.run("three") + assert result3.get_outputs() == [1] + + async def test_reset_before_first_run_is_allowed(self, simple_executor: Executor) -> None: + """Reset can be called before the first run and leaves workflow runnable.""" + workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build() + + await workflow.reset() + + result = await workflow.run("hello") + assert result.get_final_state() == WorkflowRunState.IDLE + + async def test_reset_raises_while_run_active(self, simple_executor: Executor) -> None: + """Reset must reject while a workflow run is active.""" + workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build() + + active_stream = workflow.run(WorkflowMessage(data="hi", source_id="test"), stream=True) + try: + with pytest.raises(WorkflowException, match="Cannot reset workflow while a run is active"): + await workflow.reset() + finally: + async for _ in active_stream: + pass + + +# endregion From a0d75d542d60b09acf43fe376d659f725e133f15 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:11:50 +0000 Subject: [PATCH 2/3] Make capture_checkpoint_object docstring caller-agnostic Describe the quiescence precondition (no in-flight executor messages) in terms of checkpoint coherence rather than any specific caller, so the contract does not assume the WorkflowExecutor is the only consumer. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agent_framework/_workflows/_runner.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index d63b6791163..2645fa552b2 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -348,15 +348,19 @@ async def capture_checkpoint_object( Builds a checkpoint from committed shared state, executor snapshots, and any pending request_info events, without writing to a storage backend. The caller - owns the returned object - for example, a parent ``WorkflowExecutor`` embedding - a child workflow's checkpoint in its own checkpoint payload, or a ``Workflow`` - capturing its initial baseline for :meth:`Workflow.reset`. - - Capture is only valid when the runner is quiescent with respect to message - delivery: it rejects capture while in-flight executor messages are present, - since those represent mid-superstep state that would not form a clean baseline. - This mirrors the normal per-superstep checkpoint path and makes no assumption - about which caller is taking the checkpoint. + owns the returned object and may embed, persist, or restore it however it sees + fit. + + Capture requires the runner to be quiescent with respect to message delivery: + it rejects capture while in-flight executor messages are present, because such + messages represent mid-superstep work that has not yet been delivered. The + snapshot does not carry those queued messages, so capturing mid-superstep would + either silently drop pending work or produce state that cannot be resumed + consistently. A checkpoint is only coherent at a superstep boundary where the + message queue is drained. This precondition is independent of who calls it. + + Pending request_info events are intentionally not blocked: a workflow paused + awaiting responses is itself a stable point. Args: metadata: Optional metadata to associate with the checkpoint. From 013db0ddcac690c67f5f9fef70f73494dd5f2ea3 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:15:45 +0000 Subject: [PATCH 3/3] Move in-flight-message guard into context.create_checkpoint_object Rather than the runner querying has_messages() to guard capture, put the message-quiescence check where the messages live: the context's create_checkpoint_object now rejects creation while in-flight executor messages are present. The persisting create_checkpoint path (superstep boundaries, which legitimately snapshot pending messages) shares a new private _build_checkpoint helper and is not guarded. The runner's capture_checkpoint_object just prepares executor state and delegates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agent_framework/_workflows/_runner.py | 31 +++++------- .../_workflows/_runner_context.py | 48 ++++++++++++++++++- 2 files changed, 58 insertions(+), 21 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 2645fa552b2..b942026bfa5 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -346,21 +346,17 @@ async def capture_checkpoint_object( ) -> WorkflowCheckpoint: """Capture the current runner state as an in-memory ``WorkflowCheckpoint``. - Builds a checkpoint from committed shared state, executor snapshots, and any - pending request_info events, without writing to a storage backend. The caller - owns the returned object and may embed, persist, or restore it however it sees - fit. - - Capture requires the runner to be quiescent with respect to message delivery: - it rejects capture while in-flight executor messages are present, because such - messages represent mid-superstep work that has not yet been delivered. The - snapshot does not carry those queued messages, so capturing mid-superstep would - either silently drop pending work or produce state that cannot be resumed - consistently. A checkpoint is only coherent at a superstep boundary where the - message queue is drained. This precondition is independent of who calls it. - - Pending request_info events are intentionally not blocked: a workflow paused - awaiting responses is itself a stable point. + Persists executor snapshots into committed shared state, then delegates to the + context's :meth:`RunnerContext.create_checkpoint_object` to build the snapshot + from committed shared state, executor snapshots, and any pending request_info + events, without writing to a storage backend. The caller owns the returned object + and may embed, persist, or restore it however it sees fit. + + Creation requires the runner to be quiescent with respect to message delivery; + the context rejects capture while in-flight executor messages are present (a bare + snapshot cannot carry mid-superstep queued messages coherently). Pending + request_info events are not blocked, since a workflow paused awaiting responses is + itself a stable point. Args: metadata: Optional metadata to associate with the checkpoint. @@ -371,11 +367,6 @@ async def capture_checkpoint_object( Raises: WorkflowCheckpointException: If in-flight executor messages are present. """ - if await self._ctx.has_messages(): - raise WorkflowCheckpointException( - "Cannot capture checkpoint while in-flight executor messages are present." - ) - # Persist executor snapshots into committed shared state before exporting it. await self._prepare_checkpoint_state() return await self._ctx.create_checkpoint_object( diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 5a11afd8b60..a59bf898028 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -10,6 +10,7 @@ from enum import Enum from typing import Any, Literal, Protocol, TypeVar, runtime_checkable +from ..exceptions import WorkflowCheckpointException from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint from ._const import INTERNAL_SOURCE_ID from ._events import WorkflowEvent @@ -209,6 +210,14 @@ async def create_checkpoint_object( Unlike :meth:`create_checkpoint`, this does not require checkpoint storage and does not save anything; it returns the checkpoint object for the caller to own. + The object is meant to be a self-contained, resumable snapshot, so it may only be + created when the runner is quiescent with respect to message delivery: it rejects + creation while in-flight executor messages are present. A bare checkpoint object + cannot carry mid-superstep queued messages without silently dropping work or + producing non-resumable state; a coherent snapshot exists only at a drained + superstep boundary. Pending request_info events are intentionally not blocked, + since a workflow paused awaiting responses is itself a stable point. + Args: workflow_name: The name of the workflow for which the checkpoint is being created. graph_signature_hash: Hash of the workflow graph topology to @@ -220,6 +229,9 @@ async def create_checkpoint_object( Returns: A ``WorkflowCheckpoint`` snapshot of the current context state. + + Raises: + WorkflowCheckpointException: If in-flight executor messages are present. """ ... @@ -418,6 +430,40 @@ async def create_checkpoint_object( iteration_count: int, metadata: dict[str, Any] | None = None, ) -> WorkflowCheckpoint: + # A standalone checkpoint object is only coherent when the message queue is + # drained: the snapshot is meant to be self-contained, but in-flight executor + # messages represent mid-superstep work that a bare object cannot carry without + # silently dropping it or producing non-resumable state. The persisting + # ``create_checkpoint`` path (superstep boundaries) legitimately snapshots + # pending messages, so it uses ``_build_checkpoint`` directly and is not guarded. + if self._messages: + raise WorkflowCheckpointException( + "Cannot create checkpoint object while in-flight executor messages are present." + ) + return self._build_checkpoint( + workflow_name, + graph_signature_hash, + state, + previous_checkpoint_id, + iteration_count, + metadata, + ) + + def _build_checkpoint( + self, + workflow_name: str, + graph_signature_hash: str, + state: State, + previous_checkpoint_id: CheckpointID | None, + iteration_count: int, + metadata: dict[str, Any] | None = None, + ) -> WorkflowCheckpoint: + """Build a ``WorkflowCheckpoint`` snapshot from the current context state. + + Shared by :meth:`create_checkpoint_object` (bare, quiescent snapshot) and + :meth:`create_checkpoint` (persisting, superstep-boundary snapshot). This does + not enforce message quiescence; callers that require it guard first. + """ return WorkflowCheckpoint( workflow_name=workflow_name, graph_signature_hash=graph_signature_hash, @@ -443,7 +489,7 @@ async def create_checkpoint( if not storage: raise ValueError("Checkpoint storage not configured") - checkpoint = await self.create_checkpoint_object( + checkpoint = self._build_checkpoint( workflow_name, graph_signature_hash, state,