[BREAKING] Python: Allow workflow checkpoint full replayability

python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py

@@ -922,6 +922,10 @@ def _drain_open_message() -> list[TextMessageEndEvent]:
                 else:
                     state_value = str(getattr(state, "value", state))
                 if state_value in _TERMINAL_STATES and not terminal_emitted:
+                    # Close any open assistant text message before the terminal event so
+                    # RUN_FINISHED is always the last emitted event.
+                    for end_event in _drain_open_message():
+                        yield end_event
                     if not interrupts:
                         interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow)))
                     yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)

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

@@ -2,6 +2,7 @@
 
 import logging
 import sys
+import warnings
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, Literal, cast
@@ -14,7 +15,7 @@
 from .._sessions import AgentSession
 from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream
 from ._agent_utils import resolve_agent_id
-from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
+from ._const import GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY
 from ._executor import Executor, handler
 from ._message_utils import normalize_messages_input
 from ._request_info_mixin import response_handler
@@ -251,7 +252,7 @@ async def from_str(
             executor, use ``AgentExecutorResponse.with_text(...)`` so that the message type
             stays ``AgentExecutorResponse`` and ``from_response`` is called instead.
         """
-        if not self._cache and ctx.source_executor_ids != ["Workflow"]:
+        if not self._cache and ctx.source_executor_ids != [INTERNAL_SOURCE_ID(self.id)]:
             logger.warning(
                 "AgentExecutor '%s': from_str handler invoked with an empty cache. "
                 "If you are chaining from an AgentExecutor, the upstream custom executor may be "
@@ -366,7 +367,17 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
         self._pending_responses_to_agent = pending_responses_payload or []
 
     def reset(self) -> None:
-        """Reset the internal cache of the executor."""
+        """Reset the internal cache of the executor.
+
+        .. deprecated::
+            ``AgentExecutor.reset`` is deprecated and will be removed in a future
+            version. It is unused within the framework and has no replacement.
+        """
+        warnings.warn(
+            "`AgentExecutor.reset` is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         logger.debug("AgentExecutor %s: Resetting cache", self.id)
         self._cache.clear()
 

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

@@ -59,7 +59,17 @@ class WorkflowCheckpoint:
         pending_request_info_events: Any pending request info events that have not
             yet been processed at the time of checkpointing. This allows the workflow
             to resume with the correct pending events after a restore.
-        iteration_count: Current iteration number when checkpoint was created
+        iteration_count: Current iteration number when checkpoint was created.
+            Note: iteration_count is not guaranteed to be unique across a workflow's
+            lifecycle. It marks the superstep boundary the checkpoint sits on, and the
+            same boundary can carry more than one checkpoint. For example, a run that
+            pauses at ``IDLE_WITH_PENDING_REQUESTS`` records a checkpoint after superstep
+            K; when responses are later delivered, a response-entry checkpoint is recorded
+            at the same iteration K (responses in-flight, before superstep K+1 runs).
+            Both share iteration K but are distinct checkpoints. Checkpoint ordering is
+            defined by the ``previous_checkpoint_id`` lineage chain (and ``timestamp``),
+            not by ``iteration_count``; do not use ``iteration_count`` to identify the
+            latest checkpoint in human-in-the-loop flows.
         metadata: Additional metadata (e.g., superstep info, graph signature)
         version: Checkpoint format version
 

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

@@ -80,7 +80,6 @@ def __init__(
         self._state = state
 
         # Checkpointing related attributes
-        self._resumed_from_checkpoint = False
         self._previous_checkpoint_id: CheckpointID | None = None
 
     @property
@@ -105,83 +104,77 @@ def reset_iteration_count(self) -> None:
         self._iteration = 0
 
     async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
-        """Run the workflow until no more messages are sent."""
-        try:
-            # Emit any events already produced prior to entering loop
-            if await self._ctx.has_events():
-                logger.info("Yielding pre-loop events")
-                for event in await self._ctx.drain_events():
-                    yield event
+        """Run the workflow until no more messages are sent.
 
-            # Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
-            # end of an iteration, we can think of this checkpoint as being created at the end of "superstep 0"
-            # which captures the states after which the start executor has run. Note that we execute the start
-            # executor outside of the main iteration loop.
-            if await self._ctx.has_messages() and self._iteration == 0 and not self._resumed_from_checkpoint:
-                await self.create_checkpoint_if_enabled()
-
-            while self._iteration < self._max_iterations:
-                logger.info(f"Starting superstep {self._iteration + 1}")
-                yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
-
-                # Run iteration concurrently with live event streaming: we poll
-                # for new events while the iteration coroutine progresses.
-                iteration_task = asyncio.create_task(self._run_iteration())
-                try:
-                    while not iteration_task.done():
-                        try:
-                            # Wait briefly for any new event; timeout allows progress checks
-                            event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
-                            yield event
-                        except asyncio.TimeoutError:
-                            # Periodically continue to let iteration advance
-                            continue
-                except asyncio.CancelledError:
-                    # Propagate cancellation to the iteration task to avoid orphaned work
-                    iteration_task.cancel()
-                    with contextlib.suppress(asyncio.CancelledError):
-                        await iteration_task
-                    raise
-
-                # Propagate errors from iteration, but first surface any pending events
-                try:
+        The runner's checkpoint responsibility is intentionally narrow: it creates a
+        checkpoint at the end of each superstep. The entry checkpoint that captures the
+        run's initial input (before any executor runs) is created by ``Workflow`` prior
+        to entering this loop, since only the workflow knows how the run was started
+        (fresh input vs. checkpoint resume vs. responses).
+        """
+        # Emit any events already produced prior to entering loop
+        if await self._ctx.has_events():
+            logger.info("Yielding pre-loop events")
+            for event in await self._ctx.drain_events():
+                yield event
+
+        while self._iteration < self._max_iterations:
+            logger.info(f"Starting superstep {self._iteration + 1}")
+            yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
+
+            # Run iteration concurrently with live event streaming: we poll
+            # for new events while the iteration coroutine progresses.
+            iteration_task = asyncio.create_task(self._run_iteration())
+            try:
+                while not iteration_task.done():
+                    try:
+                        # Wait briefly for any new event; timeout allows progress checks
+                        event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
+                        yield event
+                    except asyncio.TimeoutError:
+                        # Periodically continue to let iteration advance
+                        continue
+            except asyncio.CancelledError:
+                # Propagate cancellation to the iteration task to avoid orphaned work
+                iteration_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
                     await iteration_task
-                except Exception:
-                    # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
-                    if await self._ctx.has_events():
-                        for event in await self._ctx.drain_events():
-                            yield event
-                    raise
-                self._iteration += 1
-
-                # Drain any straggler events emitted at tail end
+                raise
+
+            # Propagate errors from iteration, but first surface any pending events
+            try:
+                await iteration_task
+            except Exception:
+                # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
                 if await self._ctx.has_events():
                     for event in await self._ctx.drain_events():
                         yield event
+                raise
+            self._iteration += 1
+
+            # Drain any straggler events emitted at tail end
+            if await self._ctx.has_events():
+                for event in await self._ctx.drain_events():
+                    yield event
 
-                logger.info(f"Completed superstep {self._iteration}")
+            logger.info(f"Completed superstep {self._iteration}")
 
-                # Commit pending state changes at superstep boundary
-                self._state.commit()
+            # Commit pending state changes at superstep boundary
+            self._state.commit()
 
-                # Create checkpoint after each superstep iteration
-                await self.create_checkpoint_if_enabled()
+            # Create checkpoint after each superstep iteration
+            await self.create_checkpoint_if_enabled()
 
-                yield WorkflowEvent.superstep_completed(iteration=self._iteration)
+            yield WorkflowEvent.superstep_completed(iteration=self._iteration)
 
-                # Check for convergence: no more messages to process
-                if not await self._ctx.has_messages():
-                    break
+            # Check for convergence: no more messages to process
+            if not await self._ctx.has_messages():
+                break
 
-            logger.info(f"Workflow completed after {self._iteration} supersteps")
+        logger.info(f"Workflow completed after {self._iteration} supersteps")
 
-            if self._iteration >= self._max_iterations and await self._ctx.has_messages():
-                raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
-        finally:
-            # Reset the resume flag so stale resume state never leaks into the next run on this
-            # instance - even if convergence raised before completing (e.g. an executor failure
-            # during a resumed run).
-            self._resumed_from_checkpoint = False
+        if self._iteration >= self._max_iterations and await self._ctx.has_messages():
+            raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
 
     async def _run_iteration(self) -> None:
         """Run a single iteration of the workflow.
@@ -360,7 +353,7 @@ async def build_checkpoint(self) -> WorkflowCheckpoint:
             self._iteration,
         )
 
