From 6c37657f06817df8439df80c463566026a57086e Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Wed, 31 Jul 2024 13:11:13 -0700 Subject: [PATCH] Persist chat sessions in empty windows (#224430) fix microsoft/vscode-copilot-release#1309 --- .../contrib/chat/common/chatModel.ts | 3 + .../contrib/chat/common/chatServiceImpl.ts | 96 ++++++++++++++++--- 2 files changed, 84 insertions(+), 15 deletions(-) diff --git a/src/vs/workbench/contrib/chat/common/chatModel.ts b/src/vs/workbench/contrib/chat/common/chatModel.ts index 4eb2b9f3816b5..fdd62d4be2842 100644 --- a/src/vs/workbench/contrib/chat/common/chatModel.ts +++ b/src/vs/workbench/contrib/chat/common/chatModel.ts @@ -554,6 +554,9 @@ export interface ISerializableChatData extends IExportableChatData { sessionId: string; creationDate: number; isImported: boolean; + + /** Indicates that this session was created in this window. Is cleared after the chat has been written to storage once. Needed to sync chat creations/deletions between empty windows. */ + isNew?: boolean; } export function isExportableSessionData(obj: unknown): obj is IExportableChatData { diff --git a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts index 81ac4804612ca..8355963b52a2d 100644 --- a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts @@ -80,6 +80,8 @@ export class ChatService extends Disposable implements IChatService { private readonly _pendingRequests = this._register(new DisposableMap()); private _persistedSessions: ISerializableChatsData; + /** Just for empty windows, need to enforce that a chat was deleted, even though other windows still have it */ + private _deletedChatIds = new Set(); private _transferredSessionData: IChatTransferredSessionData | undefined; public get transferredSessionData(): IChatTransferredSessionData | undefined { @@ -109,7 +111,8 @@ export class ChatService extends Disposable implements IChatService { super(); this._chatServiceTelemetry = this.instantiationService.createInstance(ChatServiceTelemetry); - const sessionData = storageService.get(serializedChatKey, StorageScope.WORKSPACE, ''); + const isEmptyWindow = !workspaceContextService.getWorkspace().folders.length; + const sessionData = storageService.get(serializedChatKey, isEmptyWindow ? StorageScope.APPLICATION : StorageScope.WORKSPACE, ''); if (sessionData) { this._persistedSessions = this.deserializeChats(sessionData); const countsForLog = Object.keys(this._persistedSessions).length; @@ -136,26 +139,85 @@ export class ChatService extends Disposable implements IChatService { } private saveState(): void { - let allSessions: (ChatModel | ISerializableChatData)[] = Array.from(this._sessionModels.values()) + const liveChats = Array.from(this._sessionModels.values()) .filter(session => session.initialLocation === ChatAgentLocation.Panel) .filter(session => session.getRequests().length > 0); - allSessions = allSessions.concat( - Object.values(this._persistedSessions) - .filter(session => !this._sessionModels.has(session.sessionId)) - .filter(session => session.requests.length)); - allSessions.sort((a, b) => (b.creationDate ?? 0) - (a.creationDate ?? 0)); - allSessions = allSessions.slice(0, maxPersistedSessions); - if (allSessions.length) { - this.trace('onWillSaveState', `Persisting ${allSessions.length} sessions`); + + const isEmptyWindow = !this.workspaceContextService.getWorkspace().folders.length; + if (isEmptyWindow) { + this.syncEmptyWindowChats(liveChats); + } else { + let allSessions: (ChatModel | ISerializableChatData)[] = liveChats; + allSessions = allSessions.concat( + Object.values(this._persistedSessions) + .filter(session => !this._sessionModels.has(session.sessionId)) + .filter(session => session.requests.length)); + allSessions.sort((a, b) => (b.creationDate ?? 0) - (a.creationDate ?? 0)); + allSessions = allSessions.slice(0, maxPersistedSessions); + if (allSessions.length) { + this.trace('onWillSaveState', `Persisting ${allSessions.length} sessions`); + } + + const serialized = JSON.stringify(allSessions); + + if (allSessions.length) { + this.trace('onWillSaveState', `Persisting ${serialized.length} chars`); + } + + this.storageService.store(serializedChatKey, serialized, StorageScope.WORKSPACE, StorageTarget.MACHINE); + } + + this._deletedChatIds.clear(); + } + + private syncEmptyWindowChats(thisWindowChats: ChatModel[]): void { + // Note- an unavoidable race condition exists here. If there are multiple empty windows open, and the user quits the application, then the focused + // window may lose active chats, because all windows are reading and writing to storageService at the same time. This can't be fixed without some + // kind of locking, but in reality, the focused window will likely have run `saveState` at some point, like on a window focus change, and it will + // generally be fine. + const sessionData = this.storageService.get(serializedChatKey, StorageScope.APPLICATION, ''); + + const originalPersistedSessions = this._persistedSessions; + let persistedSessions: ISerializableChatsData; + if (sessionData) { + persistedSessions = this.deserializeChats(sessionData); + const countsForLog = Object.keys(persistedSessions).length; + if (countsForLog > 0) { + this.trace('constructor', `Restored ${countsForLog} persisted sessions`); + } + } else { + persistedSessions = {}; } - const serialized = JSON.stringify(allSessions); + this._deletedChatIds.forEach(id => delete persistedSessions[id]); + + // Has the chat in this window been updated, and then closed? Overwrite the old persisted chats. + Object.values(originalPersistedSessions).forEach(session => { + const persistedSession = persistedSessions[session.sessionId]; + if (persistedSession && session.requests.length > persistedSession.requests.length) { + // We will add a 'modified date' at some point, but comparing the number of requests is good enough + persistedSessions[session.sessionId] = session; + } else if (!persistedSession && session.isNew) { + // This session was created in this window, and hasn't been persisted yet + session.isNew = false; + persistedSessions[session.sessionId] = session; + } + }); + + this._persistedSessions = persistedSessions; - if (allSessions.length) { - this.trace('onWillSaveState', `Persisting ${serialized.length} chars`); + // Add this window's active chat models to the set to persist. + // Having the same session open in two empty windows at the same time can lead to data loss, this is acceptable + const allSessions: Record = { ...this._persistedSessions }; + for (const chat of thisWindowChats) { + allSessions[chat.sessionId] = chat; } - this.storageService.store(serializedChatKey, serialized, StorageScope.WORKSPACE, StorageTarget.MACHINE); + let sessionsList = Object.values(allSessions); + sessionsList.sort((a, b) => (b.creationDate ?? 0) - (a.creationDate ?? 0)); + sessionsList = sessionsList.slice(0, maxPersistedSessions); + const data = JSON.stringify(sessionsList); + this.storageService.store(serializedChatKey, data, StorageScope.APPLICATION, StorageTarget.MACHINE); } notifyUserAction(action: IChatUserActionEvent): void { @@ -248,11 +310,13 @@ export class ChatService extends Disposable implements IChatService { } removeHistoryEntry(sessionId: string): void { + this._deletedChatIds.add(sessionId); delete this._persistedSessions[sessionId]; this.saveState(); } clearAllHistoryEntries(): void { + Object.values(this._persistedSessions).forEach(session => this._deletedChatIds.add(session.sessionId)); this._persistedSessions = {}; this.saveState(); } @@ -682,7 +746,9 @@ export class ChatService extends Disposable implements IChatService { if (model.initialLocation === ChatAgentLocation.Panel) { // Turn all the real objects into actual JSON, otherwise, calling 'revive' may fail when it tries to // assign values to properties that are getters- microsoft/vscode-copilot-release#1233 - this._persistedSessions[sessionId] = JSON.parse(JSON.stringify(model)); + const sessionData: ISerializableChatData = JSON.parse(JSON.stringify(model)); + sessionData.isNew = true; + this._persistedSessions[sessionId] = sessionData; } this._sessionModels.deleteAndDispose(sessionId);