Skip to content

Commit 29081ed

Browse files
authored
* keep model references for a little longer so that they can be reused * don't clear diff editor when rerendering but when its model gets disposed (a little later)
1 parent 1f28297 commit 29081ed

File tree

2 files changed

+125
-78
lines changed

2 files changed

+125
-78
lines changed

src/vs/workbench/contrib/chat/browser/chatContentParts/chatTextEditContentPart.ts

Lines changed: 111 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,22 @@
66
import * as dom from 'vs/base/browser/dom';
77
import { CancellationTokenSource } from 'vs/base/common/cancellation';
88
import { Emitter, Event } from 'vs/base/common/event';
9-
import { Disposable, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
9+
import { Disposable, IDisposable, IReference, RefCountedDisposable, toDisposable } from 'vs/base/common/lifecycle';
1010
import { Schemas } from 'vs/base/common/network';
1111
import { isEqual } from 'vs/base/common/resources';
12+
import { assertType } from 'vs/base/common/types';
1213
import { URI } from 'vs/base/common/uri';
1314
import { generateUuid } from 'vs/base/common/uuid';
1415
import { ISingleEditOperation } from 'vs/editor/common/core/editOperation';
1516
import { TextEdit } from 'vs/editor/common/languages';
1617
import { createTextBufferFactoryFromSnapshot } from 'vs/editor/common/model/textModel';
1718
import { IModelService } from 'vs/editor/common/services/model';
1819
import { DefaultModelSHA1Computer } from 'vs/editor/common/services/modelService';
19-
import { ITextModelService } from 'vs/editor/common/services/resolverService';
20+
import { IResolvedTextEditorModel, ITextModelService } from 'vs/editor/common/services/resolverService';
2021
import { localize } from 'vs/nls';
2122
import { MenuId } from 'vs/platform/actions/common/actions';
22-
import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
23+
import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
24+
import { createDecorator, IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
2325
import { IChatListItemRendererOptions } from 'vs/workbench/contrib/chat/browser/chat';
2426
import { IDisposableReference, ResourcePool } from 'vs/workbench/contrib/chat/browser/chatContentParts/chatCollections';
2527
import { IChatContentPart, IChatContentPartRenderContext } from 'vs/workbench/contrib/chat/browser/chatContentParts/chatContentParts';
@@ -28,13 +30,20 @@ import { ChatEditorOptions } from 'vs/workbench/contrib/chat/browser/chatOptions
2830
import { CodeCompareBlockPart, ICodeCompareBlockData, ICodeCompareBlockDiffData } from 'vs/workbench/contrib/chat/browser/codeBlockPart';
2931
import { IChatProgressRenderableResponseContent, IChatTextEditGroup } from 'vs/workbench/contrib/chat/common/chatModel';
3032
import { IChatService } from 'vs/workbench/contrib/chat/common/chatService';
31-
import { isResponseVM } from 'vs/workbench/contrib/chat/common/chatViewModel';
33+
import { IChatResponseViewModel, isResponseVM } from 'vs/workbench/contrib/chat/common/chatViewModel';
3234

3335
const $ = dom.$;
3436

37+
const ICodeCompareModelService = createDecorator<ICodeCompareModelService>('ICodeCompareModelService');
38+
39+
interface ICodeCompareModelService {
40+
_serviceBrand: undefined;
41+
createModel(response: IChatResponseViewModel, chatTextEdit: IChatTextEditGroup): Promise<IReference<{ originalSha1: string; original: IResolvedTextEditorModel; modified: IResolvedTextEditorModel }>>;
42+
}
43+
3544
export class ChatTextEditContentPart extends Disposable implements IChatContentPart {
3645
public readonly domNode: HTMLElement;
37-
private readonly ref: IDisposableReference<CodeCompareBlockPart> | undefined;
46+
private readonly comparePart: IDisposableReference<CodeCompareBlockPart> | undefined;
3847

3948
private readonly _onDidChangeHeight = this._register(new Emitter<void>());
4049
public readonly onDidChangeHeight = this._onDidChangeHeight.event;
@@ -45,16 +54,16 @@ export class ChatTextEditContentPart extends Disposable implements IChatContentP
4554
rendererOptions: IChatListItemRendererOptions,
4655
diffEditorPool: DiffEditorPool,
4756
currentWidth: number,
48-
@ITextModelService private readonly textModelService: ITextModelService,
49-
@IModelService private readonly modelService: IModelService,
50-
@IChatService private readonly chatService: IChatService,
57+
@ICodeCompareModelService private readonly codeCompareModelService: ICodeCompareModelService
5158
) {
5259
super();
5360
const element = context.element;
5461

62+
assertType(isResponseVM(element));
63+
5564
// TODO@jrieken move this into the CompareCodeBlock and properly say what kind of changes happen
5665
if (rendererOptions.renderTextEditsAsSummary?.(chatTextEdit.uri)) {
57-
if (isResponseVM(element) && element.response.value.every(item => item.kind === 'textEditGroup')) {
66+
if (element.response.value.every(item => item.kind === 'textEditGroup')) {
5867
this.domNode = $('.interactive-edits-summary', undefined, !element.isComplete
5968
? ''
6069
: element.isCanceled
@@ -75,14 +84,13 @@ export class ChatTextEditContentPart extends Disposable implements IChatContentP
7584
this._register(toDisposable(() => {
7685
isDisposed = true;
7786
cts.dispose(true);
78-
this.ref?.object.clearModel();
7987
}));
8088

81-
this.ref = this._register(diffEditorPool.get());
89+
this.comparePart = this._register(diffEditorPool.get());
8290

8391
// Attach this after updating text/layout of the editor, so it should only be fired when the size updates later (horizontal scrollbar, wrapping)
8492
// not during a renderElement OR a progressive render (when we will be firing this event anyway at the end of the render)
85-
this._register(this.ref.object.onDidChangeContentHeight(() => {
93+
this._register(this.comparePart.object.onDidChangeContentHeight(() => {
8694
this._onDidChangeHeight.fire();
8795
}));
8896

@@ -91,7 +99,7 @@ export class ChatTextEditContentPart extends Disposable implements IChatContentP
9199
edit: chatTextEdit,
92100
diffData: (async () => {
93101

94-
const ref = await this.textModelService.createModelReference(chatTextEdit.uri);
102+
const ref = await this.codeCompareModelService.createModel(element, chatTextEdit);
95103

96104
if (isDisposed) {
97105
ref.dispose();
@@ -100,70 +108,21 @@ export class ChatTextEditContentPart extends Disposable implements IChatContentP
100108

101109
this._register(ref);
102110

103-
const original = ref.object.textEditorModel;
104-
let originalSha1: string = '';
105-
106-
if (chatTextEdit.state) {
107-
originalSha1 = chatTextEdit.state.sha1;
108-
} else {
109-
const sha1 = new DefaultModelSHA1Computer();
110-
if (sha1.canComputeSHA1(original)) {
111-
originalSha1 = sha1.computeSHA1(original);
112-
chatTextEdit.state = { sha1: originalSha1, applied: 0 };
113-
}
114-
}
115-
116-
const modified = this.modelService.createModel(
117-
createTextBufferFactoryFromSnapshot(original.createSnapshot()),
118-
{ languageId: original.getLanguageId(), onDidChange: Event.None },
119-
URI.from({ scheme: Schemas.vscodeChatCodeBlock, path: original.uri.path, query: generateUuid() }),
120-
false
121-
);
122-
const modRef = await this.textModelService.createModelReference(modified.uri);
123-
this._register(modRef);
124-
125-
const editGroups: ISingleEditOperation[][] = [];
126-
if (isResponseVM(element)) {
127-
const chatModel = this.chatService.getSession(element.sessionId)!;
128-
129-
for (const request of chatModel.getRequests()) {
130-
if (!request.response) {
131-
continue;
132-
}
133-
for (const item of request.response.response.value) {
134-
if (item.kind !== 'textEditGroup' || item.state?.applied || !isEqual(item.uri, chatTextEdit.uri)) {
135-
continue;
136-
}
137-
for (const group of item.edits) {
138-
const edits = group.map(TextEdit.asEditOperation);
139-
editGroups.push(edits);
140-
}
141-
}
142-
if (request.response === element.model) {
143-
break;
144-
}
145-
}
146-
}
147-
148-
for (const edits of editGroups) {
149-
modified.pushEditOperations(null, edits, () => null);
150-
}
151-
152111
return {
153-
modified,
154-
original,
155-
originalSha1
112+
modified: ref.object.modified.textEditorModel,
113+
original: ref.object.original.textEditorModel,
114+
originalSha1: ref.object.originalSha1
156115
} satisfies ICodeCompareBlockDiffData;
157116
})()
158117
};
159-
this.ref.object.render(data, currentWidth, cts.token);
118+
this.comparePart.object.render(data, currentWidth, cts.token);
160119

161-
this.domNode = this.ref.object.element;
120+
this.domNode = this.comparePart.object.element;
162121
}
163122
}
164123

165124
layout(width: number): void {
166-
this.ref?.object.layout(width);
125+
this.comparePart?.object.layout(width);
167126
}
168127

169128
hasSameContent(other: IChatProgressRenderableResponseContent): boolean {
@@ -210,3 +169,87 @@ export class DiffEditorPool extends Disposable {
210169
};
211170
}
212171
}
172+
173+
class CodeCompareModelService implements ICodeCompareModelService {
174+
175+
declare readonly _serviceBrand: undefined;
176+
177+
constructor(
178+
@ITextModelService private readonly textModelService: ITextModelService,
179+
@IModelService private readonly modelService: IModelService,
180+
@IChatService private readonly chatService: IChatService,
181+
) { }
182+
183+
async createModel(element: IChatResponseViewModel, chatTextEdit: IChatTextEditGroup): Promise<IReference<{ originalSha1: string; original: IResolvedTextEditorModel; modified: IResolvedTextEditorModel }>> {
184+
185+
const original = await this.textModelService.createModelReference(chatTextEdit.uri);
186+
187+
const modified = await this.textModelService.createModelReference((this.modelService.createModel(
188+
createTextBufferFactoryFromSnapshot(original.object.textEditorModel.createSnapshot()),
189+
{ languageId: original.object.textEditorModel.getLanguageId(), onDidChange: Event.None },
190+
URI.from({ scheme: Schemas.vscodeChatCodeBlock, path: chatTextEdit.uri.path, query: generateUuid() }),
191+
false
192+
)).uri);
193+
194+
const d = new RefCountedDisposable(toDisposable(() => {
195+
original.dispose();
196+
modified.dispose();
197+
}));
198+
199+
// compute the sha1 of the original model
200+
let originalSha1: string = '';
201+
if (chatTextEdit.state) {
202+
originalSha1 = chatTextEdit.state.sha1;
203+
} else {
204+
const sha1 = new DefaultModelSHA1Computer();
205+
if (sha1.canComputeSHA1(original.object.textEditorModel)) {
206+
originalSha1 = sha1.computeSHA1(original.object.textEditorModel);
207+
chatTextEdit.state = { sha1: originalSha1, applied: 0 };
208+
}
209+
}
210+
211+
// apply edits to the "modified" model
212+
const chatModel = this.chatService.getSession(element.sessionId)!;
213+
const editGroups: ISingleEditOperation[][] = [];
214+
for (const request of chatModel.getRequests()) {
215+
if (!request.response) {
216+
continue;
217+
}
218+
for (const item of request.response.response.value) {
219+
if (item.kind !== 'textEditGroup' || item.state?.applied || !isEqual(item.uri, chatTextEdit.uri)) {
220+
continue;
221+
}
222+
for (const group of item.edits) {
223+
const edits = group.map(TextEdit.asEditOperation);
224+
editGroups.push(edits);
225+
}
226+
}
227+
if (request.response === element.model) {
228+
break;
229+
}
230+
}
231+
for (const edits of editGroups) {
232+
modified.object.textEditorModel.pushEditOperations(null, edits, () => null);
233+
}
234+
235+
// self-acquire a reference to diff models for a short while
236+
// because streaming usually means we will be using the original-model
237+
// repeatedly and thereby also should reuse the modified-model and just
238+
// update it with more edits
239+
d.acquire();
240+
setTimeout(() => d.release(), 5000);
241+
242+
return {
243+
object: {
244+
originalSha1,
245+
original: original.object,
246+
modified: modified.object
247+
},
248+
dispose() {
249+
d.release();
250+
},
251+
};
252+
}
253+
}
254+
255+
registerSingleton(ICodeCompareModelService, CodeCompareModelService, InstantiationType.Delayed);

src/vs/workbench/contrib/chat/browser/codeBlockPart.ts

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ import { renderFormattedText } from 'vs/base/browser/formattedTextRenderer';
1010
import { Button } from 'vs/base/browser/ui/button/button';
1111
import { CancellationToken } from 'vs/base/common/cancellation';
1212
import { Codicon } from 'vs/base/common/codicons';
13-
import { Emitter } from 'vs/base/common/event';
14-
import { Disposable, DisposableStore, MutableDisposable } from 'vs/base/common/lifecycle';
13+
import { Emitter, Event } from 'vs/base/common/event';
14+
import { combinedDisposable, Disposable, DisposableStore, MutableDisposable } from 'vs/base/common/lifecycle';
1515
import { Schemas } from 'vs/base/common/network';
1616
import { isEqual } from 'vs/base/common/resources';
1717
import { URI, UriComponents } from 'vs/base/common/uri';
@@ -24,7 +24,7 @@ import { CodeEditorWidget, ICodeEditorWidgetOptions } from 'vs/editor/browser/wi
2424
import { DiffEditorWidget } from 'vs/editor/browser/widget/diffEditor/diffEditorWidget';
2525
import { EDITOR_FONT_DEFAULTS, EditorOption, IEditorOptions } from 'vs/editor/common/config/editorOptions';
2626
import { IRange, Range } from 'vs/editor/common/core/range';
27-
import { IDiffEditorViewModel, ScrollType } from 'vs/editor/common/editorCommon';
27+
import { ScrollType } from 'vs/editor/common/editorCommon';
2828
import { TextEdit } from 'vs/editor/common/languages';
2929
import { EndOfLinePreference, ITextModel } from 'vs/editor/common/model';
3030
import { TextModelText } from 'vs/editor/common/model/textModelText';
@@ -67,6 +67,7 @@ import { MenuPreventer } from 'vs/workbench/contrib/codeEditor/browser/menuPreve
6767
import { SelectionClipboardContributionID } from 'vs/workbench/contrib/codeEditor/browser/selectionClipboard';
6868
import { getSimpleEditorOptions } from 'vs/workbench/contrib/codeEditor/browser/simpleEditorOptions';
6969
import { IMarkdownVulnerability } from '../common/annotations';
70+
import { assertType } from 'vs/base/common/types';
7071

7172
const $ = dom.$;
7273

@@ -483,6 +484,7 @@ export interface ICodeCompareBlockData {
483484
}
484485

485486

487+
// long-lived object that sits in the DiffPool and that gets reused
486488
export class CodeCompareBlockPart extends Disposable {
487489
protected readonly _onDidChangeContentHeight = this._register(new Emitter<void>());
488490
public readonly onDidChangeContentHeight = this._onDidChangeContentHeight.event;
@@ -494,7 +496,7 @@ export class CodeCompareBlockPart extends Disposable {
494496
readonly element: HTMLElement;
495497
private readonly messageElement: HTMLElement;
496498

497-
private readonly _lastDiffEditorViewModel = this._store.add(new MutableDisposable<IDiffEditorViewModel>());
499+
private readonly _lastDiffEditorViewModel = this._store.add(new MutableDisposable());
498500
private currentScrollWidth = 0;
499501

500502
constructor(
@@ -743,7 +745,8 @@ export class CodeCompareBlockPart extends Disposable {
743745

744746
this.element.classList.toggle('no-diff', isEditApplied);
745747

746-
if (data.edit.state?.applied) {
748+
if (isEditApplied) {
749+
assertType(data.edit.state?.applied);
747750

748751
const uriLabel = this.labelService.getUriLabel(data.edit.uri, { relative: true, noPrefix: true });
749752

@@ -787,8 +790,13 @@ export class CodeCompareBlockPart extends Disposable {
787790
return;
788791
}
789792

793+
const listener = Event.any(diffData.original.onWillDispose, diffData.modified.onWillDispose)(() => {
794+
// this a bit weird and basically duplicates https://github.com/microsoft/vscode/blob/7cbcafcbcc88298cfdcd0238018fbbba8eb6853e/src/vs/editor/browser/widget/diffEditor/diffEditorWidget.ts#L328
795+
// which cannot call `setModel(null)` without first complaining
796+
this.diffEditor.setModel(null);
797+
});
790798
this.diffEditor.setModel(viewModel);
791-
this._lastDiffEditorViewModel.value = viewModel;
799+
this._lastDiffEditorViewModel.value = combinedDisposable(listener, viewModel);
792800

793801
} else {
794802
this.diffEditor.setModel(null);
@@ -801,10 +809,6 @@ export class CodeCompareBlockPart extends Disposable {
801809
diffEditor: this.diffEditor,
802810
} satisfies ICodeCompareBlockActionContext;
803811
}
804-
805-
clearModel() {
806-
this.diffEditor.setModel(null);
807-
}
808812
}
809813

810814
export class DefaultChatTextEditor {

0 commit comments

Comments
 (0)