Refactor cloud agent (#2097)

* refactor cloud session class

* pass correct context

* throw errors and fix parameter typo

* remove stale test

* remove debug markdown

* reset confirmations commadn

* some loose ends

* tweaks while rate limited

* update test

* more accurate

* nevermind

* add back cli handleConfirmationData

* nullptr

* restore test

* polish

* this is the pr uri

* polish

extensions/copilot/package.json

@@ -2412,6 +2412,10 @@
 				"title": "%github.copilot.command.refreshAgentSessions%",
 				"icon": "$(refresh)"
 			},
+			{
+				"command": "github.copilot.cloud.resetWorkspaceConfirmations",
+				"title": "%github.copilot.command.resetCloudAgentWorkspaceConfirmations%"
+			},
 			{
 				"command": "github.copilot.cloud.sessions.openInBrowser",
 				"title": "%github.copilot.command.openCopilotAgentSessionsInBrowser%",
@@ -5245,4 +5249,4 @@
 		"string_decoder": "npm:string_decoder@1.2.0",
 		"node-gyp": "npm:node-gyp@10.3.1"
 	}
-}
+}

extensions/copilot/package.nls.json

@@ -139,6 +139,7 @@
 	"github.copilot.command.explainTerminalLastCommand": "Explain Last Terminal Command",
 	"github.copilot.command.collectWorkspaceIndexDiagnostics": "Collect Workspace Index Diagnostics",
 	"github.copilot.command.triggerPermissiveSignIn": "Login to GitHub with Full Permissions",
+	"github.copilot.command.resetCloudAgentWorkspaceConfirmations": "Reset Cloud Agent Workspace Confirmations",
 	"github.copilot.git.generateCommitMessage": "Generate Commit Message",
 	"github.copilot.git.resolveMergeConflicts": "Resolve Conflicts with AI",
 	"github.copilot.devcontainer.generateDevContainerConfig": "Generate Dev Container Configuration",

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

@@ -152,6 +152,11 @@ export class ChatSessionsContrib extends Disposable implements IExtensionContrib
 				cloudSessionsProvider.refresh();
 			})
 		);
