diff --git a/autogpt_platform/backend/backend/blocks/google/_drive.py b/autogpt_platform/backend/backend/blocks/google/_drive.py index 3e3ecb2711..cb2b52821c 100644 --- a/autogpt_platform/backend/backend/blocks/google/_drive.py +++ b/autogpt_platform/backend/backend/blocks/google/_drive.py @@ -1,16 +1,8 @@ -import asyncio -import mimetypes -import uuid -from pathlib import Path from typing import Any, Literal, Optional from pydantic import BaseModel, ConfigDict, Field from backend.data.model import SchemaField -from backend.util.file import get_exec_file_path -from backend.util.request import Requests -from backend.util.type import MediaFileType -from backend.util.virus_scanner import scan_content_safe AttachmentView = Literal[ "DOCS", @@ -30,8 +22,8 @@ ATTACHMENT_VIEWS: tuple[AttachmentView, ...] = ( ) -class GoogleDriveFile(BaseModel): - """Represents a single file/folder picked from Google Drive""" +class _GoogleDriveFileBase(BaseModel): + """Internal base class for Google Drive file representation.""" model_config = ConfigDict(populate_by_name=True) @@ -49,144 +41,115 @@ class GoogleDriveFile(BaseModel): ) -def GoogleDrivePickerField( - multiselect: bool = False, - allow_folder_selection: bool = False, - allowed_views: Optional[list[AttachmentView]] = None, - allowed_mime_types: Optional[list[str]] = None, - scopes: Optional[list[str]] = None, - title: Optional[str] = None, - description: Optional[str] = None, - placeholder: Optional[str] = None, - **kwargs, +class GoogleDriveFile(_GoogleDriveFileBase): + """ + Represents a Google Drive file/folder with optional credentials for chaining. + + Used for both inputs and outputs in Google Drive blocks. The `_credentials_id` + field enables chaining between blocks - when one block outputs a file, the + next block can use the same credentials to access it. + + When used with GoogleDriveFileField(), the frontend renders a combined + auth + file picker UI that automatically populates `_credentials_id`. + """ + + # Hidden field for credential ID - populated by frontend, preserved in outputs + credentials_id: Optional[str] = Field( + None, + alias="_credentials_id", + description="Internal: credential ID for authentication", + ) + + +def GoogleDriveFileField( + *, + title: str, + description: str | None = None, + credentials_kwarg: str = "credentials", + credentials_scopes: list[str] | None = None, + allowed_views: list[AttachmentView] | None = None, + allowed_mime_types: list[str] | None = None, + placeholder: str | None = None, + **kwargs: Any, ) -> Any: """ - Creates a Google Drive Picker input field. + Creates a Google Drive file input field with auto-generated credentials. + + This field type produces a single UI element that handles both: + 1. Google OAuth authentication + 2. File selection via Google Drive Picker + + The system automatically generates a credentials field, and the credentials + are passed to the run() method using the specified kwarg name. Args: - multiselect: Allow selecting multiple files/folders (default: False) - allow_folder_selection: Allow selecting folders (default: False) - allowed_views: List of view types to show in picker (default: ["DOCS"]) - allowed_mime_types: Filter by MIME types (e.g., ["application/pdf"]) title: Field title shown in UI description: Field description/help text + credentials_kwarg: Name of the kwarg that will receive GoogleCredentials + in the run() method (default: "credentials") + credentials_scopes: OAuth scopes required (default: drive.file) + allowed_views: List of view types to show in picker (default: ["DOCS"]) + allowed_mime_types: Filter by MIME types placeholder: Placeholder text for the button - **kwargs: Additional SchemaField arguments (advanced, hidden, etc.) + **kwargs: Additional SchemaField arguments Returns: - Field definition that produces: - - Single GoogleDriveFile when multiselect=False - - list[GoogleDriveFile] when multiselect=True + Field definition that produces GoogleDriveFile Example: >>> class MyBlock(Block): - ... class Input(BlockSchema): - ... document: GoogleDriveFile = GoogleDrivePickerField( - ... title="Select Document", - ... allowed_views=["DOCUMENTS"], + ... class Input(BlockSchemaInput): + ... spreadsheet: GoogleDriveFile = GoogleDriveFileField( + ... title="Select Spreadsheet", + ... credentials_kwarg="creds", + ... allowed_views=["SPREADSHEETS"], ... ) ... - ... files: list[GoogleDriveFile] = GoogleDrivePickerField( - ... title="Select Multiple Files", - ... multiselect=True, - ... allow_folder_selection=True, - ... ) + ... async def run( + ... self, input_data: Input, *, creds: GoogleCredentials, **kwargs + ... ): + ... # creds is automatically populated + ... file = input_data.spreadsheet """ - # Build configuration that will be sent to frontend + + # Determine scopes - drive.file is sufficient for picker-selected files + scopes = credentials_scopes or ["https://www.googleapis.com/auth/drive.file"] + + # Build picker configuration with auto_credentials embedded picker_config = { - "multiselect": multiselect, - "allow_folder_selection": allow_folder_selection, + "multiselect": False, + "allow_folder_selection": False, "allowed_views": list(allowed_views) if allowed_views else ["DOCS"], + "scopes": scopes, + # Auto-credentials config tells frontend to include _credentials_id in output + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": scopes, + "kwarg_name": credentials_kwarg, + }, } - # Add optional configurations if allowed_mime_types: picker_config["allowed_mime_types"] = list(allowed_mime_types) - # Determine required scopes based on config - base_scopes = scopes if scopes is not None else [] - picker_scopes: set[str] = set(base_scopes) - if allow_folder_selection: - picker_scopes.add("https://www.googleapis.com/auth/drive") - else: - # Use drive.file for minimal scope - only access files selected by user in picker - picker_scopes.add("https://www.googleapis.com/auth/drive.file") - - picker_config["scopes"] = sorted(picker_scopes) - - # Set appropriate default value - default_value = [] if multiselect else None - - # Use SchemaField to handle format properly return SchemaField( - default=default_value, + default=None, title=title, description=description, - placeholder=placeholder or "Choose from Google Drive", + placeholder=placeholder or "Select from Google Drive", + # Use google-drive-picker format so frontend renders existing component format="google-drive-picker", advanced=False, json_schema_extra={ "google_drive_picker_config": picker_config, + # Also keep auto_credentials at top level for backend detection + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": scopes, + "kwarg_name": credentials_kwarg, + }, **kwargs, }, ) - - -DRIVE_API_URL = "https://www.googleapis.com/drive/v3/files" -_requests = Requests(trusted_origins=["https://www.googleapis.com"]) - - -def GoogleDriveAttachmentField( - *, - title: str, - description: str | None = None, - placeholder: str | None = None, - multiselect: bool = True, - allowed_mime_types: list[str] | None = None, - **extra: Any, -) -> Any: - return GoogleDrivePickerField( - multiselect=multiselect, - allowed_views=list(ATTACHMENT_VIEWS), - allowed_mime_types=allowed_mime_types, - title=title, - description=description, - placeholder=placeholder or "Choose files from Google Drive", - **extra, - ) - - -async def drive_file_to_media_file( - drive_file: GoogleDriveFile, *, graph_exec_id: str, access_token: str -) -> MediaFileType: - if drive_file.is_folder: - raise ValueError("Google Drive selection must be a file.") - if not access_token: - raise ValueError("Google Drive access token is required for file download.") - - url = f"{DRIVE_API_URL}/{drive_file.id}?alt=media" - response = await _requests.get( - url, headers={"Authorization": f"Bearer {access_token}"} - ) - - mime_type = drive_file.mime_type or response.headers.get( - "content-type", "application/octet-stream" - ) - - MAX_FILE_SIZE = 100 * 1024 * 1024 - if len(response.content) > MAX_FILE_SIZE: - raise ValueError( - f"File too large: {len(response.content)} bytes > {MAX_FILE_SIZE} bytes" - ) - - base_path = Path(get_exec_file_path(graph_exec_id, "")) - base_path.mkdir(parents=True, exist_ok=True) - - extension = mimetypes.guess_extension(mime_type, strict=False) or ".bin" - filename = f"{uuid.uuid4()}{extension}" - target_path = base_path / filename - - await scan_content_safe(response.content, filename=filename) - await asyncio.to_thread(target_path.write_bytes, response.content) - - return MediaFileType(str(target_path.relative_to(base_path))) diff --git a/autogpt_platform/backend/backend/blocks/google/sheets.py b/autogpt_platform/backend/backend/blocks/google/sheets.py index e35749a72e..fac4e2d1aa 100644 --- a/autogpt_platform/backend/backend/blocks/google/sheets.py +++ b/autogpt_platform/backend/backend/blocks/google/sheets.py @@ -5,7 +5,7 @@ from typing import Any from google.oauth2.credentials import Credentials from googleapiclient.discovery import build -from backend.blocks.google._drive import GoogleDriveFile, GoogleDrivePickerField +from backend.blocks.google._drive import GoogleDriveFile, GoogleDriveFileField from backend.data.block import ( Block, BlockCategory, @@ -182,6 +182,28 @@ def _build_sheets_service(credentials: GoogleCredentials): return build("sheets", "v4", credentials=creds) +def _build_drive_service(credentials: GoogleCredentials): + """Build Drive service from platform credentials (with refresh token).""" + settings = Settings() + creds = Credentials( + token=( + credentials.access_token.get_secret_value() + if credentials.access_token + else None + ), + refresh_token=( + credentials.refresh_token.get_secret_value() + if credentials.refresh_token + else None + ), + token_uri="https://oauth2.googleapis.com/token", + client_id=settings.secrets.google_client_id, + client_secret=settings.secrets.google_client_secret, + scopes=credentials.scopes, + ) + return build("drive", "v3", credentials=creds) + + def _validate_spreadsheet_file(spreadsheet_file: "GoogleDriveFile") -> str | None: """Validate that the selected file is a Google Sheets spreadsheet. @@ -250,10 +272,10 @@ class BatchOperation(BlockSchemaInput): class GoogleSheetsReadBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -282,7 +304,6 @@ class GoogleSheetsReadBlock(Block): output_schema=GoogleSheetsReadBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -308,6 +329,7 @@ class GoogleSheetsReadBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -338,7 +360,7 @@ class GoogleSheetsReadBlock(Block): self._read_sheet, service, spreadsheet_id, input_data.range ) yield "result", data - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=spreadsheet_id, name=input_data.spreadsheet.name, @@ -346,6 +368,7 @@ class GoogleSheetsReadBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", _handle_sheets_api_error(str(e), "read") @@ -373,10 +396,10 @@ class GoogleSheetsReadBlock(Block): class GoogleSheetsWriteBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -408,7 +431,6 @@ class GoogleSheetsWriteBlock(Block): output_schema=GoogleSheetsWriteBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -435,6 +457,7 @@ class GoogleSheetsWriteBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -477,7 +500,7 @@ class GoogleSheetsWriteBlock(Block): input_data.values, ) yield "result", result - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -485,6 +508,7 @@ class GoogleSheetsWriteBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", _handle_sheets_api_error(str(e), "write") @@ -509,10 +533,10 @@ class GoogleSheetsWriteBlock(Block): class GoogleSheetsAppendBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -566,7 +590,6 @@ class GoogleSheetsAppendBlock(Block): output_schema=GoogleSheetsAppendBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -586,6 +609,7 @@ class GoogleSheetsAppendBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -642,7 +666,7 @@ class GoogleSheetsAppendBlock(Block): input_data.insert_data_option, ) yield "result", result - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -650,6 +674,7 @@ class GoogleSheetsAppendBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to append to Google Sheet: {str(e)}" @@ -690,10 +715,10 @@ class GoogleSheetsAppendBlock(Block): class GoogleSheetsClearBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -722,7 +747,6 @@ class GoogleSheetsClearBlock(Block): output_schema=GoogleSheetsClearBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -742,6 +766,7 @@ class GoogleSheetsClearBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -774,7 +799,7 @@ class GoogleSheetsClearBlock(Block): input_data.range, ) yield "result", result - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -782,6 +807,7 @@ class GoogleSheetsClearBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to clear Google Sheet range: {str(e)}" @@ -798,10 +824,10 @@ class GoogleSheetsClearBlock(Block): class GoogleSheetsMetadataBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -826,7 +852,6 @@ class GoogleSheetsMetadataBlock(Block): output_schema=GoogleSheetsMetadataBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -851,6 +876,7 @@ class GoogleSheetsMetadataBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -883,7 +909,7 @@ class GoogleSheetsMetadataBlock(Block): input_data.spreadsheet.id, ) yield "result", result - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -891,6 +917,7 @@ class GoogleSheetsMetadataBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to get spreadsheet metadata: {str(e)}" @@ -918,10 +945,10 @@ class GoogleSheetsMetadataBlock(Block): class GoogleSheetsManageSheetBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -955,7 +982,6 @@ class GoogleSheetsManageSheetBlock(Block): output_schema=GoogleSheetsManageSheetBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -976,6 +1002,7 @@ class GoogleSheetsManageSheetBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -1012,7 +1039,7 @@ class GoogleSheetsManageSheetBlock(Block): input_data.destination_sheet_name, ) yield "result", result - # Output the GoogleDriveFile for chaining + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -1020,6 +1047,7 @@ class GoogleSheetsManageSheetBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to manage sheet: {str(e)}" @@ -1073,10 +1101,10 @@ class GoogleSheetsManageSheetBlock(Block): class GoogleSheetsBatchOperationsBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -1104,7 +1132,6 @@ class GoogleSheetsBatchOperationsBlock(Block): output_schema=GoogleSheetsBatchOperationsBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -1135,6 +1162,7 @@ class GoogleSheetsBatchOperationsBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -1168,6 +1196,7 @@ class GoogleSheetsBatchOperationsBlock(Block): input_data.operations, ) yield "result", result + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -1175,6 +1204,7 @@ class GoogleSheetsBatchOperationsBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to perform batch operations: {str(e)}" @@ -1228,10 +1258,10 @@ class GoogleSheetsBatchOperationsBlock(Block): class GoogleSheetsFindReplaceBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -1274,7 +1304,6 @@ class GoogleSheetsFindReplaceBlock(Block): output_schema=GoogleSheetsFindReplaceBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -1297,6 +1326,7 @@ class GoogleSheetsFindReplaceBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -1331,6 +1361,7 @@ class GoogleSheetsFindReplaceBlock(Block): input_data.match_entire_cell, ) yield "result", result + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -1338,6 +1369,7 @@ class GoogleSheetsFindReplaceBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to find/replace in Google Sheet: {str(e)}" @@ -1376,10 +1408,10 @@ class GoogleSheetsFindReplaceBlock(Block): class GoogleSheetsFindBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -1434,7 +1466,6 @@ class GoogleSheetsFindBlock(Block): output_schema=GoogleSheetsFindBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -1467,6 +1498,7 @@ class GoogleSheetsFindBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -1511,6 +1543,7 @@ class GoogleSheetsFindBlock(Block): yield "count", result["count"] yield "locations", result["locations"] yield "result", {"success": True} + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -1518,6 +1551,7 @@ class GoogleSheetsFindBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to find text in Google Sheet: {str(e)}" @@ -1682,10 +1716,10 @@ class GoogleSheetsFindBlock(Block): class GoogleSheetsFormatBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -1717,7 +1751,6 @@ class GoogleSheetsFormatBlock(Block): output_schema=GoogleSheetsFormatBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -1739,6 +1772,7 @@ class GoogleSheetsFormatBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -1775,6 +1809,7 @@ class GoogleSheetsFormatBlock(Block): yield "error", result["error"] else: yield "result", result + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -1782,6 +1817,7 @@ class GoogleSheetsFormatBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", f"Failed to format Google Sheet cells: {str(e)}" @@ -1855,7 +1891,10 @@ class GoogleSheetsFormatBlock(Block): class GoogleSheetsCreateSpreadsheetBlock(Block): class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) + # Explicit credentials since this block creates a file (no file picker) + credentials: GoogleCredentialsInput = GoogleCredentialsField( + ["https://www.googleapis.com/auth/drive.file"] + ) title: str = SchemaField( description="The title of the new spreadsheet", ) @@ -1890,9 +1929,9 @@ class GoogleSheetsCreateSpreadsheetBlock(Block): output_schema=GoogleSheetsCreateSpreadsheetBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ + "credentials": TEST_CREDENTIALS_INPUT, "title": "Test Spreadsheet", "sheet_names": ["Sheet1", "Data", "Summary"], - "credentials": TEST_CREDENTIALS_INPUT, }, test_credentials=TEST_CREDENTIALS, test_output=[ @@ -1905,6 +1944,9 @@ class GoogleSheetsCreateSpreadsheetBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=TEST_CREDENTIALS_INPUT[ + "id" + ], # Preserves credential ID for chaining ), ), ("spreadsheet_id", "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms"), @@ -1926,10 +1968,12 @@ class GoogleSheetsCreateSpreadsheetBlock(Block): async def run( self, input_data: Input, *, credentials: GoogleCredentials, **kwargs ) -> BlockOutput: - service = _build_sheets_service(credentials) + drive_service = _build_drive_service(credentials) + sheets_service = _build_sheets_service(credentials) result = await asyncio.to_thread( self._create_spreadsheet, - service, + drive_service, + sheets_service, input_data.title, input_data.sheet_names, ) @@ -1939,7 +1983,7 @@ class GoogleSheetsCreateSpreadsheetBlock(Block): else: spreadsheet_id = result["spreadsheetId"] spreadsheet_url = result["spreadsheetUrl"] - # Output the full GoogleDriveFile object for easy chaining + # Output the GoogleDriveFile for chaining (includes credentials_id) yield "spreadsheet", GoogleDriveFile( id=spreadsheet_id, name=result.get("title", input_data.title), @@ -1947,40 +1991,68 @@ class GoogleSheetsCreateSpreadsheetBlock(Block): url=spreadsheet_url, iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.credentials.id, # Preserve credentials for chaining ) yield "spreadsheet_id", spreadsheet_id yield "spreadsheet_url", spreadsheet_url yield "result", {"success": True} - def _create_spreadsheet(self, service, title: str, sheet_names: list[str]) -> dict: + def _create_spreadsheet( + self, drive_service, sheets_service, title: str, sheet_names: list[str] + ) -> dict: try: - # Create the initial spreadsheet - spreadsheet_body = { - "properties": {"title": title}, - "sheets": [ - { - "properties": { - "title": sheet_names[0] if sheet_names else "Sheet1" - } - } - ], + # Create blank spreadsheet using Drive API + file_metadata = { + "name": title, + "mimeType": "application/vnd.google-apps.spreadsheet", } + result = ( + drive_service.files() + .create(body=file_metadata, fields="id, webViewLink") + .execute() + ) - result = service.spreadsheets().create(body=spreadsheet_body).execute() - spreadsheet_id = result["spreadsheetId"] - spreadsheet_url = result["spreadsheetUrl"] + spreadsheet_id = result["id"] + spreadsheet_url = result.get( + "webViewLink", + f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit", + ) + + # Rename first sheet if custom name provided (default is "Sheet1") + if sheet_names and sheet_names[0] != "Sheet1": + # Get first sheet ID and rename it + meta = ( + sheets_service.spreadsheets() + .get(spreadsheetId=spreadsheet_id) + .execute() + ) + first_sheet_id = meta["sheets"][0]["properties"]["sheetId"] + sheets_service.spreadsheets().batchUpdate( + spreadsheetId=spreadsheet_id, + body={ + "requests": [ + { + "updateSheetProperties": { + "properties": { + "sheetId": first_sheet_id, + "title": sheet_names[0], + }, + "fields": "title", + } + } + ] + }, + ).execute() # Add additional sheets if requested if len(sheet_names) > 1: - requests = [] - for sheet_name in sheet_names[1:]: - requests.append({"addSheet": {"properties": {"title": sheet_name}}}) - - if requests: - batch_body = {"requests": requests} - service.spreadsheets().batchUpdate( - spreadsheetId=spreadsheet_id, body=batch_body - ).execute() + requests = [ + {"addSheet": {"properties": {"title": name}}} + for name in sheet_names[1:] + ] + sheets_service.spreadsheets().batchUpdate( + spreadsheetId=spreadsheet_id, body={"requests": requests} + ).execute() return { "spreadsheetId": spreadsheet_id, @@ -1995,10 +2067,10 @@ class GoogleSheetsUpdateCellBlock(Block): """Update a single cell in a Google Sheets spreadsheet.""" class Input(BlockSchemaInput): - credentials: GoogleCredentialsInput = GoogleCredentialsField([]) - spreadsheet: GoogleDriveFile = GoogleDrivePickerField( + spreadsheet: GoogleDriveFile = GoogleDriveFileField( title="Spreadsheet", description="Select a Google Sheets spreadsheet", + credentials_kwarg="credentials", allowed_views=["SPREADSHEETS"], allowed_mime_types=["application/vnd.google-apps.spreadsheet"], ) @@ -2035,7 +2107,6 @@ class GoogleSheetsUpdateCellBlock(Block): output_schema=GoogleSheetsUpdateCellBlock.Output, disabled=GOOGLE_SHEETS_DISABLED, test_input={ - "credentials": TEST_CREDENTIALS_INPUT, "spreadsheet": { "id": "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms", "name": "Test Spreadsheet", @@ -2059,6 +2130,7 @@ class GoogleSheetsUpdateCellBlock(Block): url="https://docs.google.com/spreadsheets/d/1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=None, ), ), ], @@ -2096,6 +2168,7 @@ class GoogleSheetsUpdateCellBlock(Block): ) yield "result", result + # Output the GoogleDriveFile for chaining (preserves credentials_id) yield "spreadsheet", GoogleDriveFile( id=input_data.spreadsheet.id, name=input_data.spreadsheet.name, @@ -2103,6 +2176,7 @@ class GoogleSheetsUpdateCellBlock(Block): url=f"https://docs.google.com/spreadsheets/d/{input_data.spreadsheet.id}/edit", iconUrl="https://www.gstatic.com/images/branding/product/1x/sheets_48dp.png", isFolder=False, + _credentials_id=input_data.spreadsheet.credentials_id, ) except Exception as e: yield "error", _handle_sheets_api_error(str(e), "update") diff --git a/autogpt_platform/backend/backend/blocks/io.py b/autogpt_platform/backend/backend/blocks/io.py index 5a9e59ddb4..07f09eb349 100644 --- a/autogpt_platform/backend/backend/blocks/io.py +++ b/autogpt_platform/backend/backend/blocks/io.py @@ -2,6 +2,8 @@ import copy from datetime import date, time from typing import Any, Optional +# Import for Google Drive file input block +from backend.blocks.google._drive import AttachmentView, GoogleDriveFile from backend.data.block import ( Block, BlockCategory, @@ -646,6 +648,119 @@ class AgentTableInputBlock(AgentInputBlock): yield "result", input_data.value if input_data.value is not None else [] +class AgentGoogleDriveFileInputBlock(AgentInputBlock): + """ + This block allows users to select a file from Google Drive. + + It provides a Google Drive file picker UI that handles both authentication + and file selection. The selected file information (ID, name, URL, etc.) + is output for use by other blocks like Google Sheets Read. + """ + + class Input(AgentInputBlock.Input): + value: Optional[GoogleDriveFile] = SchemaField( + description="The selected Google Drive file.", + default=None, + advanced=False, + title="Selected File", + ) + allowed_views: list[AttachmentView] = SchemaField( + description="Which views to show in the file picker (DOCS, SPREADSHEETS, PRESENTATIONS, etc.).", + default_factory=lambda: ["DOCS", "SPREADSHEETS", "PRESENTATIONS"], + advanced=False, + title="Allowed Views", + ) + allow_folder_selection: bool = SchemaField( + description="Whether to allow selecting folders.", + default=False, + advanced=True, + title="Allow Folder Selection", + ) + + def generate_schema(self): + """Generate schema for the value field with Google Drive picker format.""" + schema = super().generate_schema() + + # Default scopes for drive.file access + scopes = ["https://www.googleapis.com/auth/drive.file"] + + # Build picker configuration + picker_config = { + "multiselect": False, # Single file selection only for now + "allow_folder_selection": self.allow_folder_selection, + "allowed_views": ( + list(self.allowed_views) if self.allowed_views else ["DOCS"] + ), + "scopes": scopes, + # Auto-credentials config tells frontend to include _credentials_id in output + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": scopes, + "kwarg_name": "credentials", + }, + } + + # Set format and config for frontend to render Google Drive picker + schema["format"] = "google-drive-picker" + schema["google_drive_picker_config"] = picker_config + # Also keep auto_credentials at top level for backend detection + schema["auto_credentials"] = { + "provider": "google", + "type": "oauth2", + "scopes": scopes, + "kwarg_name": "credentials", + } + + if self.value is not None: + schema["default"] = self.value.model_dump() + + return schema + + class Output(AgentInputBlock.Output): + result: GoogleDriveFile = SchemaField( + description="The selected Google Drive file with ID, name, URL, and other metadata." + ) + + def __init__(self): + test_file = GoogleDriveFile.model_validate( + { + "id": "test-file-id", + "name": "Test Spreadsheet", + "mimeType": "application/vnd.google-apps.spreadsheet", + "url": "https://docs.google.com/spreadsheets/d/test-file-id", + } + ) + super().__init__( + id="d3b32f15-6fd7-40e3-be52-e083f51b19a2", + description="Block for selecting a file from Google Drive.", + disabled=not config.enable_agent_input_subtype_blocks, + input_schema=AgentGoogleDriveFileInputBlock.Input, + output_schema=AgentGoogleDriveFileInputBlock.Output, + test_input=[ + { + "name": "spreadsheet_input", + "description": "Select a spreadsheet from Google Drive", + "allowed_views": ["SPREADSHEETS"], + "value": { + "id": "test-file-id", + "name": "Test Spreadsheet", + "mimeType": "application/vnd.google-apps.spreadsheet", + "url": "https://docs.google.com/spreadsheets/d/test-file-id", + }, + } + ], + test_output=[("result", test_file)], + ) + + async def run(self, input_data: Input, *args, **kwargs) -> BlockOutput: + """ + Yields the selected Google Drive file. + """ + if input_data.value is not None: + yield "result", input_data.value + + IO_BLOCK_IDs = [ AgentInputBlock().id, AgentOutputBlock().id, @@ -658,4 +773,5 @@ IO_BLOCK_IDs = [ AgentDropdownInputBlock().id, AgentToggleInputBlock().id, AgentTableInputBlock().id, + AgentGoogleDriveFileInputBlock().id, ] diff --git a/autogpt_platform/backend/backend/blocks/test/test_block.py b/autogpt_platform/backend/backend/blocks/test/test_block.py index 2c5313b7ab..7a1fdbcc73 100644 --- a/autogpt_platform/backend/backend/blocks/test/test_block.py +++ b/autogpt_platform/backend/backend/blocks/test/test_block.py @@ -1,8 +1,9 @@ -from typing import Type +from typing import Any, Type import pytest -from backend.data.block import Block, get_blocks +from backend.data.block import Block, BlockSchemaInput, get_blocks +from backend.data.model import SchemaField from backend.util.test import execute_block_test SKIP_BLOCK_TESTS = { @@ -132,3 +133,148 @@ async def test_block_ids_valid(block: Type[Block]): ), f"Block {block.name} ID is UUID version {parsed_uuid.version}, expected version 4" except ValueError: pytest.fail(f"Block {block.name} has invalid UUID format: {block_instance.id}") + + +class TestAutoCredentialsFieldsValidation: + """Tests for auto_credentials field validation in BlockSchema.""" + + def test_duplicate_auto_credentials_kwarg_name_raises_error(self): + """Test that duplicate kwarg_name in auto_credentials raises ValueError.""" + + class DuplicateKwargSchema(BlockSchemaInput): + """Schema with duplicate auto_credentials kwarg_name.""" + + # Both fields explicitly use the same kwarg_name "credentials" + file1: dict[str, Any] | None = SchemaField( + description="First file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + "kwarg_name": "credentials", + } + }, + ) + file2: dict[str, Any] | None = SchemaField( + description="Second file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + "kwarg_name": "credentials", # Duplicate kwarg_name! + } + }, + ) + + with pytest.raises(ValueError) as exc_info: + DuplicateKwargSchema.get_auto_credentials_fields() + + error_message = str(exc_info.value) + assert "Duplicate auto_credentials kwarg_name 'credentials'" in error_message + assert "file1" in error_message + assert "file2" in error_message + + def test_unique_auto_credentials_kwarg_names_succeed(self): + """Test that unique kwarg_name values work correctly.""" + + class UniqueKwargSchema(BlockSchemaInput): + """Schema with unique auto_credentials kwarg_name values.""" + + file1: dict[str, Any] | None = SchemaField( + description="First file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + "kwarg_name": "file1_credentials", + } + }, + ) + file2: dict[str, Any] | None = SchemaField( + description="Second file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + "kwarg_name": "file2_credentials", # Different kwarg_name + } + }, + ) + + # Should not raise + result = UniqueKwargSchema.get_auto_credentials_fields() + + assert "file1_credentials" in result + assert "file2_credentials" in result + assert result["file1_credentials"]["field_name"] == "file1" + assert result["file2_credentials"]["field_name"] == "file2" + + def test_default_kwarg_name_is_credentials(self): + """Test that missing kwarg_name defaults to 'credentials'.""" + + class DefaultKwargSchema(BlockSchemaInput): + """Schema with auto_credentials missing kwarg_name.""" + + file: dict[str, Any] | None = SchemaField( + description="File input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + # No kwarg_name specified - should default to "credentials" + } + }, + ) + + result = DefaultKwargSchema.get_auto_credentials_fields() + + assert "credentials" in result + assert result["credentials"]["field_name"] == "file" + + def test_duplicate_default_kwarg_name_raises_error(self): + """Test that two fields with default kwarg_name raises ValueError.""" + + class DefaultDuplicateSchema(BlockSchemaInput): + """Schema where both fields omit kwarg_name, defaulting to 'credentials'.""" + + file1: dict[str, Any] | None = SchemaField( + description="First file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + # No kwarg_name - defaults to "credentials" + } + }, + ) + file2: dict[str, Any] | None = SchemaField( + description="Second file input", + default=None, + json_schema_extra={ + "auto_credentials": { + "provider": "google", + "type": "oauth2", + "scopes": ["https://www.googleapis.com/auth/drive.file"], + # No kwarg_name - also defaults to "credentials" + } + }, + ) + + with pytest.raises(ValueError) as exc_info: + DefaultDuplicateSchema.get_auto_credentials_fields() + + assert "Duplicate auto_credentials kwarg_name 'credentials'" in str( + exc_info.value + ) diff --git a/autogpt_platform/backend/backend/data/block.py b/autogpt_platform/backend/backend/data/block.py index 762e9b37ef..315d63bd8f 100644 --- a/autogpt_platform/backend/backend/data/block.py +++ b/autogpt_platform/backend/backend/data/block.py @@ -266,14 +266,61 @@ class BlockSchema(BaseModel): ) } + @classmethod + def get_auto_credentials_fields(cls) -> dict[str, dict[str, Any]]: + """ + Get fields that have auto_credentials metadata (e.g., GoogleDriveFileInput). + + Returns a dict mapping kwarg_name -> {field_name, auto_credentials_config} + + Raises: + ValueError: If multiple fields have the same kwarg_name, as this would + cause silent overwriting and only the last field would be processed. + """ + result: dict[str, dict[str, Any]] = {} + schema = cls.jsonschema() + properties = schema.get("properties", {}) + + for field_name, field_schema in properties.items(): + auto_creds = field_schema.get("auto_credentials") + if auto_creds: + kwarg_name = auto_creds.get("kwarg_name", "credentials") + if kwarg_name in result: + raise ValueError( + f"Duplicate auto_credentials kwarg_name '{kwarg_name}' " + f"in fields '{result[kwarg_name]['field_name']}' and " + f"'{field_name}' on {cls.__qualname__}" + ) + result[kwarg_name] = { + "field_name": field_name, + "config": auto_creds, + } + return result + @classmethod def get_credentials_fields_info(cls) -> dict[str, CredentialsFieldInfo]: - return { - field_name: CredentialsFieldInfo.model_validate( + result = {} + + # Regular credentials fields + for field_name in cls.get_credentials_fields().keys(): + result[field_name] = CredentialsFieldInfo.model_validate( cls.get_field_schema(field_name), by_alias=True ) - for field_name in cls.get_credentials_fields().keys() - } + + # Auto-generated credentials fields (from GoogleDriveFileInput etc.) + for kwarg_name, info in cls.get_auto_credentials_fields().items(): + config = info["config"] + # Build a schema-like dict that CredentialsFieldInfo can parse + auto_schema = { + "credentials_provider": [config.get("provider", "google")], + "credentials_types": [config.get("type", "oauth2")], + "credentials_scopes": config.get("scopes"), + } + result[kwarg_name] = CredentialsFieldInfo.model_validate( + auto_schema, by_alias=True + ) + + return result @classmethod def get_input_defaults(cls, data: BlockInput) -> BlockInput: diff --git a/autogpt_platform/backend/backend/executor/manager.py b/autogpt_platform/backend/backend/executor/manager.py index bfec94176d..ee875bbf55 100644 --- a/autogpt_platform/backend/backend/executor/manager.py +++ b/autogpt_platform/backend/backend/executor/manager.py @@ -218,15 +218,53 @@ async def execute_node( # changes during execution. ⚠️ This means a set of credentials can only be used by # one (running) block at a time; simultaneous execution of blocks using same # credentials is not supported. - creds_lock = None + creds_locks: list[AsyncRedisLock] = [] input_model = cast(type[BlockSchema], node_block.input_schema) + + # Handle regular credentials fields for field_name, input_type in input_model.get_credentials_fields().items(): credentials_meta = input_type(**input_data[field_name]) - credentials, creds_lock = await creds_manager.acquire( - user_id, credentials_meta.id - ) + credentials, lock = await creds_manager.acquire(user_id, credentials_meta.id) + creds_locks.append(lock) extra_exec_kwargs[field_name] = credentials + # Handle auto-generated credentials (e.g., from GoogleDriveFileInput) + for kwarg_name, info in input_model.get_auto_credentials_fields().items(): + field_name = info["field_name"] + field_data = input_data.get(field_name) + if field_data and isinstance(field_data, dict): + # Check if _credentials_id key exists in the field data + if "_credentials_id" in field_data: + cred_id = field_data["_credentials_id"] + if cred_id: + # Credential ID provided - acquire credentials + provider = info.get("config", {}).get( + "provider", "external service" + ) + file_name = field_data.get("name", "selected file") + try: + credentials, lock = await creds_manager.acquire( + user_id, cred_id + ) + creds_locks.append(lock) + extra_exec_kwargs[kwarg_name] = credentials + except ValueError: + # Credential was deleted or doesn't exist + raise ValueError( + f"Authentication expired for '{file_name}' in field '{field_name}'. " + f"The saved {provider.capitalize()} credentials no longer exist. " + f"Please re-select the file to re-authenticate." + ) + # else: _credentials_id is explicitly None, skip credentials (for chained data) + else: + # _credentials_id key missing entirely - this is an error + provider = info.get("config", {}).get("provider", "external service") + file_name = field_data.get("name", "selected file") + raise ValueError( + f"Authentication missing for '{file_name}' in field '{field_name}'. " + f"Please re-select the file to authenticate with {provider.capitalize()}." + ) + output_size = 0 # sentry tracking nonsense to get user counts for blocks because isolation scopes don't work :( @@ -260,12 +298,17 @@ async def execute_node( # Re-raise to maintain normal error flow raise finally: - # Ensure credentials are released even if execution fails - if creds_lock and (await creds_lock.locked()) and (await creds_lock.owned()): - try: - await creds_lock.release() - except Exception as e: - log_metadata.error(f"Failed to release credentials lock: {e}") + # Ensure all credentials are released even if execution fails + for creds_lock in creds_locks: + if ( + creds_lock + and (await creds_lock.locked()) + and (await creds_lock.owned()) + ): + try: + await creds_lock.release() + except Exception as e: + log_metadata.error(f"Failed to release credentials lock: {e}") # Update execution stats if execution_stats is not None: diff --git a/autogpt_platform/backend/backend/util/test.py b/autogpt_platform/backend/backend/util/test.py index dda62e7f9f..95ea9554ed 100644 --- a/autogpt_platform/backend/backend/util/test.py +++ b/autogpt_platform/backend/backend/util/test.py @@ -144,6 +144,8 @@ async def execute_block_test(block: Block): "execution_context": ExecutionContext(), } input_model = cast(type[BlockSchema], block.input_schema) + + # Handle regular credentials fields credentials_input_fields = input_model.get_credentials_fields() if len(credentials_input_fields) == 1 and isinstance( block.test_credentials, _BaseCredentials @@ -158,6 +160,18 @@ async def execute_block_test(block: Block): if field_name in block.test_credentials: extra_exec_kwargs[field_name] = block.test_credentials[field_name] + # Handle auto-generated credentials (e.g., from GoogleDriveFileInput) + auto_creds_fields = input_model.get_auto_credentials_fields() + if auto_creds_fields and block.test_credentials: + if isinstance(block.test_credentials, _BaseCredentials): + # Single credentials object - use for all auto_credentials kwargs + for kwarg_name in auto_creds_fields.keys(): + extra_exec_kwargs[kwarg_name] = block.test_credentials + elif isinstance(block.test_credentials, dict): + for kwarg_name in auto_creds_fields.keys(): + if kwarg_name in block.test_credentials: + extra_exec_kwargs[kwarg_name] = block.test_credentials[kwarg_name] + for input_data in block.test_input: log.info(f"{prefix} in: {input_data}") diff --git a/autogpt_platform/frontend/src/app/globals.css b/autogpt_platform/frontend/src/app/globals.css index f969c9db79..1f782f753b 100644 --- a/autogpt_platform/frontend/src/app/globals.css +++ b/autogpt_platform/frontend/src/app/globals.css @@ -150,3 +150,16 @@ input[type="number"]::-webkit-inner-spin-button { input[type="number"] { -moz-appearance: textfield; } + +/* Google Drive Picker: ensure picker appears above dialogs and can receive clicks */ +[class*="picker-dialog"] { + z-index: 10000 !important; + pointer-events: auto !important; +} + +/* When Google picker is open, lower dialog z-index so picker renders on top */ +body[data-google-picker-open="true"] [data-dialog-overlay], +body[data-google-picker-open="true"] [data-dialog-content] { + z-index: 1 !important; + pointer-events: none !important; +} diff --git a/autogpt_platform/frontend/src/components/__legacy__/ui/dialog.tsx b/autogpt_platform/frontend/src/components/__legacy__/ui/dialog.tsx index 2aa089315b..4ce998b6f6 100644 --- a/autogpt_platform/frontend/src/components/__legacy__/ui/dialog.tsx +++ b/autogpt_platform/frontend/src/components/__legacy__/ui/dialog.tsx @@ -14,12 +14,20 @@ const DialogPortal = DialogPrimitive.Portal; const DialogClose = DialogPrimitive.Close; +/** + * Check if an external picker (like Google Drive) is currently open. + */ +function isExternalPickerOpen(): boolean { + return document.body.hasAttribute("data-google-picker-open"); +} + const DialogOverlay = React.forwardRef< React.ElementRef, React.ComponentPropsWithoutRef >(({ className, ...props }, ref) => ( , React.ComponentPropsWithoutRef ->(({ className, children, ...props }, ref) => ( - - - - {children} - - - Close - - - -)); +>( + ( + { + className, + children, + onPointerDownOutside, + onInteractOutside, + onFocusOutside, + ...props + }, + ref, + ) => ( + + + { + if (isExternalPickerOpen()) { + e.preventDefault(); + return; + } + onPointerDownOutside?.(e); + }} + onInteractOutside={(e) => { + if (isExternalPickerOpen()) { + e.preventDefault(); + return; + } + onInteractOutside?.(e); + }} + onFocusOutside={(e) => { + if (isExternalPickerOpen()) { + e.preventDefault(); + return; + } + onFocusOutside?.(e); + }} + className={cn( + "fixed left-[50%] top-[50%] z-50 grid w-full max-w-lg translate-x-[-50%] translate-y-[-50%] gap-4 border border-neutral-200 bg-white p-6 shadow-lg duration-200 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[state=closed]:slide-out-to-left-1/2 data-[state=closed]:slide-out-to-top-[48%] data-[state=open]:slide-in-from-left-1/2 data-[state=open]:slide-in-from-top-[48%] dark:border-neutral-800 dark:bg-neutral-950 sm:rounded-lg", + className, + )} + {...props} + > + {children} + + + Close + + + + ), +); DialogContent.displayName = DialogPrimitive.Content.displayName; const DialogHeader = ({ diff --git a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePicker.tsx b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePicker.tsx index 0500d08549..eaa44a9452 100644 --- a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePicker.tsx +++ b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePicker.tsx @@ -3,7 +3,12 @@ import { CredentialsInput } from "@/app/(platform)/library/agents/[id]/components/NewAgentLibraryView/components/modals/CredentialsInputs/CredentialsInputs"; import { Button } from "@/components/atoms/Button/Button"; import { CircleNotchIcon, FolderOpenIcon } from "@phosphor-icons/react"; -import { Props, useGoogleDrivePicker } from "./useGoogleDrivePicker"; +import { + Props as BaseProps, + useGoogleDrivePicker, +} from "./useGoogleDrivePicker"; + +export type Props = BaseProps; export function GoogleDrivePicker(props: Props) { const { diff --git a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePickerInput.tsx b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePickerInput.tsx index a1accbada5..1db9809de2 100644 --- a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePickerInput.tsx +++ b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/GoogleDrivePickerInput.tsx @@ -24,6 +24,9 @@ export function GoogleDrivePickerInput({ }: GoogleDrivePickerInputProps) { const [pickerError, setPickerError] = React.useState(null); const isMultiSelect = config.multiselect || false; + const hasAutoCredentials = !!config.auto_credentials; + + // Strip _credentials_id from value for display purposes const currentFiles = isMultiSelect ? Array.isArray(value) ? value @@ -33,25 +36,34 @@ export function GoogleDrivePickerInput({ : []; const handlePicked = useCallback( - (files: any[]) => { + (files: any[], credentialId?: string) => { // Clear any previous picker errors setPickerError(null); // Convert to GoogleDriveFile format - const convertedFiles = files.map((f) => ({ - id: f.id, - name: f.name, - mimeType: f.mimeType, - url: f.url, - iconUrl: f.iconUrl, - isFolder: f.mimeType === "application/vnd.google-apps.folder", - })); + const convertedFiles = files.map((f) => { + const file: any = { + id: f.id, + name: f.name, + mimeType: f.mimeType, + url: f.url, + iconUrl: f.iconUrl, + isFolder: f.mimeType === "application/vnd.google-apps.folder", + }; + + // Include _credentials_id when auto_credentials is configured + if (hasAutoCredentials && credentialId) { + file._credentials_id = credentialId; + } + + return file; + }); // Store based on multiselect mode const newValue = isMultiSelect ? convertedFiles : convertedFiles[0]; onChange(newValue); }, - [isMultiSelect, onChange], + [isMultiSelect, onChange, hasAutoCredentials], ); const handleRemoveFile = useCallback( @@ -79,6 +91,7 @@ export function GoogleDrivePickerInput({ views={config.allowed_views || ["DOCS"]} scopes={config.scopes || ["https://www.googleapis.com/auth/drive.file"]} disabled={false} + requirePlatformCredentials={hasAutoCredentials} onPicked={handlePicked} onCanceled={() => { // User canceled - no action needed diff --git a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/useGoogleDrivePicker.ts b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/useGoogleDrivePicker.ts index 323cd6d9d6..66386882c6 100644 --- a/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/useGoogleDrivePicker.ts +++ b/autogpt_platform/frontend/src/components/contextual/GoogleDrivePicker/useGoogleDrivePicker.ts @@ -34,7 +34,9 @@ export type Props = { disableThumbnails?: boolean; buttonText?: string; disabled?: boolean; - onPicked: (files: NormalizedPickedFile[]) => void; + /** When true, requires saved platform credentials (no consent flow fallback) */ + requirePlatformCredentials?: boolean; + onPicked: (files: NormalizedPickedFile[], credentialId?: string) => void; onCanceled: () => void; onError: (err: unknown) => void; }; @@ -65,6 +67,7 @@ export function useGoogleDrivePicker(options: Props) { const accessTokenRef = useRef(null); const tokenClientRef = useRef(null); const pickerReadyRef = useRef(false); + const usedCredentialIdRef = useRef(undefined); const credentials = useCredentials(getCredentialsSchema(requestedScopes)); const queryClient = useQueryClient(); const isReady = pickerReadyRef.current && !!tokenClientRef.current; @@ -114,6 +117,7 @@ export function useGoogleDrivePicker(options: Props) { ) { const credentialId = selectedCredential?.id || credentials.savedCredentials[0].id; + usedCredentialIdRef.current = credentialId; try { const queryOptions = getGetV1GetSpecificCredentialByIdQueryOptions( @@ -178,6 +182,20 @@ export function useGoogleDrivePicker(options: Props) { } } + // If platform credentials are required but none exist, show error + if (options?.requirePlatformCredentials) { + const error = new Error( + "Please connect your Google account in Settings before using this feature.", + ); + toast({ + title: "Google Account Required", + description: error.message, + variant: "destructive", + }); + if (onError) onError(error); + return; + } + const token = accessTokenRef.current || (await requestAccessToken()); buildAndShowPicker(token); } catch (e) { @@ -242,6 +260,24 @@ export function useGoogleDrivePicker(options: Props) { } function buildAndShowPicker(accessToken: string): void { + if (!developerKey) { + const error = new Error( + "Missing Google Drive Picker Configuration: developer key is not set", + ); + console.error("[useGoogleDrivePicker]", error.message); + onError(error); + return; + } + + if (!appId) { + const error = new Error( + "Missing Google Drive Picker Configuration: app ID is not set", + ); + console.error("[useGoogleDrivePicker]", error.message); + onError(error); + return; + } + const gp = window.google!.picker!; const builder = new gp.PickerBuilder() @@ -269,19 +305,40 @@ export function useGoogleDrivePicker(options: Props) { }); const picker = builder.build(); + + // Mark picker as open - prevents parent dialogs from closing on outside clicks + document.body.setAttribute("data-google-picker-open", "true"); + picker.setVisible(true); } function handlePickerData(data: any): void { + // Google Picker fires callback on multiple events: LOADED, PICKED, CANCEL + // Only remove the marker and process when picker is actually closed (PICKED or CANCEL) + const gp = window.google?.picker; + if (!gp || !data) return; + + const action = data[gp.Response.ACTION]; + + // Ignore LOADED action - picker is still open + // Note: gp.Action.LOADED exists at runtime but not in types + if (action === "loaded") { + return; + } + + // Remove the marker when picker closes (PICKED or CANCEL) + document.body.removeAttribute("data-google-picker-open"); + try { const files = normalizePickerResponse(data); if (files.length) { - onPicked(files); + // Pass the credential ID that was used for this picker session + onPicked(files, usedCredentialIdRef.current); } else { onCanceled(); } } catch (e) { - if (onError) onError(e); + onError(e); } } @@ -307,5 +364,6 @@ export function useGoogleDrivePicker(options: Props) { accessToken: accessTokenRef.current, selectedCredential, setSelectedCredential, + usedCredentialId: usedCredentialIdRef.current, }; } diff --git a/autogpt_platform/frontend/src/components/molecules/Dialog/components/DialogWrap.tsx b/autogpt_platform/frontend/src/components/molecules/Dialog/components/DialogWrap.tsx index 378aaa9f08..ae4dc8e065 100644 --- a/autogpt_platform/frontend/src/components/molecules/Dialog/components/DialogWrap.tsx +++ b/autogpt_platform/frontend/src/components/molecules/Dialog/components/DialogWrap.tsx @@ -4,6 +4,7 @@ import * as RXDialog from "@radix-ui/react-dialog"; import { CSSProperties, PropsWithChildren, + useCallback, useEffect, useRef, useState, @@ -20,6 +21,14 @@ interface Props extends BaseProps { withGradient?: boolean; } +/** + * Check if an external picker (like Google Drive) is currently open. + * Used to prevent dialog from closing when user interacts with the picker. + */ +function isExternalPickerOpen(): boolean { + return document.body.hasAttribute("data-google-picker-open"); +} + export function DialogWrap({ children, title, @@ -30,6 +39,30 @@ export function DialogWrap({ const scrollRef = useRef(null); const [hasVerticalScrollbar, setHasVerticalScrollbar] = useState(false); + // Prevent dialog from closing when external picker is open + const handleInteractOutside = useCallback( + (event: Event) => { + if (isExternalPickerOpen()) { + event.preventDefault(); + return; + } + handleClose(); + }, + [handleClose], + ); + + const handlePointerDownOutside = useCallback((event: Event) => { + if (isExternalPickerOpen()) { + event.preventDefault(); + } + }, []); + + const handleFocusOutside = useCallback((event: Event) => { + if (isExternalPickerOpen()) { + event.preventDefault(); + } + }, []); + useEffect(() => { function update() { const el = scrollRef.current; @@ -48,12 +81,15 @@ export function DialogWrap({ return ( - +