Use the Foyle API to report Log Events

package.json

@@ -1153,8 +1153,8 @@
     "@aws-sdk/client-eks": "^3.635.0",
     "@aws-sdk/credential-providers": "^3.635.0",
     "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240809200651-8507e5a24938.4",
-    "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-20240819224840-b69cd71876b1.1",
-    "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-20240819224840-b69cd71876b1.3",
+    "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-00000000000000-d14934cb2733.1",
+    "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-00000000000000-d14934cb2733.3",
     "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240826183545-20a8540bddaf.4",
     "@connectrpc/connect": "^1.4.0",
     "@connectrpc/connect-node": "^1.1.2",

src/extension/ai/converters.ts

@@ -3,6 +3,7 @@
 // See ../vscode_apis.md for an exlanation. It is very helpful for understanding this folder.
 
 import * as vscode from 'vscode'
+import * as parser_pb from '@buf/stateful_runme.bufbuild_es/runme/parser/v1/parser_pb'
 
 import { ServerLifecycleIdentity, getServerConfigurationValue } from '../../utils/configuration'
 import { Serializer } from '../../types'
@@ -59,3 +60,28 @@ export function cellProtosToCellData(cells: serializerTypes.Cell[]): vscode.Note
   let newCellData = serializer.SerializerBase.revive(notebook, identity)
   return newCellData
 }
+
+// vsCellsToESProto converts a VSCode NotebookCell to a RunMe Cell proto
+// Its not quite the inverse of cellProtosToCellData because this function
+// uses the ES client library as opposed to the TS client library.
+// The reason we don't rely on the serialization routines is we don't want to
+// generate an RPC just to convert the cell to a proto.
+export function vsCellsToESProtos(cells: vscode.NotebookCell[]): parser_pb.Cell[] {
+  const cellProtos: parser_pb.Cell[] = []
+
+  for (let cell of cells) {
+    const cellProto = new parser_pb.Cell()
+    if (cell.kind === vscode.NotebookCellKind.Code) {
+      cellProto.kind = parser_pb.CellKind.CODE
+    } else if (cell.kind === vscode.NotebookCellKind.Markup) {
+      cellProto.kind = parser_pb.CellKind.MARKUP
+    }
+
+    cellProto.value = cell.document.getText()
+    cellProto.metadata = cell.metadata
+
+    cellProtos.push(cellProto)
+  }
+
+  return cellProtos
+}

src/extension/ai/events.ts

@@ -0,0 +1,98 @@
+import { PromiseClient } from '@connectrpc/connect'
+import { AIService } from '@buf/jlewi_foyle.connectrpc_es/foyle/v1alpha1/agent_connect'
+import {
+  LogEventsRequest,
+  LogEventType,
+  LogEvent,
+} from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb'
+import * as vscode from 'vscode'
+import { ulid } from 'ulidx'
+
+import { RUNME_CELL_ID } from '../constants'
+import getLogger from '../logger'
+
+import { SessionManager } from './sessions'
+import * as converters from './converters'
+
+// Interface for the event reporter
+// This allows us to swap in a null op logger when AI isn't enabled
+export interface IEventReporter {
+  reportExecution(cell: vscode.NotebookCell): Promise<void>
+  reportEvents(events: LogEvent[]): Promise<void>
+}
+
+// EventReporter handles reporting events to the AI service
+export class EventReporter implements IEventReporter {
+  client: PromiseClient<typeof AIService>
+  log: ReturnType<typeof getLogger>
+
+  constructor(client: PromiseClient<typeof AIService>) {
+    this.client = client
+    this.log = getLogger('AIEventReporter')
+  }
+
+  async reportExecution(cell: vscode.NotebookCell) {
+    const contextCells: vscode.NotebookCell[] = []
+
+    // Include some previous cells as context.
+    // N.B. In principle we shouldn't need to send any additional context because we
+    // set the context id. So as soon as we put the focus on the execution cell we should
+    // start a streaming generate request which will include the entire notebook or a large portion of it.
+    // However, we still send some additional context here for two reasons
+    // 1. Help us verify that sending context ids is working correctly.
+    // 2. Its possible in the future we start rate limiting streaming generate requests and don't want to rely on it
+    // for providing the context of the cell execution.
+    let startIndex = cell.index - 1
+    if (startIndex < 0) {
+      startIndex = 0
+    }
+    for (let i = startIndex; i < cell.index; i++) {
+      contextCells.push(cell.notebook.cellAt(i))
+    }
+
+    contextCells.push(cell)
+    const cells = converters.vsCellsToESProtos(contextCells)
+    const event = new LogEvent()
+    event.selectedId = cell.metadata?.[RUNME_CELL_ID]
+    event.selectedIndex = cell.index
+    event.type = LogEventType.EXECUTE
+    event.cells = cells
+    event.contextId = SessionManager.getManager().getID()
+    return this.reportEvents([event])
+  }
+
+  async reportEvents(events: LogEvent[]) {
+    const req = new LogEventsRequest()
+    req.events = events
+    for (const event of events) {
+      if (event.eventId === '') {
+        event.eventId = ulid()
+      }
+    }
+    await this.client.logEvents(req).catch((e) => {
+      this.log.error(`Failed to log event; error: ${e}`)
+    })
+  }
+}
+
+// NullOpEventReporter is a null op implementation of the event reporter
+export class NullOpEventReporter implements IEventReporter {
+  async reportExecution(_cell: vscode.NotebookCell) {
+    // Do nothing
+  }
+
+  async reportEvents(_events: LogEvent[]) {
+    // Do nothing
+  }
+}
+
+let _globalReporter = new NullOpEventReporter()
+
+// getEventReporter returns the global event reporter
+export function getEventReporter(): IEventReporter {
+  return _globalReporter
+}
+
+export function setEventReporter(reporter: IEventReporter) {
+  _globalReporter = reporter
+}

