diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py index 5db554574f4..51c36ffdd61 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_control_flow.py @@ -219,7 +219,11 @@ async def handle_action( trigger: Any, ctx: WorkflowContext[LoopIterationResult], ) -> None: - """Advance to next item and send result.""" + """Handle loop-control signals or advance to the next item and send the result.""" + if isinstance(trigger, LoopControl): + await self.handle_loop_control(trigger, ctx) + return + state = await self._ensure_state_initialized(ctx, trigger) loop_id = self._init_executor_id diff --git a/python/packages/declarative/tests/test_graph_workflow_integration.py b/python/packages/declarative/tests/test_graph_workflow_integration.py index 830570ae0c9..eff2efde40e 100644 --- a/python/packages/declarative/tests/test_graph_workflow_integration.py +++ b/python/packages/declarative/tests/test_graph_workflow_integration.py @@ -10,6 +10,7 @@ """ import sys +from typing import Any import pytest @@ -25,13 +26,158 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) +from agent_framework import InMemoryCheckpointStorage # noqa: E402 + from agent_framework_declarative._workflows import ( # noqa: E402 ActionTrigger, DeclarativeWorkflowBuilder, ) +from agent_framework_declarative._workflows._declarative_base import LoopControl # noqa: E402 from agent_framework_declarative._workflows._factory import WorkflowFactory # noqa: E402 +class TestGraphWorkflowLoopControl: + """Exercise loop-control messages through the public workflow factory.""" + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("stop_item", [1, 2, 3]) + async def test_break_loop_stops_at_selected_item(self, stream: bool, stop_item: int) -> None: + definition = { + "kind": "Workflow", + "name": "break_selected_item", + "actions": [ + { + "kind": "Foreach", + "source": [1, 2, 3], + "actions": [ + {"kind": "SendActivity", "activity": "before"}, + {"kind": "If", "condition": f"=Local.item = {stop_item}", "then": [{"kind": "BreakLoop"}]}, + {"kind": "SendActivity", "activity": "after"}, + ], + }, + {"kind": "SendActivity", "activity": "done"}, + ], + } + workflow = WorkflowFactory().create_workflow_from_definition(definition) + if stream: + response = workflow.run({}, stream=True) + events = [event async for event in response] + result = await response.get_final_response() + assert [event.data for event in events if event.type == "output"] == result.get_outputs() + else: + result = await workflow.run({}) + assert result.get_outputs() == ["before", "after"] * (stop_item - 1) + ["before", "done"] + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("continue_item", [1, 2, 3]) + async def test_continue_loop_preserves_remaining_iterations(self, stream: bool, continue_item: int) -> None: + definition = { + "kind": "Workflow", + "name": "continue_selected_item", + "actions": [ + { + "kind": "Foreach", + "source": [1, 2, 3], + "actions": [ + {"kind": "SendActivity", "activity": "before"}, + { + "kind": "If", + "condition": f"=Local.item = {continue_item}", + "then": [{"kind": "ContinueLoop"}], + }, + {"kind": "SendActivity", "activity": "after"}, + ], + }, + {"kind": "SendActivity", "activity": "done"}, + ], + } + workflow = WorkflowFactory().create_workflow_from_definition(definition) + result = await workflow.run({}, stream=True).get_final_response() if stream else await workflow.run({}) + expected = [ + text for item in [1, 2, 3] for text in (["before"] if item == continue_item else ["before", "after"]) + ] + assert result.get_outputs() == expected + ["done"] + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("source", [[], [1, 2, 3]]) + async def test_break_loop_as_first_body_action(self, stream: bool, source: list[int]) -> None: + workflow = WorkflowFactory().create_workflow_from_definition({ + "kind": "Workflow", + "name": "break_first_action", + "actions": [ + {"kind": "Foreach", "source": source, "actions": [{"kind": "BreakLoop", "id": "stop"}]}, + {"kind": "SendActivity", "activity": "done"}, + ], + }) + result = await workflow.run({}, stream=True).get_final_response() if stream else await workflow.run({}) + assert result.get_outputs() == ["done"] + assert sum(event.type == "executor_completed" and event.executor_id == "stop" for event in result) == bool( + source + ) + + @pytest.mark.parametrize("stream", [False, True]) + async def test_nested_break_exits_only_inner_loop(self, stream: bool) -> None: + workflow = WorkflowFactory().create_workflow_from_definition({ + "kind": "Workflow", + "name": "nested_break", + "actions": [ + { + "kind": "Foreach", + "source": [1, 2], + "itemName": "outer", + "actions": [ + {"kind": "SendActivity", "activity": "outer-before"}, + { + "kind": "Foreach", + "source": [10, 20, 30], + "itemName": "inner", + "actions": [ + {"kind": "SendActivity", "activity": "inner"}, + {"kind": "BreakLoop"}, + ], + }, + {"kind": "SendActivity", "activity": "outer-after"}, + ], + }, + {"kind": "SendActivity", "activity": "done"}, + ], + }) + result = await workflow.run({}, stream=True).get_final_response() if stream else await workflow.run({}) + assert result.get_outputs() == ["outer-before", "inner", "outer-after"] * 2 + ["done"] + + @pytest.mark.parametrize( + "control_kind,expected", [("BreakLoop", ["done"]), ("ContinueLoop", ["body", "body", "done"])] + ) + async def test_loop_control_restored_into_fresh_workflow(self, control_kind: str, expected: list[str]) -> None: + storage = InMemoryCheckpointStorage() + definition: dict[str, Any] = { + "kind": "Workflow", + "name": "restored_loop_control", + "actions": [ + { + "kind": "Foreach", + "id": "loop", + "source": [1, 2, 3], + "actions": [ + {"kind": "SendActivity", "activity": "body"}, + {"kind": control_kind, "id": "control"}, + ], + }, + {"kind": "SendActivity", "activity": "done"}, + ], + } + workflow = WorkflowFactory(checkpoint_storage=storage).create_workflow_from_definition(definition) + await workflow.run({}) + checkpoint = next( + checkpoint + for checkpoint in await storage.list_checkpoints(workflow_name=workflow.name) + if any(isinstance(message.data, LoopControl) for message in checkpoint.messages.get("control", [])) + ) + restored = WorkflowFactory(checkpoint_storage=storage).create_workflow_from_definition(definition) + result = await restored.run(checkpoint_id=checkpoint.checkpoint_id) + assert result.get_outputs() == expected + + class TestGraphBasedWorkflowExecution: """Integration tests for graph-based workflow execution."""