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..b942026bfa5 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/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: 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, 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