[BREAKING] Python: Refactor Checkpointing for runner and runner context (#1645)

* Refactor Checkpointing for runner and runner context

* exception

* Fix formatting

* Comments

* rename

* Add detailed doc string

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

@@ -11,14 +11,31 @@ 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."""
+    """Represents a complete checkpoint of workflow state.
+
+    Checkpoints capture the full execution state of a workflow at a specific point,
+    enabling workflows to be paused and resumed.
+
+    Attributes:
+        checkpoint_id: Unique identifier for this checkpoint
+        workflow_id: Identifier of the workflow this checkpoint belongs to
+        timestamp: ISO 8601 timestamp when checkpoint was created
+        messages: Messages exchanged between executors
+        shared_state: Complete shared state including user data and executor states.
+                     Executor states are stored under the reserved key '_executor_state'.
+        iteration_count: Current iteration number when checkpoint was created
+        metadata: Additional metadata (e.g., superstep info, graph signature)
+        version: Checkpoint format version
+
+    Note:
+        The shared_state dict may contain reserved keys managed by the framework.
+        See SharedState class documentation for details on reserved keys.
+    """
 
     checkpoint_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     workflow_id: str = ""
@@ -27,11 +44,9 @@ class WorkflowCheckpoint:
     # 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]

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

@@ -8,6 +8,7 @@ from typing import Any
 
 from ._checkpoint import WorkflowCheckpoint
 from ._checkpoint_encoding import decode_checkpoint_value
+from ._const import EXECUTOR_STATE_KEY
 from ._request_info_executor import PendingRequestDetails, RequestInfoMessage, RequestResponse
 
 logger = logging.getLogger(__name__)
@@ -33,7 +34,7 @@ def get_checkpoint_summary(
     preview_width: int = 70,
 ) -> WorkflowCheckpointSummary:
     targets = sorted(checkpoint.messages.keys())
-    executor_ids = sorted(checkpoint.executor_states.keys())
+    executor_ids = sorted(checkpoint.shared_state.get(EXECUTOR_STATE_KEY, {}).keys())
     pending = _pending_requests_from_checkpoint(checkpoint, request_executor_ids=request_executor_ids)
 
     draft_preview: str | None = None
@@ -74,7 +75,7 @@ def _pending_requests_from_checkpoint(
 
     pending: dict[str, PendingRequestDetails] = {}
 
-    for state in checkpoint.executor_states.values():
+    for state in checkpoint.shared_state.get(EXECUTOR_STATE_KEY, {}).values():
         if not isinstance(state, Mapping):
             continue
         inner = state.get("pending_requests")

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

@@ -1,3 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-DEFAULT_MAX_ITERATIONS = 100  # Default maximum iterations for workflow execution.
+# Default maximum iterations for workflow execution.
+DEFAULT_MAX_ITERATIONS = 100
+
+# Key used to store executor state in shared state.
+EXECUTOR_STATE_KEY = "_executor_state"

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

@@ -26,6 +26,7 @@ from agent_framework import (
 from agent_framework._agents import BaseAgent
 
 from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
+from ._const import EXECUTOR_STATE_KEY
 from ._events import WorkflowEvent
 from ._executor import Executor, handler
 from ._model_utils import DictConvertible, encode_value
@@ -2136,7 +2137,7 @@ class MagenticWorkflow:
             return
 
         # At this point, checkpoint is guaranteed to be WorkflowCheckpoint
-        executor_states = checkpoint.executor_states
+        executor_states: dict[str, Any] = checkpoint.shared_state.get(EXECUTOR_STATE_KEY, {})
         orchestrator_id = getattr(orchestrator, "id", "")
         orchestrator_state = executor_states.get(orchestrator_id)
         if orchestrator_state is None:

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

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
 
 from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
 from ._checkpoint_encoding import DATACLASS_MARKER, MODEL_MARKER, decode_checkpoint_value
+from ._const import EXECUTOR_STATE_KEY
 from ._edge import EdgeGroup
 from ._edge_runner import EdgeRunner, create_edge_runner
 from ._events import WorkflowEvent
@@ -15,7 +16,6 @@ from ._executor import Executor
 from ._runner_context import (
     Message,
     RunnerContext,
-    WorkflowState,
 )
 from ._shared_state import SharedState
 
@@ -68,16 +68,9 @@ class Runner:
         """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
+    def reset_iteration_count(self) -> None:
+        """Reset the iteration count to zero."""
+        self._iteration = 0
 
     async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
         """Run the workflow until no more messages are sent."""
@@ -100,9 +93,6 @@ class Runner:
                 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}")
 
