From afc372220006870f4f4821d46491a27de1d3d97f Mon Sep 17 00:00:00 2001 From: Pankaj Singh Date: Fri, 16 Oct 2026 00:22:25 +0530 Subject: [PATCH 1/2] Add emit_message_history to record an agent transcript without seeding it Single-turn agents that want their transcript on XCom had to pass message_history=[], a parameter for seeding a conversation, not for observing a run. A dedicated flag states that intent directly, and do_xcom_push=False now also suppresses the transcript push, as it does the other XCom pushes. Co-Authored-By: Claude --- providers/common/ai/docs/message_history.rst | 9 ++++- providers/common/ai/docs/operators/agent.rst | 3 ++ .../providers/common/ai/operators/agent.py | 14 ++++++- .../unit/common/ai/operators/test_agent.py | 37 +++++++++++++++---- 4 files changed, 54 insertions(+), 9 deletions(-) diff --git a/providers/common/ai/docs/message_history.rst b/providers/common/ai/docs/message_history.rst index 453b4d7a7c37e..d0aabc58ed05a 100644 --- a/providers/common/ai/docs/message_history.rst +++ b/providers/common/ai/docs/message_history.rst @@ -30,6 +30,12 @@ turns and, after the run, pushes the full updated transcript run reads it back to resume the conversation. ``None`` (the default) keeps the single-turn behavior unchanged. +To record the transcript of a single-turn run, set ``emit_message_history=True``. +The operator starts the run with no prior turns and pushes ``result.all_messages()`` +to the same ``message_history`` XCom key, so the transcript is available to +downstream tasks for inspection, evaluation or storage. The transcript push follows +``do_xcom_push``. + The operator does **not** decide *where* a session is stored -- that keying is deployment-specific. The pattern is three tasks: load the prior transcript for the session, run the agent, store the updated transcript. The example keys a @@ -49,7 +55,8 @@ fresh session instead of trying to parse the string ``"None"``. .. note:: - ``message_history`` cannot be combined with ``enable_hitl_review`` -- the + ``message_history`` and ``emit_message_history`` cannot be combined with + ``enable_hitl_review`` -- the operator raises at construction. The post-review (human-approved) transcript is not recoverable today, so emitting the pre-review transcript would silently drop the reviewed turns. diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index fbb7b542989c4..4133e241fafb1 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -451,6 +451,9 @@ Parameters When set, the post-run transcript is pushed to XCom under the key ``message_history`` for the next run to resume. Default ``None`` (single-turn). See :doc:`../message_history`. +- ``emit_message_history``: Push the post-run transcript to XCom under the key + ``message_history`` without seeding the run, so a single-turn agent can record its + transcript. Default ``False``. See :doc:`../message_history`. - ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via ``model_dump()`` before being pushed to XCom. Default ``False`` -- the diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 16dccb90b9820..dba5c40a485f2 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -401,6 +401,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa object-storage XCom backend or trim old turns. Not supported together with ``enable_hitl_review`` (raises) -- the post-review transcript is not yet recoverable. + :param emit_message_history: When ``True``, push the post-run transcript to XCom + under the key ``message_history`` without seeding the run, so a single-turn + agent can record its transcript. ``message_history`` keeps pushing the + transcript on its own. Default ``False``. Not supported together with + ``enable_hitl_review`` (raises), for the same reason as ``message_history``. **HITL Review parameters** (requires the ``hitl_review`` plugin): @@ -502,6 +507,7 @@ def __init__( durable: bool = False, cache_prompt: bool = True, message_history: list[ModelMessage] | str | bytes | None = None, + emit_message_history: bool = False, # Agent feedback parameters enable_hitl_review: bool = False, max_hitl_iterations: int = 5, @@ -532,6 +538,7 @@ def __init__( # No validation here -- see coerce_usage_limits() docstring for why. self.usage_limits = usage_limits self.message_history = message_history + self.emit_message_history = emit_message_history self.durable = durable self.cache_prompt = cache_prompt @@ -576,6 +583,9 @@ def __init__( # surface the final message history. raise ValueError("message_history and enable_hitl_review=True cannot be used together.") + if emit_message_history and enable_hitl_review: + raise ValueError("emit_message_history and enable_hitl_review=True cannot be used together.") + if durable or enable_hitl_review: self._reject_sandbox_without_continuity(durable=durable, enable_hitl_review=enable_hitl_review) @@ -1035,7 +1045,7 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag if self._usage_budget and (run_usage := self._run_usage) is not None: self._log_cumulative_usage(run_usage) - if self.message_history is not None: + if self.message_history is not None or self.emit_message_history: self._emit_message_history(context, result) output = result.output @@ -1361,6 +1371,8 @@ def _resolve_message_history(self) -> list[ModelMessage] | None: def _emit_message_history(self, context: Context, result: Any) -> None: """Push the full post-run transcript to XCom for the next turn to resume.""" + if not self.do_xcom_push: + return transcript = ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode() context["task_instance"].xcom_push(key="message_history", value=transcript) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 5523c249b1f9b..83be2e471500c 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -2002,17 +2002,25 @@ def test_none_is_single_turn_no_history_no_emit(self, mock_hook_cls, make_mock_r assert "message_history" not in pushed_keys assert pushed_keys == {"run_id", "usage", MODEL_NAME_XCOM_KEY} + @pytest.mark.parametrize( + ("kwargs", "seeds_run"), + [ + pytest.param({"message_history": []}, True, id="message_history"), + pytest.param({"emit_message_history": True}, False, id="emit_message_history"), + ], + ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_transcript_emitted_to_xcom_when_history_set(self, mock_hook_cls, make_mock_run_result): - """When message_history is set, the post-run transcript is pushed to XCom and round-trips.""" + def test_transcript_emitted_to_xcom(self, mock_hook_cls, kwargs, seeds_run, make_mock_run_result): + """The post-run transcript is pushed to XCom and round-trips; only message_history seeds the run.""" mock_agent = _make_mock_agent("ok", make_mock_run_result) mock_agent.run_sync.return_value.all_messages.return_value = _sample_history() mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c", message_history=[]) + op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c", **kwargs) context = _make_context() op.execute(context=context) + assert ("message_history" in mock_agent.run_sync.call_args.kwargs) is seeds_run ti = context["task_instance"] pushes = {c.kwargs["key"]: c.kwargs["value"] for c in ti.xcom_push.call_args_list} assert set(pushes) == {"run_id", "usage", "message_history", MODEL_NAME_XCOM_KEY} @@ -2040,17 +2048,32 @@ def test_usage_limits_still_forwarded_with_history( @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) - def test_message_history_with_hitl_review_raises(self): - """message_history cannot be combined with HITL review (post-review transcript is lost).""" - with pytest.raises(ValueError, match="message_history and enable_hitl_review"): + @pytest.mark.parametrize("param", ["message_history", "emit_message_history"]) + def test_message_history_with_hitl_review_raises(self, param): + """Transcript emission cannot be combined with HITL review (post-review transcript is lost).""" + with pytest.raises(ValueError, match=f"{param} and enable_hitl_review"): AgentOperator( task_id="t", prompt="run", llm_conn_id="c", - message_history=[], enable_hitl_review=True, + **{param: [] if param == "message_history" else True}, ) + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + def test_transcript_not_emitted_when_do_xcom_push_disabled(self, mock_hook_cls, make_mock_run_result): + """do_xcom_push=False suppresses the transcript push as it does the other XCom pushes.""" + mock_agent = _make_mock_agent("ok", make_mock_run_result) + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = AgentOperator( + task_id="t", prompt="run", llm_conn_id="c", emit_message_history=True, do_xcom_push=False + ) + context = _make_context() + op.execute(context=context) + + context["task_instance"].xcom_push.assert_not_called() + def test_durable_path_also_seeds_message_history(self): """The durable branch forwards message_history into the run, and keeps its conversation id.""" seen: list[int] = [] From 39a4c59867749d666495455c1119af13bc834957 Mon Sep 17 00:00:00 2001 From: Pankaj Singh Date: Fri, 16 Oct 2026 12:03:59 +0530 Subject: [PATCH 2/2] Give emit_message_history its own switch independent of do_xcom_push Gating the transcript push on do_xcom_push also stopped it for message_history users, narrowing a stable feature. A tri-state emit_message_history defaults to whether the run is seeded, so existing sessions keep pushing, a single-turn run can opt in, and a seeded run can opt out. Co-Authored-By: Claude --- providers/common/ai/docs/message_history.rst | 5 +- providers/common/ai/docs/operators/agent.rst | 4 +- providers/common/ai/docs/stability.rst | 4 ++ providers/common/ai/docs/troubleshooting.rst | 2 +- .../providers/common/ai/operators/agent.py | 20 +++---- .../unit/common/ai/operators/test_agent.py | 56 ++++++++----------- 6 files changed, 43 insertions(+), 48 deletions(-) diff --git a/providers/common/ai/docs/message_history.rst b/providers/common/ai/docs/message_history.rst index d0aabc58ed05a..fa9ee78befee1 100644 --- a/providers/common/ai/docs/message_history.rst +++ b/providers/common/ai/docs/message_history.rst @@ -33,8 +33,9 @@ single-turn behavior unchanged. To record the transcript of a single-turn run, set ``emit_message_history=True``. The operator starts the run with no prior turns and pushes ``result.all_messages()`` to the same ``message_history`` XCom key, so the transcript is available to -downstream tasks for inspection, evaluation or storage. The transcript push follows -``do_xcom_push``. +downstream tasks for inspection, evaluation or storage. ``emit_message_history`` defaults +to ``True`` when ``message_history`` is set; pass ``False`` to keep a seeded run from +pushing it. The operator does **not** decide *where* a session is stored -- that keying is deployment-specific. The pattern is three tasks: load the prior transcript for diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 4133e241fafb1..9dbf792c9c9a9 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -452,8 +452,8 @@ Parameters ``message_history`` for the next run to resume. Default ``None`` (single-turn). See :doc:`../message_history`. - ``emit_message_history``: Push the post-run transcript to XCom under the key - ``message_history`` without seeding the run, so a single-turn agent can record its - transcript. Default ``False``. See :doc:`../message_history`. + ``message_history``. Defaults to ``True`` when ``message_history`` is set and ``False`` + otherwise; ``True`` records a single-turn run without seeding it. See :doc:`../message_history`. - ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via ``model_dump()`` before being pushed to XCom. Default ``False`` -- the diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index cced251110cff..d3b00c6642c61 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -69,6 +69,10 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit from can change. - Seeds the run with the given conversation, as a list of messages or its JSON form, and publishes the finished conversation under the ``message_history`` XCom key so a later run can continue it. Cannot be combined with ``enable_hitl_review``. + * - ``emit_message_history`` on ``AgentOperator`` + - Publishes the finished conversation under the ``message_history`` XCom key without + seeding the run. Defaults to ``True`` when ``message_history`` is set, and ``False`` + stops a seeded run from publishing it. Cannot be combined with ``enable_hitl_review``. * - ``durable=True`` on ``AgentOperator`` and ``@task.agent``, and :class:`~airflow.providers.common.ai.durable.AirflowDurability` (:doc:`durable_execution`) diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst index 7072e02c07b4c..4e233958dbb53 100644 --- a/providers/common/ai/docs/troubleshooting.rst +++ b/providers/common/ai/docs/troubleshooting.rst @@ -126,7 +126,7 @@ as a task failure. ``AgentOperator(durable=True)`` adds the capability to the agent, so a second one in ``capabilities`` would record every step twice. Drop one of them. -``message_history and enable_hitl_review=True cannot be used together`` +``message_history and enable_hitl_review=True cannot be used together`` and ``emit_message_history and enable_hitl_review=True cannot be used together`` The post-review transcript is not recoverable today, so the operator refuses rather than silently dropping the reviewed turns. See :doc:`message_history`. diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index dba5c40a485f2..89582de5858ad 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -401,11 +401,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa object-storage XCom backend or trim old turns. Not supported together with ``enable_hitl_review`` (raises) -- the post-review transcript is not yet recoverable. - :param emit_message_history: When ``True``, push the post-run transcript to XCom - under the key ``message_history`` without seeding the run, so a single-turn - agent can record its transcript. ``message_history`` keeps pushing the - transcript on its own. Default ``False``. Not supported together with - ``enable_hitl_review`` (raises), for the same reason as ``message_history``. + :param emit_message_history: Push the post-run transcript to XCom under the key + ``message_history``. Defaults to ``True`` when ``message_history`` is set and ``False`` + otherwise. Set ``True`` to record a single-turn run without seeding it, or ``False`` + to keep a seeded run from pushing it. Not supported together with + ``enable_hitl_review`` (raises). **HITL Review parameters** (requires the ``hitl_review`` plugin): @@ -507,7 +507,7 @@ def __init__( durable: bool = False, cache_prompt: bool = True, message_history: list[ModelMessage] | str | bytes | None = None, - emit_message_history: bool = False, + emit_message_history: bool | None = None, # Agent feedback parameters enable_hitl_review: bool = False, max_hitl_iterations: int = 5, @@ -538,7 +538,9 @@ def __init__( # No validation here -- see coerce_usage_limits() docstring for why. self.usage_limits = usage_limits self.message_history = message_history - self.emit_message_history = emit_message_history + self.emit_message_history = ( + message_history is not None if emit_message_history is None else emit_message_history + ) self.durable = durable self.cache_prompt = cache_prompt @@ -1045,7 +1047,7 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag if self._usage_budget and (run_usage := self._run_usage) is not None: self._log_cumulative_usage(run_usage) - if self.message_history is not None or self.emit_message_history: + if self.emit_message_history: self._emit_message_history(context, result) output = result.output @@ -1371,8 +1373,6 @@ def _resolve_message_history(self) -> list[ModelMessage] | None: def _emit_message_history(self, context: Context, result: Any) -> None: """Push the full post-run transcript to XCom for the next turn to resume.""" - if not self.do_xcom_push: - return transcript = ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode() context["task_instance"].xcom_push(key="message_history", value=transcript) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 83be2e471500c..be95f0f72c6aa 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -2003,15 +2003,16 @@ def test_none_is_single_turn_no_history_no_emit(self, mock_hook_cls, make_mock_r assert pushed_keys == {"run_id", "usage", MODEL_NAME_XCOM_KEY} @pytest.mark.parametrize( - ("kwargs", "seeds_run"), + ("kwargs", "emits"), [ - pytest.param({"message_history": []}, True, id="message_history"), - pytest.param({"emit_message_history": True}, False, id="emit_message_history"), + pytest.param({"message_history": []}, True, id="seeded"), + pytest.param({"emit_message_history": True}, True, id="opt-in"), + pytest.param({"message_history": [], "emit_message_history": False}, False, id="opt-out"), ], ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_transcript_emitted_to_xcom(self, mock_hook_cls, kwargs, seeds_run, make_mock_run_result): - """The post-run transcript is pushed to XCom and round-trips; only message_history seeds the run.""" + def test_transcript_emitted_to_xcom(self, mock_hook_cls, kwargs, emits, make_mock_run_result): + """The transcript push follows emit_message_history, which defaults to whether the run is seeded.""" mock_agent = _make_mock_agent("ok", make_mock_run_result) mock_agent.run_sync.return_value.all_messages.return_value = _sample_history() mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent @@ -2020,12 +2021,13 @@ def test_transcript_emitted_to_xcom(self, mock_hook_cls, kwargs, seeds_run, make context = _make_context() op.execute(context=context) - assert ("message_history" in mock_agent.run_sync.call_args.kwargs) is seeds_run + # emit_message_history never seeds the run. + assert ("message_history" in mock_agent.run_sync.call_args.kwargs) is ("message_history" in kwargs) ti = context["task_instance"] pushes = {c.kwargs["key"]: c.kwargs["value"] for c in ti.xcom_push.call_args_list} - assert set(pushes) == {"run_id", "usage", "message_history", MODEL_NAME_XCOM_KEY} - restored = ModelMessagesTypeAdapter.validate_json(pushes["message_history"]) - assert len(restored) == 2 + assert ("message_history" in pushes) is emits + if emits: + assert len(ModelMessagesTypeAdapter.validate_json(pushes["message_history"])) == 2 @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_usage_limits_still_forwarded_with_history( @@ -2048,31 +2050,19 @@ def test_usage_limits_still_forwarded_with_history( @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) - @pytest.mark.parametrize("param", ["message_history", "emit_message_history"]) - def test_message_history_with_hitl_review_raises(self, param): + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + pytest.param({"message_history": []}, "message_history and enable_hitl_review", id="seeded"), + pytest.param( + {"emit_message_history": True}, "emit_message_history and enable_hitl_review", id="opt-in" + ), + ], + ) + def test_message_history_with_hitl_review_raises(self, kwargs, match): """Transcript emission cannot be combined with HITL review (post-review transcript is lost).""" - with pytest.raises(ValueError, match=f"{param} and enable_hitl_review"): - AgentOperator( - task_id="t", - prompt="run", - llm_conn_id="c", - enable_hitl_review=True, - **{param: [] if param == "message_history" else True}, - ) - - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_transcript_not_emitted_when_do_xcom_push_disabled(self, mock_hook_cls, make_mock_run_result): - """do_xcom_push=False suppresses the transcript push as it does the other XCom pushes.""" - mock_agent = _make_mock_agent("ok", make_mock_run_result) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - - op = AgentOperator( - task_id="t", prompt="run", llm_conn_id="c", emit_message_history=True, do_xcom_push=False - ) - context = _make_context() - op.execute(context=context) - - context["task_instance"].xcom_push.assert_not_called() + with pytest.raises(ValueError, match=match): + AgentOperator(task_id="t", prompt="run", llm_conn_id="c", enable_hitl_review=True, **kwargs) def test_durable_path_also_seeds_message_history(self): """The durable branch forwards message_history into the run, and keeps its conversation id."""