diff --git a/airbyte/_util/unsupported_inputs.py b/airbyte/_util/unsupported_inputs.py new file mode 100644 index 000000000..46ddaae7b --- /dev/null +++ b/airbyte/_util/unsupported_inputs.py @@ -0,0 +1,43 @@ +# Copyright (c) 2024 Airbyte, Inc., all rights reserved. +"""Startup guards for removed PyAirbyte environment variables.""" + +import os + +from airbyte.constants import CLOUD_MCP_SAFE_MODE_ENV_VAR +from airbyte.exceptions import AirbyteLibInputError + + +UNSUPPORTED_ENV_VARS: dict[str, str] = { + CLOUD_MCP_SAFE_MODE_ENV_VAR: ( + "Remove `AIRBYTE_CLOUD_MCP_SAFE_MODE` and consider using one of these " + "alternatives:\n" + "- `AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES=0`: block changes to " + "sources, destinations, and connections.\n" + "- `AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS=0`: block SQL and API " + "passthrough tools.\n" + "- `AIRBYTE_CLOUD_MCP_READONLY_MODE=1`: only expose read-only tools.\n" + "- `AIRBYTE_CLOUD_MCP_ALLOW_DESTRUCTIVE_TOOLS=0`: hide destructive tools.\n" + "- `AIRBYTE_CLOUD_MCP_LOCAL_DEMO_MODE=1`: only allow changes to " + "resources created in this session (local mode only; not recommended " + "except for demos).\n" + ), +} +"""Removed PyAirbyte environment variables, mapped to replacement guidance. + +Checked at startup of the `pyab` CLI and both MCP servers. Startup fails if any +listed name is set. +""" + + +def check_unsupported_env_vars() -> None: + """Fail fast if a removed PyAirbyte env var is set. + + Raises `AirbyteLibInputError` on the first name in `UNSUPPORTED_ENV_VARS` + present in the environment with any value, including an empty string. + """ + for name, guidance in UNSUPPORTED_ENV_VARS.items(): + if name in os.environ: + raise AirbyteLibInputError( + message=f"`{name}` is no longer supported.", + guidance=guidance, + ) diff --git a/airbyte/cli/pyab.py b/airbyte/cli/pyab.py index 98714a717..944e74762 100644 --- a/airbyte/cli/pyab.py +++ b/airbyte/cli/pyab.py @@ -67,9 +67,10 @@ from typing import TYPE_CHECKING, Annotated, Any import yaml -from cyclopts import App, Parameter +from cyclopts import App, ArgumentCollection, Parameter from airbyte._util.destination_smoke_tests import run_destination_smoke_test +from airbyte._util.unsupported_inputs import check_unsupported_env_vars from airbyte.destinations.util import get_destination, get_noop_destination from airbyte.exceptions import AirbyteLibInputError from airbyte.secrets.util import get_secret @@ -77,6 +78,8 @@ if TYPE_CHECKING: + from collections.abc import Sequence + from airbyte.destinations.base import Destination from airbyte.sources.base import Source @@ -359,11 +362,21 @@ def _resolve_destination_job( ) +def _check_unsupported_env_vars_hook( + apps: Sequence[App], # noqa: ARG001 + commands: tuple[str, ...], # noqa: ARG001 + arguments: ArgumentCollection, # noqa: ARG001 +) -> None: + """Fail any CLI command when a removed env var is set.""" + check_unsupported_env_vars() + + cli = App( name="pyab", # pyrefly: ignore[unexpected-keyword] help="PyAirbyte CLI.", # pyrefly: ignore[unexpected-keyword] help_format="plaintext", # pyrefly: ignore[unexpected-keyword] version_flags=[], # pyrefly: ignore[unexpected-keyword] + config=[_check_unsupported_env_vars_hook], # pyrefly: ignore[unexpected-keyword] ) diff --git a/airbyte/constants.py b/airbyte/constants.py index 72fe7c9f5..e69198c31 100644 --- a/airbyte/constants.py +++ b/airbyte/constants.py @@ -306,6 +306,9 @@ def is_hosted_mcp_mode() -> bool: MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS" """Environment variable controlling access to external-access Cloud tools.""" +MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_ALLOW_DESTRUCTIVE_TOOLS" +"""Environment variable controlling access to destructive Cloud tools.""" + MCP_LOCAL_DEMO_MODE_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_LOCAL_DEMO_MODE" """Environment variable enabling opt-in local demo mode for STDIO MCP sessions.""" @@ -390,6 +393,9 @@ def is_hosted_mcp_mode() -> bool: MCP_CONFIG_ALLOW_EXTERNAL_ACCESS: str = "allow_external_access" """Config arg name for external-access permission.""" +MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS: str = "allow_destructive_tools" +"""Config arg name for the Cloud destructive-tools permission.""" + MCP_CONFIG_EXCLUDE_MODULES: str = "airbyte_exclude_modules" """Config arg name for the legacy AIRBYTE_MCP_DOMAINS_DISABLED setting.""" @@ -434,6 +440,9 @@ def is_hosted_mcp_mode() -> bool: MCP_ALLOW_EXTERNAL_ACCESS_HEADER: str = "X-MCP-Allow-External-Access" """HTTP header key for narrowing external-access permission.""" +MCP_ALLOW_DESTRUCTIVE_TOOLS_HEADER: str = "X-MCP-Allow-Destructive-Tools" +"""HTTP header key for narrowing destructive-tools 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 cc1848048..b94cefb22 100644 --- a/airbyte/exceptions.py +++ b/airbyte/exceptions.py @@ -339,25 +339,40 @@ class AirbyteLocalDemoModeError(AirbyteMCPError): """ +@dataclass +class AirbyteReadOnlyModeError(AirbyteMCPError): + """A non-read-only tool was blocked by MCP read-only mode.""" + + guidance: str | None = "Unset `AIRBYTE_CLOUD_MCP_READONLY_MODE`." + + @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 " - "`AIRBYTE_CLOUD_MCP_READONLY_MODE`, and ensure the " + "Set `AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES=1` and ensure the " "`X-MCP-Allow-Pipeline-Changes` request header is not `0`." ) +@dataclass +class AirbyteDestructiveToolsDisabledError(AirbyteMCPError): + """A destructive operation was blocked by MCP policy.""" + + guidance: str | None = ( + "Set `AIRBYTE_CLOUD_MCP_ALLOW_DESTRUCTIVE_TOOLS=1` and ensure the " + "`X-MCP-Allow-Destructive-Tools` 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 read-only mode) also disable external access when this setting is unset." + "`X-MCP-Allow-External-Access` request header is not `0`." ) diff --git a/airbyte/mcp/__init__.py b/airbyte/mcp/__init__.py index 105c56668..a8e3a0e7f 100644 --- a/airbyte/mcp/__init__.py +++ b/airbyte/mcp/__init__.py @@ -123,10 +123,7 @@ `AIRBYTE_CLOUD_MCP_READONLY_MODE` environment variable. When enabled, only read-only Cloud tools are available. Write and destructive operations are -disabled. - -This mode does allow running syncs on existing connectors, since sync operations -are not considered to be modifications of the Airbyte Cloud workspace. +disabled. Running and cancelling syncs are also disabled. Set the environment variable `AIRBYTE_CLOUD_MCP_READONLY_MODE=1` to enable read-only mode. @@ -145,16 +142,18 @@ `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. +external data or skill documentation. The `AIRBYTE_CLOUD_MCP_ALLOW_DESTRUCTIVE_TOOLS` +environment variable and `X-MCP-Allow-Destructive-Tools` request header control +destructive-annotated tools. 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. `AIRBYTE_CLOUD_MCP_READONLY_MODE=1` also disables pipeline changes. -When external access is unset, it is disabled if pipeline changes are disabled; -otherwise external-access tools remain available. Explicitly allowing external access -overrides these calculated defaults. Running or cancelling a sync is not a pipeline change. +Read-only mode blocks both pipeline changes and running or cancelling syncs. However, running +or cancelling a sync is _not_ considered a pipeline change, so those tools stay available when +pipeline changes are disabled, unless read-only mode is also enabled. ## Authentication for Remote (HTTP) Servers diff --git a/airbyte/mcp/_error_handling.py b/airbyte/mcp/_error_handling.py index f428982f5..0ada4cfe7 100644 --- a/airbyte/mcp/_error_handling.py +++ b/airbyte/mcp/_error_handling.py @@ -35,6 +35,7 @@ AirbyteConnectionSyncTimeoutError, AirbyteConnectorInUseError, AirbyteConnectorNotReadyError, + AirbyteDestructiveToolsDisabledError, AirbyteExternalAccessDisabledError, AirbyteLibError, AirbyteLibInputError, @@ -42,6 +43,7 @@ AirbyteMCPError, AirbyteMissingResourceError, AirbytePipelineChangesDisabledError, + AirbyteReadOnlyModeError, ) @@ -81,6 +83,8 @@ AirbyteMCPError, AirbyteLocalDemoModeError, AirbytePipelineChangesDisabledError, + AirbyteReadOnlyModeError, + AirbyteDestructiveToolsDisabledError, AirbyteExternalAccessDisabledError, AirbyteAgentsUnavailableError, AirbyteConnectorNotReadyError, diff --git a/airbyte/mcp/_policy_middleware.py b/airbyte/mcp/_policy_middleware.py index 1c7486f73..f1f92b7b9 100644 --- a/airbyte/mcp/_policy_middleware.py +++ b/airbyte/mcp/_policy_middleware.py @@ -3,16 +3,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from fastmcp.server.middleware import Middleware +from fastmcp_extensions.tool_filters import ( + ANNOTATION_DESTRUCTIVE_HINT, + ANNOTATION_READ_ONLY_HINT, + get_annotation, +) +from mcp.types import Tool from airbyte.exceptions import ( + AirbyteDestructiveToolsDisabledError, AirbyteExternalAccessDisabledError, AirbytePipelineChangesDisabledError, + AirbyteReadOnlyModeError, ) from airbyte.mcp._tool_utils import ( - check_legacy_safe_mode_header, + _is_cloud_tool, + airbyte_readonly_mode_enabled, + check_unsupported_mcp_headers, + destructive_tools_allowed, external_access_allowed, get_tool_policy, pipeline_changes_allowed, @@ -23,7 +34,8 @@ from collections.abc import Sequence from fastmcp.server.middleware import CallNext, MiddlewareContext - from fastmcp.tools import Tool, ToolResult + from fastmcp.tools import Tool as FastMCPTool + from fastmcp.tools import ToolResult from mcp.types import CallToolRequestParams, ListToolsRequest @@ -36,7 +48,7 @@ async def on_call_tool( call_next: CallNext[CallToolRequestParams, ToolResult], ) -> ToolResult: """Reject calls denied by effective pipeline or external-access settings.""" - check_legacy_safe_mode_header() + check_unsupported_mcp_headers() fastmcp_context = context.fastmcp_context if fastmcp_context is None: return await call_next(context) @@ -45,8 +57,21 @@ async def on_call_tool( if tool is None: return await call_next(context) + cloud_tool = _is_cloud_tool(fastmcp_context.fastmcp, cast(Tool, tool)) + if ( + cloud_tool + and airbyte_readonly_mode_enabled() + and not get_annotation(cast(Tool, tool), ANNOTATION_READ_ONLY_HINT, default=False) + ): + raise AirbyteReadOnlyModeError( + message="Read-only mode is enabled; only read-only tools are available." + ) policy = get_tool_policy(tool) - if policy.pipeline_change and pipeline_changes_allowed(fastmcp_context) is False: + if ( + cloud_tool + and policy.pipeline_change + and pipeline_changes_allowed(fastmcp_context) is False + ): raise AirbytePipelineChangesDisabledError( message="Pipeline-changing tools are disabled by the MCP server policy." ) @@ -54,13 +79,21 @@ async def on_call_tool( raise AirbyteExternalAccessDisabledError( message="External-access tools are disabled by the MCP server policy." ) + if ( + cloud_tool + and get_annotation(cast(Tool, tool), ANNOTATION_DESTRUCTIVE_HINT, default=False) + and destructive_tools_allowed(fastmcp_context) is False + ): + raise AirbyteDestructiveToolsDisabledError( + message="Destructive tools are disabled by the MCP server policy." + ) return await call_next(context) async def on_list_tools( self, context: MiddlewareContext[ListToolsRequest], - call_next: CallNext[ListToolsRequest, Sequence[Tool]], - ) -> Sequence[Tool]: - """Reject listings when the removed `X-MCP-Safe-Mode` header is present.""" - check_legacy_safe_mode_header() + call_next: CallNext[ListToolsRequest, Sequence[FastMCPTool]], + ) -> Sequence[FastMCPTool]: + """Reject listings when a removed MCP request header is present.""" + check_unsupported_mcp_headers() return await call_next(context) diff --git a/airbyte/mcp/_tool_utils.py b/airbyte/mcp/_tool_utils.py index d0e01ca2a..2473e0e92 100644 --- a/airbyte/mcp/_tool_utils.py +++ b/airbyte/mcp/_tool_utils.py @@ -31,6 +31,7 @@ _REGISTERED_TOOLS, # noqa: PLC2701 ) from fastmcp_extensions.tool_filters import ( + ANNOTATION_DESTRUCTIVE_HINT, ANNOTATION_READ_ONLY_HINT, CONFIG_INCLUDE_MODULES, CONFIG_TRUSTED_EXECUTION, @@ -45,14 +46,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_DESTRUCTIVE_TOOLS_ENV_VAR, + MCP_ALLOW_DESTRUCTIVE_TOOLS_HEADER, 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_DESTRUCTIVE_TOOLS, MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, MCP_CONFIG_ALLOW_PIPELINE_CHANGES, MCP_CONFIG_API_URL, @@ -144,59 +147,37 @@ def check_guid_created_in_session(guid: str) -> None: ) -def check_if_demo_mode_guards() -> None: - """Fail fast if the removed `AIRBYTE_CLOUD_MCP_SAFE_MODE` env var is set. - - Raises `AirbyteLibInputError` if the variable is present in the environment - with any value, including an empty string. - """ - if CLOUD_MCP_SAFE_MODE_ENV_VAR not in os.environ: - return - demo_mode_line = ( - "" - if is_hosted_mcp_mode() - else ( - "- `AIRBYTE_CLOUD_MCP_LOCAL_DEMO_MODE=1`: only allow changes to " - "resources created in this session (not recommended except for demos).\n" - ) - ) - raise AirbyteLibInputError( - message="`AIRBYTE_CLOUD_MCP_SAFE_MODE` is no longer supported.", - guidance=( - "Remove `AIRBYTE_CLOUD_MCP_SAFE_MODE` and consider using one of these " - "alternatives:\n" - "- `AIRBYTE_CLOUD_MCP_ALLOW_PIPELINE_CHANGES=0`: block changes to " - "sources, destinations, and connections.\n" - "- `AIRBYTE_CLOUD_MCP_ALLOW_EXTERNAL_ACCESS=0`: block SQL and API " - "passthrough tools.\n" - "- `AIRBYTE_CLOUD_MCP_READONLY_MODE=1`: only expose read-only tools.\n" - "- `MCP_NO_DESTRUCTIVE_TOOLS=1`: hide destructive tools.\n" - f"{demo_mode_line}" - ), - ) +UNSUPPORTED_MCP_HEADERS: dict[str, str] = { + MCP_LEGACY_SAFE_MODE_HEADER: ( + "Remove the `X-MCP-Safe-Mode` header and consider using one of these " + "alternatives:\n" + "- `X-MCP-Allow-Pipeline-Changes: 0`: block changes to sources, " + "destinations, and connections.\n" + "- `X-MCP-Allow-External-Access: 0`: block SQL and API passthrough tools.\n" + "- `X-MCP-Readonly-Mode: 1`: only expose read-only tools.\n" + "- `X-MCP-Allow-Destructive-Tools: 0`: hide destructive tools.\n" + ), +} +"""Removed MCP request headers, mapped to replacement guidance. + +Tool listing and calls fail if any is sent. +""" -def check_legacy_safe_mode_header() -> None: - """Fail fast if the removed `X-MCP-Safe-Mode` request header is present. +def check_unsupported_mcp_headers() -> None: + """Fail fast if a removed MCP request header is present. - Raises `AirbyteLibInputError` if the header is present with any value. - Outside an HTTP request there are no headers, so this is a no-op. + Raises `AirbyteLibInputError` on the first name in `UNSUPPORTED_MCP_HEADERS` + sent with any value. Outside an HTTP request there are no headers, so this + is a no-op. """ headers = get_http_headers() - if MCP_LEGACY_SAFE_MODE_HEADER.lower() not in headers: - return - raise AirbyteLibInputError( - message="The `X-MCP-Safe-Mode` header is no longer supported.", - guidance=( - "Remove the `X-MCP-Safe-Mode` header and consider using one of these " - "alternatives:\n" - "- `X-MCP-Allow-Pipeline-Changes: 0`: block changes to sources, " - "destinations, and connections.\n" - "- `X-MCP-Allow-External-Access: 0`: block SQL and API passthrough tools.\n" - "- `X-MCP-Readonly-Mode: 1`: only expose read-only tools.\n" - "- `X-MCP-No-Destructive-Tools: 1`: hide destructive tools.\n" - ), - ) + for name, guidance in UNSUPPORTED_MCP_HEADERS.items(): + if name.lower() in headers: + raise AirbyteLibInputError( + message=f"The `{name}` header is no longer supported.", + guidance=guidance, + ) # ============================================================================= @@ -262,6 +243,15 @@ def check_legacy_safe_mode_header() -> None: ) """Config arg for permission to run external-access Cloud tools.""" +ALLOW_DESTRUCTIVE_TOOLS_CONFIG_ARG = MCPServerConfigArg( + name=MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS, + http_header_key=MCP_ALLOW_DESTRUCTIVE_TOOLS_HEADER, + env_var=MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, + default="", + required=False, +) +"""Config arg for permission to run destructive Cloud tools.""" + TRUSTED_EXECUTION_CONFIG_ARG = MCPServerConfigArg( name=CONFIG_TRUSTED_EXECUTION, env_var=MCP_TRUSTED_EXECUTION_ENV_VAR, @@ -541,9 +531,33 @@ def _normalize_mcp_module(mcp_module: str) -> str: return mcp_module +_CLOUD_MCP_MODULES = frozenset({"cloud", "interactive"}) +"""MCP tool modules backed by Airbyte Cloud APIs. + +Cloud-scoped policies (pipeline changes, read-only mode, destructive tools) +apply to these modules only; local-operation tools are unaffected. +""" + + +def _is_cloud_tool(app: FastMCP, tool: Tool) -> bool: + """Return whether the tool belongs to a Cloud-backed MCP module.""" + mcp_module = get_tool_traits(app, tool.name).mcp_module + return _normalize_mcp_module(str(mcp_module or "")) in _CLOUD_MCP_MODULES + + def airbyte_readonly_mode_filter(tool: Tool, app: FastMCP) -> bool: - """Hide pipeline-changing tools when pipeline changes are disabled.""" - return not (pipeline_changes_allowed(app) is False and get_tool_policy(tool).pipeline_change) + """Hide non-read-only Cloud tools and pipeline-changing tools when disabled.""" + if ( + airbyte_readonly_mode_enabled() + and _is_cloud_tool(app, tool) + and not get_annotation(tool, ANNOTATION_READ_ONLY_HINT, default=False) + ): + return False + return not ( + _is_cloud_tool(app, tool) + and pipeline_changes_allowed(app) is False + and get_tool_policy(tool).pipeline_change + ) def airbyte_external_access_filter(tool: Tool, app: FastMCP) -> bool: @@ -551,6 +565,15 @@ def airbyte_external_access_filter(tool: Tool, app: FastMCP) -> bool: return not (get_tool_policy(tool).external_access and external_access_allowed(app) is False) +def airbyte_destructive_tools_filter(tool: Tool, app: FastMCP) -> bool: + """Hide destructive-annotated Cloud tools when destructive tools are disabled.""" + return not ( + _is_cloud_tool(app, tool) + and get_annotation(tool, ANNOTATION_DESTRUCTIVE_HINT, default=False) + and destructive_tools_allowed(app) is False + ) + + def _resolve_policy( app_or_ctx: FastMCP | Context, config_name: str, @@ -565,8 +588,13 @@ def _resolve_policy( return request_value +def airbyte_readonly_mode_enabled() -> bool: + """Return whether `AIRBYTE_CLOUD_MCP_READONLY_MODE` is set (env-only today).""" + return _str_to_bool(os.environ.get(MCP_READONLY_MODE_ENV_VAR)) is True + + 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: + if airbyte_readonly_mode_enabled(): return False return _resolve_policy( app_or_ctx, @@ -576,19 +604,21 @@ def pipeline_changes_allowed(app_or_ctx: FastMCP | Context) -> bool | None: def external_access_allowed(app_or_ctx: FastMCP | Context) -> bool | None: - """Resolve explicit permission first, then derive the external-access default. - - If external access is not explicitly configured, it is denied when pipeline - changes are disabled and left unset otherwise. - """ - explicit_value = _resolve_policy( + """Resolve the external-access setting.""" + return _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 - return False if pipeline_changes_allowed(app_or_ctx) is False else None + + +def destructive_tools_allowed(app_or_ctx: FastMCP | Context) -> bool | None: + """Resolve the destructive-tools setting.""" + return _resolve_policy( + app_or_ctx, + MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS, + MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, + ) def _insiders_mode(app: FastMCP) -> bool | None: diff --git a/airbyte/mcp/http_main.py b/airbyte/mcp/http_main.py index 918da85cb..341fea101 100644 --- a/airbyte/mcp/http_main.py +++ b/airbyte/mcp/http_main.py @@ -152,6 +152,7 @@ run_mcp_http_server, ) +from airbyte._util.unsupported_inputs import check_unsupported_env_vars from airbyte.constants import set_hosted_mcp_mode from airbyte.mcp._client_credentials import ( client_credentials_enabled, @@ -163,7 +164,6 @@ resolve_log_format, ) from airbyte.mcp._telemetry import McpRequestTelemetryMiddleware -from airbyte.mcp._tool_utils import check_if_demo_mode_guards from airbyte.mcp._transport_security import ( HTTP_HOST_ENV, HostOriginGuardMiddleware, @@ -292,7 +292,7 @@ def main() -> None: log_format = resolve_log_format() uvicorn_config = configure_logging(log_format) set_hosted_mcp_mode() - check_if_demo_mode_guards() + check_unsupported_env_vars() app.instructions = build_mcp_server_instructions(hosted=True) install(app) if log_format == "json": diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index dadb8304d..238a037df 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -79,6 +79,7 @@ from airbyte._util.meta import set_mcp_mode from airbyte._util.telemetry import DO_NOT_TRACK, PYAIRBYTE_MCP_TRACKING_KEY +from airbyte._util.unsupported_inputs import check_unsupported_env_vars from airbyte.constants import AIRBYTE_OFFLINE_MODE, _str_to_bool, is_hosted_mcp_mode from airbyte.mcp._config import load_secrets_to_env_vars from airbyte.mcp._error_handling import ( @@ -94,6 +95,7 @@ AIRBYTE_EXCLUDE_MODULES_CONFIG_ARG, AIRBYTE_INCLUDE_MODULES_CONFIG_ARG, AIRBYTE_READONLY_MODE_CONFIG_ARG, + ALLOW_DESTRUCTIVE_TOOLS_CONFIG_ARG, ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, ALLOW_PIPELINE_CHANGES_CONFIG_ARG, API_URL_CONFIG_ARG, @@ -105,10 +107,10 @@ ORGANIZATION_ID_CONFIG_ARG, TRUSTED_EXECUTION_CONFIG_ARG, WORKSPACE_ID_CONFIG_ARG, + airbyte_destructive_tools_filter, airbyte_external_access_filter, airbyte_module_filter, airbyte_readonly_mode_filter, - check_if_demo_mode_guards, validate_airbyte_domains, ) from airbyte.mcp._user_identity import ( @@ -534,6 +536,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 AIRBYTE_READONLY_MODE_CONFIG_ARG, ALLOW_PIPELINE_CHANGES_CONFIG_ARG, ALLOW_EXTERNAL_ACCESS_CONFIG_ARG, + ALLOW_DESTRUCTIVE_TOOLS_CONFIG_ARG, AIRBYTE_EXCLUDE_MODULES_CONFIG_ARG, AIRBYTE_INCLUDE_MODULES_CONFIG_ARG, INSIDERS_CONFIG_ARG, @@ -554,6 +557,7 @@ async def _mcp_mode_lifespan( # noqa: RUF029 tool_filters=[ airbyte_readonly_mode_filter, airbyte_external_access_filter, + airbyte_destructive_tools_filter, airbyte_module_filter, ], auth=_create_auth(), @@ -613,7 +617,7 @@ def main() -> None: It should not be called directly; instead, consult the MCP client documentation for instructions on how to connect to the server. """ - check_if_demo_mode_guards() + check_unsupported_env_vars() print("Starting Airbyte MCP server.", file=sys.stderr) try: asyncio.run(app.run_stdio_async()) diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index d9c4e0ff7..3513c3458 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -175,3 +175,15 @@ Inside each module page, primitives are grouped by kind (`## Tools`, (``) above its H3 so links like `cloud.md#deploy_connector_to_cloud` resolve in both pdoc and Docusaurus. Regenerate after any change to MCP tool signatures, descriptions, or schemas. + +## Removing Environment Variables and MCP Headers + +Never silently ignore a removed env var or MCP request header. Add it to +`UNSUPPORTED_ENV_VARS` in `airbyte/_util/unsupported_inputs.py` (checked at +startup of the `pyab` CLI and both MCP servers) or `UNSUPPORTED_MCP_HEADERS` in +`airbyte/mcp/_tool_utils.py`, mapped to guidance that names the replacement settings. + +- Env vars fail at startup. Headers fail on tool listing and tool calls. +- Any value counts as set, including empty or `0`. +- Remove all other docs references to the removed setting. The map entry is the only + place it stays. diff --git a/tests/unit_tests/test_mcp_policy_middleware.py b/tests/unit_tests/test_mcp_policy_middleware.py index b272ad2ba..b457eee61 100644 --- a/tests/unit_tests/test_mcp_policy_middleware.py +++ b/tests/unit_tests/test_mcp_policy_middleware.py @@ -10,17 +10,22 @@ import pytest from fastmcp import Client from fastmcp.exceptions import ToolError -from fastmcp_extensions import mcp_server +from fastmcp_extensions import ToolTraits, mcp_server from airbyte.constants import ( + MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS, + MCP_CONFIG_ALLOW_PIPELINE_CHANGES, MCP_READONLY_MODE_ENV_VAR, ) from airbyte.exceptions import ( + AirbyteDestructiveToolsDisabledError, AirbyteExternalAccessDisabledError, AirbyteLibInputError, AirbytePipelineChangesDisabledError, + AirbyteReadOnlyModeError, ) from airbyte.mcp import _tool_utils from airbyte.mcp._error_handling import ( @@ -40,15 +45,26 @@ async def get_tool(self, name: str) -> Any: return self.tool +_MW_TOOL_MODULES: dict[str, str] = {} +"""Tool-name -> mcp_module map for the `get_tool_traits` stub, cleared per test.""" + + @pytest.fixture(autouse=True) def clear_policy_environment(monkeypatch: pytest.MonkeyPatch) -> None: for env_var in ( + MCP_ALLOW_DESTRUCTIVE_TOOLS_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", {}) + _MW_TOOL_MODULES.clear() + monkeypatch.setattr( + _tool_utils, + "get_tool_traits", + lambda _app, name: ToolTraits(mcp_module=_MW_TOOL_MODULES.get(name, "cloud")), + ) @pytest.mark.parametrize( @@ -94,8 +110,26 @@ def clear_policy_environment(monkeypatch: pytest.MonkeyPatch) -> None: False, {MCP_READONLY_MODE_ENV_VAR: "1"}, {}, - AirbytePipelineChangesDisabledError, - id="legacy-read-only-denial", + AirbyteReadOnlyModeError, + id="readonly-mode-denial", + ), + pytest.param( + "sync_control", + ToolPolicy(pipeline_change=False), + False, + {MCP_READONLY_MODE_ENV_VAR: "1"}, + {}, + AirbyteReadOnlyModeError, + id="readonly-mode-blocks-sync-control", + ), + pytest.param( + "readonly_tool", + None, + True, + {MCP_READONLY_MODE_ENV_VAR: "1"}, + {}, + None, + id="readonly-mode-allows-read-only", ), pytest.param( "external_query", @@ -130,8 +164,26 @@ def clear_policy_environment(monkeypatch: pytest.MonkeyPatch) -> None: False, {MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR: "0"}, {}, - AirbyteExternalAccessDisabledError, - id="pipeline-changes-disabled-default-denial", + None, + id="pipeline-changes-disabled-pass-through", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + True, + {MCP_READONLY_MODE_ENV_VAR: "1"}, + {}, + None, + id="readonly-mode-external-pass-through", + ), + pytest.param( + "external_query", + ToolPolicy(pipeline_change=False, external_access=True), + False, + {}, + {MCP_CONFIG_ALLOW_PIPELINE_CHANGES: "0"}, + None, + id="pipeline-header-deny-pass-through", ), ], ) @@ -177,10 +229,131 @@ async def invoke() -> str: asyncio.run(invoke()) -def test_legacy_safe_mode_header_blocks_list_and_call( +@pytest.mark.parametrize( + ("environment", "request_config", "destructive_hint", "expected_error"), + [ + pytest.param( + {MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR: "0"}, + {}, + True, + AirbyteDestructiveToolsDisabledError, + id="env-deny", + ), + pytest.param( + {}, + {MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS: "0"}, + True, + AirbyteDestructiveToolsDisabledError, + id="header-deny", + ), + pytest.param( + {MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR: "1"}, + {MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS: "0"}, + True, + AirbyteDestructiveToolsDisabledError, + id="header-narrows-env-allow", + ), + pytest.param({}, {}, True, None, id="unset-default-allows"), + pytest.param( + {MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR: "0"}, + {}, + False, + None, + id="non-destructive-unaffected", + ), + ], +) +def test_destructive_tools_guard( + monkeypatch: pytest.MonkeyPatch, + environment: dict[str, str], + request_config: dict[str, str], + destructive_hint: bool, + expected_error: type[Exception] | None, +) -> None: + """Destructive-annotated tools are blocked when destructive tools are denied.""" + 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) + + tool_name = "destructive_tool" + tool = SimpleNamespace( + name=tool_name, + annotations=SimpleNamespace( + read_only_hint=True, + destructive_hint=destructive_hint, + ), + ) + context = SimpleNamespace( + fastmcp_context=SimpleNamespace(fastmcp=_FakeFastMCP(tool)), + 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_local_tools_bypass_cloud_policy_guards( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Local-module tools are not blocked by Cloud read-only/destructive policies.""" + monkeypatch.setattr( + _tool_utils, + "get_mcp_config", + lambda _app, key, **_kwargs: None, + ) + + async def call_next(_context: Any) -> str: + return "called" + + async def call(tool_name: str, annotations: SimpleNamespace) -> str: + tool = SimpleNamespace(name=tool_name, annotations=annotations) + context = SimpleNamespace( + fastmcp_context=SimpleNamespace(fastmcp=_FakeFastMCP(tool)), + message=SimpleNamespace(name=tool_name), + ) + return await PolicyGuardMiddleware().on_call_tool(context, call_next) + + _MW_TOOL_MODULES["destination_smoke_test"] = "local" + _MW_TOOL_MODULES["sync_connector"] = "local" + # No `ToolPolicy` registration: the tools exercise the real default + # (non-read-only -> pipeline_change=True). + + monkeypatch.setenv(MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, "0") + monkeypatch.setenv(MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, "0") + assert ( + asyncio.run( + call( + "destination_smoke_test", + SimpleNamespace(destructive_hint=True), + ) + ) + == "called" + ) + + monkeypatch.setenv(MCP_READONLY_MODE_ENV_VAR, "1") + assert ( + asyncio.run(call("sync_connector", SimpleNamespace(read_only_hint=False))) + == "called" + ) + + +def test_unsupported_mcp_header_blocks_list_and_call( monkeypatch: pytest.MonkeyPatch, ) -> None: - """The removed `X-MCP-Safe-Mode` header fails both list_tools and call_tool.""" + """A removed request header fails both list_tools and call_tool.""" monkeypatch.setattr( _tool_utils, "get_http_headers", diff --git a/tests/unit_tests/test_mcp_tool_filters.py b/tests/unit_tests/test_mcp_tool_filters.py index ead7e37dd..318953b9b 100644 --- a/tests/unit_tests/test_mcp_tool_filters.py +++ b/tests/unit_tests/test_mcp_tool_filters.py @@ -11,8 +11,10 @@ from fastmcp_extensions.tool_filters import CONFIG_INCLUDE_MODULES from airbyte.constants import ( + MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, + MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS, MCP_CONFIG_ALLOW_EXTERNAL_ACCESS, MCP_CONFIG_ALLOW_PIPELINE_CHANGES, MCP_CONFIG_EXCLUDE_MODULES, @@ -32,6 +34,9 @@ APP = cast(FastMCP, object()) """Stand-in for the app; the filter only passes it to `get_mcp_config`, which is patched.""" +_TOOL_MODULES: dict[str, str] = {} +"""Tool-name -> mcp_module map for the `get_tool_traits` stub, cleared per test.""" + def _tool(name: str) -> Tool: """Return a tool-like object named for the module it belongs to.""" @@ -43,6 +48,7 @@ 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 ( + MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, MCP_ALLOW_EXTERNAL_ACCESS_ENV_VAR, MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, MCP_READONLY_MODE_ENV_VAR, @@ -57,11 +63,12 @@ def mcp_config(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: ) # `mcp_module` is internal state (never on the wire), so stub the traits # registry lookup with a tool-name -> module map. + _TOOL_MODULES.clear() monkeypatch.setattr( _tool_utils, "get_tool_traits", lambda _app, name: ToolTraits( - mcp_module=name if name != "unannotated" else None + mcp_module=_TOOL_MODULES.get(name, name if name != "unannotated" else None) ), ) return config @@ -174,13 +181,32 @@ def test_sync_tools_stay_visible_when_pipeline_changes_are_disabled( assert _tool_utils.airbyte_readonly_mode_filter(tool, APP) +def test_non_read_only_tools_hidden_in_readonly_mode( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], +) -> None: + """Read-only mode hides non-read-only tools but keeps read-only-annotated tools.""" + assert not mcp_config + monkeypatch.setenv(MCP_READONLY_MODE_ENV_VAR, "1") + for tool_name in ("run_cloud_sync", "cancel_cloud_sync"): + tool = _policy_tool(tool_name, pipeline_change=False) + assert not _tool_utils.airbyte_readonly_mode_filter(tool, APP) + read_only_tool = _policy_tool( + "list_cloud_workspaces", pipeline_change=False, read_only_hint=True + ) + assert _tool_utils.airbyte_readonly_mode_filter(read_only_tool, APP) + + def _policy_tool( name: str, *, pipeline_change: bool | None, external_access: bool = False, read_only_hint: bool = False, + destructive_hint: bool = False, + module: str = "cloud", ) -> Tool: + _TOOL_MODULES[name] = module if pipeline_change is not None or external_access: _tool_utils._TOOL_POLICIES[name] = _tool_utils.ToolPolicy( pipeline_change=bool(pipeline_change), @@ -190,7 +216,10 @@ def _policy_tool( Tool, SimpleNamespace( name=name, - annotations=SimpleNamespace(read_only_hint=read_only_hint), + annotations=SimpleNamespace( + read_only_hint=read_only_hint, + destructive_hint=destructive_hint, + ), ), ) @@ -260,8 +289,10 @@ def test_pipeline_change_filter( pytest.param("0", "1", None, None, False, id="external-env-deny-wins"), pytest.param(None, "0", None, None, False, id="external-header-deny"), pytest.param("1", None, None, "1", True, id="explicit-external-allow"), - pytest.param(None, None, "0", None, False, id="pipeline-deny-default"), - pytest.param(None, None, None, "1", False, id="legacy-readonly-default"), + pytest.param( + None, None, "0", None, True, id="pipeline-deny-external-unaffected" + ), + pytest.param(None, None, None, "1", True, id="readonly-external-unaffected"), pytest.param(None, None, None, None, True, id="unset-default-allows"), ], ) @@ -286,7 +317,94 @@ def test_external_access_filter( 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 + "external-access-tool", + pipeline_change=False, + external_access=True, + read_only_hint=True, ) assert _tool_utils.airbyte_external_access_filter(tool, APP) is visible + + +@pytest.mark.parametrize( + ("env_value", "header", "destructive", "visible"), + [ + pytest.param("0", None, True, False, id="env-deny"), + pytest.param(None, "0", True, False, id="header-deny"), + pytest.param("1", "0", True, False, id="header-narrows-env-allow"), + pytest.param(None, None, True, True, id="unset-default-allows"), + pytest.param("0", None, False, True, id="non-destructive-unaffected"), + ], +) +def test_destructive_tools_filter( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], + env_value: str | None, + header: str | None, + destructive: bool, + visible: bool, +) -> None: + if env_value is None: + monkeypatch.delenv(MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, raising=False) + else: + monkeypatch.setenv(MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, env_value) + if header is not None: + mcp_config[MCP_CONFIG_ALLOW_DESTRUCTIVE_TOOLS] = header + tool = _policy_tool( + "destructive-tool", + pipeline_change=False, + destructive_hint=destructive, + ) + + assert _tool_utils.airbyte_destructive_tools_filter(tool, APP) is visible + + +def test_local_tools_unaffected_by_readonly_and_destructive_policies( + monkeypatch: pytest.MonkeyPatch, + mcp_config: dict[str, str], +) -> None: + """Local-module tools stay visible under Cloud read-only and destructive policies.""" + assert not mcp_config + monkeypatch.setenv(MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, "0") + monkeypatch.setenv(MCP_READONLY_MODE_ENV_VAR, "1") + monkeypatch.setenv(MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, "0") + + # No `ToolPolicy` registration: the tools exercise the real default + # (non-read-only -> pipeline_change=True). + destructive_local_tool = _policy_tool( + "destination_smoke_test", + pipeline_change=None, + destructive_hint=True, + module="local", + ) + assert _tool_utils.airbyte_destructive_tools_filter(destructive_local_tool, APP) + assert _tool_utils.airbyte_readonly_mode_filter(destructive_local_tool, APP) + + non_readonly_local_tool = _policy_tool( + "sync_connector", + pipeline_change=None, + module="local", + ) + assert _tool_utils.airbyte_readonly_mode_filter(non_readonly_local_tool, APP) + + +def test_real_server_keeps_local_tools_listed_under_cloud_denials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Real registrations: local tools survive Cloud denials; Cloud tools do not.""" + import asyncio + + from airbyte.mcp.server import app as airbyte_app + + monkeypatch.setenv(MCP_ALLOW_PIPELINE_CHANGES_ENV_VAR, "0") + monkeypatch.setenv(MCP_READONLY_MODE_ENV_VAR, "1") + monkeypatch.setenv(MCP_ALLOW_DESTRUCTIVE_TOOLS_ENV_VAR, "0") + + def visible(name: str) -> bool: + tool = cast(Tool, asyncio.run(airbyte_app.get_tool(name))) + return _tool_utils.airbyte_readonly_mode_filter( + tool, airbyte_app + ) and _tool_utils.airbyte_destructive_tools_filter(tool, airbyte_app) + + assert visible("destination_smoke_test") + assert not visible("run_cloud_sync") diff --git a/tests/unit_tests/test_mcp_tool_utils.py b/tests/unit_tests/test_mcp_tool_utils.py index 7b593c707..78221219f 100644 --- a/tests/unit_tests/test_mcp_tool_utils.py +++ b/tests/unit_tests/test_mcp_tool_utils.py @@ -10,7 +10,6 @@ from fastmcp_extensions import MCPServerConfigArg from airbyte.constants import ( - CLOUD_MCP_SAFE_MODE_ENV_VAR, MCP_LOCAL_DEMO_MODE_ENV_VAR, ) from airbyte.exceptions import AirbyteLibInputError, AirbyteLocalDemoModeError @@ -27,8 +26,7 @@ ToolPolicy, _resolve_transport_bearer_token, check_guid_created_in_session, - check_if_demo_mode_guards, - check_legacy_safe_mode_header, + check_unsupported_mcp_headers, mcp_tool, register_guid_created_in_session, ) @@ -152,64 +150,24 @@ def test_check_guid_created_in_session_passes_when_hosted( check_guid_created_in_session("any-guid-at-all") -def test_check_if_demo_mode_guards_passes_when_unset( +def test_check_unsupported_mcp_headers_passes_when_absent( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Test that startup check passes when the legacy env var is unset.""" - monkeypatch.delenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, raising=False) - check_if_demo_mode_guards() - - -@pytest.mark.parametrize("value", ["1", "0", ""]) -def test_check_if_demo_mode_guards_raises_on_stdio( - monkeypatch: pytest.MonkeyPatch, - value: str, -) -> None: - """Test that any value of the legacy env var fails startup on STDIO.""" - monkeypatch.setenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, value) - monkeypatch.setattr("airbyte.mcp._tool_utils.is_hosted_mcp_mode", lambda: False) - with pytest.raises(AirbyteLibInputError) as exc_info: - check_if_demo_mode_guards() - assert "no longer supported" in str(exc_info.value) - guidance = exc_info.value.guidance - assert isinstance(guidance, str) - assert guidance.startswith( - "Remove `AIRBYTE_CLOUD_MCP_SAFE_MODE` and consider using one of these alternatives:\n" - ) - bullets = [line for line in guidance.splitlines() if line.startswith("- ")] - assert "AIRBYTE_CLOUD_MCP_LOCAL_DEMO_MODE" in bullets[-1] - assert "not recommended except for demos" in bullets[-1] - - -def test_check_if_demo_mode_guards_raises_when_hosted( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test that the hosted guidance omits the STDIO-only demo mode bullet.""" - monkeypatch.setenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, "1") - monkeypatch.setattr("airbyte.mcp._tool_utils.is_hosted_mcp_mode", lambda: True) - with pytest.raises(AirbyteLibInputError) as exc_info: - check_if_demo_mode_guards() - assert "DEMO" not in (exc_info.value.guidance or "") - - -def test_check_legacy_safe_mode_header_passes_when_absent( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test that the header check passes when the legacy header is absent.""" + """Test that the header check passes when no listed header is sent.""" monkeypatch.setattr("airbyte.mcp._tool_utils.get_http_headers", lambda **_: {}) - check_legacy_safe_mode_header() + check_unsupported_mcp_headers() -def test_check_legacy_safe_mode_header_raises_when_present( +def test_check_unsupported_mcp_headers_raises_when_present( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Test that the removed header fails the request with migration guidance.""" + """Test that a listed header fails the request with mapped guidance.""" monkeypatch.setattr( "airbyte.mcp._tool_utils.get_http_headers", lambda **_: {"x-mcp-safe-mode": "1"}, ) with pytest.raises(AirbyteLibInputError) as exc_info: - check_legacy_safe_mode_header() + check_unsupported_mcp_headers() assert "X-MCP-Safe-Mode" in str(exc_info.value) guidance = exc_info.value.guidance assert isinstance(guidance, str) @@ -217,7 +175,24 @@ def test_check_legacy_safe_mode_header_raises_when_present( "Remove the `X-MCP-Safe-Mode` header and consider using one of these alternatives:\n" ) assert "X-MCP-Allow-Pipeline-Changes" in guidance - assert "DEMO" not in guidance + + +def test_check_unsupported_mcp_headers_is_generic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that the header guard works for any name added to the map.""" + monkeypatch.setattr( + "airbyte.mcp._tool_utils.UNSUPPORTED_MCP_HEADERS", + {"X-Custom-Removed": "use Y"}, + ) + monkeypatch.setattr( + "airbyte.mcp._tool_utils.get_http_headers", + lambda **_: {"x-custom-removed": "1"}, + ) + with pytest.raises(AirbyteLibInputError) as exc_info: + check_unsupported_mcp_headers() + assert "X-Custom-Removed" in str(exc_info.value) + assert exc_info.value.guidance == "use Y" def test_multiple_guids_can_be_registered() -> None: diff --git a/tests/unit_tests/test_unsupported_inputs.py b/tests/unit_tests/test_unsupported_inputs.py new file mode 100644 index 000000000..243f060d4 --- /dev/null +++ b/tests/unit_tests/test_unsupported_inputs.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024 Airbyte, Inc., all rights reserved. +"""Unit tests for the removed-env-var startup guard and the `pyab` CLI.""" + +from __future__ import annotations + +import pytest + +from airbyte._util.unsupported_inputs import check_unsupported_env_vars +from airbyte.cli import pyab +from airbyte.constants import CLOUD_MCP_SAFE_MODE_ENV_VAR +from airbyte.exceptions import AirbyteLibInputError + + +def test_check_unsupported_env_vars_passes_when_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that the startup check passes when no listed env var is set.""" + monkeypatch.delenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, raising=False) + check_unsupported_env_vars() + + +@pytest.mark.parametrize("value", ["1", "0", ""]) +def test_check_unsupported_env_vars_raises_when_set( + monkeypatch: pytest.MonkeyPatch, + value: str, +) -> None: + """Test that any value of a listed env var fails startup with mapped guidance.""" + monkeypatch.setenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, value) + with pytest.raises(AirbyteLibInputError) as exc_info: + check_unsupported_env_vars() + assert CLOUD_MCP_SAFE_MODE_ENV_VAR in str(exc_info.value) + guidance = exc_info.value.guidance + assert isinstance(guidance, str) + assert guidance.startswith( + "Remove `AIRBYTE_CLOUD_MCP_SAFE_MODE` and consider using one of these alternatives:\n" + ) + + +def test_check_unsupported_env_vars_is_generic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that the env-var guard works for any name added to the map.""" + monkeypatch.setattr( + "airbyte._util.unsupported_inputs.UNSUPPORTED_ENV_VARS", + {"FAKE_REMOVED_VAR": "use X"}, + ) + monkeypatch.setenv("FAKE_REMOVED_VAR", "1") + with pytest.raises(AirbyteLibInputError) as exc_info: + check_unsupported_env_vars() + assert "FAKE_REMOVED_VAR" in str(exc_info.value) + assert exc_info.value.guidance == "use X" + + +def test_pyab_cli_fails_on_unsupported_env_var( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that every `pyab` CLI command runs the env-var guard first.""" + monkeypatch.setenv(CLOUD_MCP_SAFE_MODE_ENV_VAR, "1") + with pytest.raises(AirbyteLibInputError): + pyab.cli(["benchmark"], exit_on_error=False)