feat(mcp): let agents discover their own Cloud org and workspace IDs - #1094
feat(mcp): let agents discover their own Cloud org and workspace IDs#1094Aaron ("AJ") Steers (aaronsteers) wants to merge 2 commits into
Conversation
Co-Authored-By: AJ Steers <aj@airbyte.io>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
👋 Greetings, Airbyte Team Member!Here are some helpful tips and reminders for your convenience. 💡 Show Tips and TricksTesting This PyAirbyte VersionYou can test this version of PyAirbyte using the following: # Run PyAirbyte CLI from this branch:
uvx --from 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1785180153-cloud-mcp-org-discovery' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1785180153-cloud-mcp-org-discovery'PR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
📚 Show Repo GuidanceHelpful ResourcesCommunity SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
There was a problem hiding this comment.
Pull request overview
This PR exposes Airbyte Cloud organization/workspace discovery to MCP agents so they can determine their own org/workspace IDs using existing Cloud API capabilities, instead of requiring users to manually provide UUIDs.
Changes:
- Add
CloudClient.list_organizations()and refactorget_organization()to reuse it. - Add a new read-only MCP tool
list_cloud_organizations, and enhancelist_cloud_workspacesto support org-less discovery plus structured{workspaces, message}results. - Update server/auth/workspace-missing guidance text and unit tests to direct agents to discovery-first flows and to avoid auto-selecting among multiple candidates.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
airbyte/cloud/client.py |
Adds list_organizations() and reuses it in get_organization(). |
airbyte/mcp/cloud.py |
Introduces result models for discovery, adds list_cloud_organizations, and updates list_cloud_workspaces behavior + auth degradation. |
airbyte/mcp/server.py |
Updates MCP server instructions to recommend discovery-first and explicit user selection. |
airbyte/exceptions.py |
Updates missing-workspace guidance to instruct agents to run discovery tools first. |
tests/unit_tests/test_cloud_credentials.py |
Adds unit coverage for organization discovery and new workspace discovery behavior/error degradation. |
tests/unit_tests/test_exceptions.py |
Updates assertions to match the revised guidance strings. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| CloudWorkspaceResult( | ||
| workspace_id=ws.workspace_id, | ||
| workspace_name=ws.name, | ||
| organization_id=ws.organization_id or "", | ||
| ) |
There was a problem hiding this comment.
Agreed, fixing. Making CloudWorkspaceResult.organization_id nullable and returning None when the org is unknown, rather than coercing to "". An independent end-to-end run against Cloud flagged the same thing: the no-arg public-API path returned organization_id: "" while the org-scoped path returned the real ID, which is exactly the ambiguity you describe.
| message=( | ||
| "Workspace discovery is unavailable because these credentials do not " | ||
| "have permission to list workspaces. Provide a workspace ID or use " | ||
| "credentials with workspace access." | ||
| ), |
There was a problem hiding this comment.
Good catch, fixing. That except block does wrap _resolve_organization_id(), so an org-resolution 403 currently reports as a workspace-permission failure. Broadening the message to cover both organization and workspace discovery, and folding it into a shared helper so the two discovery tools can't drift.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughCloudClient now exposes typed organization discovery. MCP organization and workspace tools return structured results with selection guidance and permission-error messages, while server instructions and workspace-context errors require explicit discovery before choosing a workspace. ChangesCloud discovery flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MCPClient
participant CloudClient
participant OrganizationsTool
participant WorkspacesTool
MCPClient->>OrganizationsTool: discover organizations
OrganizationsTool->>CloudClient: list_organizations()
CloudClient-->>OrganizationsTool: CloudOrganization objects
OrganizationsTool-->>MCPClient: organizations and optional message
MCPClient->>WorkspacesTool: discover workspaces
WorkspacesTool->>CloudClient: list workspaces across organizations
CloudClient-->>WorkspacesTool: workspace results or permission error
WorkspacesTool-->>MCPClient: workspaces and selection guidance
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
airbyte/mcp/cloud.py (1)
1406-1417: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate permission-error handling between the two discovery tools.
list_cloud_workspacesandlist_cloud_organizationsrepeat the sameexcept AirbyteError/ status-code check / empty-result construction pattern. Might be worth extracting a small shared helper (e.g. returning the permission message orNone) to avoid re-copying this for future discovery tools — wdyt?Also applies to: 1450-1461
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@airbyte/mcp/cloud.py` around lines 1406 - 1417, Extract the duplicated AirbyteError permission handling from list_cloud_workspaces and list_cloud_organizations into a small shared helper that checks for HTTPStatus.UNAUTHORIZED or HTTPStatus.FORBIDDEN and returns the common permission-unavailable result or message. Update both discovery tools to reuse that helper while preserving re-raising for other AirbyteError statuses and the existing empty-result response.airbyte/mcp/server.py (1)
116-119: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMinor: awkward phrasing "and then ... first".
Small wording nit in agent-facing instructions — could read more cleanly. wdyt?
✏️ Suggested rewording
- AIRBYTE_CLOUD_WORKSPACE_ID for local or stdio connections). When no organization - or workspace ID is configured, call list_cloud_organizations and then - list_cloud_workspaces first. If multiple organizations or workspaces are - returned, ask the user to choose explicitly; never select automatically. + 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.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@airbyte/mcp/server.py` around lines 116 - 119, Rewrite the agent-facing instruction around list_cloud_organizations and list_cloud_workspaces to state clearly that when neither ID is configured, call list_cloud_organizations followed by list_cloud_workspaces. Preserve the requirement to ask the user when multiple results exist and never select automatically.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@airbyte/mcp/cloud.py`:
- Around line 1384-1389: Update the docstring for the workspace-listing tool to
avoid promising cross-organization results whenever organization arguments are
omitted. Clarify that omitted arguments use the client's configured default
organization when one exists, while the public cross-organization listing is
used only when neither explicit nor default organization configuration is
present. Keep the existing behavior and references to
CloudClient.list_workspaces() unchanged.
---
Nitpick comments:
In `@airbyte/mcp/cloud.py`:
- Around line 1406-1417: Extract the duplicated AirbyteError permission handling
from list_cloud_workspaces and list_cloud_organizations into a small shared
helper that checks for HTTPStatus.UNAUTHORIZED or HTTPStatus.FORBIDDEN and
returns the common permission-unavailable result or message. Update both
discovery tools to reuse that helper while preserving re-raising for other
AirbyteError statuses and the existing empty-result response.
In `@airbyte/mcp/server.py`:
- Around line 116-119: Rewrite the agent-facing instruction around
list_cloud_organizations and list_cloud_workspaces to state clearly that when
neither ID is configured, call list_cloud_organizations followed by
list_cloud_workspaces. Preserve the requirement to ask the user when multiple
results exist and never select automatically.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 58bc6da3-791f-4587-bb90-677f28723264
📒 Files selected for processing (6)
airbyte/cloud/client.pyairbyte/exceptions.pyairbyte/mcp/cloud.pyairbyte/mcp/server.pytests/unit_tests/test_cloud_credentials.pytests/unit_tests/test_exceptions.py
End-to-end MCP verification (headless, real server over stdio)Ran the actual MCP server ( Bare credentials → org → workspace → workspace-scoped call, no human inputBaseline contrast on
|
Code Coverage OverviewLanguages: Python Python / code-coverage/pytest-fastThe overall coverage in commit 7c9d9a7 in the Show a code coverage summary of the most impacted files.
Python / code-coverage/pytest-no-credsThe overall coverage in commit 7c9d9a7 in the Show a code coverage summary of the most impacted files.
Python / code-coverage/pytestThe overall coverage in commit 7c9d9a7 in the Show a code coverage summary of the most impacted files.
Updated |
Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
airbyte/mcp/cloud.py:261
CloudWorkspaceListResult.messageis used both when discovery returns no results and when discovery is unavailable due to 401/403 (permission) errors, but the field docstring only mentions the "no results" case. Update the docstring so it matches actual behavior and helps tool consumers interpret permission-denied responses.
message: str | None = None
"""Additional guidance when discovery returns no results."""
airbyte/mcp/cloud.py:237
- Now that
CloudWorkspaceResult.organization_idis nullable, returning placeholder strings (e.g. incheck_airbyte_cloud_workspaceit currently uses "[unavailable - requires ORGANIZATION_READER permission]") conflicts with the field’s meaning and can look like a real org ID to agents/consumers. Consider returningNonewhen the org is unavailable and usingmessage/separate fields for the unavailability reason.
organization_id: str | None
"""ID of the organization, if known and available."""
airbyte/mcp/cloud.py:88
- In
_handle_discovery_permission_error, usingraise errorwill re-raise the exception but can lose the original traceback context compared to a bareraise. Preferraisehere so unexpected errors preserve the original stack trace for debugging.
status_code = (error.context or {}).get("status_code")
if status_code not in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}:
raise error
Summary
Agents connecting to Cloud MCP with only client credentials had no way to learn their own organization or workspace ID, so they had to stop and ask the user for a UUID. The API could already answer the question —
GET /organizationsandGET /workspacesboth succeed with no ID supplied — but the MCP surface never exposed it:list_cloud_workspaceshard-required an org selector and its docstring said it "will NOT list workspaces across all organizations", and there was no organization-listing tool at all.This PR exposes the existing capability rather than adding new API calls:
CloudClient.list_organizations()returning typedCloudOrganizations, backed by the existingapi_util.list_organizations_for_user().get_organization()now filters that list instead of duplicating the call — behavior is unchanged sinceCloudOrganizationalready scopes its own credentials toorganization_idin its constructor.list_cloud_organizationsMCP tool taking no inputs.list_cloud_workspacesnow takes the org-less public-API path when neitherorganization_idnororganization_nameis given, and keeps the Config APIlist_by_organization_idpath when one is:Both discovery tools return a result model with a
messagefield rather than raising on the two expected dead ends — zero results, and 401/403 from under-permissioned (e.g. workspace-scoped) credentials that can't read organizations. Anything else still raises.Guidance in
MCP_SERVER_INSTRUCTIONS,CLOUD_AUTH_TIP_TEXT, andAirbyteMissingWorkspaceContextErrornow tells agents to discover first and to ask the user to choose when discovery returns multiple candidates — never auto-select. No session/selection state is added here; this is read-only discovery only.Verified live against Airbyte Cloud with sandbox client credentials and no org/workspace configured:
list_cloud_organizations({})returnedAirbyte Team, andlist_cloud_workspaces({})returnedDevin's Sandbox.Link to Devin session: https://app.devin.ai/sessions/0f9b254d63ef4705a5dbdcc22bc40d7f
Requested by: Aaron ("AJ") Steers (@aaronsteers)
Summary by CodeRabbit
New Features
Bug Fixes
Tests