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():

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -17,6 +17,7 @@ from typing import Protocol, cast
 
 from agent_framework import (
     ChatOptions,
+    CheckpointID,
     Content,
     ContextProvider,
     FileCheckpointStorage,
@@ -343,6 +344,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
 
     # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally
     CHECKPOINT_STORAGE_PATH = "/.checkpoints"
+    INITIAL_CHECKPOINT_STORAGE_NAME = "initial"
     FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json"
 
     def __init__(
@@ -386,7 +388,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             )
 
         self._is_workflow_agent = False
-        self._checkpoint_storage_path = None
         if isinstance(agent, WorkflowAgent):
             if agent.workflow._runner_context.has_checkpointing():  # pyright: ignore[reportPrivateUsage]
                 raise RuntimeError(
@@ -399,6 +400,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
                 else os.path.join(os.getcwd(), self.CHECKPOINT_STORAGE_PATH.lstrip("/"))
             )
             self._is_workflow_agent = True
+            # The initial checkpoint storage that stores the workflow's initial state. We will use this checkpoint
+            # to restore the workflow when no conversation_id or previous_response_id is supplied in a request.
+            self._initial_checkpoint_storage = _checkpoint_storage_for_context(
+                self._checkpoint_storage_path, self.INITIAL_CHECKPOINT_STORAGE_NAME
+            )
+            self._initial_checkpoint_id: CheckpointID | None = None
 
         self._agent = agent
         self._approval_storage = (
@@ -580,8 +587,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
 
         # The following should never happen due to the checks above.
         # This is for type safety and defensive programming.
-        if self._checkpoint_storage_path is None:
-            raise RuntimeError("Checkpoint storage path is not configured for workflow agent.")
         if not isinstance(self._agent, WorkflowAgent):
             raise RuntimeError("Agent is not a workflow agent.")
 
@@ -590,6 +595,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         # any future async resources owned by the workflow are entered here.
         await self._ensure_agent_ready()
 
+        # Create a checkpoint to store the initial state of the workflow, if it doesn't already exist.
+        # This allows us to restore to a clean slate when no conversation_id or previous_response_id
+        # is supplied in a request.
+        if self._initial_checkpoint_id is None:
+            self._initial_checkpoint_id = await self._agent.workflow.create_checkpoint(self._initial_checkpoint_storage)
+
         # Determine the latest checkpoint (if any) so we can resume the
         # workflow's prior state for this turn. The directory is keyed by
         # the inbound context id (conversation_id when set, otherwise
@@ -599,15 +610,36 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         # the only place that state lives is the workflow checkpoint, so
         # on every turn we restore the latest checkpoint and feed the new
         # input back into the start executor as a continuation rather than
-        # a fresh run.
-        latest_checkpoint_id: str | None = None
-        restore_storage: FileCheckpointStorage | None = None
+        # a fresh run. If no conversation_id or previous_response_id is
+        # supplied, the workflow will be restored to the initial checkpoint
+        # to avoid context bleed between requests.
+        latest_checkpoint_id: str = self._initial_checkpoint_id
+        restore_storage: FileCheckpointStorage = self._initial_checkpoint_storage
         if context_id is not None:
             restore_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id)
             latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name)
             if latest_checkpoint is not None:
                 latest_checkpoint_id = latest_checkpoint.checkpoint_id
 
+        # Restore the workflow to the latest checkpoint and run it with the
+        # new input. Events (including request info events) will not be emitted
+        # during restoration (in streaming) or after restoration (in non-streaming)
+        # since we assume the client had already seen those events and we don't want
+        # to emit duplicates.
+        if is_streaming_request:
+            async for _ in self._agent.run(
+                stream=True,
+                checkpoint_id=latest_checkpoint_id,
+                checkpoint_storage=restore_storage,
+            ):
+                pass
+        else:
+            await self._agent.run(
+                stream=False,
+                checkpoint_id=latest_checkpoint_id,
+                checkpoint_storage=restore_storage,
+            )
+
         # Storage that will receive checkpoints written during this turn.
         # When the caller chains with previous_response_id, the next turn
         # will reference the current response_id as its previous_response_id,
@@ -619,43 +651,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         write_context_id = context.conversation_id or context.response_id
         write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id)
 
-        # Multi-turn pattern: when we have a prior checkpoint, restore it
-        # first (drive the workflow back to idle with prior state intact),
-        # then make a separate call that delivers the new user input. This
-        # depends on Workflow.run preserving shared state across calls. The
-        # restore-only call may yield events from any pending in-flight
-        # work in the checkpoint; we consume those internally here so they
-        # don't surface to the response stream as duplicates.
-        #
-        # If the restored checkpoint had pending request_info events, the
-        # restore-only call replays them through
-        # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates``
-        # and populates ``self._agent.pending_requests``. That is the correct
-        # state: those requests are genuinely outstanding, and the next
-        # ``run(input_messages, ...)`` call may contain ``function_call_output``
-        # items (carried as FunctionResult/FunctionApprovalResponse content)
-        # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`.
-        if latest_checkpoint_id is not None:
-            if is_streaming_request:
-                async for _ in self._agent.run(
-                    stream=True,
-                    checkpoint_id=latest_checkpoint_id,
-                    checkpoint_storage=restore_storage,
-                ):
-                    pass
-            else:
-                await self._agent.run(
-                    stream=False,
-                    checkpoint_id=latest_checkpoint_id,
-                    checkpoint_storage=restore_storage,
-                )
-        else:
-            # We reset the workflow if neither conversation_id nor previous_response_id
-            # was supplied, because this implies there's no prior state to restore and
-            # we want to ensure a clean slate. Workflow may contain in-memory state that
-            # needs to be cleared on new conversations.
-            await self._agent.workflow.reset_for_new_run()
-
         # Now run the agent with the latest input
         response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model)