feat: Implement session-based safe mode for Cloud ops (#844)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: aaronsteers

airbyte/mcp/_tool_utils.py

@@ -27,6 +27,7 @@
 AIRBYTE_CLOUD_MCP_SAFE_MODE = os.environ.get("AIRBYTE_CLOUD_MCP_SAFE_MODE", "").strip() == "1"
 
 _REGISTERED_TOOLS: list[tuple[Callable[..., Any], dict[str, Any]]] = []
+_GUIDS_CREATED_IN_SESSION: set[str] = set()
 
 
 class SafeModeError(Exception):
@@ -35,6 +36,32 @@ class SafeModeError(Exception):
     pass
 
 
+def register_guid_created_in_session(guid: str) -> None:
+    """Register a GUID as created in this session.
+
+    Args:
+        guid: The GUID to register
+    """
+    _GUIDS_CREATED_IN_SESSION.add(guid)
+
+
+def check_guid_created_in_session(guid: str) -> None:
+    """Check if a GUID was created in this session.
+
+    Raises SafeModeError if the GUID was not created in this session and
+    AIRBYTE_CLOUD_MCP_SAFE_MODE is set to 1.
+
+    Args:
+        guid: The GUID to check
+    """
+    if AIRBYTE_CLOUD_MCP_SAFE_MODE and guid not in _GUIDS_CREATED_IN_SESSION:
+        raise SafeModeError(
+            f"Cannot perform destructive operation on '{guid}': "
+            f"Object was not created in this session. "
+            f"AIRBYTE_CLOUD_MCP_SAFE_MODE is set to '1'."
+        )
+
+
 def should_register_tool(annotations: dict[str, Any]) -> bool:
     """Check if a tool should be registered based on safe mode settings.
 

airbyte/mcp/cloud_ops.py

@@ -18,7 +18,12 @@
 from airbyte.cloud.connectors import CloudDestination, CloudSource, CustomCloudSourceDefinition
 from airbyte.cloud.workspaces import CloudWorkspace
 from airbyte.destinations.util import get_noop_destination
-from airbyte.mcp._tool_utils import mcp_tool, register_tools
+from airbyte.mcp._tool_utils import (
+    check_guid_created_in_session,
+    mcp_tool,
+    register_guid_created_in_session,
+    register_tools,
+)
 from airbyte.mcp._util import resolve_config, resolve_list_of_strings
 
 
@@ -96,6 +101,7 @@ def deploy_source_to_cloud(
     except Exception as ex:
         return f"Failed to deploy source '{source_name}': {ex}"
     else:
+        register_guid_created_in_session(deployed_source.connector_id)
         return (
             f"Successfully deployed source '{source_name}' with ID '{deployed_source.connector_id}'"
             f" and URL: {deployed_source.connector_url}"
@@ -166,6 +172,7 @@ def deploy_destination_to_cloud(
     except Exception as ex:
         return f"Failed to deploy destination '{destination_name}': {ex}"
     else:
+        register_guid_created_in_session(deployed_destination.connector_id)
         return (
             f"Successfully deployed destination '{destination_name}' "
             f"with ID: {deployed_destination.connector_id}"
@@ -227,6 +234,7 @@ def create_connection_on_cloud(
     except Exception as ex:
         return f"Failed to create connection '{connection_name}': {ex}"
     else:
+        register_guid_created_in_session(deployed_connection.connection_id)
         return (
             f"Successfully created connection '{connection_name}' "
             f"with ID '{deployed_connection.connection_id}' and "
@@ -342,6 +350,7 @@ def deploy_noop_destination_to_cloud(
     except Exception as ex:
         return f"Failed to deploy No-op Destination: {ex}"
     else:
+        register_guid_created_in_session(deployed_destination.connector_id)
         return (
             f"Successfully deployed No-op Destination "
             f"with ID '{deployed_destination.connector_id}' and "
@@ -617,6 +626,7 @@ def publish_custom_source_definition(
     except Exception as ex:
         return f"Failed to publish custom source definition '{name}': {ex}"
     else:
+        register_guid_created_in_session(custom_source.definition_id)
         return (
             "Successfully published custom YAML source definition:\n"
             + _get_custom_source_definition_description(
@@ -684,6 +694,7 @@ def update_custom_source_definition(
     Note: Only YAML (declarative) connectors are currently supported.
     Docker-based custom sources are not yet available.
     """
+    check_guid_created_in_session(definition_id)
     try:
         processed_manifest = manifest_yaml
         if isinstance(manifest_yaml, str) and "\n" not in manifest_yaml:
@@ -733,6 +744,7 @@ def permanently_delete_custom_source_definition(
     Note: Only YAML (declarative) connectors are currently supported.
     Docker-based custom sources are not yet available.
     """
+    check_guid_created_in_session(definition_id)
     workspace: CloudWorkspace = _get_cloud_workspace()
     definition = workspace.get_custom_source_definition(
         definition_id=definition_id,
@@ -806,6 +818,7 @@ def update_cloud_source_config(
     and `AIRBYTE_API_ROOT` environment variables will be used to authenticate with the
     Airbyte Cloud API.
     """
+    check_guid_created_in_session(source_id)
     workspace: CloudWorkspace = _get_cloud_workspace()
     source = workspace.get_source(source_id=source_id)
 
@@ -881,6 +894,7 @@ def update_cloud_destination_config(
     and `AIRBYTE_API_ROOT` environment variables will be used to authenticate with the
     Airbyte Cloud API.
     """
+    check_guid_created_in_session(destination_id)
     workspace: CloudWorkspace = _get_cloud_workspace()
     destination = workspace.get_destination(destination_id=destination_id)
 
@@ -983,6 +997,7 @@ def set_cloud_connection_selected_streams(
     and `AIRBYTE_API_ROOT` environment variables will be used to authenticate with the
     Airbyte Cloud API.
     """
+    check_guid_created_in_session(connection_id)
     workspace: CloudWorkspace = _get_cloud_workspace()
     connection = workspace.get_connection(connection_id=connection_id)