Cancel the agent run when a common.ai LLM or agent task is killed

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

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()

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,

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
 

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
 

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
 

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
 

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)
 

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)

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
+        )

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)

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
+        )

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)