Strip cell markers from Notebook Cell Next Edits

src/extension/inlineEdits/node/nextEditProviderTelemetry.ts

@@ -112,6 +112,7 @@ export interface INextEditProviderTelemetry extends ILlmNESTelemetry, IDiagnosti
 	readonly postProcessingOutcome: string | undefined;
 	readonly isNESForAnotherDoc: boolean;
 	readonly notebookCellMarkerCount: number;
+	readonly notebookCellMarkerOutcome: string | undefined;
 	readonly notebookCellMarkerIndex: number;
 	readonly isActiveDocument?: boolean;
 	readonly isNaturalLanguageDominated: boolean;
@@ -443,6 +444,7 @@ export class NextEditProviderTelemetryBuilder extends Disposable {
 			supersededByOpportunityId: this._supersededByOpportunityId,
 			pickedNES: this._nesTypePicked,
 			hadLlmNES: this._hadLlmNES,
+			notebookCellMarkerOutcome: this._notebookCellMarkerOutcome,
 			isActiveDocument: this._isActiveDocument,
 			isNESForAnotherDoc: this._isNESForAnotherDoc,
 			notebookCellMarkerCount: this._notebookCellMarkerCount,
@@ -529,6 +531,12 @@ export class NextEditProviderTelemetryBuilder extends Disposable {
 		return this;
 	}
 
+	private _notebookCellMarkerOutcome?: string;
+	public setNotebookCellMarkerOutcome(outcome: string): this {
+		this._notebookCellMarkerOutcome = outcome;
+		return this;
+	}
+
 	private _notebookCellMarkerIndex: number = -1;
 	public setNotebookCellMarkerIndex(index: number): this {
 		this._notebookCellMarkerIndex = index;
@@ -649,6 +657,7 @@ export class TelemetrySender implements IDisposable {
 		const {
 			opportunityId,
 			headerRequestId,
+			notebookCellMarkerOutcome,
 			requestN,
 			providerId,
 			modelName,
@@ -737,6 +746,7 @@ export class TelemetrySender implements IDisposable {
 		"fetchError": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Fetch error message" },
 		"pickedNES": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Whether the request had picked NES" },
 		"diagnosticType": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Type of diagnostics" },
+		"notebookCellMarkerOutcome": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Outcome of the notebook cell marker processing" },
 		"diagnosticDroppedReasons": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Reasons for dropping diagnostics NES suggestions" },
 		"requestN": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Request number", "isMeasurement": true },
 		"hadStatelessNextEditProviderCall": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Whether the request had a stateless next edit provider call", "isMeasurement": true },
@@ -806,6 +816,7 @@ export class TelemetrySender implements IDisposable {
 				diagnosticDroppedReasons,
 				pickedNES,
 				notebookType,
+				notebookCellMarkerOutcome
 			},
 			{
 				requestN,

src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts

@@ -244,7 +244,7 @@ export class InlineCompletionProviderImpl implements InlineCompletionItemProvide
 				isInlineCompletion = allowInlineCompletions && isInlineSuggestion(position, document, range, result.edit.newText);
 				completionItem = serveAsCompletionsProvider && !isInlineCompletion ?
 					undefined :
-					this.createCompletionItem(doc, document, position, range, result);
+					this.createCompletionItem(doc, document, position, range, result, telemetryBuilder);
 			} else {
 				// nes is for a different document.
 				telemetryBuilder.setNotebookCellMarkerIndex((result.edit.newText || '').indexOf('#%% vscode.cell [id='));
@@ -351,6 +351,7 @@ export class InlineCompletionProviderImpl implements InlineCompletionItemProvide
 		position: Position,
 		range: Range,
 		result: NonNullable<(NextEditResult | DiagnosticsNextEditResult)['result']>,
+		telemetryBuilder: NextEditProviderTelemetryBuilder
 	): Omit<NesCompletionItem, 'telemetryBuilder' | 'info' | 'showInlineEditMenu' | 'action' | 'wasShown' | 'isInlineEdit'> | undefined {
 
 		// Only show edit when the cursor is max 4 lines away from the edit
@@ -370,9 +371,13 @@ export class InlineCompletionProviderImpl implements InlineCompletionItemProvide
 		} : undefined;
 
 
+		const telemetryInfo = { cellMarkerOutcome: '' };
+		const insertText = doc.adjustTextEdit(document, result.edit.newText, telemetryInfo);
+		telemetryBuilder.setNotebookCellMarkerOutcome(telemetryInfo.cellMarkerOutcome);
+
 		return {
 			range,
-			insertText: result.edit.newText,
+			insertText,
 			showRange,
 			displayLocation,
 		};

src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts

@@ -13,7 +13,7 @@ import { LanguageId } from '../../../../platform/inlineEdits/common/dataTypes/la
 import { EditReason } from '../../../../platform/inlineEdits/common/editReason';
 import { IObservableDocument, ObservableWorkspace, StringEditWithReason } from '../../../../platform/inlineEdits/common/observableWorkspace';
 import { createAlternativeNotebookDocument, IAlternativeNotebookDocument, toAltNotebookCellChangeEdit, toAltNotebookChangeEdit } from '../../../../platform/notebook/common/alternativeNotebookTextDocument';
-import { getDefaultLanguage } from '../../../../platform/notebook/common/helpers';
+import { EOL, getDefaultLanguage } from '../../../../platform/notebook/common/helpers';
 import { IExperimentationService } from '../../../../platform/telemetry/common/nullExperimentationService';
 import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry';
 import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
@@ -430,6 +430,16 @@ export interface IVSCodeObservableDocument extends IObservableDocument {
 	 * @param itemResolveCount Number of code actions to resolve (too large numbers slow down code actions)
 	 */
 	getCodeActions(range: OffsetRange, itemResolveCount: number, token: CancellationToken): Promise<CodeActionData[] | undefined>;
+	/**
+	 * This only applies to Notebook documents.
+	 * For text documents, the same value is returned.
+	 *
+	 * With notebook documents, cells are concatenated and seprated using special cell markers.
+	 * If the cell markers appear in the suggestions, then those end up being turned into text edits
+	 * that include the cell markers. This confuses users as they did not type those cell markers.
+	 * This method allows adjusting the text edit to remove those cell markers.
+	 */
+	adjustTextEdit(textDocument: TextDocument, newText: string, notebookTelemetry: { cellMarkerOutcome: string }): string;
 }
 
 export interface IVSCodeObservableTextDocument extends IObservableDocument, IVSCodeObservableDocument {
@@ -546,6 +556,9 @@ class VSCodeObservableTextDocument extends AbstractVSCodeObservableDocument impl
 
 		return actions?.map(action => toCodeActionData(action, this, range => this.toOffsetRange(this.textDocument, range)));
 	}
+	adjustTextEdit(_textDocument: TextDocument, newText: string): string {
+		return newText;
+	}
 }
 
 class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument implements IVSCodeObservableDocument {
@@ -637,6 +650,78 @@ class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument
 			}));
 		})).then(results => coalesce(results.flat()));
 	}
