Opens all changes after open in worktree flow finished for PRs (#3701)

* Opens all changes after open in worktree flow finished for PRs

* Update src/uris/deepLinks/deepLinkService.ts

Co-authored-by: Ramin Tadayon <[email protected]>

* Refinements

 - Only shows changes for "open in worktree" action
- Removes unnecessary state transition
- Only shows changes if worktree is created in flow
- Applies to branches in views and graph

---------

Co-authored-by: Ramin Tadayon <[email protected]>
Co-authored-by: Ramin Tadayon <[email protected]>

Author: nzaytsev

src/commands/git/switch.ts

@@ -34,7 +34,7 @@ interface Context {
 
 interface State {
 	repos: string | string[] | Repository | Repository[];
-	onWorkspaceChanging?: (() => Promise<void>) | (() => void);
+	onWorkspaceChanging?: ((isNewWorktree?: boolean) => Promise<void>) | ((isNewWorktree?: boolean) => void);
 	reference: GitReference;
 	createBranch?: string;
 	fastForwardTo?: GitReference;

src/commands/git/worktree.ts

@@ -102,7 +102,7 @@ interface CreateState {
 		title?: string;
 	};
 
-	onWorkspaceChanging?: (() => Promise<void>) | (() => void);
+	onWorkspaceChanging?: ((isNewWorktree?: boolean) => Promise<void>) | ((isNewWorktree?: boolean) => void);
 	skipWorktreeConfirmations?: boolean;
 }
 
@@ -139,7 +139,8 @@ interface OpenState {
 		};
 	};
 
-	onWorkspaceChanging?: (() => Promise<void>) | (() => void);
+	onWorkspaceChanging?: ((isNewWorktree?: boolean) => Promise<void>) | ((isNewWorktree?: boolean) => void);
+	isNewWorktree?: boolean;
 	skipWorktreeConfirmations?: boolean;
 }
 
@@ -628,6 +629,7 @@ export class WorktreeGitCommand extends QuickCommand<State> {
 						confirm: action === 'prompt',
 						openOnly: true,
 						overrides: { disallowBack: true },
+						isNewWorktree: true,
 						skipWorktreeConfirmations: state.skipWorktreeConfirmations,
 						onWorkspaceChanging: state.onWorkspaceChanging,
 					} satisfies OpenStepState,
