Python: Support checkpoint hydration and new input in a single workflow run (#7863)

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)
 

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,35 +490,27 @@ 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:
             if streaming:
                 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:

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), "

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), "

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