diff --git a/airbyte/constants.py b/airbyte/constants.py index e69198c31..ab4788995 100644 --- a/airbyte/constants.py +++ b/airbyte/constants.py @@ -361,25 +361,27 @@ def is_hosted_mcp_mode() -> bool: """MCP tool modules that are hidden unless insiders mode is enabled. Enable them with `AIRBYTE_MCP_INSIDERS` / `X-MCP-Insiders`, or by naming the module in -the include list. +the include list. Individual tools can also opt in with `mcp_tool(insiders_only=True)`; those +are shown only when insiders mode is on, not by an include list. """ MCP_INSIDERS_ENV_VAR: str = "AIRBYTE_MCP_INSIDERS" """Environment variable that advertises insiders MCP tools. Off by default. -Set to `1`/`true`/`yes` to advertise the tools in `MCP_INSIDERS_MODULES` to every -caller, or to `0`/`false`/`no` to hide them from every caller. Either value overrides -`MCP_INSIDERS_HEADER`; any other value, including an empty string, leaves the decision -to that header. +Set to `0`/`false`/`no` to hide insiders tools (see `MCP_INSIDERS_MODULES`) from every +caller, regardless of `MCP_INSIDERS_HEADER`. Set to `1`/`true`/`yes` to advertise them by +default; a caller can still opt out with `MCP_INSIDERS_HEADER: 0`. Any other value, +including an empty string, leaves the decision to that header. """ MCP_INSIDERS_HEADER: str = "X-MCP-Insiders" """HTTP header key that advertises insiders MCP tools, per request. -Set to `1`/`true`/`yes` to add the tools in `MCP_INSIDERS_MODULES` to the advertised +Set to `1`/`true`/`yes` to add insiders tools (see `MCP_INSIDERS_MODULES`) to the advertised tool surface. This selects which tools are advertised and is not an access-control boundary: every insiders tool authorizes each call against the Airbyte API. -`MCP_INSIDERS_ENV_VAR` overrides this header when explicitly set. +A falsy `MCP_INSIDERS_ENV_VAR` overrides this header; a truthy one makes insiders tools +the default, which this header can only narrow by setting it to `0`. """ # MCP Config Arg Names (used with get_mcp_config) diff --git a/airbyte/mcp/_tool_utils.py b/airbyte/mcp/_tool_utils.py index 2473e0e92..f4d3ae4cd 100644 --- a/airbyte/mcp/_tool_utils.py +++ b/airbyte/mcp/_tool_utils.py @@ -427,6 +427,7 @@ def _parse_csv_config(value: str) -> list[str]: class ToolPolicy: pipeline_change: bool external_access: bool = False + insiders_only: bool = False _TOOL_POLICIES: dict[str, ToolPolicy] = {} @@ -457,11 +458,14 @@ def mcp_tool( # noqa: PLR0913 - mirrors the upstream decorator's kwargs pipeline_change: bool | None = None, external_access: bool = False, tracing: bool | Callable[[Mapping[str, object]], Mapping[str, object]] = True, + insiders_only: bool = False, force_disabled: bool = False, ) -> Callable[[_MCP_TOOL_FUNC], _MCP_TOOL_FUNC]: """Decorate an MCP tool with deferred Airbyte registration metadata. - When `force_disabled` is true, the function is returned unchanged and never registered. + When `insiders_only` is true, the tool is hidden unless insiders mode is on + (`AIRBYTE_MCP_INSIDERS` / `X-MCP-Insiders`). When `force_disabled` is true, the function + is returned unchanged and never registered. """ base_decorator = _mcp_tool( read_only=read_only, @@ -492,6 +496,7 @@ def decorator(func: _MCP_TOOL_FUNC) -> _MCP_TOOL_FUNC: _TOOL_POLICIES[decorated.__name__] = ToolPolicy( pipeline_change=not read_only if pipeline_change is None else pipeline_change, external_access=external_access, + insiders_only=insiders_only, ) return decorated @@ -574,6 +579,11 @@ def airbyte_destructive_tools_filter(tool: Tool, app: FastMCP) -> bool: ) +def airbyte_insiders_tools_filter(tool: Tool, app: FastMCP) -> bool: + """Hide insiders-flagged tools unless insiders mode is on.""" + return not get_tool_policy(tool).insiders_only or _insiders_mode(app) is True + + def _resolve_policy( app_or_ctx: FastMCP | Context, config_name: str, diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index e9bfb4ffd..8b6ca9753 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -108,9 +108,6 @@ from airbyte.cloud.sync_results import SyncResult -SEARCH_TOOLS_DISABLED = True -"""Hide the search tools until the search backend is enabled.""" - DELETE_NAME_GUARD_TIP_TEXT = ( 'IMPORTANT: This operation requires the resource name to contain "delete-me" or ' '"deleteme" (case insensitive). Otherwise, the deletion is rejected. Do not rename the ' @@ -1878,7 +1875,7 @@ def execute_external_sql_query( open_world=True, external_access=True, tracing=lambda args: agent_action_attributes("execute_external_search_query", args), - force_disabled=SEARCH_TOOLS_DISABLED, + insiders_only=True, ) def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the connector API. ctx: Context, @@ -2007,7 +2004,7 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the read_only=True, idempotent=True, external_access=True, - force_disabled=SEARCH_TOOLS_DISABLED, + insiders_only=True, ) def get_cloud_search_status( ctx: Context, diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index 238a037df..d3a236891 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -109,6 +109,7 @@ WORKSPACE_ID_CONFIG_ARG, airbyte_destructive_tools_filter, airbyte_external_access_filter, + airbyte_insiders_tools_filter, airbyte_module_filter, airbyte_readonly_mode_filter, validate_airbyte_domains, @@ -559,6 +560,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 airbyte_external_access_filter, airbyte_destructive_tools_filter, airbyte_module_filter, + airbyte_insiders_tools_filter, ], auth=_create_auth(), lifespan=_mcp_mode_lifespan, diff --git a/scripts/generate_mcp_markdown.py b/scripts/generate_mcp_markdown.py index 5979d4489..dd5d882e4 100755 --- a/scripts/generate_mcp_markdown.py +++ b/scripts/generate_mcp_markdown.py @@ -100,6 +100,7 @@ MCP_INSIDERS_MODULES, MCP_TRUSTED_EXECUTION_ENV_VAR, ) +from airbyte.mcp._tool_utils import _TOOL_POLICIES # noqa: PLC2701 DEFAULT_OUTPUT = Path("docs/mcp-generated") @@ -240,10 +241,11 @@ def _annotate_modules( ) -> None: """Preserve module/capability metadata in the inspect report. - `mcp_module` and `requiresClientFilesystem` never appear on the wire in - `mcp` 2, so the docs pipeline synthesizes them here from the internal - `ToolTraits` registry (`get_tool_traits`) and the `_REGISTERED_*` - fallback map — fields on the *report*, not on the wire tool. + `mcp_module`, `requiresClientFilesystem` and `insidersOnly` never appear on + the wire in `mcp` 2, so the docs pipeline synthesizes them here from the + internal `ToolTraits` registry (`get_tool_traits`), the `_TOOL_POLICIES` + registry, and the `_REGISTERED_*` fallback map — fields on the *report*, + not on the wire tool. """ for section in ("tools", "prompts", "resources", "templates"): for item in report.get(section) or []: @@ -257,6 +259,12 @@ def _annotate_modules( **(item.get("annotations") or {}), "requiresClientFilesystem": True, } + policy = _TOOL_POLICIES.get(str(key)) + if policy and policy.insiders_only: + item["annotations"] = { + **(item.get("annotations") or {}), + "insidersOnly": True, + } def _merge_reports(report_a: dict[str, Any], report_b: dict[str, Any]) -> dict[str, Any]: @@ -426,6 +434,12 @@ def _render_hint_badges( "servers, or name the module in the include-modules setting; " f"`{MCP_INSIDERS_ENV_VAR}=0` disables it regardless)." ) + elif annotations.get("insidersOnly") is True: + lines.append( + "**Availability:** experimental, insiders only " + f"(`{MCP_INSIDERS_ENV_VAR}=1` for stdio, `{MCP_INSIDERS_HEADER}: 1` for hosted " + f"servers; `{MCP_INSIDERS_ENV_VAR}=0` disables it regardless)." + ) if title := annotations.get("title"): lines.append(f"**Title:** {title}") return ("\n\n".join(lines) + "\n\n") if lines else "" diff --git a/tests/unit_tests/fixtures/arg_trace_plan.json b/tests/unit_tests/fixtures/arg_trace_plan.json index 90abfd011..f685e4fca 100644 --- a/tests/unit_tests/fixtures/arg_trace_plan.json +++ b/tests/unit_tests/fixtures/arg_trace_plan.json @@ -78,6 +78,20 @@ "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", @@ -100,6 +114,13 @@ "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", diff --git a/tests/unit_tests/test_mcp_arg_trace_parity.py b/tests/unit_tests/test_mcp_arg_trace_parity.py index 5302276c6..2fb14b483 100644 --- a/tests/unit_tests/test_mcp_arg_trace_parity.py +++ b/tests/unit_tests/test_mcp_arg_trace_parity.py @@ -22,6 +22,7 @@ from opentelemetry.instrumentation.requests import RequestsInstrumentor from opentelemetry.sdk.trace import ReadableSpan +from airbyte.constants import MCP_INSIDERS_ENV_VAR from airbyte.mcp import _otel as observability from airbyte.mcp import cloud from airbyte.mcp.server import app as server_app @@ -49,15 +50,20 @@ "page_size": 10, "cursor": SENTINEL, }, - "list_cloud_connections": {"name_contains": SENTINEL, "limit": 5}, -} -if not cloud.SEARCH_TOOLS_DISABLED: - CALLS["execute_external_search_query"] = { + "execute_external_search_query": { "connector_id": CONNECTOR_ID, "prompt": f"find {SENTINEL}", "streams": [{"name": SENTINEL}], "limit": 3, - } + }, + "list_cloud_connections": {"name_contains": SENTINEL, "limit": 5}, +} + + +@pytest.fixture(autouse=True) +def _insiders_mode(monkeypatch: pytest.MonkeyPatch) -> None: + """Advertise the insiders-only search tools so their args are covered.""" + monkeypatch.setenv(MCP_INSIDERS_ENV_VAR, "1") def _plan() -> dict[str, dict[str, str]]: @@ -168,10 +174,6 @@ def test_records_match_fixture_and_exclude_raw_values( assert SENTINEL not in caplog.text -@pytest.mark.skipif( - cloud.SEARCH_TOOLS_DISABLED, - reason="`execute_external_search_query` is hidden until the search backend is enabled.", -) def test_prompt_records_carry_similarity_and_sql_stays_equality_only( traced_app: FastMCP, ) -> None: diff --git a/tests/unit_tests/test_mcp_cloud.py b/tests/unit_tests/test_mcp_cloud.py index 62e425a22..a44f8f845 100644 --- a/tests/unit_tests/test_mcp_cloud.py +++ b/tests/unit_tests/test_mcp_cloud.py @@ -36,6 +36,7 @@ ConnectionStatus, JobStatusEnum, ) +from airbyte.constants import MCP_INSIDERS_ENV_VAR from airbyte.mcp import cloud as cloud_mcp from airbyte.mcp import guidance as guidance_mcp from airbyte.mcp._arg_resolvers import resolve_list_of_dicts @@ -2154,13 +2155,29 @@ def test_search_tools_use_connector_type_without_probe( assert [kind for kind, _ in connector.calls] == ["search", "search_status"] -def test_search_tools_are_advertised() -> None: - """The search tools register unless `SEARCH_TOOLS_DISABLED` hides them.""" +@pytest.mark.parametrize( + ("insiders_env", "expected"), + [ + pytest.param(None, False, id="hidden_by_default"), + pytest.param("1", True, id="shown_with_insiders_env"), + pytest.param("0", False, id="hidden_with_insiders_env_off"), + ], +) +def test_search_tools_are_insiders_only( + monkeypatch: pytest.MonkeyPatch, + insiders_env: str | None, + expected: bool, +) -> None: + """The search tools are advertised only in insiders mode.""" from airbyte.mcp import server + if insiders_env is None: + monkeypatch.delenv(MCP_INSIDERS_ENV_VAR, raising=False) + else: + monkeypatch.setenv(MCP_INSIDERS_ENV_VAR, insiders_env) names = {tool.name for tool in asyncio.run(server.app.list_tools())} for tool_name in ("execute_external_search_query", "get_cloud_search_status"): - assert (tool_name in names) is not cloud_mcp.SEARCH_TOOLS_DISABLED + assert (tool_name in names) is expected @pytest.mark.parametrize( diff --git a/tests/unit_tests/test_mcp_tool_annotations.py b/tests/unit_tests/test_mcp_tool_annotations.py index aee3c206a..107d62e3e 100644 --- a/tests/unit_tests/test_mcp_tool_annotations.py +++ b/tests/unit_tests/test_mcp_tool_annotations.py @@ -11,6 +11,7 @@ OPEN_WORLD_TOOL_NAMES = { "execute_external_api_query", "execute_external_sql_query", + "execute_external_search_query", } diff --git a/tests/unit_tests/test_mcp_tool_filters.py b/tests/unit_tests/test_mcp_tool_filters.py index 318953b9b..9dd7ae66f 100644 --- a/tests/unit_tests/test_mcp_tool_filters.py +++ b/tests/unit_tests/test_mcp_tool_filters.py @@ -170,6 +170,42 @@ def test_insiders_gate_is_empty() -> None: assert config_arg.env_var == MCP_INSIDERS_ENV_VAR +@pytest.mark.parametrize( + ("env_value", "caller_value", "expected"), + [ + pytest.param(None, None, False, id="hidden_by_default"), + pytest.param(None, "1", True, id="caller_header_enables"), + pytest.param(None, "0", False, id="caller_header_disables"), + pytest.param("1", None, True, id="env_enables"), + pytest.param("1", "0", False, id="caller_narrows_env"), + pytest.param("0", "1", False, id="env_off_wins"), + ], +) +def test_insiders_tool_visibility( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], + env_value: str | None, + caller_value: str | None, + expected: bool, +) -> None: + """Tools registered with `insiders_only=True` follow the insiders env var and header.""" + if env_value is not None: + monkeypatch.setenv(MCP_INSIDERS_ENV_VAR, env_value) + if caller_value is not None: + mcp_config[MCP_CONFIG_INSIDERS] = caller_value + _tool_utils._TOOL_POLICIES["insiders_tool"] = _tool_utils.ToolPolicy( + pipeline_change=False, + insiders_only=True, + ) + _tool_utils._TOOL_POLICIES["regular_tool"] = _tool_utils.ToolPolicy( + pipeline_change=False, + ) + + insiders_filter = _tool_utils.airbyte_insiders_tools_filter + assert insiders_filter(_tool("insiders_tool"), APP) is expected + assert insiders_filter(_tool("regular_tool"), APP) is True + + def test_sync_tools_stay_visible_when_pipeline_changes_are_disabled( monkeypatch: pytest.MonkeyPatch, mcp_config: dict[str, str], diff --git a/tests/unit_tests/test_mcp_tool_utils.py b/tests/unit_tests/test_mcp_tool_utils.py index 78221219f..4b145e98c 100644 --- a/tests/unit_tests/test_mcp_tool_utils.py +++ b/tests/unit_tests/test_mcp_tool_utils.py @@ -261,11 +261,23 @@ def policy_explicit_override() -> None: ) +def test_mcp_tool_force_disabled_skips_registration() -> None: + def force_disabled_tool() -> None: + pass + + assert mcp_tool(read_only=True, force_disabled=True)(force_disabled_tool) is ( + force_disabled_tool + ) + assert force_disabled_tool.__name__ not in _TOOL_POLICIES + + def test_server_tools_register_expected_policies() -> None: assert server_module.app external_tools = ( "execute_external_api_query", "execute_external_sql_query", + "execute_external_search_query", + "get_cloud_search_status", "get_agent_skill_docs", ) for tool_name in external_tools: @@ -274,6 +286,14 @@ def test_server_tools_register_expected_policies() -> None: for tool_name in ("run_cloud_sync", "cancel_cloud_sync"): assert _TOOL_POLICIES[tool_name].pipeline_change is False + insiders_tools = { + name for name, policy in _TOOL_POLICIES.items() if policy.insiders_only + } + assert insiders_tools == { + "execute_external_search_query", + "get_cloud_search_status", + } + @pytest.mark.parametrize( ("config_arg", "expected_env", "expected_header", "expected_name"),