Add dropdown to create 1-4 concurrent PR variants from cloud agent sessions

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

@@ -22,7 +22,9 @@ type ConfirmationResult = { step: string; accepted: boolean; metadata?: CreatePr
 interface CreatePromptMetadata {
 	prompt: string;
 	history?: string;
-	references?: vscode.ChatPromptReference[];
+	references?: readonly vscode.ChatPromptReference[];
+	variationsCount?: number;
+	originalSessionItem?: vscode.ChatSessionItem;
 }
 
 interface UncommittedChangesMetadata {
@@ -76,6 +78,8 @@ export interface ICommentResult {
 
 const AGENTS_OPTION_GROUP_ID = 'agents';
 const DEFAULT_AGENT_ID = '___vscode_default___';
+const VARIATIONS_OPTION_GROUP_ID = 'variations';
+const DEFAULT_VARIATIONS_COUNT = '1';
 
 export class CopilotChatSessionsProvider extends Disposable implements vscode.ChatSessionContentProvider, vscode.ChatSessionItemProvider {
 	public static readonly TYPE = 'copilot-cloud-agent';
@@ -87,6 +91,7 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 	private chatSessions: Map<number, PullRequestSearchItem> = new Map();
 	private chatSessionItemsPromise: Promise<vscode.ChatSessionItem[]> | undefined;
 	private sessionAgentMap: Map<Uri, string> = new Map();
+	private sessionVariationsMap: Map<string, string> = new Map(); // Use string keys (Uri.toString()) for proper equality comparison
 	public chatParticipant = vscode.chat.createChatParticipant(CopilotChatSessionsProvider.TYPE, async (request, context, stream, token) =>
 		await this.chatParticipantImpl(request, context, stream, token)
 	);
@@ -112,6 +117,17 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 			return { optionGroups: [] };
 		}
 
+		// Always provide variations option group if we have a valid repo
+		const variationItems: vscode.ChatSessionProviderOptionItem[] = [
+			{ id: '1', name: vscode.l10n.t('1 variant') },
+			{ id: '2', name: vscode.l10n.t('2 variants') },
+			{ id: '3', name: vscode.l10n.t('3 variants') },
+			{ id: '4', name: vscode.l10n.t('4 variants') }
+		];
+
+		const optionGroups: vscode.ChatSessionProviderOptionGroup[] = [];
+
+		// Try to fetch custom agents, but don't fail if this errors
 		try {
 			const customAgents = await this._octoKitService.getCustomAgents(repoId.org, repoId.repo);
 			const agentItems: vscode.ChatSessionProviderOptionItem[] = [
@@ -121,23 +137,31 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 					name: agent.display_name || agent.name
 				}))
 			];
-			return {
-				optionGroups: [
-					{
-						id: AGENTS_OPTION_GROUP_ID,
-						name: vscode.l10n.t('Custom Agents'),
-						description: vscode.l10n.t('Select which agent to use'),
-						items: agentItems,
-					}
-				]
-			};
+
+			optionGroups.push({
+				id: AGENTS_OPTION_GROUP_ID,
+				name: vscode.l10n.t('Custom Agents'),
+				description: vscode.l10n.t('Select which agent to use'),
+				items: agentItems,
+			});
 		} catch (error) {
 			this.logService.error(`Error fetching custom agents: ${error}`);
-			return { optionGroups: [] };
+			// Continue without custom agents option group
 		}
+
+		// Always add variations option group
+		optionGroups.push({
+			id: VARIATIONS_OPTION_GROUP_ID,
+			name: vscode.l10n.t('PR Variations'),
+			description: vscode.l10n.t('Number of PR variants to generate'),
+			items: variationItems,
+		});
+
+		return { optionGroups };
 	}
 
 	provideHandleOptionsChange(resource: Uri, updates: ReadonlyArray<vscode.ChatSessionOptionUpdate>, token: vscode.CancellationToken): void {
+		this.logService.info(`[VARIANTS DEBUG] provideHandleOptionsChange called - Resource: ${resource}, Updates: ${JSON.stringify(updates)}`);
 		for (const update of updates) {
 			if (update.optionId === AGENTS_OPTION_GROUP_ID) {
 				if (update.value) {
@@ -147,6 +171,14 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 					this.sessionAgentMap.delete(resource);
 					this.logService.info(`Agent cleared for session ${resource}`);
 				}
+			} else if (update.optionId === VARIATIONS_OPTION_GROUP_ID) {
+				if (update.value) {
+					this.sessionVariationsMap.set(resource.toString(), update.value);
+					this.logService.info(`[VARIANTS DEBUG] Variations changed for session ${resource}: ${update.value}`);
+				} else {
+					this.sessionVariationsMap.delete(resource.toString());
+					this.logService.info(`Variations cleared for session ${resource}`);
+				}
 			}
 		}
 	}
