From 372522fb734836cfe056b19369f8633d3cb733ce Mon Sep 17 00:00:00 2001 From: ryo-whaletech <304937387+ryo-whaletech@users.noreply.github.com> Date: Tue, 22 Sep 2026 01:35:13 +0900 Subject: [PATCH 1/3] Fix nested workflow cancellation kwargs propagation --- .../_workflows/_workflow_executor.py | 66 +++++++++++-------- .../tests/workflow/test_workflow_kwargs.py | 61 +++++++++++++++++ 2 files changed, 101 insertions(+), 26 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 691dafc03b8..160ca4a226c 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -351,6 +351,42 @@ 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) + @staticmethod + def _get_child_invocation_kwargs( + ctx: WorkflowContext[Any, Any], + ) -> 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, {}) + + # Use the caller's raw kwargs so legacy per-executor mappings are resolved + # against the child workflow's executor IDs rather than the parent's. + fi_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None + ci_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None + for key in ("function_invocation_kwargs", "client_kwargs"): + raw_key = ( + RAW_FUNCTION_INVOCATION_KWARGS_KEY if key == "function_invocation_kwargs" else RAW_CLIENT_KWARGS_KEY + ) + raw_value = parent_kwargs.get(raw_key) + if raw_value is not None: + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any], raw_value) + else: + normalized: Any = parent_kwargs.get(key) + if isinstance(normalized, dict): + normalized_dict = cast(dict[str, Any], normalized) + if len(normalized_dict) == 1 and GLOBAL_KWARGS_KEY in normalized_dict: + normalized = normalized_dict[GLOBAL_KWARGS_KEY] + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any] | None, normalized) + if resolved is not None: + if key == "function_invocation_kwargs": + fi_kwargs = resolved + 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. @@ -378,32 +414,7 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id}") - # Get kwargs from parent workflow's State to propagate to subworkflow - parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - - # Use the caller's raw kwargs so legacy per-executor mappings are resolved - # against the child workflow's executor IDs rather than the parent's. - fi_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None - ci_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None - for key in ("function_invocation_kwargs", "client_kwargs"): - raw_key = ( - RAW_FUNCTION_INVOCATION_KWARGS_KEY if key == "function_invocation_kwargs" else RAW_CLIENT_KWARGS_KEY - ) - raw_value = parent_kwargs.get(raw_key) - if raw_value is not None: - resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any], raw_value) - else: - normalized: Any = parent_kwargs.get(key) - if isinstance(normalized, dict): - normalized_dict = cast(dict[str, Any], normalized) - if len(normalized_dict) == 1 and GLOBAL_KWARGS_KEY in normalized_dict: - normalized = normalized_dict[GLOBAL_KWARGS_KEY] - resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any] | None, normalized) - if resolved is not None: - if key == "function_invocation_kwargs": - fi_kwargs = resolved - else: - ci_kwargs = resolved + 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( @@ -466,9 +477,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) 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 6685c70827a..c729ca6abf1 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -1114,6 +1114,67 @@ async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, An ] +async def test_subworkflow_cancellation_replaces_child_invocation_kwargs() -> None: + """Nested cancellation applies replacement invocation kwargs to the child workflow.""" + from agent_framework import Executor, 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, {})) + + child = WorkflowBuilder(start_executor=RequestingExecutor(id="child-requester")).build() + child_executor = WorkflowExecutor(child, id="child", propagate_request=True) + parent = WorkflowBuilder(start_executor=child_executor).build() + old_function_kwargs = {"phase": "old"} + old_client_kwargs = {"model": "old"} + new_function_kwargs = WorkflowInvocationKwargs( + global_kwargs={"phase": "new"}, + executor_kwargs={"child-requester": {"request": "new"}}, + ) + new_client_kwargs = WorkflowInvocationKwargs( + global_kwargs={"model": "new"}, + executor_kwargs={"child-requester": {"timeout": 30}}, + ) + + paused = await parent.run( + "start", + function_invocation_kwargs=old_function_kwargs, + client_kwargs=old_client_kwargs, + ) + [request] = paused.get_request_info_events() + _ = await parent.cancel_pending_requests( + [request.request_id], + function_invocation_kwargs=new_function_kwargs, + client_kwargs=new_client_kwargs, + ) + + assert captured_child_kwargs == [ + { + "function_invocation_kwargs": { + "global_kwargs": {"phase": "new"}, + "executor_kwargs": {"child-requester": {"request": "new"}}, + }, + "client_kwargs": { + "global_kwargs": {"model": "new"}, + "executor_kwargs": {"child-requester": {"timeout": 30}}, + }, + } + ] + + async def test_subworkflow_kwargs_accessible_via_state() -> None: """Test that kwargs are accessible via State within subworkflow. From e4dc04d4a6c33c3ac157294eabecf6cde927a235 Mon Sep 17 00:00:00 2001 From: WhaleTech <304937387+ryo-whaletech@users.noreply.github.com> Date: Tue, 22 Sep 2026 02:40:51 +0900 Subject: [PATCH 2/3] Test nested cancellation with one kwargs channel --- .../tests/workflow/test_workflow_kwargs.py | 56 +++++++++++++------ 1 file changed, 38 insertions(+), 18 deletions(-) diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index c729ca6abf1..4d4bda0aaf9 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -1114,8 +1114,16 @@ async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, An ] -async def test_subworkflow_cancellation_replaces_child_invocation_kwargs() -> None: - """Nested cancellation applies replacement invocation kwargs to the child workflow.""" +@pytest.mark.parametrize( + ("replace_function_kwargs", "replace_client_kwargs"), + [(True, True), (True, False), (False, True)], + ids=["both", "function-only", "client-only"], +) +async def test_subworkflow_cancellation_replaces_child_invocation_kwargs( + replace_function_kwargs: bool, + replace_client_kwargs: bool, +) -> None: + """Nested cancellation replaces supplied kwargs and clears a stale omitted peer channel.""" from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler, response_handler from agent_framework._workflows._workflow_executor import WorkflowExecutor @@ -1155,24 +1163,36 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An client_kwargs=old_client_kwargs, ) [request] = paused.get_request_info_events() - _ = await parent.cancel_pending_requests( - [request.request_id], - function_invocation_kwargs=new_function_kwargs, - client_kwargs=new_client_kwargs, - ) + if replace_function_kwargs and replace_client_kwargs: + _ = await parent.cancel_pending_requests( + [request.request_id], + function_invocation_kwargs=new_function_kwargs, + client_kwargs=new_client_kwargs, + ) + elif replace_function_kwargs: + _ = await parent.cancel_pending_requests( + [request.request_id], + function_invocation_kwargs=new_function_kwargs, + ) + else: + _ = await parent.cancel_pending_requests( + [request.request_id], + client_kwargs=new_client_kwargs, + ) - assert captured_child_kwargs == [ - { - "function_invocation_kwargs": { - "global_kwargs": {"phase": "new"}, - "executor_kwargs": {"child-requester": {"request": "new"}}, - }, - "client_kwargs": { - "global_kwargs": {"model": "new"}, - "executor_kwargs": {"child-requester": {"timeout": 30}}, - }, + expected_child_kwargs: dict[str, Any] = {} + if replace_function_kwargs: + expected_child_kwargs["function_invocation_kwargs"] = { + "global_kwargs": {"phase": "new"}, + "executor_kwargs": {"child-requester": {"request": "new"}}, } - ] + if replace_client_kwargs: + expected_child_kwargs["client_kwargs"] = { + "global_kwargs": {"model": "new"}, + "executor_kwargs": {"child-requester": {"timeout": 30}}, + } + + assert captured_child_kwargs == [expected_child_kwargs] async def test_subworkflow_kwargs_accessible_via_state() -> None: From 85b8d570b70a9c7c6ea75eee460e928086aa3514 Mon Sep 17 00:00:00 2001 From: WhaleTech <304937387+ryo-whaletech@users.noreply.github.com> Date: Wed, 30 Sep 2026 22:56:34 +0900 Subject: [PATCH 3/3] Python: Preserve explicit empty kwargs during nested cancellation --- .../_workflows/_workflow_executor.py | 8 +- .../tests/workflow/test_workflow_kwargs.py | 113 ++++++++++++++---- 2 files changed, 99 insertions(+), 22 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 86117b5d8a7..2fe0ea5c145 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -394,6 +394,8 @@ def can_handle(self, message: WorkflowMessage) -> bool: 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, @@ -420,6 +422,10 @@ def _get_child_invocation_kwargs( 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) @@ -561,7 +567,7 @@ 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) + 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(), diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index b767c0958c8..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,9 @@ 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"), [ @@ -1121,15 +1124,36 @@ async def resume(self, request: str, response: str, ctx: WorkflowContext[Any, An (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", ], - ids=["both", "function-only", "client-only", "two-level", "checkpoint-restore"], ) 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 ( @@ -1178,14 +1202,28 @@ def build_parent() -> Workflow: parent = build_parent() old_function_kwargs = {"phase": "old"} old_client_kwargs = {"model": "old"} - new_function_kwargs = WorkflowInvocationKwargs( + 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( + 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", @@ -1204,41 +1242,74 @@ def build_parent() -> Workflow: [request.request_id], checkpoint_id=pending_checkpoint.checkpoint_id, checkpoint_storage=storage, - function_invocation_kwargs=new_function_kwargs, - client_kwargs=new_client_kwargs, - ) - elif replace_function_kwargs and replace_client_kwargs: - _ = await parent.cancel_pending_requests( - [request.request_id], - function_invocation_kwargs=new_function_kwargs, - client_kwargs=new_client_kwargs, - ) - elif replace_function_kwargs: - _ = await parent.cancel_pending_requests( - [request.request_id], - function_invocation_kwargs=new_function_kwargs, + 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], - client_kwargs=new_client_kwargs, + 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"}, - "executor_kwargs": {"child-requester": {"request": "new"}}, + "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"}, - "executor_kwargs": {"child-requester": {"timeout": 30}}, + "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.