From c1fa5c945b1cd12c2d80c4db8b85c8243cd01db1 Mon Sep 17 00:00:00 2001
From: Jeremy Lewi <jeremy@lewi.us>
Date: Fri, 25 Oct 2024 13:17:40 -0700
Subject: [PATCH] Report cell execution status to Foyle. (#1745)

* This updates the extension to report the status of the cell execution to Foyle when
  reporting execution events.

* This is highly valuable signal for the AI. This will help us from learning from
  broken/invalid commands.

* Fix jlewi/foyle#310

Co-authored-by: Sebastian Tiedtke <sebastiantiedtke@gmail.com>
---
 src/extension/ai/events.ts | 13 ++++++++++---
 src/extension/kernel.ts    |  8 ++++----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/extension/ai/events.ts b/src/extension/ai/events.ts
index a4924237c..71b9e4e5c 100644
--- a/src/extension/ai/events.ts
+++ b/src/extension/ai/events.ts
@@ -4,6 +4,7 @@ import {
   LogEventsRequest,
   LogEventType,
   LogEvent,
+  LogEvent_ExecuteStatus,
 } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb'
 import * as vscode from 'vscode'
 import { ulid } from 'ulidx'
@@ -17,7 +18,7 @@ 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>
+  reportExecution(cell: vscode.NotebookCell, status: boolean): Promise<void>
   reportEvents(events: LogEvent[]): Promise<void>
 }
 
@@ -31,7 +32,7 @@ export class EventReporter implements IEventReporter {
     this.log = getLogger('AIEventReporter')
   }
 
-  async reportExecution(cell: vscode.NotebookCell) {
+  async reportExecution(cell: vscode.NotebookCell, executionSuccess: boolean) {
     const contextCells: vscode.NotebookCell[] = []
 
     // Include some previous cells as context.
@@ -58,6 +59,12 @@ export class EventReporter implements IEventReporter {
     event.type = LogEventType.EXECUTE
     event.cells = cells
     event.contextId = SessionManager.getManager().getID()
+
+    if (executionSuccess) {
+      event.executeStatus = LogEvent_ExecuteStatus.SUCCEEDED
+    } else {
+      event.executeStatus = LogEvent_ExecuteStatus.FAILED
+    }
     return this.reportEvents([event])
   }
 
@@ -77,7 +84,7 @@ export class EventReporter implements IEventReporter {
 
 // NullOpEventReporter is a null op implementation of the event reporter
 export class NullOpEventReporter implements IEventReporter {
-  async reportExecution(_cell: vscode.NotebookCell) {
+  async reportExecution(_cell: vscode.NotebookCell, _status: boolean) {
     // Do nothing
   }
 
diff --git a/src/extension/kernel.ts b/src/extension/kernel.ts
index 4220ab117..76a2811b5 100644
--- a/src/extension/kernel.ts
+++ b/src/extension/kernel.ts
@@ -777,8 +777,6 @@ export class Kernel implements Disposable {
     }
 
     TelemetryReporter.sendTelemetryEvent('cell.startExecute')
-    // todo(sebastian): rewrite to use non-blocking impl
-    const execCellReport = getEventReporter().reportExecution(cell)
     runmeExec.start(Date.now())
 
     const annotations = getAnnotations(cell)
@@ -825,10 +823,12 @@ export class Kernel implements Disposable {
       successfulCellExecution = false
       log.error('Error executing cell', e.message)
       window.showErrorMessage(e.message)
-    } finally {
-      await execCellReport
     }
 
+    // todo(sebastian): rewrite to use non-blocking impl
+    const execCellReport = getEventReporter().reportExecution(cell, successfulCellExecution)
+    await execCellReport
+
     TelemetryReporter.sendTelemetryEvent('cell.endExecute', {
       'cell.success': successfulCellExecution?.toString(),
       'cell.mimeType': annotations.mimeType,