+		this.copilotCloudRegistrations.add(
+			vscode.commands.registerCommand('github.copilot.cloud.resetWorkspaceConfirmations', () => {
+				cloudSessionsProvider.resetWorkspaceContext();
+			})
+		);
 		this.copilotCloudRegistrations.add(
 			vscode.commands.registerCommand('github.copilot.cloud.sessions.openInBrowser', async (chatSessionItem: vscode.ChatSessionItem) => {
 				cloudSessionsProvider.openSessionsInBrowser(chatSessionItem);

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

@@ -26,11 +26,19 @@ import { PermissionRequest, requestPermission } from '../../agents/copilotcli/no
 import { ChatSummarizerProvider } from '../../prompt/node/summarizer';
 import { IToolsService } from '../../tools/common/toolsService';
 import { ICopilotCLITerminalIntegration } from './copilotCLITerminalIntegration';
-import { ConfirmationResult, CopilotCloudSessionsProvider, UncommittedChangesStep } from './copilotCloudSessionsProvider';
+import { CopilotCloudSessionsProvider } from './copilotCloudSessionsProvider';
 
 const MODELS_OPTION_ID = 'model';
 const ISOLATION_OPTION_ID = 'isolation';
 
+const UncommittedChangesStep = 'uncommitted-changes';
+type ConfirmationResult = { step: string; accepted: boolean; metadata?: CLIConfirmationMetadata };
+interface CLIConfirmationMetadata {
+	prompt: string;
+	references?: readonly vscode.ChatPromptReference[];
+	chatContext: vscode.ChatContext;
+}
+
 // Track model selections per session
 // TODO@rebornix: we should have proper storage for the session model preference (revisit with API)
 const _sessionModel: Map<string, vscode.ChatSessionProviderOptionItem | undefined> = new Map();
@@ -414,7 +422,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 			}
 
 			if (!isUntitled && confirmationResults.length) {
-				return await this.handleConfirmationData(session.object, request.prompt, confirmationResults, context, stream, token);
+				return await this.handleConfirmationData(request, session.object, request.prompt, confirmationResults, context, stream, token);
 			}
 
 			if (request.prompt.startsWith('/delegate')) {
@@ -499,18 +507,12 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		}
 
 		const prompt = request.prompt.substring('/delegate'.length).trim();
-		if (!await this.cloudSessionProvider.tryHandleUncommittedChanges({
-			prompt: prompt,
-			chatContext: context
-		}, stream, token)) {
-			const prInfo = await this.cloudSessionProvider.createDelegatedChatSession({
-				prompt,
-				chatContext: context
-			}, stream, token);
-			if (prInfo) {
-				await this.recordPushToSession(session, request.prompt, prInfo);
-			}
+
+		const prInfo = await this.cloudSessionProvider.delegate(request, stream, context, token, { prompt, chatContext: context });
+		if (prInfo) {
+			await this.recordPushToSession(session, request.prompt, prInfo);
 		}
+
 	}
 
 	private getAcceptedRejectedConfirmationData(request: vscode.ChatRequest): ConfirmationResult[] {
@@ -521,7 +523,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		return results;
 	}
 
-	private async handleConfirmationData(session: ICopilotCLISession, prompt: string, results: ConfirmationResult[], context: vscode.ChatContext, stream: vscode.ChatResponseStream, token: vscode.CancellationToken) {
+	private async handleConfirmationData(request: vscode.ChatRequest, session: ICopilotCLISession, prompt: string, results: ConfirmationResult[], context: vscode.ChatContext, stream: vscode.ChatResponseStream, token: vscode.CancellationToken) {
 		const uncommittedChangesData = results.find(data => data.step === UncommittedChangesStep);
 		if (!uncommittedChangesData) {
 			stream.warning(`Unknown confirmation step: ${results.map(r => r.step).join(', ')}\n\n`);
@@ -533,12 +535,7 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 			return {};
 		}
 
-		const prInfo = await this.cloudSessionProvider?.createDelegatedChatSession({
-			prompt: uncommittedChangesData.metadata.prompt,
-			references: uncommittedChangesData.metadata.references,
-			autoPushAndCommit: uncommittedChangesData.metadata.autoPushAndCommit,
-			chatContext: context
-		}, stream, token);
+		const prInfo = await this.cloudSessionProvider?.delegate(request, stream, context, token, uncommittedChangesData.metadata);
 		if (prInfo) {
 			await this.recordPushToSession(session, prompt, prInfo);
 		}
@@ -580,13 +577,13 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 	private async recordPushToSession(
 		session: ICopilotCLISession,
 		userPrompt: string,
-		prInfo: { uri: string; title: string; description: string; author: string; linkTag: string }
+		prInfo: { uri: vscode.Uri; title: string; description: string; author: string; linkTag: string }
 	): Promise<void> {
 		// Add user message event
 		session.addUserMessage(userPrompt);
 
 		// Add assistant message event with embedded PR metadata
-		const assistantMessage = `GitHub Copilot cloud agent has begun working on your request. Follow its progress in the associated chat and pull request.\n<pr_metadata uri="${prInfo.uri}" title="${escapeXml(prInfo.title)}" description="${escapeXml(prInfo.description)}" author="${escapeXml(prInfo.author)}" linkTag="${escapeXml(prInfo.linkTag)}"/>`;
+		const assistantMessage = `GitHub Copilot cloud agent has begun working on your request. Follow its progress in the associated chat and pull request.\n<pr_metadata uri="${prInfo.uri.toString()}" title="${escapeXml(prInfo.title)}" description="${escapeXml(prInfo.description)}" author="${escapeXml(prInfo.author)}" linkTag="${escapeXml(prInfo.linkTag)}"/>`;
 		session.addUserAssistantMessage(assistantMessage);
 	}
 }

extensions/copilot/src/extension/chatSessions/vscode-node/copilotCloudGitOperationsManager.ts

@@ -4,8 +4,12 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
+import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
+import { IGitExtensionService } from '../../../platform/git/common/gitExtensionService';
+import { IGitService } from '../../../platform/git/common/gitService';
 import { Repository } from '../../../platform/git/vscode/git';
 import { ILogService } from '../../../platform/log/common/logService';
+import { getRepoId } from '../vscode/copilotCodingAgentUtils';
 
 export interface GitRepoInfo {
 	repository: Repository;
@@ -14,10 +18,93 @@ export interface GitRepoInfo {
 }
 
 export class CopilotCloudGitOperationsManager {
-	constructor(private readonly logService: ILogService) { }
+	constructor(
+		private readonly logService: ILogService,
+		private readonly gitService: IGitService,
+		private readonly gitExtensionService: IGitExtensionService,
+		private readonly configurationService: IConfigurationService
+	) { }
 
-	async commitAndPushChanges(repoInfo: GitRepoInfo): Promise<string> {
-		const { repository, remoteName, baseRef } = repoInfo;
+	private get autoCommitAndPushEnabled(): boolean {
+		return this.configurationService.getConfig(ConfigKey.AgentDelegateAutoCommitAndPush);
+	}
+
+	async repoInfo(): Promise<GitRepoInfo> {
+		// TODO: support selecting remote
+		// await this.promptAndUpdatePreferredGitHubRemote(true);
+		const repoId = await getRepoId(this.gitService);
+		if (!repoId) {
+			throw new Error(vscode.l10n.t('Repository information is not available. Open a GitHub repository to continue with cloud agent.'));
+		}
+		const currentRepository = this.gitService.activeRepository.get();
+		if (!currentRepository) {
+			throw new Error(vscode.l10n.t('No active repository found. Open a GitHub repository to continue with cloud agent.'));
+		}
+		const git = this.gitExtensionService.getExtensionApi();
+		const repo = git?.getRepository(currentRepository?.rootUri);
+		// Checks if user has permission to access the repository
+		if (!repo) {
+			throw new Error(
+				vscode.l10n.t(
+					'Unable to access {0}. Please check your permissions and try again.',
+					`\`${repoId.org}/${repoId.repo}\``
+				)
+			);
+		}
+		return {
+			repository: repo,
+			remoteName: repo.state.HEAD?.upstream?.remote ?? currentRepository.upstreamRemote ?? repo.state.remotes?.[0]?.name ?? 'origin',
+			baseRef: currentRepository.headBranchName ?? 'main'
+		};
+	}
+
+	async validateRemoteHasBaseRef(stream: vscode.ChatResponseStream): Promise<void> {
+		const { repository, remoteName, baseRef } = await this.repoInfo();
+		stream.progress(vscode.l10n.t('Verifying branch \'{0}\' exists on remote \'{1}\'', baseRef, remoteName));
+		if (repository && remoteName && baseRef) {
+			try {
+				const remoteBranches =
+					(await repository.getBranches({ remote: true }))
+						.filter(b => b.remote); // Has an associated remote
+				const expectedRemoteBranch = `${remoteName}/${baseRef}`;
+				const alternateNames = new Set<string>([
+					expectedRemoteBranch,
+					`refs/remotes/${expectedRemoteBranch}`,
+					baseRef
+				]);
+				const hasRemoteBranch = remoteBranches.some(branch => {
+					if (!branch.name) {
+						return false;
+					}
+					if (branch.remote && branch.remote !== remoteName) {
+						return false;
+					}
+					const candidateName =
+						(branch.remote && branch.name.startsWith(branch.remote + '/'))
+							? branch.name
+							: `${branch.remote}/${branch.name}`;
+					return alternateNames.has(candidateName);
+				});
+
+				if (!hasRemoteBranch) {
+					if (this.autoCommitAndPushEnabled) {
+						this.logService.warn(`Base branch '${expectedRemoteBranch}' not found on remote. Auto-pushing because autoCommitAndPush is enabled.`);
+						stream.progress(vscode.l10n.t('Pushing branch \'{0}\'', baseRef));
+						await repository.push(remoteName, baseRef, true);
+					} else {
+						throw new Error('autoCommitAndPush is disabled');
+					}
+				}
+			} catch (error) {
+				this.logService.error(`Failed to verify remote branch for cloud agent: ${error instanceof Error ? error.message : String(error)}`);
+				throw new Error(vscode.l10n.t('Branch \'{0}\' does not exist on remote \'{1}\'. Push the branch manually or enable \'github.copilot.chat.agent.delegate.autoCommitAndPush\'', baseRef, remoteName));
+			}
+		}
+	}
+
+
+	async commitAndPushChanges(): Promise<string> {
+		const { repository, remoteName, baseRef } = await this.repoInfo();
 		const asyncBranch = await this.generateRandomBranchName(repository, 'copilot');
 
 		const commitMessage = vscode.l10n.t('Checkpoint from VS Code for cloud agent session');
@@ -25,7 +112,7 @@ export class CopilotCloudGitOperationsManager {
 			await repository.createBranch(asyncBranch, true);
 			await this.performCommit(asyncBranch, repository, commitMessage);
 			await repository.push(remoteName, asyncBranch, true);
-			this.showBranchSwitchNotification(repository, baseRef, asyncBranch);
+			await this.switchBackToBaseRef(repository, baseRef, asyncBranch);
 			return asyncBranch;
 		} catch (error) {
 			await this.rollbackToOriginalBranch(repository, baseRef);
@@ -41,6 +128,7 @@ export class CopilotCloudGitOperationsManager {
 				throw new Error(vscode.l10n.t('Uncommitted changes still detected.'));
 			}
 		} catch (error) {
+			// TODO: stream.progress('waiting for user to manually commit changes');
 			const commitSuccessful = await this.handleInteractiveCommit(repository);
 			if (!commitSuccessful) {
 				throw new Error(vscode.l10n.t('Failed to commit changes. Please commit or stash your changes manually before using the cloud agent.'));
@@ -115,17 +203,9 @@ export class CopilotCloudGitOperationsManager {
 		});
 	}
 
-	private showBranchSwitchNotification(repository: Repository, baseRef: string, newRef: string): void {
+	private async switchBackToBaseRef(repository: Repository, baseRef: string, newRef: string): Promise<void> {
 		if (repository.state.HEAD?.name !== baseRef) {
-			const SWAP_BACK_TO_ORIGINAL_BRANCH = vscode.l10n.t('Swap back to \'{0}\'', baseRef);
-			vscode.window.showInformationMessage(
-				vscode.l10n.t('Pending changes pushed to remote branch \'{0}\'.', newRef),
-				SWAP_BACK_TO_ORIGINAL_BRANCH,
-			).then(async (selection) => {
-				if (selection === SWAP_BACK_TO_ORIGINAL_BRANCH) {
-					await repository.checkout(baseRef);
-				}
-			});
+			await repository.checkout(baseRef);
 		}
 	}
 

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

@@ -71,8 +71,13 @@ class FakeGitService extends mock<IGitService>() {
 
 // Cloud provider fake for delegate scenario
 class FakeCloudProvider extends mock<CopilotCloudSessionsProvider>() {
-	override tryHandleUncommittedChanges = vi.fn(async () => false);
-	override createDelegatedChatSession = vi.fn(async () => ({ uri: 'pr://1', title: 'PR Title', description: 'Desc', author: 'Me', linkTag: 'tag' })) as unknown as CopilotCloudSessionsProvider['createDelegatedChatSession'];
+	override delegate = vi.fn(async () => ({
+		uri: vscode.Uri.parse('pr://1'),
+		title: 'PR Title',
+		description: 'PR Description',
+		author: 'Test Author',
+		linkTag: '#1'
+	})) as unknown as CopilotCloudSessionsProvider['delegate'];
 }
 
 
@@ -245,8 +250,7 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 		// Warning should appear (we emitted stream.warning). The mock stream only records markdown.
 		// Delegate path adds assistant PR metadata; ensure output contains PR metadata tag instead of relying on warning capture.
 		expect(sdkSession.emittedEvents[1].content).toMatch(/<pr_metadata uri="pr:\/\/1"/);
-		expect(cloudProvider.tryHandleUncommittedChanges).toHaveBeenCalled();
-		expect(cloudProvider.createDelegatedChatSession).toHaveBeenCalled();
+		expect(cloudProvider.delegate).toHaveBeenCalled();
 	});
 
 	it('handles /delegate command for new session', async () => {
@@ -261,8 +265,7 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 
 		expect(manager.sessions.size).toBe(1);
 		const sdkSession = Array.from(manager.sessions.values())[0];
-		expect(cloudProvider.tryHandleUncommittedChanges).toHaveBeenCalled();
-		expect(cloudProvider.createDelegatedChatSession).toHaveBeenCalled();
+		expect(cloudProvider.delegate).toHaveBeenCalled();
 		// PR metadata recorded
 		expect(sdkSession.emittedEvents.length).toBe(2);
 		expect(sdkSession.emittedEvents[0].event).toBe('user.message');
@@ -299,13 +302,13 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 		const sessionId = 'existing-confirm';
 		const sdkSession = new MockCliSdkSession(sessionId, new Date());
 		manager.sessions.set(sessionId, sdkSession);
-		const request = new TestChatRequest('Apply');
-		(request as any).acceptedConfirmationData = [{ step: 'uncommitted-changes', metadata: { prompt: 'delegate work' } }];
+		const request = new TestChatRequest('my prompt');
 		const context = createChatContext(sessionId, false);
+		(request as any).acceptedConfirmationData = [{ step: 'uncommitted-changes', metadata: { chatContext: context } }];
 		const stream = new MockChatResponseStream();
 		const token = disposables.add(new CancellationTokenSource()).token;
 		// Cloud provider will create delegated chat session returning prInfo
-		(cloudProvider.createDelegatedChatSession as unknown as ReturnType<typeof vi.fn>).mockResolvedValue({ uri: 'pr://2', title: 'T', description: 'D', author: 'A', linkTag: 'L' });
+		(cloudProvider.delegate as unknown as ReturnType<typeof vi.fn>).mockResolvedValue({ uri: 'pr://2', title: 'T', description: 'D', author: 'A', linkTag: 'L' });
 
 		await participant.createHandler()(request, context, stream, token);
 
@@ -317,7 +320,13 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 		expect(sdkSession.emittedEvents[1].event).toBe('assistant.message');
 		expect(sdkSession.emittedEvents[1].content).toContain('pr://2');
 		// Cloud provider used with provided metadata
-		expect(cloudProvider.createDelegatedChatSession).toHaveBeenCalledWith({ prompt: 'delegate work', chatContext: context }, expect.anything(), token);
+		expect(cloudProvider.delegate).toHaveBeenCalledWith(
+			request,
+			stream,
+			context,
+			token,
+			{ chatContext: context }
+		);
 	});
 
 	it('handleConfirmationData cancels when uncommitted-changes rejected', async () => {
@@ -334,7 +343,7 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 
 		// Should not record push or call delegate session
 		expect(sdkSession.emittedEvents.length).toBe(0);
-		expect(cloudProvider.createDelegatedChatSession).not.toHaveBeenCalled();
+		expect(cloudProvider.delegate).not.toHaveBeenCalled();
 		// Cancellation message markdown captured
 		expect(stream.output.some(o => /Cloud agent delegation request cancelled/i.test(o))).toBe(true);
 	});

extensions/copilot/src/extension/chatSessions/vscode/copilotCodingAgentUtils.ts

@@ -13,9 +13,6 @@ export const CONTINUE_TRUNCATION = vscode.l10n.t('Continue with truncation');
 export const body_suffix = vscode.l10n.t('Created from [VS Code](https://code.visualstudio.com/docs/copilot/copilot-coding-agent).');
 // https://github.com/github/sweagentd/blob/main/docs/adr/0001-create-job-api.md
 export const JOBS_API_VERSION = 'v1';
-type RemoteAgentSuccessResult = { link: string; state: 'success'; number: number; webviewUri: vscode.Uri; llmDetails: string; sessionId: string };
-type RemoteAgentErrorResult = { error: string; innerError?: string; state: 'error' };
-export type RemoteAgentResult = RemoteAgentSuccessResult | RemoteAgentErrorResult;
 
 /**
  * Truncation utility to ensure the problem statement sent to Copilot API is under the maximum length.

extensions/copilot/src/platform/github/common/githubAPI.ts

@@ -232,9 +232,13 @@ export async function makeSearchGraphQLRequest(
 		first
 	};
 
+	// TODO: Handle rate limiting
+	//       result.errors[0]
+	//         {type: 'RATE_LIMIT', code: 'graphql_rate_limit', message: 'API rate limit already exceeded for user ID xxxxxxx.'}
+
 	const result = await makeGitHubGraphQLRequest(fetcherService, logService, telemetry, host, query, token, variables);
 
-	return result ? result.data.search.nodes : [];
+	return result.data?.search?.nodes ?? [];
 }
 
 export async function getPullRequestFromGlobalId(