From 004b49fd309c08aa044d728a9e522eb55ccd07ca Mon Sep 17 00:00:00 2001 From: Jeremy Lewi Date: Thu, 24 Oct 2024 14:38:12 -0700 Subject: [PATCH] Squashed commit of the following: branch: jlewi/ghost_markup https://github.com/stateful/vscode-runme/pull/1744 Trigger GhostCells On Cell Execution --- package-lock.json | 14 +++--- package.json | 4 +- src/extension/ai/ghost.ts | 95 ++++++++++++++++++++++++++++--------- src/extension/ai/manager.ts | 14 ++++-- src/extension/ai/stream.ts | 6 ++- 5 files changed, 97 insertions(+), 36 deletions(-) diff --git a/package-lock.json b/package-lock.json index df233e1ac..736079ece 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,8 +15,8 @@ "@aws-sdk/client-eks": "^3.670.0", "@aws-sdk/credential-providers": "^3.670.0", "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240904201038-e0425cfebb28.4", - "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20241002214001-8c67df189216.1", - "@buf/jlewi_foyle.connectrpc_es": "^1.5.0-20241002214001-8c67df189216.1", + "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20241018224331-ac296bc95724.1", + "@buf/jlewi_foyle.connectrpc_es": "^1.5.0-20241018224331-ac296bc95724.1", "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240913234806-45813f39881a.4", "@connectrpc/connect": "^1.4.0", "@connectrpc/connect-node": "^1.6.1", @@ -2402,8 +2402,8 @@ } }, "node_modules/@buf/jlewi_foyle.bufbuild_es": { - "version": "1.10.0-20241002214001-8c67df189216.1", - "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.10.0-20241002214001-8c67df189216.1.tgz", + "version": "1.10.0-20241018224331-ac296bc95724.1", + "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.10.0-20241018224331-ac296bc95724.1.tgz", "dependencies": { "@buf/bufbuild_protovalidate.bufbuild_es": "1.10.0-20240212200630-3014d81c3a48.1", "@buf/stateful_runme.bufbuild_es": "1.10.0-20240913234806-45813f39881a.1" @@ -2413,11 +2413,11 @@ } }, "node_modules/@buf/jlewi_foyle.connectrpc_es": { - "version": "1.5.0-20241002214001-8c67df189216.1", - "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.connectrpc_es/-/jlewi_foyle.connectrpc_es-1.5.0-20241002214001-8c67df189216.1.tgz", + "version": "1.5.0-20241018224331-ac296bc95724.1", + "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.connectrpc_es/-/jlewi_foyle.connectrpc_es-1.5.0-20241018224331-ac296bc95724.1.tgz", "dependencies": { "@buf/bufbuild_protovalidate.connectrpc_es": "1.5.0-20240212200630-3014d81c3a48.1", - "@buf/jlewi_foyle.bufbuild_es": "1.10.0-20241002214001-8c67df189216.1", + "@buf/jlewi_foyle.bufbuild_es": "1.10.0-20241018224331-ac296bc95724.1", "@buf/stateful_runme.connectrpc_es": "1.5.0-20240913234806-45813f39881a.1" }, "peerDependencies": { diff --git a/package.json b/package.json index 61783b7b7..a822f1500 100644 --- a/package.json +++ b/package.json @@ -1266,8 +1266,8 @@ "@aws-sdk/client-eks": "^3.670.0", "@aws-sdk/credential-providers": "^3.670.0", "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240904201038-e0425cfebb28.4", - "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20241002214001-8c67df189216.1", - "@buf/jlewi_foyle.connectrpc_es": "^1.5.0-20241002214001-8c67df189216.1", + "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20241018224331-ac296bc95724.1", + "@buf/jlewi_foyle.connectrpc_es": "^1.5.0-20241018224331-ac296bc95724.1", "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240913234806-45813f39881a.4", "@connectrpc/connect": "^1.4.0", "@connectrpc/connect-node": "^1.6.1", diff --git a/src/extension/ai/ghost.ts b/src/extension/ai/ghost.ts index d5dde8e51..3ff3c8c84 100644 --- a/src/extension/ai/ghost.ts +++ b/src/extension/ai/ghost.ts @@ -1,5 +1,6 @@ import * as vscode from 'vscode' import * as agent_pb from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb' +import { StreamGenerateRequest_Trigger } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb' import getLogger from '../logger' import * as serializer from '../serializer' @@ -19,6 +20,14 @@ const log = getLogger() export const ghostKey = '_ghostCell' export const ghostCellKindKey = '_ghostCellKind' +// Schemes are defined at +// https://github.com/microsoft/vscode/blob/a56879c50db91715377005d6182d12742d1ba5c7/src/vs/base/common/network.ts#L64 +export const vsCodeCellScheme = 'vscode-notebook-cell' + +// vsCodeOutputScheme is the scheme for the output window (not cell outputs). +// The output window is where the logs for Runme are displayed. +export const vsCodeOutputScheme = 'output' + const ghostDecoration = vscode.window.createTextEditorDecorationType({ color: '#888888', // Light grey color }) @@ -88,12 +97,16 @@ export class GhostCellGenerator implements stream.CompletionHandlers { let matchedCell = notebook.cellAt(cellChangeEvent.cellIndex) // Has the cell changed since the last time we processed an event? + // TODO(https://github.com/jlewi/foyle/issues/312): I think there's an edge case where we don't + // correctly detect that the cell has changed and a new stream needs to be initiated. let newCell = true if (nbState.activeCell?.document.uri === matchedCell?.document.uri) { newCell = false } - log.info(`buildRequest: is newCell: ${newCell} , firstRequest: ${firstRequest}`) + log.info( + `buildRequest: is newCell: ${newCell} , firstRequest: ${firstRequest}, trigger ${cellChangeEvent.trigger}`, + ) // Update notebook state nbState.activeCell = matchedCell @@ -119,6 +132,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { notebookUri: notebook.uri.toString(), }), }, + trigger: cellChangeEvent.trigger, }) return request @@ -137,6 +151,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { cell: notebook.cells[0], }), }, + trigger: cellChangeEvent.trigger, }) return request @@ -236,6 +251,8 @@ export class GhostCellGenerator implements stream.CompletionHandlers { return } + // Are schemes defined here: + // https://github.com/microsoft/vscode/blob/a56879c50db91715377005d6182d12742d1ba5c7/src/vs/base/common/network.ts#L64 if (editor.document.uri.scheme !== 'vscode-notebook-cell') { // Doesn't correspond to a notebook cell so do nothing return @@ -305,8 +322,7 @@ export class CellChangeEventGenerator { } handleOnDidChangeNotebookCell = (event: vscode.TextDocumentChangeEvent) => { - if (event.document.uri.scheme !== 'vscode-notebook-cell') { - // ignore other open events + if (![vsCodeCellScheme].includes(event.document.uri.scheme)) { return } var matchedCell: vscode.NotebookCell | undefined @@ -332,31 +348,66 @@ export class CellChangeEventGenerator { } this.streamCreator.handleEvent( - new stream.CellChangeEvent(notebook.uri.toString(), matchedCell.index), + new stream.CellChangeEvent( + notebook.uri.toString(), + matchedCell.index, + StreamGenerateRequest_Trigger.CELL_TEXT_CHANGE, + ), ) } -} -// handleOnDidChangeVisibleTextEditors is called when the visible text editors change. -// This includes when a TextEditor is created. I also think it can be the result of scrolling. -export function handleOnDidChangeVisibleTextEditors(editors: readonly vscode.TextEditor[]) { - for (const editor of editors) { - log.info(`onDidChangeVisibleTextEditors Fired for editor ${editor.document.uri}`) - if (editor.document.uri.scheme !== 'vscode-notebook-cell') { - log.info(`onDidChangeVisibleTextEditors Fired fo ${editor.document.uri}`) - // Doesn't correspond to a notebook cell so do nothing - continue - } - const cell = getCellFromCellDocument(editor.document) - if (cell === undefined) { - continue - } + // handleOnDidChangeVisibleTextEditors is called when the visible text editors change. + // This includes when a TextEditor is created. I also think it can be the result of scrolling. + // When cells become visible we need to apply ghost decorations. + // + // This event is also fired when a code cell is executed and its output becomes visible. + // We use that to trigger completion generation because we want the newly rendered code + // cell output to affect the suggestions. + handleOnDidChangeVisibleTextEditors = (editors: readonly vscode.TextEditor[]) => { + for (const editor of editors) { + if (![vsCodeCellScheme].includes(editor.document.uri.scheme)) { + // Doesn't correspond to a notebook or output cell so do nothing + continue + } - if (!isGhostCell(cell)) { - continue + const cell = getCellFromCellDocument(editor.document) + if (cell === undefined) { + continue + } + + if (!isGhostCell(cell)) { + continue + } + + editorAsGhost(editor) } + } - editorAsGhost(editor) + handleOnDidChangeNotebookDocument = (event: vscode.NotebookDocumentChangeEvent) => { + event.cellChanges.forEach((change) => { + if (change.outputs !== undefined) { + // If outputs change then we want to trigger completions. + + // N.B. It looks like if you click the "configure" button associated with a cell then it will trigger + // an output change event. I don't think there's any easy way to filter those events out. To filter + // those events out we'd need to keep track of the output item with mime type + // application/vnd.code.notebook.stdout and then detect when the stdout changes. That would require + // keeping track of that state. If we trigger on the "configure" then we send a request to the Foyle + // server and we can rely on the Foyle server to do the debouncing. + + // It is the responsibility of the StreamCreator to decide whether the change should be processed.. + // In particular its possible that the cell that changed is not the active cell. Therefore + // we may not want to generate completions for it. For example, you can have multiple cells + // running. So in principle the active cell could be different from the cell that changed. + this.streamCreator.handleEvent( + new stream.CellChangeEvent( + change.cell.notebook.uri.toString(), + change.cell.index, + StreamGenerateRequest_Trigger.CELL_OUTPUT_CHANGE, + ), + ) + } + }) } } diff --git a/src/extension/ai/manager.ts b/src/extension/ai/manager.ts index a2b16727e..b0050c1f1 100644 --- a/src/extension/ai/manager.ts +++ b/src/extension/ai/manager.ts @@ -16,6 +16,7 @@ export class AIManager { subscriptions: vscode.Disposable[] = [] client: PromiseClient completionGenerator: generate.CompletionGenerator + constructor() { this.log = getLogger('AIManager') this.log.info('AI: Initializing AI Manager') @@ -52,7 +53,7 @@ export class AIManager { let eventGenerator = new ghost.CellChangeEventGenerator(creator) // onDidChangeTextDocument fires when the contents of a cell changes. - // We use this to generate completions. + // We use this to generate completions when the contents of a cell changes. this.subscriptions.push( vscode.workspace.onDidChangeTextDocument(eventGenerator.handleOnDidChangeNotebookCell), ) @@ -61,13 +62,20 @@ export class AIManager { // This can happen due to scrolling. // We need to trap this event to apply decorations to turn cells into ghost cells. this.subscriptions.push( - vscode.window.onDidChangeVisibleTextEditors(ghost.handleOnDidChangeVisibleTextEditors), + vscode.window.onDidChangeVisibleTextEditors( + eventGenerator.handleOnDidChangeVisibleTextEditors, + ), ) // When a cell is selected we want to check if its a ghost cell and if so render it a non-ghost cell. this.subscriptions.push( vscode.window.onDidChangeActiveTextEditor(cellGenerator.handleOnDidChangeActiveTextEditor), - // vscode.window.onDidChangeActiveTextEditor(localOnDidChangeActiveTextEditor), + ) + + this.subscriptions.push( + vscode.workspace.onDidChangeNotebookDocument( + eventGenerator.handleOnDidChangeNotebookDocument, + ), ) } diff --git a/src/extension/ai/stream.ts b/src/extension/ai/stream.ts index 863a33274..44a530165 100644 --- a/src/extension/ai/stream.ts +++ b/src/extension/ai/stream.ts @@ -3,6 +3,7 @@ import { AIService } from '@buf/jlewi_foyle.connectrpc_es/foyle/v1alpha1/agent_c import { StreamGenerateRequest, StreamGenerateResponse, + StreamGenerateRequest_Trigger, } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb' import getLogger from '../logger' @@ -167,10 +168,11 @@ export class StreamCreator { export class CellChangeEvent { public notebookUri: string public cellIndex: number - - constructor(notebookUri: string, cellIndex: number) { + public trigger: StreamGenerateRequest_Trigger + constructor(notebookUri: string, cellIndex: number, trigger: StreamGenerateRequest_Trigger) { this.notebookUri = notebookUri this.cellIndex = cellIndex + this.trigger = trigger } }