Trigger Completion Generation On Cell Execution

src/extension/ai/events.ts

@@ -5,6 +5,7 @@ import {
   LogEventType,
   LogEvent,
   LogEvent_ExecuteStatus,
+  StreamGenerateRequest_Trigger,
 } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb'
 import * as vscode from 'vscode'
 import { ulid } from 'ulidx'
@@ -15,6 +16,7 @@ import getLogger from '../logger'
 
 import { SessionManager } from './sessions'
 import * as converters from './converters'
+import * as stream from './stream'
 
 // Interface for the event reporter
 // This allows us to swap in a null op logger when AI isn't enabled
@@ -29,10 +31,12 @@ export class EventReporter implements vscode.Disposable, IEventReporter {
   log: ReturnType<typeof getLogger>
   queue: Subject<LogEvent[]> = new Subject()
   subscription: Subscription
+  streamCreator: stream.StreamCreator
 
-  constructor(client: PromiseClient<typeof AIService>) {
+  constructor(client: PromiseClient<typeof AIService>, streamCreator: stream.StreamCreator) {
     this.client = client
     this.log = getLogger('AIEventReporter')
+    this.streamCreator = streamCreator
 
     this.subscription = this.queue
       .pipe(
@@ -93,6 +97,13 @@ export class EventReporter implements vscode.Disposable, IEventReporter {
     } else {
       event.executeStatus = LogEvent_ExecuteStatus.FAILED
     }
+    // Fire an event to trigger the AI service
+    const cellChangeEvent = new stream.CellChangeEvent(
+      cell.notebook.uri.toString(),
+      cell.index,
+      StreamGenerateRequest_Trigger.CELL_EXECUTE,
+    )
+    this.streamCreator.handleEvent(cellChangeEvent)
     return this.reportEvents([event])
   }
 

src/extension/ai/ghost.ts

@@ -124,6 +124,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers {
       let notebookProto = await this.converter.notebookDataToProto(notebookData)
       let request = new agent_pb.StreamGenerateRequest({
         contextId: SessionManager.getManager().getID(),
+        trigger: cellChangeEvent.trigger,
         request: {
           case: 'fullContext',
           value: new agent_pb.FullContext({
@@ -142,6 +143,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers {
       let notebookProto = await this.converter.notebookDataToProto(notebookData)
       let request = new agent_pb.StreamGenerateRequest({
         contextId: SessionManager.getManager().getID(),
+        trigger: cellChangeEvent.trigger,
         request: {
           case: 'update',
           value: new agent_pb.UpdateContext({
@@ -377,40 +379,6 @@ export class CellChangeEventGenerator {
       editorAsGhost(editor)
     }
   }
-
-  handleOnDidChangeNotebookDocument = (event: vscode.NotebookDocumentChangeEvent) => {
-    // N.B. For non-interactive cells this will trigger each time the output is updated.
-    // For interactive cells this doesn't appear to trigger each time the cell output is updated.
-    // For example, if you have a long running command (e.g. a bash for loop with a sleep that
-    // echos a message on each iteration) then this won't trigger on each iteration for
-    // an interactive cell but will for non-interactive.
-    event.cellChanges.forEach(async (change) => {
-      log.info(`handleOnDidChangeNotebookDocument: change: ${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.
-        //
-        await this.streamCreator.handleEvent(
-          new stream.CellChangeEvent(
-            change.cell.notebook.uri.toString(),
-            change.cell.index,
-            StreamGenerateRequest_Trigger.CELL_OUTPUT_CHANGE,
-          ),
-        )
-      }
-    })
-  }
 }
 
 // editorAsGhost decorates an editor as a ghost cell.

src/extension/ai/manager.ts

@@ -33,12 +33,6 @@ export class AIManager implements vscode.Disposable {
     this.completionGenerator = new generate.CompletionGenerator(this.client, this.converter)
     if (autoComplete) {
       this.registerGhostCellEvents()
-
-      const reporter = new events.EventReporter(this.client)
-      this.subscriptions.push(reporter)
-
-      // Update the global event reporter to use the AI service
-      events.setEventReporter(reporter)
     }
   }
 
@@ -61,6 +55,12 @@ export class AIManager implements vscode.Disposable {
     // and turns each window into an AsyncIterable of streaming requests.
     let creator = new stream.StreamCreator(cellGenerator, this.client)
 
+    const reporter = new events.EventReporter(this.client, creator)
+    this.subscriptions.push(reporter)
+
+    // Update the global event reporter to use the AI service
+    events.setEventReporter(reporter)
+
     let eventGenerator = new ghost.CellChangeEventGenerator(creator)
     // onDidChangeTextDocument fires when the contents of a cell changes.
     // We use this to generate completions when the contents of a cell changes.
@@ -82,14 +82,6 @@ export class AIManager implements vscode.Disposable {
       vscode.window.onDidChangeActiveTextEditor(cellGenerator.handleOnDidChangeActiveTextEditor),
     )
 
-    // We use onDidChangeNotebookDocument to listen for changes to outputs.
-    // We use this to trigger updates in response to a cell's output being updated.
-    this.subscriptions.push(
-      vscode.workspace.onDidChangeNotebookDocument(
-        eventGenerator.handleOnDidChangeNotebookDocument,
-      ),
-    )
-
     // Create a new status bar item aligned to the right
     let statusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 100)
     statusBarItem.text = 'Session: <None>'