diff --git a/providers/common/ai/docs/agent_security.rst b/providers/common/ai/docs/agent_security.rst index f58a79aeb8609..8997ab57999ea 100644 --- a/providers/common/ai/docs/agent_security.rst +++ b/providers/common/ai/docs/agent_security.rst @@ -115,7 +115,7 @@ No single layer is sufficient — they work together. not stop an agent reaching connections through some other tool. It also does not sanitize what the code computes or returns. Custom images can carry secrets and a backend you add can expose its own identity. The - ``sbx`` backend leaks orphaned microVMs if the worker is killed, and its + ``sbx`` backend leaks orphaned microVMs if the worker is killed outright (SIGKILL), and its CPU allocation defaults to every host CPU; the Modal backend reclaims its sandboxes on a server-side timeout, and if you opt into ``egress_enforcement="sni"`` its hostname allowlist is enforced at the TLS diff --git a/providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py b/providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py new file mode 100644 index 0000000000000..5e40ce85b507a --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py @@ -0,0 +1,70 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Mixin that cancels an operator's in-flight pydantic-ai run when the task is killed.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +from pydantic_ai import CancellationToken + +if TYPE_CHECKING: + from pydantic_ai import Agent, AgentRunResult + + +class CancellableAgentRunMixin: + """ + Run a pydantic-ai agent synchronously with kill-time cancellation wired in. + + The wrapper holds the in-flight run's ``CancellationToken`` so :meth:`on_kill` can + cancel it. Cancelling makes ``run_sync`` raise ``RunCancelled`` and unwind, giving the + agent's toolsets a chance to exit (tearing down a provisioned sandbox, for one) before + SIGKILL rather than leaving the run to die mid-flight. + + This needs the Task SDK to call ``on_kill`` on SIGTERM, which lands in Airflow 3.0.4 and + 3.1.0. On 3.0.0 to 3.0.3 the handler is absent, so a kill runs the pre-existing path (the + run continues until SIGKILL). + """ + + # Set only while a run is in flight. Read by on_kill from the signal handler. + _cancellation_token: CancellationToken | None = None + + # Provided by BaseOperator at runtime. Declared here for the type checker. + log: Any + + def run_agent_sync( + self, agent: Agent[Any, Any], user_prompt: Any, **run_kwargs: Any + ) -> AgentRunResult[Any]: + """Call ``agent.run_sync`` under a fresh cancellation token held for :meth:`on_kill`.""" + self._cancellation_token = CancellationToken() + try: + return agent.run_sync(user_prompt, cancellation_token=self._cancellation_token, **run_kwargs) + finally: + self._cancellation_token = None + + def on_kill(self) -> None: + token = self._cancellation_token + if token is None: + return + self.log.info("Task killed, cancelling in-flight agent run") + # Cancel from a separate thread, not inline. on_kill runs in the Task SDK's + # SIGTERM handler on the same thread that drives run_sync, and cancel() only + # interrupts a blocked run when issued from a different thread. Called inline it + # defers until the in-flight await returns, so the worker is SIGKILLed at the + # grace deadline before the run unwinds. + threading.Thread(target=token.cancel, name="agent-run-cancel", daemon=True).start() 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 4038edf9f1290..03be4e1af8f5f 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 @@ -29,6 +29,7 @@ from pydantic_ai.capabilities import Toolset from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook +from airflow.providers.common.ai.mixins.cancellable_run import CancellableAgentRunMixin from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin from airflow.providers.common.ai.observability import ( build_run_identity_attributes, @@ -125,7 +126,10 @@ def _build_code_mode() -> Any: return CodeMode() -class AgentOperator(BaseOperator, HITLReviewMixin): +# CancellableAgentRunMixin must precede BaseOperator so its on_kill overrides BaseOperator's +# no-op. The other mixins only add methods, so they can trail BaseOperator. See the MRO guard +# test in tests/unit/common/ai/mixins/test_cancellable_run.py. +class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin): """ Run a pydantic-ai Agent with tools and multi-turn reasoning. @@ -555,6 +559,8 @@ def execute(self, context: Context) -> Any: storage = self._durable_storage counter = self._durable_counter + # A killed run raises RunCancelled (see run_agent_sync), which propagates to fail the + # task. The durable cache cleanup below is skipped on the raise, preserving it for retry. if self.durable and storage is not None and counter is not None: from pydantic_ai.models import infer_model @@ -565,9 +571,9 @@ def execute(self, context: Context) -> Any: resolved_model = infer_model(agent.model) caching_model = CachingModel(resolved_model, storage=storage, counter=counter) with agent.override(model=caching_model): - result = agent.run_sync(self.prompt, **run_kwargs) + result = self.run_agent_sync(agent, self.prompt, **run_kwargs) else: - result = agent.run_sync(self.prompt, **run_kwargs) + result = self.run_agent_sync(agent, self.prompt, **run_kwargs) log_run_summary(self.log, result) self._emit_run_metadata(context, result) @@ -680,7 +686,8 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu if identity: stamp_identity_on_agent_spans(agent, identity) messages = message_history or [] - result = agent.run_sync( + result = self.run_agent_sync( + agent, feedback, message_history=messages, usage_limits=usage_limits, diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py index 54832926b2cbe..df3d0e8d515ef 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py @@ -27,6 +27,7 @@ from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin +from airflow.providers.common.ai.mixins.cancellable_run import CancellableAgentRunMixin from airflow.providers.common.ai.policies.decision import DecisionPolicy from airflow.providers.common.ai.utils.decision import ( DECISION_XCOM_KEY, @@ -70,7 +71,10 @@ __all__ = ["DecisionPolicy", "LLMOperator"] -class LLMOperator(BaseOperator, LLMApprovalMixin): +# CancellableAgentRunMixin must precede BaseOperator so its on_kill overrides BaseOperator's +# no-op. The other mixins only add methods, so they can trail BaseOperator. See the MRO guard +# test in tests/unit/common/ai/mixins/test_cancellable_run.py. +class LLMOperator(CancellableAgentRunMixin, BaseOperator, LLMApprovalMixin): """ Call an LLM with a prompt and return the output. @@ -319,7 +323,7 @@ def execute(self, context: Context) -> Any: agent: Agent[object, Any] = self.llm_hook.create_agent( output_type=self.output_type, instructions=self.system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt, usage_limits=usage_limits) + result = self.run_agent_sync(agent, self.prompt, usage_limits=usage_limits) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py index 7f6a271a08234..939e3fe5b4e91 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py @@ -229,7 +229,7 @@ def execute(self, context: Context) -> str | Iterable[str] | None: instructions=self.system_prompt, **self.agent_params, ) - result = agent.run_sync(self.prompt, usage_limits=usage_limits) + result = self.run_agent_sync(agent, self.prompt, usage_limits=usage_limits) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py index a9686dc82428d..8d90bc77b907c 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py @@ -164,7 +164,7 @@ def execute(self, context: Context) -> Any: instructions=self._build_system_prompt(), **self.agent_params, ) - result = agent.run_sync(request.user_content, usage_limits=usage_limits) + result = self.run_agent_sync(agent, request.user_content, usage_limits=usage_limits) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py index b75a8645197c6..d82b3df2da011 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py @@ -340,7 +340,7 @@ def execute(self, context: Context) -> dict[str, Any]: **self.agent_params, ) self.log.info("Running LLM schema comparison...") - result = agent.run_sync(self.prompt, usage_limits=usage_limits) + result = self.run_agent_sync(agent, self.prompt, usage_limits=usage_limits) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index f744acaeba6e3..ea392282e538c 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -162,7 +162,7 @@ def execute(self, context: Context) -> str: agent = self.llm_hook.create_agent( output_type=str, instructions=full_system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt, usage_limits=usage_limits) + result = self.run_agent_sync(agent, self.prompt, usage_limits=usage_limits) log_run_summary(self.log, result) sql = self._strip_llm_output(result.output, dialect=self._resolved_dialect) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py index 00fd431b66f7c..2be4009e07831 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from pydantic import BaseModel @@ -70,7 +70,7 @@ def my_prompt(): assert result == "The top customer is Acme Corp." assert op.prompt == "Who is our top customer?" mock_agent.run_sync.assert_called_once_with( - "Who is our top customer?", usage_limits=None, run_id="ti-1" + "Who is our top customer?", usage_limits=None, run_id="ti-1", cancellation_token=ANY ) @pytest.mark.parametrize( @@ -105,7 +105,9 @@ def my_prompt(): op.execute(context=_make_context()) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None, run_id="ti-1") + mock_agent.run_sync.assert_called_once_with( + prompt, usage_limits=None, run_id="ti-1", cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_sequence_prompt_with_hitl_review_raises_before_run_sync(self, mock_hook_cls): @@ -149,7 +151,7 @@ def my_prompt(topic): assert op.prompt == "Analyze revenue trends" mock_agent.run_sync.assert_called_once_with( - "Analyze revenue trends", usage_limits=None, run_id="ti-1" + "Analyze revenue trends", usage_limits=None, run_id="ti-1", cancellation_token=ANY ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py index a66a4006bc367..1e1e656d57bba 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from pydantic_ai.messages import ImageUrl @@ -45,7 +45,9 @@ def my_prompt(): assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("Summarize this text", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Summarize this text", usage_limits=None, cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_forwards_fallback_conn_ids_to_hook(self, mock_hook_cls, make_mock_run_result): @@ -98,7 +100,7 @@ def my_prompt(): op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None, cancellation_token=ANY) @pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="require_approval needs Airflow >= 3.1.0") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @@ -137,4 +139,6 @@ def my_prompt(topic): op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Summarize quantum computing" - mock_agent.run_sync.assert_called_once_with("Summarize quantum computing", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Summarize quantum computing", usage_limits=None, cancellation_token=ANY + ) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py index 3ff6993626102..5f39b54955877 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py @@ -17,7 +17,7 @@ from __future__ import annotations from enum import Enum -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from pydantic_ai.messages import ImageUrl @@ -57,7 +57,9 @@ def my_prompt(): assert result == "positive" assert op.prompt == "Route this review" - mock_agent.run_sync.assert_called_once_with("Route this review", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Route this review", usage_limits=None, cancellation_token=ANY + ) mock_do_branch.assert_called_once() @pytest.mark.parametrize( @@ -121,7 +123,7 @@ def my_prompt(): op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None, cancellation_token=ANY) @patch.object(LLMBranchOperator, "do_branch") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py index 2346ed6a0cf67..b4bfb040bfa5b 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest @@ -72,7 +72,9 @@ def my_prompt(): assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "prepared prompt", usage_limits=None, cancellation_token=ANY + ) @pytest.mark.parametrize( "return_value", @@ -116,4 +118,6 @@ def my_prompt(topic): op.execute(context={"task_instance": MagicMock(spec=["task_id"])}) assert op.prompt == "Summarize system logs" - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "prepared prompt", usage_limits=None, cancellation_token=ANY + ) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py index 099b4118c783c..e2d1f2c54b801 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from pydantic_ai.messages import ImageUrl @@ -45,7 +45,9 @@ def my_prompt_fn(): assert result == "SELECT 1" assert op.prompt == "Get all users" - mock_agent.run_sync.assert_called_once_with("Get all users", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Get all users", usage_limits=None, cancellation_token=ANY + ) @pytest.mark.parametrize( "return_value", @@ -79,7 +81,7 @@ def my_prompt_fn(): op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None, cancellation_token=ANY) @pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="require_approval needs Airflow >= 3.1.0") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_cancellable_run.py b/providers/common/ai/tests/unit/common/ai/mixins/test_cancellable_run.py new file mode 100644 index 0000000000000..69cf184dbc806 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/mixins/test_cancellable_run.py @@ -0,0 +1,100 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import threading +import time +from unittest.mock import DEFAULT, MagicMock + +import pytest +from pydantic_ai import CancellationToken + +from airflow.providers.common.ai.mixins.cancellable_run import CancellableAgentRunMixin +from airflow.providers.common.ai.operators.agent import AgentOperator +from airflow.providers.common.ai.operators.llm import LLMOperator + + +class TestRunAgentSync: + def test_forwards_held_cancellation_token_and_clears_after_success(self): + mixin = CancellableAgentRunMixin() + agent = MagicMock(spec=["run_sync"]) + held: dict[str, object] = {} + + def capture(*args, **kwargs): + # Capture the token the mixin holds mid-run: that is the one on_kill would cancel. + held["token"] = mixin._cancellation_token + return DEFAULT + + agent.run_sync.side_effect = capture + + result = mixin.run_agent_sync(agent, "prompt", usage_limits=None) + + assert result is agent.run_sync.return_value + passed = agent.run_sync.call_args.kwargs["cancellation_token"] + assert isinstance(passed, CancellationToken) + # run_sync must receive the exact token on_kill cancels, not just some CancellationToken. + assert passed is held["token"] + agent.run_sync.assert_called_once_with("prompt", cancellation_token=passed, usage_limits=None) + # The token is dropped once the run returns so a later on_kill is a no-op. + assert mixin._cancellation_token is None + + def test_clears_token_when_run_raises(self): + mixin = CancellableAgentRunMixin() + agent = MagicMock(spec=["run_sync"]) + agent.run_sync.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError): + mixin.run_agent_sync(agent, "prompt") + + assert mixin._cancellation_token is None + + +class TestOnKill: + def test_noop_when_no_run_active(self): + mixin = CancellableAgentRunMixin() + mixin.log = MagicMock() + # No run in flight: on_kill must not raise and must not cancel anything. + mixin.on_kill() + + def test_cancels_active_token_off_the_calling_thread(self): + """cancel() only interrupts a blocked run_sync when issued from a thread other than + the one running the run. on_kill runs in the SIGTERM handler on that same thread, so + it must hand the cancel to a separate thread. An inline cancel would run on the + calling thread and fail this test.""" + mixin = CancellableAgentRunMixin() + mixin.log = MagicMock() + cancel_thread: dict[str, int] = {} + token = MagicMock(spec=CancellationToken) + token.cancel.side_effect = lambda: cancel_thread.setdefault("id", threading.get_ident()) + mixin._cancellation_token = token + + mixin.on_kill() + + deadline = time.monotonic() + 2 + while "id" not in cancel_thread and time.monotonic() < deadline: + time.sleep(0.01) + token.cancel.assert_called_once_with() + assert cancel_thread["id"] != threading.get_ident() + + +class TestOnKillMroBinding: + @pytest.mark.parametrize("operator_cls", [AgentOperator, LLMOperator]) + def test_operator_on_kill_resolves_to_mixin(self, operator_cls): + """The mixin must precede BaseOperator in the bases so its on_kill wins. Reordering it + after BaseOperator (to match the sibling mixins) would silently restore the no-op and + disable kill-time cancellation, which the direct on_kill tests above would not catch.""" + assert operator_cls.on_kill is CancellableAgentRunMixin.on_kill 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 89411f5599202..129749b149862 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 @@ -19,7 +19,7 @@ import sys from datetime import timedelta from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from pydantic import BaseModel @@ -224,7 +224,9 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls, make_moc ) op.execute(context=_make_context()) - mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits, run_id="ti-1") + mock_agent.run_sync.assert_called_once_with( + "run", usage_limits=limits, run_id="ti-1", cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_execute_coerces_usage_limits_dict_before_run_sync(self, mock_hook_cls, make_mock_run_result): @@ -351,6 +353,7 @@ def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls, mak "Add detail", message_history=[], usage_limits=limits, + cancellation_token=ANY, ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) @@ -390,7 +393,9 @@ def test_execute_creates_agent_from_hook(self, mock_hook_cls, make_mock_run_resu mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with( output_type=str, instructions="You are helpful." ) - mock_agent.run_sync.assert_called_once_with("What is the answer?", usage_limits=None, run_id="ti-1") + mock_agent.run_sync.assert_called_once_with( + "What is the answer?", usage_limits=None, run_id="ti-1", cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls, make_mock_run_result): @@ -814,6 +819,7 @@ def test_regenerate_with_feedback_calls_agent_with_feedback_and_history( "Add more detail", message_history=msg_history, usage_limits=None, + cancellation_token=ANY, ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) @@ -920,7 +926,9 @@ def test_execute_non_durable_does_not_wrap(self, mock_hook_cls, make_mock_run_re op.execute(context=_make_context()) # run_sync called directly, no override - mock_agent.run_sync.assert_called_once_with("test", usage_limits=None, run_id="ti-1") + mock_agent.run_sync.assert_called_once_with( + "test", usage_limits=None, run_id="ti-1", cancellation_token=ANY + ) def test_build_durable_capabilities_wraps_toolset_capability(self): """A ``Toolset`` capability's inner toolset is wrapped with CachingToolset; @@ -1179,6 +1187,28 @@ def test_durable_path_also_seeds_message_history( assert len(passed) == 2 +class TestAgentOperatorCancellation: + """A killed run raises RunCancelled and propagates to fail the task.""" + + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + def test_run_cancelled_propagates_without_emitting_history(self, mock_hook_cls): + """RunCancelled is not swallowed, and no partial transcript is salvaged to XCom even for a + message_history session: a retry clears the TI's XCom before it starts, so nothing reads it.""" + from pydantic_ai import RunCancelled + + mock_agent = MagicMock(spec=["run_sync", "instrument"]) + mock_agent.run_sync.side_effect = RunCancelled("killed", messages=_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=[]) + context = _make_context() + with pytest.raises(RunCancelled): + op.execute(context=context) + + pushed_keys = {c.kwargs["key"] for c in context["task_instance"].xcom_push.call_args_list} + assert "message_history" not in pushed_keys + + class TestAgentOperatorHITLArgumentChecks: """The order in which __init__ reports conflicting HITL arguments.""" diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index 827de2b0730fb..bf87665a0a6d0 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -19,7 +19,7 @@ import logging from datetime import timedelta from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest @@ -133,7 +133,9 @@ def test_execute_returns_string_output(self, mock_hook_cls, make_mock_run_result result = op.execute(context=MagicMock()) assert result == "Paris is the capital of France." - mock_agent.run_sync.assert_called_once_with("What is the capital of France?", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "What is the capital of France?", usage_limits=None, cancellation_token=ANY + ) mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with( output_type=str, instructions="" ) @@ -157,7 +159,7 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls, make_moc ) op.execute(context=MagicMock()) - mock_agent.run_sync.assert_called_once_with("Summarize", usage_limits=limits) + mock_agent.run_sync.assert_called_once_with("Summarize", usage_limits=limits, cancellation_token=ANY) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_coerces_usage_limits_dict_before_run_sync(self, mock_hook_cls, make_mock_run_result): diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py index 835a4e4911736..2f75362aaecb3 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py @@ -17,7 +17,7 @@ from __future__ import annotations from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest @@ -85,7 +85,9 @@ def test_execute_single_branch(self, mock_hook_cls, mock_do_branch, make_mock_ru assert result == "task_a" mock_do_branch.assert_called_once_with(ctx, "task_a") - mock_agent.run_sync.assert_called_once_with("Pick a branch", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Pick a branch", usage_limits=None, cancellation_token=ANY + ) @patch.object(LLMBranchOperator, "do_branch") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py index d5ff67fd88107..2617697679e52 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py @@ -18,7 +18,7 @@ from datetime import timedelta from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest @@ -128,7 +128,9 @@ def test_execute_returns_string_output(self, mock_build_request, mock_hook_cls): max_text_chars=100_000, sample_rows=10, ) - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "prepared prompt", usage_limits=None, cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch( diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py index 1443ce54661d1..8be1f8c5e7d47 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py @@ -280,7 +280,9 @@ def test_execute(self, mock_build_system_prompt, mock_build_schema_context, make instructions="system_prompt", param="value", ) - mock_agent.run_sync.assert_called_once_with("user_prompt", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "user_prompt", usage_limits=None, cancellation_token=mock.ANY + ) assert result == {"compatible": True, "mismatches": [], "summary": "All good"} @mock.patch( @@ -395,6 +397,7 @@ def test_execute_schema_comparison_mixed_conn(self, mock_get_db_hook, db_hook, m mock_agent.run_sync.assert_called_once_with( "Compare S3 Parquet schema against the Postgres table and flag breaking changes", usage_limits=None, + cancellation_token=mock.ANY, ) assert result["compatible"] is True assert result["summary"] == "S3 and Postgres schemas are compatible" diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index b373877549f45..1676cc114aeed 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py @@ -20,7 +20,7 @@ import sys from datetime import timedelta from decimal import Decimal -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import ANY, MagicMock, PropertyMock, patch from uuid import uuid4 import pytest @@ -241,7 +241,9 @@ def test_execute_with_schema_context(self, mock_hook_cls, make_mock_run_result): result = op.execute(context=MagicMock()) assert result == "SELECT id, name FROM users WHERE active = true" - mock_agent.run_sync.assert_called_once_with("Get active users", usage_limits=None) + mock_agent.run_sync.assert_called_once_with( + "Get active users", usage_limits=None, cancellation_token=ANY + ) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_coerces_usage_limits_dict_before_run_sync(self, mock_hook_cls, make_mock_run_result):