Add UsageLimits support to common.ai operators

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

@@ -230,6 +230,11 @@ Parameters
   every tool call is logged in real time. Default ``True``.
 - ``agent_params``: Additional keyword arguments passed to the pydantic-ai
   ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on every
+  agent run (initial run, durable replay, and HITL regeneration). Use it to
+  cap requests, tokens, or tool calls per task -- agents are particularly
+  prone to runaway tool loops, so ``tool_calls_limit`` is a useful guardrail.
+  See :ref:`howto/operator:llm` for an example. Default ``None``.
 - ``durable``: When ``True``, enables step-level caching of model responses and
   tool results via ObjectStorage. On retry, cached steps are replayed instead of
   re-executing expensive LLM calls. Requires the ``[common.ai] durable_cache_path``

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

@@ -66,6 +66,36 @@ the full list of supported parameters.
     :start-after: [START howto_operator_llm_agent_params]
     :end-before: [END howto_operator_llm_agent_params]
 
+Usage Limits
+------------
+
+Set ``usage_limits`` to a
+`pydantic-ai UsageLimits <https://ai.pydantic.dev/api/usage/#pydantic_ai.usage.UsageLimits>`__
+to fail the task when the run exceeds a configured budget — request count,
+input/output tokens, or tool calls. The check happens inside pydantic-ai's
+run loop, so the limit applies even when ``retries`` triggers multiple model
+calls within a single task.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+    :language: python
+    :start-after: [START howto_operator_llm_usage_limits]
+    :end-before: [END howto_operator_llm_usage_limits]
+
+Common knobs on ``UsageLimits``:
+
+- ``request_limit`` — max model requests per run (caps retry/tool-loop blow-ups).
+  pydantic-ai applies a default of ``50`` when ``UsageLimits()`` is constructed
+  without an explicit value, so passing ``UsageLimits(input_tokens_limit=4_000)``
+  silently inherits that 50-request cap. Set ``request_limit=None`` to disable
+  it explicitly when you only want a token cap.
+- ``input_tokens_limit`` / ``output_tokens_limit`` — per-run token caps.
+- ``total_tokens_limit`` — combined input + output cap.
+- ``tool_calls_limit`` — max tool invocations (``AgentOperator`` only).
+
+When the limit is hit pydantic-ai raises ``UsageLimitExceeded``, which
+propagates to Airflow as a task failure — Airflow's standard retry policy
+applies on top.
+
 TaskFlow Decorator
 ------------------
 
@@ -132,6 +162,8 @@ Parameters
   for structured output.
 - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent``
   constructor (e.g. ``retries``, ``model_settings``, ``tools``). Supports Jinja templating.
+- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on the run.  Fails
+  the task when token / request / tool-call budgets are exceeded.  Default ``None``.
 - ``require_approval``: If ``True``, the task defers after generating output and waits
   for human review.  Default ``False``.
 - ``approval_timeout``: Maximum time to wait for a review (``timedelta``).  ``None``

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py

@@ -21,6 +21,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
+from pydantic_ai.usage import UsageLimits
 
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.compat.sdk import dag, task
@@ -118,6 +119,30 @@ def extract(text: str):
 example_llm_decorator_structured()
 
 
+# [START howto_operator_llm_usage_limits]
+@dag
+def example_llm_operator_usage_limits():
+    LLMOperator(
+        task_id="capped_summary",
+        prompt="Summarize the attached design doc in three bullet points.",
+        llm_conn_id="pydanticai_default",
+        system_prompt="You are a concise technical reviewer.",
+        # Fail the task if the run exceeds 5 model requests, 4_000 input
+        # tokens, or 1_000 output tokens.  Useful for guardrails on shared
+        # connections or untrusted prompts.
+        usage_limits=UsageLimits(
+            request_limit=5,
+            input_tokens_limit=4_000,
+            output_tokens_limit=1_000,
+        ),
+    )
+
+
+# [END howto_operator_llm_usage_limits]
+
+example_llm_operator_usage_limits()
+
+
 # [START howto_operator_llm_approval]
 @dag
 def example_llm_operator_approval():

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

@@ -40,6 +40,7 @@
 if TYPE_CHECKING:
     from pydantic_ai import Agent
     from pydantic_ai.toolsets.abstract import AbstractToolset
+    from pydantic_ai.usage import UsageLimits
 
     from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
     from airflow.providers.common.ai.durable.storage import DurableStorage
