Skip to content

Commit

Permalink
Remove direct refs of original & modified models (#241134)
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne authored Feb 19, 2025
1 parent df66760 commit 7e0db1c
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 14 deletions.
11 changes: 7 additions & 4 deletions ...vs/workbench/contrib/notebook/browser/contrib/chatEdit/notebookOriginalModelRefFactory.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import { INotebookService } from '../../../common/notebookService.js';
import { bufferToStream, VSBuffer } from '../../../../../../base/common/buffer.js';
import { NotebookTextModel } from '../../../common/model/notebookTextModel.js';
import { createDecorator, IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
import { ChatEditingModifiedDocumentEntry } from '../../../../chat/browser/chatEditing/chatEditingModifiedDocumentEntry.js';
import { ITextModelService } from '../../../../../../editor/common/services/resolverService.js';


export const INotebookOriginalModelReferenceFactory = createDecorator<INotebookOriginalModelReferenceFactory>('INotebookOriginalModelReferenceFactory');
Expand All @@ -22,7 +22,9 @@ export interface INotebookOriginalModelReferenceFactory {

export class OriginalNotebookModelReferenceCollection extends ReferenceCollection<Promise<NotebookTextModel>> {
private readonly modelsToDispose = new Set<string>();
constructor(@INotebookService private readonly notebookService: INotebookService) {
constructor(@INotebookService private readonly notebookService: INotebookService,
@ITextModelService private readonly modelService: ITextModelService
) {
super();
}

Expand All @@ -33,9 +35,10 @@ export class OriginalNotebookModelReferenceCollection extends ReferenceCollectio
if (model) {
return model;
}
// TODO@DonJayamanne FIX ME, don't use `originalModel`
const bytes = VSBuffer.fromString((fileEntry as ChatEditingModifiedDocumentEntry).originalModel.getValue());
const modelRef = await this.modelService.createModelReference(uri);
const bytes = VSBuffer.fromString(modelRef.object.textEditorModel.getValue());
const stream = bufferToStream(bytes);
modelRef.dispose();

return this.notebookService.createNotebookTextModel(viewType, uri, stream);
}
Expand Down
38 changes: 28 additions & 10 deletions src/vs/workbench/contrib/notebook/browser/contrib/chatEdit/notebookSynchronizer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import { raceCancellation, ThrottledDelayer } from '../../../../../../base/commo
import { CellDiffInfo, computeDiff, prettyChanges } from '../../diff/notebookDiffViewModel.js';
import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js';
import { INotebookEditorWorkerService } from '../../../common/services/notebookWorkerService.js';
import { ChatEditingModifiedDocumentEntry } from '../../../../chat/browser/chatEditing/chatEditingModifiedDocumentEntry.js';
import { CellEditType, ICellDto2, ICellEditOperation, ICellReplaceEdit, NotebookData, NotebookSetting } from '../../../common/notebookCommon.js';
import { URI } from '../../../../../../base/common/uri.js';
import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js';
Expand All @@ -30,6 +29,7 @@ import { ITextModelService } from '../../../../../../editor/common/services/reso
import { SnapshotContext } from '../../../../../services/workingCopy/common/fileWorkingCopy.js';
import { INotebookEditorService } from '../../services/notebookEditorService.js';
import { CellEditState } from '../../notebookBrowser.js';
import { IModelService } from '../../../../../../editor/common/services/model.js';


export const INotebookModelSynchronizerFactory = createDecorator<INotebookModelSynchronizerFactory>('INotebookModelSynchronizerFactory');
Expand Down Expand Up @@ -81,6 +81,7 @@ export class NotebookModelSynchronizer extends Disposable {
@INotebookService private readonly notebookService: INotebookService,
@INotebookEditorService private readonly notebookEditorService: INotebookEditorService,
@IChatService chatService: IChatService,
@IModelService private readonly modelService: IModelService,
@ITextModelService private readonly textModelService: ITextModelService,
@INotebookLoggingService private readonly logService: INotebookLoggingService,
@IConfigurationService private readonly configurationService: IConfigurationService,
Expand Down Expand Up @@ -128,12 +129,14 @@ export class NotebookModelSynchronizer extends Disposable {
snapshotCreated = true;
}

// TODO@DonJayamanne FIX ME, don't use `modifiedModel`
const modifiedModel = (entry as ChatEditingModifiedDocumentEntry).modifiedModel;
const modifiedModel = this.modelService.getModel(entry.modifiedURI);
if (!modifiedModel) {
return;
}
let cancellationToken = store.add(new CancellationTokenSource());
store.add(modifiedModel.onDidChangeContent(async () => {
// TODO@DonJayamanne FIX ME, don't use `originalModel`
if (!this.isTextEditFromUs && !modifiedModel.isDisposed() && !(entry as ChatEditingModifiedDocumentEntry).originalModel.isDisposed() && modifiedModel.getValue() !== (entry as ChatEditingModifiedDocumentEntry).originalModel.getValue()) {
const originalModel = this.modelService.getModel(entry.originalURI);
if (originalModel && !this.isTextEditFromUs && !modifiedModel.isDisposed() && !originalModel.isDisposed() && modifiedModel.getValue() !== originalModel.getValue()) {
cancellationToken = store.add(new CancellationTokenSource());
updateNotebookModel(entry, cancellationToken.token);
}
Expand Down Expand Up @@ -258,7 +261,10 @@ export class NotebookModelSynchronizer extends Disposable {
}

private async accept(entry: IModifiedFileEntry) {
const modifiedModel = (entry as ChatEditingModifiedDocumentEntry).modifiedModel;
const modifiedModel = this.modelService.getModel(entry.modifiedURI);
if (!modifiedModel) {
return;
}
const content = modifiedModel.getValue();
await this.updateNotebook(VSBuffer.fromString(content), false);
this._diffInfo.set(undefined, undefined);
Expand All @@ -270,7 +276,10 @@ export class NotebookModelSynchronizer extends Disposable {
}

private async updateTextDocumentModel(entry: IModifiedFileEntry) {
const modifiedModel = (entry as ChatEditingModifiedDocumentEntry).modifiedModel;
const modifiedModel = this.modelService.getModel(entry.modifiedURI);
if (!modifiedModel) {
return;
}
const stream = await this.notebookService.createNotebookTextDocumentSnapshot(this.model.uri, SnapshotContext.Save, CancellationToken.None);
const buffer = await streamToBuffer(stream);
const text = new TextDecoder().decode(buffer.buffer);
Expand Down Expand Up @@ -300,7 +309,12 @@ export class NotebookModelSynchronizer extends Disposable {
}

private async updateNotebookModel(entry: IModifiedFileEntry, token: CancellationToken) {
const modifiedModelVersion = (entry as ChatEditingModifiedDocumentEntry).modifiedModel.getVersionId();
const modifiedModel = this.modelService.getModel(entry.modifiedURI);
if (!modifiedModel) {
return;
}

const modifiedModelVersion = modifiedModel.getVersionId();
const currentModel = this.model;
const modelVersion = currentModel?.versionId ?? 0;
const modelWithChatEdits = await this.getModifiedModelForDiff(entry, token);
Expand All @@ -312,7 +326,7 @@ export class NotebookModelSynchronizer extends Disposable {
const cellDiffInfo = (await this.computeDiff(originalModel, modelWithChatEdits, token))?.cellDiffInfo;
// This is the diff from the current model to the model with chat edits.
const cellDiffInfoToApplyEdits = (await this.computeDiff(currentModel, modelWithChatEdits, token))?.cellDiffInfo;
const currentVersion = (entry as ChatEditingModifiedDocumentEntry).modifiedModel.getVersionId();
const currentVersion = modifiedModel.getVersionId();
if (!cellDiffInfo || !cellDiffInfoToApplyEdits || token.isCancellationRequested || currentVersion !== modifiedModelVersion || modelVersion !== currentModel.versionId) {
return;
}
Expand Down Expand Up @@ -443,7 +457,11 @@ export class NotebookModelSynchronizer extends Disposable {
private previousUriOfModelForDiff?: URI;

private async getModifiedModelForDiff(entry: IModifiedFileEntry, token: CancellationToken): Promise<NotebookTextModel | undefined> {
const text = (entry as ChatEditingModifiedDocumentEntry).modifiedModel.getValue();
const modifiedModel = this.modelService.getModel(entry.modifiedURI);
if (!modifiedModel) {
return;
}
const text = modifiedModel.getValue();
const bytes = VSBuffer.fromString(text);
const uri = entry.modifiedURI.with({ scheme: `NotebookChatEditorController.modifiedScheme${Date.now().toString()}` });
const stream = bufferToStream(bytes);
Expand Down

0 comments on commit 7e0db1c

Please sign in to comment.