Python: Add reset() to workflow

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

python/packages/core/agent_framework/_workflows/_runner.py

@@ -340,21 +340,32 @@ 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.
+        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.
 
-        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.
         """
         # Persist executor snapshots into committed shared state before exporting it.
         await self._prepare_checkpoint_state()
@@ -364,6 +375,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:

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,

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(

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