kaxil commented on code in PR #74026: URL: https://github.com/apache/airflow/pull/74026#discussion_r4352309838
########## providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py: ########## @@ -0,0 +1,1044 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import copy +import importlib +import json +import sys +from collections.abc import AsyncIterator +from typing import Any +from unittest import mock + +import pytest + +pytest.importorskip("claude_agent_sdk") + +from claude_agent_sdk import AssistantMessage, Message, ResultMessage, TextBlock, UserMessage +from pydantic_ai.exceptions import ModelRetry + +from airflow.providers.common.ai.tools import AirflowTool, ToolCallError, ToolResult, collect_tools +from airflow.providers.common.ai.tools.claude_agent_sdk import AirflowTools +from airflow.providers.common.ai.toolsets.sql import SQLToolset +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +from unit.common.ai.tools.test__from_toolset import _scripted +from unit.common.ai.toolsets.test_sql import _make_mock_db_hook + +MODULE = "airflow.providers.common.ai.tools.claude_agent_sdk" + + +def _tool(outcome: ToolResult | Exception, ran: list[dict[str, Any]] | None = None) -> AirflowTool: + """A ``lookup`` tool that returns or raises ``outcome``, recording its calls in ``ran``.""" + + async def function(arguments: dict[str, Any]) -> ToolResult: + if ran is not None: + ran.append(arguments) + if isinstance(outcome, Exception): + raise outcome + return outcome + + return AirflowTool( + name="lookup", + description="Look something up.", + parameters={"type": "object", "properties": {"key": {"type": "string"}}, "required": ["key"]}, + function=function, + ) + + +async def _pong(arguments: dict[str, Any]) -> ToolResult: + return ToolResult("pong") + + +def _build_tools(*sources: Any, server_name: str = "airflow") -> tuple[AirflowTools, dict[str, Any]]: + """ + Build ``AirflowTools``, with ``create_sdk_mcp_server`` mocked to capture its SDK tools. + + The ``@tool`` decorator itself is not mocked: it is only a plain wrapper + (``SdkMcpTool(handler=handler, ...)``), so the handlers it returns run the same code + they would through the real MCP server. Returns the instance and its handlers by name. + """ + captured: dict[str, Any] = {} + + def _capture_server(*, name: str, version: str = "1.0.0", tools: Any = None) -> Any: + captured["sdk_tools"] = tools + return mock.sentinel.server + + with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as mock_create_server: + mock_create_server.side_effect = _capture_server + tools = AirflowTools(*sources, server_name=server_name) + handlers = {sdk_tool.name: sdk_tool.handler for sdk_tool in captured["sdk_tools"]} + return tools, handlers + + +def _assistant_message(message_id: str | None = None) -> AssistantMessage: + return AssistantMessage( + content=[TextBlock(text="")], model="claude-test", stop_reason="tool_use", message_id=message_id + ) + + +def _user_message() -> UserMessage: + return UserMessage(content="") + + +def _result_message(**overrides: Any) -> ResultMessage: + fields: dict[str, Any] = { + "subtype": "success", + "duration_ms": 1, + "duration_api_ms": 1, + "is_error": False, + "num_turns": 1, + "session_id": "session-1", + } + fields.update(overrides) + return ResultMessage(**fields) + + +async def _fake_cli( + handlers: dict[str, Any], + turns: list[list[tuple[str, dict[str, Any]]]], + *, + final: ResultMessage | None = None, + message_ids: list[str | None] | None = None, +) -> AsyncIterator[Message]: + """ + Stand in for the CLI driving the tools ``AirflowTools`` served through one or more turns. + + For each turn, yields an ``AssistantMessage`` with a ``tool_use`` stop reason, then + awaits every one of that turn's calls -- as the CLI waits for a turn's tool results + before sending the next request -- before yielding the ``UserMessage`` carrying them. + Ends with a ``ResultMessage``. ``message_ids`` gives each turn's ``AssistantMessage.message_id``. + """ + for index, turn in enumerate(turns): + yield _assistant_message(message_ids[index] if message_ids else None) + for tool_name, args in turn: + await handlers[tool_name](args) + yield _user_message() + yield final if final is not None else _result_message() + + +async def _call_real_server(tools: AirflowTools, name: str, args: dict[str, Any]) -> dict[str, Any]: + """Call a tool through the real in-process MCP server, with the JSON-RPC the SDK feeds its bridge.""" + return await _request_real_server(tools, "tools/call", {"name": name, "arguments": args}) + + +async def _request_real_server(tools: AirflowTools, method: str, params: dict[str, Any]) -> dict[str, Any]: + """Send one JSON-RPC request to the real in-process MCP server and return its ``result``.""" + moved = ( + "claude_agent_sdk's private path _internal.sdk_mcp_bridge.SdkMcpBridge has changed; " + "update _request_real_server in this test module." + ) + try: + bridge_module = importlib.import_module("claude_agent_sdk._internal.sdk_mcp_bridge") + except ImportError: + pytest.fail(moved) + bridge_cls = getattr(bridge_module, "SdkMcpBridge", None) + if bridge_cls is None: + pytest.fail(moved) + bridge = bridge_cls("airflow", tools.server["instance"]) + try: + await bridge.handle( + { + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "test", "version": "1"}, + }, + } + ) + await bridge.handle({"jsonrpc": "2.0", "method": "notifications/initialized"}) + reply = await bridge.handle( + { + "jsonrpc": "2.0", + "id": 1, + "method": method, + "params": params, + } + ) + return reply["result"] + finally: + await bridge.aclose() + + +class TestTools: + def test_each_toolset_tool_becomes_an_mcp_tool_with_its_schema(self): + tools, handlers = _build_tools(_tool(ToolResult("ok")), SQLToolset("pg_default")) + + assert tools.allowed_tools == [ + "mcp__airflow__lookup", + "mcp__airflow__list_tables", + "mcp__airflow__get_schema", + "mcp__airflow__query", + "mcp__airflow__check_query", + ] + assert set(handlers) == {"lookup", "list_tables", "get_schema", "query", "check_query"} + + def test_allowed_tools_use_a_custom_server_name(self): + tools, _ = _build_tools(_tool(ToolResult("ok")), server_name="warehouse") + + assert tools.allowed_tools == ["mcp__warehouse__lookup"] + + def test_an_mcp_tool_does_not_share_the_source_schema(self): + source = _tool(ToolResult("ok")) + before = copy.deepcopy(source.parameters) + captured: dict[str, Any] = {} + + def _capture_server(*, name: str, version: str = "1.0.0", tools: Any = None) -> Any: + captured["sdk_tools"] = tools + return mock.sentinel.server + + with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as mock_create_server: + mock_create_server.side_effect = _capture_server + AirflowTools(source) + + captured["sdk_tools"][0].input_schema["properties"]["key"]["type"] = "integer" + assert source.parameters == before + + def test_every_mcp_tool_keeps_a_large_result_inline(self): + captured: dict[str, Any] = {} + sources = [_tool(ToolResult("ok")), SQLToolset("pg_default")] + + def _capture_server(*, name: str, version: str = "1.0.0", tools: Any = None) -> Any: + captured["sdk_tools"] = tools + return mock.sentinel.server + + with mock.patch(f"{MODULE}.create_sdk_mcp_server", autospec=True) as mock_create_server: + mock_create_server.side_effect = _capture_server + AirflowTools(*sources) + + assert [t.name for t in captured["sdk_tools"]] == [t.name for t in collect_tools(sources)] + for sdk_tool in captured["sdk_tools"]: + assert sdk_tool.annotations.maxResultSizeChars == 500_000 + + @pytest.mark.asyncio + async def test_a_schema_without_properties_takes_no_arguments(self): + source = AirflowTool( + name="ping", + description="Ping.", + parameters={"type": "object"}, + function=mock.create_autospec(_pong, return_value=ToolResult("pong")), + ) + tools = AirflowTools(source) + + listed = await _request_real_server(tools, "tools/list", {}) + + assert listed["tools"][0]["inputSchema"].get("required", []) == [] + assert "type" not in listed["tools"][0]["inputSchema"]["properties"] + assert source.parameters == {"type": "object"} + + @pytest.mark.asyncio + async def test_the_cli_lists_each_tool_with_its_description_schema_and_result_limit(self): + schema = { + "type": "object", + "properties": {"key": {"type": "string", "description": "The key."}}, + "required": ["key"], + } + source = AirflowTool( + name="lookup", + description="Look a value up.", + parameters=schema, + function=mock.create_autospec(_pong, return_value=ToolResult("pong")), + ) + tools = AirflowTools(source) + + listed = await _request_real_server(tools, "tools/list", {}) + + (listed_tool,) = listed["tools"] + assert listed_tool["name"] == "lookup" + assert listed_tool["description"] == "Look a value up." + assert listed_tool["inputSchema"] == schema + assert listed_tool["_meta"] == {"anthropic/maxResultSizeChars": 500_000} + + +class TestHandler: + @pytest.mark.asyncio + async def test_the_model_reads_a_toolset_result(self): + ts = SQLToolset("pg_default") + ts._hook = _make_mock_db_hook(records=[(1, "Ada")], last_description=[("id",), ("name",)]) + tools, handlers = _build_tools(ts) + + payload = await handlers["query"]({"sql": "SELECT id, name FROM users"}) + + assert "is_error" not in payload + assert json.loads(payload["content"][0]["text"])["rows"] == [[1, "Ada"]] + + @pytest.mark.asyncio + async def test_a_correctable_failure_reaches_the_model_as_an_error(self): + tools, handlers = _build_tools(_tool(ToolResult("Unknown column 'nme'", is_error=True))) + + payload = await handlers["lookup"]({"key": "a"}) + + assert payload["is_error"] is True + assert payload["content"] == [{"type": "text", "text": "Unknown column 'nme'"}] + + @pytest.mark.asyncio + async def test_a_structured_result_reaches_the_model_as_json(self): + tools, handlers = _build_tools(_tool(ToolResult({"rows": [[1, None]]}))) + + payload = await handlers["lookup"]({"key": "a"}) + + assert json.loads(payload["content"][0]["text"]) == {"rows": [[1, None]]} + + @pytest.mark.asyncio + async def test_not_driven_by_run_hands_the_failure_to_the_model_with_a_warning(self): + tools, handlers = _build_tools(_tool(PermissionError("role cannot read orders"))) + + with mock.patch(f"{MODULE}.log", autospec=True) as mock_log: + payload = await handlers["lookup"]({"key": "a"}) + + assert payload["is_error"] is True + assert "PermissionError: role cannot read orders" in payload["content"][0]["text"] + mock_log.warning.assert_called_once() + assert "not driven by AirflowTools.run" in mock_log.warning.call_args.args[0] + + @pytest.mark.asyncio + async def test_the_toolsets_calls_are_counted_as_claude_agent_sdk(self): + ts = SQLToolset("pg_default") + ts._hook = _make_mock_db_hook() + tools, handlers = _build_tools(ts) + + with mock.patch( + "airflow.providers.common.ai.utils.tool_metrics.Stats", mock.MagicMock(spec=["incr"]) + ) as stats: + await handlers["list_tables"]({}) + + assert stats.incr.call_args.kwargs["tags"]["framework"] == "claude_agent_sdk" + + @pytest.mark.enable_redact + @pytest.mark.asyncio + async def test_a_value_mask_secrets_cannot_descend_into_is_masked_after_rendering( + self, registered_secret + ): + """ + ``mask_secrets`` leaves an opaque object (not a dict, list or dataclass) untouched, so + only masking *after* it is turned into text -- what ``serialize_for_llm`` does -- + catches a secret in its ``repr``. ``json.dumps(value, default=str)`` masks first and + renders after, so it would not catch this. + """ + + class _Row: + def __init__(self, secret: str) -> None: + self._secret = secret + + def __repr__(self) -> str: + return f"_Row(secret={self._secret!r})" + + tools, handlers = _build_tools(_tool(ToolResult(_Row(f"key={registered_secret}")))) + + payload = await handlers["lookup"]({"key": "a"}) + + text = payload["content"][0]["text"] + assert registered_secret not in text + assert "key=***" in text + + @pytest.mark.enable_redact + @pytest.mark.asyncio + async def test_a_secret_in_a_plain_string_result_is_masked(self, registered_secret): + tools, handlers = _build_tools(_tool(ToolResult(f"key={registered_secret}"))) + + payload = await handlers["lookup"]({"key": "a"}) + + assert registered_secret not in payload["content"][0]["text"] + assert payload["content"][0]["text"] == "key=***" + + +class TestRun: + @pytest.mark.asyncio + async def test_run_returns_the_result_message(self): + tools, handlers = _build_tools(_tool(ToolResult("ok"))) + final = _result_message(result="done") + + result = await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]], final=final)) + + assert result is final + + @pytest.mark.asyncio + async def test_run_raises_without_a_result_message(self): + tools, handlers = _build_tools() + + async def _no_result() -> AsyncIterator[Message]: + return + yield # pragma: no cover - makes this an async generator function + + with pytest.raises(RuntimeError, match="ended without a ResultMessage"): + await tools.run(_no_result()) + + @pytest.mark.asyncio + async def test_a_failure_before_the_stream_ends_without_a_result_message_still_raises_it(self): + """ + The CLI can end the stream right after a failed tool call without sending another + message at all -- no further ``UserMessage``, no ``ResultMessage``. The in-loop check + (see the test above) never runs again once the stream is exhausted, so the failure + has to be checked once more after the loop, before falling back to the generic + "ended without a ResultMessage" error. + """ + tools, handlers = _build_tools(_tool(PermissionError("role cannot read orders"))) + + async def _cli_ends_right_after_the_failure() -> AsyncIterator[Message]: + yield _assistant_message() + await handlers["lookup"]({"key": "a"}) + # The stream ends here: no further message, no ResultMessage. + + with pytest.raises(ToolCallError, match="PermissionError: role cannot read orders"): + await tools.run(_cli_ends_right_after_the_failure()) + + @pytest.mark.asyncio + async def test_any_other_failure_fails_the_run_at_the_next_message(self): + """ + The handler answers the model within the turn -- ``ran`` has one call, and the model + did read the masked failure -- but ``run`` still raises at the next message rather + than letting the stream reach a ``ResultMessage``. + """ + ran: list[dict[str, Any]] = [] + tools, handlers = _build_tools(_tool(PermissionError("role cannot read orders"), ran)) + + with pytest.raises(ToolCallError, match="PermissionError: role cannot read orders"): + await tools.run(_fake_cli(handlers, [[("lookup", {"key": "a"})]])) + + assert len(ran) == 1 + + @pytest.mark.asyncio + async def test_a_failure_in_one_turn_stops_the_stream_before_the_next_turn_starts(self): + ran: list[dict[str, Any]] = [] + turn_two_started: list[bool] = [] + tools, handlers = _build_tools(_tool(PermissionError("role cannot read orders"), ran)) + + async def _two_turns() -> AsyncIterator[Message]: + yield _assistant_message() + await handlers["lookup"]({"key": "a"}) + yield _user_message() + turn_two_started.append(True) + yield _assistant_message() + await handlers["lookup"]({"key": "b"}) + yield _user_message() + yield _result_message() + + with pytest.raises(ToolCallError, match="PermissionError: role cannot read orders"): + await tools.run(_two_turns()) + + assert turn_two_started == [] + assert len(ran) == 1 + + @pytest.mark.asyncio + async def test_calls_after_a_failure_in_the_same_turn_do_not_run(self): + ran: list[dict[str, Any]] = [] + tools, handlers = _build_tools(_tool(PermissionError("role cannot read orders"), ran)) + + async def _three_calls_one_turn() -> AsyncIterator[Message]: + yield _assistant_message() + for _ in range(3): + await handlers["lookup"]({"key": "a"}) + yield _user_message() + yield _result_message() + + with pytest.raises(ToolCallError, match="PermissionError"): + await tools.run(_three_calls_one_turn()) + + assert len(ran) == 1 + + @pytest.mark.asyncio + async def test_an_error_result_message_raises_instead_of_being_returned(self): + tools, _ = _build_tools() + + async def _ends_at_the_result() -> AsyncIterator[Message]: + yield _result_message( + is_error=True, + subtype="error_max_turns", + result="partial answer", + errors=["hit the turn limit"], + num_turns=3, + ) + + with pytest.raises(RuntimeError) as exc_info: + await tools.run(_ends_at_the_result()) + + message = str(exc_info.value) + assert "error_max_turns" in message + assert "partial answer" in message + assert "hit the turn limit" in message + + @pytest.mark.asyncio + async def test_an_error_turn_followed_by_a_successful_turn_still_raises(self): + tools, _ = _build_tools() + + async def _error_then_success() -> AsyncIterator[Message]: + yield _result_message(is_error=True, subtype="error_during_execution", result="first turn failed") Review Comment: With only one error result, this can't tell first from last: changing `run()` to `if message.is_error:` (keep the last error) still passes. A second error result with different text between these two, plus a check that its text isn't in the message, would pin the "keep the first" behaviour the comment in `run()` describes. ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,278 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/frameworks:claude_agent_sdk: + +Claude Agent SDK +================= + +If your agent is a loop on the `Claude Agent SDK +<https://platform.claude.com/docs/en/api/agent-sdk/overview>`__'s ``query()``, which runs +the Claude Code CLI as a subprocess, you can run it in an Airflow task as it is and give +it the toolsets this provider ships. +:class:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools` serves +``SQLToolset``, ``HookToolset`` and most other toolsets to the CLI through an in-process +MCP server. If you have no loop yet, ``AgentOperator`` runs one for you, with durable +replay and human review; see :doc:`index`. + +.. note:: + + ``AirflowTools`` is experimental: it may change in a minor release of this provider. + The framework-neutral tool interface it builds on is stable. + See :ref:`howto/stability`. + +Run a loop in a task +--------------------- + +This task answers a question about a database. The client's API key comes from an +Airflow connection: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py + :language: python + :start-after: [START example_claude_agent_sdk] + :end-before: [END example_claude_agent_sdk] + +The model can call ``list_tables``, ``get_schema`` and ``query``, the tools ``SQLToolset`` Review Comment: `SQLToolset` always serves `check_query` too, so this list is one short. Someone writing a `can_use_tool` policy from it would miss that tool. ########## providers/common/ai/docs/frameworks/claude_agent_sdk.rst: ########## @@ -0,0 +1,278 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/frameworks:claude_agent_sdk: + +Claude Agent SDK +================= + +If your agent is a loop on the `Claude Agent SDK +<https://platform.claude.com/docs/en/api/agent-sdk/overview>`__'s ``query()``, which runs +the Claude Code CLI as a subprocess, you can run it in an Airflow task as it is and give +it the toolsets this provider ships. +:class:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools` serves +``SQLToolset``, ``HookToolset`` and most other toolsets to the CLI through an in-process +MCP server. If you have no loop yet, ``AgentOperator`` runs one for you, with durable +replay and human review; see :doc:`index`. + +.. note:: + + ``AirflowTools`` is experimental: it may change in a minor release of this provider. + The framework-neutral tool interface it builds on is stable. + See :ref:`howto/stability`. + +Run a loop in a task +--------------------- + +This task answers a question about a database. The client's API key comes from an +Airflow connection: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_claude_agent_sdk.py + :language: python + :start-after: [START example_claude_agent_sdk] + :end-before: [END example_claude_agent_sdk] + +The model can call ``list_tables``, ``get_schema`` and ``query``, the tools ``SQLToolset`` +serves. Each call writes a group to the task log. This is the log of one run against a +temporary SQLite database, captured as the task logger wrote it, for the question "Inspect +the columns of each table first with get_schema, then tell me which customer name has the +highest total order amount and in which city they live." Which tools the model calls, in +what order, and how long each takes depend on the model and the question: + +.. code-block:: text + + 2026-10-19T07:56:54.324694Z [info ] ::group::Tool call: list_tables [airflow.providers.common.ai.tools] loc=__init__.py:116 + 2026-10-19T07:56:54.603210Z [info ] Tool list_tables returned in 0.28s [airflow.providers.common.ai.tools] loc=__init__.py:130 + 2026-10-19T07:56:54.603326Z [info ] ::endgroup:: [airflow.providers.common.ai.tools] loc=__init__.py:132 + 2026-10-19T07:56:56.308886Z [info ] ::group::Tool call: get_schema [airflow.providers.common.ai.tools] loc=__init__.py:116 + 2026-10-19T07:56:56.309896Z [info ] Tool get_schema returned in 0.00s [airflow.providers.common.ai.tools] loc=__init__.py:130 + 2026-10-19T07:56:56.310010Z [info ] ::endgroup:: [airflow.providers.common.ai.tools] loc=__init__.py:132 + 2026-10-19T07:56:56.476896Z [info ] ::group::Tool call: get_schema [airflow.providers.common.ai.tools] loc=__init__.py:116 + 2026-10-19T07:56:56.478181Z [info ] Tool get_schema returned in 0.00s [airflow.providers.common.ai.tools] loc=__init__.py:130 + 2026-10-19T07:56:56.478343Z [info ] ::endgroup:: [airflow.providers.common.ai.tools] loc=__init__.py:132 + 2026-10-19T07:56:58.700718Z [info ] ::group::Tool call: query [airflow.providers.common.ai.tools] loc=__init__.py:116 + 2026-10-19T07:56:58.704926Z [info ] Running statement: SELECT c.name, c.city, SUM(o.amount) AS total_amount + FROM customers c + JOIN orders o ON o.customer_id = c.id + GROUP BY c.id, c.name, c.city + ORDER BY total_amount DESC + LIMIT 5, parameters: None [airflow.task.hooks.airflow.providers.sqlite.hooks.sqlite.SqliteHook] loc=sql.py:896 + 2026-10-19T07:56:59.495567Z [info ] Tool query returned in 0.79s [airflow.providers.common.ai.tools] loc=__init__.py:130 + 2026-10-19T07:56:59.495691Z [info ] ::endgroup:: [airflow.providers.common.ai.tools] loc=__init__.py:132 + +Only ``AirflowTools`` and ``tools.run(query(...))`` come from Airflow. ``query()``, the +system prompt, the model and ``max_turns`` are the SDK's own. +:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.build_options` builds +``ClaudeAgentOptions`` with the Airflow tools wired in and safe defaults for the fields +that matter most for this adapter; see `Built-in tools run on the worker host`_. You can +also build ``ClaudeAgentOptions`` yourself and add +:attr:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.server` and +:attr:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.allowed_tools` to +it. Put the server in ``mcp_servers`` under the key ``tools.server_name`` (``"airflow"`` by +default). The SDK names a server after its dict key, so ``{"warehouse": tools.server}`` +serves ``mcp__warehouse__*`` tools while ``allowed_tools`` lists ``mcp__airflow__*`` and +the names do not match. + +What each tool does +-------------------- + +``AirflowTools`` accepts any toolset that implements +:class:`~airflow.providers.common.ai.tools.ToolProvider`, and individual +:class:`~airflow.providers.common.ai.tools.AirflowTool` objects. ``MCPToolset`` is the +exception; give the SDK the MCP server directly instead. Each tool: + +- Keeps the toolset tool's name, description and argument schema, served as an MCP tool + named ``mcp__<server_name>__<tool name>``. +- Runs through the toolset, so connection resolution, SQL validation, + ``allowed_tables``, ``allowed_methods`` and result limits behave as they do in + :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. Because the + built-in tools are off, the model cannot open a result the SDK writes to a file, so each + tool is marked with 500,000 characters as the size the SDK should keep inline. The bundled + CLI caps that value at 500,000 (observed on CLI 2.1.286) and counts it in UTF-16 code units, + so an emoji counts as two. A single result longer than that + is truncated to fit and ends with a notice asking the model to narrow the request; the + model still reads it as a normal result, not an error. Without the truncation the CLI would + spill the result to a file the model cannot read. The CLI also writes each result to stdout + twice, JSON-escaped, and the SDK fails the stream on any line over its ``max_buffer_size`` + (1 MiB by default), which a result well under 500,000 characters can exceed. + ``build_options()`` therefore sets ``max_buffer_size`` to 16 MiB. +- Passes its result through Airflow's secret masker before the model sees it. A + connection password that turns up in a query result reaches the model as ``***``. +- Writes a ``Tool call: <name>`` group to the task log. Calls to a toolset are also + counted in the ``common_ai.tool_calls`` metric with ``framework=claude_agent_sdk`` (see + :doc:`../observability`). + +When a tool fails +------------------- + +A failure the model can correct, such as a query naming a table outside +``allowed_tables`` or one the database rejects, comes back to it as an MCP tool result +with ``isError`` set. This is what the model read when it asked a ``SQLToolset`` to run a +``DELETE`` statement, which the read-only validation rejects. It is the ``tool_result`` +block from the ``UserMessage`` in the SDK stream of a real run (indented here; +``tool_use_id`` is the value from that run), so it shows the API-level message the model +reads, not the handler's return value or the MCP wire format: + +.. code-block:: json + + { + "type": "tool_result", + "tool_use_id": "toolu_01Sitdzj84kaH1CKYC8LFW2b", + "content": "The query tool failed: Statement type 'Delete' is not allowed. Allowed types: Select, Union, Intersect, Except, Describe, Show\nUse the list_tables and get_schema tools to inspect the database, then fix the query and try again.", + "is_error": true + } + +``SQLToolset``, ``HookToolset`` and ``ObjectStorageToolset`` allow one correction: a tool +that keeps failing ends the run with +``ToolCallError: query kept failing after 1 correction(s): ...``. The retry budget is +counted per run of +:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run`, and a model +turn counts as one failure however many of its calls fail, the same way the Anthropic and +Strands adapters count them. The turn is the ``AssistantMessage.message_id``; without one, +calls that run at the same time count as one failure. A call whose arguments fail the SDK's own JSON Schema validation never +reaches Airflow's tools, so it is not counted against the budget and not recorded in the +``common_ai.tool_calls`` metric; see `max_turns`_ below. + +Any other failure, such as a hook raising, ends the run. The Claude Agent SDK's own tool +dispatch catches every exception a handler raises and turns it into an error result the +model reads, so +:meth:`~airflow.providers.common.ai.tools.claude_agent_sdk.AirflowTools.run` has to work around +that to still fail the task: it raises +:class:`~airflow.providers.common.ai.tools.ToolCallError` once the model has read the +(masked) failure, rather than letting the stream continue to a later +``ResultMessage``. Airflow's own retry takes over once the task fails. + +A run can also end with a ``ResultMessage`` whose ``is_error`` is set, for example when +``max_turns`` is reached or the API call fails. ``run()`` does not return that result as +a success. Which exception you get depends on the stream: + +* ``query()``: the CLI exits non-zero right after an error result, and the SDK raises Review Comment: These bullets cover a single prompt. With a streamed prompt, the `run()` docstring says an error turn followed by a successful one raises `RuntimeError`, since the SDK drops its saved error result once the conversation moves on. A sentence after the bullets would match the docstring. -- 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]
