Store workspaceState with associated editSession

src/vs/workbench/contrib/editSessions/browser/editSessions.contribution.ts

@@ -743,7 +743,7 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 			return undefined;
 		}
 
-		const data: EditSession = { folders, version: 2 };
+		const data: EditSession = { folders, version: 2, workspaceStateId: this.editSessionsStorageService.lastWrittenResources.get('workspaceState')?.ref };
 
 		try {
 			this.logService.info(`Storing edit session...`);

src/vs/workbench/contrib/editSessions/browser/editSessionsStorageService.ts

@@ -58,6 +58,16 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 		return this._didSignOut.event;
 	}
 
+	private _lastWrittenResources = new Map<SyncResource, { ref: string; content: string }>();
+	get lastWrittenResources() {
+		return this._lastWrittenResources;
+	}
+
+	private _lastReadResources = new Map<SyncResource, { ref: string; content: string }>();
+	get lastReadResources() {
+		return this._lastReadResources;
+	}
+
 	storeClient: EditSessionsStoreClient | undefined; // TODO@joyceerhl lifecycle hack
 
 	constructor(
@@ -89,9 +99,9 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 	}
 
 	/**
-	 *
-	 * @param editSession An object representing edit session state to be restored.
-	 * @returns The ref of the stored edit session state.
+	 * @param resource: The resource to retrieve content for.
+	 * @param content An object representing resource state to be restored.
+	 * @returns The ref of the stored state.
 	 */
 	async write(resource: SyncResource, content: string | EditSession): Promise<string> {
 		await this.initialize(false);
@@ -103,14 +113,20 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 			content.machine = await this.getOrCreateCurrentMachineId();
 		}
 
-		return this.storeClient!.writeResource(resource, typeof content === 'string' ? content : JSON.stringify(content), null, undefined, createSyncHeaders(generateUuid()));
+		content = typeof content === 'string' ? content : JSON.stringify(content);
+		const ref = await this.storeClient!.writeResource(resource, content, null, undefined, createSyncHeaders(generateUuid()));
+
+		this._lastWrittenResources.set(resource, { ref, content });
+
+		return ref;
 	}
 
 	/**
+	 * @param resource: The resource to retrieve content for.
 	 * @param ref: A specific content ref to retrieve content for, if it exists.
 	 * If undefined, this method will return the latest saved edit session, if any.
 	 *
-	 * @returns An object representing the requested or latest edit session state, if any.
+	 * @returns An object representing the requested or latest state, if any.
 	 */
 	async read(resource: SyncResource, ref: string | undefined): Promise<{ ref: string; content: string } | undefined> {
 		await this.initialize(false);
@@ -133,7 +149,11 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 		}
 
 		// TODO@joyceerhl Validate session data, check schema version
-		return (content !== undefined && content !== null && ref !== undefined) ? { ref, content } : undefined;
+		if (content !== undefined && content !== null && ref !== undefined) {
+			this._lastReadResources.set(resource, { ref, content });
+			return { ref, content };
+		}
+		return undefined;
 	}
 
 	async delete(resource: SyncResource, ref: string | null) {

src/vs/workbench/contrib/editSessions/common/editSessions.ts

@@ -35,6 +35,9 @@ export interface IEditSessionsStorageService {
 
 	storeClient: EditSessionsStoreClient | undefined;
 
+	lastReadResources: Map<SyncResource, { ref: string; content: string }>;
+	lastWrittenResources: Map<SyncResource, { ref: string; content: string }>;
+
 	initialize(silent?: boolean): Promise<boolean>;
 	read(resource: SyncResource, ref: string | undefined): Promise<{ ref: string; content: string } | undefined>;
 	write(resource: SyncResource, content: string | EditSession): Promise<string>;
@@ -82,6 +85,7 @@ export const EditSessionSchemaVersion = 3;
 
 export interface EditSession {
 	version: number;
+	workspaceStateId?: string;
 	machine?: string;
 	folders: Folder[];
 }

src/vs/workbench/contrib/editSessions/common/workspaceStateSync.ts

@@ -17,7 +17,7 @@ import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity'
 import { IUserDataProfile } from 'vs/platform/userDataProfile/common/userDataProfile';
 import { AbstractSynchroniser, IAcceptResult, IMergeResult, IResourcePreview, ISyncResourcePreview } from 'vs/platform/userDataSync/common/abstractSynchronizer';
 import { IRemoteUserData, IResourceRefHandle, IUserDataSyncBackupStoreService, IUserDataSyncConfiguration, IUserDataSyncEnablementService, IUserDataSyncLogService, IUserDataSyncStoreService, IUserDataSynchroniser, IWorkspaceState, SyncResource } from 'vs/platform/userDataSync/common/userDataSync';
-import { IEditSessionsStorageService } from 'vs/workbench/contrib/editSessions/common/editSessions';
+import { EditSession, IEditSessionsStorageService } from 'vs/workbench/contrib/editSessions/common/editSessions';
 import { IWorkspaceIdentityService } from 'vs/workbench/services/workspaces/common/workspaceIdentityService';
 
 
@@ -99,11 +99,14 @@ export class WorkspaceStateSynchroniser extends AbstractSynchroniser implements
 		});
 
 		const content: IWorkspaceState = { folders, storage: contributedData, version: this.version };
-		this.editSessionsStorageService.write('workspaceState', stringify(content));
+		await this.editSessionsStorageService.write('workspaceState', stringify(content));
 	}
 
 	override async apply(): Promise<ISyncResourcePreview | null> {
-		const resource = await this.editSessionsStorageService.read('workspaceState', undefined);
+		const payload = this.editSessionsStorageService.lastReadResources.get('editSessions')?.content;
+		const workspaceStateId = payload ? (JSON.parse(payload) as EditSession).workspaceStateId : undefined;
+
+		const resource = await this.editSessionsStorageService.read('workspaceState', workspaceStateId);
 		if (!resource) {
 			return null;
 		}
@@ -143,6 +146,8 @@ export class WorkspaceStateSynchroniser extends AbstractSynchroniser implements
 			}
 			this.storageService.storeAll(storageEntries, true);
 		}
+
+		this.editSessionsStorageService.delete('workspaceState', resource.ref);
 		return null;
 	}