Preserve customization roots during provisional sessions

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts

@@ -484,14 +484,17 @@ class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizat
 			return undefined;
 		}
 		const sessionState = this._readSessionState(sessionResource);
+		const workingDirectories = sessionState === undefined
+			? this._provisionalSessionService.getProvisionalWorkingDirectories(sessionResource)?.map(uri => uri.toString())
+			: sessionState.workingDirectories;
 		const rootState = target.connection.rootState.value;
 		const channel = target.backendSession.toString();
 		return {
 			customizations: sessionState?.customizations ?? [],
 			resourceUris: target.connection.resourceUris,
 			folderPickerDecision: readSessionFolderPickerDecision(sessionState?._meta),
-			workingDirectory: sessionState?.workingDirectories?.[0],
-			workingDirectories: sessionState?.workingDirectories,
+			workingDirectory: workingDirectories?.[0],
+			workingDirectories,
 			rootConfig: rootState && !(rootState instanceof Error) ? rootState.config : undefined,
 			isBundledMcpServer: (pluginUri, serverName) => this._activeClientService.isBundledMcpServer(pluginUri, serverName),
 			authenticate: request => target.connection.authenticate(request),

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostUntitledProvisionalSessionService.ts

@@ -105,6 +105,9 @@ export interface IAgentHostUntitledProvisionalSessionService {
 	 */
 	get(sessionResource: URI): URI | undefined;
 
+	/** Working directories used to create the current provisional generation. */
+	getProvisionalWorkingDirectories(sessionResource: URI): readonly URI[] | undefined;
+
 	/**
 	 * Initial config the editor window applies to every new Agent Host session.
 	 * Returns `undefined` in the Agents window, where the sessions provider owns
@@ -373,6 +376,14 @@ export class AgentHostUntitledProvisionalSessionService extends Disposable imple
 		return this._generationMatchingDesiredState(entry)?.backendSession;
 	}
 
+	getProvisionalWorkingDirectories(sessionResource: URI): readonly URI[] | undefined {
+		const entry = this._entries.get(sessionResource);
+		if (!entry || entry.disposed) {
+			return undefined;
+		}
+		return this._generationMatchingDesiredState(entry)?.workingDirectories;
+	}
+
 	private _computeWorkingDirectories(primary: URI | undefined, provider: string): readonly URI[] | undefined {
 		return computeWorkingDirectories(primary, this._workspaceContextService.getWorkspace().folders.map(folder => folder.uri), this._agentHostService.rootState.value, provider);
 	}

src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostUntitledProvisionalSessionService.test.ts

@@ -1494,6 +1494,26 @@ suite('AgentHostUntitledProvisionalSessionService', () => {
 		});
 	});
 
+	test('retains working directories after rebinding a provisional session', async () => {
+		const folderA = URI.file('/repoA');
+		const folderB = URI.file('/repoB');
+		workspaceFolders = [folderA, folderB];
+		agentHost.rootStateAgents = [agentInfo('copilot', true)];
+		const untitled = untitledChatUri('rebind-roots');
+		const real = URI.from({ scheme: 'agent-host-copilot', path: '/real-rebind-roots' });
+
+		await provisional.getOrCreate(untitled, 'copilot', folderA);
+		await provisional.tryRebind(untitled, real, 'copilot');
+
+		assert.deepStrictEqual({
+			untitled: provisional.getProvisionalWorkingDirectories(untitled),
+			real: provisional.getProvisionalWorkingDirectories(real)?.map(directory => directory.toString()),
+		}, {
+			untitled: undefined,
+			real: [folderA.toString(), folderB.toString()],
+		});
+	});
+
 	test('sends only the primary when the provider does not advertise multiple working directories', async () => {
 		const folderA = URI.file('/repoA');
 		const folderB = URI.file('/repoB');