kaxil commented on code in PR #74026:
URL: https://github.com/apache/airflow/pull/74026#discussion_r4314046571
##########
providers/common/ai/docs/stability.rst:
##########
@@ -265,6 +265,9 @@ Everything this provider ships that is not in the table
above is experimental.
(:doc:`frameworks/openai_agents`)
- Written against OpenAI Agents SDK 0.20 and 0.23. The SDK is pre-1.0 and
changes
its tool and runner APIs between minor releases.
+ * - The Claude Agent SDK adapter (:doc:`frameworks/claude_agent_sdk`)
+ - New; the SDK is pre-1.0 and bundles the Claude Code CLI, whose tool
dispatch turns
Review Comment:
It's the SDK's in-process MCP server that turns a handler's exception into
an error result, not the CLI. The guide and the class docstring already say so;
something like "the SDK is pre-1.0, and its in-process MCP server turns a
handler's exception into ..." would match them.
##########
providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py:
##########
@@ -0,0 +1,331 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Give Airflow tools to an agent loop run by the Claude Agent SDK.
+
+.. note:: Experimental; see :mod:`airflow.providers.common.ai.tools`.
+"""
+
+from __future__ import annotations
+
+try:
+ from claude_agent_sdk import (
+ AssistantMessage,
+ ClaudeAgentOptions,
+ McpSdkServerConfig,
+ Message,
+ ResultMessage,
+ ToolAnnotations,
+ create_sdk_mcp_server,
+ tool,
+ )
+except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException(e)
+
+import copy
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
collect_tools
+from airflow.providers.common.ai.tools._from_toolset import tool_call_scope
+from airflow.providers.common.ai.utils.tool_definition import serialize_for_llm
+from airflow.providers.common.ai.utils.tool_metrics import calling_framework
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncIterable, Awaitable, Callable
+
+ from airflow.providers.common.ai.tools import ToolProvider
+
+log = logging.getLogger(__name__)
+
+__all__ = ["AirflowTools"]
+
+
+_MAX_ERROR_DETAIL_CHARS = 500
+_MAX_ERROR_ENTRIES = 5
+# The bundled CLI clamps ``maxResultSizeChars`` to 500_000 (observed on CLI
2.1.286; other
+# versions may differ), so a larger value is a no-op. With ``tools=[]`` the
model has no Read
+# tool, so a result over the ceiling that the CLI spills to a
``<persisted-output>`` file is lost.
+_MAX_INLINE_RESULT_CHARS = 500_000
Review Comment:
A result well under 500,000 can still kill the stream. The CLI puts each
tool result on one stdout line twice (`content` and `tool_use_result`),
JSON-escaped, and the SDK rejects any line over `max_buffer_size`, 1 MiB by
default.
I ran the bundled CLI from 0.2.165 against a fake Messages API with a tool
returning `serialize_for_llm`-style JSON rows:
```
350,000 chars -> OK
420,000 chars -> CLIJSONDecodeError: JSON message exceeded maximum buffer
size of 1048576 bytes
500,000 chars, max_buffer_size=8 MiB -> OK, model got all 500,000 inline
300,000 emoji -> model got a <persisted-output> stub
```
So an uncapped `HookToolset` or `ObjectStorageToolset` result fails the task
rather than reaching the model truncated. I think `build_options()` should
default `max_buffer_size` (and the guide should list it with the fields to set
by hand, and drop the "about 1 million characters"). The emoji case is because
the CLI counts the limit in UTF-16 units, so `_truncate_result_text` would need
to measure the same way.
##########
providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py:
##########
@@ -0,0 +1,331 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Give Airflow tools to an agent loop run by the Claude Agent SDK.
+
+.. note:: Experimental; see :mod:`airflow.providers.common.ai.tools`.
+"""
+
+from __future__ import annotations
+
+try:
+ from claude_agent_sdk import (
+ AssistantMessage,
+ ClaudeAgentOptions,
+ McpSdkServerConfig,
+ Message,
+ ResultMessage,
+ ToolAnnotations,
+ create_sdk_mcp_server,
+ tool,
+ )
+except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException(e)
+
+import copy
+import logging
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
collect_tools
+from airflow.providers.common.ai.tools._from_toolset import tool_call_scope
+from airflow.providers.common.ai.utils.tool_definition import serialize_for_llm
+from airflow.providers.common.ai.utils.tool_metrics import calling_framework
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncIterable, Awaitable, Callable
+
+ from airflow.providers.common.ai.tools import ToolProvider
+
+log = logging.getLogger(__name__)
+
+__all__ = ["AirflowTools"]
+
+
+_MAX_ERROR_DETAIL_CHARS = 500
+_MAX_ERROR_ENTRIES = 5
+# The bundled CLI clamps ``maxResultSizeChars`` to 500_000 (observed on CLI
2.1.286; other
+# versions may differ), so a larger value is a no-op. With ``tools=[]`` the
model has no Read
+# tool, so a result over the ceiling that the CLI spills to a
``<persisted-output>`` file is lost.
+_MAX_INLINE_RESULT_CHARS = 500_000
+
+
+def _truncate_result_text(text: str) -> str:
+ if len(text) <= _MAX_INLINE_RESULT_CHARS:
+ return text
+ notice = (
+ f"\n\n[Truncated: this result had {len(text)} characters and only the
beginning is shown. "
+ "Narrow the request, for example by adding filters, selecting fewer
columns or lowering "
+ "the LIMIT, to get a smaller result.]"
+ )
+ # The notice counts toward the ceiling: text kept up to the ceiling plus
the notice would
+ # still be spilled by the CLI.
+ return text[: _MAX_INLINE_RESULT_CHARS - len(notice)] + notice
+
+
+def _bound_detail(text: str) -> str:
+ omitted = len(text) - _MAX_ERROR_DETAIL_CHARS
+ return text if omitted <= 0 else
f"{text[:_MAX_ERROR_DETAIL_CHARS]}...(+{omitted} chars)"
+
+
+def _describe_error_result(result: ResultMessage) -> str:
+ errors = result.errors or []
+ omitted_errors = len(errors) - _MAX_ERROR_ENTRIES
+ fields: dict[str, Any] = {
+ "subtype": result.subtype,
+ "result": result.result,
+ "errors": errors[:_MAX_ERROR_ENTRIES],
+ "stop_reason": result.stop_reason,
+ "api_error_status": result.api_error_status,
+ "terminal_reason": result.terminal_reason,
+ }
+ details = []
+ for name, value in fields.items():
+ if not value:
+ continue
+ text = _bound_detail(repr(value)) if name in ("result", "errors") else
repr(value)
+ if name == "errors" and omitted_errors > 0:
+ text = f"{text} (+{omitted_errors} more)"
+ details.append(f"{name}={text}")
+ return f"The Claude Agent SDK run ended with an error result ({',
'.join(details)})."
+
+
+async def _close_stream(messages: AsyncIterable[Message]) -> None:
+ # ``AsyncIterable`` does not guarantee ``aclose``; ``query()`` and
``receive_response()`` do.
+ aclose = getattr(messages, "aclose", None)
+ if aclose is not None:
+ await aclose()
+
+
+class _Run:
+ """One run of :meth:`AirflowTools.run`: identifies its calls and collects
their failures."""
+
+ def __init__(self) -> None:
+ self.failures: list[ToolCallError] = []
+ self.turn: str | None = None
+
+
+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
+ annotations =
ToolAnnotations(maxResultSizeChars=_MAX_INLINE_RESULT_CHARS)
+ sdk_tools = [
+ tool(t.name, t.description, self._build_tool_schema(t),
annotations=annotations)(
+ self._build_handler(t)
+ )
+ for t in airflow_tools
+ ]
+ self.server: McpSdkServerConfig =
create_sdk_mcp_server(name=server_name, tools=sdk_tools)
+
+ def build_options(self, **kwargs: Any) -> ClaudeAgentOptions:
+ """
+ Build ``ClaudeAgentOptions`` wired to this instance's tools, with safe
defaults.
+
+ ``ClaudeAgentOptions()`` on its own loads every settings source on the
worker
+ (``~/.claude``, project settings, a project ``.mcp.json``) and leaves
the CLI's
+ built-in tools on, including a worker-host ``Bash``. This method
instead
+ defaults ``tools`` to ``[]`` (no built-in tools), ``setting_sources``
to ``[]``
+ (no settings files) and ``strict_mcp_config`` to ``True`` (only the
MCP servers
+ passed here), and ``verbatim_prompts`` to ``True``. It adds
:attr:`server` and
+ :attr:`allowed_tools` to whatever ``mcp_servers`` and
``allowed_tools`` the caller
+ passes. A keyword this method does not default, such as ``model`` or
``env``, is
+ passed through unchanged.
+
+ ``verbatim_prompts`` stops the CLI from expanding ``@path`` file
mentions in a prompt;
+ otherwise anyone who controls the prompt could send a file the worker
user can read to
+ the model, without passing the secret masker. It also turns off slash
command dispatch
+ and the CLI's turn-start attachment pass.
+
+ When ``can_use_tool`` is passed, the Airflow tools are not added to
``allowed_tools``,
+ because an allowed tool never reaches the callback, and
``permission_mode`` defaults to
+ ``"default"``. The callback then decides on every Airflow tool call;
the names are still
+ in :attr:`allowed_tools` if you want to allow some by hand.
+
+ Everything it builds can also be assembled by hand; see :attr:`server`
and
+ :attr:`allowed_tools`.
+
+ :raises ValueError: when ``mcp_servers`` already has a server named
+ :attr:`server_name`.
+ """
+ mcp_servers = dict(kwargs.pop("mcp_servers", {}))
+ if self.server_name in mcp_servers:
+ raise ValueError(
+ f"mcp_servers already has a server named {self.server_name!r}.
Pass a different "
+ "server_name to AirflowTools, or drop it from mcp_servers."
+ )
+ mcp_servers[self.server_name] = self.server
+ allowed_tools = list(kwargs.pop("allowed_tools", []))
+ if kwargs.get("can_use_tool") is None:
+ allowed_tools.extend(self.allowed_tools)
+ else:
+ kwargs.setdefault("permission_mode", "default")
+ kwargs.setdefault("tools", [])
+ kwargs.setdefault("setting_sources", [])
+ kwargs.setdefault("strict_mcp_config", True)
+ kwargs.setdefault("verbatim_prompts", True)
+ return ClaudeAgentOptions(mcp_servers=mcp_servers,
allowed_tools=allowed_tools, **kwargs)
+
+ async def run(self, messages: AsyncIterable[Message]) -> ResultMessage:
+ """
+ Drive ``messages``, the stream ``query()`` or
``ClaudeSDKClient.receive_response()`` returns.
+
+ ``ClaudeSDKClient.receive_messages()`` is not supported: it never ends.
+
+ ``messages`` is closed when this returns or raises. For ``query()``
that ends the CLI
+ session, so it stops asking for tool calls; a ``ClaudeSDKClient``
stays connected until
+ you disconnect it, for example by leaving its ``async with`` block.
+
+ :raises RuntimeError: when this instance is already running a query,
when
+ ``messages`` ends without a
:class:`~claude_agent_sdk.ResultMessage`, or, for
+ ``ClaudeSDKClient.receive_response()``, when the result message
has ``is_error``
+ set (for example ``max_turns`` reached or an API failure).
+ :raises ResultError: for ``query()`` when the result message has
``is_error`` set.
+ The CLI exits non-zero right after such a result, so the SDK raises
+ ``claude_agent_sdk.ResultError`` (a ``ProcessError``) from the
stream before this
+ method checks the result.
+ :raises ToolCallError: when a tool failed in a way the model cannot
fix.
+ """
+ if self._run is not None:
+ raise RuntimeError(
+ f"{self!r} is already running a query. Build a separate
AirflowTools instance "
+ "for a concurrent run, or await the first run before starting
another."
+ )
+ state = _Run()
+ self._run = state
+ raised = True
+ try:
+ result: ResultMessage | None = None
+ async for message in messages:
+ if isinstance(message, AssistantMessage):
+ # The CLI sends each turn's message before running that
turn's tools, which is
+ # what ``handle_tool_call`` relies on. Observed with the
bundled CLI, not documented.
+ state.turn = message.message_id
+ if isinstance(message, ResultMessage):
+ result = message
+ if state.failures:
+ raise state.failures[0]
+ if result is None:
+ if state.failures:
+ raise state.failures[0]
+ raise RuntimeError("The Claude Agent SDK stream ended without
a ResultMessage.")
+ if result.is_error:
+ raise RuntimeError(_describe_error_result(result))
+ raised = False
+ return result
+ finally:
+ try:
+ await _close_stream(messages)
Review Comment:
For `query()` this closes only the outer generator. `query()` is a plain
`async for ... yield` over `process_query`, so closing it doesn't close
`process_query`, and the SDK's `Query.close()` only runs later from asyncio's
async-gen finalizer, after `self._run = None` below. A `tools/call` that lands
in that gap still skips the "Not run" guard; the SDK's teardown happened to
cancel that handler before the tool body ran, but that's loop scheduling, not
something this code guarantees. Keeping the ended run visible to the handler
(say `self._last_run`) would close it without relying on teardown order. The
docstring's "so it stops asking for tool calls" is also only true after the
SDK's 5-second shutdown wait; the CLI made 5 more model requests in that window.
##########
providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py:
##########
@@ -0,0 +1,929 @@
+# 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_result_message_bounds_its_free_text_fields(self):
+ tools, _ = _build_tools()
+
+ async def _runaway() -> AsyncIterator[Message]:
+ yield _result_message(
+ is_error=True,
+ subtype="error_max_turns",
+ result="x" * 100_000,
+ errors=[f"error-{i}" for i in range(1_000)],
+ )
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_runaway())
+
+ message = str(exc_info.value)
+ assert len(message) < 2_000
+ assert "error_max_turns" in message
+ assert "(+99502 chars)" in message
+ assert "error-4" in message
+ assert "error-5" not in message
+ assert "(+995 more)" in message
+
+ @pytest.mark.asyncio
+ async def
test_an_api_failure_reported_as_a_success_subtype_still_raises(self):
+ tools, _ = _build_tools()
+
+ async def _api_failure() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True, subtype="success",
api_error_status=529)
+
+ with pytest.raises(RuntimeError, match="529"):
+ await tools.run(_api_failure())
+
+ @pytest.mark.asyncio
+ async def
test_an_error_result_message_leaves_unset_fields_out_of_the_message(self):
+ tools, _ = _build_tools()
+
+ async def _bare_error() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True,
subtype="error_during_execution")
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_bare_error())
+
+ message = str(exc_info.value)
+ assert "error_during_execution" in message
+ for unset in ("result=", "errors=", "stop_reason=",
"api_error_status=", "terminal_reason="):
+ assert unset not in message
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_failure_takes_priority_over_an_error_result_message(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ async def _failure_then_error_result() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _result_message(is_error=True, subtype="error_max_turns")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_failure_then_error_result())
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_that_keeps_failing_fails_the_run_once_its_retries_are_spent(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {})], [("step",
{})]]))
+
+ @pytest.mark.asyncio
+ async def
test_each_run_of_a_reused_instance_starts_with_a_fresh_retry_budget(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ for _ in range(2):
+ result = await tools.run(_fake_cli(handlers, [[("step", {})]]))
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_one_model_turn_count_as_one(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ result = await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=["msg_1"]))
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_different_model_turns_count_separately(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(
+ _fake_cli(handlers, [[("step", {})], [("step", {})]],
message_ids=["msg_1", "msg_2"])
+ )
+
+ @pytest.mark.asyncio
+ async def
test_sequential_failures_count_separately_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=[None]))
+
+ @pytest.mark.asyncio
+ async def
test_simultaneous_failures_count_as_one_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ async def _two_calls_at_once() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await asyncio.gather(handlers["step"]({}), handlers["step"]({}))
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_calls_at_once())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_real_server_path_sees_the_turn_the_run_recorded(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools = AirflowTools(step)
+
+ async def _two_failures_in_one_turn() -> AsyncIterator[Message]:
+ yield _assistant_message("msg_1")
+ await _call_real_server(tools, "step", {})
+ await _call_real_server(tools, "step", {})
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_failures_in_one_turn())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_stream_is_closed_while_the_run_is_still_active(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+ closed_during_run: list[bool] = []
+
+ async def _failing_cli() -> AsyncIterator[Message]:
+ try:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _user_message()
+ yield _result_message()
+ finally:
+ closed_during_run.append(tools._run is not None)
+
+ stream = _failing_cli()
+ with pytest.raises(ToolCallError):
+ await tools.run(stream)
+
+ assert closed_during_run == [True]
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_stream_without_aclose_is_driven_to_the_end(self):
+ class _PlainStream:
+ def __init__(self, messages: list[Message]) -> None:
+ self._messages = messages
+
+ def __aiter__(self) -> AsyncIterator[Message]:
+ async def _iterate() -> AsyncIterator[Message]:
+ for message in self._messages:
+ yield message
+
+ return _iterate()
+
+ final = _result_message(result="done")
+
+ result = await
AirflowTools(_tool(ToolResult("ok"))).run(_PlainStream([_assistant_message(),
final]))
+
+ assert result is final
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_does_not_hide_the_tools_failure(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_after_a_clean_run_raises(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli({}, [])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(OSError, match="close failed"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def
test_a_failing_aclose_is_logged_when_it_does_not_hide_the_tools_failure(self,
caplog):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError):
+ await tools.run(_StubbornStream())
+
+ assert {
Review Comment:
This is what fails the Compat 2.11.1 and 3.0.6 jobs: `TypeError: argument of
type 'LogCaptureFixture' is not iterable`. `x in caplog` needs the structlog
capture fixture, which falls back to the stock `LogCaptureFixture` before 3.1,
and the provider supports 2.11. Patching `log` with `autospec`, as the test at
line 304 does, works on every version.
##########
providers/common/ai/tests/unit/common/ai/tools/test_claude_agent_sdk.py:
##########
@@ -0,0 +1,929 @@
+# 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_result_message_bounds_its_free_text_fields(self):
+ tools, _ = _build_tools()
+
+ async def _runaway() -> AsyncIterator[Message]:
+ yield _result_message(
+ is_error=True,
+ subtype="error_max_turns",
+ result="x" * 100_000,
+ errors=[f"error-{i}" for i in range(1_000)],
+ )
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_runaway())
+
+ message = str(exc_info.value)
+ assert len(message) < 2_000
+ assert "error_max_turns" in message
+ assert "(+99502 chars)" in message
+ assert "error-4" in message
+ assert "error-5" not in message
+ assert "(+995 more)" in message
+
+ @pytest.mark.asyncio
+ async def
test_an_api_failure_reported_as_a_success_subtype_still_raises(self):
+ tools, _ = _build_tools()
+
+ async def _api_failure() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True, subtype="success",
api_error_status=529)
+
+ with pytest.raises(RuntimeError, match="529"):
+ await tools.run(_api_failure())
+
+ @pytest.mark.asyncio
+ async def
test_an_error_result_message_leaves_unset_fields_out_of_the_message(self):
+ tools, _ = _build_tools()
+
+ async def _bare_error() -> AsyncIterator[Message]:
+ yield _result_message(is_error=True,
subtype="error_during_execution")
+
+ with pytest.raises(RuntimeError) as exc_info:
+ await tools.run(_bare_error())
+
+ message = str(exc_info.value)
+ assert "error_during_execution" in message
+ for unset in ("result=", "errors=", "stop_reason=",
"api_error_status=", "terminal_reason="):
+ assert unset not in message
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_failure_takes_priority_over_an_error_result_message(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ async def _failure_then_error_result() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _result_message(is_error=True, subtype="error_max_turns")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_failure_then_error_result())
+
+ @pytest.mark.asyncio
+ async def
test_a_tool_that_keeps_failing_fails_the_run_once_its_retries_are_spent(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {})], [("step",
{})]]))
+
+ @pytest.mark.asyncio
+ async def
test_each_run_of_a_reused_instance_starts_with_a_fresh_retry_budget(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ for _ in range(2):
+ result = await tools.run(_fake_cli(handlers, [[("step", {})]]))
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_one_model_turn_count_as_one(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ result = await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=["msg_1"]))
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_failures_in_different_model_turns_count_separately(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(
+ _fake_cli(handlers, [[("step", {})], [("step", {})]],
message_ids=["msg_1", "msg_2"])
+ )
+
+ @pytest.mark.asyncio
+ async def
test_sequential_failures_count_separately_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ await tools.run(_fake_cli(handlers, [[("step", {}), ("step",
{})]], message_ids=[None]))
+
+ @pytest.mark.asyncio
+ async def
test_simultaneous_failures_count_as_one_when_the_turn_has_no_message_id(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools, handlers = _build_tools(step)
+
+ async def _two_calls_at_once() -> AsyncIterator[Message]:
+ yield _assistant_message()
+ await asyncio.gather(handlers["step"]({}), handlers["step"]({}))
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_calls_at_once())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_real_server_path_sees_the_turn_the_run_recorded(self):
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+ tools = AirflowTools(step)
+
+ async def _two_failures_in_one_turn() -> AsyncIterator[Message]:
+ yield _assistant_message("msg_1")
+ await _call_real_server(tools, "step", {})
+ await _call_real_server(tools, "step", {})
+ yield _user_message()
+ yield _result_message()
+
+ result = await tools.run(_two_failures_in_one_turn())
+
+ assert isinstance(result, ResultMessage)
+
+ @pytest.mark.asyncio
+ async def test_the_stream_is_closed_while_the_run_is_still_active(self):
+ ran: list[dict[str, Any]] = []
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders"), ran))
+ closed_during_run: list[bool] = []
+
+ async def _failing_cli() -> AsyncIterator[Message]:
+ try:
+ yield _assistant_message()
+ await handlers["lookup"]({"key": "a"})
+ yield _user_message()
+ yield _result_message()
+ finally:
+ closed_during_run.append(tools._run is not None)
+
+ stream = _failing_cli()
+ with pytest.raises(ToolCallError):
+ await tools.run(stream)
+
+ assert closed_during_run == [True]
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_stream_without_aclose_is_driven_to_the_end(self):
+ class _PlainStream:
+ def __init__(self, messages: list[Message]) -> None:
+ self._messages = messages
+
+ def __aiter__(self) -> AsyncIterator[Message]:
+ async def _iterate() -> AsyncIterator[Message]:
+ for message in self._messages:
+ yield message
+
+ return _iterate()
+
+ final = _result_message(result="done")
+
+ result = await
AirflowTools(_tool(ToolResult("ok"))).run(_PlainStream([_assistant_message(),
final]))
+
+ assert result is final
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_does_not_hide_the_tools_failure(self):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def test_a_failing_aclose_after_a_clean_run_raises(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli({}, [])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(OSError, match="close failed"):
+ await tools.run(_StubbornStream())
+
+ assert tools._run is None
+
+ @pytest.mark.asyncio
+ async def
test_a_failing_aclose_is_logged_when_it_does_not_hide_the_tools_failure(self,
caplog):
+ tools, handlers = _build_tools(_tool(PermissionError("role cannot read
orders")))
+
+ class _StubbornStream:
+ def __aiter__(self) -> AsyncIterator[Message]:
+ return _fake_cli(handlers, [[("lookup", {"key": "a"})]])
+
+ async def aclose(self) -> None:
+ raise OSError("close failed")
+
+ with pytest.raises(ToolCallError):
+ await tools.run(_StubbornStream())
+
+ assert {
+ "event": "Closing the Claude Agent SDK stream failed: close
failed",
+ "log_level": "warning",
+ } in caplog
+
+ @pytest.mark.asyncio
+ async def test_a_cancelled_run_closes_the_stream_and_stays_cancelled(self):
+ tools, _ = _build_tools(_tool(ToolResult("ok")))
+ started = asyncio.Event()
+ closed: list[bool] = []
+
+ async def _blocked_cli() -> AsyncIterator[Message]:
+ try:
+ yield _assistant_message()
+ started.set()
+ await asyncio.Event().wait()
+ yield _result_message()
+ finally:
+ closed.append(True)
+
+ task = asyncio.ensure_future(tools.run(_blocked_cli()))
+ await started.wait()
+ task.cancel()
+
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+ assert closed == [True]
Review Comment:
This still passes with `_close_stream` removed: cancelling `run()` throws
`CancelledError` into the generator's pending `__anext__`, so its own `finally`
runs either way. A class-based stream like `_StubbornStream` whose `aclose`
records the close would make this assertion check the close.
--
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]
