feat: Implement safe mode filtering for Cloud ops tools#840
feat: Implement safe mode filtering for Cloud ops tools#840Aaron ("AJ") Steers (aaronsteers) wants to merge 1 commit into
Conversation
Add server-side filtering based on two environment variables: - AIRBYTE_CLOUD_MCP_READONLY_MODE=1: Blocks all write operations on Cloud resources - AIRBYTE_CLOUD_MCP_SAFE_MODE=1: Blocks destructive operations (updates/deletes) on Cloud resources Implementation: - Created safe_mode.py module with filtering logic and SafeModeError exception - Added enforce_cloud_safe_mode decorator that checks annotations before tool execution - Wrapped all Cloud ops tools in register_cloud_ops_tools() with safe mode enforcement - Local ops tools are unaffected by these modes (only Cloud ops are restricted) The filtering uses the MCP annotations (readOnlyHint, destructiveHint) that were added in the previous PR to determine which tools to block. 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/1761261873-implement-safe-mode-filtering' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1761261873-implement-safe-mode-filtering'Helpful ResourcesPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
Community SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
|
Closing in favor of a cleaner approach using decorators on tool definitions |
📝 WalkthroughWalkthroughIntroduces centralized safe-mode enforcement for cloud operations tools. Creates a new module with environment-based mode detection and authorization checks, then wraps all cloud operation tools with a decorator applying read-only, idempotent, and destructive operation annotations for standardized safety enforcement. Changes
Sequence DiagramsequenceDiagram
participant User as User/MCP Client
participant Cloud_Ops as cloud_ops (wrapped tool)
participant SafeMode as enforce_cloud_safe_mode
participant Checker as check_cloud_tool_allowed
participant Original as Original Tool Logic
User->>Cloud_Ops: Call cloud operation
Cloud_Ops->>SafeMode: Execution (with annotations)
SafeMode->>Checker: check_cloud_tool_allowed(annotations, tool_name)
alt Mode Check
Checker->>Checker: Check environment modes
alt Readonly Mode + No READ_ONLY_HINT
Checker-->>SafeMode: Raise SafeModeError
SafeMode-->>Cloud_Ops: Exception
Cloud_Ops-->>User: Error: Operation blocked
else Safe Mode + DESTRUCTIVE_HINT=True
Checker-->>SafeMode: Raise SafeModeError
SafeMode-->>Cloud_Ops: Exception
Cloud_Ops-->>User: Error: Destructive operation blocked
else Allowed
Checker-->>SafeMode: Success
SafeMode->>Original: Call wrapped function
Original-->>SafeMode: Result
SafeMode-->>Cloud_Ops: Return result
Cloud_Ops-->>User: Success response
end
end
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes The safe-mode module introduces straightforward environment-based checks and exception handling with clear logic flow. Cloud ops changes follow a consistent, repetitive decorator-wrapping pattern across tools. No complex business logic or intertwined dependencies complicate the review—each tool annotation mapping is predictable and self-contained. Possibly related PRs
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
airbyte/mcp/cloud_ops.py(2 hunks)airbyte/mcp/safe_mode.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
airbyte/mcp/cloud_ops.py (1)
airbyte/mcp/safe_mode.py (1)
enforce_cloud_safe_mode(80-98)
⏰ 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.11, Windows)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (No Creds)
- GitHub Check: Pytest (Fast)
🔇 Additional comments (5)
airbyte/mcp/safe_mode.py (3)
21-24: LGTM!Clean exception definition for safe mode violations.
27-42: LGTM!Both mode detection functions follow a consistent pattern with good defensive programming using
.strip().
80-98: The decorator implementation is correct, but manual verification of tool registration is critical per PR objectives.The
@wrapsdecorator is properly imported and applied, which preserves__annotations__and function metadata. However, the core concern about FastMCP introspection remains valid:
- Decorator pattern: Correct usage with
@wraps(func)preserves the original function's annotations- Registration safeguard: Tools are registered with explicit
annotations=kwarg alongside the decorated wrapper, suggesting intentional defensive coding- Missing test coverage: No unit tests found for decorated cloud tools—exactly as noted in PR objectives
The real risk isn't the decorator itself, but whether FastMCP properly introspects the wrapper function. With
inspect.signature(), the wrapper shows(*args: Any, **kwargs: Any). FastMCP would need to use__annotations__instead, which@wrapspreserves.Have you run
poe mcp-tool-teston the decorated cloud tools to verify:
- Parameters are correctly recognized?
- Type validation works as expected?
- Tool descriptions are properly exposed?
This manual testing would confirm whether FastMCP's introspection works correctly with the decorated functions. Wdyt about adding at least one integration test for a decorated tool to prevent regression?
airbyte/mcp/cloud_ops.py (2)
27-27: LGTM!Clean import of the safe mode enforcement decorator.
749-826: Consistent decorator application pattern across all 15 tools.The pattern of wrapping each tool with
enforce_cloud_safe_mode(annotations)(function)and passing the same annotations toapp.tool()is applied consistently across all cloud operations. This ensures centralized safe mode enforcement as intended.However, this relies on FastMCP properly handling the decorated functions, which I flagged as a critical concern in
safe_mode.py. Once you verify that FastMCP can introspect wrapped functions correctly (signatures, type hints, parameter validation), this implementation should work as designed.
| check_workspace_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| deploy_source_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| deploy_destination_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| deploy_noop_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| create_connection_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| run_sync_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| get_sync_status_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| get_sync_logs_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| list_sources_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| list_destinations_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| list_connections_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| publish_custom_annotations = { | ||
| DESTRUCTIVE_HINT: False, | ||
| } | ||
| list_custom_annotations = { | ||
| READ_ONLY_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } | ||
| update_custom_annotations = { | ||
| DESTRUCTIVE_HINT: True, | ||
| } | ||
| delete_custom_annotations = { | ||
| DESTRUCTIVE_HINT: True, | ||
| IDEMPOTENT_HINT: True, | ||
| } |
There was a problem hiding this comment.
Read-only operations are missing explicit DESTRUCTIVE_HINT: False and will be incorrectly blocked in safe mode.
Most read-only annotation dictionaries (e.g., check_workspace_annotations, get_sync_status_annotations, etc.) set READ_ONLY_HINT: True and IDEMPOTENT_HINT: True but omit DESTRUCTIVE_HINT.
Since safe_mode.py line 71 defaults DESTRUCTIVE_HINT to True, these read-only operations will be treated as destructive when AIRBYTE_CLOUD_MCP_SAFE_MODE=1 is set, causing them to be incorrectly blocked.
To fix this, should all read-only annotation dictionaries explicitly set DESTRUCTIVE_HINT: False? For example:
check_workspace_annotations = {
READ_ONLY_HINT: True,
IDEMPOTENT_HINT: True,
+ DESTRUCTIVE_HINT: False,
}Also, minor question: Lines 744-747 set both DESTRUCTIVE_HINT: True and IDEMPOTENT_HINT: True for delete operations. While deletes can be idempotent (repeated deletes are safe), this combination is unusual. Can you confirm this is intentional?
wdyt?
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In airbyte/mcp/cloud_ops.py around lines 695 to 747, several read-only
annotation dicts (check_workspace_annotations, get_sync_status_annotations,
get_sync_logs_annotations, list_sources_annotations,
list_destinations_annotations, list_connections_annotations,
list_custom_annotations) currently set READ_ONLY_HINT and IDEMPOTENT_HINT but
omit DESTRUCTIVE_HINT; add DESTRUCTIVE_HINT: False to each of those read-only
dictionaries so safe mode won’t treat them as destructive, and for
delete_custom_annotations confirm intent: if deletes should be blocked in safe
mode keep DESTRUCTIVE_HINT: True, otherwise change it to DESTRUCTIVE_HINT: False
(or add an inline comment clarifying the choice).
| def check_cloud_tool_allowed(annotations: dict[str, Any], tool_name: str) -> None: | ||
| """Check if a Cloud ops tool is allowed to execute based on safe mode settings. | ||
|
|
||
| Args: | ||
| annotations: Tool annotations dict containing readOnlyHint and destructiveHint | ||
| tool_name: Name of the tool being checked | ||
|
|
||
| Raises: | ||
| SafeModeError: If the tool is blocked by safe mode restrictions | ||
| """ | ||
| readonly_mode = is_readonly_mode_enabled() | ||
| safe_mode = is_safe_mode_enabled() | ||
|
|
||
| if not readonly_mode and not safe_mode: | ||
| return | ||
|
|
||
| if readonly_mode: | ||
| is_readonly = annotations.get(READ_ONLY_HINT, False) | ||
| if not is_readonly: | ||
| raise SafeModeError( | ||
| f"Tool '{tool_name}' is blocked by AIRBYTE_CLOUD_MCP_READONLY_MODE. " | ||
| f"This tool performs write operations on Cloud resources. " | ||
| f"To allow write operations, set AIRBYTE_CLOUD_MCP_READONLY_MODE=0 or unset it." | ||
| ) | ||
|
|
||
| if safe_mode: | ||
| is_destructive = annotations.get(DESTRUCTIVE_HINT, True) # Default is True per FastMCP | ||
| if is_destructive: | ||
| raise SafeModeError( | ||
| f"Tool '{tool_name}' is blocked by AIRBYTE_CLOUD_MCP_SAFE_MODE. " | ||
| f"This tool performs destructive operations (updates/deletes) on Cloud resources. " | ||
| f"To allow destructive operations, set AIRBYTE_CLOUD_MCP_SAFE_MODE=0 or unset it." | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
Verify the default value of DESTRUCTIVE_HINT and its consistency with tool annotations.
The logic correctly blocks operations based on readonly and safe modes. However, Line 71 defaults DESTRUCTIVE_HINT to True, which means any tool without this hint explicitly set will be treated as destructive in safe mode.
Looking at cloud_ops.py, several annotation dictionaries don't explicitly set DESTRUCTIVE_HINT (e.g., check_workspace_annotations, get_sync_status_annotations, list_sources_annotations, etc.). These will default to destructive, which seems incorrect for read-only operations. Could you verify:
- Is
DESTRUCTIVE_HINTdefaulting toTrueper FastMCP spec, or should it default toFalse? - Should all read-only annotation dicts in
cloud_ops.pyexplicitly setDESTRUCTIVE_HINT: False?
wdyt?
🏁 Script executed:
#!/bin/bash
# Description: Check if read-only tools explicitly set DESTRUCTIVE_HINT to False
# Search for annotation dictionaries that set READ_ONLY_HINT but don't set DESTRUCTIVE_HINT
rg -A 5 'READ_ONLY_HINT.*True' airbyte/mcp/cloud_ops.py | rg -v 'DESTRUCTIVE_HINT'Length of output: 737
🏁 Script executed:
#!/bin/bash
# Get full context of read-only annotation dicts to see if there's any pattern
rg -B 2 -A 7 'READ_ONLY_HINT.*True' airbyte/mcp/cloud_ops.py | head -60Length of output: 1088
🏁 Script executed:
#!/bin/bash
# Search for any logic that correlates READ_ONLY_HINT with DESTRUCTIVE_HINT in the codebase
rg -n 'READ_ONLY_HINT|DESTRUCTIVE_HINT' airbyte/mcp/safe_mode.pyLength of output: 235
🏁 Script executed:
#!/bin/bash
# Check if DESTRUCTIVE_HINT is mentioned in any comments or docs
rg -B 2 -A 2 'DESTRUCTIVE_HINT' airbyte/mcp/ | head -40Length of output: 1841
🏁 Script executed:
#!/bin/bash
# Look for constant definitions
rg -n 'READ_ONLY_HINT|DESTRUCTIVE_HINT' airbyte/mcp/ --type-list | head -20Length of output: 389
Fix the semantic inconsistency between READ_ONLY_HINT and DESTRUCTIVE_HINT.
The verification confirms the original concern. All read-only tools in cloud_ops.py lack an explicit DESTRUCTIVE_HINT setting. With the default of True on line 71, these read-only tools will be incorrectly blocked in safe mode—contradicting the premise that read-only operations are non-destructive.
Two paths forward:
- Change the default to
False(line 71): This treatsDESTRUCTIVE_HINTas opt-in for destructive operations rather than opt-out for safe ones. - Explicitly add
DESTRUCTIVE_HINT: Falseto all read-only annotation dicts incloud_ops.py: This makes the intent explicit (which is generally preferable).
Which approach aligns with your intent here—deny-by-default or explicit safety marking, wdyt?
🤖 Prompt for AI Agents
In airbyte/mcp/safe_mode.py around lines 45 to 77, the current default
DESTRUCTIVE_HINT=True causes read-only tools to be treated as destructive and
blocked in safe mode; fix by making DESTRUCTIVE_HINT explicitly False for all
read-only tool annotations in cloud_ops.py (preferred), or if you want a global
behavior change instead, change the default on line 71 from True to False so
DESTRUCTIVE_HINT is opt-in for destructive operations; apply one of these two
fixes consistently and run tests.
feat: Implement safe mode filtering for Cloud ops tools
Summary
Implements server-side filtering for PyAirbyte MCP Cloud operations based on two environment variables:
AIRBYTE_CLOUD_MCP_READONLY_MODE=1: Blocks all write operations on Cloud resources (allows only read-only tools like list/get operations)AIRBYTE_CLOUD_MCP_SAFE_MODE=1: Blocks destructive operations (updates/deletes) on Cloud resources while allowing non-destructive writes (deploy, create, run syncs)Implementation details:
airbyte/mcp/safe_mode.pymodule withenforce_cloud_safe_modedecoratorregister_cloud_ops_tools()with safe mode enforcementreadOnlyHint,destructiveHint) to determine which tools to blockSafeModeErrorwith clear error messages when tools are blockedReview & Testing Checklist for Human
Test READONLY_MODE: Set
AIRBYTE_CLOUD_MCP_READONLY_MODE=1and verify:Test SAFE_MODE: Set
AIRBYTE_CLOUD_MCP_SAFE_MODE=1and verify:Test both modes together: Set both env vars to "1" and verify expected behavior (readonly mode should block all writes)
Verify local ops unaffected: Confirm that local ops tools (sync_source_to_cache, run_sql_query, etc.) work normally regardless of these env var settings
Check error messages: Verify SafeModeError messages are clear and provide guidance on how to disable restrictions
Notes
DESTRUCTIVE_HINTdefaults toTrueper FastMCP spec (line 74 in safe_mode.py) - if wrong, filtering breaksRequested by: AJ Steers (Aaron ("AJ") Steers (@aaronsteers))
Devin session: https://app.devin.ai/sessions/026dbd2898474bc2a3ce21aec839533e
Summary by CodeRabbit