Python: Apply caller-supplied run kwargs over checkpoint kwargs on restore

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,

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