diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 33d16cd9bb..1c1a78a9f0 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -242,6 +242,7 @@ from ._workflows._agent_executor import ( ) from ._workflows._agent_utils import resolve_agent_id from ._workflows._checkpoint import ( + CheckpointID, CheckpointStorage, FileCheckpointStorage, InMemoryCheckpointStorage, @@ -285,7 +286,6 @@ from ._workflows._functional import ( workflow, ) from ._workflows._request_info_mixin import response_handler -from ._workflows._runner import Runner from ._workflows._runner_context import ( InProcRunnerContext, RunnerContext, @@ -379,6 +379,7 @@ __all__ = [ "ChatResponse", "ChatResponseUpdate", "CheckResult", + "CheckpointID", "CheckpointStorage", "ClassSkill", "CompactionProvider", @@ -470,7 +471,6 @@ __all__ = [ "RoleLiteral", "RubricScore", "RunContext", - "Runner", "RunnerContext", "SecretString", "SelectiveToolCallCompactionStrategy", diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 124cc2fb57..28f1bec7d5 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -65,7 +65,7 @@ class Runner: # Checkpointing related attributes self._resumed_from_checkpoint = False - self._previous_checkpoint_id: CheckpointID | None = None + self.previous_checkpoint_id: CheckpointID | None = None @property def context(self) -> RunnerContext: @@ -113,12 +113,12 @@ class Runner: for event in await self._ctx.drain_events(): yield event - # Create the first checkpoint. Checkpoints are usually considered to be created at the end of an iteration, - # we can think of the first checkpoint as being created at the end of a "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. + # 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 a "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 not self._resumed_from_checkpoint: - await self._create_checkpoint_if_enabled() + await self.create_checkpoint_if_enabled() while self._iteration < self._max_iterations: logger.info(f"Starting superstep {self._iteration + 1}") @@ -165,7 +165,7 @@ class Runner: self._state.commit() # Create checkpoint after each superstep iteration - await self._create_checkpoint_if_enabled() + await self.create_checkpoint_if_enabled() yield WorkflowEvent.superstep_completed(iteration=self._iteration) @@ -231,7 +231,7 @@ class Runner: ] await asyncio.gather(*tasks) - async def _create_checkpoint_if_enabled(self) -> None: + async def create_checkpoint_if_enabled(self) -> None: """Create a checkpoint if checkpointing is enabled and attach a label and metadata.""" if not self._ctx.has_checkpointing(): return @@ -249,7 +249,7 @@ class Runner: self._workflow_name, self._graph_signature_hash, self._state, - self._previous_checkpoint_id, + self.previous_checkpoint_id, self._iteration, ) @@ -257,9 +257,9 @@ class Runner: "Created checkpoint: %s with parent checkpoint at iteration %d: %s", checkpoint_id, self._iteration, - self._previous_checkpoint_id, + self.previous_checkpoint_id, ) - self._previous_checkpoint_id = checkpoint_id + self.previous_checkpoint_id = checkpoint_id except Exception as e: logger.warning( "Failed to create checkpoint at iteration %d: %s. " @@ -267,7 +267,7 @@ class Runner: "The next successfully-created checkpoint will be parented to the last successful checkpoint: %s", self._iteration, e, - self._previous_checkpoint_id, + self.previous_checkpoint_id, ) async def restore_from_checkpoint( @@ -396,7 +396,7 @@ class Runner: """ self._resumed_from_checkpoint = True self._iteration = checkpoint.iteration_count - self._previous_checkpoint_id = checkpoint.checkpoint_id + self.previous_checkpoint_id = 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 6e240a1096..5c9a78dd75 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -18,9 +18,9 @@ from typing import TYPE_CHECKING, Any, Literal, overload from .._sessions import ContextProvider from .._types import ResponseStream -from ..exceptions import WorkflowException, WorkflowRunnerException +from ..exceptions import WorkflowCheckpointException, WorkflowException, WorkflowRunnerException from ..observability import OtelAttr, capture_exception, create_workflow_span -from ._checkpoint import CheckpointStorage +from ._checkpoint import CheckpointID, CheckpointStorage from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY from ._edge import ( EdgeGroup, @@ -474,6 +474,43 @@ class Workflow(DictConvertible): """Get the list of executors in the workflow.""" return list(self.executors.values()) + async def create_checkpoint(self, checkpoint_storage: CheckpointStorage | None) -> CheckpointID: + """Create a checkpoint of the current workflow state in the provided storage. + + Args: + checkpoint_storage: The CheckpointStorage instance where the checkpoint will be stored. + If None, will use the workflow's default checkpoint storage if configured, or raise + if checkpointing is not enabled. + + Notes: + - Checkpoints can only be created when the workflow is idle (not actively running). + - Checkpoints are automatically created at the end of each superstep if a checkpoint storage is configured. + Use this method only when necessary, for example to capture the initial state of the workflow prior to the + first run. + - Creating a checkpoint manually will alter the checkpoint lineage. The new checkpoint will become the + parent of the next checkpoint created automatically (if checkpointing is enabled by providing a storage). + """ + if self._is_run_active(): + raise WorkflowException( + "Cannot create checkpoint while a workflow run is active. " + "Checkpointing is only allowed between runs when the workflow is idle." + ) + + if checkpoint_storage is None and not self._runner.context.has_checkpointing(): + raise WorkflowCheckpointException( + "Checkpoint storage must be provided to create a checkpoint when checkpointing is not enabled." + ) + if checkpoint_storage is not None: + self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + + try: + await self._runner.create_checkpoint_if_enabled() + if self._runner.previous_checkpoint_id is None: + raise WorkflowCheckpointException("Failed to create checkpoint.") + return self._runner.previous_checkpoint_id + finally: + self._runner.context.clear_runtime_checkpoint_storage() + async def _run_workflow_with_tracing( self, initial_executor_fn: Callable[[], Awaitable[None]] | None = None, @@ -1165,6 +1202,15 @@ class Workflow(DictConvertible): **kwargs, ) + def _is_run_active(self) -> bool: + """Check if a workflow run is currently active. + + Returns: + True if a run is active, False otherwise. + """ + existing_stream = self._active_run() if self._active_run is not None else None + return existing_stream is not None + async def reset_for_new_run(self) -> None: """Reset the workflow for a new run that is independent from prior runs. diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 6199bcd8ec..f6b36ebe32 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -921,7 +921,7 @@ async def test_runner_mark_resumed_sets_previous_checkpoint_id(): ) # Pre-condition: nothing to chain back to - assert runner._previous_checkpoint_id is None # pyright: ignore[reportPrivateUsage] + assert runner.previous_checkpoint_id is None resumed_checkpoint = WorkflowCheckpoint( checkpoint_id="resumed-cp-id", @@ -933,7 +933,7 @@ async def test_runner_mark_resumed_sets_previous_checkpoint_id(): 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] + assert runner.previous_checkpoint_id == "resumed-cp-id" async def test_runner_post_resume_checkpoint_chains_to_resumed_checkpoint(): diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 498879861b..1a5b0e32d0 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -17,6 +17,7 @@ from typing import Protocol, cast from agent_framework import ( ChatOptions, + CheckpointID, Content, ContextProvider, FileCheckpointStorage, @@ -343,6 +344,7 @@ class ResponsesHostServer(ResponsesAgentServerHost): # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally CHECKPOINT_STORAGE_PATH = "/.checkpoints" + INITIAL_CHECKPOINT_STORAGE_NAME = "initial" FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json" def __init__( @@ -386,7 +388,6 @@ class ResponsesHostServer(ResponsesAgentServerHost): ) self._is_workflow_agent = False - self._checkpoint_storage_path = None if isinstance(agent, WorkflowAgent): if agent.workflow._runner_context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] raise RuntimeError( @@ -399,6 +400,12 @@ class ResponsesHostServer(ResponsesAgentServerHost): else os.path.join(os.getcwd(), self.CHECKPOINT_STORAGE_PATH.lstrip("/")) ) self._is_workflow_agent = True + # The initial checkpoint storage that stores the workflow's initial state. We will use this checkpoint + # to restore the workflow when no conversation_id or previous_response_id is supplied in a request. + self._initial_checkpoint_storage = _checkpoint_storage_for_context( + self._checkpoint_storage_path, self.INITIAL_CHECKPOINT_STORAGE_NAME + ) + self._initial_checkpoint_id: CheckpointID | None = None self._agent = agent self._approval_storage = ( @@ -580,8 +587,6 @@ class ResponsesHostServer(ResponsesAgentServerHost): # The following should never happen due to the checks above. # This is for type safety and defensive programming. - if self._checkpoint_storage_path is None: - raise RuntimeError("Checkpoint storage path is not configured for workflow agent.") if not isinstance(self._agent, WorkflowAgent): raise RuntimeError("Agent is not a workflow agent.") @@ -590,6 +595,12 @@ class ResponsesHostServer(ResponsesAgentServerHost): # any future async resources owned by the workflow are entered here. await self._ensure_agent_ready() + # Create a checkpoint to store the initial state of the workflow, if it doesn't already exist. + # This allows us to restore to a clean slate when no conversation_id or previous_response_id + # is supplied in a request. + if self._initial_checkpoint_id is None: + self._initial_checkpoint_id = await self._agent.workflow.create_checkpoint(self._initial_checkpoint_storage) + # Determine the latest checkpoint (if any) so we can resume the # workflow's prior state for this turn. The directory is keyed by # the inbound context id (conversation_id when set, otherwise @@ -599,15 +610,36 @@ class ResponsesHostServer(ResponsesAgentServerHost): # the only place that state lives is the workflow checkpoint, so # on every turn we restore the latest checkpoint and feed the new # input back into the start executor as a continuation rather than - # a fresh run. - latest_checkpoint_id: str | None = None - restore_storage: FileCheckpointStorage | None = None + # a fresh run. If no conversation_id or previous_response_id is + # supplied, the workflow will be restored to the initial checkpoint + # to avoid context bleed between requests. + latest_checkpoint_id: str = self._initial_checkpoint_id + restore_storage: FileCheckpointStorage = self._initial_checkpoint_storage if context_id is not None: restore_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id) latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) if latest_checkpoint is not None: latest_checkpoint_id = latest_checkpoint.checkpoint_id + # Restore the workflow to the latest checkpoint and run it with the + # new input. Events (including request info events) will not be emitted + # during restoration (in streaming) or after restoration (in non-streaming) + # since we assume the client had already seen those events and we don't want + # to emit duplicates. + if is_streaming_request: + async for _ in self._agent.run( + stream=True, + checkpoint_id=latest_checkpoint_id, + checkpoint_storage=restore_storage, + ): + pass + else: + await self._agent.run( + stream=False, + checkpoint_id=latest_checkpoint_id, + checkpoint_storage=restore_storage, + ) + # Storage that will receive checkpoints written during this turn. # When the caller chains with previous_response_id, the next turn # will reference the current response_id as its previous_response_id, @@ -619,43 +651,6 @@ class ResponsesHostServer(ResponsesAgentServerHost): write_context_id = context.conversation_id or context.response_id write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id) - # Multi-turn pattern: when we have a prior checkpoint, restore it - # first (drive the workflow back to idle with prior state intact), - # then make a separate call that delivers the new user input. This - # depends on Workflow.run preserving shared state across calls. The - # restore-only call may yield events from any pending in-flight - # work in the checkpoint; we consume those internally here so they - # don't surface to the response stream as duplicates. - # - # If the restored checkpoint had pending request_info events, the - # restore-only call replays them through - # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates`` - # and populates ``self._agent.pending_requests``. That is the correct - # state: those requests are genuinely outstanding, and the next - # ``run(input_messages, ...)`` call may contain ``function_call_output`` - # items (carried as FunctionResult/FunctionApprovalResponse content) - # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`. - if latest_checkpoint_id is not None: - if is_streaming_request: - async for _ in self._agent.run( - stream=True, - checkpoint_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, - ): - pass - else: - await self._agent.run( - stream=False, - checkpoint_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, - ) - else: - # We reset the workflow if neither conversation_id nor previous_response_id - # was supplied, because this implies there's no prior state to restore and - # we want to ensure a clean slate. Workflow may contain in-memory state that - # needs to be cleared on new conversations. - await self._agent.workflow.reset_for_new_run() - # Now run the agent with the latest input response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model)