From f9587ae2737b5cc1c15cb89f5104dfc52188a701 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 28 Jul 2026 14:53:48 -0700 Subject: [PATCH 1/4] Allow workflow checkpoint full replayability Seed the initial run input through the start executor's internal self-edge and record an entry checkpoint (iteration 0) before any executor runs, plus a response-entry checkpoint when responses are delivered, so a run is fully replayable from its checkpoints. Simplify the runner to only checkpoint after each superstep. Drop stale events in apply_checkpoint on restore, and deprecate the unused RunnerContext.reset_for_new_run. --- .../agent_framework_ag_ui/_workflow_run.py | 4 + .../packages/core/agent_framework/__init__.py | 6 +- .../core/agent_framework/__init__.pyi | 3 +- .../_workflows/_agent_executor.py | 4 +- .../agent_framework/_workflows/_checkpoint.py | 12 +- .../agent_framework/_workflows/_runner.py | 131 +++++----- .../_workflows/_runner_context.py | 24 +- .../agent_framework/_workflows/_workflow.py | 39 ++- .../core/test_function_invocation_logic.py | 1 + .../core/tests/workflow/test_checkpoint.py | 223 +++++++++++++++++ .../test_request_info_event_rehydrate.py | 58 +++++ .../core/tests/workflow/test_runner.py | 231 +++--------------- .../core/tests/workflow/test_workflow.py | 35 ++- 13 files changed, 481 insertions(+), 290 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 324933cbeb..3f6e76e5fd 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/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) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ae09f7610b..45ac531dd6 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/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", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index d4547535e8..4f8804a641 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/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", diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 22e872be07..b7787736fa 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/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 " diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 7473f403e6..2b267979e9 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/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 diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 3de7bb1614..ac5558dbe1 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/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 diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index 0917d23274..f70b9075f7 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/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 diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 4d2f471926..2feb582f61 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/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 @@ -654,16 +654,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 (iteration 0) 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 +1032,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. diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index ae01d82a22..b950aab0f0 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -2374,6 +2374,7 @@ def test_replace_approval_contents_with_results_allows_reused_call_id_after_comp ("call_reused", "second output"), ] + def test_replace_approval_contents_with_results_uses_result_call_ids_for_placeholders() -> None: from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index e7aa4f0b3a..346d041cb3 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -427,6 +427,229 @@ def _build_workflow() -> Any: assert checkpoints_by_id[chain[-1]].previous_checkpoint_id is None +async def test_workflow_entry_checkpoint_records_input_and_replays(): + """The entry checkpoint records the run's raw input, enabling a full replay from the start executor. + + The first checkpoint (iteration 0) is created before any executor runs and captures the input + message seeded onto the start executor's internal self-edge. Restoring it on a fresh instance + re-delivers that input to the start executor, so the entire run - including the start executor - + replays and produces the same output. + """ + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, handler + from agent_framework._workflows._const import INTERNAL_SOURCE_ID + from agent_framework._workflows._executor import Executor + + class StartExecutor(Executor): + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message + "-start", target_id="finish") + + class FinishExecutor(Executor): + @handler + async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.yield_output(message + "-done") + + storage = InMemoryCheckpointStorage() + + def _build_workflow() -> Any: + start = StartExecutor(id="start") + finish = FinishExecutor(id="finish") + return ( + WorkflowBuilder( + name="entry-replay-test", + max_iterations=10, + start_executor=start, + checkpoint_storage=storage, + ) + .add_edge(start, finish) + .build() + ) + + # First run. + workflow = _build_workflow() + workflow_name = workflow.name + first_outputs = (await workflow.run("hello")).get_outputs() + assert first_outputs == ["hello-start-done"] + + # The entry checkpoint (iteration 0) is created before any executor runs and begins a new lineage. + checkpoints = await storage.list_checkpoints(workflow_name=workflow_name) + entry = [cp for cp in checkpoints if cp.iteration_count == 0] + assert len(entry) == 1, "A fresh checkpointed run must create exactly one entry checkpoint at iteration 0" + entry_cp = entry[0] + assert entry_cp.previous_checkpoint_id is None + + # The raw input is recorded as an in-flight message on the start executor's internal self-edge. + seeded = entry_cp.messages.get(INTERNAL_SOURCE_ID("start")) + assert seeded is not None and len(seeded) == 1, "Entry checkpoint must record the seeded input message" + assert seeded[0].data == "hello" + + # Restore the entry checkpoint on a fresh instance: the start executor replays from the raw input + # and the whole run reproduces the original output. + replay_workflow = _build_workflow() + replay_outputs = (await replay_workflow.run(checkpoint_id=entry_cp.checkpoint_id)).get_outputs() + assert replay_outputs == first_outputs + + +async def test_workflow_one_checkpoint_per_superstep_lineage_and_replay(): + """A checkpointed run yields exactly one checkpoint per superstep plus the entry checkpoint, + forms a single unbroken lineage, and can be replayed from any pre-terminal checkpoint. + + - Count: checkpoints == (# supersteps) + 1 (the iteration-0 entry checkpoint), with contiguous + iteration counts ``0..N`` (no gaps or duplicates). + - Lineage: the entry checkpoint has no parent and every later checkpoint chains to its predecessor. + - Replay: restoring any non-terminal checkpoint on a fresh instance reproduces the final output. + """ + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, handler + from agent_framework._workflows._executor import Executor + + class StartExecutor(Executor): + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message + "-start", target_id="middle") + + class MiddleExecutor(Executor): + @handler + async def process(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message + "-middle", target_id="finish") + + class FinishExecutor(Executor): + @handler + async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.yield_output(message + "-done") + + storage = InMemoryCheckpointStorage() + + def _build_workflow() -> Any: + start = StartExecutor(id="start") + middle = MiddleExecutor(id="middle") + finish = FinishExecutor(id="finish") + return ( + WorkflowBuilder( + name="per-superstep-test", + max_iterations=10, + start_executor=start, + checkpoint_storage=storage, + ) + .add_edge(start, middle) + .add_edge(middle, finish) + .build() + ) + + # First run: count supersteps via the emitted control-plane events. + workflow = _build_workflow() + workflow_name = workflow.name + events = [event async for event in workflow.run("hello", stream=True)] + first_outputs = [event.data for event in events if event.type == "output"] + superstep_count = sum(1 for event in events if event.type == "superstep_completed") + assert first_outputs == ["hello-start-middle-done"] + assert superstep_count == 3, "start -> middle -> finish converges in three supersteps" + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow_name), + key=lambda c: c.iteration_count, + ) + + # Count: one checkpoint per superstep plus the entry checkpoint, with contiguous iteration counts. + assert len(checkpoints) == superstep_count + 1 + assert [cp.iteration_count for cp in checkpoints] == list(range(superstep_count + 1)) + + # Lineage: the entry checkpoint has no parent; every later checkpoint chains to its predecessor. + assert checkpoints[0].previous_checkpoint_id is None + for prev, cur in zip(checkpoints, checkpoints[1:]): + assert cur.previous_checkpoint_id == prev.checkpoint_id, ( + f"Checkpoint at iteration {cur.iteration_count} must chain to iteration {prev.iteration_count}" + ) + + # Replay: restoring any non-terminal checkpoint reproduces the final output. The terminal + # checkpoint (highest iteration) is the converged state with no work left, so it is excluded. + for cp in checkpoints[:-1]: + replay_workflow = _build_workflow() + replay_outputs = (await replay_workflow.run(checkpoint_id=cp.checkpoint_id)).get_outputs() + assert replay_outputs == first_outputs, ( + f"Replay from the checkpoint at iteration {cp.iteration_count} must reproduce the final output" + ) + + +async def test_workflow_response_entry_checkpoint_records_response_and_replays(): + """Delivering responses records a response-entry checkpoint that captures the responses in-flight, + making a human-in-the-loop continuation fully replayable. + + The response-entry checkpoint is created before the runner processes the responses. It sits at + the same iteration as the pending-request (IDLE) checkpoint but is a distinct checkpoint that + chains from it and carries the delivered response as an in-flight message (with no pending + requests). Resuming from it on a fresh instance re-delivers the response and reproduces the output. + """ + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowRunState, handler, response_handler + from agent_framework._workflows._executor import Executor + from agent_framework._workflows._request_info_mixin import RequestInfoMixin + + class HilExecutor(Executor, RequestInfoMixin): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, response_type=str, request_id="req1") + + @response_handler + async def on_response( + self, + original_request: str, + response: str, + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: + await ctx.yield_output(f"{original_request}->{response}") + + storage = InMemoryCheckpointStorage() + + def _build_workflow() -> Any: + hil = HilExecutor(id="hil") + return WorkflowBuilder( + name="response-entry-test", + max_iterations=10, + start_executor=hil, + checkpoint_storage=storage, + ).build() + + # First run: pauses at IDLE_WITH_PENDING_REQUESTS awaiting the response. + workflow = _build_workflow() + workflow_name = workflow.name + result = await workflow.run("approve") + request_events = result.get_request_info_events() + assert len(request_events) == 1 + request_id = request_events[0].request_id + assert result.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + + checkpoints_before = await storage.list_checkpoints(workflow_name=workflow_name) + before_ids = {cp.checkpoint_id for cp in checkpoints_before} + # The last checkpoint before responding sits at the pending-request boundary and records the request. + pending_cp = max(checkpoints_before, key=lambda c: c.iteration_count) + assert pending_cp.pending_request_info_events + + # Deliver the response; the workflow completes. + final = await workflow.run(responses={request_id: "yes"}) + assert final.get_outputs() == ["approve->yes"] + + checkpoints_after = await storage.list_checkpoints(workflow_name=workflow_name) + new_checkpoints = [cp for cp in checkpoints_after if cp.checkpoint_id not in before_ids] + + # The response-entry checkpoint sits at the SAME iteration as the pending-request checkpoint, + # chains from it, has no pending requests (they were answered), and records the response in-flight. + response_entry = next(cp for cp in new_checkpoints if cp.iteration_count == pending_cp.iteration_count) + assert response_entry.checkpoint_id != pending_cp.checkpoint_id + assert response_entry.previous_checkpoint_id == pending_cp.checkpoint_id + assert not response_entry.pending_request_info_events + assert response_entry.messages, "Response-entry checkpoint must record the delivered response in-flight" + + # Resuming from the response-entry checkpoint replays the continuation and reproduces the output. + replay_workflow = _build_workflow() + replay_outputs = (await replay_workflow.run(checkpoint_id=response_entry.checkpoint_id)).get_outputs() + assert replay_outputs == ["approve->yes"] + + async def test_memory_checkpoint_storage_roundtrip_json_native_types(): """Test that JSON-native types (str, int, float, bool, None) roundtrip correctly.""" storage = InMemoryCheckpointStorage() diff --git a/python/packages/core/tests/workflow/test_request_info_event_rehydrate.py b/python/packages/core/tests/workflow/test_request_info_event_rehydrate.py index d280a788cd..c05ab66984 100644 --- a/python/packages/core/tests/workflow/test_request_info_event_rehydrate.py +++ b/python/packages/core/tests/workflow/test_request_info_event_rehydrate.py @@ -385,3 +385,61 @@ async def test_checkpoint_restore_with_partial_responses_reemits_unhandled_reque assert final_status.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, ( f"Workflow should be IDLE_WITH_PENDING_REQUESTS, got {final_status.state}" ) + + +async def test_apply_checkpoint_drops_stale_events() -> None: + """apply_checkpoint should drop events left over from a prior run. + + After restore, only the pending request_info events carried by the checkpoint are + re-emitted onto the event queue; unrelated events queued earlier must be discarded so + they don't leak into the resumed run. + """ + # Build a checkpoint that carries a single pending request_info event. + source = InProcRunnerContext(InMemoryCheckpointStorage()) + keep_event = WorkflowEvent.request_info( + request_id="keep-me", + source_executor_id="gateway", + request_data=MockRequest(), + response_type=bool, + ) + await source.add_request_info_event(keep_event) + checkpoint_id = await source.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1) + checkpoint = await source.load_checkpoint(checkpoint_id) + assert checkpoint is not None + + # Target context has a stale event queued from a prior run. + target = InProcRunnerContext(InMemoryCheckpointStorage()) + stale_event = WorkflowEvent.request_info( + request_id="stale", + source_executor_id="old", + request_data=MockRequest(), + response_type=bool, + ) + await target.add_event(stale_event) + assert await target.has_events() + + await target.apply_checkpoint(checkpoint) + + events = await target.drain_events() + request_ids = [event.request_id for event in events] + assert request_ids == ["keep-me"] + assert "stale" not in request_ids + + +async def test_apply_checkpoint_preserves_streaming_flag() -> None: + """apply_checkpoint must not reset the streaming flag set before restore. + + The run pipeline calls set_streaming() before restoring, so applying a checkpoint must + leave the flag untouched (unlike the deprecated reset_for_new_run, which clears it). + """ + source = InProcRunnerContext(InMemoryCheckpointStorage()) + checkpoint_id = await source.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1) + checkpoint = await source.load_checkpoint(checkpoint_id) + assert checkpoint is not None + + target = InProcRunnerContext(InMemoryCheckpointStorage()) + target.set_streaming(True) + + await target.apply_checkpoint(checkpoint) + + assert target.is_streaming() is True diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 4839a8cd63..3690cdaa1f 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -557,7 +557,7 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: assert executor.count == 7 assert state.get("shared_key") == "shared_value" - assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage] + assert runner._previous_checkpoint_id == checkpoint.checkpoint_id # pyright: ignore[reportPrivateUsage] async def test_runner_build_checkpoint_includes_in_flight_messages(): @@ -731,7 +731,7 @@ async def test_runner_restore_from_checkpoint_with_external_storage(): # Restore using external storage await runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage=storage) - assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage] + assert runner._previous_checkpoint_id == checkpoint_id # pyright: ignore[reportPrivateUsage] assert runner._iteration == 5 # pyright: ignore[reportPrivateUsage] assert state.get("test_key") == "test_value" @@ -956,51 +956,6 @@ async def test_runner_restore_executor_states_no_states(): await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage] -async def test_runner_checkpoint_with_resumed_flag(): - """Test that resumed flag prevents initial checkpoint creation.""" - storage = InMemoryCheckpointStorage() - ctx = CheckpointingContext(storage) - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") - - edges = [ - SingleEdgeGroup(executor_a.id, executor_b.id), - SingleEdgeGroup(executor_b.id, executor_a.id), - ] - - executors: dict[str, Executor] = { - executor_a.id: executor_a, - executor_b.id: executor_b, - } - state = State() - - runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash") - resumed_checkpoint = WorkflowCheckpoint( - checkpoint_id="resumed-cp", - workflow_name="test_name", - graph_signature_hash="test_hash", - iteration_count=5, - ) - runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage] - - # Add a message to trigger the checkpoint creation path - await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START")) - - await executor_a.execute( - MockMessage(data=8), - ["START"], - state, - ctx, - ) - - # Run until convergence - async for _ in runner.run_until_convergence(): - pass - - # After completing, resumed flag should be reset - assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage] - - async def test_runner_mark_resumed_sets_previous_checkpoint_id(): """_mark_resumed must populate _previous_checkpoint_id so future checkpoints chain back to the resume point.""" runner = Runner( @@ -1023,7 +978,6 @@ async def test_runner_mark_resumed_sets_previous_checkpoint_id(): ) runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage] - assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage] assert runner._iteration == 3 # pyright: ignore[reportPrivateUsage] assert runner._previous_checkpoint_id == "resumed-cp-id" # pyright: ignore[reportPrivateUsage] @@ -1150,171 +1104,50 @@ async def test_runner_drains_events_on_iteration_exception(): assert len(output_events) >= 1 -async def test_runner_resumed_flag_reset_after_failed_resumed_run(): - """A failed *resumed* run must not leak the resume flag into the next run. - - The resume flag suppresses the initial "superstep 0" (entry) checkpoint when resuming from an - iteration-0 checkpoint (which already exists and must not be recreated). It used to be cleared - only on the success path, so an executor failure during a resumed run left it ``True`` and the - next fresh run wrongly skipped its entry checkpoint. The flag is now cleared in a ``finally`` so - this holds even when convergence raises. +async def test_runner_creates_exactly_one_checkpoint_per_superstep(): + """The runner creates exactly one checkpoint per completed superstep, with a consistent lineage. - This also verifies checkpoint creation on the re-run: the resumed (failed) run creates no entry - checkpoint, while the subsequent fresh run does. + Driving a loop that converges in a known number of supersteps must yield one checkpoint per + superstep - iteration counts ``1..N`` with no gaps or duplicates - and each checkpoint must + chain to its immediate predecessor, the first beginning a fresh lineage with no parent. The + runner never creates an entry (iteration-0) checkpoint; that is the ``Workflow``'s responsibility. """ storage = InMemoryCheckpointStorage() ctx = CheckpointingContext(storage) - executor_a = PassthroughExecutor(id="executor_a") - executor_b = ExecutorThatFailsWithEvents(id="executor_b", runner_ctx=ctx, fail_on_iteration=1) - - edges = [SingleEdgeGroup(executor_a.id, executor_b.id)] + executor_a = MockExecutor(id="executor_a") + executor_b = MockExecutor(id="executor_b") + edges = [ + SingleEdgeGroup(executor_a.id, executor_b.id), + SingleEdgeGroup(executor_b.id, executor_a.id), + ] executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b} state = State() - runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash") - # Simulate a resumed run; this marks the runner as resumed so the next run skips - # the superstep-0 checkpoint. - resumed_checkpoint = WorkflowCheckpoint( - checkpoint_id="resumed-cp", - workflow_name="test_name", - graph_signature_hash="test_hash", - iteration_count=0, - ) - runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage] - assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage] - - # Run the resumed turn; executor_b fails mid-iteration before any superstep - # checkpoint is created. + # executor_a seeds the loop; the runner then drives supersteps until the value reaches 10. await executor_a.execute(MockMessage(data=0), ["START"], state, ctx) - with pytest.raises(RuntimeError, match="Executor failed with pending events"): - async for _ in runner.run_until_convergence(): - pass - - # The fix: the resume flag is cleared even though the run raised. - assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage] - # The resumed (failed) run created no superstep-0 checkpoint (it was skipped). - assert await storage.list_checkpoints(workflow_name="test_name") == [] - - # Re-run as a fresh turn: with the flag correctly reset, the runner now creates - # the initial superstep-0 checkpoint (iteration_count == 0) before failing again. - runner.reset_iteration_count() - await executor_a.execute(MockMessage(data=0), ["START"], state, ctx) - with pytest.raises(RuntimeError, match="Executor failed with pending events"): - async for _ in runner.run_until_convergence(): - pass - - checkpoints = await storage.list_checkpoints(workflow_name="test_name") - assert any(cp.iteration_count == 0 for cp in checkpoints), ( - "Fresh run after a failed resumed run must create the superstep-0 checkpoint; " - "a leaked resume flag would have skipped it" - ) - - -async def test_runner_creates_entry_checkpoint_at_iteration_zero(): - """A fresh run creates the entry (superstep-0) checkpoint at iteration 0 with no parent. - - This is the baseline the lineage-consistency guard must preserve: when starting from iteration 0 - with messages queued and not resumed, the entry checkpoint is created and begins a new lineage - (``previous_checkpoint_id is None``). - """ - storage = InMemoryCheckpointStorage() - ctx = CheckpointingContext(storage) - # Terminal executor with no outgoing edges: the runner runs one superstep and converges. - source = MockExecutor(id="source") - state = State() - runner = Runner([], {source.id: source}, state, ctx, "test_name", graph_signature_hash="test_hash") - - assert runner._iteration == 0 # pyright: ignore[reportPrivateUsage] - assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage] - - await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id=source.id)) - - async for _ in runner.run_until_convergence(): - pass - - checkpoints = await storage.list_checkpoints(workflow_name="test_name") - entry_checkpoints = [cp for cp in checkpoints if cp.iteration_count == 0] - assert len(entry_checkpoints) == 1, "A fresh run must create exactly one entry checkpoint at iteration 0" - assert entry_checkpoints[0].previous_checkpoint_id is None, ( - "The entry checkpoint of a fresh run must begin a new lineage with no parent" - ) - - -async def test_runner_skips_entry_checkpoint_when_iteration_nonzero(): - """The entry (superstep-0) checkpoint must only be created at iteration 0 to keep lineage consistent. - - A re-run that did not reset the iteration count (and is not marked as resumed) must not write an - entry checkpoint carrying a non-zero ``iteration_count`` - doing so would place two checkpoints at - the same iteration in the lineage. The ``_iteration == 0`` guard suppresses the entry checkpoint in - this case while still allowing the normal per-superstep checkpoints to be created. - """ - storage = InMemoryCheckpointStorage() - ctx = CheckpointingContext(storage) - # Terminal executor with no outgoing edges: the runner runs one superstep and converges. - source = MockExecutor(id="source") - state = State() - runner = Runner([], {source.id: source}, state, ctx, "test_name", graph_signature_hash="test_hash") - - # Simulate a re-run that kept its iteration count and is not marked as resumed. - runner._iteration = 5 # pyright: ignore[reportPrivateUsage] - assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage] - - await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id=source.id)) - - async for _ in runner.run_until_convergence(): - pass - - checkpoints = await storage.list_checkpoints(workflow_name="test_name") - # No entry checkpoint at the pre-existing iteration count may be created. - assert all(cp.iteration_count != 5 for cp in checkpoints), ( - "Entry checkpoint must not be created at a non-zero iteration; lineage would have a duplicate iteration" - ) - # The normal post-superstep checkpoint is still created (iteration advanced to 6). - assert any(cp.iteration_count == 6 for cp in checkpoints) - - -async def test_runner_resumed_from_iteration_zero_skips_entry_checkpoint(): - """Resuming from an iteration-0 checkpoint must not recreate the entry checkpoint. - - Here ``_iteration == 0`` is true, so the iteration guard alone would not suppress the entry - checkpoint; the resume flag is what prevents recreating the checkpoint that already exists at - iteration 0. - """ - storage = InMemoryCheckpointStorage() - ctx = CheckpointingContext(storage) - source = MockExecutor(id="source") - state = State() - runner = Runner([], {source.id: source}, state, ctx, "test_name", graph_signature_hash="test_hash") - - # Resume from an iteration-0 checkpoint: iteration stays 0 but the run is marked as resumed. - resumed_checkpoint = WorkflowCheckpoint( - checkpoint_id="entry-cp", - workflow_name="test_name", - graph_signature_hash="test_hash", - iteration_count=0, - ) - runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage] - assert runner._iteration == 0 # pyright: ignore[reportPrivateUsage] - assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage] - - await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id=source.id)) - - async for _ in runner.run_until_convergence(): - pass + superstep_completed = 0 + async for event in runner.run_until_convergence(): + if event.type == "superstep_completed": + superstep_completed += 1 - # The pre-loop entry checkpoint is skipped; only the post-superstep checkpoint (iteration 1) is created, - # and it chains back to the resumed entry checkpoint. checkpoints = sorted( await storage.list_checkpoints(workflow_name="test_name"), - key=lambda c: c.timestamp, - ) - assert all(cp.checkpoint_id != "entry-cp" for cp in checkpoints), "Resumed entry checkpoint must not be recreated" - assert checkpoints, "The resumed run must still create its post-superstep checkpoint" - assert checkpoints[0].previous_checkpoint_id == "entry-cp", ( - "The first post-resume checkpoint must chain back to the resumed entry checkpoint" + key=lambda c: c.iteration_count, ) + # Exactly one checkpoint per completed superstep, with contiguous iteration counts 1..N and no entry checkpoint. + assert superstep_completed == runner._iteration # pyright: ignore[reportPrivateUsage] + assert len(checkpoints) == superstep_completed + assert [cp.iteration_count for cp in checkpoints] == list(range(1, superstep_completed + 1)) + + # Lineage: the first checkpoint begins a fresh chain; each subsequent one chains to its predecessor. + assert checkpoints[0].previous_checkpoint_id is None + for prev, cur in zip(checkpoints, checkpoints[1:]): + assert cur.previous_checkpoint_id == prev.checkpoint_id, ( + f"Checkpoint at iteration {cur.iteration_count} must chain to iteration {prev.iteration_count}" + ) + class SlowEventEmittingExecutor(Executor): """An executor that emits events with delays to test straggler event draining.""" diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 1d7087b6fe..2f672f591d 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -230,9 +230,9 @@ async def test_fan_out(): # and executor_completed (type='executor_completed') # executor_b will also emit an output event (type='output') # Each superstep will emit a started event (type='started') and status event (type='status') - # This workflow will converge in 2 supersteps because executor_c will send one more message - # after executor_b completes - assert len(events) == 11 + # Superstep 1 runs the start executor (executor_a) on the seeded input; the workflow then + # takes two more supersteps because executor_c sends one more message after executor_b completes. + assert len(events) == 13 assert events.get_final_state() == WorkflowRunState.IDLE outputs = events.get_outputs() @@ -255,8 +255,9 @@ async def test_fan_out_multiple_completed_events(): # and executor_completed (type='executor_completed') # executor_b and executor_c will also emit an output event (type='output') # Each superstep will emit a started event (type='started') and status event (type='status') - # This workflow will converge in 1 superstep because executor_a and executor_b will not send further messages - assert len(events) == 10 + # Superstep 1 runs the start executor (executor_a) on the seeded input; superstep 2 runs + # executor_b and executor_c, after which the workflow converges. + assert len(events) == 12 # Multiple outputs are expected from both executors outputs = events.get_outputs() @@ -283,7 +284,9 @@ async def test_fan_in(): # and executor_completed (type='executor_completed') # aggregator will also emit an output event (type='output') # Each superstep will emit a started event (type='started') and status event (type='status') - assert len(events) == 13 + # Superstep 1 runs the start executor (executor_a) on the seeded input, superstep 2 runs the + # fan-out targets, and superstep 3 runs the aggregator. + assert len(events) == 15 assert events.get_final_state() == WorkflowRunState.IDLE outputs = events.get_outputs() @@ -318,6 +321,26 @@ async def test_workflow_with_checkpointing_enabled(simple_executor: Executor): assert result is not None +async def test_run_with_unhandled_input_type_raises(simple_executor: Executor): + """Running with an input the start executor cannot handle must fail loudly, not silently drop it.""" + workflow = WorkflowBuilder(start_executor=simple_executor).build() + + # simple_executor only handles str; an int is not routable to it. + with pytest.raises(RuntimeError, match="cannot handle input of type 'int'"): + await workflow.run(42) + + +async def test_run_unwraps_workflow_message_input(simple_executor: Executor): + """A WorkflowMessage passed to run() is unwrapped (not double-wrapped) so its data reaches the handler.""" + # simple_executor handles str; passing WorkflowMessage(data=) must not be dropped as a type + # mismatch. Without unwrapping, the start executor would receive a WorkflowMessage (not a str) and + # the input would be silently dropped by the internal edge runner. + workflow = WorkflowBuilder(start_executor=simple_executor).build() + + result = await workflow.run(WorkflowMessage(data="wrapped", source_id="test", target_id=None)) + assert result.get_final_state() == WorkflowRunState.IDLE + + async def test_workflow_checkpointing_not_enabled_for_external_restore( simple_executor: Executor, ): From 624456d3c777cea44d17895f17afb33c19449c8e Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 28 Jul 2026 15:13:57 -0700 Subject: [PATCH 2/4] Fix type --- python/packages/core/agent_framework/_workflows/_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 2feb582f61..8ff1718169 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -655,7 +655,7 @@ async def _execute_with_message_or_checkpoint( # Handle initial message elif message is not 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.. + # 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( From 85fff4f4b0b2d87d2c2ae379b6bd3dbfd7dafcdc Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 28 Jul 2026 21:02:42 -0700 Subject: [PATCH 3/4] Add max iteration detailed doc string --- python/packages/core/agent_framework/_workflows/_workflow.py | 2 ++ .../core/agent_framework/_workflows/_workflow_builder.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 8ff1718169..3dfeeadfd5 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -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 diff --git a/python/packages/core/agent_framework/_workflows/_workflow_builder.py b/python/packages/core/agent_framework/_workflows/_workflow_builder.py index e092710913..de389b3eb5 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_builder.py +++ b/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 From c408c5bdfe9a1d12d684d0adba1db017dc10db94 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 29 Jul 2026 10:26:52 -0700 Subject: [PATCH 4/4] Refine comments --- python/packages/core/agent_framework/_workflows/_workflow.py | 2 +- python/packages/core/tests/workflow/test_checkpoint.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 3dfeeadfd5..77060b93a9 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -675,7 +675,7 @@ async def _execute_with_message_or_checkpoint( ) await self._runner.context.send_message(entry_message) - # Record the entry checkpoint (iteration 0) capturing the seeded input before any + # 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() diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 346d041cb3..5f3da78cd1 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -430,7 +430,7 @@ def _build_workflow() -> Any: async def test_workflow_entry_checkpoint_records_input_and_replays(): """The entry checkpoint records the run's raw input, enabling a full replay from the start executor. - The first checkpoint (iteration 0) is created before any executor runs and captures the input + The first checkpoint is created before any executor runs and captures the input message seeded onto the start executor's internal self-edge. Restoring it on a fresh instance re-delivers that input to the start executor, so the entire run - including the start executor - replays and produces the same output. @@ -473,7 +473,7 @@ def _build_workflow() -> Any: first_outputs = (await workflow.run("hello")).get_outputs() assert first_outputs == ["hello-start-done"] - # The entry checkpoint (iteration 0) is created before any executor runs and begins a new lineage. + # The entry checkpoint is created before any executor runs. checkpoints = await storage.list_checkpoints(workflow_name=workflow_name) entry = [cp for cp in checkpoints if cp.iteration_count == 0] assert len(entry) == 1, "A fresh checkpointed run must create exactly one entry checkpoint at iteration 0"