@@ -134,9 +124,6 @@ class Runner:
                     for event in await self._ctx.drain_events():
                         yield event
 
-                # Update context with current iteration state immediately
-                await self._update_context_with_shared_state()
-
                 logger.info(f"Completed superstep {self._iteration}")
 
                 # Create checkpoint after each superstep iteration
@@ -195,7 +182,6 @@ class Runner:
         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,
@@ -203,7 +189,11 @@ class Runner:
             }
             if self.graph_signature_hash:
                 metadata["graph_signature"] = self.graph_signature_hash
-            checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata)
+            checkpoint_id = await self._ctx.create_checkpoint(
+                self._shared_state,
+                self._iteration,
+                metadata=metadata,
+            )
             logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}")
             return checkpoint_id
         except Exception as e:
@@ -213,6 +203,10 @@ class Runner:
     async def _auto_snapshot_executor_states(self) -> None:
         """Populate executor state by calling snapshot hooks on executors if available.
 
+        TODO(@taochen#1614): this method is potentially problematic if executors also call
+        set_executor_state on the context directly. We should clarify the intended usage
+        pattern for executor state management.
+
         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.
@@ -234,32 +228,13 @@ class Runner:
                         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_executor_state(exec_id, state_dict)
+                    await self._set_executor_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_workflow_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_workflow_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,
@@ -304,20 +279,16 @@ class Runner:
                     checkpoint_id,
                 )
 
-            await self._restore_executor_states(checkpoint.executor_states)
-
-            state = _convert_checkpoint_to_workflow_state(checkpoint)
-            await self._ctx.set_workflow_state(state)
-
-            if checkpoint.workflow_id:
-                self._ctx.set_workflow_id(checkpoint.workflow_id)
             self._workflow_id = checkpoint.workflow_id
+            # Restore shared state
+            await self._shared_state.import_state(checkpoint.shared_state)
+            # Restore executor states using the restored shared state
+            await self._restore_executor_states()
+            # Apply the checkpoint to the context
+            await self._ctx.apply_checkpoint(checkpoint)
+            # Mark the runner as resumed
+            self._mark_resumed(checkpoint.iteration_count)
 
-            await self._restore_shared_state_from_context()
-            self.mark_resumed(
-                iteration=checkpoint.iteration_count,
-                max_iterations=checkpoint.max_iterations,
-            )
             logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
             return True
         except ValueError:
@@ -326,12 +297,24 @@ class Runner:
             logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
             return False
 
-    async def _restore_executor_states(self, executor_states: dict[str, dict[str, Any]]) -> None:
-        for exec_id, state in executor_states.items():
-            executor = self._executors.get(exec_id)
+    async def _restore_executor_states(self) -> None:
+        has_executor_states = await self._shared_state.has(EXECUTOR_STATE_KEY)
+        if not has_executor_states:
+            return
+
+        executor_states = await self._shared_state.get(EXECUTOR_STATE_KEY)
+        if not isinstance(executor_states, dict):
+            raise ValueError("Executor states in shared state is not a dictionary. Unable to restore.")
+
+        for executor_id, state in executor_states.items():
+            if not isinstance(executor_id, str):
+                raise ValueError("Executor ID in executor states is not a string. Unable to restore.")
+            if not isinstance(state, dict):
+                raise ValueError(f"Executor state for {executor_id} is not a dictionary. Unable to restore.")
+
+            executor = self._executors.get(executor_id)
             if not executor:
-                logger.debug(f"Executor {exec_id} not found during state restoration; skipping.")
-                continue
+                raise ValueError(f"Executor {executor_id} not found during state restoration.")
 
             restored = False
             restore_method = getattr(executor, "restore_state", None)
@@ -342,26 +325,10 @@ class Runner:
                         await maybe  # type: ignore[arg-type]
                     restored = True
             except Exception as ex:  # pragma: no cover - defensive
-                logger.debug(f"Executor {exec_id} restore_state failed: {ex}")
+                raise ValueError(f"Executor {executor_id} restore_state failed: {ex}") from ex
 
             if not restored:
-                logger.debug(f"Executor {exec_id} does not support state restoration; skipping.")
-
-    async def _restore_shared_state_from_context(self) -> None:
-        try:
-            restored_state = await self._ctx.get_workflow_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}")
+                logger.debug(f"Executor {executor_id} does not support state restoration; skipping.")
 
     def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]:
         """Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners.
