Python: introduce workflow checkpointing (#366)

* Add workflow checkpointing functionality.

* Reintroduce protocol that went missing during merge

* Checkpoint updates

* Fix ordering of checkpointing

* Cleanup

* Cleanup - thanks Copilot

* Cleanup - thanks Copilot

* State reset updates

* State reset updates 2

* Workflow fixes and updates. Addressed PR feedback

* A few updates

python/packages/main/agent_framework/workflow/__init__.py

@@ -28,6 +28,10 @@ _IMPORTS = [
     "RequestInfoMessage",
     "WorkflowRunResult",
     "Workflow",
+    "FileCheckpointStorage",
+    "InMemoryCheckpointStorage",
+    "CheckpointStorage",
+    "WorkflowCheckpoint",
 ]
 
 

python/packages/main/agent_framework/workflow/__init__.pyi

@@ -6,15 +6,19 @@ from agent_framework_workflow import (
     AgentExecutorResponse,
     AgentRunEvent,
     AgentRunStreamingEvent,
+    CheckpointStorage,
     Executor,
     ExecutorCompletedEvent,
     ExecutorEvent,
     ExecutorInvokeEvent,
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
     RequestInfoEvent,
     RequestInfoExecutor,
     RequestInfoMessage,
     Workflow,
     WorkflowBuilder,
+    WorkflowCheckpoint,
     WorkflowCompletedEvent,
     WorkflowContext,
     WorkflowEvent,
@@ -30,15 +34,19 @@ __all__ = [
     "AgentExecutorResponse",
     "AgentRunEvent",
     "AgentRunStreamingEvent",
+    "CheckpointStorage",
     "Executor",
     "ExecutorCompletedEvent",
     "ExecutorEvent",
     "ExecutorInvokeEvent",
+    "FileCheckpointStorage",
+    "InMemoryCheckpointStorage",
     "RequestInfoEvent",
     "RequestInfoExecutor",
     "RequestInfoMessage",
     "Workflow",
     "WorkflowBuilder",
+    "WorkflowCheckpoint",
     "WorkflowCompletedEvent",
     "WorkflowContext",
     "WorkflowEvent",

python/packages/workflow/agent_framework_workflow/__init__.py

@@ -2,6 +2,15 @@
 
 import importlib.metadata
 
+from ._checkpoint import (
+    CheckpointStorage,
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+from ._const import (
+    DEFAULT_MAX_ITERATIONS,
+)
 from ._events import (
     AgentRunEvent,
     AgentRunStreamingEvent,
@@ -22,6 +31,11 @@ from ._executor import (
     RequestInfoMessage,
     handler,
 )
+from ._runner_context import (
+    InProcRunnerContext,
+    Message,
+    RunnerContext,
+)
 from ._validation import (
     EdgeDuplicationError,
     GraphConnectivityError,
@@ -40,26 +54,34 @@ except importlib.metadata.PackageNotFoundError:
 
 
 __all__ = [
+    "DEFAULT_MAX_ITERATIONS",
     "AgentExecutor",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentRunEvent",
     "AgentRunStreamingEvent",
+    "CheckpointStorage",
     "EdgeDuplicationError",
     "Executor",
     "ExecutorCompletedEvent",
     "ExecutorEvent",
     "ExecutorInvokeEvent",
+    "FileCheckpointStorage",
     "GraphConnectivityError",
+    "InMemoryCheckpointStorage",
+    "InProcRunnerContext",
+    "Message",
     "RequestInfoEvent",
     "RequestInfoEvent",
     "RequestInfoExecutor",
     "RequestInfoExecutor",
     "RequestInfoMessage",
+    "RunnerContext",
     "TypeCompatibilityError",
     "ValidationTypeEnum",
     "Workflow",
     "WorkflowBuilder",
+    "WorkflowCheckpoint",
     "WorkflowCompletedEvent",
     "WorkflowContext",
     "WorkflowEvent",

python/packages/workflow/agent_framework_workflow/_checkpoint.py

@@ -0,0 +1,199 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import logging
+import os
+import uuid
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Protocol
+
+from ._const import DEFAULT_MAX_ITERATIONS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class WorkflowCheckpoint:
+    """Represents a complete checkpoint of workflow state."""
+
+    checkpoint_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    workflow_id: str = ""
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    # Core workflow state
+    messages: dict[str, list[dict[str, Any]]] = field(default_factory=dict)  # type: ignore[misc]
+    shared_state: dict[str, Any] = field(default_factory=dict)  # type: ignore[misc]
+    executor_states: dict[str, dict[str, Any]] = field(default_factory=dict)  # type: ignore[misc]
+
+    # Runtime state
+    iteration_count: int = 0
+    max_iterations: int = DEFAULT_MAX_ITERATIONS
+
+    # Metadata
+    metadata: dict[str, Any] = field(default_factory=dict)  # type: ignore[misc]
+    version: str = "1.0"
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "WorkflowCheckpoint":
+        return cls(**data)
+
+
+class CheckpointStorage(Protocol):
+    """Protocol for checkpoint storage backends."""
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        ...
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        ...
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+        ...
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+        ...
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        ...
+
+
+class InMemoryCheckpointStorage:
+    """In-memory checkpoint storage for testing and development."""
+
+    def __init__(self):
+        """Initialize the memory storage."""
+        self._checkpoints: dict[str, WorkflowCheckpoint] = {}
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        self._checkpoints[checkpoint.checkpoint_id] = checkpoint
+        logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory")
+        return checkpoint.checkpoint_id
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        checkpoint = self._checkpoints.get(checkpoint_id)
+        if checkpoint:
+            logger.debug(f"Loaded checkpoint {checkpoint_id} from memory")
+        return checkpoint
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+        if workflow_id is None:
+            return list(self._checkpoints.keys())
+        return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+        if workflow_id is None:
+            return list(self._checkpoints.values())
+        return [cp for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        if checkpoint_id in self._checkpoints:
+            del self._checkpoints[checkpoint_id]
+            logger.debug(f"Deleted checkpoint {checkpoint_id} from memory")
+            return True
+        return False
+
+
+class FileCheckpointStorage:
+    """File-based checkpoint storage for persistence."""
+
+    def __init__(self, storage_path: str | Path):
+        """Initialize the file storage."""
+        self.storage_path = Path(storage_path)
+        self.storage_path.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Initialized file checkpoint storage at {self.storage_path}")
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        file_path = self.storage_path / f"{checkpoint.checkpoint_id}.json"
+        checkpoint_dict = asdict(checkpoint)
+
+        def _write_atomic() -> None:
+            tmp_path = file_path.with_suffix(".json.tmp")
+            with open(tmp_path, "w") as f:
+                json.dump(checkpoint_dict, f, indent=2, ensure_ascii=False)
+            os.replace(tmp_path, file_path)
+
+        await asyncio.to_thread(_write_atomic)
+
+        logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}")
+        return checkpoint.checkpoint_id
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        file_path = self.storage_path / f"{checkpoint_id}.json"
+
+        if not file_path.exists():
+            return None
+
+        def _read() -> dict[str, Any]:
+            with open(file_path) as f:
+                return json.load(f)
+
+        checkpoint_dict = await asyncio.to_thread(_read)
+
+        checkpoint = WorkflowCheckpoint(**checkpoint_dict)
+        logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}")
+        return checkpoint
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+
+        def _list_ids() -> list[str]:
+            checkpoint_ids: list[str] = []
+            for file_path in self.storage_path.glob("*.json"):
+                try:
+                    with open(file_path) as f:
+                        data = json.load(f)
+                    if workflow_id is None or data.get("workflow_id") == workflow_id:
+                        checkpoint_ids.append(data.get("checkpoint_id", file_path.stem))
+                except Exception as e:
+                    logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
+            return checkpoint_ids
+
+        return await asyncio.to_thread(_list_ids)
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+
+        def _list_checkpoints() -> list[WorkflowCheckpoint]:
+            checkpoints: list[WorkflowCheckpoint] = []
+            for file_path in self.storage_path.glob("*.json"):
+                try:
+                    with open(file_path) as f:
+                        data = json.load(f)
+                    if workflow_id is None or data.get("workflow_id") == workflow_id:
+                        checkpoints.append(WorkflowCheckpoint.from_dict(data))
+                except Exception as e:
+                    logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
+            return checkpoints
+
+        return await asyncio.to_thread(_list_checkpoints)
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        file_path = self.storage_path / f"{checkpoint_id}.json"
+
+        def _delete() -> bool:
+            if file_path.exists():
+                file_path.unlink()
+                logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}")
+                return True
+            return False
+
+        return await asyncio.to_thread(_delete)

python/packages/workflow/agent_framework_workflow/_const.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+DEFAULT_MAX_ITERATIONS = 100  # Default maximum iterations for workflow execution.

python/packages/workflow/agent_framework_workflow/_runner.py

@@ -4,16 +4,16 @@ import asyncio
 import logging
 from collections import defaultdict
 from collections.abc import AsyncIterable
+from typing import Any
 
 from ._edge import Edge
 from ._events import WorkflowEvent
+from ._executor import Executor
 from ._runner_context import Message, RunnerContext
 from ._shared_state import SharedState
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MAX_ITERATIONS = 100
-
 
 class Runner:
     """A class to run a workflow in Pregel supersteps."""
@@ -23,8 +23,9 @@ class Runner:
         edges: list[Edge],
         shared_state: SharedState,
         ctx: RunnerContext,
-        max_iterations: int = DEFAULT_MAX_ITERATIONS,
-    ) -> None:
+        max_iterations: int = 100,
+        workflow_id: str | None = None,
+    ):
         """Initialize the runner with edges, shared state, and context.
 
         Args:
@@ -32,66 +33,107 @@ class Runner:
             shared_state: The shared state for the workflow.
             ctx: The runner context for the workflow.
             max_iterations: The maximum number of iterations to run.
+            workflow_id: The workflow ID for checkpointing.
         """
         self._edge_map = self._parse_edges(edges)
         self._ctx = ctx
         self._iteration = 0
         self._max_iterations = max_iterations
         self._shared_state = shared_state
-        self._is_running = False
+        self._workflow_id = workflow_id
+        self._running = False
+        self._resumed_from_checkpoint = False  # Track whether we resumed
+
+        # Set workflow ID in context if provided
+        if workflow_id:
+            self._ctx.set_workflow_id(workflow_id)
 
     @property
     def context(self) -> RunnerContext:
         """Get the workflow context."""
         return self._ctx
 
+    def mark_resumed(self, iteration: int | None = None, max_iterations: int | None = None) -> None:
+        """Mark the runner as having resumed from a checkpoint.
+
+        Optionally set the current iteration and max iterations.
+        """
+        self._resumed_from_checkpoint = True
+        if iteration is not None:
+            self._iteration = iteration
+        if max_iterations is not None:
+            self._max_iterations = max_iterations
+
     async def run_until_convergence(self) -> AsyncIterable[WorkflowEvent]:
         """Run the workflow until no more messages are sent."""
+        if self._running:
+            raise RuntimeError("Runner is already running.")
+
+        self._running = True
         try:
-            if self._is_running:
-                raise RuntimeError("Runner is already running.")
-            self._is_running = True
+            # Process any events from initial execution before checkpointing
+            if await self._ctx.has_events():
+                logger.info("Processing events from initial execution")
+                events = await self._ctx.drain_events()
+                for event in events:
+                    logger.info(f"Yielding initial event: {event}")
+                    yield event
+
+            # Create first checkpoint if there are messages from initial execution
+            if await self._ctx.has_messages() and self._ctx.has_checkpointing():
+                if not self._resumed_from_checkpoint:
+                    logger.info("Creating checkpoint after initial execution")
+                    await self._create_checkpoint_if_enabled("after_initial_execution")
+                else:
+                    logger.info("Skipping 'after_initial_execution' checkpoint because we resumed from a checkpoint")
+
+            # Initialize context with starting iteration state
+            await self._update_context_with_shared_state()
+
             while self._iteration < self._max_iterations:
+                logger.info(f"Starting superstep {self._iteration + 1}")
                 await self._run_iteration()
                 self._iteration += 1
 
+                # Update context with current iteration state immediately
+                await self._update_context_with_shared_state()
+
+                logger.info(f"Completed superstep {self._iteration}")
+
+                # Process events first before any checkpointing
                 if await self._ctx.has_events():
+                    logger.info("Processing events before checkpointing")
                     events = await self._ctx.drain_events()
                     for event in events:
+                        logger.debug(f"Yielding event: {event}")
                         yield event
 
+                # Create checkpoint after each superstep iteration
+                await self._create_checkpoint_if_enabled(f"superstep_{self._iteration}")
+
                 if not await self._ctx.has_messages():
                     break
-            else:
+
+            if self._iteration >= self._max_iterations and await self._ctx.has_messages():
                 raise RuntimeError(f"Runner did not converge after {self._max_iterations} iterations.")
-        finally:
-            self._is_running = False
+
+            logger.info(f"Workflow completed after {self._iteration} supersteps")
             self._iteration = 0
+            self._resumed_from_checkpoint = False  # Reset resume flag for next run
+        finally:
+            self._running = False
 
     async def _run_iteration(self):
-        """Run a superstep of the workflow execution."""
-
         async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None:
-            """Deliver messages to the executors.
-
-            Outer loop to concurrently deliver messages from all sources to their targets.
-            """
-
             async def _deliver_messages_inner(
                 edge: Edge,
                 messages: list[Message],
             ) -> None:
-                """Deliver messages to a specific target executor.
-
-                Inner loop to deliver messages to a specific target executor.
-                """
                 for message in messages:
                     if message.target_id is not None and message.target_id != edge.target_id:
                         continue
-
                     if not edge.can_handle(message.data):
                         continue
-
                     await edge.send_message(message, self._shared_state, self._ctx)
 
             associated_edges = self._edge_map.get(source_executor_id, [])
@@ -105,6 +147,127 @@ class Runner:
         ]
         await asyncio.gather(*tasks)
 
+    async def _create_checkpoint_if_enabled(self, checkpoint_type: str) -> str | None:
+        """Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
+        if not self._ctx.has_checkpointing():
+            return None
+
+        try:
+            # Auto-snapshot executor states
+            await self._auto_snapshot_executor_states()
+            await self._update_context_with_shared_state()
+            checkpoint_category = "initial" if checkpoint_type == "after_initial_execution" else "superstep"
+            metadata = {
+                "superstep": self._iteration,
+                "checkpoint_type": checkpoint_category,
+            }
+            checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata)
+            logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}")
+            return checkpoint_id
+        except Exception as e:
+            logger.warning(f"Failed to create {checkpoint_type} checkpoint: {e}")
+            return None
+
+    async def _auto_snapshot_executor_states(self) -> None:
+        """Populate executor state by calling snapshot hooks on executors if available.
+
+        Convention:
+          - If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it.
+          - Else if it has a plain attribute `state` that is a dict, use that.
+        Only JSON-serializable dicts should be provided by executors.
+        """
+        executors: dict[str, Executor] = {}
+        for edge_list in self._edge_map.values():
+            for edge in edge_list:
+                executors[edge.source.id] = edge.source
+                executors[edge.target.id] = edge.target
+        for exec_id, executor in executors.items():
+            state_dict: dict[str, Any] | None = None
+            snapshot = getattr(executor, "snapshot_state", None)
+            try:
+                if callable(snapshot):
+                    maybe = snapshot()
+                    if asyncio.iscoroutine(maybe):  # type: ignore[arg-type]
+                        maybe = await maybe  # type: ignore[assignment]
+                    if isinstance(maybe, dict):
+                        state_dict = maybe  # type: ignore[assignment]
+                else:
+                    state_attr = getattr(executor, "state", None)
+                    if isinstance(state_attr, dict):
+                        state_dict = state_attr  # type: ignore[assignment]
+            except Exception as ex:  # pragma: no cover
+                logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}")
+            if state_dict is not None:
+                try:
+                    await self._ctx.set_state(exec_id, state_dict)
+                except Exception as ex:  # pragma: no cover
+                    logger.debug(f"Failed to persist state for executor {exec_id}: {ex}")
+
+    async def _update_context_with_shared_state(self) -> None:
+        if not self._ctx.has_checkpointing():
+            return
+
+        try:
+            current_state = await self._ctx.get_checkpoint_state()
+
+            shared_state_data = {}
+            async with self._shared_state.hold():
+                if hasattr(self._shared_state, "_state"):
+                    shared_state_data = dict(self._shared_state._state)  # type: ignore[attr-defined]
+
+            current_state["shared_state"] = shared_state_data
+            current_state["iteration_count"] = self._iteration
+            current_state["max_iterations"] = self._max_iterations
+
+            await self._ctx.set_checkpoint_state(current_state)
+        except Exception as e:
+            logger.warning(f"Failed to update context with shared state: {e}")
+
+    async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
+        """Restore workflow state from a checkpoint.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from
+
+        Returns:
+            True if restoration was successful, False otherwise
+        """
+        if not self._ctx.has_checkpointing():
+            logger.warning("Context does not support checkpointing")
+            return False
+
+        try:
+            success = await self._ctx.restore_from_checkpoint(checkpoint_id)
+            if not success:
+                return False
+
+            await self._restore_shared_state_from_context()
+            self.mark_resumed()  # mark resumed; iteration/max already restored from context
+            logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
+            return False
+
+    async def _restore_shared_state_from_context(self) -> None:
+        if not self._ctx.has_checkpointing():
+            return
+
+        try:
+            restored_state = await self._ctx.get_checkpoint_state()
+
+            shared_state_data = restored_state.get("shared_state", {})
+            if shared_state_data and hasattr(self._shared_state, "_state"):
+                async with self._shared_state.hold():
+                    self._shared_state._state.clear()  # type: ignore[attr-defined]
+                    self._shared_state._state.update(shared_state_data)  # type: ignore[attr-defined]
+
+            self._iteration = restored_state.get("iteration_count", 0)
+            self._max_iterations = restored_state.get("max_iterations", self._max_iterations)
+
+        except Exception as e:
+            logger.warning(f"Failed to restore shared state from context: {e}")
+
     def _parse_edges(self, edges: list[Edge]) -> dict[str, list[Edge]]:
         """Parse the edges of the workflow into a more convenient format.
 

python/packages/workflow/agent_framework_workflow/_runner_context.py

@@ -1,11 +1,15 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import logging
+import uuid
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Protocol, TypeVar, runtime_checkable
+from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable
 
+from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
+from ._const import DEFAULT_MAX_ITERATIONS
 from ._events import WorkflowEvent
+from ._shared_state import SharedState
 
 logger = logging.getLogger(__name__)
 
@@ -21,9 +25,21 @@ class Message:
     target_id: str | None = None
 
 
+class CheckpointState(TypedDict):
+    messages: dict[str, list[dict[str, Any]]]
+    shared_state: dict[str, Any]
+    executor_states: dict[str, dict[str, Any]]
+    iteration_count: int
+    max_iterations: int
+
+
 @runtime_checkable
 class RunnerContext(Protocol):
-    """Protocol for the execution context used by the runner."""
+    """Protocol for the execution context used by the runner.
+
+    A single context that supports messaging, events, and optional checkpointing.
+    If checkpoint storage is not configured, checkpoint methods may raise.
+    """
 
     async def send_message(self, message: Message) -> None:
         """Send a message from the executor to the context.
@@ -73,43 +89,208 @@ class RunnerContext(Protocol):
         """
         ...
 
+    async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
+        """Set the state for a specific executor.
 
-class InProcRunnerContext(RunnerContext):
-    """In-process execution context for local execution of workflows."""
+        Args:
+            executor_id: The ID of the executor whose state is being set.
+            state: The state data to be set for the executor.
+        """
+        ...
 
-    def __init__(self):
-        """Initialize the in-process execution context."""
+    async def get_state(self, executor_id: str) -> dict[str, Any] | None:
+        """Get the state for a specific executor.
+
+        Args:
+            executor_id: The ID of the executor whose state is being retrieved.
+
+        Returns:
+            The state data for the executor, or None if not found.
+        """
+        ...
+
+    # Checkpointing capability
+    def has_checkpointing(self) -> bool:
+        """Check if the context supports checkpointing.
+
+        Returns:
+            True if checkpointing is supported, False otherwise.
+        """
+        ...
+
+    # Checkpointing APIs (optional, enabled by storage)
+    def set_workflow_id(self, workflow_id: str) -> None:
+        """Set the workflow ID for the context."""
+        ...
+
+    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
+        """Reset the context for a new workflow run."""
+        ...
+
+    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
+        """Create a checkpoint of the current workflow state.
+
+        Args:
+            metadata: Optional metadata to associate with the checkpoint.
+        """
+        ...
+
+    async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
+        """Restore the context from a checkpoint.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from.
+
+        Returns:
+            True if the restoration was successful, False otherwise.
+        """
+        ...
+
+    async def get_checkpoint_state(self) -> CheckpointState:
+        """Get the current state of the context suitable for checkpointing."""
+        ...
+
+    async def set_checkpoint_state(self, state: CheckpointState) -> None:
+        """Set the state of the context from a checkpoint.
+
+        Args:
+            state: The state data to set for the context.
+        """
+        ...
+
+
+class InProcRunnerContext:
+    """In-process execution context for local execution and optional checkpointing."""
+
+    def __init__(self, checkpoint_storage: CheckpointStorage | None = None):
+        """Initialize the in-process execution context.
+
+        Args:
+            checkpoint_storage: Optional storage to enable checkpointing.
+        """
         self._messages: defaultdict[str, list[Message]] = defaultdict(list)
         self._events: list[WorkflowEvent] = []
 
+        # Checkpointing configuration/state
+        self._checkpoint_storage = checkpoint_storage
+        self._workflow_id: str | None = None
+        self._shared_state: dict[str, Any] = {}
+        self._executor_states: dict[str, dict[str, Any]] = {}
+        self._iteration_count: int = 0
+        self._max_iterations: int = 100
+
     async def send_message(self, message: Message) -> None:
-        """Send a message from the executor to the context."""
         self._messages[message.source_id].append(message)
 
     async def drain_messages(self) -> dict[str, list[Message]]:
-        """Drain all messages from the context."""
         messages = dict(self._messages)
         self._messages.clear()
         return messages
 
     async def has_messages(self) -> bool:
-        """Check if there are any messages in the context."""
         return bool(self._messages)
 
     async def add_event(self, event: WorkflowEvent) -> None:
-        """Add an event to the execution context.
-
-        Args:
-            event: The event to be added.
-        """
         self._events.append(event)
 
     async def drain_events(self) -> list[WorkflowEvent]:
-        """Drain all events from the context."""
         events = self._events.copy()
         self._events.clear()
         return events
 
     async def has_events(self) -> bool:
-        """Check if there are any events in the context."""
         return bool(self._events)
+
+    async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
+        self._executor_states[executor_id] = state
+
+    async def get_state(self, executor_id: str) -> dict[str, Any] | None:
+        return self._executor_states.get(executor_id)
+
+    def has_checkpointing(self) -> bool:
+        return self._checkpoint_storage is not None
+
+    def set_workflow_id(self, workflow_id: str) -> None:
+        self._workflow_id = workflow_id
+
+    def reset_for_new_run(self, workflow_shared_state: "SharedState | None" = None) -> None:
+        self._messages.clear()
+        self._events.clear()
+        self._shared_state.clear()
+        self._executor_states.clear()
+        self._iteration_count = 0
+        if workflow_shared_state is not None and hasattr(workflow_shared_state, "_state"):
+            workflow_shared_state._state.clear()  # type: ignore[attr-defined]
+
+    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        wf_id = self._workflow_id or str(uuid.uuid4())
+        self._workflow_id = wf_id
+        state = await self.get_checkpoint_state()
+
+        checkpoint = WorkflowCheckpoint(
+            workflow_id=wf_id,
+            messages=state["messages"],
+            shared_state=state.get("shared_state", {}),
+            executor_states=state.get("executor_states", {}),
+            iteration_count=state.get("iteration_count", 0),
+            max_iterations=state.get("max_iterations", DEFAULT_MAX_ITERATIONS),
+            metadata=metadata or {},
+        )
+        checkpoint_id = await self._checkpoint_storage.save_checkpoint(checkpoint)
+        logger.info(f"Created checkpoint {checkpoint_id} for workflow {wf_id}'")
+        return checkpoint_id
+
+    async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        checkpoint = await self._checkpoint_storage.load_checkpoint(checkpoint_id)
+        if not checkpoint:
+            logger.error(f"Checkpoint {checkpoint_id} not found")
+            return False
+
+        state: CheckpointState = {
+            "messages": checkpoint.messages,
+            "shared_state": checkpoint.shared_state,
+            "executor_states": checkpoint.executor_states,
+            "iteration_count": checkpoint.iteration_count,
+            "max_iterations": checkpoint.max_iterations,
+        }
+        await self.set_checkpoint_state(state)
+        self._workflow_id = checkpoint.workflow_id
+        logger.info(f"Restored state from checkpoint {checkpoint_id}'")
+        return True
+
+    async def get_checkpoint_state(self) -> CheckpointState:
+        serializable_messages: dict[str, list[dict[str, Any]]] = {}
+        for source_id, message_list in self._messages.items():
+            serializable_messages[source_id] = [
+                {"data": msg.data, "source_id": msg.source_id, "target_id": msg.target_id} for msg in message_list
+            ]
+        return {
+            "messages": serializable_messages,
+            "shared_state": self._shared_state,
+            "executor_states": self._executor_states,
+            "iteration_count": self._iteration_count,
+            "max_iterations": self._max_iterations,
+        }
+
+    async def set_checkpoint_state(self, state: CheckpointState) -> None:
+        self._messages.clear()
+        messages_data = state.get("messages", {})
+        for source_id, message_list in messages_data.items():
+            self._messages[source_id] = [
+                Message(
+                    data=msg.get("data"),
+                    source_id=msg.get("source_id", ""),
+                    target_id=msg.get("target_id"),
+                )
+                for msg in message_list
+            ]
+        self._shared_state = state.get("shared_state", {})
+        self._executor_states = state.get("executor_states", {})
+        self._iteration_count = state.get("iteration_count", 0)
+        self._max_iterations = state.get("max_iterations", 100)

python/packages/workflow/agent_framework_workflow/_workflow.py

@@ -1,15 +1,19 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import logging
 import sys
+import uuid
 from collections.abc import AsyncIterable, Callable, Sequence
 from typing import Any
 
+from ._checkpoint import CheckpointStorage
+from ._const import DEFAULT_MAX_ITERATIONS
 from ._edge import Edge
 from ._events import RequestInfoEvent, WorkflowCompletedEvent, WorkflowEvent
 from ._executor import Executor, RequestInfoExecutor
-from ._runner import DEFAULT_MAX_ITERATIONS, Runner
-from ._runner_context import InProcRunnerContext, RunnerContext
+from ._runner import Runner
+from ._runner_context import CheckpointState, InProcRunnerContext, RunnerContext
 from ._shared_state import SharedState
 from ._validation import validate_workflow_graph
 from ._workflow_context import WorkflowContext
@@ -20,6 +24,9 @@ else:
     from typing_extensions import Self  # pragma: no cover
 
 
+logger = logging.getLogger(__name__)
+
+
 class WorkflowRunResult(list[WorkflowEvent]):
     """A list of events generated during the workflow execution in non-streaming mode."""
 
@@ -48,6 +55,9 @@ class WorkflowRunResult(list[WorkflowEvent]):
         return [event for event in self if isinstance(event, RequestInfoEvent)]
 
 
+# region Workflow
+
+
 class Workflow:
     """A class representing a workflow that can be executed.
 
@@ -77,7 +87,11 @@ class Workflow:
         }
 
         self._shared_state = SharedState()
-        self._runner = Runner(self._edges, self._shared_state, runner_context, max_iterations=max_iterations)
+
+        workflow_id = str(uuid.uuid4())
+        self._runner = Runner(
+            self._edges, self._shared_state, runner_context, max_iterations=max_iterations, workflow_id=workflow_id
+        )
 
     @property
     def edges(self) -> list[Edge]:
@@ -101,7 +115,7 @@ class Workflow:
         return list(self._executors.values())
 
     async def run_streaming(self, message: Any) -> AsyncIterable[WorkflowEvent]:
-        """Send a message to the starting executor of the workflow and stream the events generated by the workflow.
+        """Run the workflow with a starting message and stream events.
 
         Args:
             message: The message to be sent to the starting executor.
@@ -109,6 +123,9 @@ class Workflow:
         Yields:
             WorkflowEvent: The events generated during the workflow execution.
         """
+        # Reset context for a new run if supported
+        self._runner.context.reset_for_new_run(self._shared_state)
+
         executor = self._start_executor
         if isinstance(executor, str):
             executor = self._get_executor_by_id(executor)
@@ -117,15 +134,71 @@ class Workflow:
             message,
             WorkflowContext(
                 executor.id,
-                [
-                    # Using the workflow class name as the source executor ID when
-                    # delivering the first message to the starting executor
-                    self.__class__.__name__
-                ],
+                [self.__class__.__name__],
                 self._shared_state,
                 self._runner.context,
             ),
         )
+
+        async for event in self._runner.run_until_convergence():
+            yield event
+
+    async def run_streaming_from_checkpoint(
+        self,
+        checkpoint_id: str,
+        checkpoint_storage: CheckpointStorage | None = None,
+        responses: dict[str, Any] | None = None,
+    ) -> AsyncIterable[WorkflowEvent]:
+        """Resume workflow execution from a checkpoint and stream events.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from.
+            checkpoint_storage: Optional checkpoint storage to use for restoration.
+                              If not provided, the workflow must have been built with checkpointing enabled.
+            responses: Optional dictionary of responses to inject into the workflow
+                      after restoration. Keys are request IDs, values are response data.
+
+        Yields:
+            WorkflowEvent: Events generated during workflow execution.
+
+        Raises:
+            ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled.
+            RuntimeError: If checkpoint restoration fails.
+        """
+        has_checkpointing = self._runner.context.has_checkpointing()
+
+        if not has_checkpointing and not checkpoint_storage:
+            raise ValueError(
+                "Cannot restore from checkpoint: either provide checkpoint_storage parameter "
+                "or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)."
+            )
+
+        if has_checkpointing:
+            # restore via Runner so shared state and iteration are synchronized
+            restored = await self._runner.restore_from_checkpoint(checkpoint_id)
+        else:
+            if checkpoint_storage is None:
+                raise ValueError("checkpoint_storage cannot be None.")
+            restored = await self._restore_from_external_checkpoint(checkpoint_id, checkpoint_storage)
+
+        if not restored:
+            raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}")
+
+        if responses:
+            request_info_executor = self._get_executor_by_id(RequestInfoExecutor.EXECUTOR_ID)
+            if isinstance(request_info_executor, RequestInfoExecutor):
+                for request_id, response_data in responses.items():
+                    await request_info_executor.handle_response(
+                        response_data,
+                        request_id,
+                        WorkflowContext(
+                            request_info_executor.id,
+                            [self.__class__.__name__],
+                            self._shared_state,
+                            self._runner.context,
+                        ),
+                    )
+
         async for event in self._runner.run_until_convergence():
             yield event
 
