From 9af62f44b339917d7c7a40ded8531df2f5101245 Mon Sep 17 00:00:00 2001 From: Abhinav Maurya Date: Fri, 30 Oct 2026 04:03:33 +0530 Subject: [PATCH] Python: Apply caller-supplied run kwargs over checkpoint kwargs on restore --- .../agent_framework/_workflows/_workflow.py | 75 ++++++++++++------- .../tests/workflow/test_workflow_kwargs.py | 67 +++++++++++++++++ 2 files changed, 115 insertions(+), 27 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c60617dabd5..b964c3f4f07 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -743,28 +743,32 @@ async def _execute_with_message_or_checkpoint( # Handle initial message elif message is not None: - # Seed the initial input through the start executor's internal self-edge. If the caller - # passed a WorkflowMessage, unwrap it so we don't double-wrap. - start_id = self.start_executor_id - data = message.data if isinstance(message, WorkflowMessage) else message - entry_message = WorkflowMessage( - data=data, - source_id=INTERNAL_SOURCE_ID(start_id), - target_id=start_id, - ) + await self._execute_with_message(message) + + async def _execute_with_message(self, message: Any) -> None: + """Seed the initial input message to the start executor.""" + # Seed the initial input through the start executor's internal self-edge. If the caller + # passed a WorkflowMessage, unwrap it so we don't double-wrap. + start_id = self.start_executor_id + data = message.data if isinstance(message, WorkflowMessage) else message + entry_message = WorkflowMessage( + data=data, + source_id=INTERNAL_SOURCE_ID(start_id), + target_id=start_id, + ) - # Ensure that the start executor can handle the input type. - start_executor = self.get_start_executor() - if not start_executor.can_handle(entry_message): - raise RuntimeError( - f"Start executor '{start_id}' cannot handle input of type '{type(data).__name__}'. " - f"Expected input types: {start_executor.input_types}." - ) + # Ensure that the start executor can handle the input type. + start_executor = self.get_start_executor() + if not start_executor.can_handle(entry_message): + raise RuntimeError( + f"Start executor '{start_id}' cannot handle input of type '{type(data).__name__}'. " + f"Expected input types: {start_executor.input_types}." + ) - await self._runner.context.send_message(entry_message) - # Record the entry checkpoint capturing the seeded input before any - # executor runs. The runner only checkpoints after each superstep. - await self._runner.create_checkpoint_if_enabled() + await self._runner.context.send_message(entry_message) + # Record the entry checkpoint capturing the seeded input before any + # executor runs. The runner only checkpoints after each superstep. + await self._runner.create_checkpoint_if_enabled() @overload def run( @@ -985,10 +989,24 @@ async def _run_core( ), ) - if responses is not None and checkpoint_id is None: + if message is None and checkpoint_id is None and responses is None: + raise ValueError("Must provide either 'message' or 'checkpoint_id'") + + if checkpoint_id is not None: + has_checkpointing = self._runner.context.has_checkpointing() + + if not has_checkpointing and checkpoint_storage is None: + raise ValueError( + "Cannot restore from checkpoint: either provide checkpoint_storage parameter " + "or build workflow with WorkflowBuilder(checkpoint_storage=checkpoint_storage)." + ) + + await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) + + if responses is not None: await self._validate_responses_internal(responses) - initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage) + initial_executor_fn = self._resolve_execution_mode(message, responses) async with aclosing( self._run_workflow_with_tracing( @@ -1089,9 +1107,9 @@ def _resolve_execution_mode( self, message: Any | None, responses: Mapping[str, Any] | None, - checkpoint_id: str | None, - checkpoint_storage: CheckpointStorage | None, - ) -> Callable[[], Awaitable[None]]: + checkpoint_id: str | None = None, + checkpoint_storage: CheckpointStorage | None = None, + ) -> Callable[[], Awaitable[None]] | None: """Determine the initial executor function based on parameters.""" if responses is not None: if checkpoint_id is not None: @@ -1103,8 +1121,11 @@ def _resolve_execution_mode( # Send responses only (requires pending requests in workflow state) initial_executor_fn = functools.partial(self._send_responses_internal, responses) return initial_executor_fn - # Regular run or checkpoint restoration - return functools.partial(self._execute_with_message_or_checkpoint, message, checkpoint_id, checkpoint_storage) + if message is not None: + return functools.partial(self._execute_with_message, message) + if checkpoint_id is not None: + return functools.partial(self._execute_with_message_or_checkpoint, None, checkpoint_id, checkpoint_storage) + return None async def _restore_and_send_responses( self, diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index 31bd7800271..0696f780177 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -2186,4 +2186,71 @@ async def test_empty_client_kwargs_clears_previous() -> None: assert agent.captured_kwargs[-1].get("client_kwargs") == {} +@pytest.mark.parametrize("kwargs_channel", ["function_invocation_kwargs", "client_kwargs"]) +async def test_caller_kwargs_override_checkpoint_kwargs_on_restore(tmp_path: Any, kwargs_channel: str) -> None: + """Caller-provided kwargs override checkpoint-restored kwargs on continuation.""" + from agent_framework_orchestrations._orchestration_request_info import AgentRequestInfoResponse + + from agent_framework import FileCheckpointStorage + + storage = FileCheckpointStorage( + tmp_path, + allowed_checkpoint_types=[ + "agent_framework_orchestrations._orchestration_request_info:AgentRequestInfoResponse" + ], + ) + + def build_workflow() -> tuple[Any, _KwargsCapturingAgent, _KwargsCapturingAgent]: + agent = _KwargsCapturingAgent(name="agent") + sibling = _KwargsCapturingAgent(name="sibling") + workflow = ( + SequentialBuilder(participants=[agent, sibling], checkpoint_storage=storage) + .with_request_info(agents=[agent]) + .build() + ) + return workflow, agent, sibling + + workflow, agent, _ = build_workflow() + old_kwargs = {"key": "OLD"} + new_kwargs = {"key": "NEW"} + + run_kwargs = {kwargs_channel: old_kwargs} + paused = await workflow.run("test", **run_kwargs) + [request] = paused.get_request_info_events() + assert agent.captured_kwargs[0].get(kwargs_channel) == old_kwargs + + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + checkpoint = max( + (item for item in checkpoints if item.pending_request_info_events), + key=lambda item: item.timestamp, + ) + + # 1. Resuming with checkpoint_id and new kwargs: new kwargs must override checkpoint's old kwargs + resumed_workflow, _, resumed_sibling = build_workflow() + resume_kwargs = {kwargs_channel: new_kwargs} + resumed = await resumed_workflow.run( + checkpoint_id=checkpoint.checkpoint_id, + **resume_kwargs, + ) + [resumed_request] = resumed.get_request_info_events() + assert resumed_request.request_id == request.request_id + + # 2. Resuming with responses: executor should receive NEW kwargs + await resumed_workflow.run( + responses={resumed_request.request_id: AgentRequestInfoResponse.approve()}, + ) + assert resumed_sibling.captured_kwargs[0].get(kwargs_channel) == new_kwargs + + # 3. Direct combined restore (checkpoint_id + responses + new kwargs) + combined_workflow, _, combined_sibling = build_workflow() + combined_kwargs = {"key": "COMBINED"} + combined_resume_kwargs = {kwargs_channel: combined_kwargs} + await combined_workflow.run( + checkpoint_id=checkpoint.checkpoint_id, + responses={request.request_id: AgentRequestInfoResponse.approve()}, + **combined_resume_kwargs, + ) + assert combined_sibling.captured_kwargs[0].get(kwargs_channel) == combined_kwargs + + # endregion