diff --git a/airbyte/constants.py b/airbyte/constants.py index f464827c1..575b73dc4 100644 --- a/airbyte/constants.py +++ b/airbyte/constants.py @@ -298,10 +298,16 @@ def is_hosted_mcp_mode() -> bool: MCP_READONLY_MODE_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_READONLY_MODE" -"""Environment variable to enable read-only mode for the MCP server. +"""Legacy environment variable that disables pipeline-changing MCP tools when truthy.""" -When set to "1" or "true", only tools with readOnlyHint=True will be available. -""" +MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES" +"""Environment variable controlling access to Cloud pipeline-changing tools.""" + +MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS" +"""Environment variable controlling access to external-access Cloud tools.""" + +CLOUD_MCP_SAFE_MODE_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_SAFE_MODE" +"""Environment variable controlling safe mode for destructive Cloud operations.""" MCP_DOMAINS_DISABLED_ENV_VAR: str = "AIRBYTE_MCP_DOMAINS_DISABLED" """Environment variable to disable specific MCP tool domains. @@ -371,6 +377,12 @@ def is_hosted_mcp_mode() -> bool: MCP_CONFIG_READONLY_MODE: str = "airbyte_readonly_mode" """Config arg name for the legacy AIRBYTE_CLOUD_MCP_READONLY_MODE setting.""" +MCP_CONFIG_ALLOW_PIPELINE_CHANGES: str = "allow_pipeline_changes" +"""Config arg name for the Cloud pipeline-change permission.""" + +MCP_CONFIG_ALLOW_EXTERNAL_ACCESS: str = "allow_external_access" +"""Config arg name for external-access permission.""" + MCP_CONFIG_EXCLUDE_MODULES: str = "airbyte_exclude_modules" """Config arg name for the legacy AIRBYTE_MCP_DOMAINS_DISABLED setting.""" @@ -409,6 +421,12 @@ def is_hosted_mcp_mode() -> bool: MCP_EXTENSIONS_HEADER: str = "X-MCP-Extensions" """HTTP header key for client-declared MCP extension IDs.""" +MCP_ALLOW_PIPELINE_CHANGES_HEADER: str = "X-MCP-Allow-Pipeline-Changes" +"""HTTP header key for narrowing Cloud pipeline-change permission.""" + +MCP_ALLOW_EXTERNAL_ACCESS_HEADER: str = "X-MCP-Allow-External-Access" +"""HTTP header key for narrowing external-access permission.""" + # Security Note: The API root and Config API root are intentionally NOT exposed as HTTP # headers. Each hosted MCP deployment is paired to a single backend, so allowing # a caller to override these URLs per-request would let them redirect the diff --git a/airbyte/exceptions.py b/airbyte/exceptions.py index 62b843b4b..2a81ea749 100644 --- a/airbyte/exceptions.py +++ b/airbyte/exceptions.py @@ -339,6 +339,29 @@ class AirbyteSafeModeError(AirbyteMCPError): """ +@dataclass +class AirbytePipelineChangesDisabledError(AirbyteMCPError): + """A pipeline-changing operation was blocked by MCP policy.""" + + guidance: str | None = ( + "Set `AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES=1`, unset legacy " + "`AIRBYTE_CLOUD_MCP_READONLY_MODE`, and ensure the " + "`X-MCP-Allow-Pipeline-Changes` request header is not `0`." + ) + + +@dataclass +class AirbyteExternalAccessDisabledError(AirbyteMCPError): + """An external-access operation was blocked by MCP policy.""" + + guidance: str | None = ( + "Set `AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS=1` and ensure the " + "`X-MCP-Allow-External-Access` request header is not `0`. Disabled pipeline changes " + "(including legacy read-only mode) and an explicit `AIRBYTE_CLOUD_MCP_SAFE_MODE=1` " + "also disable external access when this setting is unset." + ) + + # Normalization Errors diff --git a/airbyte/mcp/__init__.py b/airbyte/mcp/__init__.py index 37c0baf9e..4423adda9 100644 --- a/airbyte/mcp/__init__.py +++ b/airbyte/mcp/__init__.py @@ -130,6 +130,10 @@ and are only allowed for objects created in the current session. Set the environment variable `AIRBYTE_CLOUD_MCP_SAFE_MODE=0` to disable safe mode. +When `AIRBYTE_CLOUD_MCP_SAFE_MODE` is explicitly set to a truthy value, external-access tools +(for example `execute_external_api_query`) are also disabled unless +`AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS=1` explicitly allows them. Safe mode's default-on +behavior does not by itself disable external access. ### Airbyte Cloud Read-Only Mode @@ -144,6 +148,26 @@ Set the environment variable `AIRBYTE_CLOUD_MCP_READONLY_MODE=1` to enable read-only mode. +## Airbyte Cloud Pipeline and External Access Permissions + +Pipeline changes and external access can be controlled independently. The +`AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES` environment variable and +`X-MCP-Allow-Pipeline-Changes` request header control pipeline-changing tools. The +`AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS` environment variable and +`X-MCP-Allow-External-Access` request header control external-access tools, such as +`execute_external_api_query`, which can run connector passthrough queries, search, and inspect +external data or skill documentation. + +Each setting accepts `1`, `true`, or `yes` to allow, and `0`, `false`, or `no` to deny. An +explicit environment denial cannot be overridden by a request header. An environment allow can +be narrowed by a request header denial. If the environment setting is unset, the request header +decides. The legacy `AIRBYTE_CLOUD_MCP_READONLY_MODE=1` also disables pipeline changes. + +When external access is unset, it is disabled if pipeline changes are disabled or +`AIRBYTE_CLOUD_MCP_SAFE_MODE=1` is explicitly set; otherwise external-access tools remain +available. Explicitly allowing external access overrides these calculated defaults. Running or +cancelling a sync is not a pipeline change. + ## Authentication for Remote (HTTP) Servers The steps above run the MCP server over **stdio** — the client launches the diff --git a/airbyte/mcp/_error_handling.py b/airbyte/mcp/_error_handling.py index 3e3a5bd93..7de1a1ecc 100644 --- a/airbyte/mcp/_error_handling.py +++ b/airbyte/mcp/_error_handling.py @@ -6,16 +6,22 @@ from airbyte.exceptions import ( AirbyteAgentsUnavailableError, AirbyteConnectorInUseError, + AirbyteExternalAccessDisabledError, AirbyteLibError, AirbyteLibInputError, AirbyteMCPError, AirbyteMissingResourceError, + AirbytePipelineChangesDisabledError, + AirbyteSafeModeError, ) MCP_TOOL_USER_FACING_ERRORS: tuple[type[AirbyteLibError], ...] = ( AirbyteLibInputError, AirbyteMCPError, + AirbyteSafeModeError, + AirbytePipelineChangesDisabledError, + AirbyteExternalAccessDisabledError, AirbyteAgentsUnavailableError, AirbyteConnectorInUseError, AirbyteMissingResourceError, diff --git a/airbyte/mcp/_policy_middleware.py b/airbyte/mcp/_policy_middleware.py new file mode 100644 index 000000000..9bc152d41 --- /dev/null +++ b/airbyte/mcp/_policy_middleware.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Runtime enforcement for MCP tool policies.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from fastmcp.server.middleware import Middleware + +from airbyte.exceptions import ( + AirbyteExternalAccessDisabledError, + AirbytePipelineChangesDisabledError, +) +from airbyte.mcp._tool_utils import ( + external_access_allowed, + get_tool_policy, + pipeline_changes_allowed, +) + + +if TYPE_CHECKING: + from fastmcp.server.middleware import CallNext, MiddlewareContext + from fastmcp.tools import ToolResult + from mcp.types import CallToolRequestParams + + +class PolicyGuardMiddleware(Middleware): + """Enforce pipeline-change and external-access policies before tool execution.""" + + async def on_call_tool( + self, + context: MiddlewareContext[CallToolRequestParams], + call_next: CallNext[CallToolRequestParams, ToolResult], + ) -> ToolResult: + """Reject calls denied by effective pipeline or external-access settings.""" + fastmcp_context = context.fastmcp_context + if fastmcp_context is None: + return await call_next(context) + + tool = await fastmcp_context.fastmcp.get_tool(context.message.name) + if tool is None: + return await call_next(context) + + policy = get_tool_policy(tool) + if policy.pipeline_change and pipeline_changes_allowed(fastmcp_context) is False: + raise AirbytePipelineChangesDisabledError( + message="Pipeline-changing tools are disabled by the MCP server policy." + ) + if policy.external_access and external_access_allowed(fastmcp_context) is False: + raise AirbyteExternalAccessDisabledError( + message="External-access tools are disabled by the MCP server policy." + ) + return await call_next(context) diff --git a/airbyte/mcp/_tool_utils.py b/airbyte/mcp/_tool_utils.py index 3f3512ca5..7684ef427 100644 --- a/airbyte/mcp/_tool_utils.py +++ b/airbyte/mcp/_tool_utils.py @@ -12,8 +12,9 @@ import inspect import os from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, TypeVar, cast from fastmcp.server.dependencies import get_access_token, get_http_headers from fastmcp_extensions import ( @@ -35,6 +36,7 @@ CONFIG_TRUSTED_EXECUTION, get_annotation, ) +from mcp.types import Tool from airbyte.constants import ( CLOUD_API_ROOT_ENV_VAR, @@ -42,9 +44,16 @@ CLOUD_CLIENT_ID_ENV_VAR, CLOUD_CLIENT_SECRET_ENV_VAR, CLOUD_CONFIG_API_ROOT_ENV_VAR, + CLOUD_MCP_SAFE_MODE_ENV_VAR, CLOUD_ORGANIZATION_ID_ENV_VAR, CLOUD_WORKSPACE_ID_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_HEADER, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_ALLOW_PIPELINE_CHANGES_HEADER, MCP_BEARER_TOKEN_HEADER, + MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, + MCP_CONFIG_ALLOW_PIPELINE_CHANGES, MCP_CONFIG_API_URL, MCP_CONFIG_BEARER_TOKEN, MCP_CONFIG_CLIENT_ID, @@ -73,7 +82,8 @@ if TYPE_CHECKING: from fastmcp import FastMCP from fastmcp.apps import AppConfig - from mcp.types import Tool + from fastmcp.server.context import Context + from fastmcp.tools import Tool as FastMCPTool _MCP_TOOL_FUNC = TypeVar("_MCP_TOOL_FUNC", bound=Callable[..., object]) @@ -172,6 +182,24 @@ def check_guid_created_in_session(guid: str) -> None: the include-list opt-in. """ +ALLOW_PIPELINE_CHANGES_CONFIG_ARG = MCPServerConfigArg( + name=MCP_CONFIG_ALLOW_PIPELINE_CHANGES, + http_header_key=MCP_ALLOW_PIPELINE_CHANGES_HEADER, + env_var=MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + default="", + required=False, +) +"""Config arg for permission to run pipeline-changing Cloud tools.""" + +ALLOW_EXTERNAL_ACCESS_CONFIG_ARG = MCPServerConfigArg( + name=MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, + http_header_key=MCP_ALLOW_EXTERNAL_ACCESS_HEADER, + env_var=MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + default="", + required=False, +) +"""Config arg for permission to run external-access Cloud tools.""" + TRUSTED_EXECUTION_CONFIG_ARG = MCPServerConfigArg( name=CONFIG_TRUSTED_EXECUTION, env_var=MCP_TRUSTED_EXECUTION_ENV_VAR, @@ -343,6 +371,24 @@ def _parse_csv_config(value: str) -> list[str]: return [item.strip() for item in value.split(",") if item.strip()] +@dataclass(frozen=True) +class ToolPolicy: + pipeline_change: bool + external_access: bool = False + + +_TOOL_POLICIES: dict[str, ToolPolicy] = {} + + +def get_tool_policy(tool: Tool | FastMCPTool) -> ToolPolicy: + """Return registered policy metadata, falling back to the tool's read-only hint.""" + policy = _TOOL_POLICIES.get(tool.name) + if policy is not None: + return policy + read_only = get_annotation(cast(Tool, tool), ANNOTATION_READ_ONLY_HINT, default=False) + return ToolPolicy(pipeline_change=not bool(read_only)) + + def mcp_tool( # noqa: PLR0913 - mirrors the upstream decorator's kwargs *, read_only: bool = False, @@ -356,6 +402,8 @@ def mcp_tool( # noqa: PLR0913 - mirrors the upstream decorator's kwargs app: AppConfig | None = None, required_capabilities: Iterable[Capability | str] | None = None, extra_help_text: str | None = None, + pipeline_change: bool | None = None, + external_access: bool = False, ) -> Callable[[_MCP_TOOL_FUNC], _MCP_TOOL_FUNC]: """Decorate an MCP tool with deferred Airbyte registration metadata.""" base_decorator = _mcp_tool( @@ -381,6 +429,10 @@ def decorator(func: _MCP_TOOL_FUNC) -> _MCP_TOOL_FUNC: if registered_func is not decorated: raise RuntimeError("Unexpected MCP tool registration state.") registered_annotations[ANNOTATION_MCP_MODULE] = _mcp_module_for_tool(decorated) + _TOOL_POLICIES[decorated.__name__] = ToolPolicy( + pipeline_change=not read_only if pipeline_change is None else pipeline_change, + external_access=external_access, + ) return decorated return decorator @@ -420,14 +472,58 @@ def _normalize_mcp_module(mcp_module: str) -> str: def airbyte_readonly_mode_filter(tool: Tool, app: FastMCP) -> bool: - """Filter tools based on legacy AIRBYTE_CLOUD_MCP_READONLY_MODE env var. + """Hide pipeline-changing tools when pipeline changes are disabled.""" + return not (pipeline_changes_allowed(app) is False and get_tool_policy(tool).pipeline_change) - When set to "1", only show tools with readOnlyHint=True. + +def airbyte_external_access_filter(tool: Tool, app: FastMCP) -> bool: + """Hide external-access tools when external access is disabled.""" + return not (get_tool_policy(tool).external_access and external_access_allowed(app) is False) + + +def _resolve_policy( + app_or_ctx: FastMCP | Context, + config_name: str, + env_var: str, +) -> bool | None: + environment_value = _str_to_bool(os.environ.get(env_var)) + if environment_value is False: + return False + request_value = _str_to_bool(get_mcp_config(app_or_ctx, config_name)) + if environment_value is True: + return request_value is not False + return request_value + + +def pipeline_changes_allowed(app_or_ctx: FastMCP | Context) -> bool | None: + if _str_to_bool(os.environ.get(MCP_READONLY_MODE_ENV_VAR)) is True: + return False + return _resolve_policy( + app_or_ctx, + MCP_CONFIG_ALLOW_PIPELINE_CHANGES, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + ) + + +def external_access_allowed(app_or_ctx: FastMCP | Context) -> bool | None: + """Resolve explicit permission first, then derive the external-access default. + + Disabled pipeline changes and explicitly enabled safe mode disable external access + when no explicit permission is configured; otherwise the result remains unset. """ - config_value = (get_mcp_config(app, MCP_CONFIG_READONLY_MODE) or "").lower() - if config_value in {"1", "true"}: - return bool(get_annotation(tool, ANNOTATION_READ_ONLY_HINT, default=False)) - return True + explicit_value = _resolve_policy( + app_or_ctx, + MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + ) + if explicit_value is not None: + return explicit_value + explicitly_enabled_safe_mode = _str_to_bool(os.environ.get(CLOUD_MCP_SAFE_MODE_ENV_VAR)) is True + return ( + False + if pipeline_changes_allowed(app_or_ctx) is False or explicitly_enabled_safe_mode + else None + ) def _insiders_mode(app: FastMCP) -> bool | None: diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 999e14d39..d47ba4efc 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -854,6 +854,7 @@ def create_connection_on_cloud( @mcp_tool( open_world=True, + pipeline_change=False, ) def run_cloud_sync( ctx: Context, @@ -1109,6 +1110,7 @@ def list_cloud_sync_jobs( @mcp_tool( destructive=True, open_world=True, + pipeline_change=False, ) def cancel_cloud_sync( ctx: Context, @@ -1514,6 +1516,7 @@ def describe_cloud_connector( idempotent=True, open_world=True, extra_help_text=SKILL_DOCS_SECTION_HINT, + external_access=True, ) def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the connector API. ctx: Context, @@ -1618,8 +1621,9 @@ def execute_external_api_query( # noqa: PLR0913 # Explicit args mirror the con # Not yet registered as an MCP tool: write actions are not supported by the backend. -# Restore the `@mcp_tool(open_world=True, extra_help_text=...)` decorator and revert the -# `_` prefix when these are live. Tracked in https://linear.app/airbyteio/issue/AGENTIC-2280 +# Restore the `@mcp_tool(open_world=True, external_access=True, +# extra_help_text=...)` decorator and revert the `_` prefix when these are live. +# Tracked in https://linear.app/airbyteio/issue/AGENTIC-2280 def _execute_external_api_action( # noqa: PLR0913 # Explicit args mirror the connector API. ctx: Context, *, @@ -1711,6 +1715,7 @@ def _execute_external_api_action( # noqa: PLR0913 # Explicit args mirror the c read_only=True, idempotent=True, open_world=True, + external_access=True, ) def execute_external_sql_query( ctx: Context, @@ -1791,6 +1796,7 @@ def execute_external_sql_query( read_only=True, idempotent=True, open_world=True, + external_access=True, ) def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the connector API. ctx: Context, @@ -1915,6 +1921,7 @@ def execute_external_search_query( # noqa: PLR0913 # Explicit args mirror the read_only=True, idempotent=True, open_world=True, + external_access=True, ) def get_cloud_search_status( ctx: Context, diff --git a/airbyte/mcp/guidance.py b/airbyte/mcp/guidance.py index 85e5a6b3d..96e4cd04a 100644 --- a/airbyte/mcp/guidance.py +++ b/airbyte/mcp/guidance.py @@ -174,6 +174,7 @@ def test_my_tools_prompt( idempotent=True, open_world=True, extra_help_text=SKILL_DOCS_SECTION_HINT, + external_access=True, ) def get_agent_skill_docs( ctx: Context, diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index b68bfac0c..4d400c3b4 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -84,6 +84,7 @@ MCP_TOOL_USER_FACING_ERRORS, format_user_facing_error, ) +from airbyte.mcp._policy_middleware import PolicyGuardMiddleware from airbyte.mcp._scope import CallScopeMiddleware, call_scope_properties from airbyte.mcp._sso_auth import SsoRealmConfig, make_sso_proxy_factory from airbyte.mcp._telemetry import ServerConnectedTelemetryMiddleware, request_properties @@ -91,6 +92,8 @@ AIRBYTE_EXCLUDE_MODULES_CONFIG_ARG, AIRBYTE_INCLUDE_MODULES_CONFIG_ARG, AIRBYTE_READONLY_MODE_CONFIG_ARG, + ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, + ALLOW_PIPELINE_CHANGES_CONFIG_ARG, API_URL_CONFIG_ARG, BEARER_TOKEN_CONFIG_ARG, CLIENT_ID_CONFIG_ARG, @@ -100,6 +103,7 @@ ORGANIZATION_ID_CONFIG_ARG, TRUSTED_EXECUTION_CONFIG_ARG, WORKSPACE_ID_CONFIG_ARG, + airbyte_external_access_filter, airbyte_module_filter, airbyte_readonly_mode_filter, validate_airbyte_domains, @@ -526,6 +530,8 @@ async def _mcp_mode_lifespan( # noqa: RUF029 include_standard_tool_filters=True, server_config_args=[ AIRBYTE_READONLY_MODE_CONFIG_ARG, + ALLOW_PIPELINE_CHANGES_CONFIG_ARG, + ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, AIRBYTE_EXCLUDE_MODULES_CONFIG_ARG, AIRBYTE_INCLUDE_MODULES_CONFIG_ARG, INSIDERS_CONFIG_ARG, @@ -545,6 +551,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 }, tool_filters=[ airbyte_readonly_mode_filter, + airbyte_external_access_filter, airbyte_module_filter, ], auth=_create_auth(), @@ -565,6 +572,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 """The Airbyte MCP Server application instance.""" app.add_middleware(ServerConnectedTelemetryMiddleware(lifecycle_telemetry_sinks)) +app.add_middleware(PolicyGuardMiddleware()) app.middleware.insert(0, CallScopeMiddleware()) app.middleware.insert(0, AirbyteUserMiddleware()) diff --git a/tests/unit_tests/test_mcp_policy_middleware.py b/tests/unit_tests/test_mcp_policy_middleware.py new file mode 100644 index 000000000..16feeef16 --- /dev/null +++ b/tests/unit_tests/test_mcp_policy_middleware.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +"""Unit tests for runtime MCP policy enforcement.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError +from fastmcp_extensions import mcp_server + +from airbyte.constants import ( + CLOUD_MCP_SAFE_MODE_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_READONLY_MODE_ENV_VAR, +) +from airbyte.exceptions import ( + AirbyteExternalAccessDisabledError, + AirbytePipelineChangesDisabledError, +) +from airbyte.mcp import _tool_utils +from airbyte.mcp._error_handling import ( + MCP_TOOL_USER_FACING_ERRORS, + format_user_facing_error, +) +from airbyte.mcp._policy_middleware import PolicyGuardMiddleware +from airbyte.mcp._tool_utils import ToolPolicy + + +class _FakeFastMCP: + def __init__(self, tool: Any) -> None: + self.tool = tool + + async def get_tool(self, name: str) -> Any: + assert name == self.tool.name + return self.tool + + +@pytest.fixture(autouse=True) +def clear_policy_environment(monkeypatch: pytest.MonkeyPatch) -> None: + for env_var in ( + CLOUD_MCP_SAFE_MODE_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_READONLY_MODE_ENV_VAR, + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(_tool_utils, "_TOOL_POLICIES", {}) + + +@pytest.mark.parametrize( + ( + "tool_name", + "policy", + "read_only_hint", + "environment", + "request_config", + "expected_error", + ), + [ + pytest.param( + "pipeline_write", + ToolPolicy(pipeline_change=True), + False, + {MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: "0"}, + {}, + AirbytePipelineChangesDisabledError, + id="pipeline-change-denied", + ), + pytest.param( + "provider_read_only", + None, + True, + {MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: "0"}, + {}, + None, + id="read-only-provider-pass-through", + ), + pytest.param( + "sync_control", + ToolPolicy(pipeline_change=False), + False, + {MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: "0"}, + {}, + None, + id="non-pipeline-change-tool-pass-through", + ), + pytest.param( + "pipeline_write", + ToolPolicy(pipeline_change=True), + False, + {MCP_READONLY_MODE_ENV_VAR: "1"}, + {}, + AirbytePipelineChangesDisabledError, + id="legacy-read-only-denial", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR: "0"}, + {}, + AirbyteExternalAccessDisabledError, + id="explicit-external-access-denial", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {CLOUD_MCP_SAFE_MODE_ENV_VAR: "1"}, + {}, + AirbyteExternalAccessDisabledError, + id="explicit-safe-mode-default-denial", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + { + CLOUD_MCP_SAFE_MODE_ENV_VAR: "1", + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR: "1", + }, + {}, + None, + id="explicit-external-access-allow", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {}, + {}, + None, + id="unset-safe-mode-default-pass-through", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {CLOUD_MCP_SAFE_MODE_ENV_VAR: "auto"}, + {}, + None, + id="auto-safe-mode-default-pass-through", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: "0"}, + {}, + AirbyteExternalAccessDisabledError, + id="pipeline-changes-disabled-default-denial", + ), + ], +) +def test_policy_guard_middleware( + monkeypatch: pytest.MonkeyPatch, + tool_name: str, + policy: ToolPolicy | None, + read_only_hint: bool, + environment: dict[str, str], + request_config: dict[str, str], + expected_error: type[Exception] | None, +) -> None: + monkeypatch.setattr( + _tool_utils, + "get_mcp_config", + lambda _app, key, **_kwargs: request_config.get(key), + ) + for env_var, value in environment.items(): + monkeypatch.setenv(env_var, value) + + if policy is not None: + _tool_utils._TOOL_POLICIES[tool_name] = policy + tool = SimpleNamespace( + name=tool_name, + annotations=SimpleNamespace(read_only_hint=read_only_hint), + ) + fastmcp_context = SimpleNamespace(fastmcp=_FakeFastMCP(tool)) + context = SimpleNamespace( + fastmcp_context=fastmcp_context, + message=SimpleNamespace(name=tool_name), + ) + + async def call_next(_context: Any) -> str: + return "called" + + async def invoke() -> str: + return await PolicyGuardMiddleware().on_call_tool(context, call_next) + + if expected_error is None: + assert asyncio.run(invoke()) == "called" + else: + with pytest.raises(expected_error): + asyncio.run(invoke()) + + +def test_blocked_policy_error_is_formatted_for_mcp_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server = mcp_server( + "test", + telemetry=False, + user_facing_errors=MCP_TOOL_USER_FACING_ERRORS, + user_facing_error_formatter=format_user_facing_error, + ) + server.add_middleware(PolicyGuardMiddleware()) + monkeypatch.setenv(MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, "0") + + @server.tool + def blocked_pipeline_tool() -> str: + return "unexpectedly executed" + + _tool_utils._TOOL_POLICIES["blocked_pipeline_tool"] = ToolPolicy( + pipeline_change=True + ) + + async def call_tool() -> None: + async with Client(server) as client: + with pytest.raises(ToolError) as error: + await client.call_tool("blocked_pipeline_tool") + assert "Pipeline-changing tools are disabled" in str(error.value) + assert "AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES=1" in str(error.value) + + asyncio.run(call_tool()) diff --git a/tests/unit_tests/test_mcp_tool_filters.py b/tests/unit_tests/test_mcp_tool_filters.py index b9028eda9..504d9b401 100644 --- a/tests/unit_tests/test_mcp_tool_filters.py +++ b/tests/unit_tests/test_mcp_tool_filters.py @@ -11,12 +11,18 @@ from fastmcp_extensions.tool_filters import CONFIG_INCLUDE_MODULES from airbyte.constants import ( + CLOUD_MCP_SAFE_MODE_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, + MCP_CONFIG_ALLOW_PIPELINE_CHANGES, MCP_CONFIG_EXCLUDE_MODULES, MCP_CONFIG_INCLUDE_MODULES, MCP_CONFIG_INSIDERS, MCP_INSIDERS_ENV_VAR, MCP_INSIDERS_HEADER, MCP_INSIDERS_MODULES, + MCP_READONLY_MODE_ENV_VAR, _str_to_bool, ) from airbyte.mcp import _tool_utils @@ -37,20 +43,28 @@ def _tool(name: str) -> Tool: def mcp_config(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: """Patch `get_mcp_config` so tests can set MCP config values directly.""" monkeypatch.delenv(MCP_INSIDERS_ENV_VAR, raising=False) + for env_var in ( + CLOUD_MCP_SAFE_MODE_ENV_VAR, + MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, + MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_READONLY_MODE_ENV_VAR, + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(_tool_utils, "_TOOL_POLICIES", {}) config: dict[str, str] = {} monkeypatch.setattr( _tool_utils, "get_mcp_config", - lambda app, key, **kwargs: config.get(key), # noqa: ARG005 + lambda _app, key, **_kwargs: config.get(key), ) # `mcp_module` is internal state (never on the wire), so stub the traits # registry lookup with a tool-name -> module map. monkeypatch.setattr( _tool_utils, "get_tool_traits", - lambda app, name: ToolTraits( + lambda _app, name: ToolTraits( mcp_module=name if name != "unannotated" else None - ), # noqa: ARG005 + ), ) return config @@ -149,3 +163,139 @@ def test_insiders_gate_is_empty() -> None: assert not config_arg.required assert config_arg.http_header_key == MCP_INSIDERS_HEADER assert config_arg.env_var == MCP_INSIDERS_ENV_VAR + + +def test_sync_tools_stay_visible_when_pipeline_changes_are_disabled( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], +) -> None: + assert not mcp_config + monkeypatch.setenv(MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, "0") + for tool_name in ("run_cloud_sync", "cancel_cloud_sync"): + tool = _policy_tool(tool_name, pipeline_change=False) + assert _tool_utils.airbyte_readonly_mode_filter(tool, APP) + + +def _policy_tool( + name: str, + *, + pipeline_change: bool | None, + external_access: bool = False, + read_only_hint: bool = False, +) -> Tool: + if pipeline_change is not None or external_access: + _tool_utils._TOOL_POLICIES[name] = _tool_utils.ToolPolicy( + pipeline_change=bool(pipeline_change), + external_access=external_access, + ) + return cast( + Tool, + SimpleNamespace( + name=name, + annotations=SimpleNamespace(read_only_hint=read_only_hint), + ), + ) + + +@pytest.mark.parametrize( + ( + "pipeline_env", + "header", + "legacy_readonly", + "policy", + "read_only_hint", + "visible", + ), + [ + pytest.param("0", None, None, True, False, False, id="env-deny"), + pytest.param("1", "0", None, True, False, False, id="header-narrows-env-allow"), + pytest.param( + "0", "1", None, True, False, False, id="env-deny-wins-over-header" + ), + pytest.param(None, "0", None, True, False, False, id="header-deny-without-env"), + pytest.param("1", None, None, True, False, True, id="env-allow"), + pytest.param("1", None, "1", True, False, False, id="legacy-readonly-deny"), + pytest.param("0", None, None, False, False, True, id="pipeline-change-false"), + pytest.param("0", None, None, None, True, True, id="read-only-hint-fallback"), + ], +) +def test_pipeline_change_filter( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], + pipeline_env: str | None, + header: str | None, + legacy_readonly: str | None, + policy: bool | None, + read_only_hint: bool, + visible: bool, +) -> None: + for env_var, value in ( + (MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, pipeline_env), + (MCP_READONLY_MODE_ENV_VAR, legacy_readonly), + ): + if value is None: + monkeypatch.delenv(env_var, raising=False) + else: + monkeypatch.setenv(env_var, value) + if header is not None: + mcp_config[MCP_CONFIG_ALLOW_PIPELINE_CHANGES] = header + tool = _policy_tool( + "pipeline-policy-tool", + pipeline_change=policy, + read_only_hint=read_only_hint, + ) + + assert _tool_utils.airbyte_readonly_mode_filter(tool, APP) is visible + + +@pytest.mark.parametrize( + ( + "external_env", + "header", + "pipeline_env", + "legacy_readonly", + "safe_mode", + "visible", + ), + [ + pytest.param("0", None, None, None, None, False, id="external-env-deny"), + pytest.param( + "1", "0", None, None, None, False, id="header-narrows-external-allow" + ), + pytest.param("0", "1", None, None, None, False, id="external-env-deny-wins"), + pytest.param(None, "0", None, None, None, False, id="external-header-deny"), + pytest.param("1", None, None, None, "1", True, id="explicit-external-allow"), + pytest.param(None, None, None, None, "1", False, id="explicit-safe-mode-deny"), + pytest.param(None, None, "0", None, None, False, id="pipeline-deny-default"), + pytest.param(None, None, None, "1", None, False, id="legacy-readonly-default"), + pytest.param(None, None, None, None, None, True, id="unset-default-allows"), + pytest.param(None, None, None, None, "auto", True, id="auto-safe-mode-allows"), + ], +) +def test_external_access_filter( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], + external_env: str | None, + header: str | None, + pipeline_env: str | None, + legacy_readonly: str | None, + safe_mode: str | None, + visible: bool, +) -> None: + for env_var, value in ( + (MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, external_env), + (MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, pipeline_env), + (MCP_READONLY_MODE_ENV_VAR, legacy_readonly), + (CLOUD_MCP_SAFE_MODE_ENV_VAR, safe_mode), + ): + if value is None: + monkeypatch.delenv(env_var, raising=False) + else: + monkeypatch.setenv(env_var, value) + if header is not None: + mcp_config[MCP_CONFIG_ALLOW_EXTERNAL_ACCESS] = header + tool = _policy_tool( + "external-access-tool", pipeline_change=False, external_access=True + ) + + assert _tool_utils.airbyte_external_access_filter(tool, APP) is visible diff --git a/tests/unit_tests/test_mcp_tool_utils.py b/tests/unit_tests/test_mcp_tool_utils.py index 1b1f76f2a..7e16ade4a 100644 --- a/tests/unit_tests/test_mcp_tool_utils.py +++ b/tests/unit_tests/test_mcp_tool_utils.py @@ -12,15 +12,21 @@ from airbyte.exceptions import AirbyteSafeModeError from airbyte.mcp._tool_utils import ( API_URL_CONFIG_ARG, + ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, + ALLOW_PIPELINE_CHANGES_CONFIG_ARG, CLIENT_ID_CONFIG_ARG, CLIENT_SECRET_CONFIG_ARG, CONFIG_API_URL_CONFIG_ARG, TRUSTED_EXECUTION_CONFIG_ARG, _GUIDS_CREATED_IN_SESSION, + _TOOL_POLICIES, + ToolPolicy, _resolve_transport_bearer_token, check_guid_created_in_session, + mcp_tool, register_guid_created_in_session, ) +from airbyte.mcp import server as server_module @dataclass @@ -166,3 +172,78 @@ def test_config_args_are_not_caller_controllable( tool surface. """ assert config_arg.http_header_key is None + + +def test_mcp_tool_registers_policy_defaults_and_overrides() -> None: + @mcp_tool(read_only=True) + def policy_read_only_default() -> None: + pass + + @mcp_tool() + def policy_write_default() -> None: + pass + + @mcp_tool(read_only=True, pipeline_change=True, external_access=True) + def policy_explicit_override() -> None: + pass + + assert _TOOL_POLICIES[policy_read_only_default.__name__] == ToolPolicy( + pipeline_change=False, + external_access=False, + ) + assert _TOOL_POLICIES[policy_write_default.__name__] == ToolPolicy( + pipeline_change=True, + external_access=False, + ) + assert _TOOL_POLICIES[policy_explicit_override.__name__] == ToolPolicy( + pipeline_change=True, + external_access=True, + ) + + +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: + assert _TOOL_POLICIES[tool_name].external_access + + for tool_name in ("run_cloud_sync", "cancel_cloud_sync"): + assert _TOOL_POLICIES[tool_name].pipeline_change is False + + +@pytest.mark.parametrize( + ("config_arg", "expected_env", "expected_header", "expected_name"), + [ + pytest.param( + ALLOW_PIPELINE_CHANGES_CONFIG_ARG, + "AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES", + "X-MCP-Allow-Pipeline-Changes", + "allow_pipeline_changes", + id="pipeline-changes", + ), + pytest.param( + ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, + "AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS", + "X-MCP-Allow-External-Access", + "allow_external_access", + id="external-access", + ), + ], +) +def test_policy_config_args( + config_arg: MCPServerConfigArg, + expected_env: str, + expected_header: str, + expected_name: str, +) -> None: + assert config_arg.name == expected_name + assert config_arg.env_var == expected_env + assert config_arg.http_header_key == expected_header + assert config_arg.default == "" + assert not config_arg.required