GhostCells: Fix removal of ghost rendering and prevent stale inserts

__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() }),

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",

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<vscode.Uri, NotebookState>
 
+  // 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<vscode.Uri, NotebookState>()
+    // 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,19 +305,14 @@ 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),
     )
   }
 }
 
 // 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
-  })
-}

src/extension/ai/manager.ts

@@ -43,14 +43,16 @@ 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),
     )
 
     // 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),
     )
   }
 

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: {