Preserve customization roots during provisional sessions

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

@@ -434,7 +434,7 @@ export function getPresentableMcpServerCustomizations(customizations: readonly C
 	return entries.filter(entry => entry.isTopLevel || !topLevelNames.has(entry.server.name));
 }
 
-class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizationService {
+export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizationService {
 
 	private readonly _sessionStateSubscriptions = this._register(new DisposableResourceMap<IDisposable & { readonly connection: IAgentConnection; readonly backendSession: URI; readonly sub: IAgentSubscription<SessionState> }>());
 
@@ -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),
@@ -527,8 +530,9 @@ class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizat
 
 	private _readSessionState(sessionResource: URI): SessionState | undefined {
 		const target = this._resolveSessionTarget(sessionResource);
-		const value = target ? this._ensureSessionStateSubscription(sessionResource, target)?.sub.value : undefined;
-		return value && !(value instanceof Error) ? value : undefined;
+		const subscription = target ? this._ensureSessionStateSubscription(sessionResource, target)?.sub : undefined;
+		const value = subscription?.value;
+		return value instanceof Error ? subscription?.verifiedValue : value;
 	}
 
 	private _ensureSessionStateSubscription(sessionResource: URI, target: IAgentHostSessionResolution): (IDisposable & { readonly connection: IAgentConnection; readonly backendSession: URI; readonly sub: IAgentSubscription<SessionState> }) | undefined {

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/agentHostCustomizationService.test.ts

@@ -4,15 +4,25 @@
  *--------------------------------------------------------------------------------------------*/
 
 import assert from 'assert';
+import { Event } from '../../../../../../base/common/event.js';
+import { IReference } from '../../../../../../base/common/lifecycle.js';
 import { ResourceMap } from '../../../../../../base/common/map.js';
 import { URI } from '../../../../../../base/common/uri.js';
+import { mock } from '../../../../../../base/test/common/mock.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
+import { IAgentHostConnectionsService } from '../../../../../../platform/agentHost/common/agentHostConnectionsService.js';
+import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js';
+import { IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js';
 import { CustomizationEnablementKind, CustomizationType, McpServerCustomization, McpServerStatus, type Customization, type CustomizationEnablement } from '../../../../../../platform/agentHost/common/state/protocol/state.js';
 import { createAgentHostResourceUriMapper, identityAgentHostResourceUriMapper, IAgentHostResourceUriMapper } from '../../../../../../platform/agentHost/common/agentHostUri.js';
+import { createSessionState, RootState, SessionState, SessionStatus, StateComponents } from '../../../../../../platform/agentHost/common/state/sessionState.js';
 import { IOutputService } from '../../../../../services/output/common/output.js';
 import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
 import { ILogService, ILoggerService, NullLogService, NullLoggerService } from '../../../../../../platform/log/common/log.js';
-import { AbstractAgentHostCustomizationService, IAgentHostCustomizationTarget } from '../../../browser/agentSessions/agentHost/agentHostCustomizationService.js';
+import { AbstractAgentHostCustomizationService, IAgentHostCustomizationTarget, WorkbenchAgentHostCustomizationService } from '../../../browser/agentSessions/agentHost/agentHostCustomizationService.js';
+import { IAgentHostUntitledProvisionalSessionService } from '../../../browser/agentSessions/agentHost/agentHostUntitledProvisionalSessionService.js';
+import { IChatService } from '../../../common/chatService/chatService.js';
+import { IAgentHostActiveClientService } from '../../../browser/agentSessions/agentHost/agentHostActiveClientService.js';
 
 class FakeTarget implements IAgentHostCustomizationTarget {
 	readonly enablementChanges: { readonly rawId: string; readonly enablement: readonly CustomizationEnablement[] }[] = [];
@@ -66,6 +76,29 @@ class TestAgentHostCustomizationService extends AbstractAgentHostCustomizationSe
 	}
 }
 
+class TestSessionSubscription extends mock<IAgentSubscription<SessionState>>() {
+	override readonly onDidChange = Event.None;
+	private current: SessionState | Error | undefined;
+	private confirmed: SessionState | undefined;
+
+	override get value(): SessionState | Error | undefined {
+		return this.current;
+	}
+
+	override get verifiedValue(): SessionState | undefined {
+		return this.confirmed;
+	}
+
+	setSnapshot(state: SessionState): void {
+		this.current = state;
+		this.confirmed = state;
+	}
+
+	setError(error: Error): void {
+		this.current = error;
+	}
+}
+
 suite('AbstractAgentHostCustomizationService', () => {
 	const store = ensureNoDisposablesAreLeakedInTestSuite();
 
@@ -268,4 +301,93 @@ suite('AbstractAgentHostCustomizationService', () => {
 			disabledReason: { source: 'scope', scope: CustomizationEnablementKind.Session },
 		});
 	});
+
+});
+
+suite('WorkbenchAgentHostCustomizationService', () => {
+	const store = ensureNoDisposablesAreLeakedInTestSuite();
+
+	test('uses provisional roots only until authoritative session state is available', () => {
+		const sessionResource = URI.parse('untitled:chat');
+		const backendSession = URI.parse('copilot:/session');
+		const provisionalRoot = URI.file('/provisional');
+		const hydratedRoot = URI.file('/hydrated');
+		const retainedRoot = URI.file('/retained');
+		const subscription = new TestSessionSubscription();
+		const connection = new class extends mock<IAgentConnection>() {
+			override readonly resourceUris = identityAgentHostResourceUriMapper;
+			override readonly onDidAction = Event.None;
+			override readonly rootState = {
+				value: undefined,
+				verifiedValue: undefined,
+				onDidChange: Event.None,
+				onWillApplyAction: Event.None,
+				onDidApplyAction: Event.None,
+			} satisfies IAgentSubscription<RootState>;
+
+			override getSubscription<T>(_kind: StateComponents): IReference<IAgentSubscription<T>> {
+				return {
+					object: subscription as unknown as IAgentSubscription<T>,
+					dispose: () => { },
+				};
+			}
+		}();
+		const instantiationService = store.add(new TestInstantiationService());
+		instantiationService.stub(ILoggerService, store.add(new NullLoggerService()));
+		instantiationService.stub(IOutputService, {
+			getChannel: () => undefined,
+			getChannelDescriptor: () => undefined,
+			showChannel: async () => { },
+		});
+		const service = store.add(new WorkbenchAgentHostCustomizationService(
+			new class extends mock<IAgentHostConnectionsService>() {
+				override readonly ambientConnection = connection;
+			}(),
+			new class extends mock<IAgentHostUntitledProvisionalSessionService>() {
+				override readonly onDidChange = Event.None;
+				override get(): URI {
+					return backendSession;
+				}
+				override getProvisionalWorkingDirectories(): readonly URI[] {
+					return [provisionalRoot];
+				}
+			}(),
+			instantiationService,
+			new NullLogService(),
+			new class extends mock<IChatService>() {
+				override readonly onDidDisposeSession = Event.None;
+			}(),
+			new class extends mock<IAgentHostActiveClientService>() { }(),
+		));
+		const createState = (workingDirectories: readonly URI[]): SessionState => createSessionState({
+			resource: backendSession.toString(),
+			provider: 'copilot',
+			title: 'Session',
+			status: SessionStatus.Idle,
+			createdAt: new Date(0).toISOString(),
+			modifiedAt: new Date(0).toISOString(),
+			workingDirectories: workingDirectories.map(uri => uri.toString()),
+		});
+
+		const beforeSnapshot = service.getWorkingDirectories(sessionResource);
+		subscription.setSnapshot(createState([hydratedRoot]));
+		const afterSnapshot = service.getWorkingDirectories(sessionResource);
+		subscription.setSnapshot(createState([]));
+		const afterEmptySnapshot = service.getWorkingDirectories(sessionResource);
+		subscription.setSnapshot(createState([retainedRoot]));
+		subscription.setError(new Error('subscription failed'));
+		const afterError = service.getWorkingDirectories(sessionResource);
+
+		assert.deepStrictEqual({
+			beforeSnapshot,
+			afterSnapshot,
+			afterEmptySnapshot,
+			afterError,
+		}, {
+			beforeSnapshot: [provisionalRoot.toString()],
+			afterSnapshot: [hydratedRoot.toString()],
+			afterEmptySnapshot: [],
+			afterError: [retainedRoot.toString()],
+		});
+	});
 });

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');