@@ -150,11 +223,7 @@ class Workflow:
                 request_id,
                 WorkflowContext(
                     request_info_executor.id,
-                    [
-                        # Using the workflow class name as the source executor ID when
-                        # delivering the first message to the starting executor
-                        self.__class__.__name__
-                    ],
+                    [self.__class__.__name__],
                     self._shared_state,
                     self._runner.context,
                 ),
@@ -177,6 +246,33 @@ class Workflow:
         events = [event async for event in self.run_streaming(message)]
         return WorkflowRunResult(events)
 
+    async def run_from_checkpoint(
+        self,
+        checkpoint_id: str,
+        checkpoint_storage: CheckpointStorage | None = None,
+        responses: dict[str, Any] | None = None,
+    ) -> WorkflowRunResult:
+        """Resume workflow execution from a checkpoint.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from.
+            checkpoint_storage: Optional checkpoint storage to use for restoration.
+                              If not provided, the workflow must have been built with checkpointing enabled.
+            responses: Optional dictionary of responses to inject into the workflow
+                      after restoration. Keys are request IDs, values are response data.
+
+        Returns:
+            A WorkflowRunResult instance containing a list of events generated during the workflow execution.
+
+        Raises:
+            ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled.
+            RuntimeError: If checkpoint restoration fails.
+        """
+        events = [
+            event async for event in self.run_streaming_from_checkpoint(checkpoint_id, checkpoint_storage, responses)
+        ]
+        return WorkflowRunResult(events)
+
     async def send_responses(self, responses: dict[str, Any]) -> WorkflowRunResult:
         """Send responses back to the workflow.
 
@@ -202,6 +298,104 @@ class Workflow:
             raise ValueError(f"Executor with ID {executor_id} not found.")
         return self._executors[executor_id]
 
+    async def _restore_from_external_checkpoint(
+        self, checkpoint_id: str, checkpoint_storage: CheckpointStorage
+    ) -> bool:
+        """Restore workflow state from an external checkpoint storage.
+
+        This method implements the state transfer pattern: load checkpoint data
+        from external storage and transfer it to the current workflow context.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from.
+            checkpoint_storage: The checkpoint storage to load from.
+
+        Returns:
+            True if restoration was successful, False otherwise.
+        """
+        try:
+            checkpoint = await checkpoint_storage.load_checkpoint(checkpoint_id)
+            if not checkpoint:
+                return False
+
+            temp_context = InProcRunnerContext(checkpoint_storage)
+            state: CheckpointState = {
+                "messages": checkpoint.messages,
+                "shared_state": checkpoint.shared_state,
+                "executor_states": checkpoint.executor_states,
+                "iteration_count": checkpoint.iteration_count,
+                "max_iterations": checkpoint.max_iterations,
+            }
+
+            await temp_context.set_checkpoint_state(state)
+            restored_state = await temp_context.get_checkpoint_state()
+            await self._transfer_state_to_context(restored_state)
+
+            # Also set runner iteration/max so superstep numbering continues
+            self._runner.mark_resumed(iteration=checkpoint.iteration_count, max_iterations=checkpoint.max_iterations)
+
+            return True
+
+        except Exception as e:
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.error(f"Failed to restore from external checkpoint {checkpoint_id}: {e}")
+            return False
+
+    async def _transfer_state_to_context(self, restored_state: CheckpointState) -> None:
+        """Transfer restored checkpoint state into the current workflow runtime.
+
+        This transfers:
+        - messages -> into the current RunnerContext so delivery can continue
+        - executor_states -> into the current RunnerContext so ctx.get_state() works after resume
+        - shared_state -> into the Workflow's SharedState so executors can read values set before the checkpoint
+        """
+        # Best-effort restoration
+        # Restore shared state so downstream executors can read values (e.g., original_input)
+        try:
+            shared_state_data = restored_state.get("shared_state", {})
+            if shared_state_data and hasattr(self._shared_state, "_state"):
+                async with self._shared_state.hold():
+                    self._shared_state._state.clear()  # type: ignore[attr-defined]
+                    self._shared_state._state.update(shared_state_data)  # type: ignore[attr-defined]
+        except Exception as exc:  # pragma: no cover
+            logger.debug("Failed to restore shared_state during external restore: %s", exc)
+
+        # Restore executor states into the context so ctx.get_state() calls after resume succeed
+        try:
+            executor_states = restored_state.get("executor_states", {})
+            for exec_id, state in executor_states.items():
+                try:
+                    await self._runner.context.set_state(exec_id, state)
+                except Exception as exc:  # pragma: no cover - ignore per-executor failures
+                    logger.debug("Failed to restore executor state for %s during external restore: %s", exec_id, exc)
+        except Exception as exc:  # pragma: no cover
+            logger.debug("Failed to iterate executor_states during external restore: %s", exc)
+
+        # Transfer pending messages into the context for delivery in the next superstep
+        messages_data = restored_state["messages"]
+        for _, message_list in messages_data.items():
+            for msg_data in message_list:
+                source_any = msg_data.get("source_id", "")
+                source_id: str = source_any if isinstance(source_any, str) else str(source_any)
+                if not source_id:
+                    source_id = ""
+                target_raw = msg_data.get("target_id")
+                target_id: str | None = (
+                    target_raw if target_raw is None or isinstance(target_raw, str) else str(target_raw)
+                )
+
+                # Build and send Message via runner context
+                from ._runner_context import Message as _Msg
+
+                await self._runner.context.send_message(
+                    _Msg(data=msg_data.get("data"), source_id=source_id, target_id=target_id)
+                )
+
+
+# region WorkflowBuilder
+
 
 class WorkflowBuilder:
     """A builder class for constructing workflows.