+	adjustTextEdit(textDocument: TextDocument, newText: string, notebookTelemetry: { cellMarkerOutcome: string }): string {
+		if (!newText.includes('%% vscode.cell [id=')) {
+			return newText;
+		}
+
+		const cell = this.altNotebook.getCell(textDocument);
+		if (!cell) {
+			// Unlikely scenario, but just in case, return the text as is.
+			notebookTelemetry.cellMarkerOutcome = 'cellNotFound';
+			return newText;
+		}
+
+		// Assume we have a notebook as follows
+		/**
+		 * #%% vscode.cell [id=111] [language=python]
+		 * print("Hello World")
+		 * #%% vscode.cell [id=222] [language=python]
+		 * print("foo")
+		 */
+		// Now assume user triggers completions in first cell after `print`.
+		// & lets assume the LLM returns a completion of `("Hello World!")\n#%% vscode.cell [id=222] [language=python]\nprint("foo bar baz")`
+		// Basically llm has returned edits to the current cell as well as the next cell.
+		// We want to trim off the cell markers and the content of other cells.
+		// So that the user only sees `("Hello World!")` as the completion text.
+		const cellMarker = this.altNotebook.getCellMarker(cell);
+		if (!cellMarker) {
+			notebookTelemetry.cellMarkerOutcome = 'noCellMarker';
+			return newText;
+		}
+
+		let numberOfMarkersRemoved = 0;
+		let updatedText = newText;
+		// If the new text starts with the cell marker, this means the entire cell is being replaced.
+		if (newText.startsWith(cellMarker)) {
+			updatedText = newText.substring(cellMarker.length + EOL.length);
+			numberOfMarkersRemoved++;
+			if (updatedText.length === 0) {
+				// If the entire cell is being replaced with empty text, this isn't possible.
+				notebookTelemetry.cellMarkerOutcome = 'emptyEditAfterRemovingFirstMarker';
+				return newText;
+			}
+		} else if (newText.includes(cellMarker)) {
+			// This is weird, possible there's some white space before the cell marker.
+			updatedText = newText.substring(newText.indexOf(cellMarker) + cellMarker.length + EOL.length);
+			numberOfMarkersRemoved++;
+			if (updatedText.length === 0) {
+				// If the entire cell is being replaced with empty text, this isn't possible.
+				notebookTelemetry.cellMarkerOutcome = 'emptyEditAfterRemovingFirstMarkerWithPrefix';
+				return newText;
+			}
+		}
+
+		// Possible we have the cell marker for the next cell in the edit (see above example).
+		const nextCell = this.notebook.cellCount > cell.index + 1 ? this.notebook.cellAt(cell.index + 1) : undefined;
+		const nextCellMarker = nextCell ? this.altNotebook.getCellMarker(nextCell) : undefined;
+		if (nextCellMarker) {
+			const idx = updatedText.indexOf(nextCellMarker);
+			if (idx !== -1) {
+				numberOfMarkersRemoved++;
+				updatedText = updatedText.substring(0, idx - EOL.length).trimEnd();
+				if (updatedText.length === 0) {
+					// This means we've got just empty text (i.e. removing all text).
+					// This isn't possible.
+					notebookTelemetry.cellMarkerOutcome = 'emptyEditAfterRemovingNextMarker';
+					return newText;
+				}
+			}
+		}
+
+		notebookTelemetry.cellMarkerOutcome = updatedText.includes('%% vscode.cell [id=') ? `stillHasCellMarker:removed${numberOfMarkersRemoved}` : `removedCellMarker${numberOfMarkersRemoved}`;
+		return updatedText;
+	}
 }
 
 function getTextDocuments(excludeNotebookCells: boolean, markdownCellUris: ResourceSet): IObservable<readonly TextDocument[]> {

src/platform/notebook/common/alternativeNotebookTextDocument.ts

@@ -30,15 +30,16 @@ class AlternativeNotebookCellSnapshot {
 		const code = cell.document.getText().replace(/\r\n|\n/g, EOL);
 		const prefix = cell.kind === NotebookCellKind.Markup ? `${cellMarker}${EOL}${blockComment[0]}${EOL}` : `${cellMarker}${EOL}`;
 		const suffix = cell.kind === NotebookCellKind.Markup ? `${EOL}${blockComment[1]}` : '';
-		return new AlternativeNotebookCellSnapshot(cell, blockComment, lineCommentStart, code, prefix, suffix);
+		return new AlternativeNotebookCellSnapshot(cell, blockComment, lineCommentStart, code, prefix, suffix, cellMarker);
 	}
 	constructor(
 		public readonly cell: NotebookCell,
 		private readonly blockComment: [string, string],
 		private readonly lineCommentStart: string,
 		private readonly code: string,
 		private readonly prefix: string,
-		private readonly suffix: string
+		private readonly suffix: string,
+		public readonly cellMarker: string
 	) {
 		this.crlfTranslator = new CrLfOffsetTranslator(cell.document.getText(), cell.document.eol);
 		this.positionTransformer = new PositionOffsetTransformer(`${prefix}${code}${suffix}`);
@@ -64,7 +65,7 @@ class AlternativeNotebookCellSnapshot {
 
 	public withTextEdit(edit: StringEdit): AlternativeNotebookCellSnapshot {
 		const newCode = edit.apply(this.code);
-		return new AlternativeNotebookCellSnapshot(this.cell, this.blockComment, this.lineCommentStart, newCode, this.prefix, this.suffix);
+		return new AlternativeNotebookCellSnapshot(this.cell, this.blockComment, this.lineCommentStart, newCode, this.prefix, this.suffix, this.cellMarker);
 	}
 
 	public get altText(): string {
@@ -177,6 +178,11 @@ abstract class AbstractAlternativeNotebookDocument {
 		return this.cellTextDocuments.get(textDocument);
 	}
 
+	public getCellMarker(cell: NotebookCell) {
+		const altCell = this.cells.find(c => c.altCell.cell === cell);
+		return altCell?.altCell.cellMarker;
+	}
+
 	public getText(range?: OffsetRange): string {
 		const altText = this.cells.map(cell => cell.altCell.altText).join(EOL);
 		return range ? range.substring(altText) : altText;