Python: Fix sub-workflow checkpoint restore to preserve sub-workflow state

python/packages/azurefunctions/agent_framework_azurefunctions/_context.py

@@ -115,6 +115,18 @@ async def create_checkpoint(
         """Checkpointing not supported in activity context."""
         raise NotImplementedError("Checkpointing is not supported in Azure Functions activity context")
 
+    async def build_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: str | None,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> WorkflowCheckpoint:
+        """Checkpointing not supported in activity context."""
+        raise NotImplementedError("Checkpointing is not supported in Azure Functions activity context")
+
     async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
         """Checkpointing not supported in activity context."""
         raise NotImplementedError("Checkpointing is not supported in Azure Functions activity context")

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

@@ -811,6 +811,24 @@ def run(
                 execution is not allowed).
         """
         self._validate_run_params(message, responses, checkpoint_id)
+        # Warn (but don't block) when a fresh message or a checkpoint restore begins while a prior
+        # run left request_info events pending. Mirrors Workflow.run. Delivering responses is the
+        # normal way to complete the pending cycle and is intentionally not warned.
+        if (message is not None or checkpoint_id is not None) and self._last_pending_request_ids:
+            logger.warning(
+                "Workflow %s received %s while %d request_info event(s) are still pending from an "
+                "unfinished request/response cycle; %s. Deliver responses (responses=...) to complete "
+                "the pending cycle before starting new input.",
+                self.name,
+                "a fresh message" if message is not None else "a checkpoint restore",
+                len(self._last_pending_request_ids),
+                (
+                    "those requests remain answerable, but this run advances workflow state, so a "
+                    "response that arrives later may apply to a workflow that has moved on"
+                    if message is not None
+                    else "those pending requests will be overwritten by the checkpoint's state"
+                ),
+            )
         if responses and checkpoint_id is None:
             # Require at least one response key to match a currently-pending
             # request; prevents silent replay against stale state while still

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

@@ -245,7 +245,11 @@ async def _prepare_checkpoint_state(self) -> None:
         self._state.commit()
 
     async def create_checkpoint_if_enabled(self) -> None:
-        """Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
+        """Create a checkpoint and save the checkpoint to the configured storage if one is configured.
+
+        Note:
+            1. This method has no effect if checkpointing is not enabled in the context.
+        """
         if not self._ctx.has_checkpointing():
             return
 
@@ -340,6 +344,57 @@ async def restore_from_checkpoint(
             logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
             raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint_id}") from e
 
+    async def build_checkpoint(self) -> WorkflowCheckpoint:
+        """Create a checkpoint object.
+
+        Returns:
+            A ``WorkflowCheckpoint``.
+        """
+        # Persist executor snapshots into committed shared state before exporting it.
+        await self._prepare_checkpoint_state()
+        return await self._ctx.build_checkpoint(
+            self._workflow_name,
+            self._graph_signature_hash,
+            self._state,
+            None,
+            self._iteration,
+        )
+
+    async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+        """Restore runner state from an in-memory ``WorkflowCheckpoint`` object.
+
+        Unlike :meth:`restore_from_checkpoint`, this does not load from a storage
+        backend; it applies a checkpoint the caller already holds - for example, a
+        child workflow checkpoint embedded in a parent ``WorkflowExecutor``'s state.
+
+        Restores shared state, executor snapshots, in-flight messages, and pending
+        request_info events, then marks the runner as resumed.
+
+        Args:
+            checkpoint: The checkpoint whose state should be restored.
+
+        Raises:
+            WorkflowCheckpointException: If the checkpoint's graph signature does not
+                match this runner's workflow, or if restoration otherwise fails.
+        """
+        if self._graph_signature_hash != checkpoint.graph_signature_hash:
+            raise WorkflowCheckpointException(
+                "Workflow graph has changed since the checkpoint was created. "
+                "Please rebuild the original workflow before resuming."
+            )
+
+        try:
+            # Clear first so import_state (which merges) does not leak stale keys from a
+            # prior run on this Workflow instance.
+            self._state.clear()
+            self._state.import_state(checkpoint.state)
+            await self._restore_executor_states()
+            await self._ctx.apply_checkpoint(checkpoint)
+            self._mark_resumed(checkpoint)
+        except Exception as e:
+            logger.error(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}: {e}")
+            raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}") from e
+
     async def _save_executor_states(self) -> None:
         """Populate executor state by calling checkpoint hooks on executors."""
         for exec_id, executor in self._executors.items():

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

