Fix telemetry (#1191)

Author: chrmarti

chat-lib/package-lock.json

@@ -1,6 +1,6 @@
 {
 	"name": "@vscode/chat-lib",
-	"version": "0.0.3",
+	"version": "0.0.4",
 	"description": "Chat and inline editing SDK extracted from VS Code Copilot Chat",
 	"main": "dist/src/main.js",
 	"types": "dist/src/main.d.ts",
@@ -15,7 +15,7 @@
 	},
 	"dependencies": {
 		"@microsoft/tiktokenizer": "^1.0.10",
-		"@vscode/copilot-api": "^0.1.8",
+		"@vscode/copilot-api": "^0.1.10",
 		"@vscode/l10n": "^0.0.18",
 		"@vscode/prompt-tsx": "^0.4.0-alpha.5",
 		"jsonc-parser": "^3.3.1",

chat-lib/package.json

@@ -18,9 +18,9 @@ import { FetchOptions, IAbortController, IHeaders, Response } from '../src/_inte
 import { IFetcher } from '../src/_internal/platform/networking/common/networking';
 import { CancellationToken } from '../src/_internal/util/vs/base/common/cancellation';
 import { URI } from '../src/_internal/util/vs/base/common/uri';
-import { StringEdit } from '../src/_internal/util/vs/editor/common/core/edits/stringEdit';
+import { StringEdit, StringReplacement } from '../src/_internal/util/vs/editor/common/core/edits/stringEdit';
 import { OffsetRange } from '../src/_internal/util/vs/editor/common/core/ranges/offsetRange';
-import { createNESProvider, ITelemetrySender } from '../src/main';
+import { createNESProvider, ILogTarget, ITelemetrySender, LogLevel } from '../src/main';
 import { ICopilotTokenManager } from '../src/_internal/platform/authentication/common/copilotTokenManager';
 import { Emitter } from '../src/_internal/util/vs/base/common/event';
 import { CopilotToken } from '../src/_internal/platform/authentication/common/copilotToken';
@@ -97,8 +97,17 @@ class TestCopilotTokenManager implements ICopilotTokenManager {
 }
 
 class TestTelemetrySender implements ITelemetrySender {
+	events: { eventName: string; properties?: Record<string, string | undefined>; measurements?: Record<string, number | undefined> }[] = [];
 	sendTelemetryEvent(eventName: string, properties?: Record<string, string | undefined>, measurements?: Record<string, number | undefined>): void {
-		// No-op
+		this.events.push({ eventName, properties, measurements });
+	}
+}
+
+class TestLogTarget implements ILogTarget {
+	logs: { level: LogLevel; message: string; metadata?: any }[] = [];
+	logIt(level: LogLevel, metadataStr: string, ...extra: any[]): void {
+		this.logs.push({ level, message: metadataStr, metadata: extra });
+		console.log(`[${LogLevel[level]}]${metadataStr}`, ...extra);
 	}
 }
 
@@ -121,20 +130,26 @@ describe('NESProvider Facade', () => {
 			const myPoint = new Point(0, 1);`.trimStart()
 		});
 		doc.setSelection([new OffsetRange(1, 1)], undefined);
+		const telemetrySender = new TestTelemetrySender();
+		const logTarget = new TestLogTarget();
 		const nextEditProvider = createNESProvider({
 			workspace,
 			fetcher: new TestFetcher({ '/chat/completions': await fs.readFile(path.join(__dirname, 'nesProvider.reply.txt'), 'utf8') }),
 			copilotTokenManager: new TestCopilotTokenManager(),
-			telemetrySender: new TestTelemetrySender(),
+			telemetrySender,
+			logTarget,
 		});
 
 		doc.applyEdit(StringEdit.insert(11, '3D'));
 
 		const result = await nextEditProvider.getNextEdit(doc.id.toUri(), CancellationToken.None);
 
-		assert(result.result?.edit);
+		assert(result.result);
 
-		doc.applyEdit(result.result.edit.toEdit());
+		const { range, newText } = result.result;
+		const offsetRange = OffsetRange.fromTo(range.start, range.endExclusive);
+		const replace = StringReplacement.replace(offsetRange, newText);
+		doc.applyEdit(replace.toEdit());
 
 		expect(doc.value.get().value).toMatchInlineSnapshot(`
 			"class Point3D {
@@ -150,5 +165,16 @@ describe('NESProvider Facade', () => {
 
 			const myPoint = new Point(0, 1);"
 		`);
+
+		nextEditProvider.handleAcceptance(result);
+		await new Promise(resolve => setTimeout(resolve, 100)); // wait for async telemetry sending
+		const event = telemetrySender.events.find(e => e.eventName === 'copilot-nes/provideInlineEdit');
+		expect(event).toBeDefined();
+		expect(event!.properties?.acceptance).toBe('accepted');
+
+		nextEditProvider.dispose();
+
+		expect(logTarget.logs.length).toBeGreaterThan(0);
+		expect(logTarget.logs.filter(l => l.level === LogLevel.Error)).toHaveLength(0);
 	});
 });

chat-lib/test/nesProvider.spec.ts

@@ -451,7 +451,7 @@ export class NextEditProviderTelemetryBuilder extends Disposable {
 
 	constructor(
 		gitExtensionService: IGitExtensionService,
-		notebookService: INotebookService,
+		notebookService: INotebookService | undefined,
 		workspaceService: IWorkspaceService,
 		providerId: string,
 		doc: IObservableDocument,

src/extension/inlineEdits/node/nextEditProviderTelemetry.ts

@@ -6,7 +6,7 @@
 import type * as vscode from 'vscode';
 import { DebugRecorder } from '../../extension/inlineEdits/node/debugRecorder';
 import { INextEditProvider, NextEditProvider } from '../../extension/inlineEdits/node/nextEditProvider';
-import { LlmNESTelemetryBuilder } from '../../extension/inlineEdits/node/nextEditProviderTelemetry';
+import { LlmNESTelemetryBuilder, NextEditProviderTelemetryBuilder, TelemetrySender } from '../../extension/inlineEdits/node/nextEditProviderTelemetry';
 import { INextEditResult } from '../../extension/inlineEdits/node/nextEditResult';
 import { ChatMLFetcherImpl } from '../../extension/prompt/node/chatMLFetcher';
 import { XtabProvider } from '../../extension/xtab/node/xtabProvider';
@@ -115,13 +115,41 @@ export interface INESProviderOptions {
 	readonly logTarget?: ILogTarget;
 }
 
-export function createNESProvider(options: INESProviderOptions): INESProvider {
+export interface INESResult {
+	readonly result?: {
+		readonly newText: string;
+		readonly range: {
+			readonly start: number;
+			readonly endExclusive: number;
+		};
+	};
+}
+
+export interface INESProvider<T extends INESResult = INESResult> {
+	getId(): string;
+	getNextEdit(documentUri: vscode.Uri, cancellationToken: CancellationToken): Promise<T>;
+	handleShown(suggestion: T): void;
+	handleAcceptance(suggestion: T): void;
+	handleRejection(suggestion: T): void;
+	handleIgnored(suggestion: T, supersededByRequestUuid: T | undefined): void;
+	dispose(): void;
+}
+
+export function createNESProvider(options: INESProviderOptions): INESProvider<INESResult> {
 	const instantiationService = setupServices(options);
 	return instantiationService.createInstance(NESProvider, options);
 }
 
-class NESProvider extends Disposable implements INESProvider {
+interface NESResult extends INESResult {
+	docId: DocumentId;
+	requestUuid: string;
+	internalResult: INextEditResult;
+	telemetryBuilder: NextEditProviderTelemetryBuilder;
+}
+
+class NESProvider extends Disposable implements INESProvider<NESResult> {
 	private readonly _nextEditProvider: INextEditProvider<INextEditResult, LlmNESTelemetryBuilder>;
+	private readonly _telemetrySender: TelemetrySender;
 	private readonly _debugRecorder: DebugRecorder;
 
 	constructor(
@@ -140,29 +168,49 @@ class NESProvider extends Disposable implements INESProvider {
 		this._debugRecorder = this._register(new DebugRecorder(this._options.workspace));
 
 		this._nextEditProvider = instantiationService.createInstance(NextEditProvider, this._options.workspace, statelessNextEditProvider, historyContextProvider, xtabHistoryTracker, this._debugRecorder);
+		this._telemetrySender = this._register(instantiationService.createInstance(TelemetrySender));
 	}
 
 	getId(): string {
 		return this._nextEditProvider.ID;
 	}
 
-	handleShown(suggestion: INextEditResult): void {
-		this._nextEditProvider.handleShown(suggestion);
+	handleShown(result: NESResult): void {
+		result.telemetryBuilder.setAsShown();
+		this._nextEditProvider.handleShown(result.internalResult);
 	}
 
-	handleAcceptance(docId: DocumentId, suggestion: INextEditResult): void {
-		this._nextEditProvider.handleAcceptance(docId, suggestion);
+	handleAcceptance(result: NESResult): void {
+		result.telemetryBuilder.setAcceptance('accepted');
+		result.telemetryBuilder.setStatus('accepted');
+		this._nextEditProvider.handleAcceptance(result.docId, result.internalResult);
+		this.handleEndOfLifetime(result);
 	}
 
-	handleRejection(docId: DocumentId, suggestion: INextEditResult): void {
-		this._nextEditProvider.handleRejection(docId, suggestion);
+	handleRejection(result: NESResult): void {
+		result.telemetryBuilder.setAcceptance('rejected');
+		result.telemetryBuilder.setStatus('rejected');
+		this._nextEditProvider.handleRejection(result.docId, result.internalResult);
+		this.handleEndOfLifetime(result);
 	}
 
-	handleIgnored(docId: DocumentId, suggestion: INextEditResult, supersededByRequestUuid: INextEditResult | undefined): void {
-		this._nextEditProvider.handleIgnored(docId, suggestion, supersededByRequestUuid);
+	handleIgnored(result: NESResult, supersededByRequestUuid: NESResult | undefined): void {
+		if (supersededByRequestUuid) {
+			result.telemetryBuilder.setSupersededBy(supersededByRequestUuid.requestUuid);
+		}
+		this._nextEditProvider.handleIgnored(result.docId, result.internalResult, supersededByRequestUuid?.internalResult);
+		this.handleEndOfLifetime(result);
 	}
 
-	async getNextEdit(documentUri: vscode.Uri, cancellationToken: CancellationToken): Promise<INextEditResult> {
+	private handleEndOfLifetime(result: NESResult): void {
+		try {
+			this._telemetrySender.sendTelemetryForBuilder(result.telemetryBuilder);
+		} finally {
+			result.telemetryBuilder.dispose();
+		}
+	}
+
+	async getNextEdit(documentUri: vscode.Uri, cancellationToken: CancellationToken): Promise<NESResult> {
 		const docId = DocumentId.create(documentUri.toString());
 
 		// Create minimal required context objects
@@ -183,30 +231,41 @@ class NESProvider extends Disposable implements INESProvider {
 		}
 
 		// Create telemetry builder - we'll need to pass null/undefined for services we don't have
-		const telemetryBuilder = new LlmNESTelemetryBuilder(
-			new NullGitExtensionService(), // IGitExtensionService
+		const telemetryBuilder = new NextEditProviderTelemetryBuilder(
+			new NullGitExtensionService(),
 			undefined, // INotebookService
-			this._workspaceService, // IWorkspaceService
-			this._nextEditProvider.ID, // providerId
-			document, // doc
-			this._debugRecorder, // debugRecorder
-			undefined // requestBookmark
+			this._workspaceService,
+			this._nextEditProvider.ID,
+			document,
+			this._debugRecorder,
+			logContext.recordingBookmark
 		);
-
-		return await this._nextEditProvider.getNextEdit(docId, context, logContext, cancellationToken, telemetryBuilder);
+		telemetryBuilder.setOpportunityId(context.requestUuid);
+
+		try {
+			const internalResult = await this._nextEditProvider.getNextEdit(docId, context, logContext, cancellationToken, telemetryBuilder.nesBuilder);
+			const result: NESResult = {
+				result: internalResult.result ? {
+					newText: internalResult.result.edit.newText,
+					range: internalResult.result.edit.replaceRange,
+				} : undefined,
+				docId,
+				requestUuid: context.requestUuid,
+				internalResult,
+				telemetryBuilder,
+			};
+			return result;
+		} catch (e) {
+			try {
+				this._telemetrySender.sendTelemetryForBuilder(telemetryBuilder);
+			} finally {
+				telemetryBuilder.dispose();
+			}
+			throw e;
+		}
 	}
 }
 
-export interface INESProvider {
-	getId(): string;
-	getNextEdit(documentUri: vscode.Uri, cancellationToken: CancellationToken): Promise<INextEditResult>;
-	handleShown(suggestion: INextEditResult): void;
-	handleAcceptance(docId: DocumentId, suggestion: INextEditResult): void;
-	handleRejection(docId: DocumentId, suggestion: INextEditResult): void;
-	handleIgnored(docId: DocumentId, suggestion: INextEditResult, supersededByRequestUuid: INextEditResult | undefined): void;
-	dispose(): void;
-}
-
 function setupServices(options: INESProviderOptions) {
 	const { fetcher, copilotTokenManager, telemetrySender, logTarget } = options;
 	const builder = new InstantiationServiceBuilder();