diff --git a/.gitignore b/.gitignore index b867eb5f0..9dc0d1d83 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ docs/*.zip # Generated MCP server docs (regenerate via `poe mcp-docs-md`) docs/mcp-generated/ +/arg_trace_plan.json # Misc .DS_Store diff --git a/AGENTS.md b/AGENTS.md index 3dc0d21b2..c695475bb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -34,3 +34,8 @@ unredacted Cloud message without guidance; this does not add it to LLM Obs metad Non-Cloud MCP trace messages are the first line of the exception class's docstring (trace-only, exported to Datadog as `error.message`), so keep those first lines fixed text with no IDs, values or URLs. + +## MCP Argument Tracing + +Trace classes come from `airbyte/mcp/arg_trace_rules.yaml`; do not put `TraceArg` markers in tool signatures. +When `tests/unit_tests/test_mcp_arg_trace_parity.py` reports an unclassified argument, add or adjust a YAML rule; `poe mcp-arg-trace-plan` writes the full plan and each decision source to gitignored `arg_trace_plan.json`. diff --git a/airbyte/mcp/_arg_trace_rules.py b/airbyte/mcp/_arg_trace_rules.py new file mode 100644 index 000000000..48614cf80 --- /dev/null +++ b/airbyte/mcp/_arg_trace_rules.py @@ -0,0 +1,11 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Load MCP argument-tracing rules from the adjacent YAML file.""" + +from pathlib import Path + +import yaml + + +ARG_TRACE_RULES: dict[str, str] = yaml.safe_load( + Path(__file__).with_name("arg_trace_rules.yaml").read_text(encoding="utf-8") +) diff --git a/airbyte/mcp/_datadog.py b/airbyte/mcp/_datadog.py index a08ff38b8..f22e1a339 100644 --- a/airbyte/mcp/_datadog.py +++ b/airbyte/mcp/_datadog.py @@ -33,6 +33,7 @@ ) from opentelemetry.instrumentation.requests import RequestsInstrumentor +from airbyte.mcp._arg_trace_rules import ARG_TRACE_RULES from airbyte.mcp._error_handling import ( MCP_TOOL_USER_FACING_ERRORS, classify_mcp_tool_error, @@ -88,7 +89,12 @@ _TOOL_ANNOTATIONS: dict[str, dict[str, Any]] = {} _ARG_PREFIX = "airbyte.mcp" # Same classification and validation as the OTel backend; intent is exported separately. -_ARG_TRACER = ArgTracer(_ARG_PREFIX, key=_arg_key, skip=(INTENT_ARG,)) +_ARG_TRACER = ArgTracer( + _ARG_PREFIX, + rules=ARG_TRACE_RULES, + 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" diff --git a/airbyte/mcp/_otel.py b/airbyte/mcp/_otel.py index 051402741..89c212199 100644 --- a/airbyte/mcp/_otel.py +++ b/airbyte/mcp/_otel.py @@ -21,6 +21,7 @@ from opentelemetry.trace import SpanKind, Status from airbyte.constants import CLOUD_API_ROOT, CLOUD_CONFIG_API_ROOT +from airbyte.mcp._arg_trace_rules import ARG_TRACE_RULES from airbyte.mcp._error_handling import classify_mcp_tool_error, mcp_tool_error_reason @@ -322,6 +323,7 @@ def install(app: FastMCP, *, environ: Mapping[str, str] | None = None) -> None: error_reason=mcp_tool_error_reason, other_spans=_http_client_span_attributes, arg_key=_arg_key, + arg_rules=ARG_TRACE_RULES, exporter=_exporter(backend, environ), ), ), diff --git a/airbyte/mcp/_trace_plan.py b/airbyte/mcp/_trace_plan.py new file mode 100644 index 000000000..8c386f2bc --- /dev/null +++ b/airbyte/mcp/_trace_plan.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Write the MCP argument trace plan and its classification sources.""" + +import asyncio +import json +from pathlib import Path + +from fastmcp_extensions.otel.middleware import trace_plan + +from airbyte.mcp import _otel as observability +from airbyte.mcp.server import app + + +def main() -> None: + """Write the trusted-execution MCP trace plan to the repository root.""" + observability.install(app, environ={"AIRBYTE_MCP_TRACING_BACKEND": "otel"}) + plan = asyncio.run(trace_plan(app, include_sources=True)) + target = Path(__file__).resolve().parents[2] / "arg_trace_plan.json" + target.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/airbyte/mcp/arg_trace_rules.yaml b/airbyte/mcp/arg_trace_rules.yaml new file mode 100644 index 000000000..8128107ae --- /dev/null +++ b/airbyte/mcp/arg_trace_rules.yaml @@ -0,0 +1,64 @@ +# How each MCP tool argument is recorded in trace spans. +# Keys are argument names or fnmatch globs; `tool:arg` keys apply to one tool. +# The first matching key wins. `bool`, `Literal`, and `Enum` arguments are +# recorded as their value, and dict or model arguments as presence, so they +# need no rule unless they should be recorded differently. + +# Tool-scoped exceptions. +get_stream_previews:streams: fingerprint +sync_source_to_cache:streams: fingerprint + +# Not recorded. +workspace_id: omit +intent: omit + +# Identifiers, numbers, and opaque inputs: keyed hash. +"*_id": hash +streams: hash +cursor: hash +sql: hash +sql_query: hash +user_email: hash +api_args: hash +manifest_yaml_file: hash +manifest_yaml_text: hash +wait_timeout: hash +min_similarity: hash +"max_*": hash +"num_*": hash +"*_hours": hash +"*_seconds": hash +"*_number": hash +"*_offset": hash + +# Free text and names: keyed hash plus a similarity fingerprint. +name: fingerprint +"*_name": fingerprint +"*_names": fingerprint +name_contains: fingerprint +selected_streams: fingerprint +"*_fields": fingerprint +scenarios: fingerprint +config_file: fingerprint +manifest_path: fingerprint +"*prefix": fingerprint +"namespace*": fingerprint +reuse_namespace: fingerprint +cron_expression: fingerprint +docker_image: fingerprint +keyword_filter: fingerprint +prompt: fingerprint +query: fingerprint +search: fingerprint +section: fingerprint + +# Small fixed vocabularies or bounded numbers: raw value. +limit: value +page_size: value +connector_type: value +entity_type: value +support_level: value +min_support_level: value +sql_dialect: value +backfill_start_time: value +defer_credentials: value # Overrides the `credential` secret-name heuristic. diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 39c1d41b0..52fc0d329 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -17,7 +17,7 @@ import requests from fastmcp import Context, FastMCP -from fastmcp_extensions import TraceArg, get_mcp_config, register_mcp_tools +from fastmcp_extensions import get_mcp_config, register_mcp_tools from pydantic import BaseModel, ConfigDict, Field from airbyte import Destination, Source, get_destination, get_source @@ -618,7 +618,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. name: Annotated[ str, Field(description="The name to use when deploying the connector."), - TraceArg.FINGERPRINT, ], connector_name: Annotated[ str | None, @@ -629,7 +628,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. ), default=None, ), - TraceArg.FINGERPRINT, ] = None, *, connector_type: Annotated[ @@ -654,7 +652,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. ), default=None, ), - TraceArg.HASH, ] = None, workspace_id: Annotated[ str | None, @@ -662,7 +659,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], config: Annotated[ dict | str | None, @@ -677,7 +673,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], unique: Annotated[ bool, @@ -692,7 +687,6 @@ def deploy_connector_to_cloud( # noqa: PLR0913 # Mirrors the API surface. description=DEFER_CREDENTIALS_TIP_TEXT, default=False, ), - TraceArg.VALUE, ] = False, ) -> str: """Deploy a source or destination connector to Airbyte Cloud. @@ -873,7 +867,6 @@ def create_connection_on_cloud( connection_name: Annotated[ str, Field(description="The name of the connection."), - TraceArg.FINGERPRINT, ], source_id: Annotated[ str, @@ -890,7 +883,6 @@ def create_connection_on_cloud( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], selected_streams: Annotated[ str | list[str] | None, @@ -903,7 +895,6 @@ def create_connection_on_cloud( ), default=None, ), - TraceArg.FINGERPRINT, ], table_prefix: Annotated[ str | None, @@ -911,7 +902,6 @@ def create_connection_on_cloud( description="Optional table prefix to use when syncing to the destination.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> str: """Create a connection between a deployed source and destination on Airbyte Cloud. @@ -985,7 +975,6 @@ def run_cloud_sync( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], wait: Annotated[ bool, @@ -1024,7 +1013,7 @@ def run_cloud_sync( @mcp_tool() def deploy_noop_destination_to_cloud( ctx: Context, - name: Annotated[str, TraceArg.FINGERPRINT] = "No-op Destination", + name: str = "No-op Destination", *, workspace_id: Annotated[ str | None, @@ -1032,7 +1021,6 @@ def deploy_noop_destination_to_cloud( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], unique: bool = True, ) -> str: @@ -1081,7 +1069,6 @@ def get_cloud_sync_status( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], include_attempts: Annotated[ bool, @@ -1147,7 +1134,6 @@ def list_cloud_sync_jobs( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], max_jobs: Annotated[ int, @@ -1250,7 +1236,6 @@ def cancel_cloud_sync( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> SyncJobResult: """Cancel a running sync job on an Airbyte Cloud connection.""" @@ -1306,7 +1291,6 @@ def list_cloud_connectors( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], connector_type: Annotated[ ConnectorType | None, @@ -1321,7 +1305,6 @@ def list_cloud_connectors( description="Optional case-insensitive substring to filter connectors by name", default=None, ), - TraceArg.FINGERPRINT, ], limit: Annotated[ int | None, @@ -1329,7 +1312,6 @@ def list_cloud_connectors( description="Optional maximum number of items to return (default: no limit)", default=None, ), - TraceArg.VALUE, ], feature_filter: Annotated[ ConnectorFeature | None, @@ -1583,7 +1565,6 @@ def describe_cloud_connector( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], with_config: Annotated[ bool, @@ -1670,7 +1651,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con "`get_agent_skill_docs` for supported entity types." ), ), - TraceArg.VALUE, ], action: Annotated[ ExternalApiReadOnlyAction, @@ -1688,7 +1668,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con ), default=None, ), - TraceArg.HASH, ] = None, select_fields: Annotated[ list[str] | str | None, @@ -1696,7 +1675,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con description="Fields to keep in the response, as a list or a CSV string.", default=None, ), - TraceArg.FINGERPRINT, ] = None, exclude_fields: Annotated[ list[str] | str | None, @@ -1704,7 +1682,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con description="Fields to drop from the response, as a list or a CSV string.", default=None, ), - TraceArg.FINGERPRINT, ] = None, page_size: Annotated[ int | None, @@ -1712,7 +1689,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con description="Maximum number of entities to return in this page.", default=None, ), - TraceArg.VALUE, ] = None, cursor: Annotated[ str | None, @@ -1741,7 +1717,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con ), default=None, ), - TraceArg.OMIT, ] = None, workspace_id: Annotated[ str | None, @@ -1749,7 +1724,6 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, ) -> ExternalApiExecuteResult: """Read data from an external system through a deployed Cloud connector's direct API. @@ -1911,7 +1885,6 @@ def execute_external_sql_query( ), default=None, ), - TraceArg.VALUE, ] = None, page_size: Annotated[ int | None, @@ -1919,7 +1892,6 @@ def execute_external_sql_query( description="Maximum number of rows to return in this page.", default=None, ), - TraceArg.VALUE, ] = None, cursor: Annotated[ str | None, @@ -1944,7 +1916,6 @@ def execute_external_sql_query( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, ) -> ExternalApiExecuteResult: """Run a read-only SQL query against a deployed SQL-passthrough destination. @@ -2031,7 +2002,6 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the prompt: Annotated[ str, Field(description="The search text, for example `refund requests from ACME`."), - TraceArg.FINGERPRINT, ], search_type: Annotated[ ExternalSearchType, @@ -2046,7 +2016,6 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the description="Maximum number of hits to return, 1 to 100. Defaults to the backend's.", default=None, ), - TraceArg.VALUE, ] = None, streams: Annotated[ list[dict[str, Any]] | str | None, @@ -2061,7 +2030,6 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the ), default=None, ), - TraceArg.HASH, ] = None, lookback_seconds: Annotated[ int | None, @@ -2110,7 +2078,6 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, ) -> ExternalSearchResult: """Search indexed data for a deployed Cloud source or destination. @@ -2181,7 +2148,6 @@ def get_cloud_search_status( description="Optional stream name to report. Omit to report every stream.", default=None, ), - TraceArg.FINGERPRINT, ] = None, namespace: Annotated[ str | None, @@ -2189,7 +2155,6 @@ def get_cloud_search_status( description="Optional stream namespace to report. Omit to report every namespace.", default=None, ), - TraceArg.FINGERPRINT, ] = None, workspace_id: Annotated[ str | None, @@ -2197,7 +2162,6 @@ def get_cloud_search_status( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, ) -> ExternalSearchStatusResult: """Report search indexing status for a deployed Cloud source or destination. @@ -2281,7 +2245,6 @@ def request_cloud_search_index_backfill( ), default=None, ), - TraceArg.HASH, ] = None, backfill_start_time: Annotated[ str | None, @@ -2294,7 +2257,6 @@ def request_cloud_search_index_backfill( ), default=None, ), - TraceArg.VALUE, ] = None, workspace_id: Annotated[ str | None, @@ -2302,7 +2264,6 @@ def request_cloud_search_index_backfill( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, ) -> ExternalSearchBackfillRequestResult: """Request a search backfill that indexes data already in a connection's destination. @@ -2365,7 +2326,6 @@ def check_cloud_connector( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> ConnectorCheckResult: """Check the configuration and credentials of a deployed source or destination.""" @@ -2397,7 +2357,6 @@ def describe_cloud_connection( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> CloudConnectionDetails: """Get detailed information about a specific deployed connection.""" @@ -2446,7 +2405,6 @@ def get_cloud_sync_logs( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], max_lines: Annotated[ int, @@ -2586,7 +2544,6 @@ def list_cloud_connections( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], name_contains: Annotated[ str | None, @@ -2594,7 +2551,6 @@ def list_cloud_connections( description="Optional case-insensitive substring to filter connections by name", default=None, ), - TraceArg.FINGERPRINT, ], limit: Annotated[ int | None, @@ -2602,7 +2558,6 @@ def list_cloud_connections( description="Optional maximum number of items to return (default: no limit)", default=None, ), - TraceArg.VALUE, ], with_connection_status: Annotated[ bool | None, @@ -2720,7 +2675,6 @@ def list_cloud_workspaces( description=("Optional organization name (exact match) to list workspaces within."), default=None, ), - TraceArg.FINGERPRINT, ], name_contains: Annotated[ str | None, @@ -2728,7 +2682,6 @@ def list_cloud_workspaces( description="Optional substring to filter workspaces by name (server-side filtering)", default=None, ), - TraceArg.FINGERPRINT, ], limit: Annotated[ int | None, @@ -2736,7 +2689,6 @@ def list_cloud_workspaces( description="Optional maximum number of items to return (default: no limit)", default=None, ), - TraceArg.VALUE, ], privilege_scope: Annotated[ WorkspacePrivilegeScope, @@ -2943,7 +2895,6 @@ def set_default_cloud_workspace( "workspaces are rejected." ), ), - TraceArg.OMIT, ], ) -> CloudDefaultWorkspaceUpdateResult: """Durably set the authenticated user's default Airbyte Cloud workspace. @@ -2986,7 +2937,6 @@ def list_cloud_organizations( description="Optional case-insensitive substring to filter organization names.", default=None, ), - TraceArg.FINGERPRINT, ] = None, limit: Annotated[ int | None, @@ -2994,7 +2944,6 @@ def list_cloud_organizations( description="Optional maximum number of organizations to return (default: 100).", default=None, ), - TraceArg.VALUE, ] = None, feature_filter: Annotated[ OrganizationFeature | None, @@ -3084,7 +3033,6 @@ def describe_cloud_workspace( ), default=None, ), - TraceArg.OMIT, ], ) -> CloudWorkspaceResult: """Get basic details about a workspace (ID, name, URL, parent organization). @@ -3133,7 +3081,6 @@ def describe_cloud_organization( description="Organization name (exact match).", default=None, ), - TraceArg.FINGERPRINT, ], ) -> CloudOrganizationResult: """Get basic details about an organization (ID, name, email, feature flags). @@ -3179,7 +3126,6 @@ def get_cloud_organization_billing_status( description="Organization name for an exact match, when ID is not provided.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> CloudOrganizationBillingStatusResult: """Get billing and account status for an organization. @@ -3231,7 +3177,6 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M name: Annotated[ str, Field(description="The name for the custom connector definition."), - TraceArg.FINGERPRINT, ], *, workspace_id: Annotated[ @@ -3240,7 +3185,6 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], manifest_yaml_text: Annotated[ str | None, @@ -3300,7 +3244,6 @@ def publish_custom_source_definition( # noqa: PLR0913 # Explicit args mirror M ), default=None, ), - TraceArg.FINGERPRINT, ], ) -> str: """Publish a custom YAML source connector definition to Airbyte Cloud. @@ -3365,7 +3308,6 @@ def list_custom_source_definitions( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> list[dict[str, Any]]: """List custom YAML source definitions in the Airbyte Cloud workspace. @@ -3406,7 +3348,6 @@ def get_custom_source_definition( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], include_draft: Annotated[ bool, @@ -3468,7 +3409,6 @@ def get_connector_builder_draft_manifest( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> dict[str, Any]: """Get the Connector Builder draft manifest for a custom source definition. @@ -3531,7 +3471,6 @@ def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MC description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], name: Annotated[ str | None, @@ -3542,7 +3481,6 @@ def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MC ), default=None, ), - TraceArg.FINGERPRINT, ] = None, pre_validate: Annotated[ bool, @@ -3577,7 +3515,6 @@ def update_custom_source_definition( # noqa: PLR0913 # Explicit args mirror MC ), default=None, ), - TraceArg.FINGERPRINT, ], ) -> str: """Update a custom YAML source definition in Airbyte Cloud. @@ -3667,7 +3604,6 @@ def permanently_delete_custom_source_definition( name: Annotated[ str, Field(description="The expected name of the custom source definition (for verification)."), - TraceArg.FINGERPRINT, ], *, workspace_id: Annotated[ @@ -3676,7 +3612,6 @@ def permanently_delete_custom_source_definition( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Permanently delete a custom YAML source definition from Airbyte Cloud. @@ -3729,7 +3664,6 @@ def permanently_delete_cloud_connector( name: Annotated[ str, Field(description="The expected name of the connector (for verification)."), - TraceArg.FINGERPRINT, ], *, connector_type: Annotated[ @@ -3745,7 +3679,6 @@ def permanently_delete_cloud_connector( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Permanently delete a deployed source or destination connector from Airbyte Cloud. @@ -3803,7 +3736,6 @@ def permanently_delete_cloud_connection( name: Annotated[ str, Field(description="The expected name of the connection (for verification)."), - TraceArg.FINGERPRINT, ], *, cascade_delete_source: Annotated[ @@ -3830,7 +3762,6 @@ def permanently_delete_cloud_connection( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Permanently delete a connection from Airbyte Cloud. @@ -3879,7 +3810,6 @@ def rename_cloud_connector( name: Annotated[ str, Field(description="New name for the connector."), - TraceArg.FINGERPRINT, ], *, connector_type: Annotated[ @@ -3895,7 +3825,6 @@ def rename_cloud_connector( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Rename a deployed source or destination connector on Airbyte Cloud.""" @@ -3929,7 +3858,6 @@ def update_cloud_connector_config( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ] = None, *, connector_type: Annotated[ @@ -3945,7 +3873,6 @@ def update_cloud_connector_config( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Update a deployed source or destination connector's configuration on Airbyte Cloud. @@ -3980,7 +3907,6 @@ def rename_cloud_connection( name: Annotated[ str, Field(description="New name for the connection."), - TraceArg.FINGERPRINT, ], *, workspace_id: Annotated[ @@ -3989,7 +3915,6 @@ def rename_cloud_connection( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Rename a connection on Airbyte Cloud.""" @@ -4014,7 +3939,6 @@ def set_cloud_connection_table_prefix( prefix: Annotated[ str, Field(description="New table prefix to use when syncing to the destination."), - TraceArg.FINGERPRINT, ], *, workspace_id: Annotated[ @@ -4023,7 +3947,6 @@ def set_cloud_connection_table_prefix( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Set the table prefix for a connection on Airbyte Cloud. @@ -4058,7 +3981,6 @@ def set_cloud_connection_selected_streams( "Must be an explicit stream name or list of streams." ) ), - TraceArg.FINGERPRINT, ], *, workspace_id: Annotated[ @@ -4067,7 +3989,6 @@ def set_cloud_connection_selected_streams( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Set the selected streams for a connection on Airbyte Cloud. @@ -4129,7 +4050,6 @@ def update_cloud_connection( ), default=None, ), - TraceArg.FINGERPRINT, ], interval_hours: Annotated[ int | None, @@ -4163,7 +4083,6 @@ def update_cloud_connection( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> str: """Update a connection's settings on Airbyte Cloud. @@ -4255,7 +4174,6 @@ def get_connection_artifact( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> dict[str, Any] | list[dict[str, Any]]: """Get a connection artifact (state or catalog) from Airbyte Cloud. diff --git a/airbyte/mcp/guidance.py b/airbyte/mcp/guidance.py index 28638f5cb..d45fc5d7a 100644 --- a/airbyte/mcp/guidance.py +++ b/airbyte/mcp/guidance.py @@ -17,7 +17,6 @@ from fastmcp import Context, FastMCP from fastmcp_extensions import ( MCPServerConfigArg, - TraceArg, get_mcp_config, mcp_prompt, register_mcp_prompts, @@ -94,7 +93,6 @@ def search_airbyte_knowledge_sources( "A single, well-formed natural-language query. " "Must be a complete sentence." ) ), - TraceArg.FINGERPRINT, ], ) -> list[dict[str, str]]: """Search Airbyte knowledge sources. @@ -209,7 +207,6 @@ def get_agent_skill_docs( ), default=None, ), - TraceArg.FINGERPRINT, ] = None, workspace_id: Annotated[ str | None, @@ -217,7 +214,6 @@ def get_agent_skill_docs( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ], ) -> AgentSkillDocsResult: """Read a deployed connector's direct-access docs (agent skill docs). @@ -256,7 +252,6 @@ def get_connector_info( connector_name: Annotated[ str, Field(description="The name of the connector to get information for."), - TraceArg.FINGERPRINT, ], ) -> ConnectorInfo | Literal["Connector not found."]: """Get metadata, documentation URL, config spec, and manifest URL for a connector. @@ -326,7 +321,6 @@ def get_api_docs_urls( "(e.g., 'source-facebook-marketing', 'destination-snowflake')" ) ), - TraceArg.FINGERPRINT, ], ) -> list[ApiDocsUrl] | Literal["Connector not found."]: """Get API documentation URLs for a connector. diff --git a/airbyte/mcp/interactive/_registry_ui.py b/airbyte/mcp/interactive/_registry_ui.py index 0dbb371c4..3ed609e29 100644 --- a/airbyte/mcp/interactive/_registry_ui.py +++ b/airbyte/mcp/interactive/_registry_ui.py @@ -11,7 +11,6 @@ from fastmcp.apps import PrefabAppConfig from fastmcp.tools.base import ToolResult -from fastmcp_extensions import TraceArg from prefab_ui.actions import OpenLink, SendMessage, SetState from prefab_ui.app import PrefabApp from prefab_ui.components import ( @@ -92,7 +91,6 @@ def show_connectors_list( "or `archived`. Empty string means no filter." ), ), - TraceArg.VALUE, ] = "", certified: Annotated[ # noqa: FBT002 - FastMCP tool parameter. bool, @@ -111,7 +109,6 @@ def show_connectors_list( "< `certified`. Empty string means no filter." ), ), - TraceArg.VALUE, ] = "", connector_type: Annotated[ str, @@ -121,7 +118,6 @@ def show_connectors_list( "Empty string means no filter." ), ), - TraceArg.VALUE, ] = "", search: Annotated[ str, @@ -131,7 +127,6 @@ def show_connectors_list( "definition ID, Docker repository, subtype, and docs URL." ), ), - TraceArg.FINGERPRINT, ] = "", limit: Annotated[ int, @@ -139,7 +134,6 @@ def show_connectors_list( description="Maximum number of connectors to return. Use `0` for no limit.", ge=0, ), - TraceArg.VALUE, ] = 0, ) -> ToolResult: """Show an interactive public connector catalog from the OSS registry.""" diff --git a/airbyte/mcp/interactive/_sync_history_ui.py b/airbyte/mcp/interactive/_sync_history_ui.py index 90f01ab89..f9da62b57 100644 --- a/airbyte/mcp/interactive/_sync_history_ui.py +++ b/airbyte/mcp/interactive/_sync_history_ui.py @@ -9,7 +9,6 @@ from fastmcp import Context # noqa: TC002 - required at runtime for FastMCP tool registration from fastmcp.apps import PrefabAppConfig from fastmcp.tools.base import ToolResult -from fastmcp_extensions import TraceArg from prefab_ui.app import PrefabApp from prefab_ui.components import ( Accordion, @@ -91,7 +90,6 @@ def show_connection_sync_history( # noqa: PLR0914 description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, max_jobs: Annotated[ int, diff --git a/airbyte/mcp/interactive/_workspace_sync_status_ui.py b/airbyte/mcp/interactive/_workspace_sync_status_ui.py index 58c97cec6..500bc59f6 100644 --- a/airbyte/mcp/interactive/_workspace_sync_status_ui.py +++ b/airbyte/mcp/interactive/_workspace_sync_status_ui.py @@ -12,7 +12,6 @@ from fastmcp import Context # noqa: TC002 - required at runtime for FastMCP tool registration from fastmcp.apps import PrefabAppConfig from fastmcp.tools.base import ToolResult -from fastmcp_extensions import TraceArg from prefab_ui.actions import OpenLink, SendMessage, SetState from prefab_ui.app import PrefabApp from prefab_ui.components import ( @@ -128,7 +127,6 @@ def show_workspace_sync_status( description=WORKSPACE_ID_TIP_TEXT, default=None, ), - TraceArg.OMIT, ] = None, max_connections: Annotated[ int, diff --git a/airbyte/mcp/local.py b/airbyte/mcp/local.py index cee0f6d49..f2f22c02d 100644 --- a/airbyte/mcp/local.py +++ b/airbyte/mcp/local.py @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Literal from fastmcp import FastMCP -from fastmcp_extensions import TraceArg, register_mcp_tools +from fastmcp_extensions import register_mcp_tools from fastmcp_extensions.otel._extras import error_stack # noqa: PLC2701 from pydantic import BaseModel, Field @@ -145,7 +145,6 @@ def validate_connector_config( connector_name: Annotated[ str, Field(description="The name of the connector to validate."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -160,7 +159,6 @@ def validate_connector_config( description="Path to a YAML or JSON file containing the connector configuration.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -168,7 +166,6 @@ def validate_connector_config( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], override_execution_mode: Annotated[ Literal["docker", "python", "yaml", "auto"], @@ -184,7 +181,6 @@ def validate_connector_config( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> tuple[bool, str]: """Validate a connector configuration. @@ -235,7 +231,6 @@ def list_connector_config_secrets( connector_name: Annotated[ str, Field(description="The name of the connector."), - TraceArg.FINGERPRINT, ], ) -> list[str]: """List all `config_secret_name` options that are known for the given connector. @@ -290,7 +285,6 @@ def list_source_streams( source_connector_name: Annotated[ str, Field(description="The name of the source connector."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -305,7 +299,6 @@ def list_source_streams( description="Path to a YAML or JSON file containing the source connector config.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -313,7 +306,6 @@ def list_source_streams( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], override_execution_mode: Annotated[ Literal["docker", "python", "yaml", "auto"], @@ -329,7 +321,6 @@ def list_source_streams( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> list[str]: """List all streams available in a source connector. @@ -361,12 +352,10 @@ def get_source_stream_json_schema( source_connector_name: Annotated[ str, Field(description="The name of the source connector."), - TraceArg.FINGERPRINT, ], stream_name: Annotated[ str, Field(description="The name of the stream."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -381,7 +370,6 @@ def get_source_stream_json_schema( description="Path to a YAML or JSON file containing the source connector config.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -389,7 +377,6 @@ def get_source_stream_json_schema( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], override_execution_mode: Annotated[ Literal["docker", "python", "yaml", "auto"], @@ -405,7 +392,6 @@ def get_source_stream_json_schema( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> dict[str, Any]: """List all properties for a specific stream in a source connector.""" @@ -433,7 +419,6 @@ def read_source_stream_records( source_connector_name: Annotated[ str, Field(description="The name of the source connector."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -448,7 +433,6 @@ def read_source_stream_records( description="Path to a YAML or JSON file containing the source connector config.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -456,13 +440,11 @@ def read_source_stream_records( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], *, stream_name: Annotated[ str, Field(description="The name of the stream to read records from."), - TraceArg.FINGERPRINT, ], max_records: Annotated[ int, @@ -485,7 +467,6 @@ def read_source_stream_records( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> list[dict[str, Any]] | str: """Get records from a source connector.""" @@ -535,7 +516,6 @@ def get_stream_previews( source_name: Annotated[ str, Field(description="The name of the source connector."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -550,7 +530,6 @@ def get_stream_previews( description="Path to a YAML or JSON file containing the source connector config.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -558,7 +537,6 @@ def get_stream_previews( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], streams: Annotated[ list[str] | str | None, @@ -569,7 +547,6 @@ def get_stream_previews( ), default=None, ), - TraceArg.FINGERPRINT, ], limit: Annotated[ int, @@ -577,7 +554,6 @@ def get_stream_previews( description="The maximum number of sample records to return per stream.", default=10, ), - TraceArg.VALUE, ], override_execution_mode: Annotated[ Literal["docker", "python", "yaml", "auto"], @@ -593,7 +569,6 @@ def get_stream_previews( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> dict[str, list[dict[str, Any]] | str]: """Get sample records (previews) from streams in a source connector. @@ -659,7 +634,6 @@ def sync_source_to_cache( source_connector_name: Annotated[ str, Field(description="The name of the source connector."), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -674,7 +648,6 @@ def sync_source_to_cache( description="Path to a YAML or JSON file containing the source connector config.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -682,7 +655,6 @@ def sync_source_to_cache( description="The name of the secret containing the configuration.", default=None, ), - TraceArg.FINGERPRINT, ], streams: Annotated[ list[str] | str, @@ -690,7 +662,6 @@ def sync_source_to_cache( description="The streams to sync.", default="suggested", ), - TraceArg.FINGERPRINT, ], override_execution_mode: Annotated[ Literal["docker", "python", "yaml", "auto"], @@ -706,7 +677,6 @@ def sync_source_to_cache( description="Path to a local YAML manifest file for declarative connectors.", default=None, ), - TraceArg.FINGERPRINT, ], ) -> str: """Run a sync from a source connector to the default DuckDB cache.""" @@ -914,7 +884,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 "(e.g. 'destination-snowflake', 'destination-motherduck')." ), ), - TraceArg.FINGERPRINT, ], config: Annotated[ dict | str | None, @@ -932,7 +901,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 description="Path to a YAML or JSON file containing the destination configuration.", default=None, ), - TraceArg.FINGERPRINT, ], config_secret_name: Annotated[ str | None, @@ -940,7 +908,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 description="The name of the secret containing the destination configuration.", default=None, ), - TraceArg.FINGERPRINT, ], scenarios: Annotated[ list[str] | str, @@ -953,7 +920,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 ), default="fast", ), - TraceArg.FINGERPRINT, ], custom_scenarios: Annotated[ list[dict[str, Any]] | None, @@ -975,7 +941,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 ), default=None, ), - TraceArg.FINGERPRINT, ], namespace_suffix: Annotated[ str | None, @@ -987,7 +952,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 ), default=None, ), - TraceArg.FINGERPRINT, ], reuse_namespace: Annotated[ str | None, @@ -999,7 +963,6 @@ def destination_smoke_test( # noqa: PLR0913, PLR0917 ), default=None, ), - TraceArg.FINGERPRINT, ], skip_preflight: Annotated[ bool, diff --git a/airbyte/mcp/registry.py b/airbyte/mcp/registry.py index ceb9b56a5..f53f722d3 100644 --- a/airbyte/mcp/registry.py +++ b/airbyte/mcp/registry.py @@ -17,7 +17,7 @@ import requests from fastmcp import FastMCP -from fastmcp_extensions import TraceArg, mcp_tool, register_mcp_tools +from fastmcp_extensions import mcp_tool, register_mcp_tools from pydantic import Field from airbyte import exceptions as exc @@ -44,7 +44,6 @@ def list_connectors( description="Filter connectors by keyword.", default=None, ), - TraceArg.FINGERPRINT, ], connector_type_filter: Annotated[ Literal["source", "destination"] | None, @@ -135,7 +134,6 @@ def get_connector_version_history( Field( description="The name of the connector (e.g., 'source-faker', 'destination-postgres')" ), - TraceArg.FINGERPRINT, ], num_versions_to_validate: Annotated[ int, @@ -157,7 +155,6 @@ def get_connector_version_history( ), default=None, ), - TraceArg.VALUE, ] = None, ) -> list[ConnectorVersionInfo] | Literal["Connector not found.", "Failed to fetch changelog."]: """Get version history for a connector. diff --git a/pyproject.toml b/pyproject.toml index e4a006c1a..cff66c6b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,8 @@ dependencies = [ "typing-extensions", "uuid7>=0.1.0,<1.0", "fastmcp>=4.0.9,<5.0", - "fastmcp-extensions>=0.35.0,<1.0.0", + # TK-TODO: Restore fastmcp-extensions>=,<1.0.0 after the rules release. + "fastmcp-extensions @ git+https://github.com/airbytehq/fastmcp-extensions.git@devin/1792462457-arg-trace-rules", "starlette", # Transitive dependency of fastmcp, imported directly "httpx", # Transitive dependency of fastmcp, imported directly "mcp", # Transitive dependency of fastmcp, imported directly @@ -116,6 +117,7 @@ allow-direct-references = true [tool.hatch.build.targets.wheel] packages = ["airbyte"] +force-include = { "airbyte/mcp/arg_trace_rules.yaml" = "airbyte/mcp/arg_trace_rules.yaml" } [tool.pytest.ini_options] # Pytest configuration as follows: @@ -188,6 +190,7 @@ mcp-serve-sse = { cmd = "python -c \"from airbyte.mcp.server import app; app.run mcp-inspect = { cmd = "fastmcp inspect airbyte/mcp/server.py:app", help = "Inspect MCP tools and resources (supports --tools, --health, etc.)" } mcp-tool-test = { cmd = "python -m fastmcp_extensions.utils.test_tool --app airbyte.mcp.server:app", help = "Test MCP tools directly with JSON arguments: poe mcp-tool-test ''" } mcp-docs-md = { cmd = "python scripts/generate_mcp_markdown.py", help = "Generate Markdown docs for the MCP server into docs/mcp-generated/ (Docusaurus- and pdoc-compatible)" } +mcp-arg-trace-plan = { shell = "AIRBYTE_MCP_TRUSTED_EXECUTION=1 python -m airbyte.mcp._trace_plan" } # Claude Code MCP Testing Tasks [tool.poe.tasks.test-my-tools] diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json deleted file mode 100644 index 758a80928..000000000 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ /dev/null @@ -1,383 +0,0 @@ -{ - "cancel_cloud_sync": { - "connection_id": "hash", - "job_id": "hash", - "workspace_id": "omit" - }, - "check_cloud_connector": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "workspace_id": "omit" - }, - "create_connection_on_cloud": { - "connection_name": "fingerprint", - "destination_id": "hash", - "selected_streams": "fingerprint list", - "source_id": "hash", - "table_prefix": "fingerprint", - "workspace_id": "omit" - }, - "deploy_connector_to_cloud": { - "config": "presence", - "config_secret_name": "fingerprint", - "connector_name": "fingerprint", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "custom_definition_id": "hash", - "defer_credentials": "value allowed=[false,true]", - "name": "fingerprint", - "unique": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "deploy_noop_destination_to_cloud": { - "name": "fingerprint", - "unique": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "describe_cloud_connection": { - "connection_id": "hash", - "workspace_id": "omit" - }, - "describe_cloud_connector": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "with_config": "value allowed=[false,true]", - "with_data_replication_docs": "value allowed=[false,true]", - "with_direct_access_guidance": "value allowed=[false,true]", - "with_replication_details": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "describe_cloud_organization": { - "organization_id": "hash", - "organization_name": "fingerprint" - }, - "describe_cloud_workspace": { - "workspace_id": "omit" - }, - "describe_default_cache": {}, - "destination_smoke_test": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "custom_scenarios": "presence", - "destination_connector_name": "fingerprint", - "docker_image": "fingerprint", - "namespace_suffix": "fingerprint", - "reuse_namespace": "fingerprint", - "scenarios": "fingerprint list", - "skip_preflight": "value allowed=[false,true]" - }, - "execute_external_api_query": { - "action": "value allowed=[\"get\",\"list\",\"search\"]", - "api_args": "hash", - "connector_id": "hash", - "cursor": "hash", - "entity_type": "value str", - "exclude_fields": "fingerprint list", - "intent": "omit", - "page_size": "value int", - "select_fields": "fingerprint list", - "skip_truncation": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "execute_external_search_query": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "destination_id": "hash", - "limit": "value int", - "lookback_seconds": "hash", - "max_context_chars": "hash", - "max_similarity_diff": "hash", - "min_similarity": "hash", - "prompt": "fingerprint", - "search_type": "value allowed=[\"hybrid\",\"keyword\",\"semantic\"]", - "streams": "hash list", - "workspace_id": "omit" - }, - "execute_external_sql_query": { - "connector_id": "hash", - "cursor": "hash", - "dry_run": "value allowed=[false,true]", - "page_size": "value int", - "sql": "hash", - "sql_dialect": "value str", - "workspace_id": "omit" - }, - "get_agent_skill_docs": { - "connector_id": "hash", - "docs_skill_id": "hash", - "section": "fingerprint", - "workspace_id": "omit" - }, - "get_api_docs_urls": { - "connector_name": "fingerprint" - }, - "get_cloud_organization_billing_status": { - "organization_id": "hash", - "organization_name": "fingerprint" - }, - "get_cloud_search_status": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "namespace": "fingerprint", - "stream_name": "fingerprint", - "workspace_id": "omit" - }, - "get_cloud_sync_logs": { - "attempt_number": "hash", - "connection_id": "hash", - "from_tail": "value allowed=[false,true]", - "job_id": "hash", - "line_offset": "hash", - "max_lines": "hash", - "workspace_id": "omit" - }, - "get_cloud_sync_status": { - "connection_id": "hash", - "include_attempts": "value allowed=[false,true]", - "job_id": "hash", - "workspace_id": "omit" - }, - "get_connection_artifact": { - "artifact_type": "value allowed=[\"catalog\",\"state\"]", - "connection_id": "hash", - "workspace_id": "omit" - }, - "get_connector_builder_draft_manifest": { - "definition_id": "hash", - "workspace_id": "omit" - }, - "get_connector_info": { - "connector_name": "fingerprint" - }, - "get_connector_version_history": { - "connector_name": "fingerprint", - "limit": "value int", - "num_versions_to_validate": "hash" - }, - "get_custom_source_definition": { - "definition_id": "hash", - "include_draft": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "get_default_cloud_context": {}, - "get_source_stream_json_schema": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "manifest_path": "fingerprint", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]", - "source_connector_name": "fingerprint", - "stream_name": "fingerprint" - }, - "get_stream_previews": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "limit": "value int", - "manifest_path": "fingerprint", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]", - "source_name": "fingerprint", - "streams": "fingerprint list" - }, - "list_cached_streams": {}, - "list_cloud_connections": { - "failing_connections_only": "value allowed=[false,true]", - "limit": "value int", - "name_contains": "fingerprint", - "with_connection_status": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "list_cloud_connectors": { - "connector_type": "value allowed=[\"destination\",\"source\"]", - "feature_filter": "value allowed=[\"direct_access\",\"direct_api_action\",\"direct_api_query\",\"direct_sql_query\",\"search_indexing\"]", - "limit": "value int", - "name_contains": "fingerprint", - "workspace_id": "omit" - }, - "list_cloud_organizations": { - "feature_filter": "value allowed=[\"direct_access\",\"search_indexing\"]", - "limit": "value int", - "name_contains": "fingerprint" - }, - "list_cloud_sync_jobs": { - "connection_id": "hash", - "from_tail": "value allowed=[false,true]", - "job_type": "value allowed=[\"clear\",\"refresh\",\"reset\",\"sync\"]", - "max_jobs": "hash", - "workspace_id": "omit" - }, - "list_cloud_workspaces": { - "limit": "value int", - "name_contains": "fingerprint", - "organization_id": "hash", - "organization_name": "fingerprint", - "privilege_scope": "value allowed=[\"any\",\"instance_admin\",\"member_of\",\"organization_admin\"]" - }, - "list_connector_config_secrets": { - "connector_name": "fingerprint" - }, - "list_connectors": { - "connector_type_filter": "value allowed=[\"destination\",\"source\"]", - "install_types": "value allowed=[\"docker\",\"java\",\"python\",\"yaml\"] list", - "keyword_filter": "fingerprint" - }, - "list_custom_source_definitions": { - "workspace_id": "omit" - }, - "list_dotenv_secrets": {}, - "list_source_streams": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "manifest_path": "fingerprint", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]", - "source_connector_name": "fingerprint" - }, - "permanently_delete_cloud_connection": { - "cascade_delete_destination": "value allowed=[false,true]", - "cascade_delete_source": "value allowed=[false,true]", - "connection_id": "hash", - "name": "fingerprint", - "workspace_id": "omit" - }, - "permanently_delete_cloud_connector": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "name": "fingerprint", - "workspace_id": "omit" - }, - "permanently_delete_custom_source_definition": { - "definition_id": "hash", - "name": "fingerprint", - "workspace_id": "omit" - }, - "publish_custom_source_definition": { - "manifest_yaml_file": "hash", - "manifest_yaml_text": "hash", - "name": "fingerprint", - "pre_validate": "value allowed=[false,true]", - "testing_values": "presence", - "testing_values_secret_name": "fingerprint", - "unique": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "read_source_stream_records": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "manifest_path": "fingerprint", - "max_records": "hash", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]", - "source_connector_name": "fingerprint", - "stream_name": "fingerprint" - }, - "rename_cloud_connection": { - "connection_id": "hash", - "name": "fingerprint", - "workspace_id": "omit" - }, - "rename_cloud_connector": { - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "name": "fingerprint", - "workspace_id": "omit" - }, - "request_cloud_search_index_backfill": { - "backfill_start_time": "value str", - "connection_id": "hash", - "streams": "hash list", - "workspace_id": "omit" - }, - "run_cloud_sync": { - "connection_id": "hash", - "wait": "value allowed=[false,true]", - "wait_timeout": "hash", - "workspace_id": "omit" - }, - "run_sql_query": { - "max_records": "hash", - "sql_query": "hash" - }, - "search_airbyte_knowledge_sources": { - "query": "fingerprint" - }, - "set_cloud_connection_selected_streams": { - "connection_id": "hash", - "stream_names": "fingerprint list", - "workspace_id": "omit" - }, - "set_cloud_connection_table_prefix": { - "connection_id": "hash", - "prefix": "fingerprint", - "workspace_id": "omit" - }, - "set_default_cloud_workspace": { - "user_email": "hash", - "workspace_id": "omit" - }, - "show_connection_sync_history": { - "agent_context": "value allowed=[\"min\",\"summary\",\"verbose\"]", - "connection_id": "hash", - "max_jobs": "hash", - "suppress_ui": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "show_connectors_list": { - "certified": "value allowed=[false,true]", - "connector_type": "value str", - "limit": "value int", - "min_support_level": "value str", - "search": "fingerprint", - "support_level": "value str" - }, - "show_workspace_sync_status": { - "agent_context": "value allowed=[\"min\",\"summary\",\"verbose\"]", - "max_connections": "hash", - "max_jobs_per_connection": "hash", - "recent_hours": "hash", - "suppress_ui": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "sync_source_to_cache": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "manifest_path": "fingerprint", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]", - "source_connector_name": "fingerprint", - "streams": "fingerprint list" - }, - "update_cloud_connection": { - "connection_id": "hash", - "cron_expression": "fingerprint", - "enabled": "value allowed=[false,true]", - "interval_hours": "hash", - "manual_schedule": "value allowed=[false,true]", - "workspace_id": "omit" - }, - "update_cloud_connector_config": { - "config": "presence", - "config_secret_name": "fingerprint", - "connector_id": "hash", - "connector_type": "value allowed=[\"destination\",\"source\"]", - "workspace_id": "omit" - }, - "update_custom_source_definition": { - "definition_id": "hash", - "manifest_yaml_file": "hash", - "manifest_yaml_text": "hash", - "name": "fingerprint", - "pre_validate": "value allowed=[false,true]", - "testing_values": "presence", - "testing_values_secret_name": "fingerprint", - "workspace_id": "omit" - }, - "validate_connector_config": { - "config": "presence", - "config_file": "fingerprint", - "config_secret_name": "fingerprint", - "connector_name": "fingerprint", - "manifest_path": "fingerprint", - "override_execution_mode": "value allowed=[\"auto\",\"docker\",\"python\",\"yaml\"]" - } -} diff --git a/tests/unit_tests/test_mcp_arg_trace_parity.py b/tests/unit_tests/test_mcp_arg_trace_parity.py index a30295b3a..60dcae7b0 100644 --- a/tests/unit_tests/test_mcp_arg_trace_parity.py +++ b/tests/unit_tests/test_mcp_arg_trace_parity.py @@ -7,27 +7,27 @@ import base64 import json import logging -import os import socket from collections.abc import Iterator -from pathlib import Path import pytest from fastmcp import Client, FastMCP from fastmcp.server.auth import AccessToken from fastmcp_extensions import capture_tool_spans from fastmcp_extensions.otel import middleware as otel_middleware +from fastmcp_extensions.otel._extras import declared_parameters # noqa: PLC2701 from fastmcp_extensions.otel.middleware import trace_plan from opentelemetry import trace from opentelemetry.instrumentation.requests import RequestsInstrumentor from opentelemetry.sdk.trace import ReadableSpan +from airbyte.mcp import _datadog from airbyte.mcp import _otel as observability +from airbyte.mcp._arg_trace_rules import ARG_TRACE_RULES from airbyte.mcp import cloud from airbyte.mcp.server import app as server_app -FIXTURE = Path(__file__).parent / "fixtures" / "arg_trace_plan.json" ARG = "airbyte.mcp.arg." SENTINEL = "SENTINELqzxprivatevalue" CONNECTOR_ID = "11111111-1111-1111-1111-111111111111" @@ -59,30 +59,6 @@ } -def _plan() -> dict[str, dict[str, str]]: - plan = asyncio.run(trace_plan(server_app)) - return {name: entry["args"] for name, entry in plan.items()} - - -def test_inventory_matches_fixture() -> None: - plan = _plan() - if os.environ.get("UPDATE_ARG_TRACE_FIXTURE"): - assert all( - "manifest_yaml_file" in plan[name] - for name in ( - "publish_custom_source_definition", - "update_custom_source_definition", - ) - ), "Enable trusted execution when regenerating the fixture." - FIXTURE.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n") - - expected = json.loads(FIXTURE.read_text()) - for name in ("publish_custom_source_definition", "update_custom_source_definition"): - if "manifest_yaml_file" not in plan[name]: - expected[name].pop("manifest_yaml_file", None) - assert plan == expected - - @pytest.fixture def traced_app(monkeypatch: pytest.MonkeyPatch) -> Iterator[FastMCP]: for name in ("OTEL_EXPORTER_OTLP_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"): @@ -132,6 +108,59 @@ def loopback_only(sock: socket.socket, address: object) -> object: observability._reset_for_tests() +def test_every_arg_is_classified_by_a_rule(traced_app: FastMCP) -> None: + plan = asyncio.run(trace_plan(traced_app, include_sources=True)) + unclassified = [ + f"{tool}.{name} uses source {source}; add or adjust a rule in " + "`airbyte/mcp/arg_trace_rules.yaml` (do not add `TraceArg` markers)." + for tool, entry in plan.items() + for name, source in entry["arg_sources"].items() + if source in {"default", "marker"} + ] + assert not unclassified, "\n".join(unclassified) + + +def test_no_dead_rules(traced_app: FastMCP) -> None: + plan = asyncio.run(trace_plan(traced_app, include_sources=True)) + used_rules = { + source.removeprefix("rule:") + for entry in plan.values() + for source in entry["arg_sources"].values() + if source.startswith("rule:") + } + unused_rules = set(ARG_TRACE_RULES) - used_rules + if not any("manifest_yaml_file" in entry["arg_sources"] for entry in plan.values()): + # The app excludes this argument during import when trusted execution is off. + unused_rules.discard("manifest_yaml_file") + assert not unused_rules, ( + "Rules in `airbyte/mcp/arg_trace_rules.yaml` matched no tool arguments: " + f"{sorted(unused_rules)}" + ) + + +def test_datadog_and_otel_classification_agree(traced_app: FastMCP) -> None: + plan = asyncio.run(trace_plan(traced_app)) + tools = asyncio.run(traced_app.list_tools(run_middleware=False)) + assert {tool.name for tool in tools} == set(plan) + parameters = {tool.name: declared_parameters(tool) for tool in tools} + assert any("intent" in names for names in parameters.values()) + for tool in tools: + classes = _datadog._ARG_TRACER.classes( + tool.name, + getattr(tool, "fn", None), + parameters[tool.name], + ) + assert "intent" not in classes, tool.name + expected = { + name: mode + for name, mode in plan[tool.name]["args"].items() + if name != "intent" + } + assert {name: cls.describe() for name, cls in classes.items()} == expected, ( + tool.name + ) + + def _call_all(app: FastMCP) -> dict[str, dict[str, object]]: async def run() -> list[ReadableSpan]: with capture_tool_spans() as spans: @@ -159,11 +188,14 @@ def _expected_keys(mode: str, value: object) -> set[str]: return keys | {"count"} if isinstance(value, list) else keys -def test_records_match_fixture_and_exclude_raw_values( +def test_records_match_plan_and_exclude_raw_values( traced_app: FastMCP, caplog: pytest.LogCaptureFixture ) -> None: caplog.set_level(logging.DEBUG) - plan = json.loads(FIXTURE.read_text()) + plan = { + name: entry["args"] + for name, entry in asyncio.run(trace_plan(traced_app)).items() + } attributes = _call_all(traced_app) for tool, arguments in CALLS.items(): attrs = attributes[tool] diff --git a/uv.lock b/uv.lock index 27b1dbfbe..9f9767ae9 100644 --- a/uv.lock +++ b/uv.lock @@ -248,7 +248,7 @@ requires-dist = [ { name = "duckdb", specifier = "==1.4.3" }, { name = "duckdb-engine", specifier = "==0.17.0" }, { name = "fastmcp", specifier = ">=4.0.9,<5.0" }, - { name = "fastmcp-extensions", specifier = ">=0.35.0,<1.0.0" }, + { name = "fastmcp-extensions", git = "https://github.com/airbytehq/fastmcp-extensions.git?rev=devin%2F1792462457-arg-trace-rules" }, { name = "google-auth", specifier = ">=2.27.0,<3.0" }, { name = "google-cloud-bigquery", specifier = ">=3.12.0,<4.0" }, { name = "google-cloud-bigquery-storage", specifier = ">=2.25.0,<3.0" }, @@ -1352,8 +1352,8 @@ wheels = [ [[package]] name = "fastmcp-extensions" -version = "0.35.0" -source = { registry = "https://pypi.org/simple" } +version = "0.35.0.post2.dev0+c69b2ed" +source = { git = "https://github.com/airbytehq/fastmcp-extensions.git?rev=devin%2F1792462457-arg-trace-rules#c69b2eddfd59afbf78d1f6961ac33335777557b5" } dependencies = [ { name = "fastmcp" }, { name = "httpx" }, @@ -1367,10 +1367,6 @@ dependencies = [ { name = "starlette" }, { name = "uvicorn" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/74/45/dcf55ab5ba27c238654c71d92af320eec2312ae3fad62572b78ea0b9f8fc/fastmcp_extensions-0.35.0.tar.gz", hash = "sha256:bad04d428433c424407277928bc44d77e356f9c2ff73b42aa9e058ad54173038", size = 357967, upload-time = "2026-10-14T05:17:25.382Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/5f/3836a5aea53632be7a7629b7514003138377cfc5e37958308839cb627d7e/fastmcp_extensions-0.35.0-py3-none-any.whl", hash = "sha256:86d31337876faec30dcc336c68b7103d84f553ce8f1555540888c717ddad9e8f", size = 156435, upload-time = "2026-10-14T05:17:23.977Z" }, -] [[package]] name = "fastmcp-slim"