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 324933cbebe..3f6e76e5fd3 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/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 22e872be07c..67388497302 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/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() diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 7473f403e65..2b267979e99 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 3de7bb16142..8e3d00636b4 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. @@ -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. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 4d2f4719260..574303bda61 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import copy import functools import hashlib import json @@ -20,8 +21,8 @@ from .._types import ResponseStream 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 ._checkpoint import CheckpointStorage, WorkflowCheckpoint +from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY from ._edge import ( EdgeGroup, FanOutEdgeGroup, @@ -35,7 +36,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 @@ -370,6 +371,11 @@ def __init__( # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # Deep-copied, in-memory snapshot of the workflow's pristine state, captured + # lazily before the first run's executors run. Held so the instance can be + # reset back to this state via ``reset()`` and reused for another run. + self._initial_checkpoint: WorkflowCheckpoint | None = None + @property def status(self) -> WorkflowRunState: """Return the current run-level status of this workflow instance. @@ -654,15 +660,18 @@ 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. + start_id = self.start_executor_id + await self._runner.context.send_message( + WorkflowMessage( + data=message, + source_id=INTERNAL_SOURCE_ID(start_id), + target_id=start_id, + ) ) + # 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( @@ -810,6 +819,10 @@ async def _run_core( self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) try: + # Capture the pristine initial checkpoint once, before any executor runs, + # so this instance can later be reset back to this state via ``reset()``. + await self._capture_initial_checkpoint_if_needed() + # Async validation: a fresh-message run is only allowed when the # runner context has fully drained from any prior run. If it still # has in-flight executor messages, the prior run didn't complete - @@ -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. @@ -1220,3 +1239,50 @@ def _is_run_active(self) -> bool: """ existing_stream = self._active_run() if self._active_run is not None else None return existing_stream is not None + + async def _capture_initial_checkpoint_if_needed(self) -> None: + """Capture a pristine in-memory checkpoint on the first run. + + The snapshot is taken before any executor runs, so it represents the + workflow's initial state. It is deep-copied so later state mutations + cannot corrupt the stored snapshot, which keeps ``reset()`` repeatable. + """ + if self._initial_checkpoint is not None: + return + checkpoint = await self._runner.build_checkpoint() + self._initial_checkpoint = copy.deepcopy(checkpoint) + + async def reset(self) -> None: + """Reset this workflow instance back to its initial state for reuse. + + Restores the shared state, in-flight messages, pending request info events, + executor checkpoint state, and iteration counter to the snapshot captured + before this instance's first run. This lets a single ``Workflow`` instance be + reused for multiple independent runs instead of building a new one each time. + + Only state captured by the checkpointing mechanism is rewound: the shared + workflow state and each executor's ``on_checkpoint_save()`` state. Executor + instance attributes that are not surfaced through + ``on_checkpoint_save()``/``on_checkpoint_restore()`` are not rewound. + + Calling ``reset()`` before the instance has ever run is a no-op, since the + instance is already in its initial state. + + Raises: + WorkflowException: If a run is currently active on this instance. + """ + if self._is_run_active(): + raise WorkflowException("Cannot reset the workflow while a run is active on the same instance.") + if self._initial_checkpoint is None: + # Never run; the instance is already in its initial state. + logger.debug("Reset called on a workflow instance that has never run; no state to restore.") + return + # Deep-copy on restore too, so the applied state never aliases the stored + # snapshot and subsequent resets remain repeatable. Start a new lineage: the + # pristine snapshot is an in-memory checkpoint that was never persisted, so + # checkpoints created by the next run must not chain to its (storage-absent) id. + await self._runner.restore_checkpoint( + copy.deepcopy(self._initial_checkpoint), + start_new_lineage=True, + ) + self._status = WorkflowRunState.IDLE 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 ae01d82a226..b950aab0f06 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 e7aa4f0b3a4..346d041cb31 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_runner.py b/python/packages/core/tests/workflow/test_runner.py index 4839a8cd631..3690cdaa1ff 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 1d7087b6fed..441efb609dd 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() diff --git a/python/packages/core/tests/workflow/test_workflow_reset.py b/python/packages/core/tests/workflow/test_workflow_reset.py new file mode 100644 index 00000000000..aa25f883535 --- /dev/null +++ b/python/packages/core/tests/workflow/test_workflow_reset.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import Any + +import pytest +from typing_extensions import Never + +from agent_framework import ( + Executor, + Workflow, + WorkflowBuilder, + WorkflowContext, + handler, +) +from agent_framework.exceptions import WorkflowException + + +class CounterExecutor(Executor): + """Executor that accumulates state in both shared state and an instance attribute. + + Shared state is rewound by Workflow.reset() automatically. The instance attribute is + rewound only because this executor surfaces it via the checkpoint hooks. + """ + + def __init__(self, id: str = "counter") -> None: + super().__init__(id) + self.internal_count = 0 + + @handler + async def run(self, msg: int, ctx: WorkflowContext[Never, int]) -> None: + shared = (ctx.get_state("count") or 0) + 1 + self.internal_count += 1 + ctx.set_state("count", shared) + await ctx.yield_output(shared) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"internal_count": self.internal_count} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.internal_count = state.get("internal_count", 0) + + +async def test_state_persists_across_runs_without_reset() -> None: + """State on a reused instance persists across runs until reset() is called.""" + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter).build() + + r1 = await wf.run(0) + assert r1.get_outputs() == [1] + + r2 = await wf.run(0) + assert r2.get_outputs() == [2] + assert counter.internal_count == 2 + + +async def test_reset_rewinds_shared_and_executor_state() -> None: + """reset() rewinds shared state and executor checkpoint state to the initial snapshot.""" + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter).build() + + await wf.run(0) + await wf.run(0) + assert counter.internal_count == 2 + + await wf.reset() + assert counter.internal_count == 0 + + r = await wf.run(0) + assert r.get_outputs() == [1] + + +async def test_reset_is_repeatable() -> None: + """The stored snapshot is isolated, so repeated reset+run cycles are deterministic.""" + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter).build() + + for _ in range(3): + await wf.reset() + r = await wf.run(0) + assert r.get_outputs() == [1] + assert counter.internal_count == 1 + + +async def test_reset_before_first_run_is_noop() -> None: + """reset() before the instance has ever run is a no-op and does not raise.""" + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter).build() + + await wf.reset() + + r = await wf.run(0) + assert r.get_outputs() == [1] + + +async def test_reset_while_run_active_raises() -> None: + """reset() is rejected while a run is active on the same instance.""" + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter).build() + + stream = wf.run(0, stream=True) + try: + with pytest.raises(WorkflowException, match="while a run is active"): + await wf.reset() + finally: + # Drain the stream so the run finalizes cleanly. + async for _ in stream: + pass + + +async def test_reset_starts_a_fresh_checkpoint_lineage() -> None: + """After reset(), the next run's checkpoints begin a fresh lineage with no dangling parent. + + The pristine snapshot restored by reset() is an in-memory checkpoint that was never + persisted, so post-reset checkpoints must not chain to its (storage-absent) id. + """ + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + counter = CounterExecutor() + wf: Workflow = WorkflowBuilder(start_executor=counter, checkpoint_storage=storage).build() + + await wf.run(0) + first_run_ids = {cp.checkpoint_id for cp in await storage.list_checkpoints(workflow_name=wf.name)} + + await wf.reset() + await wf.run(0) + + all_checkpoints = await storage.list_checkpoints(workflow_name=wf.name) + new_checkpoints = [cp for cp in all_checkpoints if cp.checkpoint_id not in first_run_ids] + assert new_checkpoints, "The post-reset run must create new checkpoints" + + # The post-reset entry checkpoint begins a fresh lineage (no parent) ... + entry = min(new_checkpoints, key=lambda cp: cp.iteration_count) + assert entry.previous_checkpoint_id is None, ( + "The post-reset entry checkpoint must start a new lineage, not chain to the in-memory reset snapshot" + ) + + # ... and no post-reset checkpoint references a parent that is absent from storage. + all_ids = {cp.checkpoint_id for cp in all_checkpoints} + for cp in new_checkpoints: + assert cp.previous_checkpoint_id is None or cp.previous_checkpoint_id in all_ids, ( + "Post-reset checkpoint chains to a parent that does not exist in storage (dangling lineage)" + )