Handles cancellation of AI commands

src/commands/explainBranch.ts

@@ -114,6 +114,8 @@ export class ExplainBranchCommand extends ExplainCommandBase {
 				},
 			);
 
+			if (result === 'cancelled') return;
+
 			if (result == null) {
 				void showGenericErrorMessage(`Unable to explain branch ${branch.name}`);
 				return;

src/commands/explainCommit.ts

@@ -90,6 +90,8 @@ export class ExplainCommitCommand extends ExplainCommandBase {
 				},
 			);
 
+			if (result === 'cancelled') return;
+
 			if (result == null) {
 				void showGenericErrorMessage('Unable to explain commit');
 				return;

src/commands/explainStash.ts

@@ -76,6 +76,8 @@ export class ExplainStashCommand extends ExplainCommandBase {
 				},
 			);
 
+			if (result === 'cancelled') return;
+
 			if (result == null) {
 				void showGenericErrorMessage('No changes found to explain for stash');
 				return;

src/commands/explainWip.ts

@@ -112,6 +112,8 @@ export class ExplainWipCommand extends ExplainCommandBase {
 				},
 			);
 
+			if (result === 'cancelled') return;
+
 			if (result == null) {
 				void showGenericErrorMessage(`Unable to explain ${label} changes`);
 				return;

src/commands/generateChangelog.ts

@@ -95,6 +95,8 @@ export async function generateChangelogAndOpenMarkdownDocument(
 ): Promise<void> {
 	const result = await container.ai.generateChangelog(changes, source, options);
 
+	if (result === 'cancelled') return;
+
 	const { range, changes: { length: count } = [] } = await changes.value;
 
 	let content = `# Changelog for ${range.head.label ?? range.head.ref}\n`;

src/commands/generateCommitMessage.ts

@@ -74,7 +74,7 @@ export class GenerateCommitMessageCommand extends ActiveEditorCommand {
 					progress: { location: ProgressLocation.Notification, title: 'Generating commit message...' },
 				},
 			);
-			if (result == null) return;
+			if (result == null || result === 'cancelled') return;
 
 			void executeCoreCommand('workbench.view.scm');
 			scmRepo.inputBox.value = `${currentMessage ? `${currentMessage}\n\n` : ''}${result.parsed.summary}${

src/commands/generateRebase.ts

@@ -343,7 +343,7 @@ export async function generateRebase(
 
 	const repo = svc.getRepository()!;
 	const result = await container.ai.generateRebase(repo, base.ref, head.ref, source, aiOptions);
-	if (result == null) return;
+	if (result == null || result === 'cancelled') return;
 
 	try {
 		// Extract the diff information from the reorganized commits

src/commands/git/stash.ts

@@ -677,6 +677,8 @@ export class StashGitCommand extends QuickCommand<State> {
 						resume?.dispose();
 						input.validationMessage = undefined;
 
+						if (result === 'cancelled') return;
+
 						const message = result?.parsed.summary;
 						if (message != null) {
 							state.message = message;

src/env/node/git/commitMessageProvider.ts

@@ -62,7 +62,8 @@ class AICommitMessageProvider implements CommitMessageProvider, Disposable {
 				},
 			);
 
-			if (result == null) return;
+			if (result == null || result === 'cancelled') return;
+
 			return `${currentMessage ? `${currentMessage}\n\n` : ''}${result.parsed.summary}${
 				result.parsed.body ? `\n\n${result.parsed.body}` : ''
 			}`;

src/git/utils/-webview/pullRequest.utils.ts

@@ -39,6 +39,8 @@ export async function describePullRequestWithAI(
 				...options,
 			},
 		);
+		if (result === 'cancelled') return undefined;
+
 		return result?.parsed ? { title: result.parsed.summary, description: result.parsed.body } : undefined;
 	} catch (ex) {
 		void window.showErrorMessage(ex.message);

src/plus/ai/aiProviderService.ts

@@ -510,7 +510,7 @@ export class AIProviderService implements Disposable {
 		commitOrRevision: GitRevisionReference | GitCommit,
 		sourceContext: AIExplainSource,
 		options?: { cancellation?: CancellationToken; progress?: ProgressOptions },
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const svc = this.container.git.getRepositoryService(commitOrRevision.repoPath);
 		return this.explainChanges(
 			async cancellation => {
@@ -543,7 +543,7 @@ export class AIProviderService implements Disposable {
 			| ((cancellationToken: CancellationToken) => Promise<PromptTemplateContext<'explain-changes'>>),
 		sourceContext: AIExplainSource,
 		options?: { cancellation?: CancellationToken; progress?: ProgressOptions },
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const { type, ...source } = sourceContext;
 
 		const result = await this.sendRequest(
@@ -588,7 +588,11 @@ export class AIProviderService implements Disposable {
 			}),
 			options,
 		);
-		return result != null ? { ...result, parsed: parseSummarizeResult(result.content) } : undefined;
+		return result === 'cancelled'
+			? result
+			: result != null
+			  ? { ...result, parsed: parseSummarizeResult(result.content) }
+			  : undefined;
 	}
 
 	async generateCommitMessage(
@@ -600,7 +604,7 @@ export class AIProviderService implements Disposable {
 			generating?: Deferred<AIModel>;
 			progress?: ProgressOptions;
 		},
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const result = await this.sendRequest(
 			'generate-commitMessage',
 			async (model, reporting, cancellation, maxInputTokens, retries) => {
@@ -639,7 +643,11 @@ export class AIProviderService implements Disposable {
 			}),
 			{ ...options, modelOptions: { outputTokens: 4096 } },
 		);
-		return result != null ? { ...result, parsed: parseSummarizeResult(result.content) } : undefined;
+		return result === 'cancelled'
+			? result
+			: result != null
+			  ? { ...result, parsed: parseSummarizeResult(result.content) }
+			  : undefined;
 	}
 
 	async generateCreatePullRequest(
@@ -653,7 +661,7 @@ export class AIProviderService implements Disposable {
 			generating?: Deferred<AIModel>;
 			progress?: ProgressOptions;
 		},
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const result = await this.sendRequest(
 			'generate-create-pullRequest',
 			async (model, reporting, cancellation, maxInputTokens, retries) => {
@@ -698,7 +706,11 @@ export class AIProviderService implements Disposable {
 			}),
 			options,
 		);
-		return result != null ? { ...result, parsed: parseSummarizeResult(result.content) } : undefined;
+		return result === 'cancelled'
+			? result
+			: result != null
+			  ? { ...result, parsed: parseSummarizeResult(result.content) }
+			  : undefined;
 	}
 
 	async generateCreateDraft(
@@ -711,7 +723,7 @@ export class AIProviderService implements Disposable {
 			progress?: ProgressOptions;
 			codeSuggestion?: boolean;
 		},
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const { type, ...source } = sourceContext;
 
 		const result = await this.sendRequest(
@@ -762,7 +774,11 @@ export class AIProviderService implements Disposable {
 			}),
 			options,
 		);
-		return result != null ? { ...result, parsed: parseSummarizeResult(result.content) } : undefined;
+		return result === 'cancelled'
+			? result
+			: result != null
+			  ? { ...result, parsed: parseSummarizeResult(result.content) }
+			  : undefined;
 	}
 
 	async generateStashMessage(
@@ -774,7 +790,7 @@ export class AIProviderService implements Disposable {
 			generating?: Deferred<AIModel>;
 			progress?: ProgressOptions;
 		},
-	): Promise<AISummarizeResult | undefined> {
+	): Promise<AISummarizeResult | 'cancelled' | undefined> {
 		const result = await this.sendRequest(
 			'generate-stashMessage',
 			async (model, reporting, cancellation, maxInputTokens, retries) => {
@@ -813,14 +829,18 @@ export class AIProviderService implements Disposable {
 			}),
 			{ ...options, modelOptions: { outputTokens: 1024 } },
 		);
-		return result != null ? { ...result, parsed: parseSummarizeResult(result.content) } : undefined;
+		return result === 'cancelled'
+			? result
+			: result != null
+			  ? { ...result, parsed: parseSummarizeResult(result.content) }
+			  : undefined;
 	}
 
 	async generateChangelog(
 		changes: Lazy<Promise<AIGenerateChangelogChanges>>,
 		source: Source,
 		options?: { cancellation?: CancellationToken; progress?: ProgressOptions },
-	): Promise<AIResult | undefined> {
+	): Promise<AIResult | 'cancelled' | undefined> {
 		const result = await this.sendRequest(
 			'generate-changelog',
 			async (model, reporting, cancellation, maxInputTokens, retries) => {
@@ -858,7 +878,7 @@ export class AIProviderService implements Disposable {
 			}),
 			options,
 		);
-		return result != null ? { ...result } : undefined;
+		return result === 'cancelled' ? result : result != null ? { ...result } : undefined;
 	}
 
 	async generateRebase(
@@ -873,7 +893,7 @@ export class AIProviderService implements Disposable {
 			progress?: ProgressOptions;
 			generateCommits?: boolean;
 		},
-	): Promise<AIRebaseResult | undefined> {
+	): Promise<AIRebaseResult | 'cancelled' | undefined> {
 		const result: Mutable<AIRebaseResult> = {
 			diff: undefined!,
 			explanation: undefined!,
@@ -984,9 +1004,13 @@ export class AIProviderService implements Disposable {
 			options,
 		);
 
+		if (rq === 'cancelled') return rq;
+
+		if (rq == null) return undefined;
+
 		try {
 			// if it is wrapped in markdown, we need to strip it
-			const content = rq!.content.replace(/^\s*```json\s*/, '').replace(/\s*```$/, '');
+			const content = rq.content.replace(/^\s*```json\s*/, '').replace(/\s*```$/, '');
 			// Parse the JSON content from the result
 			result.commits = JSON.parse(content) as AIRebaseResult['commits'];
 		} catch {
@@ -1021,12 +1045,17 @@ export class AIProviderService implements Disposable {
 			modelOptions?: { outputTokens?: number; temperature?: number };
 			progress?: ProgressOptions;
 		},
-	): Promise<AIRequestResult | undefined> {
+	): Promise<AIRequestResult | 'cancelled' | undefined> {
 		if (!(await this.ensureFeatureAccess(action, source))) {
-			return undefined;
+			return 'cancelled';
 		}
 
 		const model = await this.getModel(undefined, source);
+		if (options?.cancellation?.isCancellationRequested) {
+			options?.generating?.cancel();
+			return 'cancelled';
+		}
+
 		if (model == null || options?.cancellation?.isCancellationRequested) {
 			options?.generating?.cancel();
 			return undefined;
@@ -1053,16 +1082,26 @@ export class AIProviderService implements Disposable {
 			);
 
 			options?.generating?.cancel();
-			return undefined;
+			return 'cancelled';
 		}
 
 		const apiKey = await this._provider!.getApiKey(false);
-		if (apiKey == null || cancellation.isCancellationRequested) {
+
+		if (cancellation.isCancellationRequested) {
 			this.container.telemetry.sendEvent(
 				telementry.key,
-				cancellation.isCancellationRequested
-					? { ...telementry.data, failed: true, 'failed.reason': 'user-cancelled' }
-					: { ...telementry.data, failed: true, 'failed.reason': 'error', 'failed.error': 'Not authorized' },
+				{ ...telementry.data, failed: true, 'failed.reason': 'user-cancelled' },
+				source,
+			);
+
+			options?.generating?.cancel();
+			return 'cancelled';
+		}
+
+		if (apiKey == null) {
+			this.container.telemetry.sendEvent(
+				telementry.key,
+				{ ...telementry.data, failed: true, 'failed.reason': 'error', 'failed.error': 'Not authorized' },
 				source,
 			);
 
@@ -1121,7 +1160,7 @@ export class AIProviderService implements Disposable {
 					source,
 				);
 
-				return undefined;
+				return 'cancelled';
 			}
 			if (ex instanceof AIError) {
 				this.container.telemetry.sendEvent(

src/webviews/commitDetails/commitDetailsWebview.ts

@@ -1183,6 +1183,8 @@ export class CommitDetailsWebviewProvider
 				{ source: 'inspect', type: 'suggested_pr_change' },
 				{ progress: { location: { viewId: this.host.id } } },
 			);
+			if (result === 'cancelled') throw new Error('Operation was canceled');
+
 			if (result == null) throw new Error('Error retrieving content');
 
 			params = {

src/webviews/plus/patchDetails/patchDetailsWebview.ts

@@ -837,6 +837,8 @@ export class PatchDetailsWebviewProvider
 				{ source: 'patchDetails', type: `draft-${this._context.draft.type}` },
 				{ progress: { location: { viewId: this.host.id } } },
 			);
+			if (result === 'cancelled') throw new Error('Operation was canceled');
+
 			if (result == null) throw new Error('Error retrieving content');
 
 			params = { result: result.parsed };
@@ -879,6 +881,8 @@ export class PatchDetailsWebviewProvider
 				{ source: 'patchDetails', type: 'patch' },
 				{ progress: { location: { viewId: this.host.id } } },
 			);
+			if (result === 'cancelled') throw new Error('Operation was canceled');
+
 			if (result == null) throw new Error('Error retrieving content');
 
 			params = {