@@ -209,11 +403,12 @@ class WorkflowBuilder:
     This class provides methods to add edges and set the starting executor for the workflow.
     """
 
-    def __init__(self):
+    def __init__(self, max_iterations: int = DEFAULT_MAX_ITERATIONS):
         """Initialize the WorkflowBuilder with an empty list of edges and no starting executor."""
         self._edges: list[Edge] = []
         self._start_executor: Executor | str | None = None
-        self._max_iterations: int = DEFAULT_MAX_ITERATIONS
+        self._checkpoint_storage: CheckpointStorage | None = None
+        self._max_iterations: int = max_iterations
 
     def add_edge(
         self,
@@ -329,6 +524,15 @@ class WorkflowBuilder:
         self._max_iterations = max_iterations
         return self
 
+    def with_checkpointing(self, checkpoint_storage: CheckpointStorage) -> "Self":
+        """Enable checkpointing with the specified storage.
+
+        Args:
+            checkpoint_storage: The checkpoint storage to use.
+        """
+        self._checkpoint_storage = checkpoint_storage
+        return self
+
     def build(self) -> Workflow:
         """Build and return the constructed workflow.
 
@@ -347,4 +551,9 @@ class WorkflowBuilder:
 
         validate_workflow_graph(self._edges, self._start_executor)
 
