Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Persist chat sessions in empty windows #224430

Merged
merged 1 commit into from
Jul 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/vs/workbench/contrib/chat/common/chatModel.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
96 changes: 81 additions & 15 deletions src/vs/workbench/contrib/chat/common/chatServiceImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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);
Expand Down
Loading