Adds Undo functionality to AI generate commits and AI rebase (#4371)

* Add undo command configuration to package manifest

* Add types and interfaces for undo functionality

* Implement UndoGenerateRebaseCommand class

* Add undo functionality to generateRebase function

* Removes unused fn (and updates telemetry docs)

Author: axosoft-ramint

contributions.json

@@ -307,6 +307,10 @@
 			"label": "Switch AI Provider/Model",
 			"commandPalette": "gitlens:enabled && gitlens:gk:organization:ai:enabled"
 		},
+		"gitlens.ai.undoGenerateRebase": {
+			"label": "Undo AI Generated Changes (Preview)...",
+			"icon": "$(discard)"
+		},
 		"gitlens.annotations.nextChange": {
 			"label": "Next Change",
 			"icon": "$(arrow-down)"

docs/telemetry-events.md

@@ -6264,6 +6264,11 @@
 				"title": "Switch AI Provider/Model",
 				"category": "GitLens"
 			},
+			{
+				"command": "gitlens.ai.undoGenerateRebase",
+				"title": "Undo AI Generated Changes (Preview)...",
+				"icon": "$(discard)"
+			},
 			{
 				"command": "gitlens.annotations.nextChange",
 				"title": "Next Change",
@@ -10772,6 +10777,10 @@
 					"command": "gitlens.ai.switchProvider",
 					"when": "gitlens:enabled && gitlens:gk:organization:ai:enabled"
 				},
+				{
+					"command": "gitlens.ai.undoGenerateRebase",
+					"when": "false"
+				},
 				{
 					"command": "gitlens.annotations.nextChange",
 					"when": "false"

package.json

@@ -1,18 +1,19 @@
 import type { CancellationToken, ProgressOptions } from 'vscode';
-import { ProgressLocation } from 'vscode';
+import { ProgressLocation, window } from 'vscode';
 import type { Source } from '../constants.telemetry';
 import type { Container } from '../container';
 import type { MarkdownContentMetadata } from '../documents/markdown';
 import { getMarkdownHeaderContent } from '../documents/markdown';
 import type { GitRepositoryService } from '../git/gitRepositoryService';
-import type { GitReference } from '../git/models/reference';
+import type { GitStashCommit } from '../git/models/commit';
+import type { GitReference, GitStashReference } from '../git/models/reference';
 import { uncommitted } from '../git/models/revision';
 import { createReference } from '../git/utils/reference.utils';
 import { showGenericErrorMessage } from '../messages';
 import type { AIRebaseResult } from '../plus/ai/aiProviderService';
 import { showComparisonPicker } from '../quickpicks/comparisonPicker';
 import { getRepositoryOrShowPicker } from '../quickpicks/repositoryPicker';
-import { command } from '../system/-webview/command';
+import { command, executeCommand } from '../system/-webview/command';
 import { showMarkdownPreview } from '../system/-webview/markdown';
 import { Logger } from '../system/logger';
 import { escapeMarkdownCodeBlocks } from '../system/markdown';
@@ -36,6 +37,16 @@ export interface GenerateCommitsCommandArgs {
 	source?: Source;
 }
 
+export interface UndoGenerateRebaseCommandArgs {
+	repoPath?: string;
+	generatedHeadRef?: GitReference;
+	previousHeadRef?: GitReference;
+	generatedStashRef?: GitStashReference;
+	generatedBranchName?: string;
+	undoCommand?: `gitlens.ai.generateCommits` | `gitlens.ai.generateRebase`;
+	source?: Source;
+}
+
 /**
  * Represents a file patch with its diff header and hunk contents
  */
@@ -163,6 +174,158 @@ export class GenerateRebaseCommand extends GlCommandBase {
 	}
 }
 
