From ef0e276845924314b56f3c04ede4ebeb0ccdbebb Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Mon, 28 Jul 2025 18:43:05 +1000 Subject: [PATCH] Similar API for notebook and text docs in VSCode Workspace (#371) --- .../diagnosticsCompletions.ts | 2 +- .../diagnosticsCompletionProcessor.ts | 2 +- .../vscode-node/inlineCompletionProvider.ts | 27 +--- .../vscode-node/parts/vscodeWorkspace.ts | 124 +++++++++++------- .../common/alternativeNotebookTextDocument.ts | 6 + 5 files changed, 92 insertions(+), 69 deletions(-) diff --git a/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts b/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts index be48e9a5dd1..a3c2c1e2086 100644 --- a/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts +++ b/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsBasedCompletions/diagnosticsCompletions.ts @@ -196,7 +196,7 @@ async function getCodeActionsForNotebookDocumentDiagnostic(diagnostic: Diagnosti return []; } return Promise.all(cellRanges.map(async ([cell, range]) => { - const actions = await getCodeActionsForUriRange(cell.document.uri, range); + const actions = await getCodeActionsForUriRange(cell.uri, range); return actions.map(action => { action.diagnostics = action.diagnostics ? workspaceDocument.projectDiagnostics(cell, action.diagnostics) : undefined; return action; diff --git a/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts b/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts index 3ca7ab3ca8b..3513bccbc85 100644 --- a/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts +++ b/extensions/copilot/src/extension/inlineEdits/vscode-node/features/diagnosticsCompletionProcessor.ts @@ -313,7 +313,7 @@ export class DiagnosticsCompletionProcessor extends Disposable { .getDiagnostics(workspaceDocument.textDocument.uri) : workspaceDocument.notebook.getCells().flatMap(cell => this._languageDiagnosticsService .getDiagnostics(cell.document.uri) - .flatMap(diagnostic => workspaceDocument.projectDiagnostics(cell, [diagnostic]))); + .flatMap(diagnostic => workspaceDocument.projectDiagnostics(cell.document, [diagnostic]))); const availableDiagnostics = diagnostics .map(diagnostic => Diagnostic.fromVSCodeDiagnostic(diagnostic)) .filter(diagnostic => diagnostic.severity !== DiagnosticSeverity.Information) diff --git a/extensions/copilot/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts b/extensions/copilot/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts index c191ec9564a..5075b9b6cfd 100644 --- a/extensions/copilot/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts +++ b/extensions/copilot/src/extension/inlineEdits/vscode-node/inlineCompletionProvider.ts @@ -22,7 +22,6 @@ import { raceCancellation, timeout } from '../../../util/vs/base/common/async'; import { CancellationTokenSource } from '../../../util/vs/base/common/cancellation'; import { Event } from '../../../util/vs/base/common/event'; import { StringEdit } from '../../../util/vs/editor/common/core/edits/stringEdit'; -import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange'; import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation'; import { LineCheck } from '../../inlineChat/vscode-node/inlineChatHint'; import { NextEditProviderTelemetryBuilder, TelemetrySender } from '../node/nextEditProviderTelemetry'; @@ -205,20 +204,11 @@ export class InlineCompletionProviderImpl implements InlineCompletionItemProvide tracer.trace(`using next edit suggestion from ${suggestionInfo.source}`); - let range: Range; - - if (doc.kind === 'notebookDocument') { - // Ignore changes to other cells. - const cellRange = doc.fromOffsetRange(result.edit.replaceRange).find(([cell]) => cell.document === document); - if (cellRange) { - range = cellRange[1]; - } else { - tracer.trace('no next edit suggestion for notebook cell'); - this.telemetrySender.scheduleSendingEnhancedTelemetry(suggestionInfo.suggestion, telemetryBuilder); - return emptyList; - } - } else { - range = documentRangeFromOffsetRange(document, result.edit.replaceRange); + const range = doc.fromOffsetRange(document, result.edit.replaceRange); + if (!range) { + tracer.trace('no next edit suggestion for notebook cell'); + this.telemetrySender.scheduleSendingEnhancedTelemetry(suggestionInfo.suggestion, telemetryBuilder); + return emptyList; } // Only show edit when the cursor is max 4 lines away from the edit @@ -501,13 +491,6 @@ export function raceAndAll( return { first, all }; } -export function documentRangeFromOffsetRange(doc: TextDocument, range: OffsetRange): Range { - return new Range( - doc.positionAt(range.start), - doc.positionAt(range.endExclusive) - ); -} - function shortOpportunityId(oppId: string): string { return oppId.substring(4, 8); } diff --git a/extensions/copilot/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts b/extensions/copilot/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts index cd4540b0c5f..fc73c0d9719 100644 --- a/extensions/copilot/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts +++ b/extensions/copilot/src/extension/inlineEdits/vscode-node/parts/vscodeWorkspace.ts @@ -3,7 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { Diagnostic, DiagnosticSeverity, EndOfLine, languages, NotebookCell, NotebookDocument, Range, TextDocument, TextDocumentChangeEvent, TextDocumentContentChangeEvent, TextEditor, Uri, window, workspace } from 'vscode'; +import { Diagnostic, DiagnosticSeverity, EndOfLine, languages, NotebookDocument, Range, TextDocument, TextDocumentChangeEvent, TextDocumentContentChangeEvent, TextEditor, Uri, window, workspace } from 'vscode'; import { ConfigKey, IConfigurationService } from '../../../../platform/configuration/common/configurationService'; import { IIgnoreService } from '../../../../platform/ignore/common/ignoreService'; import { DiagnosticData } from '../../../../platform/inlineEdits/common/dataTypes/diagnosticData'; @@ -87,7 +87,7 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable doc.altNotebook.applyCellChanges(e.document, e.contentChanges); const editWithReason = new StringEditWithReason(edit.replacements, EditReason.create(e.detailedReason?.metadata as any)); transaction(tx => { - doc.value.set(new StringText(doc.altNotebook.getText()), tx, editWithReason); + doc.value.set(stringValueFromDoc(doc.altNotebook), tx, editWithReason); doc.version.set(doc.notebook.version, tx); }); } @@ -105,7 +105,7 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable doc.altNotebook.applyNotebookChanges(e.contentChanges); const editWithReason = new StringEditWithReason(edit.replacements, EditReason.unknown); transaction(tx => { - doc.value.set(new StringText(doc.altNotebook.getText()), tx, editWithReason); + doc.value.set(stringValueFromDoc(doc.altNotebook), tx, editWithReason); doc.version.set(doc.notebook.version, tx); }); })); @@ -115,11 +115,10 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable if (!doc) { return; } - if (doc.kind === 'textDocument') { - doc.selection.set(e.selections.map(s => rangeToOffsetRange(s, e.textEditor.document)), undefined); - } else { - doc.selection.set(this.getNotebookSelections(doc.notebook, e.textEditor), undefined); - } + const selections = doc.kind === 'textDocument' ? + coalesce(e.selections.map(s => doc.toOffsetRange(e.textEditor.document, s))) : + this.getNotebookSelections(doc.notebook, e.textEditor); + doc.selection.set(selections, undefined); })); this._store.add(window.onDidChangeTextEditorVisibleRanges(e => { @@ -127,11 +126,10 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable if (!doc) { return; } - if (doc.kind === 'textDocument') { - doc.visibleRanges.set(e.visibleRanges.map(r => rangeToOffsetRange(r, e.textEditor.document)), undefined); - } else { - doc.visibleRanges.set(this.getNotebookVisibleRanges(doc.notebook), undefined); - } + const visibleRanges = doc.kind === 'textDocument' ? + coalesce(e.visibleRanges.map(r => doc.toOffsetRange(e.textEditor.document, r))) : + this.getNotebookVisibleRanges(doc.notebook); + doc.visibleRanges.set(visibleRanges, undefined); })); this._store.add(languages.onDidChangeDiagnostics(e => { @@ -140,7 +138,9 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable if (!document) { return; } - const diagnostics = document.kind === 'textDocument' ? this._createTextDocumentDiagnosticData(document.textDocument) : this._createNotebookDiagnosticData(document.altNotebook); + const diagnostics = document.kind === 'textDocument' ? + this._createTextDocumentDiagnosticData(document) : + this._createNotebookDiagnosticData(document.altNotebook); document.diagnostics.set(diagnostics, undefined); }); })); @@ -183,10 +183,15 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable const documentId = DocumentId.create(doc.uri.toString()); const openedTextEditor = window.visibleTextEditors.find(e => e.document.uri.toString() === doc.uri.toString()); - const selections = openedTextEditor?.selections.map(s => rangeToOffsetRange(s, doc)); - const visibleRanges = openedTextEditor?.visibleRanges.map(r => rangeToOffsetRange(r, doc)); - const diagnostics = this._createTextDocumentDiagnosticData(doc); - const document = new VSCodeObservableTextDocument(documentId, stringValueFromDoc(doc), doc.version, selections ?? [], visibleRanges ?? [], LanguageId.create(doc.languageId), diagnostics, doc); + const document = new VSCodeObservableTextDocument(documentId, stringValueFromDoc(doc), doc.version, [], [], LanguageId.create(doc.languageId), [], doc); + + const selections = coalesce((openedTextEditor?.selections || []).map(s => document.toOffsetRange(doc, s))); + const visibleRanges = coalesce((openedTextEditor?.visibleRanges || []).map(r => document.toOffsetRange(doc, r))); + transaction(tx => { + document.selection.set(selections, tx); + document.visibleRanges.set(visibleRanges, tx); + document.diagnostics.set(this._createTextDocumentDiagnosticData(document), tx); + }); return document; }).recomputeInitiallyAndOnChange(store); @@ -235,7 +240,7 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable const visibleRanges = this.getNotebookVisibleRanges(doc); const diagnostics = this._createNotebookDiagnosticData(altNotebook); const language = getLanguage(getDefaultLanguage(altNotebook.notebook)).languageId; - const document = new VSCodeObservableNotebookDocument(documentId, new StringText(altNotebook.getText()), doc.version, selections ?? [], visibleRanges ?? [], LanguageId.create(language), diagnostics, doc, altNotebook); + const document = new VSCodeObservableNotebookDocument(documentId, stringValueFromDoc(altNotebook), doc.version, selections ?? [], visibleRanges ?? [], LanguageId.create(language), diagnostics, doc, altNotebook); return document; }).recomputeInitiallyAndOnChange(store); @@ -289,19 +294,23 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable return document; } - private _createTextDocumentDiagnosticData(document: TextDocument) { - return languages.getDiagnostics(document.uri).map(d => this._createDiagnosticData(d, document)).filter(isDefined); + private _createTextDocumentDiagnosticData(document: VSCodeObservableTextDocument) { + return languages.getDiagnostics(document.textDocument.uri).map(d => this._createDiagnosticData(d, document)).filter(isDefined); } - private _createDiagnosticData(diagnostic: Diagnostic, doc: TextDocument): DiagnosticData | undefined { + private _createDiagnosticData(diagnostic: Diagnostic, doc: VSCodeObservableTextDocument): DiagnosticData | undefined { if (!diagnostic.source || (diagnostic.severity !== DiagnosticSeverity.Error && diagnostic.severity !== DiagnosticSeverity.Warning)) { return undefined; } + const range = doc.toOffsetRange(doc.textDocument, diagnostic.range); + if (!range) { + return undefined; + } const diag: DiagnosticData = new DiagnosticData( - doc.uri, + doc.textDocument.uri, diagnostic.message, diagnostic.severity === DiagnosticSeverity.Error ? 'error' : 'warning', - rangeToOffsetRange(diagnostic.range, doc) + range ); return diag; } @@ -375,7 +384,8 @@ export class VSCodeWorkspace extends ObservableWorkspace implements IDisposable export interface IVSCodeObservableTextDocument extends IObservableDocument { kind: 'textDocument'; readonly textDocument: TextDocument; - fromOffsetRange(range: OffsetRange): Range; + fromOffsetRange(textDocument: TextDocument, range: OffsetRange): Range | undefined; + toOffsetRange(textDocument: TextDocument, range: Range): OffsetRange | undefined; } abstract class AbstractVSCodeObservableDocument { @@ -421,21 +431,35 @@ class VSCodeObservableTextDocument extends AbstractVSCodeObservableDocument impl super(id, value, versionId, selection, visibleRanges, languageId, diagnostics); } - fromOffsetRange(range: OffsetRange): Range { + fromOffsetRange(textDocument: TextDocument, range: OffsetRange): Range { + if (textDocument !== this.textDocument) { + throw new Error('TextDocument does not match the one of this observable document.'); + } return new Range( - this.textDocument.positionAt(range.start), - this.textDocument.positionAt(range.endExclusive) + textDocument.positionAt(range.start), + textDocument.positionAt(range.endExclusive) ); } + toOffsetRange(textDocument: TextDocument, range: Range): OffsetRange | undefined { + return new OffsetRange(textDocument.offsetAt(range.start), textDocument.offsetAt(range.end)); + } } export interface IVSCodeObservableNotebookDocument extends IObservableDocument { kind: 'notebookDocument'; readonly notebook: NotebookDocument; - fromOffsetRange(range: OffsetRange): [NotebookCell, Range][]; - fromRange(range: Range): [NotebookCell, Range][]; - projectRange(cell: NotebookCell, range: readonly Range[]): Range[]; - projectDiagnostics(cell: NotebookCell, diagnostics: readonly Diagnostic[]): Diagnostic[]; + /** + * Converts an offset range of Notebook Text document to a range within the provided notebook cell. + * If the range does not belong to the cell, it returns undefined. + */ + fromOffsetRange(textDocument: TextDocument, range: OffsetRange): Range | undefined; + /** + * Converts an offset range of Notebook Text document to a range within the notebook cell(s). + * The range provided could span multiple cells, so it returns an array of tuples containing the cell document and the range within that cell. + */ + fromOffsetRange(range: OffsetRange): [TextDocument, Range][]; + fromRange(range: Range): [TextDocument, Range][]; + projectDiagnostics(cell: TextDocument, diagnostics: readonly Diagnostic[]): Diagnostic[]; } class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument implements IVSCodeObservableNotebookDocument { @@ -455,26 +479,36 @@ class VSCodeObservableNotebookDocument extends AbstractVSCodeObservableDocument super(id, value, versionId, selection, visibleRanges, languageId, diagnostics); } - fromOffsetRange(range: OffsetRange): [NotebookCell, Range][] { - return this.altNotebook.fromAltOffsetRange(range); + fromOffsetRange(textDocument: TextDocument, range: OffsetRange): Range | undefined; + fromOffsetRange(range: OffsetRange): [TextDocument, Range][]; + fromOffsetRange(arg1: TextDocument | OffsetRange, range?: OffsetRange): Range | undefined | [TextDocument, Range][] { + if (arg1 instanceof OffsetRange) { + return this.altNotebook.fromAltOffsetRange(arg1).map(r => [r[0].document, r[1]]); + } else if (range !== undefined) { + const cell = this.altNotebook.getCell(arg1); + if (!cell) { + return undefined; + } + const results = this.altNotebook.fromAltOffsetRange(range); + const found = results.find(r => r[0].document === arg1); + return found ? found[1] : undefined; + } + return undefined; } - fromRange(range: Range): [NotebookCell, Range][] { - return this.altNotebook.fromAltRange(range); + fromRange(range: Range): [TextDocument, Range][] { + return this.altNotebook.fromAltRange(range).map(r => [r[0].document, r[1]]); } - projectRange(cell: NotebookCell, range: readonly Range[]): Range[] { - return this.altNotebook.toAltRange(cell, range); - } - projectDiagnostics(cell: NotebookCell, diagnostics: readonly Diagnostic[]): Diagnostic[] { + projectDiagnostics(textDocument: TextDocument, diagnostics: readonly Diagnostic[]): Diagnostic[] { + const cell = this.altNotebook.getCell(textDocument); + if (!cell) { + return []; + } return toAltDiagnostics(this.altNotebook, cell, diagnostics); } } export type IVSCodeObservableDocument = IVSCodeObservableTextDocument | IVSCodeObservableNotebookDocument; -function rangeToOffsetRange(range: Range, doc: TextDocument): OffsetRange { - return new OffsetRange(doc.offsetAt(range.start), doc.offsetAt(range.end)); -} - function getTextDocuments(excludeNotebookCells: boolean): IObservable { return observableFromEvent(undefined, e => { const d1 = workspace.onDidOpenTextDocument(e); @@ -632,7 +666,7 @@ export class VerifyTextDocumentChanges extends Disposable { } } -export function stringValueFromDoc(doc: TextDocument): StringText { +export function stringValueFromDoc(doc: TextDocument | IAlternativeNotebookDocument): StringText { return new StringText(doc.getText()); } export function editFromTextDocumentContentChangeEvents(events: readonly TextDocumentContentChangeEvent[]): StringEdit { diff --git a/extensions/copilot/src/platform/notebook/common/alternativeNotebookTextDocument.ts b/extensions/copilot/src/platform/notebook/common/alternativeNotebookTextDocument.ts index 865b674f03c..6e7ce7f3544 100644 --- a/extensions/copilot/src/platform/notebook/common/alternativeNotebookTextDocument.ts +++ b/extensions/copilot/src/platform/notebook/common/alternativeNotebookTextDocument.ts @@ -148,6 +148,12 @@ abstract class AbstractAlternativeNotebookDocument { } } + /** + * Get the cell associated with a text document. + * @param textDocument The text document to find the cell for. + * @returns The notebook cell associated with the text document, or undefined if not found. + * If a cell was inserted into the notebook and this instance hasn't been updated yet, it will return undefined. + */ public getCell(textDocument: TextDocument): NotebookCell | undefined { return this.cellTextDocuments.get(textDocument); }