You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2026/09/22 08:30:39 UTC
(airflow) branch main updated: Reorganize `common.ai` provider docs into topic-based navigation (#73523)
This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 90107b5081a Reorganize `common.ai` provider docs into topic-based navigation (#73523)
90107b5081a is described below
commit 90107b5081af5cf83dc406f1283fc41eb6d5fc81
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Tue Sep 22 09:30:29 2026 +0100
Reorganize `common.ai` provider docs into topic-based navigation (#73523)
The provider docs had a single flat "Guides" list of 20 entries, with the
seven connection pages sitting between the quick start and everything else
and the operators at position 13. Several features (durable execution,
guardrails, code mode, message history, approval gates, structured output)
existed only as subsections of the AgentOperator and LLMOperator pages, and
three pages exceeded 600 lines.
The sidebar is now grouped into captions: Getting started, Models and
providers, Operators, Toolsets, Running agents, Document and RAG pipelines,
Reliability and operations, Examples and References. The toolsets and
choosing-a-toolset pages become one overview plus a page per toolset, the
sandbox page splits into overview, configuration and backends, and the
agent and LLM operator feature sections become pages of their own. New
pages cover installation, core concepts and troubleshooting. A redirects
file covers the removed pages, and the AgentOperator page keeps a durable
execution heading because pydantic-ai's docs deep-link that anchor.
---
providers/common/ai/docs/agent_security.rst | 267 +++++
providers/common/ai/docs/approval_gates.rst | 118 ++
providers/common/ai/docs/choosing_a_toolset.rst | 611 -----------
providers/common/ai/docs/code_mode.rst | 87 ++
providers/common/ai/docs/concepts.rst | 135 +++
providers/common/ai/docs/connections/langchain.rst | 2 +-
.../common/ai/docs/connections/llamaindex.rst | 4 +-
providers/common/ai/docs/connections/mcp.rst | 2 +-
.../common/ai/docs/connections/pydantic_ai.rst | 2 +-
.../ai/docs/connections/pydantic_ai_azure.rst | 4 +-
.../ai/docs/connections/pydantic_ai_bedrock.rst | 4 +-
.../ai/docs/connections/pydantic_ai_vertex.rst | 4 +-
providers/common/ai/docs/durable_execution.rst | 159 +++
providers/common/ai/docs/end_to_end_pipelines.rst | 4 +-
providers/common/ai/docs/examples.rst | 6 -
providers/common/ai/docs/guardrails.rst | 62 ++
providers/common/ai/docs/hitl_review.rst | 38 +-
providers/common/ai/docs/hooks/index.rst | 62 --
providers/common/ai/docs/hooks/langchain.rst | 6 +
providers/common/ai/docs/hooks/mcp.rst | 2 +-
providers/common/ai/docs/hooks/pydantic_ai.rst | 10 +-
providers/common/ai/docs/index.rst | 192 ++--
providers/common/ai/docs/installation.rst | 86 ++
providers/common/ai/docs/message_history.rst | 60 +
providers/common/ai/docs/operators/agent.rst | 343 +-----
providers/common/ai/docs/operators/index.rst | 49 +-
providers/common/ai/docs/operators/llm.rst | 147 +--
providers/common/ai/docs/operators/llm_batch.rst | 2 +
providers/common/ai/docs/quickstart.rst | 10 +-
providers/common/ai/docs/redirects.txt | 22 +
providers/common/ai/docs/retry_policies.rst | 9 +-
providers/common/ai/docs/sandbox.rst | 821 --------------
providers/common/ai/docs/sandbox/backends.rst | 221 ++++
providers/common/ai/docs/sandbox/configuration.rst | 305 ++++++
providers/common/ai/docs/sandbox/index.rst | 425 ++++++++
providers/common/ai/docs/structured_output.rst | 68 ++
providers/common/ai/docs/toolsets.rst | 1145 --------------------
providers/common/ai/docs/toolsets/datafusion.rst | 166 +++
providers/common/ai/docs/toolsets/hook.rst | 98 ++
providers/common/ai/docs/toolsets/index.rst | 269 +++++
providers/common/ai/docs/toolsets/langchain.rst | 81 ++
providers/common/ai/docs/toolsets/logging.rst | 36 +
.../common/ai/docs/toolsets/managed_agent.rst | 291 +++++
providers/common/ai/docs/toolsets/mcp.rst | 166 +++
providers/common/ai/docs/toolsets/skills.rst | 169 +++
providers/common/ai/docs/toolsets/sql.rst | 225 ++++
providers/common/ai/docs/troubleshooting.rst | 143 +++
47 files changed, 3888 insertions(+), 3250 deletions(-)
diff --git a/providers/common/ai/docs/agent_security.rst b/providers/common/ai/docs/agent_security.rst
new file mode 100644
index 00000000000..f58a79aeb86
--- /dev/null
+++ b/providers/common/ai/docs/agent_security.rst
@@ -0,0 +1,267 @@
+ .. 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.
+
+Securing agent tools
+====================
+
+LLM agents call tools based on natural-language reasoning. This makes them
+powerful but introduces risks that don't exist with deterministic operators.
+
+What the agent can and cannot reach
+-------------------------------------
+
+An agent's reach is exactly the set of tools you register on it, and nothing
+more. The model never executes arbitrary code: it can only request one of the
+tools you provided, and pydantic-ai rejects any tool name outside that set
+before it runs. If no registered tool can read the environment, the
+filesystem, or other connections, the model cannot reach them, regardless of
+what the prompt instructs it to do.
+
+This is what "untrusted" means in this context. The Dag file itself is
+author-written and trusted, exactly like any other Dag. What is untrusted is
+the model's *output*: the tool-call requests and text it generates. That output
+is confined to your registered tools and bounded by the tool-call budget. An
+agent cannot create a new connection, read another connection's credentials, or
+run a shell command unless a tool you registered exposes that capability.
+
+The corollary is that every tool you add widens the blast radius, and a custom
+toolset is only as safe as you make it. A tool that returns ``os.environ`` or
+runs shell commands hands the model whatever that tool can reach. Audit any
+custom toolset, and any MCP server you connect through ``MCPToolset``, against
+the same standard the bundled toolsets below are built to.
+
+.. _toolset-defense-layers:
+
+Defense layers
+--------------
+
+No single layer is sufficient — they work together.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 40 40
+
+ * - Layer
+ - What it does
+ - What it does NOT do
+ * - **Airflow Connections**
+ - Credentials are stored in Airflow's secret backend, never in Dag code.
+ The LLM agent cannot see API keys or database passwords.
+ - Does not prevent the agent from using the connection to access data
+ the connection has access to.
+ * - **HookToolset: explicit allow-list**
+ - Only methods listed in ``allowed_methods`` are exposed as tools.
+ Auto-discovery is not supported. Methods are validated at Dag parse
+ time.
+ - Does not restrict what arguments the agent passes to allowed methods.
+ * - **SQLToolset: read-only by default**
+ - ``allow_writes=False`` (default) validates every SQL query through
+ ``validate_sql()``: SELECT-family and read-only metadata
+ (``DESCRIBE``/``SHOW``) statements pass; INSERT, UPDATE, DELETE, DROP,
+ and writes hidden behind ``EXPLAIN`` are rejected.
+ - Does not prevent the agent from reading sensitive data that the
+ database user has SELECT access to.
+ * - **DataFusionToolset: read-only by default**
+ - ``allow_writes=False`` (default) validates every SQL query through
+ ``validate_sql()`` and rejects CREATE TABLE, CREATE VIEW, INSERT
+ INTO, and other non-SELECT statements.
+ - Does not prevent the agent from reading any registered data source.
+ * - **SQLToolset: allowed_tables**
+ - Restricts the agent to listed tables across ``list_tables``,
+ ``get_schema``, ``query``, and ``check_query``. Queries are parsed and
+ every referenced table (including via subqueries, CTEs, JOINs, and
+ ``DESCRIBE``) is checked against the list before execution.
+ - Rejects ``COPY`` and every function sqlglot cannot type (the channel for
+ ``pg_read_file`` / ``query_to_xml`` / ``dblink``) unless named in
+ ``allowed_functions``. Fail-closed, but only as exact as the SQL parser. Not a
+ security boundary -- always pair it with least-privilege database grants. See
+ :ref:`allowed-tables-enforcement` below.
+ * - **SQLToolset: max_rows / max_result_bytes**
+ - Bounds a query result by rows (default 50) and by serialized size
+ (default 64 KiB), preventing the agent from pulling entire tables into
+ context.
+ - Does not limit the number of queries the agent can make, and each result
+ stays in message history for the rest of the run. Rows past ``max_rows``
+ are not read out of the cursor, but a client-buffering driver has already
+ transferred them -- this bounds context, not database or network load.
+ * - **MCPToolset: external server**
+ - Connects the agent to tools exposed by an MCP server, authenticated
+ through an Airflow connection.
+ - Does **not** constrain what those tools do. An MCP server can expose
+ shell, filesystem, or network access. Run only trusted servers and
+ audit the tools they expose.
+ * - **SandboxToolset: off-worker execution**
+ - Runs the agent's commands and file operations in a disposable microVM,
+ never in the worker process. Airflow injects nothing of its own; only
+ what ``SandboxSpec`` names goes in, and egress is denied by default. A
+ backend that cannot enforce a spec field raises rather than ignoring it.
+ Commands are bounded by a timeout ceiling and per-stream output caps.
+ - **Does not contain the agent.** The agent loop and every other toolset on
+ the same agent still run in the worker with its credentials, so this does
+ 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
+ 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
+ layer and leaves DNS resolution open.
+ * - **pydantic-ai: tool call budget**
+ - pydantic-ai's ``max_result_retries`` and ``model_settings`` control
+ how many tool-call rounds the agent can make before stopping.
+ - Requires explicit configuration — the default allows many rounds.
+
+.. _allowed-tables-enforcement:
+
+How ``allowed_tables`` is enforced
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+When ``allowed_tables`` is set it governs every tool, not just discovery:
+
+- ``list_tables`` and ``get_schema`` only reveal listed tables.
+- ``query`` and ``check_query`` parse the SQL with `sqlglot
+ <https://github.com/tobymao/sqlglot>`_ and reject it before execution if it
+ references any table that is not on the list. Tables reached indirectly are
+ caught too -- through subqueries, CTEs, JOINs, set operations (``UNION`` etc.),
+ ``DESCRIBE``, catalog views such as ``information_schema``, and DML. CTE
+ references are excluded by lexical scope, so a same-named CTE in another scope
+ cannot hide a real table, and the database/catalog is part of the match, so a
+ cross-database reference like ``otherdb.public.orders`` is refused.
+- Constructs the list cannot describe are rejected outright while it is active:
+ table-valued functions (``dblink``), ``TABLE('name')`` row sources, the
+ ``TABLE <name>`` shorthand, ``SHOW``, dynamic SQL (``EXEC``), ``COPY``
+ (file/program I/O), and **inline comments** -- because parser-vs-engine differences
+ hide in comments (MySQL executes ``/*! ... */`` while sqlglot and other engines
+ ignore it).
+- **Any function sqlglot does not recognize is rejected (fail-closed).** A function
+ whose string argument reaches data outside the table graph --
+ ``pg_read_file('/etc/passwd')`` (a file), ``query_to_xml('SELECT * FROM other_table', ...)``
+ (SQL over another table), a scalar ``dblink`` (a remote database) -- carries no table
+ reference for the parser to catch. Rather than maintain a denylist of such functions
+ (unbounded, engine-specific, and it would fail *open* on anything missed), the toolset
+ rejects every function sqlglot cannot type. Ordinary builtins (``count``, ``lower``,
+ ``sum``) are recognized and pass. A legitimate function sqlglot does not type
+ (``json_build_object``, ``jsonb_agg``) or a project UDF is rejected until you list it
+ in ``allowed_functions``:
+
+ .. code-block:: python
+
+ SQLToolset(
+ db_conn_id="analytics_db",
+ allowed_tables=["orders"],
+ allowed_functions=["json_build_object"], # opt in per function you trust
+ )
+
+So ``SELECT * FROM secrets`` with ``allowed_tables=["orders"]`` is refused, and
+the rejection is handed back to the agent so it can re-target an allowed table.
+
+.. warning::
+
+ This is a strong **application-level guardrail, not a security boundary.** The
+ fail-closed function check raises the bar, but any query the engine parses
+ differently from sqlglot is a residual gap, and ``allowed_functions`` is a trust
+ decision you own. **Always** point the connection at a least-privilege database
+ role -- that is the boundary that holds even when the parser cannot see through a
+ function, and it is what actually keeps an agent (which may be under prompt
+ injection) away from data and files you have not granted it:
+
+ .. code-block:: sql
+
+ -- Create a read-only role with access to specific tables only
+ CREATE ROLE airflow_agent_reader;
+ GRANT SELECT ON orders, customers TO airflow_agent_reader;
+ -- Use this role's credentials in the Airflow connection
+
+Defense in depth: the allow-list contains the agent's *intent* (and gives it a
+correctable error), while the database role is the boundary that holds even if
+the agent reaches data the parser cannot see. The connection should use a
+database user with the minimum privileges required.
+
+``HookToolset`` guidelines
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- List only the methods the agent needs. Never expose ``run()`` or
+ ``get_connection()`` — these give broad access.
+- Prefer read-only methods (``list_*``, ``get_*``, ``describe_*``).
+- The agent controls arguments. If a method accepts a ``path`` parameter,
+ the agent can pass any path the hook has access to.
+
+.. code-block:: python
+
+ # Good: expose only list and read
+ HookToolset(
+ s3_hook,
+ allowed_methods=["list_keys", "read_key"],
+ tool_name_prefix="s3_",
+ )
+
+ # Bad: exposes delete and write operations
+ HookToolset(
+ s3_hook,
+ allowed_methods=["list_keys", "read_key", "delete_object", "load_string"],
+ )
+
+Recommended configuration
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+**Read-only analytics** (the most common pattern):
+
+.. code-block:: python
+
+ SQLToolset(
+ db_conn_id="analytics_readonly", # Connection with SELECT-only grants
+ allowed_tables=["orders", "customers"], # Hide other tables from agent
+ allow_writes=False, # Default — validates SQL
+ max_rows=50, # Default — cap rows
+ max_result_bytes=65536, # Default — cap bytes; lower it for wide tables
+ )
+
+**Agents that need to modify data** (use with caution):
+
+.. code-block:: python
+
+ SQLToolset(
+ db_conn_id="app_db",
+ allowed_tables=["user_preferences"],
+ allow_writes=True, # Disables SQL validation — agent can INSERT/UPDATE
+ max_rows=100,
+ )
+
+Production checklist
+^^^^^^^^^^^^^^^^^^^^
+
+Before deploying an agent task to production:
+
+1. **Connection credentials**: Use Airflow's secret backend. Never hardcode
+ API keys in Dag files.
+2. **Database permissions**: Create a dedicated database user with minimum
+ required grants. Don't reuse the admin connection.
+3. **Tool allow-list**: Review ``allowed_methods`` / ``allowed_tables``. The
+ agent can call any exposed tool with any arguments.
+4. **Read-only default**: Keep ``allow_writes=False`` unless the task
+ specifically requires writes.
+5. **Result limits**: Set ``max_rows`` and ``max_result_bytes`` appropriate to
+ the use case. ``max_rows`` alone does not bound size -- on wide tables it is
+ ``max_result_bytes`` that keeps a result from dominating the context window for
+ the rest of the run.
+6. **Model budget**: Configure pydantic-ai's ``model_settings`` (e.g.
+ ``max_tokens``) and ``retries`` to bound cost and prevent runaway loops.
+7. **System prompt**: Include safety instructions in ``system_prompt`` (e.g.
+ "Only query tables related to the question. Never modify data.").
+8. **Prompt injection**: Be cautious when the prompt includes untrusted data
+ (user input, external API responses, upstream XCom). Consider sanitizing
+ inputs before passing them to the agent.
diff --git a/providers/common/ai/docs/approval_gates.rst b/providers/common/ai/docs/approval_gates.rst
new file mode 100644
index 00000000000..411a13bd022
--- /dev/null
+++ b/providers/common/ai/docs/approval_gates.rst
@@ -0,0 +1,118 @@
+ .. 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.
+
+.. _approval-gates:
+
+Approval gates for LLM operators
+================================
+
+.. seealso::
+ This page covers the one-shot approve, edit or reject gate on ``LLMOperator`` and its
+ subclasses. For a multi-round review loop on ``AgentOperator`` with a chat UI and REST
+ API, see :doc:`hitl_review`.
+
+Set ``require_approval=True`` to pause the task after the LLM generates its
+output and wait for a human reviewer to approve or reject it via the Airflow
+HITL interface. Optionally allow the reviewer to edit the output before
+approving with ``allow_modifications=True``, and set a deadline with
+``approval_timeout``.
+
+Human-in-the-loop review needs Airflow 3.1+, whether ``require_approval`` or a
+``decision_policy`` with ``on_uncertain="review"`` opens it. On an older core the
+operator raises ``AirflowOptionalProviderFeatureException`` when it is constructed, so the Dag file
+fails to import, and with it every Dag defined in that file. A dynamically mapped
+task (``.expand()``) is only constructed when it runs, so there the same error
+surfaces as a task failure -- still before the model is called.
+
+When ``approval_timeout`` expires without a review, the task fails by default.
+Set ``on_approval_timeout="approve"`` to return the generated output instead, so
+an unattended pipeline keeps moving. ``"reject"`` answers the review with a
+rejection, which still fails this operator; only
+:class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator`
+turns a rejection into a downstream skip. The chosen option is also
+pre-highlighted as the default in the review form, so ``"reject"`` makes
+Reject the primary button:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+ :language: python
+ :start-after: [START howto_operator_llm_approval]
+ :end-before: [END howto_operator_llm_approval]
+
+A pending review is not surfaced as a notification. Pass
+``approval_notifiers`` to tell the reviewers about it through any Airflow
+notifier (Slack, email, ...), the way
+:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
+``notifiers``. The notifiers run once the review is open and can reference
+the review ``{{ task.subject }}`` and ``{{ task.body }}`` in their templates,
+as the example above does. The ``@task.llm`` decorator and the operator
+subclasses accept the same parameter. A notifier whose delivery fails is
+logged and the task still waits for the review; a template error fails the
+task. A retry re-runs the LLM and re-notifies with the regenerated output,
+while the open review keeps the original subject and body.
+
+The default ``body`` contains the rendered prompt and the output. Where either
+is sensitive, template only ``{{ task.subject }}`` and a link to the review
+into channels outside Airflow's auth boundary.
+
+By default any user with the permission can answer the review. Pass
+``approval_assigned_users=[{"id": "<auth-manager-user-id>", "name": "<user-name>"}]``
+to restrict it to named reviewers, the way
+:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
+``assigned_users``. ``id`` is the user id reported by the auth manager: with
+the default ``SimpleAuthManager`` it is the username from
+``simple_auth_manager_users``; under the FAB auth manager it is the numeric
+user row id as a string, not the username. This needs Airflow 3.1+. On Airflow 3.1.0 through 3.1.5 both
+``id`` and ``name`` must match what the auth manager reports, so a wrong
+``name`` blocks the assigned reviewer as well as everyone else; from 3.1.6 only
+``id`` is compared. The list is stored when the review is first created:
+clearing the task re-runs it against the existing review row, so a changed
+list does not take effect.
+
+Reviewing uncertain output
+--------------------------
+
+A classifier model such as TypeSafe's reports a confidence for every field
+of a structured output, in ``provider_details`` on the model response. It is a
+summary of how concentrated the model's probability distribution was, not the
+probability that the field is right. ``decision_policy=DecisionPolicy(min_confidence=0.7)``
+(import ``DecisionPolicy`` from ``airflow.providers.common.ai.operators.llm``)
+is the bar the least confident field has to clear for the operator to return
+the output by itself. Only fields that reported a confidence are compared: a
+field whose type reports none (a bounded float, where the probability is the
+answer) is not gated, and the record's ``confidence`` map shows which fields
+were. Under the bar, the policy's ``on_uncertain`` applies: ``"review"``
+(default) sends the output to human review through the same approval flow as
+``require_approval``, with the same ``approval_timeout``,
+``on_approval_timeout`` and notifier settings; ``"fail"`` fails the task with
+``LowConfidenceError`` (from ``airflow.providers.common.ai.exceptions``), which
+Airflow retries like any other failure unless a retry rule says otherwise. A
+text model reports no confidence for any field, which counts as uncertain, so
+switching the connection does not silently switch off a control you set.
+``require_approval=True`` keeps its meaning and always asks, and
+``on_uncertain="review"`` needs Airflow 3.1+ like it does. The policy is
+honoured by ``LLMOperator`` and ``LLMBranchOperator``; the SQL, schema-compare
+and file-analysis operators run their own ``execute`` and reject a policy with a
+bar at construction.
+
+With or without a bar, the operator pushes a ``decision`` XCom carrying the
+model name, the per-field ``confidence`` and ``probabilities`` (empty for a text
+model), the bar that applied, why the output went to review if it did, who
+decided, and the gate configuration in force.
+A pending record is checkpointed with the paused task and finalized from that
+copy on resume, not from the XCom. See :ref:`LLMBranchOperator <howto/operator:llm_branch>` for the
+record's fields; there ``proposed`` and ``action`` name the branches, while
+here they are ``null`` and the output itself is the return value.
diff --git a/providers/common/ai/docs/choosing_a_toolset.rst b/providers/common/ai/docs/choosing_a_toolset.rst
deleted file mode 100644
index cc7f8a54a2a..00000000000
--- a/providers/common/ai/docs/choosing_a_toolset.rst
+++ /dev/null
@@ -1,611 +0,0 @@
- .. 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.
-
-.. _howto/choosing-a-toolset:
-
-Choosing a Toolset
-==================
-
-:doc:`toolsets` documents how to configure each toolset. This page answers the
-question that comes before that one: you have a system you want an agent to
-reach, so which route do you take, and what does each route give up?
-
-Read the table below by what you already have, not by what a toolset is called.
-When two routes both work, the deciding factor is rarely what each one can do —
-it is what each one cannot do, and every route has a short list.
-
-More than one row can be true at once, and the rows are not exclusive: one agent
-can carry several toolsets. Two questions break the ties. *Whose credential is
-it?* — prefer the route whose credential is an Airflow connection somebody on
-your side already reviewed. *Whose tool list is it?* — prefer the route whose
-exposed surface you chose rather than inherited. The pair that most often
-overlaps is an Airflow hook and a vendor MCP server reaching the same target;
-both questions point at the hook, because its credential is the connection and
-``allowed_methods`` is a list you write. Reach for the server when its tools
-cover work the hook does not expose, or when the alternative is re-wrapping that
-API by hand.
-
-Those two questions do not separate ``HookToolset`` from ``SQLToolset`` when the
-target is a DBAPI database, because both answer them the same way. A third one
-does: *is the work a fixed operation or an open-ended question?* A named method
-you can enumerate in advance is a hook. A question the agent has to express as
-SQL is a query, and ``SQLToolset`` answers it with schema discovery, bounded
-results and an ``allowed_tables`` walk you can switch on, none of which
-``HookToolset`` has an equivalent of.
-
-Start with what you have
-------------------------
-
-.. list-table::
- :widths: 50 50
- :header-rows: 1
-
- * - What you have
- - Route
- * - A target that already has an Airflow connection, and a hook method that
- already does the thing
- - ``HookToolset``
- * - A question that is a query, against a DBAPI database
- - ``SQLToolset``
- * - Files on an object store — Parquet, CSV, Avro — or a catalog-managed
- table format such as Iceberg, rather than rows in a database
- - ``DataFusionToolset``
- * - A vendor that already ships a server built for agents, whose tools you
- would otherwise re-wrap by hand
- - ``MCPToolset``
- * - Procedural knowledge — how to carry out a task — rather than an endpoint
- to call
- - ``AgentSkillsToolset``
- * - Work that means running code the model wrote, not calling a tool you chose
- - ``SandboxToolset``
- * - Reasoning that should happen on the vendor's own infrastructure
- - A subclass of ``BaseManagedAgentToolset`` that you write
-
-The rest of this page takes those seven in turn. Each entry gives the case for
-choosing it, what it cannot do, an example that exists in this repository, and
-where its credentials and its work come from.
-
-``HookToolset``
----------------
-
-**Choose it when** the target already has an Airflow connection and a hook, and
-what you want the agent to do is already a method on that hook. This is the
-cheapest route — no new server, no new credential, no new query dialect — and
-the only one that reaches any provider hook with synchronous methods without
-anyone writing an adapter first. :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` is a
-reflection-based adapter, so the work is choosing the method list.
-
-**What it cannot do**
-
-- It allow-lists method *names*, not arguments. Once ``read_key`` is exposed,
- the agent picks the key; the :ref:`defense-layer table <toolset-defense-layers>`
- states this outright. Choose methods whose worst case you accept, not methods
- you intend to constrain later.
-- Its calls act as barriers. The tools are registered with ``sequential=True``
- because hook methods perform synchronous I/O, so a slow call holds up every
- other tool the model emitted in that step, not only this toolset's. This is
- not specific to ``HookToolset`` — see :ref:`toolset-call-barriers`.
-- It returns exactly one shape. Every result goes through ``serialize_for_llm``
- and comes back as a JSON-encoded string; there is no structured error type and
- no ``ModelRetry`` wrapper, so a hook exception fails the agent run, and the
- task with it, instead of giving the model something it can correct.
- ``SQLToolset``, by contrast, hands the database's own error back as a retry.
-- Its ``call_tool`` calls the method and serializes what comes back. The code
- contains no path that awaits a coroutine result, and none that checks for one,
- so an ``async def`` hook method is not a case this adapter is written to
- handle. Treat synchronous methods as the supported set.
-
-**A real example.** The read-only S3 pair from the ``HookToolset`` guidance in
-:doc:`toolsets`:
-
-.. code-block:: python
-
- HookToolset(
- s3_hook,
- allowed_methods=["list_keys", "read_key"],
- tool_name_prefix="s3_",
- )
-
-**Credentials and where it runs.** The hook instance is yours, so the credential
-is whatever connection that hook resolves — the toolset never looks one up
-itself. Calls run in the Airflow worker process.
-
-``SQLToolset``
---------------
-
-**Choose it when** the question is a query and the data is in a DBAPI database.
-:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` gives the agent
-four tools — list tables, get schema, query, check query. Set ``allowed_tables``
-and that allow-list is enforced by parsing the SQL rather than by matching
-strings; see :ref:`allowed-tables-enforcement` for how the walk handles CTEs,
-subqueries and joins.
-
-**What it cannot do**
-
-- ``allowed_tables`` is an application-level guardrail, not a replacement for
- database permissions. Its own docstring says so, and names the residual gap:
- an engine or query the parser reads differently. Point ``db_conn_id`` at a
- least-privilege role whose grants match the allow-list.
-- It cannot bound the fetch for every driver. Hooks that hand their handler
- something other than a DBAPI cursor — ``ExasolHook`` and its pyexasol
- statement, for instance — fall back to a full fetch. The payload handed to the
- model is still bounded; the transfer is not. See :ref:`bounded-query-results`.
-- Its parser-level closure is opt-in, not the default. The table walk returns
- immediately while ``allowed_tables`` is unset, so out of the box the agent
- reaches every table the connection can see. Only omitting the argument grants
- that: an explicit ``None`` or empty list is rejected at construction.
- ``DESCRIBE`` and ``SHOW`` both pass, on dialects that parse them, while
- ``allowed_tables`` stays unset. Set ``allowed_tables`` and the walk
- turns fail-closed for ``SHOW`` — but not for ``DESCRIBE``: it instead
- becomes an ordinary table reference, allowed only when the table it names
- is on the list. See :ref:`allowed-tables-enforcement` for what else the
- walk rejects once ``allowed_tables`` is active. Statements that modify
- data, ``COPY`` among them, are rejected either way while ``allow_writes``
- is ``False``.
-- It does not classify failures. A connection error or a typo in a column
- name reaching ``list_tables``, ``get_schema`` or ``query`` becomes one
- ``ModelRetry``, so the two are treated the same way until the retry budget
- runs out and the task fails for Airflow to retry. Two paths do not raise:
- ``check_query`` catches its own errors and reports them back as a normal
- ``{"valid": false, ...}`` result, and ``get_schema`` returns a normal
- ``{"error": ...}`` result instead of raising when the requested table is
- outside ``allowed_tables`` — other ``get_schema`` failures still raise and
- still become a ``ModelRetry``.
-
-**A real example.** ``example_pydantic_ai_hook.py`` builds an agent around
-``SQLToolset`` inside a plain ``@task`` function, with no operator involved:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
- :language: python
- :start-after: [START howto_task_with_toolsets]
- :end-before: [END howto_task_with_toolsets]
-
-**Credentials and where it runs.** ``db_conn_id`` is resolved through
-``BaseHook.get_connection``, and the connection must supply a ``DbApiHook``.
-Queries run from the Airflow worker process against the database. Its tool
-calls act as barriers, as they do for the other routes that build their own
-tools; see :ref:`toolset-call-barriers`.
-
-``DataFusionToolset``
----------------------
-
-**Choose it when** the data is files on an object store rather than rows in a
-database — Parquet, CSV or Avro — or a table in a catalog such as Iceberg, and
-you want the agent to ask SQL questions of them without loading them anywhere
-first. (This route needs the ``datafusion`` extra of
-``apache-airflow-providers-common-sql``.) Each ``DataSourceConfig`` registers
-one table, and several can be registered so the agent can join across them.
-The two shapes take different fields: an object-store format needs a
-``uri``, while a catalog format like
-Iceberg is looked up by ``db_name`` instead, and ``DataSourceConfig`` raises
-``ValueError`` at construction if a catalog format is missing one.
-
-**What it cannot do**
-
-- It has no table allow-list. ``allow_writes=False`` is the only guard, and it
- blocks non-SELECT statements, not reach: the defense-layer table records that
- this toolset "does not prevent the agent from reading any registered data
- source". The registration list is therefore the whole boundary — register
- exactly what the agent may read.
-- It bounds the payload, not the scan. DataFusion has already materialized the
- full result before ``max_rows`` and ``max_result_bytes`` apply, so those limits
- protect the model's context, not the cost of the query.
-- It cannot tell failure kinds apart precisely. The DataFusion Python bindings
- expose no native exception types, so the retry decision is made by matching the
- error message against regular expressions — which a wording change upstream can
- quietly defeat.
-
-**A real example.** The same bucket as the ``HookToolset`` entry above, reached
-the other way. Rather than exposing ``list_keys`` and ``read_key`` and leaving
-the agent to reassemble files, this registers the prefix as a table and lets it
-write SQL:
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
- from airflow.providers.common.sql.config import DataSourceConfig
-
- toolset = DataFusionToolset(
- datasource_configs=[
- DataSourceConfig(
- conn_id="aws_default",
- table_name="sales",
- uri="s3://my-bucket/data/sales/",
- format="parquet",
- ),
- ],
- max_rows=100,
- )
-
-Which of the two fits depends on the question. "Read me this object" is a hook
-method. "What were last quarter's returns by region" is a query, and expressing
-it through ``list_keys`` and ``read_key`` means the model does the aggregation in
-its context window instead of the engine doing it.
-
-An Iceberg table is registered differently. (Iceberg support needs the
-``apache.iceberg`` extra of ``apache-airflow-providers-common-sql``; without
-it, registration raises ``AirflowOptionalProviderFeatureException``.) There is
-no ``uri`` to read files from; the catalog resolves the table by name, so the
-config carries a ``db_name`` instead, following the same ``DataSourceConfig``
-shape that ``example_analytics.py`` in the ``common.sql`` provider uses:
-
-.. code-block:: python
-
- toolset = DataFusionToolset(
- datasource_configs=[
- DataSourceConfig(
- conn_id="iceberg_default",
- table_name="users_data",
- db_name="demo",
- format="iceberg",
- ),
- ],
- max_rows=100,
- )
-
-**Credentials and where it runs.** Each ``DataSourceConfig`` carries its own
-``conn_id``, so object-store access is an Airflow connection. DataFusion is an
-embedded engine: the query runs inside the worker process, not on a remote
-cluster. Its tool calls act as barriers, as they do for the other routes that
-build their own tools; see :ref:`toolset-call-barriers`.
-
-``MCPToolset``
---------------
-
-**Choose it when** someone already publishes a server built for agents that
-covers your target. You inherit a tool surface that was designed to be called by
-a model — retry semantics and error wording are decided upstream, and a
-destructive tool can simply be absent — instead of maintaining a per-API wrapper
-yourself.
-
-**What it cannot do**
-
-- Its allow-list is client-side and opt-in, not server-side and required.
- :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` forwards
- ``get_tools`` and ``call_tool`` straight to the underlying server, so
- whatever the server exposes, the agent gets by default. Because
- ``MCPToolset`` is itself built on the same ``AbstractToolset`` base every
- toolset in this provider extends, a Dag author can call ``.filtered()`` to
- subset the advertised tool list using a filter function that inspects each
- tool's definition. That filtering happens on the client: it narrows what
- the agent is offered, it does not revoke or authorize anything on the
- server, and unlike ``allowed_methods`` on ``HookToolset``, which is required
- and rejects an empty list, nothing here requires you to set a filter. The
- defense-layer table is explicit that a server can expose shell, filesystem
- or network access.
-- It cannot guarantee the credential came from a connection. ``mcp_conn_id`` is
- the default path, but ``token_provider`` and ``env_provider`` are your own
- callables and are free to read an environment variable, a file, or an entirely
- different secret store. That is the point of them — it also means the
- connection is no longer the whole story for anyone auditing the Dag.
-- ``stdio`` transport is not isolation. It starts a child process on the worker
- host. It is not a sandbox, and when no environment is supplied the child
- inherits a small allowlist of variables rather than the full parent
- environment — so it is both less contained and less predictable than it looks.
-
-**A real example.** ``example_mcp.py`` drives setup from a connection:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
- :language: python
- :start-after: [START howto_toolset_mcp_connection]
- :end-before: [END howto_toolset_mcp_connection]
-
-and puts several servers on one agent, prefixed so their tool names stay apart:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
- :language: python
- :start-after: [START howto_toolset_mcp_multiple]
- :end-before: [END howto_toolset_mcp_multiple]
-
-No MCP server for object storage ships with this provider. If one exists for
-your target, it is a third route alongside the two above — and the two questions
-from the start of this page settle it. Its tool list is the server's rather than
-yours, and its token comes from wherever ``mcp_conn_id`` or your own callable
-says. So where a hook already reaches the same target, the hook wins; the server
-wins where it covers work the hook does not expose.
-
-**Credentials and where it runs.** ``mcp_conn_id`` supplies host, credentials and
-transport, unless a provider callable overrides that. ``http`` and ``sse`` reach
-a remote server; ``stdio`` runs a child process on the worker host.
-
-``AgentSkillsToolset``
-----------------------
-
-**Choose it when** what the agent is missing is procedural knowledge rather than
-an endpoint — how this team writes a report, which checks run before a release,
-what the house conventions are. A skill is a directory of instructions and
-optional scripts, and
-:class:`~airflow.providers.common.ai.toolsets.skills.AgentSkillsToolset` makes it
-discoverable. See :ref:`agent-skills` for the layout.
-
-**What it cannot do**
-
-- ``exclude_resources`` does not hide a file from the skill's own scripts. It
- keeps matches out of resource discovery and out of ``read_skill_resource``,
- and the parameter's documentation says plainly that it does not stop
- ``run_skill_script`` from reading them off disk. For genuinely sensitive files,
- pair it with ``exclude_tools={"run_skill_script"}``. (The parameter needs
- ``pydantic-ai-skills>=1.2.0``, which the ``skills`` extra already pins.)
-- It does not move script execution anywhere safer. The toolset's own wording for
- ``exclude_tools`` calls ``run_skill_script`` "on-worker script execution", and
- nothing in this toolset routes those scripts into a sandbox — so unless you
- exclude the tool, a skill's scripts run in the worker process with the worker's
- reach. If that is not acceptable, exclude the tool or put the work behind
- ``SandboxToolset`` instead.
-- It re-fetches a Git source on every run. ``GitSkills`` is resolved and
- shallow-cloned on the worker when the run starts, and the checkout is deleted
- when it ends; nothing is kept between runs, so a large or slow repository pays
- that clone once per run. A local directory is read in place and
- costs nothing.
-
-**A real example.** ``example_agent_skills.py`` loads skills from a local
-directory:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
- :language: python
- :start-after: [START howto_operator_agent_skills_local]
- :end-before: [END howto_operator_agent_skills_local]
-
-The two skills it ships, ``aip-tracker`` and ``sql-reporting``, are procedural by
-nature: neither adds an endpoint the agent could not already reach. That is the
-signal you are on the right route.
-
-**Credentials and where it runs.** A local directory needs no credential. A
-private repository goes through
-:class:`~airflow.providers.common.ai.skills.GitSkills` and its ``conn_id``,
-resolved by the Git provider's ``GitHook``; plain ``http://`` is refused when a
-``conn_id`` is set, so a credential is never sent in the clear. Cloning, reading
-and any script execution happen on the worker.
-
-``SandboxToolset``
-------------------
-
-**Choose it when** the work is running code the model wrote, rather than calling
-a tool you picked in advance — exploratory analysis, installing a package for one
-task, a script the model writes, runs, and fixes from its own traceback. Every
-other route on this page answers "call this thing"; this one answers "here is
-somewhere to work". It has its own page, :doc:`sandbox`, with worked scenarios
-and the limitations to read before designing a Dag around it.
-
-Before following this row, check whether the actual need is narrower than that:
-``code_mode=True`` is a flag on ``AgentOperator``, not a toolset on this table
-— it changes how the model invokes the tools it already has, letting it write
-code to call several of them instead of emitting one call per step. It does not
-give the agent somewhere to run arbitrary code of its own, and it avoids the
-``sbx`` backend's production-readiness, network-isolation, and reclamation
-caveats below and needs no backend at all — but not the reachability one: the generated code runs in
-Monty's deny-by-default sandbox, but the tools it calls still run in the
-worker, so a credential-bearing toolset on the same agent stays within reach
-whether or not code mode is on. See :ref:`code-mode` and
-:ref:`sandbox-boundaries`.
-
-**What it cannot do**
-
-- Only one of its two backends runs on Kubernetes. ``SbxSandboxBackend`` drives
- Docker Sandboxes on the worker host, and its own documentation says to use it
- for local development: it wants the ``sbx`` binary on the host, an
- authenticated Docker account, a one-time ``sbx policy init``, and on Linux KVM
- or nested virtualization — which an unprivileged container cannot provide.
- Production and Kubernetes use
- :class:`~airflow.providers.common.ai.sandbox.modal.ModalSandboxBackend`, a
- hosted backend behind the ``modal`` extra that installs nothing on the worker
- and reclaims a sandbox at its own lifetime if the worker dies. Both implement
- :class:`~airflow.providers.common.ai.sandbox.SandboxBackend`, and a third
- vendor can too.
-- It does not contain the agent. Only what these tools do runs in the sandbox;
- the agent loop, the model calls, and every other toolset on the same agent stay
- in the worker with the worker's credentials. It contains model-written code, so
- pairing it with a credential-bearing toolset on the same agent puts the
- credential back within reach. :ref:`sandbox-boundaries` sets this out in full.
-- The ``sbx`` backend cannot enforce network isolation on its own. It applies
- ``allow_egress_to`` as a per-sandbox policy rule, but only on top of a
- ``deny-all`` host policy, since a local rule can narrow egress and never
- widen it; ``block_network`` has no per-sandbox enforcement at all. Ask for
- either against a host policy that is not already ``deny-all`` and the
- backend raises rather than silently leaving the sandbox less restricted
- than the spec asked for. ``block_network`` defaults to ``True``, so a bare
- ``SandboxSpec()`` with no arguments already asks for it and is refused
- under the default ``host_network_policy="unknown"``. On Modal the same default
- maps onto the sandbox's own ``block_network`` and is enforced exactly; a
- hostname allowlist there is matched on the TLS handshake name and has to be
- opted into, for the reasons set out on :doc:`sandbox`.
-- Reclamation depends on the backend. A failed teardown is logged as a warning
- rather than raised, deliberately, so that a teardown blip cannot fail a
- finished run. On ``sbx`` nothing else picks up the slack: there is no
- server-side TTL, so a worker killed outright leaves the microVM and its
- workspace directory behind, named ``airflow-sandbox-*`` so an operator can find
- and remove them. On Modal the sandbox ends at its own ``sandbox_timeout``
- whatever became of the worker.
-- Nothing survives the run, and a file the agent built can leave only through
- the model's context, which is text-only and capped. Producing an artifact for a
- downstream task is a ``@task`` driving a backend directly today, not the agent;
- :doc:`sandbox` has the example.
-
-**A real example.** ``example_sandbox_toolset.py`` in this provider's example
-Dags has an agent investigating a revenue anomaly on the Modal backend beside a
-``SQLToolset``, the same agent shape on ``sbx`` for a laptop, and a ``@task``
-producing a file through a sandbox; all three are on :doc:`sandbox`. The two
-system tests, ``example_sandbox_toolset_sbx.py`` and
-``example_sandbox_toolset_modal.py``, run against a real backend and are
-reachable from the System Tests entry in the sidebar.
-
-**Credentials and where it runs.** This route does not end at an
-Airflow connection. Airflow puts none of its context, connections, variables or
-worker environment into the sandbox; only what you pass through
-:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` goes in, and the
-credential that provisions the sandbox never enters it. Authorization to the
-backend sits outside Airflow — ``sbx login`` and ``sbx policy init`` on the
-machine for ``sbx``, or ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` on the
-worker for Modal. Work runs in a per-run microVM on the worker host with
-``sbx``, or off the worker entirely in Modal's infrastructure. Its tool calls
-act as barriers, as they do for the other
-routes that build their own tools; see :ref:`toolset-call-barriers`.
-
-``BaseManagedAgentToolset``
----------------------------
-
-**Choose it when** the reasoning itself belongs on the vendor's infrastructure —
-the agent is already deployed there, grounded in data that never leaves, and
-Airflow's job is to submit one request and read one answer.
-:ref:`managed-agent-toolsets` covers the shape.
-
-Read the first bullet before planning around this route.
-
-**What it cannot do**
-
-- It is an extension point, not a toolset you can use as-is. This provider ships
- the base class only:
- :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
- declares ``agent_ref`` abstract and raises ``TypeError`` at construction unless
- a subclass implements ``invoke_sync`` or ``invoke``. No vendor subclass exists
- in this repository — the Snowflake Cortex, Bedrock AgentCore, Azure AI Foundry
- and Vertex AI Agent Engine names in its docstring describe the shape it expects,
- not implementations that ship. Using this route means writing that subclass.
-- It has no allow-list to offer. The whole toolset is one tool: a prompt goes in,
- an answer comes out. Whatever governs what the remote agent may touch lives on
- the vendor's side, which is the trade you are making.
-- It does not define where the credential comes from. The base class leaves
- that decision to the subclass; :ref:`managed-agent-toolsets` frames the
- intended shape as routing authentication through the provider's own hook.
-- Durable replay is off by default. ``replayable`` is ``False`` because a managed
- agent may act on systems Airflow cannot observe, so replaying from the cache
- could skip a side effect. Read-only agents can opt in.
-- Its failover variant relies on two preconditions the code cannot check.
- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
- requires members that are genuinely interchangeable — the same agent deployed
- twice, not two specialists over different data — and one-shot exchanges, because
- each member is invoked with a bare prompt and no thread reference, so a failover
- silently starts a fresh conversation rather than resuming the old one.
-
-**A real example.** There is none. No example Dag, no system test, and the only
-code in the documentation is an illustrative subclass in :doc:`toolsets`. Budget
-for writing and testing the subclass yourself.
-
-**Credentials and where it runs.** Both are the subclass's decision. Reasoning
-runs on the vendor's infrastructure; the worker sends a request and waits.
-
-Where the credentials come from
--------------------------------
-
-Airflow connections and hooks are the access-governance machinery this provider
-already has: a connection lives in a secret backend rather than in Dag code, and
-its scope is set outside the Dag. A hook turns that scope into typed methods.
-Where a route ends up getting its credential is therefore a decision worth
-making on purpose rather than inheriting.
-
-.. list-table::
- :widths: 28 44 28
- :header-rows: 1
-
- * - Route
- - Credential source
- - Where the tool call runs
- * - ``HookToolset``
- - Whatever connection the hook you supply resolves
- - Worker process
- * - ``SQLToolset``
- - ``db_conn_id``, via ``BaseHook.get_connection``
- - Worker process, against the database
- * - ``DataFusionToolset``
- - ``conn_id`` on each ``DataSourceConfig``
- - Worker process (embedded engine)
- * - ``MCPToolset``
- - ``mcp_conn_id``, unless ``token_provider`` or ``env_provider`` supplies
- one from elsewhere
- - Remote server, or a child process on the worker host for ``stdio``
- * - ``AgentSkillsToolset``
- - ``GitSkills.conn_id`` for private repositories; none for local directories
- - Worker process
- * - ``SandboxToolset``
- - None from Airflow. Host-level ``sbx login`` and ``sbx policy init``
- for ``sbx``; ambient Modal token on the worker for Modal
- - A microVM on the worker host (``sbx``), or Modal's infrastructure off
- the worker (Modal)
- * - ``BaseManagedAgentToolset``
- - Undefined by the base class; the subclass decides
- - The vendor's infrastructure
-
-Two rows are worth pausing on. ``SandboxToolset`` deliberately takes no Airflow
-credential — that is the whole point of it, and it substitutes a backend-level
-boundary, on the host or at the vendor, for the connection-level one. ``BaseManagedAgentToolset`` does not
-substitute anything; it simply leaves the question to whoever writes the
-subclass. Neither is a defect, but in both cases the access decision has moved
-somewhere Airflow cannot see it, and somebody has to make that decision again in
-the new place. The :ref:`defense-layer table <toolset-defense-layers>` is the
-right companion when you do.
-
-.. _toolset-call-barriers:
-
-Tool calls as barriers
-----------------------
-
-One behaviour cuts across these routes rather than telling them apart.
-``HookToolset``, ``SQLToolset``, ``DataFusionToolset`` and ``SandboxToolset``
-each build their own tool definitions and set ``sequential=True`` on them, which
-pydantic-ai treats as a barrier: the tool runs alone, tools the model emitted
-before it finish first, and tools emitted after it start only once it returns.
-A slow call on any of those four therefore holds up the rest of that step, not
-just its own toolset. ``BaseManagedAgentToolset`` sets ``sequential=False``
-deliberately, because the wait it introduces is remote. ``MCPToolset`` and
-``AgentSkillsToolset`` define no tools of their own — they pass through whatever
-the upstream toolset declares — so the setting is not theirs to make. Do not read
-this as a reason to choose one route over another; read it as something to expect
-from all four.
-
-Layering
---------
-
-:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` and the
-durable-execution ``CachingToolset`` are not alternatives to anything above. Both
-are wrappers: they take a toolset and return a toolset, adding per-call logging
-or replay from a durable cache. Choose a route first, then decide whether to wrap
-it. The same applies to
-:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`,
-which converts a chosen toolset for a different agent framework rather than
-offering another way to reach a system; :doc:`toolsets` covers both.
-
-Running the agent outside the data system
------------------------------------------
-
-A recurring alternative to everything on this page is to push the reasoning into
-the system that holds the data and let it run there. The trade is worth stating
-plainly, because it is the same trade the last row of the table makes.
-
-Running the agent where the data lives keeps the loop short, and it is the right
-answer when the work never leaves that system. What it costs is coupling: the
-agent can only reason over what that system holds, its credentials are governed
-by that system's model rather than Airflow's, and failures are that system's to
-explain. Running the agent in the worker instead means the model calls, the agent
-loop and every tool are orchestrated by Airflow, where they get retries, logging
-and observability like any other task. It also means one agent can hold a
-database, an object store and a vendor API in the same run, with each route's
-credential resolved through a connection that was reviewed once.
-
-This provider does not treat either as the default. The rule of thumb in
-:doc:`index` is the dividing line: if Airflow should *run* the AI step, and the
-model should stay swappable, use ``common.ai``; if the Dag *submits work to* a
-vendor-managed service and waits for the result, use that vendor's provider —
-and ``BaseManagedAgentToolset`` exists for the case where you want the second
-behaviour from inside an agent that is otherwise doing the first.
-
-Next
-----
-
-- :doc:`toolsets` — how to configure each toolset, and the full
- :ref:`defense-layer table <toolset-defense-layers>`.
-- :doc:`security` — the provider's security guidance.
-- :doc:`examples` — the example Dags referenced above.
diff --git a/providers/common/ai/docs/code_mode.rst b/providers/common/ai/docs/code_mode.rst
new file mode 100644
index 00000000000..32256d88fb1
--- /dev/null
+++ b/providers/common/ai/docs/code_mode.rst
@@ -0,0 +1,87 @@
+ .. 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.
+
+.. _code-mode:
+
+Code mode
+=========
+
+Set ``code_mode=True`` to collapse the agent's tools into a single ``run_code``
+tool powered by the `Monty <https://github.com/pydantic/monty>`__ sandbox (via
+pydantic-ai-harness). Instead of one model round-trip per tool call, the model
+writes a single Python snippet that calls the tools as functions -- with loops,
+conditionals, and ``asyncio.gather`` -- in one turn. For multi-tool workflows
+this cuts round-trips and token use.
+
+The generated code runs in Monty's deny-by-default sandbox: it cannot read the
+filesystem, the network, or environment variables. It can only call the tools
+you registered. Code mode therefore does not widen what the agent can reach --
+the tools it calls still run in the worker -- it only changes how the model
+invokes them. See :doc:`agent_security` for the tool
+boundary.
+
+When to use it
+--------------
+
+Code mode pays off for **orchestration-heavy, computation-light** workflows:
+calling several tools, looping over their results, filtering, and combining them.
+Collapsing many sequential tool calls into one turn is where the round-trip and
+token savings come from -- the example below answers a per-customer question in a
+single ``run_code`` block instead of one model round-trip per customer.
+
+It is **not a general-purpose code runtime**. The generated code is only the glue
+between tool calls; every real capability must come from a tool. Monty runs a
+subset of Python and **cannot import third-party libraries** (pandas, numpy,
+requests, boto3, ...) and has no filesystem or network access. If a task needs to
+crunch data inline with a library, you have two options, both better than code
+mode:
+
+- **Push the work into a tool.** Do the aggregation in SQL (``SQLToolset``), or
+ expose a hook method that returns the processed result (``HookToolset``). The
+ tool runs in the full worker environment with all its dependencies, and code
+ mode just orchestrates it.
+- **Give the agent a real environment**, with
+ :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`. It hands
+ the model a shell and a filesystem in a disposable sandbox off the worker, so
+ third-party packages, a real interpreter and installed binaries are all
+ available. It costs a sandbox per run and a second or so to provision, against
+ well under a millisecond for Monty, so reach for it when inline library code is
+ genuinely required rather than by default. See :doc:`sandbox/index` for the
+ backends and their limitations.
+
+The two are not exclusive: ``code_mode=True`` and a ``SandboxToolset`` can be
+enabled together, and the file tools fold into ``run_code`` while ``run_command``
+stays a tool of its own.
+
+Requires the ``code-mode`` extra::
+
+ pip install "apache-airflow-providers-common-ai[code-mode]"
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
+ :language: python
+ :start-after: [START howto_operator_agent_code_mode]
+ :end-before: [END howto_operator_agent_code_mode]
+
+Unlike passing a capability through ``agent_params`` (see
+:ref:`capabilities-passthrough`), ``code_mode`` is a plain boolean and is
+serialization-safe: the ``CodeMode`` capability is built at execution time, not
+stored on the serialized operator.
+
+.. note::
+
+ Monty is pre-1.0. The ``code-mode`` extra is opt-in so its dependency churn
+ never affects the base provider install.
diff --git a/providers/common/ai/docs/concepts.rst b/providers/common/ai/docs/concepts.rst
new file mode 100644
index 00000000000..7b6f99b9f23
--- /dev/null
+++ b/providers/common/ai/docs/concepts.rst
@@ -0,0 +1,135 @@
+ .. 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.
+
+.. _howto/concepts:
+
+Core concepts
+=============
+
+Five ideas carry the whole provider. Each has a page of its own; this page is the map.
+
+Connections choose the model
+----------------------------
+
+``common.ai`` is built on `pydantic-ai <https://ai.pydantic.dev/>`__, so the model vendor
+(OpenAI, Anthropic, Google, Bedrock, …) is picked by the connection ``llm_conn_id`` points
+at — switching providers later is a connection change, not a Dag rewrite. Most connections
+use the generic ``pydanticai`` type, but Azure OpenAI, Bedrock, and Vertex AI also have their
+own connection types (``pydanticai_azure``, ``pydanticai_bedrock``, ``pydanticai_vertex``) for
+provider-specific authentication.
+
+The model name lives on the connection in ``provider:model`` form, and ``model_id`` on an
+operator overrides it. :doc:`connections/pydantic_ai` has the full resolution order, and
+:doc:`provider_fallback` explains how one connection can name others to fail over to.
+
+Operators and decorators do the work
+------------------------------------
+
+Every operator ships with a matching ``@task`` decorator, so a Dag can use whichever style
+it already uses. ``LLMOperator`` sends one prompt and returns one answer. ``AgentOperator``
+runs a multi-turn loop in which the model calls tools until it is done. The other
+operators are specializations of the first: branching on the answer, analyzing a file,
+generating SQL, comparing schemas, or submitting many prompts as one batch.
+:doc:`operators/index` has the selection table.
+
+The AI step is orchestrated by Airflow: the model calls, the agent loop, and any tools
+run in the Airflow worker by default, where they get retries, logging, and observability like
+any other task. The exception is :ref:`SandboxToolset <sandbox-limitations>`, which exists so
+that code the *model* writes runs somewhere else.
+
+Toolsets give agents reach
+--------------------------
+
+A toolset is what an agent is allowed to call. The provider ships toolsets that wrap
+Airflow hooks, SQL databases, files through DataFusion, MCP servers, Agent Skills, a
+sandboxed shell, and vendor-managed agents. An agent's reach is exactly the toolsets you
+register on it. :doc:`toolsets/index` compares them and :doc:`agent_security` explains
+the defense layers.
+
+Existing LangChain tools are not locked out either: pydantic-ai ships
+``pydantic_ai.ext.langchain.LangChainToolset`` upstream, which wraps LangChain tools for a
+common.ai agent, and the provider's own
+:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`
+converts the other way — Airflow-managed toolsets into LangChain tools (see
+:doc:`toolsets/langchain`).
+
+Hooks are the plumbing underneath
+---------------------------------
+
+The provider's hooks bridge an Airflow connection to a specific framework's model objects.
+Each hook is a thin adapter: it reads credentials and config from the connection, then
+returns native framework objects (a ``pydantic_ai`` ``Agent`` / ``Model``, a LangChain
+``BaseChatModel`` or ``Embeddings``, an MCP client, ...). Operators and ``@task`` decorators
+use these hooks internally, and you reach for one directly only when you want the
+framework object in a plain ``@task``.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 75
+
+ * - Hook
+ - When to use
+ * - :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
+ - Default for ``common.ai`` operators (``LLMOperator``, ``AgentOperator``,
+ ``LLMBranchOperator``, ...). Returns a pydantic-ai ``Agent`` / ``Model``.
+ See :doc:`hooks/pydantic_ai`.
+ * - :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook`
+ - Direct LangChain access for tasks that compose ``Runnable``\\s, use the
+ LangChain agent surface, or need LangChain-native chat / embedding model
+ objects. Independent of the pydantic-ai-backed operators.
+ See :doc:`hooks/langchain`.
+ * - :class:`~airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook`
+ - Backs the LlamaIndex ``LlamaIndexEmbeddingOperator`` and
+ ``LlamaIndexRetrievalOperator``.
+ Returns LlamaIndex-native ``BaseEmbedding`` / ``LLM`` objects (OpenAI
+ by default). For non-OpenAI vendors, pass a pre-built
+ ``BaseEmbedding`` / ``LLM`` instance straight to the operator and
+ bypass the hook. See :doc:`hooks/llamaindex`.
+ * - :class:`~airflow.providers.common.ai.hooks.mcp.MCPHook`
+ - Backs ``MCPToolset`` (see :doc:`toolsets/mcp`) for agent tasks that call
+ tools on a remote MCP server. Configure the connection via
+ :doc:`connections/mcp`. See :doc:`hooks/mcp`.
+
+Results flow through XCom
+-------------------------
+
+Every operator pushes its result to XCom like any other task. A plain string arrives as a
+string. A Pydantic ``output_type`` arrives as the model instance, typed, so a downstream
+task can use attribute access. :doc:`structured_output` explains how the class is
+registered for deserialization and where that stops working.
+
+Related providers
+-----------------
+
+Use a vendor's own provider instead when the Dag needs that vendor's **native API surface**, a
+service the vendor runs for you, which no vendor-neutral operator wraps:
+
+* :doc:`apache-airflow-providers-openai:index` — the Embeddings and Responses APIs, and Batch
+ jobs built from a pre-uploaded JSONL file of raw request bodies.
+* :doc:`apache-airflow-providers-anthropic:index` — Message Batches built from raw Messages
+ API request bodies (multi-turn, images, tools), and Managed Agents sessions where the agent
+ loop runs on Anthropic's infrastructure rather than in the Airflow worker.
+* :doc:`apache-airflow-providers-cohere:index` — Cohere's own Embed API.
+* :doc:`apache-airflow-providers-google:index` — Vertex AI's Batch Prediction jobs
+ (``CreateBatchPredictionJobOperator``), a managed batch service like OpenAI's Batch API.
+* :doc:`apache-airflow-providers-amazon:index` — Bedrock's Batch Inference
+ (``BedrockBatchInferenceOperator``), and Bedrock AgentCore's managed agent runtime
+ (``BedrockCreateAgentRuntimeOperator`` / ``BedrockInvokeAgentRuntimeOperator``), where the
+ agent loop runs on AWS's infrastructure rather than in the Airflow worker.
+
+The :doc:`landing page <index>` has the decision table and the rule of thumb that separates
+the two.
diff --git a/providers/common/ai/docs/connections/langchain.rst b/providers/common/ai/docs/connections/langchain.rst
index f70e165f758..8d2cb3f1f03 100644
--- a/providers/common/ai/docs/connections/langchain.rst
+++ b/providers/common/ai/docs/connections/langchain.rst
@@ -17,7 +17,7 @@
.. _howto/connection:langchain:
-LangChain Connection
+LangChain connection
====================
The ``langchain`` connection type configures access to LLM providers via
diff --git a/providers/common/ai/docs/connections/llamaindex.rst b/providers/common/ai/docs/connections/llamaindex.rst
index 82482ce388e..ac162b904cb 100644
--- a/providers/common/ai/docs/connections/llamaindex.rst
+++ b/providers/common/ai/docs/connections/llamaindex.rst
@@ -17,8 +17,8 @@
.. _howto/connection:llamaindex:
-LlamaIndex Connection
-======================
+LlamaIndex connection
+=====================
The ``llamaindex`` connection type configures access to LLM and embedding
providers for `LlamaIndex <https://docs.llamaindex.ai/>`__. It backs
diff --git a/providers/common/ai/docs/connections/mcp.rst b/providers/common/ai/docs/connections/mcp.rst
index 1ddd17b5bcd..2707f3f88f9 100644
--- a/providers/common/ai/docs/connections/mcp.rst
+++ b/providers/common/ai/docs/connections/mcp.rst
@@ -17,7 +17,7 @@
.. _howto/connection:mcp:
-MCP Server Connection
+MCP server connection
=====================
The MCP connection type configures access to
diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst b/providers/common/ai/docs/connections/pydantic_ai.rst
index ba1447a5bed..41d866262d2 100644
--- a/providers/common/ai/docs/connections/pydantic_ai.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai.rst
@@ -17,7 +17,7 @@
.. _howto/connection:pydanticai:
-Pydantic AI Connection
+Pydantic AI connection
======================
The `Pydantic AI <https://ai.pydantic.dev/>`__ connection type configures access
diff --git a/providers/common/ai/docs/connections/pydantic_ai_azure.rst b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
index 692bb5909ba..3b20e25493d 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_azure.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
@@ -17,8 +17,8 @@
.. _howto/connection:pydanticai_azure:
-Pydantic AI (Azure OpenAI) Connection
-======================================
+Pydantic AI (Azure OpenAI) connection
+=====================================
The ``pydanticai_azure`` connection type configures access to
`Azure OpenAI <https://azure.microsoft.com/en-us/products/ai-services/openai-service>`__
diff --git a/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst b/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
index 4c106b8490c..8be3a4ca8b3 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
@@ -17,8 +17,8 @@
.. _howto/connection:pydanticai_bedrock:
-Pydantic AI (AWS Bedrock) Connection
-=======================================
+Pydantic AI (AWS Bedrock) connection
+====================================
The ``pydanticai_bedrock`` connection type configures access to
`AWS Bedrock <https://aws.amazon.com/bedrock/>`__ via the pydantic-ai framework.
diff --git a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
index bbf210e79b8..c1b0d9716a2 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
@@ -17,8 +17,8 @@
.. _howto/connection:pydanticai_vertex:
-Pydantic AI (Google Vertex AI) Connection
-============================================
+Pydantic AI (Google Vertex AI) connection
+=========================================
The ``pydanticai_vertex`` connection type configures access to
`Google Vertex AI <https://cloud.google.com/vertex-ai>`__ via the pydantic-ai
diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst
new file mode 100644
index 00000000000..307daf95b6e
--- /dev/null
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -0,0 +1,159 @@
+ .. 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.
+
+.. _durable-execution:
+
+Durable execution
+=================
+
+Agent tasks can involve multiple LLM calls and tool invocations. If a task
+fails mid-run (network error, timeout, transient API failure), a plain retry
+re-executes every LLM call and tool call from scratch -- repeating work that
+already succeeded and incurring additional cost.
+
+Setting ``durable=True`` caches each LLM response and tool result as it
+completes. On retry, completed steps are replayed from the cache and only the
+remaining steps run against the live model and tools. The cache is deleted
+after successful completion.
+
+Durable execution only helps when the task has retries configured. Without
+retries there is nothing to replay.
+
+This page is about making an ``AgentOperator`` retry cheap. Deciding *whether* a task
+should retry at all is :doc:`retry_policies`; a retried ``LLMBatchOperator`` re-attaches
+to its running batch instead of resubmitting (:ref:`llm-batch-reattach`).
+
+**Configuration**
+
+On **Airflow >= 3.3** the cache is stored in the
+:doc:`task state store <apache-airflow:core-concepts/task-state-store>`,
+scoped to the task instance. No configuration is required; the store handles
+persistence across retries.
+
+By default each cached step is written to the Airflow metadata database. Model
+responses and large tool results can be sizable, so for agents with large
+payloads configure ``[workers] state_store_backend`` to offload step values to
+external storage (e.g. object storage) instead of the metadata database; the
+provider then stores only a reference in the database.
+
+On **Airflow < 3.3** the cache is persisted to ObjectStorage and the location
+must be set in ``airflow.cfg``. The task raises ``ValueError`` at runtime if
+``durable=True`` and the option is missing.
+
+.. code-block:: ini
+
+ [common.ai]
+ # Local filesystem -- suitable for development
+ durable_cache_path = file:///tmp/airflow_durable_cache
+
+The value is an ObjectStorage URI, so any supported backend works. For
+production, use a shared store so retries on a different worker can read the
+cache:
+
+.. code-block:: ini
+
+ [common.ai]
+ durable_cache_path = s3://my-bucket/airflow/durable-cache
+
+**Operator example**
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_operator_agent_durable]
+ :end-before: [END howto_operator_agent_durable]
+
+**Decorator example**
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_decorator_agent_durable]
+ :end-before: [END howto_decorator_agent_durable]
+
+**How it works**
+
+1. On first execution, each LLM response and tool result is saved as the agent
+ progresses, together with a fingerprint of the request that produced it
+ (model, message history, settings, and tools for LLM steps; tool name,
+ arguments, and call id for tool steps).
+2. If the task fails and Airflow retries it, completed steps are loaded from
+ the cache and returned without calling the model or tool. Steps not yet in
+ the cache proceed normally.
+3. Before a step is replayed, its stored fingerprint is compared against the
+ current request. If anything changed between attempts -- the system
+ prompt, the model, the toolset, model settings, or the conversation so
+ far -- the stale entry is discarded, a warning is logged, and the step
+ re-runs live. A divergence also invalidates the steps after it: re-running
+ an LLM step produces fresh tool call ids, so tool results recorded under
+ the old conversation no longer match. A changed agent costs a re-run; it
+ never replays responses that belong to a different conversation.
+4. After successful completion, the cached steps are deleted.
+
+Replay verification compares the **requests** sent to models and tools, not
+the code behind them. Editing a tool's implementation between attempts does
+not invalidate an already-cached result for an identical call, and pointing
+``llm_conn_id`` at a different endpoint serving the same model name does not
+invalidate cached responses -- clear the cache to force a fully fresh run.
+
+After the run, a single INFO summary line reports how many steps were
+replayed vs executed fresh. Per-step detail is available at DEBUG level.
+
+The cache is scoped to a single task instance (Dag id, run id, task id, and
+map index), so each run replays only its own steps. On Airflow >= 3.3 the cache
+lives in the task state store and is removed when the Dag run is cleaned up; on
+Airflow < 3.3 it is a JSON file named ``{dag_id}_{task_id}_{run_id}.json`` (with
+``_{map_index}`` appended for mapped tasks) under the configured
+``durable_cache_path``.
+
+.. note::
+
+ Runs that fail permanently (exhaust all retries) leave their cached steps
+ behind. These do not affect future Dag runs (each run is scoped separately).
+ On Airflow >= 3.3 they are reclaimed when the Dag run is removed; on Airflow
+ < 3.3 the orphaned JSON files consume storage until cleaned up, so add a
+ lifecycle policy to the storage backend or remove them periodically.
+
+**Side effects and idempotency**
+
+Durable execution caches **return values**, not side effects. When a step is
+replayed, the tool's code does not run -- only the stored return value is
+returned. Two things follow from this:
+
+- If a tool completed successfully and its result was cached, the tool will
+ **not** run again on retry. Any side effect it produced (writing a file,
+ sending a message) already happened during the original run and is not
+ repeated.
+- If a tool fails *before* its result is cached, it **will** run again on
+ retry. A tool that partially completed (e.g. sent an email then raised an
+ exception) may produce the side effect a second time.
+
+All built-in toolsets (``SQLToolset`` with ``allow_writes=False``,
+``HookToolset`` in read-only mode) are read-only and replay safely. For custom
+tools with non-idempotent side effects, design the tool to be idempotent. For
+example, check whether the operation already completed before acting, or
+use database constraints to prevent duplicate writes.
+
+Tool results must be JSON-serializable to be cached. If a tool returns a
+non-serializable value (e.g. ``BinaryContent`` from MCP tools), that step is
+skipped with a warning and will re-execute on retry instead of replaying from
+cache. The task itself still succeeds.
+
+See also
+--------
+
+- :doc:`operators/agent` — the operator these settings apply to.
+- :doc:`retry_policies` — let a model decide whether a failure is worth retrying.
+- :doc:`troubleshooting` — the ``durable=True`` combinations the operator rejects.
diff --git a/providers/common/ai/docs/end_to_end_pipelines.rst b/providers/common/ai/docs/end_to_end_pipelines.rst
index 4ddc06a621e..3c38d07891a 100644
--- a/providers/common/ai/docs/end_to_end_pipelines.rst
+++ b/providers/common/ai/docs/end_to_end_pipelines.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-End-to-End Pipelines
-=====================
+End-to-end pipelines
+====================
The Dags in this guide combine several patterns from the operator and hook guides into one
production-shaped pipeline. Each section explains the architecture -- how the Dags are split,
diff --git a/providers/common/ai/docs/examples.rst b/providers/common/ai/docs/examples.rst
index af69eb85e8d..0ac806e16a2 100644
--- a/providers/common/ai/docs/examples.rst
+++ b/providers/common/ai/docs/examples.rst
@@ -164,9 +164,3 @@ Reliability
- Routing a failure with a model that answers typed questions instead of writing
text, and escalating when its confidence is low. Source:
`example_classifier_model.py <https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_classifier_model.py>`__.
-
-.. toctree::
- :hidden:
- :maxdepth: 1
-
- end_to_end_pipelines
diff --git a/providers/common/ai/docs/guardrails.rst b/providers/common/ai/docs/guardrails.rst
new file mode 100644
index 00000000000..52bfad9fd8c
--- /dev/null
+++ b/providers/common/ai/docs/guardrails.rst
@@ -0,0 +1,62 @@
+ .. 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.
+
+.. _capabilities-passthrough:
+
+Guardrails and capabilities
+===========================
+
+pydantic-ai `capabilities <https://ai.pydantic.dev/capabilities/>`__ bundle
+tools, lifecycle hooks, instructions, and model settings into composable units.
+Common ones include ``Thinking`` (reasoning at a configurable effort level),
+``WebSearch``, ``WebFetch``, ``ImageGeneration``, and ``MCP``.
+For the current capability catalog and package-specific installation notes, see
+the pydantic-ai documentation and the
+`pydantic-ai-harness capability matrix <https://github.com/pydantic/pydantic-ai-harness#capability-matrix>`__.
+
+``AgentOperator`` does not yet expose a first-class ``capabilities=`` kwarg,
+but anything passed through ``agent_params`` is forwarded to the underlying
+``Agent(...)`` constructor.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
+ :language: python
+ :start-after: [START howto_operator_agent_capabilities_thinking]
+ :end-before: [END howto_operator_agent_capabilities_thinking]
+
+Capabilities compose with toolsets -- pydantic-ai merges tools from both.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
+ :language: python
+ :start-after: [START howto_operator_agent_capabilities_composed]
+ :end-before: [END howto_operator_agent_capabilities_composed]
+
+Guardrail capabilities use the same passthrough pattern. This example uses
+``InputGuard`` from ``pydantic-ai-shields`` to reject a prompt before the agent
+run starts.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
+ :language: python
+ :start-after: [START howto_operator_agent_capabilities_input_guard]
+ :end-before: [END howto_operator_agent_capabilities_input_guard]
+
+.. warning::
+
+ ``agent_params`` is a templated field, which Airflow serializes by calling
+ ``str()`` on values it doesn't natively understand. Capability instances
+ are not yet round-trip-safe through Dag serialization, so the examples above construct them inside the ``@dag`` function -- not at module level.
+ First-class ``capabilities=`` support on ``AgentOperator`` (with proper
+ serializer hooks) is tracked as a follow-up.
diff --git a/providers/common/ai/docs/hitl_review.rst b/providers/common/ai/docs/hitl_review.rst
index 699d51c5e56..76c0140841c 100644
--- a/providers/common/ai/docs/hitl_review.rst
+++ b/providers/common/ai/docs/hitl_review.rst
@@ -17,8 +17,8 @@
.. _howto:hitl_review:
-Human-in-the-Loop (HITL) Review For Agentic Operators
-#####################################################
+Human-in-the-loop (HITL) review for agents
+==========================================
HITL Review adds an interactive feedback loop to agentic operators. After the
LLM Agent produces an initial output, a human reviewer can **approve**, **reject**, or
@@ -27,9 +27,13 @@ terminal action, or until a timeout is reached or max_iterations reached.
This document describes the architecture, workflow, API, XCom schema, and usage.
-
Overview
-========
+--------
+
+.. seealso::
+ This page covers the iterative review loop on ``AgentOperator``. For a one-shot
+ approve, edit or reject gate on ``LLMOperator`` and its subclasses, see
+ :doc:`approval_gates`.
**Components**
@@ -58,9 +62,8 @@ during review but adds cross-process coordination complexity: the agent state
be serialized and restored across defer/resume. XCom polling keeps the flow
simple and keeps all agent context in-process. Future versions may have different approaches.
-
Workflow
-========
+--------
.. code-block:: text
@@ -89,9 +92,8 @@ Workflow
| 7e. hitl_timeout elapsed |
| → push status timeout_exceeded, raise HITLTimeoutError
-
-Using HITL Review with AgentOperator
-====================================
+Using HITL review with ``AgentOperator``
+----------------------------------------
Enable the review loop with ``enable_hitl_review=True``:
@@ -138,9 +140,8 @@ navigate to
:start-after: [START howto_operator_agent_hitl_review]
:end-before: [END howto_operator_agent_hitl_review]
-
REST API
-========
+--------
The plugin exposes a FastAPI app at ``/hitl-review``. Base URL:
@@ -155,9 +156,8 @@ The plugin exposes a FastAPI app at ``/hitl-review``. Base URL:
- ``task_id`` — Task ID.
- ``map_index`` — Map index for mapped tasks. Use ``-1`` for non-mapped tasks or index for dynamic mapping.
-
Endpoints
----------
+^^^^^^^^^
.. list-table::
:header-rows: 1
@@ -184,9 +184,8 @@ Endpoints
- ``/sessions/reject``
- Reject the output. Session must be ``pending_review``.
-
Response model: HITLReviewResponse
-----------------------------------
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. code-block:: python
@@ -208,9 +207,8 @@ Response model: HITLReviewResponse
"task_completed": bool,
}
-
XCom keys and storage
-=====================
+---------------------
All keys use the prefix ``airflow_hitl_review_``.
@@ -235,9 +233,8 @@ All keys use the prefix ``airflow_hitl_review_``.
- Plugin
- Per-iteration human feedback text
-
Session lifecycle
------------------
+^^^^^^^^^^^^^^^^^
- **pending_review** — Awaiting human action. Plugin accepts approve, reject,
or feedback.
@@ -246,9 +243,8 @@ Session lifecycle
worker pushes a new output and status returns to ``pending_review``.
- **approved** / **rejected** — Terminal. Worker has exited the loop.
-
Chat UI
-=======
+-------
The plugin provides an interactive chat UI that loads in the task instance page.
The UI:
diff --git a/providers/common/ai/docs/hooks/index.rst b/providers/common/ai/docs/hooks/index.rst
deleted file mode 100644
index 13455cdf896..00000000000
--- a/providers/common/ai/docs/hooks/index.rst
+++ /dev/null
@@ -1,62 +0,0 @@
- .. 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.
-
-Common AI Hooks
-===============
-
-The common-ai provider ships hooks that bridge an Airflow connection to a specific
-LLM framework's model objects. Each hook is a thin adapter: it reads credentials and
-config from the connection, then returns native framework objects (a ``pydantic_ai``
-``Agent`` / ``Model``, a LangChain ``BaseChatModel`` or ``Embeddings``, an MCP client,
-...). Operators and ``@task`` decorators in this provider use these hooks internally.
-
-Choosing a hook
----------------
-
-.. list-table::
- :header-rows: 1
- :widths: 25 75
-
- * - Hook
- - When to use
- * - :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
- - Default for ``common.ai`` operators (``LLMOperator``, ``AgentOperator``,
- ``LLMBranchOperator``, ...). Returns a pydantic-ai ``Agent`` / ``Model``.
- * - :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook`
- - Direct LangChain access for tasks that compose ``Runnable``\\s, use the
- LangChain agent surface, or need LangChain-native chat / embedding model
- objects. Independent of the pydantic-ai-backed operators.
- * - :class:`~airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook`
- - Backs the LlamaIndex ``LlamaIndexEmbeddingOperator`` and
- ``LlamaIndexRetrievalOperator``.
- Returns LlamaIndex-native ``BaseEmbedding`` / ``LLM`` objects (OpenAI
- by default). For non-OpenAI vendors, pass a pre-built
- ``BaseEmbedding`` / ``LLM`` instance straight to the operator and
- bypass the hook.
- * - :class:`~airflow.providers.common.ai.hooks.mcp.MCPHook`
- - Backs ``MCPToolset`` (see :doc:`../toolsets`) for agent tasks that call
- tools on a remote MCP server. Configure the connection via
- :doc:`../connections/mcp`.
-
-Hook guides
------------
-
-.. toctree::
- :maxdepth: 1
- :glob:
-
- *
diff --git a/providers/common/ai/docs/hooks/langchain.rst b/providers/common/ai/docs/hooks/langchain.rst
index 5485f39ab2a..21134c8b3d3 100644
--- a/providers/common/ai/docs/hooks/langchain.rst
+++ b/providers/common/ai/docs/hooks/langchain.rst
@@ -20,6 +20,12 @@
``LangChainHook``
=================
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+
+ LangChain connection <../connections/langchain>
+
Use :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook` to
bridge an Airflow connection to `LangChain <https://python.langchain.com/>`__
chat and embedding models. The hook reads credentials (API key, optional base
diff --git a/providers/common/ai/docs/hooks/mcp.rst b/providers/common/ai/docs/hooks/mcp.rst
index 783ec6a0543..db861f8f7e8 100644
--- a/providers/common/ai/docs/hooks/mcp.rst
+++ b/providers/common/ai/docs/hooks/mcp.rst
@@ -35,7 +35,7 @@ Role in a Dag: use ``MCPToolset``, not the hook directly
Most Dags should not instantiate ``MCPHook`` directly. The recommended entry
point is :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` (see
-:doc:`../toolsets`), which resolves the hook lazily from an Airflow connection
+:doc:`../toolsets/mcp`), which resolves the hook lazily from an Airflow connection
and manages the underlying session lifecycle for you:
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
diff --git a/providers/common/ai/docs/hooks/pydantic_ai.rst b/providers/common/ai/docs/hooks/pydantic_ai.rst
index ef4070d72c5..79c9548ad71 100644
--- a/providers/common/ai/docs/hooks/pydantic_ai.rst
+++ b/providers/common/ai/docs/hooks/pydantic_ai.rst
@@ -29,7 +29,7 @@ The hook manages API credentials from an Airflow connection and creates pydantic
.. seealso::
:ref:`Connection configuration <howto/connection:pydanticai>`
-Basic Usage
+Basic usage
-----------
Use the hook in a ``@task`` function to call an LLM:
@@ -39,7 +39,7 @@ Use the hook in a ``@task`` function to call an LLM:
:start-after: [START howto_hook_pydantic_ai_basic]
:end-before: [END howto_hook_pydantic_ai_basic]
-Overriding the Model
+Overriding the model
--------------------
The model can be specified at three levels (highest priority first):
@@ -56,7 +56,7 @@ The model can be specified at three levels (highest priority first):
# Override with a specific model
hook = PydanticAIHook(llm_conn_id="my_llm", model_id="anthropic:claude-opus-4-6")
-Structured Output
+Structured output
-----------------
Pydantic-ai's structured output works naturally through the hook.
@@ -67,8 +67,8 @@ Define a Pydantic model for the expected output shape, then pass it as ``output_
:start-after: [START howto_hook_pydantic_ai_structured_output]
:end-before: [END howto_hook_pydantic_ai_structured_output]
-Loading Agent Config from a Spec File
---------------------------------------
+Loading agent config from a spec file
+-------------------------------------
Instead of hard-coding model name, instructions, and settings in Python, you can
store them in a YAML or JSON `AgentSpec
diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst
index a98b1df502f..12ef3431d5c 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -49,67 +49,9 @@ When to use this provider
- The vendor's own provider
- e.g. :doc:`apache-airflow-providers-anthropic:index`
-``common.ai`` is built on `pydantic-ai <https://ai.pydantic.dev/>`__, so the model vendor
-(OpenAI, Anthropic, Google, Bedrock, …) is picked by the connection ``llm_conn_id`` points
-at — switching providers later is a connection change, not a Dag rewrite. Most connections
-use the generic ``pydanticai`` type, but Azure OpenAI, Bedrock, and Vertex AI also have their
-own connection types (``pydanticai_azure``, ``pydanticai_bedrock``, ``pydanticai_vertex``) for
-provider-specific authentication. Existing LangChain
-tools aren't locked out either: pydantic-ai ships ``pydantic_ai.ext.langchain.LangChainToolset``
-upstream, which wraps LangChain tools for a common.ai agent, and the provider's own
-:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`
-converts the other way — Airflow-managed toolsets into LangChain tools (see :doc:`toolsets`).
-The AI step is orchestrated by Airflow: the model calls, the agent loop, and any tools
-run in the Airflow worker by default, where they get retries, logging, and observability like
-any other task. The exception is :ref:`SandboxToolset <sandbox-limitations>`, which exists so
-that code the *model* writes runs somewhere else.
-
-Use it when a Dag needs:
-
-* **Generation, classification, summarization, or structured extraction** —
- :doc:`LLMOperator and @task.llm <operators/llm>`, with Pydantic-typed output pushed to XCom.
-* **Many prompts at half the price** — :doc:`LLMBatchOperator and @task.llm_batch <operators/llm_batch>`
- submit prompts as one provider batch job (OpenAI or Anthropic), poll in deferrable mode for up to 24
- hours, re-attach on retry instead of paying twice, and land results as JSONL on object storage.
-* **Branching on a model's decision** — :doc:`LLMBranchOperator <operators/llm_branch>`.
-* **Agents with tools** — :doc:`AgentOperator <operators/agent>` runs a multi-turn agent loop
- in the worker, calling Airflow-defined :doc:`toolsets <toolsets>` (SQL, hooks, MCP servers,
- a sandboxed shell and filesystem,
- :ref:`Agent Skills <agent-skills>`), optionally collapsed into a single sandboxed
- :ref:`code mode <code-mode>` call, with optional human-in-the-loop review and durable step
- replay — if the task retries after a failure, completed steps are replayed from cache
- instead of re-executing. Guardrails from the upstream ``pydantic-ai-shields`` package
- (``InputGuard``, ``OutputGuard``, ``ToolGuard``, ``CostTracking``) plug into the same agent
- loop (see :doc:`operators/agent`).
-* **Analyzing files or comparing schemas with an LLM** —
- :doc:`LLMFileAnalysisOperator <operators/llm_file_analysis>` reads a file (object storage or
- local) into a prompt; :doc:`LLMSchemaCompareOperator <operators/llm_schema_compare>` diffs
- schemas across systems and flags drift a plain equality check would miss.
-* **Generating SQL from natural language** —
- :doc:`LLMSQLQueryOperator <operators/llm_sql>` returns the generated query via XCom for
- ``SQLExecuteQueryOperator`` or a downstream task to run; it does not execute the query itself.
-* **Document pipelines for RAG** —
- :doc:`DocumentLoaderOperator <operators/document_loader>` parses files into structured text
- and metadata, :doc:`LlamaIndexEmbeddingOperator <operators/llamaindex_embedding>` embeds it,
- and :doc:`LlamaIndexRetrievalOperator <operators/llamaindex_retrieval>` retrieves the closest
- chunks for an :doc:`LLMOperator <operators/llm>` prompt (see :doc:`operators/index` for the
- full set).
-
-Use a vendor's own provider instead when the Dag needs that vendor's **native API surface** —
-a service the vendor runs for you, which no vendor-neutral operator wraps:
-
-* :doc:`apache-airflow-providers-openai:index` — the Embeddings and Responses APIs, and Batch
- jobs built from a pre-uploaded JSONL file of raw request bodies.
-* :doc:`apache-airflow-providers-anthropic:index` — Message Batches built from raw Messages
- API request bodies (multi-turn, images, tools), and Managed Agents sessions where the agent
- loop runs on Anthropic's infrastructure rather than in the Airflow worker.
-* :doc:`apache-airflow-providers-cohere:index` — Cohere's own Embed API.
-* :doc:`apache-airflow-providers-google:index` — Vertex AI's Batch Prediction jobs
- (``CreateBatchPredictionJobOperator``), a managed batch service like OpenAI's Batch API.
-* :doc:`apache-airflow-providers-amazon:index` — Bedrock's Batch Inference
- (``BedrockBatchInferenceOperator``), and Bedrock AgentCore's managed agent runtime
- (``BedrockCreateAgentRuntimeOperator`` / ``BedrockInvokeAgentRuntimeOperator``), where the
- agent loop runs on AWS's infrastructure rather than in the Airflow worker.
+``common.ai`` is built on `pydantic-ai <https://ai.pydantic.dev/>`__: the connection picks the
+model vendor, and Airflow runs the AI step like any other task. :doc:`concepts` explains the
+ideas behind the provider in one page; :doc:`operators/index` lists what each operator is for.
As a rule of thumb: if Airflow should *run* the AI step (and the model should stay
swappable), use ``common.ai``; if the Dag *submits work to* a vendor-managed service and
@@ -123,28 +65,12 @@ OpenAI, Anthropic, or other pydantic-ai-supported connection:
:start-after: [START howto_operator_llm_basic]
:end-before: [END howto_operator_llm_basic]
-Choosing extras
-----------------
+Getting started
+---------------
-The provider's extras split into a few groups:
-
-* **Model providers** — ``openai``, ``anthropic``, ``google``, ``bedrock``, ``typesafe``:
- pick the one matching your ``llm_conn_id`` connection. ``typesafe`` differs from the rest
- in kind: it installs a classifier model that answers typed questions and cannot write
- text (see :doc:`classifier_models`). Each extra name mirrors the identically named
- ``pydantic-ai-slim`` optional dependency group; pydantic-ai supports more model providers
- than these, each under its own extra name, so check the
- `pydantic-ai install docs <https://ai.pydantic.dev/install/#slim-install>`__ for the full list.
-* **Agent tooling** — ``mcp``, ``skills``, ``code-mode``, ``shields``: MCP servers, Agent
- Skills, code-mode tool execution, and shield capabilities (input/output guards, tool
- guards, cost tracking).
-* **Document loading** — ``pdf``, ``docx``, ``avro``, ``parquet``: file formats for
- document pipelines.
-* **Retrieval / SQL** — ``sql``, ``common.sql``, ``langchain``, ``llamaindex``: RAG and
- SQL-schema tooling.
-* **Git-backed content** — ``git``: pulling Agent Skills or documents from a git connection.
-
-See the Optional dependencies table below for the exact package each extra installs.
+* :doc:`installation` — which extra to install for your model vendor.
+* :doc:`quickstart` — a connection and a first ``@task.llm`` in three steps.
+* :doc:`concepts` — connections, operators, toolsets, hooks and XCom in one page.
.. toctree::
:hidden:
@@ -158,28 +84,98 @@ See the Optional dependencies table below for the exact package each extra insta
.. toctree::
:hidden:
:maxdepth: 1
- :caption: Guides
+ :caption: Getting started
+ Installation <installation>
Quick start <quickstart>
+ Core concepts <concepts>
+ Structured output <structured_output>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Models and providers
+
Pydantic AI connection <connections/pydantic_ai>
- Pydantic AI (Azure OpenAI) connection <connections/pydantic_ai_azure>
- Pydantic AI (AWS Bedrock) connection <connections/pydantic_ai_bedrock>
- Pydantic AI (Google Vertex AI) connection <connections/pydantic_ai_vertex>
- MCP connection <connections/mcp>
- LangChain connection <connections/langchain>
- LlamaIndex connection <connections/llamaindex>
- Hooks <hooks/index>
- Toolsets <toolsets>
- Sandboxed execution <sandbox>
- Choosing a toolset <choosing_a_toolset>
- Operators <operators/index>
- Examples <examples>
- Provider fallback <provider_fallback>
- Retry Policies <retry_policies>
+ Azure OpenAI <connections/pydantic_ai_azure>
+ AWS Bedrock <connections/pydantic_ai_bedrock>
+ Google Vertex AI <connections/pydantic_ai_vertex>
Self-hosted models <self_hosted_models>
Classifier models <classifier_models>
- HITL Review <hitl_review>
+ Provider fallback <provider_fallback>
+ PydanticAIHook <hooks/pydantic_ai>
+ LangChainHook <hooks/langchain>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Operators
+
+ Choosing an operator <operators/index>
+ LLMOperator <operators/llm>
+ LLMBranchOperator <operators/llm_branch>
+ LLMFileAnalysisOperator <operators/llm_file_analysis>
+ LLMSQLQueryOperator <operators/llm_sql>
+ LLMSchemaCompareOperator <operators/llm_schema_compare>
+ LLMBatchOperator <operators/llm_batch>
+ AgentOperator <operators/agent>
+ Approval gates <approval_gates>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Toolsets
+
+ Choosing a toolset <toolsets/index>
+ HookToolset <toolsets/hook>
+ SQLToolset <toolsets/sql>
+ DataFusionToolset <toolsets/datafusion>
+ LoggingToolset <toolsets/logging>
+ MCPToolset <toolsets/mcp>
+ AgentSkillsToolset <toolsets/skills>
+ Managed agent toolsets <toolsets/managed_agent>
+ LangChain bridge <toolsets/langchain>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Running agents
+
+ Message history <message_history>
+ Guardrails <guardrails>
+ Code mode <code_mode>
+ Sandboxed execution <sandbox/index>
+ HITL review <hitl_review>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Document and RAG pipelines
+
+ DocumentLoaderOperator <operators/document_loader>
+ LlamaIndexEmbeddingOperator <operators/llamaindex_embedding>
+ LlamaIndexRetrievalOperator <operators/llamaindex_retrieval>
+ LlamaIndex connection <connections/llamaindex>
+ LlamaIndexHook <hooks/llamaindex>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Reliability and operations
+
+ Durable execution <durable_execution>
+ LLM retry policies <retry_policies>
Observability <observability>
+ Securing agent tools <agent_security>
+ Troubleshooting <troubleshooting>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Examples
+
+ Examples by scenario <examples>
+ End-to-end pipelines <end_to_end_pipelines>
.. toctree::
:hidden:
diff --git a/providers/common/ai/docs/installation.rst b/providers/common/ai/docs/installation.rst
new file mode 100644
index 00000000000..794dbe94e15
--- /dev/null
+++ b/providers/common/ai/docs/installation.rst
@@ -0,0 +1,86 @@
+ .. 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.
+
+.. _howto/installation:
+
+Installation
+============
+
+The provider needs Airflow 3.0 or later. Install it with the extra that matches the
+model vendor your connection will point at:
+
+.. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[openai]"
+
+Quote the package name so the brackets survive your shell. Swap ``openai`` for
+``anthropic``, ``google`` or ``bedrock`` as needed, or combine several extras with a
+comma. The minimum versions of ``apache-airflow`` and ``pydantic-ai-slim`` are listed in
+the requirements table on the :doc:`landing page <index>`.
+
+Choosing extras
+---------------
+
+The provider's extras split into a few groups:
+
+* **Model providers** — ``openai``, ``anthropic``, ``google``, ``bedrock``, ``typesafe``:
+ pick the one matching your ``llm_conn_id`` connection. ``typesafe`` differs from the rest
+ in kind: it installs a classifier model that answers typed questions and cannot write
+ text (see :doc:`classifier_models`). The first four mirror the identically named
+ ``pydantic-ai-slim`` optional dependency groups, and ``typesafe`` adds the ``typesafe-sdk``
+ the built-in adapter talks to; pydantic-ai supports more model providers
+ than these, each under its own extra name, so check the
+ `pydantic-ai install docs <https://ai.pydantic.dev/install/#slim-install>`__ for the full list.
+* **Agent tooling** — ``mcp``, ``skills``, ``code-mode``, ``shields``, ``modal``: MCP servers,
+ Agent Skills, code-mode tool execution, shield capabilities (input/output guards, tool
+ guards, cost tracking), and the hosted Modal backend for :doc:`sandboxed execution <sandbox/index>`.
+* **Document loading** — ``pdf``, ``docx``, ``avro``, ``parquet``: file formats for
+ document pipelines.
+* **Retrieval / SQL** — ``sql``, ``common.sql``, ``langchain``, ``llamaindex``: RAG and
+ SQL-schema tooling.
+* **Git-backed content** — ``git``: pulling Agent Skills or documents from a git connection.
+
+The ``Optional dependencies`` table on the :doc:`landing page <index>` lists the exact
+package each extra installs.
+
+Features gated on the Airflow version
+-------------------------------------
+
+The provider runs on Airflow 3.0, but some features need a newer core:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 60 40
+
+ * - Feature
+ - Needs
+ * - :doc:`Approval gates <approval_gates>` and :doc:`HITL review <hitl_review>`
+ - Airflow 3.1
+ * - :doc:`LLM retry policies <retry_policies>`
+ - Airflow 3.3
+ * - :doc:`Durable execution <durable_execution>` without configuring
+ ``[common.ai] durable_cache_path`` (the task state store)
+ - Airflow 3.3
+
+Next steps
+----------
+
+* :doc:`quickstart` creates a connection and runs a first ``@task.llm``.
+* :doc:`connections/pydantic_ai` covers every connection field and the vendor-specific
+ connection types.
+* :doc:`installing-providers-from-sources` covers verifying and installing a release
+ artifact by hand.
diff --git a/providers/common/ai/docs/message_history.rst b/providers/common/ai/docs/message_history.rst
new file mode 100644
index 00000000000..f75f8d6b2b9
--- /dev/null
+++ b/providers/common/ai/docs/message_history.rst
@@ -0,0 +1,60 @@
+ .. 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.
+
+.. _message-history:
+
+Multi-turn sessions and message history
+=======================================
+
+By default each agent run is a cold, single-turn conversation. To carry a
+conversation across runs -- a chat or iterative agent where "and the third one?"
+must resolve against an earlier answer -- pass ``message_history``.
+
+When ``message_history`` is set, the operator seeds the run with those prior
+turns and, after the run, pushes the full updated transcript
+(``result.all_messages()``) to XCom under the key ``message_history``. The next
+run reads it back to resume the conversation. ``None`` (the default) keeps the
+single-turn behavior unchanged.
+
+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
+JSON file in object storage by ``session_id`` (use ``s3://`` / ``gs://`` in a
+deployment); the first run starts from an empty ``"[]"``.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
+ :language: python
+ :start-after: [START howto_agent_session]
+ :end-before: [END howto_agent_session]
+
+``message_history`` accepts a list of pydantic-ai ``ModelMessage`` objects or
+their JSON form (``str`` / ``bytes``), so the value emitted to XCom feeds
+straight back in on the next run. When pulling it via a template, pass
+``default='[]'`` (as above) so the first run -- which has no XCom yet -- starts a
+fresh session instead of trying to parse the string ``"None"``.
+
+The transcript is **cumulative**: each turn appends to it, so it grows for the
+life of the session. For long sessions, configure an object-storage XCom backend
+or trim older turns before the next run rather than feeding the whole history
+back unbounded.
+
+.. note::
+
+ ``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 0c040a26c56..acce8184a95 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -33,8 +33,7 @@ tool-call loop where the LLM decides which tools to call and when to stop.
.. seealso::
:ref:`Connection configuration <howto/connection:pydanticai>`
-
-SQL Agent
+SQL agent
---------
The most common pattern: give an agent access to a database so it can answer
@@ -62,8 +61,7 @@ The ``SQLToolset`` provides four tools to the agent:
* - ``check_query``
- Validates SQL syntax without executing it
-
-Hook-based Tools
+Hook-based tools
----------------
Wrap any Airflow Hook's methods as agent tools using ``HookToolset``. Only
@@ -74,8 +72,7 @@ methods you explicitly list are exposed — there is no auto-discovery.
:start-after: [START howto_operator_agent_hook]
:end-before: [END howto_operator_agent_hook]
-
-TaskFlow Decorator
+TaskFlow decorator
------------------
The ``@task.agent`` decorator wraps ``AgentOperator``. The function returns
@@ -86,7 +83,6 @@ the prompt string; all other parameters are passed to the operator.
:start-after: [START howto_decorator_agent]
:end-before: [END howto_decorator_agent]
-
.. _howto/operator:agent-multimodal:
Multimodal prompts
@@ -114,8 +110,7 @@ to the model. This mirrors the input types accepted by pydantic-ai's
string, so a ``Sequence`` prompt will raise at the review boundary.
Widening HITL review to multimodal prompts is tracked as a follow-up.
-
-Structured Output
+Structured output
-----------------
Set ``output_type`` to a Pydantic ``BaseModel`` subclass to get structured data
@@ -123,15 +118,8 @@ back. The model instance is pushed to XCom unchanged so downstream tasks can
type-hint the class directly (``def downstream(result: MyModel)``) and use
attribute access (``result.field``).
-The declared ``output_type`` (and any ``BaseModel`` reachable from
-``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization by
-the worker when it loads the Dag, before any task runs. The Pydantic class must
-be defined at **module scope** and bound to an attribute matching its
-``__name__``. Same-Dag downstream tasks need no configuration. The UI's XCom
-viewer renders the value via the ``stringify`` path (no configuration needed;
-see the ``LLMOperator`` guide for the exact representation). Cross-Dag
-``xcom_pull`` consumers still need the class ``qualname`` added to
-``[core] allowed_deserialization_classes``.
+:doc:`../structured_output` explains the XCom deserialization rules, the cross-Dag gap and
+``serialize_output``.
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
:language: python
@@ -143,9 +131,8 @@ see the ``LLMOperator`` guide for the exact representation). Cross-Dag
:start-after: [START howto_decorator_agent_structured]
:end-before: [END howto_decorator_agent_structured]
-
-Chaining with Downstream Tasks
--------------------------------
+Chaining with downstream tasks
+------------------------------
The agent's output is pushed to XCom like any other operator, so downstream
tasks can consume it.
@@ -155,11 +142,10 @@ tasks can consume it.
:start-after: [START howto_agent_chain]
:end-before: [END howto_agent_chain]
-
.. _howto/operator:agent-dynamic-system-prompt:
-Dynamic System Prompt
-----------------------
+Dynamic system prompt
+---------------------
``system_prompt`` is a templated field, so instead of a static string it
can be a Jinja expression that reads a value an earlier task already
@@ -175,300 +161,25 @@ Open the **Rendered Template** tab on the task instance to see the
substituted ``system_prompt`` after Jinja fills in ``classify``'s XCom
values.
+Agent features
+--------------
-Multi-turn Sessions
--------------------
-
-By default each agent run is a cold, single-turn conversation. To carry a
-conversation across runs -- a chat or iterative agent where "and the third one?"
-must resolve against an earlier answer -- pass ``message_history``.
-
-When ``message_history`` is set, the operator seeds the run with those prior
-turns and, after the run, pushes the full updated transcript
-(``result.all_messages()``) to XCom under the key ``message_history``. The next
-run reads it back to resume the conversation. ``None`` (the default) keeps the
-single-turn behavior unchanged.
-
-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
-JSON file in object storage by ``session_id`` (use ``s3://`` / ``gs://`` in a
-deployment); the first run starts from an empty ``"[]"``.
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
- :language: python
- :start-after: [START howto_agent_session]
- :end-before: [END howto_agent_session]
-
-``message_history`` accepts a list of pydantic-ai ``ModelMessage`` objects or
-their JSON form (``str`` / ``bytes``), so the value emitted to XCom feeds
-straight back in on the next run. When pulling it via a template, pass
-``default='[]'`` (as above) so the first run -- which has no XCom yet -- starts a
-fresh session instead of trying to parse the string ``"None"``.
-
-The transcript is **cumulative**: each turn appends to it, so it grows for the
-life of the session. For long sessions, configure an object-storage XCom backend
-or trim older turns before the next run rather than feeding the whole history
-back unbounded.
-
-.. note::
-
- ``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.
-
-
-Durable Execution
------------------
-
-Agent tasks can involve multiple LLM calls and tool invocations. If a task
-fails mid-run (network error, timeout, transient API failure), a plain retry
-re-executes every LLM call and tool call from scratch -- repeating work that
-already succeeded and incurring additional cost.
-
-Setting ``durable=True`` caches each LLM response and tool result as it
-completes. On retry, completed steps are replayed from the cache and only the
-remaining steps run against the live model and tools. The cache is deleted
-after successful completion.
-
-Durable execution only helps when the task has retries configured. Without
-retries there is nothing to replay.
-
-**Configuration**
-
-On **Airflow >= 3.3** the cache is stored in the
-:doc:`task state store <apache-airflow:core-concepts/task-state-store>`,
-scoped to the task instance. No configuration is required; the store handles
-persistence across retries.
-
-By default each cached step is written to the Airflow metadata database. Model
-responses and large tool results can be sizable, so for agents with large
-payloads configure ``[workers] state_store_backend`` to offload step values to
-external storage (e.g. object storage) instead of the metadata database; the
-provider then stores only a reference in the database.
-
-On **Airflow < 3.3** the cache is persisted to ObjectStorage and the location
-must be set in ``airflow.cfg``. The task raises ``ValueError`` at runtime if
-``durable=True`` and the option is missing.
-
-.. code-block:: ini
-
- [common.ai]
- # Local filesystem -- suitable for development
- durable_cache_path = file:///tmp/airflow_durable_cache
-
-The value is an ObjectStorage URI, so any supported backend works. For
-production, use a shared store so retries on a different worker can read the
-cache:
-
-.. code-block:: ini
-
- [common.ai]
- durable_cache_path = s3://my-bucket/airflow/durable-cache
+Four features have pages of their own:
-**Operator example**
+- :doc:`../message_history` — pass ``message_history`` to carry a conversation across runs.
+- :doc:`../durable_execution` — set ``durable=True`` to replay completed model and tool steps
+ on retry instead of paying for them again.
+- :doc:`../guardrails` — pass pydantic-ai capabilities and ``pydantic-ai-shields`` guardrails
+ through ``agent_params``.
+- :doc:`../code_mode` — set ``code_mode=True`` to collapse the agent's tools into a single
+ ``run_code`` tool the model drives by writing Python.
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
- :language: python
- :start-after: [START howto_operator_agent_durable]
- :end-before: [END howto_operator_agent_durable]
-
-**Decorator example**
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
- :language: python
- :start-after: [START howto_decorator_agent_durable]
- :end-before: [END howto_decorator_agent_durable]
-
-**How it works**
-
-1. On first execution, each LLM response and tool result is saved as the agent
- progresses, together with a fingerprint of the request that produced it
- (model, message history, settings, and tools for LLM steps; tool name,
- arguments, and call id for tool steps).
-2. If the task fails and Airflow retries it, completed steps are loaded from
- the cache and returned without calling the model or tool. Steps not yet in
- the cache proceed normally.
-3. Before a step is replayed, its stored fingerprint is compared against the
- current request. If anything changed between attempts -- the system
- prompt, the model, the toolset, model settings, or the conversation so
- far -- the stale entry is discarded, a warning is logged, and the step
- re-runs live. A divergence also invalidates the steps after it: re-running
- an LLM step produces fresh tool call ids, so tool results recorded under
- the old conversation no longer match. A changed agent costs a re-run; it
- never replays responses that belong to a different conversation.
-4. After successful completion, the cached steps are deleted.
-
-Replay verification compares the **requests** sent to models and tools, not
-the code behind them. Editing a tool's implementation between attempts does
-not invalidate an already-cached result for an identical call, and pointing
-``llm_conn_id`` at a different endpoint serving the same model name does not
-invalidate cached responses -- clear the cache to force a fully fresh run.
-
-After the run, a single INFO summary line reports how many steps were
-replayed vs executed fresh. Per-step detail is available at DEBUG level.
-
-The cache is scoped to a single task instance (Dag id, run id, task id, and
-map index), so each run replays only its own steps. On Airflow >= 3.3 the cache
-lives in the task state store and is removed when the Dag run is cleaned up; on
-Airflow < 3.3 it is a JSON file named ``{dag_id}_{task_id}_{run_id}.json`` (with
-``_{map_index}`` appended for mapped tasks) under the configured
-``durable_cache_path``.
-
-.. note::
-
- Runs that fail permanently (exhaust all retries) leave their cached steps
- behind. These do not affect future Dag runs (each run is scoped separately).
- On Airflow >= 3.3 they are reclaimed when the Dag run is removed; on Airflow
- < 3.3 the orphaned JSON files consume storage until cleaned up, so add a
- lifecycle policy to the storage backend or remove them periodically.
-
-**Side effects and idempotency**
-
-Durable execution caches **return values**, not side effects. When a step is
-replayed, the tool's code does not run -- only the stored return value is
-returned. Two things follow from this:
+.. _agent-durable-execution:
-- If a tool completed successfully and its result was cached, the tool will
- **not** run again on retry. Any side effect it produced (writing a file,
- sending a message) already happened during the original run and is not
- repeated.
-- If a tool fails *before* its result is cached, it **will** run again on
- retry. A tool that partially completed (e.g. sent an email then raised an
- exception) may produce the side effect a second time.
-
-All built-in toolsets (``SQLToolset`` with ``allow_writes=False``,
-``HookToolset`` in read-only mode) are read-only and replay safely. For custom
-tools with non-idempotent side effects, design the tool to be idempotent. For
-example, check whether the operation already completed before acting, or
-use database constraints to prevent duplicate writes.
-
-Tool results must be JSON-serializable to be cached. If a tool returns a
-non-serializable value (e.g. ``BinaryContent`` from MCP tools), that step is
-skipped with a warning and will re-execute on retry instead of replaying from
-cache. The task itself still succeeds.
-
-
-.. _capabilities-passthrough:
-
-Capabilities (pydantic-ai)
---------------------------
-
-pydantic-ai `capabilities <https://ai.pydantic.dev/capabilities/>`__ bundle
-tools, lifecycle hooks, instructions, and model settings into composable units.
-Common ones include ``Thinking`` (reasoning at a configurable effort level),
-``WebSearch``, ``WebFetch``, ``ImageGeneration``, and ``MCP``.
-For the current capability catalog and package-specific installation notes, see
-the pydantic-ai documentation and the
-`pydantic-ai-harness capability matrix <https://github.com/pydantic/pydantic-ai-harness#capability-matrix>`__.
-
-``AgentOperator`` does not yet expose a first-class ``capabilities=`` kwarg,
-but anything passed through ``agent_params`` is forwarded to the underlying
-``Agent(...)`` constructor.
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
- :language: python
- :start-after: [START howto_operator_agent_capabilities_thinking]
- :end-before: [END howto_operator_agent_capabilities_thinking]
-
-Capabilities compose with toolsets -- pydantic-ai merges tools from both.
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
- :language: python
- :start-after: [START howto_operator_agent_capabilities_composed]
- :end-before: [END howto_operator_agent_capabilities_composed]
-
-Guardrail capabilities use the same passthrough pattern. This example uses
-``InputGuard`` from ``pydantic-ai-shields`` to reject a prompt before the agent
-run starts.
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py
- :language: python
- :start-after: [START howto_operator_agent_capabilities_input_guard]
- :end-before: [END howto_operator_agent_capabilities_input_guard]
-
-.. warning::
-
- ``agent_params`` is a templated field, which Airflow serializes by calling
- ``str()`` on values it doesn't natively understand. Capability instances
- are not yet round-trip-safe through Dag serialization, so the examples
- below construct them inside the ``@dag`` function -- not at module level.
- First-class ``capabilities=`` support on ``AgentOperator`` (with proper
- serializer hooks) is tracked as a follow-up.
-
-
-.. _code-mode:
-
-Code Mode (Monty sandbox)
--------------------------
-
-Set ``code_mode=True`` to collapse the agent's tools into a single ``run_code``
-tool powered by the `Monty <https://github.com/pydantic/monty>`__ sandbox (via
-pydantic-ai-harness). Instead of one model round-trip per tool call, the model
-writes a single Python snippet that calls the tools as functions -- with loops,
-conditionals, and ``asyncio.gather`` -- in one turn. For multi-tool workflows
-this cuts round-trips and token use.
-
-The generated code runs in Monty's deny-by-default sandbox: it cannot read the
-filesystem, the network, or environment variables. It can only call the tools
-you registered. Code mode therefore does not widen what the agent can reach --
-the tools it calls still run in the worker -- it only changes how the model
-invokes them. See :ref:`Toolsets security <howto/toolsets>` for the tool
-boundary.
-
-When to use it
-^^^^^^^^^^^^^^
-
-Code mode pays off for **orchestration-heavy, computation-light** workflows:
-calling several tools, looping over their results, filtering, and combining them.
-Collapsing many sequential tool calls into one turn is where the round-trip and
-token savings come from -- the example above answers a per-customer question in a
-single ``run_code`` block instead of one model round-trip per customer.
-
-It is **not a general-purpose code runtime**. The generated code is only the glue
-between tool calls; every real capability must come from a tool. Monty runs a
-subset of Python and **cannot import third-party libraries** (pandas, numpy,
-requests, boto3, ...) and has no filesystem or network access. If a task needs to
-crunch data inline with a library, you have two options, both better than code
-mode:
-
-- **Push the work into a tool.** Do the aggregation in SQL (``SQLToolset``), or
- expose a hook method that returns the processed result (``HookToolset``). The
- tool runs in the full worker environment with all its dependencies, and code
- mode just orchestrates it.
-- **Give the agent a real environment**, with
- :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`. It hands
- the model a shell and a filesystem in a disposable sandbox off the worker, so
- third-party packages, a real interpreter and installed binaries are all
- available. It costs a sandbox per run and a second or so to provision, against
- well under a millisecond for Monty, so reach for it when inline library code is
- genuinely required rather than by default. See :doc:`../toolsets` for the
- backends and their limitations.
-
-The two are not exclusive: ``code_mode=True`` and a ``SandboxToolset`` can be
-enabled together, and the file tools fold into ``run_code`` while ``run_command``
-stays a tool of its own.
-
-Requires the ``code-mode`` extra::
-
- pip install "apache-airflow-providers-common-ai[code-mode]"
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
- :language: python
- :start-after: [START howto_operator_agent_code_mode]
- :end-before: [END howto_operator_agent_code_mode]
-
-Unlike passing a capability through ``agent_params`` (see
-:ref:`capabilities-passthrough`), ``code_mode`` is a plain boolean and is
-serialization-safe: the ``CodeMode`` capability is built at execution time, not
-stored on the serialized operator.
-
-.. note::
-
- Monty is pre-1.0. The ``code-mode`` extra is opt-in so its dependency churn
- never affects the base provider install.
+Durable execution
+^^^^^^^^^^^^^^^^^
+Moved to :doc:`../durable_execution`.
Parameters
----------
@@ -525,7 +236,7 @@ Parameters
of pydantic-ai ``ModelMessage`` objects or their JSON form (``str`` / ``bytes``).
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 `Multi-turn Sessions`_.
+ 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
@@ -550,7 +261,6 @@ Parameters
- ``hitl_poll_interval``: Seconds between XCom polls while waiting for a
human response. Default ``10``.
-
Logging
-------
@@ -598,11 +308,10 @@ To disable real-time tool logging while keeping the post-run summary:
enable_tool_logging=False,
)
-
Security
--------
.. seealso::
- :ref:`Toolsets — Security <howto/toolsets>` for defense layers,
+ :doc:`../agent_security` for defense layers,
``allowed_tables`` limitations, ``HookToolset`` guidelines, recommended
configurations, and the production checklist.
diff --git a/providers/common/ai/docs/operators/index.rst b/providers/common/ai/docs/operators/index.rst
index e11398540b0..fa411a54713 100644
--- a/providers/common/ai/docs/operators/index.rst
+++ b/providers/common/ai/docs/operators/index.rst
@@ -15,11 +15,11 @@
specific language governing permissions and limitations
under the License.
-Common AI Operators
-===================
+Choosing an operator
+====================
-Choosing the right operator
----------------------------
+By use case
+-----------
The common-ai provider ships several operators (and matching ``@task`` decorators). Use this table
to pick the one that fits your use case:
@@ -88,11 +88,36 @@ reasoning — but if you don't need tools, ``LLMOperator`` is simpler and more e
(text, CSV, JSON, PDF, DOCX) into ``list[dict(text, metadata)]`` for downstream embedding.
No AI framework dependency.
-Operator guides
----------------
-
-.. toctree::
- :maxdepth: 1
- :glob:
-
- *
+What the operators cover
+------------------------
+
+Use this provider when a Dag needs:
+
+* **Generation, classification, summarization, or structured extraction** —
+ :doc:`LLMOperator and @task.llm <llm>`, with Pydantic-typed output pushed to XCom.
+* **Many prompts at half the price** — :doc:`LLMBatchOperator and @task.llm_batch <llm_batch>`
+ submit prompts as one provider batch job (OpenAI or Anthropic), poll in deferrable mode for up to 24
+ hours, re-attach on retry instead of paying twice, and land results as JSONL on object storage.
+* **Branching on a model's decision** — :doc:`LLMBranchOperator <llm_branch>`.
+* **Agents with tools** — :doc:`AgentOperator <agent>` runs a multi-turn agent loop
+ in the worker, calling Airflow-defined :doc:`toolsets <../toolsets/index>` (SQL, hooks, MCP servers,
+ a sandboxed shell and filesystem,
+ :ref:`Agent Skills <agent-skills>`), optionally collapsed into a single sandboxed
+ :ref:`code mode <code-mode>` call, with optional human-in-the-loop review and durable step
+ replay — if the task retries after a failure, completed steps are replayed from cache
+ instead of re-executing. Guardrails from the upstream ``pydantic-ai-shields`` package
+ (``InputGuard``, ``OutputGuard``, ``ToolGuard``, ``CostTracking``) plug into the same agent
+ loop (see :doc:`agent`).
+* **Analyzing files or comparing schemas with an LLM** —
+ :doc:`LLMFileAnalysisOperator <llm_file_analysis>` reads a file (object storage or
+ local) into a prompt; :doc:`LLMSchemaCompareOperator <llm_schema_compare>` diffs
+ schemas across systems and flags drift a plain equality check would miss.
+* **Generating SQL from natural language** —
+ :doc:`LLMSQLQueryOperator <llm_sql>` returns the generated query via XCom for
+ ``SQLExecuteQueryOperator`` or a downstream task to run; it does not execute the query itself.
+* **Document pipelines for RAG** —
+ :doc:`DocumentLoaderOperator <document_loader>` parses files into structured text
+ and metadata, :doc:`LlamaIndexEmbeddingOperator <llamaindex_embedding>` embeds it,
+ and :doc:`LlamaIndexRetrievalOperator <llamaindex_retrieval>` retrieves the closest
+ chunks for an :doc:`LLMOperator <llm>` prompt (see the table above for the
+ full set).
diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst
index 4cbc45f5e3b..12225e2a271 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -31,7 +31,7 @@ returns the output as XCom.
.. seealso::
:ref:`Connection configuration <howto/connection:pydanticai>`
-Basic Usage
+Basic usage
-----------
Provide a ``prompt`` and the operator returns the LLM's response as a string:
@@ -41,23 +41,13 @@ Provide a ``prompt`` and the operator returns the LLM's response as a string:
:start-after: [START howto_operator_llm_basic]
:end-before: [END howto_operator_llm_basic]
-Structured Output
+Structured output
-----------------
-Set ``output_type`` to a Pydantic ``BaseModel`` subclass. The LLM is instructed
-to return structured data, and the model instance is pushed to XCom unchanged
-so downstream tasks can type-hint the class directly
-(``def downstream(result: MyModel)``) and use attribute access (``result.field``).
-
-The declared ``output_type`` (and any ``BaseModel`` reachable from
-``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization by
-the worker when it loads the Dag, before any task runs -- so no edit to
-``[core] allowed_deserialization_classes`` is needed. The Pydantic class must be
-defined at **module scope** and bound to an attribute matching its ``__name__``;
-classes nested inside a function or ``@dag``-decorated body, parameterized
-generics, and dynamically-built classes whose ``__name__`` does not match the
-attribute they are bound to cannot be re-imported, so they are skipped with a
-warning at worker startup and the value fails to deserialize at the consumer.
+Set ``output_type`` to a Pydantic ``BaseModel`` subclass and the model instance is pushed to
+XCom unchanged, so downstream tasks can type-hint the class and use attribute access.
+:doc:`../structured_output` explains the XCom deserialization rules, the cross-Dag gap and
+``serialize_output``.
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
:language: python
@@ -69,29 +59,7 @@ warning at worker startup and the value fails to deserialize at the consumer.
:start-after: [START howto_operator_llm_structured]
:end-before: [END howto_operator_llm_structured]
-Registration covers downstream tasks in the **same Dag**: every worker walks the
-loaded Dag's tasks at startup and registers each declared class, so it also works
-for mapped producers (``.expand(...)``) and for workers that load Dags from a
-cache that bypasses operator construction.
-
-The Airflow UI's XCom viewer renders Pydantic instances via the
-``stringify`` path, which produces a representation like
-``my_module.MyModel@version=1(field=value,...)`` without consulting the
-allow-list. It is not pretty (no field-by-field rendering today), but the value
-shows up; no configuration is required.
-
-The remaining gap is **cross-Dag** ``xcom_pull`` -- a task in a different Dag
-that pulls this XCom only parses its own Dag file, not the producer's, so the
-class is not auto-registered. Add the class qualified name to
-``[core] allowed_deserialization_classes`` (or a glob that matches it) to make
-that pattern work.
-
-If a downstream consumer needs the dict shape (e.g. forwarding to an external
-system that expects JSON-style payloads), pass ``serialize_output=True`` and the
-operator calls ``model_dump()`` before pushing to XCom. The pre-PR behavior is
-available on demand without giving up the typed default.
-
-Agent Parameters
+Agent parameters
----------------
Pass additional keyword arguments to the pydantic-ai ``Agent`` constructor
@@ -106,7 +74,7 @@ the full list of supported parameters.
.. _howto/operator:llm_usage_limits:
-Usage Limits
+Usage limits
------------
Set ``usage_limits`` to a
@@ -176,7 +144,7 @@ Airflow task retry re-renders ``usage_limits`` and starts a fresh count, and for
``Decimal("0.50")`` caps one run, so it is not a bound on what the task spends in
total.
-TaskFlow Decorator
+TaskFlow decorator
------------------
The ``@task.llm`` decorator wraps ``LLMOperator``. The function returns the
@@ -205,7 +173,6 @@ the full example. ``require_approval=True`` is not currently supported with a
``Sequence`` prompt -- the approval session model expects a string -- and will
raise at the approval boundary; widening that path is tracked as a follow-up.
-
Classification with ``Literal``
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -233,100 +200,18 @@ to process a list of items in parallel:
``system_prompt`` is templated identically on ``@task.llm``, so the same
upstream-XCom pattern applies here.
-Human-in-the-Loop Approval
+Human-in-the-loop approval
--------------------------
-Set ``require_approval=True`` to pause the task after the LLM generates its
-output and wait for a human reviewer to approve or reject it via the Airflow
-HITL interface. Optionally allow the reviewer to edit the output before
-approving with ``allow_modifications=True``, and set a deadline with
-``approval_timeout``.
-
-Human-in-the-loop review needs Airflow 3.1+, whether ``require_approval`` or a
-``decision_policy`` with ``on_uncertain="review"`` opens it. On an older core the
-operator raises ``AirflowOptionalProviderFeatureException`` when it is constructed, so the Dag file
-fails to import, and with it every Dag defined in that file. A dynamically mapped
-task (``.expand()``) is only constructed when it runs, so there the same error
-surfaces as a task failure -- still before the model is called.
-
-When ``approval_timeout`` expires without a review, the task fails by default.
-Set ``on_approval_timeout="approve"`` to return the generated output instead, so
-an unattended pipeline keeps moving. ``"reject"`` answers the review with a
-rejection, which still fails this operator; only
-:class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator`
-turns a rejection into a downstream skip. The chosen option is also
-pre-highlighted as the default in the review form, so ``"reject"`` makes
-Reject the primary button:
+Set ``require_approval=True`` to pause the task after the model answers and wait for a human
+reviewer. The full guide, including timeouts, notifiers and assigned reviewers, is
+:doc:`../approval_gates`.
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
- :language: python
- :start-after: [START howto_operator_llm_approval]
- :end-before: [END howto_operator_llm_approval]
-
-A pending review is not surfaced as a notification. Pass
-``approval_notifiers`` to tell the reviewers about it through any Airflow
-notifier (Slack, email, ...), the way
-:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
-``notifiers``. The notifiers run once the review is open and can reference
-the review ``{{ task.subject }}`` and ``{{ task.body }}`` in their templates,
-as the example above does. The ``@task.llm`` decorator and the operator
-subclasses accept the same parameter. A notifier whose delivery fails is
-logged and the task still waits for the review; a template error fails the
-task. A retry re-runs the LLM and re-notifies with the regenerated output,
-while the open review keeps the original subject and body.
-
-The default ``body`` contains the rendered prompt and the output. Where either
-is sensitive, template only ``{{ task.subject }}`` and a link to the review
-into channels outside Airflow's auth boundary.
-
-By default any user with the permission can answer the review. Pass
-``approval_assigned_users=[{"id": "<auth-manager-user-id>", "name": "<user-name>"}]``
-to restrict it to named reviewers, the way
-:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
-``assigned_users``. ``id`` is the user id reported by the auth manager: with
-the default ``SimpleAuthManager`` it is the username from
-``simple_auth_manager_users``; under the FAB auth manager it is the numeric
-user row id as a string, not the username. This needs Airflow 3.1+. On Airflow 3.1.0 through 3.1.5 both
-``id`` and ``name`` must match what the auth manager reports, so a wrong
-``name`` blocks the assigned reviewer as well as everyone else; from 3.1.6 only
-``id`` is compared. The list is stored when the review is first created:
-clearing the task re-runs it against the existing review row, so a changed
-list does not take effect.
-
-Reviewing Uncertain Output
+Reviewing uncertain output
--------------------------
-A classifier model such as TypeSafe's reports a confidence for every field
-of a structured output, in ``provider_details`` on the model response. It is a
-summary of how concentrated the model's probability distribution was, not the
-probability that the field is right. ``decision_policy=DecisionPolicy(min_confidence=0.7)``
-(import ``DecisionPolicy`` from ``airflow.providers.common.ai.operators.llm``)
-is the bar the least confident field has to clear for the operator to return
-the output by itself. Only fields that reported a confidence are compared: a
-field whose type reports none (a bounded float, where the probability is the
-answer) is not gated, and the record's ``confidence`` map shows which fields
-were. Under the bar, the policy's ``on_uncertain`` applies: ``"review"``
-(default) sends the output to human review through the same approval flow as
-``require_approval``, with the same ``approval_timeout``,
-``on_approval_timeout`` and notifier settings; ``"fail"`` fails the task with
-``LowConfidenceError`` (from ``airflow.providers.common.ai.exceptions``), which
-Airflow retries like any other failure unless a retry rule says otherwise. A
-text model reports no confidence for any field, which counts as uncertain, so
-switching the connection does not silently switch off a control you set.
-``require_approval=True`` keeps its meaning and always asks, and
-``on_uncertain="review"`` needs Airflow 3.1+ like it does. The policy is
-honoured by ``LLMOperator`` and ``LLMBranchOperator``; the SQL, schema-compare
-and file-analysis operators run their own ``execute`` and reject a policy with a
-bar at construction.
-
-With or without a bar, the operator pushes a ``decision`` XCom carrying the
-model name, the per-field ``confidence`` and ``probabilities`` (empty for a text
-model), the bar that applied, why the output went to review if it did, who
-decided, and the gate configuration in force.
-A pending record is checkpointed with the paused task and finalized from that
-copy on resume, not from the XCom. See :ref:`LLMBranchOperator <howto/operator:llm_branch>` for the
-record's fields; there ``proposed`` and ``action`` name the branches, while
-here they are ``null`` and the output itself is the return value.
+A ``decision_policy`` with a confidence bar sends output the model is unsure about to the
+same review flow. See :doc:`../approval_gates`.
Parameters
----------
diff --git a/providers/common/ai/docs/operators/llm_batch.rst b/providers/common/ai/docs/operators/llm_batch.rst
index 6a04393b2bf..b31a72c082c 100644
--- a/providers/common/ai/docs/operators/llm_batch.rst
+++ b/providers/common/ai/docs/operators/llm_batch.rst
@@ -204,6 +204,8 @@ resolve to the same model and the operator rejects a mixed batch before submitti
``request_params`` and a request's ``params`` cannot override the keys the operator manages
(the model, the messages, the token cap and the structured-output directive).
+.. _llm-batch-reattach:
+
Retries re-attach instead of re-submitting
--------------------------------------------
diff --git a/providers/common/ai/docs/quickstart.rst b/providers/common/ai/docs/quickstart.rst
index 450932e3cb2..3c23bd323d5 100644
--- a/providers/common/ai/docs/quickstart.rst
+++ b/providers/common/ai/docs/quickstart.rst
@@ -32,7 +32,7 @@ plan to use, and step 3 below makes a real, billed API call to that provider.
Install the provider together with the extra matching the model SDK you plan
to use — ``openai``, ``anthropic``, ``google``, or ``bedrock`` (see
-:doc:`index` for the full list of available extras). Replace ``<extra>``
+:doc:`installation` for the full list of available extras). Replace ``<extra>``
below with the one you need:
.. code-block:: bash
@@ -76,16 +76,16 @@ Structured output
^^^^^^^^^^^^^^^^^^
Need typed data instead of a string? Set ``output_type`` to a Pydantic
-``BaseModel`` and the model instance is pushed to XCom unchanged. See the
-"Structured Output" section of the :ref:`howto/operator:llm` guide for the
-full example and its XCom-deserialization requirements.
+``BaseModel`` and the model instance is pushed to XCom unchanged. See
+:doc:`structured_output` for the full example and its XCom-deserialization
+requirements.
Where to go next
-----------------
- :doc:`operators/index` — the full set of operators and ``@task`` decorators
(file analysis, SQL, branching, schema comparison).
-- :doc:`toolsets` — give an agent tools built from Airflow hooks, SQL
+- :doc:`toolsets/index` — give an agent tools built from Airflow hooks, SQL
databases, or MCP servers.
- :ref:`howto/operator:agent` — run a multi-turn agent that reasons and calls
tools instead of a single prompt-response call.
diff --git a/providers/common/ai/docs/redirects.txt b/providers/common/ai/docs/redirects.txt
new file mode 100644
index 00000000000..b7f775cf404
--- /dev/null
+++ b/providers/common/ai/docs/redirects.txt
@@ -0,0 +1,22 @@
+# 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.
+
+# Pages moved in the documentation reorganization.
+toolsets.rst toolsets/index.rst
+choosing_a_toolset.rst toolsets/index.rst
+sandbox.rst sandbox/index.rst
+hooks/index.rst concepts.rst
diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst
index f513479bc85..e6e165f1594 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-LLM Retry Policies
-===================
+LLM retry policies
+==================
.. note::
Requires Airflow >= 3.3.0.
@@ -62,6 +62,11 @@ model the rest, and the rules catch what neither decides. See
`Escalating to an LLM`_ below. Both work with any LLM provider supported by
pydantic-ai (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.).
+This page is about deciding *whether* a failed task should retry. To make an agent's
+retries cheap by replaying the model and tool calls that already succeeded, see
+:doc:`durable_execution`; a retried ``LLMBatchOperator`` re-attaches to its running batch
+instead of paying for it again (:ref:`llm-batch-reattach`).
+
For the core retry policy concepts, see :doc:`apache-airflow:core-concepts/tasks`.
If the task also needs to survive a worker crash without losing its progress,
see :ref:`apache-airflow:concepts-resumable-tasks-retry-policies`.
diff --git a/providers/common/ai/docs/sandbox.rst b/providers/common/ai/docs/sandbox.rst
deleted file mode 100644
index d88731f1d10..00000000000
--- a/providers/common/ai/docs/sandbox.rst
+++ /dev/null
@@ -1,821 +0,0 @@
- .. 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.
-
-Sandboxed execution for agents
-==============================
-
-An agent that is asked to do open-ended work writes code, and then something has
-to run that code. By default that something is the Airflow worker: a skill
-script, a hand-written tool that shells out, or generated glue in
-:ref:`code mode <code-mode>` all execute on the host that also holds your
-connections, your filesystem and your network position, and the code they run
-was generated a moment ago from inputs you do not control.
-
-:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` gives the
-model a disposable workspace for that code instead. It exposes four tools:
-
-.. list-table::
- :widths: 25 75
- :header-rows: 1
-
- * - Tool
- - What it does
- * - ``run_command``
- - Runs a shell command. Pipes, redirection, ``&&`` and globs work. A
- non-zero exit is reported as output, not raised, so the model reads
- ``stderr`` and corrects itself.
- * - ``read_file``
- - Reads a text file, head-first, and reports the next ``offset`` so the
- model can page through a long file.
- * - ``write_file``
- - Writes text to a file, creating parent directories.
- * - ``list_directory``
- - Lists a directory. Directories are shown with a trailing ``/``.
-
-The sandbox is provisioned by a
-:class:`~airflow.providers.common.ai.sandbox.SandboxBackend` on the model's first
-tool call and destroyed when the agent run ends. Two backends ship: a hosted one on
-`Modal <https://modal.com/docs/guide/sandbox>`__ for production and Kubernetes,
-and a local microVM one on `Docker Sandboxes <https://docs.docker.com/ai/sandboxes/>`__
-for development. The four tool names and shapes match pydantic-ai's own sandbox
-capabilities, so a model that has seen one already knows this one.
-
-**Adding this toolset gives the agent shell and file operations in a separate
-workspace. Every other tool keeps its existing permissions and runs where it ran
-before.** Whether that is a new capability depends on what the agent already had:
-an agent with a skill script or a shell tool already ran generated code on the
-worker, and this is where that code should run instead. An agent that only had
-named database operations is being granted a general shell for the first time, in
-a workspace with less authority than the worker, and you are choosing to grant it.
-
-.. _sandbox-quick-start:
-
-Quick start
------------
-
-An agent investigating a revenue anomaly. The warehouse is reached through a
-:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`, so the credential
-stays in the task and the model only ever sees rows. The arithmetic happens in a
-hosted sandbox with pandas baked into the image, so the sandbox needs no network,
-and a bad pivot or a runaway loop costs a small sandbox rather than a worker slot:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
- :language: python
- :start-after: [START howto_sandbox_agent_investigation]
- :end-before: [END howto_sandbox_agent_investigation]
-
-The same toolset on the local ``sbx`` backend, for developing on a laptop. Here
-the agent maps a vendor file whose columns drift onto a staging schema: it writes a
-script, runs it, reads the traceback and fixes it, which is the loop nobody can
-write down in advance. Swapping ``SbxSandboxBackend`` for ``ModalSandboxBackend``
-is the only change between this Dag and production:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
- :language: python
- :start-after: [START howto_sandbox_agent_local]
- :end-before: [END howto_sandbox_agent_local]
-
-Install the Modal extra and authenticate as you would for Modal's own CLI. On a
-worker, set ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` in the environment
-instead; nothing is read until the first sandbox is created, so a Dag file that
-constructs the backend parses without credentials present:
-
-.. code-block:: bash
-
- pip install 'apache-airflow-providers-common-ai[modal]'
- modal token new # writes ~/.modal.toml
-
-What it is for, in practice
----------------------------
-
-The rule is: a model inside a running task writes code you cannot name in advance,
-and that code needs a real shell, a real interpreter or packages. These are the
-jobs that rule describes, each paired with the version where a sandbox is the wrong
-call.
-
-**A vendor drops a CSV nobody has a schema for.** Three date formats, an encoding
-error at row 41,000, pandas needed. Every step depends on the previous step's
-output, and a bad read of the whole file into memory should exhaust a small
-sandbox, not a worker slot shared with other tasks. If the schema is stable and
-the mapping is known, this is a plain task with no agent.
-
-**Yesterday's revenue moved and nobody knows why.** The agent queries through a
-SQL toolset with a table allowlist, writes the rows into the sandbox, and pivots
-with pandas rather than doing arithmetic in tokens. The warehouse credential never
-enters the sandbox. If it is always the same five queries, that is
-:class:`~airflow.providers.common.ai.operators.llm.LLMOperator` over a fixed SQL
-task; if the rows number in the millions, compute in the warehouse rather than
-carrying them through the model's context.
-
-**Migrate forty dbt models between warehouse dialects.** The model files go into
-the workspace, the image carries ``sqlglot`` and ``sqlfluff``, and the loop of
-convert, lint, read errors, fix runs for as long as it takes, fully offline. The
-rewritten files are the deliverable, and getting forty files out is the gap
-described under :ref:`Getting a result out <sandbox-results>`.
-
-**A failing task's traceback, handed to an agent to propose a fix.** The agent
-reproduces the crash against a sample, tries a fix, reruns, and posts a diff for a
-person to review. That code ran with the worker's credentials the first time; the
-reproduction should not. A written diagnosis with no execution is
-``LLMOperator`` reading the log.
-
-**Not a sandbox: flag 200 customers with an open ticket** using two existing
-tools. That is :ref:`code mode <code-mode>`: a short loop over registered tools,
-no packages, no shell, sub-millisecond start.
-
-**Not a sandbox: run a vendor's PII-scanner image over an S3 prefix.** No model
-writes anything. That is ``KubernetesPodOperator``, a fixed payload with templated
-arguments in its own pod.
-
-**Not only a sandbox: one team's agent tasks must run off the shared workers.**
-That is a where-does-the-task-run requirement, so ``KubernetesExecutor`` with a
-pod override, and the sandbox toolset inside the pod if those agents also write
-code. The pod protects the cluster from the task; the sandbox protects the task
-from its model.
-
-.. _sandbox-boundaries:
-
-Choosing the boundary
----------------------
-
-"Sandboxed" means different things at different layers, and picking the wrong
-layer is the most common way to end up with less protection than you think. Four
-boundaries exist, from smallest to largest:
-
-.. list-table::
- :widths: 22 33 45
- :header-rows: 1
-
- * - Boundary
- - What moves inside it
- - What that protects you from
- * - **A tool call**
- - What ``run_command`` and the file tools do. **This toolset.**
- - Model-written code damaging the worker host, reading its files, or
- reaching the network from it.
- * - **The glue between tools**
- - Generated orchestration code, via :ref:`code mode <code-mode>` and the
- Monty interpreter.
- - Generated code touching anything other than the tools you registered.
- Monty runs it in a worker subprocess on the worker host, a language-level
- sandbox rather than an OS one, and every tool it calls executes in the
- task process with the task's authority.
- * - **The agent process**
- - The whole agent loop, its LLM credentials and its message history.
- - The agent's own credentials leaking, and any *other* toolset on the same
- agent. Not available today.
- * - **The whole task**
- - The complete Airflow task, supervisor included, as
- ``KubernetesExecutor`` does.
- - Everything outside the task. Not the task from its own code: the
- agent and what it runs are both inside this boundary.
-
-``SandboxToolset`` is the first row. It relocates exactly its four tools and
-nothing else. Boundary size is not a security level on its own: the image, the
-credentials you inject, the network policy and the resource limits decide the
-actual isolation. Choose the smallest boundary that fits, then configure it.
-
-.. list-table::
- :widths: 48 52
- :header-rows: 1
-
- * - The agent needs to
- - Reach for
- * - Call operations you can name in advance
- - :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` or
- :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset`. Narrow,
- bounded, and the credential stays in the worker rather than reaching the
- model. If you can name the operations, name them.
- * - Chain those tools with glue logic
- - :ref:`code mode <code-mode>`. Monty confines generated code to the tools
- you registered and starts in well under a millisecond, against roughly a
- second to provision a hosted sandbox. Reach past it for a sandbox when the
- model needs something Monty does not have: a shell, ``pip install``, a
- compiled library.
- * - Write and run open-ended code: reshape data with no known schema, install
- a package, or fix its own failing script by reading the traceback
- - ``SandboxToolset``
- * - Produce a large artifact for a downstream task
- - Today, a ``@task`` driving a backend directly, not the agent. That is a
- missing seam in the toolset rather than a recommendation; see
- :ref:`Getting a result out <sandbox-results>`.
- * - A whole task's worth of untrusted work isolated, with no agent involved
- - ``KubernetesPodOperator``
- * - Airflow's own credentials kept away from the agent
- - Not this. Scope the connections the task can see, and give the agent only
- the toolsets it needs. See :ref:`sandbox-security`.
-
-**Code mode and the sandbox compose.** With both enabled, the three file tools
-fold into ``run_code`` as callables the generated code can loop over, and
-``run_command`` stays a tool the model calls directly, so a shell command is
-written as a shell command rather than quoted inside generated Python. Both paths
-reach the same sandbox: there is one per agent run whichever way a call arrives.
-
-Why not ``KubernetesPodOperator`` or the ``KubernetesExecutor``?
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Because they answer a different question. Both isolate a task from the rest of
-your platform. Neither puts a boundary between an agent and the code that agent
-writes.
-
-Run an ``AgentOperator`` under the ``KubernetesExecutor`` and the whole task,
-agent loop included, gets its own pod. Model-written code still executes inside
-that pod, with its service account, its mounted secrets, its connections and its
-network position. The pod protects the cluster from the task; it does not protect
-the task from what the model decided to run, because from the pod's point of view
-that code is the task.
-
-``KubernetesPodOperator`` does contain the work, and its image, command, arguments
-and environment are all templated, so it can launch a different payload per run,
-and that payload can be a complete agent. The difference is what each one is. The
-pod is the environment that hosts the workload, and everything inside it runs with
-the pod's authority. ``SandboxToolset`` is a separate workspace that an
-already-running workload reaches into repeatedly, with less authority than the
-workload itself. Expressing the workspace as a pod means one pod per turn, with
-scheduling latency and a lost filesystem between each, or letting the model run
-whatever it wants inside the pod that also holds your credentials.
-
-.. list-table::
- :widths: 30 35 35
- :header-rows: 1
-
- * -
- - ``KubernetesExecutor`` / ``KubernetesPodOperator``
- - ``SandboxToolset``
- * - Protects
- - The cluster, from the task
- - The task and its credentials, from the model's output
- * - Boundary sits
- - Around the whole task, agent included
- - Between the agent and the code it writes
- * - What runs inside is decided
- - Before the process starts
- - By the model, mid-run
- * - Credentials inside
- - The task's own: service account, secrets, connections
- - Only what ``SandboxSpec.env`` names, which is nothing by default
-
-Running an agent in a pod *and* giving it a sandbox is the setup a Kubernetes
-deployment usually wants: the pod bounds the task, the sandbox bounds the model.
-The ``sbx`` backend cannot be the second half of that, because it needs KVM on the
-worker host and an unprivileged pod cannot provide it; that is what the hosted
-backend is for.
-
-.. _sandbox-configuring:
-
-Configuring a sandbox
----------------------
-
-:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` says what a sandbox is
-provisioned with. The default denies outbound network access and injects no
-environment, and a backend that cannot enforce a field it was given **raises
-instead of ignoring it**, so a spec never gives you a false sense of a restriction
-being in force.
-
-.. code-block:: python
-
- from airflow.providers.common.ai.sandbox import SandboxSpec, SbxSandboxBackend
- from airflow.providers.common.ai.toolsets import SandboxToolset
-
- SandboxToolset(
- SbxSandboxBackend(host_network_policy="allow-all"),
- spec=SandboxSpec(
- env={"HF_TOKEN": "..."}, # only what the generated code legitimately needs
- block_network=False, # this sandbox may reach the internet
- ),
- )
-
-**Network.** Three modes exist. ``block_network=True``, the default, drops all
-outbound traffic including name resolution. On Modal it maps onto the sandbox's
-own ``block_network`` flag. ``sbx`` has no per-sandbox enforcement of it at all:
-egress there is a host-level ``sbx policy``, so the backend honors the default by
-refusing to provision unless the Deployment Manager has declared
-``host_network_policy="deny-all"``, which means a bare ``SandboxSpec()`` is refused
-under the default ``host_network_policy="unknown"`` rather than silently getting
-an open sandbox. ``allow_egress_to`` names hosts the sandbox may reach, and both
-backends refuse it until you have said something that makes it enforceable:
-``sbx`` applies it as a per-sandbox rule on top of that ``deny-all`` host policy,
-since a local rule can narrow egress and never widen it, and Modal matches
-hostnames in the TLS handshake, which is weaker than it sounds and has to be
-opted into with ``egress_enforcement="sni"`` (see
-:ref:`the Modal backend <sandbox-backend-modal>` for exactly what it does and does
-not stop). ``block_network=False`` opens outbound access; on Modal the cloud
-metadata endpoint and private address ranges stay unreachable even then.
-
-**Packages.** A default sandbox has the Python standard library and no network,
-so an agent that reaches for ``pip install`` gets a DNS failure in a few seconds.
-Bake what it needs into the image, which needs no egress at all:
-
-.. code-block:: python
-
- import modal
-
- ModalSandboxBackend(
- image=modal.Image.from_registry("python:3.12-slim").pip_install("pandas", "pyarrow"),
- )
-
-Or, on Modal, allow exactly the two hosts a Python install needs and accept the
-handshake-name caveats: with ``pypi.org`` and ``files.pythonhosted.org`` allowed,
-installing pandas took about ten seconds when measured and ``github.com`` stayed
-blocked. Tell the model in its instructions what the image already has; the tool
-description says whether the sandbox has network access and where its working
-directory is, but not which packages are installed.
-
-**Resources.** A default Modal sandbox is a fraction of a core (``nproc`` reports
-1) with the standard library and no ``curl``, ``git`` or ``wget``. Set ``cpu`` for
-anything heavier. ``memory`` on Modal is a scheduling request, not a ceiling: a
-sandbox created with ``memory=512`` allocated 1.5 GiB without complaint when
-measured, so do not reach for it to bound what model-written code can consume.
-``sbx`` gives a sandbox every host CPU by default and enforces its ``memory`` as a
-limit.
-
-**Timeouts.** Three clocks apply. A command has a budget: ``default_command_timeout``
-(60 s) when the model does not ask, up to ``max_command_timeout`` (300 s) when it
-does. The sandbox has a lifetime: ``sandbox_timeout`` on Modal, 3600 s by default,
-and the clock runs through the model's thinking time between tool calls, so size it
-against the whole agent run. And Modal can reclaim an idle sandbox after
-``idle_timeout``, which defaults to ``None`` because a gap for model generation
-looks idle and would take the agent's files with it.
-
-**Two sandboxes on one agent** need ``tool_prefix``, since tool names must be
-unique within an agent:
-
-.. code-block:: python
-
- toolsets = [
- SandboxToolset(
- SbxSandboxBackend(image="python:3.12-slim", host_network_policy="deny-all"),
- tool_prefix="py",
- ),
- SandboxToolset(
- SbxSandboxBackend(image="node:22-slim", host_network_policy="deny-all"),
- tool_prefix="node",
- ),
- ]
-
-That yields ``py_run_command``, ``node_run_command`` and so on. Without a prefix
-on at least one of them the run fails at startup. Give the model instructions on
-which one to use for what, or it will guess.
-
-.. _sandbox-credentials:
-
-Credentials
------------
-
-``SandboxSpec.env`` is the only way in. Airflow never populates it: no connection,
-variable or worker environment variable reaches a sandbox unless you name it
-there, and the credential that *provisions* the sandbox never enters it either.
-Modal's token stays on the worker and is used by the client, so code running
-inside cannot call Modal as you or create further sandboxes.
-
-Before you put a real secret in ``env``, four things are true of it.
-
-**The model can read it.** Model-generated code can print the environment, and a
-tool result travels into the model's context, the task log and whatever the agent
-returns to XCom. Treat a credential given to a sandbox as disclosed to the model
-and to everything that records the run, and scope it to that one job. If the model
-only needs to *use* a system rather than hold its credential, a
-:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` or
-:class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` runs the hook in
-the task and returns only the result, which is the better shape whenever the
-operations can be named.
-
-**The value is fixed when the Dag file is parsed.** ``toolsets`` is not a
-templated field, so a spec cannot carry a Jinja expression, and a connection
-lookup written beside it would run in the Dag processor on every parse. An agent's
-sandbox therefore cannot take a credential from an Airflow connection or a secrets
-backend today. If it has to, the job cannot be an agent's sandbox: drive a backend
-from a ``@task``, where you are in ordinary Python at run time and the lookup runs
-where it should.
-
-.. code-block:: python
-
- @task
- def run_with_credential():
- from airflow.providers.common.ai.sandbox import ModalSandboxBackend, SandboxSpec
- from airflow.providers.common.compat.sdk import BaseHook
-
- token = BaseHook.get_connection("my_api").password # resolved in the task, not at parse
- backend = ModalSandboxBackend()
- sandbox = backend.create(spec=SandboxSpec(env={"API_TOKEN": token}, block_network=False))
- try:
- ... # write_file, run_command, read_file
- finally:
- backend.destroy(sandbox)
-
-**On** ``sbx`` **the value is written to the guest filesystem.** The CLI has no
-create-time environment flag, so the backend appends ``export`` lines to
-``/etc/profile`` inside the microVM, readable by anything in that sandbox for its
-whole life. Modal passes the environment at creation and writes no file. Prefer the
-hosted backend when the sandbox needs a secret at all.
-
-**Non-string values are refused** at provisioning, with the offending keys named,
-rather than failing somewhere less obvious.
-
-.. _sandbox-results:
-
-Getting a result out
---------------------
-
-An agent's result is whatever the model returns: findings, a mapping, a table
-small enough to read. That reaches XCom through ``output_type`` like any other
-agent output, and the examples above end that way.
-
-What does not come out is a *file* the agent built. The sandbox is destroyed when
-the run ends; ``read_file`` is text-only, replacing any byte it cannot decode, so a
-parquet or an image cannot survive the round trip at any size; and everything the
-model reads is capped, 50 KiB per stream for ``run_command`` and ``max_read_bytes``
-(5 MiB) per ``read_file``. Those caps are a budget rather than a transport limit:
-a 200 MB file reads out of a live sandbox in under ten seconds on the same path,
-but raising the cap costs roughly three times the file size in worker memory to
-show the model 50 KiB of it. The toolset does not yet have a seam for an author to
-collect what the agent produced; that is a known gap rather than the intended
-design.
-
-When the deliverable is a file and the Dag already knows the job, do not use an
-agent for it. Drive a backend from a ``@task``: the input goes in through
-``write_file``, the output comes back through ``read_file`` with a budget the
-worker can afford rather than a model's, and it lands in object storage with only
-the location passed downstream. The sandbox needs no network and no credentials
-for this, and the task owns its teardown:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
- :language: python
- :start-after: [START howto_sandbox_task_artifact]
- :end-before: [END howto_sandbox_task_artifact]
-
-If the agent genuinely has to be the one producing the file, the shape that works
-is to give the sandboxed code egress to your bucket and a scoped credential in the
-spec, and have it write the object itself, passing the key back through its
-output. That widens the default deliberately and is worth a sentence in the Dag
-saying why.
-
-.. _sandbox-lifecycle:
-
-Lifecycle, failures and retries
--------------------------------
-
-The sandbox is created lazily on the first tool call, shared by every call in that
-agent run, and destroyed when the run ends. A run that never calls a tool never
-provisions one, and concurrent runs never share a sandbox. Files written by one
-call are visible to later calls in the same run; each ``run_command`` is a fresh
-shell, so shell variables and background jobs do not survive between calls.
-
-**A failed command is output, not a failure.** A non-zero exit comes back to the
-model with both streams and ``[exit code: N]`` appended, and the model corrects
-itself. A recoverable sandbox error, such as a missing path or a file over the
-read budget, becomes a bounded retry prompt. Only a terminal failure, such as
-rejected credentials or a sandbox that has gone, fails the task, so Airflow's own
-retry handles it rather than the model burning its retry budget.
-
-**A command timeout** reports ``[timed out after Ns]`` with the budget that was
-actually applied. On Modal the sandbox survives and files written earlier are
-still there for the model to inspect. On ``sbx`` the backend has to destroy the
-sandbox to be certain the command stopped, and the result then says
-``[sandbox was replaced; files from earlier calls are gone]`` so the model knows
-to start over. Any other failure that loses the sandbox is reported the same way.
-
-**A run that outlives its sandbox fails the task.** When ``sandbox_timeout`` or
-``idle_timeout`` passes mid-run, the next tool call reaches a sandbox that is gone,
-which is terminal. This is the normal end of any run longer than its sandbox's
-lifetime. Human output review is not part of the run: it starts after the agent
-has finished and the sandbox has been destroyed, so a review pause costs no
-sandbox time and keeps no files.
-
-**Nothing survives the run.** A task retry starts from an empty sandbox, and so
-does every other attempt. Two operator features assume otherwise and must not be
-combined with a sandbox today, because neither is rejected:
-
-- ``durable=True`` caches each tool result and replays it on a retry without
- calling the backend, so a replayed ``write_file`` reports success while no
- sandbox exists, and the first call that misses the cache runs against a fresh
- empty one. The model is handed a filesystem that does not match what it was just
- told, and nothing raises.
-- ``enable_hitl_review=True`` regenerates after reviewer feedback by starting a
- second agent run, and the first run's sandbox was destroyed when that run ended.
- The regenerated agent gets an empty sandbox while its own history describes
- files it wrote earlier.
-
-Cost and operations
--------------------
-
-**One sandbox spans a whole agent run**, so you pay for the model's thinking time
-between tool calls, not only for the seconds your commands run. A default Modal
-sandbox is around two cents an hour, so this rarely matters; a data-sized one
-(``cpu=8, memory=32768``) is around two dollars an hour, and Modal bills the larger
-of what you requested and what you used, which is the second reason ``memory`` is
-worth setting deliberately. With ``idle_timeout`` unset, ``sandbox_timeout`` is
-the worst-case bill for one run. Fan-out is not the thing to worry about: 25
-sandboxes created concurrently from one process all came up in under three
-seconds.
-
-**Teardown runs inside the task.** A graceful end and an ordinary exception both
-destroy the sandbox. A worker killed outright, an OOM or a lost node never run that
-code, and what happens next is the backend's business alone: Modal reclaims the
-sandbox at ``sandbox_timeout`` whatever became of the worker; ``sbx`` has no
-server-side lifetime, so the microVM and its workspace directory survive with no
-owner. A teardown that fails is logged with the sandbox identity and does not fail
-the task.
-
-**Finding what a Dag left behind.** Modal sandboxes are named ``airflow-sandbox-*``
-and carry an ``airflow_sandbox`` tag plus whatever ``tags`` you pass, and
-``modal.Sandbox.list(app_id=..., tags={"dag_id": "my_dag"})`` returns exactly those
-still running. Tags are fixed when the backend is constructed, so every sandbox in
-a mapped task carries the same ones. When checking whether anything is actually
-still running, trust ``Sandbox.list``, which returns only live sandboxes, or
-``poll()`` on one you hold: ``None`` means running, and any integer means gone,
-including ``137``, which is a stopped sandbox rather than a sick one. Modal's
-dashboard lists stopped sandboxes alongside live ones, and the ``Tasks`` count in
-``modal app list`` lags by up to about a minute. ``sbx`` sandboxes are named the
-same way for an operator sweep.
-
-Backends
---------
-
-.. _sandbox-backend-modal:
-
-Modal (hosted)
-^^^^^^^^^^^^^^
-
-:class:`~airflow.providers.common.ai.sandbox.modal.ModalSandboxBackend` runs each
-sandbox in Modal, provisioned over the API. Of the backends that ship with the
-provider, **this is the one to use in production**, and the only one that runs on
-Kubernetes: nothing has to be installed on the worker, model-written code never
-executes on the worker host, and Modal reclaims a sandbox at its own lifetime
-whether or not the worker survives. It needs the ``modal`` extra and ambient
-credentials, as under :ref:`Quick start <sandbox-quick-start>`.
-
-Constructor parameters:
-
-- ``image``: Registry tag for the sandbox image, or a prepared ``modal.Image``
- carrying pre-installed packages. Default ``"python:3.12-slim"``.
-- ``app_name``: Modal app the sandboxes are created under. Default
- ``"airflow-sandbox"``.
-- ``create_app_if_missing``: Create that app if it does not exist. Default ``True``.
-- ``sandbox_timeout``: Maximum lifetime of a sandbox in seconds. Default ``3600``.
- Modal's own default is 300, which is below a plausible agent run.
-- ``idle_timeout``: Seconds of inactivity after which Modal reclaims the sandbox,
- or ``None`` to rely on ``sandbox_timeout`` alone. Default ``None``; see
- :ref:`Configuring a sandbox <sandbox-configuring>`.
-- ``workdir``: Working directory for commands, created if the image lacks it.
- Default ``"/workspace"``. It is a starting directory, not a jail: absolute paths
- and ``..`` are passed through, and commands run as root, so the model can read
- and write anywhere in the sandbox filesystem. The sandbox boundary is what
- contains that.
-- ``cpu``, ``memory``, ``gpu``, ``region``, ``cloud``: passed through to Modal.
- ``None`` lets Modal choose. An unrecognized ``region`` or ``cloud`` fails the
- task rather than falling back.
-- ``tags``: Extra Modal tags on every sandbox, e.g. ``{"dag_id": "my_dag"}``.
-- ``egress_enforcement``: ``"strict"`` (default) or ``"sni"``. See below.
-
-**Network policy.** ``block_network=True`` maps exactly onto Modal's own
-``block_network``, which drops all outbound traffic including DNS. A spec that
-names ``allow_egress_to`` is **refused by default**, because Modal cannot combine
-an allowlist with ``block_network`` at all, and its hostname allowlist is enforced
-by matching the name in the TLS handshake, which means:
-
-- TLS on port 443 to a listed host connects; any other host is refused at once.
-- **Non-TLS traffic is dropped, not refused.** A plain HTTP connection to a listed
- host stalls until the client gives up, around two minutes of TCP retries for one
- address, so with the 60 s default command budget the model reads
- ``[timed out after 60s]`` and concludes its command was slow, never that the
- network stopped it.
-- **The destination address is not part of the decision.** A connection opened to
- an unrelated address while presenting a listed name is routed to the listed host
- and answered by it: connecting to ``8.8.8.8:443`` with ``pypi.org`` in the
- handshake returns the certificate and content of ``pypi.org``. The allowlist is a
- name-routed egress proxy, not a filter on where packets may go.
-- **DNS resolution stays open for every hostname**, listed or not, against
- authoritative servers outside Modal. A freshly generated label under a domain the
- operator controls resolves and returns its answer, so this is a two-way channel.
-- A host that shares a **TLS endpoint** with a listed one can be reached by
- presenting the listed name in the handshake and the other in the request. With
- ``pypi.org`` as the only allowed host, a TLS session opened to
- ``files.pythonhosted.org`` while presenting ``pypi.org`` was allowed through and
- answered. Other tenants of the same CDN returned ``421 Misdirected Request``; it is
- co-tenancy of the same TLS endpoint that matters, which you cannot check from
- outside and which can change without notice.
-
-So the allowlist says which name a TLS session may be routed to, and nothing else.
-Pass ``ModalSandboxBackend(egress_enforcement="sni")`` to say you accept that and
-have the allowlist applied:
-
-.. code-block:: python
-
- SandboxToolset(
- ModalSandboxBackend(egress_enforcement="sni"),
- spec=SandboxSpec(block_network=True, allow_egress_to=["pypi.org", "files.pythonhosted.org"]),
- )
-
-Entries must be bare hostnames or one leading ``*.`` label; a URL, a ``host:port``,
-an address or a single-label name is refused, because Modal applies the list without
-checking it and any of those would silently match nothing.
-
-**What the image needs.** ``write_file`` and ``list_directory`` use Modal's own
-filesystem API, served by a helper Modal injects into the sandbox, so they need
-nothing from the image. ``read_file`` deliberately does not: Modal's read API takes
-no length, so it cannot honor a read budget, and a single call on a file that
-streams without end (``/dev/zero``, a FIFO, a procfs entry) would pull it into
-worker memory unbounded. That tool runs the base class's shell implementation,
-which caps the read inside the guest, and the image needs ``stat``, ``head`` and
-``base64``. Any Debian or Ubuntu based image, including ``python:*-slim``, has
-them.
-
-**Symlinks.** Because ``write_file`` goes through the native API, writing to a path
-that is a symlink replaces the link with a regular file and leaves the original
-target untouched, where a shell redirect would follow the link.
-
-sbx (Docker Sandboxes, local)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-:class:`~airflow.providers.common.ai.sandbox.SbxSandboxBackend` runs each sandbox
-in a Docker Sandboxes microVM by driving the ``sbx`` CLI. Each sandbox is a real
-microVM with its own kernel.
-
-.. warning::
-
- **Use this backend for local development, not production.** Docker Sandboxes
- is built for running coding agents against a checkout on your own machine, and
- driving it from an Airflow worker is off-label use. A production worker would
- need the ``sbx`` binary on the host, an authenticated Docker account
- (``sbx login``), a one-time ``sbx policy init``, and on Linux, KVM or nested
- virtualization, which a worker in an unprivileged container cannot provide.
-
- **Orphans are not reclaimed.** There is no server-side lifetime. If the worker is
- killed outright, the microVM and its workspace directory survive; sandboxes are
- named ``airflow-sandbox-*`` so an operator can find and remove them.
-
-Installing the CLI is a Deployment Manager prerequisite (``brew install
-docker/tap/sbx`` or ``winget install Docker.sbx``); the backend needs no Python
-dependency. The template image must provide GNU coreutils ``timeout``, ``base64``,
-``stat``, ``head``, ``find``, ``mkdir`` and ``dirname``, which any Debian or Ubuntu
-based image has.
-
-Constructor parameters:
-
-- ``image``: Container image for the sandbox. Default ``"python:3.12-slim"``.
-- ``memory``: Memory limit in binary units. ``sbx`` enforces a 1 GiB minimum.
- Default ``"2g"``.
-- ``cpus``: CPUs to allocate. ``None`` (default) uses the ``sbx`` default, which
- is every host CPU.
-- ``sbx_path``: Path to the ``sbx`` binary. Default ``"sbx"``.
-- ``create_timeout``: Seconds allowed for provisioning; a first-run microVM boot
- plus an image pull can be slow. Default ``600``.
-- ``host_network_policy``: What ``sbx policy`` is set to on this host.
- ``"unknown"`` (default) makes ``create`` refuse any spec asking for a network
- guarantee this backend cannot make, and since ``block_network`` defaults to
- ``True`` that includes a bare ``SandboxSpec()``. Set ``"deny-all"`` after running
- ``sbx policy init deny-all``, or ``"allow-all"`` to state that egress is open
- and pass ``SandboxSpec(block_network=False)`` to match.
-
-What differs between the two
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Swapping the backend is one constructor argument, and tool names, spec and prompt
-do not change. Four behaviours do, so read them before assuming the same Dag
-behaves identically in both places:
-
-- **CPU.** ``sbx`` gives a sandbox every host CPU; Modal defaults to a fraction of
- one, so set ``cpu``.
-- **Egress allowlists.** ``sbx`` enforces ``allow_egress_to`` at the host policy
- layer; Modal matches TLS handshake names, which is weaker and has to be opted
- into.
-- **Command timeouts.** A timeout destroys an ``sbx`` sandbox and its files; a
- Modal sandbox survives with its files intact.
-- **Symlinks.** ``write_file`` through a symlink follows the link on ``sbx`` and
- replaces it on Modal.
-
-Bringing your own backend
-^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Any vendor that can create a sandbox, run a command in it and destroy it can plug
-in. Subclass :class:`~airflow.providers.common.ai.sandbox.SandboxBackend` in your
-own package and pass an instance to ``SandboxToolset``.
-
-**Three methods are required**: ``create``, ``run_command`` and ``destroy``. The
-three file operations ship as defaults implemented over ``run_command``, because
-reading, writing and listing a file are all expressible as shell commands.
-Override them only when the vendor has a native file API:
-
-.. code-block:: python
-
- from airflow.providers.common.ai.sandbox import (
- SandboxBackend,
- SandboxExecResult,
- SandboxSpec,
- )
-
-
- class AcmeSandboxBackend(SandboxBackend):
- name = "acme"
-
- def create(self, *, spec: SandboxSpec | None = None) -> str:
- return acme_sdk.create_sandbox().id
-
- def run_command(self, sandbox, command, *, timeout, max_output_bytes):
- r = acme_sdk.exec(sandbox, command, timeout=timeout)
- return SandboxExecResult(exit_code=r.exit_code, stdout=r.stdout, stderr=r.stderr)
-
- def destroy(self, sandbox) -> None:
- acme_sdk.delete_sandbox(sandbox)
-
- # Optional: inherited from SandboxBackend unless the vendor has
- # something better than shelling out.
- def read_file(self, sandbox, path, *, max_bytes) -> bytes:
- return acme_sdk.download(sandbox, path, limit=max_bytes)
-
-Four rules for an implementation:
-
-- Constructors run at Dag-parse time, so resolve credentials lazily, on first use.
-- ``destroy`` must be idempotent; destroying an already-gone sandbox is not an error.
-- Raise ``SandboxTerminalError`` when retrying cannot help and ``SandboxError``
- when it might. The first fails the task for Airflow to retry; the second
- becomes a bounded prompt back to the model.
-- If you cannot enforce something the ``SandboxSpec`` asks for, **raise**. Never
- provision a weaker sandbox than the Dag author asked for.
-
-.. _sandbox-limitations:
-
-Limitations
------------
-
-These apply to every backend. Each is explained in the section it belongs to; this
-is the list to read before designing a Dag around an agent with a sandbox.
-
-- **Nothing survives the run**, including across task retries.
- :ref:`Lifecycle <sandbox-lifecycle>`.
-- **A file the agent built cannot leave** except through the model's context,
- which is text-only and capped. :ref:`Getting a result out <sandbox-results>`.
-- **A credential cannot come from a connection or a secrets backend**; the spec is
- fixed at parse time, and anything injected is readable by the model.
- :ref:`Credentials <sandbox-credentials>`.
-- **Do not combine with** ``durable=True`` **or** ``enable_hitl_review=True``.
- Neither is rejected today. :ref:`Lifecycle <sandbox-lifecycle>`.
-- **A run that outlives** ``sandbox_timeout`` **fails the task.**
- :ref:`Lifecycle <sandbox-lifecycle>`.
-- **The hostname allowlist is a weak control** and refused unless opted into.
- :ref:`Modal <sandbox-backend-modal>`.
-- **Commands run as root and** ``workdir`` **is not a jail.**
- :ref:`Modal <sandbox-backend-modal>`.
-- **It does not contain the agent.** :ref:`sandbox-security`.
-
-Parameters
-----------
-
-- ``backend``: The backend that provisions and drives the sandbox.
-- ``spec``: What to provision it with. Defaults to no environment and no egress.
-- ``default_command_timeout``: Seconds for a ``run_command`` the model did not
- put a timeout on. Default ``60``.
-- ``max_command_timeout``: Hard ceiling for any single command, including a
- model-supplied ``timeout_seconds``. Default ``300``.
-- ``max_output_lines`` / ``max_output_bytes``: Caps per output stream and per
- file read, whichever is hit first. Defaults ``2000`` and 50 KiB. Command
- output keeps the **tail**, where errors and the exit status live; file reads
- keep the head and report a continuation offset.
-- ``max_read_bytes``: Largest file ``read_file`` will transfer. Default 5 MiB;
- larger files are refused with a hint to slice them in the shell.
-- ``tool_prefix``: Prefix for the four tool names. Needed when one agent has
- more than one ``SandboxToolset``.
-
-.. _sandbox-security:
-
-Security boundary
------------------
-
-**The sandbox does not contain the agent.** The agent loop, the model calls and
-every other toolset on the same agent still run in the worker process with the
-worker's credentials. An agent that has ``SandboxToolset`` *and* a toolset that
-can reach connections has a contained code tool sitting beside a credential path
-that is not contained: a prompt injection arriving in a queried row can steer the
-model into calling the credentialed tool, and that call executes from the worker
-with the connection's full grants. Moving the sandbox to a hosted backend changes
-none of this.
-
-The controls for that risk are different controls. Give the agent only the
-toolsets it needs. Point every connection at a least-privilege role. Prefer a
-capability toolset, which lets the model *use* a connection without ever holding
-it, over a credential in ``SandboxSpec.env``, which the model can read. And treat
-a table allowlist as a guardrail that contains intent rather than a boundary that
-contains access.
-
-What the sandbox does contain, verified against a live worker environment carrying
-sentinel credentials: nothing Airflow-shaped reaches it. No connection, variable,
-Fernet key or cloud credential appeared in the guest's environment, in any process
-environment or on its filesystem, and the provisioning token is not usable from
-inside even when the vendor's API is on the egress allowlist. Sandboxes cannot reach each
-other, so two sandboxes on one agent are isolated from one another as well as from
-the worker.
diff --git a/providers/common/ai/docs/sandbox/backends.rst b/providers/common/ai/docs/sandbox/backends.rst
new file mode 100644
index 00000000000..0843cf69270
--- /dev/null
+++ b/providers/common/ai/docs/sandbox/backends.rst
@@ -0,0 +1,221 @@
+ .. 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.
+
+Sandbox backends
+================
+
+.. _sandbox-backend-modal:
+
+Modal (hosted)
+--------------
+
+:class:`~airflow.providers.common.ai.sandbox.modal.ModalSandboxBackend` runs each
+sandbox in Modal, provisioned over the API. Of the backends that ship with the
+provider, **this is the one to use in production**, and the only one that runs on
+Kubernetes: nothing has to be installed on the worker, model-written code never
+executes on the worker host, and Modal reclaims a sandbox at its own lifetime
+whether or not the worker survives. It needs the ``modal`` extra and ambient
+credentials, as under :ref:`Quick start <sandbox-quick-start>`.
+
+Constructor parameters:
+
+- ``image``: Registry tag for the sandbox image, or a prepared ``modal.Image``
+ carrying pre-installed packages. Default ``"python:3.12-slim"``.
+- ``app_name``: Modal app the sandboxes are created under. Default
+ ``"airflow-sandbox"``.
+- ``create_app_if_missing``: Create that app if it does not exist. Default ``True``.
+- ``sandbox_timeout``: Maximum lifetime of a sandbox in seconds. Default ``3600``.
+ Modal's own default is 300, which is below a plausible agent run.
+- ``idle_timeout``: Seconds of inactivity after which Modal reclaims the sandbox,
+ or ``None`` to rely on ``sandbox_timeout`` alone. Default ``None``; see
+ :ref:`Configuring a sandbox <sandbox-configuring>`.
+- ``workdir``: Working directory for commands, created if the image lacks it.
+ Default ``"/workspace"``. It is a starting directory, not a jail: absolute paths
+ and ``..`` are passed through, and commands run as root, so the model can read
+ and write anywhere in the sandbox filesystem. The sandbox boundary is what
+ contains that.
+- ``cpu``, ``memory``, ``gpu``, ``region``, ``cloud``: passed through to Modal.
+ ``None`` lets Modal choose. An unrecognized ``region`` or ``cloud`` fails the
+ task rather than falling back.
+- ``tags``: Extra Modal tags on every sandbox, e.g. ``{"dag_id": "my_dag"}``.
+- ``egress_enforcement``: ``"strict"`` (default) or ``"sni"``. See below.
+
+**Network policy.** ``block_network=True`` maps exactly onto Modal's own
+``block_network``, which drops all outbound traffic including DNS. A spec that
+names ``allow_egress_to`` is **refused by default**, because Modal cannot combine
+an allowlist with ``block_network`` at all, and its hostname allowlist is enforced
+by matching the name in the TLS handshake, which means:
+
+- TLS on port 443 to a listed host connects; any other host is refused at once.
+- **Non-TLS traffic is dropped, not refused.** A plain HTTP connection to a listed
+ host stalls until the client gives up, around two minutes of TCP retries for one
+ address, so with the 60 s default command budget the model reads
+ ``[timed out after 60s]`` and concludes its command was slow, never that the
+ network stopped it.
+- **The destination address is not part of the decision.** A connection opened to
+ an unrelated address while presenting a listed name is routed to the listed host
+ and answered by it: connecting to ``8.8.8.8:443`` with ``pypi.org`` in the
+ handshake returns the certificate and content of ``pypi.org``. The allowlist is a
+ name-routed egress proxy, not a filter on where packets may go.
+- **DNS resolution stays open for every hostname**, listed or not, against
+ authoritative servers outside Modal. A freshly generated label under a domain the
+ operator controls resolves and returns its answer, so this is a two-way channel.
+- A host that shares a **TLS endpoint** with a listed one can be reached by
+ presenting the listed name in the handshake and the other in the request. With
+ ``pypi.org`` as the only allowed host, a TLS session opened to
+ ``files.pythonhosted.org`` while presenting ``pypi.org`` was allowed through and
+ answered. Other tenants of the same CDN returned ``421 Misdirected Request``; it is
+ co-tenancy of the same TLS endpoint that matters, which you cannot check from
+ outside and which can change without notice.
+
+So the allowlist says which name a TLS session may be routed to, and nothing else.
+Pass ``ModalSandboxBackend(egress_enforcement="sni")`` to say you accept that and
+have the allowlist applied:
+
+.. code-block:: python
+
+ SandboxToolset(
+ ModalSandboxBackend(egress_enforcement="sni"),
+ spec=SandboxSpec(block_network=True, allow_egress_to=["pypi.org", "files.pythonhosted.org"]),
+ )
+
+Entries must be bare hostnames or one leading ``*.`` label; a URL, a ``host:port``,
+an address or a single-label name is refused, because Modal applies the list without
+checking it and any of those would silently match nothing.
+
+**What the image needs.** ``write_file`` and ``list_directory`` use Modal's own
+filesystem API, served by a helper Modal injects into the sandbox, so they need
+nothing from the image. ``read_file`` deliberately does not: Modal's read API takes
+no length, so it cannot honor a read budget, and a single call on a file that
+streams without end (``/dev/zero``, a FIFO, a procfs entry) would pull it into
+worker memory unbounded. That tool runs the base class's shell implementation,
+which caps the read inside the guest, and the image needs ``stat``, ``head`` and
+``base64``. Any Debian or Ubuntu based image, including ``python:*-slim``, has
+them.
+
+**Symlinks.** Because ``write_file`` goes through the native API, writing to a path
+that is a symlink replaces the link with a regular file and leaves the original
+target untouched, where a shell redirect would follow the link.
+
+sbx (Docker Sandboxes, local)
+-----------------------------
+
+:class:`~airflow.providers.common.ai.sandbox.SbxSandboxBackend` runs each sandbox
+in a Docker Sandboxes microVM by driving the ``sbx`` CLI. Each sandbox is a real
+microVM with its own kernel.
+
+.. warning::
+
+ **Use this backend for local development, not production.** Docker Sandboxes
+ is built for running coding agents against a checkout on your own machine, and
+ driving it from an Airflow worker is off-label use. A production worker would
+ need the ``sbx`` binary on the host, an authenticated Docker account
+ (``sbx login``), a one-time ``sbx policy init``, and on Linux, KVM or nested
+ virtualization, which a worker in an unprivileged container cannot provide.
+
+ **Orphans are not reclaimed.** There is no server-side lifetime. If the worker is
+ killed outright, the microVM and its workspace directory survive; sandboxes are
+ named ``airflow-sandbox-*`` so an operator can find and remove them.
+
+Installing the CLI is a Deployment Manager prerequisite (``brew install
+docker/tap/sbx`` or ``winget install Docker.sbx``); the backend needs no Python
+dependency. The template image must provide GNU coreutils ``timeout``, ``base64``,
+``stat``, ``head``, ``find``, ``mkdir`` and ``dirname``, which any Debian or Ubuntu
+based image has.
+
+Constructor parameters:
+
+- ``image``: Container image for the sandbox. Default ``"python:3.12-slim"``.
+- ``memory``: Memory limit in binary units. ``sbx`` enforces a 1 GiB minimum.
+ Default ``"2g"``.
+- ``cpus``: CPUs to allocate. ``None`` (default) uses the ``sbx`` default, which
+ is every host CPU.
+- ``sbx_path``: Path to the ``sbx`` binary. Default ``"sbx"``.
+- ``create_timeout``: Seconds allowed for provisioning; a first-run microVM boot
+ plus an image pull can be slow. Default ``600``.
+- ``host_network_policy``: What ``sbx policy`` is set to on this host.
+ ``"unknown"`` (default) makes ``create`` refuse any spec asking for a network
+ guarantee this backend cannot make, and since ``block_network`` defaults to
+ ``True`` that includes a bare ``SandboxSpec()``. Set ``"deny-all"`` after running
+ ``sbx policy init deny-all``, or ``"allow-all"`` to state that egress is open
+ and pass ``SandboxSpec(block_network=False)`` to match.
+
+What differs between the two
+----------------------------
+
+Swapping the backend is one constructor argument, and tool names, spec and prompt
+do not change. Four behaviours do, so read them before assuming the same Dag
+behaves identically in both places:
+
+- **CPU.** ``sbx`` gives a sandbox every host CPU; Modal defaults to a fraction of
+ one, so set ``cpu``.
+- **Egress allowlists.** ``sbx`` enforces ``allow_egress_to`` at the host policy
+ layer; Modal matches TLS handshake names, which is weaker and has to be opted
+ into.
+- **Command timeouts.** A timeout destroys an ``sbx`` sandbox and its files; a
+ Modal sandbox survives with its files intact.
+- **Symlinks.** ``write_file`` through a symlink follows the link on ``sbx`` and
+ replaces it on Modal.
+
+Bringing your own backend
+-------------------------
+
+Any vendor that can create a sandbox, run a command in it and destroy it can plug
+in. Subclass :class:`~airflow.providers.common.ai.sandbox.SandboxBackend` in your
+own package and pass an instance to ``SandboxToolset``.
+
+**Three methods are required**: ``create``, ``run_command`` and ``destroy``. The
+three file operations ship as defaults implemented over ``run_command``, because
+reading, writing and listing a file are all expressible as shell commands.
+Override them only when the vendor has a native file API:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.sandbox import (
+ SandboxBackend,
+ SandboxExecResult,
+ SandboxSpec,
+ )
+
+
+ class AcmeSandboxBackend(SandboxBackend):
+ name = "acme"
+
+ def create(self, *, spec: SandboxSpec | None = None) -> str:
+ return acme_sdk.create_sandbox().id
+
+ def run_command(self, sandbox, command, *, timeout, max_output_bytes):
+ r = acme_sdk.exec(sandbox, command, timeout=timeout)
+ return SandboxExecResult(exit_code=r.exit_code, stdout=r.stdout, stderr=r.stderr)
+
+ def destroy(self, sandbox) -> None:
+ acme_sdk.delete_sandbox(sandbox)
+
+ # Optional: inherited from SandboxBackend unless the vendor has
+ # something better than shelling out.
+ def read_file(self, sandbox, path, *, max_bytes) -> bytes:
+ return acme_sdk.download(sandbox, path, limit=max_bytes)
+
+Four rules for an implementation:
+
+- Constructors run at Dag-parse time, so resolve credentials lazily, on first use.
+- ``destroy`` must be idempotent; destroying an already-gone sandbox is not an error.
+- Raise ``SandboxTerminalError`` when retrying cannot help and ``SandboxError``
+ when it might. The first fails the task for Airflow to retry; the second
+ becomes a bounded prompt back to the model.
+- If you cannot enforce something the ``SandboxSpec`` asks for, **raise**. Never
+ provision a weaker sandbox than the Dag author asked for.
diff --git a/providers/common/ai/docs/sandbox/configuration.rst b/providers/common/ai/docs/sandbox/configuration.rst
new file mode 100644
index 00000000000..379fa064f45
--- /dev/null
+++ b/providers/common/ai/docs/sandbox/configuration.rst
@@ -0,0 +1,305 @@
+ .. 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.
+
+Sandbox configuration and lifecycle
+===================================
+
+.. _sandbox-configuring:
+
+Configuring a sandbox
+---------------------
+
+:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` says what a sandbox is
+provisioned with. The default denies outbound network access and injects no
+environment, and a backend that cannot enforce a field it was given **raises
+instead of ignoring it**, so a spec never gives you a false sense of a restriction
+being in force.
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.sandbox import SandboxSpec, SbxSandboxBackend
+ from airflow.providers.common.ai.toolsets import SandboxToolset
+
+ SandboxToolset(
+ SbxSandboxBackend(host_network_policy="allow-all"),
+ spec=SandboxSpec(
+ env={"HF_TOKEN": "..."}, # only what the generated code legitimately needs
+ block_network=False, # this sandbox may reach the internet
+ ),
+ )
+
+**Network.** Three modes exist. ``block_network=True``, the default, drops all
+outbound traffic including name resolution. On Modal it maps onto the sandbox's
+own ``block_network`` flag. ``sbx`` has no per-sandbox enforcement of it at all:
+egress there is a host-level ``sbx policy``, so the backend honors the default by
+refusing to provision unless the Deployment Manager has declared
+``host_network_policy="deny-all"``, which means a bare ``SandboxSpec()`` is refused
+under the default ``host_network_policy="unknown"`` rather than silently getting
+an open sandbox. ``allow_egress_to`` names hosts the sandbox may reach, and both
+backends refuse it until you have said something that makes it enforceable:
+``sbx`` applies it as a per-sandbox rule on top of that ``deny-all`` host policy,
+since a local rule can narrow egress and never widen it, and Modal matches
+hostnames in the TLS handshake, which is weaker than it sounds and has to be
+opted into with ``egress_enforcement="sni"`` (see
+:ref:`the Modal backend <sandbox-backend-modal>` for exactly what it does and does
+not stop). ``block_network=False`` opens outbound access; on Modal the cloud
+metadata endpoint and private address ranges stay unreachable even then.
+
+**Packages.** A default sandbox has the Python standard library and no network,
+so an agent that reaches for ``pip install`` gets a DNS failure in a few seconds.
+Bake what it needs into the image, which needs no egress at all:
+
+.. code-block:: python
+
+ import modal
+
+ ModalSandboxBackend(
+ image=modal.Image.from_registry("python:3.12-slim").pip_install("pandas", "pyarrow"),
+ )
+
+Or, on Modal, allow exactly the two hosts a Python install needs and accept the
+handshake-name caveats: with ``pypi.org`` and ``files.pythonhosted.org`` allowed,
+installing pandas took about ten seconds when measured and ``github.com`` stayed
+blocked. Tell the model in its instructions what the image already has; the tool
+description says whether the sandbox has network access and where its working
+directory is, but not which packages are installed.
+
+**Resources.** A default Modal sandbox is a fraction of a core (``nproc`` reports
+1) with the standard library and no ``curl``, ``git`` or ``wget``. Set ``cpu`` for
+anything heavier. ``memory`` on Modal is a scheduling request, not a ceiling: a
+sandbox created with ``memory=512`` allocated 1.5 GiB without complaint when
+measured, so do not reach for it to bound what model-written code can consume.
+``sbx`` gives a sandbox every host CPU by default and enforces its ``memory`` as a
+limit.
+
+**Timeouts.** Three clocks apply. A command has a budget: ``default_command_timeout``
+(60 s) when the model does not ask, up to ``max_command_timeout`` (300 s) when it
+does. The sandbox has a lifetime: ``sandbox_timeout`` on Modal, 3600 s by default,
+and the clock runs through the model's thinking time between tool calls, so size it
+against the whole agent run. And Modal can reclaim an idle sandbox after
+``idle_timeout``, which defaults to ``None`` because a gap for model generation
+looks idle and would take the agent's files with it.
+
+**Two sandboxes on one agent** need ``tool_prefix``, since tool names must be
+unique within an agent:
+
+.. code-block:: python
+
+ toolsets = [
+ SandboxToolset(
+ SbxSandboxBackend(image="python:3.12-slim", host_network_policy="deny-all"),
+ tool_prefix="py",
+ ),
+ SandboxToolset(
+ SbxSandboxBackend(image="node:22-slim", host_network_policy="deny-all"),
+ tool_prefix="node",
+ ),
+ ]
+
+That yields ``py_run_command``, ``node_run_command`` and so on. Without a prefix
+on at least one of them the run fails at startup. Give the model instructions on
+which one to use for what, or it will guess.
+
+.. _sandbox-credentials:
+
+Credentials
+-----------
+
+``SandboxSpec.env`` is the only way in. Airflow never populates it: no connection,
+variable or worker environment variable reaches a sandbox unless you name it
+there, and the credential that *provisions* the sandbox never enters it either.
+Modal's token stays on the worker and is used by the client, so code running
+inside cannot call Modal as you or create further sandboxes.
+
+Before you put a real secret in ``env``, four things are true of it.
+
+**The model can read it.** Model-generated code can print the environment, and a
+tool result travels into the model's context, the task log and whatever the agent
+returns to XCom. Treat a credential given to a sandbox as disclosed to the model
+and to everything that records the run, and scope it to that one job. If the model
+only needs to *use* a system rather than hold its credential, a
+:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` or
+:class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` runs the hook in
+the task and returns only the result, which is the better shape whenever the
+operations can be named.
+
+**The value is fixed when the Dag file is parsed.** ``toolsets`` is not a
+templated field, so a spec cannot carry a Jinja expression, and a connection
+lookup written beside it would run in the Dag processor on every parse. An agent's
+sandbox therefore cannot take a credential from an Airflow connection or a secrets
+backend today. If it has to, the job cannot be an agent's sandbox: drive a backend
+from a ``@task``, where you are in ordinary Python at run time and the lookup runs
+where it should.
+
+.. code-block:: python
+
+ @task
+ def run_with_credential():
+ from airflow.providers.common.ai.sandbox import ModalSandboxBackend, SandboxSpec
+ from airflow.providers.common.compat.sdk import BaseHook
+
+ token = BaseHook.get_connection("my_api").password # resolved in the task, not at parse
+ backend = ModalSandboxBackend()
+ sandbox = backend.create(spec=SandboxSpec(env={"API_TOKEN": token}, block_network=False))
+ try:
+ ... # write_file, run_command, read_file
+ finally:
+ backend.destroy(sandbox)
+
+**On** ``sbx`` **the value is written to the guest filesystem.** The CLI has no
+create-time environment flag, so the backend appends ``export`` lines to
+``/etc/profile`` inside the microVM, readable by anything in that sandbox for its
+whole life. Modal passes the environment at creation and writes no file. Prefer the
+hosted backend when the sandbox needs a secret at all.
+
+**Non-string values are refused** at provisioning, with the offending keys named,
+rather than failing somewhere less obvious.
+
+.. _sandbox-results:
+
+Getting a result out
+--------------------
+
+An agent's result is whatever the model returns: findings, a mapping, a table
+small enough to read. That reaches XCom through ``output_type`` like any other
+agent output, and the examples on :doc:`index` end that way.
+
+What does not come out is a *file* the agent built. The sandbox is destroyed when
+the run ends; ``read_file`` is text-only, replacing any byte it cannot decode, so a
+parquet or an image cannot survive the round trip at any size; and everything the
+model reads is capped, 50 KiB per stream for ``run_command`` and ``max_read_bytes``
+(5 MiB) per ``read_file``. Those caps are a budget rather than a transport limit:
+a 200 MB file reads out of a live sandbox in under ten seconds on the same path,
+but raising the cap costs roughly three times the file size in worker memory to
+show the model 50 KiB of it. The toolset does not yet have a seam for an author to
+collect what the agent produced; that is a known gap rather than the intended
+design.
+
+When the deliverable is a file and the Dag already knows the job, do not use an
+agent for it. Drive a backend from a ``@task``: the input goes in through
+``write_file``, the output comes back through ``read_file`` with a budget the
+worker can afford rather than a model's, and it lands in object storage with only
+the location passed downstream. The sandbox needs no network and no credentials
+for this, and the task owns its teardown:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
+ :language: python
+ :start-after: [START howto_sandbox_task_artifact]
+ :end-before: [END howto_sandbox_task_artifact]
+
+If the agent genuinely has to be the one producing the file, the shape that works
+is to give the sandboxed code egress to your bucket and a scoped credential in the
+spec, and have it write the object itself, passing the key back through its
+output. That widens the default deliberately and is worth a sentence in the Dag
+saying why.
+
+.. _sandbox-lifecycle:
+
+Lifecycle, failures and retries
+-------------------------------
+
+The sandbox is created lazily on the first tool call, shared by every call in that
+agent run, and destroyed when the run ends. A run that never calls a tool never
+provisions one, and concurrent runs never share a sandbox. Files written by one
+call are visible to later calls in the same run; each ``run_command`` is a fresh
+shell, so shell variables and background jobs do not survive between calls.
+
+**A failed command is output, not a failure.** A non-zero exit comes back to the
+model with both streams and ``[exit code: N]`` appended, and the model corrects
+itself. A recoverable sandbox error, such as a missing path or a file over the
+read budget, becomes a bounded retry prompt. Only a terminal failure, such as
+rejected credentials or a sandbox that has gone, fails the task, so Airflow's own
+retry handles it rather than the model burning its retry budget.
+
+**A command timeout** reports ``[timed out after Ns]`` with the budget that was
+actually applied. On Modal the sandbox survives and files written earlier are
+still there for the model to inspect. On ``sbx`` the backend has to destroy the
+sandbox to be certain the command stopped, and the result then says
+``[sandbox was replaced; files from earlier calls are gone]`` so the model knows
+to start over. Any other failure that loses the sandbox is reported the same way.
+
+**A run that outlives its sandbox fails the task.** When ``sandbox_timeout`` or
+``idle_timeout`` passes mid-run, the next tool call reaches a sandbox that is gone,
+which is terminal. This is the normal end of any run longer than its sandbox's
+lifetime. Human output review is not part of the run: it starts after the agent
+has finished and the sandbox has been destroyed, so a review pause costs no
+sandbox time and keeps no files.
+
+**Nothing survives the run.** A task retry starts from an empty sandbox, and so
+does every other attempt. Two operator features assume otherwise and must not be
+combined with a sandbox today, because neither is rejected:
+
+- ``durable=True`` caches each tool result and replays it on a retry without
+ calling the backend, so a replayed ``write_file`` reports success while no
+ sandbox exists, and the first call that misses the cache runs against a fresh
+ empty one. The model is handed a filesystem that does not match what it was just
+ told, and nothing raises.
+- ``enable_hitl_review=True`` regenerates after reviewer feedback by starting a
+ second agent run, and the first run's sandbox was destroyed when that run ended.
+ The regenerated agent gets an empty sandbox while its own history describes
+ files it wrote earlier.
+
+Cost and operations
+-------------------
+
+**One sandbox spans a whole agent run**, so you pay for the model's thinking time
+between tool calls, not only for the seconds your commands run. A default Modal
+sandbox is around two cents an hour, so this rarely matters; a data-sized one
+(``cpu=8, memory=32768``) is around two dollars an hour, and Modal bills the larger
+of what you requested and what you used, which is the second reason ``memory`` is
+worth setting deliberately. With ``idle_timeout`` unset, ``sandbox_timeout`` is
+the worst-case bill for one run. Fan-out is not the thing to worry about: 25
+sandboxes created concurrently from one process all came up in under three
+seconds.
+
+**Teardown runs inside the task.** A graceful end and an ordinary exception both
+destroy the sandbox. A worker killed outright, an OOM or a lost node never run that
+code, and what happens next is the backend's business alone: Modal reclaims the
+sandbox at ``sandbox_timeout`` whatever became of the worker; ``sbx`` has no
+server-side lifetime, so the microVM and its workspace directory survive with no
+owner. A teardown that fails is logged with the sandbox identity and does not fail
+the task.
+
+**Finding what a Dag left behind.** Modal sandboxes are named ``airflow-sandbox-*``
+and carry an ``airflow_sandbox`` tag plus whatever ``tags`` you pass, and
+``modal.Sandbox.list(app_id=..., tags={"dag_id": "my_dag"})`` returns exactly those
+still running. Tags are fixed when the backend is constructed, so every sandbox in
+a mapped task carries the same ones. When checking whether anything is actually
+still running, trust ``Sandbox.list``, which returns only live sandboxes, or
+``poll()`` on one you hold: ``None`` means running, and any integer means gone,
+including ``137``, which is a stopped sandbox rather than a sick one. Modal's
+dashboard lists stopped sandboxes alongside live ones, and the ``Tasks`` count in
+``modal app list`` lags by up to about a minute. ``sbx`` sandboxes are named the
+same way for an operator sweep.
+
+``SandboxToolset`` parameters
+-----------------------------
+
+- ``backend``: The backend that provisions and drives the sandbox.
+- ``spec``: What to provision it with. Defaults to no environment and no egress.
+- ``default_command_timeout``: Seconds for a ``run_command`` the model did not
+ put a timeout on. Default ``60``.
+- ``max_command_timeout``: Hard ceiling for any single command, including a
+ model-supplied ``timeout_seconds``. Default ``300``.
+- ``max_output_lines`` / ``max_output_bytes``: Caps per output stream and per
+ file read, whichever is hit first. Defaults ``2000`` and 50 KiB. Command
+ output keeps the **tail**, where errors and the exit status live; file reads
+ keep the head and report a continuation offset.
+- ``max_read_bytes``: Largest file ``read_file`` will transfer. Default 5 MiB;
+ larger files are refused with a hint to slice them in the shell.
+- ``tool_prefix``: Prefix for the four tool names. Needed when one agent has
+ more than one ``SandboxToolset``.
diff --git a/providers/common/ai/docs/sandbox/index.rst b/providers/common/ai/docs/sandbox/index.rst
new file mode 100644
index 00000000000..32908cfc61e
--- /dev/null
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -0,0 +1,425 @@
+ .. 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.
+
+.. _howto/sandbox:
+
+Sandboxed execution for agents
+==============================
+
+An agent that is asked to do open-ended work writes code, and then something has
+to run that code. By default that something is the Airflow worker: a skill
+script, a hand-written tool that shells out, or generated glue in
+:ref:`code mode <code-mode>` all execute on the host that also holds your
+connections, your filesystem and your network position, and the code they run
+was generated a moment ago from inputs you do not control.
+
+:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` gives the
+model a disposable workspace for that code instead. It exposes four tools:
+
+.. list-table::
+ :widths: 25 75
+ :header-rows: 1
+
+ * - Tool
+ - What it does
+ * - ``run_command``
+ - Runs a shell command. Pipes, redirection, ``&&`` and globs work. A
+ non-zero exit is reported as output, not raised, so the model reads
+ ``stderr`` and corrects itself.
+ * - ``read_file``
+ - Reads a text file, head-first, and reports the next ``offset`` so the
+ model can page through a long file.
+ * - ``write_file``
+ - Writes text to a file, creating parent directories.
+ * - ``list_directory``
+ - Lists a directory. Directories are shown with a trailing ``/``.
+
+The sandbox is provisioned by a
+:class:`~airflow.providers.common.ai.sandbox.SandboxBackend` on the model's first
+tool call and destroyed when the agent run ends. Two backends ship: a hosted one on
+`Modal <https://modal.com/docs/guide/sandbox>`__ for production and Kubernetes,
+and a local microVM one on `Docker Sandboxes <https://docs.docker.com/ai/sandboxes/>`__
+for development. The four tool names and shapes match pydantic-ai's own sandbox
+capabilities, so a model that has seen one already knows this one.
+
+**Adding this toolset gives the agent shell and file operations in a separate
+workspace. Every other tool keeps its existing permissions and runs where it ran
+before.** Whether that is a new capability depends on what the agent already had:
+an agent with a skill script or a shell tool already ran generated code on the
+worker, and this is where that code should run instead. An agent that only had
+named database operations is being granted a general shell for the first time, in
+a workspace with less authority than the worker, and you are choosing to grant it.
+
+Pages in this section
+---------------------
+
+.. toctree::
+ :maxdepth: 1
+
+ Configuration and lifecycle <configuration>
+ Backends <backends>
+
+.. _sandbox-quick-start:
+
+Quick start
+-----------
+
+An agent investigating a revenue anomaly. The warehouse is reached through a
+:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`, so the credential
+stays in the task and the model only ever sees rows. The arithmetic happens in a
+hosted sandbox with pandas baked into the image, so the sandbox needs no network,
+and a bad pivot or a runaway loop costs a small sandbox rather than a worker slot:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
+ :language: python
+ :start-after: [START howto_sandbox_agent_investigation]
+ :end-before: [END howto_sandbox_agent_investigation]
+
+The same toolset on the local ``sbx`` backend, for developing on a laptop. Here
+the agent maps a vendor file whose columns drift onto a staging schema: it writes a
+script, runs it, reads the traceback and fixes it, which is the loop nobody can
+write down in advance. Swapping ``SbxSandboxBackend`` for ``ModalSandboxBackend``
+is the only change between this Dag and production:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
+ :language: python
+ :start-after: [START howto_sandbox_agent_local]
+ :end-before: [END howto_sandbox_agent_local]
+
+Install the Modal extra and authenticate as you would for Modal's own CLI. On a
+worker, set ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` in the environment
+instead; nothing is read until the first sandbox is created, so a Dag file that
+constructs the backend parses without credentials present:
+
+.. code-block:: bash
+
+ pip install 'apache-airflow-providers-common-ai[modal]'
+ modal token new # writes ~/.modal.toml
+
+What it is for, in practice
+---------------------------
+
+The rule is: a model inside a running task writes code you cannot name in advance,
+and that code needs a real shell, a real interpreter or packages. These are the
+jobs that rule describes, each paired with the version where a sandbox is the wrong
+call.
+
+**A vendor drops a CSV nobody has a schema for.** Three date formats, an encoding
+error at row 41,000, pandas needed. Every step depends on the previous step's
+output, and a bad read of the whole file into memory should exhaust a small
+sandbox, not a worker slot shared with other tasks. If the schema is stable and
+the mapping is known, this is a plain task with no agent.
+
+**Yesterday's revenue moved and nobody knows why.** The agent queries through a
+SQL toolset with a table allowlist, writes the rows into the sandbox, and pivots
+with pandas rather than doing arithmetic in tokens. The warehouse credential never
+enters the sandbox. If it is always the same five queries, that is
+:class:`~airflow.providers.common.ai.operators.llm.LLMOperator` over a fixed SQL
+task; if the rows number in the millions, compute in the warehouse rather than
+carrying them through the model's context.
+
+**Migrate forty dbt models between warehouse dialects.** The model files go into
+the workspace, the image carries ``sqlglot`` and ``sqlfluff``, and the loop of
+convert, lint, read errors, fix runs for as long as it takes, fully offline. The
+rewritten files are the deliverable, and getting forty files out is the gap
+described under :ref:`Getting a result out <sandbox-results>`.
+
+**A failing task's traceback, handed to an agent to propose a fix.** The agent
+reproduces the crash against a sample, tries a fix, reruns, and posts a diff for a
+person to review. That code ran with the worker's credentials the first time; the
+reproduction should not. A written diagnosis with no execution is
+``LLMOperator`` reading the log.
+
+**Not a sandbox: flag 200 customers with an open ticket** using two existing
+tools. That is :ref:`code mode <code-mode>`: a short loop over registered tools,
+no packages, no shell, sub-millisecond start.
+
+**Not a sandbox: run a vendor's PII-scanner image over an S3 prefix.** No model
+writes anything. That is ``KubernetesPodOperator``, a fixed payload with templated
+arguments in its own pod.
+
+**Not only a sandbox: one team's agent tasks must run off the shared workers.**
+That is a where-does-the-task-run requirement, so ``KubernetesExecutor`` with a
+pod override, and the sandbox toolset inside the pod if those agents also write
+code. The pod protects the cluster from the task; the sandbox protects the task
+from its model.
+
+.. _sandbox-boundaries:
+
+Choosing the boundary
+---------------------
+
+"Sandboxed" means different things at different layers, and picking the wrong
+layer is the most common way to end up with less protection than you think. Four
+boundaries exist, from smallest to largest:
+
+.. list-table::
+ :widths: 22 33 45
+ :header-rows: 1
+
+ * - Boundary
+ - What moves inside it
+ - What that protects you from
+ * - **A tool call**
+ - What ``run_command`` and the file tools do. **This toolset.**
+ - Model-written code damaging the worker host, reading its files, or
+ reaching the network from it.
+ * - **The glue between tools**
+ - Generated orchestration code, via :ref:`code mode <code-mode>` and the
+ Monty interpreter.
+ - Generated code touching anything other than the tools you registered.
+ Monty runs it in a worker subprocess on the worker host, a language-level
+ sandbox rather than an OS one, and every tool it calls executes in the
+ task process with the task's authority.
+ * - **The agent process**
+ - The whole agent loop, its LLM credentials and its message history.
+ - The agent's own credentials leaking, and any *other* toolset on the same
+ agent. Not available today.
+ * - **The whole task**
+ - The complete Airflow task, supervisor included, as
+ ``KubernetesExecutor`` does.
+ - Everything outside the task. Not the task from its own code: the
+ agent and what it runs are both inside this boundary.
+
+``SandboxToolset`` is the first row. It relocates exactly its four tools and
+nothing else. Boundary size is not a security level on its own: the image, the
+credentials you inject, the network policy and the resource limits decide the
+actual isolation. Choose the smallest boundary that fits, then configure it.
+
+.. list-table::
+ :widths: 48 52
+ :header-rows: 1
+
+ * - The agent needs to
+ - Reach for
+ * - Call operations you can name in advance
+ - :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` or
+ :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset`. Narrow,
+ bounded, and the credential stays in the worker rather than reaching the
+ model. If you can name the operations, name them.
+ * - Chain those tools with glue logic
+ - :ref:`code mode <code-mode>`. Monty confines generated code to the tools
+ you registered and starts in well under a millisecond, against roughly a
+ second to provision a hosted sandbox. Reach past it for a sandbox when the
+ model needs something Monty does not have: a shell, ``pip install``, a
+ compiled library.
+ * - Write and run open-ended code: reshape data with no known schema, install
+ a package, or fix its own failing script by reading the traceback
+ - ``SandboxToolset``
+ * - Produce a large artifact for a downstream task
+ - Today, a ``@task`` driving a backend directly, not the agent. That is a
+ missing seam in the toolset rather than a recommendation; see
+ :ref:`Getting a result out <sandbox-results>`.
+ * - A whole task's worth of untrusted work isolated, with no agent involved
+ - ``KubernetesPodOperator``
+ * - Airflow's own credentials kept away from the agent
+ - Not this. Scope the connections the task can see, and give the agent only
+ the toolsets it needs. See :ref:`sandbox-security`.
+
+**Code mode and the sandbox compose.** With both enabled, the three file tools
+fold into ``run_code`` as callables the generated code can loop over, and
+``run_command`` stays a tool the model calls directly, so a shell command is
+written as a shell command rather than quoted inside generated Python. Both paths
+reach the same sandbox: there is one per agent run whichever way a call arrives.
+
+Why not ``KubernetesPodOperator`` or the ``KubernetesExecutor``?
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Because they answer a different question. Both isolate a task from the rest of
+your platform. Neither puts a boundary between an agent and the code that agent
+writes.
+
+Run an ``AgentOperator`` under the ``KubernetesExecutor`` and the whole task,
+agent loop included, gets its own pod. Model-written code still executes inside
+that pod, with its service account, its mounted secrets, its connections and its
+network position. The pod protects the cluster from the task; it does not protect
+the task from what the model decided to run, because from the pod's point of view
+that code is the task.
+
+``KubernetesPodOperator`` does contain the work, and its image, command, arguments
+and environment are all templated, so it can launch a different payload per run,
+and that payload can be a complete agent. The difference is what each one is. The
+pod is the environment that hosts the workload, and everything inside it runs with
+the pod's authority. ``SandboxToolset`` is a separate workspace that an
+already-running workload reaches into repeatedly, with less authority than the
+workload itself. Expressing the workspace as a pod means one pod per turn, with
+scheduling latency and a lost filesystem between each, or letting the model run
+whatever it wants inside the pod that also holds your credentials.
+
+.. list-table::
+ :widths: 30 35 35
+ :header-rows: 1
+
+ * -
+ - ``KubernetesExecutor`` / ``KubernetesPodOperator``
+ - ``SandboxToolset``
+ * - Protects
+ - The cluster, from the task
+ - The task and its credentials, from the model's output
+ * - Boundary sits
+ - Around the whole task, agent included
+ - Between the agent and the code it writes
+ * - What runs inside is decided
+ - Before the process starts
+ - By the model, mid-run
+ * - Credentials inside
+ - The task's own: service account, secrets, connections
+ - Only what ``SandboxSpec.env`` names, which is nothing by default
+
+Running an agent in a pod *and* giving it a sandbox is the setup a Kubernetes
+deployment usually wants: the pod bounds the task, the sandbox bounds the model.
+The ``sbx`` backend cannot be the second half of that, because it needs KVM on the
+worker host and an unprivileged pod cannot provide it; that is what the hosted
+backend is for.
+
+When to choose it
+-----------------
+
+**Choose it when** the work is running code the model wrote, rather than calling
+a tool you picked in advance — exploratory analysis, installing a package for one
+task, a script the model writes, runs, and fixes from its own traceback. Every
+other toolset route answers "call this thing"; this one answers "here is
+somewhere to work". This page has the worked scenarios above and
+the limitations below to read before designing a Dag around it.
+
+Before reaching for it, check whether the actual need is narrower than that:
+``code_mode=True`` is a flag on ``AgentOperator``, not a toolset
+— it changes how the model invokes the tools it already has, letting it write
+code to call several of them instead of emitting one call per step. It does not
+give the agent somewhere to run arbitrary code of its own, and it avoids the
+``sbx`` backend's production-readiness, network-isolation, and reclamation
+caveats on :doc:`backends` and needs no backend at all — but not the reachability one: the generated code runs in
+Monty's deny-by-default sandbox, but the tools it calls still run in the
+worker, so a credential-bearing toolset on the same agent stays within reach
+whether or not code mode is on. See :ref:`code-mode` and
+:ref:`sandbox-boundaries`.
+
+**What it cannot do**
+
+- Only one of its two backends runs on Kubernetes. ``SbxSandboxBackend`` drives
+ Docker Sandboxes on the worker host, and its own documentation says to use it
+ for local development: it wants the ``sbx`` binary on the host, an
+ authenticated Docker account, a one-time ``sbx policy init``, and on Linux KVM
+ or nested virtualization — which an unprivileged container cannot provide.
+ Production and Kubernetes use
+ :class:`~airflow.providers.common.ai.sandbox.modal.ModalSandboxBackend`, a
+ hosted backend behind the ``modal`` extra that installs nothing on the worker
+ and reclaims a sandbox at its own lifetime if the worker dies. Both implement
+ :class:`~airflow.providers.common.ai.sandbox.SandboxBackend`, and a third
+ vendor can too.
+- It does not contain the agent. Only what these tools do runs in the sandbox;
+ the agent loop, the model calls, and every other toolset on the same agent stay
+ in the worker with the worker's credentials. It contains model-written code, so
+ pairing it with a credential-bearing toolset on the same agent puts the
+ credential back within reach. :ref:`sandbox-boundaries` sets this out in full.
+- The ``sbx`` backend cannot enforce network isolation on its own. It applies
+ ``allow_egress_to`` as a per-sandbox policy rule, but only on top of a
+ ``deny-all`` host policy, since a local rule can narrow egress and never
+ widen it; ``block_network`` has no per-sandbox enforcement at all. Ask for
+ either against a host policy that is not already ``deny-all`` and the
+ backend raises rather than silently leaving the sandbox less restricted
+ than the spec asked for. ``block_network`` defaults to ``True``, so a bare
+ ``SandboxSpec()`` with no arguments already asks for it and is refused
+ under the default ``host_network_policy="unknown"``. On Modal the same default
+ maps onto the sandbox's own ``block_network`` and is enforced exactly; a
+ hostname allowlist there is matched on the TLS handshake name and has to be
+ opted into, for the reasons set out on :doc:`backends`.
+- Reclamation depends on the backend. A failed teardown is logged as a warning
+ rather than raised, deliberately, so that a teardown blip cannot fail a
+ finished run. On ``sbx`` nothing else picks up the slack: there is no
+ server-side TTL, so a worker killed outright leaves the microVM and its
+ workspace directory behind, named ``airflow-sandbox-*`` so an operator can find
+ and remove them. On Modal the sandbox ends at its own ``sandbox_timeout``
+ whatever became of the worker.
+- Nothing survives the run, and a file the agent built can leave only through
+ the model's context, which is text-only and capped. Producing an artifact for a
+ downstream task is a ``@task`` driving a backend directly today, not the agent;
+ :doc:`configuration` has the example.
+
+**A real example.** ``example_sandbox_toolset.py`` in this provider's example
+Dags has an agent investigating a revenue anomaly on the Modal backend beside a
+``SQLToolset``, the same agent shape on ``sbx`` for a laptop, and a ``@task``
+producing a file through a sandbox; all three are on this page and :doc:`configuration`. The two
+system tests, ``example_sandbox_toolset_sbx.py`` and
+``example_sandbox_toolset_modal.py``, run against a real backend and are
+reachable from the System Tests entry in the sidebar.
+
+**Credentials and where it runs.** This route does not end at an
+Airflow connection. Airflow puts none of its context, connections, variables or
+worker environment into the sandbox; only what you pass through
+:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` goes in, and the
+credential that provisions the sandbox never enters it. Authorization to the
+backend sits outside Airflow — ``sbx login`` and ``sbx policy init`` on the
+machine for ``sbx``, or ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` on the
+worker for Modal. Work runs in a per-run microVM on the worker host with
+``sbx``, or off the worker entirely in Modal's infrastructure. Its tool calls
+act as barriers, as they do for the other
+routes that build their own tools; see :ref:`toolset-call-barriers`.
+
+.. _sandbox-limitations:
+
+Limitations
+-----------
+
+These apply to every backend. Each is explained in the section it belongs to; this
+is the list to read before designing a Dag around an agent with a sandbox.
+
+- **Nothing survives the run**, including across task retries.
+ :ref:`Lifecycle <sandbox-lifecycle>`.
+- **A file the agent built cannot leave** except through the model's context,
+ which is text-only and capped. :ref:`Getting a result out <sandbox-results>`.
+- **A credential cannot come from a connection or a secrets backend**; the spec is
+ fixed at parse time, and anything injected is readable by the model.
+ :ref:`Credentials <sandbox-credentials>`.
+- **Do not combine with** ``durable=True`` **or** ``enable_hitl_review=True``.
+ Neither is rejected today. :ref:`Lifecycle <sandbox-lifecycle>`.
+- **A run that outlives** ``sandbox_timeout`` **fails the task.**
+ :ref:`Lifecycle <sandbox-lifecycle>`.
+- **The hostname allowlist is a weak control** and refused unless opted into.
+ :ref:`Modal <sandbox-backend-modal>`.
+- **Commands run as root and** ``workdir`` **is not a jail.**
+ :ref:`Modal <sandbox-backend-modal>`.
+- **It does not contain the agent.** :ref:`sandbox-security`.
+
+.. _sandbox-security:
+
+Security boundary
+-----------------
+
+**The sandbox does not contain the agent.** The agent loop, the model calls and
+every other toolset on the same agent still run in the worker process with the
+worker's credentials. An agent that has ``SandboxToolset`` *and* a toolset that
+can reach connections has a contained code tool sitting beside a credential path
+that is not contained: a prompt injection arriving in a queried row can steer the
+model into calling the credentialed tool, and that call executes from the worker
+with the connection's full grants. Moving the sandbox to a hosted backend changes
+none of this.
+
+The controls for that risk are different controls. Give the agent only the
+toolsets it needs. Point every connection at a least-privilege role. Prefer a
+capability toolset, which lets the model *use* a connection without ever holding
+it, over a credential in ``SandboxSpec.env``, which the model can read. And treat
+a table allowlist as a guardrail that contains intent rather than a boundary that
+contains access.
+
+What the sandbox does contain, verified against a live worker environment carrying
+sentinel credentials: nothing Airflow-shaped reaches it. No connection, variable,
+Fernet key or cloud credential appeared in the guest's environment, in any process
+environment or on its filesystem, and the provisioning token is not usable from
+inside even when the vendor's API is on the egress allowlist. Sandboxes cannot reach each
+other, so two sandboxes on one agent are isolated from one another as well as from
+the worker.
diff --git a/providers/common/ai/docs/structured_output.rst b/providers/common/ai/docs/structured_output.rst
new file mode 100644
index 00000000000..2fbcfdc6e9b
--- /dev/null
+++ b/providers/common/ai/docs/structured_output.rst
@@ -0,0 +1,68 @@
+ .. 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.
+
+.. _structured-output:
+
+Structured output and XCom
+==========================
+
+Set ``output_type`` to a Pydantic ``BaseModel`` subclass. The LLM is instructed
+to return structured data, and the model instance is pushed to XCom unchanged
+so downstream tasks can type-hint the class directly
+(``def downstream(result: MyModel)``) and use attribute access (``result.field``).
+
+The declared ``output_type`` (and any ``BaseModel`` reachable from
+``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization by
+the worker when it loads the Dag, before any task runs -- so no edit to
+``[core] allowed_deserialization_classes`` is needed. The Pydantic class must be
+defined at **module scope** and bound to an attribute matching its ``__name__``;
+classes nested inside a function or ``@dag``-decorated body, parameterized
+generics, and dynamically-built classes whose ``__name__`` does not match the
+attribute they are bound to cannot be re-imported, so they are skipped with a
+warning at worker startup and the value fails to deserialize at the consumer.
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+ :language: python
+ :start-after: [START howto_operator_llm_structured_output_class]
+ :end-before: [END howto_operator_llm_structured_output_class]
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+ :language: python
+ :start-after: [START howto_operator_llm_structured]
+ :end-before: [END howto_operator_llm_structured]
+
+Registration covers downstream tasks in the **same Dag**: every worker walks the
+loaded Dag's tasks at startup and registers each declared class, so it also works
+for mapped producers (``.expand(...)``) and for workers that load Dags from a
+cache that bypasses operator construction.
+
+The Airflow UI's XCom viewer renders Pydantic instances via the
+``stringify`` path, which produces a representation like
+``my_module.MyModel@version=1(field=value,...)`` without consulting the
+allow-list. It is not pretty (no field-by-field rendering today), but the value
+shows up; no configuration is required.
+
+The remaining gap is **cross-Dag** ``xcom_pull`` -- a task in a different Dag
+that pulls this XCom only parses its own Dag file, not the producer's, so the
+class is not auto-registered. Add the class qualified name to
+``[core] allowed_deserialization_classes`` (or a glob that matches it) to make
+that pattern work.
+
+If a downstream consumer needs the dict shape (e.g. forwarding to an external
+system that expects JSON-style payloads), pass ``serialize_output=True`` and the
+operator calls ``model_dump()`` before pushing to XCom. The pre-PR behavior is
+available on demand without giving up the typed default.
diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst
deleted file mode 100644
index 63cbfa01399..00000000000
--- a/providers/common/ai/docs/toolsets.rst
+++ /dev/null
@@ -1,1145 +0,0 @@
- .. 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.
-
-.. _howto/toolsets:
-
-Toolsets — Airflow Hooks as AI Agent Tools
-==========================================
-
-Airflow's 350+ provider hooks already have typed methods, rich docstrings,
-and managed credentials. Toolsets expose them as pydantic-ai tools so that
-LLM agents can call them during multi-turn reasoning.
-
-If you have not settled on a toolset yet, start with
-:doc:`choosing_a_toolset`, which compares the routes and what each one gives up.
-
-Six toolsets are exported directly from the ``airflow.providers.common.ai.toolsets``
-package root:
-
-- :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic
- adapter for any Airflow Hook.
-- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated
- 4-tool database toolset.
-- :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` — connect to
- `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow
- connections.
-- :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` — give
- the agent a shell and a filesystem inside an isolated sandbox, off the
- Airflow worker. See :doc:`sandbox`.
-- :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
- — base class that provider packages subclass to expose a **vendor-managed
- agent**, one whose reasoning loop runs on a cloud provider's infrastructure.
-- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
- — composes several interchangeable managed agents behind a single tool.
- See :ref:`managed-agent-toolsets` below.
-
-Three more toolsets are documented later on this page. They are not re-exported from the
-package root, so import each of them from its own submodule::
-
- from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
- from airflow.providers.common.ai.toolsets.logging import LoggingToolset
- from airflow.providers.common.ai.toolsets.skills import AgentSkillsToolset
-
-All of the toolsets on this page implement pydantic-ai's
-`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be
-passed to any pydantic-ai ``Agent``, including via
-:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
-
-.. note::
-
- ``AgentOperator`` accepts **any** ``AbstractToolset`` implementation — not
- just the Airflow-native toolsets above. PydanticAI's own ``MCPToolset``
- (built over a FastMCP transport) and third-party toolsets work too. The
- Airflow-native toolsets add connection management, secret backend
- integration, and the connection UI, but you are not locked in.
-
-
-Using Toolsets Directly with PydanticAI
----------------------------------------
-
-Toolsets are standard pydantic-ai ``AbstractToolset`` implementations with no
-dependency on ``AgentOperator`` or ``@task.agent``. You can use them anywhere
-you can run Python within Airflow -- ``@task`` functions, ``PythonOperator``
-callables, or any custom operator's ``execute()`` method -- by creating a
-``pydantic_ai.Agent`` yourself:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
- :language: python
- :start-after: [START howto_task_with_toolsets]
- :end-before: [END howto_task_with_toolsets]
-
-This works because toolsets resolve Airflow connections lazily via
-``BaseHook.get_connection()``, which is available in any task execution
-context.
-
-This approach gives you full control over the agent lifecycle -- you can call
-``agent.run_sync()`` multiple times, swap models at runtime, or combine
-results from several agents in a single task. The tradeoff is that you lose
-the durable execution (step-level caching with retry replay), HITL review
-integration, and automatic tool call logging that ``AgentOperator`` provides.
-
-
-``HookToolset``
----------------
-
-Generic adapter that exposes selected methods of any Airflow Hook as
-pydantic-ai tools via introspection. Requires an explicit ``allowed_methods``
-list — there is no auto-discovery.
-
-.. code-block:: python
-
- from airflow.providers.http.hooks.http import HttpHook
- from airflow.providers.common.ai.toolsets.hook import HookToolset
-
- http_hook = HttpHook(http_conn_id="my_api")
-
- toolset = HookToolset(
- http_hook,
- allowed_methods=["run"],
- tool_name_prefix="http_",
- )
-
-For each listed method, the introspection engine:
-
-1. Builds a JSON Schema from the method signature (``inspect.signature`` +
- ``get_type_hints``).
-2. Extracts the description from the first paragraph of the docstring.
-3. Enriches parameter descriptions from Sphinx ``:param:`` or Google
- ``Args:`` blocks.
-
-Parameters
-^^^^^^^^^^
-
-- ``hook``: An instantiated Airflow Hook.
-- ``allowed_methods``: Method names to expose as tools. Required. Methods
- are validated with ``hasattr`` + ``callable`` at instantiation time.
-- ``tool_name_prefix``: Optional prefix prepended to each tool name
- (e.g. ``"s3_"`` produces ``"s3_list_keys"``).
-
-
-``SQLToolset``
---------------
-
-Curated toolset wrapping
-:class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` with four tools:
-
-.. list-table::
- :header-rows: 1
- :widths: 20 50
-
- * - Tool
- - Description
- * - ``list_tables``
- - Lists available table names (filtered by ``allowed_tables`` if set)
- * - ``get_schema``
- - Returns column names and types for a table
- * - ``query``
- - Executes a SQL query and returns bounded, columnar JSON (see
- :ref:`bounded-query-results`)
- * - ``check_query``
- - Validates SQL syntax without executing it
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets.sql import SQLToolset
-
- toolset = SQLToolset(
- db_conn_id="postgres_default",
- allowed_tables=["customers", "orders"],
- max_rows=20,
- )
-
-The ``DbApiHook`` is resolved lazily from ``db_conn_id`` on first tool call
-via ``BaseHook.get_connection(conn_id).get_hook()``.
-
-In read-only mode (``allow_writes=False``, the default) the ``query`` tool also
-accepts read-only metadata statements -- ``DESCRIBE``/``DESC`` and ``SHOW`` --
-in addition to SELECT-family queries. Agents commonly open with ``DESCRIBE`` to
-learn a table's columns, so permitting it keeps runs deterministic instead of
-hard-failing on schema discovery. The toolset passes the connection's dialect to
-the validator, so ``SHOW`` is recognized on databases that support it (Snowflake,
-MySQL, etc.); on databases without ``SHOW`` it stays rejected. Data-modifying
-statements remain blocked -- including ones hidden behind ``DESCRIBE``/``EXPLAIN``
-(e.g. ``EXPLAIN DELETE ...``, ``DESCRIBE DROP TABLE ...``), which the validator
-rejects by scanning the parsed statement for write operations. When
-``allowed_tables`` is set it scopes these statements too: a ``DESCRIBE`` names a
-table, so its target must be on the list, while ``SHOW`` enumerates objects beyond
-any single table and is rejected outright (see :ref:`allowed-tables-enforcement`).
-
-Multi-schema warehouses
-^^^^^^^^^^^^^^^^^^^^^^^^^
-
-When an agent's tables live in several schemas of one database -- common on
-Snowflake -- list them with schema-qualified ``allowed_tables`` entries:
-
-.. code-block:: python
-
- SQLToolset(
- db_conn_id="snowflake_hq",
- allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"],
- )
-
-``list_tables`` then introspects each referenced schema and returns the matching
-tables fully qualified (e.g. ``MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS``), and
-``get_schema`` routes each qualified name to its own schema. Without this, a
-single ``schema`` only covers one namespace, and leaving ``schema`` unset made
-introspection query a literal ``"None"`` schema and fail. Unqualified entries
-fall back to ``schema``, and table-name matching is case-insensitive (databases
-reflect identifiers in their own case). For tables in a different *database*, use
-a separate toolset whose connection points at that database.
-
-Parameters
-^^^^^^^^^^
-
-- ``db_conn_id``: Airflow connection ID for the database.
-- ``allowed_tables``: Restrict the agent to a fixed set of tables. Omit the
- argument (the default) to expose all tables in ``schema``. No value means
- allow-all: ``None`` and an empty list both raise ``ValueError``, so an allow-list
- built at runtime that resolves to nothing fails at import instead of silently
- exposing every table. Entries may be schema-qualified
- (``"SCHEMA.TABLE"``) to span multiple schemas; see above. Matching is
- case-insensitive. When set, the list is enforced on ``query`` and
- ``check_query`` as well as discovery -- every table a query references must be
- on it. See :ref:`allowed-tables-enforcement` for what this does and does not
- guarantee.
-- ``allowed_functions``: Names of functions that sqlglot does not recognize as
- builtins but are safe to run while ``allowed_tables`` is active (e.g.
- ``["json_build_object"]`` or a project UDF). ``None`` (default) rejects every
- unrecognized function. Matching is case-insensitive. Only consulted when
- ``allowed_tables`` is set.
-- ``schema``: Default schema/namespace for unqualified table listing and
- introspection. Schema-qualified ``allowed_tables`` entries override it per table.
-- ``allow_writes``: Allow data-modifying SQL (INSERT, UPDATE, DELETE, etc.).
- Default ``False`` -- only SELECT-family and read-only metadata
- (``DESCRIBE``/``SHOW``) statements are permitted.
-- ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
- Rows beyond it are not read out of a DBAPI cursor; what the driver has already
- transferred is its own call. See :ref:`bounded-query-results`.
-- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
- See :ref:`bounded-query-results`.
-
-.. _bounded-query-results:
-
-Bounded query results
-^^^^^^^^^^^^^^^^^^^^^
-
-A tool result stays in the model's message history for the rest of the run, so its
-cost is re-paid on every subsequent model request. The ``query`` tool of both
-``SQLToolset`` and ``DataFusionToolset`` bounds that in three ways.
-
-**The result is columnar.** Column names appear once, not once per row:
-
-.. code-block:: json
-
- {"columns": ["id", "name"], "rows": [[1, "Alice"], [2, "Bob"]], "row_count": 2}
-
-On a table with thousands of columns the repeated names, not the values, are the bulk
-of a row-of-dicts payload. Positional rows also keep columns that share a name --
-``SELECT o.id, c.id`` -- which a dict per row silently collapsed to one.
-
-**Rows are fetched, not filtered.** ``max_rows`` bounds what leaves the cursor, so a
-query matching a whole table costs the worker roughly what one matching ``max_rows``
-costs. How much is saved depends on the driver: with a server-side cursor the
-remaining rows are never sent, while a client-buffering driver (psycopg2's default
-cursor, MySQLdb) has already received them and only the per-row conversion is skipped.
-Hooks whose cursor is not DBAPI 2.0 (``ExasolHook`` passes a pyexasol statement) fall
-back to a full fetch, and ``DataFusionToolset`` materializes the full result in the
-engine before the toolset sees it; in both the payload is bounded but the transfer is
-not.
-
-**A byte budget bounds the payload.** ``max_rows`` caps rows, which says nothing about
-size -- one row of a 3000-column table is larger than a thousand rows of a narrow one.
-``max_result_bytes`` is what actually bounds context. Rows are returned as a contiguous
-prefix: the result stops at the first row that does not fit the remaining budget rather
-than skipping it and packing later ones, so a single wide row early in the result ends
-it. The result says which limit it hit:
-
-.. code-block:: json
-
- {"columns": ["..."], "rows": ["..."], "row_count": 3,
- "truncated": true, "truncated_by": "max_result_bytes"}
-
-``truncated_by`` is ``max_rows`` or ``max_result_bytes``. When not even one row fits,
-or the column names alone exceed the budget, the result carries a ``hint`` telling the
-agent to narrow its projection -- the only move that helps. ``total_rows`` is present
-when the driver reports a row count for the query; several (SQLite, some warehouse
-drivers) do not, and it is then omitted rather than guessed.
-
-The default budget is deliberately generous: the columnar shape alone shrinks a wide
-result several-fold, so results that fit before still fit. Lower ``max_result_bytes``
-when an agent makes many queries in one run, since every result is re-paid on every
-later request.
-
-``DataFusionToolset``
----------------------
-
-Curated toolset wrapping
-:class:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine`
-with three tools — ``list_tables``, ``get_schema``, and ``query`` — for
-querying files on object stores (S3, GCS, local filesystem, Iceberg) via Apache DataFusion.
-
-.. list-table::
- :header-rows: 1
- :widths: 20 50
-
- * - Tool
- - Description
- * - ``list_tables``
- - Lists registered table names
- * - ``get_schema``
- - Returns column names and types for a table (Arrow schema)
- * - ``query``
- - Executes a SQL query and returns bounded, columnar JSON (see
- :ref:`bounded-query-results`)
-
-Each :class:`~airflow.providers.common.sql.config.DataSourceConfig` entry
-registers a table backed by Parquet, CSV, Avro, or Iceberg data. Multiple
-configs can be registered so that SQL queries can join across tables.
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
- from airflow.providers.common.sql.config import DataSourceConfig
-
- toolset = DataFusionToolset(
- datasource_configs=[
- DataSourceConfig(
- conn_id="aws_default",
- table_name="sales",
- uri="s3://my-bucket/data/sales/",
- format="parquet",
- ),
- DataSourceConfig(
- conn_id="aws_default",
- table_name="returns",
- uri="s3://my-bucket/data/returns/",
- format="csv",
- ),
- ],
- max_rows=100,
- )
-
-The ``DataFusionEngine`` is created lazily on the first tool call. This
-toolset requires the ``datafusion`` extra of
-``apache-airflow-providers-common-sql``.
-
-Parameters
-^^^^^^^^^^
-
-- ``datasource_configs``: One or more
- :class:`~airflow.providers.common.sql.config.DataSourceConfig` entries.
- Requires ``apache-airflow-providers-common-sql[datafusion]``.
-- ``allow_writes``: Allow data-modifying SQL (CREATE TABLE, CREATE VIEW,
- INSERT INTO, etc.). Default ``False`` — only SELECT-family statements are
- permitted. DataFusion on object stores is mostly read-only, but it does
- support DDL for in-memory tables; this guard blocks those by default.
-- ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
-- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
- See :ref:`bounded-query-results`.
-
-``LoggingToolset``
-------------------
-
-:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` is a
-``WrapperToolset`` that intercepts ``call_tool()`` to log each tool invocation
-in real time. ``AgentOperator`` applies it automatically (see
-``enable_tool_logging``), but you can also use it directly with any pydantic-ai
-``Agent``:
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets.logging import LoggingToolset
- from airflow.providers.common.ai.toolsets.sql import SQLToolset
-
- sql_toolset = SQLToolset(db_conn_id="my_db")
- logged_toolset = LoggingToolset(wrapped=sql_toolset, logger=my_logger)
-
-Each tool call produces two INFO log lines (name + timing) and optional
-DEBUG-level argument logging. Exceptions are logged and re-raised.
-
-
-``MCPToolset``
---------------
-
-Connects to an `MCP (Model Context Protocol) <https://modelcontextprotocol.io/>`__
-server configured via an Airflow connection. MCP is an open protocol that lets
-LLMs interact with external tools and data sources through a standardized
-interface.
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets.mcp import MCPToolset
-
- toolset = MCPToolset(
- mcp_conn_id="my_mcp_server",
- tool_prefix="weather",
- )
-
-The MCP server is resolved lazily from the Airflow connection on the first
-tool call. See :ref:`howto/connection:mcp` for connection configuration.
-
-Requires the ``mcp`` extra: ``pip install "apache-airflow-providers-common-ai[mcp]"``
-
-Parameters
-^^^^^^^^^^
-
-- ``mcp_conn_id``: Airflow connection ID for the MCP server.
-- ``tool_prefix``: Optional prefix prepended to tool names to avoid
- collisions when using multiple MCP servers (e.g. ``"weather"`` produces
- ``"weather_get_forecast"``).
-- ``token_provider``: Optional zero-argument callable returning a bearer token.
- When set, it overrides the connection's static ``password`` for the
- ``Authorization`` header. Called once, the first time this toolset
- establishes a connection -- use it for short-lived or minted tokens (e.g. a
- Snowflake managed MCP server authenticated with a key-pair JWT). See
- :ref:`howto/connection:mcp`.
-- ``env_provider``: Optional zero-argument callable returning a
- ``dict[str, str]`` merged over the connection's ``Extra.env`` (winning on key
- conflicts) for the ``stdio`` subprocess environment -- use it when the
- credential a local stdio MCP server needs lives in a different connection, or
- is minted fresh per call (e.g. a Splunk/Vault token), rather than storing it
- statically here. Called once, the first time this toolset establishes a
- connection. See :ref:`howto/connection:mcp`.
-
-Using Multiple MCP Servers
-^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-.. code-block:: python
-
- AgentOperator(
- task_id="multi_mcp",
- prompt="Get the weather in London and run a calculation",
- llm_conn_id="pydanticai_default",
- toolsets=[
- MCPToolset(mcp_conn_id="weather_mcp", tool_prefix="weather"),
- MCPToolset(mcp_conn_id="code_runner_mcp", tool_prefix="code"),
- ],
- )
-
-Direct PydanticAI MCP Toolsets
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-For prototyping or when you want full PydanticAI control, you can pass
-``MCPToolset`` instances directly — no Airflow connection needed:
-
-.. code-block:: python
-
- from fastmcp.client.transports import StdioTransport
- from pydantic_ai.mcp import MCPToolset
-
- AgentOperator(
- task_id="direct_mcp",
- prompt="What tools are available?",
- llm_conn_id="pydanticai_default",
- toolsets=[
- MCPToolset("http://localhost:3001/mcp"),
- MCPToolset(StdioTransport(command="uvx", args=["mcp-run-python"])),
- ],
- )
-
-This works because PydanticAI's ``MCPToolset`` implements ``AbstractToolset``.
-The tradeoff: URLs and credentials are hardcoded in Dag code instead of being
-managed through Airflow connections and secret backends.
-
-
-.. _agent-skills:
-
-``AgentSkillsToolset``
-----------------------
-
-:class:`~airflow.providers.common.ai.toolsets.skills.AgentSkillsToolset` loads
-`Agent Skills <https://agentskills.io>`__ -- ``SKILL.md`` bundles (instructions,
-and optionally scripts and resources) that the model discovers and loads *on
-demand*. Only a compact catalog of skill names and descriptions sits in the
-prompt until the model decides it needs one, so a large skill library costs few
-tokens until used (progressive disclosure).
-
-It is backed by the community `pydantic-ai-skills
-<https://github.com/DougTrajano/pydantic-ai-skills>`__ package (MIT); native
-progressive disclosure is in flight upstream in `pydantic/pydantic-ai#5230
-<https://github.com/pydantic/pydantic-ai/pull/5230>`__. Install the optional
-extra to use it:
-
-.. code-block:: bash
-
- pip install "apache-airflow-providers-common-ai[skills]"
-
-Each source is a local directory or a connection-resolved
-:class:`~airflow.providers.common.ai.skills.GitSkills`. Sources are resolved when
-the agent enters the toolset, on the worker -- never while the Dag processor
-parses the file -- so a Git token is never baked into the serialized Dag, and
-cloned repositories are removed when the run ends.
-
-A local directory of ``SKILL.md`` bundles:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
- :language: python
- :start-after: [START howto_operator_agent_skills_local]
- :end-before: [END howto_operator_agent_skills_local]
-
-A Git repository, with credentials from an Airflow connection:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
- :language: python
- :start-after: [START howto_operator_agent_skills_git]
- :end-before: [END howto_operator_agent_skills_git]
-
-For a private repository, point ``conn_id`` at a
-:doc:`git connection <apache-airflow-providers-git:connections/git>`; credentials
-are resolved through the Git provider's ``GitHook`` (an HTTPS token in the
-connection password, or an SSH key in the connection's extra). A plain ``http://``
-URL with ``conn_id`` is rejected so a credential is never sent in cleartext, and a
-``repo_url`` that embeds a username/password is rejected (use ``conn_id``). After
-cloning, the credential is stripped from the checkout's ``.git/config``. As with
-any ``git clone``, the worker's own git configuration (credential helpers, SSH
-agent) may still apply, so run workers without ambient git credentials if you
-need strict isolation.
-
-.. warning::
-
- Skill bundles can contain scripts that the agent may run on the worker via
- the ``run_skill_script`` tool. For a remote source, anyone who can modify the
- repository can introduce code that executes on your worker, outside Dag
- review and versioning. Point ``GitSkills`` at a trusted repository, pin
- ``branch`` to a trusted ref, and treat skill contents as code that runs in
- your environment.
-
-Parameters
-^^^^^^^^^^
-
-- ``sources``: List of skill sources -- local directory paths and/or
- :class:`~airflow.providers.common.ai.skills.GitSkills`.
-- ``exclude_tools``: Optional set of skill tool names to hide from the agent
- (e.g. ``{"run_skill_script"}`` to disable on-worker script execution).
-- ``exclude_resources``: Optional glob patterns to exclude from resource
- discovery, added on top of the built-in defaults (``__pycache__``, ``*.pyc``,
- ``*.pyo``, ``.DS_Store``, ``.git``). A skill exposes every readable text file
- it contains as a resource; these patterns keep matched files out of the
- resource list and the ``read_skill_resource`` tool (e.g.
- ``["*.env", "secrets/*"]``). Each pattern matches the full skill-relative path
- or any single path component. This hides files from resource discovery only --
- it does not stop a skill's ``run_skill_script`` from reading them off disk, so
- pair it with ``exclude_tools={"run_skill_script"}`` when the files are
- genuinely sensitive.
-
-Using Agent Skills with other frameworks
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-``AgentSkillsToolset`` is a standard pydantic-ai toolset, so it also works with a
-plain ``pydantic_ai.Agent`` you build yourself, not just ``AgentOperator``.
-
-Because Agent Skills is a cross-framework format, the connection handling is also
-reusable through :func:`~airflow.providers.common.ai.skills.resolve_skills`, which
-resolves sources to local ``SKILL.md`` directories that any loader accepts:
-
-.. code-block:: python
-
- from airflow.providers.common.ai.skills import GitSkills, resolve_skills
-
- sources = ["./skills", GitSkills(repo_url="https://github.com/org/skills", conn_id="github_skills")]
- with resolve_skills(sources) as dirs:
- # LangChain DeepAgents
- agent = create_deep_agent(model="openai:gpt-5.4", skills=dirs)
- # ...or Strands
- agent = Agent(plugins=[AgentSkills(skills=dirs)])
-
-``resolve_skills`` needs the Git provider (for ``GitSkills``) but not pydantic-ai,
-and removes any cloned directories when the ``with`` block exits.
-
-
-``SandboxToolset``
-------------------
-
-:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` gives the
-agent a shell and a filesystem inside a disposable sandbox, provisioned off the
-Airflow worker by a :class:`~airflow.providers.common.ai.sandbox.SandboxBackend`
-and destroyed when the agent run ends. It exposes four tools, ``run_command``,
-``read_file``, ``write_file`` and ``list_directory``, and ships with a hosted
-backend on Modal for production and a local microVM backend for development.
-
-It has its own page: :doc:`sandbox`, covering when to reach for it against code
-mode, ``KubernetesPodOperator`` and the ``KubernetesExecutor``, how to configure
-egress, packages and credentials, how a result gets out, what it costs, and the
-limitations to read before designing a Dag around it.
-
-.. code-block:: python
-
- from airflow.providers.common.ai.operators.agent import AgentOperator
- from airflow.providers.common.ai.sandbox import ModalSandboxBackend
- from airflow.providers.common.ai.toolsets import SandboxToolset
-
- AgentOperator(
- task_id="sandboxed_analyst",
- prompt="Estimate pi with a Monte Carlo simulation of one million points.",
- llm_conn_id="pydanticai_default",
- toolsets=[SandboxToolset(ModalSandboxBackend())],
- )
-
-
-Working with LangChain
-----------------------
-
-Tools bridge in both directions between common.ai's toolsets and LangChain.
-
-**LangChain tools → ``AgentOperator``.** No Airflow code is needed. pydantic-ai
-ships `pydantic_ai.ext.langchain.LangChainToolset
-<https://ai.pydantic.dev/toolsets/>`__ upstream, which wraps existing LangChain
-tools as an ``AbstractToolset``. Drop it straight into ``AgentOperator``:
-
-.. code-block:: python
-
- from pydantic_ai.ext.langchain import LangChainToolset
-
- AgentOperator(
- task_id="agent_with_langchain_tools",
- prompt="Research the question and summarise.",
- llm_conn_id="pydanticai_default",
- toolsets=[LangChainToolset([my_langchain_tool])],
- )
-
-**common.ai toolsets → LangChain.** The reverse direction is what
-:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`
-provides. It converts any pydantic-ai toolset -- including ``SQLToolset``,
-``HookToolset``, and ``MCPToolset`` -- into a list of LangChain
-``StructuredTool`` objects, so a LangChain agent or chain can call Airflow's
-curated, connection-managed tools:
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_langchain_toolset_bridge.py
- :language: python
- :start-after: [START example_langchain_toolset_bridge]
- :end-before: [END example_langchain_toolset_bridge]
-
-Each generated tool keeps the source tool's name, description, and argument
-schema, and routes calls back through the original toolset, so the toolset's own
-behavior (connection resolution, ``SQLToolset``'s SQL validation, and
-``allowed_tables`` filtering) still applies. ``get_tools`` runs eagerly at
-conversion time to enumerate the tools.
-
-When a toolset raises pydantic-ai's ``ModelRetry`` to ask the model to correct
-its input (``SQLToolset`` does this on, for example, an unknown column), the
-bridge returns that message as the tool's output so the model sees it and tries
-again. ``ModelRetry`` is a feed-the-model-and-retry signal rather than a
-failure, so returning it preserves the self-correction the toolset was written
-for and works no matter how the agent is configured to handle tool errors
-(raising would abort the run under ``create_agent``'s default handling).
-
-The bridge does not hold a toolset session open across calls: ``get_tools`` and
-every tool call each run under their own event loop, so for ``MCPToolset`` the
-connection is opened and torn down around each call. It reconnects per call,
-which is fine for stateless tools but unsuitable for ``stdio`` MCP servers (or
-any server that keeps state between calls), since each call starts a fresh
-session.
-
-.. note::
-
- Outside an agent run there is no live ``RunContext``, so the bridge builds a
- minimal one with an inert placeholder model. The bundled toolsets ignore the
- context, so this is transparent for them. A custom toolset that reads live
- run state (``ctx.model``, ``ctx.messages``, ``ctx.usage``) will not behave
- correctly when bridged standalone.
-
-Requires the ``langchain`` extra:
-``pip install "apache-airflow-providers-common-ai[langchain]"``
-
-
-Security
---------
-
-LLM agents call tools based on natural-language reasoning. This makes them
-powerful but introduces risks that don't exist with deterministic operators.
-
-What the agent can and cannot reach
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-An agent's reach is exactly the set of tools you register on it, and nothing
-more. The model never executes arbitrary code: it can only request one of the
-tools you provided, and pydantic-ai rejects any tool name outside that set
-before it runs. If no registered tool can read the environment, the
-filesystem, or other connections, the model cannot reach them, regardless of
-what the prompt instructs it to do.
-
-This is what "untrusted" means in this context. The Dag file itself is
-author-written and trusted, exactly like any other Dag. What is untrusted is
-the model's *output*: the tool-call requests and text it generates. That output
-is confined to your registered tools and bounded by the tool-call budget. An
-agent cannot create a new connection, read another connection's credentials, or
-run a shell command unless a tool you registered exposes that capability.
-
-The corollary is that every tool you add widens the blast radius, and a custom
-toolset is only as safe as you make it. A tool that returns ``os.environ`` or
-runs shell commands hands the model whatever that tool can reach. Audit any
-custom toolset, and any MCP server you connect through ``MCPToolset``, against
-the same standard the bundled toolsets below are built to.
-
-.. _toolset-defense-layers:
-
-Defense Layers
-^^^^^^^^^^^^^^
-
-No single layer is sufficient — they work together.
-
-.. list-table::
- :header-rows: 1
- :widths: 20 40 40
-
- * - Layer
- - What it does
- - What it does NOT do
- * - **Airflow Connections**
- - Credentials are stored in Airflow's secret backend, never in Dag code.
- The LLM agent cannot see API keys or database passwords.
- - Does not prevent the agent from using the connection to access data
- the connection has access to.
- * - **HookToolset: explicit allow-list**
- - Only methods listed in ``allowed_methods`` are exposed as tools.
- Auto-discovery is not supported. Methods are validated at Dag parse
- time.
- - Does not restrict what arguments the agent passes to allowed methods.
- * - **SQLToolset: read-only by default**
- - ``allow_writes=False`` (default) validates every SQL query through
- ``validate_sql()``: SELECT-family and read-only metadata
- (``DESCRIBE``/``SHOW``) statements pass; INSERT, UPDATE, DELETE, DROP,
- and writes hidden behind ``EXPLAIN`` are rejected.
- - Does not prevent the agent from reading sensitive data that the
- database user has SELECT access to.
- * - **DataFusionToolset: read-only by default**
- - ``allow_writes=False`` (default) validates every SQL query through
- ``validate_sql()`` and rejects CREATE TABLE, CREATE VIEW, INSERT
- INTO, and other non-SELECT statements.
- - Does not prevent the agent from reading any registered data source.
- * - **SQLToolset: allowed_tables**
- - Restricts the agent to listed tables across ``list_tables``,
- ``get_schema``, ``query``, and ``check_query``. Queries are parsed and
- every referenced table (including via subqueries, CTEs, JOINs, and
- ``DESCRIBE``) is checked against the list before execution.
- - Rejects ``COPY`` and every function sqlglot cannot type (the channel for
- ``pg_read_file`` / ``query_to_xml`` / ``dblink``) unless named in
- ``allowed_functions``. Fail-closed, but only as exact as the SQL parser. Not a
- security boundary -- always pair it with least-privilege database grants. See
- :ref:`allowed-tables-enforcement` below.
- * - **SQLToolset: max_rows / max_result_bytes**
- - Bounds a query result by rows (default 50) and by serialized size
- (default 64 KiB), preventing the agent from pulling entire tables into
- context.
- - Does not limit the number of queries the agent can make, and each result
- stays in message history for the rest of the run. Rows past ``max_rows``
- are not read out of the cursor, but a client-buffering driver has already
- transferred them -- this bounds context, not database or network load.
- * - **MCPToolset: external server**
- - Connects the agent to tools exposed by an MCP server, authenticated
- through an Airflow connection.
- - Does **not** constrain what those tools do. An MCP server can expose
- shell, filesystem, or network access. Run only trusted servers and
- audit the tools they expose.
- * - **SandboxToolset: off-worker execution**
- - Runs the agent's commands and file operations in a disposable microVM,
- never in the worker process. Airflow injects nothing of its own; only
- what ``SandboxSpec`` names goes in, and egress is denied by default. A
- backend that cannot enforce a spec field raises rather than ignoring it.
- Commands are bounded by a timeout ceiling and per-stream output caps.
- - **Does not contain the agent.** The agent loop and every other toolset on
- the same agent still run in the worker with its credentials, so this does
- 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
- 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
- layer and leaves DNS resolution open.
- * - **pydantic-ai: tool call budget**
- - pydantic-ai's ``max_result_retries`` and ``model_settings`` control
- how many tool-call rounds the agent can make before stopping.
- - Requires explicit configuration — the default allows many rounds.
-
-
-.. _allowed-tables-enforcement:
-
-How ``allowed_tables`` Is Enforced
-""""""""""""""""""""""""""""""""""
-
-When ``allowed_tables`` is set it governs every tool, not just discovery:
-
-- ``list_tables`` and ``get_schema`` only reveal listed tables.
-- ``query`` and ``check_query`` parse the SQL with `sqlglot
- <https://github.com/tobymao/sqlglot>`_ and reject it before execution if it
- references any table that is not on the list. Tables reached indirectly are
- caught too -- through subqueries, CTEs, JOINs, set operations (``UNION`` etc.),
- ``DESCRIBE``, catalog views such as ``information_schema``, and DML. CTE
- references are excluded by lexical scope, so a same-named CTE in another scope
- cannot hide a real table, and the database/catalog is part of the match, so a
- cross-database reference like ``otherdb.public.orders`` is refused.
-- Constructs the list cannot describe are rejected outright while it is active:
- table-valued functions (``dblink``), ``TABLE('name')`` row sources, the
- ``TABLE <name>`` shorthand, ``SHOW``, dynamic SQL (``EXEC``), ``COPY``
- (file/program I/O), and **inline comments** -- because parser-vs-engine differences
- hide in comments (MySQL executes ``/*! ... */`` while sqlglot and other engines
- ignore it).
-- **Any function sqlglot does not recognize is rejected (fail-closed).** A function
- whose string argument reaches data outside the table graph --
- ``pg_read_file('/etc/passwd')`` (a file), ``query_to_xml('SELECT * FROM other_table', ...)``
- (SQL over another table), a scalar ``dblink`` (a remote database) -- carries no table
- reference for the parser to catch. Rather than maintain a denylist of such functions
- (unbounded, engine-specific, and it would fail *open* on anything missed), the toolset
- rejects every function sqlglot cannot type. Ordinary builtins (``count``, ``lower``,
- ``sum``) are recognized and pass. A legitimate function sqlglot does not type
- (``json_build_object``, ``jsonb_agg``) or a project UDF is rejected until you list it
- in ``allowed_functions``:
-
- .. code-block:: python
-
- SQLToolset(
- db_conn_id="analytics_db",
- allowed_tables=["orders"],
- allowed_functions=["json_build_object"], # opt in per function you trust
- )
-
-So ``SELECT * FROM secrets`` with ``allowed_tables=["orders"]`` is refused, and
-the rejection is handed back to the agent so it can re-target an allowed table.
-
-.. warning::
-
- This is a strong **application-level guardrail, not a security boundary.** The
- fail-closed function check raises the bar, but any query the engine parses
- differently from sqlglot is a residual gap, and ``allowed_functions`` is a trust
- decision you own. **Always** point the connection at a least-privilege database
- role -- that is the boundary that holds even when the parser cannot see through a
- function, and it is what actually keeps an agent (which may be under prompt
- injection) away from data and files you have not granted it:
-
- .. code-block:: sql
-
- -- Create a read-only role with access to specific tables only
- CREATE ROLE airflow_agent_reader;
- GRANT SELECT ON orders, customers TO airflow_agent_reader;
- -- Use this role's credentials in the Airflow connection
-
-Defense in depth: the allow-list contains the agent's *intent* (and gives it a
-correctable error), while the database role is the boundary that holds even if
-the agent reaches data the parser cannot see. The connection should use a
-database user with the minimum privileges required.
-
-
-HookToolset Guidelines
-""""""""""""""""""""""
-
-- List only the methods the agent needs. Never expose ``run()`` or
- ``get_connection()`` — these give broad access.
-- Prefer read-only methods (``list_*``, ``get_*``, ``describe_*``).
-- The agent controls arguments. If a method accepts a ``path`` parameter,
- the agent can pass any path the hook has access to.
-
-.. code-block:: python
-
- # Good: expose only list and read
- HookToolset(
- s3_hook,
- allowed_methods=["list_keys", "read_key"],
- tool_name_prefix="s3_",
- )
-
- # Bad: exposes delete and write operations
- HookToolset(
- s3_hook,
- allowed_methods=["list_keys", "read_key", "delete_object", "load_string"],
- )
-
-
-Recommended Configuration
-"""""""""""""""""""""""""
-
-**Read-only analytics** (the most common pattern):
-
-.. code-block:: python
-
- SQLToolset(
- db_conn_id="analytics_readonly", # Connection with SELECT-only grants
- allowed_tables=["orders", "customers"], # Hide other tables from agent
- allow_writes=False, # Default — validates SQL
- max_rows=50, # Default — cap rows
- max_result_bytes=65536, # Default — cap bytes; lower it for wide tables
- )
-
-**Agents that need to modify data** (use with caution):
-
-.. code-block:: python
-
- SQLToolset(
- db_conn_id="app_db",
- allowed_tables=["user_preferences"],
- allow_writes=True, # Disables SQL validation — agent can INSERT/UPDATE
- max_rows=100,
- )
-
-
-Production Checklist
-""""""""""""""""""""
-
-Before deploying an agent task to production:
-
-1. **Connection credentials**: Use Airflow's secret backend. Never hardcode
- API keys in Dag files.
-2. **Database permissions**: Create a dedicated database user with minimum
- required grants. Don't reuse the admin connection.
-3. **Tool allow-list**: Review ``allowed_methods`` / ``allowed_tables``. The
- agent can call any exposed tool with any arguments.
-4. **Read-only default**: Keep ``allow_writes=False`` unless the task
- specifically requires writes.
-5. **Result limits**: Set ``max_rows`` and ``max_result_bytes`` appropriate to
- the use case. ``max_rows`` alone does not bound size -- on wide tables it is
- ``max_result_bytes`` that keeps a result from dominating the context window for
- the rest of the run.
-6. **Model budget**: Configure pydantic-ai's ``model_settings`` (e.g.
- ``max_tokens``) and ``retries`` to bound cost and prevent runaway loops.
-7. **System prompt**: Include safety instructions in ``system_prompt`` (e.g.
- "Only query tables related to the question. Never modify data.").
-8. **Prompt injection**: Be cautious when the prompt includes untrusted data
- (user input, external API responses, upstream XCom). Consider sanitizing
- inputs before passing them to the agent.
-
-.. _managed-agent-toolsets:
-
-Managed Agent Toolsets
-----------------------
-
-Cloud vendors now run agents on your behalf — Snowflake Cortex Agents, Amazon
-Bedrock AgentCore runtimes, Azure AI Foundry hosted agents, Vertex AI Agent
-Engine. Their reasoning loops execute on the vendor's infrastructure, so they
-are not something ``AgentOperator`` runs; they are something an Airflow task
-*consults*.
-
-:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
-is the contract for exposing one of those as a tool. Each provider package
-ships its own subclass, so credentials keep flowing through that provider's
-existing hook and no new connection types are needed.
-
-A subclass implements two members:
-
-``agent_ref``
- Normalized identity of the remote agent — ``platform`` and ``name`` — logged
- on every call so a run can be audited for which agents it consulted.
-
-``invoke_sync(prompt)``
- Send the prompt, return the agent's answer. Return the *answer*, not the
- transport envelope.
-
-Tool naming, argument validation, result serialisation, and logging are handled
-by the base class, so every provider's implementation presents the same surface
-to the calling model.
-
-``tool_name`` is the required identifier — it is what the model emits when it
-calls the tool, and the Dag author chooses it. ``description`` is optional and
-falls back to the tool name rendered as prose, the same way ``HookToolset``
-derives one from a method name when there is no docstring.
-
-.. note::
-
- Writing a description is still worth the line. It is what tells the model to
- consult the agent rather than answer from its own knowledge, and it is the
- only place to record a scope limit the name cannot carry — "cannot see
- revenue figures". Because the argument schema is always a bare prompt, the
- name and the description are the whole of what the model knows about the
- agent.
-
-Sync or async?
-^^^^^^^^^^^^^^
-
-An agent run is driven by a single event loop, shared by the model requests and
-every toolset's tool calls. A call that blocks that loop stalls all of them for
-as long as the remote agent reasons -- seconds to minutes, not milliseconds --
-and it does so silently, with no error to point at.
-
-Implement ``invoke_sync`` when the vendor call is blocking. The base class runs
-it in a worker thread, keeping the loop free::
-
- class MyManagedAgentToolset(BaseManagedAgentToolset):
- @property
- def agent_ref(self) -> dict[str, str]:
- return {"platform": "example.cloud", "name": self._agent_id}
-
- def invoke_sync(self, prompt: str) -> str:
- return self._hook.conn.invoke(agent=self._agent_id, prompt=prompt)
-
-Override ``invoke(prompt)`` -- the async coroutine -- when the call is already
-asynchronous. Implement one hook or the other; a subclass that implements
-neither is rejected when it is constructed.
-
-.. warning::
-
- A thread cannot be cancelled. A caller that stops waiting for
- ``invoke_sync`` does not stop the call -- it keeps a worker thread and its
- socket until the call returns. Set a timeout on the underlying request so
- that is bounded.
-
-Toolset or operator?
-^^^^^^^^^^^^^^^^^^^^
-
-Most managed-agent platforms do not offer a plain one-request-one-answer API. Some
-require polling a job; others require creating a session and tearing it down around
-each exchange. A toolset can do either, but only by blocking inside
-``invoke()`` — it cannot defer to the Triggerer, and it has no post-task hook to
-clean up with if the worker dies mid-call.
-
-That draws a boundary worth respecting:
-
-.. list-table::
- :header-rows: 1
- :widths: 45 55
-
- * - Shape
- - Surface to use
- * - A short consultation *inside* an agent's reasoning, where failing the task
- would discard the calling agent's accumulated context
- - A managed agent toolset
- * - Long-running submitted work as a pipeline step in its own right
- - That provider's own operator, with deferral or
- :class:`~airflow.sdk.bases.resumablejobmixin.ResumableJobMixin`
-
-
-``ResumableJobMixin`` exists for exactly the second case: it persists the external
-job ID to the task state store before polling, so a worker crash reconnects to the
-running job instead of submitting a duplicate. A toolset cannot offer that, because
-the retry boundary is the task, not the tool call — on retry the agent loop restarts
-and re-issues the call. Durable execution covers the *completed* call (see
-``replayable`` below); it does not cover a call that was still in flight.
-
-Error handling
-^^^^^^^^^^^^^^
-
-Failures sort into three buckets, and conflating them is the most common way an
-implementation goes wrong:
-
-.. list-table::
- :header-rows: 1
- :widths: 22 30 48
-
- * - Raise
- - When
- - Who recovers
- * - ``ModelRetry``
- - The agent rejected the request in a way rephrasing could fix.
- - The calling model, bounded by its ``usage_limits``.
- * - ``ManagedAgentInvocationError``
- - Terminal: bad credentials, missing agent, revoked quota.
- - Nobody — the task fails fast instead of burning retries.
- * - *let it propagate*
- - Transient: 429, 5xx, connection reset, read timeout.
- - Airflow's task-level retry. A rephrase does nothing for a 503.
-
-Durable execution
-^^^^^^^^^^^^^^^^^
-
-``replayable`` is ``False`` by default. A managed agent may act on systems
-Airflow cannot observe, so replaying a cached answer on retry could skip a side
-effect. Implementations whose agent is read-only should set it to ``True`` to
-avoid paying for the same invocation twice.
-
-Deferral
-^^^^^^^^
-
-A toolset call runs in the worker and cannot defer to the Triggerer — it blocks
-for the duration of the call. See `Toolset or operator?`_ above for when that is
-acceptable and when the provider's own deferrable operator is the right surface
-instead.
-
-Failover between interchangeable agents
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-:class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
-composes several managed agents into one tool, trying them in order until one
-answers. It is itself a ``BaseManagedAgentToolset``, so the calling model sees a
-single tool and has no say in which provider serves the request — the policy
-stays deterministic Python rather than a prompt instruction a model may ignore.
-Groups nest.
-
-.. code-block:: python
-
- from airflow.providers.common.ai.toolsets import FailoverManagedAgentToolset
-
- resilient = FailoverManagedAgentToolset(
- tool_name="ask_claims_agent",
- description="Reviews an insurance claim and returns a coverage determination.",
- members=[bedrock_claims_agent, foundry_claims_agent], # same image, two clouds
- )
-
-Members must satisfy two preconditions the class cannot check.
-
-**Substitutability.** The same agent deployed twice, not two specialists with
-different data. Two containerised agents built from one image qualify; agents
-bound to one platform's own objects — a Cortex Agent over Snowflake semantic
-models — do not, because there is nothing equivalent to fail over *to*.
-
-**Statelessness per invocation.** Server-side conversation state is the norm
-across managed-agent platforms, not the exception — optional on some (Cortex
-``thread_id``), mandatory on others where a session is created and torn down
-around each exchange. Each member is invoked with a bare prompt and no thread
-reference, so a failover silently starts a fresh conversation on the standby:
-correct for a one-shot consultation, wrong for a multi-turn one. Treat one-shot
-as a restriction a group is deliberately held to, not a safe default.
-
-The three error buckets do real work here:
-
-- ``ManagedAgentInvocationError`` and transient failures move to the next member.
-- ``ModelRetry`` is re-raised immediately and never triggers failover. A prompt
- the primary could not parse will not parse on the standby either, so failing
- over would spend the standby's budget reproducing the same error.
-- The last member's exception propagates unchanged, so a total outage still fails
- the task rather than returning something misleading.
-
-``failover_on`` defaults to ``Exception`` because ``common.ai`` cannot enumerate
-the cloud SDKs' exception trees — ``requests``, ``botocore`` and the Azure SDK
-share no common base. It can be narrowed when the members' exception types are
-known.
-
-``replayable`` on a group is ``True`` only when every member is, because the
-durable cache cannot know which member produced the answer it holds.
-
-.. note::
-
- For a **standalone** agent call, prefer plain Airflow task-level failover:
- two tasks, the second with ``trigger_rule=TriggerRule.ALL_FAILED``. That keeps
- which provider served the request visible in the grid at no code cost, and
- makes failover rate a task metric. This class is for the case a task boundary
- cannot express — a managed agent consulted as a tool *inside* a longer agent
- run, where failing the task would discard the calling agent's accumulated
- context and re-run every earlier tool call.
-
-Two counters make failover visible, because a failover is a *success-shaped*
-event — without them a primary that has been down for a week looks identical to a
-healthy one:
-
-.. list-table::
- :header-rows: 1
- :widths: 30 70
-
- * - Metric
- - Tags
- * - ``managed_agent.failover``
- - ``from_platform``, ``to_platform`` — one per failover transition
- * - ``managed_agent.served``
- - ``platform``, ``role`` (``primary`` / ``standby``) — one per answer
-
-The standby-served fraction is a ratio over ``managed_agent.served`` alone, so
-"are we quietly running on the standby?" is a dashboard question rather than a log
-grep. Both are tagged by platform rather than agent name to keep cardinality
-bounded.
-
-One limitation remains: which member served a *particular* answer is in the task
-log but not in XCom. ``agent_ref`` on a group describes the group, not the
-responder, because the responder is not known until after the call. The counters
-cover the operational question; per-answer provenance for an audit trail would
-need ``AgentOperator`` to collect per-toolset metadata.
diff --git a/providers/common/ai/docs/toolsets/datafusion.rst b/providers/common/ai/docs/toolsets/datafusion.rst
new file mode 100644
index 00000000000..380689a3783
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/datafusion.rst
@@ -0,0 +1,166 @@
+ .. 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.
+
+``DataFusionToolset``
+=====================
+
+Curated toolset wrapping
+:class:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine`
+with three tools — ``list_tables``, ``get_schema``, and ``query`` — for
+querying files on object stores (S3, GCS, local filesystem, Iceberg) via Apache DataFusion.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 50
+
+ * - Tool
+ - Description
+ * - ``list_tables``
+ - Lists registered table names
+ * - ``get_schema``
+ - Returns column names and types for a table (Arrow schema)
+ * - ``query``
+ - Executes a SQL query and returns bounded, columnar JSON (see
+ :ref:`bounded-query-results`)
+
+Each :class:`~airflow.providers.common.sql.config.DataSourceConfig` entry
+registers a table backed by Parquet, CSV, Avro, or Iceberg data. Multiple
+configs can be registered so that SQL queries can join across tables.
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+ from airflow.providers.common.sql.config import DataSourceConfig
+
+ toolset = DataFusionToolset(
+ datasource_configs=[
+ DataSourceConfig(
+ conn_id="aws_default",
+ table_name="sales",
+ uri="s3://my-bucket/data/sales/",
+ format="parquet",
+ ),
+ DataSourceConfig(
+ conn_id="aws_default",
+ table_name="returns",
+ uri="s3://my-bucket/data/returns/",
+ format="csv",
+ ),
+ ],
+ max_rows=100,
+ )
+
+The ``DataFusionEngine`` is created lazily on the first tool call. This
+toolset requires the ``datafusion`` extra of
+``apache-airflow-providers-common-sql``.
+
+Parameters
+----------
+
+- ``datasource_configs``: One or more
+ :class:`~airflow.providers.common.sql.config.DataSourceConfig` entries.
+ Requires ``apache-airflow-providers-common-sql[datafusion]``.
+- ``allow_writes``: Allow data-modifying SQL (CREATE TABLE, CREATE VIEW,
+ INSERT INTO, etc.). Default ``False`` — only SELECT-family statements are
+ permitted. DataFusion on object stores is mostly read-only, but it does
+ support DDL for in-memory tables; this guard blocks those by default.
+- ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
+- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
+ See :ref:`bounded-query-results`.
+
+When to choose it
+-----------------
+
+**Choose it when** the data is files on an object store rather than rows in a
+database — Parquet, CSV or Avro — or a table in a catalog such as Iceberg, and
+you want the agent to ask SQL questions of them without loading them anywhere
+first. (This route needs the ``datafusion`` extra of
+``apache-airflow-providers-common-sql``.) Each ``DataSourceConfig`` registers
+one table, and several can be registered so the agent can join across them.
+The two shapes take different fields: an object-store format needs a
+``uri``, while a catalog format like
+Iceberg is looked up by ``db_name`` instead, and ``DataSourceConfig`` raises
+``ValueError`` at construction if a catalog format is missing one.
+
+**What it cannot do**
+
+- It has no table allow-list. ``allow_writes=False`` is the only guard, and it
+ blocks non-SELECT statements, not reach: the defense-layer table records that
+ this toolset "does not prevent the agent from reading any registered data
+ source". The registration list is therefore the whole boundary — register
+ exactly what the agent may read.
+- It bounds the payload, not the scan. DataFusion has already materialized the
+ full result before ``max_rows`` and ``max_result_bytes`` apply, so those limits
+ protect the model's context, not the cost of the query.
+- It cannot tell failure kinds apart precisely. The DataFusion Python bindings
+ expose no native exception types, so the retry decision is made by matching the
+ error message against regular expressions — which a wording change upstream can
+ quietly defeat.
+
+**A real example.** The same bucket as the ``HookToolset`` example on :doc:`hook`, reached
+the other way. Rather than exposing ``list_keys`` and ``read_key`` and leaving
+the agent to reassemble files, this registers the prefix as a table and lets it
+write SQL:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+ from airflow.providers.common.sql.config import DataSourceConfig
+
+ toolset = DataFusionToolset(
+ datasource_configs=[
+ DataSourceConfig(
+ conn_id="aws_default",
+ table_name="sales",
+ uri="s3://my-bucket/data/sales/",
+ format="parquet",
+ ),
+ ],
+ max_rows=100,
+ )
+
+Which of the two fits depends on the question. "Read me this object" is a hook
+method. "What were last quarter's returns by region" is a query, and expressing
+it through ``list_keys`` and ``read_key`` means the model does the aggregation in
+its context window instead of the engine doing it.
+
+An Iceberg table is registered differently. (Iceberg support needs the
+``apache.iceberg`` extra of ``apache-airflow-providers-common-sql``; without
+it, registration raises ``AirflowOptionalProviderFeatureException``.) There is
+no ``uri`` to read files from; the catalog resolves the table by name, so the
+config carries a ``db_name`` instead, following the same ``DataSourceConfig``
+shape that ``example_analytics.py`` in the ``common.sql`` provider uses:
+
+.. code-block:: python
+
+ toolset = DataFusionToolset(
+ datasource_configs=[
+ DataSourceConfig(
+ conn_id="iceberg_default",
+ table_name="users_data",
+ db_name="demo",
+ format="iceberg",
+ ),
+ ],
+ max_rows=100,
+ )
+
+**Credentials and where it runs.** Each ``DataSourceConfig`` carries its own
+``conn_id``, so object-store access is an Airflow connection. DataFusion is an
+embedded engine: the query runs inside the worker process, not on a remote
+cluster. Its tool calls act as barriers, as they do for the other routes that
+build their own tools; see :ref:`toolset-call-barriers`.
diff --git a/providers/common/ai/docs/toolsets/hook.rst b/providers/common/ai/docs/toolsets/hook.rst
new file mode 100644
index 00000000000..a1a65c98128
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/hook.rst
@@ -0,0 +1,98 @@
+ .. 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.
+
+``HookToolset``
+===============
+
+Generic adapter that exposes selected methods of any Airflow Hook as
+pydantic-ai tools via introspection. Requires an explicit ``allowed_methods``
+list — there is no auto-discovery.
+
+.. code-block:: python
+
+ from airflow.providers.http.hooks.http import HttpHook
+ from airflow.providers.common.ai.toolsets.hook import HookToolset
+
+ http_hook = HttpHook(http_conn_id="my_api")
+
+ toolset = HookToolset(
+ http_hook,
+ allowed_methods=["run"],
+ tool_name_prefix="http_",
+ )
+
+For each listed method, the introspection engine:
+
+1. Builds a JSON Schema from the method signature (``inspect.signature`` +
+ ``get_type_hints``).
+2. Extracts the description from the first paragraph of the docstring.
+3. Enriches parameter descriptions from Sphinx ``:param:`` or Google
+ ``Args:`` blocks.
+
+Parameters
+----------
+
+- ``hook``: An instantiated Airflow Hook.
+- ``allowed_methods``: Method names to expose as tools. Required. Methods
+ are validated with ``hasattr`` + ``callable`` at instantiation time.
+- ``tool_name_prefix``: Optional prefix prepended to each tool name
+ (e.g. ``"s3_"`` produces ``"s3_list_keys"``).
+
+When to choose it
+-----------------
+
+**Choose it when** the target already has an Airflow connection and a hook, and
+what you want the agent to do is already a method on that hook. This is the
+cheapest route — no new server, no new credential, no new query dialect — and
+the only one that reaches any provider hook with synchronous methods without
+anyone writing an adapter first. :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` is a
+reflection-based adapter, so the work is choosing the method list.
+
+**What it cannot do**
+
+- It allow-lists method *names*, not arguments. Once ``read_key`` is exposed,
+ the agent picks the key; the :ref:`defense-layer table <toolset-defense-layers>`
+ states this outright. Choose methods whose worst case you accept, not methods
+ you intend to constrain later.
+- Its calls act as barriers. The tools are registered with ``sequential=True``
+ because hook methods perform synchronous I/O, so a slow call holds up every
+ other tool the model emitted in that step, not only this toolset's. This is
+ not specific to ``HookToolset`` — see :ref:`toolset-call-barriers`.
+- It returns exactly one shape. Every result goes through ``serialize_for_llm``
+ and comes back as a JSON-encoded string; there is no structured error type and
+ no ``ModelRetry`` wrapper, so a hook exception fails the agent run, and the
+ task with it, instead of giving the model something it can correct.
+ ``SQLToolset``, by contrast, hands the database's own error back as a retry.
+- Its ``call_tool`` calls the method and serializes what comes back. The code
+ contains no path that awaits a coroutine result, and none that checks for one,
+ so an ``async def`` hook method is not a case this adapter is written to
+ handle. Treat synchronous methods as the supported set.
+
+**A real example.** The read-only S3 pair from the ``HookToolset`` guidance in
+:doc:`../agent_security`:
+
+.. code-block:: python
+
+ HookToolset(
+ s3_hook,
+ allowed_methods=["list_keys", "read_key"],
+ tool_name_prefix="s3_",
+ )
+
+**Credentials and where it runs.** The hook instance is yours, so the credential
+is whatever connection that hook resolves — the toolset never looks one up
+itself. Calls run in the Airflow worker process.
diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst
new file mode 100644
index 00000000000..e40653608fe
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -0,0 +1,269 @@
+ .. 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.
+
+.. _howto/toolsets:
+
+Toolsets
+========
+
+Choosing a toolset
+------------------
+
+Each toolset's guide documents how to configure it. This section answers the
+question that comes before that one: you have a system you want an agent to
+reach, so which route do you take, and what does each route give up?
+
+Read the table below by what you already have, not by what a toolset is called.
+When two routes both work, the deciding factor is rarely what each one can do —
+it is what each one cannot do, and every route has a short list.
+
+More than one row can be true at once, and the rows are not exclusive: one agent
+can carry several toolsets. Two questions break the ties. *Whose credential is
+it?* — prefer the route whose credential is an Airflow connection somebody on
+your side already reviewed. *Whose tool list is it?* — prefer the route whose
+exposed surface you chose rather than inherited. The pair that most often
+overlaps is an Airflow hook and a vendor MCP server reaching the same target;
+both questions point at the hook, because its credential is the connection and
+``allowed_methods`` is a list you write. Reach for the server when its tools
+cover work the hook does not expose, or when the alternative is re-wrapping that
+API by hand.
+
+Those two questions do not separate ``HookToolset`` from ``SQLToolset`` when the
+target is a DBAPI database, because both answer them the same way. A third one
+does: *is the work a fixed operation or an open-ended question?* A named method
+you can enumerate in advance is a hook. A question the agent has to express as
+SQL is a query, and ``SQLToolset`` answers it with schema discovery, bounded
+results and an ``allowed_tables`` walk you can switch on, none of which
+``HookToolset`` has an equivalent of.
+
+Start with what you have
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. list-table::
+ :widths: 50 50
+ :header-rows: 1
+
+ * - What you have
+ - Route
+ * - A target that already has an Airflow connection, and a hook method that
+ already does the thing
+ - ``HookToolset``
+ * - A question that is a query, against a DBAPI database
+ - ``SQLToolset``
+ * - Files on an object store — Parquet, CSV, Avro — or a catalog-managed
+ table format such as Iceberg, rather than rows in a database
+ - ``DataFusionToolset``
+ * - A vendor that already ships a server built for agents, whose tools you
+ would otherwise re-wrap by hand
+ - ``MCPToolset``
+ * - Procedural knowledge — how to carry out a task — rather than an endpoint
+ to call
+ - ``AgentSkillsToolset``
+ * - Work that means running code the model wrote, not calling a tool you chose
+ - ``SandboxToolset``
+ * - Reasoning that should happen on the vendor's own infrastructure
+ - A subclass of ``BaseManagedAgentToolset`` that you write
+
+The hook, SQL, DataFusion, MCP, Agent Skills and managed-agent guides each have a
+*When to choose it* section giving the case for choosing it, what it cannot do, an
+example that exists in this repository, and where its credentials and its work come
+from. :doc:`../sandbox/index` carries the same section for ``SandboxToolset``.
+
+The toolsets
+------------
+
+Airflow's 350+ provider hooks already have typed methods, rich docstrings,
+and managed credentials. Toolsets expose them as pydantic-ai tools so that
+LLM agents can call them during multi-turn reasoning.
+
+Six toolsets are exported directly from the ``airflow.providers.common.ai.toolsets``
+package root:
+
+- :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic
+ adapter for any Airflow Hook. Guide: :doc:`hook`.
+- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated
+ 4-tool database toolset. Guide: :doc:`sql`.
+- :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` — connect to
+ `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow
+ connections. Guide: :doc:`mcp`.
+- :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` — give
+ the agent a shell and a filesystem inside an isolated sandbox, off the
+ Airflow worker. Guide: :doc:`../sandbox/index`.
+- :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
+ — base class that provider packages subclass to expose a **vendor-managed
+ agent**, one whose reasoning loop runs on a cloud provider's infrastructure.
+ Guide: :doc:`managed_agent`.
+- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
+ — composes several interchangeable managed agents behind a single tool.
+ See :ref:`managed-agent-toolsets`.
+
+Three more toolsets (:doc:`datafusion`, :doc:`logging`, :doc:`skills`) are not
+re-exported from the package root, so import each of them from its own submodule::
+
+ from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
+ from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+ from airflow.providers.common.ai.toolsets.skills import AgentSkillsToolset
+
+All of these toolsets implement pydantic-ai's
+`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be
+passed to any pydantic-ai ``Agent``, including via
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
+
+.. note::
+
+ ``AgentOperator`` accepts **any** ``AbstractToolset`` implementation — not
+ just the Airflow-native toolsets above. PydanticAI's own ``MCPToolset``
+ (built over a FastMCP transport) and third-party toolsets work too. The
+ Airflow-native toolsets add connection management, secret backend
+ integration, and the connection UI, but you are not locked in.
+
+Where the credentials come from
+-------------------------------
+
+Airflow connections and hooks are the access-governance machinery this provider
+already has: a connection lives in a secret backend rather than in Dag code, and
+its scope is set outside the Dag. A hook turns that scope into typed methods.
+Where a route ends up getting its credential is therefore a decision worth
+making on purpose rather than inheriting.
+
+.. list-table::
+ :widths: 28 44 28
+ :header-rows: 1
+
+ * - Route
+ - Credential source
+ - Where the tool call runs
+ * - ``HookToolset``
+ - Whatever connection the hook you supply resolves
+ - Worker process
+ * - ``SQLToolset``
+ - ``db_conn_id``, via ``BaseHook.get_connection``
+ - Worker process, against the database
+ * - ``DataFusionToolset``
+ - ``conn_id`` on each ``DataSourceConfig``
+ - Worker process (embedded engine)
+ * - ``MCPToolset``
+ - ``mcp_conn_id``, unless ``token_provider`` or ``env_provider`` supplies
+ one from elsewhere
+ - Remote server, or a child process on the worker host for ``stdio``
+ * - ``AgentSkillsToolset``
+ - ``GitSkills.conn_id`` for private repositories; none for local directories
+ - Worker process
+ * - ``SandboxToolset``
+ - None from Airflow. Host-level ``sbx login`` and ``sbx policy init``
+ for ``sbx``; ambient Modal token on the worker for Modal
+ - A microVM on the worker host (``sbx``), or Modal's infrastructure off
+ the worker (Modal)
+ * - ``BaseManagedAgentToolset``
+ - Undefined by the base class; the subclass decides
+ - The vendor's infrastructure
+
+Two rows are worth pausing on. ``SandboxToolset`` deliberately takes no Airflow
+credential — that is the whole point of it, and it substitutes a backend-level
+boundary, on the host or at the vendor, for the connection-level one. ``BaseManagedAgentToolset`` does not
+substitute anything; it simply leaves the question to whoever writes the
+subclass. Neither is a defect, but in both cases the access decision has moved
+somewhere Airflow cannot see it, and somebody has to make that decision again in
+the new place. The :ref:`defense-layer table <toolset-defense-layers>` is the
+right companion when you do.
+
+.. _toolset-call-barriers:
+
+Tool calls as barriers
+----------------------
+
+One behaviour cuts across these routes rather than telling them apart.
+``HookToolset``, ``SQLToolset``, ``DataFusionToolset`` and ``SandboxToolset``
+each build their own tool definitions and set ``sequential=True`` on them, which
+pydantic-ai treats as a barrier: the tool runs alone, tools the model emitted
+before it finish first, and tools emitted after it start only once it returns.
+A slow call on any of those four therefore holds up the rest of that step, not
+just its own toolset. ``BaseManagedAgentToolset`` sets ``sequential=False``
+deliberately, because the wait it introduces is remote. ``MCPToolset`` and
+``AgentSkillsToolset`` define no tools of their own — they pass through whatever
+the upstream toolset declares — so the setting is not theirs to make. Do not read
+this as a reason to choose one route over another; read it as something to expect
+from all four.
+
+Layering
+--------
+
+:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` and the
+durable-execution ``CachingToolset`` are not alternatives to anything above. Both
+are wrappers: they take a toolset and return a toolset, adding per-call logging
+or replay from a durable cache. Choose a route first, then decide whether to wrap
+it. The same applies to
+:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`,
+which converts a chosen toolset for a different agent framework rather than
+offering another way to reach a system; :doc:`logging` and :doc:`langchain` cover
+both.
+
+Using toolsets outside ``AgentOperator``
+----------------------------------------
+
+Toolsets are standard pydantic-ai ``AbstractToolset`` implementations with no
+dependency on ``AgentOperator`` or ``@task.agent``. You can use them anywhere
+you can run Python within Airflow -- ``@task`` functions, ``PythonOperator``
+callables, or any custom operator's ``execute()`` method -- by creating a
+``pydantic_ai.Agent`` yourself:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
+ :language: python
+ :start-after: [START howto_task_with_toolsets]
+ :end-before: [END howto_task_with_toolsets]
+
+This works because toolsets resolve Airflow connections lazily via
+``BaseHook.get_connection()``, which is available in any task execution
+context.
+
+This approach gives you full control over the agent lifecycle -- you can call
+``agent.run_sync()`` multiple times, swap models at runtime, or combine
+results from several agents in a single task. The tradeoff is that you lose
+the durable execution (step-level caching with retry replay), HITL review
+integration, and automatic tool call logging that ``AgentOperator`` provides.
+
+Running the agent outside the data system
+-----------------------------------------
+
+A recurring alternative to everything on this page is to push the reasoning into
+the system that holds the data and let it run there. The trade is worth stating
+plainly, because it is the same trade the last row of the table makes.
+
+Running the agent where the data lives keeps the loop short, and it is the right
+answer when the work never leaves that system. What it costs is coupling: the
+agent can only reason over what that system holds, its credentials are governed
+by that system's model rather than Airflow's, and failures are that system's to
+explain. Running the agent in the worker instead means the model calls, the agent
+loop and every tool are orchestrated by Airflow, where they get retries, logging
+and observability like any other task. It also means one agent can hold a
+database, an object store and a vendor API in the same run, with each route's
+credential resolved through a connection that was reviewed once.
+
+This provider does not treat either as the default. The rule of thumb in
+:doc:`../index` is the dividing line: if Airflow should *run* the AI step, and the
+model should stay swappable, use ``common.ai``; if the Dag *submits work to* a
+vendor-managed service and waits for the result, use that vendor's provider —
+and ``BaseManagedAgentToolset`` exists for the case where you want the second
+behaviour from inside an agent that is otherwise doing the first.
+
+See also
+--------
+
+- :doc:`../agent_security` — defense layers, ``allowed_tables`` enforcement, ``HookToolset``
+ guidelines and the production checklist.
+- :doc:`../security` — the provider's security policy.
+- :doc:`../examples` — the example Dags referenced above.
diff --git a/providers/common/ai/docs/toolsets/langchain.rst b/providers/common/ai/docs/toolsets/langchain.rst
new file mode 100644
index 00000000000..e83dcecd109
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/langchain.rst
@@ -0,0 +1,81 @@
+ .. 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.
+
+LangChain bridge
+================
+
+Tools bridge in both directions between common.ai's toolsets and LangChain.
+
+**LangChain tools → ``AgentOperator``.** No Airflow code is needed. pydantic-ai
+ships `pydantic_ai.ext.langchain.LangChainToolset
+<https://ai.pydantic.dev/toolsets/>`__ upstream, which wraps existing LangChain
+tools as an ``AbstractToolset``. Drop it straight into ``AgentOperator``:
+
+.. code-block:: python
+
+ from pydantic_ai.ext.langchain import LangChainToolset
+
+ AgentOperator(
+ task_id="agent_with_langchain_tools",
+ prompt="Research the question and summarise.",
+ llm_conn_id="pydanticai_default",
+ toolsets=[LangChainToolset([my_langchain_tool])],
+ )
+
+**common.ai toolsets → LangChain.** The reverse direction is what
+:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`
+provides. It converts any pydantic-ai toolset -- including ``SQLToolset``,
+``HookToolset``, and ``MCPToolset`` -- into a list of LangChain
+``StructuredTool`` objects, so a LangChain agent or chain can call Airflow's
+curated, connection-managed tools:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_langchain_toolset_bridge.py
+ :language: python
+ :start-after: [START example_langchain_toolset_bridge]
+ :end-before: [END example_langchain_toolset_bridge]
+
+Each generated tool keeps the source tool's name, description, and argument
+schema, and routes calls back through the original toolset, so the toolset's own
+behavior (connection resolution, ``SQLToolset``'s SQL validation, and
+``allowed_tables`` filtering) still applies. ``get_tools`` runs eagerly at
+conversion time to enumerate the tools.
+
+When a toolset raises pydantic-ai's ``ModelRetry`` to ask the model to correct
+its input (``SQLToolset`` does this on, for example, an unknown column), the
+bridge returns that message as the tool's output so the model sees it and tries
+again. ``ModelRetry`` is a feed-the-model-and-retry signal rather than a
+failure, so returning it preserves the self-correction the toolset was written
+for and works no matter how the agent is configured to handle tool errors
+(raising would abort the run under ``create_agent``'s default handling).
+
+The bridge does not hold a toolset session open across calls: ``get_tools`` and
+every tool call each run under their own event loop, so for ``MCPToolset`` the
+connection is opened and torn down around each call. It reconnects per call,
+which is fine for stateless tools but unsuitable for ``stdio`` MCP servers (or
+any server that keeps state between calls), since each call starts a fresh
+session.
+
+.. note::
+
+ Outside an agent run there is no live ``RunContext``, so the bridge builds a
+ minimal one with an inert placeholder model. The bundled toolsets ignore the
+ context, so this is transparent for them. A custom toolset that reads live
+ run state (``ctx.model``, ``ctx.messages``, ``ctx.usage``) will not behave
+ correctly when bridged standalone.
+
+Requires the ``langchain`` extra:
+``pip install "apache-airflow-providers-common-ai[langchain]"``
diff --git a/providers/common/ai/docs/toolsets/logging.rst b/providers/common/ai/docs/toolsets/logging.rst
new file mode 100644
index 00000000000..f91d3e5b0e1
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/logging.rst
@@ -0,0 +1,36 @@
+ .. 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.
+
+``LoggingToolset``
+==================
+
+:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` is a
+``WrapperToolset`` that intercepts ``call_tool()`` to log each tool invocation
+in real time. ``AgentOperator`` applies it automatically (see
+``enable_tool_logging``), but you can also use it directly with any pydantic-ai
+``Agent``:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ sql_toolset = SQLToolset(db_conn_id="my_db")
+ logged_toolset = LoggingToolset(wrapped=sql_toolset, logger=my_logger)
+
+Each tool call produces two INFO log lines (name + timing) and optional
+DEBUG-level argument logging. Exceptions are logged and re-raised.
diff --git a/providers/common/ai/docs/toolsets/managed_agent.rst b/providers/common/ai/docs/toolsets/managed_agent.rst
new file mode 100644
index 00000000000..e2993a8553b
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/managed_agent.rst
@@ -0,0 +1,291 @@
+ .. 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.
+
+.. _managed-agent-toolsets:
+
+Managed agent toolsets
+======================
+
+Cloud vendors now run agents on your behalf — Snowflake Cortex Agents, Amazon
+Bedrock AgentCore runtimes, Azure AI Foundry hosted agents, Vertex AI Agent
+Engine. Their reasoning loops execute on the vendor's infrastructure, so they
+are not something ``AgentOperator`` runs; they are something an Airflow task
+*consults*.
+
+:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
+is the contract for exposing one of those as a tool. Each provider package
+ships its own subclass, so credentials keep flowing through that provider's
+existing hook and no new connection types are needed.
+
+A subclass implements two members:
+
+``agent_ref``
+ Normalized identity of the remote agent — ``platform`` and ``name`` — logged
+ on every call so a run can be audited for which agents it consulted.
+
+``invoke_sync(prompt)``
+ Send the prompt, return the agent's answer. Return the *answer*, not the
+ transport envelope.
+
+Tool naming, argument validation, result serialisation, and logging are handled
+by the base class, so every provider's implementation presents the same surface
+to the calling model.
+
+``tool_name`` is the required identifier — it is what the model emits when it
+calls the tool, and the Dag author chooses it. ``description`` is optional and
+falls back to the tool name rendered as prose, the same way ``HookToolset``
+derives one from a method name when there is no docstring.
+
+.. note::
+
+ Writing a description is still worth the line. It is what tells the model to
+ consult the agent rather than answer from its own knowledge, and it is the
+ only place to record a scope limit the name cannot carry — "cannot see
+ revenue figures". Because the argument schema is always a bare prompt, the
+ name and the description are the whole of what the model knows about the
+ agent.
+
+Sync or async?
+--------------
+
+An agent run is driven by a single event loop, shared by the model requests and
+every toolset's tool calls. A call that blocks that loop stalls all of them for
+as long as the remote agent reasons -- seconds to minutes, not milliseconds --
+and it does so silently, with no error to point at.
+
+Implement ``invoke_sync`` when the vendor call is blocking. The base class runs
+it in a worker thread, keeping the loop free::
+
+ class MyManagedAgentToolset(BaseManagedAgentToolset):
+ @property
+ def agent_ref(self) -> dict[str, str]:
+ return {"platform": "example.cloud", "name": self._agent_id}
+
+ def invoke_sync(self, prompt: str) -> str:
+ return self._hook.conn.invoke(agent=self._agent_id, prompt=prompt)
+
+Override ``invoke(prompt)`` -- the async coroutine -- when the call is already
+asynchronous. Implement one hook or the other; a subclass that implements
+neither is rejected when it is constructed.
+
+.. warning::
+
+ A thread cannot be cancelled. A caller that stops waiting for
+ ``invoke_sync`` does not stop the call -- it keeps a worker thread and its
+ socket until the call returns. Set a timeout on the underlying request so
+ that is bounded.
+
+Toolset or operator?
+--------------------
+
+Most managed-agent platforms do not offer a plain one-request-one-answer API. Some
+require polling a job; others require creating a session and tearing it down around
+each exchange. A toolset can do either, but only by blocking inside
+``invoke()`` — it cannot defer to the Triggerer, and it has no post-task hook to
+clean up with if the worker dies mid-call.
+
+That draws a boundary worth respecting:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 45 55
+
+ * - Shape
+ - Surface to use
+ * - A short consultation *inside* an agent's reasoning, where failing the task
+ would discard the calling agent's accumulated context
+ - A managed agent toolset
+ * - Long-running submitted work as a pipeline step in its own right
+ - That provider's own operator, with deferral or
+ :class:`~airflow.sdk.bases.resumablejobmixin.ResumableJobMixin`
+
+``ResumableJobMixin`` exists for exactly the second case: it persists the external
+job ID to the task state store before polling, so a worker crash reconnects to the
+running job instead of submitting a duplicate. A toolset cannot offer that, because
+the retry boundary is the task, not the tool call — on retry the agent loop restarts
+and re-issues the call. Durable execution covers the *completed* call (see
+``replayable`` below); it does not cover a call that was still in flight.
+
+Error handling
+--------------
+
+Failures sort into three buckets, and conflating them is the most common way an
+implementation goes wrong:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 22 30 48
+
+ * - Raise
+ - When
+ - Who recovers
+ * - ``ModelRetry``
+ - The agent rejected the request in a way rephrasing could fix.
+ - The calling model, bounded by its ``usage_limits``.
+ * - ``ManagedAgentInvocationError``
+ - Terminal: bad credentials, missing agent, revoked quota.
+ - Nobody — the task fails fast instead of burning retries.
+ * - *let it propagate*
+ - Transient: 429, 5xx, connection reset, read timeout.
+ - Airflow's task-level retry. A rephrase does nothing for a 503.
+
+Durable execution
+-----------------
+
+``replayable`` is ``False`` by default. A managed agent may act on systems
+Airflow cannot observe, so replaying a cached answer on retry could skip a side
+effect. Implementations whose agent is read-only should set it to ``True`` to
+avoid paying for the same invocation twice.
+
+Deferral
+--------
+
+A toolset call runs in the worker and cannot defer to the Triggerer — it blocks
+for the duration of the call. See `Toolset or operator?`_ above for when that is
+acceptable and when the provider's own deferrable operator is the right surface
+instead.
+
+Failover between interchangeable agents
+---------------------------------------
+
+:class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
+composes several managed agents into one tool, trying them in order until one
+answers. It is itself a ``BaseManagedAgentToolset``, so the calling model sees a
+single tool and has no say in which provider serves the request — the policy
+stays deterministic Python rather than a prompt instruction a model may ignore.
+Groups nest.
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets import FailoverManagedAgentToolset
+
+ resilient = FailoverManagedAgentToolset(
+ tool_name="ask_claims_agent",
+ description="Reviews an insurance claim and returns a coverage determination.",
+ members=[bedrock_claims_agent, foundry_claims_agent], # same image, two clouds
+ )
+
+Members must satisfy two preconditions the class cannot check.
+
+**Substitutability.** The same agent deployed twice, not two specialists with
+different data. Two containerised agents built from one image qualify; agents
+bound to one platform's own objects — a Cortex Agent over Snowflake semantic
+models — do not, because there is nothing equivalent to fail over *to*.
+
+**Statelessness per invocation.** Server-side conversation state is the norm
+across managed-agent platforms, not the exception — optional on some (Cortex
+``thread_id``), mandatory on others where a session is created and torn down
+around each exchange. Each member is invoked with a bare prompt and no thread
+reference, so a failover silently starts a fresh conversation on the standby:
+correct for a one-shot consultation, wrong for a multi-turn one. Treat one-shot
+as a restriction a group is deliberately held to, not a safe default.
+
+The three error buckets do real work here:
+
+- ``ManagedAgentInvocationError`` and transient failures move to the next member.
+- ``ModelRetry`` is re-raised immediately and never triggers failover. A prompt
+ the primary could not parse will not parse on the standby either, so failing
+ over would spend the standby's budget reproducing the same error.
+- The last member's exception propagates unchanged, so a total outage still fails
+ the task rather than returning something misleading.
+
+``failover_on`` defaults to ``Exception`` because ``common.ai`` cannot enumerate
+the cloud SDKs' exception trees — ``requests``, ``botocore`` and the Azure SDK
+share no common base. It can be narrowed when the members' exception types are
+known.
+
+``replayable`` on a group is ``True`` only when every member is, because the
+durable cache cannot know which member produced the answer it holds.
+
+.. note::
+
+ For a **standalone** agent call, prefer plain Airflow task-level failover:
+ two tasks, the second with ``trigger_rule=TriggerRule.ALL_FAILED``. That keeps
+ which provider served the request visible in the grid at no code cost, and
+ makes failover rate a task metric. This class is for the case a task boundary
+ cannot express — a managed agent consulted as a tool *inside* a longer agent
+ run, where failing the task would discard the calling agent's accumulated
+ context and re-run every earlier tool call.
+
+Two counters make failover visible, because a failover is a *success-shaped*
+event — without them a primary that has been down for a week looks identical to a
+healthy one:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 70
+
+ * - Metric
+ - Tags
+ * - ``managed_agent.failover``
+ - ``from_platform``, ``to_platform`` — one per failover transition
+ * - ``managed_agent.served``
+ - ``platform``, ``role`` (``primary`` / ``standby``) — one per answer
+
+The standby-served fraction is a ratio over ``managed_agent.served`` alone, so
+"are we quietly running on the standby?" is a dashboard question rather than a log
+grep. Both are tagged by platform rather than agent name to keep cardinality
+bounded.
+
+One limitation remains: which member served a *particular* answer is in the task
+log but not in XCom. ``agent_ref`` on a group describes the group, not the
+responder, because the responder is not known until after the call. The counters
+cover the operational question; per-answer provenance for an audit trail would
+need ``AgentOperator`` to collect per-toolset metadata.
+
+When to choose it
+-----------------
+
+**Choose it when** the reasoning itself belongs on the vendor's infrastructure —
+the agent is already deployed there, grounded in data that never leaves, and
+Airflow's job is to submit one request and read one answer.
+:ref:`managed-agent-toolsets` covers the shape.
+
+Read the first bullet before planning around this route.
+
+**What it cannot do**
+
+- It is an extension point, not a toolset you can use as-is. This provider ships
+ the base class only:
+ :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
+ declares ``agent_ref`` abstract and raises ``TypeError`` at construction unless
+ a subclass implements ``invoke_sync`` or ``invoke``. No vendor subclass exists
+ in this repository — the Snowflake Cortex, Bedrock AgentCore, Azure AI Foundry
+ and Vertex AI Agent Engine names in its docstring describe the shape it expects,
+ not implementations that ship. Using this route means writing that subclass.
+- It has no allow-list to offer. The whole toolset is one tool: a prompt goes in,
+ an answer comes out. Whatever governs what the remote agent may touch lives on
+ the vendor's side, which is the trade you are making.
+- It does not define where the credential comes from. The base class leaves
+ that decision to the subclass; :ref:`managed-agent-toolsets` frames the
+ intended shape as routing authentication through the provider's own hook.
+- Durable replay is off by default. ``replayable`` is ``False`` because a managed
+ agent may act on systems Airflow cannot observe, so replaying from the cache
+ could skip a side effect. Read-only agents can opt in.
+- Its failover variant relies on two preconditions the code cannot check.
+ :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
+ requires members that are genuinely interchangeable — the same agent deployed
+ twice, not two specialists over different data — and one-shot exchanges, because
+ each member is invoked with a bare prompt and no thread reference, so a failover
+ silently starts a fresh conversation rather than resuming the old one.
+
+**A real example.** There is none. No example Dag, no system test, and the only
+code in the documentation is an illustrative subclass on this page. Budget
+for writing and testing the subclass yourself.
+
+**Credentials and where it runs.** Both are the subclass's decision. Reasoning
+runs on the vendor's infrastructure; the worker sends a request and waits.
diff --git a/providers/common/ai/docs/toolsets/mcp.rst b/providers/common/ai/docs/toolsets/mcp.rst
new file mode 100644
index 00000000000..2c7f378cc5a
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/mcp.rst
@@ -0,0 +1,166 @@
+ .. 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.
+
+``MCPToolset``
+==============
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+
+ MCP connection <../connections/mcp>
+ MCPHook <../hooks/mcp>
+
+Connects to an `MCP (Model Context Protocol) <https://modelcontextprotocol.io/>`__
+server configured via an Airflow connection. MCP is an open protocol that lets
+LLMs interact with external tools and data sources through a standardized
+interface.
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets.mcp import MCPToolset
+
+ toolset = MCPToolset(
+ mcp_conn_id="my_mcp_server",
+ tool_prefix="weather",
+ )
+
+The MCP server is resolved lazily from the Airflow connection on the first
+tool call. See :ref:`howto/connection:mcp` for connection configuration.
+
+Requires the ``mcp`` extra: ``pip install "apache-airflow-providers-common-ai[mcp]"``
+
+Parameters
+----------
+
+- ``mcp_conn_id``: Airflow connection ID for the MCP server.
+- ``tool_prefix``: Optional prefix prepended to tool names to avoid
+ collisions when using multiple MCP servers (e.g. ``"weather"`` produces
+ ``"weather_get_forecast"``).
+- ``token_provider``: Optional zero-argument callable returning a bearer token.
+ When set, it overrides the connection's static ``password`` for the
+ ``Authorization`` header. Called once, the first time this toolset
+ establishes a connection -- use it for short-lived or minted tokens (e.g. a
+ Snowflake managed MCP server authenticated with a key-pair JWT). See
+ :ref:`howto/connection:mcp`.
+- ``env_provider``: Optional zero-argument callable returning a
+ ``dict[str, str]`` merged over the connection's ``Extra.env`` (winning on key
+ conflicts) for the ``stdio`` subprocess environment -- use it when the
+ credential a local stdio MCP server needs lives in a different connection, or
+ is minted fresh per call (e.g. a Splunk/Vault token), rather than storing it
+ statically here. Called once, the first time this toolset establishes a
+ connection. See :ref:`howto/connection:mcp`.
+
+Using multiple MCP servers
+--------------------------
+
+.. code-block:: python
+
+ AgentOperator(
+ task_id="multi_mcp",
+ prompt="Get the weather in London and run a calculation",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ MCPToolset(mcp_conn_id="weather_mcp", tool_prefix="weather"),
+ MCPToolset(mcp_conn_id="code_runner_mcp", tool_prefix="code"),
+ ],
+ )
+
+Direct pydantic-ai MCP toolsets
+-------------------------------
+
+For prototyping or when you want full PydanticAI control, you can pass
+``MCPToolset`` instances directly — no Airflow connection needed:
+
+.. code-block:: python
+
+ from fastmcp.client.transports import StdioTransport
+ from pydantic_ai.mcp import MCPToolset
+
+ AgentOperator(
+ task_id="direct_mcp",
+ prompt="What tools are available?",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ MCPToolset("http://localhost:3001/mcp"),
+ MCPToolset(StdioTransport(command="uvx", args=["mcp-run-python"])),
+ ],
+ )
+
+This works because PydanticAI's ``MCPToolset`` implements ``AbstractToolset``.
+The tradeoff: URLs and credentials are hardcoded in Dag code instead of being
+managed through Airflow connections and secret backends.
+
+When to choose it
+-----------------
+
+**Choose it when** someone already publishes a server built for agents that
+covers your target. You inherit a tool surface that was designed to be called by
+a model — retry semantics and error wording are decided upstream, and a
+destructive tool can simply be absent — instead of maintaining a per-API wrapper
+yourself.
+
+**What it cannot do**
+
+- Its allow-list is client-side and opt-in, not server-side and required.
+ :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` forwards
+ ``get_tools`` and ``call_tool`` straight to the underlying server, so
+ whatever the server exposes, the agent gets by default. Because
+ ``MCPToolset`` is itself built on the same ``AbstractToolset`` base every
+ toolset in this provider extends, a Dag author can call ``.filtered()`` to
+ subset the advertised tool list using a filter function that inspects each
+ tool's definition. That filtering happens on the client: it narrows what
+ the agent is offered, it does not revoke or authorize anything on the
+ server, and unlike ``allowed_methods`` on ``HookToolset``, which is required
+ and rejects an empty list, nothing here requires you to set a filter. The
+ defense-layer table is explicit that a server can expose shell, filesystem
+ or network access.
+- It cannot guarantee the credential came from a connection. ``mcp_conn_id`` is
+ the default path, but ``token_provider`` and ``env_provider`` are your own
+ callables and are free to read an environment variable, a file, or an entirely
+ different secret store. That is the point of them — it also means the
+ connection is no longer the whole story for anyone auditing the Dag.
+- ``stdio`` transport is not isolation. It starts a child process on the worker
+ host. It is not a sandbox, and when no environment is supplied the child
+ inherits a small allowlist of variables rather than the full parent
+ environment — so it is both less contained and less predictable than it looks.
+
+**A real example.** ``example_mcp.py`` drives setup from a connection:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
+ :language: python
+ :start-after: [START howto_toolset_mcp_connection]
+ :end-before: [END howto_toolset_mcp_connection]
+
+and puts several servers on one agent, prefixed so their tool names stay apart:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_mcp.py
+ :language: python
+ :start-after: [START howto_toolset_mcp_multiple]
+ :end-before: [END howto_toolset_mcp_multiple]
+
+No MCP server for object storage ships with this provider. If one exists for
+your target, it is a third route alongside the hook and DataFusion routes (:doc:`hook`,
+:doc:`datafusion`) — and the two questions in :ref:`Choosing a toolset <howto/toolsets>` settle
+it. Its tool list is the server's rather than
+yours, and its token comes from wherever ``mcp_conn_id`` or your own callable
+says. So where a hook already reaches the same target, the hook wins; the server
+wins where it covers work the hook does not expose.
+
+**Credentials and where it runs.** ``mcp_conn_id`` supplies host, credentials and
+transport, unless a provider callable overrides that. ``http`` and ``sse`` reach
+a remote server; ``stdio`` runs a child process on the worker host.
diff --git a/providers/common/ai/docs/toolsets/skills.rst b/providers/common/ai/docs/toolsets/skills.rst
new file mode 100644
index 00000000000..1ce37c1391b
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/skills.rst
@@ -0,0 +1,169 @@
+ .. 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.
+
+.. _agent-skills:
+
+``AgentSkillsToolset``
+======================
+
+:class:`~airflow.providers.common.ai.toolsets.skills.AgentSkillsToolset` loads
+`Agent Skills <https://agentskills.io>`__ -- ``SKILL.md`` bundles (instructions,
+and optionally scripts and resources) that the model discovers and loads *on
+demand*. Only a compact catalog of skill names and descriptions sits in the
+prompt until the model decides it needs one, so a large skill library costs few
+tokens until used (progressive disclosure).
+
+It is backed by the community `pydantic-ai-skills
+<https://github.com/DougTrajano/pydantic-ai-skills>`__ package (MIT); native
+progressive disclosure is in flight upstream in `pydantic/pydantic-ai#5230
+<https://github.com/pydantic/pydantic-ai/pull/5230>`__. Install the optional
+extra to use it:
+
+.. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[skills]"
+
+Each source is a local directory or a connection-resolved
+:class:`~airflow.providers.common.ai.skills.GitSkills`. Sources are resolved when
+the agent enters the toolset, on the worker -- never while the Dag processor
+parses the file -- so a Git token is never baked into the serialized Dag, and
+cloned repositories are removed when the run ends.
+
+A local directory of ``SKILL.md`` bundles:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
+ :language: python
+ :start-after: [START howto_operator_agent_skills_local]
+ :end-before: [END howto_operator_agent_skills_local]
+
+A Git repository, with credentials from an Airflow connection:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
+ :language: python
+ :start-after: [START howto_operator_agent_skills_git]
+ :end-before: [END howto_operator_agent_skills_git]
+
+For a private repository, point ``conn_id`` at a
+:doc:`git connection <apache-airflow-providers-git:connections/git>`; credentials
+are resolved through the Git provider's ``GitHook`` (an HTTPS token in the
+connection password, or an SSH key in the connection's extra). A plain ``http://``
+URL with ``conn_id`` is rejected so a credential is never sent in cleartext, and a
+``repo_url`` that embeds a username/password is rejected (use ``conn_id``). After
+cloning, the credential is stripped from the checkout's ``.git/config``. As with
+any ``git clone``, the worker's own git configuration (credential helpers, SSH
+agent) may still apply, so run workers without ambient git credentials if you
+need strict isolation.
+
+.. warning::
+
+ Skill bundles can contain scripts that the agent may run on the worker via
+ the ``run_skill_script`` tool. For a remote source, anyone who can modify the
+ repository can introduce code that executes on your worker, outside Dag
+ review and versioning. Point ``GitSkills`` at a trusted repository, pin
+ ``branch`` to a trusted ref, and treat skill contents as code that runs in
+ your environment.
+
+Parameters
+----------
+
+- ``sources``: List of skill sources -- local directory paths and/or
+ :class:`~airflow.providers.common.ai.skills.GitSkills`.
+- ``exclude_tools``: Optional set of skill tool names to hide from the agent
+ (e.g. ``{"run_skill_script"}`` to disable on-worker script execution).
+- ``exclude_resources``: Optional glob patterns to exclude from resource
+ discovery, added on top of the built-in defaults (``__pycache__``, ``*.pyc``,
+ ``*.pyo``, ``.DS_Store``, ``.git``). A skill exposes every readable text file
+ it contains as a resource; these patterns keep matched files out of the
+ resource list and the ``read_skill_resource`` tool (e.g.
+ ``["*.env", "secrets/*"]``). Each pattern matches the full skill-relative path
+ or any single path component. This hides files from resource discovery only --
+ it does not stop a skill's ``run_skill_script`` from reading them off disk, so
+ pair it with ``exclude_tools={"run_skill_script"}`` when the files are
+ genuinely sensitive.
+
+Using Agent Skills with other frameworks
+----------------------------------------
+
+``AgentSkillsToolset`` is a standard pydantic-ai toolset, so it also works with a
+plain ``pydantic_ai.Agent`` you build yourself, not just ``AgentOperator``.
+
+Because Agent Skills is a cross-framework format, the connection handling is also
+reusable through :func:`~airflow.providers.common.ai.skills.resolve_skills`, which
+resolves sources to local ``SKILL.md`` directories that any loader accepts:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.skills import GitSkills, resolve_skills
+
+ sources = ["./skills", GitSkills(repo_url="https://github.com/org/skills", conn_id="github_skills")]
+ with resolve_skills(sources) as dirs:
+ # LangChain DeepAgents
+ agent = create_deep_agent(model="openai:gpt-5.4", skills=dirs)
+ # ...or Strands
+ agent = Agent(plugins=[AgentSkills(skills=dirs)])
+
+``resolve_skills`` needs the Git provider (for ``GitSkills``) but not pydantic-ai,
+and removes any cloned directories when the ``with`` block exits.
+
+When to choose it
+-----------------
+
+**Choose it when** what the agent is missing is procedural knowledge rather than
+an endpoint — how this team writes a report, which checks run before a release,
+what the house conventions are. A skill is a directory of instructions and
+optional scripts, and
+:class:`~airflow.providers.common.ai.toolsets.skills.AgentSkillsToolset` makes it
+discoverable. See :ref:`agent-skills` for the layout.
+
+**What it cannot do**
+
+- ``exclude_resources`` does not hide a file from the skill's own scripts. It
+ keeps matches out of resource discovery and out of ``read_skill_resource``,
+ and the parameter's documentation says plainly that it does not stop
+ ``run_skill_script`` from reading them off disk. For genuinely sensitive files,
+ pair it with ``exclude_tools={"run_skill_script"}``. (The parameter needs
+ ``pydantic-ai-skills>=1.2.0``, which the ``skills`` extra already pins.)
+- It does not move script execution anywhere safer. The toolset's own wording for
+ ``exclude_tools`` calls ``run_skill_script`` "on-worker script execution", and
+ nothing in this toolset routes those scripts into a sandbox — so unless you
+ exclude the tool, a skill's scripts run in the worker process with the worker's
+ reach. If that is not acceptable, exclude the tool or put the work behind
+ ``SandboxToolset`` instead.
+- It re-fetches a Git source on every run. ``GitSkills`` is resolved and
+ shallow-cloned on the worker when the run starts, and the checkout is deleted
+ when it ends; nothing is kept between runs, so a large or slow repository pays
+ that clone once per run. A local directory is read in place and
+ costs nothing.
+
+**A real example.** ``example_agent_skills.py`` loads skills from a local
+directory:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
+ :language: python
+ :start-after: [START howto_operator_agent_skills_local]
+ :end-before: [END howto_operator_agent_skills_local]
+
+The two skills it ships, ``aip-tracker`` and ``sql-reporting``, are procedural by
+nature: neither adds an endpoint the agent could not already reach. That is the
+signal you are on the right route.
+
+**Credentials and where it runs.** A local directory needs no credential. A
+private repository goes through
+:class:`~airflow.providers.common.ai.skills.GitSkills` and its ``conn_id``,
+resolved by the Git provider's ``GitHook``; plain ``http://`` is refused when a
+``conn_id`` is set, so a credential is never sent in the clear. Cloning, reading
+and any script execution happen on the worker.
diff --git a/providers/common/ai/docs/toolsets/sql.rst b/providers/common/ai/docs/toolsets/sql.rst
new file mode 100644
index 00000000000..eacaedbd7d5
--- /dev/null
+++ b/providers/common/ai/docs/toolsets/sql.rst
@@ -0,0 +1,225 @@
+ .. 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.
+
+``SQLToolset``
+==============
+
+Curated toolset wrapping
+:class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` with four tools:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 50
+
+ * - Tool
+ - Description
+ * - ``list_tables``
+ - Lists available table names (filtered by ``allowed_tables`` if set)
+ * - ``get_schema``
+ - Returns column names and types for a table
+ * - ``query``
+ - Executes a SQL query and returns bounded, columnar JSON (see
+ :ref:`bounded-query-results`)
+ * - ``check_query``
+ - Validates SQL syntax without executing it
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ toolset = SQLToolset(
+ db_conn_id="postgres_default",
+ allowed_tables=["customers", "orders"],
+ max_rows=20,
+ )
+
+The ``DbApiHook`` is resolved lazily from ``db_conn_id`` on first tool call
+via ``BaseHook.get_connection(conn_id).get_hook()``.
+
+In read-only mode (``allow_writes=False``, the default) the ``query`` tool also
+accepts read-only metadata statements -- ``DESCRIBE``/``DESC`` and ``SHOW`` --
+in addition to SELECT-family queries. Agents commonly open with ``DESCRIBE`` to
+learn a table's columns, so permitting it keeps runs deterministic instead of
+hard-failing on schema discovery. The toolset passes the connection's dialect to
+the validator, so ``SHOW`` is recognized on databases that support it (Snowflake,
+MySQL, etc.); on databases without ``SHOW`` it stays rejected. Data-modifying
+statements remain blocked -- including ones hidden behind ``DESCRIBE``/``EXPLAIN``
+(e.g. ``EXPLAIN DELETE ...``, ``DESCRIBE DROP TABLE ...``), which the validator
+rejects by scanning the parsed statement for write operations. When
+``allowed_tables`` is set it scopes these statements too: a ``DESCRIBE`` names a
+table, so its target must be on the list, while ``SHOW`` enumerates objects beyond
+any single table and is rejected outright (see :ref:`allowed-tables-enforcement`).
+
+Multi-schema warehouses
+-------------------------
+
+When an agent's tables live in several schemas of one database -- common on
+Snowflake -- list them with schema-qualified ``allowed_tables`` entries:
+
+.. code-block:: python
+
+ SQLToolset(
+ db_conn_id="snowflake_hq",
+ allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"],
+ )
+
+``list_tables`` then introspects each referenced schema and returns the matching
+tables fully qualified (e.g. ``MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS``), and
+``get_schema`` routes each qualified name to its own schema. Without this, a
+single ``schema`` only covers one namespace, and leaving ``schema`` unset made
+introspection query a literal ``"None"`` schema and fail. Unqualified entries
+fall back to ``schema``, and table-name matching is case-insensitive (databases
+reflect identifiers in their own case). For tables in a different *database*, use
+a separate toolset whose connection points at that database.
+
+Parameters
+----------
+
+- ``db_conn_id``: Airflow connection ID for the database.
+- ``allowed_tables``: Restrict the agent to a fixed set of tables. Omit the
+ argument (the default) to expose all tables in ``schema``. No value means
+ allow-all: ``None`` and an empty list both raise ``ValueError``, so an allow-list
+ built at runtime that resolves to nothing fails at import instead of silently
+ exposing every table. Entries may be schema-qualified
+ (``"SCHEMA.TABLE"``) to span multiple schemas; see above. Matching is
+ case-insensitive. When set, the list is enforced on ``query`` and
+ ``check_query`` as well as discovery -- every table a query references must be
+ on it. See :ref:`allowed-tables-enforcement` for what this does and does not
+ guarantee.
+- ``allowed_functions``: Names of functions that sqlglot does not recognize as
+ builtins but are safe to run while ``allowed_tables`` is active (e.g.
+ ``["json_build_object"]`` or a project UDF). ``None`` (default) rejects every
+ unrecognized function. Matching is case-insensitive. Only consulted when
+ ``allowed_tables`` is set.
+- ``schema``: Default schema/namespace for unqualified table listing and
+ introspection. Schema-qualified ``allowed_tables`` entries override it per table.
+- ``allow_writes``: Allow data-modifying SQL (INSERT, UPDATE, DELETE, etc.).
+ Default ``False`` -- only SELECT-family and read-only metadata
+ (``DESCRIBE``/``SHOW``) statements are permitted.
+- ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
+ Rows beyond it are not read out of a DBAPI cursor; what the driver has already
+ transferred is its own call. See :ref:`bounded-query-results`.
+- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
+ See :ref:`bounded-query-results`.
+
+.. _bounded-query-results:
+
+Bounded query results
+---------------------
+
+A tool result stays in the model's message history for the rest of the run, so its
+cost is re-paid on every subsequent model request. The ``query`` tool of both
+``SQLToolset`` and ``DataFusionToolset`` bounds that in three ways.
+
+**The result is columnar.** Column names appear once, not once per row:
+
+.. code-block:: json
+
+ {"columns": ["id", "name"], "rows": [[1, "Alice"], [2, "Bob"]], "row_count": 2}
+
+On a table with thousands of columns the repeated names, not the values, are the bulk
+of a row-of-dicts payload. Positional rows also keep columns that share a name --
+``SELECT o.id, c.id`` -- which a dict per row silently collapsed to one.
+
+**Rows are fetched, not filtered.** ``max_rows`` bounds what leaves the cursor, so a
+query matching a whole table costs the worker roughly what one matching ``max_rows``
+costs. How much is saved depends on the driver: with a server-side cursor the
+remaining rows are never sent, while a client-buffering driver (psycopg2's default
+cursor, MySQLdb) has already received them and only the per-row conversion is skipped.
+Hooks whose cursor is not DBAPI 2.0 (``ExasolHook`` passes a pyexasol statement) fall
+back to a full fetch, and ``DataFusionToolset`` materializes the full result in the
+engine before the toolset sees it; in both the payload is bounded but the transfer is
+not.
+
+**A byte budget bounds the payload.** ``max_rows`` caps rows, which says nothing about
+size -- one row of a 3000-column table is larger than a thousand rows of a narrow one.
+``max_result_bytes`` is what actually bounds context. Rows are returned as a contiguous
+prefix: the result stops at the first row that does not fit the remaining budget rather
+than skipping it and packing later ones, so a single wide row early in the result ends
+it. The result says which limit it hit:
+
+.. code-block:: json
+
+ {"columns": ["..."], "rows": ["..."], "row_count": 3,
+ "truncated": true, "truncated_by": "max_result_bytes"}
+
+``truncated_by`` is ``max_rows`` or ``max_result_bytes``. When not even one row fits,
+or the column names alone exceed the budget, the result carries a ``hint`` telling the
+agent to narrow its projection -- the only move that helps. ``total_rows`` is present
+when the driver reports a row count for the query; several (SQLite, some warehouse
+drivers) do not, and it is then omitted rather than guessed.
+
+The default budget is deliberately generous: the columnar shape alone shrinks a wide
+result several-fold, so results that fit before still fit. Lower ``max_result_bytes``
+when an agent makes many queries in one run, since every result is re-paid on every
+later request.
+
+When to choose it
+-----------------
+
+**Choose it when** the question is a query and the data is in a DBAPI database.
+:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` gives the agent
+four tools — list tables, get schema, query, check query. Set ``allowed_tables``
+and that allow-list is enforced by parsing the SQL rather than by matching
+strings; see :ref:`allowed-tables-enforcement` for how the walk handles CTEs,
+subqueries and joins.
+
+**What it cannot do**
+
+- ``allowed_tables`` is an application-level guardrail, not a replacement for
+ database permissions. Its own docstring says so, and names the residual gap:
+ an engine or query the parser reads differently. Point ``db_conn_id`` at a
+ least-privilege role whose grants match the allow-list.
+- It cannot bound the fetch for every driver. Hooks that hand their handler
+ something other than a DBAPI cursor — ``ExasolHook`` and its pyexasol
+ statement, for instance — fall back to a full fetch. The payload handed to the
+ model is still bounded; the transfer is not. See :ref:`bounded-query-results`.
+- Its parser-level closure is opt-in, not the default. The table walk returns
+ immediately while ``allowed_tables`` is unset, so out of the box the agent
+ reaches every table the connection can see. Only omitting the argument grants
+ that: an explicit ``None`` or empty list is rejected at construction.
+ ``DESCRIBE`` and ``SHOW`` both pass, on dialects that parse them, while
+ ``allowed_tables`` stays unset. Set ``allowed_tables`` and the walk
+ turns fail-closed for ``SHOW`` — but not for ``DESCRIBE``: it instead
+ becomes an ordinary table reference, allowed only when the table it names
+ is on the list. See :ref:`allowed-tables-enforcement` for what else the
+ walk rejects once ``allowed_tables`` is active. Statements that modify
+ data, ``COPY`` among them, are rejected either way while ``allow_writes``
+ is ``False``.
+- It does not classify failures. A connection error or a typo in a column
+ name reaching ``list_tables``, ``get_schema`` or ``query`` becomes one
+ ``ModelRetry``, so the two are treated the same way until the retry budget
+ runs out and the task fails for Airflow to retry. Two paths do not raise:
+ ``check_query`` catches its own errors and reports them back as a normal
+ ``{"valid": false, ...}`` result, and ``get_schema`` returns a normal
+ ``{"error": ...}`` result instead of raising when the requested table is
+ outside ``allowed_tables`` — other ``get_schema`` failures still raise and
+ still become a ``ModelRetry``.
+
+**A real example.** ``example_pydantic_ai_hook.py`` builds an agent around
+``SQLToolset`` inside a plain ``@task`` function, with no operator involved:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
+ :language: python
+ :start-after: [START howto_task_with_toolsets]
+ :end-before: [END howto_task_with_toolsets]
+
+**Credentials and where it runs.** ``db_conn_id`` is resolved through
+``BaseHook.get_connection``, and the connection must supply a ``DbApiHook``.
+Queries run from the Airflow worker process against the database. Its tool
+calls act as barriers, as they do for the other routes that build their own
+tools; see :ref:`toolset-call-barriers`.
diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst
new file mode 100644
index 00000000000..763fbbb54bc
--- /dev/null
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -0,0 +1,143 @@
+ .. 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.
+
+.. _howto/troubleshooting:
+
+Troubleshooting
+===============
+
+The errors below are the ones a first Dag most often hits. Each names the message you
+see, the cause, and the fix.
+
+Model and connection errors
+---------------------------
+
+``No model specified for connection '...'``
+ The connection has no ``model`` in its extra and the operator did not pass
+ ``model_id``. Set the **Model** field on the connection (``provider:model`` form) or
+ pass ``model_id`` to the operator. See :doc:`connections/pydantic_ai`.
+
+``Connection '...' has no default model provider, so the bare model name '...' cannot be resolved``
+ The generic ``pydanticai`` connection type needs a ``provider:`` prefix on the model
+ name, for example ``openai:gpt-5`` rather than ``gpt-5``. Add the prefix, or use a
+ vendor connection type (:doc:`connections/pydantic_ai_azure`,
+ :doc:`connections/pydantic_ai_bedrock`, :doc:`connections/pydantic_ai_vertex`), which
+ supply their own platform.
+
+``'...' is not a provider pydantic-ai recognizes``
+ The text before the first ``:`` in the model name is not a pydantic-ai provider.
+ Check it for a typo. If the vendor's own model id contains a ``:`` (Bedrock-style
+ version suffixes, for example), use the matching vendor connection type instead of
+ the generic one.
+
+An ``ImportError`` for ``pydantic_ai.models.<vendor>`` or the vendor SDK
+ The provider is installed without the extra for that vendor. Install it, quoting the
+ package name so the brackets survive the shell:
+
+ .. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[openai]"
+
+ :doc:`installation` lists the extras.
+
+``A fallback chain is configured for '...' but no model is set``
+ ``fallback_conn_ids`` on the connection needs an explicit primary model. Set the
+ **Model** field on the primary connection or ``model_id`` on the operator. A model
+ taken from an agent spec file cannot be wrapped in a fallback chain. See
+ :doc:`provider_fallback`.
+
+``Fallback connection '...' resolves to ..., which is not a PydanticAIHook``
+ Every entry in ``fallback_conn_ids`` must be one of the ``pydanticai`` connection
+ types. Chains are also not resolved recursively, so a fallback connection may not
+ declare its own ``fallback_conn_ids``; list every vendor directly on the primary.
+
+Operator construction errors
+----------------------------
+
+These raise while the Dag file is parsed, so the whole file fails to import. A mapped
+task (``.expand()``) is constructed at run time instead, so there the same error surfaces
+as a task failure.
+
+``require_approval=True needs Airflow 3.1+`` / ``DecisionPolicy(on_uncertain='review') needs Airflow 3.1+`` / ``approval_assigned_users needs Airflow 3.1+`` / ``Human in the loop functionality needs Airflow 3.1+``
+ Human-in-the-loop review, whether through ``require_approval``,
+ ``DecisionPolicy(on_uncertain="review")`` or ``enable_hitl_review``, needs Airflow 3.1
+ or later. Upgrade the core, or use ``on_uncertain="fail"`` and drop the review flags
+ on an older core. See :doc:`approval_gates` and :doc:`hitl_review`.
+
+``durable=True and enable_hitl_review=True cannot be used together`` / ``durable=True and code_mode=True cannot be used together``
+ Durable replay assumes a stable step order across attempts, which neither a human
+ review loop nor code mode provides. Pick one. See :doc:`durable_execution`.
+
+``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`.
+
+``code_mode=True requires the 'code-mode' extra``
+ Install ``apache-airflow-providers-common-ai[code-mode]``. See :doc:`code_mode`.
+
+``... does not support decision_policy yet``
+ Only ``LLMOperator`` and ``LLMBranchOperator`` honor a ``DecisionPolicy`` with a
+ confidence bar. The SQL, schema-compare and file-analysis operators run their own
+ ``execute`` and reject one at construction; use ``require_approval=True`` there for an
+ unconditional review.
+
+``on_approval_timeout=... needs a review path ... and a positive approval_timeout to fire``
+ ``on_approval_timeout`` other than ``"fail"`` only makes sense when a review can open
+ (``require_approval=True`` or a reviewing ``decision_policy``) **and**
+ ``approval_timeout`` is set to a positive ``timedelta``. Set both or leave the default.
+
+Run-time errors
+---------------
+
+``Agent model must be set when durable=True``
+ The agent was built without a model, usually because the connection has no
+ ``model`` and no ``model_id`` was passed. Durable execution needs the model resolved
+ up front so that replayed steps can be matched against their fingerprints. Fix the connection as described
+ above.
+
+``durable=True`` on Airflow below 3.3 fails with a ``ValueError`` about ``durable_cache_path``
+ On cores older than 3.3 the step cache lives in object storage and
+ ``[common.ai] durable_cache_path`` must be set. On 3.3 and later the task state store
+ is used and the option is ignored. See :doc:`durable_execution`.
+
+A structured ``output_type`` arrives downstream as a string or fails to deserialize
+ The Pydantic class must be defined at module scope under its own ``__name__`` so the
+ worker can register it for XCom deserialization, and a consumer in a *different* Dag
+ needs the class added to ``[core] allowed_deserialization_classes``. Pass
+ ``serialize_output=True`` to receive a plain ``dict`` instead. See
+ :doc:`structured_output`.
+
+A review task waits for a long time
+ That is expected: the task is waiting for a reviewer. An approval gate on an LLM operator
+ releases its worker slot while it waits (it pauses as awaiting input on Airflow 3.3+, and
+ defers to the triggerer on older cores); a HITL review on ``AgentOperator`` polls from the
+ worker and holds its slot. Set ``approval_timeout`` or ``hitl_timeout`` so an unattended
+ review cannot wait forever. See :doc:`approval_gates` and :doc:`hitl_review`.
+
+The provider returns a rate limit or is down
+ Three features answer this at different layers. Airflow's own ``retries`` re-run the
+ task. :doc:`provider_fallback` fails over to another vendor inside one attempt.
+ :doc:`retry_policies` lets a model classify the failure and decide whether a retry is
+ worth it. They compose; each page says where it sits among the others.
+
+Still stuck
+-----------
+
+Every operator logs a post-run summary with the model name, token usage and the tool
+call sequence, and ``AgentOperator`` logs each tool call as it happens. Turn the task
+log level to ``DEBUG`` to see tool arguments and the model output. :doc:`observability`
+covers exporting the same information as OpenTelemetry traces.