diff --git a/airbyte/mcp/_error_handling.py b/airbyte/mcp/_error_handling.py index 4568f05ce..22bcc28d9 100644 --- a/airbyte/mcp/_error_handling.py +++ b/airbyte/mcp/_error_handling.py @@ -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, @@ -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( @@ -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"" - if secret and len(leaf) >= 4: # noqa: PLR2004 - text = text.replace(leaf, placeholder) - else: - text = re.sub( - rf"(? str: + """Redact input_value echoes, bearer tokens, URL credentials, and URL query strings.""" text = _INPUT_VALUE.sub("input_value=", text) text = _BEARER_TOKEN.sub("Bearer ", text) text = _URL_CREDENTIALS.sub(r"\1@", text) return _URL_QUERY.sub(r"\1", text) -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. @@ -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 diff --git a/tests/unit_tests/test_mcp_datadog.py b/tests/unit_tests/test_mcp_datadog.py index d26bbe3a4..6106e23f1 100644 --- a/tests/unit_tests/test_mcp_datadog.py +++ b/tests/unit_tests/test_mcp_datadog.py @@ -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 '': 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() ) diff --git a/tests/unit_tests/test_mcp_tool_errors.py b/tests/unit_tests/test_mcp_tool_errors.py index 7f5be4f9c..fbe346a88 100644 --- a/tests/unit_tests/test_mcp_tool_errors.py +++ b/tests/unit_tests/test_mcp_tool_errors.py @@ -5,7 +5,6 @@ import asyncio import json -from types import SimpleNamespace from typing import Literal import pytest @@ -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 @@ -357,11 +356,11 @@ def fail(connection_id: str) -> str: text = _call_error(server, "fail", {"connection_id": "SENTINEL-arg-value-1234"}) - assert "boom with " in text - assert "SENTINEL-arg-value-1234" not in text + assert "boom with SENTINEL-arg-value-1234" in text + assert " None: +def test_user_facing_error_echoed_id_passes_through() -> None: server = _server_with_agent_error_text() @server.tool @@ -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 "" in text + assert "00000000-1111-2222-3333-444444444444" in text + assert " None: +def test_secret_valued_argument_passes_through_error_text() -> None: server = _server_with_agent_error_text() @server.tool @@ -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 "" 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 " 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" @@ -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 "" in text + assert " 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 " - "'': Object was not created.", - id="non-secret-uuid-echo-masked", - ), - pytest.param( - {"password": "abc123secret"}, - "boom in xxabc123secretyy field", - "boom in xxyy 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 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=;" - ) + assert redact_agent_text(text) == text