-        return Workflow(self._edges, self._start_executor, InProcRunnerContext(), self._max_iterations)
+        context = InProcRunnerContext(self._checkpoint_storage)
+
+        return Workflow(self._edges, self._start_executor, context, self._max_iterations)
+
+
+# endregion

python/packages/workflow/agent_framework_workflow/_workflow_context.py

@@ -89,3 +89,18 @@ class WorkflowContext:
     def shared_state(self) -> SharedState:
         """Get the shared state."""
         return self._shared_state
+
+    async def set_state(self, state: dict[str, Any]) -> None:
+        """Persist this executor's state into the checkpointable context.
+
+        Executors call this with a JSON-serializable dict capturing the minimal
+        state needed to resume. It replaces any previously stored state.
+        """
+        if hasattr(self._runner_context, "set_state"):
+            await self._runner_context.set_state(self._executor_id, state)  # type: ignore[arg-type]
+
+    async def get_state(self) -> dict[str, Any] | None:
+        """Retrieve previously persisted state for this executor, if any."""
+        if hasattr(self._runner_context, "get_state"):
+            return await self._runner_context.get_state(self._executor_id)  # type: ignore[return-value]
+        return None

python/packages/workflow/tests/test_checkpoint.py

@@ -0,0 +1,334 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+
+from agent_framework_workflow._checkpoint import (
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+
+
+def test_workflow_checkpoint_default_values():
+    checkpoint = WorkflowCheckpoint()
+
+    assert checkpoint.checkpoint_id != ""
+    assert checkpoint.workflow_id == ""
+    assert checkpoint.timestamp != ""
+    assert checkpoint.messages == {}
+    assert checkpoint.shared_state == {}
+    assert checkpoint.executor_states == {}
+    assert checkpoint.iteration_count == 0
+    assert checkpoint.max_iterations == 100
+    assert checkpoint.metadata == {}
+    assert checkpoint.version == "1.0"
+
+
+def test_workflow_checkpoint_custom_values():
+    custom_timestamp = datetime.now(timezone.utc).isoformat()
+    checkpoint = WorkflowCheckpoint(
+        checkpoint_id="test-checkpoint-123",
+        workflow_id="test-workflow-456",
+        timestamp=custom_timestamp,
+        messages={"executor1": [{"data": "test"}]},
+        shared_state={"key": "value"},
+        executor_states={"executor1": {"state": "active"}},
+        iteration_count=5,
+        max_iterations=50,
+        metadata={"test": True},
+        version="2.0",
+    )
+
+    assert checkpoint.checkpoint_id == "test-checkpoint-123"
+    assert checkpoint.workflow_id == "test-workflow-456"
+    assert checkpoint.timestamp == custom_timestamp
+    assert checkpoint.messages == {"executor1": [{"data": "test"}]}
+    assert checkpoint.shared_state == {"key": "value"}
+    assert checkpoint.executor_states == {"executor1": {"state": "active"}}
+    assert checkpoint.iteration_count == 5
+    assert checkpoint.max_iterations == 50
+    assert checkpoint.metadata == {"test": True}
+    assert checkpoint.version == "2.0"
+
+
+async def test_memory_checkpoint_storage_save_and_load():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow", messages={"executor1": [{"data": "hello"}]})
+
+    # Save checkpoint
+    saved_id = await storage.save_checkpoint(checkpoint)
+    assert saved_id == checkpoint.checkpoint_id
+
+    # Load checkpoint
+    loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+    assert loaded_checkpoint is not None
+    assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+    assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+    assert loaded_checkpoint.messages == checkpoint.messages
+
+
+async def test_memory_checkpoint_storage_load_nonexistent():
+    storage = InMemoryCheckpointStorage()
+
+    result = await storage.load_checkpoint("nonexistent-id")
+    assert result is None
+
+
+async def test_memory_checkpoint_storage_list_checkpoints():
+    storage = InMemoryCheckpointStorage()
+
+    # Create checkpoints for different workflows
+    checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+    await storage.save_checkpoint(checkpoint1)
+    await storage.save_checkpoint(checkpoint2)
+    await storage.save_checkpoint(checkpoint3)
+
+    # Test list_checkpoint_ids for workflow-1
+    workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+    assert len(workflow1_checkpoint_ids) == 2
+    assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+    assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+    # Test list_checkpoints for workflow-1 (returns objects)
+    workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+    assert len(workflow1_checkpoints) == 2
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+    assert {cp.checkpoint_id for cp in workflow1_checkpoints} == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+    # Test list_checkpoint_ids for workflow-2
+    workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+    assert len(workflow2_checkpoint_ids) == 1
+    assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+    # Test list_checkpoints for workflow-2 (returns objects)
+    workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+    assert len(workflow2_checkpoints) == 1
+    assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+    # Test list_checkpoint_ids for non-existent workflow
+    empty_checkpoint_ids = await storage.list_checkpoint_ids("nonexistent-workflow")
+    assert len(empty_checkpoint_ids) == 0
+
+    # Test list_checkpoints for non-existent workflow
+    empty_checkpoints = await storage.list_checkpoints("nonexistent-workflow")
+    assert len(empty_checkpoints) == 0
+
+    # Test list_checkpoint_ids without workflow filter (all checkpoints)
+    all_checkpoint_ids = await storage.list_checkpoint_ids()
+    assert len(all_checkpoint_ids) == 3
+    expected_ids = {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id, checkpoint3.checkpoint_id}
+    assert expected_ids.issubset(set(all_checkpoint_ids))
+
+    # Test list_checkpoints without workflow filter (all checkpoints)
+    all_checkpoints = await storage.list_checkpoints()
+    assert len(all_checkpoints) == 3
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_memory_checkpoint_storage_delete():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+    # Save checkpoint
+    await storage.save_checkpoint(checkpoint)
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is not None
+
+    # Delete checkpoint
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is True
+
+    # Verify deletion
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is None
+
+    # Try to delete again
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is False
+
+
+async def test_file_checkpoint_storage_save_and_load():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
+            shared_state={"key": "value"},
+        )
+
+        # Save checkpoint
+        saved_id = await storage.save_checkpoint(checkpoint)
+        assert saved_id == checkpoint.checkpoint_id
+
+        # Verify file was created
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Load checkpoint
+        loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+        assert loaded_checkpoint is not None
+        assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+        assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+        assert loaded_checkpoint.messages == checkpoint.messages
+        assert loaded_checkpoint.shared_state == checkpoint.shared_state
+
+
+async def test_file_checkpoint_storage_load_nonexistent():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        result = await storage.load_checkpoint("nonexistent-id")
+        assert result is None
+
+
+async def test_file_checkpoint_storage_list_checkpoints():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoints for different workflows
+        checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+        await storage.save_checkpoint(checkpoint1)
+        await storage.save_checkpoint(checkpoint2)
+        await storage.save_checkpoint(checkpoint3)
+
+        # Test list_checkpoint_ids for workflow-1
+        workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+        assert len(workflow1_checkpoint_ids) == 2
+        assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+        assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+        # Test list_checkpoints for workflow-1 (returns objects)
+        workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+        assert len(workflow1_checkpoints) == 2
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+        checkpoint_ids = {cp.checkpoint_id for cp in workflow1_checkpoints}
+        assert checkpoint_ids == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+        # Test list_checkpoint_ids for workflow-2
+        workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+        assert len(workflow2_checkpoint_ids) == 1
+        assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+        # Test list_checkpoints for workflow-2 (returns objects)
+        workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+        assert len(workflow2_checkpoints) == 1
+        assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+        # Test list all checkpoints
+        all_checkpoint_ids = await storage.list_checkpoint_ids()
+        assert len(all_checkpoint_ids) == 3
+
+        all_checkpoints = await storage.list_checkpoints()
+        assert len(all_checkpoints) == 3
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_file_checkpoint_storage_delete():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+        # Save checkpoint
+        await storage.save_checkpoint(checkpoint)
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Delete checkpoint
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is True
+        assert not file_path.exists()
+
+        # Try to delete again
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is False
+
+
+async def test_file_checkpoint_storage_directory_creation():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        nested_path = Path(temp_dir) / "nested" / "checkpoint" / "storage"
+        storage = FileCheckpointStorage(nested_path)
+
+        # Directory should be created
+        assert nested_path.exists()
+        assert nested_path.is_dir()
+
+        # Should be able to save checkpoints
+        checkpoint = WorkflowCheckpoint(workflow_id="test")
+        await storage.save_checkpoint(checkpoint)
+
+        file_path = nested_path / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+
+async def test_file_checkpoint_storage_corrupted_file():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a corrupted JSON file
+        corrupted_file = Path(temp_dir) / "corrupted.json"
+        with open(corrupted_file, "w") as f:  # noqa: ASYNC230
+            f.write("{ invalid json }")
+
+        # list_checkpoints should handle the corrupted file gracefully
+        checkpoints = await storage.list_checkpoints("any-workflow")
+        assert checkpoints == []
+
+
+async def test_file_checkpoint_storage_json_serialization():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoint with complex nested data
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="complex-workflow",
+            messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
+            shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
+            executor_states={"executor1": {"state": "active", "config": {"timeout": 30, "retries": 3}}},
+        )
+
+        # Save and load
+        await storage.save_checkpoint(checkpoint)
+        loaded = await storage.load_checkpoint(checkpoint.checkpoint_id)
+
+        assert loaded is not None
+        assert loaded.messages == checkpoint.messages
+        assert loaded.shared_state == checkpoint.shared_state
+        assert loaded.executor_states == checkpoint.executor_states
+
+        # Verify the JSON file is properly formatted
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        with open(file_path) as f:  # noqa: ASYNC230
+            data = json.load(f)
+
+        assert data["messages"]["executor1"][0]["data"]["nested"]["value"] == 42
+        assert data["shared_state"]["list"] == [1, 2, 3]
+        assert data["shared_state"]["bool"] is True
+        assert data["shared_state"]["null"] is None
+
+
+def test_checkpoint_storage_protocol_compliance():
+    # This test ensures both implementations have all required methods
+    memory_storage = InMemoryCheckpointStorage()
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        file_storage = FileCheckpointStorage(temp_dir)
+
+        for storage in [memory_storage, file_storage]:
+            # Test that all protocol methods exist and are callable
+            assert hasattr(storage, "save_checkpoint")
+            assert callable(storage.save_checkpoint)
+            assert hasattr(storage, "load_checkpoint")
+            assert callable(storage.load_checkpoint)
+            assert hasattr(storage, "list_checkpoint_ids")
+            assert callable(storage.list_checkpoint_ids)
+            assert hasattr(storage, "list_checkpoints")
+            assert callable(storage.list_checkpoints)
+            assert hasattr(storage, "delete_checkpoint")
+            assert callable(storage.delete_checkpoint)

