automations: fix: honor worktree and provider session configuration

src/vs/platform/agentHost/node/agentHostAutomationService.ts

@@ -754,10 +754,12 @@ export class AgentHostAutomationService extends Disposable implements IAgentHost
 				await this._execution.cancelSession(session);
 				return;
 			}
-			// Clients restore the last turn's model configuration, not the SDK's creation defaults.
-			const message: Message = definition.message.model === undefined && definition.session.model !== undefined
-				? { ...definition.message, model: definition.session.model }
-				: definition.message;
+			// Turn selections override creation defaults in both the provider and restored clients.
+			const message: Message = {
+				...definition.message,
+				...(definition.message.model === undefined && definition.session.model !== undefined ? { model: definition.session.model } : {}),
+				...(definition.message.agent === undefined && definition.session.agent !== undefined ? { agent: definition.session.agent } : {}),
+			};
 			await this._execution.startSession(session, message);
 		} catch (error) {
 			try {

src/vs/platform/agentHost/node/agentService.ts

@@ -2926,10 +2926,7 @@ export class AgentService extends Disposable implements IAgentService {
 			config = { ...config, importConversation: { ...config.importConversation, turns: importedTurns } };
 		}
 
-		// Resolve host-owned isolation before provider creation. Providers such as
-		// Codex may schedule eager prewarming from createSession; marking a
-		// client-chosen worktree session pending first prevents that prewarm from
-		// materializing in the picked folder before the host creates the worktree.
+		// Providers must see pending isolation before creation can schedule eager prewarming.
 		const initializeSideEffects = this._sideEffects.initialize();
 		const sessionConfig = await this._resolveCreatedSessionConfig(provider, config);
 		const deferWorktreeCreation = sessionConfig?.values?.[SessionConfigKey.Isolation] === 'worktree' && !config?.importConversation;
@@ -3445,15 +3442,15 @@ export class AgentService extends Disposable implements IAgentService {
 	}
 
 	private async _createProviderSession(provider: IAgent, config: IAgentCreateSessionConfig | undefined, deferWorktreeCreation: boolean): Promise<IAgentCreateSessionResult> {
-		const requestedSessionId = deferWorktreeCreation && config?.session ? AgentSession.id(config.session) : undefined;
-		if (requestedSessionId) {
-			this._worktree.notePending(requestedSessionId);
+		const session = config?.session ?? this._mintSessionUri(provider);
+		const pendingSessionId = deferWorktreeCreation ? AgentSession.id(session) : undefined;
+		if (pendingSessionId) {
+			this._worktree.notePending(pendingSessionId);
 		}
 
 		let created: IAgentCreateSessionResult | undefined;
 		try {
 			const providerConfig = config ? this._toProviderConfig(config) : undefined;
-			const session = config?.session ?? this._mintSessionUri(provider);
 			const defaultChatUri = URI.parse(buildDefaultChatUri(session));
 			const boundConfig: IAgentCreateSessionConfig = { ...(providerConfig ?? {}), session };
 			const result = await provider.chats.createChat(defaultChatUri, this._chatContext(session, defaultChatUri), this._toCreateChatOptions(boundConfig));
@@ -3464,9 +3461,6 @@ export class AgentService extends Disposable implements IAgentService {
 				...(result?.provisional ? { provisional: true } : {}),
 				...(result ? { chat: result } : {}),
 			};
-			if (deferWorktreeCreation && created.provisional) {
-				this._worktree.notePending(AgentSession.id(created.session));
-			}
 			await this._persistDefaultChatBacking(created);
 			return created;
 		} catch (err) {
@@ -3475,9 +3469,8 @@ export class AgentService extends Disposable implements IAgentService {
 			}
 			throw err;
 		} finally {
-			const returnedPendingSessionId = created?.provisional ? AgentSession.id(created.session) : undefined;
-			if (requestedSessionId && requestedSessionId !== returnedPendingSessionId) {
-				this._worktree.clearPending(requestedSessionId);
+			if (pendingSessionId && !created?.provisional) {
+				this._worktree.clearPending(pendingSessionId);
 			}
 		}
 	}
@@ -3493,6 +3486,8 @@ export class AgentService extends Disposable implements IAgentService {
 			await provider.chats.disposeChat(defaultChatUri, this._chatContext(session, defaultChatUri));
 		} catch (disposeError) {
 			this._logService.error(disposeError, `[AgentService] Failed to roll back default chat of provider session ${session.toString()}`);
+		} finally {
+			this._worktree.clearPending(AgentSession.id(session));
 		}
 	}
 

src/vs/platform/agentHost/node/codex/codexAgent.ts

@@ -69,10 +69,11 @@ import { IAgentHostProxyResolver } from '../agentHostProxyResolver.js';
 import { MODEL_REFRESH_BASE_DELAY_MS, MODEL_REFRESH_MAX_ATTEMPTS, MODEL_REFRESH_MAX_DELAY_MS, modelRefreshBackoff } from '../shared/modelRefreshRetry.js';
 import { AGENT_HOST_WORKSPACELESS_INSTRUCTIONS } from '../shared/workspacelessInstructions.js';
 import { IAgentHostCheckpointService } from '../../common/agentHostCheckpointService.js';
+import { IAgentHostGitService, tryResolvePrimaryWorktreeRoot } from '../../common/agentHostGitService.js';
 import { ISessionDataService } from '../../common/sessionDataService.js';
 import { ICopilotApiService } from '../shared/copilotApiService.js';
 import { extractForwardedErrorInfo } from '../shared/proxyChatError.js';
-import { IAgentHostWorktreeIsolation, type IAgentHostWorktreePendingState } from '../shared/worktreeIsolation.js';
+import { IAgentHostWorktreeIsolation } from '../shared/worktreeIsolation.js';
 import { getServerToolDisplay } from '../shared/serverToolGroups.js';
 import { IAgentSdkDownloader, IAgentSdkPackage } from '../agentSdkDownloader.js';
 import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
@@ -1252,7 +1253,6 @@ export class CodexAgent extends Disposable implements IAgent {
 	private readonly _metadataStore: CodexSessionMetadataStore;
 	private _lastSignInRequest: string | undefined;
 	private _lastSignOutRequest: string | undefined;
-	private readonly _worktree: IAgentHostWorktreePendingState;
 
 	/**
 	 * The agent host's server-tool host (feedback "comments" today, more in the
@@ -1280,12 +1280,12 @@ export class CodexAgent extends Disposable implements IAgent {
 		@IAgentHostOTelService private readonly _otelService: IAgentHostOTelService,
 		@IAgentHostCustomizationEnablementService private readonly _customizationEnablementService: IAgentHostCustomizationEnablementService,
 		@IAgentHostSessionTitleSignal sessionTitleSignal: IAgentHostSessionTitleSignal,
-		@IAgentHostWorktreeIsolation worktree: IAgentHostWorktreeIsolation,
+		@IAgentHostWorktreeIsolation private readonly _worktree: IAgentHostWorktreeIsolation,
 		@ISessionDataService private readonly _sessionDataService: ISessionDataService,
 		@ITelemetryService private readonly _telemetryService: ITelemetryService,
+		@IAgentHostGitService private readonly _gitService: IAgentHostGitService,
 	) {
 		super();
-		this._worktree = worktree;
 		this._metadataStore = this._instantiationService.createInstance(CodexSessionMetadataStore);
 		this._githubMcpServerEnabled = this._isGitHubMcpServerEnabled();
 		this._publishAccountInfo({ status: 'unknown' });
@@ -1927,13 +1927,50 @@ export class CodexAgent extends Disposable implements IAgent {
 		return resolved.filter(candidate => candidate !== undefined);
 	}
 
+	/** Resolve native workspace agents in the host-owned worktree without rewriting their persisted selection identity. */
+	private async _resolveSelectedAgent(session: ICodexSession): Promise<AgentSelection | undefined> {
+		const agent = session.agent;
+		if (!agent || !session.workingDirectory) {
+			return agent;
+		}
+		const agentUri = URI.parse(agent.uri);
+		const agentsDirectory = extUriBiasedIgnorePathCase.dirname(agentUri);
+		const sourceRoot = extUriBiasedIgnorePathCase.dirname(extUriBiasedIgnorePathCase.dirname(agentsDirectory));
+		if (!extUriBiasedIgnorePathCase.isEqual(agentsDirectory, URI.joinPath(sourceRoot, '.github', 'agents'))
+			|| extUriBiasedIgnorePathCase.isEqual(sourceRoot, session.workingDirectory)) {
+			return agent;
+		}
+		const worktree = await this._worktree.readWorktreeMetadata(session.configurationResource);
+		if (!worktree?.repositoryRoot || !worktree.worktreePath
+			|| !extUriBiasedIgnorePathCase.isEqual(session.workingDirectory, worktree.worktreePath)) {
+			return agent;
+		}
+		if (!extUriBiasedIgnorePathCase.isEqual(sourceRoot, worktree.repositoryRoot)) {
+			try {
+				const checkoutRoot = await this._gitService.getRepositoryRoot(sourceRoot);
+				if (!checkoutRoot || !extUriBiasedIgnorePathCase.isEqual(checkoutRoot, sourceRoot)
+					|| !extUriBiasedIgnorePathCase.isEqual(await tryResolvePrimaryWorktreeRoot(this._gitService, checkoutRoot), worktree.repositoryRoot)) {
+					return agent;
+				}
+			} catch (error) {
+				this._logService.warn('[Codex] Failed to resolve the selected workspace agent repository', error);
+				return agent;
+			}
+		}
+		return {
+			...agent,
+			uri: URI.joinPath(worktree.worktreePath, '.github', 'agents', extUriBiasedIgnorePathCase.basename(agentUri)).toString(),
+		};
+	}
+
 	private async _buildCustomizationLaunch(session: ICodexSession): Promise<ICodexCustomizationLaunch> {
 		const plugins = this._enabledClientPlugins(session);
-		const [workspaceAgents, workspaceSkills] = await Promise.all([
+		const [workspaceAgents, workspaceSkills, selectedAgent] = await Promise.all([
 			discoverCodexWorkspaceAgents(this._customizationWorkingDirectories(session), this._fileService),
 			discoverCodexWorkspaceSkills(this._customizationWorkingDirectories(session), this._fileService),
+			this._resolveSelectedAgent(session),
 		]);
-		const customization = await codexCustomizationConfig(workspaceAgents.agents, plugins, session.agent, this._fileService);
+		const customization = await codexCustomizationConfig(workspaceAgents.agents, plugins, selectedAgent, this._fileService);
 		const developerInstructions = [
 			customization.developerInstructions,
 			session.managedWorkingDirectory ? AGENT_HOST_WORKSPACELESS_INSTRUCTIONS : '',
@@ -1965,7 +2002,7 @@ export class CodexAgent extends Disposable implements IAgent {
 			})),
 		];
 		const signature = JSON.stringify({
-			agent: session.agent?.uri,
+			agent: selectedAgent?.uri,
 			agentRoles: customization.agentRoles,
 			developerInstructions,
 			selectedCapabilityRoots: selectedCapabilityRoots.map(root => root.location.path),

src/vs/platform/agentHost/test/node/agentHostAutomationService.test.ts

@@ -725,13 +725,16 @@ suite('AgentHostAutomationService', () => {
 		})));
 	});
 
-	for (const hasMessageModel of [false, true]) {
-		test(hasMessageModel ? 'preserves an explicit Automation message model' : 'records the Automation model configuration on its first turn', async () => {
+	for (const messageOverride of ['none', 'model', 'agent', 'both']) {
+		test(`records the Automation session selections on its first turn with ${messageOverride} message overrides`, async () => {
 			const session = URI.parse('mock:/model-configuration-run');
 			const model = { id: 'mock-model', config: { thinkingLevel: 'low', contextSize: 272_000 } };
-			const messageModel = hasMessageModel ? { id: 'other-model', config: { thinkingLevel: 'high' } } : undefined;
+			const agent = { uri: 'file:///workspace/.github/agents/reviewer.agent.md' };
+			const messageModel = messageOverride === 'model' || messageOverride === 'both' ? { id: 'other-model', config: { thinkingLevel: 'high' } } : undefined;
+			const messageAgent = messageOverride === 'agent' || messageOverride === 'both' ? { uri: 'file:///other/agents/reviewer.agent.md' } : undefined;
 			const completed = new DeferredPromise<void>();
 			let createdModel: AutomationDefinition['session']['model'];
+			let createdAgent: AutomationDefinition['session']['agent'];
 			disposables.add(stateManager.onDidEmitEnvelope(envelope => {
 				if (envelope.action.type === ActionType.AutomationRunLifecycleChanged && envelope.action.lifecycle.status === AutomationRunStatus.Completed) {
 					void completed.complete();
@@ -740,6 +743,7 @@ suite('AgentHostAutomationService', () => {
 			const service = createService({
 				createSession: async template => {
 					createdModel = template.model;
+					createdAgent = template.agent;
 					stateManager.createSession({
 						resource: session.toString(),
 						provider: 'mock',
@@ -767,9 +771,13 @@ suite('AgentHostAutomationService', () => {
 			});
 			const automation = definition();
 			automation.session.model = model;
+			automation.session.agent = agent;
 			if (messageModel) {
 				automation.message.model = messageModel;
 			}
+			if (messageAgent) {
+				automation.message.agent = messageAgent;
+			}
 			await service.completeMigration();
 			await service.handleCreate({ ...createAction(), definition: automation });
 			await service.runAutomation({
@@ -781,12 +789,24 @@ suite('AgentHostAutomationService', () => {
 
 			assert.deepStrictEqual({
 				createdModel,
+				createdAgent,
 				recordedModel: stateManager.getChatState(buildDefaultChatUri(session))?.turns[0]?.message.model,
+				recordedAgent: stateManager.getChatState(buildDefaultChatUri(session))?.turns[0]?.message.agent,
 				savedModel: stateManager.getAutomationCatalogState()?.entries[0].definition.session.model,
+				savedAgent: stateManager.getAutomationCatalogState()?.entries[0].definition.session.agent,
+				savedMessage: stateManager.getAutomationCatalogState()?.entries[0].definition.message,
 			}, {
 				createdModel: model,
+				createdAgent: agent,
 				recordedModel: messageModel ?? model,
+				recordedAgent: messageAgent ?? agent,
 				savedModel: model,
+				savedAgent: agent,
+				savedMessage: {
+					...definition().message,
+					...(messageModel ? { model: messageModel } : {}),
+					...(messageAgent ? { agent: messageAgent } : {}),
+				},
 			});
 		});
 	}