+@command()
+export class UndoGenerateRebaseCommand extends GlCommandBase {
+	constructor(private readonly container: Container) {
+		super('gitlens.ai.undoGenerateRebase');
+	}
+
+	async execute(args?: UndoGenerateRebaseCommandArgs): Promise<void> {
+		try {
+			if (!args?.undoCommand) {
+				Logger.error(undefined, 'UndoGenerateRebaseCommand', 'execute', 'Missing undoCommand parameter');
+				void window.showErrorMessage('Unable to undo: Missing command information');
+				return;
+			}
+
+			if (args.undoCommand === 'gitlens.ai.generateRebase') {
+				await this.undoGenerateRebase(args);
+			} else if (args.undoCommand === 'gitlens.ai.generateCommits') {
+				await this.undoGenerateCommits(args);
+			} else {
+				const unknownCommand = args.undoCommand as string;
+				Logger.error(
+					undefined,
+					'UndoGenerateRebaseCommand',
+					'execute',
+					`Unknown undoCommand: ${unknownCommand}`,
+				);
+				void window.showErrorMessage(`Unable to undo: Unknown command ${unknownCommand}`);
+			}
+		} catch (ex) {
+			Logger.error(ex, 'UndoGenerateRebaseCommand', 'execute');
+			void showGenericErrorMessage('Unable to undo operation');
+		}
+	}
+
+	private async undoGenerateRebase(args: UndoGenerateRebaseCommandArgs): Promise<void> {
+		// Check required parameters
+		if (!args.repoPath || !args.generatedBranchName) {
+			Logger.error(
+				undefined,
+				'UndoGenerateRebaseCommand',
+				'undoGenerateRebase',
+				'Missing required parameters: repoPath or generatedBranchName',
+			);
+			void window.showErrorMessage('Unable to undo rebase: Missing required information');
+			return;
+		}
+
+		const svc = this.container.git.getRepositoryService(args.repoPath);
+
+		// Warn user and ask for confirmation
+		const confirm = { title: 'Delete Branch' };
+		const cancel = { title: 'Cancel', isCloseAffordance: true };
+		const result = await window.showWarningMessage(
+			`This will delete the branch '${args.generatedBranchName}'. This action cannot be undone.\n\nAre you sure you want to continue?`,
+			{ modal: true },
+			confirm,
+			cancel,
+		);
+
+		if (result !== confirm) return;
+
+		try {
+			// Try to delete the branch
+			await svc.branches.deleteLocalBranch?.(args.generatedBranchName, { force: true });
+			void window.showInformationMessage(
+				`Successfully deleted branch '${args.generatedBranchName}'. Undo completed.`,
+			);
+		} catch (ex) {
+			Logger.error(ex, 'UndoGenerateRebaseCommand', 'undoGenerateRebase');
+
+			// Check if it's because the user is on the branch or other specific errors
+			const errorMessage = ex instanceof Error ? ex.message : String(ex);
+			if (errorMessage.includes('checked out') || errorMessage.includes('current branch')) {
+				void window.showErrorMessage(
+					`Cannot delete branch '${args.generatedBranchName}' because it is currently checked out.`,
+				);
+			} else {
+				void window.showErrorMessage(`Failed to delete branch '${args.generatedBranchName}': ${errorMessage}`);
+			}
+		}
+	}
+
+	private async undoGenerateCommits(args: UndoGenerateRebaseCommandArgs): Promise<void> {
+		// Check required parameters
+		if (!args.repoPath || !args.generatedHeadRef || !args.previousHeadRef || !args.generatedStashRef) {
+			Logger.error(
+				undefined,
+				'UndoGenerateRebaseCommand',
+				'undoGenerateCommits',
+				'Missing required parameters: repoPath, generatedHeadRef, previousHeadRef, or generatedStashRef',
+			);
+			void window.showErrorMessage('Unable to undo commits: Missing required information');
+			return;
+		}
+
+		const svc = this.container.git.getRepositoryService(args.repoPath);
+
+		try {
+			// Check if current HEAD matches the generated HEAD
+			const log = await svc.commits.getLog(undefined, { limit: 1 });
+			const currentCommit = log?.commits.values().next().value;
+			if (!currentCommit || currentCommit.sha !== args.generatedHeadRef.ref) {
+				void window.showErrorMessage(
+					'Cannot undo commits: Your HEAD reference has changed since the commits were generated. Please ensure you are on the correct commit.',
+				);
+				return;
+			}
+
+			// Warn user and ask for confirmation
+			const confirm = { title: 'Undo Commits' };
+			const cancel = { title: 'Cancel', isCloseAffordance: true };
+			const result = await window.showWarningMessage(
+				`This will reset your current branch to ${args.previousHeadRef.ref} and restore your previous working changes. Any work done after generating commits will be lost.\n\nAre you sure you want to continue?`,
+				{ modal: true },
+				confirm,
+				cancel,
+			);
+
+			if (result !== confirm) return;
+
+			// Check if there are working tree changes and stash them
+			const status = await svc.status.getStatus();
+			if (status?.files && status.files.length > 0) {
+				await svc.stash?.saveStash(undefined, undefined, { includeUntracked: true });
+			}
+
+			// Reset hard to the previous HEAD
+			await svc.reset(args.previousHeadRef.ref, { hard: true });
+
+			// Apply the generated stash
+			try {
+				await svc.stash?.applyStash(args.generatedStashRef.ref);
+			} catch (ex) {
+				Logger.error(ex, 'UndoGenerateRebaseCommand', 'undoGenerateCommits', 'Failed to apply stash');
+				void window.showWarningMessage(
+					`Reset completed, but failed to apply the original stash: ${
+						ex instanceof Error ? ex.message : String(ex)
+					}`,
+				);
+				return;
+			}
+
+			void window.showInformationMessage(
+				'Successfully undid the generated commits and restored your previous working changes. Undo completed.',
+			);
+		} catch (ex) {
+			Logger.error(ex, 'UndoGenerateRebaseCommand', 'undoGenerateCommits');
+			void window.showErrorMessage(`Failed to undo commits: ${ex instanceof Error ? ex.message : String(ex)}`);
+		}
+	}
+}
+
 export async function generateRebase(
 	container: Container,
 	svc: GitRepositoryService,
@@ -191,18 +354,62 @@ export async function generateRebase(
 
 		let generateType: 'commits' | 'rebase' = 'rebase';
 		let headRefSlug = head.ref;
+		let generatedBranchName: string | undefined;
+		let previousHeadRef: GitReference | undefined;
+		let generatedHeadRef: GitReference | undefined;
+		let generatedStashRef: GitStashReference | undefined;
+		let stashCommit: GitStashCommit | undefined;
+		let previousStashCommit: GitStashCommit | undefined;
+
 		const shas = await repo.git.patch?.createUnreachableCommitsFromPatches(base.ref, diffInfo);
 		if (shas?.length) {
 			if (head.ref === uncommitted) {
 				generateType = 'commits';
 				headRefSlug = 'uncommitted';
+
+				// Capture the current HEAD before making changes
+				const log = await svc.commits.getLog(undefined, { limit: 1 });
+				if (log?.commits.size) {
+					const currentCommit = log.commits.values().next().value;
+					if (currentCommit) {
+						previousHeadRef = createReference(currentCommit.sha, svc.path, { refType: 'revision' });
+					}
+				}
+
+				let stash = await svc.stash?.getStash();
+				if (stash?.stashes.size) {
+					const latestStash = stash.stashes.values().next().value;
+					if (latestStash) {
+						previousStashCommit = latestStash;
+					}
+				}
+
 				// stash the working changes
 				await svc.stash?.saveStash(undefined, undefined, { includeUntracked: true });
-				// await repo.git.checkout?.(shas[shas.length - 1]);
+
+				// Get the latest stash reference
+				stash = await svc.stash?.getStash();
+				if (stash?.stashes.size) {
+					stashCommit = stash.stashes.values().next().value;
+					if (stashCommit) {
+						generatedStashRef = createReference(stashCommit.ref, svc.path, {
+							refType: 'stash',
+							name: stashCommit.stashName,
+							number: stashCommit.stashNumber,
+							message: stashCommit.message,
+							stashOnRef: stashCommit.stashOnRef,
+						});
+					}
+				}
+
 				// reset the current branch to the new shas
 				await svc.reset(shas[shas.length - 1], { hard: true });
+
+				// Capture the new HEAD after reset
+				generatedHeadRef = createReference(shas[shas.length - 1], svc.path, { refType: 'revision' });
 			} else {
-				await svc.branches.createBranch?.(`rebase/${head.ref}-${Date.now()}`, shas[shas.length - 1]);
+				generatedBranchName = `rebase/${head.ref}-${Date.now()}`;
+				await svc.branches.createBranch?.(generatedBranchName, shas[shas.length - 1]);
 			}
 		}
 
@@ -214,6 +421,40 @@ export async function generateRebase(
 		);
 
 		showMarkdownPreview(documentUri);
+
+		// Show success notification with Undo button
+		const undoButton = { title: 'Undo' };
+		const resultNotification = await window.showInformationMessage(
+			generateType === 'commits'
+				? 'Successfully generated commits from your working changes.'
+				: 'Successfully generated rebase branch.',
+			undoButton,
+		);
+
+		if (resultNotification === undoButton) {
+			if (generateType === 'commits') {
+				// Undo GenerateCommitsCommand
+				void executeCommand('gitlens.ai.undoGenerateRebase', {
+					undoCommand: 'gitlens.ai.generateCommits',
+					repoPath: svc.path,
+					generatedHeadRef: generatedHeadRef,
+					previousHeadRef: previousHeadRef,
+					generatedStashRef:
+						stashCommit != null && stashCommit.ref !== previousStashCommit?.ref
+							? generatedStashRef
+							: undefined,
+					source: source,
+				} satisfies UndoGenerateRebaseCommandArgs);
+			} else {
+				// Undo GenerateRebaseCommand
+				void executeCommand('gitlens.ai.undoGenerateRebase', {
+					undoCommand: 'gitlens.ai.generateRebase',
+					repoPath: svc.path,
+					generatedBranchName: generatedBranchName,
+					source: source,
+				} satisfies UndoGenerateRebaseCommandArgs);
+			}
+		}
 	} catch (ex) {
 		Logger.error(ex, 'GenerateRebaseCommand', 'execute');
 		void showGenericErrorMessage('Unable to parse rebase result');

src/commands/generateRebase.ts

@@ -21,6 +21,7 @@ export type ContributedCommands =
 	| 'gitlens.ai.generateCommits:views'
 	| 'gitlens.ai.rebaseOntoCommit:graph'
 	| 'gitlens.ai.rebaseOntoCommit:views'
+	| 'gitlens.ai.undoGenerateRebase'
 	| 'gitlens.annotations.nextChange'
 	| 'gitlens.annotations.previousChange'
 	| 'gitlens.computingFileAnnotations'