Add methods to manage additional copilot cli session metadata (#4300)

* Add methods to manage additional copilot cli session metadata

* Updates

Author: DonJayamanne

src/extension/chatSessions/common/chatSessionMetadataStore.ts

@@ -7,6 +7,7 @@ import type * as vscode from 'vscode';
 import type { Uri } from 'vscode';
 import { createServiceIdentifier } from '../../../util/common/services';
 import { ChatSessionWorktreeProperties } from './chatSessionWorktreeService';
+import type { IWorkspaceInfo } from './workspaceInfo';
 
 export interface WorkspaceFolderEntry {
 	readonly folderPath: string;
@@ -16,6 +17,10 @@ export interface WorkspaceFolderEntry {
 export interface ChatSessionMetadataFile {
 	worktreeProperties?: ChatSessionWorktreeProperties;
 	workspaceFolder?: WorkspaceFolderEntry;
+	additionalWorkspaces?: {
+		worktreeProperties?: ChatSessionWorktreeProperties;
+		workspaceFolder?: WorkspaceFolderEntry;
+	}[];
 	/**
 	 * Whether the session metadata has been written to the Copilot CLI session state directory.
 	 */
@@ -34,4 +39,6 @@ export interface IChatSessionMetadataStore {
 	getWorktreeProperties(folder: Uri): Promise<ChatSessionWorktreeProperties | undefined>;
 	getSessionWorkspaceFolder(sessionId: string): Promise<vscode.Uri | undefined>;
 	getUsedWorkspaceFolders(): Promise<WorkspaceFolderEntry[]>;
+	getAdditionalWorkspaces(sessionId: string): Promise<IWorkspaceInfo[]>;
+	setAdditionalWorkspaces(sessionId: string, workspaces: IWorkspaceInfo[]): Promise<void>;
 }

src/extension/chatSessions/vscode-node/chatSessionMetadataStoreImpl.ts

@@ -15,6 +15,7 @@ import { ResourceMap } from '../../../util/vs/base/common/map';
 import { dirname, isEqual } from '../../../util/vs/base/common/resources';
 import { ChatSessionMetadataFile, IChatSessionMetadataStore, WorkspaceFolderEntry } from '../common/chatSessionMetadataStore';
 import { ChatSessionWorktreeData, ChatSessionWorktreeProperties } from '../common/chatSessionWorktreeService';
+import { IWorkspaceInfo } from '../common/workspaceInfo';
 import { getCopilotCLISessionDir } from '../copilotcli/node/cliHelpers';
 
 const WORKSPACE_FOLDER_MEMENTO_KEY = 'github.copilot.cli.sessionWorkspaceFolders';
@@ -53,7 +54,7 @@ export class ChatSessionMetadataStore extends Disposable implements IChatSession
 					continue;
 				}
 				if (!metadata.writtenToDisc) {
-					if ((metadata.workspaceFolder || metadata.worktreeProperties)) {
+					if ((metadata.workspaceFolder || metadata.worktreeProperties || metadata.additionalWorkspaces?.length)) {
 						this.updateSessionMetadata(sessionId, metadata, false).catch(ex => {
 							this.logService.error(ex, `[ChatSessionMetadataStore] Failed to write metadata for session ${sessionId} to session state: `);
 						});
@@ -116,7 +117,7 @@ export class ChatSessionMetadataStore extends Disposable implements IChatSession
 			// These promises can run in background and no need to wait for them.
 			// Even if user exits early we have all the data in the global storage and we'll restore from that next time.
 			if (!metadata.writtenToDisc) {
-				if ((metadata.workspaceFolder || metadata.worktreeProperties)) {
+				if ((metadata.workspaceFolder || metadata.worktreeProperties || metadata.additionalWorkspaces?.length)) {
 					this.updateSessionMetadata(sessionId, metadata, false).catch(ex => {
 						this.logService.error(ex, `[ChatSessionMetadataStore] Failed to write metadata for session ${sessionId} to session state: `);
 					});
@@ -216,6 +217,33 @@ export class ChatSessionMetadataStore extends Disposable implements IChatSession
 		}
 		return Array.from(entries.entries()).map(([folderUri, timestamp]) => ({ folderPath: folderUri.fsPath, timestamp }));
 	}
+
+	async getAdditionalWorkspaces(sessionId: string): Promise<IWorkspaceInfo[]> {
+		const metadata = await this.getSessionMetadata(sessionId);
+		if (!metadata?.additionalWorkspaces?.length) {
+			return [];
+		}
+		return metadata.additionalWorkspaces.map(ws => ({
+			folder: !ws.worktreeProperties && ws.workspaceFolder?.folderPath ? Uri.file(ws.workspaceFolder.folderPath) : undefined,
+			repository: ws.worktreeProperties?.repositoryPath ? Uri.file(ws.worktreeProperties.repositoryPath) : undefined,
+			worktree: ws.worktreeProperties?.worktreePath ? Uri.file(ws.worktreeProperties.worktreePath) : undefined,
+			worktreeProperties: ws.worktreeProperties,
+		}));
+	}
+
+	async setAdditionalWorkspaces(sessionId: string, workspaces: IWorkspaceInfo[]): Promise<void> {
+		await this._intialize.value;
+		const existing = this._cache[sessionId] ?? {};
+		const additionalWorkspaces = workspaces.map(ws => ({
+			worktreeProperties: ws.worktreeProperties,
+			workspaceFolder: !ws.worktreeProperties && ws.folder ? { folderPath: ws.folder.fsPath, timestamp: Date.now() } : undefined,
+		}));
+		const metadata: ChatSessionMetadataFile = { ...existing, additionalWorkspaces };
+		this._cache[sessionId] = metadata;
+		await this.updateSessionMetadata(sessionId, metadata);
+		this.updateGlobalStorage();
+	}
+
 	private async getSessionMetadata(sessionId: string): Promise<ChatSessionMetadataFile | undefined> {
 		await this._intialize.value;
 		if (sessionId in this._cache) {

src/extension/chatSessions/vscode-node/test/chatSessionMetadataStoreImpl.spec.ts

@@ -14,6 +14,7 @@ import { Emitter } from '../../../../util/vs/base/common/event';
 import { URI } from '../../../../util/vs/base/common/uri';
 import { eventToPromise } from '../../../completions-core/vscode-node/lib/src/prompt/asyncUtils';
 import { ChatSessionWorktreeData, ChatSessionWorktreeProperties } from '../../common/chatSessionWorktreeService';
+import { IWorkspaceInfo } from '../../common/workspaceInfo';
 import { getCopilotCLISessionDir } from '../../copilotcli/node/cliHelpers';
 import { ChatSessionMetadataStore } from '../chatSessionMetadataStoreImpl';
 
@@ -339,7 +340,7 @@ describe('ChatSessionMetadataStore', () => {
 			store.dispose();
 		});
 
-		it('should not retry entries with no workspaceFolder or worktreeProperties', async () => {
+		it('should not retry entries with no workspaceFolder, worktreeProperties, or additionalWorkspaces', async () => {
 			const existingData = {
 				'session-empty': {},
 				'session-folder': { workspaceFolder: { folderPath: Uri.file('/workspace/a').fsPath, timestamp: 100 } },
@@ -358,6 +359,32 @@ describe('ChatSessionMetadataStore', () => {
 			expect(sessionEmptyStatCalls).toHaveLength(0);
 			store.dispose();
 		});
+
+		it('should retry entries that only have additionalWorkspaces (not delete as invalid data)', async () => {
+			// A session with only additionalWorkspaces and writtenToDisc: false
+			// must be retried, not deleted from cache — otherwise data is lost after a crash.
+			const existingData = {
+				'session-only-additional': {
+					additionalWorkspaces: [
+						{ workspaceFolder: { folderPath: Uri.file('/extra/workspace').fsPath, timestamp: 100 } },
+					],
+				},
+			};
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify(existingData));
+
+			// Pre-create the session directory so the recovery write can succeed
+			await mockFs.createDirectory(sessionDirectoryUri('session-only-additional'));
+			const fileCreated = eventToPromise(mockFs.onDidCreateFile.event);
+
+			const store = await createStore();
+			await fileCreated;
+
+			const fileUri = sessionMetadataFileUri('session-only-additional');
+			const rawContent = await mockFs.readFile(fileUri);
+			const written = JSON.parse(new TextDecoder().decode(rawContent));
+			expect(written.additionalWorkspaces).toHaveLength(1);
+			store.dispose();
+		});
 	});
 
 	// ──────────────────────────────────────────────────────────────────────────
@@ -1409,6 +1436,186 @@ describe('ChatSessionMetadataStore', () => {
 		});
 	});
 
+	// ──────────────────────────────────────────────────────────────────────────
+	// setAdditionalWorkspaces / getAdditionalWorkspaces
+	// ──────────────────────────────────────────────────────────────────────────
+	describe('setAdditionalWorkspaces / getAdditionalWorkspaces', () => {
+		it('should store and retrieve workspace-folder type additional workspaces', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+
+			const workspaces: IWorkspaceInfo[] = [
+				{ folder: Uri.file('/extra/a'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+				{ folder: Uri.file('/extra/b'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			];
+			await store.setAdditionalWorkspaces('session-1', workspaces);
+
+			const result = await store.getAdditionalWorkspaces('session-1');
+			expect(result).toHaveLength(2);
+			expect(result[0].folder?.fsPath).toBe(Uri.file('/extra/a').fsPath);
+			expect(result[1].folder?.fsPath).toBe(Uri.file('/extra/b').fsPath);
+			store.dispose();
+		});
+
+		it('should store and retrieve worktree type additional workspaces', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+			const props = makeWorktreeV1Props();
+
+			const workspaces: IWorkspaceInfo[] = [
+				{ folder: undefined, repository: Uri.file('/repo'), worktree: Uri.file('/repo/.worktrees/wt'), worktreeProperties: props },
+			];
+			await store.setAdditionalWorkspaces('session-wt', workspaces);
+
+			const result = await store.getAdditionalWorkspaces('session-wt');
+			expect(result).toHaveLength(1);
+			expect(result[0].worktreeProperties?.branchName).toBe(props.branchName);
+			expect(result[0].worktree?.fsPath).toBe(Uri.file('/repo/.worktrees/wt').fsPath);
+			// worktreeProperties present → folder should be undefined per getAdditionalWorkspaces logic
+			expect(result[0].folder).toBeUndefined();
+			store.dispose();
+		});
+
+		it('should return empty array when no additional workspaces are set', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({
+				'session-1': { workspaceFolder: { folderPath: Uri.file('/a').fsPath, timestamp: 1 } },
+			}));
+			const store = await createStore();
+
+			const result = await store.getAdditionalWorkspaces('session-1');
+			expect(result).toEqual([]);
+			store.dispose();
+		});
+
+		it('should return empty array for unknown session', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+
+			const result = await store.getAdditionalWorkspaces('nonexistent');
+			expect(result).toEqual([]);
+			store.dispose();
+		});
+
+		it('should write additionalWorkspaces to per-session file', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+
+			await store.setAdditionalWorkspaces('session-1', [
+				{ folder: Uri.file('/extra/a'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			]);
+
+			const fileUri = sessionMetadataFileUri('session-1');
+			const rawContent = await mockFs.readFile(fileUri);
+			const written = JSON.parse(new TextDecoder().decode(rawContent));
+			expect(written.additionalWorkspaces).toHaveLength(1);
+			expect(written.additionalWorkspaces[0].workspaceFolder?.folderPath).toBe(Uri.file('/extra/a').fsPath);
+			store.dispose();
+		});
+
+		it('should preserve existing workspaceFolder when setting additionalWorkspaces', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({
+				'session-1': { workspaceFolder: { folderPath: Uri.file('/primary').fsPath, timestamp: 100 } },
+			}));
+			const store = await createStore();
+
+			await store.setAdditionalWorkspaces('session-1', [
+				{ folder: Uri.file('/extra/a'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			]);
+
+			// Primary workspace folder should still be accessible
+			const folder = await store.getSessionWorkspaceFolder('session-1');
+			expect(folder?.fsPath).toBe(Uri.file('/primary').fsPath);
+
+			// Additional workspaces should also be present
+			const result = await store.getAdditionalWorkspaces('session-1');
+			expect(result).toHaveLength(1);
+			expect(result[0].folder?.fsPath).toBe(Uri.file('/extra/a').fsPath);
+			store.dispose();
+		});
+
+		it('should replace previous additionalWorkspaces on subsequent call', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+
+			await store.setAdditionalWorkspaces('session-1', [
+				{ folder: Uri.file('/old'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			]);
+			await store.setAdditionalWorkspaces('session-1', [
+				{ folder: Uri.file('/new/a'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+				{ folder: Uri.file('/new/b'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			]);
+
+			const result = await store.getAdditionalWorkspaces('session-1');
+			expect(result).toHaveLength(2);
+			expect(result[0].folder?.fsPath).toBe(Uri.file('/new/a').fsPath);
+			store.dispose();
+		});
+
+		it('should trigger debounced bulk storage update', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({}));
+			const store = await createStore();
+
+			await store.setAdditionalWorkspaces('session-1', [
+				{ folder: Uri.file('/extra'), repository: undefined, worktree: undefined, worktreeProperties: undefined },
+			]);
+			await vi.advanceTimersByTimeAsync(1_100);
+
+			const rawContent = await mockFs.readFile(BULK_METADATA_FILE);
+			const written = JSON.parse(new TextDecoder().decode(rawContent));
+			expect(written['session-1']?.additionalWorkspaces).toBeDefined();
+			store.dispose();
+		});
+
+		it('should restore additionalWorkspaces from bulk file on startup', async () => {
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({
+				'session-1': {
+					additionalWorkspaces: [
+						{ workspaceFolder: { folderPath: Uri.file('/restored/a').fsPath, timestamp: 100 } },
+					],
+					writtenToDisc: true,
+				},
+			}));
+			const store = await createStore();
+
+			const result = await store.getAdditionalWorkspaces('session-1');
+			expect(result).toHaveLength(1);
+			expect(result[0].folder?.fsPath).toBe(Uri.file('/restored/a').fsPath);
+			store.dispose();
+		});
+
+		it('should survive crash recovery: entry with only additionalWorkspaces is re-persisted not deleted', async () => {
+			// Simulate VS Code crash: bulk file has the entry but writtenToDisc is falsy
+			// (updateSessionMetadata never completed before the crash).
+			mockFs.mockFile(BULK_METADATA_FILE, JSON.stringify({
+				'session-crash': {
+					additionalWorkspaces: [
+						{ workspaceFolder: { folderPath: Uri.file('/extra/workspace').fsPath, timestamp: 100 } },
+					],
+					// writtenToDisc intentionally absent (falsy) — simulates crash before write completed
+				},
+			}));
+
+			// Pre-create session directory so recovery write can succeed
+			await mockFs.createDirectory(sessionDirectoryUri('session-crash'));
+			const fileCreated = eventToPromise(mockFs.onDidCreateFile.event);
+
+			const store = await createStore();
+			await fileCreated;
+
+			// Entry should have been re-persisted to per-session file
+			const fileUri = sessionMetadataFileUri('session-crash');
+			const rawContent = await mockFs.readFile(fileUri);
+			const written = JSON.parse(new TextDecoder().decode(rawContent));
+			expect(written.additionalWorkspaces).toHaveLength(1);
+			expect(written.additionalWorkspaces[0].workspaceFolder?.folderPath).toBe(Uri.file('/extra/workspace').fsPath);
+
+			// And still readable via the API
+			const result = await store.getAdditionalWorkspaces('session-crash');
+			expect(result).toHaveLength(1);
+			store.dispose();
+		});
+	});
+
 	// ──────────────────────────────────────────────────────────────────────────
 	// Constructor & edge cases
 	// ──────────────────────────────────────────────────────────────────────────