@@ -294,9 +326,14 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 			// Query for the sub-agent that the remote reports for this session
 			|| undefined; /* TODO: Needs API to support this. */
 
+		const selectedVariations = this.sessionVariationsMap.get(resource.toString()) || DEFAULT_VARIATIONS_COUNT;
+
 		return {
 			history,
-			options: selectedAgent ? { [AGENTS_OPTION_GROUP_ID]: selectedAgent } : undefined,
+			options: {
+				...(selectedAgent ? { [AGENTS_OPTION_GROUP_ID]: selectedAgent } : {}),
+				[VARIATIONS_OPTION_GROUP_ID]: selectedVariations
+			},
 			activeResponseCallback: this.findActiveResponseCallback(sessions, pr),
 			requestHandler: undefined
 		};
@@ -350,14 +387,18 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 
 	private createEmptySession(resource: Uri): vscode.ChatSession {
 		const sessionId = resource ? resource.path.slice(1) : undefined;
+		const variationsValue = this.sessionVariationsMap.get(resource.toString());
+		this.logService.info(`[VARIANTS DEBUG] createEmptySession called - Resource: ${resource}, sessionId: ${sessionId}, variationsValue: ${variationsValue}`);
+
 		return {
 			history: [],
 			...(sessionId && sessionId.startsWith('untitled-')
 				? {
 					options: {
 						[AGENTS_OPTION_GROUP_ID]:
-							this.sessionAgentMap.get(resource)
-							?? (this.sessionAgentMap.set(resource, DEFAULT_AGENT_ID), DEFAULT_AGENT_ID)
+							this.sessionAgentMap.get(resource) || DEFAULT_AGENT_ID,
+						[VARIATIONS_OPTION_GROUP_ID]:
+							variationsValue || DEFAULT_VARIATIONS_COUNT
 					}
 				}
 				: {}),
@@ -477,40 +518,135 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 	}
 
 	async createDelegatedChatSession(metadata: CreatePromptMetadata, stream: vscode.ChatResponseStream, token: vscode.CancellationToken): Promise<PullRequestInfo | undefined> {
-		const { prompt, history, references } = metadata;
-		const number = await this.startSession(stream, token, 'chat', prompt, history, references);
-		if (!number) {
-			return undefined;
+		const { prompt, history, references, variationsCount = 1, originalSessionItem } = metadata;
+
+		this.logService.info(`[VARIANTS DEBUG] createDelegatedChatSession called - variationsCount: ${variationsCount}`);
+
+		// Notify user about creating multiple variants
+		if (variationsCount > 1) {
+			this.logService.info(`[VARIANTS DEBUG] Creating ${variationsCount} PR variants...`);
+			stream.markdown(vscode.l10n.t('Creating {0} PR variants...', variationsCount));
+			stream.markdown('\n\n');
 		}
-		const pullRequest = await this.findPR(number);
-		if (!pullRequest) {
-			stream.warning(vscode.l10n.t('Could not find the associated pull request {0} for this chat session.', number));
-			return undefined;
+
+		const prInfos: PullRequestInfo[] = [];
+
+		// Create each variant
+		for (let i = 0; i < variationsCount; i++) {
+			if (token.isCancellationRequested) {
+				stream.warning(vscode.l10n.t('PR creation cancelled.'));
+				break;
+			}
+
+			// Add variant indicator to prompt if creating multiple variants
+			const variantPrompt = variationsCount > 1
+				? `${prompt}\n\n[Variant ${i + 1} of ${variationsCount}]`
+				: prompt;
+
+			stream.progress(vscode.l10n.t('Creating variant {0} of {1}', i + 1, variationsCount));
+
+			const number = await this.startSession(stream, token, 'chat', variantPrompt, history, references);
+			if (!number) {
+				stream.warning(vscode.l10n.t('Failed to create variant {0}', i + 1));
+				continue;
+			}
+
+			const pullRequest = await this.findPR(number);
+			if (!pullRequest) {
+				stream.warning(vscode.l10n.t('Could not find the associated pull request {0} for variant {1}.', number, i + 1));
+				continue;
+			}
+
+			const uri = await toOpenPullRequestWebviewUri({ owner: pullRequest.repository.owner.login, repo: pullRequest.repository.name, pullRequestNumber: pullRequest.number });
+			const card = new vscode.ChatResponsePullRequestPart(uri, pullRequest.title, pullRequest.body, getAuthorDisplayName(pullRequest.author), `#${pullRequest.number}`);
+			stream.push(card);
+
+			const prInfo = {
+				uri: uri.toString(),
+				title: pullRequest.title,
+				description: pullRequest.body,
+				author: getAuthorDisplayName(pullRequest.author),
+				linkTag: `#${pullRequest.number}`
+			};
+			prInfos.push(prInfo);
 		}
 
-		const uri = await toOpenPullRequestWebviewUri({ owner: pullRequest.repository.owner.login, repo: pullRequest.repository.name, pullRequestNumber: pullRequest.number });
-		const card = new vscode.ChatResponsePullRequestPart(uri, pullRequest.title, pullRequest.body, getAuthorDisplayName(pullRequest.author), `#${pullRequest.number}`);
-		stream.push(card);
-		stream.markdown(vscode.l10n.t('GitHub Copilot cloud agent has begun working on your request. Follow its progress in the associated chat and pull request.'));
-		await vscode.commands.executeCommand('vscode.open', vscode.Uri.from({ scheme: CopilotChatSessionsProvider.TYPE, path: '/' + number }));
-		// Return PR info for embedding in session history
-		return {
-			uri: uri.toString(),
-			title: pullRequest.title,
-			description: pullRequest.body,
-			author: getAuthorDisplayName(pullRequest.author),
-			linkTag: `#${pullRequest.number}`
-		};
+		// After creating all PRs, fire the commit event to switch from untitled to first PR session
+		if (variationsCount > 1 && originalSessionItem && prInfos.length > 0) {
+			const firstPrNumber = parseInt(prInfos[0].linkTag.substring(1), 10);
+			this._onDidCommitChatSessionItem.fire({
+				original: originalSessionItem,
+				modified: {
+					resource: vscode.Uri.from({ scheme: CopilotChatSessionsProvider.TYPE, path: '/' + firstPrNumber }),
+					label: `Pull Request ${firstPrNumber}`
+				}
+			});
+		}
+
+		// Show summary message
+		if (prInfos.length === variationsCount) {
+			if (variationsCount > 1) {
+				stream.markdown(vscode.l10n.t('Successfully created {0} PR variants. GitHub Copilot cloud agent has begun working on your requests. Follow their progress in the associated chats and pull requests.', variationsCount));
+			} else {
+				stream.markdown(vscode.l10n.t('GitHub Copilot cloud agent has begun working on your request. Follow its progress in the associated chat and pull request.'));
+			}
+		} else if (prInfos.length > 0) {
+			stream.warning(vscode.l10n.t('Created {0} of {1} requested variants.', prInfos.length, variationsCount));
+		}
+
+		// Open the first PR - but only for single variant or after all variants created
+		// For multiple variants, don't auto-open to avoid interrupting the user experience
+		if (prInfos.length > 0 && variationsCount === 1) {
+			const firstPrNumber = parseInt(prInfos[0].linkTag.substring(1), 10);
+			await vscode.commands.executeCommand('vscode.open', vscode.Uri.from({ scheme: CopilotChatSessionsProvider.TYPE, path: '/' + firstPrNumber }));
+			return prInfos[0];
+		}
+
+		return prInfos.length > 0 ? prInfos[0] : undefined;
 	}
 
 	private async chatParticipantImpl(request: vscode.ChatRequest, context: vscode.ChatContext, stream: vscode.ChatResponseStream, token: vscode.CancellationToken) {
+		this.logService.info(`[VARIANTS DEBUG] chatParticipantImpl called - hasSessionContext: ${!!context.chatSessionContext}, isUntitled: ${context.chatSessionContext?.isUntitled}, hasConfirmation: ${!!(request.acceptedConfirmationData || request.rejectedConfirmationData)}`);
+
 		if (request.acceptedConfirmationData || request.rejectedConfirmationData) {
 			return await this.handleConfirmationData(request, stream, token);
 		}
 
 		if (context.chatSessionContext?.isUntitled) {
 			/* Generate new cloud agent session from an 'untitled' session */
 			const selectedAgent = this.sessionAgentMap.get(context.chatSessionContext.chatSessionItem.resource);
+			const variationsCount = parseInt(this.sessionVariationsMap.get(context.chatSessionContext.chatSessionItem.resource.toString()) || DEFAULT_VARIATIONS_COUNT, 10);
+
+			// Debug logging
+			this.logService.info(`[VARIANTS DEBUG] Untitled session - Resource: ${context.chatSessionContext.chatSessionItem.resource}, Variations from map: ${this.sessionVariationsMap.get(context.chatSessionContext.chatSessionItem.resource.toString())}, Parsed count: ${variationsCount}`);
+			this.logService.info(`[VARIANTS DEBUG] Map contents: ${JSON.stringify(Array.from(this.sessionVariationsMap.entries()).map(([k, v]) => ({ key: k, value: v })))}`);
+			this.logService.info(`[VARIANTS DEBUG] Resource toString: ${context.chatSessionContext.chatSessionItem.resource.toString()}, path: ${context.chatSessionContext.chatSessionItem.resource.path}`);
+
+			// For untitled sessions with multiple variants, show confirmation first
+			if (variationsCount > 1) {
+				this.logService.info(`[VARIANTS DEBUG] Showing confirmation for ${variationsCount} variants`);
+				// Show confirmation modal with premium request cost warning
+				const confirmationDetails = vscode.l10n.t('The agent will work asynchronously to create {0} pull request variants with your requested changes. This will use {0} premium requests.', variationsCount);
+
+				stream.confirmation(
+					vscode.l10n.t('Delegate to cloud agent'),
+					confirmationDetails,
+					{
+						step: 'create',
+						metadata: {
+							prompt: context.chatSummary?.prompt ?? request.prompt,
+							history: context.chatSummary?.history,
+							references: request.references,
+							variationsCount,
+							originalSessionItem: context.chatSessionContext.chatSessionItem
+						}
+					},
+					['Delegate', 'Cancel']
+				);
+				return {};
+			}
+
+			// Single variant - use the simple flow
 			const number = await this.startSession(
 				stream,
 				token,
@@ -545,6 +681,30 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 					return {};
 				}
 
+				// Check if user wants to create multiple variants from this session
+				const variationsCount = parseInt(this.sessionVariationsMap.get(context.chatSessionContext.chatSessionItem.resource.toString()) || DEFAULT_VARIATIONS_COUNT, 10);
+
+				if (variationsCount > 1) {
+					// Show confirmation for creating multiple variants from existing session
+					const confirmationDetails = vscode.l10n.t('The agent will work asynchronously to create {0} pull request variants with your requested changes. This will use {0} premium requests.', variationsCount);
+
+					stream.confirmation(
+						vscode.l10n.t('Create PR variants'),
+						confirmationDetails,
+						{
+							step: 'create',
+							metadata: {
+								prompt: context.chatSummary?.prompt ?? request.prompt,
+								history: context.chatSummary?.history,
+								references: request.references,
+								variationsCount,
+							}
+						},
+						['Create', 'Cancel']
+					);
+					return {};
+				}
+
 				stream.progress(vscode.l10n.t('Preparing'));
 				const session = SessionIdForPr.parse(context.chatSessionContext.chatSessionItem.resource);
 				let prNumber = session?.prNumber;
@@ -594,15 +754,26 @@ export class CopilotChatSessionsProvider extends Disposable implements vscode.Ch
 			}
 		} else {
 			/* @copilot invoked from a 'normal' chat or 'cloud button' */
+			// Get variations count - default to 1 if not set (for non-untitled sessions)
+			// Since we're in a normal chat without session context, default to 1
+			const variationsCount = 1;
+
+			// Construct confirmation message
+			let confirmationDetails = this.DELEGATE_MODAL_DETAILS;
+			if (variationsCount > 1) {
+				confirmationDetails = vscode.l10n.t('The agent will work asynchronously to create {0} pull request variants with your requested changes. This will use {0} premium requests.', variationsCount);
+			}
+
 			stream.confirmation(
 				vscode.l10n.t('Delegate to cloud agent'),
-				this.DELEGATE_MODAL_DETAILS,
+				confirmationDetails,
 				{
 					step: 'create',
 					metadata: {
 						prompt: context.chatSummary?.prompt ?? request.prompt,
 						history: context.chatSummary?.history,
 						references: request.references,
+						variationsCount,
 					}
 				},
 				['Delegate', 'Cancel']