src/extension/ai/ghost.ts

@@ -1,13 +1,15 @@
 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'
+import { RUNME_CELL_ID } from '../constants'
 
 import * as converters from './converters'
 import * as stream from './stream'
 import * as protos from './protos'
+import { SessionManager } from './sessions'
+import { getEventReporter } from './events'
 
 const log = getLogger()
 
@@ -39,13 +41,11 @@ export class GhostCellGenerator implements stream.CompletionHandlers {
   // 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
@@ -109,7 +109,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers {
 
       let notebookProto = serializer.GrpcSerializer.marshalNotebook(notebookData)
       let request = new agent_pb.StreamGenerateRequest({
-        contextId: this.contextID,
+        contextId: SessionManager.getManager().getID(),
         request: {
           case: 'fullContext',
           value: new agent_pb.FullContext({
@@ -129,7 +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,
+        contextId: SessionManager.getManager().getID(),
         request: {
           case: 'update',
           value: new agent_pb.UpdateContext({
@@ -144,10 +144,10 @@ 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) {
+    if (response.contextId !== SessionManager.getManager().getID()) {
       // 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}`,
+        `Ignoring response with contextID ${response.contextId} because it doesn't match the current contextID ${SessionManager.getManager().getID()}`,
       )
       return
     }
@@ -215,11 +215,11 @@ 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
+    const oldCID = SessionManager.getManager().getID()
     // We need to generate a new context ID because the context has changed.
-    this.contextID = ulid()
+    const contextID = SessionManager.getManager().newID()
     log.info(
-      `onDidChangeActiveTextEditor fired: editor: ${editor?.document.uri}; new contextID: ${this.contextID}; old contextID: ${oldCID}`,
+      `onDidChangeActiveTextEditor fired: editor: ${editor?.document.uri}; new contextID: ${contextID}; old contextID: ${oldCID}`,
     )
     if (editor === undefined) {
       return
@@ -253,6 +253,13 @@ export class GhostCellGenerator implements stream.CompletionHandlers {
     // If the cell is a ghost cell we want to remove the decoration
     // and replace it with a non-ghost cell.
     editorAsNonGhost(editor)
+
+    const event = new agent_pb.LogEvent()
+    event.type = agent_pb.LogEventType.ACCEPTED
+    event.contextId = oldCID
+    event.selectedId = cell.metadata?.[RUNME_CELL_ID]
+    event.selectedIndex = cell.index
+    getEventReporter().reportEvents([event])
   }
 
   shutdown(): void {