diff --git a/airbyte/cloud/client.py b/airbyte/cloud/client.py index ced3c7ba7..797faf250 100644 --- a/airbyte/cloud/client.py +++ b/airbyte/cloud/client.py @@ -285,6 +285,27 @@ def name_filter(workspace_name: str) -> bool: ) ] + def list_organizations(self) -> list[CloudOrganization]: + """List all organizations available to this client.""" + return [ + CloudOrganization( + organization_id=organization.organization_id, + organization_name=organization.organization_name, + email=organization.email, + client_id=self.client_id, + client_secret=self.client_secret, + bearer_token=self.bearer_token, + public_api_root=self.public_api_root, + config_api_root=self.config_api_root, + ) + for organization in api_util.list_organizations_for_user( + api_root=self.public_api_root, + client_id=self.client_id, + client_secret=self.client_secret, + bearer_token=self.bearer_token, + ) + ] + def get_organization( self, organization_id: str | None = None, @@ -302,12 +323,7 @@ def get_organization( message="Organization ID or organization name is required." ) - organizations = api_util.list_organizations_for_user( - api_root=self.public_api_root, - client_id=self.client_id, - client_secret=self.client_secret, - bearer_token=self.bearer_token, - ) + organizations = self.list_organizations() if resolved_organization_id: matching_organizations = [ organization @@ -332,18 +348,4 @@ def get_organization( context={"organization_name": organization_name}, ) - organization = matching_organizations[0] - - organization_credentials = self._credentials.with_organization_id( - organization.organization_id - ) - return CloudOrganization( - organization_id=organization.organization_id, - organization_name=organization.organization_name, - email=organization.email, - client_id=organization_credentials.client_id, - client_secret=organization_credentials.client_secret, - bearer_token=organization_credentials.bearer_token, - public_api_root=organization_credentials.public_api_root, - config_api_root=organization_credentials.config_api_root, - ) + return matching_organizations[0] diff --git a/airbyte/exceptions.py b/airbyte/exceptions.py index 96117b140..413fbbfe9 100644 --- a/airbyte/exceptions.py +++ b/airbyte/exceptions.py @@ -287,13 +287,17 @@ def __post_init__(self) -> 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." + "Call `list_cloud_organizations` and `list_cloud_workspaces` first. " + "If discovery returns exactly one workspace, provide its ID via the " + f"`{MCP_WORKSPACE_ID_HEADER}` header or the `workspace_id` parameter; " + "otherwise ask the user to choose." ) else: self.guidance = ( - f"Set the `{CLOUD_WORKSPACE_ID_ENV_VAR}` environment variable, or pass " - "the `workspace_id` parameter." + "Call `list_cloud_organizations` and `list_cloud_workspaces` first. " + "If discovery returns exactly one workspace, set its ID in " + f"`{CLOUD_WORKSPACE_ID_ENV_VAR}` or pass the `workspace_id` parameter; " + "otherwise ask the user to choose." ) diff --git a/airbyte/mcp/cloud.py b/airbyte/mcp/cloud.py index eb6a1925f..094874d1e 100644 --- a/airbyte/mcp/cloud.py +++ b/airbyte/mcp/cloud.py @@ -10,8 +10,10 @@ # tool / helper definitions as a redundant "API Documentation" list. __all__: list[str] = [] +from collections.abc import Callable +from http import HTTPStatus from pathlib import Path -from typing import Annotated, Any, Literal, cast +from typing import Annotated, Any, Literal, TypeVar, cast from fastmcp import Context, FastMCP from fastmcp_extensions import get_mcp_config, mcp_tool, register_mcp_tools @@ -42,6 +44,7 @@ ) from airbyte.destinations.util import get_noop_destination from airbyte.exceptions import ( + AirbyteError, AirbyteMissingResourceError, AirbyteMissingWorkspaceContextError, PyAirbyteInputError, @@ -57,12 +60,13 @@ CLOUD_AUTH_TIP_TEXT = ( 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"`{MCP_CLIENT_ID_HEADER}` and `{MCP_CLIENT_SECRET_HEADER}` headers. To discover " + f"available organizations and workspaces, call `list_cloud_organizations` and " + f"`list_cloud_workspaces` before asking the user for an ID. 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." + f"`{CLOUD_CLIENT_SECRET_ENV_VAR}`. If discovery returns multiple candidates, " + f"ask the user to choose one; do not select automatically." ) WORKSPACE_ID_TIP_TEXT = ( f"Workspace ID. Hosted MCP connections pass it via the " @@ -70,6 +74,24 @@ f"`{CLOUD_WORKSPACE_ID_ENV_VAR}` environment variable." ) +_DiscoveryResult = TypeVar("_DiscoveryResult") + + +def _handle_discovery_permission_error( + error: AirbyteError, + *, + make_result: Callable[[str], _DiscoveryResult], +) -> _DiscoveryResult: + """Return a graceful result for discovery permission errors.""" + status_code = (error.context or {}).get("status_code") + if status_code not in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: + raise error + return make_result( + "Organization or workspace discovery is unavailable because these credentials " + "do not have the required permission or access. Provide an organization or " + "workspace ID, or use credentials with the needed access." + ) + class CloudSourceResult(BaseModel): """Information about a deployed source connector in Airbyte Cloud.""" @@ -192,6 +214,16 @@ class CloudOrganizationResult(BaseModel): Defaults to False unless we have affirmative evidence of a locked state.""" +class CloudOrganizationListResult(BaseModel): + """Result of discovering organizations in Airbyte Cloud.""" + + organizations: list[CloudOrganizationResult] + """Organizations visible to the authenticated credentials.""" + + message: str | None = None + """Additional guidance when discovery returns no results or is unavailable.""" + + class CloudWorkspaceResult(BaseModel): """Information about a workspace in Airbyte Cloud.""" @@ -201,8 +233,8 @@ class CloudWorkspaceResult(BaseModel): """Display name of the workspace.""" workspace_url: str | None = None """URL to access the workspace in Airbyte Cloud.""" - organization_id: str - """ID of the organization (requires ORGANIZATION_READER permission).""" + organization_id: str | None + """ID of the organization, if known and available.""" organization_name: str | None = None """Name of the organization (requires ORGANIZATION_READER permission).""" payment_status: str | None = None @@ -219,6 +251,16 @@ class CloudWorkspaceResult(BaseModel): Requires ORGANIZATION_READER permission.""" +class CloudWorkspaceListResult(BaseModel): + """Result of discovering workspaces in Airbyte Cloud.""" + + workspaces: list[CloudWorkspaceResult] + """Workspaces visible to the authenticated credentials.""" + + message: str | None = None + """Additional guidance when discovery returns no results.""" + + class LogReadResult(BaseModel): """Result of reading sync logs with pagination support.""" @@ -1332,16 +1374,14 @@ def list_cloud_workspaces( organization_id: Annotated[ str | None, Field( - description="Organization ID. Required if organization_name is not provided.", + description="Optional organization ID to list workspaces within.", default=None, ), ], organization_name: Annotated[ str | None, Field( - description=( - "Organization name (exact match). " "Required if organization_id is not provided." - ), + description=("Optional organization name (exact match) to list workspaces within."), default=None, ), ], @@ -1359,35 +1399,99 @@ def list_cloud_workspaces( default=None, ), ], -) -> list[CloudWorkspaceResult]: - """List all workspaces in a specific organization. +) -> CloudWorkspaceListResult: + """List all workspaces visible to the authenticated credentials. - Requires either organization_id OR organization_name (exact match) to be provided. - This tool will NOT list workspaces across all organizations - you must specify - which organization to list workspaces from. + When an organization ID or exact organization name is provided, the Config API + lists workspaces in that organization. When neither is provided and the client + has no default organization, the public API lists workspaces across organizations + visible to the current credentials. Otherwise, results are scoped to the client's + default organization. """ client = _get_cloud_client(ctx) - resolved_org_id = _resolve_organization_id( - organization_id=organization_id, - organization_name=organization_name, - client=client, - ) - - workspaces = client.list_workspaces( - organization_id=resolved_org_id, - name_contains=name_contains, - limit=limit, - ) + try: + workspaces = client.list_workspaces( + organization_id=( + _resolve_organization_id( + organization_id=organization_id, + organization_name=organization_name, + client=client, + ) + if organization_id is not None or organization_name is not None + else None + ), + name_contains=name_contains, + limit=limit, + ) + except AirbyteError as error: + return _handle_discovery_permission_error( + error, + make_result=lambda message: CloudWorkspaceListResult( + workspaces=[], + message=message, + ), + ) - return [ + results = [ CloudWorkspaceResult( workspace_id=ws.workspace_id, workspace_name=ws.name, - organization_id=ws.organization_id or "", + organization_id=ws.organization_id, ) for ws in workspaces ] + return CloudWorkspaceListResult( + workspaces=results, + message=( + "No workspaces were returned for these credentials. Verify the " + "credentials or ask the user to provide a workspace ID." + if not results + else None + ), + ) + + +@mcp_tool( + read_only=True, + idempotent=True, + open_world=True, + extra_help_text=CLOUD_AUTH_TIP_TEXT, +) +def list_cloud_organizations( + ctx: Context, +) -> CloudOrganizationListResult: + """List organizations visible to the authenticated Airbyte Cloud credentials.""" + try: + organizations = _get_cloud_client(ctx).list_organizations() + except AirbyteError as error: + return _handle_discovery_permission_error( + error, + make_result=lambda message: CloudOrganizationListResult( + organizations=[], + message=message, + ), + ) + + if not organizations: + return CloudOrganizationListResult( + organizations=[], + message=( + "No organizations were returned for these credentials. Verify the " + "credentials or ask the user to provide an organization ID." + ), + ) + + return CloudOrganizationListResult( + organizations=[ + CloudOrganizationResult( + id=organization.organization_id, + name=organization.organization_name or "", + email=organization.email or "", + ) + for organization in organizations + ] + ) @mcp_tool( diff --git a/airbyte/mcp/server.py b/airbyte/mcp/server.py index 0f6474193..07f6c806e 100644 --- a/airbyte/mcp/server.py +++ b/airbyte/mcp/server.py @@ -113,7 +113,10 @@ - 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) + AIRBYTE_CLOUD_WORKSPACE_ID for local or stdio connections). When no organization + or workspace ID is configured, first call list_cloud_organizations, then + list_cloud_workspaces. If multiple organizations or workspaces are + returned, ask the user to choose explicitly; never select automatically. - 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 b23ba0390..a05ddb2ae 100644 --- a/tests/unit_tests/test_cloud_credentials.py +++ b/tests/unit_tests/test_cloud_credentials.py @@ -11,7 +11,11 @@ from airbyte.cloud.models import CloudWorkspaceInfo from airbyte.cloud.organizations import CloudOrganization from airbyte.cloud.workspaces import CloudWorkspace -from airbyte.exceptions import AirbyteMissingResourceError, PyAirbyteInputError +from airbyte.exceptions import ( + AirbyteError, + AirbyteMissingResourceError, + PyAirbyteInputError, +) from airbyte.mcp import cloud as mcp_cloud from airbyte.secrets.base import SecretString @@ -484,6 +488,191 @@ def test_cloud_client_get_organization_uses_default_organization_id( assert organization.organization_id == "default-org" +@pytest.mark.parametrize( + "organizations", + [ + pytest.param([], id="empty"), + pytest.param( + [ + models.OrganizationResponse( + organization_id="organization-id", + organization_name="Organization", + email="test@example.com", + ) + ], + id="single", + ), + pytest.param( + [ + models.OrganizationResponse( + organization_id="organization-id-1", + organization_name="Organization 1", + email="one@example.com", + ), + models.OrganizationResponse( + organization_id="organization-id-2", + organization_name="Organization 2", + email="two@example.com", + ), + ], + id="multiple", + ), + ], +) +def test_cloud_client_list_organizations_returns_typed_resources( + monkeypatch: pytest.MonkeyPatch, + organizations: list[models.OrganizationResponse], +) -> None: + monkeypatch.setattr( + api_util, + "list_organizations_for_user", + lambda **_: organizations, + ) + + result = CloudClient(bearer_token="token").list_organizations() + + assert all(isinstance(organization, CloudOrganization) for organization in result) + assert [organization.organization_id for organization in result] == [ + organization.organization_id for organization in organizations + ] + + +def test_cloud_client_get_organization_reuses_list_organizations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + organizations = [ + CloudOrganization( + organization_id="organization-id", + organization_name="Organization", + email="test@example.com", + ) + ] + client = CloudClient(bearer_token="token") + monkeypatch.setattr(client, "list_organizations", lambda: organizations) + + result = client.get_organization(organization_id="organization-id") + + assert result is organizations[0] + + +def test_mcp_list_cloud_organizations_returns_empty_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class EmptyClient: + def list_organizations(self) -> list[CloudOrganization]: + return [] + + monkeypatch.setattr(mcp_cloud, "_get_cloud_client", lambda _: EmptyClient()) + + result = mcp_cloud.list_cloud_organizations(None) + + assert result.organizations == [] + assert result.message + + +@pytest.mark.parametrize("organization_count", [1, 2]) +def test_mcp_list_cloud_organizations_returns_all_results( + monkeypatch: pytest.MonkeyPatch, + organization_count: int, +) -> None: + class OrganizationClient: + def list_organizations(self) -> list[CloudOrganization]: + return [ + CloudOrganization( + organization_id=f"organization-id-{index}", + organization_name=f"Organization {index}", + email=f"test-{index}@example.com", + ) + for index in range(organization_count) + ] + + monkeypatch.setattr(mcp_cloud, "_get_cloud_client", lambda _: OrganizationClient()) + + result = mcp_cloud.list_cloud_organizations(None) + + assert len(result.organizations) == organization_count + assert result.message is None + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_mcp_list_cloud_organizations_degrades_on_auth_error( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +) -> None: + class ForbiddenClient: + def list_organizations(self) -> list[CloudOrganization]: + raise AirbyteError(context={"status_code": status_code}) + + monkeypatch.setattr(mcp_cloud, "_get_cloud_client", lambda _: ForbiddenClient()) + + result = mcp_cloud.list_cloud_organizations(None) + + assert result.organizations == [] + assert "permission" in (result.message or "") + + +def test_mcp_list_cloud_workspaces_uses_public_api_without_organization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured_organization_id: str | None = "unset" + + class WorkspaceClient: + def list_workspaces( + self, *, organization_id: str | None = None, **_: object + ) -> list[CloudWorkspaceInfo]: + nonlocal captured_organization_id + captured_organization_id = organization_id + return [ + CloudWorkspaceInfo( + workspaceId="workspace-id", + name="Workspace", + organizationId="organization-id", + ), + CloudWorkspaceInfo( + workspaceId="workspace-without-org", + name="Workspace without organization", + organizationId=None, + ), + ] + + monkeypatch.setattr(mcp_cloud, "_get_cloud_client", lambda _: WorkspaceClient()) + + result = mcp_cloud.list_cloud_workspaces( + None, + organization_id=None, + organization_name=None, + name_contains=None, + limit=None, + ) + + assert captured_organization_id is None + assert result.workspaces[0].workspace_id == "workspace-id" + assert result.workspaces[1].organization_id is None + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_mcp_list_cloud_workspaces_degrades_on_auth_error( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +) -> None: + class ForbiddenClient: + def list_workspaces(self, **_: object) -> list[CloudWorkspaceInfo]: + raise AirbyteError(context={"status_code": status_code}) + + monkeypatch.setattr(mcp_cloud, "_get_cloud_client", lambda _: ForbiddenClient()) + + result = mcp_cloud.list_cloud_workspaces( + None, + organization_id=None, + organization_name=None, + name_contains=None, + limit=None, + ) + + assert result.workspaces == [] + assert "permission" in (result.message or "") + + def test_mcp_resolve_organization_id_skips_lookup_when_id_provided() -> None: class FailingClient: def get_organization(self, **_: object) -> object: diff --git a/tests/unit_tests/test_exceptions.py b/tests/unit_tests/test_exceptions.py index 8dd2c84bd..633a04f99 100644 --- a/tests/unit_tests/test_exceptions.py +++ b/tests/unit_tests/test_exceptions.py @@ -109,14 +109,18 @@ def test_cloud_credentials_error_guidance( [ pytest.param( True, - "Provide the workspace ID via the `X-Airbyte-Workspace-Id` header, or pass " - "the `workspace_id` parameter.", + "Call `list_cloud_organizations` and `list_cloud_workspaces` first. If " + "discovery returns exactly one workspace, provide its ID via the " + "`X-Airbyte-Workspace-Id` header or the `workspace_id` parameter; " + "otherwise ask the user to choose.", id="hosted", ), pytest.param( False, - "Set the `AIRBYTE_CLOUD_WORKSPACE_ID` environment variable, or pass the " - "`workspace_id` parameter.", + "Call `list_cloud_organizations` and `list_cloud_workspaces` first. If " + "discovery returns exactly one workspace, set its ID in " + "`AIRBYTE_CLOUD_WORKSPACE_ID` or pass the `workspace_id` parameter; " + "otherwise ask the user to choose.", id="local", ), ],