chore(mcp): migrate to fastmcp-extensions library #949
Conversation
- Replace custom @mcp_tool decorator with fastmcp_extensions.mcp_tool - Replace custom @mcp_prompt decorator with fastmcp_extensions.mcp_prompt - Replace custom register_tools() with fastmcp_extensions.register_mcp_tools() - Replace custom register_prompts() with fastmcp_extensions.register_mcp_prompts() - Use mcp_server() factory for FastMCP app creation - Delete _annotations.py (use annotations from fastmcp_extensions) - Refactor _tool_utils.py to only contain safe mode logic - Remove obsolete MCP domain constants from constants.py - Update tests to only test remaining safe mode functionality - Add fastmcp-extensions>=0.2.0 dependency 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/1768864254-migrate-mcp-to-fastmcp-extensions' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1768864254-migrate-mcp-to-fastmcp-extensions'Helpful ResourcesPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
Community SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
PyTest Results (Fast Tests Only, No Creds)326 tests - 22 326 ✅ - 22 5m 39s ⏱️ -7s Results for commit a856bf5. ± Comparison against base commit a7b8c61. This pull request removes 28 and adds 6 tests. Note that renamed tests count towards both.♻️ This comment has been updated with latest results. |
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughReplaces static MCP domain lists and header-based credential resolution with environment/config-driven constants and MCP config APIs; removes FastMCP annotation hints; migrates MCP server, tool, and prompt registration to fastmcp-extensions; adds session-scoped safe-mode GUID tracking, new arg resolvers, secret-loading, renamed exports, tests, and pyproject dependency. (39 words) Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant MCP_Server as mcp_server
participant ToolModule as ToolModule
participant SafeUtils as _tool_utils
Client->>MCP_Server: start (mcp_server init with AIRBYTE_CONFIG_ARGS)
MCP_Server->>ToolModule: register_mcp_tools(mcp_module=...)
ToolModule-->>MCP_Server: tools & prompts registered
Client->>MCP_Server: invoke create-resource
MCP_Server->>ToolModule: call create handler
ToolModule->>SafeUtils: register_guid_created_in_session(guid)
SafeUtils-->>ToolModule: ack
Client->>MCP_Server: invoke destructive-action(target_guid)
MCP_Server->>SafeUtils: check_guid_created_in_session(target_guid)
alt GUID in session or safe-mode disabled
SafeUtils-->>MCP_Server: ok
MCP_Server->>ToolModule: call destructive handler
ToolModule-->>Client: success
else GUID not in session and safe-mode enabled
SafeUtils-->>MCP_Server: SafeModeError
MCP_Server-->>Client: deny operation (error)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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}
|
Devin, please rename the Python files so that we preserve the original module names used for filtering. |
Rename files to preserve backward compatibility with domain filtering: - cloud_ops.py -> cloud.py (domain: 'cloud') - local_ops.py -> local.py (domain: 'local') - connector_registry.py -> registry.py (domain: 'registry') Update all imports and references in: - server.py - __init__.py - test_mcp_connector_registry.py Co-Authored-By: AJ Steers <aj@airbyte.io>
|
@aaronsteers Done! Renamed the files to preserve the original domain names:
Updated all imports and references in |
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
airbyte/mcp/registry.py (1)
2-2: Minor docstring inconsistency.The docstring says "Airbyte Cloud MCP operations" but this file is
registry.pyfor connector registry operations. Should this be "Airbyte Connector Registry MCP operations" instead, wdyt?Suggested fix
-"""Airbyte Cloud MCP operations.""" +"""Airbyte Connector Registry MCP operations."""
🧹 Nitpick comments (1)
airbyte/mcp/server.py (1)
9-12: Missing noqa comments for side-effect imports.Lines 9-11 import modules for their side effects (tool/prompt registration) but only line 12 has the
# noqa: F401comment. To be consistent and avoid future linter warnings, should we add similar comments to lines 9-11, wdyt?Suggested fix
-import airbyte.mcp.cloud -import airbyte.mcp.local -import airbyte.mcp.prompts -import airbyte.mcp.registry # noqa: F401 - Import to register tools +import airbyte.mcp.cloud # noqa: F401 - Import to register tools +import airbyte.mcp.local # noqa: F401 - Import to register tools +import airbyte.mcp.prompts # noqa: F401 - Import to register prompts +import airbyte.mcp.registry # noqa: F401 - Import to register tools
devin-ai-integration
bot
changed the title
There was a problem hiding this comment.
Pull request overview
This PR refactors PyAirbyte's MCP implementation to use the fastmcp-extensions library, eliminating ~450 lines of custom infrastructure code. The migration replaces custom decorators and registration logic with library-provided equivalents while preserving safe mode functionality.
Changes:
- Migrated to
fastmcp-extensionslibrary for MCP tool and prompt decorators - Simplified tool registration using
mcp_server()factory andregister_mcp_tools()/register_mcp_prompts() - Removed custom annotation constants and domain filtering logic (now handled by library)
Reviewed changes
Copilot reviewed 12 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
pyproject.toml |
Added fastmcp-extensions>=0.2.0,<1.0.0 dependency |
airbyte/mcp/server.py |
Replaced custom registration with library's mcp_server() and registration functions |
airbyte/mcp/cloud.py |
Removed domain parameter from @mcp_tool decorators and registration function |
airbyte/mcp/local.py |
Removed domain parameter from @mcp_tool decorators and registration function |
airbyte/mcp/registry.py |
Removed domain parameter from @mcp_tool decorators and registration function |
airbyte/mcp/prompts.py |
Replaced manual prompt registration with @mcp_prompt decorator |
airbyte/mcp/_tool_utils.py |
Retained only safe mode logic; removed custom decorator and registration code |
airbyte/mcp/_annotations.py |
Deleted entire file (annotations now from library) |
airbyte/constants.py |
Removed obsolete MCP domain filtering constants |
airbyte/mcp/__init__.py |
Updated module names from *_ops to match new file names |
tests/unit_tests/test_mcp_tool_utils.py |
Replaced domain filtering tests with safe mode functionality tests |
tests/unit_tests/test_mcp_connector_registry.py |
Updated import path from connector_registry to registry |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
… env vars Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
PyTest Results (Full)396 tests - 22 379 ✅ - 22 23m 58s ⏱️ - 1m 6s Results for commit a856bf5. ± Comparison against base commit a7b8c61. This pull request removes 28 and adds 6 tests. Note that renamed tests count towards both.♻️ This comment has been updated with latest results. |
Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/mcp/server.py`:
- Around line 28-31: Remove the now-unnecessary "# noqa: F401" comments on the
side-effect import lines so Ruff no longer flags them as unused: edit the import
statements referencing airbyte.mcp.cloud, airbyte.mcp.local, and
airbyte.mcp.prompts to remove their trailing "# noqa: F401" directives while
leaving the imports themselves intact (the import of airbyte.mcp.registry can
remain unchanged); ensure these are preserved as side-effect imports in
server.py so registration still occurs.
…ct imports Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 13 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Devin, check latest bot feedback |
Co-Authored-By: AJ Steers <aj@airbyte.io>
…ions utility Co-Authored-By: AJ Steers <aj@airbyte.io>
…fig() directly - Delete header extraction functions from _util.py (resolve_cloud_credentials, resolve_workspace_id, get_*_from_headers, _get_header_value) - Delete header constants from _util.py (HEADER_CLIENT_ID, HEADER_CLIENT_SECRET, HEADER_WORKSPACE_ID, HEADER_API_URL) - Add config args for bearer_token, client_id, client_secret, api_url to server.py - Add MCP config constants and header constants to constants.py - Update _get_cloud_workspace to use get_mcp_config() with Context parameter - Update all cloud.py tools to accept ctx: Context parameter - Update list_cloud_workspaces and describe_cloud_organization to use get_mcp_config() This removes ~215 lines of custom header extraction code in favor of the fastmcp-extensions config arg system. Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/mcp/cloud.py`:
- Around line 1408-1428: The api_url obtained in list_cloud_workspaces (and
similarly in describe_cloud_organization) can be None but is passed to
_resolve_organization_id and api_util.list_workspaces_in_organization which
expect a str; add an explicit fallback like setting api_url = api_url or
api_util.CLOUD_API_ROOT immediately after fetching it so callers
(_resolve_organization_id, api_util.list_workspaces_in_organization) always
receive a non-None api_root consistent with CloudWorkspace's default.
🧹 Nitpick comments (1)
airbyte/mcp/server.py (1)
154-174: Minor consideration: exclude/include filter precedence.The current logic applies
exclude_modulesfirst (lines 166-168), theninclude_modules(lines 170-172). This means if both are set, onlyexclude_modulestakes effect andinclude_modulesis ignored.Is this the intended behavior? Some systems apply both filters (i.e., include first, then exclude from the included set). If this is intentional, might be worth adding a brief comment to clarify the precedence, wdyt?
💡 Optional: Add clarifying comment
def airbyte_module_filter(tool: Tool, app: FastMCP) -> bool: """Filter tools based on legacy AIRBYTE_MCP_DOMAINS and AIRBYTE_MCP_DOMAINS_DISABLED. When AIRBYTE_MCP_DOMAINS_DISABLED is set, hide tools from those modules. When AIRBYTE_MCP_DOMAINS is set, only show tools from those modules. + + Note: If both are set, exclude_modules takes precedence and include_modules is ignored. """
…tils.py Tidies up server.py by moving the config args and filter functions to _tool_utils.py: - Move AIRBYTE_CONFIG_ARGS and all individual config args to _tool_utils.py - Move airbyte_readonly_mode_filter and airbyte_module_filter to _tool_utils.py - Move _parse_csv_config helper function to _tool_utils.py - Update server.py to import these from _tool_utils.py This keeps server.py focused on server initialization while _tool_utils.py contains all the config and filtering logic. Co-Authored-By: AJ Steers <aj@airbyte.io>
The new name better describes the function's side effects - it loads secrets from dotenv files and secret managers into environment variables. This must run before mcp_server() so config args can resolve from the loaded env vars. Co-Authored-By: AJ Steers <aj@airbyte.io>
- Rename resolve_config to resolve_connector_config for clarity - Move resolve_connector_config and resolve_list_of_strings to new _arg_resolvers.py - Update imports in cloud.py, local.py, and registry.py - Simplify _util.py to only contain secret loading functions This separates argument resolution logic from secret loading logic. Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/mcp/_arg_resolvers.py`:
- Around line 30-73: In resolve_list_of_strings, the JSON-array branch currently
lets non-list[str] JSON values fall through to CSV parsing and the docstring
incorrectly states None/empty returns None; change behavior to raise a
ValueError when json.loads(value) successfully parses but the result is not a
list of strings (e.g., dict, list with non-strings, etc.), and update the
docstring text to accurately state that None returns None and empty/blank
strings return an empty list; keep existing CSV fallback and preserve TypeError
for unsupported types.
♻️ Duplicate comments (1)
airbyte/mcp/cloud.py (1)
1359-1489: Guard against missingapi_urlin org/workspace lookups.
When the MCP config arg isn’t set,api_urlcan beNone, but_resolve_organization_idandlist_workspaces_in_organizationexpect a non‑optionalapi_root. Would you be open to defaulting toapi_util.CLOUD_API_ROOThere to avoid runtime errors? wdyt?🛠️ Suggested fix
- api_url = get_mcp_config(ctx, MCP_CONFIG_API_URL) + api_url = get_mcp_config(ctx, MCP_CONFIG_API_URL) + api_url = api_url or api_util.CLOUD_API_ROOT ... - api_url = get_mcp_config(ctx, MCP_CONFIG_API_URL) + api_url = get_mcp_config(ctx, MCP_CONFIG_API_URL) + api_url = api_url or api_util.CLOUD_API_ROOT
🧹 Nitpick comments (1)
airbyte/mcp/_arg_resolvers.py (1)
82-96: Docstring mismatch on empty config inputs.
The function returns{}when all inputs areNone, but the Raises section says ValueError. Could we update the docstring to match? wdyt?✏️ Docstring tweak
- Returns: - Resolved configuration dictionary + Returns: + Resolved configuration dictionary (empty if no inputs provided) @@ - Raises: - ValueError: If no configuration provided or if JSON parsing fails + Raises: + ValueError: If JSON parsing fails or a provided input is invalid
… formatting - Remove AIRBYTE_CONFIG_ARGS list from _tool_utils.py - Import and pass config args directly inline in server.py - Use multiline list formatting for tool_filters Per PR review feedback for cleaner diffs when items are added/removed. Co-Authored-By: AJ Steers <aj@airbyte.io>
…id_is_set internally - Remove exclude_workspace_id_arg parameter from register_cloud_tools - Import and evaluate AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET directly in cloud.py - Simplify server.py by removing the passthrough parameter Per PR review feedback to avoid unnecessary parameter passthrough. Co-Authored-By: AJ Steers <aj@airbyte.io>
The file now only contains secret loading configuration functions, so _config.py is a more accurate name. Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/mcp/server.py`:
- Around line 91-95: The flag AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET is evaluated at
import time and can be stale if load_secrets_to_env_vars() sets
AIRBYTE_CLOUD_WORKSPACE_ID later; before calling register_cloud_tools(app, ...),
recompute the presence of the env var (e.g., re-check
os.environ.get("AIRBYTE_CLOUD_WORKSPACE_ID") or similar) and pass that fresh
boolean into register_cloud_tools instead of the module-level
AIRBYTE_CLOUD_WORKSPACE_ID_IS_SET so the workspace_id arg visibility reflects
secrets loaded at runtime.
- Add api_url guard to default to CLOUD_API_ROOT when get_mcp_config returns None
- Fix docstring in _arg_resolvers.py to match actual behavior (returns {} when no inputs)
Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
airbyte/mcp/cloud.py (1)
1888-1912: Potential inconsistency: Missingworkspace_idparameter on destructive delete operations?I noticed that
permanently_delete_cloud_source,permanently_delete_cloud_destination, andpermanently_delete_cloud_connectiondon't have the optionalworkspace_idparameter, whilepermanently_delete_custom_source_definitiondoes have it.Other operations on the same resources (like
rename_cloud_source,update_cloud_source_config) include theworkspace_idparameter for flexibility.Is this intentional as an extra safety measure (forcing users to rely solely on the MCP config workspace), or should these functions also accept an optional
workspace_idoverride for consistency with the rest of the API, wdyt?
♻️ Duplicate comments (2)
airbyte/mcp/_arg_resolvers.py (2)
20-27: Replace overload ellipsis withpassto avoid no‑op statements?
Line 22 and Line 27 still use...as a statement body, which static analyzers flag as “statement has no effect.” Would you be open to swapping inpasshere, wdyt?🔧 Proposed change
`@overload` -def resolve_list_of_strings(value: None) -> None: ... +def resolve_list_of_strings(value: None) -> None: + pass @@ `@overload` -def resolve_list_of_strings(value: str | list[str] | set[str]) -> list[str]: ... +def resolve_list_of_strings(value: str | list[str] | set[str]) -> list[str]: + pass
31-73: Tighten JSON‑array handling + docstring alignment?
The docstring contradicts behavior (Line 36 vs Line 39), and non‑list[str]JSON arrays fall through to CSV parsing (Lines 63‑69). Would you be open to clarifying the docstring and raising on non‑list[str]JSON arrays, wdyt?🛠️ Proposed change
- 2. None or empty input will return None. + 2. None input will return None. @@ - 5. If the input is empty or None, an empty list will be returned. + 5. Empty/blank string input will return an empty list. @@ if value.startswith("[") and value.endswith("]"): # Try to parse as JSON array: try: - parsed = json.loads(value) - if isinstance(parsed, list) and all(isinstance(item, str) for item in parsed): - return parsed + parsed = json.loads(value) except json.JSONDecodeError as ex: raise ValueError(f"Invalid JSON array: {value}") from ex + if not (isinstance(parsed, list) and all(isinstance(item, str) for item in parsed)): + raise ValueError(f"JSON array must be a list of strings: {value}") + return parsed
🧹 Nitpick comments (1)
airbyte/mcp/cloud.py (1)
1406-1428: Past review concern addressed — nice fix!The
api_urlfallback toapi_util.CLOUD_API_ROOTon line 1409 properly handles the case whenMCP_CONFIG_API_URLis not set. This matches the defensive pattern used in_get_cloud_workspace.One small observation: the credential resolution pattern (lines 1406-1409) and
SecretStringwrapping (lines 1415-1417, 1423-1425) is duplicated here and indescribe_cloud_organization. Would it be worth extracting a small helper like_get_mcp_credentials(ctx)that returns a tuple/dataclass with the wrapped secrets? This would reduce the duplication and keep credential handling in one place, wdyt?
aaronsteers
deleted the
devin/1768864254-migrate-mcp-to-fastmcp-extensions
branch
Summary
Migrates PyAirbyte's MCP implementation to use the new
fastmcp-extensionslibrary, deleting ~450 lines of custom MCP infrastructure code that is now provided by the library.Key changes:
@mcp_tooldecorator withfastmcp_extensions.mcp_tool@mcp_promptwithfastmcp_extensions.mcp_promptmcp_server()factory andregister_mcp_tools()/register_mcp_prompts()for registration_annotations.pyentirely (annotations now from library)_tool_utils.pyto only contain safe mode logic (62 lines, down from 265)constants.pyBreaking changes:
AIRBYTE_MCP_DOMAINSandAIRBYTE_MCP_DOMAINS_DISABLEDenv vars are no longer supportedReview & Testing Checklist for Human
poe mcp-serve-localand confirm tools are registered without errorspoe mcp-tool-test <tool_name> '<json_args>'to verify at least one tool from each module (cloud_ops, local_ops, connector_registry) worksinclude_standard_tool_filters=Trueinmcp_server()provides equivalent readonly mode functionalityAIRBYTE_MCP_DOMAINS*env varsRecommended test plan:
Notes
poe fix-and-check)Link to Devin run: https://app.devin.ai/sessions/a1d776294421421e88a60f61e20e1670
Requested by: @aaronsteers
Summary by CodeRabbit
New Features
Refactor
Tests
Chores
✏️ 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.