Skip to content

fix(mcp): stop masking the agent's own arguments in tool error text #1366

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Aaron ("AJ") Steers (aaronsteers) merged 1 commit into from
Oct 13, 2026
Merged

fix(mcp): stop masking the agent's own arguments in tool error text #1366

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 8 additions & 140 deletions airbyte/mcp/_error_handling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,15 @@
from __future__ import annotations

import asyncio
import json
import logging
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING
from urllib.parse import urlsplit

import requests
from fastmcp.exceptions import DisabledError, NotFoundError, ToolError, ValidationError
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp_extensions.otel._arg_digests import classify_tool # noqa: PLC2701
from fastmcp_extensions.otel._extras import _chain, declared_parameters # noqa: PLC2701
from fastmcp_extensions.otel.models import TraceArg
from fastmcp_extensions.otel._extras import _chain # noqa: PLC2701

from airbyte._util.api_util import (
SDKError,
Expand Down Expand Up @@ -49,12 +45,10 @@


if TYPE_CHECKING:
from fastmcp.tools import Tool, ToolResult
from fastmcp.tools import ToolResult
from mcp import types as mt


logger = logging.getLogger(__name__)

_UPSTREAM_METHODS = frozenset({"GET", "POST", "PUT", "PATCH", "DELETE"})
# Literal segments come from path templates in airbyte_api.
_CLOUD_API_ROUTE_SEGMENTS = frozenset(
Expand Down Expand Up @@ -283,144 +277,22 @@ def format_validation_error(error: ValidationError) -> str:
return " ".join(parts) if parts else "Invalid arguments."


_JSON_LITERALS = frozenset({"true", "false", "null"})
_INPUT_VALUE = re.compile(r"input_value=(?:'[^']*'|\"[^\"]*\"|[^\s,\])]+)")
_BEARER_TOKEN = re.compile(r"Bearer\s+[^\s'\")]+")
_URL_CREDENTIALS = re.compile(r"([a-zA-Z][a-zA-Z0-9+.-]*://)[^\s/?#'\")]+@")
_URL_QUERY = re.compile(r"([a-zA-Z][a-zA-Z0-9+.-]*://[^\s?#'\")]*)\?[^\s'\")]*")


def _secret_leaf(leaf: str) -> bool:
"""Whether a secret-class leaf is masked: everything but JSON literals."""
return bool(leaf) and not (
leaf in _JSON_LITERALS or (leaf.isdigit() and len(leaf) < 4) # noqa: PLR2004
)


def _collect_leaves(
value: Any, # noqa: ANN401
path: str,
*,
secret: bool,
leaves: list[tuple[str, str, bool]],
) -> None:
"""Collect (placeholder path, leaf text, secret) tuples for values to hide."""
if isinstance(value, bool):
return # `true`/`false` are JSON literals, never masked.
if isinstance(value, (int, float)):
leaf = str(value)
if (secret and _secret_leaf(leaf)) or (not secret and len(leaf) >= 8): # noqa: PLR2004
leaves.append((path, leaf, secret))
return
if isinstance(value, str):
if (secret and _secret_leaf(value)) or (not secret and len(value) >= 8): # noqa: PLR2004
leaves.append((path, value, secret))
# A JSON-string config may hold nested secrets; index its leaves too.
if value[:1] in {"{", "["}:
try:
nested = json.loads(value)
except ValueError:
return
if isinstance(nested, (dict, list)):
_collect_leaves(nested, path, secret=secret, leaves=leaves)
return
if isinstance(value, Mapping):
for key, item in value.items():
if isinstance(key, str):
_collect_leaves(item, f"{path}.{key}", secret=secret, leaves=leaves)
return
if isinstance(value, (list, tuple)):
for item in value:
_collect_leaves(item, path, secret=secret, leaves=leaves)


def _allowed_choices(
tool: Tool | None,
name: str,
arg_class: Any, # noqa: ANN401
) -> tuple[Any, ...]:
"""Return the declared `enum`/`const` choices of a parameter, if any."""
if tool is not None:
schema = declared_parameters(tool).get(name)
if isinstance(schema, Mapping):
if isinstance(schema.get("enum"), list):
return tuple(schema["enum"])
if "const" in schema:
return (schema["const"],)
allowed = getattr(arg_class, "allowed", None)
return tuple(allowed) if allowed is not None else ()


def redact_agent_text(
text: str,
arguments: Mapping[str, Any] | None,
tool: Tool | None,
) -> str:
"""Mask every echoed argument value in agent-facing error text.

Each parameter is classified by the same rules arg tracing uses, so error
text never shows a value tracing would not record in clear: secret-class
parameters lose every leaf, other parameters lose leaves of eight or more
characters, and `VALUE` (enum-like) parameters stay clear only when they
hold one of the declared choices. Secret leaves shorter than four
characters and non-secret leaves are masked as whole tokens.
"""
if arguments:
func = getattr(tool, "fn", None) if tool is not None else None
try:
classes = classify_tool(func, sorted(arguments), tool=getattr(tool, "name", ""))
except Exception:
classes = {}
leaves: list[tuple[str, str, bool]] = []
for name, value in arguments.items():
arg_class = classes.get(name)
mode = arg_class.mode if arg_class is not None else TraceArg.PRESENCE
if mode is TraceArg.VALUE:
if any(value == choice for choice in _allowed_choices(tool, name, arg_class)):
continue
secret = False
else:
secret = mode in {TraceArg.OMIT, TraceArg.PRESENCE}
_collect_leaves(value, name, secret=secret, leaves=leaves)
seen: set[str] = set()
for path, leaf, secret in sorted(leaves, key=lambda item: -len(item[1])):
if leaf in seen:
continue
seen.add(leaf)
placeholder = f"<value of {path.split('.', maxsplit=1)[0]}>"
if secret and len(leaf) >= 4: # noqa: PLR2004
text = text.replace(leaf, placeholder)
else:
text = re.sub(
rf"(?<![A-Za-z0-9_-]){re.escape(leaf)}(?![A-Za-z0-9_-])",
placeholder.replace("\\", "\\\\"),
text,
)

def redact_agent_text(text: str) -> str:
"""Redact input_value echoes, bearer tokens, URL credentials, and URL query strings."""
text = _INPUT_VALUE.sub("input_value=<redacted>", text)
text = _BEARER_TOKEN.sub("Bearer <redacted>", text)
text = _URL_CREDENTIALS.sub(r"\1<redacted>@", text)
return _URL_QUERY.sub(r"\1", text)
Comment thread
aaronsteers marked this conversation as resolved.


async def _call_arguments(
context: MiddlewareContext[mt.CallToolRequestParams],
) -> tuple[Mapping[str, Any] | None, Tool | None]:
"""Return the call's arguments and tool, either `None` when unavailable."""
message = context.message
arguments = getattr(message, "arguments", None)
tool = None
fastmcp_context = context.fastmcp_context
if fastmcp_context is not None:
try:
tool = await fastmcp_context.fastmcp.get_tool(message.name)
except Exception:
logger.debug("Could not resolve the called tool for error redaction")
return arguments if isinstance(arguments, Mapping) else None, tool


class AgentErrorTextMiddleware(Middleware):
"""Rewrite tool error text for the agent and mask echoed argument values.
"""Rewrite tool error text for the agent and redact credentials.

Runs outside `UserFacingErrorMiddleware` and telemetry: telemetry has
already classified the original exception when this middleware replaces it.
Expand All @@ -434,10 +306,6 @@ async def on_call_tool(
try:
return await call_next(context)
except ValidationError as error:
arguments, tool = await _call_arguments(context)
raise ValidationError(
redact_agent_text(format_validation_error(error), arguments, tool)
) from None
raise ValidationError(redact_agent_text(format_validation_error(error))) from None
except ToolError as error:
arguments, tool = await _call_arguments(context)
raise ToolError(redact_agent_text(agent_error_text(error), arguments, tool)) from None
raise ToolError(redact_agent_text(agent_error_text(error))) from None
4 changes: 2 additions & 2 deletions tests/unit_tests/test_mcp_datadog.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1462,10 +1462,10 @@ async def run():
assert private_text not in json.dumps(span.get_tags())
local_demo_text = responses[8][2]["result"]["content"][0]["text"]
assert local_demo_text == (
"Cannot modify '<value of guid>': local demo mode only allows "
"Cannot modify 'abc-secret-id-123': local demo mode only allows "
"changes to resources created in this session."
)
assert "abc-secret-id-123" not in local_demo_text
assert "abc-secret-id-123" in local_demo_text
assert str(ValueError(private_payload)) not in json.dumps(
value_error_span.get_tags()
)
Expand Down
110 changes: 20 additions & 90 deletions tests/unit_tests/test_mcp_tool_errors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

import asyncio
import json
from types import SimpleNamespace
from typing import Literal

import pytest
Expand Down Expand Up @@ -348,7 +347,7 @@ def broken() -> str:
assert "api.airbyte.com" not in text


def test_unexpected_error_echoed_argument_is_masked() -> None:
def test_unexpected_error_echoed_argument_passes_through() -> None:
server = _server_with_agent_error_text()

@server.tool
Expand All @@ -357,11 +356,11 @@ def fail(connection_id: str) -> str:

text = _call_error(server, "fail", {"connection_id": "SENTINEL-arg-value-1234"})

assert "boom with <value of connection_id>" in text
assert "SENTINEL-arg-value-1234" not in text
assert "boom with SENTINEL-arg-value-1234" in text
assert "<value of" not in text


def test_user_facing_error_echoed_id_is_masked() -> None:
def test_user_facing_error_echoed_id_passes_through() -> None:
server = _server_with_agent_error_text()

@server.tool
Expand All @@ -377,11 +376,11 @@ def unsafe(connection_id: str) -> str:
)

assert "Cannot perform destructive operation" in text
assert "00000000-1111-2222-3333-444444444444" not in text
assert "<value of connection_id>" in text
assert "00000000-1111-2222-3333-444444444444" in text
assert "<value of" not in text


def test_secret_valued_argument_never_shows_a_prefix() -> None:
def test_secret_valued_argument_passes_through_error_text() -> None:
server = _server_with_agent_error_text()

@server.tool
Expand All @@ -400,13 +399,13 @@ def read_records(config: dict | str | None, stream_name: str) -> str:
},
)

