feat(mcp): Add optional workspace_id parameter to cloud tools with conditional schema hiding#876
Conversation
…nditional schema hiding - Add workspace_id parameter to all 22 cloud MCP tools - Modify _get_cloud_workspace() to accept optional workspace_id override - Implement exclude_args pattern in register_tools() to hide workspace_id when AIRBYTE_CLOUD_WORKSPACE_ID env var is set - When env var is set, workspace_id param is hidden from MCP schema (uses env var) - When env var is not set, workspace_id param is exposed (must be provided in tool calls) Closes #875 Co-Authored-By: AJ Steers <aj@airbyte.io>
Original prompt from AJ Steers |
🤖 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. Testing 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/1764107552-mcp-workspace-id-param' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1764107552-mcp-workspace-id-param'Helpful ResourcesPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
Community SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
📝 WalkthroughWalkthroughAdded optional workspace_id support across cloud MCP operations and updated tool registration to conditionally hide the workspace_id argument for cloud tools when an AIRBYTE_CLOUD_WORKSPACE_ID environment variable is set. Changes
Sequence Diagram(s)sequenceDiagram
participant Registrar as register_tools
participant Inspect as inspect
participant App as app.tool
rect rgb(240,248,255)
Note right of Registrar: Registering cloud tools\n(per-tool signature inspection)
end
Registrar->>Inspect: inspect.signature(callable_fn)
Inspect-->>Registrar: parameter list
alt AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET == True
Registrar->>Registrar: exclude_args = ["workspace_id"]
else AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET == False
Registrar->>Registrar: exclude_args = None
end
Registrar->>App: app.tool(callable_fn, annotations=..., exclude_args=exclude_args)
App-->>Registrar: registration complete
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Would you like me to suggest a small test matrix (env var set/unset × a representative cloud tool) to validate registration and runtime behavior, wdyt? Possibly related PRs
Pre-merge checks and finishing touches✅ Passed checks (5 passed)
✨ Finishing touches
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
airbyte/mcp/cloud_ops.py (2)
540-565: Consider adding error handling for workspace retrieval in list functions.I noticed that
list_deployed_cloud_source_connectors(and similar list functions) don't wrap the_get_cloud_workspacecall in a try/except block, unlike the deploy functions. If the workspace retrieval fails (e.g., invalid workspace_id), it would raise an unhandled exception.This might be intentional for read-only operations, but would you consider adding consistent error handling for a better user experience, wdyt?
410-421: Minor style inconsistency in parameter annotations.I noticed that
nameanduniqueparameters lack theAnnotated[..., Field(...)]wrapper that other tools use consistently, while the newworkspace_idparameter does have it. This is pre-existing, but would you consider aligning the style for consistency across tools, wdyt?def deploy_noop_destination_to_cloud( - name: str = "No-op Destination", + name: Annotated[ + str, + Field( + description="The name to use when deploying the destination.", + default="No-op Destination", + ), + ] = "No-op Destination", *, workspace_id: Annotated[ str | None, Field( description="Workspace ID. Defaults to AIRBYTE_CLOUD_WORKSPACE_ID env var.", default=None, ), ] = None, - unique: bool = True, + unique: Annotated[ + bool, + Field( + description="Whether to require a unique name.", + default=True, + ), + ] = True, ) -> str:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
airbyte/mcp/_tool_utils.py(3 hunks)airbyte/mcp/cloud_ops.py(41 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
airbyte/mcp/cloud_ops.py (3)
tests/integration_tests/cloud/conftest.py (1)
workspace_id(49-50)airbyte/cloud/workspaces.py (2)
CloudWorkspace(64-612)list_custom_source_definitions(554-582)airbyte/cloud/auth.py (1)
resolve_cloud_workspace_id(36-41)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (All, Python 3.11, Windows)
- GitHub Check: Pytest (Fast)
- GitHub Check: Pytest (No Creds)
🔇 Additional comments (9)
airbyte/mcp/_tool_utils.py (3)
10-10: LGTM!The
inspectimport is appropriately added to support the new signature inspection logic for excluding parameters.
29-29: LGTM!Clean boolean flag derived from the environment variable. The
strip()andbool()pattern correctly handles empty strings and whitespace.
167-179: Excellent! I have completed the verification. The web search and codebase analysis confirm that:
- ✅
exclude_argsis a valid FastMCP parameter - Confirmed by FastMCP documentation- ✅ FastMCP version supports it - Project uses
fastmcp >=2.11.3,<3.0.0(exclude_args is deprecated in 2.14+, still works)- ✅
workspace_idhas a default value - Foundworkspace_id: str | None = None, which satisfies FastMCP's constraint that only arguments with defaults can be excluded- ✅ Implementation is correct - The code properly excludes workspace_id only when domain is "cloud" AND the environment variable is set
exclude_argsparameter is valid and properly implemented.The verification confirms that FastMCP's
app.tool()API fully supports theexclude_argsparameter. The code correctly passesNonewhen there are no args to exclude, and properly excludesworkspace_idonly when both conditions are met (domain == "cloud" and AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET). Theworkspace_idparameter has a default value (None), satisfying FastMCP's constraint that only arguments with defaults can be excluded.Minor note: The list comprehension on line 171 iterates over a single-element list
["workspace_id"], which works but could be simplified to a direct membership check if preferred.airbyte/mcp/cloud_ops.py (6)
66-78: LGTM on the helper function update!The
_get_cloud_workspacefunction cleanly handles the optionalworkspace_idby delegating toresolve_cloud_workspace_id, which per the relevant code snippets handles the env var fallback. Good encapsulation.
95-101: Consistent parameter pattern.The
workspace_idparameter follows the established pattern with properAnnotatedtype hint andFielddescription. Nice work keeping it uniform across all the tools!
173-179: Consistent workspace_id propagation across deployment and sync tools.The pattern is uniformly applied to
deploy_destination_to_cloud,create_connection_on_cloud, andrun_cloud_sync. Each correctly passes the optionalworkspace_idthrough to_get_cloud_workspace.Also applies to: 265-271, 318-324
374-383: Nice addition of workspace override capability to the health check.Adding
workspace_idtocheck_airbyte_cloud_workspaceallows users to verify connectivity to a specific workspace without relying on the environment variable. The keyword-only signature with*is appropriate here.
989-990: Safe mode check ordering is correct.The
check_guid_created_in_sessionis called before_get_cloud_workspace, which is the right order - it validates permission before incurring the cost of workspace authentication. Good defensive pattern!
1316-1326: Overall implementation looks solid!The
workspace_idparameter has been consistently added across all cloud operations, and theregister_cloud_ops_toolsfunction correctly delegates toregister_toolswhich now handles the conditional parameter exclusion. This is a clean implementation of per-workspace context support.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
airbyte/mcp/cloud_ops.py (1)
95-101: Consider extracting a type alias to reduce repetition, wdyt?The
workspace_idannotation is repeated identically 22 times. You could define a type alias at the module level to DRY this up:WorkspaceIdParam = Annotated[ str | None, Field( description="Workspace ID. Defaults to AIRBYTE_CLOUD_WORKSPACE_ID env var.", default=None, ), ]Then each function signature becomes simply
workspace_id: WorkspaceIdParam. This would make future changes (e.g., updating the description) a single-line edit rather than 22.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
airbyte/mcp/cloud_ops.py(41 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
airbyte/mcp/cloud_ops.py (3)
tests/integration_tests/cloud/conftest.py (1)
workspace_id(49-50)airbyte/cloud/workspaces.py (2)
CloudWorkspace(64-612)list_custom_source_definitions(554-582)airbyte/cloud/auth.py (1)
resolve_cloud_workspace_id(36-41)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Windows)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (Fast)
- GitHub Check: Pytest (No Creds)
🔇 Additional comments (3)
airbyte/mcp/cloud_ops.py (3)
66-78: LGTM! Clean implementation of workspace_id threading.The helper function correctly passes the optional
workspace_idtoresolve_cloud_workspace_id, which handles the fallback to the environment variable. Nice and clean!
374-403: Nice backward-compatible addition!The
check_airbyte_cloud_workspacefunction signature change is clean - it now accepts an optionalworkspace_idwhile maintaining backward compatibility for existing callers who rely on the environment variable.
1316-1326: Theexclude_argsimplementation in_tool_utils.pyis correctly handling workspace_id filtering.The code at lines 167-172 properly implements the conditional hiding of
workspace_idfrom the MCP schema:
- Line 29 defines
AIRBYTE_CLOUD_WORKSPACE_ID_IS_SETwhich checks if theAIRBYTE_CLOUD_WORKSPACE_IDenvironment variable is set- Lines 169-172 conditionally build an
exclude_argslist that filters outworkspace_idonly when the env var is present and the domain is "cloud"- Line 178 passes
exclude_argstoapp.tool(), which FastMCP uses to hide the parameter from the schemaThis implementation ensures that when
AIRBYTE_CLOUD_WORKSPACE_IDis configured, theworkspace_idparameter is automatically excluded from the MCP tool schema, while still allowing it to be passed programmatically when the function is called. The logic correctly handles all 20+ cloud operation functions that now have theworkspace_idparameter.
Aaron ("AJ") Steers (aaronsteers)
marked this pull request as ready for review
Aaron ("AJ") Steers (aaronsteers)
changed the title
PyTest Results (Full)389 tests ±0 372 ✅ - 1 25m 0s ⏱️ - 3m 14s For more details on these failures, see this check. Results for commit 4203f57. ± Comparison against base commit f1cb667. |
Aaron ("AJ") Steers (aaronsteers)
merged commit a8cf82c
into
main
Aaron ("AJ") Steers (aaronsteers)
deleted the
devin/1764107552-mcp-workspace-id-param
branch
1 participant
feat(mcp): Add optional workspace_id parameter to cloud tools with conditional schema hiding
Summary
This PR adds an optional
workspace_idparameter to all 22 cloud MCP tools, enabling workspace context switching for admin/ops use cases. The parameter is conditionally hidden from the MCP schema based on whether theAIRBYTE_CLOUD_WORKSPACE_IDenvironment variable is set:workspace_idparam is hidden from MCP schema (uses env var value)workspace_idparam is exposed and must be provided in tool callsThis follows the same
exclude_argspattern used in connector-builder-mcp for themanifestparameter.Changes:
_get_cloud_workspace()to accept optionalworkspace_idoverrideworkspace_idparameter to all 22 cloud tools incloud_ops.pyexclude_argslogic inregister_tools()in_tool_utils.pyCloses #875
Review & Testing Checklist for Human
exclude_argsworks with FastMCP: Test that whenAIRBYTE_CLOUD_WORKSPACE_IDis set, theworkspace_idparameter does NOT appear in the MCP tool schemaworkspace_idin a tool call correctly overrides the workspace contextresolve_cloud_workspace_id()accepts input: Confirm the underlying function inairbyte/cloud/auth.pycorrectly handles the optionalinput_valueparameterworkspace_idis not providedRecommended test plan:
AIRBYTE_CLOUD_WORKSPACE_IDenv var, start MCP server, verifyworkspace_idis NOT in tool schemasworkspace_idIS in tool schemasworkspace_idand verify it worksNotes
workspace_idparameter annotation is repeated 22 times with identical structureSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.
Important
Auto-merge enabled.
This PR is set to merge automatically when all requirements are met.
Note
Auto-merge may have been disabled. Please check the PR status to confirm.