ensure fast new chat session swap (#1327)

* ensure fast new chat session swap

* move pending session management to session service.

Author: rebornix

src/extension/agents/copilotcli/node/copilotcliSessionService.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { Session, SessionManager } from '@github/copilot/sdk';
-import type { CancellationToken } from 'vscode';
+import type { CancellationToken, ChatContext, ChatRequest } from 'vscode';
 import { IEnvService } from '../../../../platform/env/common/envService';
 import { IVSCodeExtensionContext } from '../../../../platform/extContext/common/extensionContext';
 import { ILogService } from '../../../../platform/log/common/logService';
@@ -20,6 +20,8 @@ export interface ICopilotCLISession {
 	readonly timestamp: Date;
 }
 
+export type ExtendedChatRequest = ChatRequest & { prompt: string };
+
 export interface ICopilotCLISessionService {
 	readonly _serviceBrand: undefined;
 
@@ -35,6 +37,11 @@ export interface ICopilotCLISessionService {
 	// Session wrapper tracking
 	trackSessionWrapper<T extends IDisposable>(sessionId: string, wrapper: T): void;
 	findSessionWrapper<T extends IDisposable>(sessionId: string): T | undefined;
+
+	// Pending request tracking (for untitled sessions)
+	setPendingRequest(sessionId: string, request: ExtendedChatRequest, context: ChatContext): void;
+	getPendingRequest(sessionId: string): { request: ExtendedChatRequest; context: ChatContext } | undefined;
+	clearPendingRequest(sessionId: string): void;
 }
 
 export const ICopilotCLISessionService = createServiceIdentifier<ICopilotCLISessionService>('ICopilotCLISessionService');
@@ -45,6 +52,7 @@ export class CopilotCLISessionService implements ICopilotCLISessionService {
 	private _sessionManager: SessionManager | undefined;
 	private _sessionWrappers = new DisposableMap<string, IDisposable>();
 	private _sessions = new Map<string, ICopilotCLISession>();
+	private _pendingRequests = new Map<string, { request: any; context: any }>();
 
 	constructor(
 		@ILogService private readonly logService: ILogService,
@@ -223,4 +231,16 @@ export class CopilotCLISessionService implements ICopilotCLISessionService {
 		// Fallback to session ID
 		return `Session ${sdkSession.sessionId.slice(0, 8)}`;
 	}
+
+	public setPendingRequest(sessionId: string, request: ExtendedChatRequest, context: ChatContext): void {
+		this._pendingRequests.set(sessionId, { request, context });
+	}
+
+	public getPendingRequest(sessionId: string): { request: ExtendedChatRequest; context: ChatContext } | undefined {
+		return this._pendingRequests.get(sessionId);
+	}
+
+	public clearPendingRequest(sessionId: string): void {
+		this._pendingRequests.delete(sessionId);
+	}
 }

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

@@ -58,8 +58,8 @@ export class ChatSessionsContrib extends Disposable implements IExtensionContrib
 		const copilotcliSessionItemProvider = this._register(copilotcliAgentInstaService.createInstance(CopilotCLIChatSessionItemProvider));
 		this._register(vscode.chat.registerChatSessionItemProvider(this.copilotcliSessionType, copilotcliSessionItemProvider));
 		const copilotcliAgentManager = this._register(copilotcliAgentInstaService.createInstance(CopilotCLIAgentManager));
-		const copilotcliChatSessionContentProvider = copilotcliAgentInstaService.createInstance(CopilotCLIChatSessionContentProvider);
-		const copilotcliChatSessionParticipant = new CopilotCLIChatSessionParticipant(this.copilotcliSessionType, copilotcliAgentManager, copilotcliSessionItemProvider);
+		const copilotcliChatSessionContentProvider = copilotcliAgentInstaService.createInstance(CopilotCLIChatSessionContentProvider, copilotcliAgentManager);
+		const copilotcliChatSessionParticipant = new CopilotCLIChatSessionParticipant(this.copilotcliSessionType, copilotcliAgentManager, copilotCLISessionService, copilotcliSessionItemProvider);
 		const copilotcliParticipant = vscode.chat.createChatParticipant(this.copilotcliSessionType, copilotcliChatSessionParticipant.createHandler());
 		this._register(vscode.chat.registerChatSessionContentProvider(this.copilotcliSessionType, copilotcliChatSessionContentProvider, copilotcliParticipant));
 		this._register(registerCLIChatCommands(copilotcliSessionItemProvider, copilotCLISessionService));

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

@@ -17,8 +17,8 @@ import * as path from '../../../util/vs/base/common/path';
 import { URI } from '../../../util/vs/base/common/uri';
 import { localize } from '../../../util/vs/nls';
 import { CopilotCLIAgentManager } from '../../agents/copilotcli/node/copilotcliAgentManager';
-import { ICopilotCLISessionService } from '../../agents/copilotcli/node/copilotcliSessionService';
-import { buildChatHistoryFromEvents, parseChatMessagesToEvents } from '../../agents/copilotcli/node/copilotcliToolInvocationFormatter';
+import { ExtendedChatRequest, ICopilotCLISessionService } from '../../agents/copilotcli/node/copilotcliSessionService';
+import { buildChatHistoryFromEvents, parseChatMessagesToEvents, stripSystemReminders } from '../../agents/copilotcli/node/copilotcliToolInvocationFormatter';
 
 export class CopilotCLIChatSessionItemProvider extends Disposable implements vscode.ChatSessionItemProvider {
 	private readonly _onDidChangeChatSessionItems = this._register(new Emitter<void>());
@@ -169,6 +169,7 @@ export class CopilotCLIChatSessionItemProvider extends Disposable implements vsc
 export class CopilotCLIChatSessionContentProvider implements vscode.ChatSessionContentProvider {
 
 	constructor(
+		private readonly copilotcliAgentManager: CopilotCLIAgentManager,
 		@ICopilotCLISessionService private readonly sessionService: ICopilotCLISessionService,
 	) { }
 
@@ -180,9 +181,39 @@ export class CopilotCLIChatSessionContentProvider implements vscode.ChatSessionC
 
 		const history = existingSession ? buildChatHistoryFromEvents(events) : [];
 
+		// Check if there's a pending request for this new session
+		const pendingRequest = this.sessionService.getPendingRequest(copilotcliSessionId);
+
+		const activeResponseCallback = pendingRequest
+			? async (stream: vscode.ChatResponseStream, token: vscode.CancellationToken) => {
+				this.sessionService.clearPendingRequest(copilotcliSessionId);
+				await this.copilotcliAgentManager.handleRequest(
+					copilotcliSessionId,
+					pendingRequest.request,
+					pendingRequest.context,
+					stream,
+					token
+				);
+			}
+			: undefined;
+
+		// If there's a pending request, add it to the history as the first request turn
+		if (pendingRequest) {
+			const request = pendingRequest.request;
+			const requestTurn = new vscode.ChatRequestTurn2(
+				stripSystemReminders(request.prompt),
+				undefined,
+				[...request.references],
+				'',
+				[],
+				undefined
+			);
+			history.push(requestTurn);
+		}
+
 		return {
 			history,
-			activeResponseCallback: undefined,
+			activeResponseCallback,
 			requestHandler: undefined,
 		};
 	}
@@ -192,6 +223,7 @@ export class CopilotCLIChatSessionParticipant {
 	constructor(
 		private readonly sessionType: string,
 		private readonly copilotcliAgentManager: CopilotCLIAgentManager,
+		private readonly sessionService: ICopilotCLISessionService,
 		private readonly sessionItemProvider: CopilotCLIChatSessionItemProvider,
 	) { }
 
@@ -202,23 +234,24 @@ export class CopilotCLIChatSessionParticipant {
 	private async handleRequest(request: vscode.ChatRequest, context: vscode.ChatContext, stream: vscode.ChatResponseStream, token: vscode.CancellationToken): Promise<vscode.ChatResult | void> {
 		// Resolve the prompt with references before processing
 		const resolvedPrompt = this.resolvePrompt(request);
-		const processedRequest = { ...request, prompt: resolvedPrompt };
+		const processedRequest: ExtendedChatRequest = { ...request, prompt: resolvedPrompt };
 
-		const create = async () => {
-			const { copilotcliSessionId } = await this.copilotcliAgentManager.handleRequest(undefined, processedRequest, context, stream, token);
-			if (!copilotcliSessionId) {
-				stream.warning(localize('copilotcli.failedToCreateSession', "Failed to create a new CopilotCLI session."));
-				return undefined;
-			}
-			return copilotcliSessionId;
-		};
 		const { chatSessionContext } = context;
 		if (chatSessionContext) {
 			if (chatSessionContext.isUntitled) {
-				const copilotcliSessionId = await create();
-				if (copilotcliSessionId) {
-					this.sessionItemProvider.swap(chatSessionContext.chatSessionItem, { id: copilotcliSessionId, label: processedRequest.prompt ?? 'CopilotCLI' });
+				// Create a new session (just get the session ID, don't run the request yet)
+				const newSdkSession = await this.sessionService.getOrCreateSDKSession(undefined, processedRequest.prompt);
+				const copilotcliSessionId = newSdkSession?.sessionId;
+				if (!copilotcliSessionId) {
+					stream.warning(localize('copilotcli.failedToCreateSession', "Failed to create a new CopilotCLI session."));
+					return {};
 				}
+
+				// Store the pending request that will be executed by activeResponseCallback
+				this.sessionService.setPendingRequest(copilotcliSessionId, processedRequest, context);
+
+				// Immediately swap to the new session (this will trigger provideChatSessionContent)
+				this.sessionItemProvider.swap(chatSessionContext.chatSessionItem, { id: copilotcliSessionId, label: processedRequest.prompt ?? 'CopilotCLI' });
 				return {};
 			}