From adf8c987f3b9944f9553925f605f104edbc96898 Mon Sep 17 00:00:00 2001 From: Shivani Bhandari Date: Thu, 27 Aug 2026 16:42:51 +0530 Subject: [PATCH] Python: Allow checkpoint hydrate and new input in one workflow.run Restore and validate the checkpoint before seeding start-executor input so hosts can resume multi-turn workflows without a separate hydrate round trip (#7863). Keep the prior two-step path working. --- .../agent_framework_ag_ui/_workflow_run.py | 14 ++- .../core/agent_framework/_workflows/_agent.py | 84 ++++++++++---- .../agent_framework/_workflows/_functional.py | 8 +- .../agent_framework/_workflows/_workflow.py | 52 +++++---- .../core/tests/workflow/test_checkpoint.py | 86 ++++++++++++++ .../workflow/test_functional_workflow.py | 27 +++-- .../core/tests/workflow/test_workflow.py | 13 +-- .../checkpoint_hydrate_with_input.py | 105 ++++++++++++++++++ 8 files changed, 314 insertions(+), 75 deletions(-) create mode 100644 python/samples/03-workflows/checkpoint/checkpoint_hydrate_with_input.py 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 35b142df5f..a82f075386 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 @@ -1116,11 +1116,15 @@ def _drain_open_message() -> list[TextMessageEndEvent]: telemetry_conversation_id = str(supplied_thread_id) if supplied_thread_id is not None else None telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id) with telemetry_context(): - if responses or checkpoint_id is not None: - # ``message`` is mutually exclusive with both ``responses`` and - # ``checkpoint_id`` in the core API; ``responses`` + ``checkpoint_id`` - # restores the checkpoint and delivers the responses in a single call. - event_stream = workflow.run(stream=True, responses=responses or None, **checkpoint_kwargs, **fwd_kwargs) + if responses: + # HITL: restore (optional) and deliver responses in one call. + event_stream = workflow.run(stream=True, responses=responses, **checkpoint_kwargs, **fwd_kwargs) + elif checkpoint_id is not None: + # Pure checkpoint restore. Incoming chat messages are not start-executor + # input here; AG-UI maps HITL replies through ``responses`` above. + # Hosts that need hydrate+new-user-input can pass both via core + # ``Workflow.run(message=..., checkpoint_id=...)`` (#7863). + event_stream = workflow.run(stream=True, **checkpoint_kwargs, **fwd_kwargs) else: event_stream = workflow.run(message=messages, stream=True, **checkpoint_kwargs, **fwd_kwargs) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index f0b51969c0..57a8f5a3d2 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -407,27 +407,73 @@ async def _run_core( Yields: WorkflowEvent objects from the workflow execution. """ - # Restore the workflow state if a checkpoint is provided + run_kwargs = { + "checkpoint_storage": checkpoint_storage, + "function_invocation_kwargs": function_invocation_kwargs, + "client_kwargs": client_kwargs, + } + + # Checkpoint + input: restore and continue in one workflow.run (#7863). if checkpoint_id is not None: if checkpoint_storage is None: raise AgentInvalidRequestException("checkpoint_storage must be provided when checkpoint_id is provided") logger.debug(f"Restoring workflow from checkpoint {checkpoint_id}") - # Restore the workflow from checkpoint + + if not input_messages: + # Hydrate only; no new turn. + if streaming: + async for _ in self.workflow.run( + stream=True, + checkpoint_id=checkpoint_id, + checkpoint_storage=checkpoint_storage, + ): + pass + else: + _ = await self.workflow.run( + checkpoint_id=checkpoint_id, + checkpoint_storage=checkpoint_storage, + ) + logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") + return + + checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint.pending_request_info_events: + # HITL continuation: restore and deliver responses together. + function_responses = self._extract_function_responses(input_messages) + if streaming: + async for event in self.workflow.run( + responses=function_responses, + stream=True, + checkpoint_id=checkpoint_id, + **run_kwargs, + ): + yield event + else: + for event in await self.workflow.run( + responses=function_responses, + checkpoint_id=checkpoint_id, + **run_kwargs, + ): + yield event + return + + # Idle checkpoint: restore and seed the new user turn together. if streaming: - async for _ in self.workflow.run( + async for event in self.workflow.run( + message=input_messages, stream=True, checkpoint_id=checkpoint_id, - checkpoint_storage=checkpoint_storage, + **run_kwargs, ): - pass + yield event else: - _ = await self.workflow.run( + for event in await self.workflow.run( + message=input_messages, checkpoint_id=checkpoint_id, - checkpoint_storage=checkpoint_storage, - ) - if not input_messages: - logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") - return + **run_kwargs, + ): + yield event + return final_state = self._workflow.status logger.debug(f"Workflow state: {final_state}") @@ -444,17 +490,13 @@ async def _run_core( async for event in self.workflow.run( responses=function_responses, stream=True, - checkpoint_storage=checkpoint_storage, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, + **run_kwargs, ): yield event else: for event in await self.workflow.run( responses=function_responses, - checkpoint_storage=checkpoint_storage, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, + **run_kwargs, ): yield event elif final_state == WorkflowRunState.IDLE: @@ -462,17 +504,13 @@ async def _run_core( async for event in self.workflow.run( message=input_messages, stream=True, - checkpoint_storage=checkpoint_storage, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, + **run_kwargs, ): yield event else: for event in await self.workflow.run( message=input_messages, - checkpoint_storage=checkpoint_storage, - function_invocation_kwargs=function_invocation_kwargs, - client_kwargs=client_kwargs, + **run_kwargs, ): yield event else: diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 9e4deb340b..a02f546f4b 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -824,8 +824,9 @@ def run( HITL interruption; *checkpoint_id* restores from a previously saved checkpoint. *responses* may be combined with *checkpoint_id* to restore a checkpoint and inject HITL responses in a single call. - *message* is mutually exclusive with both *responses* and - *checkpoint_id*. + *message* may be combined with *checkpoint_id* to restore a checkpoint + and apply new start input in a single call (#7863). *message* remains + mutually exclusive with *responses*. Args: message: Input data passed as the first positional argument to @@ -1282,9 +1283,6 @@ def _validate_run_params( if message is not None and responses is not None: raise ValueError("Cannot provide both 'message' and 'responses'. Use one or the other.") - if message is not None and checkpoint_id is not None: - raise ValueError("Cannot provide both 'message' and 'checkpoint_id'. Use one or the other.") - if message is None and responses is None and checkpoint_id is None: raise ValueError( "Must provide at least one of: 'message' (new run), 'responses' (send responses), " diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 77060b93a9..bbb4281ebd 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -628,11 +628,17 @@ async def _execute_with_message_or_checkpoint( checkpoint_id: str | None, checkpoint_storage: CheckpointStorage | None, ) -> None: - """Internal handler for executing workflow with either initial message or checkpoint restoration. + """Restore from a checkpoint and/or seed a start-executor message. + + Checkpoint restoration always runs first when ``checkpoint_id`` is set so + validation/hydration failures never process ``message``. When both are + provided, the restored state is hydrated and then ``message`` is applied + in the same run (#7863). Args: - message: Initial message for the start executor (for new runs). - checkpoint_id: ID of checkpoint to restore from (for resuming runs). + message: Initial message for the start executor (new turn, optionally + after restore). + checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. Raises: @@ -642,7 +648,7 @@ async def _execute_with_message_or_checkpoint( if message is None and checkpoint_id is None: raise ValueError("Must provide either 'message' or 'checkpoint_id'") - # Handle checkpoint restoration + # Handle checkpoint restoration first so failures never execute new input. if checkpoint_id is not None: has_checkpointing = self._runner.context.has_checkpointing() @@ -654,8 +660,8 @@ async def _execute_with_message_or_checkpoint( await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) - # Handle initial message - elif message is not None: + # Seed start-executor input for a new turn (alone or after restore). + if 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. start_id = self.start_executor_id @@ -725,7 +731,8 @@ def run( Args: message: Initial message for the start executor. Required for new workflow runs. - Mutually exclusive with responses. + Mutually exclusive with responses. Can be combined with checkpoint_id to + restore a checkpoint and apply new input in a single call (#7863). stream: If True, returns a ResponseStream of events with ``get_final_response()`` for the final WorkflowRunResult. If False (default), returns an awaitable WorkflowRunResult. @@ -734,8 +741,8 @@ def run( exclusive with message. Can be combined with checkpoint_id to restore a checkpoint and send responses in a single call. checkpoint_id: ID of checkpoint to restore from. Can be used alone (resume - from checkpoint), with message (not allowed), or with responses - (restore then send responses). + from checkpoint), with message (restore then seed start-executor input), + or with responses (restore then send responses). checkpoint_storage: Runtime checkpoint storage. include_status_events: Whether to include status events (non-streaming only). function_invocation_kwargs: Keyword arguments forwarded to tool invocations in @@ -825,17 +832,17 @@ async def _run_core( self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) try: - # 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 - - # the caller must either resume from a checkpoint or wait for the - # prior run to drain. Pending request_info events are intentionally - # NOT blocked here (they are answered via a follow-up ``responses=...`` - # run); the warning below surfaces the abandon/overwrite cases instead. - # NOTE: _validate_run_params already enforces that ``message`` is - # mutually exclusive with both ``checkpoint_id`` and ``responses``, - # so we don't need to re-check those here. - if message is not None and await self._runner.context.has_messages(): + # Async validation: a fresh-message run (without checkpoint restore) 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 - the caller must either resume from a checkpoint (alone or + # combined with message) or wait for the prior run to drain. When + # ``checkpoint_id`` is also provided, restore replaces the context before + # the new message is seeded, so leftover in-flight messages are allowed. + # Pending request_info events are intentionally NOT blocked here (they + # are answered via a follow-up ``responses=...`` run); the warning below + # surfaces the abandon/overwrite cases instead. + if message is not None and checkpoint_id is None and await self._runner.context.has_messages(): raise RuntimeError( "Cannot start a new run with 'message' while in-flight executor " "messages remain from a prior run. Resume from a checkpoint " @@ -947,16 +954,13 @@ def _validate_run_params( Rules: - message and responses are mutually exclusive - - message and checkpoint_id are mutually exclusive - At least one of message, responses, or checkpoint_id must be provided + - message + checkpoint_id is allowed (restore then seed start-executor input) - responses + checkpoint_id is allowed (restore then send) """ if message is not None and responses is not None: raise ValueError("Cannot provide both 'message' and 'responses'. Use one or the other.") - if message is not None and checkpoint_id is not None: - raise ValueError("Cannot provide both 'message' and 'checkpoint_id'. Use one or the other.") - if message is None and responses is None and checkpoint_id is None: raise ValueError( "Must provide at least one of: 'message' (new run), 'responses' (send responses), " diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5f3da78cd1..2796c44fc7 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -336,6 +336,92 @@ async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: ) +async def test_workflow_run_restores_checkpoint_and_applies_message_in_one_call(): + """Hydrate from a completed-turn checkpoint and seed new start input atomically (#7863).""" + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, handler + from agent_framework._workflows._executor import Executor + + class AccumulatorExecutor(Executor): + def __init__(self, id: str) -> None: + super().__init__(id=id) + self.history: list[str] = [] + + @handler + async def accumulate(self, message: str, ctx: WorkflowContext[Never, list[str]]) -> None: # type: ignore[valid-type] + self.history.append(message) + await ctx.yield_output(list(self.history)) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"history": list(self.history)} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.history = list(state.get("history", [])) + + storage = InMemoryCheckpointStorage() + + def _build() -> tuple[Any, AccumulatorExecutor]: + start = AccumulatorExecutor(id="start") + workflow = WorkflowBuilder( + name="hydrate-and-input", + max_iterations=5, + start_executor=start, + checkpoint_storage=storage, + ).build() + return workflow, start + + workflow, _ = _build() + first = await workflow.run("turn-1") + assert first.get_outputs() == [["turn-1"]] + latest = await storage.get_latest(workflow_name=workflow.name) + assert latest is not None + + # Fresh instance: restore + new input in a single run (streaming and non-streaming). + resumed, resumed_start = _build() + second = await resumed.run("turn-2", checkpoint_id=latest.checkpoint_id) + assert second.get_outputs() == [["turn-1", "turn-2"]] + assert resumed_start.history == ["turn-1", "turn-2"] + + after_second = await storage.get_latest(workflow_name=workflow.name) + assert after_second is not None + streamed, streamed_start = _build() + events = [event async for event in streamed.run("turn-3", checkpoint_id=after_second.checkpoint_id, stream=True)] + assert any(e.type == "output" and e.data == ["turn-1", "turn-2", "turn-3"] for e in events) + assert streamed_start.history == ["turn-1", "turn-2", "turn-3"] + + +async def test_workflow_run_invalid_checkpoint_does_not_apply_message(): + """Checkpoint validation failures must occur before the new message is seeded.""" + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, handler + from agent_framework._workflows._executor import Executor + + class CountingExecutor(Executor): + def __init__(self, id: str) -> None: + super().__init__(id=id) + self.calls = 0 + + @handler + async def run(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + self.calls += 1 + await ctx.yield_output(message) + + storage = InMemoryCheckpointStorage() + executor = CountingExecutor(id="start") + workflow = WorkflowBuilder( + name="hydrate-fail", + start_executor=executor, + checkpoint_storage=storage, + ).build() + + with pytest.raises(WorkflowCheckpointException): + await workflow.run("should-not-run", checkpoint_id="missing-checkpoint-id") + + assert executor.calls == 0 + + async def test_workflow_checkpoint_ancestry_preserved_after_resume(): """Resuming from a checkpoint must preserve ancestry: future checkpoints chain back to the resumed one.""" from typing_extensions import Never diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 96b7e6edf6..1b5650c9cd 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -467,14 +467,6 @@ async def wf(x: int) -> None: with pytest.raises(ValueError, match="Cannot provide both"): await wf.run("hello", responses={"r1": "val"}) - async def test_invalid_params_message_and_checkpoint(self): - @built_workflow - async def wf(x: int) -> None: - pass - - with pytest.raises(ValueError, match="Cannot provide both"): - await wf.run("hello", checkpoint_id="abc") - async def test_invalid_params_nothing(self): @built_workflow async def wf(x: int) -> None: @@ -670,6 +662,25 @@ async def hitl_wf(doc: str, ctx: RunContext) -> str: result2 = await hitl_wf.run(checkpoint_id=ckpt_id, responses={"req1": "Approved!"}) assert result2.get_outputs() == ["Done: Approved!"] + async def test_checkpoint_and_message_in_one_call(self): + """Restore from a completed run and apply a new message without a second hydrate call.""" + storage = InMemoryCheckpointStorage() + + @built_workflow(checkpoint_storage=storage) + async def echo_wf(text: str, ctx: RunContext) -> str: + history = ctx.get_state("history") or [] + history = [*history, text] + ctx.set_state("history", history) + return "|".join(history) + + first = await echo_wf.run("one") + assert first.get_outputs() == ["one"] + latest = await storage.get_latest(workflow_name="echo_wf") + assert latest is not None + + second = await echo_wf.run("two", checkpoint_id=latest.checkpoint_id) + assert second.get_outputs() == ["one|two"] + async def test_checkpoint_without_storage_raises(self): @built_workflow async def wf(x: int) -> int: diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 2f672f591d..439bd8b11f 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1387,16 +1387,9 @@ async def test_workflow_run_parameter_validation(simple_executor: Executor) -> N result = await workflow.run(test_message) assert result.get_final_state() == WorkflowRunState.IDLE - # Invalid: message + checkpoint_id (mutually exclusive). Multi-turn - # state preservation is handled by Workflow.run preserving state across - # calls, so the host pattern is two separate calls (restore-then-run), - # not a single combined call. - with pytest.raises(ValueError, match="Cannot provide both 'message' and 'checkpoint_id'"): - await workflow.run(test_message, checkpoint_id="some-checkpoint") - - with pytest.raises(ValueError, match="Cannot provide both 'message' and 'checkpoint_id'"): - async for _ in workflow.run(test_message, checkpoint_id="some-checkpoint", stream=True): - pass + # Invalid: message + responses remain mutually exclusive + with pytest.raises(ValueError, match="Cannot provide both 'message' and 'responses'"): + await workflow.run(test_message, responses={"req": True}) # Invalid: none of message or checkpoint_id with pytest.raises(ValueError, match="Must provide at least one of"): diff --git a/python/samples/03-workflows/checkpoint/checkpoint_hydrate_with_input.py b/python/samples/03-workflows/checkpoint/checkpoint_hydrate_with_input.py new file mode 100644 index 0000000000..5d0cf07d3b --- /dev/null +++ b/python/samples/03-workflows/checkpoint/checkpoint_hydrate_with_input.py @@ -0,0 +1,105 @@ +# Copyright (c) Microsoft. All rights reserved. + +""" +Sample: Checkpoint hydration and new input in one workflow.run + +Purpose: +Show how to restore a persisted checkpoint and apply a new user turn in a +single ``workflow.run`` call instead of the older two-step hydrate-then-run +pattern (#7863). + +What you learn: +- How ``message`` and ``checkpoint_id`` can be combined on ``Workflow.run`` +- That checkpoint restore runs before the new message is seeded +- That streaming and non-streaming share the same combined API +- That the prior two-step pattern remains valid for hosts that prefer it + +Pipeline: +1) A stateful start executor accumulates each turn and yields the history +2) Turn 1 runs to completion and persists checkpoints +3) A fresh workflow instance resumes with ``run(message=..., checkpoint_id=...)`` + +Prerequisites: +- Basic understanding of workflow executors, edges, and checkpoint storage +""" + +import asyncio +import sys +from typing import Any + +from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + WorkflowBuilder, + WorkflowContext, + handler, +) + +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + + +class HistoryExecutor(Executor): + """Accumulates each turn's text and yields the full history.""" + + def __init__(self, id: str) -> None: + super().__init__(id=id) + self._history: list[str] = [] + + @handler + async def accumulate(self, text: str, ctx: WorkflowContext[Any, list[str]]) -> None: + self._history.append(text) + print(f"HistoryExecutor: recorded {text!r}; history={self._history}") + await ctx.yield_output(list(self._history)) + + @override + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"history": list(self._history)} + + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self._history = list(state.get("history", [])) + + +def _build(storage: InMemoryCheckpointStorage): + start = HistoryExecutor(id="history") + return WorkflowBuilder( + name="hydrate-with-input-sample", + start_executor=start, + checkpoint_storage=storage, + ).build() + + +async def main() -> None: + storage = InMemoryCheckpointStorage() + + # Turn 1 — fresh run + workflow = _build(storage) + result1 = await workflow.run("hello") + print(f"Turn 1 outputs: {result1.get_outputs()}") + + latest = await storage.get_latest(workflow_name=workflow.name) + if latest is None: + raise RuntimeError("Expected a checkpoint after turn 1") + + # Turn 2 — restore + new input in one call (non-streaming) + resumed = _build(storage) + result2 = await resumed.run("again", checkpoint_id=latest.checkpoint_id) + print(f"Turn 2 outputs: {result2.get_outputs()}") + + latest = await storage.get_latest(workflow_name=workflow.name) + if latest is None: + raise RuntimeError("Expected a checkpoint after turn 2") + + # Turn 3 — same combined API with streaming + streamed = _build(storage) + print("Turn 3 (streaming):") + async for event in streamed.run("once more", checkpoint_id=latest.checkpoint_id, stream=True): + if event.type == "output": + print(f" output event: {event.data}") + + +if __name__ == "__main__": + asyncio.run(main())