Skip to content

Commit 4867ea3

Browse files
roblourensroblourens
authored
Clean up chat model management and add tests (#278896)
* Clean up async chatmodel disposal enable reusing the same model when it is re-acquired during disposal * Move ChatModelStore, add tests * More tests * Clean up tests * test * fix
1 parent a8e321d commit 4867ea3

File tree

8 files changed

+438
-130
lines changed

8 files changed

+438
-130
lines changed

‎src/vs/workbench/contrib/chat/common/chatModelStore.ts‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { CancellationToken } from '../../../../base/common/cancellation.js';
7+
import { IDisposable, IReference, ReferenceCollection } from '../../../../base/common/lifecycle.js';
8+
import { ObservableMap } from '../../../../base/common/observable.js';
9+
import { URI } from '../../../../base/common/uri.js';
10+
import { ILogService } from '../../../../platform/log/common/log.js';
11+
import { IChatEditingSession } from './chatEditingService.js';
12+
import { ChatModel, IExportableChatData, ISerializableChatData } from './chatModel.js';
13+
import { ChatAgentLocation } from './constants.js';
14+
15+
export interface IStartSessionProps {
16+
readonly initialData?: IExportableChatData | ISerializableChatData;
17+
readonly location: ChatAgentLocation;
18+
readonly token: CancellationToken;
19+
readonly sessionResource: URI;
20+
readonly sessionId?: string;
21+
readonly canUseTools: boolean;
22+
readonly transferEditingSession?: IChatEditingSession;
23+
}
24+
25+
export interface ChatModelStoreDelegate {
26+
createModel: (props: IStartSessionProps) => ChatModel;
27+
willDisposeModel: (model: ChatModel) => Promise<void>;
28+
}
29+
30+
export class ChatModelStore extends ReferenceCollection<ChatModel> implements IDisposable {
31+
private readonly _models = new ObservableMap<string, ChatModel>();
32+
private readonly _modelsToDispose = new Set<string>();
33+
private readonly _pendingDisposals = new Set<Promise<void>>();
34+
35+
constructor(
36+
private readonly delegate: ChatModelStoreDelegate,
37+
@ILogService private readonly logService: ILogService,
38+
) {
39+
super();
40+
}
41+
42+
public get observable() {
43+
return this._models.observable;
44+
}
45+
46+
public values(): Iterable<ChatModel> {
47+
return this._models.values();
48+
}
49+
50+
/**
51+
* Get a ChatModel directly without acquiring a reference.
52+
*/
53+
public get(uri: URI): ChatModel | undefined {
54+
return this._models.get(this.toKey(uri));
55+
}
56+
57+
public has(uri: URI): boolean {
58+
return this._models.has(this.toKey(uri));
59+
}
60+
61+
public acquireExisting(uri: URI): IReference<ChatModel> | undefined {
62+
const key = this.toKey(uri);
63+
if (!this._models.has(key)) {
64+
return undefined;
65+
}
66+
return this.acquire(key);
67+
}
68+
69+
public acquireOrCreate(props: IStartSessionProps): IReference<ChatModel> {
70+
return this.acquire(this.toKey(props.sessionResource), props);
71+
}
72+
73+
protected createReferencedObject(key: string, props?: IStartSessionProps): ChatModel {
74+
this._modelsToDispose.delete(key);
75+
const existingModel = this._models.get(key);
76+
if (existingModel) {
77+
return existingModel;
78+
}
79+
80+
if (!props) {
81+
throw new Error(`No start session props provided for chat session ${key}`);
82+
}
83+
84+
this.logService.trace(`Creating chat session ${key}`);
85+
const model = this.delegate.createModel(props);
86+
if (model.sessionResource.toString() !== key) {
87+
throw new Error(`Chat session key mismatch for ${key}`);
88+
}
89+
this._models.set(key, model);
90+
return model;
91+
}
92+
93+
protected destroyReferencedObject(key: string, object: ChatModel): void {
94+
this._modelsToDispose.add(key);
95+
const promise = this.doDestroyReferencedObject(key, object);
96+
this._pendingDisposals.add(promise);
97+
promise.finally(() => {
98+
this._pendingDisposals.delete(promise);
99+
});
100+
}
101+
102+
private async doDestroyReferencedObject(key: string, object: ChatModel): Promise<void> {
103+
try {
104+
await this.delegate.willDisposeModel(object);
105+
} catch (error) {
106+
this.logService.error(error);
107+
} finally {
108+
if (this._modelsToDispose.has(key)) {
109+
this.logService.trace(`Disposing chat session ${key}`);
110+
this._models.delete(key);
111+
object.dispose();
112+
}
113+
this._modelsToDispose.delete(key);
114+
}
115+
}
116+
117+
/**
118+
* For test use only
119+
*/
120+
async waitForModelDisposals(): Promise<void> {
121+
await Promise.all(this._pendingDisposals);
122+
}
123+
124+
private toKey(uri: URI): string {
125+
return uri.toString();
126+
}
127+
128+
dispose(): void {
129+
this._models.forEach(model => model.dispose());
130+
}
131+
}

‎src/vs/workbench/contrib/chat/common/chatService.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,6 +1001,11 @@ export interface IChatService {
10011001
readonly edits2Enabled: boolean;
10021002

10031003
readonly requestInProgressObs: IObservable<boolean>;
1004+
1005+
/**
1006+
* For tests only!
1007+
*/
1008+
waitForModelDisposals(): Promise<void>;
10041009
}
10051010

10061011
export interface IChatSessionContext {

‎src/vs/workbench/contrib/chat/common/chatServiceImpl.ts‎

Lines changed: 19 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ import { BugIndicatingError, ErrorNoTelemetry } from '../../../../base/common/er
1010
import { Emitter, Event } from '../../../../base/common/event.js';
1111
import { MarkdownString } from '../../../../base/common/htmlContent.js';
1212
import { Iterable } from '../../../../base/common/iterator.js';
13-
import { Disposable, DisposableMap, DisposableStore, IDisposable, IReference, MutableDisposable, ReferenceCollection } from '../../../../base/common/lifecycle.js';
13+
import { Disposable, DisposableMap, DisposableStore, IDisposable, MutableDisposable } from '../../../../base/common/lifecycle.js';
1414
import { revive } from '../../../../base/common/marshalling.js';
1515
import { Schemas } from '../../../../base/common/network.js';
16-
import { autorun, derived, IObservable, ObservableMap } from '../../../../base/common/observable.js';
16+
import { autorun, derived, IObservable } from '../../../../base/common/observable.js';
1717
import { StopWatch } from '../../../../base/common/stopwatch.js';
1818
import { isDefined } from '../../../../base/common/types.js';
1919
import { URI } from '../../../../base/common/uri.js';
@@ -29,8 +29,8 @@ import { IWorkspaceContextService } from '../../../../platform/workspace/common/
2929
import { IExtensionService } from '../../../services/extensions/common/extensions.js';
3030
import { IMcpService } from '../../mcp/common/mcpTypes.js';
3131
import { IChatAgentCommand, IChatAgentData, IChatAgentHistoryEntry, IChatAgentRequest, IChatAgentResult, IChatAgentService } from './chatAgents.js';
32-
import { IChatEditingSession } from './chatEditingService.js';
3332
import { ChatModel, ChatRequestModel, ChatRequestRemovalReason, IChatModel, IChatRequestModel, IChatRequestVariableData, IChatResponseModel, IExportableChatData, ISerializableChatData, ISerializableChatDataIn, ISerializableChatsData, normalizeSerializableChatData, toChatHistoryContent, updateRanges } from './chatModel.js';
33+
import { ChatModelStore, IStartSessionProps } from './chatModelStore.js';
3434
import { chatAgentLeader, ChatRequestAgentPart, ChatRequestAgentSubcommandPart, ChatRequestSlashCommandPart, ChatRequestTextPart, chatSubcommandLeader, getPromptText, IParsedChatRequest } from './chatParserTypes.js';
3535
import { ChatRequestParser } from './chatRequestParser.js';
3636
import { ChatMcpServersStarting, IChatCompleteResponse, IChatDetail, IChatFollowup, IChatModelReference, IChatProgress, IChatSendRequestData, IChatSendRequestOptions, IChatSendRequestResponseState, IChatService, IChatSessionContext, IChatTransferredSessionData, IChatUserActionEvent } from './chatService.js';
@@ -71,91 +71,7 @@ class CancellableRequest implements IDisposable {
7171
}
7272
}
7373

74-
interface IStartSessionProps {
75-
readonly initialData?: IExportableChatData | ISerializableChatData;
76-
readonly location: ChatAgentLocation;
77-
readonly token: CancellationToken;
78-
readonly sessionResource: URI;
79-
readonly sessionId?: string;
80-
readonly canUseTools: boolean;
81-
readonly transferEditingSession?: IChatEditingSession;
82-
}
83-
84-
interface ChatModelStoreDelegate {
85-
createModel: (props: IStartSessionProps) => ChatModel;
86-
willDisposeModel: (model: ChatModel) => Promise<void>;
87-
}
88-
89-
class ChatModelStore extends ReferenceCollection<ChatModel> implements IDisposable {
90-
private readonly _models = new ObservableMap<string, ChatModel>();
91-
92-
constructor(
93-
private readonly delegate: ChatModelStoreDelegate,
94-
@ILogService private readonly logService: ILogService,
95-
) {
96-
super();
97-
}
98-
99-
public get observable() {
100-
return this._models.observable;
101-
}
102-
103-
public values(): Iterable<ChatModel> {
104-
return this._models.values();
105-
}
106-
107-
public get(uri: URI): ChatModel | undefined {
108-
return this._models.get(this.toKey(uri));
109-
}
110-
111-
public has(uri: URI): boolean {
112-
return this._models.has(this.toKey(uri));
113-
}
114-
115-
public acquireExisting(uri: URI): IReference<ChatModel> | undefined {
116-
const key = this.toKey(uri);
117-
if (!this._models.has(key)) {
118-
return undefined;
119-
}
120-
return this.acquire(key);
121-
}
122-
123-
public acquireOrCreate(props: IStartSessionProps): IReference<ChatModel> {
124-
return this.acquire(this.toKey(props.sessionResource), props);
125-
}
126-
127-
protected createReferencedObject(key: string, props?: IStartSessionProps): ChatModel {
128-
if (!props) {
129-
throw new Error(`No start session props provided for chat session ${key}`);
130-
}
131-
132-
this.logService.trace(`Creating chat session ${key}`);
133-
const model = this.delegate.createModel(props);
134-
if (model.sessionResource.toString() !== key) {
135-
throw new Error(`Chat session key mismatch for ${key}`);
136-
}
137-
this._models.set(key, model);
138-
return model;
139-
}
140-
141-
protected async destroyReferencedObject(key: string, object: ChatModel): Promise<void> {
142-
try {
143-
await this.delegate.willDisposeModel(object);
144-
} finally {
145-
this.logService.trace(`Disposing chat session ${key}`);
146-
this._models.delete(key);
147-
object.dispose();
148-
}
149-
}
15074

151-
private toKey(uri: URI): string {
152-
return uri.toString();
153-
}
154-
155-
dispose(): void {
156-
this._models.forEach(model => model.dispose());
157-
}
158-
}
15975

16076
class DisposableResourceMap<T extends IDisposable> extends Disposable {
16177

@@ -214,6 +130,13 @@ export class ChatService extends Disposable implements IChatService {
214130

215131
readonly requestInProgressObs: IObservable<boolean>;
216132

133+
/**
134+
* For test use only
135+
*/
136+
waitForModelDisposals(): Promise<void> {
137+
return this._sessionModels.waitForModelDisposals();
138+
}
139+
217140
public get edits2Enabled(): boolean {
218141
return this.configurationService.getValue(ChatConfiguration.Edits2Enabled);
219142
}
@@ -239,17 +162,15 @@ export class ChatService extends Disposable implements IChatService {
239162
super();
240163

241164
this._sessionModels = this._register(instantiationService.createInstance(ChatModelStore, {
242-
createModel: props => this._startSession(props),
243-
willDisposeModel: async model => {
244-
if (this._persistChats) {
245-
const localSessionId = LocalChatSessionUri.parseLocalSessionId(model.sessionResource);
246-
if (localSessionId && (model.initialLocation === ChatAgentLocation.Chat)) {
247-
// Always preserve sessions that have custom titles, even if empty
248-
if (model.getRequests().length === 0 && !model.customTitle) {
249-
this._chatSessionStore.deleteSession(localSessionId);
250-
} else {
251-
this._chatSessionStore.storeSessions([model]);
252-
}
165+
createModel: (props: IStartSessionProps) => this._startSession(props),
166+
willDisposeModel: async (model: ChatModel) => {
167+
const localSessionId = LocalChatSessionUri.parseLocalSessionId(model.sessionResource);
168+
if (localSessionId && (model.initialLocation === ChatAgentLocation.Chat)) {
169+
// Always preserve sessions that have custom titles, even if empty
170+
if (model.getRequests().length === 0 && !model.customTitle) {
171+
await this._chatSessionStore.deleteSession(localSessionId);
172+
} else {
173+
await this._chatSessionStore.storeSessions([model]);
253174
}
254175
}
255176
}
@@ -295,14 +216,6 @@ export class ChatService extends Disposable implements IChatService {
295216
});
296217
}
297218

298-
private _persistChats = true;
299-
/**
300-
* For test only
301-
*/
302-
setChatPersistanceEnabled(enabled: boolean): void {
303-
this._persistChats = enabled;
304-
}
305-
306219
public get editingSessions() {
307220
return [...this._sessionModels.values()].map(v => v.editingSession).filter(isDefined);
308221
}
@@ -312,10 +225,6 @@ export class ChatService extends Disposable implements IChatService {
312225
}
313226

314227
private saveState(): void {
315-
if (!this._persistChats) {
316-
return;
317-
}
318-
319228
const liveChats = Array.from(this._sessionModels.values())
320229
.filter(session => {
321230
if (!LocalChatSessionUri.parseLocalSessionId(session.sessionResource)) {

‎src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,6 @@ suite('ChatEditingService', function () {
105105
editingService = value;
106106

107107
chatService = insta.get(IChatService);
108-
(chatService as ChatService).setChatPersistanceEnabled(false);
109108

110109
store.add(insta.get(IChatSessionsService) as ChatSessionsService); // Needs to be disposed in between test runs to clear extensionPoint contribution
111110
store.add(chatService as ChatService);
@@ -131,8 +130,9 @@ suite('ChatEditingService', function () {
131130
}));
132131
});
133132

134-
teardown(() => {
133+
teardown(async () => {
135134
store.clear();
135+
await chatService.waitForModelDisposals();
136136
});
137137

138138
ensureNoDisposablesAreLeakedInTestSuite();

0 commit comments

Comments
 (0)