python/packages/workflow/tests/test_workflow.py

@@ -1,10 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import tempfile
 from dataclasses import dataclass
 
 import pytest
 from agent_framework.workflow import (
     Executor,
+    FileCheckpointStorage,
     RequestInfoEvent,
     RequestInfoExecutor,
     RequestInfoMessage,
@@ -15,6 +17,8 @@ from agent_framework.workflow import (
     handler,
 )
 
+from agent_framework_workflow import Message
+
 
 @dataclass
 class MockMessage:
@@ -275,3 +279,278 @@ async def test_fan_in():
 
     completed_event = events.get_completed_event()
     assert completed_event is not None and completed_event.data == 4
+
+
+@pytest.fixture
+def simple_executor() -> Executor:
+    class SimpleExecutor(Executor):
+        @handler
+        async def handle_message(self, message: Message, context: WorkflowContext) -> None:
+            pass
+
+    return SimpleExecutor("test_executor")
+
+
+async def test_workflow_with_checkpointing_enabled(simple_executor: Executor):
+    """Test that a workflow can be built with checkpointing enabled."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Build workflow with checkpointing - should not raise any errors
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Verify workflow was created and can run
+        test_message = Message(data="test message", source_id="test", target_id=None)
+        result = await workflow.run(test_message)
+        assert result is not None
+
+
+async def test_workflow_checkpointing_not_enabled_for_external_restore(simple_executor: Executor):
+    """Test that external checkpoint restoration fails when workflow doesn't support checkpointing."""
+    # Build workflow WITHOUT checkpointing
+    workflow = (
+        WorkflowBuilder()
+        .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+        .set_start_executor(simple_executor)
+        .build()
+    )
+
+    # Attempt to restore from checkpoint without providing external storage should fail
+    try:
+        [event async for event in workflow.run_streaming_from_checkpoint("fake-checkpoint-id")]
+        raise AssertionError("Expected ValueError to be raised")
+    except ValueError as e:
+        assert "Cannot restore from checkpoint" in str(e)
+        assert "either provide checkpoint_storage parameter" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(simple_executor: Executor):
+    # Build workflow WITHOUT checkpointing
+    workflow = (
+        WorkflowBuilder()
+        .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+        .set_start_executor(simple_executor)
+        .build()
+    )
+
+    # Attempt to run from checkpoint should fail
+    try:
+        async for _ in workflow.run_streaming_from_checkpoint("fake_checkpoint_id"):
+            pass
+        raise AssertionError("Expected ValueError to be raised")
+    except ValueError as e:
+        assert "Cannot restore from checkpoint" in str(e)
+        assert "either provide checkpoint_storage parameter" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(simple_executor: Executor):
+    """Test that attempting to restore from a non-existent checkpoint fails appropriately."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Attempt to run from non-existent checkpoint should fail
+        try:
+            async for _ in workflow.run_streaming_from_checkpoint("nonexistent_checkpoint_id"):
+                pass
+            raise AssertionError("Expected RuntimeError to be raised")
+        except RuntimeError as e:
+            assert "Failed to restore from checkpoint" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_with_external_storage(simple_executor: Executor):
+    """Test that external checkpoint storage can be provided for restoration."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Create a workflow WITHOUT checkpointing
+        workflow_without_checkpointing = (
+            WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
+        )
+
+        # Resume from checkpoint using external storage parameter
+        try:
+            events: list[WorkflowEvent] = []
+            async for event in workflow_without_checkpointing.run_streaming_from_checkpoint(
+                checkpoint_id, checkpoint_storage=storage
+            ):
+                events.append(event)
+                if len(events) >= 2:  # Limit to avoid infinite loops
+                    break
+        except Exception:
+            # Expected since we have minimal setup, but method should accept the parameters
+            pass
+
+
+async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Executor):
+    """Test the non-streaming run_from_checkpoint method."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Test non-streaming run_from_checkpoint method
+        result = await workflow.run_from_checkpoint(checkpoint_id)
+        assert isinstance(result, list)  # Should return WorkflowRunResult which extends list
+        assert hasattr(result, "get_completed_event")  # Should have WorkflowRunResult methods
+
+
+async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor):
+    """Test that run_streaming_from_checkpoint accepts responses parameter."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Test that run_stream_from_checkpoint accepts responses parameter
+        responses = {"request_123": {"data": "test_response"}}
+
+        try:
+            events: list[WorkflowEvent] = []
+            async for event in workflow.run_streaming_from_checkpoint(checkpoint_id, responses=responses):
+                events.append(event)
+                if len(events) >= 2:  # Limit to avoid infinite loops
+                    break
+        except Exception:
+            # Expected since we have minimal setup, but method should accept the parameters
+            pass
+
+
+@dataclass
+class StateTrackingMessage:
+    """A message that tracks state for testing context reset behavior."""
+
+    data: str
+    run_id: str
+
+
+class StateTrackingExecutor(Executor):
+    """An executor that tracks state in shared state to test context reset behavior."""
+
+    @handler(output_types=[])
+    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext) -> None:
+        """Handle the message and track it in shared state."""
+        # Get existing messages from shared state
+        try:
+            existing_messages = await ctx.get_shared_state("processed_messages")
+        except KeyError:
+            existing_messages = []
+
+        # Record this message
+        message_record = f"{message.run_id}:{message.data}"
+        existing_messages.append(message_record)  # type: ignore
+
+        # Update shared state
+        await ctx.set_shared_state("processed_messages", existing_messages)
+
+        # Complete workflow with current shared state
+        await ctx.add_event(WorkflowCompletedEvent(data=existing_messages.copy()))  # type: ignore
+
+
+async def test_workflow_multiple_runs_no_state_collision():
+    """Test that running the same workflow instance multiple times doesn't have state collision."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create executor that tracks state in shared state
+        state_executor = StateTrackingExecutor("state_executor")
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(state_executor, state_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(state_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Run 1: Should only see messages from run 1
+        result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1"))
+        completed1 = result1.get_completed_event()
+        assert completed1 is not None
+        assert completed1.data == ["run1:message1"]
+
+        # Run 2: Should only see messages from run 2, not run 1
+        result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2"))
+        completed2 = result2.get_completed_event()
+        assert completed2 is not None
+        assert completed2.data == ["run2:message2"]  # Should NOT contain run1 data
+
+        # Run 3: Should only see messages from run 3
+        result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3"))
+        completed3 = result3.get_completed_event()
+        assert completed3 is not None
+        assert completed3.data == ["run3:message3"]  # Should NOT contain run1 or run2 data
+
+        # Verify that each run only processed its own message
+        # This confirms that the checkpointable context properly resets between runs
+        assert completed1.data != completed2.data
+        assert completed2.data != completed3.data
+        assert completed1.data != completed3.data

python/samples/getting_started/workflow/step_02_simple_workflow_condition.py

@@ -3,13 +3,7 @@
 import asyncio
 from dataclasses import dataclass
 
-from agent_framework.workflow import (
-    Executor,
-    WorkflowBuilder,
-    WorkflowCompletedEvent,
-    WorkflowContext,
-    handler,
-)
+from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler
 
 """
 The following sample demonstrates a basic workflow with two executors

