Render multi-branch review choices as a multi-select

providers/common/ai/docs/operators/llm_branch.rst

@@ -83,8 +83,8 @@ branch(es) and wait for a human reviewer to approve the choice before any
 downstream task is skipped. The review form shows the LLM's choice and the
 valid downstream task IDs. When ``allow_modifications=True``, the reviewer
 can also change the choice — rendered as a dropdown of the downstream task
-IDs, or a free-text JSON list of task IDs (e.g. ``["task_a", "task_b"]``)
-with ``allow_multiple_branches=True``. The reviewed branch(es) are validated
+IDs, or a multi-select of them with ``allow_multiple_branches=True``. The
+reviewed branch(es) are validated
 against the downstream task IDs before branching:
 
 .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py

providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Protocol
@@ -109,14 +110,19 @@ def defer_for_approval(
         :param modification_schema: JSON schema for the editable ``output`` param
             when ``allow_modifications=True``. Defaults to ``{"type": "string"}``.
             Pass e.g. ``{"type": "string", "enum": [...]}`` to render a dropdown
-            of valid values in the review form.
+            of valid values in the review form, or ``{"type": "array", "items":
+            {"type": "string", "enum": [...]}, "examples": [...]}`` to render a
+            multi-select (JSON Schema forbids ``enum`` at the array level, so the
+            options come from ``examples``); a list submitted by the reviewer is
+            returned from ``execute_complete`` re-serialized as a compact JSON string.
         """
         from airflow.providers.standard.triggers.hitl import HITLTrigger
         from airflow.sdk.execution_time.hitl import upsert_hitl_detail
         from airflow.sdk.timezone import utcnow
 
         self.validate_approval_prompt()
 
+        raw_output = output
         if isinstance(output, BaseModel):
             output = output.model_dump_json()
         elif not isinstance(output, str):
@@ -133,9 +139,17 @@ def defer_for_approval(
 
         hitl_params: dict[str, dict[str, Any]] = {}
         if self.allow_modifications:
+            # The multi-select rendered for an array schema needs the list, not its JSON string
+            param_value: Any = output
+            if (
+                modification_schema is not None
+                and modification_schema.get("type") == "array"
+                and isinstance(raw_output, list)
+            ):
+                param_value = raw_output
             hitl_params = {
                 "output": {
-                    "value": output,
+                    "value": param_value,
                     "description": "Edit the output before approving (optional).",
                     "schema": modification_schema or {"type": "string"},
                 },
@@ -215,13 +229,33 @@ def execute_complete(self, context: Context, generated_output: str, event: dict[
         # when allow_modifications=False, bypassing the read-only approval flow.
         if getattr(self, "allow_modifications", False) and params_input:
             modified = params_input.get("output")
+            if "output" in params_input and modified is None:
+                raise HITLTriggerEventError(
+                    {
+                        "error": "Modified output must not be empty; edit it or reject instead.",
+                        "error_type": "validation",
+                    }
+                )
+            if isinstance(modified, list):
+                for item in modified:
+                    if not isinstance(item, str):
+                        raise HITLTriggerEventError(
+                            {
+                                "error": f"Modified output list items must be strings, "
+                                f"got {type(item).__name__}.",
+                                "error_type": "validation",
+                            }
+                        )
+                # Compact so an unchanged selection compares equal to generated_output
+                modified = json.dumps(modified, separators=(",", ":"))
             if modified is not None and not isinstance(modified, str):
                 # On the awaiting_input path nothing upstream schema-validates params_input
                 # (HITLTrigger did on the legacy path), so enforce the string contract here
                 # rather than returning a non-string as the task's output.
                 raise HITLTriggerEventError(
                     {
-                        "error": f"Modified output must be a string, got {type(modified).__name__}.",
+                        "error": f"Modified output must be a string or a list of strings, "
+                        f"got {type(modified).__name__}.",
                         "error_type": "validation",
                     }
                 )

providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py

@@ -56,7 +56,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
     unselected downstream tasks once a reviewer approves. The review form
     lists the valid downstream task IDs; with ``allow_modifications=True``
     the editable choice is rendered as a dropdown of those IDs (single-branch
-    mode) or a free-text JSON list (``allow_multiple_branches=True``), and
+    mode) or a multi-select of them (``allow_multiple_branches=True``), and
     the reviewed branch(es) are validated against the downstream task IDs
     before branching.
     """
@@ -121,7 +121,9 @@ def execute(self, context: Context) -> str | Iterable[str] | None:
                 f"```\nPrompt: {self.prompt}\n\nChosen branch(es): {chosen}\n```"
             )
             modification_schema = (
-                None if self.allow_multiple_branches else {"type": "string", "enum": choices}
+                {"type": "array", "items": {"type": "string", "enum": choices}, "examples": choices}
+                if self.allow_multiple_branches
+                else {"type": "string", "enum": choices}
             )
             self.defer_for_approval(  # type: ignore[misc]
                 context, branches, body=body, modification_schema=modification_schema

providers/common/ai/tests/unit/common/ai/mixins/test_approval.py

@@ -158,6 +158,22 @@ def test_modification_schema_overrides_output_param_schema(
         param = mock_upsert.call_args[1]["params"]["output"]
         assert param["schema"] == schema
 
+    @patch(HITL_TRIGGER_PATH, autospec=True)
+    @patch(UPSERT_HITL_PATH)
+    def test_array_schema_passes_list_param_value(
+        self, mock_upsert, mock_trigger_cls, approval_op_with_modifications, context
+    ):
+        choices = ["task_a", "task_b"]
+        schema = {"type": "array", "items": {"type": "string", "enum": choices}, "examples": choices}
+
+        approval_op_with_modifications.defer_for_approval(context, ["task_a"], modification_schema=schema)
+
+        param = mock_upsert.call_args[1]["params"]["output"]
+        assert param["value"] == ["task_a"]
+        assert param["schema"] == schema
+        defer_kwargs = approval_op_with_modifications.defer.call_args[1]
+        assert defer_kwargs["kwargs"]["generated_output"] == '["task_a"]'
+
     @patch(HITL_TRIGGER_PATH, autospec=True)
     @patch(UPSERT_HITL_PATH)
     def test_no_modifications_params_empty(self, mock_upsert, mock_trigger_cls, approval_op, context):
@@ -292,6 +308,52 @@ def test_approved_with_non_string_modified_output_raises(self, approval_op_with_
                 {}, generated_output="original output", event=event
             )
 
+    def test_approved_with_list_modified_output_is_serialized(self, approval_op_with_modifications):
+        event = {
+            "chosen_options": ["Approve"],
+            "responded_by_user": "editor",
+            "params_input": {"output": ["task_b", "task_c"]},
+        }
+
+        result = approval_op_with_modifications.execute_complete(
+            {}, generated_output='["task_a"]', event=event
+        )
+
+        assert result == '["task_b","task_c"]'
+
+    def test_approved_with_unmodified_list_output_returns_original(self, approval_op_with_modifications):
+        event = {
+            "chosen_options": ["Approve"],
+            "responded_by_user": "editor",
+            "params_input": {"output": ["task_a"]},
+        }
+
+        result = approval_op_with_modifications.execute_complete(
+            {}, generated_output='["task_a"]', event=event
+        )
+
+        assert result == '["task_a"]'
+
+    def test_approved_with_non_string_list_items_raises(self, approval_op_with_modifications):
+        event = {
+            "chosen_options": ["Approve"],
+            "responded_by_user": "editor",
+            "params_input": {"output": ["task_a", 2]},
+        }
+
+        with pytest.raises(HITLTriggerEventError, match="items must be strings, got int"):
+            approval_op_with_modifications.execute_complete({}, generated_output='["task_a"]', event=event)
+
+    def test_approved_with_cleared_output_raises(self, approval_op_with_modifications):
+        event = {
+            "chosen_options": ["Approve"],
+            "responded_by_user": "editor",
+            "params_input": {"output": None},
+        }
+
+        with pytest.raises(HITLTriggerEventError, match="must not be empty"):
+            approval_op_with_modifications.execute_complete({}, generated_output='["task_a"]', event=event)
+
     def test_approved_with_unmodified_output(self, approval_op_with_modifications):
         event = {
             "chosen_options": ["Approve"],

providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py

@@ -25,7 +25,7 @@
 from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
-from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.common.compat.sdk import Param, ParamValidationError, TaskDeferred
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
 
@@ -313,10 +313,10 @@ def test_review_form_lists_choices_and_renders_enum_dropdown(
     @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True)
     @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")
     @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True)
-    def test_review_form_multi_branch_keeps_string_schema(
+    def test_review_form_multi_branch_renders_multiselect(
         self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
     ):
-        """With allow_multiple_branches the editable param stays free-text (JSON list)."""
+        """With allow_multiple_branches the editable param is an array enum (multi-select)."""
         downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a", "task_b": "task_b"})
 
         mock_agent = MagicMock(spec=["run_sync"])
@@ -338,7 +338,17 @@ def test_review_form_multi_branch_keeps_string_schema(
 
         call_kwargs = mock_upsert.call_args.kwargs
         assert "Valid branches: `task_a`, `task_b`" in call_kwargs["body"]
-        assert call_kwargs["params"]["output"]["schema"] == {"type": "string"}
+        assert call_kwargs["params"]["output"]["schema"] == {
+            "type": "array",
+            "items": {"type": "string", "enum": ["task_a", "task_b"]},
+            "examples": ["task_a", "task_b"],
+        }
+        assert call_kwargs["params"]["output"]["value"] == ["task_a"]
+
+        schema = call_kwargs["params"]["output"]["schema"]
+        assert Param(schema=schema).resolve(["task_a"]) == ["task_a"]
+        with pytest.raises(ParamValidationError):
+            Param(schema=schema).resolve(["task_x"])
 
     @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True)
     def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_cls):
@@ -405,6 +415,30 @@ def test_execute_complete_with_modified_branch(self, mock_do_branch):
         assert result == "task_b"
         mock_do_branch.assert_called_once_with(ctx, "task_b")
 
+    @patch.object(LLMBranchOperator, "do_branch")
+    def test_execute_complete_with_multiselect_modified_branches(self, mock_do_branch):
+        """A list submitted by the multi-select review form branches into those tasks."""
+        mock_do_branch.return_value = ["task_b", "task_c"]
+        op = LLMBranchOperator(
+            task_id="t",
+            prompt="p",
+            llm_conn_id="c",
+            allow_multiple_branches=True,
+            allow_modifications=True,
+        )
+        op.downstream_task_ids = {"task_a", "task_b", "task_c"}
+        event = {
+            "chosen_options": ["Approve"],
+            "responded_by_user": "admin",
+            "params_input": {"output": ["task_b", "task_c"]},
+        }
+        ctx = _make_context()
+
+        result = op.execute_complete(ctx, generated_output='["task_a"]', event=event)
+
+        assert result == ["task_b", "task_c"]
+        mock_do_branch.assert_called_once_with(ctx, ["task_b", "task_c"])
+
     @patch.object(LLMBranchOperator, "do_branch")
     def test_execute_complete_rejects_invalid_modified_branch(self, mock_do_branch):
         """A reviewer-modified branch outside downstream_task_ids fails validation."""