Path: blob/main/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookEntry.ts
5220 views
/*---------------------------------------------------------------------------------------------1* Copyright (c) Microsoft Corporation. All rights reserved.2* Licensed under the MIT License. See License.txt in the project root for license information.3*--------------------------------------------------------------------------------------------*/45import { streamToBuffer } from '../../../../../base/common/buffer.js';6import { CancellationToken } from '../../../../../base/common/cancellation.js';7import { StringSHA1 } from '../../../../../base/common/hash.js';8import { DisposableStore, IReference, thenRegisterOrDispose } from '../../../../../base/common/lifecycle.js';9import { ResourceMap, ResourceSet } from '../../../../../base/common/map.js';10import { Schemas } from '../../../../../base/common/network.js';11import { autorun, IObservable, ITransaction, ObservablePromise, observableValue, transaction } from '../../../../../base/common/observable.js';12import { isEqual } from '../../../../../base/common/resources.js';13import { assertType } from '../../../../../base/common/types.js';14import { URI } from '../../../../../base/common/uri.js';15import { generateUuid } from '../../../../../base/common/uuid.js';16import { Range } from '../../../../../editor/common/core/range.js';17import { LineRange } from '../../../../../editor/common/core/ranges/lineRange.js';18import { nullDocumentDiff } from '../../../../../editor/common/diff/documentDiffProvider.js';19import { DetailedLineRangeMapping, RangeMapping } from '../../../../../editor/common/diff/rangeMapping.js';20import { TextEdit } from '../../../../../editor/common/languages.js';21import { ITextModel } from '../../../../../editor/common/model.js';22import { IModelService } from '../../../../../editor/common/services/model.js';23import { ITextModelService } from '../../../../../editor/common/services/resolverService.js';24import { localize } from '../../../../../nls.js';25import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';26import { IFileService } from '../../../../../platform/files/common/files.js';27import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';28import { IUndoRedoElement, IUndoRedoService, UndoRedoElementType } from '../../../../../platform/undoRedo/common/undoRedo.js';29import { IEditorPane, SaveReason } from '../../../../common/editor.js';30import { IFilesConfigurationService } from '../../../../services/filesConfiguration/common/filesConfigurationService.js';31import { SnapshotContext } from '../../../../services/workingCopy/common/fileWorkingCopy.js';32import { IAiEditTelemetryService } from '../../../editTelemetry/browser/telemetry/aiEditTelemetry/aiEditTelemetryService.js';33import { NotebookTextDiffEditor } from '../../../notebook/browser/diff/notebookDiffEditor.js';34import { INotebookTextDiffEditor } from '../../../notebook/browser/diff/notebookDiffEditorBrowser.js';35import { CellDiffInfo } from '../../../notebook/browser/diff/notebookDiffViewModel.js';36import { getNotebookEditorFromEditorPane } from '../../../notebook/browser/notebookBrowser.js';37import { NotebookCellTextModel } from '../../../notebook/common/model/notebookCellTextModel.js';38import { NotebookTextModel } from '../../../notebook/common/model/notebookTextModel.js';39import { CellEditType, ICellDto2, ICellEditOperation, ICellReplaceEdit, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookSetting, NotebookTextModelChangedEvent, TransientOptions } from '../../../notebook/common/notebookCommon.js';40import { computeDiff } from '../../../notebook/common/notebookDiff.js';41import { INotebookEditorModelResolverService } from '../../../notebook/common/notebookEditorModelResolverService.js';42import { INotebookLoggingService } from '../../../notebook/common/notebookLoggingService.js';43import { INotebookService } from '../../../notebook/common/notebookService.js';44import { INotebookEditorWorkerService } from '../../../notebook/common/services/notebookWorkerService.js';45import { IChatService } from '../../common/chatService/chatService.js';46import { ChatEditKind, IModifiedEntryTelemetryInfo, IModifiedFileEntryEditorIntegration, ISnapshotEntry, ModifiedFileEntryState } from '../../common/editing/chatEditingService.js';47import { IChatResponseModel } from '../../common/model/chatModel.js';48import { AbstractChatEditingModifiedFileEntry } from './chatEditingModifiedFileEntry.js';49import { createSnapshot, deserializeSnapshot, getNotebookSnapshotFileURI, restoreSnapshot, SnapshotComparer } from './notebook/chatEditingModifiedNotebookSnapshot.js';50import { ChatEditingNewNotebookContentEdits } from './notebook/chatEditingNewNotebookContentEdits.js';51import { ChatEditingNotebookCellEntry } from './notebook/chatEditingNotebookCellEntry.js';52import { ChatEditingNotebookDiffEditorIntegration, ChatEditingNotebookEditorIntegration } from './notebook/chatEditingNotebookEditorIntegration.js';53import { ChatEditingNotebookFileSystemProvider } from './notebook/chatEditingNotebookFileSystemProvider.js';54import { adjustCellDiffAndOriginalModelBasedOnCellAddDelete, adjustCellDiffAndOriginalModelBasedOnCellMovements, adjustCellDiffForKeepingAnInsertedCell, adjustCellDiffForRevertingADeletedCell, adjustCellDiffForRevertingAnInsertedCell, calculateNotebookRewriteRatio, getCorrespondingOriginalCellIndex, isTransientIPyNbExtensionEvent } from './notebook/helpers.js';55import { countChanges, ICellDiffInfo, sortCellChanges } from './notebook/notebookCellChanges.js';565758const SnapshotLanguageId = 'VSCodeChatNotebookSnapshotLanguage';5960export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifiedFileEntry {61static NewModelCounter: number = 0;62private readonly modifiedModel: NotebookTextModel;63private readonly originalModel: NotebookTextModel;64override originalURI: URI;65/**66* JSON stringified version of the original notebook.67*/68override initialContent: string;69/**70* Whether we're still generating diffs from a response.71*/72private _isProcessingResponse = observableValue<boolean>('isProcessingResponse', false);73get isProcessingResponse(): IObservable<boolean> {74return this._isProcessingResponse;75}76private _isEditFromUs: boolean = false;77/**78* Whether all edits are from us, e.g. is possible a user has made edits, then this will be false.79*/80private _allEditsAreFromUs: boolean = true;81private readonly _changesCount = observableValue<number>(this, 0);82override changesCount: IObservable<number> = this._changesCount;8384private readonly cellEntryMap = new ResourceMap<ChatEditingNotebookCellEntry>();85private modifiedToOriginalCell = new ResourceMap<URI>();86private readonly _cellsDiffInfo = observableValue<ICellDiffInfo[]>('diffInfo', []);8788get cellsDiffInfo(): IObservable<ICellDiffInfo[]> {89return this._cellsDiffInfo;90}9192get viewType() {93return this.modifiedModel.viewType;94}9596/**97* List of Cell URIs that are edited,98* Will be cleared once all edits have been accepted.99* I.e. this will only contain URIS while acceptAgentEdits is being called & before `isLastEdit` is sent.100* I.e. this is populated only when edits are being streamed.101*/102private readonly editedCells = new ResourceSet();103104public static async create(uri: URI, _multiDiffEntryDelegate: { collapse: (transaction: ITransaction | undefined) => void }, telemetryInfo: IModifiedEntryTelemetryInfo, chatKind: ChatEditKind, initialContent: string | undefined, instantiationService: IInstantiationService): Promise<AbstractChatEditingModifiedFileEntry> {105return instantiationService.invokeFunction(async accessor => {106const notebookService = accessor.get(INotebookService);107const resolver = accessor.get(INotebookEditorModelResolverService);108const configurationServie = accessor.get(IConfigurationService);109const resourceRef: IReference<IResolvedNotebookEditorModel> = await resolver.resolve(uri);110const notebook = resourceRef.object.notebook;111const originalUri = getNotebookSnapshotFileURI(telemetryInfo.sessionResource, telemetryInfo.requestId, generateUuid(), notebook.uri.scheme === Schemas.untitled ? `/${notebook.uri.path}` : notebook.uri.path, notebook.viewType);112const [options, buffer] = await Promise.all([113notebookService.withNotebookDataProvider(resourceRef.object.notebook.notebookType),114notebookService.createNotebookTextDocumentSnapshot(notebook.uri, SnapshotContext.Backup, CancellationToken.None).then(s => streamToBuffer(s))115]);116const disposables = new DisposableStore();117// Register so that we can load this from file system.118disposables.add(ChatEditingNotebookFileSystemProvider.registerFile(originalUri, buffer));119const originalRef = await resolver.resolve(originalUri, notebook.viewType);120if (initialContent !== undefined) {121try {122restoreSnapshot(originalRef.object.notebook, initialContent);123} catch (ex) {124console.error(`Error restoring snapshot: ${initialContent}`, ex);125initialContent = createSnapshot(notebook, options.serializer.options, configurationServie);126}127} else {128initialContent = createSnapshot(notebook, options.serializer.options, configurationServie);129// Both models are the same, ensure the cell ids are the same, this way we get a perfect diffing.130// No need to generate edits for this.131// We want to ensure they are identitcal, possible original notebook was open and got modified.132// Or something gets changed between serialization & deserialization of the snapshot into the original.133// E.g. in jupyter notebooks the metadata contains transient data that gets updated after deserialization.134restoreSnapshot(originalRef.object.notebook, initialContent);135const edits: ICellEditOperation[] = [];136notebook.cells.forEach((cell, index) => {137const internalId = generateCellHash(cell.uri);138edits.push({ editType: CellEditType.PartialInternalMetadata, index, internalMetadata: { internalId } });139});140resourceRef.object.notebook.applyEdits(edits, true, undefined, () => undefined, undefined, false);141originalRef.object.notebook.applyEdits(edits, true, undefined, () => undefined, undefined, false);142}143const instance = instantiationService.createInstance(ChatEditingModifiedNotebookEntry, resourceRef, originalRef, _multiDiffEntryDelegate, options.serializer.options, telemetryInfo, chatKind, initialContent);144instance._register(disposables);145return instance;146});147}148149public static canHandleSnapshotContent(initialContent: string | undefined): boolean {150if (!initialContent) {151return false;152}153154try {155deserializeSnapshot(initialContent);156return true;157} catch (ex) {158// not a valid snapshot159return false;160}161}162163public static canHandleSnapshot(snapshot: ISnapshotEntry): boolean {164if (snapshot.languageId === SnapshotLanguageId && ChatEditingModifiedNotebookEntry.canHandleSnapshotContent(snapshot.current)) {165return true;166}167return false;168}169170private readonly initialContentComparer: SnapshotComparer;171172constructor(173private readonly modifiedResourceRef: IReference<IResolvedNotebookEditorModel>,174originalResourceRef: IReference<IResolvedNotebookEditorModel>,175private readonly _multiDiffEntryDelegate: { collapse: (transaction: ITransaction | undefined) => void },176private readonly transientOptions: TransientOptions | undefined,177telemetryInfo: IModifiedEntryTelemetryInfo,178kind: ChatEditKind,179initialContent: string,180@IConfigurationService private readonly configurationService: IConfigurationService,181@IFilesConfigurationService fileConfigService: IFilesConfigurationService,182@IChatService chatService: IChatService,183@IFileService fileService: IFileService,184@IInstantiationService instantiationService: IInstantiationService,185@ITextModelService private readonly textModelService: ITextModelService,186@IModelService private readonly modelService: IModelService,187@IUndoRedoService undoRedoService: IUndoRedoService,188@INotebookEditorWorkerService private readonly notebookEditorWorkerService: INotebookEditorWorkerService,189@INotebookLoggingService private readonly loggingService: INotebookLoggingService,190@INotebookEditorModelResolverService private readonly notebookResolver: INotebookEditorModelResolverService,191@IAiEditTelemetryService aiEditTelemetryService: IAiEditTelemetryService,192) {193super(modifiedResourceRef.object.notebook.uri, telemetryInfo, kind, configurationService, fileConfigService, chatService, fileService, undoRedoService, instantiationService, aiEditTelemetryService);194this.initialContentComparer = new SnapshotComparer(initialContent);195this.modifiedModel = this._register(modifiedResourceRef).object.notebook;196this.originalModel = this._register(originalResourceRef).object.notebook;197this.originalURI = this.originalModel.uri;198this.initialContent = initialContent;199this.initializeModelsFromDiff();200this._register(this.modifiedModel.onDidChangeContent(this.mirrorNotebookEdits, this));201}202203initializeModelsFromDiffImpl(cellsDiffInfo: CellDiffInfo[]) {204this.cellEntryMap.forEach(entry => entry.dispose());205this.cellEntryMap.clear();206const diffs = cellsDiffInfo.map((cellDiff, i) => {207switch (cellDiff.type) {208case 'delete':209return this.createDeleteCellDiffInfo(cellDiff.originalCellIndex);210case 'insert':211return this.createInsertedCellDiffInfo(cellDiff.modifiedCellIndex);212default:213return this.createModifiedCellDiffInfo(cellDiff.modifiedCellIndex, cellDiff.originalCellIndex);214}215});216this._cellsDiffInfo.set(diffs, undefined);217this._changesCount.set(countChanges(diffs), undefined);218}219220getIndexOfCellHandle(handle: number) {221return this.modifiedModel.cells.findIndex(c => c.handle === handle);222}223224private computeRequestId: number = 0;225async initializeModelsFromDiff() {226const id = ++this.computeRequestId;227if (this._areOriginalAndModifiedIdenticalImpl()) {228const cellsDiffInfo: CellDiffInfo[] = this.modifiedModel.cells.map((_, index) => {229return { type: 'unchanged', originalCellIndex: index, modifiedCellIndex: index } satisfies CellDiffInfo;230});231this.initializeModelsFromDiffImpl(cellsDiffInfo);232return;233}234const cellsDiffInfo: CellDiffInfo[] = [];235try {236this._isProcessingResponse.set(true, undefined);237const notebookDiff = await this.notebookEditorWorkerService.computeDiff(this.originalURI, this.modifiedURI);238if (id !== this.computeRequestId || this._store.isDisposed) {239return;240}241const result = computeDiff(this.originalModel, this.modifiedModel, notebookDiff);242if (result.cellDiffInfo.length) {243cellsDiffInfo.push(...result.cellDiffInfo);244}245} catch (ex) {246this.loggingService.error('Notebook Chat', 'Error computing diff:\n' + ex);247} finally {248this._isProcessingResponse.set(false, undefined);249}250this.initializeModelsFromDiffImpl(cellsDiffInfo);251}252updateCellDiffInfo(cellsDiffInfo: ICellDiffInfo[], transcation: ITransaction | undefined) {253this._cellsDiffInfo.set(sortCellChanges(cellsDiffInfo), transcation);254this._changesCount.set(countChanges(cellsDiffInfo), transcation);255}256257mirrorNotebookEdits(e: NotebookTextModelChangedEvent) {258if (this._isEditFromUs || this._isExternalEditInProgress || Array.from(this.cellEntryMap.values()).some(entry => entry.isEditFromUs)) {259return;260}261262// Possible user reverted the changes from SCM or the like.263// Or user just reverted the changes made via edits (e.g. edit made a change in a cell and user undid that change either by typing over or other).264// Computing snapshot is too slow, as this event gets triggered for every key stroke in a cell,265// const didResetToOriginalContent = createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService) === this.initialContent;266let didResetToOriginalContent = this.initialContentComparer.isEqual(this.modifiedModel);267const currentState = this._stateObs.get();268if (currentState === ModifiedFileEntryState.Modified && didResetToOriginalContent) {269this._stateObs.set(ModifiedFileEntryState.Rejected, undefined);270this.updateCellDiffInfo([], undefined);271this.initializeModelsFromDiff();272this._notifySessionAction('rejected');273return;274}275276if (!e.rawEvents.length) {277return;278}279280if (currentState === ModifiedFileEntryState.Rejected) {281return;282}283284if (isTransientIPyNbExtensionEvent(this.modifiedModel.notebookType, e)) {285return;286}287288this._allEditsAreFromUs = false;289this._userEditScheduler.schedule();290291// Changes to cell text is sync'ed and handled separately.292// See ChatEditingNotebookCellEntry._mirrorEdits293for (const event of e.rawEvents.filter(event => event.kind !== NotebookCellsChangeType.ChangeCellContent)) {294switch (event.kind) {295case NotebookCellsChangeType.ChangeDocumentMetadata: {296const edit: ICellEditOperation = {297editType: CellEditType.DocumentMetadata,298metadata: this.modifiedModel.metadata299};300this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);301break;302}303case NotebookCellsChangeType.ModelChange: {304let cellDiffs = sortCellChanges(this._cellsDiffInfo.get());305// Ensure the new notebook cells have internalIds306this._applyEditsSync(() => {307event.changes.forEach(change => {308change[2].forEach((cell, i) => {309if (cell.internalMetadata.internalId) {310return;311}312const index = change[0] + i;313const internalId = generateCellHash(cell.uri);314const edits: ICellEditOperation[] = [{ editType: CellEditType.PartialInternalMetadata, index, internalMetadata: { internalId } }];315this.modifiedModel.applyEdits(edits, true, undefined, () => undefined, undefined, false);316cell.internalMetadata ??= {};317cell.internalMetadata.internalId = internalId;318});319});320});321event.changes.forEach(change => {322cellDiffs = adjustCellDiffAndOriginalModelBasedOnCellAddDelete(change,323cellDiffs,324this.modifiedModel.cells.length,325this.originalModel.cells.length,326this.originalModel.applyEdits.bind(this.originalModel),327this.createModifiedCellDiffInfo.bind(this));328});329this.updateCellDiffInfo(cellDiffs, undefined);330this.disposeDeletedCellEntries();331break;332}333case NotebookCellsChangeType.ChangeCellLanguage: {334const index = getCorrespondingOriginalCellIndex(event.index, this._cellsDiffInfo.get());335if (typeof index === 'number') {336const edit: ICellEditOperation = {337editType: CellEditType.CellLanguage,338index,339language: event.language340};341this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);342}343break;344}345case NotebookCellsChangeType.ChangeCellMetadata: {346// ipynb and other extensions can alter metadata, ensure we update the original model in the corresponding cell.347const index = getCorrespondingOriginalCellIndex(event.index, this._cellsDiffInfo.get());348if (typeof index === 'number') {349const edit: ICellEditOperation = {350editType: CellEditType.Metadata,351index,352metadata: event.metadata353};354this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);355}356break;357}358case NotebookCellsChangeType.ChangeCellMime:359break;360case NotebookCellsChangeType.ChangeCellInternalMetadata: {361const index = getCorrespondingOriginalCellIndex(event.index, this._cellsDiffInfo.get());362if (typeof index === 'number') {363const edit: ICellEditOperation = {364editType: CellEditType.PartialInternalMetadata,365index,366internalMetadata: event.internalMetadata367};368this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);369}370break;371}372case NotebookCellsChangeType.Output: {373// User can run cells.374const index = getCorrespondingOriginalCellIndex(event.index, this._cellsDiffInfo.get());375if (typeof index === 'number') {376const edit: ICellEditOperation = {377editType: CellEditType.Output,378index,379append: event.append,380outputs: event.outputs381};382this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);383}384break;385}386case NotebookCellsChangeType.OutputItem: {387// outputs are shared between original and modified model, so the original model is already updated.388break;389}390case NotebookCellsChangeType.Move: {391const result = adjustCellDiffAndOriginalModelBasedOnCellMovements(event, this._cellsDiffInfo.get().slice());392if (result) {393this.originalModel.applyEdits(result[1], true, undefined, () => undefined, undefined, false);394this._cellsDiffInfo.set(result[0], undefined);395}396break;397}398default: {399break;400}401}402}403404didResetToOriginalContent = this.initialContentComparer.isEqual(this.modifiedModel);405if (currentState === ModifiedFileEntryState.Modified && didResetToOriginalContent) {406this._stateObs.set(ModifiedFileEntryState.Rejected, undefined);407this.updateCellDiffInfo([], undefined);408this.initializeModelsFromDiff();409return;410}411}412413protected override async _doAccept(): Promise<void> {414this.updateCellDiffInfo([], undefined);415const snapshot = createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService);416restoreSnapshot(this.originalModel, snapshot);417this.initializeModelsFromDiff();418await this._collapse(undefined);419420const config = this._fileConfigService.getAutoSaveConfiguration(this.modifiedURI);421if (this.modifiedModel.uri.scheme !== Schemas.untitled && (!config.autoSave || !this.notebookResolver.isDirty(this.modifiedURI))) {422// SAVE after accept for manual-savers, for auto-savers423// trigger explict save to get save participants going424await this._applyEdits(async () => {425try {426await this.modifiedResourceRef.object.save({427reason: SaveReason.EXPLICIT,428force: true,429});430} catch {431// ignored432}433});434}435}436437protected override async _doReject(): Promise<void> {438this.updateCellDiffInfo([], undefined);439if (this.createdInRequestId === this._telemetryInfo.requestId) {440await this._applyEdits(async () => {441await this.modifiedResourceRef.object.revert({ soft: true });442await this._fileService.del(this.modifiedURI);443});444this._onDidDelete.fire();445} else {446await this._applyEdits(async () => {447const snapshot = createSnapshot(this.originalModel, this.transientOptions, this.configurationService);448this.restoreSnapshotInModifiedModel(snapshot);449if (this._allEditsAreFromUs && Array.from(this.cellEntryMap.values()).every(entry => entry.allEditsAreFromUs)) {450// save the file after discarding so that the dirty indicator goes away451// and so that an intermediate saved state gets reverted452await this.modifiedResourceRef.object.save({ reason: SaveReason.EXPLICIT, skipSaveParticipants: true });453}454});455this.initializeModelsFromDiff();456await this._collapse(undefined);457}458}459460private async _collapse(transaction: ITransaction | undefined): Promise<void> {461this._multiDiffEntryDelegate.collapse(transaction);462}463464protected override _createEditorIntegration(editor: IEditorPane): IModifiedFileEntryEditorIntegration {465const notebookEditor = getNotebookEditorFromEditorPane(editor);466if (!notebookEditor && editor.getId() === NotebookTextDiffEditor.ID) {467const diffEditor = (editor.getControl() as INotebookTextDiffEditor);468return this._instantiationService.createInstance(ChatEditingNotebookDiffEditorIntegration, diffEditor, this._cellsDiffInfo);469}470assertType(notebookEditor);471return this._instantiationService.createInstance(ChatEditingNotebookEditorIntegration, this, editor, this.modifiedModel, this.originalModel, this._cellsDiffInfo);472}473474protected override _resetEditsState(tx: ITransaction): void {475super._resetEditsState(tx);476this.cellEntryMap.forEach(entry => !entry.isDisposed && entry.clearCurrentEditLineDecoration());477}478479protected override _createUndoRedoElement(response: IChatResponseModel): IUndoRedoElement | undefined {480const request = response.session.getRequests().find(req => req.id === response.requestId);481const label = request?.message.text ? localize('chatNotebookEdit1', "Chat Edit: '{0}'", request.message.text) : localize('chatNotebookEdit2', "Chat Edit");482const transientOptions = this.transientOptions;483const outputSizeLimit = this.configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;484485// create a snapshot of the current state of the model, before the next set of edits486let initial = createSnapshot(this.modifiedModel, transientOptions, outputSizeLimit);487let last = '';488let redoState = ModifiedFileEntryState.Rejected;489490return {491type: UndoRedoElementType.Resource,492resource: this.modifiedURI,493label,494code: 'chat.edit',495confirmBeforeUndo: false,496undo: async () => {497last = createSnapshot(this.modifiedModel, transientOptions, outputSizeLimit);498this._isEditFromUs = true;499try {500restoreSnapshot(this.modifiedModel, initial);501restoreSnapshot(this.originalModel, initial);502} finally {503this._isEditFromUs = false;504}505redoState = this._stateObs.get() === ModifiedFileEntryState.Accepted ? ModifiedFileEntryState.Accepted : ModifiedFileEntryState.Rejected;506this._stateObs.set(ModifiedFileEntryState.Rejected, undefined);507this.updateCellDiffInfo([], undefined);508this.initializeModelsFromDiff();509this._notifySessionAction('userModified');510},511redo: async () => {512initial = createSnapshot(this.modifiedModel, transientOptions, outputSizeLimit);513this._isEditFromUs = true;514try {515restoreSnapshot(this.modifiedModel, last);516restoreSnapshot(this.originalModel, last);517} finally {518this._isEditFromUs = false;519}520this._stateObs.set(redoState, undefined);521this.updateCellDiffInfo([], undefined);522this.initializeModelsFromDiff();523this._notifySessionAction('userModified');524}525};526}527528protected override async _areOriginalAndModifiedIdentical(): Promise<boolean> {529return this._areOriginalAndModifiedIdenticalImpl();530}531532private _areOriginalAndModifiedIdenticalImpl(): boolean {533const snapshot = createSnapshot(this.originalModel, this.transientOptions, this.configurationService);534return new SnapshotComparer(snapshot).isEqual(this.modifiedModel);535}536537private newNotebookEditGenerator?: ChatEditingNewNotebookContentEdits;538override async acceptAgentEdits(resource: URI, edits: (TextEdit | ICellEditOperation)[], isLastEdits: boolean, responseModel: IChatResponseModel | undefined): Promise<void> {539const isCellUri = resource.scheme === Schemas.vscodeNotebookCell;540const cell = isCellUri && this.modifiedModel.cells.find(cell => isEqual(cell.uri, resource));541let cellEntry: ChatEditingNotebookCellEntry | undefined;542if (cell) {543const index = this.modifiedModel.cells.indexOf(cell);544const entry = this._cellsDiffInfo.get().slice().find(entry => entry.modifiedCellIndex === index);545if (!entry) {546// Not possible.547console.error('Original cell model not found');548return;549}550551cellEntry = this.getOrCreateModifiedTextFileEntryForCell(cell, await entry.modifiedModel.promise, await entry.originalModel.promise);552}553554// For all cells that were edited, send the `isLastEdits` flag.555const finishPreviousCells = async () => {556await Promise.all(Array.from(this.editedCells).map(async (uri) => {557const cell = this.modifiedModel.cells.find(cell => isEqual(cell.uri, uri));558const cellEntry = cell && this.cellEntryMap.get(cell.uri);559await cellEntry?.acceptAgentEdits([], true, responseModel);560}));561this.editedCells.clear();562};563564await this._applyEdits(async () => {565await Promise.all(edits.map(async (edit, idx) => {566const last = isLastEdits && idx === edits.length - 1;567if (TextEdit.isTextEdit(edit)) {568// Possible we're getting the raw content for the notebook.569if (isEqual(resource, this.modifiedModel.uri)) {570this.newNotebookEditGenerator ??= this._instantiationService.createInstance(ChatEditingNewNotebookContentEdits, this.modifiedModel);571this.newNotebookEditGenerator.acceptTextEdits([edit]);572} else {573// If we get cell edits, its impossible to get text edits for the notebook uri.574this.newNotebookEditGenerator = undefined;575if (!this.editedCells.has(resource)) {576await finishPreviousCells();577this.editedCells.add(resource);578}579await cellEntry?.acceptAgentEdits([edit], last, responseModel);580}581} else {582// If we notebook edits, its impossible to get text edits for the notebook uri.583this.newNotebookEditGenerator = undefined;584this.acceptNotebookEdit(edit);585}586}));587});588589// If the last edit for a cell was sent, then handle it590if (isLastEdits) {591await finishPreviousCells();592}593594// isLastEdits can be true for cell Uris, but when its true for Cells edits.595// It cannot be true for the notebook itself.596isLastEdits = !isCellUri && isLastEdits;597598// If this is the last edit and & we got regular text edits for generating new notebook content599// Then generate notebook edits from those text edits & apply those notebook edits.600if (isLastEdits && this.newNotebookEditGenerator) {601const notebookEdits = await this.newNotebookEditGenerator.generateEdits();602this.newNotebookEditGenerator = undefined;603notebookEdits.forEach(edit => this.acceptNotebookEdit(edit));604}605606transaction((tx) => {607this._stateObs.set(ModifiedFileEntryState.Modified, tx);608if (!isLastEdits) {609const newRewriteRation = Math.max(this._rewriteRatioObs.get(), calculateNotebookRewriteRatio(this._cellsDiffInfo.get(), this.originalModel, this.modifiedModel));610this._rewriteRatioObs.set(Math.min(1, newRewriteRation), tx);611} else {612this.editedCells.clear();613this._resetEditsState(tx);614this._rewriteRatioObs.set(1, tx);615}616});617}618619private disposeDeletedCellEntries() {620const cellsUris = new ResourceSet(this.modifiedModel.cells.map(cell => cell.uri));621Array.from(this.cellEntryMap.keys()).forEach(uri => {622if (cellsUris.has(uri)) {623return;624}625this.cellEntryMap.get(uri)?.dispose();626this.cellEntryMap.delete(uri);627});628}629630acceptNotebookEdit(edit: ICellEditOperation): void {631// make the actual edit632this.modifiedModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);633this.disposeDeletedCellEntries();634635636if (edit.editType !== CellEditType.Replace) {637return;638}639// Ensure cells have internal Ids.640edit.cells.forEach((_, i) => {641const index = edit.index + i;642const cell = this.modifiedModel.cells[index];643if (cell.internalMetadata.internalId) {644return;645}646const internalId = generateCellHash(cell.uri);647const edits: ICellEditOperation[] = [{ editType: CellEditType.PartialInternalMetadata, index, internalMetadata: { internalId } }];648this.modifiedModel.applyEdits(edits, true, undefined, () => undefined, undefined, false);649});650651let diff: ICellDiffInfo[] = [];652if (edit.count === 0) {653// All existing indexes are shifted by number of cells added.654diff = sortCellChanges(this._cellsDiffInfo.get());655diff.forEach(d => {656if (d.type !== 'delete' && d.modifiedCellIndex >= edit.index) {657d.modifiedCellIndex += edit.cells.length;658}659});660const diffInsert = edit.cells.map((_, i) => this.createInsertedCellDiffInfo(edit.index + i));661diff.splice(edit.index, 0, ...diffInsert);662} else {663// All existing indexes are shifted by number of cells removed.664// And unchanged cells should be converted to deleted cells.665diff = sortCellChanges(this._cellsDiffInfo.get()).map((d) => {666if (d.type === 'unchanged' && d.modifiedCellIndex >= edit.index && d.modifiedCellIndex <= (edit.index + edit.count - 1)) {667return this.createDeleteCellDiffInfo(d.originalCellIndex);668}669if (d.type !== 'delete' && d.modifiedCellIndex >= (edit.index + edit.count)) {670d.modifiedCellIndex -= edit.count;671return d;672}673return d;674});675}676this.updateCellDiffInfo(diff, undefined);677}678679private computeStateAfterAcceptingRejectingChanges(accepted: boolean) {680const currentSnapshot = createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService);681if (new SnapshotComparer(currentSnapshot).isEqual(this.originalModel)) {682const state = accepted ? ModifiedFileEntryState.Accepted : ModifiedFileEntryState.Rejected;683this._stateObs.set(state, undefined);684this._notifySessionAction(accepted ? 'accepted' : 'rejected');685}686}687688createModifiedCellDiffInfo(modifiedCellIndex: number, originalCellIndex: number): ICellDiffInfo {689const modifiedCell = this.modifiedModel.cells[modifiedCellIndex];690const originalCell = this.originalModel.cells[originalCellIndex];691this.modifiedToOriginalCell.set(modifiedCell.uri, originalCell.uri);692const modifiedCellModelPromise = this.resolveCellModel(modifiedCell.uri);693const originalCellModelPromise = this.resolveCellModel(originalCell.uri);694695Promise.all([modifiedCellModelPromise, originalCellModelPromise]).then(([modifiedCellModel, originalCellModel]) => {696this.getOrCreateModifiedTextFileEntryForCell(modifiedCell, modifiedCellModel, originalCellModel);697});698699const diff = observableValue('diff', nullDocumentDiff);700const unchangedCell: ICellDiffInfo = {701type: 'unchanged',702modifiedCellIndex,703originalCellIndex,704keep: async (changes: DetailedLineRangeMapping) => {705const [modifiedCellModel, originalCellModel] = await Promise.all([modifiedCellModelPromise, originalCellModelPromise]);706const entry = this.getOrCreateModifiedTextFileEntryForCell(modifiedCell, modifiedCellModel, originalCellModel);707return entry ? entry.keep(changes) : false;708},709undo: async (changes: DetailedLineRangeMapping) => {710const [modifiedCellModel, originalCellModel] = await Promise.all([modifiedCellModelPromise, originalCellModelPromise]);711const entry = this.getOrCreateModifiedTextFileEntryForCell(modifiedCell, modifiedCellModel, originalCellModel);712return entry ? entry.undo(changes) : false;713},714modifiedModel: new ObservablePromise(modifiedCellModelPromise),715originalModel: new ObservablePromise(originalCellModelPromise),716diff717};718719return unchangedCell;720721}722createInsertedCellDiffInfo(modifiedCellIndex: number): ICellDiffInfo {723const cell = this.modifiedModel.cells[modifiedCellIndex];724const lines = cell.getValue().split(/\r?\n/);725const originalRange = new Range(1, 0, 1, 0);726const modifiedRange = new Range(1, 0, lines.length, lines[lines.length - 1].length);727const innerChanges = new RangeMapping(originalRange, modifiedRange);728const changes = [new DetailedLineRangeMapping(new LineRange(1, 1), new LineRange(1, lines.length), [innerChanges])];729// When a new cell is inserted, we use the ChatEditingCodeEditorIntegration to handle the edits.730// & to also display undo/redo and decorations.731// However that needs a modified and original model.732// For inserted cells there's no original model, so we create a new empty text model and pass that as the original.733const originalModelUri = this.modifiedModel.uri.with({ query: (ChatEditingModifiedNotebookEntry.NewModelCounter++).toString(), scheme: 'emptyCell' });734const originalModel = this.modelService.getModel(originalModelUri) || this._register(this.modelService.createModel('', null, originalModelUri));735this.modifiedToOriginalCell.set(cell.uri, originalModelUri);736const keep = async () => {737this._applyEditsSync(() => this.keepPreviouslyInsertedCell(cell));738this.computeStateAfterAcceptingRejectingChanges(true);739return true;740};741const undo = async () => {742this._applyEditsSync(() => this.undoPreviouslyInsertedCell(cell));743this.computeStateAfterAcceptingRejectingChanges(false);744return true;745};746this.resolveCellModel(cell.uri).then(modifiedModel => {747if (this._store.isDisposed) {748return;749}750// We want decorators for the cell just as we display decorators for modified cells.751// This way we have the ability to accept/reject the entire cell.752this.getOrCreateModifiedTextFileEntryForCell(cell, modifiedModel, originalModel);753});754return {755type: 'insert' as const,756originalCellIndex: undefined,757modifiedCellIndex: modifiedCellIndex,758keep,759undo,760modifiedModel: new ObservablePromise(this.resolveCellModel(cell.uri)),761originalModel: new ObservablePromise(Promise.resolve(originalModel)),762diff: observableValue('deletedCellDiff', {763changes,764identical: false,765moves: [],766quitEarly: false,767})768} satisfies ICellDiffInfo;769}770createDeleteCellDiffInfo(originalCellIndex: number): ICellDiffInfo {771const originalCell = this.originalModel.cells[originalCellIndex];772const lines = new Array(originalCell.textBuffer.getLineCount()).fill(0).map((_, i) => originalCell.textBuffer.getLineContent(i + 1));773const originalRange = new Range(1, 0, lines.length, lines[lines.length - 1].length);774const modifiedRange = new Range(1, 0, 1, 0);775const innerChanges = new RangeMapping(modifiedRange, originalRange);776const changes = [new DetailedLineRangeMapping(new LineRange(1, lines.length), new LineRange(1, 1), [innerChanges])];777const modifiedModelUri = this.modifiedModel.uri.with({ query: (ChatEditingModifiedNotebookEntry.NewModelCounter++).toString(), scheme: 'emptyCell' });778const modifiedModel = this.modelService.getModel(modifiedModelUri) || this._register(this.modelService.createModel('', null, modifiedModelUri));779const keep = async () => {780this._applyEditsSync(() => this.keepPreviouslyDeletedCell(this.originalModel.cells.indexOf(originalCell)));781this.computeStateAfterAcceptingRejectingChanges(true);782return true;783};784const undo = async () => {785this._applyEditsSync(() => this.undoPreviouslyDeletedCell(this.originalModel.cells.indexOf(originalCell), originalCell));786this.computeStateAfterAcceptingRejectingChanges(false);787return true;788};789790// This will be deleted.791return {792type: 'delete' as const,793modifiedCellIndex: undefined,794originalCellIndex,795originalModel: new ObservablePromise(this.resolveCellModel(originalCell.uri)),796modifiedModel: new ObservablePromise(Promise.resolve(modifiedModel)),797keep,798undo,799diff: observableValue('cellDiff', {800changes,801identical: false,802moves: [],803quitEarly: false,804})805} satisfies ICellDiffInfo;806}807808private undoPreviouslyInsertedCell(cell: NotebookCellTextModel) {809let diffs: ICellDiffInfo[] = [];810this._applyEditsSync(() => {811const index = this.modifiedModel.cells.indexOf(cell);812diffs = adjustCellDiffForRevertingAnInsertedCell(index,813this._cellsDiffInfo.get(),814this.modifiedModel.applyEdits.bind(this.modifiedModel));815});816this.disposeDeletedCellEntries();817this.updateCellDiffInfo(diffs, undefined);818}819820private keepPreviouslyInsertedCell(cell: NotebookCellTextModel) {821const modifiedCellIndex = this.modifiedModel.cells.indexOf(cell);822if (modifiedCellIndex === -1) {823// Not possible.824return;825}826const cellToInsert: ICellDto2 = {827cellKind: cell.cellKind,828language: cell.language,829metadata: cell.metadata,830outputs: cell.outputs,831source: cell.getValue(),832mime: cell.mime,833internalMetadata: {834internalId: cell.internalMetadata.internalId835}836};837this.cellEntryMap.get(cell.uri)?.dispose();838this.cellEntryMap.delete(cell.uri);839const cellDiffs = adjustCellDiffForKeepingAnInsertedCell(840modifiedCellIndex,841this._cellsDiffInfo.get().slice(),842cellToInsert,843this.originalModel.applyEdits.bind(this.originalModel),844this.createModifiedCellDiffInfo.bind(this)845);846this.updateCellDiffInfo(cellDiffs, undefined);847}848849private undoPreviouslyDeletedCell(deletedOriginalIndex: number, originalCell: NotebookCellTextModel) {850const cellToInsert: ICellDto2 = {851cellKind: originalCell.cellKind,852language: originalCell.language,853metadata: originalCell.metadata,854outputs: originalCell.outputs,855source: originalCell.getValue(),856mime: originalCell.mime,857internalMetadata: {858internalId: originalCell.internalMetadata.internalId859}860};861let cellDiffs: ICellDiffInfo[] = [];862this._applyEditsSync(() => {863cellDiffs = adjustCellDiffForRevertingADeletedCell(864deletedOriginalIndex,865this._cellsDiffInfo.get(),866cellToInsert,867this.modifiedModel.applyEdits.bind(this.modifiedModel),868this.createModifiedCellDiffInfo.bind(this)869);870});871this.updateCellDiffInfo(cellDiffs, undefined);872}873874875private keepPreviouslyDeletedCell(deletedOriginalIndex: number) {876// Delete this cell from original as well.877const edit: ICellReplaceEdit = { cells: [], count: 1, editType: CellEditType.Replace, index: deletedOriginalIndex, };878this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, false);879const diffs = sortCellChanges(this._cellsDiffInfo.get())880.filter(d => !(d.type === 'delete' && d.originalCellIndex === deletedOriginalIndex))881.map(diff => {882if (diff.type !== 'insert' && diff.originalCellIndex > deletedOriginalIndex) {883return {884...diff,885originalCellIndex: diff.originalCellIndex - 1,886};887}888return diff;889});890this.updateCellDiffInfo(diffs, undefined);891}892893private async _applyEdits(operation: () => Promise<void>) {894// make the actual edit895this._isEditFromUs = true;896try {897await operation();898} finally {899this._isEditFromUs = false;900}901}902903private _applyEditsSync(operation: () => void) {904// make the actual edit905this._isEditFromUs = true;906try {907operation();908} finally {909this._isEditFromUs = false;910}911}912913public getCurrentSnapshot() {914return createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService);915}916917override createSnapshot(chatSessionResource: URI, requestId: string | undefined, undoStop: string | undefined): ISnapshotEntry {918return {919resource: this.modifiedURI,920languageId: SnapshotLanguageId,921snapshotUri: getNotebookSnapshotFileURI(chatSessionResource, requestId, undoStop, this.modifiedURI.path, this.modifiedModel.viewType),922original: createSnapshot(this.originalModel, this.transientOptions, this.configurationService),923current: createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService),924state: this.state.get(),925telemetryInfo: this.telemetryInfo,926};927}928929override equalsSnapshot(snapshot: ISnapshotEntry | undefined): boolean {930return !!snapshot &&931isEqual(this.modifiedURI, snapshot.resource) &&932this.state.get() === snapshot.state &&933new SnapshotComparer(snapshot.original).isEqual(this.originalModel) &&934new SnapshotComparer(snapshot.current).isEqual(this.modifiedModel);935936}937938override async restoreFromSnapshot(snapshot: ISnapshotEntry, restoreToDisk = true): Promise<void> {939this.updateCellDiffInfo([], undefined);940this._stateObs.set(snapshot.state, undefined);941restoreSnapshot(this.originalModel, snapshot.original);942if (restoreToDisk) {943this.restoreSnapshotInModifiedModel(snapshot.current);944}945this.initializeModelsFromDiff();946}947948override async resetToInitialContent(): Promise<void> {949this.updateCellDiffInfo([], undefined);950this.restoreSnapshotInModifiedModel(this.initialContent);951this.initializeModelsFromDiff();952}953954public restoreModifiedModelFromSnapshot(snapshot: string) {955this.restoreSnapshotInModifiedModel(snapshot);956return this.initializeModelsFromDiff();957}958959private restoreSnapshotInModifiedModel(snapshot: string) {960if (snapshot === createSnapshot(this.modifiedModel, this.transientOptions, this.configurationService)) {961return;962}963964this._applyEditsSync(() => {965// See private _setDocValue in chatEditingModifiedDocumentEntry.ts966this.modifiedModel.pushStackElement();967restoreSnapshot(this.modifiedModel, snapshot);968this.modifiedModel.pushStackElement();969});970}971972private readonly cellTextModelMap = new ResourceMap<ITextModel>();973974private async resolveCellModel(cellURI: URI): Promise<ITextModel> {975const cell = this.originalModel.cells.concat(this.modifiedModel.cells).find(cell => isEqual(cell.uri, cellURI));976if (!cell) {977throw new Error('Cell not found');978}979const model = this.cellTextModelMap.get(cell.uri);980if (model) {981this.cellTextModelMap.set(cell.uri, model);982return model;983} else {984const textEditorModel = await thenRegisterOrDispose(this.textModelService.createModelReference(cell.uri), this._store);985const model = textEditorModel.object.textEditorModel;986this.cellTextModelMap.set(cell.uri, model);987return model;988}989}990991getOrCreateModifiedTextFileEntryForCell(cell: NotebookCellTextModel, modifiedCellModel: ITextModel, originalCellModel: ITextModel): ChatEditingNotebookCellEntry | undefined {992let cellEntry = this.cellEntryMap.get(cell.uri);993if (cellEntry) {994return cellEntry;995}996if (this._store.isDisposed) {997return;998}999const disposables = new DisposableStore();1000cellEntry = this._register(this._instantiationService.createInstance(ChatEditingNotebookCellEntry, this.modifiedResourceRef.object.resource, cell, modifiedCellModel, originalCellModel, () => this._isExternalEditInProgress, disposables));1001this.cellEntryMap.set(cell.uri, cellEntry);1002disposables.add(autorun(r => {1003if (this.modifiedModel.cells.indexOf(cell) === -1) {1004return;1005}1006const diffs = this.cellsDiffInfo.read(undefined).slice();1007const index = this.modifiedModel.cells.indexOf(cell);1008let entry = diffs.find(entry => entry.modifiedCellIndex === index);1009if (!entry) {1010// Not possible.1011return;1012}1013const entryIndex = diffs.indexOf(entry);1014entry.diff.set(cellEntry.diffInfo.read(r), undefined);1015if (cellEntry.diffInfo.read(undefined).identical && entry.type === 'modified') {1016entry = {1017...entry,1018type: 'unchanged',1019};1020}1021if (!cellEntry.diffInfo.read(undefined).identical && entry.type === 'unchanged') {1022entry = {1023...entry,1024type: 'modified',1025};1026}1027diffs.splice(entryIndex, 1, { ...entry });10281029transaction(tx => {1030this.updateCellDiffInfo(diffs, tx);1031});1032}));10331034disposables.add(autorun(r => {1035if (this.modifiedModel.cells.indexOf(cell) === -1) {1036return;1037}10381039const cellState = cellEntry.state.read(r);1040if (cellState === ModifiedFileEntryState.Accepted) {1041this.computeStateAfterAcceptingRejectingChanges(true);1042} else if (cellState === ModifiedFileEntryState.Rejected) {1043this.computeStateAfterAcceptingRejectingChanges(false);1044}1045}));10461047return cellEntry;1048}10491050async computeEditsFromSnapshots(beforeSnapshot: string, afterSnapshot: string): Promise<(TextEdit | ICellEditOperation)[]> {1051// For notebooks, we restore the snapshot and compute the cell-level edits1052// This is a simplified approach that replaces cells as needed10531054const beforeData = deserializeSnapshot(beforeSnapshot);1055const afterData = deserializeSnapshot(afterSnapshot);10561057const edits: ICellEditOperation[] = [];10581059// Simple approach: replace all cells1060// A more sophisticated approach would diff individual cells1061if (beforeData.data.cells.length > 0) {1062edits.push({1063editType: CellEditType.Replace,1064index: 0,1065count: beforeData.data.cells.length,1066cells: afterData.data.cells1067});1068} else if (afterData.data.cells.length > 0) {1069edits.push({1070editType: CellEditType.Replace,1071index: 0,1072count: 0,1073cells: afterData.data.cells1074});1075}10761077return edits;1078}10791080async save(): Promise<void> {1081if (this.modifiedModel.uri.scheme === Schemas.untitled) {1082return;1083}10841085// Save the notebook if dirty1086if (this.notebookResolver.isDirty(this.modifiedModel.uri)) {1087await this.modifiedResourceRef.object.save({1088reason: SaveReason.EXPLICIT,1089skipSaveParticipants: true1090});1091}1092}10931094async revertToDisk(): Promise<void> {1095if (this.modifiedModel.uri.scheme === Schemas.untitled) {1096return;1097}10981099// Revert to reload from disk1100await this.modifiedResourceRef.object.revert({ soft: false });1101}1102}110311041105function generateCellHash(cellUri: URI) {1106const hash = new StringSHA1();1107hash.update(cellUri.toString());1108return hash.digest().substring(0, 8);1109}111011111112