@@ -1081,7 +1083,7 @@ export class WorktreeGitCommand extends QuickCommand<State> {
 
 				const location = convertOpenFlagsToLocation(state.flags);
 				if (location === 'currentWindow' || location === 'newWindow') {
-					await state.onWorkspaceChanging?.();
+					await state.onWorkspaceChanging?.(state.isNewWorktree);
 				}
 
 				openWorkspace(state.worktree.uri, { location: convertOpenFlagsToLocation(state.flags), name: name });

src/plus/launchpad/launchpadProvider.ts

@@ -540,7 +540,13 @@ export class LaunchpadProvider implements Disposable {
 				? item.openRepository.localBranch.name
 				: item.headRef.name;
 
-		return getPullRequestBranchDeepLink(this.container, branchName, item.repoIdentity.remote.url, action);
+		return getPullRequestBranchDeepLink(
+			this.container,
+			branchName,
+			item.repoIdentity.remote.url,
+			action,
+			item.underlyingPullRequest,
+		);
 	}
 
 	private async getMatchingOpenRepository(
@@ -1023,17 +1029,31 @@ export function getPullRequestBranchDeepLink(
 	headRefBranchName: string,
 	remoteUrl: string,
 	action?: DeepLinkActionType,
+	pr?: PullRequest,
 ) {
 	const schemeOverride = configuration.get('deepLinks.schemeOverride');
 	const scheme = typeof schemeOverride === 'string' ? schemeOverride : env.uriScheme;
+
+	const searchParams = new URLSearchParams({
+		url: ensureRemoteUrl(remoteUrl),
+	});
+	if (action) {
+		searchParams.set('action', action);
+	}
+	if (pr) {
+		searchParams.set('prId', pr.id);
+		searchParams.set('prTitle', pr.title);
+	}
+	if (pr?.refs) {
+		searchParams.set('prBaseRef', pr.refs.base.sha);
+		searchParams.set('prHeadRef', pr.refs.head.sha);
+	}
 	// TODO: Get the proper pull URL from the provider, rather than tacking .git at the end of the
 	// url from the head ref.
 	return Uri.parse(
 		`${scheme}://${container.context.extension.id}/${'link' satisfies UriTypes}/${DeepLinkType.Repository}/-/${
 			DeepLinkType.Branch
-		}/${encodeURIComponent(headRefBranchName)}?url=${encodeURIComponent(ensureRemoteUrl(remoteUrl))}${
-			action != null ? `&action=${action}` : ''
-		}`,
+		}/${encodeURIComponent(headRefBranchName)}?${searchParams.toString()}`,
 	);
 }
 

src/plus/webviews/graph/graphWebview.ts

@@ -62,7 +62,11 @@ import { GitContributor } from '../../../git/models/contributor';
 import type { GitGraph, GitGraphRowType } from '../../../git/models/graph';
 import { getGkProviderThemeIconString } from '../../../git/models/graph';
 import type { PullRequest } from '../../../git/models/pullRequest';
-import { getComparisonRefsForPullRequest, serializePullRequest } from '../../../git/models/pullRequest';
+import {
+	getComparisonRefsForPullRequest,
+	getRepositoryIdentityForPullRequest,
+	serializePullRequest,
+} from '../../../git/models/pullRequest';
 import type {
 	GitBranchReference,
 	GitReference,
@@ -115,6 +119,7 @@ import { getContext, onDidChangeContext } from '../../../system/vscode/context';
 import type { OpenWorkspaceLocation } from '../../../system/vscode/utils';
 import { isDarkTheme, isLightTheme, openWorkspace } from '../../../system/vscode/utils';
 import { isWebviewItemContext, isWebviewItemGroupContext, serializeWebviewItemContext } from '../../../system/webview';
+import { DeepLinkActionType } from '../../../uris/deepLinks/deepLink';
 import { RepositoryFolderNode } from '../../../views/nodes/abstract/repositoryFolderNode';
 import type { IpcCallMessageType, IpcMessage, IpcNotification } from '../../../webviews/protocol';
 import type { WebviewHost, WebviewProvider, WebviewShowingArgs } from '../../../webviews/webviewProvider';
@@ -123,6 +128,7 @@ import { isSerializedState } from '../../../webviews/webviewsController';
 import type { SubscriptionChangeEvent } from '../../gk/account/subscriptionService';
 import type { ConnectionStateChangeEvent } from '../../integrations/integrationService';
 import { remoteProviderIdToIntegrationId } from '../../integrations/integrationService';
+import { getPullRequestBranchDeepLink } from '../../launchpad/launchpadProvider';
 import type {
 	BranchState,
 	DidChangeRefsVisibilityParams,
@@ -3643,6 +3649,24 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 	private async openInWorktree(item?: GraphItemContext) {
 		if (isGraphItemRefContext(item, 'branch')) {
 			const { ref } = item.webviewItemValue;
+			const repo = this.container.git.getRepository(ref.repoPath);
+			const branch = await repo?.git.getBranch(ref.name);
+			const pr = await branch?.getAssociatedPullRequest();
+			if (branch != null && repo != null && pr != null) {
+				const remoteUrl = (await branch.getRemote())?.url ?? getRepositoryIdentityForPullRequest(pr).remote.url;
+				if (remoteUrl != null) {
+					const deepLink = getPullRequestBranchDeepLink(
+						this.container,
+						branch.getNameWithoutRemote(),
+						remoteUrl,
+						DeepLinkActionType.SwitchToPullRequestWorktree,
+						pr,
+					);
+
+					return this.container.deepLinks.processDeepLinkUri(deepLink, false, repo);
+				}
+			}
+
 			await executeGitCommand({
 				command: 'switch',
 				state: {

src/uris/deepLinks/deepLink.ts

@@ -100,6 +100,7 @@ export interface DeepLink {
 	secondaryTargetId?: string;
 	secondaryRemoteUrl?: string;
 	action?: string;
+	prId?: string;
 	params?: URLSearchParams;
 }
 
@@ -178,6 +179,7 @@ export function parseDeepLinkUri(uri: Uri): DeepLink | undefined {
 				secondaryRemoteUrl: secondaryRemoteUrl,
 				action: action,
 				params: urlParams,
+				prId: urlParams.get('prId') ?? undefined,
 			};
 		}
 		case DeepLinkType.Draft: {
@@ -239,6 +241,7 @@ export const enum DeepLinkServiceState {
 	OpenInspect,
 	SwitchToRef,
 	RunCommand,
+	OpenAllPrChanges,
 }
 
 export const enum DeepLinkServiceAction {
@@ -271,6 +274,7 @@ export const enum DeepLinkServiceAction {
 	OpenFile,
 	OpenInspect,
 	OpenSwitch,
+	OpenAllPrChanges,
 }
 
 export type DeepLinkRepoOpenType = 'clone' | 'folder' | 'workspace' | 'current';
@@ -417,6 +421,12 @@ export const deepLinkStateTransitionTable: Record<string, Record<string, DeepLin
 		[DeepLinkServiceAction.DeepLinkCancelled]: DeepLinkServiceState.Idle,
 	},
 	[DeepLinkServiceState.OpenInspect]: {
+		[DeepLinkServiceAction.OpenAllPrChanges]: DeepLinkServiceState.OpenAllPrChanges,
+		[DeepLinkServiceAction.DeepLinkResolved]: DeepLinkServiceState.Idle,
+		[DeepLinkServiceAction.DeepLinkErrored]: DeepLinkServiceState.Idle,
+		[DeepLinkServiceAction.DeepLinkCancelled]: DeepLinkServiceState.Idle,
+	},
+	[DeepLinkServiceState.OpenAllPrChanges]: {
 		[DeepLinkServiceAction.DeepLinkResolved]: DeepLinkServiceState.Idle,
 		[DeepLinkServiceAction.DeepLinkErrored]: DeepLinkServiceState.Idle,
 		[DeepLinkServiceAction.DeepLinkCancelled]: DeepLinkServiceState.Idle,

src/uris/deepLinks/deepLinkService.ts

@@ -4,7 +4,7 @@ import { Commands } from '../../constants.commands';
 import type { StoredDeepLinkContext, StoredNamedRef } from '../../constants.storage';
 import type { Container } from '../../container';
 import { executeGitCommand } from '../../git/actions';
-import { openFileAtRevision } from '../../git/actions/commit';
+import { openComparisonChanges, openFileAtRevision } from '../../git/actions/commit';
 import type { GitBranch } from '../../git/models/branch';
 import { getBranchNameWithoutRemote } from '../../git/models/branch';
 import type { GitCommit } from '../../git/models/commit';
@@ -19,6 +19,7 @@ import { missingRepositoryId } from '../../gk/models/repositoryIdentities';
 import { ensureAccount, ensurePaidPlan } from '../../plus/utils';
 import type { ShowInCommitGraphCommandArgs } from '../../plus/webviews/graph/protocol';
 import { createQuickPickSeparator } from '../../quickpicks/items/common';
+import { debug } from '../../system/decorators/log';
 import { once } from '../../system/event';
 import { Logger } from '../../system/logger';
 import { normalizePath } from '../../system/path';
@@ -68,10 +69,7 @@ export class DeepLinkService implements Disposable {
 		this._disposables.push(container.uri.onDidReceiveUri(async (uri: Uri) => this.processDeepLinkUri(uri)));
 
 		const pendingDeepLink = this.container.storage.get('deepLinks:pending');
-		if (pendingDeepLink != null) {
-			void this.container.storage.delete('deepLinks:pending');
-			void this.processPendingDeepLink(pendingDeepLink);
-		}
+		void this.processPendingDeepLink(pendingDeepLink);
 	}
 
 	dispose() {
@@ -241,9 +239,11 @@ export class DeepLinkService implements Disposable {
 		}
 	}
 
-	private async processPendingDeepLink(pendingDeepLink: StoredDeepLinkContext) {
-		if (pendingDeepLink.url == null) return;
-
+	@debug()
+	private async processPendingDeepLink(pendingDeepLink: StoredDeepLinkContext | undefined) {
+		if (pendingDeepLink == null) return;
+		void this.container.storage.delete('deepLinks:pending');
+		if (pendingDeepLink?.url == null) return;
 		const link = parseDeepLinkUri(Uri.parse(pendingDeepLink.url));
 		if (link == null) return;
 
@@ -509,6 +509,7 @@ export class DeepLinkService implements Disposable {
 	// TODO @axosoft-ramint: Move all the logic for matching a repo, prompting to add repo, matching remote, etc. for a target (branch, PR, etc.)
 	// to a separate service where it can be used outside of the context of deep linking. Then the deep link service should leverage it,
 	// and we should stop using deep links to process things like Launchpad switch actions, Open in Worktree command, etc.
+	@debug()
 	private async processDeepLink(
 		initialAction: DeepLinkServiceAction = DeepLinkServiceAction.DeepLinkEventFired,
 		useProgress: boolean = true,
@@ -1078,7 +1079,6 @@ export class DeepLinkService implements Disposable {
 				}
 				case DeepLinkServiceState.GoToTarget: {
 					// Need to re-fetch the remotes in case we opened in a new window
-
 					if (targetType === DeepLinkType.Repository) {
 						if (
 							this._context.action === DeepLinkActionType.Switch ||
@@ -1321,9 +1321,25 @@ export class DeepLinkService implements Disposable {
 							state: this._context.state,
 						};
 
+						// Form a new link URL with PR info stripped out in case we are opening an existing PR worktree,
+						// in which case we do not want to advance to the  "Open All PR Changes" step in the flow.
+						// We should only advance to that step if a worktree is newly created in the flow.
+						const oldUrl = Uri.parse(this._context.url ?? '').toString(true);
+						const urlParams = new URL(oldUrl).searchParams;
+						urlParams.delete('prId');
+						urlParams.delete('prTitle');
+						urlParams.delete('prBaseRef');
+						urlParams.delete('prHeadRef');
+						const newUrlParams = urlParams.toString();
+						const nonPrUrl =
+							newUrlParams.length > 0 ? `${oldUrl.split('?')[0]}?${newUrlParams}` : oldUrl.split('?')[0];
+
 						// Storing link info in case the switch causes a new window to open
-						const onWorkspaceChanging = async () =>
-							this.container.storage.store('deepLinks:pending', pendingDeepLink);
+						const onWorkspaceChanging = async (isNewWorktree?: boolean) =>
+							this.container.storage.store(
+								'deepLinks:pending',
+								isNewWorktree ? pendingDeepLink : { ...pendingDeepLink, url: nonPrUrl },
+							);
 
 						await executeGitCommand({
 							command: 'switch',
@@ -1373,7 +1389,6 @@ export class DeepLinkService implements Disposable {
 				case DeepLinkServiceState.OpenInspect: {
 					// If we arrive at this step, clear any stored data used for the "new window" option
 					await this.container.storage.delete('deepLinks:pending');
-
 					if (!repo) {
 						action = DeepLinkServiceAction.DeepLinkErrored;
 						message = 'Missing repository.';
@@ -1386,6 +1401,46 @@ export class DeepLinkService implements Disposable {
 						repository: repo,
 						source: 'launchpad',
 					} satisfies ShowWipArgs);
+					const { params } = this._context;
+					if (
+						this._context.action === DeepLinkActionType.SwitchToPullRequestWorktree &&
+						params != null &&
+						(params.get('prId') != null || params.get('prTitle') != null) &&
+						params.get('prBaseRef') != null &&
+						params.get('prHeadRef') != null
+					) {
+						action = DeepLinkServiceAction.OpenAllPrChanges;
+						break;
+					}
+
+					action = DeepLinkServiceAction.DeepLinkResolved;
+					break;
+				}
+				case DeepLinkServiceState.OpenAllPrChanges: {
+					const prId = this._context.params?.get('prId');
+					const prHeadRef = this._context.params?.get('prHeadRef');
+					const prBaseRef = this._context.params?.get('prBaseRef');
+					const prTitle = this._context.params?.get('prTitle');
+					if (!repoPath || (!prId && !prTitle) || !prHeadRef || !prBaseRef) {
+						action = DeepLinkServiceAction.DeepLinkErrored;
+						if (!repoPath) {
+							message = 'No repository path was provided.';
+						} else if (!prId) {
+							message = 'No pull request id provided.';
+						} else {
+							message = 'No pull request refs was provided.';
+						}
+						break;
+					}
+					await openComparisonChanges(
+						this.container,
+						{
+							repoPath: repoPath,
+							lhs: prBaseRef,
+							rhs: prHeadRef,
+						},
+						{ title: `Changes in Pull Request ${prTitle ? `"${prTitle}"` : `#${prId}`}` },
+					);
 					action = DeepLinkServiceAction.DeepLinkResolved;
 					break;
 				}

src/views/viewCommands.ts

@@ -786,6 +786,23 @@ export class ViewCommands {
 		if (!node.is('branch') && !node.is('pullrequest') && !node.is('launchpad-item')) return;
 
 		if (node.is('branch')) {
+			const pr = await node.branch.getAssociatedPullRequest();
+			if (pr != null) {
+				const remoteUrl =
+					(await node.branch.getRemote())?.url ?? getRepositoryIdentityForPullRequest(pr).remote.url;
+				if (remoteUrl != null) {
+					const deepLink = getPullRequestBranchDeepLink(
+						this.container,
+						node.branch.getNameWithoutRemote(),
+						remoteUrl,
+						DeepLinkActionType.SwitchToPullRequestWorktree,
+						pr,
+					);
+
+					return this.container.deepLinks.processDeepLinkUri(deepLink, false, node.repo);
+				}
+			}
+
 			return executeGitCommand({
 				command: 'switch',
 				state: {
@@ -808,6 +825,7 @@ export class ViewCommands {
 				pr.refs.head.branch,
 				repoIdentity.remote.url,
 				DeepLinkActionType.SwitchToPullRequestWorktree,
+				pr,
 			);
 
 			const prRepo = await getOrOpenPullRequestRepository(this.container, pr, {