Optimize diff storage and improve Open Diff button behavior

- Stop storing full file contents in conversation history to reduce memory
- Compute and store only startLine/startLines for accurate line numbers on reload
- Open Diff button now only shows on last pending edit request
- Button uses stored diff data directly instead of re-reading file
- Hide button when edit result arrives (edit no longer pending)
- Show simple completion messages for Edit/MultiEdit/Write results
- Use virtual document scheme (claude-diff:) for read-only diff viewer
- Simplify tool result handling for Read/TodoWrite

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

src/extension.ts

@@ -6,6 +6,17 @@ import getHtml from './ui';
 
 const exec = util.promisify(cp.exec);
 
+// Storage for diff content (used by DiffContentProvider)
+const diffContentStore = new Map<string, string>();
+
+// Custom TextDocumentContentProvider for read-only diff views
+class DiffContentProvider implements vscode.TextDocumentContentProvider {
+	provideTextDocumentContent(uri: vscode.Uri): string {
+		const content = diffContentStore.get(uri.path);
+		return content || '';
+	}
+}
+
 export function activate(context: vscode.ExtensionContext) {
 	console.log('Claude Code Chat extension is being activated!');
 	const provider = new ClaudeChatProvider(context.extensionUri, context);
@@ -23,6 +34,10 @@ export function activate(context: vscode.ExtensionContext) {
 	const webviewProvider = new ClaudeChatWebviewProvider(context.extensionUri, context, provider);
 	vscode.window.registerWebviewViewProvider('claude-code-chat.chat', webviewProvider);
 
+	// Register custom content provider for read-only diff views
+	const diffProvider = new DiffContentProvider();
+	context.subscriptions.push(vscode.workspace.registerTextDocumentContentProvider('claude-diff', diffProvider));
+
 	// Listen for configuration changes
 	const configChangeDisposable = vscode.workspace.onDidChangeConfiguration(event => {
 		if (event.affectsConfiguration('claudeCodeChat.wsl')) {
@@ -300,6 +315,9 @@ class ClaudeChatProvider {
 			case 'openDiff':
 				this._openDiffEditor(message.oldContent, message.newContent, message.filePath);
 				return;
+			case 'openDiffByIndex':
+				this._openDiffByMessageIndex(message.messageIndex);
+				return;
 			case 'createImageFile':
 				this._createImageFile(message.imageData, message.imageType);
 				return;
@@ -657,7 +675,7 @@ class ClaudeChatProvider {
 		});
 	}
 
-	private _processJsonStreamData(jsonData: any) {
+	private async _processJsonStreamData(jsonData: any) {
 		switch (jsonData.type) {
 			case 'system':
 				if (jsonData.subtype === 'init') {
@@ -717,6 +735,7 @@ class ClaudeChatProvider {
 							// Show tool execution with better formatting
 							const toolInfo = `🔧 Executing: ${content.name}`;
 							let toolInput = '';
+							let fileContentBefore: string | undefined;
 
 							if (content.input) {
 								// Special formatting for TodoWrite to make it more readable
@@ -731,6 +750,44 @@ class ClaudeChatProvider {
 									// Send raw input to UI for formatting
 									toolInput = '';
 								}
+
+								// For Edit/MultiEdit/Write, read current file content (before state)
+								if ((content.name === 'Edit' || content.name === 'MultiEdit' || content.name === 'Write') && content.input.file_path) {
+									try {
+										const fileUri = vscode.Uri.file(content.input.file_path);
+										const fileData = await vscode.workspace.fs.readFile(fileUri);
+										fileContentBefore = Buffer.from(fileData).toString('utf8');
+									} catch {
+										// File might not exist yet (for Write), that's ok
+										fileContentBefore = '';
+									}
+								}
+							}
+
+							// Compute startLine(s) while we have the file content
+							let startLine: number | undefined;
+							let startLines: number[] | undefined;
+							if (fileContentBefore !== undefined) {
+								if (content.name === 'Edit' && content.input.old_string) {
+									const position = fileContentBefore.indexOf(content.input.old_string);
+									if (position !== -1) {
+										const textBefore = fileContentBefore.substring(0, position);
+										startLine = (textBefore.match(/\n/g) || []).length + 1;
+									} else {
+										startLine = 1;
+									}
+								} else if (content.name === 'MultiEdit' && content.input.edits) {
+									startLines = content.input.edits.map((edit: any) => {
+										if (edit.old_string) {
+											const position = fileContentBefore!.indexOf(edit.old_string);
+											if (position !== -1) {
+												const textBefore = fileContentBefore!.substring(0, position);
+												return (textBefore.match(/\n/g) || []).length + 1;
+											}
+										}
+										return 1;
+									});
+								}
 							}
 
 							// Show tool use and save to conversation
@@ -740,7 +797,10 @@ class ClaudeChatProvider {
 									toolInfo: toolInfo,
 									toolInput: toolInput,
 									rawInput: content.input,
-									toolName: content.name
+									toolName: content.name,
+									fileContentBefore: fileContentBefore,
+									startLine: startLine,
+									startLines: startLines
 								}
 							});
 						}
@@ -762,11 +822,25 @@ class ClaudeChatProvider {
 
 							const isError = content.is_error || false;
 
-							// Find the last tool use to get the tool name and input
+							// Find the last tool use to get the tool name, input, and computed startLine
 							const lastToolUse = this._currentConversation[this._currentConversation.length - 1]
 
 							const toolName = lastToolUse?.data?.toolName;
 							const rawInput = lastToolUse?.data?.rawInput;
+							const startLine = lastToolUse?.data?.startLine;
+							const startLines = lastToolUse?.data?.startLines;
+
+							// For Edit/MultiEdit/Write, read current file content (after state)
+							let fileContentAfter: string | undefined;
+							if ((toolName === 'Edit' || toolName === 'MultiEdit' || toolName === 'Write') && rawInput?.file_path && !isError) {
+								try {
+									const fileUri = vscode.Uri.file(rawInput.file_path);
+									const fileData = await vscode.workspace.fs.readFile(fileUri);
+									fileContentAfter = Buffer.from(fileData).toString('utf8');
+								} catch {
+									// File read failed, that's ok
+								}
+							}
 
 							// Don't send tool result for Read and TodoWrite tools unless there's an error
 							if ((toolName === 'Read' || toolName === 'TodoWrite') && !isError) {
@@ -791,7 +865,10 @@ class ClaudeChatProvider {
 										isError: isError,
 										toolUseId: content.tool_use_id,
 										toolName: toolName,
-										rawInput: rawInput
+										rawInput: rawInput,
+										fileContentAfter: fileContentAfter,
+										startLine: startLine,
+										startLines: startLines
 									}
 								});
 							}
@@ -1818,14 +1895,30 @@ class ClaudeChatProvider {
 			this._conversationStartTime = new Date().toISOString();
 		}
 
+		// The message index will be the current length (0-indexed position after push)
+		const messageIndex = this._currentConversation.length;
+
+		// For tool messages that support diff, include the message index
+		const messageToSend = (message.type === 'toolUse' || message.type === 'toolResult')
+			? { ...message, data: { ...message.data, messageIndex } }
+			: message;
+
 		// Send to UI using the helper method
-		this._postMessage(message);
+		this._postMessage(messageToSend);
+
+		// Strip fileContentBefore/fileContentAfter from saved data to reduce storage
+		// Keep startLine/startLines which are small and needed for accurate line numbers on reload
+		let dataToSave = message.data;
+		if (message.type === 'toolUse' || message.type === 'toolResult') {
+			const { fileContentBefore, fileContentAfter, ...rest } = message.data || {};
+			dataToSave = rest; // startLine and startLines are preserved in rest
+		}
 
 		// Save to conversation
 		this._currentConversation.push({
 			timestamp: new Date().toISOString(),
 			messageType: message.type,
-			data: message.data
+			data: dataToSave
 		});
 
 		// Persist conversation
@@ -2099,9 +2192,14 @@ class ClaudeChatProvider {
 							}
 						}
 
+						// For tool messages, include the message index so Open Diff buttons work
+						const messageData = (message.messageType === 'toolUse' || message.messageType === 'toolResult')
+							? { ...message.data, messageIndex: i }
+							: message.data;
+
 						this._postMessage({
 							type: message.messageType,
-							data: message.data
+							data: messageData
 						});
 						if (message.messageType === 'userInput') {
 							try {
@@ -2343,35 +2441,77 @@ class ClaudeChatProvider {
 		}
 	}
 
-	private async _openDiffEditor(oldContent: string, newContent: string, filePath: string) {
+	private async _openDiffByMessageIndex(messageIndex: number) {
 		try {
-			const storageUri = this._context.storageUri;
+			const message = this._currentConversation[messageIndex];
-			if (!storageUri) {
+			if (!message) {
-				vscode.window.showErrorMessage('No storage location available');
+				console.error('Message not found at index:', messageIndex);
 				return;
 			}
 
-			const baseName = path.basename(filePath);
+			const data = message.data;
-			const ext = path.extname(filePath);
+			const toolName = data.toolName;
-			const nameWithoutExt = baseName.slice(0, -ext.length) || baseName;
+			const rawInput = data.rawInput;
-			const timestamp = Date.now();
+			let filePath = rawInput?.file_path || '';
+			let oldContent = '';
+			let newContent = '';
 
-			// Create temp files in extension's storage directory
+			if (!filePath) {
-			const tempDirUri = vscode.Uri.joinPath(storageUri, 'diff-temp');
+				console.error('No file path found for message at index:', messageIndex);
-
+				return;
-			// Ensure temp directory exists
-			try {
-				await vscode.workspace.fs.createDirectory(tempDirUri);
-			} catch {
-				// Directory might already exist, ignore error
 			}
 
-			const oldUri = vscode.Uri.joinPath(tempDirUri, `${nameWithoutExt}.old.${timestamp}${ext}`);
+			// Read current file from disk - this is the "before" state since edit hasn't been applied yet
-			const newUri = vscode.Uri.joinPath(tempDirUri, `${nameWithoutExt}.new.${timestamp}${ext}`);
+			try {
+				const fileUri = vscode.Uri.file(filePath);
+				const fileData = await vscode.workspace.fs.readFile(fileUri);
+				oldContent = Buffer.from(fileData).toString('utf8');
+			} catch {
+				// File might not exist yet (for Write creating new file)
+				oldContent = '';
+			}
 
-			// Write content to temp files using VS Code filesystem API
+			// Compute "after" state by applying the edit to current file
-			await vscode.workspace.fs.writeFile(oldUri, Buffer.from(oldContent, 'utf8'));
+			if (toolName === 'Edit' && rawInput?.old_string && rawInput?.new_string) {
-			await vscode.workspace.fs.writeFile(newUri, Buffer.from(newContent, 'utf8'));
+				newContent = oldContent.replace(rawInput.old_string, rawInput.new_string);
+			} else if (toolName === 'MultiEdit' && rawInput?.edits) {
+				newContent = oldContent;
+				for (const edit of rawInput.edits) {
+					if (edit.old_string && edit.new_string) {
+						newContent = newContent.replace(edit.old_string, edit.new_string);
+					}
+				}
+			} else if (toolName === 'Write' && rawInput?.content) {
+				newContent = rawInput.content;
+			}
+
+			if (oldContent !== newContent) {
+				await this._openDiffEditor(oldContent, newContent, filePath);
+			} else {
+				vscode.window.showInformationMessage('No changes to show - the edit may have already been applied.');
+			}
+		} catch (error) {
+			console.error('Error opening diff by message index:', error);
+		}
+	}
+
+	private async _openDiffEditor(oldContent: string, newContent: string, filePath: string) {
+		try {
+			// oldContent and newContent are now full file contents passed from the webview
+			const baseName = path.basename(filePath);
+			const timestamp = Date.now();
+
+			// Create unique paths for the virtual documents
+			const oldPath = `/${timestamp}/old/${baseName}`;
+			const newPath = `/${timestamp}/new/${baseName}`;
+
+			// Store content in the global store for the content provider
+			diffContentStore.set(oldPath, oldContent);
+			diffContentStore.set(newPath, newContent);
+
+			// Create URIs with our custom scheme
+			const oldUri = vscode.Uri.parse(`claude-diff:${oldPath}`);
+			const newUri = vscode.Uri.parse(`claude-diff:${newPath}`);
 
 			// Ensure side-by-side diff mode is enabled
 			const diffConfig = vscode.workspace.getConfiguration('diffEditor');
@@ -2383,27 +2523,20 @@ class ClaudeChatProvider {
 			// Open diff editor
 			await vscode.commands.executeCommand('vscode.diff', oldUri, newUri, `${baseName} (Changes)`);
 
-			// Track which files need to be cleaned up
+			// Clean up stored content when documents are closed
-			const filesToCleanup = new Set([oldUri.toString(), newUri.toString()]);
+			const closeListener = vscode.workspace.onDidCloseTextDocument((doc) => {
-
+				if (doc.uri.toString() === oldUri.toString()) {
-			// Listen for document close events to clean up temp files
+					diffContentStore.delete(oldPath);
-			const closeListener = vscode.workspace.onDidCloseTextDocument(async (doc) => {
+				}
-				if (filesToCleanup.has(doc.uri.toString())) {
+				if (doc.uri.toString() === newUri.toString()) {
-					filesToCleanup.delete(doc.uri.toString());
+					diffContentStore.delete(newPath);
-					try {
+				}
-						await vscode.workspace.fs.delete(doc.uri, { useTrash: false });
+				// Dispose listener when both are cleaned up
-					} catch {
+				if (!diffContentStore.has(oldPath) && !diffContentStore.has(newPath)) {
-						// File might already be deleted, ignore
+					closeListener.dispose();
-					}
-
-					// Dispose listener when both files are cleaned up
-					if (filesToCleanup.size === 0) {
-						closeListener.dispose();
-					}
 				}
 			});
 
-			// Also add to disposables to clean up on extension deactivate
 			this._disposables.push(closeListener);
 		} catch (error) {
 			vscode.window.showErrorMessage(`Failed to open diff editor: ${error}`);

src/script.ts

@@ -15,22 +15,19 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 		let selectedFileIndex = -1;
 		let planModeEnabled = false;
 		let thinkingModeEnabled = false;
+		let lastPendingEditIndex = -1; // Track the last Edit/MultiEdit/Write toolUse without result
+		let lastPendingEditData = null; // Store diff data for the pending edit { filePath, oldContent, newContent }
 
-		// Storage for diff data to enable "Open Diff" functionality
+		// Open diff using stored data (no file read needed)
-		const diffDataStore = {};
+		function openDiffEditor() {
-
+			if (lastPendingEditData) {
-		function openDiffEditor(diffId) {
+				vscode.postMessage({
-			const data = diffDataStore[diffId];
+					type: 'openDiff',
-			if (!data) {
+					filePath: lastPendingEditData.filePath,
-				console.error('Diff data not found for id:', diffId);
+					oldContent: lastPendingEditData.oldContent,
-				return;
+					newContent: lastPendingEditData.newContent
+				});
 			}
-			vscode.postMessage({
-				type: 'openDiff',
-				oldContent: data.oldString,
-				newContent: data.newString,
-				filePath: data.filePath
-			});
 		}
 
 		function shouldAutoScroll(messagesDiv) {
@@ -179,12 +176,52 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 				} else {
 					// Format raw input with expandable content for long values
 					// Use diff format for Edit, MultiEdit, and Write tools, regular format for others
-					if (data.toolName === 'Edit') {
+					if (data.toolName === 'Edit' || data.toolName === 'MultiEdit' || data.toolName === 'Write') {
-						contentDiv.innerHTML = formatEditToolDiff(data.rawInput);
+						// Only show Open Diff button if we have fileContentBefore (live session, not reload)
-					} else if (data.toolName === 'MultiEdit') {
+						const showButton = data.fileContentBefore !== undefined && data.messageIndex >= 0;
-						contentDiv.innerHTML = formatMultiEditToolDiff(data.rawInput);
+
-					} else if (data.toolName === 'Write') {
+						// Hide any existing pending edit button before showing new one
-						contentDiv.innerHTML = formatWriteToolDiff(data.rawInput);
+						if (showButton && lastPendingEditIndex >= 0) {
+							const prevContent = document.querySelector('[data-edit-message-index="' + lastPendingEditIndex + '"]');
+							if (prevContent) {
+								const btn = prevContent.querySelector('.diff-open-btn');
+								if (btn) btn.style.display = 'none';
+							}
+							lastPendingEditData = null;
+						}
+
+						if (showButton) {
+							lastPendingEditIndex = data.messageIndex;
+							contentDiv.setAttribute('data-edit-message-index', data.messageIndex);
+
+							// Compute and store diff data for when button is clicked
+							const oldContent = data.fileContentBefore || '';
+							let newContent = oldContent;
+							if (data.toolName === 'Edit' && data.rawInput.old_string && data.rawInput.new_string) {
+								newContent = oldContent.replace(data.rawInput.old_string, data.rawInput.new_string);
+							} else if (data.toolName === 'MultiEdit' && data.rawInput.edits) {
+								for (const edit of data.rawInput.edits) {
+									if (edit.old_string && edit.new_string) {
+										newContent = newContent.replace(edit.old_string, edit.new_string);
+									}
+								}
+							} else if (data.toolName === 'Write' && data.rawInput.content) {
+								newContent = data.rawInput.content;
+							}
+							lastPendingEditData = {
+								filePath: data.rawInput.file_path,
+								oldContent: oldContent,
+								newContent: newContent
+							};
+						}
+
+						if (data.toolName === 'Edit') {
+							contentDiv.innerHTML = formatEditToolDiff(data.rawInput, data.fileContentBefore, showButton, data.startLine);
+						} else if (data.toolName === 'MultiEdit') {
+							contentDiv.innerHTML = formatMultiEditToolDiff(data.rawInput, data.fileContentBefore, showButton, data.startLines);
+						} else {
+							contentDiv.innerHTML = formatWriteToolDiff(data.rawInput, data.fileContentBefore, showButton);
+						}
 					} else {
 						contentDiv.innerHTML = formatToolInputUI(data.rawInput);
 					}
@@ -243,94 +280,42 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 			const messagesDiv = document.getElementById('messages');
 			const shouldScroll = shouldAutoScroll(messagesDiv);
 
-			// Show detailed diff for Edit, MultiEdit, and Write tools
+			// When result comes in for Edit/MultiEdit/Write, hide the Open Diff button on the request
-			if (!data.isError && data.rawInput) {
+			// since the edit has now been applied (no longer pending)
-				if (data.toolName === 'Edit' && data.rawInput.file_path && data.rawInput.old_string && data.rawInput.new_string) {
+			if (lastPendingEditIndex >= 0) {
-					const parsed = parseToolResult(data.content);
+				// Find and hide the button on the corresponding toolUse
-					const diffHTML = generateUnifiedDiffHTML(
+				const toolUseContent = document.querySelector('[data-edit-message-index="' + lastPendingEditIndex + '"]');
-						data.rawInput.old_string,
+				if (toolUseContent) {
-						data.rawInput.new_string,
+					const btn = toolUseContent.querySelector('.diff-open-btn');
-						data.rawInput.file_path,
+					if (btn) {
-						parsed.startLine
+						btn.style.display = 'none';
-					);
-
-					const messageDiv = document.createElement('div');
-					messageDiv.className = 'message tool-result';
-					messageDiv.innerHTML = diffHTML;
-					messagesDiv.appendChild(messageDiv);
-
-					scrollToBottomIfNeeded(messagesDiv, shouldScroll);
-					return;
-				} else if (data.toolName === 'MultiEdit' && data.rawInput.file_path && data.rawInput.edits) {
-					const parsed = parseToolResult(data.content);
-					let html = '';
-					const formattedPath = formatFilePath(data.rawInput.file_path);
-					html += '<div class="diff-file-path" onclick="openFileInEditor(\\\'' + escapeHtml(data.rawInput.file_path) + '\\\')">' + formattedPath + '</div>\\n';
-					html += '<div class="diff-container">';
-					html += '<div class="diff-header">' + data.rawInput.edits.length + ' edit' + (data.rawInput.edits.length > 1 ? 's' : '') + '</div>';
-
-					for (let i = 0; i < data.rawInput.edits.length; i++) {
-						const edit = data.rawInput.edits[i];
-						if (i > 0) html += '<div class="diff-edit-separator"></div>';
-						html += '<div class="edit-number">Edit #' + (i + 1) + '</div>';
-						const editDiff = generateUnifiedDiffHTML(edit.old_string, edit.new_string, data.rawInput.file_path, parsed.startLine);
-						html += editDiff;
 					}
-
-					html += '</div>';
-
-					const messageDiv = document.createElement('div');
-					messageDiv.className = 'message tool-result';
-					messageDiv.innerHTML = html;
-					messagesDiv.appendChild(messageDiv);
-
-					scrollToBottomIfNeeded(messagesDiv, shouldScroll);
-					return;
-				} else if (data.toolName === 'Write' && data.rawInput.file_path && data.rawInput.content) {
-					const parsed = parseToolResult(data.content);
-					const newLines = data.rawInput.content.split('\\n');
-					let html = '';
-					const formattedPath = formatFilePath(data.rawInput.file_path);
-					html += '<div class="diff-file-path" onclick="openFileInEditor(\\\'' + escapeHtml(data.rawInput.file_path) + '\\\')">' + formattedPath + '</div>\\n';
-					html += '<div class="diff-container">';
-					html += '<div class="diff-header">New file: Lines 1-' + newLines.length + '</div>';
-
-					for (let i = 0; i < newLines.length; i++) {
-						const lineNumStr = (i + 1).toString().padStart(3);
-						html += '<div class="diff-line added">+' + lineNumStr + '  ' + escapeHtml(newLines[i]) + '</div>';
-					}
-
-					html += '</div>';
-					html += '<div class="diff-summary">Summary: +' + newLines.length + ' line' + (newLines.length > 1 ? 's' : '') + ' added</div>';
-
-					const messageDiv = document.createElement('div');
-					messageDiv.className = 'message tool-result';
-					messageDiv.innerHTML = html;
-					messagesDiv.appendChild(messageDiv);
-
-					scrollToBottomIfNeeded(messagesDiv, shouldScroll);
-					return;
 				}
+				lastPendingEditIndex = -1;
+				lastPendingEditData = null;
 			}
 
-			// For Read and TodoWrite tools with hidden flag, just hide loading state and show completion message
+			// For Read and TodoWrite tools, just hide loading state (no result message needed)
-			if (data.hidden && (data.toolName === 'Read' || data.toolName === 'TodoWrite') && !data.isError) {
+			if ((data.toolName === 'Read' || data.toolName === 'TodoWrite') && !data.isError) {
-				return
+				return;
-				// Show completion message
+			}
-				const toolName = data.toolName;
+
+			// For Edit/MultiEdit/Write, show simple completion message (diff is already shown on request)
+			if ((data.toolName === 'Edit' || data.toolName === 'MultiEdit' || data.toolName === 'Write') && !data.isError) {
 				let completionText;
-				if (toolName === 'Read') {
+				if (data.toolName === 'Edit') {
-					completionText = '✅ Read completed';
+					completionText = '✅ Edit completed';
-				} else if (toolName === 'TodoWrite') {
+				} else if (data.toolName === 'MultiEdit') {
-					completionText = '✅ Update Todos completed';
+					completionText = '✅ MultiEdit completed';
 				} else {
-					completionText = '✅ ' + toolName + ' completed';
+					completionText = '✅ Write completed';
 				}
 				addMessage(completionText, 'system');
-				return; // Don't show the result message
+				scrollToBottomIfNeeded(messagesDiv, shouldScroll);
+				return;
 			}
 			
-			if(data.isError && data.content === "File has not been read yet. Read it first before writing to it."){
+			if(data.isError && data.content?.includes("File has not been read yet. Read it first before writing to it.")){
 				return addMessage("File has not been read yet. Let me read it first before writing to it.", 'system');
 			}
 
@@ -521,21 +506,15 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 		}
 
 		// Generate unified diff HTML with line numbers
-		function generateUnifiedDiffHTML(oldString, newString, filePath, startLine = 1) {
+		// showButton controls whether to show the "Open Diff" button
+		function generateUnifiedDiffHTML(oldString, newString, filePath, startLine = 1, showButton = false) {
 			const oldLines = oldString.split('\\n');
 			const newLines = newString.split('\\n');
 			const diff = computeLineDiff(oldLines, newLines);
 
-			// Generate unique ID for this diff
+			// Generate unique ID for this diff (used for truncation)
 			const diffId = 'diff_' + Math.random().toString(36).substr(2, 9);
 
-			// Store diff data for "Open Diff" functionality
-			diffDataStore[diffId] = {
-				oldString: oldString,
-				newString: newString,
-				filePath: filePath
-			};
-
 			let html = '';
 			const formattedPath = formatFilePath(filePath);
 
@@ -630,16 +609,18 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 			if (summary) {
 				html += '<div class="diff-summary-row">';
 				html += '<span class="diff-summary">Summary: ' + summary + '</span>';
-				html += '<button class="diff-open-btn" onclick="openDiffEditor(\\\'' + diffId + '\\\')" title="Open side-by-side diff in VS Code">';
+				if (showButton) {
-				html += '<svg width="14" height="14" viewBox="0 0 16 16"><rect x="1" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><rect x="9" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><line x1="2.5" y1="4" x2="5.5" y2="4" stroke="#f85149" stroke-width="1.5"/><line x1="2.5" y1="7" x2="5.5" y2="7" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="2.5" y1="10" x2="5.5" y2="10" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="4" x2="13.5" y2="4" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="7" x2="13.5" y2="7" stroke="#3fb950" stroke-width="1.5"/><line x1="10.5" y1="10" x2="13.5" y2="10" stroke="#3fb950" stroke-width="1.5"/></svg>';
+					html += '<button class="diff-open-btn" onclick="openDiffEditor()" title="Open side-by-side diff in VS Code">';
-				html += 'Open Diff</button>';
+					html += '<svg width="14" height="14" viewBox="0 0 16 16"><rect x="1" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><rect x="9" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><line x1="2.5" y1="4" x2="5.5" y2="4" stroke="#f85149" stroke-width="1.5"/><line x1="2.5" y1="7" x2="5.5" y2="7" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="2.5" y1="10" x2="5.5" y2="10" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="4" x2="13.5" y2="4" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="7" x2="13.5" y2="7" stroke="#3fb950" stroke-width="1.5"/><line x1="10.5" y1="10" x2="13.5" y2="10" stroke="#3fb950" stroke-width="1.5"/></svg>';
+					html += 'Open Diff</button>';
+				}
 				html += '</div>';
 			}
 
 			return html;
 		}
 
-		function formatEditToolDiff(input) {
+		function formatEditToolDiff(input, fileContentBefore, showButton = false, providedStartLine = null) {
 			if (!input || typeof input !== 'object') {
 				return formatToolInputUI(input);
 			}
@@ -649,11 +630,21 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 				return formatToolInputUI(input);
 			}
 
-			// Show full diff (line numbers will be approximate until result comes back)
+			// Use provided startLine if available (from saved data), otherwise compute from fileContentBefore
-			return generateUnifiedDiffHTML(input.old_string, input.new_string, input.file_path, 1);
+			let startLine = providedStartLine || 1;
+			if (!providedStartLine && fileContentBefore) {
+				const position = fileContentBefore.indexOf(input.old_string);
+				if (position !== -1) {
+					// Count newlines before the match to get line number
+					const textBefore = fileContentBefore.substring(0, position);
+					startLine = (textBefore.match(/\\n/g) || []).length + 1;
+				}
+			}
+
+			return generateUnifiedDiffHTML(input.old_string, input.new_string, input.file_path, startLine, showButton);
 		}
 
-		function formatMultiEditToolDiff(input) {
+		function formatMultiEditToolDiff(input, fileContentBefore, showButton = false, providedStartLines = null) {
 			if (!input || typeof input !== 'object') {
 				return formatToolInputUI(input);
 			}
@@ -674,39 +665,60 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 					if (index > 0) {
 						html += '<div class="diff-edit-separator"></div>';
 					}
+
+					// Use provided startLine if available, otherwise compute from fileContentBefore
+					let startLine = (providedStartLines && providedStartLines[index]) || 1;
+					if (!providedStartLines && fileContentBefore) {
+						const position = fileContentBefore.indexOf(edit.old_string);
+						if (position !== -1) {
+							const textBefore = fileContentBefore.substring(0, position);
+							startLine = (textBefore.match(/\\n/g) || []).length + 1;
+						}
+					}
+
 					const oldLines = edit.old_string.split('\\n');
 					const newLines = edit.new_string.split('\\n');
 					const diff = computeLineDiff(oldLines, newLines);
 
 					html += '<div class="diff-container">';
-					html += '<div class="diff-header">Edit ' + (index + 1) + '</div>';
+					html += '<div class="diff-header">Edit ' + (index + 1) + ' (Line ' + startLine + ')</div>';
 
-					let addedCount = 0;
+					let lineNum = startLine;
-					let removedCount = 0;
 					for (const change of diff) {
 						let prefix, cssClass;
 						if (change.type === 'context') {
 							prefix = ' ';
 							cssClass = 'context';
+							lineNum++;
 						} else if (change.type === 'added') {
 							prefix = '+';
 							cssClass = 'added';
-							addedCount++;
+							lineNum++;
 						} else {
 							prefix = '-';
 							cssClass = 'removed';
-							removedCount++;
 						}
-						html += '<div class="diff-line ' + cssClass + '">' + prefix + '    ' + escapeHtml(change.content) + '</div>';
+						const lineNumStr = (change.type === 'removed' ? '' : lineNum - 1).toString().padStart(3);
+						html += '<div class="diff-line ' + cssClass + '">' + prefix + lineNumStr + '  ' + escapeHtml(change.content) + '</div>';
 					}
 					html += '</div>';
 				}
 			});
 
+			// Add summary row with Open Diff button
+			html += '<div class="diff-summary-row">';
+			html += '<span class="diff-summary">Summary: ' + input.edits.length + ' edit' + (input.edits.length > 1 ? 's' : '') + '</span>';
+			if (showButton) {
+				html += '<button class="diff-open-btn" onclick="openDiffEditor()" title="Open side-by-side diff in VS Code">';
+				html += '<svg width="14" height="14" viewBox="0 0 16 16"><rect x="1" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><rect x="9" y="1" width="6" height="14" rx="1" fill="none" stroke="currentColor" stroke-opacity="0.5"/><line x1="2.5" y1="4" x2="5.5" y2="4" stroke="#f85149" stroke-width="1.5"/><line x1="2.5" y1="7" x2="5.5" y2="7" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="2.5" y1="10" x2="5.5" y2="10" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="4" x2="13.5" y2="4" stroke="currentColor" stroke-opacity="0.4" stroke-width="1.5"/><line x1="10.5" y1="7" x2="13.5" y2="7" stroke="#3fb950" stroke-width="1.5"/><line x1="10.5" y1="10" x2="13.5" y2="10" stroke="#3fb950" stroke-width="1.5"/></svg>';
+				html += 'Open Diff</button>';
+			}
+			html += '</div>';
+
 			return html;
 		}
 
-		function formatWriteToolDiff(input) {
+		function formatWriteToolDiff(input, fileContentBefore, showButton = false) {
 			if (!input || typeof input !== 'object') {
 				return formatToolInputUI(input);
 			}
@@ -716,8 +728,11 @@ const getScript = (isTelemetryEnabled: boolean) => `<script>
 				return formatToolInputUI(input);
 			}
 
-			// Show full content as added lines (new file)
+			// fileContentBefore may be empty string if new file, or existing content if overwriting
-			return generateUnifiedDiffHTML('', input.content, input.file_path, 1);
+			const fullFileBefore = fileContentBefore || '';
+
+			// Show full content as added lines (new file or replacement)
+			return generateUnifiedDiffHTML(fullFileBefore, input.content, input.file_path, 1, showButton);
 		}
 
 		function escapeHtml(text) {