Fix output type loss in AgentOperator's human-in-the-loop review

When enable_hitl_review=True and output_type is a non-str, non-BaseModel
type (e.g. list[str], bool), the reviewed output was serialized with
str(output), producing a Python repr rather than valid JSON. Rehydrating
it back into output_type then silently failed and returned the raw repr
string instead of the original value. Serialize with TypeAdapter instead,
and reuse the existing rehydrate_pydantic_output helper (already used by
the require_approval path) instead of a separate, broken json.loads
fallback.

Author: ColtenOuO

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

@@ -21,7 +21,7 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Protocol
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
 from airflow.providers.common.ai.utils.hitl_review import (
@@ -271,5 +271,5 @@ def _to_string(output: Any) -> str:
         if isinstance(output, BaseModel):
             return output.model_dump_json()
         if not isinstance(output, str):
-            return str(output)
+            return TypeAdapter(type(output)).dump_json(output).decode()
         return output

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

@@ -18,14 +18,13 @@
 
 from __future__ import annotations
 
-import json
 from collections.abc import Sequence
 from dataclasses import replace
 from datetime import timedelta
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
 from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
@@ -493,16 +492,11 @@ def execute(self, context: Context) -> Any:
                 output,
                 message_history=result.all_messages(),
             )
-            if isinstance(self.output_type, type) and issubclass(self.output_type, BaseModel):
-                return rehydrate_pydantic_output(
-                    self.output_type,
-                    result_str,
-                    serialize_output=self._serialize_model_output,
-                )
-            try:
-                return json.loads(result_str)
-            except (ValueError, TypeError):
-                return result_str
+            return rehydrate_pydantic_output(
+                self.output_type,
+                result_str,
+                serialize_output=self._serialize_model_output,
+            )
 
         if self._serialize_model_output and isinstance(output, BaseModel):
             output = output.model_dump()
@@ -556,4 +550,6 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu
         output = result.output
         if isinstance(output, BaseModel):
             output = output.model_dump_json()
-        return str(output), result.all_messages()
+        elif not isinstance(output, str):
+            output = TypeAdapter(type(output)).dump_json(output).decode()
+        return output, result.all_messages()

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

@@ -23,9 +23,12 @@
 if not AIRFLOW_V_3_1_PLUS:
     pytest.skip("Human in the loop is only compatible with Airflow >= 3.1.0", allow_module_level=True)
 
+import json
 from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
+from pydantic import BaseModel
+
 from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
 from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
 from airflow.providers.common.ai.utils.hitl_review import (
@@ -107,6 +110,30 @@ def context(mock_ti):
     return {"task_instance": mock_ti}
 
 
+class TestToString:
+
+    def test_str_output_passes_through_unchanged(self):
+        assert HITLReviewMixin._to_string("plain text") == "plain text"
+
+    def test_base_model_output_uses_model_dump_json(self):
+        class Summary(BaseModel):
+            text: str
+
+        assert HITLReviewMixin._to_string(Summary(text="hi")) == '{"text":"hi"}'
+
+    def test_list_output_is_valid_json_not_python_repr(self):
+        result = HITLReviewMixin._to_string(["tag-a", "tag-b"])
+
+        assert result == '["tag-a","tag-b"]'
+        assert json.loads(result) == ["tag-a", "tag-b"]
+
+    def test_bool_output_is_valid_json_not_python_repr(self):
+        result = HITLReviewMixin._to_string(True)
+
+        assert result == "true"
+        assert json.loads(result) is True
+
+
 class TestHITLReviewMixin:
     @patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep", autospec=True)
     def test_pushes_session_and_output_on_start(

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

@@ -432,6 +432,31 @@ def test_execute_with_hitl_rehydrates_base_model(self, mock_hook_cls, mock_run_h
         assert result.text == "Approved summary"
         assert result.score == 0.9
 
+    @pytest.mark.skipif(
+        not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0"
+    )
+    @patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True)
+    @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+    def test_execute_with_hitl_rehydrates_non_base_model_output(self, mock_hook_cls, mock_run_hitl):
+        mock_result = _make_mock_run_result(["tag-a", "tag-b"])
+        mock_agent = MagicMock(spec=["run_sync"])
+        mock_agent.run_sync.return_value = mock_result
+        mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
+        mock_run_hitl.return_value = '["tag-a", "tag-b"]'
+
+        op = AgentOperator(
+            task_id="test",
+            prompt="Summarize",
+            llm_conn_id="my_llm",
+            output_type=list[str],
+            enable_hitl_review=True,
+            hitl_timeout=timedelta(minutes=5),
+        )
+        context = MagicMock()
+        result = op.execute(context=context)
+
+        assert result == ["tag-a", "tag-b"]
+
     @pytest.mark.skipif(
         not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0"
     )
@@ -584,6 +609,27 @@ def test_regenerate_with_feedback_serializes_base_model_output(self, mock_hook_c
 
         assert output == '{"text":"Revised","score":0.0}'
 
+    @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+    def test_regenerate_with_feedback_serializes_non_base_model_output(self, mock_hook_cls):
+        mock_result = _make_mock_run_result(["tag-a", "tag-b"])
+        mock_result.all_messages.return_value = []
+        mock_agent = MagicMock(spec=["run_sync"])
+        mock_agent.run_sync.return_value = mock_result
+        mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
+
+        op = AgentOperator(
+            task_id="test",
+            prompt="test",
+            llm_conn_id="my_llm",
+            output_type=list[str],
+        )
+        output, _ = op.regenerate_with_feedback(
+            feedback="Expand",
+            message_history=[],
+        )
+
+        assert output == '["tag-a","tag-b"]'
+
 
 class TestAgentOperatorDurable:
     def test_durable_param_stored(self):