diff --git a/src/notebooks/deepnote/deepnoteCellCopyHandler.ts b/src/notebooks/deepnote/deepnoteCellCopyHandler.ts new file mode 100644 index 0000000000..e5889dfa55 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteCellCopyHandler.ts @@ -0,0 +1,401 @@ +import { injectable, inject } from 'inversify'; +import { + workspace, + NotebookDocumentChangeEvent, + NotebookEdit, + WorkspaceEdit, + commands, + window, + NotebookCellData, + NotebookRange, + env, + NotebookCellOutputItem, + NotebookCellOutput +} from 'vscode'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IDisposableRegistry } from '../../platform/common/types'; +import { logger } from '../../platform/logging'; +import { generateBlockId, generateSortingKey } from './dataConversionUtils'; + +/** + * Marker prefix for clipboard data to identify Deepnote cell metadata + */ +const CLIPBOARD_MARKER = '___DEEPNOTE_CELL_METADATA___'; + +/** + * Interface for cell metadata stored in clipboard + */ +interface ClipboardCellMetadata { + metadata: Record; + kind: number; + languageId: string; + value: string; +} + +/** + * Handles cell copy operations in Deepnote notebooks to ensure metadata is preserved. + * + * VSCode's built-in copy commands don't preserve custom cell metadata, so this handler + * intercepts copy/cut/paste commands and stores metadata in the clipboard as JSON. + * This allows metadata to be preserved across copy/paste and cut/paste operations. + */ +@injectable() +export class DeepnoteCellCopyHandler implements IExtensionSyncActivationService { + private processingChanges = false; + + constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {} + + public activate(): void { + // Override built-in notebook copy/cut commands to preserve metadata for Deepnote notebooks + this.disposables.push(commands.registerCommand('notebook.cell.copyDown', () => this.copyCellDownInterceptor())); + this.disposables.push(commands.registerCommand('notebook.cell.copyUp', () => this.copyCellUpInterceptor())); + this.disposables.push(commands.registerCommand('notebook.cell.copy', () => this.copyCellInterceptor())); + this.disposables.push(commands.registerCommand('notebook.cell.cut', () => this.cutCellInterceptor())); + this.disposables.push(commands.registerCommand('notebook.cell.paste', () => this.pasteCellInterceptor())); + + // Listen for notebook document changes to detect when cells are added without metadata + this.disposables.push(workspace.onDidChangeNotebookDocument((e) => this.onDidChangeNotebookDocument(e))); + } + + /** + * Interceptor for the built-in notebook.cell.copyDown command. + * Routes to our custom implementation for Deepnote notebooks. + */ + private async copyCellDownInterceptor(): Promise { + const editor = window.activeNotebookEditor; + if (editor && editor.notebook && editor.notebook.notebookType === 'deepnote') { + await this.copyCellDown(); + } else { + await commands.executeCommand('default:notebook.cell.copyDown'); + } + } + + /** + * Interceptor for the built-in notebook.cell.copyUp command. + * Routes to our custom implementation for Deepnote notebooks. + */ + private async copyCellUpInterceptor(): Promise { + const editor = window.activeNotebookEditor; + if (editor && editor.notebook && editor.notebook.notebookType === 'deepnote') { + await this.copyCellUp(); + } else { + await commands.executeCommand('default:notebook.cell.copyUp'); + } + } + + /** + * Interceptor for the built-in notebook.cell.copy command. + * Stores cell metadata in clipboard for Deepnote notebooks. + */ + private async copyCellInterceptor(): Promise { + const editor = window.activeNotebookEditor; + if (editor && editor.notebook && editor.notebook.notebookType === 'deepnote') { + await this.copyCellToClipboard({ isCut: false }); + } else { + await commands.executeCommand('default:notebook.cell.copy'); + } + } + + /** + * Interceptor for the built-in notebook.cell.cut command. + * Stores cell metadata in clipboard for Deepnote notebooks. + */ + private async cutCellInterceptor(): Promise { + const editor = window.activeNotebookEditor; + if (editor && editor.notebook && editor.notebook.notebookType === 'deepnote') { + await this.copyCellToClipboard({ isCut: true }); + } else { + await commands.executeCommand('default:notebook.cell.cut'); + } + } + + /** + * Interceptor for the built-in notebook.cell.paste command. + * Restores cell metadata from clipboard for Deepnote notebooks. + */ + private async pasteCellInterceptor(): Promise { + const editor = window.activeNotebookEditor; + if (editor && editor.notebook && editor.notebook.notebookType === 'deepnote') { + await this.pasteCellFromClipboard(); + } else { + await commands.executeCommand('default:notebook.cell.paste'); + } + } + + private async copyCellDown(): Promise { + await this.copyCellAtOffset(1); + } + + private async copyCellUp(): Promise { + await this.copyCellAtOffset(-1); + } + + /** + * Copy a cell at a specific offset from the current cell. + * @param offset -1 for copy up, 1 for copy down + */ + private async copyCellAtOffset(offset: number): Promise { + const editor = window.activeNotebookEditor; + + if (!editor || !editor.notebook || editor.notebook.notebookType !== 'deepnote') { + logger.debug(`copyCellAtOffset called for non-Deepnote notebook`); + return; + } + + const selection = editor.selection; + if (!selection) { + return; + } + + const cellToCopy = editor.notebook.cellAt(selection.start); + const insertIndex = offset > 0 ? selection.start + 1 : selection.start; + + // Create a new cell with the same content and metadata + const newCell = new NotebookCellData( + cellToCopy.kind, + cellToCopy.document.getText(), + cellToCopy.document.languageId + ); + + // Copy all metadata (ID and sortingKey will be generated by onDidChangeNotebookDocument) + if (cellToCopy.metadata) { + const copiedMetadata = structuredClone(cellToCopy.metadata); + newCell.metadata = copiedMetadata; + logger.debug('DeepnoteCellCopyHandler: Copying cell with metadata preserved'); + } + + // Copy outputs if present + if (cellToCopy.outputs.length > 0) { + newCell.outputs = cellToCopy.outputs.map( + (o) => + new NotebookCellOutput( + o.items.map((i) => new NotebookCellOutputItem(i.data, i.mime)), + o.metadata ? structuredClone(o.metadata) : undefined + ) + ); + } + + // Insert the new cell + const edit = new WorkspaceEdit(); + edit.set(editor.notebook.uri, [NotebookEdit.insertCells(insertIndex, [newCell])]); + + const success = await workspace.applyEdit(edit); + + if (success) { + // Move selection to the new cell + editor.selection = new NotebookRange(insertIndex, insertIndex + 1); + logger.debug(`DeepnoteCellCopyHandler: Successfully copied cell to index ${insertIndex}`); + } else { + logger.warn('DeepnoteCellCopyHandler: Failed to copy cell'); + } + } + + private async onDidChangeNotebookDocument(e: NotebookDocumentChangeEvent): Promise { + // Only process Deepnote notebooks + if (e.notebook.notebookType !== 'deepnote') { + return; + } + + // Avoid recursive processing + if (this.processingChanges) { + return; + } + + // Check for cell additions (which includes copies) + for (const change of e.contentChanges) { + if (change.addedCells.length === 0) { + continue; + } + + // When cells are copied, VSCode should preserve metadata automatically. + // However, we need to ensure that: + // 1. Each cell has a unique ID + // 2. The sortingKey is updated based on the new position + // 3. All other metadata (including sql_integration_id) is preserved + + const cellsNeedingMetadataFix: Array<{ index: number; metadata: Record }> = []; + + for (const cell of change.addedCells) { + const metadata = cell.metadata || {}; + + // Log the metadata to see what's actually being copied + logger.debug('DeepnoteCellCopyHandler: Cell added with metadata'); + + // Only process Deepnote cells (cells with type or pocket metadata) + if (!metadata.type && !metadata.__deepnotePocket) { + continue; + } + + const cellIndex = e.notebook.getCells().indexOf(cell); + + if (cellIndex === -1) { + continue; + } + + // Check if this cell needs metadata updates + // We update the ID and sortingKey for all added Deepnote cells to ensure uniqueness + const updatedMetadata = { ...metadata }; + + // Generate new ID for the cell (important for copied cells) + updatedMetadata.id = generateBlockId(); + + // Update sortingKey based on the new position + if (updatedMetadata.__deepnotePocket) { + updatedMetadata.__deepnotePocket = { + ...updatedMetadata.__deepnotePocket, + sortingKey: generateSortingKey(cellIndex) + }; + } else if (updatedMetadata.sortingKey) { + updatedMetadata.sortingKey = generateSortingKey(cellIndex); + } + + // All other metadata (including sql_integration_id) is preserved from the original metadata + cellsNeedingMetadataFix.push({ + index: cellIndex, + metadata: updatedMetadata + }); + + logger.debug( + `DeepnoteCellCopyHandler: Updated metadata for ${metadata.type} cell at index ${cellIndex}` + ); + } + + // Apply metadata fixes if needed + if (cellsNeedingMetadataFix.length > 0) { + await this.applyMetadataFixes(e.notebook.uri, cellsNeedingMetadataFix); + } + } + } + + private async applyMetadataFixes( + notebookUri: import('vscode').Uri, + fixes: Array<{ index: number; metadata: Record }> + ): Promise { + try { + this.processingChanges = true; + + const edit = new WorkspaceEdit(); + + // Create all the edits at once instead of calling set() multiple times + const edits = fixes.map((fix) => NotebookEdit.updateCellMetadata(fix.index, fix.metadata)); + edit.set(notebookUri, edits); + + const success = await workspace.applyEdit(edit); + + if (success) { + logger.debug(`DeepnoteCellCopyHandler: Successfully updated metadata for ${fixes.length} cell(s)`); + } else { + logger.warn(`DeepnoteCellCopyHandler: Failed to apply metadata fixes for ${fixes.length} cell(s)`); + } + } catch (error) { + logger.error('DeepnoteCellCopyHandler: Error applying metadata fixes', error); + } finally { + this.processingChanges = false; + } + } + + /** + * Copy or cut a cell to the clipboard with metadata preserved. + * @param isCut Whether this is a cut operation (will delete the cell after copying) + */ + private async copyCellToClipboard(params: { isCut: boolean }): Promise { + const editor = window.activeNotebookEditor; + + if (!editor || !editor.notebook || editor.notebook.notebookType !== 'deepnote') { + logger.debug(`copyCellToClipboard called for non-Deepnote notebook`); + return; + } + + const selection = editor.selection; + if (!selection) { + return; + } + + const cellToCopy = editor.notebook.cellAt(selection.start); + + // Create clipboard data with all cell information + const clipboardData: ClipboardCellMetadata = { + metadata: cellToCopy.metadata || {}, + kind: cellToCopy.kind, + languageId: cellToCopy.document.languageId, + value: cellToCopy.document.getText() + }; + + // Store in clipboard as JSON with marker + const clipboardText = `${CLIPBOARD_MARKER}${JSON.stringify(clipboardData)}`; + await env.clipboard.writeText(clipboardText); + + logger.debug(`DeepnoteCellCopyHandler: ${params.isCut ? 'Cut' : 'Copied'} cell to clipboard with metadata`); + + // If this is a cut operation, delete the cell + if (params.isCut) { + const edit = new WorkspaceEdit(); + edit.set(editor.notebook.uri, [ + NotebookEdit.deleteCells(new NotebookRange(selection.start, selection.start + 1)) + ]); + await workspace.applyEdit(edit); + logger.debug(`DeepnoteCellCopyHandler: Deleted cell after cut operation`); + } + } + + /** + * Paste a cell from the clipboard, restoring metadata if available. + */ + private async pasteCellFromClipboard(): Promise { + const editor = window.activeNotebookEditor; + + if (!editor || !editor.notebook || editor.notebook.notebookType !== 'deepnote') { + logger.debug(`pasteCellFromClipboard called for non-Deepnote notebook`); + return; + } + + const selection = editor.selection; + if (!selection) { + return; + } + + // Read from clipboard + const clipboardText = await env.clipboard.readText(); + + // Check if clipboard contains our metadata marker + if (!clipboardText.startsWith(CLIPBOARD_MARKER)) { + await commands.executeCommand('default:notebook.cell.paste'); + return; + } + + try { + // Parse clipboard data + const jsonText = clipboardText.substring(CLIPBOARD_MARKER.length); + const clipboardData: ClipboardCellMetadata = JSON.parse(jsonText); + + // Create new cell with preserved metadata + const newCell = new NotebookCellData(clipboardData.kind, clipboardData.value, clipboardData.languageId); + + const insertIndex = selection.start; + + // Copy metadata (ID and sortingKey will be generated by onDidChangeNotebookDocument) + const copiedMetadata = { ...clipboardData.metadata }; + newCell.metadata = copiedMetadata; + + logger.debug('DeepnoteCellCopyHandler: Copying cell with metadata preserved'); + + // Insert the new cell + const edit = new WorkspaceEdit(); + edit.set(editor.notebook.uri, [NotebookEdit.insertCells(insertIndex, [newCell])]); + + const success = await workspace.applyEdit(edit); + + if (success) { + // Move selection to the new cell + editor.selection = new NotebookRange(insertIndex, insertIndex + 1); + logger.debug(`DeepnoteCellCopyHandler: Successfully pasted cell at index ${insertIndex}`); + } else { + logger.warn('DeepnoteCellCopyHandler: Failed to paste cell'); + } + } catch (error) { + logger.error('DeepnoteCellCopyHandler: Error parsing clipboard data', error); + await commands.executeCommand('default:notebook.cell.paste'); + } + } +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index 3bd4f9d236..1d51e66450 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -67,6 +67,7 @@ import { DeepnoteServerProvider } from '../kernels/deepnote/deepnoteServerProvid import { DeepnoteInitNotebookRunner, IDeepnoteInitNotebookRunner } from './deepnote/deepnoteInitNotebookRunner.node'; import { DeepnoteRequirementsHelper, IDeepnoteRequirementsHelper } from './deepnote/deepnoteRequirementsHelper.node'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; +import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { registerControllerTypes(serviceManager, isDevMode); @@ -152,6 +153,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlIntegrationStartupCodeProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteCellCopyHandler + ); // Deepnote kernel services serviceManager.addSingleton(IDeepnoteToolkitInstaller, DeepnoteToolkitInstaller);