@@ -102,6 +103,12 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
         arguments at DEBUG level. Set to ``False`` to disable.
     :param agent_params: Additional keyword arguments passed to the pydantic-ai
         ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+    :param usage_limits: Optional pydantic-ai
+        :class:`~pydantic_ai.usage.UsageLimits` enforced on every agent run
+        (initial run, durable replay, and HITL regeneration). Pass
+        ``UsageLimits(request_limit=..., total_tokens_limit=..., tool_calls_limit=..., ...)``
+        to fail the task when the agent exceeds the configured token, request,
+        or tool budget. ``None`` (default) means no enforcement.
     :param durable: When ``True``, enables step-level caching of model
         responses and tool results for durable execution.  On retry, cached
         steps are replayed instead of re-executing.  Default ``False``.
@@ -147,6 +154,7 @@ def __init__(
         toolsets: list[AbstractToolset] | None = None,
         enable_tool_logging: bool = True,
         agent_params: dict[str, Any] | None = None,
+        usage_limits: UsageLimits | None = None,
         durable: bool = False,
         # Agent feedback parameters
         enable_hitl_review: bool = False,
@@ -165,6 +173,7 @@ def __init__(
         self.toolsets = toolsets
         self.enable_tool_logging = enable_tool_logging
         self.agent_params = agent_params or {}
+        self.usage_limits = usage_limits
 
         self.durable = durable
 
@@ -246,9 +255,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)
+                result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         else:
-            result = agent.run_sync(self.prompt)
+            result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
 
         log_run_summary(self.log, result)
 
@@ -293,7 +302,7 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu
         """Re-run the agent with *feedback* appended to the conversation history."""
         agent = self._build_agent()
         messages = message_history or []
-        result = agent.run_sync(feedback, message_history=messages)
+        result = agent.run_sync(feedback, message_history=messages, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
 
         output = result.output

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

@@ -32,6 +32,7 @@
 
 if TYPE_CHECKING:
     from pydantic_ai import Agent
+    from pydantic_ai.usage import UsageLimits
 
     from airflow.sdk import Context
 
@@ -56,6 +57,11 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``tools``).
         See `pydantic-ai Agent docs <https://ai.pydantic.dev/api/agent/>`__
         for the full list.
+    :param usage_limits: Optional pydantic-ai
+        :class:`~pydantic_ai.usage.UsageLimits` enforced on the run. Pass
+        ``UsageLimits(request_limit=..., total_tokens_limit=..., ...)`` to fail
+        the task when the agent exceeds the configured token, request, or tool
+        budget. ``None`` (default) means no enforcement.
     :param require_approval: If ``True``, the task defers after generating
         output and waits for a human reviewer to approve or reject via the
         HITL interface.  Default ``False``.
@@ -83,6 +89,7 @@ def __init__(
         system_prompt: str = "",
         output_type: type = str,
         agent_params: dict[str, Any] | None = None,
+        usage_limits: UsageLimits | None = None,
         require_approval: bool = False,
         approval_timeout: timedelta | None = None,
         allow_modifications: bool = False,
@@ -95,6 +102,7 @@ def __init__(
         self.system_prompt = system_prompt
         self.output_type = output_type
         self.agent_params = agent_params or {}
+        self.usage_limits = usage_limits
         self.require_approval = require_approval
         self.approval_timeout = approval_timeout
         self.allow_modifications = allow_modifications
@@ -118,7 +126,7 @@ def execute(self, context: Context) -> Any:
         agent: Agent[None, Any] = self.llm_hook.create_agent(
             output_type=self.output_type, instructions=self.system_prompt, **self.agent_params
         )
-        result = agent.run_sync(self.prompt)
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
         output = result.output
 

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

@@ -81,7 +81,7 @@ def execute(self, context: Context) -> str | Iterable[str] | None:
             instructions=self.system_prompt,
             **self.agent_params,
         )
-        result = agent.run_sync(self.prompt)
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
         output = result.output
 

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

@@ -134,7 +134,7 @@ def execute(self, context: Context) -> Any:
             instructions=self._build_system_prompt(),
             **self.agent_params,
         )
-        result = agent.run_sync(request.user_content)
+        result = agent.run_sync(request.user_content, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
         output = result.output
 

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

@@ -309,7 +309,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)
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
 
         output_result = result.output.model_dump()

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

@@ -147,7 +147,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)
+        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
         log_run_summary(self.log, result)
         sql = self._strip_llm_output(result.output)
 

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