@@ -413,13 +380,28 @@ class Runner:
                 return True
         return False
 
+    def _mark_resumed(self, iteration: int) -> None:
+        """Mark the runner as having resumed from a checkpoint.
 
-def _convert_checkpoint_to_workflow_state(checkpoint: WorkflowCheckpoint) -> WorkflowState:
-    """Helper function to convert a WorkflowCheckpoint to a WorkflowState."""
-    return {
-        "messages": checkpoint.messages,
-        "shared_state": checkpoint.shared_state,
-        "executor_states": checkpoint.executor_states,
-        "iteration_count": checkpoint.iteration_count,
-        "max_iterations": checkpoint.max_iterations,
-    }
+        Optionally set the current iteration and max iterations.
+        """
+        self._resumed_from_checkpoint = True
+        self._iteration = iteration
+
+    async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
+        """Store executor state in shared state under a reserved key.
+
+        Executors call this with a JSON-serializable dict capturing the minimal
+        state needed to resume. It replaces any previously stored state.
+        """
+        has_existing_states = await self._shared_state.has(EXECUTOR_STATE_KEY)
+        if has_existing_states:
+            existing_states = await self._shared_state.get(EXECUTOR_STATE_KEY)
+        else:
+            existing_states = {}
+
+        if not isinstance(existing_states, dict):
+            raise ValueError("Existing executor states in shared state is not a dictionary.")
+
+        existing_states[executor_id] = state
+        await self._shared_state.set(EXECUTOR_STATE_KEY, existing_states)

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

@@ -5,11 +5,10 @@ import logging
 import uuid
 from copy import copy
 from dataclasses import dataclass
-from typing import Any, Protocol, TypedDict, TypeVar, cast, runtime_checkable
+from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable
 
 from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
 from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
-from ._const import DEFAULT_MAX_ITERATIONS
 from ._events import WorkflowEvent
 from ._shared_state import SharedState
 
@@ -43,7 +42,7 @@ class Message:
         return self.source_span_ids[0] if self.source_span_ids else None
 
 
-class WorkflowState(TypedDict):
+class _WorkflowState(TypedDict):
     """TypedDict representing the serializable state of a workflow execution.
 
     This includes all state data needed for checkpointing and restoration.
@@ -51,9 +50,7 @@ class WorkflowState(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
@@ -116,26 +113,6 @@ class RunnerContext(Protocol):
         """Wait for and return the next event emitted by the workflow run."""
         ...
 
-    async def set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
-        """Set the state for a specific executor.
-
-        Args:
-            executor_id: The ID of the executor whose state is being set.
-            state: The state data to be set for the executor.
-        """
-        ...
-
-    async def get_executor_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.
@@ -150,7 +127,7 @@ class RunnerContext(Protocol):
         """Set the workflow ID for the context."""
         ...
 
-    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
+    def reset_for_new_run(self) -> None:
         """Reset the context for a new workflow run."""
         ...
 
@@ -170,27 +147,42 @@ class RunnerContext(Protocol):
         """
         ...
 
