This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3259e8960f8 Mask tool output from every capability in Common AI agents
(#73993)
3259e8960f8 is described below
commit 3259e8960f8e7a0a4b257fad3b6e321c2a1fffbf
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 12:39:22 2026 +0100
Mask tool output from every capability in Common AI agents (#73993)
Agents built by Common AI masked tool results only for toolsets
AgentOperator
could find and wrap: toolsets=, agent_params toolsets and a Toolset
capability
holding a toolset. Tools that arrive through any other capability (MCP,
PrefixTools, CombinedCapability, a Toolset built per run), function tools in
agent_params, and every tool given to LLMOperator and the other LLM
operators
reached the model unmasked. PydanticAIHook.create_agent now adds a
MaskingCapability that masks around every tool call, innermost so other
capabilities see only masked results. It also masks the retry or failure
message pydantic-ai builds from a tool's ModelRetry or ToolFailed, and a
ToolReturn's metadata.
Masking a dataclass result now copies it and masks the copy's fields instead
of rebuilding it with dataclasses.replace, which failed on an InitVar
without
a default and reset init=False fields.
* Fix review findings on Common AI tool-output masking
Running out of tool retries is no longer logged as the tool failing: the
UnexpectedModelBehavior pydantic-ai raises below MaskingCapability is logged
at debug and stripped like a retry. mask_secrets keeps a Pydantic model that
holds no registered secret as the model, so a ToolReturn's metadata or a
model a tool returns still has its attributes in the message history.
The comment on why AgentOperator keeps toolset-level masking now names the
real reason, toolset wrappers a capability adds below the capability hooks;
the durable journal masks what it records itself, and its test now covers a
leaf toolset AirflowDurability journals as well as a durable unit. The
MaskingCapability docstring credits create_agent, shows how to add it to an
agent built with Agent(...), and names what still sees a result before it is
masked. The docs list output functions as not masked, and the Pydantic AI
framework page and the capture_content option no longer contradict the
masking create_agent now does.
* Simplify Common AI tool-output masking
Mask a ToolReturn and the message part inside a ToolRetryError or
ToolFailedError through mask_secrets' own dataclass handling instead of
listing fields by hand, which also masks a failed ToolReturnPart's
metadata. Move the comment on why AgentOperator keeps toolset-level
masking next to that code, and drop an operator test the hook test
already covers.
---
providers/common/ai/docs/agent_security.rst | 18 +-
providers/common/ai/docs/capabilities.rst | 6 +-
.../common/ai/docs/frameworks/pydantic_ai.rst | 22 ++-
providers/common/ai/docs/observability.rst | 10 +-
providers/common/ai/provider.yaml | 8 +-
.../providers/common/ai/durable/fingerprint.py | 7 +-
.../providers/common/ai/get_provider_info.py | 2 +-
.../providers/common/ai/hooks/pydantic_ai.py | 10 ++
.../airflow/providers/common/ai/operators/agent.py | 4 +
.../airflow/providers/common/ai/utils/masking.py | 26 ++-
.../providers/common/ai/utils/toolset_base.py | 83 +++++++++-
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 45 +++++
.../tests/unit/common/ai/operators/test_agent.py | 182 +++++++++++++++++++--
.../ai/tests/unit/common/ai/utils/test_masking.py | 51 +++++-
.../unit/common/ai/utils/test_toolset_base.py | 42 ++++-
15 files changed, 459 insertions(+), 57 deletions(-)
diff --git a/providers/common/ai/docs/agent_security.rst
b/providers/common/ai/docs/agent_security.rst
index 58a26dc17b9..9ce4a350a7f 100644
--- a/providers/common/ai/docs/agent_security.rst
+++ b/providers/common/ai/docs/agent_security.rst
@@ -80,16 +80,22 @@ No single layer is sufficient on its own. They work
together.
- What a tool returns, and the error text handed to the model so it can
correct a
call, pass through Airflow's secret masker first. A connection password
that shows
up in a database error or a hook's return value reaches the model, the
model
- provider and any trace as ``***``. ``AgentOperator`` applies this to
the toolsets
- you pass in ``toolsets``, in ``agent_params["toolsets"]`` and in a
``Toolset``
- capability, including your own. The SQL, hook, object storage,
DataFusion, MCP,
+ provider and any trace as ``***``. Every agent this provider builds
applies this
+ to what each tool returns or raises while it runs, however the tool
reaches the
+ agent: ``toolsets``, ``agent_params``, or any capability, including
your own.
+ That covers ``AgentOperator``, the other LLM operators, an agent you
build with
+ ``PydanticAIHook.create_agent``, and one you build with ``Agent(...)``
and give
+ ``MaskingCapability`` (see :doc:`frameworks/pydantic_ai`). The SQL,
hook, object storage, DataFusion, MCP,
sandbox and managed-agent toolsets apply it wherever they run,
including in a Pydantic AI agent
you build yourself.
- Masks only secrets Airflow has registered, such as connection passwords
and
sensitive connection extras. A credential that exists only in the data
itself is
- not recognized. Not masked: prompts, model output, function tools
passed as
- ``agent_params["tools"]``, a ``Toolset`` capability built by a
function, a
- framework's own tools, and MCP servers a framework connects to itself.
In a
+ not recognized. Not masked: prompts, model output, errors raised while
validating
+ a tool's arguments or the agent's output, what an output function
(``ToolOutput``
+ or a function ``output_type``) returns or raises, a result a capability
returns
+ in place of running the tool, tools the model provider runs itself
(native ``WebSearch``,
+ ``WebFetch`` or ``ImageGeneration``, and ``MCP(native=True)``), a
framework's own
+ tools, and MCP servers a framework connects to itself. In a
sandbox, the model writes the commands, so it can print a secret in a
form the
masker does not recognize; masking there guards against accidents only.
* - **ObjectStorageToolset: one read-only root**
diff --git a/providers/common/ai/docs/capabilities.rst
b/providers/common/ai/docs/capabilities.rst
index 12bf47565b6..43106a7cf78 100644
--- a/providers/common/ai/docs/capabilities.rst
+++ b/providers/common/ai/docs/capabilities.rst
@@ -85,9 +85,9 @@ templated inside a ``Toolset`` capability the same way as in
``toolsets=``:
A ``Toolset`` capability built from a function is resolved when the run
starts, so its
connection IDs are not templated.
-Tool results from a ``Toolset`` capability are masked like those from
``toolsets=``.
-When ``enable_tool_logging=True`` (the default), ``AgentOperator`` logs calls
to
-these toolsets and other tools contributed by capabilities, including MCP
+Tool results from any capability whose tools run in the worker are masked like
those
+from ``toolsets=``. When ``enable_tool_logging=True`` (the default),
``AgentOperator``
+logs calls to these toolsets and other tools contributed by capabilities,
including MCP
toolsets. Output tools and provider-native tools, including native MCP, are not
covered. See :doc:`toolsets/logging` for configuration details and the complete
limitations. Put a toolset in the capability list when you need to order it
diff --git a/providers/common/ai/docs/frameworks/pydantic_ai.rst
b/providers/common/ai/docs/frameworks/pydantic_ai.rst
index f88e3d30887..0ae4bd8cf98 100644
--- a/providers/common/ai/docs/frameworks/pydantic_ai.rst
+++ b/providers/common/ai/docs/frameworks/pydantic_ai.rst
@@ -42,12 +42,23 @@ connections and self-hosted endpoints as ``AgentOperator``;
see :doc:`../model_p
Every toolset this provider ships is a Pydantic AI toolset, so it goes into
``toolsets=`` as it is, next to any toolsets and tools of your own.
-For the OpenTelemetry spans ``AgentOperator`` emits, build the agent with the
hook's
+For the OpenTelemetry spans and secret masking ``AgentOperator`` gets, build
the agent
+with the hook's
:meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.create_agent`
instead of ``Agent(...)``. It takes the same arguments, and under ``[common.ai]
otel_export_enabled`` it sends the agent's spans through Airflow's tracing,
without prompt
text unless ``capture_content`` is on; see :doc:`../observability`. Pydantic
AI's own
-instrumentation records prompt and completion text by default.
+instrumentation records prompt and completion text by default. It also passes
what every
+tool returns or raises through Airflow's secret masker, whichever toolset or
capability
+supplied the tool. To get that masking on an agent you build with
``Agent(...)``, add
+:class:`~airflow.providers.common.ai.utils.toolset_base.MaskingCapability`
last in its
+``capabilities``:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.utils.toolset_base import
MaskingCapability
+
+ agent = Agent(model, toolsets=[my_toolset],
capabilities=[MaskingCapability()])
What you get, and what you do not
---------------------------------
@@ -57,8 +68,8 @@ as they do inside ``AgentOperator``: SQL validation,
``allowed_tables``, result
object-storage path checks, and the secret masker on everything a tool
returns, including
the text of a failure the model is asked to correct. Their calls are counted
in the
``common_ai.tool_calls`` metric described in :doc:`../observability`. The
Agent Skills
-toolset is the exception: its results are masked only inside
``AgentOperator``, and its
-calls are not counted.
+toolset is the exception: its results are masked only in an agent built with
+``create_agent`` or given ``MaskingCapability``, and its calls are not counted.
``AgentOperator`` adds things on top of the agent that a task running your own
agent
does not get:
@@ -66,9 +77,6 @@ does not get:
- Durable replay of model and tool steps across task retries
(``durable=True``).
- Human review of the output (``enable_hitl_review``), and a pause for a
person to
approve marked tool calls before they run (:doc:`../tool_approval`).
-- Masking of the results of the Agent Skills toolset and of toolsets you wrote
yourself.
- There is no public masking wrapper, so run the agent through
``AgentOperator`` if their
- results can carry a secret.
- Rendering of templated connection IDs in toolsets, such as ``SQLToolset("{{
... }}")``.
In your own task, pass the connection ID itself.
- Tool call logging, and the task's identity (``airflow.dag_id``,
``airflow.task_id``
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index bd964c7dd12..b55aeb112b4 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -169,12 +169,14 @@ outputs (``gen_ai.input.messages`` /
``gen_ai.output.messages``), set:
.. warning::
- With ``capture_content`` enabled, prompts, completions, and tool IO are
+ With ``capture_content`` enabled, prompts, completions, and tool arguments
are
exported to your tracing backend **without redaction**. Airflow's secret
masking
applies to logs and rendered template fields, not to OpenTelemetry span
- attributes, so it does not scrub this content. Enable it only for debugging
- in a trusted environment. It has no effect unless ``otel_export_enabled``
is
- ``True``.
+ attributes. The one exception is what a function tool returns or raises
while it
+ runs in an agent this provider builds, which is masked before the span
records it.
+ An output function's result is not (see :doc:`agent_security`). Enable it
only for
+ debugging in a trusted environment. It has no effect unless
``otel_export_enabled``
+ is ``True``.
Agents built with other frameworks
----------------------------------
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index ee1e188fc1b..fc902b7341b 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -159,9 +159,11 @@ config:
finish reason are recorded, never message text. Turning this on
exports
model inputs and outputs to your tracing backend without redaction.
Airflow's
secret masking applies to logs and rendered template fields, not to
- OpenTelemetry span attributes, so it does not scrub this content.
- Enable it only for debugging in a trusted environment. Has no effect
- unless ``otel_export_enabled`` is ``True``.
+ OpenTelemetry span attributes. The exception is what a function tool
+ returns or raises while it runs in an agent this provider builds,
which
+ is masked before the span records it. Enable it only for debugging in
+ a trusted environment. Has no effect unless ``otel_export_enabled``
is
+ ``True``.
version_added: 0.4.0
type: boolean
example: "False"
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
index a7462216d99..3082f4df64f 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
@@ -156,9 +156,10 @@ def fingerprint_model_request(
return _digest(
{
"model": model_identifier,
- # The journal records tool results with secrets masked, while
a live run hands
- # the model the unmasked result, so a replayed history holds
``***`` where the
- # live one held the secret. Masking the history here as well
makes both match.
+ # The journal records tool results with secrets masked, and
not everything in a
+ # live run's history is (an output function's result, for
one), so a replayed
+ # history can hold ``***`` where the live one held the secret.
Masking the
+ # history here as well makes both match.
"messages": _strip_volatile(mask_secrets(dumped)),
"settings": _content_settings(model_settings),
"params": _normalize_params(params),
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
index 07bba3befff..3af5ce60876 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
@@ -133,7 +133,7 @@ def get_provider_info():
"default": "False",
},
"capture_content": {
- "description": "Capture prompt, completion, and
tool-call content on the emitted\nGenAI spans (``gen_ai.input.messages`` /
``gen_ai.output.messages``).\n\nOff by default: only token counts, model id,
latency, tool names, and\nfinish reason are recorded, never message text.
Turning this on exports\nmodel inputs and outputs to your tracing backend
without redaction. Airflow's\nsecret masking applies to logs and rendered
template fields, not to\nOpenTelemetry span at [...]
+ "description": "Capture prompt, completion, and
tool-call content on the emitted\nGenAI spans (``gen_ai.input.messages`` /
``gen_ai.output.messages``).\n\nOff by default: only token counts, model id,
latency, tool names, and\nfinish reason are recorded, never message text.
Turning this on exports\nmodel inputs and outputs to your tracing backend
without redaction. Airflow's\nsecret masking applies to logs and rendered
template fields, not to\nOpenTelemetry span at [...]
"version_added": "0.4.0",
"type": "boolean",
"example": "False",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
index f44ae25ce54..3a6a91404d8 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
@@ -27,6 +27,7 @@ from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.ai.observability import
genai_instrumentation_settings
+from airflow.providers.common.ai.utils.toolset_base import MaskingCapability
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import BaseHook
@@ -611,6 +612,12 @@ class PydanticAIHook(BaseHook):
GenAI spans through Airflow's tracing pipeline. See
:mod:`airflow.providers.common.ai.observability`.
+ What each tool returns or raises while it runs passes through
Airflow's secret
+ masker before the model sees it, whichever toolset or capability
supplied the
+ tool: the agent gets a
+
:class:`~airflow.providers.common.ai.utils.toolset_base.MaskingCapability` after
+ any ``capabilities`` passed in.
+
:param output_type: The expected output type from the agent (default:
``str``).
:param instructions: System-level instructions for the agent.
Required when *spec_file* is not given. When *spec_file* is given,
@@ -633,6 +640,9 @@ class PydanticAIHook(BaseHook):
# still wins over the provider's auto-instrumentation.
caller_instrument = agent_kwargs.pop("instrument", _UNSET)
+ # Last, so it is the innermost capability; see MaskingCapability.
+ agent_kwargs["capabilities"] = [*(agent_kwargs.get("capabilities") or
()), MaskingCapability()]
+
if spec_file is not None:
from_file_kwargs = dict(agent_kwargs)
model = self._get_conn_if_model_configured()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 9078c2e9f85..8bc303e6494 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -782,6 +782,10 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
# pydantic-ai wraps capability hooks in list order, so merging the
two lists
# would pick an order the Dag author never wrote down.
raise ValueError("Pass capabilities either as capabilities=... or
in agent_params, not both.")
+ # create_agent adds a MaskingCapability that masks every tool call.
The toolsets are
+ # masked here as well because toolset wrappers a capability adds
(get_wrapper_toolset,
+ # as ToolLoggingCapability does) run below the capability hooks and
would otherwise see
+ # the unmasked result. The durable journal masks what it records on
its own.
if self.toolsets:
toolsets: list[AbstractToolset] = [ensure_masked(ts) for ts in
self.toolsets]
extra_kwargs["toolsets"] = toolsets
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/masking.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/masking.py
index 9fe2de9ebdc..90916a03a39 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/masking.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/masking.py
@@ -18,6 +18,7 @@
from __future__ import annotations
+import copy
import dataclasses
import json
from typing import Any, overload
@@ -46,7 +47,8 @@ def mask_secrets(value: Any) -> Any:
Strings are masked wherever they sit in nested dicts, lists, tuples, sets
and dataclasses,
dict keys included, and the shape and types are kept. Bytes are masked as
UTF-8 text. A
- Pydantic model is turned into the JSON-compatible data the model would be
shown. Images,
+ Pydantic model that holds a registered secret is turned into the
JSON-compatible data the
+ model would be shown, with the secret masked; one that holds none is
returned as it is. Images,
documents and other multimodal content, and any other object, pass through
as they are.
Registered secrets are the ones Airflow knows about, such as connection
passwords and
sensitive connection extras; a credential that only appears in the data
itself is not
@@ -78,6 +80,9 @@ def dumps_masked(value: Any, **kwargs: Any) -> str:
def _json_default(value: object) -> Any:
if isinstance(value, bytes):
return value.decode("utf-8", "replace")
+ if isinstance(value, BaseModel):
+ # Only reached for a model mask_secrets kept as it was, so one holding
no secret.
+ return to_jsonable_python(value)
if dataclasses.is_dataclass(value) and not isinstance(value, type):
return {field.name: getattr(value, field.name) for field in
dataclasses.fields(value)}
return mask_secrets(str(value))
@@ -91,7 +96,11 @@ def _mask(value: Any, seen: frozenset[int]) -> Any:
masked = mask_secrets(text)
return value if masked == text else masked.encode("utf-8",
"surrogateescape")
if isinstance(value, BaseModel):
- return _mask(to_jsonable_python(value), seen)
+ data = to_jsonable_python(value)
+ masked = _mask(data, seen)
+ # Kept as the model when there was nothing to mask, so code that reads
its attributes
+ # later, such as a ToolReturn's metadata in the message history, still
can.
+ return value if masked == data else masked
if isinstance(value, (dict, list, tuple, set, frozenset)) or
_is_masked_dataclass(value):
if id(value) in seen:
return _CYCLE
@@ -105,10 +114,15 @@ def _mask(value: Any, seen: frozenset[int]) -> Any:
if isinstance(value, (set, frozenset)):
return type(value)(_mask(item, seen) for item in value)
if _is_masked_dataclass(value):
- # Rebuilt rather than turned into a dict, so a type the framework acts
on, such as
- # pydantic-ai's TextContent, still is one.
- fields = [field for field in dataclasses.fields(value) if field.init]
- return dataclasses.replace(value, **{f.name: _mask(getattr(value,
f.name), seen) for f in fields})
+ # A copy with its fields masked, rather than a dict, so a type the
framework acts on,
+ # such as pydantic-ai's TextContent, still is one. Not
dataclasses.replace, which fails
+ # on an InitVar without a default and resets fields declared with
init=False.
+ masked = copy.copy(value)
+ for field in dataclasses.fields(value):
+ if hasattr(value, field.name):
+ # object.__setattr__ so a frozen dataclass can be masked too.
+ object.__setattr__(masked, field.name, _mask(getattr(value,
field.name), seen))
+ return masked
return value
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
index 10caa772cc4..4a90d956d6c 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
@@ -19,15 +19,22 @@
from __future__ import annotations
import asyncio
-import dataclasses
import logging
import threading
from abc import abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, TypeVar
-from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry,
ToolFailed
-from pydantic_ai.messages import ToolReturn
+from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering
+from pydantic_ai.exceptions import (
+ ApprovalRequired,
+ CallDeferred,
+ ModelRetry,
+ ToolFailed,
+ ToolFailedError,
+ ToolRetryError,
+ UnexpectedModelBehavior,
+)
from pydantic_ai.toolsets import DynamicToolset
from pydantic_ai.toolsets.abstract import AbstractToolset
from pydantic_ai.toolsets.wrapper import WrapperToolset
@@ -41,6 +48,9 @@ if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from pydantic_ai._run_context import RunContext
+ from pydantic_ai.capabilities.abstract import ValidatedToolArgs,
WrapToolExecuteHandler
+ from pydantic_ai.messages import ToolCallPart
+ from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets import ToolsetFunc
from pydantic_ai.toolsets.abstract import ToolsetTool
@@ -93,6 +103,17 @@ async def _mask_call(name: str, call: Awaitable[Any], *,
count_as: str | None =
except (ModelRetry, ToolFailed) as e:
log.debug("Tool %s returned an error for the model", name, exc_info=e)
error = _strip(e)
+ except (ToolRetryError, ToolFailedError) as e:
+ # The form a ModelRetry or ToolFailed takes by the time a capability's
+ # wrap_tool_execute sees it. The model is sent the message part it
carries, not the
+ # exception's text, so that part is what gets masked.
+ log.debug("Tool %s returned an error for the model", name, exc_info=e)
+ error = _mask_tool_error_part(e)
+ except UnexpectedModelBehavior as e:
+ # What pydantic-ai raises, below a capability's wrap_tool_execute,
once a tool has
+ # used up its retries. The run fails with it, so it is not logged as
the tool failing.
+ log.debug("Tool %s ran out of retries", name, exc_info=e)
+ error = _strip(e)
except Exception as e:
if not getattr(e, _STRIPPED, False):
log.warning("Tool %s failed", name, exc_info=e)
@@ -103,13 +124,18 @@ async def _mask_call(name: str, call: Awaitable[Any], *,
count_as: str | None =
if error is not None:
# Raised outside the except blocks, so Python does not chain the
original back on.
raise error
- if isinstance(result, ToolReturn):
- return dataclasses.replace(
- result, return_value=mask_secrets(result.return_value),
content=mask_secrets(result.content)
- )
+ # A ToolReturn is masked field by field, its metadata included: not sent
to the model,
+ # but kept in the message history and recorded in traces.
return mask_secrets(result)
+def _mask_tool_error_part(error: ToolRetryError | ToolFailedError) ->
ToolRetryError | ToolFailedError:
+ """Rebuild ``error`` around a masked copy of the message part it
carries."""
+ if isinstance(error, ToolRetryError):
+ return ToolRetryError(mask_secrets(error.tool_retry))
+ return ToolFailedError(mask_secrets(error.tool_failed))
+
+
def _strip(error: Exception) -> Exception:
"""
Mask what ``error`` would print and drop its cause chain.
@@ -259,6 +285,49 @@ class MaskingToolset(WrapperToolset[Any]):
return await _mask_call(name, self.wrapped.call_tool(name, tool_args,
ctx, tool))
+@dataclass
+class MaskingCapability(AbstractCapability[Any]):
+ """
+ Mask what every tool hands the model, however the tool reached the agent.
+
+ :class:`MaskingToolset` only covers a toolset the operator can find and
wrap. A
+ capability can supply tools the operator never sees as a toolset: ``MCP``,
a
+ ``Toolset`` built per run, or a ``Toolset`` inside ``PrefixTools`` or a
+ ``CombinedCapability``. This masks around every tool call instead.
+
+
:meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.create_agent`
+ adds it to every agent it builds, so add it yourself only to an agent you
build with
+ ``Agent(...)``, last in its ``capabilities``:
+
+ .. code-block:: python
+
+ agent = Agent(model, toolsets=[my_toolset],
capabilities=[*my_capabilities, MaskingCapability()])
+
+ It declares the innermost position, and pydantic-ai breaks ties between
capabilities
+ that declare it by list order, the last one innermost, so the agent's
other capabilities
+ hooking tool execution see a result only after it is masked. Two kinds
still see it
+ first: a capability passed to ``run(capabilities=...)``, which wraps
inside the agent's
+ own, and a toolset wrapper a capability adds with ``get_wrapper_toolset``,
which runs
+ below every capability's hooks.
+ """
+
+ def get_ordering(self) -> CapabilityOrdering:
+ """Join the innermost capabilities; list order decides which of them
is closest to the tool."""
+ return CapabilityOrdering(position="innermost")
+
+ async def wrap_tool_execute(
+ self,
+ ctx: RunContext[Any],
+ *,
+ call: ToolCallPart,
+ tool_def: ToolDefinition,
+ args: ValidatedToolArgs,
+ handler: WrapToolExecuteHandler,
+ ) -> Any:
+ """Run the tool and mask its result, or the error it raised, before
anything else sees it."""
+ return await _mask_call(call.tool_name, handler(args))
+
+
def ensure_masked(toolset: AbstractToolset[Any] | ToolsetFunc[Any]) ->
AbstractToolset[Any]:
"""
Return ``toolset`` wrapped in :class:`MaskingToolset`, unless it already
masks its own output.
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
index 4b703301b87..098408e02d6 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@@ -26,7 +26,9 @@ from unittest.mock import MagicMock, call, patch
import pytest
from pydantic_ai import Agent
+from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.exceptions import FallbackExceptionGroup, ModelAPIError
+from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart,
ToolReturnPart
from pydantic_ai.models import Model
from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.models.function import FunctionModel
@@ -44,6 +46,7 @@ from airflow.providers.common.ai.hooks.pydantic_ai import (
_has_recognized_provider_prefix,
_looks_like_unrecognized_provider_prefix,
)
+from airflow.providers.common.ai.utils.toolset_base import MaskingCapability
from airflow.providers.common.compat.module_loading import import_string,
qualname
from airflow.providers.common.compat.sdk import AirflowNotFoundException,
BaseHook
from airflow.providers_manager import HookInfo
@@ -1273,6 +1276,7 @@ class TestPydanticAIHookCreateAgent:
mock_model,
output_type=str,
instructions="You are a helpful assistant.",
+ capabilities=[MaskingCapability()],
)
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
@@ -1298,8 +1302,42 @@ class TestPydanticAIHookCreateAgent:
output_type=dict,
instructions="Be helpful.",
retries=3,
+ capabilities=[MaskingCapability()],
)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ def test_create_agent_adds_masking_after_the_callers_capabilities(self,
mock_agent_cls, mock_infer_model):
+ """Last, so it is innermost among capabilities that also ask to be."""
+ mock_infer_model.return_value = MagicMock(spec=Model)
+ callers = [MagicMock(spec=AbstractCapability)]
+ hook = PydanticAIHook(llm_conn_id="test_conn",
model_id="openai:gpt-5.6-sol")
+ conn = Connection(conn_id="test_conn", conn_type="pydanticai")
+ with patch.object(hook, "get_connection", autospec=True,
return_value=conn):
+ hook.create_agent(instructions="hi", capabilities=callers)
+
+ assert mock_agent_cls.call_args.kwargs["capabilities"] == [*callers,
MaskingCapability()]
+ assert len(callers) == 1
+
+ @pytest.mark.enable_redact
+ def test_create_agent_masks_what_a_tool_returns(self, registered_secret):
+ """Operators other than AgentOperator pass agent_params tools straight
to create_agent."""
+
+ def read_setting() -> str:
+ return f"api key: {registered_secret}"
+
+ def call_then_echo(messages, info):
+ returns = [p for m in messages for p in m.parts if isinstance(p,
ToolReturnPart)]
+ if returns:
+ return
ModelResponse(parts=[TextPart(content=str(returns[-1].content))])
+ return ModelResponse(parts=[ToolCallPart(tool_name="read_setting",
args={}, tool_call_id="c1")])
+
+ hook = PydanticAIHook(llm_conn_id="test_conn")
+ with patch.object(hook, "get_conn", autospec=True,
return_value=FunctionModel(call_then_echo)):
+ agent = hook.create_agent(instructions="hi", tools=[read_setting])
+
+ assert agent.run_sync("hi").output == "api key: ***"
+
def test_create_agent_without_instructions_or_spec_file_raises(self):
hook = PydanticAIHook(llm_conn_id="test_conn",
model_id="openai:gpt-5.6-sol")
with pytest.raises(ValueError, match="instructions is required"):
@@ -1321,6 +1359,7 @@ class TestPydanticAIHookCreateAgent:
"/path/to/agent.yaml",
model=mock_model,
output_type=str,
+ capabilities=[MaskingCapability()],
)
mock_agent_cls.assert_not_called()
@@ -1339,6 +1378,7 @@ class TestPydanticAIHookCreateAgent:
mock_agent_cls.from_file.assert_called_once_with(
"/path/to/agent.yaml",
output_type=str,
+ capabilities=[MaskingCapability()],
)
mock_agent_cls.assert_not_called()
@@ -1377,6 +1417,7 @@ class TestPydanticAIHookCreateAgent:
spec_path,
model=mock_model,
output_type=str,
+ capabilities=[MaskingCapability()],
)
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
@@ -1401,6 +1442,7 @@ class TestPydanticAIHookCreateAgent:
model=mock_model,
output_type=str,
instructions="Override instructions.",
+ capabilities=[MaskingCapability()],
)
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
@@ -1419,6 +1461,7 @@ class TestPydanticAIHookCreateAgent:
"/path/to/agent.yaml",
model=mock_model,
output_type=dict,
+ capabilities=[MaskingCapability()],
)
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
@@ -1443,6 +1486,7 @@ class TestPydanticAIHookCreateAgent:
output_type=str,
retries=5,
end_strategy="early",
+ capabilities=[MaskingCapability()],
)
@@ -1493,6 +1537,7 @@ class TestPydanticAIHookCreateAgentInstrumentation:
mock_model,
output_type=str,
instructions="hi",
+ capabilities=[MaskingCapability()],
)
assert agent.instrument is False
mock_settings.assert_not_called()
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index e45611283a7..27f5cda1ba2 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -34,18 +34,20 @@ from pydantic import BaseModel
from pydantic_ai import Agent, DeferredToolRequests, RunContext, Tool,
ToolOutput
from pydantic_ai.capabilities import (
AbstractCapability,
+ CapabilityOrdering,
CombinedCapability,
PrefixTools,
Thinking,
Toolset,
WebSearch,
)
-from pydantic_ai.exceptions import UsageLimitExceeded
+from pydantic_ai.exceptions import ModelRetry, ToolFailed, UsageLimitExceeded
from pydantic_ai.messages import (
ModelMessage,
ModelMessagesTypeAdapter,
ModelRequest,
ModelResponse,
+ RetryPromptPart,
TextPart,
ToolCallPart,
ToolReturnPart,
@@ -64,6 +66,7 @@ from airflow.providers.common.ai.durable import
AirflowDurability
from airflow.providers.common.ai.durable.base import RUN_ID_KEY, build_step_key
from airflow.providers.common.ai.durable.journal import DurableJournal,
DurableRun, journal_scope
from airflow.providers.common.ai.durable.storage import DurableStorage
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.operators.agent import AgentOperator,
HITLReviewLink
from airflow.providers.common.ai.sandbox.base import (
HOLDER_TAG,
@@ -3511,11 +3514,61 @@ class TestAgentOperatorDurableUsageBudgetEndToEnd:
def _echo_tool_result(messages, info: AgentInfo) -> ModelResponse:
- """Call ``read_setting`` once, then answer with whatever it returned."""
- returns = [p for m in messages for p in m.parts if isinstance(p,
ToolReturnPart)]
- if returns:
- return
ModelResponse(parts=[TextPart(content=str(returns[-1].content))])
- return ModelResponse(parts=[ToolCallPart(tool_name="read_setting",
args={}, tool_call_id="c1")])
+ """
+ Call the first tool the agent offers once, whatever its name, then answer
with what came back.
+
+ That is the tool's result, or the retry prompt pydantic-ai sends when the
tool raises.
+ """
+ replies = [p for m in messages for p in m.parts if isinstance(p,
(ToolReturnPart, RetryPromptPart))]
+ if replies:
+ return
ModelResponse(parts=[TextPart(content=str(replies[-1].content))])
+ return ModelResponse(
+ parts=[ToolCallPart(tool_name=info.function_tools[0].name, args={},
tool_call_id="c1")]
+ )
+
+
[email protected]
+class _ToolsFromGetToolset(AbstractCapability):
+ """Provides tools the way ``MCP`` and other built-in capabilities do,
through ``get_toolset``."""
+
+ toolset: FunctionToolset =
dataclasses.field(default_factory=FunctionToolset)
+
+ def get_toolset(self):
+ return self.toolset
+
+
[email protected]
+class _RecordToolResults(AbstractCapability):
+ """A Dag author's capability that looks at every tool result, as a logging
or guard capability would."""
+
+ seen: list = dataclasses.field(default_factory=list)
+
+ def get_ordering(self):
+ # Asks for the same spot the masking capability takes, next to the
tool.
+ return CapabilityOrdering(position="innermost")
+
+ async def wrap_tool_execute(self, ctx, *, call, tool_def, args, handler):
+ result = await handler(args)
+ self.seen.append(result)
+ return result
+
+
+class _OwnLeafToolset(AbstractToolset):
+ """A Dag author's own toolset, which pydantic-ai's durable backend does
not run as a durable unit."""
+
+ def __init__(self, inner: FunctionToolset):
+ self.inner = inner
+
+ @property
+ def id(self):
+ return self.inner.id
+
+ async def get_tools(self, ctx):
+ tools = await self.inner.get_tools(ctx)
+ return {name: dataclasses.replace(tool, toolset=self) for name, tool
in tools.items()}
+
+ async def call_tool(self, name, tool_args, ctx, tool):
+ return await self.inner.call_tool(name, tool_args, ctx,
dataclasses.replace(tool, toolset=self.inner))
@pytest.mark.enable_redact
@@ -3523,12 +3576,17 @@ class TestAgentOperatorMasksToolOutput:
"""What any tool hands the model is masked, however the toolset reaches
the agent."""
@staticmethod
- def _run(op: AgentOperator, storage=None) -> str:
- hook = MagicMock(spec=["create_agent"])
- hook.create_agent.side_effect = lambda **kw:
Agent(FunctionModel(_echo_tool_result), **kw)
+ def _agent(op: AgentOperator) -> Agent:
+ # The real create_agent, which adds the masking capability, with a
scripted model.
+ hook = PydanticAIHook(llm_conn_id="c")
op.llm_hook = hook
+ with patch.object(hook, "get_conn", autospec=True,
return_value=FunctionModel(_echo_tool_result)):
+ return op._build_agent()
+
+ def _run(self, op: AgentOperator, storage=None) -> str:
+ agent = self._agent(op)
with journal_scope(DurableJournal(storage)) if storage is not None
else nullcontext():
- return op._build_agent().run_sync("hi").output
+ return agent.run_sync("hi").output
@staticmethod
def _dag_authors_toolset(secret: str) -> FunctionToolset:
@@ -3565,20 +3623,120 @@ class TestAgentOperatorMasksToolOutput:
assert self._run(op) == "api key: ***"
- def
test_a_toolset_capability_is_masked_before_the_durable_journal_stores_it(self,
registered_secret):
+ @pytest.mark.parametrize(
+ "passed_as",
+ [
+ pytest.param(lambda ts: {"toolsets": [ts]}, id="toolsets"),
+ pytest.param(lambda ts: {"capabilities": [Toolset(ts)]},
id="toolset-capability"),
+ ],
+ )
+ @pytest.mark.parametrize(
+ "toolset",
+ [
+ pytest.param(lambda ts: ts, id="durable-unit"),
+ pytest.param(_OwnLeafToolset, id="own-leaf"),
+ ],
+ )
+ def test_the_durable_journal_never_stores_a_registered_secret(
+ self, registered_secret, passed_as, toolset
+ ):
+ """
+ The journal masks what it records itself, whichever path journals the
call.
+
+ It sits inside every masking layer: pydantic-ai runs a function
toolset as a durable
+ unit, and AirflowDurability journals any other leaf toolset, both
below the masking
+ wrapper around the toolset and the masking capability.
+ """
storage = MemoryStorage()
op = AgentOperator(
task_id="t",
prompt="hi",
llm_conn_id="c",
durable=True,
- agent_params={"capabilities":
[Toolset(self._dag_authors_toolset(registered_secret))]},
+ **passed_as(toolset(self._dag_authors_toolset(registered_secret))),
)
assert self._run(op, storage) == "api key: ***"
recorded = [step["payload"] for step in storage.steps() if
step["kind"] == "tool"]
assert recorded == [{"kind": "tool_return", "result": "api key: ***"}]
+ @pytest.mark.parametrize(
+ "capability",
+ [
+ pytest.param(lambda ts: PrefixTools(wrapped=Toolset(ts),
prefix="settings"), id="prefixed"),
+ pytest.param(lambda ts: CombinedCapability([Toolset(ts)]),
id="combined"),
+ pytest.param(lambda ts: Toolset(lambda ctx: ts),
id="toolset-built-per-run"),
+ pytest.param(lambda ts: _ToolsFromGetToolset(ts),
id="capability-providing-tools"),
+ ],
+ )
+ def test_tools_from_any_capability(self, registered_secret, capability):
+ op = AgentOperator(
+ task_id="t",
+ prompt="hi",
+ llm_conn_id="c",
+
capabilities=[capability(self._dag_authors_toolset(registered_secret))],
+ )
+
+ assert self._run(op) == "api key: ***"
+
+ def test_other_capabilities_only_see_the_masked_result(self,
registered_secret):
+ seen = []
+ op = AgentOperator(
+ task_id="t",
+ prompt="hi",
+ llm_conn_id="c",
+ capabilities=[
+ _RecordToolResults(seen),
+
_ToolsFromGetToolset(self._dag_authors_toolset(registered_secret)),
+ ],
+ )
+
+ self._run(op)
+
+ assert seen == ["api key: ***"]
+
+ @pytest.mark.parametrize(
+ ("error", "reply_is"),
+ [
+ pytest.param(ModelRetry, lambda part: isinstance(part,
RetryPromptPart), id="model-retry"),
+ pytest.param(
+ ToolFailed,
+ lambda part: isinstance(part, ToolReturnPart) and part.outcome
== "failed",
+ id="tool-failed",
+ ),
+ ],
+ )
+ def test_an_error_from_a_capabilitys_tool_is_masked(self,
registered_secret, error, reply_is, caplog):
+ """
+ pydantic-ai hands a capability both as its own wrapper errors; the
model sees the part inside.
+
+ A ModelRetry comes back to the model as a retry prompt and a
ToolFailed as a failed tool
+ result, which does not spend the retry budget; masking must keep them
apart. Neither is
+ the tool failing, so neither is logged as one.
+ """
+
+ def read_setting() -> str:
+ raise error(f"bad key {registered_secret}")
+
+ op = AgentOperator(
+ task_id="t",
+ prompt="hi",
+ llm_conn_id="c",
+
capabilities=[_ToolsFromGetToolset(FunctionToolset(tools=[read_setting]))],
+ )
+
+ result = self._agent(op).run_sync("hi")
+
+ (reply,) = [
+ p
+ for m in result.all_messages()
+ for p in m.parts
+ if isinstance(p, (ToolReturnPart, RetryPromptPart))
+ ]
+ assert reply_is(reply)
+ assert reply.content == "bad key ***"
+ assert "Tool read_setting failed" not in caplog.messages
+
def test_a_toolset_that_masks_its_own_output_is_not_wrapped_again(self):
sql = SQLToolset("pg_default")
op = AgentOperator(
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_masking.py
b/providers/common/ai/tests/unit/common/ai/utils/test_masking.py
index e6dffecaeec..5c821d5b10f 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_masking.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_masking.py
@@ -17,7 +17,7 @@
from __future__ import annotations
import json
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass, field
import pytest
from pydantic import BaseModel
@@ -55,6 +55,16 @@ class TestMaskSecrets:
assert mask_secrets([Setting(value=registered_secret)]) == [{"value":
"***"}]
+ def test_keeps_a_pydantic_model_that_holds_no_secret(self,
registered_secret):
+ """Code that reads the model's attributes later, such as from the
message history, still can."""
+
+ class Setting(BaseModel):
+ value: str
+
+ setting = Setting(value="plain")
+
+ assert mask_secrets(setting) is setting
+
def test_leaves_other_objects_such_as_images_alone(self,
registered_secret):
image = BinaryContent(data=b"\x89PNG", media_type="image/png")
@@ -69,6 +79,37 @@ class TestMaskSecrets:
assert mask_secrets(Credentials("svc", registered_secret)) ==
Credentials("svc", "***")
assert mask_secrets(TextContent(content=f"key={registered_secret}"))
== TextContent(content="key=***")
+ def test_masks_a_dataclass_it_could_not_construct_again(self,
registered_secret):
+ """An InitVar has no value to pass back in, and an init=False field is
not an argument."""
+
+ @dataclass
+ class Report:
+ rows: list
+ source: InitVar[str]
+ label: str = field(init=False, default="")
+ # Never set: a field the copy must not try to read.
+ cache: str = field(init=False)
+
+ def __post_init__(self, source: str) -> None:
+ self.label = f"from {source}"
+
+ report = Report([registered_secret], source=registered_secret)
+
+ masked = mask_secrets(report)
+
+ assert type(masked) is Report
+ assert masked.rows == ["***"]
+ assert masked.label == "from ***"
+ assert not hasattr(masked, "cache")
+ assert report.rows == [registered_secret]
+
+ def test_masks_a_slotted_dataclass(self, registered_secret):
+ @dataclass(slots=True)
+ class Row:
+ value: str
+
+ assert mask_secrets(Row(registered_secret)) == Row("***")
+
def test_masks_bytes_as_text(self, registered_secret):
assert mask_secrets(f"key={registered_secret}".encode()) == b"key=***"
@@ -115,3 +156,11 @@ class TestDumpsMasked:
dumped = dumps_masked({"row": Row(f"pg://svc:{secret}@db"), "blob":
f"key={secret}".encode()})
assert json.loads(dumped) == {"row": {"dsn": "pg://svc:***@db"},
"blob": "key=***"}
+
+ def test_dumps_a_pydantic_model_that_holds_no_secret_as_its_data(self,
registered_secret):
+ class Setting(BaseModel):
+ value: str
+
+ assert json.loads(dumps_masked({"setting": Setting(value="plain")}))
== {
+ "setting": {"value": "plain"}
+ }
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
b/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
index df0f364ede8..fa5b339293d 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
@@ -26,7 +26,7 @@ from unittest.mock import MagicMock
import pytest
from pydantic_ai import Agent
from pydantic_ai._run_context import RunContext
-from pydantic_ai.exceptions import ApprovalRequired, ModelRetry, ToolFailed
+from pydantic_ai.exceptions import ApprovalRequired, ModelRetry, ToolFailed,
UnexpectedModelBehavior
from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart,
ToolReturn, ToolReturnPart
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.models.test import TestModel
@@ -34,7 +34,12 @@ from pydantic_ai.toolsets.abstract import ToolsetTool
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.usage import RunUsage
-from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
MaskingToolset, ensure_masked
+from airflow.providers.common.ai.utils.toolset_base import (
+ AirflowToolset,
+ MaskingCapability,
+ MaskingToolset,
+ ensure_masked,
+)
class _ScriptedToolset(AirflowToolset):
@@ -146,11 +151,15 @@ class TestMasking:
def test_masks_a_tool_return_and_keeps_it_one(self, registered_secret):
result = _call(
_ScriptedToolset(
- ToolReturn(return_value=f"key={registered_secret}",
content=f"for {registered_secret}")
+ ToolReturn(
+ return_value=f"key={registered_secret}",
+ content=f"for {registered_secret}",
+ metadata={"raw": registered_secret},
+ )
)
)
- assert result == ToolReturn(return_value="key=***", content="for ***")
+ assert result == ToolReturn(return_value="key=***", content="for ***",
metadata={"raw": "***"})
def test_an_approval_request_keeps_its_type_and_is_not_logged_as_a_failure(
self, registered_secret, caplog
@@ -228,6 +237,31 @@ class TestMaskingToolset:
assert asyncio.run(call()) == "api key: ***"
[email protected]_redact
+class TestMaskingCapability:
+ def test_running_out_of_tool_retries_is_not_logged_as_a_tool_failure(self,
registered_secret, caplog):
+ """pydantic-ai raises UnexpectedModelBehavior below the capability
once the retries are used up."""
+
+ def read_setting() -> str:
+ raise ModelRetry(f"bad key {registered_secret}")
+
+ def keep_calling(messages, info) -> ModelResponse:
+ return ModelResponse(
+ parts=[ToolCallPart(tool_name="read_setting", args={},
tool_call_id=f"c{len(messages)}")]
+ )
+
+ agent = Agent(
+ FunctionModel(keep_calling),
+ toolsets=[FunctionToolset([read_setting], max_retries=1)],
+ capabilities=[MaskingCapability()],
+ )
+
+ with pytest.raises(UnexpectedModelBehavior, match="exceeded max
retries"):
+ agent.run_sync("hi")
+
+ assert "Tool read_setting failed" not in caplog.messages
+
+
class _ApiToolset(FunctionToolset):
"""A Dag author's toolset whose one tool returns a connection string."""