@@ -56,7 +56,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?")
+        mock_agent.run_sync.assert_called_once_with("Who is our top customer?", usage_limits=None)
 
     @pytest.mark.parametrize(
         "return_value",
@@ -92,7 +92,7 @@ def my_prompt(topic):
         op.execute(context={"task_instance": MagicMock()})
 
         assert op.prompt == "Analyze revenue trends"
-        mock_agent.run_sync.assert_called_once_with("Analyze revenue trends")
+        mock_agent.run_sync.assert_called_once_with("Analyze revenue trends", usage_limits=None)
 
     @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
     def test_execute_passes_toolsets_through(self, mock_hook_cls):

providers/common/ai/tests/unit/common/ai/decorators/test_llm.py

@@ -54,7 +54,7 @@ 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")
+        mock_agent.run_sync.assert_called_once_with("Summarize this text", usage_limits=None)
 
     @pytest.mark.parametrize(
         "return_value",
@@ -90,4 +90,4 @@ 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")
+        mock_agent.run_sync.assert_called_once_with("Summarize quantum computing", usage_limits=None)

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

@@ -66,7 +66,7 @@ def my_prompt():
 
         assert result == "positive"
         assert op.prompt == "Route this review"
-        mock_agent.run_sync.assert_called_once_with("Route this review")
+        mock_agent.run_sync.assert_called_once_with("Route this review", usage_limits=None)
         mock_do_branch.assert_called_once()
 
     @pytest.mark.parametrize(

providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py

@@ -71,7 +71,7 @@ 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")
+        mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None)
 
     @pytest.mark.parametrize(
         "return_value",
@@ -115,4 +115,4 @@ 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")
+        mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None)

providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py

@@ -54,7 +54,7 @@ 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")
+        mock_agent.run_sync.assert_called_once_with("Get all users", usage_limits=None)
 
     @pytest.mark.parametrize(
         "return_value",

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

@@ -21,6 +21,7 @@
 
 import pytest
 from pydantic import BaseModel
+from pydantic_ai.usage import UsageLimits
 
 from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink
 from airflow.providers.common.ai.toolsets.logging import LoggingToolset
@@ -79,6 +80,44 @@ def test_template_fields(self):
 
 
 class TestAgentOperatorExecute:
+    @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+    def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls):
+        """``usage_limits`` is forwarded to ``agent.run_sync`` on the non-durable path."""
+        mock_agent = _make_mock_agent("ok")
+        mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
+
+        limits = UsageLimits(request_limit=3, tool_calls_limit=5)
+        op = AgentOperator(
+            task_id="test",
+            prompt="run",
+            llm_conn_id="my_llm",
+            usage_limits=limits,
+        )
+        op.execute(context=MagicMock())
+
+        mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits)
+
+    @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+    def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls):
+        """``usage_limits`` is also forwarded by ``regenerate_with_feedback``."""
+        mock_agent = _make_mock_agent("revised")
+        mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
+
+        limits = UsageLimits(request_limit=1)
+        op = AgentOperator(
+            task_id="test",
+            prompt="run",
+            llm_conn_id="my_llm",
+            usage_limits=limits,
+        )
+        op.regenerate_with_feedback(feedback="Add detail", message_history=[])
+
+        mock_agent.run_sync.assert_called_once_with(
+            "Add detail",
+            message_history=[],
+            usage_limits=limits,
+        )
+
     @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
     def test_execute_creates_agent_from_hook(self, mock_hook_cls):
         mock_agent = _make_mock_agent("The answer is 42.")
@@ -97,7 +136,7 @@ def test_execute_creates_agent_from_hook(self, mock_hook_cls):
         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?")
+        mock_agent.run_sync.assert_called_once_with("What is the answer?", usage_limits=None)
 
     @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
     def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls):
@@ -378,6 +417,7 @@ def test_regenerate_with_feedback_calls_agent_with_feedback_and_history(self, mo
         mock_agent.run_sync.assert_called_once_with(
             "Add more detail",
             message_history=msg_history,
+            usage_limits=None,
         )
 
     @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
@@ -462,4 +502,4 @@ def test_execute_non_durable_does_not_wrap(self, mock_hook_cls):
         op.execute(context=MagicMock())
 
         # run_sync called directly, no override
-        mock_agent.run_sync.assert_called_once_with("test")
+        mock_agent.run_sync.assert_called_once_with("test", usage_limits=None)