-    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
+    async def create_checkpoint(
+        self,
+        shared_state: SharedState,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
         """Create a checkpoint of the current workflow state.
 
         Args:
+            shared_state: The shared state to include in the checkpoint.
+                          This is needed to capture the full state of the workflow.
+                          The shared state is not managed by the context itself.
+            iteration_count: The current iteration count of the workflow.
             metadata: Optional metadata to associate with the checkpoint.
+
+        Returns:
+            The ID of the created checkpoint.
         """
         ...
 
     async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
-        """Load a checkpoint without mutating the current context state."""
-        ...
-
-    async def get_workflow_state(self) -> WorkflowState:
-        """Get the current state of the workflow suitable for checkpointing."""
-        ...
-
-    async def set_workflow_state(self, state: WorkflowState) -> None:
-        """Set the state of the workflow from a checkpoint.
+        """Load a checkpoint without mutating the current context state.
 
         Args:
-            state: The state data to set for the workflow.
+            checkpoint_id: The ID of the checkpoint to load.
+
+        Returns:
+            The loaded checkpoint, or None if it does not exist.
+        """
+        ...
+
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+        """Apply a checkpoint to the current context, mutating its state.
+
+        Args:
+            checkpoint: The checkpoint whose state is to be applied.
         """
         ...
 
@@ -211,14 +203,11 @@ class InProcRunnerContext:
         # 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
 
         # Streaming flag - set by workflow's run_stream() vs run()
         self._streaming: bool = False
 
+    # region Messaging and Events
     async def send_message(self, message: Message) -> None:
         self._messages.setdefault(message.source_id, [])
         self._messages[message.source_id].append(message)
@@ -259,15 +248,70 @@ class InProcRunnerContext:
         """
         return await self._event_queue.get()
 
-    async def set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
-        self._executor_states[executor_id] = state
+    # endregion Messaging and Events
 
-    async def get_executor_state(self, executor_id: str) -> dict[str, Any] | None:
-        return self._executor_states.get(executor_id)
+    # region Checkpointing
 
     def has_checkpointing(self) -> bool:
         return self._checkpoint_storage is not None
 
+    async def create_checkpoint(
+        self,
+        shared_state: SharedState,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        self._workflow_id = self._workflow_id or str(uuid.uuid4())
+        state = await self._get_serialized_workflow_state(shared_state, iteration_count)
+
+        checkpoint = WorkflowCheckpoint(
+            workflow_id=self._workflow_id,
+            messages=state["messages"],
+            shared_state=state["shared_state"],
+            iteration_count=state["iteration_count"],
+            metadata=metadata or {},
+        )
+        checkpoint_id = await self._checkpoint_storage.save_checkpoint(checkpoint)
+        logger.info(f"Created checkpoint {checkpoint_id} for workflow {self._workflow_id}")
+        return checkpoint_id
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+        return await self._checkpoint_storage.load_checkpoint(checkpoint_id)
+
+    def reset_for_new_run(self) -> None:
+        """Reset the context for a new workflow run.
+
+        This clears messages, events, and resets streaming flag.
+        """
+        self._messages.clear()
+        # Clear any pending events (best-effort) by recreating the queue
+        self._event_queue = asyncio.Queue()
+        self._streaming = False  # Reset streaming flag
+
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+        self._messages.clear()
+        messages_data = checkpoint.messages
+        for source_id, message_list in messages_data.items():
+            self._messages[source_id] = [
+                Message(
+                    data=decode_checkpoint_value(msg.get("data")),
+                    source_id=msg.get("source_id", ""),
+                    target_id=msg.get("target_id"),
+                    trace_contexts=msg.get("trace_contexts"),
+                    source_span_ids=msg.get("source_span_ids"),
+                )
+                for msg in message_list
+            ]
+
+        self._workflow_id = checkpoint.workflow_id
+
+    # endregion Checkpointing
+
     def set_workflow_id(self, workflow_id: str) -> None:
         self._workflow_id = workflow_id
 
@@ -287,44 +331,7 @@ class InProcRunnerContext:
         """
         return self._streaming
 
-    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
-        self._messages.clear()
-        # Clear any pending events (best-effort) by recreating the queue
-        self._event_queue = asyncio.Queue()
-        self._shared_state.clear()
-        self._executor_states.clear()
-        self._iteration_count = 0
-        self._streaming = False  # Reset streaming flag
-        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_workflow_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 load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
-        if not self._checkpoint_storage:
-            raise ValueError("Checkpoint storage not configured")
-        return await self._checkpoint_storage.load_checkpoint(checkpoint_id)
-
-    async def get_workflow_state(self) -> WorkflowState:
+    async def _get_serialized_workflow_state(self, shared_state: SharedState, iteration_count: int) -> _WorkflowState:
         serializable_messages: dict[str, list[dict[str, Any]]] = {}
         for source_id, message_list in self._messages.items():
             serializable_messages[source_id] = [
@@ -340,48 +347,6 @@ class InProcRunnerContext:
 
         return {
             "messages": serializable_messages,
-            "shared_state": encode_checkpoint_value(self._shared_state),
-            "executor_states": encode_checkpoint_value(self._executor_states),
-            "iteration_count": self._iteration_count,
-            "max_iterations": self._max_iterations,
+            "shared_state": encode_checkpoint_value(await shared_state.export_state()),
+            "iteration_count": iteration_count,
         }
-
-    async def set_workflow_state(self, state: WorkflowState) -> None:
-        self._messages.clear()
-        messages_data = state.get("messages", {})
-        for source_id, message_list in messages_data.items():
-            self._messages[source_id] = [
-                Message(
-                    data=decode_checkpoint_value(msg.get("data")),
-                    source_id=msg.get("source_id", ""),
-                    target_id=msg.get("target_id"),
-                    trace_contexts=msg.get("trace_contexts"),
-                    source_span_ids=msg.get("source_span_ids"),
-                )
-                for msg in message_list
-            ]
-        # Restore shared_state
-        decoded_shared_raw = decode_checkpoint_value(state.get("shared_state", {}))
-        if isinstance(decoded_shared_raw, dict):
-            self._shared_state = cast(dict[str, Any], decoded_shared_raw)
-        else:  # fallback to empty dict if corrupted
-            self._shared_state = {}
-
-        # Restore executor_states ensuring value types are dicts
-        decoded_exec_raw = decode_checkpoint_value(state.get("executor_states", {}))
-        if isinstance(decoded_exec_raw, dict):
-            typed_exec: dict[str, dict[str, Any]] = {}
-            for k_raw, v_raw in decoded_exec_raw.items():  # type: ignore[assignment]
-                if isinstance(k_raw, str) and isinstance(v_raw, dict):
-                    # Filter inner dict to string keys only (best-effort)
-                    inner: dict[str, Any] = {}
-                    for inner_k, inner_v in v_raw.items():  # type: ignore[assignment]
-                        if isinstance(inner_k, str):
-                            inner[inner_k] = inner_v
-                    typed_exec[k_raw] = inner
-            self._executor_states = typed_exec
-        else:
-            self._executor_states = {}
-
-        self._iteration_count = state.get("iteration_count", 0)
-        self._max_iterations = state.get("max_iterations", 100)

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

@@ -7,7 +7,21 @@ from typing import Any
 
 
 class SharedState:
-    """A class to manage shared state in a workflow."""
+    """A class to manage shared state in a workflow.
+
+    SharedState provides thread-safe access to workflow state data that needs to be
+    shared across executors during workflow execution.
+
+    Reserved Keys:
+        The following keys are reserved for internal framework use and should not be
+        modified by user code:
+
+        - `_executor_state`: Stores executor state for checkpointing (managed by Runner)
+
+    Warning:
+        Do not use keys starting with underscore (_) as they may be reserved for
+        internal framework operations.
+    """
 
     def __init__(self) -> None:
         """Initialize the shared state."""
@@ -34,6 +48,24 @@ class SharedState:
         async with self._shared_state_lock:
             await self.delete_within_hold(key)
 
+    async def clear(self) -> None:
+        """Clear the entire shared state."""
+        async with self._shared_state_lock:
+            self._state.clear()
+
+    async def export_state(self) -> dict[str, Any]:
+        """Get a serialized copy of the entire shared state."""
+        async with self._shared_state_lock:
+            return dict(self._state)
+
+    async def import_state(self, state: dict[str, Any]) -> None:
+        """Populate the shared state from a serialized state dictionary.
+
+        This replaces the entire current state with the provided state.
+        """
+        async with self._shared_state_lock:
+            self._state.update(state)
+
     @asynccontextmanager
     async def hold(self) -> AsyncIterator["SharedState"]:
         """Context manager to hold the shared state lock for multiple operations.

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

@@ -183,13 +183,11 @@ class Workflow(DictConvertible):
         # Convert start_executor to string ID if it's an Executor instance
         start_executor_id = start_executor.id if isinstance(start_executor, Executor) else start_executor
 
-        id = str(uuid.uuid4())
-
         self.edge_groups = list(edge_groups)
         self.executors = dict(executors)
         self.start_executor_id = start_executor_id
         self.max_iterations = max_iterations
-        self.id = id
+        self.id = str(uuid.uuid4())
         self.name = name
         self.description = description
 
@@ -202,7 +200,7 @@ class Workflow(DictConvertible):
             self._shared_state,
             runner_context,
             max_iterations=max_iterations,
-            workflow_id=id,
+            workflow_id=self.id,
         )
 
         # Flag to prevent concurrent workflow executions
@@ -317,7 +315,9 @@ class Workflow(DictConvertible):
 
                 # Reset context for a new run if supported
                 if reset_context:
-                    self._runner.context.reset_for_new_run(self._shared_state)
+                    self._runner.reset_iteration_count()
+                    self._runner.context.reset_for_new_run()
+                    await self._shared_state.clear()
 
                 # Set streaming mode after reset
                 self._runner_context.set_streaming(streaming)

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

@@ -11,6 +11,7 @@ from opentelemetry.trace import SpanKind
 from typing_extensions import Never, TypeVar
 
 from ..observability import OtelAttr, create_workflow_span
+from ._const import EXECUTOR_STATE_KEY
 from ._events import (
     WorkflowEvent,
     WorkflowEventSource,
@@ -436,16 +437,34 @@ class WorkflowContext(Generic[T_Out, T_W_Out]):
         return self._shared_state
 
     async def set_executor_state(self, state: dict[str, Any]) -> None:
-        """Persist this executor's state into the checkpointable context.
+        """Store executor state in shared state under a reserved key.
 
         Executors call this with a JSON-serializable dict capturing the minimal
         state needed to resume. It replaces any previously stored state.
         """
-        await self._runner_context.set_executor_state(self._executor_id, state)
+        has_existing_states = await self._shared_state.has(EXECUTOR_STATE_KEY)
+        if has_existing_states:
+            existing_states = await self._shared_state.get(EXECUTOR_STATE_KEY)
+        else:
+            existing_states = {}
+
+        if not isinstance(existing_states, dict):
+            raise ValueError("Existing executor states in shared state is not a dictionary.")
+
+        existing_states[self._executor_id] = state
+        await self._shared_state.set(EXECUTOR_STATE_KEY, existing_states)
 
     async def get_executor_state(self) -> dict[str, Any] | None:
         """Retrieve previously persisted state for this executor, if any."""
-        return await self._runner_context.get_executor_state(self._executor_id)
+        has_existing_states = await self._shared_state.has(EXECUTOR_STATE_KEY)
+        if not has_existing_states:
+            return None
+
+        existing_states = await self._shared_state.get(EXECUTOR_STATE_KEY)
+        if not isinstance(existing_states, dict):
+            raise ValueError("Existing executor states in shared state is not a dictionary.")
+
+        return existing_states.get(self._executor_id)
 
     def is_streaming(self) -> bool:
         """Check if the workflow is running in streaming mode.

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

@@ -20,9 +20,7 @@ def test_workflow_checkpoint_default_values():
     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"
 
@@ -35,9 +33,7 @@ def test_workflow_checkpoint_custom_values():
         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",
     )
@@ -47,9 +43,7 @@ def test_workflow_checkpoint_custom_values():
     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"
 
@@ -290,7 +284,6 @@ async def test_file_checkpoint_storage_json_serialization():
             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
@@ -300,7 +293,6 @@ async def test_file_checkpoint_storage_json_serialization():
         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"

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

@@ -8,6 +8,7 @@ from typing import Any
 from agent_framework._workflows._checkpoint import WorkflowCheckpoint
 from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
 from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
+from agent_framework._workflows._const import EXECUTOR_STATE_KEY
 from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent
 from agent_framework._workflows._request_info_executor import (
     PendingRequestDetails,
@@ -16,10 +17,7 @@ from agent_framework._workflows._request_info_executor import (
     RequestInfoMessage,
     RequestResponse,
 )
-from agent_framework._workflows._runner_context import (
-    Message,
-    WorkflowState,
-)
+from agent_framework._workflows._runner_context import Message
 from agent_framework._workflows._shared_state import SharedState
 from agent_framework._workflows._workflow_context import WorkflowContext
 
@@ -29,9 +27,6 @@ PENDING_STATE_KEY = RequestInfoExecutor._PENDING_SHARED_STATE_KEY  # pyright: ig
 class _StubRunnerContext:
     """Minimal runner context stub for exercising WorkflowContext helpers."""
 
-    def __init__(self, stored_state: dict[str, Any] | None = None) -> None:
-        self._state = stored_state or {}
-
     async def send_message(self, message: Message) -> None:  # pragma: no cover - unused in tests
         return None
 
@@ -53,31 +48,27 @@ class _StubRunnerContext:
     async def next_event(self) -> WorkflowEvent:  # pragma: no cover - unused
         raise RuntimeError("Not implemented in stub context")
 
-    async def get_executor_state(self, executor_id: str) -> dict[str, Any] | None:  # pragma: no cover - trivial
-        return self._state
-
-    async def set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:  # pragma: no cover - unused
-        self._state = state
-
     def has_checkpointing(self) -> bool:  # pragma: no cover - unused
         return False
 
     def set_workflow_id(self, workflow_id: str) -> None:  # pragma: no cover - unused
         pass
 
-    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:  # pragma: no cover - unused
+    def reset_for_new_run(self) -> None:  # pragma: no cover - unused
         pass
 
-    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:  # pragma: no cover - unused
+    async def create_checkpoint(
+        self,
+        shared_state: SharedState,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:  # pragma: no cover - unused
         raise RuntimeError("Checkpointing not supported in stub context")
 
     async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:  # pragma: no cover - unused
         return None
 
-    async def get_workflow_state(self) -> WorkflowState:  # pragma: no cover - unused
-        return {}  # type: ignore[return-value]
-
-    async def set_workflow_state(self, state: WorkflowState) -> None:  # pragma: no cover - unused
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:  # pragma: no cover - unused
         pass
 
     def set_streaming(self, streaming: bool) -> None:  # pragma: no cover - unused
@@ -120,8 +111,8 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None:
         },
     )
 
-    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -143,8 +134,8 @@ async def test_has_pending_request_detects_snapshot() -> None:
         },
     )
 
