Refactor Copilot CLI session request handling to ensure commits happen on last request (#4467)

* Refactor Copilot CLI session request handling to ensure commits happen on last request

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: DonJayamanne

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

@@ -1050,11 +1050,23 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 	/**
 	 * Map to track pending requests for untitled sessions.
 	 * Key = Untitled Session Id
-	 * Vlaue = Map of Request Id to the Promise of the request being handled
+	 * Value = Map of Request Id to the Promise of the request being handled
 	 * So if we have multiple requests (can happen when steering) for the same untitled session.
 	 */
 	private readonly pendingRequestsForUntitledSessions = new Map<string, Map<string, Promise<vscode.ChatResult | void>>>();
-	private readonly pendingCommitWorktreeChangesBySession = new Map<string, Promise<void>>();
+
+	/**
+	 * Tracks in-flight requests per session so we can coordinate worktree
+	 * commit / PR handling and cleanup.
+	 *
+	 * We generally cannot have parallel requests for the same session, but when
+	 * steering is involved there can be multiple requests in flight for a
+	 * single session (the original request continues running while steering
+	 * requests are processed). This map records all active requests for each
+	 * session so that any worktree-related actions are deferred until the last
+	 * in-flight request for that session has completed.
+	 */
+	private readonly pendingRequestBySession = new Map<string, Set<vscode.ChatRequest>>();
 
 	/**
 	 * Outer request handler that supports *yielding* for session steering.
@@ -1065,7 +1077,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 	 *    previous request (status is `InProgress` or `NeedsInput`).
 	 * 2. VS Code signals this by setting `context.yieldRequested = true` on the
 	 *    *previous* request's context object.
-	 * 3. This handler polls `context.yieldRequested` every 500 ms. Once detected
+	 * 3. This handler polls `context.yieldRequested` every 100 ms. Once detected
 	 *    the outer `Promise.race` resolves, returning control to VS Code so it
 	 *    can dispatch the new (steering) request.
 	 * 4. Crucially, the inner `handleRequestImpl` promise is **not** cancelled
@@ -1106,7 +1118,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		}
 	}
 
-	private sendTelemetryForHandleRequest(request: vscode.ChatRequest, context: vscode.ChatContext, result: vscode.ChatResult | void, error: unknown): void {
+	private sendTelemetryForHandleRequest(request: vscode.ChatRequest, context: vscode.ChatContext): void {
 		const { chatSessionContext } = context;
 		const hasChatSessionItem = String(!!chatSessionContext?.chatSessionItem);
 		let isUntitled = String(chatSessionContext?.isUntitled);
@@ -1138,11 +1150,12 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		let { chatSessionContext } = context;
 		const disposables = new DisposableStore();
 		let sessionId: string | undefined = undefined;
+		let sdkSessionId: string | undefined = undefined;
 		try {
 
 			const initialOptions = chatSessionContext?.initialSessionOptions;
 			if (initialOptions && chatSessionContext) {
-				if (initialOptions && initialOptions.length > 0) {
+				if (initialOptions.length > 0) {
 					const sessionResource = chatSessionContext.chatSessionItem.resource;
 					const sessionId = SessionIdForCLI.parse(sessionResource);
 					for (const opt of initialOptions) {
@@ -1186,7 +1199,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 				}
 			}
 
-			this.sendTelemetryForHandleRequest(request, context, undefined, undefined);
+			this.sendTelemetryForHandleRequest(request, context);
 
 			const [authInfo,] = await Promise.all([this.copilotCLISDK.getAuthInfo().catch((ex) => this.logService.error(ex, 'Authorization failed')), this.lockRepoOptionForSession(context, token)]);
 			if (!authInfo) {
@@ -1242,7 +1255,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 				}
 				return {};
 			}
-
+			sdkSessionId = session.object.sessionId;
 			this.copilotCLIAgents.trackSessionAgent(session.object.sessionId, agent?.name);
 			if (isUntitled && !this.useController) {
 				disposables.add(toDisposable(() => this.sessionItemProvider.untitledSessionIdMapping.delete(session.object.sessionId)));
@@ -1259,6 +1272,10 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 				// No need to track the temproary workspace folders as its been persisted.
 				this.folderRepositoryManager.deleteUntitledSessionFolder(id);
 			}
+			const requestsForSession = this.pendingRequestBySession.get(session.object.sessionId) ?? new Set<vscode.ChatRequest>();
+			requestsForSession.add(request);
+			this.pendingRequestBySession.set(session.object.sessionId, requestsForSession);
+
 			// Check if we have context stored for this request (created in createCLISessionAndSubmitRequest, work around)
 			const contextForRequest = this.contextForRequest.get(session.object.sessionId);
 			this.contextForRequest.delete(session.object.sessionId);
@@ -1269,22 +1286,22 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 				const { prompt, attachments } = contextForRequest;
 				this.contextForRequest.delete(session.object.sessionId);
 				await session.object.handleRequest(request, { prompt }, attachments, model, authInfo, token);
-				await this.commitWorktreeChangesIfNeeded(session.object, token);
+				await this.commitWorktreeChangesIfNeeded(request, session.object, token);
 			} else if (request.command && !request.prompt && !isUntitled) {
 				const input = (copilotCLICommands as readonly string[]).includes(request.command)
 					? { command: request.command as CopilotCLICommand }
 					: { prompt: `/${request.command}` };
 				await session.object.handleRequest(request, input, [], model, authInfo, token);
-				await this.commitWorktreeChangesIfNeeded(session.object, token);
+				await this.commitWorktreeChangesIfNeeded(request, session.object, token);
 			} else if (request.prompt && Object.values(builtinSlashSCommands).includes(request.prompt)) {
 				await session.object.handleRequest(request, { prompt: request.prompt }, [], model, authInfo, token);
-				await this.commitWorktreeChangesIfNeeded(session.object, token);
+				await this.commitWorktreeChangesIfNeeded(request, session.object, token);
 			} else {
 				// Construct the full prompt with references to be sent to CLI.
 				const plan = request.modeInstructions2 ? isCopilotCLIPlanAgent(request.modeInstructions2) : false;
 				const { prompt, attachments } = await this.promptResolver.resolvePrompt(request, undefined, [], session.object.workspace, [], token);
 				await session.object.handleRequest(request, { prompt, plan }, attachments, model, authInfo, token);
-				await this.commitWorktreeChangesIfNeeded(session.object, token);
+				await this.commitWorktreeChangesIfNeeded(request, session.object, token);
 			}
 
 			if (isUntitled && !token.isCancellationRequested) {
@@ -1322,6 +1339,15 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 			throw ex;
 		}
 		finally {
+			if (sdkSessionId) {
+				const requestsForSession = this.pendingRequestBySession.get(sdkSessionId);
+				if (requestsForSession) {
+					requestsForSession.delete(request);
+					if (requestsForSession.size === 0) {
+						this.pendingRequestBySession.delete(sdkSessionId);
+					}
+				}
+			}
 			if (chatSessionContext?.chatSessionItem.resource) {
 				this.sessionItemProvider.notifySessionsChange();
 			}
@@ -1408,17 +1434,21 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		}
 	}
 
-	private async commitWorktreeChangesIfNeeded(session: ICopilotCLISession, token: vscode.CancellationToken): Promise<void> {
-		if (token.isCancellationRequested) {
+	private async commitWorktreeChangesIfNeeded(request: vscode.ChatRequest, session: ICopilotCLISession, token: vscode.CancellationToken): Promise<void> {
+		const pendingRequests = this.pendingRequestBySession.get(session.sessionId);
+		if (pendingRequests && pendingRequests.size > 1) {
+			// We still have pending requests for this session, which means the user has done some steering.
+			// Wait for all requests to complete, the last request to complete will handle the commit.
+			pendingRequests.delete(request);
 			return;
 		}
 
-		const existingPromise = this.pendingCommitWorktreeChangesBySession.get(session.sessionId);
-		if (existingPromise) {
-			return existingPromise;
+		if (token.isCancellationRequested) {
+			pendingRequests?.delete(request);
+			return;
 		}
 
-		const commitPromise = (async () => {
+		try {
 			if (session.status === vscode.ChatSessionStatus.Completed) {
 				const workingDirectory = getWorkingDirectory(session.workspace);
 				if (isIsolationEnabled(session.workspace)) {
@@ -1434,14 +1464,9 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 			}
 
 			await this.handlePullRequestCreated(session);
-		})().finally(() => {
-			if (this.pendingCommitWorktreeChangesBySession.get(session.sessionId) === commitPromise) {
-				this.pendingCommitWorktreeChangesBySession.delete(session.sessionId);
-			}
-		});
-
-		this.pendingCommitWorktreeChangesBySession.set(session.sessionId, commitPromise);
-		return commitPromise;
+		} finally {
+			pendingRequests?.delete(request);
+		}
 	}
 
 	private async handlePullRequestCreated(session: ICopilotCLISession): Promise<void> {
@@ -1698,7 +1723,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 			// Now that we've delegated it to a session, we can get out of here.
 			// Else if the request takes say 10 minutes, the caller would be blocked for that long.
 			session.object.handleRequest(request, { prompt }, attachments, model, authInfo, token)
-				.then(() => this.commitWorktreeChangesIfNeeded(session.object, token))
+				.then(() => this.commitWorktreeChangesIfNeeded(request, session.object, token))
 				.catch(error => {
 					this.logService.error(`Failed to handle CLI session request: ${error}`);
 					// Optionally: stream.error(error) to notify the user

src/extension/chatSessions/vscode-node/test/copilotCLIChatSessionParticipant.spec.ts

@@ -216,6 +216,10 @@ class TestCopilotCLISession extends CopilotCLISession {
 	public requests: Array<{ input: CopilotCLISessionInput; attachments: Attachment[]; modelId: string | undefined; authInfo: NonNullable<SessionOptions['authInfo']>; token: vscode.CancellationToken }> = [];
 	public static nextHandleRequestResult: Promise<void> | undefined;
 	public static handleRequestHook: ((request: { id: string; toolInvocationToken: vscode.ChatParticipantToolToken }, input: CopilotCLISessionInput) => Promise<void>) | undefined;
+	public static statusOverride?: vscode.ChatSessionStatus;
+	override get status(): vscode.ChatSessionStatus | undefined {
+		return TestCopilotCLISession.statusOverride;
+	}
 	override handleRequest(request: { id: string; toolInvocationToken: vscode.ChatParticipantToolToken }, input: CopilotCLISessionInput, attachments: Attachment[], modelId: string | undefined, authInfo: NonNullable<SessionOptions['authInfo']>, token: vscode.CancellationToken): Promise<void> {
 		this.requests.push({ input, attachments, modelId, authInfo, token });
 		if (TestCopilotCLISession.handleRequestHook) {
@@ -270,6 +274,7 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 		cliSessions.length = 0;
 		TestCopilotCLISession.nextHandleRequestResult = undefined;
 		TestCopilotCLISession.handleRequestHook = undefined;
+		TestCopilotCLISession.statusOverride = undefined;
 		// By default, simulate the command not being available so that
 		// handleDelegationFromAnotherChat falls into its catch block and
 		// calls handleRequest directly. The workaround tests override this.
@@ -541,6 +546,59 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 		await handlerPromise;
 	});
 
+	it('defers worktree handleRequestCompleted until all steering requests complete', async () => {
+		// Use an existing (non-untitled) session so both concurrent requests are guaranteed to
+		// resolve to the same SDK session and share the same pendingRequestBySession entry.
+		const sessionId = 'existing-worktree-session';
+		const sdkSession = new MockCliSdkSession(sessionId, new Date());
+		manager.sessions.set(sessionId, sdkSession);
+
+		const worktreeProperties = {
+			autoCommit: true,
+			baseCommit: 'deadbeef',
+			branchName: 'test',
+			repositoryPath: `${sep}repo`,
+			worktreePath: `${sep}worktree`,
+			version: 1
+		} satisfies ChatSessionWorktreeProperties;
+		// FolderRepositoryManagerImpl.getFolderRepository checks worktreeService.getWorktreeProperties(sessionId)
+		// when the session ID is not an untitled ID.
+		(worktree.getWorktreeProperties as unknown as ReturnType<typeof vi.fn>).mockResolvedValue(worktreeProperties);
+		// Simulate the session completing so the worktree commit path runs
+		TestCopilotCLISession.statusOverride = vscode.ChatSessionStatus.Completed;
+
+		let resolveFirst!: () => void;
+		const firstDeferred = new Promise<void>(resolve => { resolveFirst = resolve; });
+		let resolveSecond!: () => void;
+		const secondDeferred = new Promise<void>(resolve => { resolveSecond = resolve; });
+
+		TestCopilotCLISession.handleRequestHook = vi.fn((_request, input) => {
+			if (input.prompt === 'First') { return firstDeferred; }
+			return secondDeferred;
+		});
+
+		const context = createChatContext(sessionId, false);
+		const stream = new MockChatResponseStream();
+
+		const firstRequest = new TestChatRequest('First');
+		const firstToken = disposables.add(new CancellationTokenSource()).token;
+		const firstPromise = participant.createHandler()(firstRequest, context, stream, firstToken);
+
+		const secondRequest = new TestChatRequest('Second');
+		const secondToken = disposables.add(new CancellationTokenSource()).token;
+		const secondPromise = participant.createHandler()(secondRequest, context, stream, secondToken);
+
+		// Second (steering) request completes first — commit must NOT fire while first is still pending
+		resolveSecond();
+		await secondPromise;
+		expect(worktree.handleRequestCompleted).not.toHaveBeenCalled();
+
+		// First request completes last — commit fires exactly once
+		resolveFirst();
+		await firstPromise;
+		expect(worktree.handleRequestCompleted).toHaveBeenCalledTimes(1);
+	});
+
 	it('defers untitled session swap while a steering request is still pending', async () => {
 		(itemProvider.isNewSession as ReturnType<typeof vi.fn>).mockImplementation((sessionId: string) => sessionId.startsWith('untitled:'));
 		let resolveFirstRequest!: () => void;
@@ -1290,7 +1348,7 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 			);
 			expect(lockCalls.length).toBeGreaterThan(0);
 
-			// Verify unlock was called on failure
+			// Verify unlock was NOT called on failure (session creation failed but workspace was trusted)
 			const unlockCalls = allCalls.filter(
 				call => call[1].some((update: any) => update.optionId === 'repository' && typeof update.value === 'string')
 			);