Agents - create branch when needed during a pull request creation (#333163)

* Agents - create branch when needed during a pull request creation

* Pull request feedback

Author: lszomoru

src/vs/platform/agentHost/common/agentHostGitService.ts

@@ -254,6 +254,8 @@ export interface IAgentHostGitService {
 	 * recreating the worktree.
 	 */
 	branchExists(repositoryRoot: URI, branchName: string): Promise<boolean>;
+	/** Creates a new branch and optionally checks it out while preserving the working tree. */
+	createBranch(workingDirectory: URI, branchName: string, options?: { readonly checkout?: boolean }): Promise<void>;
 	/**
 	 * Returns true when the working tree has any tracked, staged, or
 	 * untracked changes. Used by archive cleanup to skip removing a

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

@@ -356,6 +356,14 @@ export class AgentHostGitService implements IAgentHostGitService {
 		return output !== undefined;
 	}
 
+	async createBranch(workingDirectory: URI, branchName: string, options?: { readonly checkout?: boolean }): Promise<void> {
+		const args = options?.checkout
+			? ['checkout', '-q', '-b', branchName, '--no-track']
+			: ['branch', '-q', branchName];
+
+		await this._runGit(workingDirectory, args, { throwOnError: true });
+	}
+
 	async hasUncommittedChanges(workingDirectory: URI): Promise<boolean> {
 		const output = await this._runGitStatus(workingDirectory, ['--porcelain']);
 		return !!output && output.trim().length > 0;

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

@@ -18,6 +18,8 @@ import { type AutoMergeMethod, type CreatedPullRequest, IAgentHostOctoKitService
 import type { InvokeChangesetOperationParams, InvokeChangesetOperationResult } from '../common/state/protocol/channels-changeset/commands.js';
 import { ICopilotApiService, type ICopilotUtilityChatMessage } from './shared/copilotApiService.js';
 import { buildConversationContext } from '../common/agentHostConversationContext.js';
+import { IAgentBranchNameGenerator } from './shared/agentBranchNameGenerator.js';
+import { SessionConfigKey } from '../common/sessionConfigKeys.js';
 
 /**
  * Soft upper bound, in characters, for the conversation context fed to the
@@ -51,13 +53,14 @@ export interface PullRequestCreatedEvent {
  *
  * 1. Resolve session → working directory + current/base branch from
  *    {@link ISessionGitState}.
- * 2. Commit any uncommitted working-tree changes.
- * 3. Push the current branch to its GitHub upstream remote (with `--set-upstream` when missing).
- * 4. Resolve `owner` / `repo` from {@link ISessionGitState.githubOwner}
+ * 2. If the current branch is the base branch, create a generated session branch.
+ * 3. Commit any uncommitted working-tree changes.
+ * 4. Push the current branch to its GitHub upstream remote (with `--set-upstream` when missing).
+ * 5. Resolve `owner` / `repo` from {@link ISessionGitState.githubOwner}
  *    / {@link ISessionGitState.githubRepo} (populated by the git probe).
- * 5. Reuse an existing PR for the branch, or POST `/repos/{owner}/{repo}/pulls`
+ * 6. Reuse an existing PR for the branch, or POST `/repos/{owner}/{repo}/pulls`
  *    via {@link IAgentHostOctoKitService}.
- * 6. Return the PR URL as an {@link InvokeChangesetOperationResult.followUp}.
+ * 7. Return the PR URL as an {@link InvokeChangesetOperationResult.followUp}.
  */
 export class AgentHostPullRequestOperationHandler implements IChangesetOperationHandler {
 
@@ -78,6 +81,7 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		@IAgentHostOctoKitService private readonly _octoKitService: IAgentHostOctoKitService,
 		@IAgentHostGitHubEndpointService private readonly _gitHubEndpointService: IAgentHostGitHubEndpointService,
 		@ICopilotApiService private readonly _copilotApiService: ICopilotApiService,
+		@IAgentBranchNameGenerator private readonly _branchNameGenerator: IAgentBranchNameGenerator,
 		@ILogService private readonly _logService: ILogService,
 	) { }
 
@@ -100,8 +104,8 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 			throw new ProtocolError(JsonRpcErrorCodes.InvalidParams, `Not a changeset URI: ${params.channel}`);
 		}
 		this._throwIfCancelled(token);
-		const sessionUri = parsed.sessionUri;
 
+		const sessionUri = parsed.sessionUri;
 		const sessionState = this._getSessionState(sessionUri);
 		if (!sessionState) {
 			throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Session not found: ${sessionUri}`);
@@ -123,17 +127,18 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		const workingDirectory = URI.parse(workingDirectoryStr);
 		const storedGitState = readSessionGitState(sessionState._meta);
 		const effectiveBaseBranch = await this._resolveBaseBranchName(sessionUri);
-		const gitState = await this._gitService.getSessionGitState(workingDirectory, effectiveBaseBranch) ?? storedGitState;
-		const branchName = gitState?.branchName ?? await this._gitService.getCurrentBranch(workingDirectory);
+
+		let gitState = await this._gitService.getSessionGitState(workingDirectory, effectiveBaseBranch) ?? storedGitState;
+		let branchName = gitState?.branchName ?? await this._gitService.getCurrentBranch(workingDirectory);
 		if (!branchName) {
 			throw new ProtocolError(JsonRpcErrorCodes.InternalError, `Could not determine current branch for ${workingDirectory}`);
 		}
 
-		const baseBranchName = effectiveBaseBranch ?? gitState?.baseBranchName ?? (await this._gitService.getDefaultBranch(workingDirectory))?.name;
+		const defaultBranch = await this._gitService.getDefaultBranch(workingDirectory);
+		const baseBranchName = effectiveBaseBranch ?? gitState?.baseBranchName ?? defaultBranch?.name;
 		if (!baseBranchName) {
 			throw new ProtocolError(JsonRpcErrorCodes.InternalError, `Could not determine base branch for ${workingDirectory}`);
 		}
-		const base = baseBranchName;
 
 		const repoResource = this._gitHubEndpointService.getRepoResource();
 		const authToken = this._authenticationService.getAuthToken({
@@ -149,9 +154,39 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		}
 
 		const hasUncommitted = await this._gitService.hasUncommittedChanges(workingDirectory);
+
+		// Create a new branch if the current branch is the same
+		// as the base branch and there are uncommitted changes
+		if (hasUncommitted && branchName === baseBranchName) {
+			const branchPrefix = sessionState.config?.values[SessionConfigKey.WorktreeBranchPrefix];
+
+			try {
+				const generatedBranchName = await this._branchNameGenerator.generateBranchName({
+					sessionId: URI.parse(sessionUri).path.split('/').filter(Boolean).pop() ?? sessionUri,
+					message: sessionState.turns.find(turn => turn.message.text.trim())?.message.text,
+					githubToken: authToken,
+					signal,
+					branchPrefix: typeof branchPrefix === 'string' ? branchPrefix : undefined,
+					branchNameCollides: candidate => this._gitService.branchExists(workingDirectory, candidate).catch(() => true),
+				});
+
+				this._throwIfCancelled(token);
+				this._logService.info(`[AgentHostPullRequestOperationHandler] Creating branch ${generatedBranchName} for session ${sessionUri}`);
+
+				await this._gitService.createBranch(workingDirectory, generatedBranchName, { checkout: true });
+				branchName = generatedBranchName;
+
+				gitState = await this._gitService.getSessionGitState(workingDirectory, effectiveBaseBranch);
+			} catch (err) {
+				this._throwIfCancelled(token);
+				throw new ProtocolError(JsonRpcErrorCodes.InternalError, `Failed to create a branch before creating a pull request: ${err instanceof Error ? err.message : String(err)}`);
+			}
+		}
+
 		if (hasUncommitted) {
 			this._throwIfCancelled(token);
 			this._logService.info(`[AgentHostPullRequestOperationHandler] Committing uncommitted changes for session ${sessionUri}`);
+
 			try {
 				await this._gitService.commitAll(workingDirectory, this._formatCommitMessage(branchName));
 			} catch (err) {
@@ -161,7 +196,7 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		}
 		this._throwIfCancelled(token);
 
-		const branchChanges = await this._gitService.computeSessionFileDiffs(workingDirectory, { sessionUri, baseBranch: base });
+		const branchChanges = await this._gitService.computeSessionFileDiffs(workingDirectory, { sessionUri, baseBranch: baseBranchName });
 		if (branchChanges === undefined) {
 			throw new ProtocolError(JsonRpcErrorCodes.InternalError, localize('agentHost.changeset.pr.computeChangesFailed', "Could not compute branch changes to create a pull request."));
 		}
@@ -171,7 +206,7 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		this._throwIfCancelled(token);
 
 		const githubHeadOwner = gitState?.githubHeadOwner;
-		const upstreamBranch = githubHeadOwner ? parseUpstreamBranchName(gitState.upstreamBranchName) : undefined;
+		const upstreamBranch = githubHeadOwner ? parseUpstreamBranchName(gitState?.upstreamBranchName) : undefined;
 		const headOwner = upstreamBranch && githubHeadOwner ? githubHeadOwner : gitHubState.owner;
 		const headBranch = upstreamBranch?.branch ?? branchName;
 		const pushRef = headBranch === branchName ? branchName : `${branchName}:${headBranch}`;
@@ -195,12 +230,12 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 		}
 		this._throwIfCancelled(token);
 
-		const generated = await this._generateTitleAndDescription(sessionState, branchName, base, branchChanges, signal, token);
-		this._throwIfCancelled(token);
+		const generated = await this._generateTitleAndDescription(sessionState, branchName, baseBranchName, branchChanges, signal, token);
 		const title = generated?.title ?? this._formatTitle(branchName);
-		const body = generated?.description ?? this._formatBody(branchName, base);
+		const body = generated?.description ?? this._formatBody(branchName, baseBranchName);
+		this._throwIfCancelled(token);
 
-		this._logService.info(`[AgentHostPullRequestOperationHandler] Creating ${this._draft ? 'draft ' : ''}PR ${gitHubState.owner}/${gitHubState.repo} ${createHead} -> ${base}`);
+		this._logService.info(`[AgentHostPullRequestOperationHandler] Creating ${this._draft ? 'draft ' : ''}PR ${gitHubState.owner}/${gitHubState.repo} ${createHead} -> ${baseBranchName}`);
 		let created: CreatedPullRequest;
 		try {
 			created = await this._octoKitService.createPullRequest(
@@ -209,7 +244,7 @@ export class AgentHostPullRequestOperationHandler implements IChangesetOperation
 				title,
 				body,
 				createHead,
-				base,
+				baseBranchName,
 				this._draft,
 				authToken,
 				signal,

src/vs/platform/agentHost/node/shared/agentBranchNameGenerator.ts

@@ -40,10 +40,12 @@ export interface IAgentBranchNameGeneratorRequest {
 export const IAgentBranchNameGenerator = createDecorator<IAgentBranchNameGenerator>('agentBranchNameGenerator');
 
 export interface IAgentBranchNameGenerator {
+	readonly _serviceBrand: undefined;
 	generateBranchName(request: IAgentBranchNameGeneratorRequest): Promise<string>;
 }
 
 export class AgentBranchNameGenerator implements IAgentBranchNameGenerator {
+	declare readonly _serviceBrand: undefined;
 
 	constructor(
 		@ICopilotApiService private readonly _copilotApiService: ICopilotApiService,

src/vs/platform/agentHost/test/common/sessionTestHelpers.ts

@@ -350,6 +350,7 @@ export function createNoopGitService(): import('../../common/agentHostGitService
 		addExistingWorktree: async () => { },
 		removeWorktree: async () => { },
 		branchExists: async () => false,
+		createBranch: async () => { },
 		hasUncommittedChanges: async () => false,
 		commitAll: async () => { },
 		mergeBranch: async () => '',

src/vs/platform/agentHost/test/node/agentHostCommitOperationHandler.test.ts

@@ -47,6 +47,7 @@ class TestGitService implements IAgentHostGitService {
 	async addExistingWorktree(): Promise<void> { }
 	async removeWorktree(): Promise<void> { }
 	async branchExists(): Promise<boolean> { return false; }
+	async createBranch(): Promise<void> { }
 	async hasUncommittedChanges(): Promise<boolean> {
 		this.calls.push('hasUncommittedChanges');
 		return this.uncommitted;

src/vs/platform/agentHost/test/node/agentHostDiscardChangesOperationHandler.test.ts

@@ -35,6 +35,7 @@ class TestGitService implements IAgentHostGitService {
 	async addExistingWorktree(): Promise<void> { }
 	async removeWorktree(): Promise<void> { }
 	async branchExists(): Promise<boolean> { return false; }
+	async createBranch(): Promise<void> { }
 	async hasUncommittedChanges(): Promise<boolean> { return true; }
 	async commitAll(): Promise<void> { }
 	async mergeBranch(): Promise<string> { return ''; }

src/vs/platform/agentHost/test/node/agentHostGitService.integrationTest.ts

@@ -581,6 +581,25 @@ suite('AgentHostGitService - worktree helpers (real git)', () => {
 		assert.strictEqual(await svc!.branchExists(URI.file(dir), 'does-not-exist'), false);
 	});
 
+	(hasGit ? test : test.skip)('createBranch preserves dirty changes and leaves the base branch unchanged', async () => {
+		const dir = initRepo();
+		const fs = await import('fs/promises');
+		await fs.writeFile(join(dir, 'dirty.txt'), 'session changes');
+		const baseHead = cp.execFileSync('git', ['rev-parse', 'main'], { cwd: dir, env, encoding: 'utf8' }).trim();
+
+		await svc!.createBranch(URI.file(dir), 'agents/session', { checkout: true });
+
+		const branchName = cp.execFileSync('git', ['branch', '--show-current'], { cwd: dir, env, encoding: 'utf8' }).trim();
+		const currentHead = cp.execFileSync('git', ['rev-parse', 'HEAD'], { cwd: dir, env, encoding: 'utf8' }).trim();
+		const status = cp.execFileSync('git', ['status', '--porcelain'], { cwd: dir, env, encoding: 'utf8' }).trim();
+		assert.deepStrictEqual({ branchName, currentHead, baseHead, status }, {
+			branchName: 'agents/session',
+			currentHead: baseHead,
+			baseHead,
+			status: '?? dirty.txt',
+		});
+	});
+
 	(hasGit ? test : test.skip)('hasUncommittedChanges flips with untracked and committed work', async () => {
 		const dir = initRepo();
 		assert.strictEqual(await svc!.hasUncommittedChanges(URI.file(dir)), false);