diff --git a/airbyte/cloud/_credentials.py b/airbyte/cloud/_credentials.py index a842f6c73..6fc347cda 100644 --- a/airbyte/cloud/_credentials.py +++ b/airbyte/cloud/_credentials.py @@ -15,7 +15,7 @@ CLOUD_ORGANIZATION_ID_ENV_VAR, CLOUD_WORKSPACE_ID_ENV_VAR, ) -from airbyte.exceptions import PyAirbyteInputError +from airbyte.exceptions import AirbyteNoCloudCredentialsError, PyAirbyteInputError from airbyte.secrets.base import SecretString from airbyte.secrets.util import try_get_secret @@ -86,15 +86,7 @@ def from_auth( guidance="Provide both client ID and client secret, or use a bearer token.", ) if not resolved_bearer_token and not resolved_client_id: - guidance = ( - "Set Airbyte Cloud credentials in environment variables." - if env_vars - else "Provide either bearer_token or both client_id and client_secret." - ) - raise PyAirbyteInputError( - message="No Airbyte credentials found.", - guidance=guidance, - ) + raise AirbyteNoCloudCredentialsError(_env_vars=env_vars) return cls( client_id=SecretString(resolved_client_id) if resolved_client_id else None, diff --git a/airbyte/constants.py b/airbyte/constants.py index af5d03429..68a4520bb 100644 --- a/airbyte/constants.py +++ b/airbyte/constants.py @@ -261,6 +261,21 @@ def _str_to_bool(value: str) -> bool: # MCP (Model Context Protocol) Constants +_HOSTED_MCP_MODE_ENABLED: bool = False +"""Whether the process is serving MCP over hosted HTTP transport.""" + + +def set_hosted_mcp_mode() -> None: + """Set the flag indicating the process serves MCP over hosted HTTP transport.""" + global _HOSTED_MCP_MODE_ENABLED + _HOSTED_MCP_MODE_ENABLED = True + + +def is_hosted_mcp_mode() -> bool: + """Return True if the process serves MCP over hosted HTTP transport.""" + return _HOSTED_MCP_MODE_ENABLED + + MCP_READONLY_MODE_ENV_VAR: str = "AIRBYTE_CLOUD_MCP_READONLY_MODE" """Environment variable to enable read-only mode for the MCP server. diff --git a/airbyte/exceptions.py b/airbyte/exceptions.py index 4b5e8aa52..96117b140 100644 --- a/airbyte/exceptions.py +++ b/airbyte/exceptions.py @@ -42,14 +42,24 @@ from dataclasses import dataclass from pathlib import Path from textwrap import indent -from typing import TYPE_CHECKING, Any - -from airbyte.constants import AIRBYTE_PRINT_FULL_ERROR_LOGS +from typing import TYPE_CHECKING, Any, Protocol + +from airbyte.constants import ( + AIRBYTE_PRINT_FULL_ERROR_LOGS, + CLOUD_BEARER_TOKEN_ENV_VAR, + CLOUD_CLIENT_ID_ENV_VAR, + CLOUD_CLIENT_SECRET_ENV_VAR, + CLOUD_WORKSPACE_ID_ENV_VAR, + MCP_BEARER_TOKEN_HEADER, + MCP_CLIENT_ID_HEADER, + MCP_CLIENT_SECRET_HEADER, + MCP_WORKSPACE_ID_HEADER, + is_hosted_mcp_mode, +) if TYPE_CHECKING: from airbyte._util.api_duck_types import AirbyteApiResponseDuckType - from airbyte.cloud.workspaces import CloudWorkspace NEW_ISSUE_URL = "https://github.com/airbytehq/airbyte/issues/new/choose" @@ -222,6 +232,71 @@ class PyAirbyteNoStreamsSelectedError(PyAirbyteInputError): available_streams: list[str] | None = None +@dataclass +class AirbyteNoCloudCredentialsError(PyAirbyteInputError): + """No Airbyte credentials found.""" + + guidance: str | None = None + _allow_bearer: bool = True + _env_vars: bool = True + + def __post_init__(self) -> None: + """Set guidance for the current execution mode.""" + if self.guidance is not None: + return + if is_hosted_mcp_mode(): + if self._allow_bearer: + self.guidance = ( + f"Provide a bearer token via the `{MCP_BEARER_TOKEN_HEADER}` header, " + f"or client credentials via the `{MCP_CLIENT_ID_HEADER}` and " + f"`{MCP_CLIENT_SECRET_HEADER}` headers." + ) + else: + self.guidance = ( + f"Provide client credentials via the `{MCP_CLIENT_ID_HEADER}` and " + f"`{MCP_CLIENT_SECRET_HEADER}` headers." + ) + elif self._allow_bearer and self._env_vars: + self.guidance = ( + f"Provide `bearer_token`, or both `client_id` and `client_secret`, as " + f"arguments or via the `{CLOUD_BEARER_TOKEN_ENV_VAR}`, " + f"`{CLOUD_CLIENT_ID_ENV_VAR}`, and `{CLOUD_CLIENT_SECRET_ENV_VAR}` " + "environment variables." + ) + elif self._allow_bearer: + self.guidance = "Provide `bearer_token`, or both `client_id` and `client_secret`." + elif self._env_vars: + self.guidance = ( + f"Provide both `client_id` and `client_secret`, as arguments or via the " + f"`{CLOUD_CLIENT_ID_ENV_VAR}` and `{CLOUD_CLIENT_SECRET_ENV_VAR}` " + "environment variables." + ) + else: + self.guidance = "Provide both `client_id` and `client_secret`." + + +@dataclass +class AirbyteMissingWorkspaceContextError(PyAirbyteInputError): + """Workspace ID is required but not provided.""" + + guidance: str | None = None + + def __post_init__(self) -> None: + """Set guidance for the current execution mode.""" + if self.guidance is not None: + return + if is_hosted_mcp_mode(): + self.guidance = ( + f"Provide the workspace ID via the `{MCP_WORKSPACE_ID_HEADER}` header, " + "or pass the `workspace_id` parameter." + ) + else: + self.guidance = ( + f"Set the `{CLOUD_WORKSPACE_ID_ENV_VAR}` environment variable, or pass " + "the `workspace_id` parameter." + ) + + # MCP Server Errors @@ -456,6 +531,19 @@ class PyAirbyteSecretNotFoundError(PyAirbyteError): # Airbyte API Errors +class _WorkspaceWithUrl(Protocol): + """Structural type for a workspace that exposes a `workspace_url`. + + Declared locally so `exceptions` does not need to import `airbyte.cloud`, which + would create an import cycle. Any object with a `workspace_url` attribute (e.g. + `CloudWorkspace`) satisfies this via structural (duck) typing. + """ + + @property + def workspace_url(self) -> str | None: + """The web URL of the workspace.""" + + @dataclass class AirbyteError(PyAirbyteError): """An error occurred while communicating with the hosted Airbyte instance.""" @@ -463,7 +551,7 @@ class AirbyteError(PyAirbyteError): response: AirbyteApiResponseDuckType | None = None """The API response from the failed request.""" - workspace: CloudWorkspace | None = None + workspace: _WorkspaceWithUrl | None = None """The workspace where the error occurred.""" @property diff --git a/airbyte/mcp/_tool_utils.py b/airbyte/mcp/_tool_utils.py index 94816df6a..46e4089bd 100644 --- a/airbyte/mcp/_tool_utils.py +++ b/airbyte/mcp/_tool_utils.py @@ -16,7 +16,11 @@ from typing import TYPE_CHECKING, TypeVar from fastmcp.apps import UI_EXTENSION_ID -from fastmcp.server.dependencies import get_access_token, get_context, get_http_headers +from fastmcp.server.dependencies import ( + get_access_token, + get_context, + get_http_headers, +) from fastmcp_extensions import MCPServerConfigArg, get_mcp_config from fastmcp_extensions import mcp_tool as _mcp_tool from fastmcp_extensions.decorators import ( diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index 0beab1b41..eb6a1925f 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -25,15 +25,27 @@ from airbyte.cloud.models import JobTypeEnum from airbyte.cloud.workspaces import CloudWorkspace from airbyte.constants import ( + CLOUD_BEARER_TOKEN_ENV_VAR, + CLOUD_CLIENT_ID_ENV_VAR, + CLOUD_CLIENT_SECRET_ENV_VAR, + CLOUD_WORKSPACE_ID_ENV_VAR, + MCP_BEARER_TOKEN_HEADER, + MCP_CLIENT_ID_HEADER, + MCP_CLIENT_SECRET_HEADER, MCP_CONFIG_API_URL, MCP_CONFIG_BEARER_TOKEN, MCP_CONFIG_CLIENT_ID, MCP_CONFIG_CLIENT_SECRET, MCP_CONFIG_CONFIG_API_URL, MCP_CONFIG_WORKSPACE_ID, + MCP_WORKSPACE_ID_HEADER, ) from airbyte.destinations.util import get_noop_destination -from airbyte.exceptions import AirbyteMissingResourceError, PyAirbyteInputError +from airbyte.exceptions import ( + AirbyteMissingResourceError, + AirbyteMissingWorkspaceContextError, + PyAirbyteInputError, +) from airbyte.mcp._arg_resolvers import resolve_connector_config, resolve_list_of_strings from airbyte.mcp._tool_utils import ( AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET, @@ -43,11 +55,20 @@ CLOUD_AUTH_TIP_TEXT = ( - "By default, the `AIRBYTE_CLOUD_CLIENT_ID`, `AIRBYTE_CLOUD_CLIENT_SECRET`, " - "and `AIRBYTE_CLOUD_WORKSPACE_ID` environment variables " - "will be used to authenticate with the Airbyte Cloud API." + f"When connecting to a hosted MCP server, provide a bearer token via the " + f"`{MCP_BEARER_TOKEN_HEADER}` header, or client credentials via the " + f"`{MCP_CLIENT_ID_HEADER}` and `{MCP_CLIENT_SECRET_HEADER}` headers, plus " + f"the workspace ID via the `{MCP_WORKSPACE_ID_HEADER}` header. For local or " + f"stdio connections, set the `{CLOUD_BEARER_TOKEN_ENV_VAR}` environment " + f"variable, or both `{CLOUD_CLIENT_ID_ENV_VAR}` and " + f"`{CLOUD_CLIENT_SECRET_ENV_VAR}`, plus `{CLOUD_WORKSPACE_ID_ENV_VAR}` for the " + f"workspace." +) +WORKSPACE_ID_TIP_TEXT = ( + f"Workspace ID. Hosted MCP connections pass it via the " + f"`{MCP_WORKSPACE_ID_HEADER}` header; local or stdio connections use the " + f"`{CLOUD_WORKSPACE_ID_ENV_VAR}` environment variable." ) -WORKSPACE_ID_TIP_TEXT = "Workspace ID. Defaults to `AIRBYTE_CLOUD_WORKSPACE_ID` env var." class CloudSourceResult(BaseModel): @@ -259,10 +280,7 @@ def _get_cloud_workspace( """ resolved_workspace_id = workspace_id or get_mcp_config(ctx, MCP_CONFIG_WORKSPACE_ID) if not resolved_workspace_id: - raise PyAirbyteInputError( - message="Workspace ID is required but not provided.", - guidance="Set AIRBYTE_CLOUD_WORKSPACE_ID env var or pass workspace_id parameter.", - ) + raise AirbyteMissingWorkspaceContextError return _get_cloud_client(ctx).get_workspace(resolved_workspace_id) diff --git a/airbyte/mcp/http_main.py b/airbyte/mcp/http_main.py index 5f51cdc30..0ca1e61ca 100644 --- a/airbyte/mcp/http_main.py +++ b/airbyte/mcp/http_main.py @@ -50,6 +50,7 @@ register_landing_page, ) +from airbyte.constants import set_hosted_mcp_mode from airbyte.mcp.server import ( DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT, @@ -80,6 +81,7 @@ def _get_server_url() -> str: def main() -> None: """Start the Airbyte MCP server with HTTP transport.""" logging.basicConfig(level=logging.INFO) + set_hosted_mcp_mode() # When deployed behind a path-stripping LB (MCP_SERVER_URL has a path # component like /cloud-mcp), serve the MCP endpoint at root so the diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index 249cff11a..88bd21e58 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -110,8 +110,10 @@ - Listing and describing environment variables for connector configuration Operational modes: -- Cloud operations: Deploy and manage connectors on Airbyte Cloud (requires - AIRBYTE_CLOUD_CLIENT_ID, AIRBYTE_CLOUD_CLIENT_SECRET, AIRBYTE_CLOUD_WORKSPACE_ID) +- Cloud operations: Deploy and manage connectors on Airbyte Cloud (use request + headers when connecting to a hosted MCP server, or AIRBYTE_CLOUD_CLIENT_ID + + AIRBYTE_CLOUD_CLIENT_SECRET (or AIRBYTE_CLOUD_BEARER_TOKEN) plus + AIRBYTE_CLOUD_WORKSPACE_ID for local or stdio connections) - Local operations: Run connectors locally for data extraction (requires AIRBYTE_PROJECT_DIR for artifact storage) diff --git a/tests/unit_tests/test_cloud_credentials.py b/tests/unit_tests/test_cloud_credentials.py index e3438f99a..67f511b84 100644 --- a/tests/unit_tests/test_cloud_credentials.py +++ b/tests/unit_tests/test_cloud_credentials.py @@ -69,12 +69,17 @@ def fake_try_get_secret( [ pytest.param( False, - "Provide either bearer_token or both client_id and client_secret.", + "Provide `bearer_token`, or both `client_id` and `client_secret`.", id="explicit_inputs", ), pytest.param( True, - "Set Airbyte Cloud credentials in environment variables.", + ( + "Provide `bearer_token`, or both `client_id` and `client_secret`, as " + "arguments or via the `AIRBYTE_CLOUD_BEARER_TOKEN`, " + "`AIRBYTE_CLOUD_CLIENT_ID`, and `AIRBYTE_CLOUD_CLIENT_SECRET` " + "environment variables." + ), id="env_vars", ), ], diff --git a/tests/unit_tests/test_exceptions.py b/tests/unit_tests/test_exceptions.py index 57545b4cb..8dd2c84bd 100644 --- a/tests/unit_tests/test_exceptions.py +++ b/tests/unit_tests/test_exceptions.py @@ -2,8 +2,13 @@ from __future__ import annotations import inspect -import pytest import airbyte.exceptions as exceptions_module +import pytest + +from airbyte.exceptions import ( + AirbyteMissingWorkspaceContextError, + AirbyteNoCloudCredentialsError, +) def test_exceptions(): @@ -27,5 +32,109 @@ def test_exceptions(): assert name.endswith("Error") +@pytest.mark.parametrize( + ("hosted", "allow_bearer", "env_vars", "expected_guidance"), + [ + pytest.param( + True, + True, + True, + "Provide a bearer token via the `Authorization` header, or client credentials " + "via the `X-Airbyte-Cloud-Client-Id` and `X-Airbyte-Cloud-Client-Secret` headers.", + id="hosted_with_bearer", + ), + pytest.param( + True, + False, + False, + "Provide client credentials via the `X-Airbyte-Cloud-Client-Id` and " + "`X-Airbyte-Cloud-Client-Secret` headers.", + id="hosted_client_credentials_only", + ), + pytest.param( + False, + True, + True, + "Provide `bearer_token`, or both `client_id` and `client_secret`, as arguments " + "or via the `AIRBYTE_CLOUD_BEARER_TOKEN`, `AIRBYTE_CLOUD_CLIENT_ID`, and " + "`AIRBYTE_CLOUD_CLIENT_SECRET` environment variables.", + id="local_with_bearer", + ), + pytest.param( + False, + False, + True, + "Provide both `client_id` and `client_secret`, as arguments or via the " + "`AIRBYTE_CLOUD_CLIENT_ID` and `AIRBYTE_CLOUD_CLIENT_SECRET` environment variables.", + id="local_client_credentials_only", + ), + pytest.param( + False, + True, + False, + "Provide `bearer_token`, or both `client_id` and `client_secret`.", + id="local_without_env_vars", + ), + pytest.param( + False, + False, + False, + "Provide both `client_id` and `client_secret`.", + id="local_client_credentials_only_without_env_vars", + ), + ], +) +def test_cloud_credentials_error_guidance( + hosted: bool, + allow_bearer: bool, + env_vars: bool, + expected_guidance: str, +) -> None: + """Render cloud credential guidance for each supported mode.""" + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(exceptions_module, "is_hosted_mcp_mode", lambda: hosted) + error = AirbyteNoCloudCredentialsError( + _allow_bearer=allow_bearer, + _env_vars=env_vars, + ) + + assert error.get_message() == "No Airbyte credentials found." + assert error.guidance == expected_guidance + assert "Allow Bearer" not in str(error) + assert "Env Vars" not in str(error) + + +@pytest.mark.parametrize( + ("hosted", "expected_guidance"), + [ + pytest.param( + True, + "Provide the workspace ID via the `X-Airbyte-Workspace-Id` header, or pass " + "the `workspace_id` parameter.", + id="hosted", + ), + pytest.param( + False, + "Set the `AIRBYTE_CLOUD_WORKSPACE_ID` environment variable, or pass the " + "`workspace_id` parameter.", + id="local", + ), + ], +) +def test_missing_workspace_context_error_guidance( + hosted: bool, + expected_guidance: str, +) -> None: + """Render workspace guidance for each supported mode.""" + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(exceptions_module, "is_hosted_mcp_mode", lambda: hosted) + error = AirbyteMissingWorkspaceContextError() + + assert error.get_message() == "Workspace ID is required but not provided." + assert error.guidance == expected_guidance + assert "Allow Bearer" not in str(error) + assert "Env Vars" not in str(error) + + if __name__ == "__main__": pytest.main()