Add create checkpoint to workflow interface

python/packages/core/agent_framework/__init__.py

@@ -242,6 +242,7 @@ from ._workflows._agent_executor import (
 )
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
+    CheckpointID,
     CheckpointStorage,
     FileCheckpointStorage,
     InMemoryCheckpointStorage,
@@ -285,7 +286,6 @@ from ._workflows._functional import (
     workflow,
 )
 from ._workflows._request_info_mixin import response_handler
-from ._workflows._runner import Runner
 from ._workflows._runner_context import (
     InProcRunnerContext,
     RunnerContext,
@@ -379,6 +379,7 @@ __all__ = [
     "ChatResponse",
     "ChatResponseUpdate",
     "CheckResult",
+    "CheckpointID",
     "CheckpointStorage",
     "ClassSkill",
     "CompactionProvider",
@@ -470,7 +471,6 @@ __all__ = [
     "RoleLiteral",
     "RubricScore",
     "RunContext",
-    "Runner",
     "RunnerContext",
     "SecretString",
     "SelectiveToolCallCompactionStrategy",

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

@@ -65,7 +65,7 @@ class Runner:
 
         # Checkpointing related attributes
         self._resumed_from_checkpoint = False
-        self._previous_checkpoint_id: CheckpointID | None = None
+        self.previous_checkpoint_id: CheckpointID | None = None
 
     @property
     def context(self) -> RunnerContext:
@@ -113,12 +113,12 @@ class Runner:
             for event in await self._ctx.drain_events():
                 yield event
 
-        # Create the first checkpoint. Checkpoints are usually considered to be created at the end of an iteration,
-        # we can think of the first checkpoint as being created at the end of a "superstep 0" which captures the
-        # states after which the start executor has run.  Note that we execute the start executor outside of the
-        # main iteration loop.
+        # Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
+        # end of an iteration, we can think of this checkpoint as being created at the end of a "superstep 0"
+        # which captures the states after which the start executor has run. Note that we execute the start
+        # executor outside of the main iteration loop.
         if await self._ctx.has_messages() and not self._resumed_from_checkpoint:
-            await self._create_checkpoint_if_enabled()
+            await self.create_checkpoint_if_enabled()
 
         while self._iteration < self._max_iterations:
             logger.info(f"Starting superstep {self._iteration + 1}")
@@ -165,7 +165,7 @@ class Runner:
             self._state.commit()
 
             # Create checkpoint after each superstep iteration
-            await self._create_checkpoint_if_enabled()
+            await self.create_checkpoint_if_enabled()
 
             yield WorkflowEvent.superstep_completed(iteration=self._iteration)
 
@@ -231,7 +231,7 @@ class Runner:
         ]
         await asyncio.gather(*tasks)
 
-    async def _create_checkpoint_if_enabled(self) -> None:
+    async def create_checkpoint_if_enabled(self) -> None:
         """Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
         if not self._ctx.has_checkpointing():
             return
@@ -249,7 +249,7 @@ class Runner:
                 self._workflow_name,
                 self._graph_signature_hash,
                 self._state,
-                self._previous_checkpoint_id,
+                self.previous_checkpoint_id,
                 self._iteration,
             )
 
@@ -257,9 +257,9 @@ class Runner:
                 "Created checkpoint: %s with parent checkpoint at iteration %d: %s",
                 checkpoint_id,
                 self._iteration,
-                self._previous_checkpoint_id,
+                self.previous_checkpoint_id,
             )
-            self._previous_checkpoint_id = checkpoint_id
+            self.previous_checkpoint_id = checkpoint_id
         except Exception as e:
             logger.warning(
                 "Failed to create checkpoint at iteration %d: %s. "
@@ -267,7 +267,7 @@ class Runner:
                 "The next successfully-created checkpoint will be parented to the last successful checkpoint: %s",
                 self._iteration,
                 e,
-                self._previous_checkpoint_id,
+                self.previous_checkpoint_id,
             )
 
     async def restore_from_checkpoint(
@@ -396,7 +396,7 @@ class Runner:
         """
         self._resumed_from_checkpoint = True
         self._iteration = checkpoint.iteration_count
-        self._previous_checkpoint_id = checkpoint.checkpoint_id
+        self.previous_checkpoint_id = checkpoint.checkpoint_id
 
     async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
         """Store executor state in state under a reserved key.

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

@@ -18,9 +18,9 @@ from typing import TYPE_CHECKING, Any, Literal, overload
 
 from .._sessions import ContextProvider
 from .._types import ResponseStream
-from ..exceptions import WorkflowException, WorkflowRunnerException
+from ..exceptions import WorkflowCheckpointException, WorkflowException, WorkflowRunnerException
 from ..observability import OtelAttr, capture_exception, create_workflow_span
-from ._checkpoint import CheckpointStorage
+from ._checkpoint import CheckpointID, CheckpointStorage
 from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
 from ._edge import (
     EdgeGroup,
@@ -474,6 +474,43 @@ class Workflow(DictConvertible):
         """Get the list of executors in the workflow."""
         return list(self.executors.values())
 
+    async def create_checkpoint(self, checkpoint_storage: CheckpointStorage | None) -> CheckpointID:
+        """Create a checkpoint of the current workflow state in the provided storage.
+
+        Args:
+            checkpoint_storage: The CheckpointStorage instance where the checkpoint will be stored.
+                If None, will use the workflow's default checkpoint storage if configured, or raise
+                if checkpointing is not enabled.
+
+        Notes:
+            - Checkpoints can only be created when the workflow is idle (not actively running).
+            - Checkpoints are automatically created at the end of each superstep if a checkpoint storage is configured.
+              Use this method only when necessary, for example to capture the initial state of the workflow prior to the
+              first run.
+            - Creating a checkpoint manually will alter the checkpoint lineage. The new checkpoint will become the
+              parent of the next checkpoint created automatically (if checkpointing is enabled by providing a storage).
+        """
+        if self._is_run_active():
+            raise WorkflowException(
+                "Cannot create checkpoint while a workflow run is active. "
+                "Checkpointing is only allowed between runs when the workflow is idle."
+            )
+
+        if checkpoint_storage is None and not self._runner.context.has_checkpointing():
+            raise WorkflowCheckpointException(
+                "Checkpoint storage must be provided to create a checkpoint when checkpointing is not enabled."
+            )
+        if checkpoint_storage is not None:
+            self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
+
+        try:
+            await self._runner.create_checkpoint_if_enabled()
+            if self._runner.previous_checkpoint_id is None:
+                raise WorkflowCheckpointException("Failed to create checkpoint.")
+            return self._runner.previous_checkpoint_id
+        finally:
+            self._runner.context.clear_runtime_checkpoint_storage()
+
     async def _run_workflow_with_tracing(
         self,
         initial_executor_fn: Callable[[], Awaitable[None]] | None = None,
@@ -1165,6 +1202,15 @@ class Workflow(DictConvertible):
             **kwargs,
         )
 
+    def _is_run_active(self) -> bool:
+        """Check if a workflow run is currently active.
+
+        Returns:
+            True if a run is active, False otherwise.
+        """
+        existing_stream = self._active_run() if self._active_run is not None else None
+        return existing_stream is not None
+
     async def reset_for_new_run(self) -> None:
         """Reset the workflow for a new run that is independent from prior runs.
 

python/packages/core/tests/workflow/test_runner.py

@@ -921,7 +921,7 @@ async def test_runner_mark_resumed_sets_previous_checkpoint_id():
     )
 
     # Pre-condition: nothing to chain back to
-    assert runner._previous_checkpoint_id is None  # pyright: ignore[reportPrivateUsage]
+    assert runner.previous_checkpoint_id is None
 
     resumed_checkpoint = WorkflowCheckpoint(
         checkpoint_id="resumed-cp-id",
@@ -933,7 +933,7 @@ async def test_runner_mark_resumed_sets_previous_checkpoint_id():
 
     assert runner._resumed_from_checkpoint is True  # pyright: ignore[reportPrivateUsage]
     assert runner._iteration == 3  # pyright: ignore[reportPrivateUsage]
-    assert runner._previous_checkpoint_id == "resumed-cp-id"  # pyright: ignore[reportPrivateUsage]
+    assert runner.previous_checkpoint_id == "resumed-cp-id"
 
 
 async def test_runner_post_resume_checkpoint_chains_to_resumed_checkpoint():