python/samples/getting_started/workflow/step_06_map_reduce.py

@@ -3,24 +3,11 @@
 import ast
 import asyncio
 import os
-import sys
 from collections import defaultdict
 from dataclasses import dataclass
 
 import aiofiles
-from agent_framework.workflow import (
-    Executor,
-    WorkflowBuilder,
-    WorkflowCompletedEvent,
-    WorkflowContext,
-    handler,
-)
-
-if sys.version_info >= (3, 12):
-    pass  # pragma: no cover
-else:
-    pass  # pragma: no cover
-
+from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler
 
 """
 The following sample demonstrates a basic map reduce workflow that
@@ -119,7 +106,8 @@ class Map(Executor):
             data: An instance of SplitCompleted signaling the map step can be started.
             ctx: The execution context containing the shared state and other information.
         """
-        # Retrieve the data to be processed from the shared state.# Define a key for the shared state to store the data to be processed
+        # Retrieve the data to be processed from the shared state.
+        # Define a key for the shared state to store the data to be processed
         data_to_be_processed: list[str] = await ctx.get_shared_state(SHARED_STATE_DATA_KEY)
         chunk_start, chunk_end = await ctx.get_shared_state(self.id)
 

python/samples/getting_started/workflow/step_07_checkpoint_resume.py

