Persist chat sessions in empty windows

fix microsoft/vscode-copilot-release#1309

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 {

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<string, CancellationTokenSource>());
 	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<string>();
 
 	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<string, ISerializableChatData | ChatModel> = { ...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);