-    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -152,9 +143,8 @@ async def test_has_pending_request_detects_snapshot() -> None:
 
 
 async def test_has_pending_request_false_when_snapshot_absent() -> None:
-    shared_state = SharedState()
-    runner_ctx = _StubRunnerContext({"pending_requests": {}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -196,7 +186,6 @@ def test_pending_requests_from_checkpoint_and_summary() -> None:
                 }
             }
         },
-        executor_states={},
         iteration_count=1,
     )
 
@@ -284,10 +273,13 @@ async def test_run_persists_pending_requests_in_runner_state() -> None:
     await executor.execute(approval, ctx.source_executor_ids, shared_state, runner_ctx)
 
     # Runner state should include both pending snapshot and serialized request events
-    assert PENDING_STATE_KEY in runner_ctx._state  # pyright: ignore[reportPrivateUsage]
-    assert approval.request_id in runner_ctx._state[PENDING_STATE_KEY]  # pyright: ignore[reportPrivateUsage]
+    assert await shared_state.has(EXECUTOR_STATE_KEY)
+    executor_state = await shared_state.get(EXECUTOR_STATE_KEY)
+    assert executor.id in executor_state
+    assert PENDING_STATE_KEY in executor_state[executor.id]
+    assert approval.request_id in executor_state[executor.id][PENDING_STATE_KEY]
 
     response_ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx)
     await executor.handle_response("approved", approval.request_id, response_ctx)  # type: ignore
 
