Sessions: fixes for detect PR (#4629)

* Sessions: fixes for detect PR

* Test and comment

* Fix

Author: osortega

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

@@ -1546,30 +1546,63 @@ export class CopilotCLIChatSessionParticipant extends Disposable {
 		}
 	}
 
+	private static readonly _PR_DETECTION_RETRY_COUNT = 3;
+	private static readonly _PR_DETECTION_INITIAL_DELAY_MS = 2_000;
+
 	private async handlePullRequestCreated(session: ICopilotCLISession): Promise<void> {
+		const sessionId = session.sessionId;
 		let prUrl = session.createdPullRequestUrl;
 
 		if (!prUrl) {
-			prUrl = await this.detectPullRequestForSession(session.sessionId);
+			// Only attempt retry detection if the session has v2 worktree properties
+			// with branch info — v1 worktrees can't store PR URLs, and sessions
+			// without worktree properties have nothing to look up.
+			const worktreeProperties = await this.copilotCLIWorktreeManagerService.getWorktreeProperties(sessionId);
+			if (worktreeProperties?.version === 2 && worktreeProperties.branchName && worktreeProperties.repositoryPath) {
+				prUrl = await this.detectPullRequestWithRetry(sessionId);
+			}
 		}
 
 		if (!prUrl) {
 			return;
 		}
 
 		try {
-			const worktreeProperties = await this.copilotCLIWorktreeManagerService.getWorktreeProperties(session.sessionId);
+			const worktreeProperties = await this.copilotCLIWorktreeManagerService.getWorktreeProperties(sessionId);
 			if (worktreeProperties && worktreeProperties.version === 2) {
-				await this.copilotCLIWorktreeManagerService.setWorktreeProperties(session.sessionId, {
+				await this.copilotCLIWorktreeManagerService.setWorktreeProperties(sessionId, {
 					...worktreeProperties,
 					pullRequestUrl: prUrl,
 					changes: undefined,
 				});
 				this.sessionItemProvider.notifySessionsChange();
 			}
 		} catch (error) {
-			this.logService.error(`Failed to persist pull request metadata: ${error instanceof Error ? error.message : String(error)}`);
+			this.logService.error(`Failed to persist pull request metadata for session ${sessionId}: ${error instanceof Error ? error.message : String(error)}`);
+		}
+	}
+
+	/**
+	 * Attempts to detect a pull request for a freshly-completed session using
+	 * exponential backoff. The GitHub API may not have indexed the PR immediately
+	 * after `gh pr create` returns, so we retry with increasing delays:
+	 * attempt 1: 2s, attempt 2: 4s, attempt 3: 8s.
+	 */
+	private async detectPullRequestWithRetry(sessionId: string): Promise<string | undefined> {
+		const maxRetries = CopilotCLIChatSessionParticipant._PR_DETECTION_RETRY_COUNT;
+		const initialDelay = CopilotCLIChatSessionParticipant._PR_DETECTION_INITIAL_DELAY_MS;
+
+		for (let attempt = 0; attempt < maxRetries; attempt++) {
+			const delay = initialDelay * Math.pow(2, attempt);
+			await new Promise<void>(resolve => setTimeout(resolve, delay));
+
+			const prUrl = await this.detectPullRequestForSession(sessionId);
+			if (prUrl) {
+				return prUrl;
+			}
 		}
+
+		return undefined;
 	}
 
 	/**

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

@@ -40,7 +40,7 @@ import { IAgentSessionsWorkspace } from '../../common/agentSessionsWorkspace';
 import { IChatCustomAgentsService } from '../../common/chatCustomAgentsService';
 import { IChatSessionWorkspaceFolderService } from '../../common/chatSessionWorkspaceFolderService';
 import { IChatSessionWorktreeCheckpointService } from '../../common/chatSessionWorktreeCheckpointService';
-import { IChatSessionWorktreeService, type ChatSessionWorktreeFile, type ChatSessionWorktreeProperties } from '../../common/chatSessionWorktreeService';
+import { IChatSessionWorktreeService, type ChatSessionWorktreeFile, type ChatSessionWorktreeProperties, type ChatSessionWorktreePropertiesV2 } from '../../common/chatSessionWorktreeService';
 import { MockChatSessionMetadataStore } from '../../common/test/mockChatSessionMetadataStore';
 import { getWorkingDirectory, IWorkspaceInfo } from '../../common/workspaceInfo';
 import { IChatDelegationSummaryService } from '../../copilotcli/common/delegationSummaryService';
@@ -1971,4 +1971,141 @@ describe('CopilotCLIChatSessionParticipant.handleRequest', () => {
 			expect(agent?.tools).toBeNull();
 		});
 	});
+
+	describe('PR detection with retry', () => {
+		let octoKitService: IOctoKitService;
+
+		const v2WorktreeProperties: ChatSessionWorktreePropertiesV2 = {
+			version: 2,
+			baseCommit: 'abc123',
+			branchName: 'copilot/test-branch',
+			baseBranchName: 'main',
+			repositoryPath: `${sep}repo`,
+			worktreePath: `${sep}worktree`,
+		};
+
+		const repoContext: RepoContext = {
+			rootUri: Uri.file(`${sep}repo`),
+			kind: 'repository',
+			remotes: ['origin'],
+			remoteFetchUrls: ['https://github.com/testowner/testrepo.git'],
+		} as unknown as RepoContext;
+
+		beforeEach(() => {
+			vi.useFakeTimers();
+			octoKitService = {
+				findPullRequestByHeadBranch: vi.fn(async () => undefined),
+			} as unknown as IOctoKitService;
+
+			// Set up folder & git repo so session creation succeeds with worktree isolation
+			folderRepositoryManager.setUntitledSessionFolder('untitled:pr-test', Uri.file(`${sep}repo`));
+			git.setRepo(repoContext);
+			(worktree.createWorktree as unknown as ReturnType<typeof vi.fn>).mockResolvedValue(v2WorktreeProperties);
+			// After session creation, getWorktreeProperties returns v2 for any session
+			(worktree.getWorktreeProperties as unknown as ReturnType<typeof vi.fn>).mockResolvedValue(v2WorktreeProperties);
+			TestCopilotCLISession.statusOverride = vscode.ChatSessionStatus.Completed;
+
+			// Recreate participant with the controllable octoKitService
+			participant = new CopilotCLIChatSessionParticipant(
+				contentProvider,
+				promptResolver,
+				itemProvider,
+				cloudProvider,
+				repositoryTracker,
+				git,
+				models as unknown as ICopilotCLIModels,
+				new NullCopilotCLIAgents(),
+				sessionService,
+				worktree,
+				worktreeCheckpointService,
+				workspaceFolderService,
+				telemetry,
+				logService,
+				new PromptsServiceImpl(new NullWorkspaceService()),
+				new (mock<IChatDelegationSummaryService>())(),
+				folderRepositoryManager,
+				configurationService,
+				sdk,
+				new MockChatSessionMetadataStore(),
+				customSessionTitleService,
+				octoKitService,
+			);
+		});
+
+		afterEach(() => {
+			vi.useRealTimers();
+		});
+
+		it('retries PR detection with exponential backoff and succeeds on second attempt', async () => {
+			const findPr = octoKitService.findPullRequestByHeadBranch as ReturnType<typeof vi.fn>;
+			findPr
+				.mockResolvedValueOnce(undefined) // attempt 1: not found
+				.mockResolvedValueOnce({ url: 'https://github.com/testowner/testrepo/pull/42' }); // attempt 2: found
+
+			const request = new TestChatRequest('Create a PR');
+			const context = createChatContext('untitled:pr-test', true);
+			const stream = new MockChatResponseStream();
+			const token = disposables.add(new CancellationTokenSource()).token;
+
+			const handlerPromise = participant.createHandler()(request, context, stream, token);
+			await vi.runAllTimersAsync();
+			await handlerPromise;
+
+			// Should have been called twice (after 2s delay, then after 4s delay)
+			expect(findPr).toHaveBeenCalledTimes(2);
+			// Should have persisted the PR URL
+			expect(worktree.setWorktreeProperties).toHaveBeenCalledWith(
+				expect.any(String),
+				expect.objectContaining({ pullRequestUrl: 'https://github.com/testowner/testrepo/pull/42' })
+			);
+		});
+
+		it('stops retrying once all attempts are exhausted', async () => {
+			const findPr = octoKitService.findPullRequestByHeadBranch as ReturnType<typeof vi.fn>;
+			findPr.mockResolvedValue(undefined); // always returns not found
+
+			const request = new TestChatRequest('Create something');
+			const context = createChatContext('untitled:pr-test', true);
+			const stream = new MockChatResponseStream();
+			const token = disposables.add(new CancellationTokenSource()).token;
+
+			const handlerPromise = participant.createHandler()(request, context, stream, token);
+			await vi.runAllTimersAsync();
+			await handlerPromise;
+
+			// 3 attempts total (after 2s, 4s, and 8s delays)
+			expect(findPr).toHaveBeenCalledTimes(3);
+			// Should NOT have persisted any PR URL since all attempts failed
+			const setPropsCallsWithPrUrl = (worktree.setWorktreeProperties as ReturnType<typeof vi.fn>).mock.calls
+				.filter((args: unknown[]) => (args[1] as { pullRequestUrl?: string })?.pullRequestUrl !== undefined);
+			expect(setPropsCallsWithPrUrl).toHaveLength(0);
+		});
+
+		it('skips retry when session already has createdPullRequestUrl', async () => {
+			const findPr = octoKitService.findPullRequestByHeadBranch as ReturnType<typeof vi.fn>;
+
+			// Make the session report a PR URL directly
+			TestCopilotCLISession.handleRequestHook = vi.fn(async () => {
+				const session = cliSessions[cliSessions.length - 1];
+				(session as any)._createdPullRequestUrl = 'https://github.com/testowner/testrepo/pull/99';
+			});
+
+			const request = new TestChatRequest('Create a PR via MCP');
+			const context = createChatContext('untitled:pr-test', true);
+			const stream = new MockChatResponseStream();
+			const token = disposables.add(new CancellationTokenSource()).token;
+
+			const handlerPromise = participant.createHandler()(request, context, stream, token);
+			await vi.runAllTimersAsync();
+			await handlerPromise;
+
+			// Should NOT have called the GitHub API since session had the URL
+			expect(findPr).not.toHaveBeenCalled();
+			// Should have persisted the session's PR URL
+			expect(worktree.setWorktreeProperties).toHaveBeenCalledWith(
+				expect.any(String),
+				expect.objectContaining({ pullRequestUrl: 'https://github.com/testowner/testrepo/pull/99' })
+			);
+		});
+	});
 });