@@ -0,0 +1,215 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import os
+from pathlib import Path
+
+from agent_framework.workflow import (
+    Executor,
+    FileCheckpointStorage,
+    WorkflowBuilder,
+    WorkflowCompletedEvent,
+    WorkflowContext,
+    handler,
+)
+
+"""
+Demonstrates workflow checkpointing, shared state, and resumption at superstep boundaries.
+
+Flow:
+1) UpperCaseExecutor: "hello world" -> "HELLO WORLD" (writes shared_state: original_input, upper_output)
+2) ReverseTextExecutor: "HELLO WORLD" -> "DLROW OLLEH"
+3) LowerCaseExecutor: "DLROW OLLEH" -> "dlrow olleh" (reads shared_state, emits WorkflowCompletedEvent)
+
+Initial run checkpoints:
+- after_initial_execution: messages from upper_case_executor
+- superstep_1: messages from reverse_text_executor
+- superstep_2: no messages (final events only)
+
+Resume:
+- Resume from the checkpoint containing "DLROW OLLEH" (superstep_1); only LowerCaseExecutor runs.
+- Iteration continues from the checkpoint; one checkpoint is created after the resumed superstep.
+"""
+
+# Define the temporary directory for storing checkpoints
+DIR = os.path.dirname(__file__)
+TEMP_DIR = os.path.join(DIR, "tmp", "checkpoints")
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+
+class UpperCaseExecutor(Executor):
+    @handler(output_types=[str])
+    async def to_upper_case(self, text: str, ctx: WorkflowContext) -> None:
+        result = text.upper()
+        print(f"UpperCaseExecutor: '{text}' -> '{result}'")
+        # Persist executor state into checkpointable context
+        prev = await ctx.get_state() or {}
+        count = int(prev.get("count", 0)) + 1
+        await ctx.set_state({
+            "count": count,
+            "last_input": text,
+            "last_output": result,
+        })
+        # Write to shared_state so downstream executors (and checkpoints) can see it
+        await ctx.set_shared_state("original_input", text)
+        await ctx.set_shared_state("upper_output", result)
+        await ctx.send_message(result)
+
+
+class LowerCaseExecutor(Executor):
+    @handler(output_types=[str])
+    async def to_lower_case(self, text: str, ctx: WorkflowContext) -> None:
+        result = text.lower()
+        print(f"LowerCaseExecutor: '{text}' -> '{result}'")
+        # Read from shared_state written by UpperCaseExecutor
+        orig = await ctx.get_shared_state("original_input")
+        upper = await ctx.get_shared_state("upper_output")
+        print(f"LowerCaseExecutor (shared_state): original_input='{orig}', upper_output='{upper}'")
+        # Persist executor state into checkpointable context
+        prev = await ctx.get_state() or {}
+        count = int(prev.get("count", 0)) + 1
+        await ctx.set_state({
+            "count": count,
+            "last_input": text,
+            "last_output": result,
+            "final": True,
+        })
+        await ctx.add_event(WorkflowCompletedEvent(result))
+
+
+class ReverseTextExecutor(Executor):
+    def __init__(self, id: str):
+        """Initialize the executor with an ID."""
+        super().__init__(id=id)
+
+    @handler(output_types=[str])
+    async def reverse_text(self, text: str, ctx: WorkflowContext) -> None:
+        result = text[::-1]
+        print(f"ReverseTextExecutor: '{text}' -> '{result}'")
+        # Persist executor state into checkpointable context
+        prev = await ctx.get_state() or {}
+        count = int(prev.get("count", 0)) + 1
+        await ctx.set_state({
+            "count": count,
+            "last_input": text,
+            "last_output": result,
+        })
+        await ctx.send_message(result)
+
+
+async def find_checkpoint_with_message(
+    checkpoint_storage: FileCheckpointStorage, workflow_id: str, needle: str
+) -> str | None:
+    """Find the checkpoint that contains a message data value exactly equal to 'needle'."""
+    checkpoints = await checkpoint_storage.list_checkpoints(workflow_id=workflow_id)
+    # Sort by timestamp ascending so earlier checkpoints appear first
+    checkpoints.sort(key=lambda cp: cp.timestamp)
+    for checkpoint in checkpoints:
+        for executor_messages in checkpoint.messages.values():
+            for message in executor_messages:
+                if message.get("data") == needle:
+                    return checkpoint.checkpoint_id
+    return None
+
+
+async def main():
+    # Clear existing checkpoints in this sample directory
+    checkpoint_dir = Path(TEMP_DIR)
+    for file in checkpoint_dir.glob("*.json"):
+        file.unlink()
+
+    upper_case_executor = UpperCaseExecutor(id="upper_case_executor")
+    reverse_text_executor = ReverseTextExecutor(id="reverse_text_executor")
+    lower_case_executor = LowerCaseExecutor(id="lower_case_executor")
+
+    checkpoint_storage = FileCheckpointStorage(storage_path=TEMP_DIR)
+
+    workflow = (
+        WorkflowBuilder(max_iterations=5)
+        .add_edge(upper_case_executor, reverse_text_executor)
+        .add_edge(reverse_text_executor, lower_case_executor)
+        .set_start_executor(upper_case_executor)
+        .with_checkpointing(checkpoint_storage=checkpoint_storage)
+        .build()
+    )
+
+    print("Running workflow with initial message...")
+    async for event in workflow.run_streaming(message="hello world"):
+        print(f"Event: {event}")
+
+    # Inspect checkpoints
+    all_checkpoints = await checkpoint_storage.list_checkpoints()
+    if not all_checkpoints:
+        print("No checkpoints found!")
+        return
+
+    # All checkpoints from this run share one workflow_id
+    workflow_id = all_checkpoints[0].workflow_id
+
+    # Dump a quick summary including shared_state keys of interest
+    print("\nCheckpoint summary:")
+    for cp in sorted(all_checkpoints, key=lambda c: c.timestamp):
+        msg_count = sum(len(v) for v in cp.messages.values())
+        state_keys = sorted(list(cp.executor_states.keys())) if hasattr(cp, "executor_states") else []
+        orig = cp.shared_state.get("original_input") if hasattr(cp, "shared_state") else None
+        upper = cp.shared_state.get("upper_output") if hasattr(cp, "shared_state") else None
+        print(
+            f"- {cp.checkpoint_id} | "
+            f"iter={cp.iteration_count} | messages={msg_count} | states={state_keys} | "
+            f"shared_state: original_input='{orig}', upper_output='{upper}'"
+        )
+
+    # Find the checkpoint with DLROW OLLEH
+    # This will have us resume at the third (last) executor (node)
+    checkpoint_id = await find_checkpoint_with_message(checkpoint_storage, workflow_id, "DLROW OLLEH")
+    if not checkpoint_id:
+        print("Could not find checkpoint with 'DLROW OLLEH'!")
+        return
+
+    # The previous workflow can also be used.
+    # Showing that the workflow can run from a previous checkpoint,
+    # when checkpointing is not enabled for the particular instance.
+    new_workflow = (
+        WorkflowBuilder(max_iterations=5)
+        .add_edge(upper_case_executor, reverse_text_executor)
+        .add_edge(reverse_text_executor, lower_case_executor)
+        .set_start_executor(upper_case_executor)
+        .build()
+    )
+
+    print(f"\nResuming from checkpoint: {checkpoint_id}")
+    async for event in new_workflow.run_streaming_from_checkpoint(checkpoint_id, checkpoint_storage=checkpoint_storage):
+        print(f"Resumed Event: {event}")
+
+    """
+    Sample Output:
+
+    Running workflow with initial message...
+    UpperCaseExecutor: 'hello world' -> 'HELLO WORLD'
+    Event: ExecutorInvokeEvent(executor_id=upper_case_executor)
+    Event: ExecutorCompletedEvent(executor_id=upper_case_executor)
+    ReverseTextExecutor: 'HELLO WORLD' -> 'DLROW OLLEH'
+    Event: ExecutorInvokeEvent(executor_id=reverse_text_executor)
+    Event: ExecutorCompletedEvent(executor_id=reverse_text_executor)
+    LowerCaseExecutor: 'DLROW OLLEH' -> 'dlrow olleh'
+    LowerCaseExecutor (shared_state): original_input='hello world', upper_output='HELLO WORLD'
+    Event: ExecutorInvokeEvent(executor_id=lower_case_executor)
+    Event: WorkflowCompletedEvent(data=dlrow olleh)
+    Event: ExecutorCompletedEvent(executor_id=lower_case_executor)
+
+    Checkpoint summary:
+    - dfc63e72-8e8d-454f-9b6d-0d740b9062e6 | label='after_initial_execution' | iter=0 | messages=1 | states=['upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD'
+    - a78c345a-e5d9-45ba-82c0-cb725452d91b | label='superstep_1' | iter=1 | messages=1 | states=['reverse_text_executor', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD'
+    - 637c1dbd-a525-4404-9583-da03980537a2 | label='superstep_2' | iter=2 | messages=0 | states=['lower_case_executor', 'reverse_text_executor', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD'
+
+    Resuming from checkpoint: a78c345a-e5d9-45ba-82c0-cb725452d91b
+    LowerCaseExecutor: 'DLROW OLLEH' -> 'dlrow olleh'
+    LowerCaseExecutor (shared_state): original_input='hello world', upper_output='HELLO WORLD'
+    Resumed Event: ExecutorInvokeEvent(executor_id=lower_case_executor)
+    Resumed Event: WorkflowCompletedEvent(data=dlrow olleh)
+    Resumed Event: ExecutorCompletedEvent(executor_id=lower_case_executor)
+    """  # noqa: E501
+
+
+if __name__ == "__main__":
+    asyncio.run(main())