Skip to content

Commit

Permalink
Fix accept/reject in Notebook Chat (#241429)
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne authored Feb 21, 2025
1 parent f5f0a61 commit ad74b0a
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 43 deletions.
92 changes: 52 additions & 40 deletions src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookEntry.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,8 @@ import { ChatEditingNotebookEditorIntegration, countChanges, ICellDiffInfo } fro
import { ChatEditingSnapshotTextModelContentProvider } from './chatEditingTextModelContentProviders.js';



const noopKeep = () => Promise.resolve(true);
const noopUndo = () => Promise.resolve(true);

export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifiedFileEntry {
private readonly modifiedModel: NotebookTextModel;
Expand Down Expand Up @@ -115,26 +116,23 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
notebookService.withNotebookDataProvider(resourceRef.object.notebook.notebookType),
notebookService.createNotebookTextDocumentSnapshot(notebook.uri, SnapshotContext.Backup, CancellationToken.None).then(s => streamToBuffer(s))
]);
const originalDisposables = new DisposableStore();
originalDisposables.add(ChatEditingNotebookFileSystemProvider.registerFile(originalUri, buffer));

const disposables = new DisposableStore();
disposables.add(ChatEditingNotebookFileSystemProvider.registerFile(originalUri, buffer));
const originalRef = await resolver.resolve(originalUri, notebook.viewType);
originalDisposables.add(originalRef);
if (initialContent) {
restoreSnapshot(originalRef.object.notebook, initialContent);
}
initialContent = initialContent || createSnapshot(originalRef.object.notebook, options.serializer.options, configurationServie);
const dispsoables = new DisposableStore();
const modifiedCells = new ResourceMap<ITextModel>();
const originalCells = new ResourceMap<ITextModel>();
await Promise.all(resourceRef.object.notebook.cells.map(async cell => {
modifiedCells.set(cell.uri, dispsoables.add(await textModelService.createModelReference(cell.uri)).object.textEditorModel);
modifiedCells.set(cell.uri, disposables.add(await textModelService.createModelReference(cell.uri)).object.textEditorModel);
}).concat(originalRef.object.notebook.cells.map(async cell => {
originalCells.set(cell.uri, dispsoables.add(await textModelService.createModelReference(cell.uri)).object.textEditorModel);
originalCells.set(cell.uri, disposables.add(await textModelService.createModelReference(cell.uri)).object.textEditorModel);
})));

