Fix code duplication issues in streaming diff animation

- Eliminated duplicate temp file creation and diff view opening logic in StreamingDiffController
- Consolidated session creation logic in DiffAnimationHandler
- Created reusable cleanup helper function in Messages
- Removed duplicate getCursorState function
- Reduced codebase by 53+ lines of duplicate code

Author: MarcoWang3

packages/amazonq/src/lsp/chat/diffAnimation/diffAnimationHandler.ts

@@ -252,21 +252,28 @@ export class DiffAnimationHandler implements vscode.Disposable {
         }
     }
 
-    public async startStreamingDiffSession(toolUseId: string, filePath: string): Promise<void> {
+    public async startStreamingDiffSession(
+        toolUseId: string,
+        filePath: string,
+        providedOriginalContent?: string
+    ): Promise<void> {
         try {
-            let originalContent = ''
-            const pendingWrite = Array.from(this.pendingWrites.values()).find(
-                (write) => write.toolUseId === toolUseId && write.filePath === filePath
-            )
+            let originalContent = providedOriginalContent || ''
 
-            if (pendingWrite) {
-                originalContent = pendingWrite.originalContent
-            } else {
-                try {
-                    const document = await vscode.workspace.openTextDocument(vscode.Uri.file(filePath))
-                    originalContent = document.getText()
-                } catch {
-                    originalContent = ''
+            if (!providedOriginalContent) {
+                const pendingWrite = Array.from(this.pendingWrites.values()).find(
+                    (write) => write.toolUseId === toolUseId && write.filePath === filePath
+                )
+
+                if (pendingWrite) {
+                    originalContent = pendingWrite.originalContent
+                } else {
+                    try {
+                        const document = await vscode.workspace.openTextDocument(vscode.Uri.file(filePath))
+                        originalContent = document.getText()
+                    } catch {
+                        originalContent = ''
+                    }
                 }
             }
 
@@ -288,17 +295,7 @@ export class DiffAnimationHandler implements vscode.Disposable {
         filePath: string,
         originalContent: string
     ): Promise<void> {
-        try {
-            this.streamingSessions.set(toolUseId, {
-                toolUseId,
-                filePath,
-                originalContent,
-                startTime: Date.now(),
-            })
-            await this.streamingDiffController.openStreamingDiffView(toolUseId, filePath, originalContent)
-        } catch (error) {
-            getLogger().error(`[DiffAnimationHandler] ❌ Failed to start streaming session for ${toolUseId}: ${error}`)
-        }
+        return this.startStreamingDiffSession(toolUseId, filePath, originalContent)
     }
 
     public async streamContentUpdate(

packages/amazonq/src/lsp/chat/diffAnimation/streamingDiffController.ts

@@ -376,6 +376,51 @@ export class StreamingDiffController implements vscode.Disposable {
         }
     }
 
+    /**
+     * Helper method to create temp file and open diff view - eliminates code duplication
+     */
+    private async createTempFileAndOpenDiff(
+        filePath: string,
+        originalContent: string,
+        toolUseId: string,
+        diffTitle: string
+    ): Promise<{ tempFilePath: string; activeDiffEditor: vscode.TextEditor }> {
+        const fileName = path.basename(filePath)
+        const tempFilePath = path.join(path.dirname(filePath), `.amazonq-temp-${toolUseId}-${fileName}`)
+        const tempFileUri = vscode.Uri.file(tempFilePath)
+
+        // Create virtual URI for original content
+        const originalUri = vscode.Uri.parse(`${diffViewUriScheme}:${fileName}`).with({
+            query: Buffer.from(originalContent).toString('base64'),
+        })
+
+        // Create temporary file with original content
+        await this.createTempFile(tempFilePath, originalContent)
+
+        // Open diff view between virtual original and temp file
+        const activeDiffEditor = await new Promise<vscode.TextEditor>((resolve, reject) => {
+            const disposable = vscode.window.onDidChangeActiveTextEditor((editor) => {
+                if (editor && editor.document.uri.fsPath === tempFilePath) {
+                    disposable.dispose()
+                    resolve(editor)
+                }
+            })
+
+            void vscode.commands.executeCommand('vscode.diff', originalUri, tempFileUri, diffTitle, {
+                preserveFocus: true,
+                preview: false,
+            })
+
+            // Timeout after 10 seconds
+            setTimeout(() => {
+                disposable.dispose()
+                reject(new Error('Failed to open diff editor within timeout'))
+            }, 10000)
+        })
+
+        return { tempFilePath, activeDiffEditor }
+    }
+
     /**
      * Configure VSCode's diff editor settings for better focus on changed regions
      */
@@ -528,54 +573,21 @@ export class StreamingDiffController implements vscode.Disposable {
             const originalContent = originalDocument.getText()
             const fileName = path.basename(fsReplaceComplete.filePath)
 
-            // **CRITICAL FIX: Create temporary file for animation (like fsWrite)**
-            const tempFilePath = path.join(
-                path.dirname(fsReplaceComplete.filePath),
-                `.amazonq-temp-${fsReplaceComplete.toolUseId}-${fileName}`
+            // **STEP 1 & 2: Use helper method to create temp file and open diff view**
+            const { tempFilePath, activeDiffEditor } = await this.createTempFileAndOpenDiff(
+                fsReplaceComplete.filePath,
+                originalContent,
+                fsReplaceComplete.toolUseId,
+                `${fileName}: Original ↔ Amazon Q fsReplace Changes`
             )
             const tempFileUri = vscode.Uri.file(tempFilePath)
 
-            // Create virtual URI for original content
-            const originalUri = vscode.Uri.parse(`${diffViewUriScheme}:${fileName}`).with({
-                query: Buffer.from(originalContent).toString('base64'),
-            })
-
-            // **STEP 1: Create temp file with original content**
-            await this.createTempFile(tempFilePath, originalContent)
-
-            // **STEP 2: Apply all diffs to create final content**
+            // **STEP 3: Apply all diffs to create final content**
             let finalContent = originalContent
             for (const { oldStr, newStr } of structuredDiffs) {
                 finalContent = finalContent.replace(oldStr, newStr)
             }
 
-            // **STEP 3: Open diff view between original (virtual) and temp file**
-            const activeDiffEditor = await new Promise<vscode.TextEditor>((resolve, reject) => {
-                const disposable = vscode.window.onDidChangeActiveTextEditor((editor) => {
-                    if (editor && editor.document.uri.fsPath === tempFilePath) {
-                        disposable.dispose()
-                        resolve(editor)
-                    }
-                })
-
-                void vscode.commands.executeCommand(
-                    'vscode.diff',
-                    originalUri,
-                    tempFileUri,
-                    `${fileName}: Original ↔ Amazon Q fsReplace Changes`,
-                    {
-                        preserveFocus: true,
-                        preview: false,
-                    }
-                )
-
-                // Timeout after 10 seconds
-                setTimeout(() => {
-                    disposable.dispose()
-                    reject(new Error('Failed to open diff editor within timeout'))
-                }, 10000)
-            })
-
             // **STEP 4: Configure diff editor settings**
             await this.configureDiffEditorSettings(activeDiffEditor)
 

packages/amazonq/src/lsp/chat/messages.ts

@@ -79,6 +79,7 @@ import { decryptResponse, encryptRequest } from '../encryption'
 import { focusAmazonQPanel } from './commands'
 import { DiffAnimationHandler } from './diffAnimation/diffAnimationHandler'
 import { getLogger } from 'aws-core-vscode/shared'
+import { getCursorState } from '../utils'
 
 // Create a singleton instance of DiffAnimationHandler
 let diffAnimationHandler: DiffAnimationHandler | undefined
@@ -130,21 +131,6 @@ export function registerLanguageServerEventListener(languageClient: LanguageClie
     })
 }
 
-function getCursorState(selection: readonly vscode.Selection[]) {
-    return selection.map((s) => ({
-        range: {
-            start: {
-                line: s.start.line,
-                character: s.start.character,
-            },
-            end: {
-                line: s.end.line,
-                character: s.end.character,
-            },
-        },
-    }))
-}
-
 // Initialize DiffAnimationHandler on first use
 function getDiffAnimationHandler(): DiffAnimationHandler {
     if (!diffAnimationHandler) {
@@ -153,6 +139,20 @@ function getDiffAnimationHandler(): DiffAnimationHandler {
     return diffAnimationHandler
 }
 
+// Helper function to clean up temp files - eliminates code duplication
+async function cleanupTempFiles(context: string): Promise<void> {
+    try {
+        const animationHandler = getDiffAnimationHandler()
+        const streamingController = (animationHandler as any).streamingDiffController
+        if (streamingController && streamingController.cleanupChatSession) {
+            await streamingController.cleanupChatSession()
+            getLogger().info(`[VSCode Client] 🧹 Cleaned up temp files ${context}`)
+        }
+    } catch (error) {
+        getLogger().warn(`[VSCode Client] ⚠️ Failed to cleanup temp files ${context}: ${error}`)
+    }
+}
+
 export function registerMessageListeners(
     languageClient: LanguageClient,
     provider: AmazonQChatViewProvider,
@@ -285,16 +285,7 @@ export function registerMessageListeners(
 
                 // **CRITICAL FIX**: Clean up temp files when chat is stopped
                 // This ensures temp files are cleaned up when users stop ongoing operations
-                try {
-                    const animationHandler = getDiffAnimationHandler()
-                    const streamingController = (animationHandler as any).streamingDiffController
-                    if (streamingController && streamingController.cleanupChatSession) {
-                        await streamingController.cleanupChatSession()
-                        getLogger().info(`[VSCode Client] 🧹 Cleaned up temp files after stopping chat`)
-                    }
-                } catch (error) {
-                    getLogger().warn(`[VSCode Client] ⚠️ Failed to cleanup temp files after stopping chat: ${error}`)
-                }
+                await cleanupTempFiles('after stopping chat')
                 break
             }
             case chatRequestType.method: {
@@ -306,16 +297,7 @@ export function registerMessageListeners(
 
                 // **CRITICAL FIX**: Clean up temp files from previous chat sessions
                 // This ensures temp files don't accumulate when users start new conversations
-                try {
-                    const animationHandler = getDiffAnimationHandler()
-                    const streamingController = (animationHandler as any).streamingDiffController
-                    if (streamingController && streamingController.cleanupChatSession) {
-                        await streamingController.cleanupChatSession()
-                        getLogger().info(`[VSCode Client] 🧹 Cleaned up temp files before starting new chat`)
-                    }
-                } catch (error) {
-                    getLogger().warn(`[VSCode Client] ⚠️ Failed to cleanup temp files before new chat: ${error}`)
-                }
+                await cleanupTempFiles('before starting new chat')
 
                 const chatDisposable = languageClient.onProgress(
                     chatRequestType,