@@ -195,6 +195,34 @@ def is_streaming(self) -> bool:
         """
         ...
 
+    async def build_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: CheckpointID | None,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> WorkflowCheckpoint:
+        """Build a checkpoint and return it for the caller to own.
+
+        The checkpoint is constructed in memory and handed back to the caller; nothing is
+        persisted and no checkpoint storage is required.
+
+        Args:
+            workflow_name: The name of the workflow for which the checkpoint is being created.
+            graph_signature_hash: Hash of the workflow graph topology to
+                validate checkpoint compatibility during restore.
+            state: The state to include in the checkpoint.
+            previous_checkpoint_id: The ID of the previous checkpoint, if any, to form a checkpoint chain.
+            iteration_count: The current iteration count of the workflow.
+            metadata: Optional metadata to associate with the checkpoint.
+
+        Returns:
+            A ``WorkflowCheckpoint`` of the current context state.
+        """
+        ...
+
     async def create_checkpoint(
         self,
         workflow_name: str,
@@ -204,7 +232,7 @@ async def create_checkpoint(
         iteration_count: int,
         metadata: dict[str, Any] | None = None,
     ) -> CheckpointID:
-        """Create a checkpoint of the current workflow state.
+        """Persist a checkpoint of the current workflow state to configured storage and return its ID.
 
         Args:
             workflow_name: The name of the workflow for which the checkpoint is being created.
@@ -219,6 +247,9 @@ async def create_checkpoint(
 
         Returns:
             The ID of the created checkpoint.
+
+        Raises:
+            ValueError: If checkpoint storage is not configured.
         """
         ...
 
@@ -381,29 +412,48 @@ def clear_runtime_checkpoint_storage(self) -> None:
     def has_checkpointing(self) -> bool:
         return self._get_effective_checkpoint_storage() is not None
 
-    async def create_checkpoint(
+    async def build_checkpoint(
         self,
         workflow_name: str,
         graph_signature_hash: str,
         state: State,
         previous_checkpoint_id: CheckpointID | None,
         iteration_count: int,
         metadata: dict[str, Any] | None = None,
-    ) -> CheckpointID:
-        storage = self._get_effective_checkpoint_storage()
-        if not storage:
-            raise ValueError("Checkpoint storage not configured")
-
-        checkpoint = WorkflowCheckpoint(
+    ) -> WorkflowCheckpoint:
+        return WorkflowCheckpoint(
             workflow_name=workflow_name,
             graph_signature_hash=graph_signature_hash,
             previous_checkpoint_id=previous_checkpoint_id,
-            messages=dict(self._messages),
+            # Copy the per-source lists so the snapshot is isolated from later context mutations.
+            messages={source_id: list(messages) for source_id, messages in self._messages.items()},
             state=state.export_state(),
             pending_request_info_events=dict(self._pending_request_info_events),
             iteration_count=iteration_count,
             metadata=metadata or {},
         )
+
+    async def create_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: CheckpointID | None,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> CheckpointID:
+        storage = self._get_effective_checkpoint_storage()
+        if not storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        checkpoint = await self.build_checkpoint(
+            workflow_name,
+            graph_signature_hash,
+            state,
+            previous_checkpoint_id,
+            iteration_count,
+            metadata,
+        )
         checkpoint_id = await storage.save(checkpoint)
         logger.debug(f"Created checkpoint {checkpoint_id}")
         return checkpoint_id

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

@@ -814,10 +814,9 @@ async def _run_core(
             # runner context has fully drained from any prior run. If it still
             # has in-flight executor messages, the prior run didn't complete -
             # the caller must either resume from a checkpoint or wait for the
-            # prior run to drain. (Pending request_info events are intentionally
-            # NOT blocked here: a follow-up run with message=... is the normal
-            # way to deliver a response to those pending requests, e.g. via
-            # WorkflowAgent._process_pending_requests.)
+            # prior run to drain. Pending request_info events are intentionally
+            # NOT blocked here (they are answered via a follow-up ``responses=...``
+            # run); the warning below surfaces the abandon/overwrite cases instead.
             # NOTE: _validate_run_params already enforces that ``message`` is
             # mutually exclusive with both ``checkpoint_id`` and ``responses``,
             # so we don't need to re-check those here.
@@ -830,6 +829,33 @@ async def _run_core(
                     "checkpointing; there is no in-process recovery path."
                 )
 
+            # Warn (but don't block) when a fresh message or a checkpoint restore begins while the
+            # workflow still has pending request_info events from an unfinished request/response
+            # cycle. A fresh ``message`` does NOT drop those pending requests - they remain pending and
+            # can still be answered later - but the new run advances executor and shared state, so when
+            # a response for an earlier request eventually arrives the workflow may have moved on,
+            # yielding inconsistent results. A ``checkpoint_id`` restore instead replaces the context's
+            # pending requests with the checkpoint's state. Delivering ``responses`` is the normal way to
+            # answer pending requests and is intentionally not warned. Mirrors the WorkflowExecutor
+            # warning for overlapping sub-workflow executions.
+            if message is not None or checkpoint_id is not None:
+                pending_request_info_events = await self._runner.context.get_pending_request_info_events()
+                if pending_request_info_events:
+                    logger.warning(
+                        "Workflow %s received %s while %d request_info event(s) are still pending from an "
+                        "unfinished request/response cycle; %s. Deliver responses (responses=...) to complete "
+                        "the pending cycle before starting new input.",
+                        self.id,
+                        "a fresh message" if message is not None else "a checkpoint restore",
+                        len(pending_request_info_events),
+                        (
+                            "those requests remain pending, but this run advances executor and shared state, "
+                            "so a response that arrives later may apply to a workflow that has moved on"
+                            if message is not None
+                            else "those pending requests will be overwritten by the checkpoint's state"
+                        ),
+                    )
+
             initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage)
 
             async for event in self._run_workflow_with_tracing(