diff --git a/src/notebooks/deepnote/staleOutputStatusBarProvider.ts b/src/notebooks/deepnote/staleOutputStatusBarProvider.ts new file mode 100644 index 000000000..7fb099a6e --- /dev/null +++ b/src/notebooks/deepnote/staleOutputStatusBarProvider.ts @@ -0,0 +1,173 @@ +import { + CancellationToken, + EventEmitter, + NotebookCell, + NotebookCellKind, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookDocumentChangeEvent, + NotebookEdit, + ProviderResult, + WorkspaceEdit, + l10n, + notebooks, + workspace +} from 'vscode'; +import { inject, injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { computeHash } from '../../platform/common/crypto'; +import { IDisposableRegistry } from '../../platform/common/types'; +import { Pocket } from '../../platform/deepnote/pocket'; +import { notebookCellExecutions, NotebookCellExecutionState } from '../../platform/notebooks/cellExecutionStateService'; + +/** + * Provides status bar items for cells with stale outputs. + * An output is considered stale when the cell content has changed since the output was generated. + */ +@injectable() +export class StaleOutputStatusBarProvider + implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService +{ + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {} + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + // Refresh when any Deepnote notebook changes (e.g., cell content edited) + this.disposables.push( + workspace.onDidChangeNotebookDocument((e: NotebookDocumentChangeEvent) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + // Listen for cell execution completion to update contentHash + this.disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell.notebook.notebookType === 'deepnote' && e.state === NotebookCellExecutionState.Idle) { + void this.handleCellExecutionComplete(e.cell); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): ProviderResult { + if (token?.isCancellationRequested) { + return; + } + + if (cell.kind !== NotebookCellKind.Code) { + return; + } + + if (cell.outputs.length === 0) { + return; + } + + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + const storedHash = pocket?.contentHash; + + // If no contentHash exists, the cell was never executed (by us), don't show indicator + if (!storedHash) { + return; + } + + return this.checkAndCreateStatusBarItem(cell, storedHash); + } + + private async checkAndCreateStatusBarItem( + cell: NotebookCell, + storedHash: string + ): Promise { + const currentContent = cell.document.getText(); + const currentHash = `sha256:${await computeHash(currentContent, 'SHA-256')}`; + + const cellKey = this.getCellKey(cell); + const cachedExecutedHash = this.executedContentHashes.get(cellKey); + + if (cachedExecutedHash && cachedExecutedHash === currentHash) { + return; + } + + if (currentHash === storedHash) { + return; + } + + return { + text: `$(warning) ${l10n.t('Output may be stale')}`, + alignment: 1, // NotebookCellStatusBarAlignment.Left + priority: 85, + tooltip: l10n.t( + 'Cell content has changed since outputs were generated.\nRe-run the cell to update outputs.' + ), + command: { + title: l10n.t('Re-run cell'), + command: 'notebook.cell.execute', + arguments: [{ ranges: [{ start: cell.index, end: cell.index + 1 }], document: cell.notebook.uri }] + } + }; + } + + /** + * In-memory cache of content hashes at the time of execution. + * This provides immediate access to the hash without waiting for metadata persistence. + * Key format: cell document URI string + */ + private readonly executedContentHashes = new Map(); + + private getCellKey(cell: NotebookCell): string { + return cell.document.uri.toString(); + } + + /** + * Updates the cell's contentHash when execution completes successfully. + * This ensures the stale indicator is removed after re-running a cell. + */ + private async handleCellExecutionComplete(cell: NotebookCell): Promise { + if (cell.kind !== NotebookCellKind.Code) { + return; + } + + if (cell.executionSummary?.success === false) { + return; + } + + try { + const content = cell.document.getText(); + const hash = await computeHash(content, 'SHA-256'); + const newContentHash = `sha256:${hash}`; + + const cellKey = this.getCellKey(cell); + + this.executedContentHashes.set(cellKey, newContentHash); + this._onDidChangeCellStatusBarItems.fire(); + + // Also persist the hash to cell metadata for cross-session persistence + const existingPocket = (cell.metadata?.__deepnotePocket as Record) || {}; + const updatedPocket = { ...existingPocket, contentHash: newContentHash }; + + const edit = new WorkspaceEdit(); + const updatedMetadata = { + ...cell.metadata, + __deepnotePocket: updatedPocket + }; + + edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]); + + await workspace.applyEdit(edit); + } catch { + // Silently ignore errors - the stale indicator will just remain + } + } +} diff --git a/src/notebooks/deepnote/staleOutputStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/staleOutputStatusBarProvider.unit.test.ts new file mode 100644 index 000000000..a92355e5d --- /dev/null +++ b/src/notebooks/deepnote/staleOutputStatusBarProvider.unit.test.ts @@ -0,0 +1,245 @@ +import { expect } from 'chai'; +import { anything, verify, when } from 'ts-mockito'; + +import { CancellationToken, NotebookCell, NotebookCellKind, NotebookCellOutput, NotebookDocument, Uri } from 'vscode'; + +import { computeHash } from '../../platform/common/crypto'; +import type { IDisposableRegistry } from '../../platform/common/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { StaleOutputStatusBarProvider } from './staleOutputStatusBarProvider'; + +suite('StaleOutputStatusBarProvider', () => { + let provider: StaleOutputStatusBarProvider; + let mockDisposables: IDisposableRegistry; + let mockToken: CancellationToken; + + setup(() => { + mockDisposables = [] as any; + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new StaleOutputStatusBarProvider(mockDisposables); + }); + + function createMockCell(options: { + kind?: NotebookCellKind; + outputs?: NotebookCellOutput[]; + metadata?: Record; + text?: string; + }): NotebookCell { + const notebookUri = Uri.file('/test/notebook.deepnote'); + + return { + index: 0, + notebook: { + uri: notebookUri, + notebookType: 'deepnote' + } as NotebookDocument, + kind: options.kind ?? NotebookCellKind.Code, + document: { + uri: Uri.file('/test/notebook.deepnote#cell0'), + fileName: '/test/notebook.deepnote#cell0', + isUntitled: false, + languageId: 'python', + version: 1, + isDirty: false, + isClosed: false, + getText: () => options.text ?? 'print("hello")', + save: async () => true, + eol: 1, + lineCount: 1, + lineAt: () => ({ text: '' }) as any, + offsetAt: () => 0, + positionAt: () => ({}) as any, + validateRange: () => ({}) as any, + validatePosition: () => ({}) as any + } as any, + metadata: options.metadata || {}, + outputs: options.outputs ?? [], + executionSummary: undefined + } as any; + } + + function createMockOutput(): NotebookCellOutput { + return { + items: [{ mime: 'text/plain', data: new Uint8Array() }], + metadata: undefined + } as NotebookCellOutput; + } + + suite('provideCellStatusBarItems', () => { + test('should return undefined when cancellation token is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + + const cell = createMockCell({ + outputs: [createMockOutput()], + metadata: { __deepnotePocket: { contentHash: 'sha256:abc123' } } + }); + + const result = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(result).to.be.undefined; + }); + + test('should return undefined for markdown cells', () => { + const cell = createMockCell({ + kind: NotebookCellKind.Markup, + outputs: [createMockOutput()], + metadata: { __deepnotePocket: { contentHash: 'sha256:abc123' } } + }); + + const result = provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.be.undefined; + }); + + test('should return undefined for cells without outputs', () => { + const cell = createMockCell({ + outputs: [], + metadata: { __deepnotePocket: { contentHash: 'sha256:abc123' } } + }); + + const result = provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.be.undefined; + }); + + test('should return undefined when no contentHash in pocket (never executed)', () => { + const cell = createMockCell({ + outputs: [createMockOutput()], + metadata: { __deepnotePocket: {} } + }); + + const result = provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.be.undefined; + }); + + test('should return undefined when no pocket exists', () => { + const cell = createMockCell({ + outputs: [createMockOutput()], + metadata: {} + }); + + const result = provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.be.undefined; + }); + + test('should return undefined when hashes match (not stale)', async () => { + const cellContent = 'print("hello")'; + const hash = await computeHash(cellContent, 'SHA-256'); + const storedHash = `sha256:${hash}`; + + const cell = createMockCell({ + text: cellContent, + outputs: [createMockOutput()], + metadata: { __deepnotePocket: { contentHash: storedHash } } + }); + + const result = await provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.be.undefined; + }); + + test('should return warning indicator when hashes differ (stale)', async () => { + const originalContent = 'print("original")'; + const originalHash = await computeHash(originalContent, 'SHA-256'); + const storedHash = `sha256:${originalHash}`; + + // Cell content has changed since execution + const cell = createMockCell({ + text: 'print("modified")', + outputs: [createMockOutput()], + metadata: { __deepnotePocket: { contentHash: storedHash } } + }); + + const result = await provider.provideCellStatusBarItems(cell, mockToken); + + expect(result).to.not.be.undefined; + expect((result as any).text).to.include('Output may be stale'); + expect((result as any).alignment).to.equal(1); // Left alignment + expect((result as any).priority).to.equal(85); + expect((result as any).tooltip).to.include('Cell content has changed'); + expect((result as any).command.command).to.equal('notebook.cell.execute'); + }); + }); + + suite('activate', () => { + setup(() => { + resetVSCodeMocks(); + when( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider(anything(), anything()) + ).thenReturn({ dispose: () => undefined } as any); + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything())).thenReturn({ + dispose: () => undefined + } as any); + }); + + teardown(() => { + resetVSCodeMocks(); + }); + + test('registers notebook cell status bar provider for deepnote notebooks', () => { + provider.activate(); + + verify( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider('deepnote', provider) + ).once(); + }); + + test('registers workspace.onDidChangeNotebookDocument listener', () => { + provider.activate(); + + verify(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything())).once(); + }); + + test('fires status bar update event when deepnote notebook changes', () => { + let changeHandler: ((e: any) => void) | undefined; + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything())).thenCall((handler) => { + changeHandler = handler; + + return { dispose: () => undefined }; + }); + + let eventFired = false; + provider.onDidChangeCellStatusBarItems(() => { + eventFired = true; + }); + + provider.activate(); + expect(changeHandler).to.not.be.undefined; + + // Fire the event with a deepnote notebook + changeHandler!({ notebook: { notebookType: 'deepnote' } }); + + expect(eventFired).to.be.true; + }); + + test('does not fire status bar update event for non-deepnote notebooks', () => { + let changeHandler: ((e: any) => void) | undefined; + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything())).thenCall((handler) => { + changeHandler = handler; + + return { dispose: () => undefined }; + }); + + let eventFired = false; + provider.onDidChangeCellStatusBarItems(() => { + eventFired = true; + }); + + provider.activate(); + expect(changeHandler).to.not.be.undefined; + + // Fire the event with a jupyter notebook (not deepnote) + changeHandler!({ notebook: { notebookType: 'jupyter-notebook' } }); + + expect(eventFired).to.be.false; + }); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index 53ac77c50..2b881a4f8 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -58,6 +58,7 @@ import { IPlatformDeepnoteNotebookManager } from '../platform/notebooks/deepnote/types'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; +import { StaleOutputStatusBarProvider } from './deepnote/staleOutputStatusBarProvider'; import { IDeepnoteToolkitInstaller, IDeepnoteServerStarter, @@ -179,6 +180,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + StaleOutputStatusBarProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, SqlIntegrationStartupCodeProvider