-    assert runner_ctx._state[PENDING_STATE_KEY] == {}  # pyright: ignore[reportPrivateUsage]
+    assert executor_state[executor.id][PENDING_STATE_KEY] == {}

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

@@ -414,9 +414,7 @@ async def test_workflow_run_stream_from_checkpoint_with_external_storage(simple_
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -451,9 +449,7 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -484,9 +480,7 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executo
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -525,7 +519,7 @@ class StateTrackingExecutor(Executor):
     """An executor that tracks state in shared state to test context reset behavior."""
 
     @handler
-    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list]) -> None:
+    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list[Any]]) -> None:
         """Handle the message and track it in shared state."""
         # Get existing messages from shared state
         try:

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

@@ -6,7 +6,7 @@ import pytest
 from opentelemetry import trace
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-from agent_framework import WorkflowBuilder
+from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder
 from agent_framework._workflows._executor import Executor, handler
 from agent_framework._workflows._runner_context import InProcRunnerContext, Message
 from agent_framework._workflows._shared_state import SharedState
@@ -426,7 +426,7 @@ async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExp
 @pytest.mark.parametrize("enable_otel", [False], indirect=True)
 async def test_message_trace_context_serialization(span_exporter: InMemorySpanExporter) -> None:
     """Test that message trace context is properly serialized/deserialized."""