const instance = instantiationService.createInstance(ChatEditingModifiedNotebookEntry, resourceRef, originalRef, modifiedCells, originalCells, _multiDiffEntryDelegate, options.serializer.options, telemetryInfo, chatKind, initialContent);
instance._register(dispsoables);
instance._register(disposables);
return instance;
});
}
Expand All @@ -161,15 +159,13 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
super(modifiedResourceRef.object.notebook.uri, telemetryInfo, kind, configurationService, fileConfigService, chatService, fileService, instantiationService);
this._register(modifiedResourceRef);
this._register(originalResourceRef);
modifiedCellModels.forEach(cell => this._register(cell));
originalCellModels.forEach(cell => this._register(cell));
this.modifiedModel = modifiedResourceRef.object.notebook;
this.originalModel = originalResourceRef.object.notebook;
this.originalURI = this.originalModel.uri;
this.initialContent = initialContent;
this._register(this.modifiedModel.onDidChangeContent(this.mirrorNotebookEdits, this));
this._cellDiffInfo.set(this.modifiedModel.cells.map((cell, i) => ({ modifiedCellIndex: i, originalCellIndex: i, diff: nullDocumentDiff, type: 'unchanged' })), undefined);
this._maxModifiedLineNumbers.set(this.modifiedModel.cells.map(() => 0), undefined);
this.createEmptyDiffs();
this.modifiedModel.cells.forEach((cell, index) => {
const originalModel = this.originalCellModels.get(this.originalModel.cells[index].uri)!;
this.modifiedToOriginalCellMap.set(cell.uri, originalModel);
Expand All @@ -178,6 +174,22 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
this.emptyTextModel = this._register(model.createModel('', null, URI.parse('vscode-chat:///empty')));
}

createEmptyDiffs() {
this._cellDiffInfo.set(this.modifiedModel.cells.map((cell, i) => {
return { modifiedCellIndex: i, originalCellIndex: i, diff: this.getDiffForUnchangedCell(cell), type: 'unchanged' };
}), undefined);
}

getDiffForUnchangedCell(cell: NotebookCellTextModel): IDocumentDiff2 {
return {
...nullDocumentDiff,
keep: noopKeep,
undo: noopUndo,
originalModel: this.modifiedToOriginalCellMap.get(cell.uri) || this.emptyTextModel,
modifiedModel: this.modifiedCellModels.get(cell.uri) || this.emptyTextModel,
};
}

mirrorNotebookEdits(e: NotebookTextModelChangedEvent) {
if (this._isEditFromUs || Array.from(this.cellEntryMap.values()).some(entry => entry.isEditFromUs)) {
// TODO@DonJayamanne Apply this same edit to the original notebook.
Expand All @@ -204,24 +216,29 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
const snapshot = this.modifiedModel.createSnapshot({ context: SnapshotContext.Backup, outputSizeLimit, transientOptions: this.transientOptions });
this.originalModel.restoreSnapshot(snapshot, this.transientOptions);
this._changesCount.set(0, tx);
this._cellDiffInfo.set(this.modifiedModel.cells.map((cell, i) => ({ modifiedCellIndex: i, originalCellIndex: i, diff: nullDocumentDiff, type: 'unchanged' })), undefined);
this.createEmptyDiffs();

await this._collapse(tx);
}

protected override async _doReject(tx: ITransaction | undefined): Promise<void> {
if (this.createdInRequestId === this._telemetryInfo.requestId) {
await this.modifiedResourceRef.object.revert({ soft: true });
await this._fileService.del(this.modifiedURI);
await this._applyEdits(async () => {
await this.modifiedResourceRef.object.revert({ soft: true });
await this._fileService.del(this.modifiedURI);
});
this._onDidDelete.fire();
} else {
const outputSizeLimit = this.configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;
const snapshot = this.originalModel.createSnapshot({ context: SnapshotContext.Backup, outputSizeLimit, transientOptions: this.transientOptions });
this.modifiedModel.restoreSnapshot(snapshot, this.transientOptions);
if (this._allEditsAreFromUs && this._entries.get().every(entry => entry.allEditsAreFromUs)) {
// save the file after discarding so that the dirty indicator goes away
// and so that an intermediate saved state gets reverted
await this.modifiedResourceRef.object.save({ reason: SaveReason.EXPLICIT, skipSaveParticipants: true });
}
await this._applyEdits(async () => {
const outputSizeLimit = this.configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;
const snapshot = this.originalModel.createSnapshot({ context: SnapshotContext.Backup, outputSizeLimit, transientOptions: this.transientOptions });
this.modifiedModel.restoreSnapshot(snapshot, this.transientOptions);
if (this._allEditsAreFromUs && this._entries.get().every(entry => entry.allEditsAreFromUs)) {
// save the file after discarding so that the dirty indicator goes away
// and so that an intermediate saved state gets reverted
await this.modifiedResourceRef.object.save({ reason: SaveReason.EXPLICIT, skipSaveParticipants: true });
}
});
await this._collapse(tx);
}

Expand All @@ -238,7 +255,7 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie

protected override _resetEditsState(tx: ITransaction): void {
super._resetEditsState(tx);
this.cellEntryMap.forEach(entry => entry.clearCurrentEditLineDecoration());
this.cellEntryMap.forEach(entry => !entry.disposed && entry.clearCurrentEditLineDecoration());
}

override async acceptAgentEdits(resource: URI, edits: (TextEdit | ICellEditOperation)[], isLastEdits: boolean, responseModel: IChatResponseModel): Promise<void> {
Expand All @@ -255,7 +272,7 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
});
};

this._applyEdits(async () => {
await this._applyEdits(async () => {
await Promise.all(edits.map(async edit => {
if (TextEdit.isTextEdit(edit)) {
if (!this.editedCells.has(resource)) {
Expand Down Expand Up @@ -323,21 +340,11 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
this.modifiedCellModels.set(cell.uri, modifiedModel);
this.modifiedToOriginalCellMap.set(cell.uri, originalModel);
const keep = async () => {
try {
this._isEditFromUs = true;
this.keepPreviouslyInsertedCell(cell);
} finally {
this._isEditFromUs = false;
}
await this._applyEdits(async () => this.keepPreviouslyInsertedCell(cell));
return true;
};
const undo = async () => {
try {
this._isEditFromUs = true;
this.undoPreviouslyInsertedCell(cell);
} finally {
this._isEditFromUs = false;
}
await this._applyEdits(async () => this.undoPreviouslyInsertedCell(cell));
return true;
};
this.getOrCreateModifiedTextFileEntryForCell(cell, keep, undo);
Expand Down Expand Up @@ -513,7 +520,6 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
if (cellEntry) {
return cellEntry;
}
const index = this.modifiedModel.cells.indexOf(cell);
const originalCellModel = this.modifiedToOriginalCellMap.get(cell.uri);
const modifiedCellModel = this.modifiedCellModels.get(cell.uri);
if (!modifiedCellModel || !originalCellModel) {
Expand All @@ -526,13 +532,16 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
this._register(autorun(r => {
const cellDiff = cellEntry.diffInfo.read(r);
let diffs = this.cellDiffInfo.get().slice();
const index = this.modifiedModel.cells.indexOf(cell);
const entry = diffs.find(entry => entry.modifiedCellIndex === index);
if (!entry) {
// Not possible.
return;
}
entry.diff = cellDiff;
entry.diff = cellDiff.identical ? nullDocumentDiff : cellDiff;
entry.diff = { ...entry.diff, ...cellDiff };
if (cellDiff.identical) {
entry.diff = { ...entry.diff, ...nullDocumentDiff };
}
if (entry.type === 'unchanged' || entry.type === 'modified') {
entry.type = cellDiff.identical ? 'unchanged' : 'modified';
}
Expand Down Expand Up @@ -696,6 +705,9 @@ class ChatEditingNotebookCellEntry extends ObservableDisposable {
}

public clearCurrentEditLineDecoration() {
if (this.modifiedModel.isDisposed()) {
return;
}
this._editDecorations = this.modifiedModel.deltaDecorations(this._editDecorations, []);
}

Expand Down
6 changes: 3 additions & 3 deletions src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingNotebookEditorIntegration.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ export type ICellDiffInfo = {
originalCellIndex: number;
modifiedCellIndex: number;
type: 'unchanged';
diff: IDocumentDiff; // Null diff Change (property to be consistent with others, also we have a list of all line numbers)
diff: IDocumentDiff2; // Null diff Change (property to be consistent with others, also we have a list of all line numbers)
} | {
originalCellIndex: number;
modifiedCellIndex: number;
Expand Down Expand Up @@ -131,12 +131,12 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
this.cellEditorIntegrations.set(cell, { integration, diff: diff2 });
this._register(integration);
this._register(editor.onDidDispose(() => {
this.cellEditorIntegrations.get(cell)!.integration.dispose();
this.cellEditorIntegrations.get(cell)?.integration.dispose();
this.cellEditorIntegrations.delete(cell);
}));
this._register(editor.onDidChangeModel(() => {
if (editor.getModel() !== cell.textModel) {
this.cellEditorIntegrations.get(cell)!.integration.dispose();
this.cellEditorIntegrations.get(cell)?.integration.dispose();
this.cellEditorIntegrations.delete(cell);
}
}));
Expand Down

0 comments on commit ad74b0a

Please sign in to comment.