From 1a402106c5ae672b9ce40416560668c5cbbaaa2f Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Wed, 27 Aug 2025 18:20:48 +1000 Subject: [PATCH] Strip cell markers from Notebook Cell Next Edits --- .../node/nextEditProviderTelemetry.ts | 11 +++ .../vscode-node/inlineCompletionProvider.ts | 9 +- .../vscode-node/parts/vscodeWorkspace.ts | 87 ++++++++++++++++++- .../common/alternativeNotebookTextDocument.ts | 12 ++- 4 files changed, 113 insertions(+), 6 deletions(-) diff --git a/src/extension/inlineEdits/node/nextEditProviderTelemetry.ts b/src/extension/inlineEdits/node/nextEditProviderTelemetry.ts index e6f7ae4ee..de2938075 100644 --- a/src/extension/inlineEdits/node/nextEditProviderTelemetry.ts +++ b/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, diff --git a/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts b/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts index 19624acc4..f16f45a05 100644 --- a/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts +++ b/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 | 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, }; diff --git a/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts b/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts index 9f59c6aee..eafb7db36 100644 --- a/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts +++ b/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; + /** + * 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 { diff --git a/src/platform/notebook/common/alternativeNotebookTextDocument.ts b/src/platform/notebook/common/alternativeNotebookTextDocument.ts index 79a83e76e..552a88f72 100644 --- a/src/platform/notebook/common/alternativeNotebookTextDocument.ts +++ b/src/platform/notebook/common/alternativeNotebookTextDocument.ts @@ -30,7 +30,7 @@ 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, @@ -38,7 +38,8 @@ class AlternativeNotebookCellSnapshot { 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;