This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 320b0c5d03d Let Claude Agent SDK agent loops use common.ai toolsets
through a new AirflowTools adapter (#74026)
320b0c5d03d is described below
commit 320b0c5d03d66dd4a2dc743bbad776a4eb3852cf
Author: Wei Lee <[email protected]>
AuthorDate: Wed Oct 28 23:34:25 2026 +0900
Let Claude Agent SDK agent loops use common.ai toolsets through a new
AirflowTools adapter (#74026)
---
providers/common/ai/docs/features.rst | 8 +-
.../common/ai/docs/frameworks/claude_agent_sdk.rst | 282 ++++++
providers/common/ai/docs/frameworks/index.rst | 27 +-
providers/common/ai/docs/index.rst | 55 +-
providers/common/ai/docs/observability.rst | 4 +-
providers/common/ai/docs/stability.rst | 3 +
providers/common/ai/docs/toolsets/index.rst | 2 +-
providers/common/ai/pyproject.toml | 3 +
.../ai/example_dags/example_claude_agent_sdk.py | 90 ++
.../airflow/providers/common/ai/tools/__init__.py | 5 +-
.../providers/common/ai/tools/claude_agent_sdk.py | 369 +++++++
.../providers/common/ai/utils/tool_metrics.py | 4 +-
.../unit/common/ai/tools/test_claude_agent_sdk.py | 1060 ++++++++++++++++++++
.../observability/metrics/metrics_template.yaml | 4 +-
uv.lock | 27 +-
15 files changed, 1898 insertions(+), 45 deletions(-)
diff --git a/providers/common/ai/docs/features.rst
b/providers/common/ai/docs/features.rst
index e7f7149c191..7a461dec8ae 100644
--- a/providers/common/ai/docs/features.rst
+++ b/providers/common/ai/docs/features.rst
@@ -247,8 +247,8 @@ Recover from failures
Bring your own agent
--------------------
-The Strands, ADK, Anthropic SDK, OpenAI Agents SDK and other-framework routes
get every
-toolset except
+The Strands, ADK, Anthropic SDK, OpenAI Agents SDK, Claude Agent SDK and
other-framework
+routes get every toolset except
``MCPToolset`` and ``AgentSkillsToolset``.
.. list-table::
@@ -275,6 +275,10 @@ toolset except
- Gives an OpenAI Agents SDK agent the toolsets. Run through
``run_durably``, a retry
replays its completed calls.
- None; install ``openai-agents``
+ * - :doc:`Claude Agent SDK <frameworks/claude_agent_sdk>`
+ - Gives the Claude Agent SDK's agent loop the toolsets through an
in-process MCP
+ server.
+ - ``claude-agent-sdk``
* - :doc:`LangChain <hooks/langchain>`
- Builds LangChain chat and embedding models from a connection, and
:doc:`turns a toolset into LangChain tools <toolsets/langchain>`.
diff --git a/providers/common/ai/docs/frameworks/claude_agent_sdk.rst
b/providers/common/ai/docs/frameworks/claude_agent_sdk.rst
new file mode 100644
index 00000000000..6b27e4f8dbe
--- /dev/null
+++ b/providers/common/ai/docs/frameworks/claude_agent_sdk.rst
@@ -0,0 +1,282 @@
+ .. 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/frameworks:claude_agent_sdk:
+
+Claude Agent SDK
+=================
+
+If your agent is a loop on the `Claude Agent SDK
+<https://platform.claude.com/docs/en/api/agent-sdk/overview>`__'s ``query()``,
which runs
+the Claude Code CLI as a subprocess, you can run it in an Airflow task as it
is and give
+it the toolsets this provider ships.
+:class:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools`
serves
+``SQLToolset``, ``HookToolset`` and most other toolsets to the CLI through an
in-process
+MCP server. If you have no loop yet, ``AgentOperator`` runs one for you, with
durable
+replay and human review; see :doc:`index`.
+
+.. note::
+
+ ``AirflowTools`` is experimental: it may change in a minor release of this
provider.
+ The framework-neutral tool interface it builds on is stable.
+ See :ref:`howto/stability`.
+
+Run a loop in a task
+---------------------
+
+This task answers a question about a database. The client's API key comes from
an
+Airflow connection:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py
+ :language: python
+ :start-after: [START example_claude_agent_sdk]
+ :end-before: [END example_claude_agent_sdk]
+
+The model can call ``list_tables``, ``get_schema``, ``query`` and
``check_query``, the
+tools ``SQLToolset`` serves. Each call writes a group to the task log. This is
the log of
+one run against a temporary SQLite database, captured as the task logger wrote
it, for
+the question "Inspect the columns of each table first with get_schema, then
tell me
+which customer name has the highest total order amount and in which city they
live."
+Which tools the model calls, in what order, and how long each takes depend on
the model
+and the question:
+
+.. code-block:: text
+
+ 2026-10-19T07:56:54.324694Z [info ] ::group::Tool call: list_tables
[airflow.providers.common.ai.tools] loc=__init__.py:116
+ 2026-10-19T07:56:54.603210Z [info ] Tool list_tables returned in 0.28s
[airflow.providers.common.ai.tools] loc=__init__.py:130
+ 2026-10-19T07:56:54.603326Z [info ] ::endgroup::
[airflow.providers.common.ai.tools] loc=__init__.py:132
+ 2026-10-19T07:56:56.308886Z [info ] ::group::Tool call: get_schema
[airflow.providers.common.ai.tools] loc=__init__.py:116
+ 2026-10-19T07:56:56.309896Z [info ] Tool get_schema returned in 0.00s
[airflow.providers.common.ai.tools] loc=__init__.py:130
+ 2026-10-19T07:56:56.310010Z [info ] ::endgroup::
[airflow.providers.common.ai.tools] loc=__init__.py:132
+ 2026-10-19T07:56:56.476896Z [info ] ::group::Tool call: get_schema
[airflow.providers.common.ai.tools] loc=__init__.py:116
+ 2026-10-19T07:56:56.478181Z [info ] Tool get_schema returned in 0.00s
[airflow.providers.common.ai.tools] loc=__init__.py:130
+ 2026-10-19T07:56:56.478343Z [info ] ::endgroup::
[airflow.providers.common.ai.tools] loc=__init__.py:132
+ 2026-10-19T07:56:58.700718Z [info ] ::group::Tool call: query
[airflow.providers.common.ai.tools] loc=__init__.py:116
+ 2026-10-19T07:56:58.704926Z [info ] Running statement: SELECT c.name,
c.city, SUM(o.amount) AS total_amount
+ FROM customers c
+ JOIN orders o ON o.customer_id = c.id
+ GROUP BY c.id, c.name, c.city
+ ORDER BY total_amount DESC
+ LIMIT 5, parameters: None
[airflow.task.hooks.airflow.providers.sqlite.hooks.sqlite.SqliteHook]
loc=sql.py:896
+ 2026-10-19T07:56:59.495567Z [info ] Tool query returned in 0.79s
[airflow.providers.common.ai.tools] loc=__init__.py:130
+ 2026-10-19T07:56:59.495691Z [info ] ::endgroup::
[airflow.providers.common.ai.tools] loc=__init__.py:132
+
+Only ``AirflowTools`` and ``tools.run(query(...))`` come from Airflow.
``query()``, the
+system prompt, the model and ``max_turns`` are the SDK's own.
+:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.build_options`
builds
+``ClaudeAgentOptions`` with the Airflow tools wired in and safe defaults for
the fields
+that matter most for this adapter; see `Built-in tools run on the worker
host`_. You can
+also build ``ClaudeAgentOptions`` yourself and add
+:attr:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.server`
and
+:attr:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.allowed_tools`
to
+it. Put the server in ``mcp_servers`` under the key ``tools.server_name``
(``"airflow"`` by
+default). The SDK names a server after its dict key, so ``{"warehouse":
tools.server}``
+serves ``mcp__warehouse__*`` tools while ``allowed_tools`` lists
``mcp__airflow__*`` and
+the names do not match.
+
+What each tool does
+--------------------
+
+``AirflowTools`` accepts any toolset that implements
+:class:`~airflow.providers.common.ai.tools.ToolProvider`, and individual
+:class:`~airflow.providers.common.ai.tools.AirflowTool` objects.
``MCPToolset`` is the
+exception; give the SDK the MCP server directly instead. Each tool:
+
+- Keeps the toolset tool's name, description and argument schema, served as an
MCP tool
+ named ``mcp__<server_name>__<tool name>``.
+- Runs through the toolset, so connection resolution, SQL validation,
+ ``allowed_tables``, ``allowed_methods`` and result limits behave as they do
in
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. Because
the
+ built-in tools are off, the model cannot open a result the SDK writes to a
file, so each
+ tool is marked with 500,000 characters as the size the SDK should keep
inline. The bundled
+ CLI caps that value at 500,000 (observed on CLI 2.1.286) and counts it in
UTF-16 code units,
+ so an emoji counts as two. A single result longer than that
+ is truncated to fit and ends with a notice asking the model to narrow the
request; the
+ model still reads it as a normal result, not an error. Without the
truncation the CLI would
+ spill the result to a file the model cannot read. The CLI also writes each
result to stdout
+ twice, JSON-escaped, and the SDK fails the stream on any line over its
``max_buffer_size``
+ (1 MiB by default), which a result well under 500,000 characters can exceed.
+ ``build_options()`` therefore sets ``max_buffer_size`` to 16 MiB.
+- Passes its result through Airflow's secret masker before the model sees it. A
+ connection password that turns up in a query result reaches the model as
``***``.
+- Writes a ``Tool call: <name>`` group to the task log. Calls to a toolset are
also
+ counted in the ``common_ai.tool_calls`` metric with
``framework=claude_agent_sdk`` (see
+ :doc:`../observability`).
+
+When a tool fails
+-------------------
+
+A failure the model can correct, such as a query naming a table outside
+``allowed_tables`` or one the database rejects, comes back to it as an MCP
tool result
+with ``isError`` set. This is what the model read when it asked a
``SQLToolset`` to run a
+``DELETE`` statement, which the read-only validation rejects. It is the
``tool_result``
+block from the ``UserMessage`` in the SDK stream of a real run (indented here;
+``tool_use_id`` is the value from that run), so it shows the API-level message
the model
+reads, not the handler's return value or the MCP wire format:
+
+.. code-block:: json
+
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01Sitdzj84kaH1CKYC8LFW2b",
+ "content": "The query tool failed: Statement type 'Delete' is not
allowed. Allowed types: Select, Union, Intersect, Except, Describe, Show\nUse
the list_tables and get_schema tools to inspect the database, then fix the
query and try again.",
+ "is_error": true
+ }
+
+``SQLToolset``, ``HookToolset`` and ``ObjectStorageToolset`` allow one
correction: a tool
+that keeps failing ends the run with
+``ToolCallError: query kept failing after 1 correction(s): ...``. The retry
budget is
+counted per run of
+:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run`,
and a model
+turn counts as one failure however many of its calls fail, the same way the
Anthropic and
+Strands adapters count them. The turn is the ``AssistantMessage.message_id``;
without one,
+calls that run at the same time count as one failure. A call whose arguments
fail the SDK's own JSON Schema validation never
+reaches Airflow's tools, so it is not counted against the budget and not
recorded in the
+``common_ai.tool_calls`` metric; see `max_turns`_ below.
+
+Any other failure, such as a hook raising, ends the run. The Claude Agent
SDK's own tool
+dispatch catches every exception a handler raises and turns it into an error
result the
+model reads, so
+:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run`
has to work around
+that to still fail the task: it raises
+:class:`~airflow.providers.common.ai.tools.ToolCallError` once the model has
read the
+(masked) failure, rather than letting the stream continue to a later
+``ResultMessage``. Airflow's own retry takes over once the task fails.
+
+A run can also end with a ``ResultMessage`` whose ``is_error`` is set, for
example when
+``max_turns`` is reached or the API call fails. ``run()`` does not return that
result as
+a success. Which exception you get depends on the stream:
+
+* ``query()``: the CLI exits non-zero right after an error result, and the SDK
raises
+ ``claude_agent_sdk.ResultError`` (a ``ProcessError``) from the stream on the
next
+ iteration, before ``run()`` reaches its own check. ``ResultError`` carries
the result's
+ ``subtype``, ``terminal_reason`` and the rest of its payload.
+* ``ClaudeSDKClient.receive_response()``: the stream ends at the
``ResultMessage``, so
+ ``run()`` raises ``RuntimeError`` with the result's ``subtype``, ``result``,
+ ``errors``, ``stop_reason``, ``api_error_status`` and ``terminal_reason``
where they are
+ set, with long values shortened.
+
+For a streamed prompt, ``run()`` raises ``RuntimeError`` whenever any result
message has
+``is_error`` set, even if a later turn succeeds.
+
+If a tool failure of the kind above was also recorded, ``ToolCallError`` is
raised first
+on either stream. ``RuntimeError`` is also raised when the stream ends without
any
+``ResultMessage``.
+
+.. _max_turns:
+
+Set ``max_turns``
+-------------------
+
+Because an argument-validation failure from the SDK itself is invisible to
Airflow's
+tools, a model that keeps sending malformed arguments is not stopped by the
retry
+budget above. Always set ``max_turns`` on ``build_options()`` (or on
+``ClaudeAgentOptions`` directly) so such a loop still ends.
+
+Drive the loop with ``tools.run``
+-----------------------------------
+
+``run()`` takes the stream that ``query()`` or
``ClaudeSDKClient.receive_response()``
+returns. ``ClaudeSDKClient.receive_messages()`` is not supported, because it
never ends.
+
+The Claude Agent SDK's tool dispatch hands every handler exception to the
model as an
+error result. If you iterate ``query()`` yourself instead of using
+:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run`,
+a hook that raises becomes one more thing for the model to read, and the task
can succeed
+with an answer written around the failure. Airflow's tools log a warning when
a failure
+goes to the model this way.
+
+An ``AirflowTools`` instance drives one query at a time: build a new instance
per task,
+or call
+:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run`
again
+only once the previous call has returned. A tool call that arrives after a run
has ended
+is refused with an ``is_error`` result until the next run starts, including a
handler
+called directly between two runs.
+
+Built-in tools run on the worker host
+----------------------------------------
+
+``ClaudeAgentOptions()`` on its own loads every settings source on the worker
+(``~/.claude``, project settings, a project ``.mcp.json``) and leaves the
CLI's built-in
+tools on, including a worker-host ``Bash`` -- the bundled CLI runs as a
subprocess on the
+same worker that runs the task, with its environment variables.
``build_options()`` closes all
+of that by default: no built-in tools, no settings sources, and
``strict_mcp_config=True``
+so only the MCP servers passed in the options are used: the one
``AirflowTools`` builds
+plus any ``mcp_servers`` you pass to ``build_options()``. It also sets
+``verbatim_prompts=True``, so the CLI does not expand ``@path`` file mentions
in a prompt.
+Otherwise anyone who controls the prompt could send any file the worker user
can read to the
+model, and that content would not pass through the secret masker. It sets
+``max_buffer_size`` to 16 MiB (the SDK's default is 1 MiB) so a result near
the 500,000
+character limit does not break the stream; a value you pass wins.
+
+``verbatim_prompts=True`` has side effects:
+
+- Slash commands in the prompt are not dispatched.
+- The CLI's turn-start attachment pass is off.
+- The SDK expects CLI 2.1.248 or later; an older CLI only logs a warning.
+
+When you pass ``can_use_tool`` to ``build_options()``:
+
+- The Airflow tools are not added to ``allowed_tools``, because an allowed
tool never reaches
+ the callback. Your callback decides on every Airflow tool call.
+- ``permission_mode`` defaults to ``"default"``; a value you pass wins.
+- ``AirflowTools.allowed_tools`` still lists the tool names if you want to
allow some by hand.
+
+If you build ``ClaudeAgentOptions`` yourself instead, set those fields the
same way.
+
+Tracing
+---------
+
+:func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`
does not cover
+the Claude Agent SDK: it instruments spans created in this Python process, but
the agent
+loop runs in the bundled CLI's own subprocess. The SDK passes the active
span's W3C trace
+context to the CLI, but whether the CLI emits its own telemetry from it is not
something
+this provider verifies or relies on.
+
+Differences from ``AgentOperator``
+-------------------------------------
+
+The loop is the SDK's, so what ``AgentOperator`` builds on top of Pydantic AI
does not
+apply:
+
+- A task retry runs the loop again from the start. Tools that change data must
be safe
+ to repeat, because there is no step-level replay as with
+ ``AgentOperator(durable=True)``.
+- There is no human review step like
``AgentOperator(enable_hitl_review=True)``.
+- A toolset's connection ID is used as written. It is not rendered as a
template, so
+ pass the value rather than a Jinja expression.
+
+Installation
+--------------
+
+Install this provider with its ``claude-agent-sdk`` extra:
+
+.. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[claude-agent-sdk]"
+
+The extra installs the SDK, which bundles the Claude Code CLI. The adapter
needs
+``claude-agent-sdk`` 0.2.160 or later, which the extra requires, and is tested
with
+0.2.168.
+
+.. seealso::
+
+ - :doc:`../toolsets/index`, for what each toolset exposes and its limits.
+ - :doc:`index`, for the framework-neutral tool interface ``AirflowTools``
is built on.
+ - `Claude Agent SDK overview
+ <https://platform.claude.com/docs/en/api/agent-sdk/overview>`__, in
Anthropic's
+ documentation, for the SDK and the CLI it bundles.
diff --git a/providers/common/ai/docs/frameworks/index.rst
b/providers/common/ai/docs/frameworks/index.rst
index ff9b6169563..661e3ca98d1 100644
--- a/providers/common/ai/docs/frameworks/index.rst
+++ b/providers/common/ai/docs/frameworks/index.rst
@@ -50,7 +50,7 @@ AI because it does not tie a Dag to either:
Pydantic AI is the recommendation, not a requirement. For a new agent, start
with
``AgentOperator``, which gets the most from this provider. If your team has
standardized
-on Strands, ADK, LangChain, the Anthropic SDK or the OpenAI Agents SDK, keep
that framework: the table below
+on Strands, ADK, LangChain, the Anthropic SDK, the OpenAI Agents SDK or the
Claude Agent SDK, keep that framework: the table below
shows what this provider gives an agent written in each one.
Supported frameworks
@@ -103,6 +103,12 @@ is Airflow's and which part stays yours.
the model cannot fix fail the task. Guide: :doc:`openai_agents`.
- You
- ``openai-agents``
+ * - Claude Agent SDK
+ -
:class:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools` gives
the
+ SDK's agent loop the same tools through an in-process MCP server, and
makes a tool
+ failure the model cannot fix fail the task. Guide:
:doc:`claude_agent_sdk`.
+ - You
+ - ``[claude-agent-sdk]`` extra
* - LangChain
- :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook`
returns chat
and embedding models configured from a connection, and
@@ -119,9 +125,9 @@ is Airflow's and which part stays yours.
- No agent
- ``[llamaindex]`` extra
-The Strands, ADK, Anthropic and OpenAI Agents integrations and the tracing
helper are experimental: they
-can change or be removed in a minor release of this provider. The
framework-neutral tool
-interface under them is stable. See :ref:`howto/stability`.
+The Strands, ADK, Anthropic, OpenAI Agents and Claude Agent SDK integrations
and the tracing
+helper are experimental: they can change or be removed in a minor release of
this provider. The
+framework-neutral tool interface under them is stable. See
:ref:`howto/stability`.
Tested versions
---------------
@@ -145,11 +151,15 @@ Tested versions
- 0.20.0 and 0.23.1
- 0.21.0 and later require ``openai>=3``, while Airflow's constraints
file pins
``openai`` below 3, so installing with that file gives 0.20.
+ * - Claude Agent SDK
+ - 0.2.168
+ - Bundles the Claude Code CLI as the subprocess the SDK talks to.
These are the versions the adapter tests have been run against. Strands and
ADK cap
packages below the versions in Airflow's dependency lock, and current OpenAI
Agents SDK
releases need ``openai`` 3, so CI runs their adapter tests in a job of their
own, one
-framework at a time.
+framework at a time. The Claude Agent SDK adapter has no such conflict, so CI
runs its
+tests.
Outside canary runs, the job holds every package already in the CI image at its
installed version with ``--override``, so Strands' and ADK's own dependency
caps are
@@ -175,6 +185,8 @@ Start from the agent you already have:
``AirflowTools`` to its tool runner. See :doc:`anthropic`.
- **An existing OpenAI Agents SDK agent**: keep it, and add Airflow's toolsets
to its
``tools`` with ``AirflowTools``. See :doc:`openai_agents`.
+- **A loop written on the Claude Agent SDK**: keep it, and hand it Airflow's
toolsets with
+ ``AirflowTools``. See :doc:`claude_agent_sdk`.
- **An existing LangChain agent**: build its model with ``LangChainHook`` and
add
Airflow's toolsets with ``airflow_toolset_to_langchain_tools``.
- **No agent yet, or no preference**: use ``AgentOperator`` or
``@task.agent``. That route
@@ -194,7 +206,7 @@ Airflow also depends on, such as OpenTelemetry or a vendor
SDK.
To give such a framework Airflow's toolsets rather than hand-written tools,
build on
the framework-neutral interface in :mod:`airflow.providers.common.ai.tools`,
which is
-what the Strands, ADK, Anthropic and OpenAI Agents integrations are written
against:
+what the Strands, ADK, Anthropic, OpenAI Agents and Claude Agent SDK
integrations are written against:
- :class:`~airflow.providers.common.ai.tools.AirflowTool` is one operation: a
name, a
description, a JSON Schema for its arguments and an async function. Call it
through
@@ -214,7 +226,7 @@ what the Strands, ADK, Anthropic and OpenAI Agents
integrations are written agai
An adapter maps these onto the framework's own tool type and error status, and
makes
sure a ``ToolCallError`` ends the run rather than reaching the model, as the
Strands
plugin does. :func:`~airflow.providers.common.ai.tools.collect_tools` turns
the toolsets and
-tools an adapter is given into one list, as the Strands, ADK, Anthropic and
OpenAI Agents
+tools an adapter is given into one list, as the Strands, ADK, Anthropic,
OpenAI Agents and Claude Agent SDK
adapters do.
Blocking hook calls, such as a SQL query, run one at a time in the task's
process, so an
@@ -235,3 +247,4 @@ on; see :doc:`../observability`. For an OpenAI Agents SDK
agent, also register
Google ADK <adk>
Anthropic SDK <anthropic>
OpenAI Agents SDK <openai_agents>
+ Claude Agent SDK <claude_agent_sdk>
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 68152731e5e..d94b5b20574 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -224,33 +224,34 @@ Install them when installing from PyPI. For example:
pip install apache-airflow-providers-common-ai[anthropic]
-===============
=======================================================================================================================================
-Extra Dependencies
-===============
=======================================================================================================================================
-``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``,
``anthropic>=1.1.0,<2``
-``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0``
-``google`` ``pydantic-ai-slim[google]>=2.36.0``
-``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0``
-``typesafe`` ``typesafe-sdk>=0.6.0``
-``mcp`` ``pydantic-ai-slim[mcp]>=2.36.0``
-``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2``
-``opensandbox`` ``opensandbox>=1.1.0``
-``boat`` ``boat-sdk>=1.9.3``
-``openshell`` ``openshell>=0.1.2,<0.2``, ``grpcio>=1.78.0``,
``protobuf>=6.31.1``
-``code-mode`` ``pydantic-ai-harness[codemode]>=0.35.0``
-``shields`` ``pydantic-ai-shields>=0.3.4``
-``skills`` ``apache-airflow-providers-git>=0.4.0``,
``pydantic-ai-skills>=1.2.0,<2``
-``islo`` ``islo>=0.3.19``
-``avro`` ``fastavro>=1.10.0; python_version < "3.14"``,
``fastavro>=1.12.1; python_version >= "3.14"``
-``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``,
``pyarrow>=22.0.0; python_version >= '3.14'``
-``sql`` ``apache-airflow-providers-common-sql>=2.2.0``,
``sqlglot>=30.0.0``
-``common.sql`` ``apache-airflow-providers-common-sql>=2.2.0``
-``langchain`` ``langchain>=1.0.0``
-``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.14.5``,
``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.8``
-``pdf`` ``pypdf>=4.0.0``
-``docx`` ``python-docx>=1.0.0``
-``git`` ``apache-airflow-providers-git``
-===============
=======================================================================================================================================
+====================
=======================================================================================================================================
+Extra Dependencies
+====================
=======================================================================================================================================
+``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``,
``anthropic>=1.1.0,<2``
+``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0``
+``google`` ``pydantic-ai-slim[google]>=2.36.0``
+``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0``
+``typesafe`` ``typesafe-sdk>=0.6.0``
+``mcp`` ``pydantic-ai-slim[mcp]>=2.36.0``
+``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2``
+``opensandbox`` ``opensandbox>=1.1.0``
+``boat`` ``boat-sdk>=1.9.3``
+``openshell`` ``openshell>=0.1.2,<0.2``, ``grpcio>=1.78.0``,
``protobuf>=6.31.1``
+``code-mode`` ``pydantic-ai-harness[codemode]>=0.35.0``
+``shields`` ``pydantic-ai-shields>=0.3.4``
+``skills`` ``apache-airflow-providers-git>=0.4.0``,
``pydantic-ai-skills>=1.2.0,<2``
+``islo`` ``islo>=0.3.19``
+``avro`` ``fastavro>=1.10.0; python_version < "3.14"``,
``fastavro>=1.12.1; python_version >= "3.14"``
+``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``,
``pyarrow>=22.0.0; python_version >= '3.14'``
+``sql`` ``apache-airflow-providers-common-sql>=2.2.0``,
``sqlglot>=30.0.0``
+``common.sql`` ``apache-airflow-providers-common-sql>=2.2.0``
+``langchain`` ``langchain>=1.0.0``
+``llamaindex`` ``dataclasses-json>=0.6.7``,
``llama-index-core>=0.14.5``, ``llama-index-embeddings-openai>=0.6.0``,
``llama-index-llms-openai>=0.6.8``
+``pdf`` ``pypdf>=4.0.0``
+``docx`` ``python-docx>=1.0.0``
+``git`` ``apache-airflow-providers-git``
+``claude-agent-sdk`` ``claude-agent-sdk>=0.2.160``
+====================
=======================================================================================================================================
Downloading official packages
-----------------------------
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index ad5bbb769c8..8068bdf7a4b 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -277,8 +277,8 @@ carries three tags:
- The toolset class, such as ``SQLToolset`` or ``ObjectStorageToolset``.
* - ``framework``
- ``pydantic_ai`` for ``AgentOperator`` and your own Pydantic AI agents;
- ``strands``, ``adk``, ``langchain``, ``anthropic`` or ``openai_agents``
for a call
- through that framework's adapter;
+ ``strands``, ``adk``, ``langchain``, ``anthropic``, ``openai_agents``
or ``claude_agent_sdk``
+ for a call through that framework's adapter;
``none`` for a call to a toolset's ``airflow_tools()`` made without an
adapter.
* - ``outcome``
- ``executed`` when the call returned, ``failed`` when it raised
(including a
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index df3ee5898b3..4ab1ac43d81 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -272,6 +272,9 @@ Everything this provider ships that is not in the table
above is experimental.
(:doc:`frameworks/openai_agents`)
- Written against OpenAI Agents SDK 0.20 and 0.23. The SDK is pre-1.0 and
changes
its tool and runner APIs between minor releases.
+ * - The Claude Agent SDK adapter (:doc:`frameworks/claude_agent_sdk`)
+ - New; the SDK is pre-1.0, and its in-process MCP server turns a
handler's exception
+ into an error result the adapter has to work around.
* - The ``common_ai.tool_calls`` metric and
:func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`
(:doc:`observability`)
diff --git a/providers/common/ai/docs/toolsets/index.rst
b/providers/common/ai/docs/toolsets/index.rst
index 2291ead1a45..230b6e9c0fd 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -299,7 +299,7 @@ used to allow exactly one correction whatever ``retries``
said, so an agent that
query, and a large integer ``retries`` meant for output validation also lets a
failing
database be queried that many times. Pass ``max_retries=1`` to a toolset to
keep the old
behaviour. Outside a pydantic-ai agent (the LangChain, Strands, Google ADK,
Anthropic
-tool runner and OpenAI Agents SDK bridges) there is no agent budget, so each
tool gets one correction unless its
+tool runner, OpenAI Agents SDK and Claude Agent SDK bridges) there is no agent
budget, so each tool gets one correction unless its
toolset sets ``max_retries``.
Layering
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 14f46d326b5..d41a04879b2 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -179,6 +179,8 @@ dependencies = [
"git" = [
"apache-airflow-providers-git"
]
+# The Claude Agent SDK adapter in
airflow.providers.common.ai.tools.claude_agent_sdk.
+"claude-agent-sdk" = ["claude-agent-sdk>=0.2.160"]
[dependency-groups]
dev = [
@@ -211,6 +213,7 @@ dev = [
"openshell>=0.1.2,<0.2",
"grpcio>=1.78.0",
"protobuf>=6.31.1",
+ "claude-agent-sdk>=0.2.160",
]
# To build docs:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py
new file mode 100644
index 00000000000..d455954436f
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py
@@ -0,0 +1,90 @@
+# 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.
+"""
+Answer a question about a database with the Claude Agent SDK's agent loop, in
an Airflow task.
+
+The agent loop is the SDK's own ``query()``, which runs the bundled Claude
Code CLI as a
+subprocess. Airflow supplies the API key, from a connection, and the tools:
+:class:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools`
serves a
+``SQLToolset`` to the CLI through an in-process MCP server, with read-only SQL
validation,
+bounded results and Airflow's secret masker applied to every result.
``AirflowTools.build_options``
+also turns off the CLI's built-in tools and every settings source on the
worker, so the only
+tools the model can call are Airflow's.
+
+Before running:
+
+1. Install this provider's ``claude-agent-sdk`` extra:
+ ``pip install "apache-airflow-providers-common-ai[claude-agent-sdk]"``.
+2. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
password is
+ an Anthropic API key.
+3. Create a database connection (``DB_CONN_ID``, default ``sql_default``)
whose hook is a
+ ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+"""
+
+from __future__ import annotations
+
+import os
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "anthropic_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "claude-sonnet-5")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_claude_agent_sdk]
+@dag(tags=["example"])
+def example_claude_agent_sdk():
+ """Answer a question about a database with the Claude Agent SDK's agent
loop."""
+
+ @task
+ def ask_the_warehouse(question: str = DEFAULT_QUESTION) -> str | None:
+ import asyncio
+
+ from claude_agent_sdk import query
+
+ from airflow.providers.common.ai.tools.claude_agent_sdk import
AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ llm = BaseHook.get_connection(LLM_CONN_ID)
+ tools = AirflowTools(SQLToolset(db_conn_id=DB_CONN_ID))
+ options = tools.build_options(
+ model=LLM_MODEL,
+ system_prompt=(
+ "You are a SQL analyst. Use list_tables and get_schema to
explore the "
+ "database, then run read-only queries to answer the question."
+ ),
+ # Without a limit the model can keep asking for tool calls
indefinitely if it
+ # never settles on an answer; see "When a tool fails" in the guide.
+ max_turns=10,
+ env={"ANTHROPIC_API_KEY": llm.password},
+ )
+
+ async def ask():
+ return await tools.run(query(prompt=question, options=options))
+
+ return asyncio.run(ask()).result
+
+ ask_the_warehouse()
+
+
+# [END example_claude_agent_sdk]
+
+
+example_claude_agent_sdk()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
index 349a11081c1..5f116b5569d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
@@ -20,8 +20,9 @@ Framework-neutral tools backed by Airflow connections.
An :class:`AirflowTool` is one operation an agent can call: a name, a
description, a JSON Schema for its arguments, and an async function. Nothing in
it belongs to a particular agent framework, so the same tool can be handed to
-Strands Agents, Google ADK, LangChain, the Anthropic SDK's tool runner or the
OpenAI
-Agents SDK through a small adapter, while the agent itself stays native to
that framework.
+Strands Agents, Google ADK, LangChain, the Anthropic SDK's tool runner, the
OpenAI Agents SDK
+or the Claude Agent SDK through a small adapter, while the agent itself stays
native to that
+framework.
The toolsets this provider ships, such as
:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py
new file mode 100644
index 00000000000..49cae35fda9
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py
@@ -0,0 +1,369 @@
+# 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.
+"""
+Give Airflow tools to an agent loop run by the Claude Agent SDK.
+
+.. note:: Experimental; see :ref:`howto/stability`.
+"""
+
+from __future__ import annotations
+
+try:
+ from claude_agent_sdk import (
+ AssistantMessage,
+ ClaudeAgentOptions,
+ McpSdkServerConfig,
+ Message,
+ ResultMessage,
+ ToolAnnotations,
+ create_sdk_mcp_server,
+ tool,
+ )
+except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException(e)
+
+import copy
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
collect_tools
+from airflow.providers.common.ai.tools._from_toolset import tool_call_scope
+from airflow.providers.common.ai.utils.tool_definition import serialize_for_llm
+from airflow.providers.common.ai.utils.tool_metrics import calling_framework
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncIterable, Awaitable, Callable
+
+ from airflow.providers.common.ai.tools import ToolProvider
+
+log = logging.getLogger(__name__)
+
+__all__ = ["AirflowTools"]
+
+
+_MAX_ERROR_DETAIL_CHARS = 500
+_MAX_ERROR_ENTRIES = 5
+# The bundled CLI clamps ``maxResultSizeChars`` to 500_000 (observed on CLI
2.1.286; other
+# versions may differ), so a larger value is a no-op. With ``tools=[]`` the
model has no Read
+# tool, so a result over the ceiling that the CLI spills to a
``<persisted-output>`` file is lost.
+_MAX_INLINE_RESULT_CHARS = 500_000
+# The CLI writes each tool result to stdout twice, JSON-escaped, and the SDK
fails the stream on any
+# line over ``max_buffer_size`` (1 MiB by default). Escaping can multiply a
result by up to ~6, so the
+# buffer must sit far above ``_MAX_INLINE_RESULT_CHARS``. Sizes observed with
the CLI bundled in
+# claude-agent-sdk 0.2.165, not a documented guarantee.
+_MAX_BUFFER_BYTES = 16 * 1024 * 1024
+
+
+def _count_utf16_units(text: str) -> int:
+ return len(text.encode("utf-16-le", errors="surrogatepass")) // 2
+
+
+def _truncate_result_text(text: str) -> str:
+ # The CLI measures the ceiling in UTF-16 code units, so a code point
outside the BMP (most
+ # emoji) counts twice; ``len(text)`` would let such a result through and
the CLI would spill it.
+ if _count_utf16_units(text) <= _MAX_INLINE_RESULT_CHARS:
+ return text
+ notice = (
+ f"\n\n[Truncated: this result had {len(text)} characters and only the
beginning is shown. "
+ "Narrow the request, for example by adding filters, selecting fewer
columns or lowering "
+ "the LIMIT, to get a smaller result.]"
+ )
+ # The notice counts toward the ceiling: text kept up to the ceiling plus
the notice would
+ # still be spilled by the CLI.
+ keep_units = _MAX_INLINE_RESULT_CHARS - _count_utf16_units(notice)
+ units = text.encode("utf-16-le", errors="surrogatepass")[: keep_units * 2]
+ # A cut that lands between the two halves of a surrogate pair would leave
a lone high surrogate.
+ if units and 0xD8 <= units[-1] <= 0xDB:
+ units = units[:-2]
+ return units.decode("utf-16-le", errors="surrogatepass") + notice
+
+
+def _bound_detail(text: str) -> str:
+ omitted = len(text) - _MAX_ERROR_DETAIL_CHARS
+ return text if omitted <= 0 else
f"{text[:_MAX_ERROR_DETAIL_CHARS]}...(+{omitted} chars)"
+
+
+def _describe_error_result(result: ResultMessage) -> str:
+ errors = result.errors or []
+ omitted_errors = len(errors) - _MAX_ERROR_ENTRIES
+ fields: dict[str, Any] = {
+ "subtype": result.subtype,
+ "result": result.result,
+ "errors": errors[:_MAX_ERROR_ENTRIES],
+ "stop_reason": result.stop_reason,
+ "api_error_status": result.api_error_status,
+ "terminal_reason": result.terminal_reason,
+ }
+ details = []
+ for name, value in fields.items():
+ if not value:
+ continue
+ text = _bound_detail(repr(value)) if name in ("result", "errors") else
repr(value)
+ if name == "errors" and omitted_errors > 0:
+ text = f"{text} (+{omitted_errors} more)"
+ details.append(f"{name}={text}")
+ return f"The Claude Agent SDK run ended with an error result ({',
'.join(details)})."
+
+
+async def _close_stream(messages: AsyncIterable[Message]) -> None:
+ # ``AsyncIterable`` does not guarantee ``aclose``; ``query()`` and
``receive_response()`` do.
+ aclose = getattr(messages, "aclose", None)
+ if aclose is not None:
+ await aclose()
+
+
+class _Run:
+ """
+ One run of :meth:`AirflowTools.run`.
+
+ Identifies its calls, collects their failures, and records whether it has
ended.
+ """
+
+ def __init__(self) -> None:
+ self.failures: list[ToolCallError] = []
+ self.turn: str | None = None
+ self.ended = False
+
+
+class AirflowTools:
+ """
+ Airflow's tools for an agent loop run by the Claude Agent SDK.
+
+ Pass toolsets that implement
+ :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as
``SQLToolset``
+ and ``HookToolset``, or individual
+ :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. They are
served
+ through an in-process MCP server; add it and :attr:`allowed_tools` to
+ ``ClaudeAgentOptions``, or build them with :meth:`build_options`, and
drive the SDK's
+ ``query()`` with :meth:`run`:
+
+ .. code-block:: python
+
+ from claude_agent_sdk import query
+
+ from airflow.providers.common.ai.tools.claude_agent_sdk import
AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ tools = AirflowTools(SQLToolset(db_conn_id="warehouse",
allowed_tables=["orders"]))
+ options = tools.build_options(model="claude-sonnet-5", max_turns=10)
+ result = await tools.run(query(prompt="...", options=options))
+
+ Each tool keeps the source tool's name, description and argument schema,
and every
+ result passes through Airflow's secret masker first. A failure the model
can
+ correct, such as a query naming a missing column, reaches it as an
``is_error``
+ tool result. A failure the model cannot fix, such as a hook raising, makes
+ :meth:`run` raise
:class:`~airflow.providers.common.ai.tools.ToolCallError`, so the
+ task fails and Airflow retries it. The SDK's own tool dispatch catches
every
+ exception a handler raises and turns it into an ``is_error`` result the
model
+ reads, so drive the query with :meth:`run` rather than iterating it
yourself.
+
+ An instance runs one query at a time: build a new ``AirflowTools`` per
task, or
+ call :meth:`run` again once the previous run has finished. A tool call
that arrives
+ after a run has ended is refused until the next :meth:`run` starts. A call
that lands
+ after the next run has started cannot be told apart and is served by that
run.
+
+ :param sources: Toolsets and tools to serve, in order.
+ :param server_name: The MCP server name the tools are served under. Also
the
+ prefix of each tool's name in :attr:`allowed_tools`
+ (``mcp__<server_name>__<tool name>``).
+ """
+
+ def __init__(self, *sources: ToolProvider | AirflowTool, server_name: str
= "airflow") -> None:
+ self.server_name = server_name
+ airflow_tools = collect_tools(sources)
+ self.allowed_tools: list[str] = [f"mcp__{server_name}__{t.name}" for t
in airflow_tools]
+ # ``_run`` is the run in progress (guards re-entry); ``_last_run`` is
the one that most
+ # recently ended (refuses late calls). It is only read while ``_run``
is None.
+ self._run: _Run | None = None
+ self._last_run: _Run | None = None
+ annotations =
ToolAnnotations(maxResultSizeChars=_MAX_INLINE_RESULT_CHARS)
+ sdk_tools = [
+ tool(t.name, t.description, self._build_tool_schema(t),
annotations=annotations)(
+ self._build_handler(t)
+ )
+ for t in airflow_tools
+ ]
+ self.server: McpSdkServerConfig =
create_sdk_mcp_server(name=server_name, tools=sdk_tools)
+
+ def build_options(self, **kwargs: Any) -> ClaudeAgentOptions:
+ """
+ Build ``ClaudeAgentOptions`` wired to this instance's tools, with safe
defaults.
+
+ ``ClaudeAgentOptions()`` on its own loads every settings source on the
worker
+ (``~/.claude``, project settings, a project ``.mcp.json``) and leaves
the CLI's
+ built-in tools on, including a worker-host ``Bash``. This method
instead
+ defaults ``tools`` to ``[]`` (no built-in tools), ``setting_sources``
to ``[]``
+ (no settings files) and ``strict_mcp_config`` to ``True`` (only the
MCP servers
+ passed here), and ``verbatim_prompts`` to ``True``. It adds
:attr:`server` and
+ :attr:`allowed_tools` to whatever ``mcp_servers`` and
``allowed_tools`` the caller
+ passes. A keyword this method does not default, such as ``model`` or
``env``, is
+ passed through unchanged.
+
+ ``verbatim_prompts`` stops the CLI from expanding ``@path`` file
mentions in a prompt;
+ otherwise anyone who controls the prompt could send a file the worker
user can read to
+ the model, without passing the secret masker. It also turns off slash
command dispatch
+ and the CLI's turn-start attachment pass.
+
+ When ``can_use_tool`` is passed, the Airflow tools are not added to
``allowed_tools``,
+ because an allowed tool never reaches the callback, and
``permission_mode`` defaults to
+ ``"default"``. The callback then decides on every Airflow tool call;
the names are still
+ in :attr:`allowed_tools` if you want to allow some by hand.
+
+ Everything it builds can also be assembled by hand; see :attr:`server`
and
+ :attr:`allowed_tools`.
+
+ :raises ValueError: when ``mcp_servers`` already has a server named
+ :attr:`server_name`.
+ """
+ mcp_servers = dict(kwargs.pop("mcp_servers", {}))
+ if self.server_name in mcp_servers:
+ raise ValueError(
+ f"mcp_servers already has a server named {self.server_name!r}.
Pass a different "
+ "server_name to AirflowTools, or drop it from mcp_servers."
+ )
+ mcp_servers[self.server_name] = self.server
+ allowed_tools = list(kwargs.pop("allowed_tools", []))
+ if kwargs.get("can_use_tool") is None:
+ allowed_tools.extend(self.allowed_tools)
+ else:
+ kwargs.setdefault("permission_mode", "default")
+ kwargs.setdefault("tools", [])
+ kwargs.setdefault("setting_sources", [])
+ kwargs.setdefault("strict_mcp_config", True)
+ kwargs.setdefault("verbatim_prompts", True)
+ kwargs.setdefault("max_buffer_size", _MAX_BUFFER_BYTES)
+ return ClaudeAgentOptions(mcp_servers=mcp_servers,
allowed_tools=allowed_tools, **kwargs)
+
+ async def run(self, messages: AsyncIterable[Message]) -> ResultMessage:
+ """
+ Drive ``messages``, the stream ``query()`` or
``ClaudeSDKClient.receive_response()`` returns.
+
+ ``ClaudeSDKClient.receive_messages()`` is not supported: it never ends.
+
+ ``messages`` is closed when this returns or raises. For ``query()``
that ends the CLI
+ session, but only once the SDK finishes tearing down, so a tool call
can still arrive
+ after this returns or raises; it is answered with an ``is_error`` "Not
run" result.
+ A ``ClaudeSDKClient`` stays connected until you disconnect it, for
example by leaving
+ its ``async with`` block.
+
+ :raises RuntimeError: when this instance is already running a query,
when
+ ``messages`` ends without a
:class:`~claude_agent_sdk.ResultMessage`, or, for
+ ``ClaudeSDKClient.receive_response()``, when the result message
has ``is_error``
+ set (for example ``max_turns`` reached or an API failure).
+ For a streamed prompt, any result message with ``is_error`` set
raises, even when a
+ later turn succeeds.
+ :raises ResultError: for ``query()`` when the result message has
``is_error`` set.
+ The CLI exits non-zero right after such a result, so the SDK raises
+ ``claude_agent_sdk.ResultError`` (a ``ProcessError``) from the
stream before this
+ method checks the result.
+ :raises ToolCallError: when a tool failed in a way the model cannot
fix.
+ """
+ if self._run is not None:
+ raise RuntimeError(
+ f"{self!r} is already running a query. Build a separate
AirflowTools instance "
+ "for a concurrent run, or await the first run before starting
another."
+ )
+ state = _Run()
+ self._run = state
+ raised = True
+ try:
+ result: ResultMessage | None = None
+ # Each turn yields its own ResultMessage; keep the first error so
a later success
+ # cannot mask it.
+ error_result: ResultMessage | None = None
+ async for message in messages:
+ if isinstance(message, AssistantMessage):
+ # The CLI sends each turn's message before running that
turn's tools, which is
+ # what ``handle_tool_call`` relies on. Observed with the
bundled CLI, not documented.
+ state.turn = message.message_id
+ if isinstance(message, ResultMessage):
+ result = message
+ if message.is_error and error_result is None:
+ error_result = message
+ if state.failures:
+ raise state.failures[0]
+ if result is None:
+ if state.failures:
+ raise state.failures[0]
+ raise RuntimeError("The Claude Agent SDK stream ended without
a ResultMessage.")
+ if error_result is not None:
+ raise RuntimeError(_describe_error_result(error_result))
+ raised = False
+ return result
+ finally:
+ try:
+ await _close_stream(messages)
+ except Exception as e:
+ if not raised:
+ raise
+ log.warning("Closing the Claude Agent SDK stream failed: %s",
e)
+ finally:
+ state.ended = True
+ self._last_run = state
+ self._run = None
+
+ @staticmethod
+ def _build_tool_schema(airflow_tool: AirflowTool) -> dict[str, Any]:
+ schema = copy.deepcopy(airflow_tool.parameters)
+ schema.setdefault("properties", {})
+ return schema
+
+ def _build_handler(
+ self, airflow_tool: AirflowTool
+ ) -> Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]:
+ async def handle_tool_call(args: dict[str, Any]) -> dict[str, Any]:
+ state = self._run if self._run is not None else self._last_run
+ if state is not None and (state.failures or state.ended):
+ reason = (
+ "an earlier Airflow tool call in this run failed"
+ if state.failures
+ else "the Airflow tool run has ended"
+ )
+ return {
+ "content": [{"type": "text", "text": f"Not run:
{reason}."}],
+ "is_error": True,
+ }
+ with (
+ calling_framework("claude_agent_sdk"),
+ tool_call_scope(run=state, turn=state.turn if state is not
None else None),
+ ):
+ try:
+ result = await airflow_tool.call(args)
+ except ToolCallError as e:
+ if state is None:
+ log.warning(
+ "Tool %s failed, but the Claude Agent SDK query
was not driven by "
+ "AirflowTools.run, so the model reads this failure
instead of the "
+ "task failing: %s",
+ airflow_tool.name,
+ e,
+ )
+ else:
+ state.failures.append(e)
+ return {"content": [{"type": "text", "text": str(e)}],
"is_error": True}
+ payload: dict[str, Any] = {
+ "content": [
+ {"type": "text", "text":
_truncate_result_text(serialize_for_llm(result.content))}
+ ]
+ }
+ if result.is_error:
+ payload["is_error"] = True
+ return payload
+
+ return handle_tool_call
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
index 81eee0ca318..910892bf9d1 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
@@ -28,7 +28,9 @@ if TYPE_CHECKING:
from collections.abc import Iterator
#: The agent framework a tool call came through, as the ``framework`` tag
reports it.
-Framework = Literal["pydantic_ai", "strands", "adk", "langchain", "anthropic",
"openai_agents", "none"]
+Framework = Literal[
+ "pydantic_ai", "strands", "adk", "langchain", "anthropic",
"openai_agents", "claude_agent_sdk", "none"
+]
# Set by each framework adapter around the calls it makes.
_framework: ContextVar[Framework | None] =
ContextVar("common_ai_tool_framework", default=None)
diff --git
a/providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py
b/providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py
new file mode 100644
index 00000000000..8788f02cb07
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py
@@ -0,0 +1,1060 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import copy
+import importlib
+import json
+import sys
+from collections.abc import AsyncIterator
+from typing import Any
+from unittest import mock
+
+import pytest
+
+pytest.importorskip("claude_agent_sdk")
+
+from claude_agent_sdk import AssistantMessage, Message, ResultMessage,
TextBlock, UserMessage
+from pydantic_ai.exceptions import ModelRetry
+
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
ToolResult, collect_tools
+from airflow.providers.common.ai.tools.claude_agent_sdk import AirflowTools
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+from unit.common.ai.tools.test__from_toolset import _scripted
+from unit.common.ai.toolsets.test_sql import _make_mock_db_hook
+
+MODULE = "airflow.providers.common.ai.tools.claude_agent_sdk"
+
+
+def _tool(outcome: ToolResult | Exception, ran: list[dict[str, Any]] | None =
None) -> AirflowTool:
+ """A ``lookup`` tool that returns or raises ``outcome``, recording its
calls in ``ran``."""
+
+ async def function(arguments: dict[str, Any]) -> ToolResult:
+ if ran is not None:
+ ran.append(arguments)
+ if isinstance(outcome, Exception):
+ raise outcome
+ return outcome
+
+ return AirflowTool(
+ name="lookup",
+ description="Look something up.",
+ parameters={"type": "object", "properties": {"key": {"type":
"string"}}, "required": ["key"]},
+ function=function,
+ )
+
+
+async def _pong(arguments: dict[str, Any]) -> ToolResult:
+ return ToolResult("pong")
+
+
+def _build_tools(*sources: Any, server_name: str = "airflow") ->
tuple[AirflowTools, dict[str, Any]]:
+ """
+ Build ``AirflowTools``, with ``create_sdk_mcp_server`` mocked to capture
its SDK tools.
+
+ The ``@tool`` decorator itself is not mocked: it is only a plain wrapper
+ (``SdkMcpTool(handler=handler, ...)``), so the handlers it returns run the
same code
+ they would through the real MCP server. Returns the instance and its
handlers by name.
+ """
+ captured: dict[str, Any] = {}
+
+ def _capture_server(*, name: str, version: str = "1.0.0", tools: Any =
None) -> Any:
+ captured["sdk_tools"] = tools
+ return mock.sentinel.server
+
+ with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as
mock_create_server:
+ mock_create_server.side_effect = _capture_server
+ tools = AirflowTools(*sources, server_name=server_name)
+ handlers = {sdk_tool.name: sdk_tool.handler for sdk_tool in
captured["sdk_tools"]}
+ return tools, handlers
+
+
+def _assistant_message(message_id: str | None = None) -> AssistantMessage:
+ return AssistantMessage(
+ content=[TextBlock(text="")], model="claude-test",
stop_reason="tool_use", message_id=message_id
+ )
+
+
+def _user_message() -> UserMessage:
+ return UserMessage(content="")
+
+
+def _result_message(**overrides: Any) -> ResultMessage:
+ fields: dict[str, Any] = {
+ "subtype": "success",
+ "duration_ms": 1,
+ "duration_api_ms": 1,
+ "is_error": False,
+ "num_turns": 1,
+ "session_id": "session-1",
+ }
+ fields.update(overrides)
+ return ResultMessage(**fields)
+
+
+async def _fake_cli(
+ handlers: dict[str, Any],
+ turns: list[list[tuple[str, dict[str, Any]]]],
+ *,
+ final: ResultMessage | None = None,
+ message_ids: list[str | None] | None = None,
+) -> AsyncIterator[Message]:
+ """
+ Stand in for the CLI driving the tools ``AirflowTools`` served through one
or more turns.
+
+ For each turn, yields an ``AssistantMessage`` with a ``tool_use`` stop
reason, then
+ awaits every one of that turn's calls -- as the CLI waits for a turn's
tool results
+ before sending the next request -- before yielding the ``UserMessage``
carrying them.
+ Ends with a ``ResultMessage``. ``message_ids`` gives each turn's
``AssistantMessage.message_id``.
+ """
+ for index, turn in enumerate(turns):
+ yield _assistant_message(message_ids[index] if message_ids else None)
+ for tool_name, args in turn:
+ await handlers[tool_name](args)
+ yield _user_message()
+ yield final if final is not None else _result_message()
+
+
+async def _call_real_server(tools: AirflowTools, name: str, args: dict[str,
Any]) -> dict[str, Any]:
+ """Call a tool through the real in-process MCP server, with the JSON-RPC
the SDK feeds its bridge."""
+ return await _request_real_server(tools, "tools/call", {"name": name,
"arguments": args})
+
+
+async def _request_real_server(tools: AirflowTools, method: str, params:
dict[str, Any]) -> dict[str, Any]:
+ """Send one JSON-RPC request to the real in-process MCP server and return
its ``result``."""
+ moved = (
+ "claude_agent_sdk's private path _internal.sdk_mcp_bridge.SdkMcpBridge
has changed; "
+ "update _request_real_server in this test module."
+ )
+ try:
+ bridge_module =
importlib.import_module("claude_agent_sdk._internal.sdk_mcp_bridge")
+ except ImportError:
+ pytest.fail(moved)
+ bridge_cls = getattr(bridge_module, "SdkMcpBridge", None)
+ if bridge_cls is None:
+ pytest.fail(moved)
+ bridge = bridge_cls("airflow", tools.server["instance"])
+ try:
+ await bridge.handle(
+ {
+ "jsonrpc": "2.0",
+ "id": 0,
+ "method": "initialize",
+ "params": {
+ "protocolVersion": "2025-06-18",
+ "capabilities": {},
+ "clientInfo": {"name": "test", "version": "1"},
+ },
+ }
+ )
+ await bridge.handle({"jsonrpc": "2.0", "method":
"notifications/initialized"})
+ reply = await bridge.handle(
+ {
+ "jsonrpc": "2.0",
+ "id": 1,
+ "method": method,
+ "params": params,
+ }
+ )
+ return reply["result"]
+ finally:
+ await bridge.aclose()
+
+
+class TestTools:
+ def test_each_toolset_tool_becomes_an_mcp_tool_with_its_schema(self):
+ tools, handlers = _build_tools(_tool(ToolResult("ok")),
SQLToolset("pg_default"))
+
+ assert tools.allowed_tools == [
+ "mcp__airflow__lookup",
+ "mcp__airflow__list_tables",
+ "mcp__airflow__get_schema",
+ "mcp__airflow__query",
+ "mcp__airflow__check_query",
+ ]
+ assert set(handlers) == {"lookup", "list_tables", "get_schema",
"query", "check_query"}
+
+ def test_allowed_tools_use_a_custom_server_name(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")),
server_name="warehouse")
+
+ assert tools.allowed_tools == ["mcp__warehouse__lookup"]
+
+ def test_an_mcp_tool_does_not_share_the_source_schema(self):
+ source = _tool(ToolResult("ok"))
+ before = copy.deepcopy(source.parameters)
+ captured: dict[str, Any] = {}
+
+ def _capture_server(*, name: str, version: str = "1.0.0", tools: Any =
None) -> Any:
+ captured["sdk_tools"] = tools
+ return mock.sentinel.server
+
+ with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as
mock_create_server:
+ mock_create_server.side_effect = _capture_server
+ AirflowTools(source)
+
+ captured["sdk_tools"][0].input_schema["properties"]["key"]["type"] =
"integer"
+ assert source.parameters == before
+
+ def test_every_mcp_tool_keeps_a_large_result_inline(self):
+ captured: dict[str, Any] = {}
+ sources = [_tool(ToolResult("ok")), SQLToolset("pg_default")]
+
+ def _capture_server(*, name: str, version: str = "1.0.0", tools: Any =
None) -> Any:
+ captured["sdk_tools"] = tools
+ return mock.sentinel.server
+
+ with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as
mock_create_server:
+ mock_create_server.side_effect = _capture_server
+ AirflowTools(*sources)
+
+ assert [t.name for t in captured["sdk_tools"]] == [t.name for t in
collect_tools(sources)]
+ for sdk_tool in captured["sdk_tools"]:
+ assert sdk_tool.annotations.maxResultSizeChars == 500_000
+
+ @pytest.mark.asyncio
+ async def test_a_schema_without_properties_takes_no_arguments(self):
+ source = AirflowTool(
+ name="ping",
+ description="Ping.",
+ parameters={"type": "object"},
+ function=mock.create_autospec(_pong,
return_value=ToolResult("pong")),
+ )
+ tools = AirflowTools(source)
+
+ listed = await _request_real_server(tools, "tools/list", {})
+
+ assert listed["tools"][0]["inputSchema"].get("required", []) == []
+ assert "type" not in listed["tools"][0]["inputSchema"]["properties"]
+ assert source.parameters == {"type": "object"}
+
+ @pytest.mark.asyncio
+ async def
test_the_cli_lists_each_tool_with_its_description_schema_and_result_limit(self):
+ schema = {
+ "type": "object",
+ "properties": {"key": {"type": "string", "description": "The
key."}},
+ "required": ["key"],
+ }
+ source = AirflowTool(
+ name="lookup",
+ description="Look a value up.",
+ parameters=schema,
+ function=mock.create_autospec(_pong,
return_value=ToolResult("pong")),
+ )
+ tools = AirflowTools(source)
+
+ listed = await _request_real_server(tools, "tools/list", {})
+
+ (listed_tool,) = listed["tools"]
+ assert listed_tool["name"] == "lookup"
+ assert listed_tool["description"] == "Look a value up."
+ assert listed_tool["inputSchema"] == schema
+ assert listed_tool["_meta"] == {"anthropic/maxResultSizeChars":
500_000}
+
+
+class TestHandler:
+ @pytest.mark.asyncio
+ async def test_the_model_reads_a_toolset_result(self):
+ ts = SQLToolset("pg_default")
+ ts._hook = _make_mock_db_hook(records=[(1, "Ada")],
last_description=[("id",), ("name",)])
+ tools, handlers = _build_tools(ts)
+
+ payload = await handlers["query"]({"sql": "SELECT id, name FROM
users"})
+
+ assert "is_error" not in payload
+ assert json.loads(payload["content"][0]["text"])["rows"] == [[1,
"Ada"]]
+
+ @pytest.mark.asyncio
+ async def test_a_correctable_failure_reaches_the_model_as_an_error(self):
+ tools, handlers = _build_tools(_tool(ToolResult("Unknown column
'nme'", is_error=True)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert payload["is_error"] is True
+ assert payload["content"] == [{"type": "text", "text": "Unknown column
'nme'"}]
+
+ @pytest.mark.asyncio
+ async def test_a_structured_result_reaches_the_model_as_json(self):
+ tools, handlers = _build_tools(_tool(ToolResult({"rows": [[1,
None]]})))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert json.loads(payload["content"][0]["text"]) == {"rows": [[1,
None]]}
+
+ @pytest.mark.asyncio
+ async def
test_not_driven_by_run_hands_the_failure_to_the_model_with_a_warning(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ with mock.patch(f"{MODULE}.log", autospec=True) as mock_log:
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert payload["is_error"] is True
+ assert "PermissionError: role cannot read orders" in
payload["content"][0]["text"]
+ mock_log.warning.assert_called_once()
+ assert "not driven by AirflowTools.run" in
mock_log.warning.call_args.args[0]
+
+ @pytest.mark.asyncio
+ async def test_the_toolsets_calls_are_counted_as_claude_agent_sdk(self):
+ ts = SQLToolset("pg_default")
+ ts._hook = _make_mock_db_hook()
+ tools, handlers = _build_tools(ts)
+
+ with mock.patch(
+ "airflow.providers.common.ai.utils.tool_metrics.Stats",
mock.MagicMock(spec=["incr"])
+ ) as stats:
+ await handlers["list_tables"]({})
+
+ assert stats.incr.call_args.kwargs["tags"]["framework"] ==
"claude_agent_sdk"
+
+ @pytest.mark.enable_redact
+ @pytest.mark.asyncio
+ async def
test_a_value_mask_secrets_cannot_descend_into_is_masked_after_rendering(
+ self, registered_secret
+ ):
+ """
+ ``mask_secrets`` leaves an opaque object (not a dict, list or
dataclass) untouched, so
+ only masking *after* it is turned into text -- what
``serialize_for_llm`` does --
+ catches a secret in its ``repr``. ``json.dumps(value, default=str)``
masks first and
+ renders after, so it would not catch this.
+ """
+
+ class _Row:
+ def __init__(self, secret: str) -> None:
+ self._secret = secret
+
+ def __repr__(self) -> str:
+ return f"_Row(secret={self._secret!r})"
+
+ tools, handlers =
_build_tools(_tool(ToolResult(_Row(f"key={registered_secret}"))))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ text = payload["content"][0]["text"]
+ assert registered_secret not in text
+ assert "key=***" in text
+
+ @pytest.mark.enable_redact
+ @pytest.mark.asyncio
+ async def test_a_secret_in_a_plain_string_result_is_masked(self,
registered_secret):
+ tools, handlers =
_build_tools(_tool(ToolResult(f"key={registered_secret}")))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert registered_secret not in payload["content"][0]["text"]
+ assert payload["content"][0]["text"] == "key=***"
+
+
+class TestRun:
+ @pytest.mark.asyncio
+ async def test_run_returns_the_result_message(self):
+ tools, handlers = _build_tools(_tool(ToolResult("ok")))
+ final = _result_message(result="done")
+
+ result = await tools.run(_fake_cli(handlers, [[("lookup", {"key":
"a"})]], final=final))
+
+ assert result is final
+
+ @pytest.mark.asyncio
+ async def test_run_raises_without_a_result_message(self):
+ tools, handlers = _build_tools()
+
+ async def _no_result() -> AsyncIterator[Message]:
+ return
+ yield # pragma: no cover - makes this an async generator function
+
+ with pytest.raises(RuntimeError, match="ended without a
ResultMessage"):
+ await tools.run(_no_result())
+
+ @pytest.mark.asyncio
+ async def
test_a_failure_before_the_stream_ends_without_a_result_message_still_raises_it(self):
+ """
+ The CLI can end the stream right after a failed tool call without
sending another
+ message at all -- no further ``UserMessage``, no ``ResultMessage``.
The in-loop check
+ (see the test above) never runs again once the stream is exhausted, so
the failure
+ has to be checked once more after the loop, before falling back to the
generic
+ "ended without a ResultMessage" error.
+ """
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ async def _cli_ends_right_after_the_failure() ->
AsyncIterator[Message]:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ # The stream ends here: no further message, no ResultMessage.
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_cli_ends_right_after_the_failure())
+
+ @pytest.mark.asyncio
+ async def test_any_other_failure_fails_the_run_at_the_next_message(self):
+ """
+ The handler answers the model within the turn -- ``ran`` has one call,
and the model
+ did read the masked failure -- but ``run`` still raises at the next
message rather
+ than letting the stream reach a ``ResultMessage``.
+ """
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]]))
+
+ assert len(ran) == 1
+
+ @pytest.mark.asyncio
+ async def
test_a_failure_in_one_turn_stops_the_stream_before_the_next_turn_starts(self):
+ ran: list[dict[str, Any]] = []
+ turn_two_started: list[bool] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+
+ async def _two_turns() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _user_message()
+ turn_two_started.append(True)
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "b"})
+ yield _user_message()
+ yield _result_message()
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_two_turns())
+
+ assert turn_two_started == []
+ assert len(ran) == 1
+
+ @pytest.mark.asyncio
+ async def test_calls_after_a_failure_in_the_same_turn_do_not_run(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+
+ async def _three_calls_one_turn() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ for _ in range(3):
+ await handlers["lookup"]({"key": "a"})
+ yield _user_message()
+ yield _result_message()
+
+ with pytest.raises(ToolCallError, match="PermissionError"):
+ await tools.run(_three_calls_one_turn())
+
+ assert len(ran) == 1
+
+ @pytest.mark.asyncio
+ async def
test_an_error_result_message_raises_instead_of_being_returned(self):
+ tools, _ = _build_tools()
+
+ async def _ends_at_the_result() -> AsyncIterator[Message]:
+ yield _result_message(
+ is_error=True,
+ subtype="error_max_turns",
+ result="partial answer",
+ errors=["hit the turn limit"],
+ num_turns=3,
+ )
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_ends_at_the_result())
+
+ message = str(exc_info.value)
+ assert "error_max_turns" in message
+ assert "partial answer" in message
+ assert "hit the turn limit" in message
+
+ @pytest.mark.asyncio
+ async def
test_an_error_turn_followed_by_a_successful_turn_still_raises(self):
+ tools, _ = _build_tools()
+
+ async def _error_then_success() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True,
subtype="error_during_execution", result="first turn failed")
+ yield _result_message()
+
+ with pytest.raises(RuntimeError, match="first turn failed"):
+ await tools.run(_error_then_success())
+
+ @pytest.mark.asyncio
+ async def
test_the_first_error_turn_is_reported_when_several_turns_fail(self):
+ tools, _ = _build_tools()
+
+ async def _two_errors_then_success() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True,
subtype="error_during_execution", result="first turn failed")
+ yield _result_message(
+ is_error=True, subtype="error_during_execution",
result="second turn failed"
+ )
+ yield _result_message()
+
+ with pytest.raises(RuntimeError, match="first turn failed") as
exc_info:
+ await tools.run(_two_errors_then_success())
+
+ assert "second turn failed" not in str(exc_info.value)
+
+ @pytest.mark.asyncio
+ async def test_an_error_result_message_bounds_its_free_text_fields(self):
+ tools, _ = _build_tools()
+
+ async def _runaway() -> AsyncIterator[Message]:
+ yield _result_message(
+ is_error=True,
+ subtype="error_max_turns",
+ result="x" * 100_000,
+ errors=[f"error-{i}" for i in range(1_000)],
+ )
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_runaway())
+
+ message = str(exc_info.value)
+ assert len(message) < 2_000
+ assert "error_max_turns" in message
+ assert "(+99502 chars)" in message
+ assert "error-4" in message
+ assert "error-5" not in message
+ assert "(+995 more)" in message
+
+ @pytest.mark.asyncio
+ async def
test_an_api_failure_reported_as_a_success_subtype_still_raises(self):
+ tools, _ = _build_tools()
+
+ async def _api_failure() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True, subtype="success",
api_error_status=529)
+
+ with pytest.raises(RuntimeError, match="529"):
+ await tools.run(_api_failure())
+
+ @pytest.mark.asyncio
+ async def
test_an_error_result_message_leaves_unset_fields_out_of_the_message(self):
+ tools, _ = _build_tools()
+
+ async def _bare_error() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True,
subtype="error_during_execution")
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_bare_error())
+
+ message = str(exc_info.value)
+ assert "error_during_execution" in message
+ for unset in ("result=", "errors=", "stop_reason=",
"api_error_status=", "terminal_reason="):
+ assert unset not in message
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_failure_takes_priority_over_an_error_result_message(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ async def _failure_then_error_result() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _result_message(is_error=True, subtype="error_max_turns")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_failure_then_error_result())
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_that_keeps_failing_fails_the_run_once_its_retries_are_spent(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {})], [("step",
{})]]))
+
+ @pytest.mark.asyncio
+ async def
test_each_run_of_a_reused_instance_starts_with_a_fresh_retry_budget(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ for _ in range(2):
+ result = await tools.run(_fake_cli(handlers, [[("step", {})]]))
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_one_model_turn_count_as_one(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ result = await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=["msg_1"]))
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_different_model_turns_count_separately(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(
+ _fake_cli(handlers, [[("step", {})], [("step", {})]],
message_ids=["msg_1", "msg_2"])
+ )
+
+ @pytest.mark.asyncio
+ async def
test_sequential_failures_count_separately_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=[None]))
+
+ @pytest.mark.asyncio
+ async def
test_simultaneous_failures_count_as_one_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ async def _two_calls_at_once() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await asyncio.gather(handlers["step"]({}), handlers["step"]({}))
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_calls_at_once())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_real_server_path_sees_the_turn_the_run_recorded(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools = AirflowTools(step)
+
+ async def _two_failures_in_one_turn() -> AsyncIterator[Message]:
+ yield _assistant_message("msg_1")
+ await _call_real_server(tools, "step", {})
+ await _call_real_server(tools, "step", {})
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_failures_in_one_turn())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_stream_is_closed_while_the_run_is_still_active(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+ closed_during_run: list[bool] = []
+
+ async def _failing_cli() -> AsyncIterator[Message]:
+ try:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _user_message()
+ yield _result_message()
+ finally:
+ closed_during_run.append(tools._run is not None)
+
+ stream = _failing_cli()
+ with pytest.raises(ToolCallError):
+ await tools.run(stream)
+
+ assert closed_during_run == [True]
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_call_after_a_successful_run_has_ended_does_not_run(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(ToolResult("ok"), ran))
+ await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]]))
+ assert len(ran) == 1
+
+ payload = await handlers["lookup"]({"key": "b"})
+
+ assert payload["is_error"] is True
+ assert payload["content"][0]["text"].startswith("Not run")
+ assert "failed" not in payload["content"][0]["text"]
+ assert len(ran) == 1
+
+ @pytest.mark.asyncio
+ async def test_a_call_after_a_failed_run_has_ended_does_not_run(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+ with pytest.raises(ToolCallError):
+ await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]]))
+
+ payload = await handlers["lookup"]({"key": "b"})
+
+ assert payload["is_error"] is True
+ assert payload["content"][0]["text"] == "Not run: an earlier Airflow
tool call in this run failed."
+ assert len(ran) == 1
+
+ @pytest.mark.asyncio
+ async def test_a_new_run_is_not_blocked_by_the_previous_ended_run(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(ToolResult("ok"), ran))
+ await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]]))
+
+ await tools.run(_fake_cli(handlers, [[("lookup", {"key": "b"})]]))
+
+ assert ran == [{"key": "a"}, {"key": "b"}]
+
+ @pytest.mark.asyncio
+ async def test_a_call_landing_after_the_stream_is_closed_is_refused(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(ToolResult("ok"), ran))
+ late_calls: list[asyncio.Task[dict[str, Any]]] = []
+
+ async def _cli_whose_teardown_calls_a_tool_later() ->
AsyncIterator[Message]:
+ try:
+ yield _assistant_message()
+ yield _result_message()
+ finally:
+
late_calls.append(asyncio.create_task(handlers["lookup"]({"key": "late"})))
+
+ await tools.run(_cli_whose_teardown_calls_a_tool_later())
+ assert tools._run is None
+
+ payload = await late_calls[0]
+
+ assert payload["is_error"] is True
+ assert payload["content"][0]["text"].startswith("Not run")
+ assert ran == []
+
+ @pytest.mark.asyncio
+ async def test_a_stream_without_aclose_is_driven_to_the_end(self):
+ class _PlainStream:
+ def __init__(self, messages: list[Message]) -> None:
+ self._messages = messages
+
+ def __aiter__(self) -> AsyncIterator[Message]:
+ async def _iterate() -> AsyncIterator[Message]:
+ for message in self._messages:
+ yield message
+
+ return _iterate()
+
+ final = _result_message(result="done")
+
+ result = await
AirflowTools(_tool(ToolResult("ok"))).run(_PlainStream([_assistant_message(),
final]))
+
+ assert result is final
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_does_not_hide_the_tools_failure(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_after_a_clean_run_raises(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli({}, [])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(OSError, match="close failed"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.log", autospec=True)
+ async def
test_a_failing_aclose_is_logged_when_it_does_not_hide_the_tools_failure(self,
mock_log):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError):
+ await tools.run(_StubbornStream())
+
+ mock_log.warning.assert_called_once()
+ assert "Closing the Claude Agent SDK stream failed" in
mock_log.warning.call_args.args[0]
+ assert str(mock_log.warning.call_args.args[1]) == "close failed"
+
+ @pytest.mark.asyncio
+ async def test_a_cancelled_run_closes_the_stream_and_stays_cancelled(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+ started = asyncio.Event()
+ aclose_calls: list[bool] = []
+
+ class _BlockedStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return self._iterate()
+
+ async def _iterate(self) -> AsyncIterator[Message]:
+ yield _assistant_message()
+ started.set()
+ await asyncio.Event().wait()
+ yield _result_message()
+
+ async def aclose(self) -> None:
+ aclose_calls.append(True)
+
+ task = asyncio.ensure_future(tools.run(_BlockedStream()))
+ await started.wait()
+ task.cancel()
+
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+ assert aclose_calls == [True]
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_concurrent_runs_on_one_instance_are_rejected(self):
+ tools, handlers = _build_tools(_tool(ToolResult("ok")))
+ gate = asyncio.Event()
+
+ async def _blocked_cli() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await gate.wait()
+ yield _result_message()
+
+ first = asyncio.ensure_future(tools.run(_blocked_cli()))
+ await asyncio.sleep(0)
+
+ with pytest.raises(RuntimeError, match="already running"):
+ await tools.run(_fake_cli(handlers, []))
+
+ gate.set()
+ await first
+
+
+class TestRealMcpServer:
+ @pytest.mark.asyncio
+ async def
test_a_tool_result_marked_as_an_error_reaches_the_model_as_isError(self):
+ tools = AirflowTools(_tool(ToolResult("not found", is_error=True)))
+
+ result = await _call_real_server(tools, "lookup", {"key": "a"})
+
+ assert result["isError"] is True
+ assert "not found" in result["content"][0]["text"]
+
+ @pytest.mark.asyncio
+ async def test_a_successful_tool_result_is_not_marked_isError(self):
+ tools = AirflowTools(_tool(ToolResult("ok")))
+
+ result = await _call_real_server(tools, "lookup", {"key": "a"})
+
+ assert result["isError"] is False
+
+ @pytest.mark.asyncio
+ async def
test_the_sdk_returns_an_airflow_tool_failure_to_the_model_instead_of_raising(self):
+ tools = AirflowTools(_tool(PermissionError("role cannot read orders")))
+
+ result = await _call_real_server(tools, "lookup", {"key": "a"})
+
+ assert result["isError"] is True
+ assert "PermissionError: role cannot read orders" in
result["content"][0]["text"]
+
+ @pytest.mark.asyncio
+ async def test_the_sdk_catches_an_exception_escaping_the_handler(self):
+ tools = AirflowTools(_tool(ToolResult("ok")))
+
+ with mock.patch.object(AirflowTool, "call", autospec=True,
side_effect=RuntimeError("boom")):
+ result = await _call_real_server(tools, "lookup", {"key": "a"})
+
+ assert result["isError"] is True
+ assert "boom" in result["content"][0]["text"]
+
+
+class TestResultTruncation:
+ @pytest.mark.asyncio
+ async def test_a_result_over_the_ceiling_is_cut_to_fit_with_a_notice(self):
+ _, handlers = _build_tools(_tool(ToolResult("x" * 1_200_000)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ text = payload["content"][0]["text"]
+ # Booleans keep a failing assertion from rendering a 500k-character
diff.
+ assert len(text) <= 500_000
+ assert ("Truncated: this result had 1200000 characters" in text) is
True
+ assert text.startswith("x" * 1000) is True
+ assert "is_error" not in payload
+
+ @pytest.mark.asyncio
+ async def test_a_result_at_the_ceiling_is_returned_as_is(self):
+ _, handlers = _build_tools(_tool(ToolResult("x" * 500_000)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ text = payload["content"][0]["text"]
+ assert len(text) == 500_000
+ assert ("Truncated" in text) is False
+
+ @pytest.mark.asyncio
+ async def test_a_result_one_over_the_ceiling_is_cut(self):
+ _, handlers = _build_tools(_tool(ToolResult("x" * 500_001)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ text = payload["content"][0]["text"]
+ assert len(text) <= 500_000
+ assert ("Truncated" in text) is True
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("text", "is_cut"),
+ [
+ pytest.param("\N{GRINNING FACE}" * 250_000, False,
id="at-the-ceiling-in-utf16-units"),
+ pytest.param("\N{GRINNING FACE}" * 250_001, True,
id="one-emoji-over-the-ceiling"),
+ pytest.param("\N{GRINNING FACE}" * 300_000, True,
id="under-the-ceiling-in-code-points"),
+ ],
+ )
+ async def test_the_ceiling_is_measured_in_utf16_units(self, text, is_cut):
+ _, handlers = _build_tools(_tool(ToolResult(text)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ result = payload["content"][0]["text"]
+ assert (("Truncated" in result) is is_cut) is True
+ assert (len(result.encode("utf-16-le")) // 2 <= 500_000) is True
+
+ @pytest.mark.asyncio
+ async def test_a_cut_never_splits_a_surrogate_pair(self):
+ # The leading "a" shifts every emoji to an odd UTF-16 offset, so the
cut falls between the
+ # two halves of a pair unless it is moved back to a code point
boundary.
+ _, handlers = _build_tools(_tool(ToolResult("a" + "\N{GRINNING FACE}"
* 300_000)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ result = payload["content"][0]["text"]
+ kept = result[: result.index("\n\n[Truncated")]
+ # Comparing the whole string would make pytest diff ~250k characters
on failure.
+ assert kept[0] == "a"
+ assert set(kept[1:]) == {"\N{GRINNING FACE}"}
+ assert (len(result.encode("utf-16-le")) // 2 <= 500_000) is True
+
+ @pytest.mark.asyncio
+ async def test_a_small_result_is_returned_as_is(self):
+ _, handlers = _build_tools(_tool(ToolResult("short")))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert payload == {"content": [{"type": "text", "text": "short"}]}
+
+ @pytest.mark.asyncio
+ async def test_a_truncated_error_result_stays_an_error(self):
+ _, handlers = _build_tools(_tool(ToolResult("x" * 600_000,
is_error=True)))
+
+ payload = await handlers["lookup"]({"key": "a"})
+
+ assert len(payload["content"][0]["text"]) <= 500_000
+ assert payload["is_error"] is True
+
+
+async def _decide_tool_use(tool_name: str, input_data: dict[str, Any],
context: Any) -> Any:
+ """Shape of a ``can_use_tool`` callback, for ``create_autospec``."""
+
+
+class TestOptions:
+ def test_safe_defaults(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options()
+
+ assert options.tools == []
+ assert options.setting_sources == []
+ assert options.strict_mcp_config is True
+ assert options.mcp_servers == {"airflow": tools.server}
+ assert options.allowed_tools == tools.allowed_tools
+ assert options.verbatim_prompts is True
+ assert options.permission_mode is None
+ assert options.max_buffer_size == 16 * 1024 * 1024
+
+ def test_explicit_max_buffer_size_is_kept(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options(max_buffer_size=1024)
+
+ assert options.max_buffer_size == 1024
+
+ def test_explicit_kwargs_override_the_defaults(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options(
+ tools=["Bash"], setting_sources=["project"],
strict_mcp_config=False, verbatim_prompts=False
+ )
+
+ assert options.verbatim_prompts is False
+ assert options.tools == ["Bash"]
+ assert options.setting_sources == ["project"]
+ assert options.strict_mcp_config is False
+
+ def test_other_kwargs_pass_through(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options(model="claude-sonnet-5", max_turns=10)
+
+ assert options.model == "claude-sonnet-5"
+ assert options.max_turns == 10
+
+ def test_mcp_servers_and_allowed_tools_merge_with_the_callers(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+ other_server = mock.sentinel.other_server
+
+ options = tools.build_options(mcp_servers={"other": other_server},
allowed_tools=["other_tool"])
+
+ assert options.mcp_servers == {"other": other_server, "airflow":
tools.server}
+ assert options.allowed_tools == ["other_tool", *tools.allowed_tools]
+
+ def
test_a_can_use_tool_callback_keeps_the_airflow_tools_out_of_allowed_tools(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+ can_use_tool = mock.create_autospec(_decide_tool_use)
+
+ options = tools.build_options(can_use_tool=can_use_tool,
allowed_tools=["other_tool"])
+
+ assert options.allowed_tools == ["other_tool"]
+ assert options.can_use_tool is can_use_tool
+ assert options.permission_mode == "default"
+ assert tools.allowed_tools == ["mcp__airflow__lookup"]
+
+ def
test_a_can_use_tool_callback_does_not_override_an_explicit_permission_mode(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options(
+ can_use_tool=mock.create_autospec(_decide_tool_use),
permission_mode="acceptEdits"
+ )
+
+ assert options.permission_mode == "acceptEdits"
+
+ def test_a_none_can_use_tool_is_the_same_as_not_passing_one(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ options = tools.build_options(can_use_tool=None)
+
+ assert options.allowed_tools == tools.allowed_tools
+ assert options.permission_mode is None
+
+ def test_a_colliding_server_name_raises(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ with pytest.raises(ValueError, match="already has a server named
'airflow'"):
+ tools.build_options(mcp_servers={"airflow":
mock.sentinel.other_server})
+
+
+def test_missing_sdk_raises_optional_feature_exception() -> None:
+ with mock.patch.dict(sys.modules, {"claude_agent_sdk": None}):
+ sys.modules.pop(MODULE, None)
+ try:
+ with pytest.raises(AirflowOptionalProviderFeatureException):
+ importlib.import_module(MODULE)
+ finally:
+ sys.modules.pop(MODULE, None)
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 0ede098457c..4a70cc9eb1b 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -511,8 +511,8 @@ metrics:
- name: "common_ai.tool_calls"
description: "Number of calls to a toolset of the common.ai provider.
Metric with toolset
(the toolset class), framework (``pydantic_ai``, ``strands``, ``adk``,
``langchain``,
- ``anthropic``, ``openai_agents`` or ``none``) and outcome (``executed``,
``failed`` or
- ``replayed``) tagging."
+ ``anthropic``, ``openai_agents``, ``claude_agent_sdk`` or ``none``) and
outcome (``executed``,
+ ``failed`` or ``replayed``) tagging."
type: "counter"
legacy_name: "-"
name_variables: ["toolset", "framework", "outcome"]
diff --git a/uv.lock b/uv.lock
index 1e9303eb957..754e3d378b0 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4635,6 +4635,9 @@ bedrock = [
boat = [
{ name = "boat-sdk" },
]
+claude-agent-sdk = [
+ { name = "claude-agent-sdk" },
+]
code-mode = [
{ name = "pydantic-ai-harness", extra = ["codemode"] },
]
@@ -4714,6 +4717,7 @@ dev = [
{ name = "apache-airflow-providers-standard" },
{ name = "apache-airflow-task-sdk" },
{ name = "boat-sdk" },
+ { name = "claude-agent-sdk" },
{ name = "grpcio" },
{ name = "islo" },
{ name = "langchain" },
@@ -4745,6 +4749,7 @@ requires-dist = [
{ name = "apache-airflow-providers-modal", marker = "extra == 'modal'",
editable = "providers/modal" },
{ name = "apache-airflow-providers-standard", editable =
"providers/standard" },
{ name = "boat-sdk", marker = "extra == 'boat'", specifier = ">=1.9.3" },
+ { name = "claude-agent-sdk", marker = "extra == 'claude-agent-sdk'",
specifier = ">=0.2.160" },
{ name = "dataclasses-json", marker = "extra == 'llamaindex'", specifier =
">=0.6.7" },
{ name = "fastavro", marker = "python_full_version >= '3.14' and extra ==
'avro'", specifier = ">=1.12.1" },
{ name = "fastavro", marker = "python_full_version < '3.14' and extra ==
'avro'", specifier = ">=1.10.0" },
@@ -4776,7 +4781,7 @@ requires-dist = [
{ name = "structlog", specifier = ">=24.2.0" },
{ name = "typesafe-sdk", marker = "extra == 'typesafe'", specifier =
">=0.6.0" },
]
-provides-extras = ["anthropic", "bedrock", "google", "openai", "typesafe",
"mcp", "modal", "opensandbox", "boat", "openshell", "code-mode", "shields",
"skills", "islo", "avro", "parquet", "sql", "common-sql", "langchain",
"llamaindex", "pdf", "docx", "git"]
+provides-extras = ["anthropic", "bedrock", "google", "openai", "typesafe",
"mcp", "modal", "opensandbox", "boat", "openshell", "code-mode", "shields",
"skills", "islo", "avro", "parquet", "sql", "common-sql", "langchain",
"llamaindex", "pdf", "docx", "git", "claude-agent-sdk"]
[package.metadata.requires-dev]
dev = [
@@ -4791,6 +4796,7 @@ dev = [
{ name = "apache-airflow-providers-standard", editable =
"providers/standard" },
{ name = "apache-airflow-task-sdk", editable = "task-sdk" },
{ name = "boat-sdk", specifier = ">=1.9.3" },
+ { name = "claude-agent-sdk", specifier = ">=0.2.160" },
{ name = "grpcio", specifier = ">=1.78.0" },
{ name = "islo", specifier = ">=0.3.19" },
{ name = "langchain", specifier = ">=1.0.0" },
@@ -11606,6 +11612,25 @@ wheels = [
{ url =
"https://files.pythonhosted.org/packages/22/fc/976d9c4b79e28cbda95b1acf574b00f811d9aec0fce55b63d573d6fa446b/ciso8601-2.3.3-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl";,
hash =
"sha256:58799673ffdf621fe138fb8af6a89daf4ddefdf7ca4a10777ad8d55f3f171b6e", size
= 18284, upload-time = "2025-08-20T16:31:28.43Z" },
]
+[[package]]
+name = "claude-agent-sdk"
+version = "0.2.168"
+source = { registry = "https://pypi.org/simple"; }
+dependencies = [
+ { name = "anyio" },
+ { name = "jsonschema" },
+ { name = "mcp" },
+ { name = "sniffio" },
+]
+sdist = { url =
"https://files.pythonhosted.org/packages/9a/ee/31e159daf956f0c5503eac67d02743da692c0999fef046cdde3cc52702b9/claude_agent_sdk-0.2.168.tar.gz";,
hash =
"sha256:80388054a09c137d76fc29caa6e46bc90d2ae2ef09fac5ee1c38fc36d2762fda", size
= 371591, upload-time = "2026-10-15T20:20:43.683Z" }
+wheels = [
+ { url =
"https://files.pythonhosted.org/packages/15/31/9a205fc0b7ae1b66d5964774bb534ab378047c9bf87578f443bd1fb2205e/claude_agent_sdk-0.2.168-py3-none-macosx_11_0_arm64.whl";,
hash =
"sha256:2ef1e7fc442c5abdc1e698027bf1b5faa8ea2130134a14364e2ae3c53bf31fdd", size
= 104355641, upload-time = "2026-10-15T20:20:48.105Z" },
+ { url =
"https://files.pythonhosted.org/packages/9f/ba/12bbe92d5e9145c7aaa94aa349a0859b77a9b58ccc40ae9e69cf7f078e5b/claude_agent_sdk-0.2.168-py3-none-macosx_11_0_x86_64.whl";,
hash =
"sha256:653fc7bbe56630e6ddd9601a7cfe50586667e74102da3fb90571456c63655a62", size
= 109438860, upload-time = "2026-10-15T20:20:52.274Z" },
+ { url =
"https://files.pythonhosted.org/packages/ad/5a/26e587846f14f67f3353f96811b86931263a9c7e264d17a691d9577e4da5/claude_agent_sdk-0.2.168-py3-none-manylinux_2_17_aarch64.whl";,
hash =
"sha256:77ae66b10edfd07b6ea3dec2af106ad47a9d616f8eb52d4f9b385e1b00bb2a27", size
= 114030032, upload-time = "2026-10-15T20:20:56.813Z" },
+ { url =
"https://files.pythonhosted.org/packages/00/03/09d347dcddbbab4becd41fef3a811967eddf24135b2d62dac9e321c81537/claude_agent_sdk-0.2.168-py3-none-manylinux_2_17_x86_64.whl";,
hash =
"sha256:be197be92549909ec62a2cee04c27f4b0a38c80f723735538a3efc352507f81d", size
= 114443626, upload-time = "2026-10-15T20:21:01.794Z" },
+ { url =
"https://files.pythonhosted.org/packages/48/93/520a18ab853b688afcc31d504344e541a7dc68bc2b515ea8a5fef1c868a1/claude_agent_sdk-0.2.168-py3-none-win_amd64.whl";,
hash =
"sha256:ed1814b0dd1c113bdaf3e51e6a67aa5b96b8f5c5839a43aeb6ae1e6ed23a5b35", size
= 117292475, upload-time = "2026-10-15T20:21:07.053Z" },
+]
+
[[package]]
name = "click"
version = "8.5.0"
