Lee-W commented on code in PR #74026:
URL: https://github.com/apache/airflow/pull/74026#discussion_r4398448160
##########
providers/common/ai/provider.yaml:
##########
@@ -192,6 +192,7 @@ connection-types:
- Heroku
- Hugging Face
- LiteLLM
+ - Meta
Review Comment:
Reset `uv.lock` to main and ran a plain `uv lock`; it now only adds
`claude-agent-sdk`. The Meta entry and the floor bump are no longer part of
this PR, since main has since added Meta itself.
##########
providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py:
##########
@@ -0,0 +1,362 @@
+# 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 :mod:`airflow.providers.common.ai.tools`.
+"""
+
+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).
+ :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
+ 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 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 result.is_error:
+ raise RuntimeError(_describe_error_result(result))
Review Comment:
`run()` now keeps the first `is_error` result and raises on it after the
loop, so an error turn followed by a successful one still fails.
##########
providers/common/ai/docs/frameworks/claude_agent_sdk.rst:
##########
@@ -0,0 +1,275 @@
+ .. 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`` and ``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
Review Comment:
Added a sentence: the server must be keyed by `tools.server_name` in
`mcp_servers`, with the mismatched-key example.
##########
providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py:
##########
@@ -0,0 +1,362 @@
+# 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 :mod:`airflow.providers.common.ai.tools`.
Review Comment:
The note now points at :ref:`howto/stability`, like the sibling adapters.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
