kaxil commented on code in PR #74026: URL: https://github.com/apache/airflow/pull/74026#discussion_r4290027055
########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,219 @@ + .. 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`` and the framework-neutral tool interface it builds on are Review Comment: This calls the framework-neutral tool interface experimental too, but `stability.rst`, `frameworks/index.rst` and every sibling guide say it's stable and only `AirflowTools` is experimental. The sibling wording would fit here. ########## providers/common/ai/docs/index.rst: ########## @@ -251,6 +250,35 @@ Extra Dependencies ``docx`` ``python-docx>=1.0.0`` ``git`` ``apache-airflow-providers-git`` =============== ======================================================================================================================================= +======= Review Comment: This `=======` and the table above it look like leftover merge residue: the old table lost its top border, and with this line in between docutils renders the whole extras section as a paragraph (0 tables, and no warning to catch it). Keeping only the new table below fixes it. ########## providers/common/ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py: ########## @@ -0,0 +1,88 @@ +# 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 asyncio +import os + +from claude_agent_sdk import query Review Comment: The sibling examples import the SDK and the adapter inside the task. With them at module level, the guide's `exampleinclude` (START/END markers below the imports) shows `query`, `AirflowTools`, `SQLToolset` and `asyncio` with no imports, and parsing the file without the extra fails with `ModuleNotFoundError`. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,266 @@ +# 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 ( + 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.query_results import DEFAULT_MAX_RESULT_BYTES +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 + + +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)})." + + +class _Run: + """One run of :meth:`AirflowTools.run`: identifies its calls and collects their failures.""" + + def __init__(self) -> None: + self.failures: list[ToolCallError] = [] + + +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. + + :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] + self._run: _Run | None = None + # With ``tools=[]`` the model has no Read tool, so a result the SDK spills to a file is lost to + # it. ``maxResultSizeChars`` counts characters, so using SQLToolset's byte cap as the limit + # raises the inline ceiling to at least the size of any result it caps. + annotations = ToolAnnotations(maxResultSizeChars=DEFAULT_MAX_RESULT_BYTES) Review Comment: 65,536 is the budget `build_query_result` gives the column names and rows, and its docstring says the `truncated`/`hint` envelope goes on top of that. A `SQLToolset` with a larger `max_result_bytes`, or a `HookToolset` result (no cap), can also be bigger. Anything past the annotation still spills: a 66,000-char result came back to the model as `<persisted-output> Output too large ... saved to ...`, which it can't open with `tools=[]`. A ceiling well above the default cap would make the "at least the size of any result it caps" comment true. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,266 @@ +# 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 ( + 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.query_results import DEFAULT_MAX_RESULT_BYTES +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 + + +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)})." + + +class _Run: + """One run of :meth:`AirflowTools.run`: identifies its calls and collects their failures.""" + + def __init__(self) -> None: + self.failures: list[ToolCallError] = [] + + +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. + + :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] + self._run: _Run | None = None + # With ``tools=[]`` the model has no Read tool, so a result the SDK spills to a file is lost to + # it. ``maxResultSizeChars`` counts characters, so using SQLToolset's byte cap as the limit + # raises the inline ceiling to at least the size of any result it caps. + annotations = ToolAnnotations(maxResultSizeChars=DEFAULT_MAX_RESULT_BYTES) + 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 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. + 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 = [*kwargs.pop("allowed_tools", []), *self.allowed_tools] + kwargs.setdefault("tools", []) + kwargs.setdefault("setting_sources", []) + kwargs.setdefault("strict_mcp_config", True) Review Comment: These defaults leave `verbatim_prompts` at `False`, so the CLI still expands `@path` mentions in the prompt. I ran the bundled CLI against a fake Messages API with exactly these options and the prompt `Summarise @/tmp/x/secret.txt please`: the file's contents were in the request body, even with `tools=[]`. With `verbatim_prompts=True` they weren't. The example takes `question` as a task argument, so anyone who can trigger the Dag with conf can send any file the worker user can read (`airflow.cfg`, for one) to the model, and that content never goes through the masker. Can `build_options()` default `verbatim_prompts=True` as well, with the "Built-in tools run on the worker host" list in the guide and `test_safe_defaults` updated to match? ########## providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py: ########## @@ -0,0 +1,637 @@ +# 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.ai.utils.query_results import DEFAULT_MAX_RESULT_BYTES + +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, + ) + + +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() -> AssistantMessage: + return AssistantMessage(content=[TextBlock(text="")], model="claude-test", stop_reason="tool_use") + + +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, +) -> 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``. + """ + for turn in turns: + yield _assistant_message() + 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): Review Comment: This only checks names, and nothing in the file reads a description back: passing `""` instead of `t.description` in the adapter still leaves all 36 tests green. A `tools/list` through `_request_real_server` could assert name, description, `inputSchema` and the `_meta` max-size entry in one go. Smaller things: the import at line 629 can go to the top, and the `AsyncMock` at line 230 has no spec. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,266 @@ +# 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 ( + 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.query_results import DEFAULT_MAX_RESULT_BYTES +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 + + +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)})." + + +class _Run: + """One run of :meth:`AirflowTools.run`: identifies its calls and collects their failures.""" + + def __init__(self) -> None: + self.failures: list[ToolCallError] = [] + + +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. + + :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] + self._run: _Run | None = None + # With ``tools=[]`` the model has no Read tool, so a result the SDK spills to a file is lost to + # it. ``maxResultSizeChars`` counts characters, so using SQLToolset's byte cap as the limit + # raises the inline ceiling to at least the size of any result it caps. + annotations = ToolAnnotations(maxResultSizeChars=DEFAULT_MAX_RESULT_BYTES) + 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 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. + 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 = [*kwargs.pop("allowed_tools", []), *self.allowed_tools] + kwargs.setdefault("tools", []) + kwargs.setdefault("setting_sources", []) + kwargs.setdefault("strict_mcp_config", True) + 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. + + :raises RuntimeError: when this instance is already running a query, when + ``messages`` ends without a :class:`~claude_agent_sdk.ResultMessage`, or when + the result message has ``is_error`` set (for example ``max_turns`` reached or + an API failure). + :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 + try: + result: ResultMessage | None = None + async for message in messages: + 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)) + return result + finally: + 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 state is not None and state.failures: + return { + "content": [ + { + "type": "text", + "text": "Not run: an earlier Airflow tool call in this run failed.", + } + ], + "is_error": True, + } + with calling_framework("claude_agent_sdk"), tool_call_scope(run=state): Review Comment: With no `turn` here, `_RetryBudget.failed` relies on "calls that started before the last counted failure belong to the same turn". That never applies with this SDK: the bundled CLI runs a turn's MCP calls one after another unless the tool sets `readOnlyHint` (three calls in one turn ran at 0.0s, 0.5s and 1.0s, and overlapped once I set the hint). So two bad `query` calls in one turn raise "kept failing after 1 correction(s)" and fail the task, where the Anthropic and Strands adapters count that turn once. `run()` sees each `AssistantMessage`, which carries `message_id`, so recording it on the run state and passing it as `turn=` would line up with them, though the ordering against the handler isn't documented and would want a test against the real server. If you'd rather not depend on that, the guide's "calls that run at the same time count as one failure" (line 115) needs to say each failed call counts. ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,219 @@ + .. 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`` and the framework-neutral tool interface it builds on are + experimental. They may change in a minor release of this provider. + 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 like this one to the task log. This output is +illustrative, not captured from a real run; the tool name and timing vary with what the +model decides to do: + +.. code-block:: text + + [info] ::group::Tool call: list_tables + [info] Tool list_tables returned in 0.00s + [info] ::endgroup:: + +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. + +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 ``SQLToolset``'s 64 KiB result limit as the size the SDK should + keep inline. +- 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. For a ``query`` call the database rejected, the model would read a +block shaped like this one (illustrative, not captured from a real run; 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_example_1", + "content": "The query tool failed: <database error>\nUse the list_tables and get_schema ... 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`, not per model +turn: the SDK's handler does not identify which model turn a call belongs to, so calls +that run at the same time count as one failure, the same way the ADK adapter counts them +(see :doc:`adk`). 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: it raises ``RuntimeError`` with the result's ``subtype``, ``result``, Review Comment: For `query()` this raises `ResultError`, not `RuntimeError`: the CLI exits non-zero right after an `is_error` result, and the SDK raises `ResultError` (a `ProcessError`) from the stream on the next iteration, before `run()` reaches its own check. Only `receive_response()` gets to the `RuntimeError`. Can this paragraph and the `run()` docstring say which one each stream raises? ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,266 @@ +# 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 ( + 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.query_results import DEFAULT_MAX_RESULT_BYTES +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 + + +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)})." + + +class _Run: + """One run of :meth:`AirflowTools.run`: identifies its calls and collects their failures.""" + + def __init__(self) -> None: + self.failures: list[ToolCallError] = [] + + +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. + + :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] + self._run: _Run | None = None + # With ``tools=[]`` the model has no Read tool, so a result the SDK spills to a file is lost to + # it. ``maxResultSizeChars`` counts characters, so using SQLToolset's byte cap as the limit + # raises the inline ceiling to at least the size of any result it caps. + annotations = ToolAnnotations(maxResultSizeChars=DEFAULT_MAX_RESULT_BYTES) + 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 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. + 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 = [*kwargs.pop("allowed_tools", []), *self.allowed_tools] + kwargs.setdefault("tools", []) + kwargs.setdefault("setting_sources", []) + kwargs.setdefault("strict_mcp_config", True) + 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. + + :raises RuntimeError: when this instance is already running a query, when + ``messages`` ends without a :class:`~claude_agent_sdk.ResultMessage`, or when + the result message has ``is_error`` set (for example ``max_turns`` reached or + an API failure). + :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 + try: + result: ResultMessage | None = None + async for message in messages: Review Comment: When `run()` raises out of this loop, nothing closes `messages`, so the CLI keeps going until the generator is finalized. With a fake transport, a caller that catches `ToolCallError` and keeps the event loop alive saw the CLI send another `tools/call` 0.4s later, and since the `finally` had already set `self._run = None`, the handler skipped the "Not run" guard and ran the tool. In the example's `asyncio.run` shape it does get cleaned up at loop shutdown, but the SDK's own `inner.aclose()` fails there and logs a `RuntimeError: aclose(): asynchronous generator is already running` traceback. Closing the stream in the `finally` before clearing `_run` (`contextlib.aclosing`, or `await messages.aclose()` when it has one) fixes both; I checked the first one that way. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,266 @@ +# 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 ( + 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.query_results import DEFAULT_MAX_RESULT_BYTES +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 + + +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)})." + + +class _Run: + """One run of :meth:`AirflowTools.run`: identifies its calls and collects their failures.""" + + def __init__(self) -> None: + self.failures: list[ToolCallError] = [] + + +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. + + :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] + self._run: _Run | None = None + # With ``tools=[]`` the model has no Read tool, so a result the SDK spills to a file is lost to + # it. ``maxResultSizeChars`` counts characters, so using SQLToolset's byte cap as the limit + # raises the inline ceiling to at least the size of any result it caps. + annotations = ToolAnnotations(maxResultSizeChars=DEFAULT_MAX_RESULT_BYTES) + 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 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. + 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 = [*kwargs.pop("allowed_tools", []), *self.allowed_tools] Review Comment: If someone passes `can_use_tool` here to approve or deny the SQL calls, every Airflow tool is still added to `allowed_tools`, so the CLI auto-approves them and the callback never sees one. The SDK only emits `CanUseToolShadowedWarning`. Should `build_options()` leave its tools out of `allowed_tools` when `can_use_tool` is passed, or at least say so in the docstring? ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,219 @@ + .. 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`` and the framework-neutral tool interface it builds on are + experimental. They may change in a minor release of this provider. + 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 like this one to the task log. This output is +illustrative, not captured from a real run; the tool name and timing vary with what the Review Comment: This block and the `is_error` one at line 97 still say illustrative, and the log shows only `list_tables`. The reply on the earlier thread had screenshots of a real run calling `list_tables`, `get_schema` and `query`; can that output go in here? -- 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]
