fix: Preserve tab state (pinning/position) during diff

Author: seedlord

src/integrations/editor/DiffViewProvider.ts

@@ -19,6 +19,7 @@ export class DiffViewProvider {
 	private documentWasOpen = false
 	private originalViewColumn?: vscode.ViewColumn // Store the original view column
 	private userFocusedEditorInfo?: { uri: vscode.Uri; viewColumn: vscode.ViewColumn } // Store user's focus before diff
+	private originalTabState?: { uri: vscode.Uri; isPinned: boolean; viewColumn: vscode.ViewColumn; index: number } // Store original tab state if open
 	private relPath?: string
 	private newContent?: string
 	private activeDiffEditor?: vscode.TextEditor
@@ -67,10 +68,23 @@ export class DiffViewProvider {
 				(tab) => tab.input instanceof vscode.TabInputText && arePathsEqual(tab.input.uri.fsPath, absolutePath),
 			)
 		for (const tab of tabs) {
+			// Store state BEFORE closing
+			if (tab.input instanceof vscode.TabInputText) {
+				// Ensure it's a text tab to access URI safely
+				this.originalTabState = {
+					uri: tab.input.uri,
+					isPinned: tab.isPinned,
+					viewColumn: tab.group.viewColumn,
+					index: tab.group.tabs.indexOf(tab), // Correct way to get the index
+				}
+				this.documentWasOpen = true // Set flag indicating we found an open tab state
+				console.log("Original tab state saved:", this.originalTabState) // Optional: for debugging
+			}
+
 			if (!tab.isDirty) {
 				await vscode.window.tabGroups.close(tab)
 			}
-			this.documentWasOpen = true
+			// Removed this.documentWasOpen = true from here as it's set when state is saved
 		}
 
 		// Store the currently focused editor before opening the diff
@@ -174,9 +188,13 @@ export class DiffViewProvider {
 		await vscode.window.showTextDocument(vscode.Uri.file(absolutePath), { preview: false })
 		await this.closeAllDiffViews()
 
-		// If the original document was open, try to focus it.
-		// VS Code should handle showing the updated content automatically since the file was saved.
-		await this._focusOriginalDocument(absolutePath, this.originalViewColumn)
+		// Restore original tab state if it existed, otherwise handle focus normally
+		if (this.originalTabState) {
+			await this._restoreOriginalTabState()
+		} else {
+			// Fallback to old focus logic only if no original tab state was saved (e.g., new file)
+			await this._focusOriginalDocument(absolutePath, undefined) // Pass undefined as viewColumn is part of originalTabState now
+		}
 
 		/*
 		Getting diagnostics before and after the file edit is a better approach than
@@ -260,10 +278,13 @@ export class DiffViewProvider {
 			// Close the diff view first
 			await this.closeAllDiffViews()
 
-			// If the document was originally open, ensure it's focused.
-			// The revert logic already applied the original content and saved.
-			await this._focusOriginalDocument(absolutePath, this.originalViewColumn)
-
+			// Restore original tab state if it existed, otherwise handle focus normally
+			if (this.originalTabState) {
+				await this._restoreOriginalTabState()
+			} else {
+				// Fallback to old focus logic only if no original tab state was saved (e.g., new file)
+				await this._focusOriginalDocument(absolutePath, undefined) // Pass undefined as viewColumn is part of originalTabState now
+			}
 		}
 		// edit is done
 		await this.reset()
@@ -310,16 +331,61 @@ export class DiffViewProvider {
 			const disposable = vscode.window.onDidChangeActiveTextEditor((editor) => {
 				if (editor && arePathsEqual(editor.document.uri.fsPath, uri.fsPath)) {
 					disposable.dispose()
-					// Diff editor is now active, resolve the promise
-					resolve(editor)
+					// Diff editor is now active
+					// Pin the diff editor if the original tab was pinned
+					const pinAndMoveIfNeeded = async () => {
+						if (this.originalTabState?.isPinned) {
+							try {
+								await vscode.commands.executeCommand("workbench.action.pinEditor")
+								console.log("Diff editor pinned.")
+								// Add a small delay after pinning before moving
+								await new Promise((resolve) => setTimeout(resolve, 50))
+
+								// Move the pinned diff editor to the original index
+								const targetGroup = vscode.window.tabGroups.all.find(
+									(group) => group.viewColumn === this.originalTabState?.viewColumn,
+								)
+								const index = this.originalTabState.index
+								if (targetGroup && index >= 0) {
+									// Check if index is still valid after potential async operations
+									if (index < targetGroup.tabs.length) {
+										await vscode.commands.executeCommand("moveActiveEditor", {
+											to: "position",
+											value: index + 1, // 1-based index
+										})
+										console.log(`Diff editor moved to index ${index}.`)
+									} else {
+										console.warn(
+											`Diff editor move skipped: Index ${index} out of bounds after pinning/delay.`,
+										)
+									}
+								} else {
+									console.warn(
+										`Could not move diff editor: Invalid index (${index}) or target group not found.`,
+									)
+								}
+							} catch (err) {
+								console.error("Failed to pin or move diff editor:", err)
+							}
+						}
+						// Resolve the promise regardless of pin/move success
+						resolve(editor)
+					}
+					pinAndMoveIfNeeded() // Execute async pin/move logic
 				}
 			})
 			const options: vscode.TextDocumentShowOptions = {
 				// preserveFocus: true, // Removed to prevent focus issues
 			}
-			if (this.originalViewColumn !== undefined) {
+			// Use viewColumn from originalTabState if available
+			if (this.originalTabState?.viewColumn !== undefined) {
+				options.viewColumn = this.originalTabState.viewColumn
+			} else if (this.originalViewColumn !== undefined) {
+				// Fallback to originalViewColumn if originalTabState is not set (e.g., file wasn't open)
 				options.viewColumn = this.originalViewColumn
 			}
+
+			// Execute the diff command first
 			vscode.commands.executeCommand(
 				"vscode.diff",
 				vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${fileName}`).with({
@@ -381,7 +447,7 @@ export class DiffViewProvider {
 
 	private async _focusOriginalDocument(
 		absolutePath: string,
-		viewColumn: vscode.ViewColumn | undefined, // Note: viewColumn here is the original view column of the *modified* file's tab, if it was open
+		_viewColumn: vscode.ViewColumn | undefined, // Prefixed as unused now
 	): Promise<void> {
 		let focusRestoredOrHandled = false
 
@@ -421,34 +487,86 @@ export class DiffViewProvider {
 			}
 		}
 
-		// Fallback logic for *existing* documents (only runs if focus restoration failed AND it wasn't a new file)
-		if (!focusRestoredOrHandled && this.documentWasOpen && viewColumn) {
-			console.log("Executing fallback logic for existing document as primary focus restore failed.")
-			// Fallback 1: Try to focus the editor tab corresponding to the *modified* file
-			const originalEditor = vscode.window.visibleTextEditors.find(
-				(editor) => arePathsEqual(editor.document.uri.fsPath, absolutePath) && editor.viewColumn === viewColumn,
+		// Fallback logic for existing documents is now handled by _restoreOriginalTabState
+		// Keep the final log message for clarity
+		else if (!focusRestoredOrHandled) {
+			console.log(
+				"No specific focus action taken by _focusOriginalDocument (focus restore might have succeeded, or it was a new file handled above).",
 			)
-			if (originalEditor) {
-				const position = new vscode.Position(0, 0)
-				originalEditor.revealRange(new vscode.Range(position, position), vscode.TextEditorRevealType.AtTop)
-				console.log("Focus set to modified file's original editor (fallback 1):", absolutePath)
+		}
+	}
+
+	private async _restoreOriginalTabState(): Promise<void> {
+		if (!this.originalTabState) {
+			console.log("No original tab state to restore.")
+			return
+		}
+
+		console.log("Attempting to restore original tab state:", this.originalTabState)
+		const { uri, viewColumn, isPinned, index } = this.originalTabState
+
+		try {
+			// 1. Show the document in the correct view column
+			// Prefix 'editor' as it's not directly used after assignment (focus happens implicitly)
+			const _editor = await vscode.window.showTextDocument(uri, {
+				viewColumn: viewColumn,
+				preview: false, // Ensure it's not a preview tab
+				preserveFocus: false, // Ensure this editor gets focus initially for commands
+			})
+			console.log("Document shown:", uri.fsPath)
+
+			// Small delay to allow VS Code to potentially update tab state after showing
+			await new Promise((resolve) => setTimeout(resolve, 100)) // 100ms delay
+
+			// 2. Pin the editor if necessary
+			if (isPinned) {
+				await vscode.commands.executeCommand("workbench.action.pinEditor")
+				console.log("Editor pinned.")
+				// Another small delay might be needed after pinning before moving
+				await new Promise((resolve) => setTimeout(resolve, 50))
+			}
+
+			// 3. Move the editor to the original index
+			// Ensure the target index is valid. VS Code's move command is 1-based.
+			// We need to find the current group to check tab count.
+			const targetGroup = vscode.window.tabGroups.all.find((group) => group.viewColumn === viewColumn)
+			if (targetGroup && index >= 0 && index < targetGroup.tabs.length) {
+				// The 'moveActiveEditor' command uses 1-based index for 'value'
+				await vscode.commands.executeCommand("moveActiveEditor", { to: "position", value: index + 1 })
+				console.log(`Editor moved to index ${index}.`)
 			} else {
-				// Fallback 2: Open the modified file if its editor wasn't found
+				console.warn(`Could not move editor: Invalid index (${index}) or target group not found.`)
+			}
+
+			// 4. Restore original user focus if it was different from the restored tab
+			if (this.userFocusedEditorInfo && !arePathsEqual(this.userFocusedEditorInfo.uri.fsPath, uri.fsPath)) {
 				try {
-					await vscode.window.showTextDocument(vscode.Uri.file(absolutePath), {
-						preview: false,
-						viewColumn: viewColumn,
-						preserveFocus: false, // Force focus
+					await vscode.window.showTextDocument(this.userFocusedEditorInfo.uri, {
+						viewColumn: this.userFocusedEditorInfo.viewColumn,
+						preserveFocus: false, // Force focus back
 					})
-					console.log("Opened modified file's editor (fallback 2):", absolutePath)
-				} catch (error) {
-					console.error("Failed to show modified document (fallback 2):", error)
+					console.log(
+						"Focus restored to originally focused editor (after tab state restore):",
+						this.userFocusedEditorInfo.uri.fsPath,
+					)
+				} catch (focusError) {
+					console.warn("Failed to restore original user focus after tab state restore:", focusError)
+					// If restoring original focus fails, at least the target tab should be focused.
+					await vscode.window.showTextDocument(uri, { viewColumn: viewColumn, preserveFocus: false })
 				}
+			} else {
+				// Ensure the restored tab keeps focus if no other editor was focused or if it was the focused one
+				await vscode.window.showTextDocument(uri, { viewColumn: viewColumn, preserveFocus: false })
+				console.log("Focus kept on restored tab:", uri.fsPath)
+			}
+		} catch (error) {
+			console.error("Error restoring original tab state:", error)
+			// Fallback: Just try to show the document without state restoration
+			try {
+				await vscode.window.showTextDocument(uri, { viewColumn: viewColumn, preview: false })
+			} catch (fallbackError) {
+				console.error("Fallback showTextDocument also failed:", fallbackError)
 			}
-		} else if (!focusRestoredOrHandled) {
-			console.log(
-				"No specific focus action taken (focus restore might have succeeded, or it was a new file handled above, or fallbacks for existing file didn't apply).",
-			)
 		}
 	}
 
@@ -459,7 +577,8 @@ export class DiffViewProvider {
 		this.originalContent = undefined
 		this.createdDirs = []
 		this.documentWasOpen = false
-		this.originalViewColumn = undefined // Reset stored view column
+		this.originalViewColumn = undefined // Reset stored view column - Keep for potential fallback? Replaced by originalTabState.viewColumn mostly.
+		this.originalTabState = undefined // Reset stored tab state
 		this.userFocusedEditorInfo = undefined // Reset stored user focus info
 		this.activeDiffEditor = undefined
 		this.fadedOverlayController = undefined