-    ctx = InProcRunnerContext()
+    ctx = InProcRunnerContext(InMemoryCheckpointStorage())
 
     # Create message with trace context
     message = Message(
@@ -439,16 +439,18 @@ async def test_message_trace_context_serialization(span_exporter: InMemorySpanEx
 
     await ctx.send_message(message)
 
-    # Get context state (which serializes messages)
-    state = await ctx.get_workflow_state()
+    # Create a checkpoint that includes the message
+    checkpoint_id = await ctx.create_checkpoint(SharedState(), 0)
+    checkpoint = await ctx.load_checkpoint(checkpoint_id)
+    assert checkpoint is not None
 
     # Check serialized message includes trace context
-    serialized_msg = state["messages"]["source"][0]
+    serialized_msg = checkpoint.messages["source"][0]
     assert serialized_msg["trace_contexts"] == [{"traceparent": "00-trace-span-01"}]
     assert serialized_msg["source_span_ids"] == ["span123"]
 
     # Test deserialization
-    await ctx.set_workflow_state(state)
+    await ctx.apply_checkpoint(checkpoint)
     restored_messages = await ctx.drain_messages()
 
     restored_msg = list(restored_messages.values())[0][0]

python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py

@@ -196,7 +196,7 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None:
     for cp in sorted(checkpoints, key=lambda c: c.timestamp):
         summary = get_checkpoint_summary(cp)
         msg_count = sum(len(v) for v in cp.messages.values())
-        state_keys = sorted(cp.executor_states.keys())
+        state_keys = sorted(summary.executor_ids)
         orig = cp.shared_state.get("original_input")
         upper = cp.shared_state.get("upper_output")