[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/__init__.py

@@ -289,7 +289,10 @@
         "InMemoryCheckpointStorage",
         "WorkflowCheckpoint",
     ),
-    "._workflows._const": ("DEFAULT_MAX_ITERATIONS",),
+    "._workflows._const": (
+        "DEFAULT_MAX_ITERATIONS",
+        "INTERNAL_SOURCE_ID",
+    ),
     "._workflows._edge": (
         "Case",
         "Default",
@@ -369,6 +372,7 @@
     "GROUP_INDEX_KEY",
     "GROUP_KIND_KEY",
     "GROUP_TOKEN_COUNT_KEY",
+    "INTERNAL_SOURCE_ID",
     "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY",
     "SKIP_PARSING",
     "SUMMARIZED_BY_SUMMARY_ID_KEY",

python/packages/core/agent_framework/__init__.pyi

@@ -255,7 +255,7 @@ from ._workflows._checkpoint import (
     InMemoryCheckpointStorage,
     WorkflowCheckpoint,
 )
-from ._workflows._const import DEFAULT_MAX_ITERATIONS
+from ._workflows._const import DEFAULT_MAX_ITERATIONS, INTERNAL_SOURCE_ID
 from ._workflows._edge import (
     Case,
     Default,
@@ -336,6 +336,7 @@ __all__ = [
     "GROUP_INDEX_KEY",
     "GROUP_KIND_KEY",
     "GROUP_TOKEN_COUNT_KEY",
+    "INTERNAL_SOURCE_ID",
     "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY",
     "SKIP_PARSING",
     "SUMMARIZED_BY_SUMMARY_ID_KEY",

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

@@ -14,7 +14,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 +251,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 "

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.
@@ -453,11 +446,11 @@ 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.
+        """Restore per-run checkpoint bookkeeping from a resumed checkpoint.
 
-        Optionally set the current iteration and max iterations.
+        Sets the iteration counter and previous-checkpoint pointer so that
+        checkpoints created after the resume continue the restored lineage.
         """
-        self._resumed_from_checkpoint = True
         self._iteration = checkpoint.iteration_count
         self._previous_checkpoint_id = checkpoint.checkpoint_id
 

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

@@ -4,6 +4,7 @@
 
 import asyncio
 import logging
+import warnings
 from collections.abc import Callable
 from copy import copy
 from dataclasses import dataclass
@@ -176,7 +177,12 @@ def clear_runtime_checkpoint_storage(self) -> None:
         ...
 
     def reset_for_new_run(self) -> None:
-        """Reset the context for a new workflow run."""
+        """Reset the context for a new workflow run.
+
+        .. deprecated::
+            ``reset_for_new_run`` is deprecated and will be removed in a future version.
+            ``apply_checkpoint`` should reset the context prior to applying a checkpoint.
+        """
         ...
 
     def set_streaming(self, streaming: bool) -> None:
@@ -469,7 +475,19 @@ def reset_for_new_run(self) -> None:
 
         This clears messages, events, and resets streaming flag.
         Runtime checkpoint storage is NOT cleared here as it's managed at the workflow level.
+
+        .. deprecated::
+            ``reset_for_new_run`` is deprecated and will be removed in a future version.
+            ``apply_checkpoint`` should reset the context prior to applying a checkpoint.
         """
+        warnings.warn(
+            (
+                "`reset_for_new_run` is deprecated and will be removed in a future version. "
+                "`apply_checkpoint` should reset the context prior to applying a checkpoint."
+            ),
+            DeprecationWarning,
+            stacklevel=2,
+        )
         self._messages.clear()
         # Drop any pending events. The queue and its loop marker are cleared so the queue
         # rebinds lazily under the running loop on next use.
@@ -479,6 +497,10 @@ def reset_for_new_run(self) -> None:
 
     async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
         """Apply a checkpoint to the current context, mutating its state."""
+        # Drop any events left over from a prior run so the restored state starts clean.
+        self._event_queue = None
+        self._event_queue_loop = None
+
         # Restore messages
         self._messages.clear()
         messages_data = checkpoint.messages

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

@@ -21,7 +21,7 @@
 from ..exceptions import WorkflowException
 from ..observability import OtelAttr, capture_exception, create_workflow_span
 from ._checkpoint import CheckpointStorage
-from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
+from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY
 from ._edge import (
     EdgeGroup,
     FanOutEdgeGroup,
@@ -35,7 +35,7 @@
 from ._executor import Executor
 from ._model_utils import DictConvertible
 from ._runner import RunnerImpl
-from ._runner_context import RunnerContext
+from ._runner_context import RunnerContext, WorkflowMessage
 from ._state import State
 from ._typing_utils import is_instance_of, try_coerce_to_type
 from ._validation import ValidationTypeEnum, WorkflowValidationError
@@ -299,6 +299,8 @@ def __init__(
             start_executor: The starting executor for the workflow.
             runner_context: The RunnerContext instance to be used during workflow execution.
             max_iterations: The maximum number of iterations the workflow will run for convergence.
+                The first iteration is the initial run of the start executor, and each subsequent
+                iteration is a superstep.
             name: A human-readable name for the workflow. This can be used to identify the workflow in
                 checkpoints, and telemetry. If the workflow is built using WorkflowBuilder, this will be the
                 name of the builder. This name should be unique across different workflow definitions for
@@ -654,16 +656,29 @@ async def _execute_with_message_or_checkpoint(
 
         # Handle initial message
         elif message is not None:
-            executor = self.get_start_executor()
-            await executor.execute(
-                message,
-                [self.__class__.__name__],
-                self._runner.state,
-                self._runner.context,
-                trace_contexts=None,
-                source_span_ids=None,
+            # Seed the initial input through the start executor's internal self-edge. If the caller
+            # passed a WorkflowMessage, unwrap it so we don't double-wrap.
+            start_id = self.start_executor_id
+            data = message.data if isinstance(message, WorkflowMessage) else message
+            entry_message = WorkflowMessage(
+                data=data,
+                source_id=INTERNAL_SOURCE_ID(start_id),
+                target_id=start_id,
             )
 
+            # Ensure that the start executor can handle the input type.
+            start_executor = self.get_start_executor()
+            if not start_executor.can_handle(entry_message):
+                raise RuntimeError(
+                    f"Start executor '{start_id}' cannot handle input of type '{type(data).__name__}'. "
+                    f"Expected input types: {start_executor.input_types}."
+                )
+
+            await self._runner.context.send_message(entry_message)
+            # Record the entry checkpoint capturing the seeded input before any
+            # executor runs. The runner only checkpoints after each superstep.
+            await self._runner.create_checkpoint_if_enabled()
+
     @overload
     def run(
         self,
@@ -1019,6 +1034,12 @@ async def _send_responses_internal(self, responses: Mapping[str, Any]) -> None:
             for request_id, response in coerced_responses.items()
         ])
 
+        # Record a response-entry checkpoint capturing the delivered responses in-flight, before the
+        # runner processes them in the next superstep. This mirrors the initial-message entry
+        # checkpoint so a human-in-the-loop continuation is fully replayable: resuming from this
+        # checkpoint re-delivers the responses and replays the superstep that consumes them.
+        await self._runner.create_checkpoint_if_enabled()
+
     def _get_executor_by_id(self, executor_id: str) -> Executor:
         """Get an executor by its ID.
 

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

@@ -100,7 +100,9 @@ def __init__(
         """Initialize the WorkflowBuilder.
 
         Args:
-            max_iterations: Maximum number of iterations for workflow convergence. Default is 100.
+            max_iterations: Maximum number of iterations for workflow convergence. The first
+                iteration is the initial run of the start executor, and each subsequent iteration
+                is a superstep. Default is 100.
             name: A human-readable name for the workflow builder. This name will be the identifier
                 for all workflow instances created from this builder. If not provided, a unique name
                 will be generated. This will be useful for versioning, monitoring, checkpointing, and