Add local agent host test

src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts

@@ -740,14 +740,16 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC
 		if (config?.activeClient?.customizations) {
 			this._grantImplicitReadsForCustomizations(config.activeClient.customizations);
 		}
-		await this._sendRequest('createSession', {
+		const inflight = this._sendRequest('createSession', {
 			channel: session.toString(),
 			provider,
 			model: config?.model,
 			workingDirectory: config?.workingDirectory ? fromAgentHostUri(config.workingDirectory).toString() : undefined,
 			config: config?.config,
 			activeClient: config?.activeClient,
 		});
+		this._subscriptionManager.trackSessionCreate(session, inflight);
+		await inflight;
 		return session;
 	}
 

src/vs/platform/agentHost/common/state/agentSubscription.ts

@@ -474,6 +474,7 @@ type ManagedSubscriptionEntry = { sub: ManagedSubscription; kind: StateComponent
 export class AgentSubscriptionManager extends Disposable {
 
 	private readonly _subscriptions = new ResourceMap<ManagedSubscriptionEntry>();
+	private readonly _inflightCreates = new ResourceMap<Promise<unknown>>();
 	private _referenceOwnerIds = 0;
 	private readonly _rootState: RootStateSubscription;
 	private readonly _clientId: string;
@@ -520,6 +521,20 @@ export class AgentSubscriptionManager extends Disposable {
 		return entry?.sub as IAgentSubscription<T> | undefined;
 	}
 
+	/**
+	 * Register an in-flight `createSession` Promise for a session URI. Any
+	 * subscribe issued for this resource while the create is pending waits
+	 * for the Promise before issuing the wire-level subscribe.
+	 */
+	trackSessionCreate(resource: URI, promise: Promise<unknown>): void {
+		this._inflightCreates.set(resource, promise);
+		void promise.finally(() => {
+			if (this._inflightCreates.get(resource) === promise) {
+				this._inflightCreates.delete(resource);
+			}
+		});
+	}
+
 	/**
 	 * Get or create a refcounted subscription to any resource. Disposing
 	 * the returned reference decrements the refcount; when it reaches zero
@@ -553,15 +568,28 @@ export class AgentSubscriptionManager extends Disposable {
 		// Kick off server subscription asynchronously.
 		// Capture the entry reference so we can validate it hasn't been
 		// replaced by a new subscription for the same key (race guard).
-		this._subscribe(resource).then(snapshot => {
-			if (this._subscriptions.get(resource) === entry) {
-				sub.handleSnapshot(snapshot.state as never, snapshot.fromSeq);
+		void (async () => {
+			const inflight = this._inflightCreates.get(resource);
+			if (inflight) {
+				try {
+					await inflight;
+				} catch {
+					// Swallow — fall through to subscribe so the error
+					// surfaces consistently via setError() on the
+					// subscription, matching the no-inflight path.
+				}
 			}
-		}).catch(err => {
-			if (this._subscriptions.get(resource) === entry) {
-				sub.setError(err instanceof Error ? err : new Error(String(err)));
+			try {
+				const snapshot = await this._subscribe(resource);
+				if (this._subscriptions.get(resource) === entry) {
+					sub.handleSnapshot(snapshot.state as never, snapshot.fromSeq);
+				}
+			} catch (err) {
+				if (this._subscriptions.get(resource) === entry) {
+					sub.setError(err instanceof Error ? err : new Error(String(err)));
+				}
 			}
-		});
+		})();
 
 		return this._acquireReference<T>(resource, entry, owner);
 	}

src/vs/platform/agentHost/electron-browser/localAgentHostService.ts

@@ -199,7 +199,17 @@ export class LocalAgentHostServiceClient extends Disposable implements IAgentHos
 		return this._proxy.listSessions();
 	}
 	createSession(config?: IAgentCreateSessionConfig): Promise<URI> {
-		return this._proxy.createSession(config);
+		const promise = this._proxy.createSession(config);
+		// When the caller pre-specifies the session URI, a subscribe for
+		// that URI can race the in-flight create. Register the promise so
+		// `AgentSubscriptionManager.getSubscription` gates the wire-level
+		// subscribe on it (avoids transient `AHP_SESSION_NOT_FOUND`).
+		// When the server assigns the URI, no caller can subscribe to it
+		// ahead of `await createSession()`, so there's no race to track.
+		if (config?.session) {
+			this._subscriptionManager.trackSessionCreate(config.session, promise);
+		}
+		return promise;
 	}
 	resolveSessionConfig(params: IAgentResolveSessionConfigParams): Promise<ResolveSessionConfigResult> {
 		return this._proxy.resolveSessionConfig(params);

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

@@ -141,6 +141,10 @@ export class AgentHostStateManager extends Disposable {
 		return [...this._sessionStates.keys()];
 	}
 
+	getAnnouncedSessionSummaries(): SessionSummary[] {
+		return [...this._lastNotifiedSummaries.values()];
+	}
+
 	/**
 	 * Returns all session URIs whose keys start with the given prefix.
 	 * Used to discover subagent sessions for a given parent.

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

@@ -418,8 +418,39 @@ export class AgentService extends Disposable implements IAgentService {
 			return s;
 		});
 
-		this._logService.trace(`[AgentService] listSessions returned ${withStatus.length} sessions`);
-		return withStatus;
+		// Overlay any session that has been announced via `sessionAdded`
+		// but is missing from the providers' `listSessions` snapshot.
+		// Providers can briefly drop a just-materialized session (e.g.
+		// between firing `sessionAdded` and the SDK's session DB becoming
+		// visible to the next `listSessions` call), and immediately after
+		// `session/turnComplete` we've observed `CopilotAgent.listSessions`
+		// return an empty array transiently. Without this overlay,
+		// renderer-side session caches evict the live session, which
+		// closes the chat view holding the in-flight response bubble.
+		const known = new Set(withStatus.map(s => s.session.toString()));
+		const additions: IAgentSessionMetadata[] = [];
+		for (const summary of this._stateManager.getAnnouncedSessionSummaries()) {
+			if (known.has(summary.resource)) {
+				continue;
+			}
+			additions.push({
+				session: URI.parse(summary.resource),
+				startTime: summary.createdAt,
+				modifiedTime: summary.modifiedAt,
+				summary: summary.title,
+				status: summary.status,
+				activity: summary.activity,
+				model: summary.model,
+				agent: summary.agent,
+				workingDirectory: typeof summary.workingDirectory === 'string' ? URI.parse(summary.workingDirectory) : undefined,
+				...(summary.project ? { project: { uri: URI.parse(summary.project.uri), displayName: summary.project.displayName } } : {}),
+				changesets: summary.changesets,
+			});
+		}
+		const combined = additions.length > 0 ? [...withStatus, ...additions] : withStatus;
+
+		this._logService.trace(`[AgentService] listSessions returned ${combined.length} sessions (${additions.length} state-manager fallback)`);
+		return combined;
 	}
 
 	async createSession(config?: IAgentCreateSessionConfig): Promise<URI> {

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -24,7 +24,7 @@ import { localize } from '../../../../nls.js';
 import { IParsedPlugin, parseAgentFile, parsePlugin, parseSkillFile } from '../../../agentPlugins/common/pluginParsers.js';
 import { IFileService } from '../../../files/common/files.js';
 import { IInstantiationService } from '../../../instantiation/common/instantiation.js';
-import { ILogService } from '../../../log/common/log.js';
+import { ILogService, LogLevel } from '../../../log/common/log.js';
 import { AgentHostConfigKey, agentHostCustomizationConfigSchema, toContainerCustomization } from '../../common/agentHostCustomizationConfig.js';
 import { AgentHostSessionSyncEnabledConfigKey, AutoApproveLevel, ISchemaProperty, SessionMode, createSchema, platformRootSchema, platformSessionSchema, schemaProperty } from '../../common/agentHostSchema.js';
 import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js';
@@ -51,6 +51,24 @@ import { SessionPluginBundler } from '../shared/sessionPluginBundler.js';
 import { CopilotSlashCommandCompletionProvider } from './copilotSlashCommandCompletionProvider.js';
 import { getEffectiveAgents } from '../../common/customAgents.js';
 
+/**
+ * Maps a VS Code {@link LogLevel} to the Copilot CLI runtime's `logLevel`
+ * option so the spawned CLI logs (written to `~/.copilot/logs/process-*.log`)
+ * match the agent host's configured verbosity. `Trace` maps to the CLI's most
+ * verbose `'all'` level so renderer-side trace logging surfaces the CLI's
+ * internal diagnostics.
+ */
+function copilotCliLogLevelFor(level: LogLevel): NonNullable<CopilotClientOptions['logLevel']> {
+	switch (level) {
+		case LogLevel.Off: return 'none';
+		case LogLevel.Trace: return 'all';
+		case LogLevel.Debug: return 'debug';
+		case LogLevel.Info: return 'info';
+		case LogLevel.Warning: return 'warning';
+		case LogLevel.Error: return 'error';
+	}
+}
+
 interface ICreatedWorktree {
 	readonly repositoryRoot: URI;
 	readonly worktree: URI;
@@ -529,6 +547,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 				connection: RuntimeConnection.forStdio({ path: cliPath }),
 				env,
 				telemetry,
+				logLevel: copilotCliLogLevelFor(this._logService.getLevel()),
 				enableRemoteSessions: this._isSessionSyncEnabled(),
 			};
 			const client = this._createCopilotClient(clientOptions);

src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts

@@ -25,7 +25,7 @@ import type { IAgentSubscription } from '../../../../../platform/agentHost/commo
 import { ResolveSessionConfigResult } from '../../../../../platform/agentHost/common/state/protocol/commands.js';
 import { AgentCustomization, AgentSelection, Customization, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionActiveClient, SessionState, SessionSummary, type ChangesetSummary } from '../../../../../platform/agentHost/common/state/protocol/state.js';
 import { ActionType, isSessionAction, NotificationType } from '../../../../../platform/agentHost/common/state/sessionActions.js';
-import { readSessionGitState, ROOT_STATE_URI, SessionMeta, StateComponents, type ISessionGitState } from '../../../../../platform/agentHost/common/state/sessionState.js';
+import { AgentInfo, readSessionGitState, ROOT_STATE_URI, SessionMeta, StateComponents, type ISessionGitState } from '../../../../../platform/agentHost/common/state/sessionState.js';
 import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
 import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
 import { ILogService } from '../../../../../platform/log/common/log.js';
@@ -1077,6 +1077,8 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 	get sessionTypes(): readonly ISessionType[] { return this._sessionTypes; }
 	protected _sessionTypes: ISessionType[] = [];
 
+	private _lastAgents: AgentInfo[] | undefined;
+
 	protected readonly _onDidChangeSessionTypes = this._register(new Emitter<void>());
 	readonly onDidChangeSessionTypes: Event<void> = this._onDidChangeSessionTypes.event;
 
@@ -1272,8 +1274,11 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 	 * id/label set actually changed.
 	 */
 	protected _syncSessionTypesFromRootState(rootState: RootState): void {
-		this._onDidChangeCustomAgents.fire();
-		this._onDidChangeCustomizations.fire();
+		if (this._lastAgents !== rootState.agents) {
+			this._lastAgents = rootState.agents;
+			this._onDidChangeCustomAgents.fire();
+			this._onDidChangeCustomizations.fire();
+		}
 		const next = rootState.agents.map((agent): ISessionType => ({
 			id: agent.provider,
 			label: this._formatSessionTypeLabel(agent.displayName?.trim() || agent.provider),

src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts

@@ -206,6 +206,20 @@ class MockAgentHostService extends mock<IAgentHostService>() {
 		this._onDidRootStateChange.fire(this._rootStateValue);
 	}
 
+	/**
+	 * Fires a root state change that preserves the current `agents` reference,
+	 * simulating non-agent root deltas (e.g. `RootActiveSessionsChanged` on
+	 * every turn start/complete) that the real reducer emits without
+	 * replacing the `agents` slice.
+	 */
+	fireNonAgentRootStateChange(): void {
+		if (!this._rootStateValue || this._rootStateValue instanceof Error) {
+			throw new Error('rootState not initialized; call setAgents first');
+		}
+		this._rootStateValue = { ...this._rootStateValue };
+		this._onDidRootStateChange.fire(this._rootStateValue);
+	}
+
 	clearRootState(): void {
 		this._rootStateValue = undefined;
 	}
@@ -945,12 +959,21 @@ suite('LocalAgentHostSessionsProvider', () => {
 		let fired = 0;
 		disposables.add(provider.onDidChangeCustomAgents(() => { fired++; }));
 
-		// Root state change should fire the event.
+		// A root state change that replaces the agents reference should
+		// fire the event. This is the only path that mutates agents in the
+		// real reducer (`RootAgentsChanged`).
 		agentHost.setAgents([
 			{ provider: 'copilotcli', displayName: 'Copilot', description: '', models: [] } as AgentInfo,
 		]);
 		const afterRoot = fired;
-		assert.ok(afterRoot > 0, 'expected event to fire on root state change');
+		assert.ok(afterRoot > 0, 'expected event to fire when the agents reference is replaced');
+
+		// A subsequent root state change that preserves the agents reference
+		// (e.g. `activeSessionsChanged` on every turn start/complete) must
+		// NOT fire — firing on those caused chat session bubbles to be
+		// re-hydrated mid-turn, dropping streamed responses.
+		agentHost.fireNonAgentRootStateChange();
+		assert.strictEqual(fired, afterRoot, 'expected event NOT to fire on non-agent root deltas (preserved agents reference)');
 
 		// Session-state update with new customizations should fire it again.
 		provider.getSessionConfig(session!.sessionId);

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

@@ -54,6 +54,7 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr
 	private readonly _authTokenCache = new AgentHostAuthTokenCache();
 
 	private readonly _isSessionsWindow: boolean;
+	private readonly _enableSmokeTestDriver: boolean;
 
 	constructor(
 		@IAgentHostService private readonly _agentHostService: IAgentHostService,
@@ -64,16 +65,17 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr
 		@ILanguageModelsService private readonly _languageModelsService: ILanguageModelsService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@IAgentHostFileSystemService _agentHostFileSystemService: IAgentHostFileSystemService,
-		@IConfigurationService configurationService: IConfigurationService,
+		@IConfigurationService private readonly _configurationService: IConfigurationService,
 		@ICustomizationHarnessService private readonly _customizationHarnessService: ICustomizationHarnessService,
 		@IWorkbenchEnvironmentService environmentService: IWorkbenchEnvironmentService,
 		@IAgentHostActiveClientService private readonly _activeClientService: IAgentHostActiveClientService,
 	) {
 		super();
 
 		this._isSessionsWindow = environmentService.isSessionsWindow;
+		this._enableSmokeTestDriver = !!environmentService.enableSmokeTestDriver;
 
-		if (!configurationService.getValue<boolean>(AgentHostEnabledSettingId)) {
+		if (!this._configurationService.getValue<boolean>(AgentHostEnabledSettingId)) {
 			return;
 		}
 
@@ -235,6 +237,11 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr
 	private async _authenticateWithServer(agents: readonly AgentInfo[]): Promise<void> {
 		this._agentHostService.setAuthenticationPending(true);
 		try {
+			const testToken = this._getScenarioAutomationToken();
+			if (testToken !== undefined) {
+				await this._seedTestToken(agents, testToken);
+				return;
+			}
 			await authenticateProtectedResources(agents, {
 				authTokenCache: this._authTokenCache,
 				authenticationService: this._authenticationService,
@@ -256,6 +263,14 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr
 	 * to the server. Returns true if authentication succeeded.
 	 */
 	private async _resolveAuthenticationInteractively(protectedResources: ProtectedResourceMetadata[]): Promise<boolean> {
+		const testToken = this._getScenarioAutomationToken();
+		if (testToken !== undefined) {
+			for (const resource of protectedResources) {
+				await this._agentHostService.authenticate({ resource: resource.resource, token: testToken });
+				this._authTokenCache.updateAndIsChanged(resource.resource, resource.scopes_supported, testToken);
+			}
+			return protectedResources.length > 0;
+		}
 		try {
 			return await resolveAuthenticationInteractively(protectedResources, {
 				authTokenCache: this._authTokenCache,
@@ -269,4 +284,29 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr
 		}
 		return false;
 	}
+
+	private async _seedTestToken(agents: readonly AgentInfo[], token: string): Promise<void> {
+		for (const agent of agents) {
+			for (const resource of agent.protectedResources ?? []) {
+				if (!this._authTokenCache.updateAndIsChanged(resource.resource, resource.scopes_supported, token)) {
+					continue;
+				}
+				try {
+					await this._agentHostService.authenticate({ resource: resource.resource, token });
+				} catch (err) {
+					this._authTokenCache.clear(resource.resource);
+					throw err;
+				}
+			}
+		}
+	}
+
+	private _getScenarioAutomationToken(): string | undefined {
+		// Smoke-test escape hatch.
+		if (!this._enableSmokeTestDriver) {
+			return undefined;
+		}
+		const token = this._configurationService.getValue('chat.agentHost.unsafeTestToken');
+		return typeof token === 'string' && token.length > 0 ? token : undefined;
+	}
 }