kaxil commented on code in PR #74026: URL: https://github.com/apache/airflow/pull/74026#discussion_r4237644786
########## providers/common/ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py: ########## @@ -0,0 +1,84 @@ +# 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 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 +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 + async def ask_the_warehouse(question: str = DEFAULT_QUESTION) -> str | None: Review Comment: An `async def` task needs Airflow 3.2+, but the provider supports 2.11, and this is the guide's only recipe. A sync `@task` that wraps the loop in `asyncio.run(...)`, as `example_adk_agent.py` does, runs on every supported version. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,215 @@ +# 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, + 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"] + + +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 + sdk_tools = [ + tool(t.name, t.description, copy.deepcopy(t.parameters))(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 ``claude_agent_sdk.query()`` (or a client) returns. + + :raises RuntimeError: when this instance is already running a query, or when + ``messages`` ends without a :class:`~claude_agent_sdk.ResultMessage`. + :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.") + return result Review Comment: The docstring says `messages` can come from a client too, but `ClaudeSDKClient.receive_response()` stops right after the `ResultMessage`, so an error result (max_turns hit, API failure) is returned here as a success and the task passes with an empty `result.result`. `query()` only raises `ResultError` for that case because it waits for the CLI to exit non-zero. Raising when `result.is_error` is set covers both, and `receive_messages()` never ends, so the docstring should say it isn't supported. ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,195 @@ + .. 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 called one tool, ``lookup``. Here is the task log's tool-call group, Review Comment: This log, and the `is_error` block at line 92, come from the unit test's `lookup` tool, not the `SQLToolset` example just above, which would call `list_tables`, `get_schema` or `query`. Can you capture both from a real run of the example? ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,195 @@ + .. 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 called one tool, ``lookup``. Here is the task log's tool-call group, +timestamps trimmed: + +.. code-block:: text + + [info] ::group::Tool call: lookup + [info] Tool lookup 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`. +- 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``, comes back to it as an MCP tool result with ``isError`` set. The +model received this block when ``lookup`` answered that a column did not exist: + +.. code-block:: json + + { + "type": "tool_result", + "tool_use_id": "toolu_lookup_1", + "content": "Unknown column 'nme'. Did you mean 'name'?", + "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:`~...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:`~...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. + +.. _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`` +----------------------------------- + +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:`~...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:`run` again only once the previous call has returned. + +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 server ``AirflowTools`` builds is used. If you build ``ClaudeAgentOptions`` Review Comment: `build_options()` also keeps any `mcp_servers` the caller passes, so it is the servers passed to `build_options()` plus Airflow's, not only the one `AirflowTools` builds. ########## providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py: ########## @@ -0,0 +1,417 @@ +# 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 +from airflow.providers.common.ai.tools.claude_agent_sdk import AirflowTools +from airflow.providers.common.ai.toolsets.sql import SQLToolset + +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() + + +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 + + +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 Review Comment: If I remove the in-loop raise at `claude_agent_sdk.py:168` and only check for failures after the stream ends, this test still passes. `_fake_cli` yields nothing but the `ResultMessage` after the failing call, so there's no later message to fail on. Adding a second turn after the failure and asserting it never runs would actually test "at the next message". Every test also mocks `create_sdk_mcp_server`, so nothing checks the `is_error` to `isError` mapping or that the SDK catches handler exceptions. A test against the real server instance would cover both. ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,195 @@ + .. 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 called one tool, ``lookup``. Here is the task log's tool-call group, +timestamps trimmed: + +.. code-block:: text + + [info] ::group::Tool call: lookup + [info] Tool lookup 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`. +- 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``, comes back to it as an MCP tool result with ``isError`` set. The +model received this block when ``lookup`` answered that a column did not exist: + +.. code-block:: json + + { + "type": "tool_result", + "tool_use_id": "toolu_lookup_1", + "content": "Unknown column 'nme'. Did you mean 'name'?", + "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:`~...tools.claude_agent_sdk.AirflowTools.run`, not per model Review Comment: `~...tools.claude_agent_sdk.AirflowTools.run` doesn't resolve, since Sphinx strips only one leading dot, so this ref and the ones at 115, 136 and the bare `:meth:`run`` at 141 render as unlinked text. The full `airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run` path, as used for `build_options` above, links. ########## providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py: ########## @@ -0,0 +1,215 @@ +# 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, + 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"] + + +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 + sdk_tools = [ + tool(t.name, t.description, copy.deepcopy(t.parameters))(self._build_handler(t)) Review Comment: The SDK's `tool()` docstring says a result over Claude Code's inline size limit gets written to a file, and the model only sees a preview. With `tools=[]` there is no `Read` tool, so the model can't open that file. `SQLToolset` results go up to 64 KiB, so I think these tools need `ToolAnnotations(maxResultSizeChars=...)`. Separately: `_build_input_schema` only passes a dict schema through when it has `properties`. A hand-built `AirflowTool` with `{"type": "object"}` ends up with a required `type` argument. Calling `setdefault("properties", {})` on the copy would fix it. -- 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]