-    async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+    async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint, *, start_new_lineage: bool = False) -> None:
         """Restore runner state from an in-memory ``WorkflowCheckpoint`` object.
 
         Unlike :meth:`restore_from_checkpoint`, this does not load from a storage
@@ -372,6 +365,9 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
 
         Args:
             checkpoint: The checkpoint whose state should be restored.
+            start_new_lineage: When True, clear the previous-checkpoint pointer after restoring
+                so checkpoints created by the next run begin a fresh lineage with no parent. Use
+                this when the restored checkpoint is not part of the persisted lineage.
 
         Raises:
             WorkflowCheckpointException: If the checkpoint's graph signature does not
@@ -390,7 +386,7 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
             self._state.import_state(checkpoint.state)
             await self._restore_executor_states()
             await self._ctx.apply_checkpoint(checkpoint)
-            self._mark_resumed(checkpoint)
+            self._mark_resumed(checkpoint, start_new_lineage=start_new_lineage)
         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
@@ -452,14 +448,20 @@ def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[
 
         return parsed
 
-    def _mark_resumed(self, checkpoint: WorkflowCheckpoint) -> None:
-        """Mark the runner as having resumed from a checkpoint.
+    def _mark_resumed(self, checkpoint: WorkflowCheckpoint, *, start_new_lineage: bool = False) -> None:
+        """Restore per-run checkpoint bookkeeping from a resumed checkpoint.
+
+        Sets the iteration counter and previous-checkpoint pointer so that
+        checkpoints created after the resume continue the restored lineage.
 
-        Optionally set the current iteration and max iterations.
+        When ``start_new_lineage`` is True, the previous-checkpoint pointer is cleared
+        instead, so checkpoints created by the next run begin a fresh lineage with no
+        parent. This is used when restoring a checkpoint that is not part of the
+        persisted lineage (e.g. an in-memory reset snapshot) to avoid chaining later
+        persisted checkpoints to a parent id that does not exist in storage.
         """
-        self._resumed_from_checkpoint = True
         self._iteration = checkpoint.iteration_count
-        self._previous_checkpoint_id = checkpoint.checkpoint_id
+        self._previous_checkpoint_id = None if start_new_lineage else checkpoint.checkpoint_id
 
     async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
         """Store executor state in state under a reserved key.