diff --git a/__mocks__/vscode.ts b/__mocks__/vscode.ts index ce828ce9a..32cd00253 100644 --- a/__mocks__/vscode.ts +++ b/__mocks__/vscode.ts @@ -93,6 +93,7 @@ export const window = { showInformationMessage: vi.fn(), showErrorMessage: vi.fn(), createTerminal: vi.fn().mockReturnValue(terminal), + createTextEditorDecorationType: vi.fn(), showNotebookDocument: vi.fn(), showTextDocument: vi.fn(), onDidChangeActiveNotebookEditor: vi.fn().mockReturnValue({ dispose: vi.fn() }), diff --git a/package-lock.json b/package-lock.json index beace0b79..59bf04109 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,8 +15,8 @@ "@aws-sdk/client-eks": "^3.630.0", "@aws-sdk/credential-providers": "^3.630.0", "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240709201032-5be6b6661793.4", - "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20240801044500-fba3b5ceb26b.1", - "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-20240801044500-fba3b5ceb26b.3", + "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20240818004946-86c4c8cbe96c.1", + "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-20240818004946-86c4c8cbe96c.3", "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240731204114-2df61af8e022.4", "@connectrpc/connect": "^1.4.0", "@connectrpc/connect-node": "^1.1.2", @@ -2418,8 +2418,8 @@ } }, "node_modules/@buf/jlewi_foyle.bufbuild_es": { - "version": "1.10.0-20240801044500-fba3b5ceb26b.1", - "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.10.0-20240801044500-fba3b5ceb26b.1.tgz", + "version": "1.10.0-20240818004946-86c4c8cbe96c.1", + "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.10.0-20240818004946-86c4c8cbe96c.1.tgz", "dependencies": { "@buf/bufbuild_protovalidate.bufbuild_es": "1.10.0-20240212200630-3014d81c3a48.1", "@buf/googleapis_googleapis.bufbuild_es": "1.10.0-20240729201550-8bc2c51e08c4.1", @@ -2430,12 +2430,12 @@ } }, "node_modules/@buf/jlewi_foyle.connectrpc_es": { - "version": "1.4.0-20240801044500-fba3b5ceb26b.3", - "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.connectrpc_es/-/jlewi_foyle.connectrpc_es-1.4.0-20240801044500-fba3b5ceb26b.3.tgz", + "version": "1.4.0-20240818004946-86c4c8cbe96c.3", + "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.connectrpc_es/-/jlewi_foyle.connectrpc_es-1.4.0-20240818004946-86c4c8cbe96c.3.tgz", "dependencies": { "@buf/bufbuild_protovalidate.connectrpc_es": "1.4.0-20240212200630-3014d81c3a48.3", "@buf/googleapis_googleapis.connectrpc_es": "1.4.0-20240729201550-8bc2c51e08c4.3", - "@buf/jlewi_foyle.bufbuild_es": "1.7.2-20240801044500-fba3b5ceb26b.2", + "@buf/jlewi_foyle.bufbuild_es": "1.7.2-20240818004946-86c4c8cbe96c.2", "@buf/stateful_runme.connectrpc_es": "1.4.0-20240731204114-2df61af8e022.3" }, "peerDependencies": { @@ -2457,8 +2457,8 @@ } }, "node_modules/@buf/jlewi_foyle.connectrpc_es/node_modules/@buf/jlewi_foyle.bufbuild_es": { - "version": "1.7.2-20240801044500-fba3b5ceb26b.2", - "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.7.2-20240801044500-fba3b5ceb26b.2.tgz", + "version": "1.7.2-20240818004946-86c4c8cbe96c.2", + "resolved": "https://buf.build/gen/npm/v1/@buf/jlewi_foyle.bufbuild_es/-/jlewi_foyle.bufbuild_es-1.7.2-20240818004946-86c4c8cbe96c.2.tgz", "dependencies": { "@buf/bufbuild_protovalidate.bufbuild_es": "1.7.2-20240212200630-3014d81c3a48.2", "@buf/googleapis_googleapis.bufbuild_es": "1.7.2-20240729201550-8bc2c51e08c4.2", diff --git a/package.json b/package.json index 01589569c..3539b2e95 100644 --- a/package.json +++ b/package.json @@ -1146,8 +1146,8 @@ "@aws-sdk/client-eks": "^3.630.0", "@aws-sdk/credential-providers": "^3.630.0", "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240709201032-5be6b6661793.4", - "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20240801044500-fba3b5ceb26b.1", - "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-20240801044500-fba3b5ceb26b.3", + "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20240818004946-86c4c8cbe96c.1", + "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-20240818004946-86c4c8cbe96c.3", "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240731204114-2df61af8e022.4", "@connectrpc/connect": "^1.4.0", "@connectrpc/connect-node": "^1.1.2", diff --git a/src/extension/ai/ghost.ts b/src/extension/ai/ghost.ts index 002139653..d30e6ae15 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 { ulid } from 'ulidx' import getLogger from '../logger' import * as serializer from '../serializer' @@ -15,6 +16,10 @@ const log = getLogger() // the ghost metadata is not persisted to the markdown file. export const ghostKey = '_ghostCell' +const ghostDecoration = vscode.window.createTextEditorDecorationType({ + color: '#888888', // Light grey color +}) + // TODO(jeremy): How do we handle multiple notebooks? Arguably you should only be generating // completions for the active notebook. So as soon as the active notebook changes we should // stop generating completions for the old notebook and start generating completions for the new notebook. @@ -31,8 +36,16 @@ export const ghostKey = '_ghostCell' export class GhostCellGenerator implements stream.CompletionHandlers { private notebookState: Map + // contextID is the ID of the context we are generating completions for. + // It is used to detect whether a completion response is stale and should be + // discarded because the context has changed. + private contextID: string + constructor() { this.notebookState = new Map() + // Generate a random context ID. This should be unnecessary because presumable the event to change + // the active cell will be sent before any requests are sent but it doesn't hurt to be safe. + this.contextID = ulid() } // Updated method to check and initialize notebook state @@ -96,6 +109,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { let notebookProto = serializer.GrpcSerializer.marshalNotebook(notebookData) let request = new agent_pb.StreamGenerateRequest({ + contextId: this.contextID, request: { case: 'fullContext', value: new agent_pb.FullContext({ @@ -115,6 +129,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { let notebook = protos.notebookTSToES(notebookProto) // Generate an update request let request = new agent_pb.StreamGenerateRequest({ + contextId: this.contextID, request: { case: 'update', value: new agent_pb.UpdateContext({ @@ -129,6 +144,14 @@ export class GhostCellGenerator implements stream.CompletionHandlers { // processResponse applies the changes from the response to the notebook. processResponse(response: agent_pb.StreamGenerateResponse) { + if (response.contextId !== this.contextID) { + // TODO(jeremy): Is this logging too verbose? + log.info( + `Ignoring response with contextID ${response.contextId} because it doesn't match the current contextID ${this.contextID}`, + ) + return + } + let cellsTs = protos.cellsESToTS(response.cells) let newCellData = converters.cellProtosToCellData(cellsTs) @@ -189,6 +212,49 @@ export class GhostCellGenerator implements stream.CompletionHandlers { }) } + // handleOnDidChangeActiveTextEditor updates the ghostKey cell decoration and rendering + // when it is selected + handleOnDidChangeActiveTextEditor = (editor: vscode.TextEditor | undefined) => { + const oldCID = this.contextID + // We need to generate a new context ID because the context has changed. + this.contextID = ulid() + log.info( + `onDidChangeActiveTextEditor fired: editor: ${editor?.document.uri}; new contextID: ${this.contextID}; old contextID: ${oldCID}`, + ) + if (editor === undefined) { + return + } + + if (editor.document.uri.scheme !== 'vscode-notebook-cell') { + // Doesn't correspond to a notebook cell so do nothing + return + } + + const cell = getCellFromCellDocument(editor.document) + if (cell === undefined) { + return + } + + if (!isGhostCell(cell)) { + return + } + // ...cell.metadata creates a shallow copy of the metadata object + const updatedMetadata = { ...cell.metadata, [ghostKey]: false } + const update = vscode.NotebookEdit.updateCellMetadata(cell.index, updatedMetadata) + const edit = new vscode.WorkspaceEdit() + edit.set(editor.document.uri, [update]) + vscode.workspace.applyEdit(edit).then((result: boolean) => { + log.trace(`updateCellMetadata to deactivate ghostcell resolved with ${result}`) + if (!result) { + log.error('applyEdit failed') + return + } + }) + // If the cell is a ghost cell we want to remove the decoration + // and replace it with a non-ghost cell. + editorAsNonGhost(editor) + } + shutdown(): void { log.info('Shutting down') } @@ -239,11 +305,6 @@ export class CellChangeEventGenerator { return } - log.info( - `onDidChangeTextDocument Fired for notebook ${event.document.uri}` + - 'this should fire when a cell is added to a notebook', - ) - this.streamCreator.handleEvent( new stream.CellChangeEvent(notebook.uri.toString(), matchedCell.index), ) @@ -251,7 +312,7 @@ export class CellChangeEventGenerator { } // handleOnDidChangeVisibleTextEditors is called when the visible text editors change. -// This includes when a TextEditor is created. +// 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}`) @@ -266,7 +327,6 @@ export function handleOnDidChangeVisibleTextEditors(editors: readonly vscode.Tex } if (!isGhostCell(cell)) { - log.info(`Not a ghost editor doc:${cell.document.uri}`) continue } @@ -282,14 +342,18 @@ function editorAsGhost(editor: vscode.TextEditor) { textDoc.positionAt(textDoc.getText().length), ) - editor.setDecorations(getGhostDecoration(), [range]) + editor.setDecorations(ghostDecoration, [range]) } function editorAsNonGhost(editor: vscode.TextEditor) { // To remove the decoration we set the range to an empty range and pass in a reference // to the original decoration // https://github.com/microsoft/vscode-extension-samples/blob/main/decorator-sample/USAGE.md#tips - editor.setDecorations(getGhostDecoration(), []) + // + // Important: ghostDecoration must be a reference to the same object that was used to create the decoration. + // that's how VSCode knows which decoration to remove. If you use a "copy" (i.e. a decoration with the same value) + // the decoration won't get removed. + editor.setDecorations(ghostDecoration, []) } function isGhostCell(cell: vscode.NotebookCell): boolean { @@ -313,51 +377,3 @@ function getCellFromCellDocument(textDoc: vscode.TextDocument): vscode.NotebookC }) return matchedCell } - -// handleOnDidChangeActiveTextEditor updates the ghostKey cell decoration and rendering -// when it is selected -export function handleOnDidChangeActiveTextEditor(editor: vscode.TextEditor | undefined) { - log.info(`onDidChangeActiveTextEditor Fired for editor ${editor?.document.uri}`) - if (editor === undefined) { - return - } - - if (editor.document.uri.scheme !== 'vscode-notebook-cell') { - // Doesn't correspond to a notebook cell so do nothing - return - } - - const cell = getCellFromCellDocument(editor.document) - if (cell === undefined) { - return - } - - if (!isGhostCell(cell)) { - return - } - // ...cell.metadata creates a shallow copy of the metadata object - const updatedMetadata = { ...cell.metadata, [ghostKey]: false } - const update = vscode.NotebookEdit.updateCellMetadata(cell.index, updatedMetadata) - const edit = new vscode.WorkspaceEdit() - edit.set(editor.document.uri, [update]) - vscode.workspace.applyEdit(edit).then((result: boolean) => { - log.trace(`updateCellMetadata to deactivate ghostcell resolved with ${result}`) - if (!result) { - log.error('applyEdit failed') - return - } - }) - // If the cell is a ghost cell we want to remove the decoration - // and replace it with a non-ghost cell. - editorAsNonGhost(editor) -} - -// n.b. this is a function and not a top level const because that causes problems with the vitest -// mocking framework. -// N.B. I think we could potentially have solved that by doing something like -// https://github.com/stateful/vscode-runme/pull/1475#issuecomment-2278636467 -function getGhostDecoration(): vscode.TextEditorDecorationType { - return vscode.window.createTextEditorDecorationType({ - color: '#888888', // Light grey color - }) -} diff --git a/src/extension/ai/manager.ts b/src/extension/ai/manager.ts index d1ca36429..7fffece8f 100644 --- a/src/extension/ai/manager.ts +++ b/src/extension/ai/manager.ts @@ -43,6 +43,7 @@ export class AIManager { ) // onDidChangeVisibleTextEditors fires when the visible text editors change. + // 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), @@ -50,7 +51,8 @@ export class AIManager { // 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(ghost.handleOnDidChangeActiveTextEditor), + vscode.window.onDidChangeActiveTextEditor(cellGenerator.handleOnDidChangeActiveTextEditor), + // vscode.window.onDidChangeActiveTextEditor(localOnDidChangeActiveTextEditor), ) } diff --git a/tests/extension/serializer.test.ts b/tests/extension/serializer.test.ts index 57c2bd771..fb1dade68 100644 --- a/tests/extension/serializer.test.ts +++ b/tests/extension/serializer.test.ts @@ -29,6 +29,7 @@ vi.mock('vscode', () => ({ activeNotebookEditor: undefined, showErrorMessage: vi.fn().mockResolvedValue({}), createOutputChannel: vi.fn(), + createTextEditorDecorationType: vi.fn().mockResolvedValue({}), }, Uri: { joinPath: vi.fn().mockReturnValue('/foo/bar') }, workspace: {