diff --git a/airbyte/_util/api_util.py b/airbyte/_util/api_util.py index 2a8ac858b..8002227fe 100644 --- a/airbyte/_util/api_util.py +++ b/airbyte/_util/api_util.py @@ -71,6 +71,9 @@ "The source or destination is still a draft: a person has not finished its setup in " "Airbyte Cloud (HTTP 409 actor-not-ready)." ) +FORBIDDEN_RESOURCE_MESSAGE = ( + "The requested resource was not found, or these credentials can't access it (HTTP 403)." +) ACTOR_NOT_READY_MCP_GUIDANCE = ( "This is expected right after creating a connector with `defer_credentials=True`. " "Don't retry yet. Ask the user to open the connector's settings page, complete " @@ -197,6 +200,16 @@ def _is_actor_not_ready_error(error: SDKError) -> bool: ) +def sdk_error_message(error: SDKError) -> str: + """Return the fixed message an SDKError is reported with; never Cloud body text.""" + if _is_actor_not_ready_error(error): + return ACTOR_NOT_READY_ERROR_MESSAGE + if error.status_code == HTTPStatus.FORBIDDEN: + return FORBIDDEN_RESOURCE_MESSAGE + message, _ = describe_cloud_error(parse_cloud_error(error.status_code, error.body)) + return message + + def _wrap_sdk_error( error: SDKError, base_context: dict[str, Any] | None = None ) -> AirbyteCloudError: @@ -223,19 +236,14 @@ def _wrap_sdk_error( status_code = sdk_context.get("status_code") is_forbidden = status_code == HTTPStatus.FORBIDDEN - message, guidance = describe_cloud_error(problem) + _, guidance = describe_cloud_error(problem) error_type = ( AirbyteMissingResourceError if is_forbidden or status_code == HTTPStatus.NOT_FOUND else AirbyteCloudError ) return error_type( - message=( - "The requested resource was not found, or these credentials can't access it " - "(HTTP 403)." - if is_forbidden - else message - ), + message=sdk_error_message(error), guidance=FORBIDDEN_RESOURCE_GUIDANCE if is_forbidden else guidance, context=merged_context, ) diff --git a/airbyte/mcp/_datadog.py b/airbyte/mcp/_datadog.py index f8c846732..7dcee1f4c 100644 --- a/airbyte/mcp/_datadog.py +++ b/airbyte/mcp/_datadog.py @@ -36,6 +36,8 @@ from airbyte.mcp._error_handling import ( MCP_TOOL_USER_FACING_ERRORS, classify_mcp_tool_error, + cloud_error_route, + cloud_error_trace_message, mcp_tool_error_reason, ) from airbyte.mcp._otel import _arg_key, _env, _flag @@ -83,6 +85,8 @@ _ARG_TRACER = ArgTracer(_ARG_PREFIX, key=_arg_key, skip=(INTENT_ARG,)) _TOOL_ERROR_MESSAGE = "tool resulted in an error" _STATUS_CODE = "airbyte.mcp.upstream.status_code" +_UPSTREAM_METHOD = "airbyte.mcp.upstream.method" +_UPSTREAM_ROUTE = "airbyte.mcp.upstream.route" def _build_tool_maps() -> None: @@ -202,6 +206,13 @@ def _exception_attributes(error: BaseException) -> dict[str, str]: } if (stack := error_stack(cause)) is not None: attrs["error.stack"] = stack + if route := cloud_error_route(cause): + attrs[_UPSTREAM_METHOD], attrs[_UPSTREAM_ROUTE] = route + try: + if message := cloud_error_trace_message(cause): + attrs["error.message"] = message + except Exception: + logger.debug("Datadog Cloud error message unavailable") return attrs @@ -246,6 +257,15 @@ def _tool_error_message(span: Span) -> str: return f"{_TOOL_ERROR_MESSAGE} ({details})" if details else _TOOL_ERROR_MESSAGE +def _raised_tool_error_message(error: BaseException, span: Span) -> str: + """Return a fixed Cloud message when available, falling back to the span summary.""" + try: + message = cloud_error_trace_message(error) + except Exception: + message = None + return message if message is not None else _tool_error_message(span) + + class _StripMetaTraceContextMiddleware: async def __call__( self, @@ -520,14 +540,14 @@ async def _trace_call( logger.debug("Datadog tool attributes unavailable") -def _annotate_attributes(span: Span, source: Mapping[str, str | bool]) -> None: +def _annotate_attributes( + span: Span, source: Mapping[str, str | bool], *, llmobs: bool = True +) -> None: """Keep native APM attributes and LLM metadata consistent. Argument records pass the OTel backend's export validation: ints become metrics, and the validated records also reach LLM metadata. """ - from ddtrace.llmobs import LLMObs # noqa: PLC0415 - attrs = dict(source) arg_attrs = {key: attrs.pop(key) for key in list(attrs) if is_arg_key(_ARG_PREFIX, key)} tool = attrs.get("gen_ai.tool.name") @@ -553,13 +573,14 @@ def _annotate_attributes(span: Span, source: Mapping[str, str | bool]) -> None: metadata.update({key.removeprefix("airbyte.mcp."): value for key, value in accepted.items()}) if "gen_ai.tool.call.id" in attrs: metadata["tool_id"] = attrs["gen_ai.tool.call.id"] - LLMObs.annotate(span, metadata=metadata) + if llmobs: + from ddtrace.llmobs import LLMObs # noqa: PLC0415 + LLMObs.annotate(span, metadata=metadata) -def _annotate_request(span: Span, ctx: ServerRequestContext[Any]) -> None: - """Annotate protocol identity without copying tool arguments into the SDK.""" - from ddtrace.llmobs import LLMObs # noqa: PLC0415 +def _annotate_request(span: Span, ctx: ServerRequestContext[Any], *, llmobs: bool = True) -> None: + """Annotate protocol identity without copying tool arguments into the SDK.""" request_types = { "initialize": InitializeRequest, "tools/list": ListToolsRequest, @@ -605,26 +626,31 @@ def _annotate_request(span: Span, ctx: ServerRequestContext[Any]) -> None: ) except Exception: logger.debug("Early tool scope unavailable") - _annotate_attributes(span, attrs) + _annotate_attributes(span, attrs, llmobs=llmobs) session = attrs.get("airbyte.mcp.session_id") if session: tags["mcp_session_id"] = session - LLMObs.annotate( - span, - tags=tags, - metadata={"pyairbyte.version": get_version()}, - # Seed scalar payloads so the processor can replace Input even if the - # tool fails before the metadata middleware runs. Never seed raw data. - input_data=REDACTED if tool_call else None, - output_data=REDACTED if tool_call else None, - ) - if not tool_call: - request = ( - request_types[ctx.method] - .model_validate({"method": ctx.method, "params": params}) - .model_dump(exclude={"params": {"meta": "_dd_trace_context"}}) + if llmobs: + from ddtrace.llmobs import LLMObs # noqa: PLC0415 + + LLMObs.annotate( + span, + tags=tags, + metadata={"pyairbyte.version": get_version()}, + # Seed scalar payloads so the processor can replace Input even if the + # tool fails before the metadata middleware runs. Never seed raw data. + input_data=REDACTED if tool_call else None, + output_data=REDACTED if tool_call else None, ) - LLMObs.annotate(span, input_data=request) + if not tool_call: + request = ( + request_types[ctx.method] + .model_validate({"method": ctx.method, "params": params}) + .model_dump(exclude={"params": {"meta": "_dd_trace_context"}}) + ) + LLMObs.annotate(span, input_data=request) + else: + span.set_tags(tags) class _DatadogRequestMiddleware: @@ -667,8 +693,8 @@ async def __call__( # noqa: PLR0912, PLR0915 from ddtrace import tracer # noqa: PLC0415 from ddtrace.llmobs import LLMObs # noqa: PLC0415 - # Native server instrumentation covered initialize/call. Listing is an - # explicit addition; other operations retain deployment HTTP tracing. + # Use APM spans for lifecycle methods and LLM Obs spans for tool calls. + # Other operations retain deployment HTTP tracing. if ctx.method not in {"initialize", "tools/list", "tools/call"}: return await call_next(ctx) span = None @@ -696,17 +722,16 @@ async def __call__( # noqa: PLR0912, PLR0915 except Exception: logger.debug("Datadog MCP distributed context unavailable") name = str(params.get("name", "unknown_tool")) if tool_call else f"mcp.{ctx.method}" - start = LLMObs.tool if tool_call else LLMObs.task - span = start(name=name) - # LLMObs retains its display name independently of these APM fields. - span.name = f"mcp.{ctx.method}" + if tool_call: + span = LLMObs.tool(name=name) + span.name = f"mcp.{ctx.method}" + else: + span = tracer.trace(f"mcp.{ctx.method}", resource="server_request") if tool_call: # Per-tool APM trace metrics group by resource; keep it bounded. span.resource = name if name in _TOOL_MODULES else "unknown_tool" - else: - span.resource = "server_request" span.set_metric("_dd.measured", 1) - _annotate_request(span, ctx) + _annotate_request(span, ctx, llmobs=tool_call) except Exception: logger.debug("Datadog request attributes unavailable") try: @@ -719,7 +744,7 @@ async def __call__( # noqa: PLR0912, PLR0915 "tools/list": ListToolsResult, "tools/call": CallToolResult, }[ctx.method].model_validate(result) - output = REDACTED if tool_call else response.model_dump(mode="json") + output = REDACTED if tool_call and getattr(response, "is_error", False): span.error = 1 error_type = span.get_tag("airbyte.mcp.error_type") or "ToolError" @@ -732,11 +757,15 @@ async def __call__( # noqa: PLR0912, PLR0915 # A returned error result raised nothing to classify. if span.get_tag("airbyte.mcp.error.category") is None: error_attrs.update(_error_classification(None)) - _annotate_attributes(span, error_attrs) - span.set_tag("error.message", _tool_error_message(span)) + _annotate_attributes(span, error_attrs, llmobs=tool_call) + if span.get_tag("error.message") is None: + span.set_tag("error.message", _tool_error_message(span)) else: - _annotate_attributes(span, {"airbyte.mcp.outcome": "success"}) - LLMObs.annotate(span, output_data=output) + _annotate_attributes( + span, {"airbyte.mcp.outcome": "success"}, llmobs=tool_call + ) + if tool_call: + LLMObs.annotate(span, output_data=output) except Exception: logger.debug("Datadog response attributes unavailable") return result @@ -747,9 +776,9 @@ async def __call__( # noqa: PLR0912, PLR0915 # and traceback, which can contain arguments and results. if error := sys.exc_info()[1]: span.error = 1 - _annotate_attributes(span, _exception_attributes(error)) + _annotate_attributes(span, _exception_attributes(error), llmobs=tool_call) if tool_call: - span.set_tag("error.message", _tool_error_message(span)) + span.set_tag("error.message", _raised_tool_error_message(error, span)) span.__exit__(None, None, None) except Exception: logger.debug("Datadog span completion failed") diff --git a/airbyte/mcp/_error_handling.py b/airbyte/mcp/_error_handling.py index ca80f9d83..bcbbac45a 100644 --- a/airbyte/mcp/_error_handling.py +++ b/airbyte/mcp/_error_handling.py @@ -8,14 +8,21 @@ import re from collections.abc import Mapping from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit +import requests from fastmcp.exceptions import 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 airbyte._util.api_util import error_response_body, sdk_error_response +from airbyte._util.api_util import ( + SDKError, + error_response_body, + sdk_error_message, + sdk_error_response, +) from airbyte._util.cloud_errors import ( describe_cloud_error, is_valid_problem_slug, @@ -45,6 +52,29 @@ 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( + { + "connections", + "destinations", + "health", + "jobs", + "organizations", + "oauthCredentials", + "permissions", + "sources", + "initiateOAuth", + "streams", + "tags", + "users", + "workspaces", + "definitions", + "declarative_sources", + } +) +_MAX_UPSTREAM_ROUTE_LENGTH = 200 + MCP_TOOL_USER_FACING_ERRORS: tuple[type[AirbyteLibError], ...] = ( AirbyteLibInputError, @@ -60,6 +90,78 @@ """Expected errors returned to MCP clients as concise message and guidance text.""" +def cloud_error_trace_message(error: BaseException) -> str | None: + """Return the fixed Cloud error message for a failed call's span, if any. + + Only the message half: guidance can carry Cloud's error ID. + """ + cloud_error: AirbyteCloudApiError | None = None + for cause in _chain(error): + if ( + isinstance(cause, AirbyteLibError) + and isinstance(cause.__cause__, requests.HTTPError) + and "problem_type" in (cause.context or {}) + ): + return cause.message + if isinstance(cause, AirbyteCloudApiError): + body = error_response_body(cause) + if body is not None: + message, _ = describe_cloud_error(parse_cloud_error(cause.status_code, body)) + return message + cloud_error = cause + elif isinstance(cause, SDKError): + return sdk_error_message(cause) + + if cloud_error is not None: + message, _ = describe_cloud_error(parse_cloud_error(cloud_error.status_code, None)) + return message + return None + + +def cloud_error_route(error: BaseException) -> tuple[str, str] | None: + """Return a bounded HTTP method and sanitized path for an upstream SDK error.""" + for cause in _chain(error): + if not isinstance(cause, SDKError): + continue + try: + response = getattr(cause, "raw_response", None) + request = getattr(response, "request", None) + if request is None: + continue + + method = request.method.upper() + if method not in _UPSTREAM_METHODS: + return None + + segments = [ + segment for segment in urlsplit(str(request.url)).path.split("/") if segment + ] + version_index = next( + ( + index + for index in range(len(segments) - 1, -1, -1) + if re.fullmatch(r"v[0-9]{1,3}", segments[index]) + ), + None, + ) + route_segments: list[str] = [] + for index, segment in enumerate(segments): + if version_index is None or index < version_index: + route_segments.append("{id}") + elif index == version_index or segment in _CLOUD_API_ROUTE_SEGMENTS: + route_segments.append(segment) + else: + route_segments.append("{id}") + route = "/" + "/".join(route_segments) + if len(route) > _MAX_UPSTREAM_ROUTE_LENGTH: + return None + except Exception: + return None + else: + return method, route + return None + + def _cloud_api_error_text(error: AirbyteCloudApiError) -> str: """Return fixed Cloud error text, retaining direct-access guidance when present.""" message, table_guidance = describe_cloud_error( diff --git a/tests/unit_tests/test_cloud_api_util.py b/tests/unit_tests/test_cloud_api_util.py index 093690de3..947c29836 100644 --- a/tests/unit_tests/test_cloud_api_util.py +++ b/tests/unit_tests/test_cloud_api_util.py @@ -2921,21 +2921,18 @@ def test_wrap_sdk_error_describes_cloud_problem( def test_wrap_sdk_error_403_keeps_fixed_forbidden_text() -> None: - wrapped = api_util._wrap_sdk_error( - _sdk_error( - 403, - _problem_body( - "https://reference.airbyte.com/reference/errors#forbidden", - data={"message": _INTERNAL_MARKER}, - ), - ) + error = _sdk_error( + 403, + _problem_body( + "https://reference.airbyte.com/reference/errors#forbidden", + data={"message": _INTERNAL_MARKER}, + ), ) + wrapped = api_util._wrap_sdk_error(error) assert type(wrapped) is AirbyteMissingResourceError - assert wrapped.get_message() == ( - "The requested resource was not found, or these credentials can't " - "access it (HTTP 403)." - ) + assert wrapped.get_message() == api_util.FORBIDDEN_RESOURCE_MESSAGE + assert api_util.sdk_error_message(error) == api_util.FORBIDDEN_RESOURCE_MESSAGE assert wrapped.guidance == api_util.FORBIDDEN_RESOURCE_GUIDANCE assert wrapped.context["problem_type"] == "forbidden" assert _INTERNAL_MARKER not in str(wrapped) diff --git a/tests/unit_tests/test_mcp_datadog.py b/tests/unit_tests/test_mcp_datadog.py index 8f0e4e1e7..3494a978c 100644 --- a/tests/unit_tests/test_mcp_datadog.py +++ b/tests/unit_tests/test_mcp_datadog.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Airbyte, Inc., all rights reserved. """Native Datadog payload policy and real HTTP/LLM span contracts.""" +import asyncio import json import os import subprocess @@ -69,7 +70,7 @@ def test_datadog_initialize_tags_use_client_info( monkeypatch.setattr(meta, "_MCP_MODE_ENABLED", True) monkeypatch.setattr(meta, "get_http_headers", lambda: headers) monkeypatch.setattr(_datadog, "_request_trace_attributes", lambda: {}) - monkeypatch.setattr(_datadog, "_annotate_attributes", lambda *_: None) + monkeypatch.setattr(_datadog, "_annotate_attributes", lambda *_, **__: None) annotations: list[dict[str, object]] = [] monkeypatch.setattr( "ddtrace.llmobs.LLMObs.annotate", @@ -198,6 +199,138 @@ def _problem(problem_type: str, title: str) -> str: }) +def _wrapped_sdk_error_with_request( + method: str, url: str, status_code: int = 403 +) -> ToolError: + error = _wrapped_sdk_error(status_code, _problem("forbidden", "Forbidden")) + sdk_error = error.__cause__.__cause__ + assert isinstance(sdk_error, SDKError) + sdk_error.raw_response.request = requests.Request(method, url).prepare() + return error + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + ( + "https://user:pw@api.airbyte.com/v1/connections/" + "123e4567-e89b-42d3-a456-426614174000?x=secret", + ("GET", "/v1/connections/{id}"), + ), + ( + "https://api.airbyte.com/v1/jobs/12345/status", + ("GET", "/v1/jobs/{id}/{id}"), + ), + ( + "https://api.airbyte.com/v1/Connections/MyConnection", + ("GET", "/v1/{id}/{id}"), + ), + ( + "https://h/customers/users/api/public/v1/connections/abc-def", + ("GET", "/{id}/{id}/{id}/{id}/v1/connections/{id}"), + ), + ( + "https://h/v1/customers/v2/connections/abc", + ("GET", "/{id}/{id}/v2/connections/{id}"), + ), + ( + "https://api.airbyte.com/v1/workspaces/" + "123e4567-e89b-42d3-a456-426614174000/definitions/sources/" + "123e4567-e89b-42d3-a456-426614174001", + ("GET", "/v1/workspaces/{id}/definitions/sources/{id}"), + ), + ( + "https://api.airbyte.com/v1/sources/initiateOAuth", + ("GET", "/v1/sources/initiateOAuth"), + ), + ( + "https://api.airbyte.com/connections/abc", + ("GET", "/{id}/{id}"), + ), + ], +) +def test_cloud_error_route_sanitizes_request_url( + url: str, expected: tuple[str, str] +) -> None: + error = _wrapped_sdk_error_with_request("GET", url) + + assert _datadog.cloud_error_route(error) == expected + assert "user:pw" not in str(expected) + assert "secret" not in str(expected) + + +def test_cloud_error_route_omits_missing_or_invalid_requests() -> None: + without_request = _wrapped_sdk_error(403, _problem("forbidden", "Forbidden")) + weird_method = _wrapped_sdk_error_with_request( + "CONNECT", "https://api.airbyte.com/v1/connections/123" + ) + + assert _datadog.cloud_error_route(without_request) is None + assert _datadog.cloud_error_route(RuntimeError("private")) is None + assert _datadog.cloud_error_route(weird_method) is None + assert "airbyte.mcp.upstream.method" not in _datadog._exception_attributes( + RuntimeError("private") + ) + assert "airbyte.mcp.upstream.method" not in _datadog._exception_attributes( + without_request + ) + assert "airbyte.mcp.upstream.route" not in _datadog._exception_attributes( + weird_method + ) + + +@pytest.mark.parametrize( + ("status_code", "problem_type", "expected_message"), + [ + ( + 404, + "resource-not-found", + "The resource was not found. (Cloud error: resource-not-found, HTTP 404)", + ), + ( + 403, + "forbidden", + "The requested resource was not found, or these credentials can't access it " + "(HTTP 403).", + ), + ], +) +def test_cloud_error_trace_message_for_config_api_errors( + monkeypatch: pytest.MonkeyPatch, + status_code: int, + problem_type: str, + expected_message: str, +) -> None: + response = requests.Response() + response.status_code = status_code + response.url = "https://config.airbyte.com/v1/workspaces/get" + response.request = requests.Request("POST", response.url).prepare() + response._content = json.dumps({"type": problem_type, "title": "SENTINEL"}).encode() + monkeypatch.setattr(api_util.requests, "request", lambda **_: response) + + with pytest.raises(api_util.AirbyteCloudError) as exc_info: + api_util._make_config_api_request( + path="/workspaces/get", + json={"workspaceId": "SENTINEL"}, + api_root="https://api.airbyte.com/v1", + config_api_root="https://config.airbyte.com/v1", + client_id=None, + client_secret=None, + bearer_token=api_util.SecretString("token"), + ) + + assert isinstance(exc_info.value.__cause__, requests.HTTPError) + assert _datadog.cloud_error_trace_message(exc_info.value) == expected_message + + +def test_cloud_error_trace_message_ignores_missing_resource_without_http_error() -> ( + None +): + error = api_util.AirbyteMissingResourceError(message="resource SENTINEL") + + assert _datadog.cloud_error_trace_message(error) is None + + @pytest.mark.parametrize( ("error", "expected"), [ @@ -358,12 +491,15 @@ def test_exception_attributes_classify_without_error_text( error: BaseException, expected: dict[str, str] ) -> None: attrs = _datadog._exception_attributes(error) - - assert attrs == { + expected_attrs = { "airbyte.mcp.outcome": "exception", "error.type": expected["airbyte.mcp.error_type"], **expected, } + if message := _datadog.cloud_error_trace_message(error): + expected_attrs["error.message"] = message + + assert attrs == expected_attrs assert "SENTINEL" not in json.dumps(attrs) @@ -417,6 +553,92 @@ def test_returned_tool_error_is_classified_as_tool_error() -> None: } +@pytest.mark.parametrize( + ("method", "params", "result"), + [ + ( + "initialize", + { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "clientInfo": {"name": "agent", "version": "1"}, + }, + { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "serverInfo": {"name": "airbyte", "version": "1"}, + }, + ), + ("tools/list", {}, {"tools": []}), + ], +) +def test_non_tool_requests_create_apm_spans_without_llmobs( + monkeypatch: pytest.MonkeyPatch, + method: str, + params: dict[str, object], + result: dict[str, object], +) -> None: + from ddtrace import tracer + from ddtrace.llmobs import LLMObs + + class Span: + def __init__(self) -> None: + self.tags: dict[str, str] = {} + self.metrics: dict[str, int] = {} + self.error = 0 + + def set_tag(self, key: str, value: str) -> None: + self.tags[key] = value + + def set_tags(self, tags: dict[str, str]) -> None: + self.tags.update(tags) + + def get_tag(self, key: str) -> str | None: + return self.tags.get(key) + + def set_metric(self, key: str, value: int) -> None: + self.metrics[key] = value + + def get_metric(self, key: str) -> int | None: + return self.metrics.get(key) + + def __exit__(self, *_args: object) -> None: + return None + + span = Span() + trace_calls: list[tuple[str, str]] = [] + annotate_calls: list[object] = [] + monkeypatch.setattr(LLMObs, "_instance", None) + monkeypatch.setattr( + LLMObs, "annotate", lambda *args, **kwargs: annotate_calls.append(args) + ) + + def trace(name: str, *, resource: str) -> Span: + trace_calls.append((name, resource)) + span.name = name + span.resource = resource + return span + + monkeypatch.setattr(tracer, "trace", trace) + + async def call_next(_ctx: object) -> dict[str, object]: + return result + + returned = asyncio.run( + _datadog._DatadogRequestMiddleware()( + SimpleNamespace(method=method, params=params), call_next + ) + ) + + assert returned == result + assert trace_calls == [(f"mcp.{method}", "server_request")] + assert span.name == f"mcp.{method}" + assert span.resource == "server_request" + assert span.tags["airbyte.mcp.outcome"] == "success" + assert span.tags["mcp_method"] == method + assert annotate_calls == [] + + @pytest.mark.parametrize( ("tags", "expected"), [ @@ -822,6 +1044,18 @@ def fail_tool() -> str: private_payload, response=response ) + @app.tool() + def raised_cloud_sdk_error() -> str: + request_url = ( + "https://api.airbyte.com/v1/connections/" + "123e4567-e89b-42d3-a456-426614174000?x=1" + ) + response = requests.Response() + response.status_code = 403 + response.url = request_url + response.request = requests.Request("GET", request_url).prepare() + raise SDKError("private response text", 403, "{}", response) + @app.tool() async def handle_error() -> str: from fastmcp.exceptions import ToolError @@ -871,6 +1105,7 @@ async def run(): }, ), ("tools/list", {}), + ("tools/call", {"name": "raised_cloud_sdk_error", "arguments": {}}), *[ ( "tools/call", @@ -907,6 +1142,65 @@ async def run(): ) assert response.status_code == 200 responses.append((method, params, response.json())) + lifecycle_spans = [ + span + for span in spans + if span.name in {"mcp.initialize", "mcp.tools/list"} + ] + assert {span.name for span in lifecycle_spans} == { + "mcp.initialize", + "mcp.tools/list", + } + for lifecycle_span in lifecycle_spans: + assert lifecycle_span.resource == "server_request" + assert lifecycle_span.span_type != "llm" + assert lifecycle_span.get_tag("airbyte.mcp.outcome") == "success" + assert lifecycle_span._get_ctx_item("_llmobs.cached_event") is None + initialize_span = next( + span for span in lifecycle_spans if span.name == "mcp.initialize" + ) + assert initialize_span.get_tag("client_name") == "custom-agent" + assert initialize_span.get_tag("client_version") == "custom-agent_1.2" + assert ( + initialize_span.get_tag("airbyte.mcp.mcp_protocol_version") + == "2025-11-25" + ) + assert ( + initialize_span.get_tag("mcp_session_id") + == (client_properties["session_id"]) + ) + assert responses[2][2]["result"]["isError"] + cloud_span = next( + span + for span in reversed(spans) + if span.get_tag("airbyte.mcp.upstream.route") + == "/v1/connections/{id}" + ) + cloud_metadata = cloud_span._get_ctx_item("_llmobs.cached_event")[ + "meta" + ]["metadata"] + cloud_uuid = "123e4567-e89b-42d3-a456-426614174000" + assert cloud_span.get_tag("error.message") == ( + api_util.FORBIDDEN_RESOURCE_MESSAGE + ) + assert cloud_span.get_tag("airbyte.mcp.upstream.method") == "GET" + assert cloud_span.get_tag("airbyte.mcp.upstream.route") == ( + "/v1/connections/{id}" + ) + assert cloud_uuid not in json.dumps(cloud_span.get_tags()) + assert "x=1" not in json.dumps(cloud_span.get_tags()) + assert cloud_uuid not in json.dumps(cloud_metadata) + assert "x=1" not in json.dumps(cloud_metadata) + assert "error.message" not in cloud_metadata + returned_error_span = next( + span + for span in reversed(spans) + if span.get_tag("airbyte.mcp.outcome") == "tool_error" + and span.get_tag("error.type") == "ToolError" + ) + assert returned_error_span.get_tag("error.message").startswith( + "tool resulted in an error (" + ) # Telemetry failures must not abort dispatch or replace its result. for target in ("tool", "annotate"): before = len(calls) @@ -1062,6 +1356,85 @@ async def unavailable_tool(_ctx): "tool resulted in an error (tool_unavailable)" ) + async def raised_tool_message(error: BaseException) -> str | None: + async def fail_tool(_ctx): + raise error + + with pytest.raises(type(error)): + await _datadog._DatadogRequestMiddleware()( + SimpleNamespace( + method="tools/call", + params={"name": "failed_cloud_tool", "arguments": {}}, + ), + fail_tool, + ) + span = next(span for span in reversed(spans) if span.span_type == "llm") + return span.get_tag("error.message") + + forbidden_message = await raised_tool_message( + _wrapped_sdk_error( + 403, + _problem( + "https://reference.airbyte.com/reference/errors#forbidden", + "forbidden", + ), + ) + ) + assert forbidden_message == api_util.FORBIDDEN_RESOURCE_MESSAGE + + upstream_uuid = "123e4567-e89b-42d3-a456-426614174000" + upstream_query = "x=secret-query" + upstream_message = await raised_tool_message( + _wrapped_sdk_error_with_request( + "GET", + f"https://user:pw@api.airbyte.com/v1/connections/{upstream_uuid}" + f"?{upstream_query}", + ) + ) + assert upstream_message == api_util.FORBIDDEN_RESOURCE_MESSAGE + upstream_span = next( + span for span in reversed(spans) if span.span_type == "llm" + ) + assert upstream_span.get_tag("airbyte.mcp.upstream.method") == "GET" + assert upstream_span.get_tag("airbyte.mcp.upstream.route") == ( + "/v1/connections/{id}" + ) + upstream_metadata = upstream_span._get_ctx_item("_llmobs.cached_event")["meta"][ + "metadata" + ] + assert upstream_uuid not in json.dumps(upstream_span.get_tags()) + assert upstream_uuid not in json.dumps(upstream_metadata) + assert upstream_query not in json.dumps(upstream_span.get_tags()) + assert upstream_query not in json.dumps(upstream_metadata) + + unexpected_problem = json.loads( + _problem(_GENERIC_PROBLEM_TYPE, "unexpected-problem") + ) + error_id = "123e4567-e89b-42d3-a456-426614174000" + unexpected_problem["errorId"] = error_id + unexpected_message = await raised_tool_message( + _wrapped_sdk_error(500, json.dumps(unexpected_problem)) + ) + assert unexpected_message is not None + assert unexpected_message.endswith( + "(Cloud error: unexpected-problem, HTTP 500)" + ) + assert error_id not in unexpected_message + assert "Error ID" not in unexpected_message + + hostile_title = "password=SENTINEL-cloud-title" + hostile_message = await raised_tool_message( + _wrapped_sdk_error(500, _problem(_GENERIC_PROBLEM_TYPE, hostile_title)) + ) + assert hostile_message is not None + assert hostile_message.endswith("(HTTP 500)") + assert hostile_title not in hostile_message + + runtime_message = await raised_tool_message( + RuntimeError("private runtime message") + ) + assert runtime_message == "tool resulted in an error (unclassified)" + async def fail_request(_ctx): raise RuntimeError("private request failure") @@ -1071,8 +1444,11 @@ async def fail_request(_ctx): fail_request, ) failed_request = next( - span for span in reversed(spans) if span.span_type == "llm" + span for span in reversed(spans) if span.name == "mcp.tools/list" ) + assert failed_request.error + assert failed_request.get_tag("airbyte.mcp.outcome") == "exception" + assert failed_request.get_tag("error.type") == "RuntimeError" assert failed_request.get_tag("error.message") is None # Cancellation must propagate and restore the previous Datadog context. @@ -1182,8 +1558,8 @@ async def dispatch(_ctx): "private synthetic result", ): assert secret not in serialized_native - primary = native[:7] - assert len(primary) == 7 + primary = native[:6] + assert len(primary) == 6 lookup = {span.span_id: span for span in spans} for span in primary: assert lookup[span.parent_id].name == "starlette.request" @@ -1194,23 +1570,22 @@ async def dispatch(_ctx): assert metadata["auth_method"] == "client_credentials" assert metadata["mcp_protocol_version"] == "2025-11-25" assert metadata["session_id"] == client_properties["session_id"] - assert [event["meta"]["span"]["kind"] for event in events] == [ - "task", - "task", - "tool", - "tool", - "tool", - "tool", - "tool", + assert [event["meta"]["span"]["kind"] for event in events] == ["tool"] * 6 + assert [event["name"] for event in events] == [ + "raised_cloud_sdk_error", + "execute_external_api_query", + "execute_external_api_query", + "execute_external_api_query", + "run_sql_query", + "unknown_tool", ] - assert events[0]["name"] == "mcp.initialize" - assert "client_name:custom-agent" in events[0]["tags"] - assert "client_version:custom-agent_1.2" in events[0]["tags"] assert [span.resource for span in primary] == [ - "server_request", - "server_request", - ] + ["execute_external_api_query"] * 3 + ["run_sql_query", "unknown_tool"] - for index in (2, 3, 4): + "unknown_tool", + *(["execute_external_api_query"] * 3), + "run_sql_query", + "unknown_tool", + ] + for index in (1, 2, 3): span, event = primary[index], events[index] assert json.loads(event["meta"]["input"]["value"]) == { "method": "tools/call", @@ -1229,7 +1604,7 @@ async def dispatch(_ctx): entity = "Custom Entities/東京" * 12 # Only bounded approved telemetry leaves the process; execution # receives the original argument. - assert entities[index - 2] == entity + assert entities[index - 1] == entity assert event["meta"]["metadata"]["agent.entity_type"] == entity assert span.get_tag("airbyte.mcp.agent.entity_type") == entity assert ( @@ -1247,7 +1622,7 @@ async def dispatch(_ctx): assert span.get_tag("airbyte.mcp.client_version") == "v" * 256 assert ( event["meta"]["metadata"]["tool_id"] - == hashlib.sha256(str(index + 1).encode()).hexdigest() + == hashlib.sha256(str(index + 3).encode()).hexdigest() ) assert "mcp_tool_kind:server" in event["tags"] assert ( @@ -1265,15 +1640,11 @@ async def dispatch(_ctx): and s.name == "urllib3.request" for s in spans ) - assert ( - json.loads(responses[2][2]["result"]["content"][0]["text"]) - == private_arguments - ) assert ( json.loads(responses[3][2]["result"]["content"][0]["text"]) == private_arguments ) - assert primary[2].error == 0 + assert primary[2].error == 1 tool_records = [ record for record in records if record.name == "execute_external_api_query" ] @@ -1286,14 +1657,14 @@ async def dispatch(_ctx): == "44444444-4444-4444-4444-444444444444" ) assert ( - events[5]["meta"]["metadata"]["organization_id"] + events[4]["meta"]["metadata"]["organization_id"] == "33333333-3333-3333-3333-333333333333" ) - assert "workspace_id" not in events[5]["meta"]["metadata"] + assert "workspace_id" not in events[4]["meta"]["metadata"] for index, outcome, error_type in ( - (2, "success", None), - (3, "tool_error", "ToolError"), - (4, "exception", "ValueError"), + (1, "success", None), + (2, "tool_error", "ToolError"), + (3, "exception", "ValueError"), ): span, event = primary[index], events[index] assert bool(span.error) is (error_type is not None) @@ -1301,18 +1672,18 @@ async def dispatch(_ctx): assert span.get_tag("airbyte.mcp.outcome") == outcome assert event["meta"]["metadata"]["outcome"] == outcome assert event["meta"]["metadata"].get("error_type") == error_type - assert events[4]["meta"]["metadata"]["error_type"] == "ValueError" - assert events[3]["meta"]["metadata"]["error.category"] == "tool_error" - assert events[4]["meta"]["metadata"]["error.category"] == "unclassified" - assert primary[3].get_tag("error.message") == ( + assert events[3]["meta"]["metadata"]["error_type"] == "ValueError" + assert events[2]["meta"]["metadata"]["error.category"] == "tool_error" + assert events[3]["meta"]["metadata"]["error.category"] == "unclassified" + assert primary[2].get_tag("error.message") == ( "tool resulted in an error (tool_error)" ) - assert events[5]["meta"]["output"]["value"] == "[REDACTED]" + assert events[4]["meta"]["output"]["value"] == "[REDACTED]" assert ( - responses[5][2]["result"]["content"][0]["text"] + responses[6][2]["result"]["content"][0]["text"] == "private synthetic result" ) - assert events[6]["name"] == "unknown_tool" and primary[6].error + assert events[5]["name"] == "unknown_tool" and primary[5].error for span_id, correlation in correlations[:3]: assert str(span_id) == correlation["dd.span_id"] asyncio.run(_native_distributed_context_contract(app))