diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 7c2217edc00..2fe0ea5c145 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -391,34 +391,16 @@ def can_handle(self, message: WorkflowMessage) -> bool: # For other messages, only handle if the wrapped workflow can accept them as input return any(is_instance_of(message.data, input_type) for input_type in self.workflow.input_types) - @handler - async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, Any]) -> None: - """Execute the sub-workflow with raw input data. - - This handler starts a new sub-workflow execution. When the sub-workflow - needs external information, it pauses and sends a request to the parent. - - Args: - input_data: The input data to send to the sub-workflow. - ctx: The workflow context from the parent. - """ - # The sub-workflow is a single shared instance. If it still has pending request_info events - # from an unfinished request/response cycle, a new input advances its shared state and can - # interfere with that cycle - a response arriving later may apply to a sub-workflow that has - # moved on. We allow it (the sub-workflow may be stateless) but warn so the risk is visible. - pending_requests = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] - if pending_requests: - logger.warning( - f"WorkflowExecutor {self.id} received a new input message while its sub-workflow " - f"({self.workflow.id}) still has {len(pending_requests)} pending request(s) from an " - f"unfinished request/response cycle. The sub-workflow is a single shared instance, so the " - f"new input advances shared state and can interfere with the in-flight cycle. Ensure the " - f"sub-workflow is stateless, or complete the pending cycle before sending new input." - ) - - logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id}") - - # Get kwargs from parent workflow's State to propagate to subworkflow + def _get_child_invocation_kwargs( + self, + ctx: WorkflowContext[Any, Any], + *, + preserve_empty: bool = False, + ) -> tuple[ + WorkflowInvocationKwargs | Mapping[str, Any] | None, + WorkflowInvocationKwargs | Mapping[str, Any] | None, + ]: + """Extract invocation kwargs from the parent context for the child workflow.""" parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) parent_resolved_kwargs: dict[str, Any] = ctx.get_state(RESOLVED_WORKFLOW_RUN_KWARGS_KEY, {}) parent_routed_kwargs: dict[str, list[str]] = ctx.get_state(ROUTED_WORKFLOW_RUN_KWARGS_KEY, {}) @@ -440,6 +422,10 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A self._parent_routed_keys_to_exclude(routed_keys), cast(Mapping[str, Any] | None, parent_resolved_kwargs.get(key)), ) + # Continuations must preserve an explicitly supplied channel after scoping; + # None would preserve the child's stale kwargs instead of clearing them. + if preserve_empty and resolved is None: + resolved = dict[str, Any]() if ( transparent_alias is not None and isinstance(resolved, Mapping) @@ -463,6 +449,37 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A else: ci_kwargs = resolved + return fi_kwargs, ci_kwargs + + @handler + async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, Any]) -> None: + """Execute the sub-workflow with raw input data. + + This handler starts a new sub-workflow execution. When the sub-workflow + needs external information, it pauses and sends a request to the parent. + + Args: + input_data: The input data to send to the sub-workflow. + ctx: The workflow context from the parent. + """ + # The sub-workflow is a single shared instance. If it still has pending request_info events + # from an unfinished request/response cycle, a new input advances its shared state and can + # interfere with that cycle - a response arriving later may apply to a sub-workflow that has + # moved on. We allow it (the sub-workflow may be stateless) but warn so the risk is visible. + pending_requests = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] + if pending_requests: + logger.warning( + f"WorkflowExecutor {self.id} received a new input message while its sub-workflow " + f"({self.workflow.id}) still has {len(pending_requests)} pending request(s) from an " + f"unfinished request/response cycle. The sub-workflow is a single shared instance, so the " + f"new input advances shared state and can interfere with the in-flight cycle. Ensure the " + f"sub-workflow is stateless, or complete the pending cycle before sending new input." + ) + + logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id}") + + fi_kwargs, ci_kwargs = self._get_child_invocation_kwargs(ctx) + # Run the sub-workflow and collect all events, passing parent kwargs result = await self.workflow.run( input_data, @@ -550,9 +567,12 @@ async def handle_propagated_request_response( @override async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: """Propagate cancellation into the wrapped workflow.""" + fi_kwargs, ci_kwargs = self._get_child_invocation_kwargs(ctx, preserve_empty=True) result = await self.workflow.cancel_pending_requests( [request_id], tools=ctx.get_runtime_tools(), + function_invocation_kwargs=fi_kwargs, + client_kwargs=ci_kwargs, ) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index ba91dc601f8..31bd7800271 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -1114,6 +1114,202 @@ async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, An ] +@pytest.mark.parametrize( + "kwargs_shape", ["child-values", "wrapper-only", "legacy-wrapper-only", "empty-mapping", "empty-typed"] +) +@pytest.mark.parametrize( + ("replace_function_kwargs", "replace_client_kwargs", "nested", "restore_checkpoint"), + [ + (True, True, False, False), + (True, False, False, False), + (False, True, False, False), + (True, True, True, False), + (True, False, True, False), + (False, True, True, False), + (True, True, False, True), + (True, False, False, True), + (False, True, False, True), + (False, False, False, False), + (False, False, True, False), + (False, False, False, True), + ], + ids=[ + "both", + "function-only", + "client-only", + "two-level-both", + "two-level-function-only", + "two-level-client-only", + "checkpoint-both", + "checkpoint-function-only", + "checkpoint-client-only", + "omitted", + "two-level-omitted", + "checkpoint-omitted", + ], +) +async def test_subworkflow_cancellation_replaces_child_invocation_kwargs( + replace_function_kwargs: bool, + replace_client_kwargs: bool, + nested: bool, + restore_checkpoint: bool, + kwargs_shape: str, +) -> None: + """Nested cancellation replaces supplied kwargs and clears a stale omitted peer channel.""" + from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + Workflow, + WorkflowBuilder, + WorkflowContext, + handler, + response_handler, + ) + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + captured_child_kwargs: list[dict[str, Any]] = [] + + class RequestingExecutor(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[Any, Any]) -> None: + del message + await ctx.request_info("Continue?", str, request_id="child-request") + + @response_handler + async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, Any]) -> None: + del request, response, ctx + + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: + del request_id + captured_child_kwargs.append(ctx.get_state(RESOLVED_WORKFLOW_RUN_KWARGS_KEY, {})) + + storage = InMemoryCheckpointStorage() + + def build_parent() -> Workflow: + child = WorkflowBuilder(start_executor=RequestingExecutor(id="child-requester")).build() + child_executor = WorkflowExecutor(child, id="child", propagate_request=True) + start_executor = ( + WorkflowExecutor( + WorkflowBuilder(start_executor=child_executor).build(), id="middle", propagate_request=True + ) + if nested + else child_executor + ) + return WorkflowBuilder( + name="nested-cancellation", start_executor=start_executor, checkpoint_storage=storage + ).build() + + parent = build_parent() + old_function_kwargs = {"phase": "old"} + old_client_kwargs = {"model": "old"} + new_function_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] = WorkflowInvocationKwargs( + global_kwargs={"phase": "new"}, + executor_kwargs={"child": {"parent": "new"}, "child-requester": {"request": "new"}}, + ) + new_client_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] = WorkflowInvocationKwargs( + global_kwargs={"model": "new"}, + executor_kwargs={"child": {"parent": "new"}, "child-requester": {"timeout": 30}}, + ) + if kwargs_shape == "wrapper-only": + wrapper_id = "middle" if nested else "child" + new_function_kwargs = WorkflowInvocationKwargs(executor_kwargs={wrapper_id: {"parent": "new"}}) + new_client_kwargs = WorkflowInvocationKwargs(executor_kwargs={wrapper_id: {"parent": "new"}}) + elif kwargs_shape == "legacy-wrapper-only": + wrapper_id = "middle" if nested else "child" + new_function_kwargs = {wrapper_id: {"parent": "new"}} + new_client_kwargs = {wrapper_id: {"parent": "new"}} + elif kwargs_shape == "empty-mapping": + new_function_kwargs = {} + new_client_kwargs = {} + elif kwargs_shape == "empty-typed": + new_function_kwargs = WorkflowInvocationKwargs() + new_client_kwargs = WorkflowInvocationKwargs() + + paused = await parent.run( + "start", + function_invocation_kwargs=old_function_kwargs, + client_kwargs=old_client_kwargs, + ) + [request] = paused.get_request_info_events() + if restore_checkpoint: + checkpoints = await storage.list_checkpoints(workflow_name=parent.name) + pending_checkpoint = max( + (checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events), + key=lambda checkpoint: checkpoint.timestamp, + ) + parent = build_parent() + _ = await parent.cancel_pending_requests( + [request.request_id], + checkpoint_id=pending_checkpoint.checkpoint_id, + checkpoint_storage=storage, + function_invocation_kwargs=new_function_kwargs if replace_function_kwargs else None, + client_kwargs=new_client_kwargs if replace_client_kwargs else None, + ) + else: + _ = await parent.cancel_pending_requests( + [request.request_id], + function_invocation_kwargs=new_function_kwargs if replace_function_kwargs else None, + client_kwargs=new_client_kwargs if replace_client_kwargs else None, + ) + + expected_child_kwargs: dict[str, Any] = {} + if replace_function_kwargs: + expected_child_kwargs["function_invocation_kwargs"] = { + "global_kwargs": {"phase": "new"} if kwargs_shape == "child-values" else {}, + "executor_kwargs": {"child-requester": {"request": "new"}} if kwargs_shape == "child-values" else {}, + } + if replace_client_kwargs: + expected_child_kwargs["client_kwargs"] = { + "global_kwargs": {"model": "new"} if kwargs_shape == "child-values" else {}, + "executor_kwargs": {"child-requester": {"timeout": 30}} if kwargs_shape == "child-values" else {}, + } + if not replace_function_kwargs and not replace_client_kwargs: + expected_child_kwargs = { + "function_invocation_kwargs": {"global_kwargs": old_function_kwargs, "executor_kwargs": {}}, + "client_kwargs": {"global_kwargs": old_client_kwargs, "executor_kwargs": {}}, + } + + assert captured_child_kwargs == [expected_child_kwargs] + + +@pytest.mark.parametrize("kwargs_shape", ["typed", "legacy", "omitted"]) +async def test_subworkflow_cancellation_kwargs_reach_downstream_agent(kwargs_shape: str) -> None: + """Cancellation clears scoped-out kwargs before a real downstream AgentExecutor runs.""" + from agent_framework import AgentExecutor, Executor, WorkflowBuilder, WorkflowContext, handler, response_handler + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + class RequestingExecutor(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[Any, list[Message]]) -> None: + del message + await ctx.request_info("Continue?", str, request_id="child-request") + + @response_handler + async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, list[Message]]) -> None: + del request, response, ctx + + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, list[Message]]) -> None: + del request_id + await ctx.send_message([Message("user", ["Continue after cancellation"])]) + + agent = _KwargsCapturingAgent(name="downstream") + requester = RequestingExecutor(id="child-requester") + child = WorkflowBuilder(start_executor=requester).add_edge(requester, AgentExecutor(agent)).build() + parent = WorkflowBuilder(start_executor=WorkflowExecutor(child, id="child", propagate_request=True)).build() + paused = await parent.run("start", function_invocation_kwargs={"phase": "old"}, client_kwargs={"model": "old"}) + [request] = paused.get_request_info_events() + replacement_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None + if kwargs_shape == "typed": + replacement_kwargs = WorkflowInvocationKwargs(executor_kwargs={"child": {"parent": "new"}}) + elif kwargs_shape == "legacy": + replacement_kwargs = {"child": {"parent": "new"}} + await parent.cancel_pending_requests([request.request_id], function_invocation_kwargs=replacement_kwargs) + assert agent.captured_kwargs[0].get("function_invocation_kwargs") == ( + {"phase": "old"} if kwargs_shape == "omitted" else {} + ) + assert agent.captured_kwargs[0].get("client_kwargs") == ({"model": "old"} if kwargs_shape == "omitted" else None) + + async def test_subworkflow_kwargs_accessible_via_state() -> None: """Test that kwargs are accessible via State within subworkflow.