diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d47287237f..75e6b70451 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/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] diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py b/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py index a102055f62..06f4a025fb 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py +++ b/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") diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index b426692725..0adf0b0978 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/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" diff --git a/python/packages/core/agent_framework/_workflows/_magentic.py b/python/packages/core/agent_framework/_workflows/_magentic.py index 55a2c29a71..2b132633aa 100644 --- a/python/packages/core/agent_framework/_workflows/_magentic.py +++ b/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: diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 16c841c279..bf1ab93d8a 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/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) diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 05ff147f12..d91e73a69a 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/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) diff --git a/python/packages/core/agent_framework/_workflows/_shared_state.py b/python/packages/core/agent_framework/_workflows/_shared_state.py index 22225518ea..93057021fb 100644 --- a/python/packages/core/agent_framework/_workflows/_shared_state.py +++ b/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. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 7cce9a1429..e24d25c8a2 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/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) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index a91d5af14e..1a9562fca6 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/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. diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 0515a7ca35..236696e79c 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/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" diff --git a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py b/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py index e0412b4c1c..91bc716829 100644 --- a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py +++ b/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] == {} diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index c87dcfaf9f..f66c7048e8 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/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: diff --git a/python/packages/core/tests/workflow/test_workflow_observability.py b/python/packages/core/tests/workflow/test_workflow_observability.py index e8c0f809e6..d7ceb18f90 100644 --- a/python/packages/core/tests/workflow/test_workflow_observability.py +++ b/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] diff --git a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py index 93307bf0e6..17fb44c87b 100644 --- a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py +++ b/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")