assert secret not in text
assert "sk_live" not in text
assert "<value of config>" in text
assert "stream_with_secret_value" not in text
assert secret in text
assert "sk_live" in text
assert "stream_with_secret_value" in text
assert "<value of" not in text


def test_redaction_placeholder_does_not_expose_nested_mapping_key() -> None:
def test_nested_argument_echo_passes_through_without_mapping_key() -> None:
server = _server_with_agent_error_text()
hostile_key = "password-or-injected-text"
secret = "sk_live_FAKE1234567890"
Expand All @@ -421,85 +420,16 @@ def read_records(config: dict[str, str]) -> str:
{"config": {hostile_key: secret}},
)

assert secret not in text
assert secret in text
assert hostile_key not in text
assert "<value of config>" in text
assert "<value of" not in text


def _tool_stub(fn):
return SimpleNamespace(fn=fn, name="test-tool")


def _classify_args(stream_name: str, connection_id: str, password: str) -> None:
"""Hints so `classify_tool` gives HASH/PRESENCE modes like real tools."""


@pytest.mark.parametrize(
("arguments", "text", "expected"),
[
pytest.param(
{"stream_name": "user"},
"Ask the user; don't retry.",
"Ask the user; don't retry.",
id="non-secret-word-user",
),
pytest.param(
{"stream_name": "retry"},
"Ask the user; don't retry.",
"Ask the user; don't retry.",
id="non-secret-word-retry",
),
pytest.param(
{"stream_name": "users"},
"Ask the user; don't retry.",
"Ask the user; don't retry.",
id="non-secret-word-users",
),
pytest.param(
{"stream_name": "customer"},
"Check the customers table first.",
"Check the customers table first.",
id="non-secret-token-inside-longer-token",
),
pytest.param(
{"connection_id": "f47ac10b-58cc-4372-a567-0e02b2c3d479"},
"Cannot perform destructive operation on "
"'f47ac10b-58cc-4372-a567-0e02b2c3d479': Object was not created.",
"Cannot perform destructive operation on "
"'<value of connection_id>': Object was not created.",
id="non-secret-uuid-echo-masked",
),
pytest.param(
{"password": "abc123secret"},
"boom in xxabc123secretyy field",
"boom in xx<value of password>yy field",
id="secret-substring-masked",
),
],
)
def test_redact_agent_text_masks_whole_tokens_only(
arguments: dict, text: str, expected: str
) -> None:
assert redact_agent_text(text, arguments, _tool_stub(_classify_args)) == expected


def test_short_secret_is_masked_as_a_whole_token() -> None:
assert (
redact_agent_text(
"Ask the user for a valid value",
{"password": "a"},
_tool_stub(_classify_args),
)
== "Ask the user for <value of password> valid value"
def test_redact_agent_text_keeps_connection_error_id_and_url() -> None:
text = (
"Connection 'c0bfcb08-bd3c-411e-8076-3d80fa18c137' not found. See "
"https://cloud.airbyte.com/workspaces/266ebdfe-0d7b-4540-9817-de7e4505ba61/"
"source/c0bfcb08-bd3c-411e-8076-3d80fa18c137"
)


def test_long_secret_is_masked_inside_a_token() -> None:
assert (
redact_agent_text(
"Request failed with pw=hunter2xyz;",
{"password": "hunter2xyz"},
_tool_stub(_classify_args),
)
== "Request failed with pw=<value of password>;"
)
assert redact_agent_text(text) == text
Loading