diff --git a/packages/amazonq/.changes/next-release/Feature-5fb3c926-5b03-4d8b-b573-ba6692bf79ed.json b/packages/amazonq/.changes/next-release/Feature-5fb3c926-5b03-4d8b-b573-ba6692bf79ed.json new file mode 100644 index 00000000000..fba91903ad1 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Feature-5fb3c926-5b03-4d8b-b573-ba6692bf79ed.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "Real time diff animation will be rendered during code generation" +} diff --git a/packages/amazonq/src/lsp/chat/diffAnimation/diffAnimationHandler.ts b/packages/amazonq/src/lsp/chat/diffAnimation/diffAnimationHandler.ts new file mode 100644 index 00000000000..904f2419a9a --- /dev/null +++ b/packages/amazonq/src/lsp/chat/diffAnimation/diffAnimationHandler.ts @@ -0,0 +1,243 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as vscode from 'vscode' +import { getLogger } from 'aws-core-vscode/shared' +import { StreamingDiffController } from './streamingDiffController' + +export class DiffAnimationHandler implements vscode.Disposable { + private streamingDiffController: StreamingDiffController + private streamingSessions = new Map< + string, + { toolUseId: string; filePath: string; originalContent: string; startTime: number } + >() + + constructor() { + this.streamingDiffController = new StreamingDiffController() + } + + public async startStreamingDiffSession( + toolUseId: string, + filePath: string, + providedOriginalContent?: string + ): Promise { + try { + let originalContent = providedOriginalContent || '' + + if (!providedOriginalContent) { + try { + const document = await vscode.workspace.openTextDocument(vscode.Uri.file(filePath)) + originalContent = document.getText() + } catch { + originalContent = '' + } + } + + this.streamingSessions.set(toolUseId, { + toolUseId, + filePath, + originalContent, + startTime: Date.now(), + }) + + await this.streamingDiffController.openStreamingDiffView(toolUseId, filePath, originalContent) + } catch (error) { + getLogger().error(`Failed to start streaming session for ${toolUseId}: ${error}`) + } + } + + public async startStreamingWithOriginalContent( + toolUseId: string, + filePath: string, + originalContent: string + ): Promise { + return this.startStreamingDiffSession(toolUseId, filePath, originalContent) + } + + public async streamContentUpdate( + toolUseId: string, + partialContent: string, + isFinal: boolean = false + ): Promise { + const session = this.streamingSessions.get(toolUseId) + if (!session) { + return + } + + if (!isFinal && partialContent.trim() === '') { + return + } + + try { + await this.streamingDiffController.streamContentUpdate(toolUseId, partialContent, isFinal) + + if (isFinal) { + this.streamingSessions.delete(toolUseId) + } + } catch (error) { + getLogger().error(`Failed to stream content for ${toolUseId}: ${error}`) + this.streamingSessions.delete(toolUseId) + } + } + + public isStreamingActive(toolUseId: string): boolean { + return this.streamingSessions.has(toolUseId) && this.streamingDiffController.isStreamingActive(toolUseId) + } + + public getStreamingStats(toolUseId: string): any { + const session = this.streamingSessions.get(toolUseId) + const streamingStats = this.streamingDiffController.getStreamingStats(toolUseId) + return { + sessionExists: !!session, + sessionDuration: session ? Date.now() - session.startTime : 0, + filePath: session?.filePath, + ...streamingStats, + } + } + + /** + * Common logic for initializing streaming sessions + */ + private async initializeStreamingSession( + streamingChunk: any, + initializingStreamsByFile: Map>, + errorPrefix: string + ): Promise { + const filePath = streamingChunk.filePath + const isAlreadyInitializing = + filePath && + initializingStreamsByFile.has(filePath) && + initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId) + + if (!this.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) { + if (!initializingStreamsByFile.has(filePath)) { + initializingStreamsByFile.set(filePath, new Set()) + } + initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId) + + try { + await this.startStreamingDiffSession(streamingChunk.toolUseId, filePath) + } catch (error) { + getLogger().error(`${errorPrefix} ${streamingChunk.toolUseId}: ${error}`) + if (errorPrefix.includes('fsReplace')) { + // Don't rethrow for fsReplace + } else { + throw error + } + } finally { + if (filePath && initializingStreamsByFile.has(filePath)) { + const toolUseIds = initializingStreamsByFile.get(filePath)! + toolUseIds.delete(streamingChunk.toolUseId) + if (toolUseIds.size === 0) { + initializingStreamsByFile.delete(filePath) + } + } + } + } + } + + /** + * Common logic for processing streaming content and cleanup + */ + private async processStreamingContent( + streamingChunk: any, + initializingStreamsByFile: Map> + ): Promise { + if (streamingChunk.fsWriteParams) { + if (this.streamingDiffController && (this.streamingDiffController as any).updateFsWriteParams) { + ;(this.streamingDiffController as any).updateFsWriteParams( + streamingChunk.toolUseId, + streamingChunk.fsWriteParams + ) + } + } + + await this.streamContentUpdate( + streamingChunk.toolUseId, + streamingChunk.content || '', + streamingChunk.isComplete || false + ) + + if (!streamingChunk.isComplete || !streamingChunk.filePath) { + return + } + + const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath) + if (!toolUseIds) { + return + } + + toolUseIds.delete(streamingChunk.toolUseId) + + if (toolUseIds.size === 0) { + initializingStreamsByFile.delete(streamingChunk.filePath) + } + } + + /** + * Handle streaming chunk processing for diff animations + */ + public async handleStreamingChunk( + streamingChunk: any, + initializingStreamsByFile: Map>, + processedChunks: Map> + ): Promise { + // Handle fsReplace streaming chunks separately + if (streamingChunk.toolName === 'fsReplace') { + try { + const contentHash = streamingChunk.content + ? `${streamingChunk.content.substring(0, 50)}-${streamingChunk.content.length}` + : 'empty' + const chunkHash = `${streamingChunk.toolUseId}-${contentHash}-${streamingChunk.fsWriteParams?.pairIndex || 0}-${streamingChunk.isComplete}` + + if (!processedChunks.has(streamingChunk.toolUseId)) { + processedChunks.set(streamingChunk.toolUseId, new Set()) + } + + const toolChunks = processedChunks.get(streamingChunk.toolUseId)! + + if (streamingChunk.fsWriteParams?.command === 'fsReplace_diffPair') { + if (toolChunks.has(chunkHash)) { + return + } + } else { + const simpleHash = `${streamingChunk.toolUseId}-${streamingChunk.content?.length || 0}` + if (toolChunks.has(simpleHash) && streamingChunk.isComplete) { + return + } + toolChunks.add(simpleHash) + } + + toolChunks.add(chunkHash) + + await this.initializeStreamingSession( + streamingChunk, + initializingStreamsByFile, + 'Failed to initialize fsReplace streaming session for' + ) + await this.processStreamingContent(streamingChunk, initializingStreamsByFile) + } catch (error) { + getLogger().error(`Failed to process fsReplace streaming chunk: ${error}`) + initializingStreamsByFile.delete(streamingChunk.toolUseId) + } + return + } + + try { + await this.initializeStreamingSession( + streamingChunk, + initializingStreamsByFile, + 'Failed to initialize streaming session for' + ) + await this.processStreamingContent(streamingChunk, initializingStreamsByFile) + } catch (error) { + getLogger().error(`Failed to process streaming chunk: ${error}`) + } + } + + public async dispose(): Promise { + this.streamingSessions.clear() + this.streamingDiffController.dispose() + } +} diff --git a/packages/amazonq/src/lsp/chat/diffAnimation/streamingDiffController.ts b/packages/amazonq/src/lsp/chat/diffAnimation/streamingDiffController.ts new file mode 100644 index 00000000000..13eb48803e1 --- /dev/null +++ b/packages/amazonq/src/lsp/chat/diffAnimation/streamingDiffController.ts @@ -0,0 +1,675 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as vscode from 'vscode' +import * as path from 'path' +import { getLogger } from 'aws-core-vscode/shared' +export const diffViewUriScheme = 'amazonq-diff' +interface FsWriteParams { + command?: string + insertLine?: number + oldStr?: string + newStr?: string + fileText?: string + explanation?: string + pairIndex?: number + totalPairs?: number +} + +type StreamingSession = { + filePath: string + tempFilePath: string + originalContent: string + activeDiffEditor: vscode.TextEditor + fadedOverlayController: DecorationController + activeLineController: DecorationController + streamedLines: string[] + disposed: boolean + fsWriteParams?: FsWriteParams +} + +/** + * Streaming Diff Controller using temporary files for animations + */ +export class StreamingDiffController implements vscode.Disposable { + private activeStreamingSessions = new Map() + + private fsReplaceSessionsByFile = new Map< + string, + { + toolUseIds: Set + totalExpectedPairs: number + completedPairs: number + tempFilePath: string + lastActivity: number + } + >() + + private contentProvider: DiffContentProvider + + constructor() { + this.contentProvider = new DiffContentProvider() + vscode.workspace.registerTextDocumentContentProvider(diffViewUriScheme, this.contentProvider) + } + + private disposeSession(session: StreamingSession): void { + session.disposed = true + session.fadedOverlayController.clear() + session.activeLineController.clear() + } + + private logError(context: string, toolUseId: string, error: any): void { + getLogger().error(`[StreamingDiffController] ❌ ${context} for ${toolUseId}: ${error}`) + } + + updateFsWriteParams(toolUseId: string, fsWriteParams: FsWriteParams): void { + const session = this.activeStreamingSessions.get(toolUseId) + if (session) { + session.fsWriteParams = fsWriteParams + + if (fsWriteParams?.command === 'fsReplace_diffPair') { + const filePath = session.filePath + const { totalPairs = 1 } = fsWriteParams + + if (!this.fsReplaceSessionsByFile.has(filePath)) { + this.fsReplaceSessionsByFile.set(filePath, { + toolUseIds: new Set([toolUseId]), + totalExpectedPairs: totalPairs, + completedPairs: 0, + tempFilePath: session.tempFilePath, + lastActivity: Date.now(), + }) + } else { + const fsReplaceSession = this.fsReplaceSessionsByFile.get(filePath)! + fsReplaceSession.toolUseIds.add(toolUseId) + fsReplaceSession.lastActivity = Date.now() + } + } + } + } + async openStreamingDiffView(toolUseId: string, filePath: string, originalContent: string): Promise { + try { + const fileName = path.basename(filePath) + + let tempFilePath: string + let shouldCreateNewTempFile = true + + // Check if there's already an fsReplace session for this file + const existingFsReplaceSession = this.fsReplaceSessionsByFile.get(filePath) + if (existingFsReplaceSession) { + tempFilePath = existingFsReplaceSession.tempFilePath + shouldCreateNewTempFile = false + + // Add this toolUseId to the existing session + existingFsReplaceSession.toolUseIds.add(toolUseId) + existingFsReplaceSession.lastActivity = Date.now() + } else { + // Create new temp file path + tempFilePath = path.join(path.dirname(filePath), `.amazonq-temp-${toolUseId}-${fileName}`) + } + + const tempFileUri = vscode.Uri.file(tempFilePath) + + const originalUri = vscode.Uri.parse(`${diffViewUriScheme}:${fileName}`).with({ + query: Buffer.from(originalContent).toString('base64'), + }) + if (shouldCreateNewTempFile) { + await this.createTempFile(tempFilePath, originalContent) + } + const activeDiffEditor = await new Promise((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 Changes`, + { + preserveFocus: true, + preview: false, + } + ) + + setTimeout(() => { + disposable.dispose() + reject(new Error('Failed to open diff editor within timeout')) + }, 10000) + }) + + const fadedOverlayController = new DecorationController('fadedOverlay', activeDiffEditor) + const activeLineController = new DecorationController('activeLine', activeDiffEditor) + + // Apply faded overlay to all lines initially + fadedOverlayController.addLines(0, activeDiffEditor.document.lineCount) + + // Store the streaming session with temp file path + this.activeStreamingSessions.set(toolUseId, { + filePath, + tempFilePath, + originalContent, + activeDiffEditor, + fadedOverlayController, + activeLineController, + streamedLines: [], + disposed: false, + }) + } catch (error) { + getLogger().error(`Failed to open diff view for ${toolUseId}: ${error}`) + throw error + } + } + + /** + * Stream content updates to temporary file for animation - handles different fsWrite and fsReplace operation types + */ + async streamContentUpdate(toolUseId: string, partialContent: string, isFinal: boolean = false): Promise { + const session = this.activeStreamingSessions.get(toolUseId) + + if (!session || session.disposed) { + return + } + + try { + const command = session.fsWriteParams?.command + + if (command === 'fsReplace_diffPair') { + await this.handleFsReplaceDiffPair(session, partialContent, isFinal) + return + } else if (command === 'fsReplace_completion') { + await this.handleFsReplaceCompletionSignal(session) + return + } + + let contentToAnimate = partialContent + + if (session.fsWriteParams?.command === 'append') { + try { + const needsNewline = session.originalContent.length !== 0 && !session.originalContent.endsWith('\n') + contentToAnimate = session.originalContent + (needsNewline ? '\n' : '') + partialContent + } catch (error) { + contentToAnimate = partialContent + } + } else if (session.fsWriteParams?.command === 'create') { + contentToAnimate = partialContent + } + const accumulatedLines = contentToAnimate.split('\n') + if (!isFinal) { + accumulatedLines.pop() + } + + const diffEditor = session.activeDiffEditor + const document = diffEditor.document + + if (!diffEditor || !document) { + getLogger().warn( + `[StreamingDiffController] Diff editor or document unavailable for ${toolUseId}, skipping animation update` + ) + return + } + + const beginningOfDocument = new vscode.Position(0, 0) + diffEditor.selection = new vscode.Selection(beginningOfDocument, beginningOfDocument) + const newLines = accumulatedLines.slice(session.streamedLines.length) + + for (let i = 0; i < newLines.length; i++) { + const lineIndex = session.streamedLines.length + i + const lineContent = newLines[i] + const edit = new vscode.WorkspaceEdit() + + if (lineIndex < document.lineCount) { + const lineRange = new vscode.Range(lineIndex, 0, lineIndex, document.lineAt(lineIndex).text.length) + edit.replace(document.uri, lineRange, lineContent) + } else { + const insertPosition = new vscode.Position(document.lineCount, 0) + const contentToInsert = (lineIndex > 0 ? '\n' : '') + lineContent + edit.insert(document.uri, insertPosition, contentToInsert) + } + + await vscode.workspace.applyEdit(edit) + + session.activeLineController.setActiveLine(lineIndex) + session.fadedOverlayController.updateOverlayAfterLine(lineIndex, document.lineCount) + + this.scrollEditorToLine(diffEditor, lineIndex) + } + session.streamedLines = accumulatedLines + + if (!isFinal) { + return + } + + // Final cleanup when streaming is complete + if (session.streamedLines.length < document.lineCount) { + const edit = new vscode.WorkspaceEdit() + edit.delete(document.uri, new vscode.Range(session.streamedLines.length, 0, document.lineCount, 0)) + await vscode.workspace.applyEdit(edit) + } + + try { + await document.save() + } catch (saveError) { + getLogger().error(`Failed to save temp file ${session.tempFilePath}: ${saveError}`) + } + + session.fadedOverlayController.clear() + session.activeLineController.clear() + + setTimeout(async () => { + try { + await this.cleanupTempFile(session.tempFilePath) + this.disposeSession(session) + this.activeStreamingSessions.delete(toolUseId) + } catch (error) { + getLogger().warn(`Failed to auto-cleanup temp file ${session.tempFilePath}: ${error}`) + } + }, 500) + } catch (error) { + getLogger().error( + `[StreamingDiffController] ❌ Failed to stream animation update for ${toolUseId}: ${error}` + ) + } + } + + async handleFsReplaceDiffPair(session: StreamingSession, partialContent: string, isFinal: boolean): Promise { + try { + const diffEditor = session.activeDiffEditor + const document = diffEditor.document + + if (!diffEditor || !document) { + getLogger().warn( + `[StreamingDiffController] Diff editor or document unavailable for fsReplace diffPair, skipping operation` + ) + return + } + await new Promise((resolve) => setTimeout(resolve, 10)) + + // Extract diff pair parameters from fsWriteParams (removed startLine - calculate dynamically) + const { oldStr, newStr, pairIndex, totalPairs } = session.fsWriteParams || {} + + if (!oldStr || !newStr) { + return + } + const currentContent = document.getText() + if (document.uri.fsPath !== session.tempFilePath) { + return + } + try { + const correctDocument = await vscode.workspace.openTextDocument(vscode.Uri.file(session.tempFilePath)) + if (correctDocument) { + const correctEditor = vscode.window.visibleTextEditors.find( + (editor) => editor.document.uri.fsPath === session.tempFilePath + ) + if (correctEditor) { + session.activeDiffEditor = correctEditor + } + } + } catch (error) { + getLogger().error(`[StreamingDiffController] ❌ Failed to correct document path: ${error}`) + return + } + + // Find the location of oldStr in the current content + const oldStrIndex = currentContent.indexOf(oldStr) + if (oldStrIndex === -1) { + return + } + + const beforeOldStr = currentContent.substring(0, oldStrIndex) + const startLineNumber = beforeOldStr.split('\n').length - 1 + const oldStrLines = oldStr.split('\n') + const endLineNumber = startLineNumber + oldStrLines.length - 1 + this.scrollEditorToLine(diffEditor, startLineNumber) + for (let lineNum = startLineNumber; lineNum <= endLineNumber; lineNum++) { + session.activeLineController.setActiveLine(lineNum) + await new Promise((resolve) => setTimeout(resolve, 50)) + } + const edit = new vscode.WorkspaceEdit() + const oldStrStartPos = document.positionAt(oldStrIndex) + const oldStrEndPos = document.positionAt(oldStrIndex + oldStr.length) + const replaceRange = new vscode.Range(oldStrStartPos, oldStrEndPos) + edit.replace(document.uri, replaceRange, newStr) + await vscode.workspace.applyEdit(edit) + const newStrLines = newStr.split('\n') + const newEndLineNumber = startLineNumber + newStrLines.length - 1 + for (let lineNum = startLineNumber; lineNum <= newEndLineNumber; lineNum++) { + session.activeLineController.setActiveLine(lineNum) + session.fadedOverlayController.updateOverlayAfterLine(lineNum, document.lineCount) + // Small delay to create smooth line-by-line animation effect for visual feedback + await new Promise((resolve) => setTimeout(resolve, 20)) + } + // Clear active line highlighting after a brief delay to allow user to see the final result + // before moving to the next diff pair or completing the animation + setTimeout(() => { + session.activeLineController.clear() + }, 500) + + try { + await document.save() + } catch (saveError) { + getLogger().error( + `[StreamingDiffController] ❌ Failed to save fsReplace diffPair temp file: ${saveError}` + ) + } + if (!isFinal) { + return + } + await this.handleFsReplaceCompletion(session, pairIndex || 0, totalPairs || 1) + } catch (error) { + getLogger().error(`[StreamingDiffController] ❌ Failed to handle fsReplace diffPair: ${error}`) + } + } + /** + * Create temporary file for animation + */ + private async createTempFile(tempFilePath: string, initialContent: string): Promise { + try { + const edit = new vscode.WorkspaceEdit() + edit.createFile(vscode.Uri.file(tempFilePath), { overwrite: true }) + await vscode.workspace.applyEdit(edit) + const document = await vscode.workspace.openTextDocument(vscode.Uri.file(tempFilePath)) + const fullEdit = new vscode.WorkspaceEdit() + fullEdit.replace(document.uri, new vscode.Range(0, 0, document.lineCount, 0), initialContent) + await vscode.workspace.applyEdit(fullEdit) + + await document.save() + } catch (error) { + getLogger().error(`[StreamingDiffController] ❌ Failed to create temp file ${tempFilePath}: ${error}`) + throw error + } + } + + /** + * Clean up temporary file after animation + */ + private async cleanupTempFile(tempFilePath: string): Promise { + try { + const edit = new vscode.WorkspaceEdit() + edit.deleteFile(vscode.Uri.file(tempFilePath), { ignoreIfNotExists: true }) + await vscode.workspace.applyEdit(edit) + } catch (error) { + getLogger().warn(`[StreamingDiffController] ⚠️ Failed to cleanup temp file ${tempFilePath}: ${error}`) + } + } + /** + * Scroll editor to line + */ + private scrollEditorToLine(editor: vscode.TextEditor, line: number): void { + editor.revealRange(new vscode.Range(line, 0, line, 0), vscode.TextEditorRevealType.InCenter) + } + + isStreamingActive(toolUseId: string): boolean { + const session = this.activeStreamingSessions.get(toolUseId) + return session !== undefined && !session.disposed + } + + /** + * Get streaming stats + */ + getStreamingStats(toolUseId: string): { isActive: boolean; contentLength: number } | undefined { + const session = this.activeStreamingSessions.get(toolUseId) + if (!session) { + return undefined + } + + return { + isActive: this.isStreamingActive(toolUseId), + contentLength: session.streamedLines.join('\n').length, + } + } + + /** + * Close streaming session + */ + async closeDiffView(toolUseId: string): Promise { + const session = this.activeStreamingSessions.get(toolUseId) + if (!session) { + return + } + + try { + this.disposeSession(session) + + // Clean up temp file immediately when session is closed + if (session.tempFilePath) { + await this.cleanupTempFile(session.tempFilePath) + } + + this.activeStreamingSessions.delete(toolUseId) + } catch (error) { + this.logError('Failed to close streaming session', toolUseId, error) + } + } + + /** + * Handle fsReplace completion signal from parser - triggers immediate cleanup + */ + private async handleFsReplaceCompletionSignal(session: StreamingSession): Promise { + const filePath = session.filePath + try { + // Clear decorations immediately + session.fadedOverlayController.clear() + session.activeLineController.clear() + const document = session.activeDiffEditor?.document + if (document) { + try { + await document.save() + } catch (saveError) { + getLogger().error(`[StreamingDiffController] ❌ Failed to save fsReplace temp file: ${saveError}`) + } + } + // Delay cleanup to allow final UI updates and user to see completion state + setTimeout(async () => { + try { + await this.cleanupTempFile(session.tempFilePath) + session.disposed = true + const sessionsToRemove: string[] = [] + for (const [toolUseId, sessionData] of this.activeStreamingSessions.entries()) { + if (sessionData.filePath === filePath) { + sessionsToRemove.push(toolUseId) + } + } + + for (const toolUseId of sessionsToRemove) { + this.activeStreamingSessions.delete(toolUseId) + } + + this.fsReplaceSessionsByFile.delete(filePath) + } catch (error) { + getLogger().warn( + `[StreamingDiffController] ⚠️ Failed to cleanup fsReplace session for ${filePath}: ${error}` + ) + } + }, 500) + } catch (error) { + getLogger().error(`[StreamingDiffController] ❌ Failed to handle fsReplace completion signal: ${error}`) + } + } + + /** + * Clean up multiple streaming sessions by their tool use IDs + */ + private async cleanupSessions(toolUseIds: Set): Promise { + for (const toolUseId of toolUseIds) { + const sessionToCleanup = this.activeStreamingSessions.get(toolUseId) + if (!sessionToCleanup) { + continue + } + sessionToCleanup.disposed = true + this.activeStreamingSessions.delete(toolUseId) + } + } + + /** + * Handle fsReplace completion - properly track and cleanup when all diff pairs for a file are done + */ + private async handleFsReplaceCompletion( + session: StreamingSession, + pairIndex: number, + totalPairs: number + ): Promise { + const filePath = session.filePath + const fsReplaceSession = this.fsReplaceSessionsByFile.get(filePath) + + if (!fsReplaceSession) { + return + } + + fsReplaceSession.completedPairs++ + fsReplaceSession.lastActivity = Date.now() + const allPairsComplete = fsReplaceSession.completedPairs >= fsReplaceSession.totalExpectedPairs + const isLastPairInSequence = pairIndex === totalPairs - 1 + if (allPairsComplete && isLastPairInSequence) { + session.fadedOverlayController.clear() + session.activeLineController.clear() + setTimeout(async () => { + try { + await this.cleanupTempFile(fsReplaceSession.tempFilePath) + await this.cleanupSessions(fsReplaceSession.toolUseIds) + this.fsReplaceSessionsByFile.delete(filePath) + } catch (error) { + getLogger().warn( + `[StreamingDiffController] ⚠️ Failed to cleanup fsReplace session for ${filePath}: ${error}` + ) + } + }, 1000) // 1 second delay to ensure all operations complete + } + } + + /** + * Clean up all temporary files for a chat session + */ + async cleanupChatSession(): Promise { + const tempFilePaths: string[] = [] + + // Collect from active streaming sessions + for (const session of this.activeStreamingSessions.values()) { + if (session.tempFilePath) { + tempFilePaths.push(session.tempFilePath) + } + } + + // Collect from fs replace sessions + for (const session of this.fsReplaceSessionsByFile.values()) { + if (session.tempFilePath) { + tempFilePaths.push(session.tempFilePath) + } + } + + for (const tempFilePath of tempFilePaths) { + try { + await this.cleanupTempFile(tempFilePath) + } catch (error) { + getLogger().warn(`[StreamingDiffController] ⚠️ Failed to cleanup temp file ${tempFilePath}: ${error}`) + } + } + this.fsReplaceSessionsByFile.clear() + } + + /** + * Dispose all resources + */ + dispose(): void { + void this.cleanupChatSession() + + for (const [toolUseId, session] of this.activeStreamingSessions.entries()) { + try { + this.disposeSession(session) + } catch (error) { + this.logError('Error disposing session', toolUseId, error) + } + } + + this.activeStreamingSessions.clear() + } +} + +class DiffContentProvider implements vscode.TextDocumentContentProvider { + provideTextDocumentContent(uri: vscode.Uri): string { + try { + return Buffer.from(uri.query, 'base64').toString('utf8') + } catch { + return '' + } + } +} + +const fadedOverlayDecorationType = vscode.window.createTextEditorDecorationType({ + backgroundColor: 'rgba(255, 255, 0, 0.1)', + opacity: '0.4', + isWholeLine: true, +}) + +const activeLineDecorationType = vscode.window.createTextEditorDecorationType({ + backgroundColor: 'rgba(255, 255, 0, 0.3)', + opacity: '1', + isWholeLine: true, + border: '1px solid rgba(255, 255, 0, 0.5)', +}) + +type DecorationType = 'fadedOverlay' | 'activeLine' + +class DecorationController { + private decorationType: DecorationType + private editor: vscode.TextEditor + private ranges: vscode.Range[] = [] + + constructor(decorationType: DecorationType, editor: vscode.TextEditor) { + this.decorationType = decorationType + this.editor = editor + } + + getDecoration() { + switch (this.decorationType) { + case 'fadedOverlay': + return fadedOverlayDecorationType + case 'activeLine': + return activeLineDecorationType + } + } + + addLines(startIndex: number, numLines: number) { + if (startIndex < 0 || numLines <= 0) { + return + } + + const lastRange = this.ranges[this.ranges.length - 1] + if (lastRange && lastRange.end.line === startIndex - 1) { + this.ranges[this.ranges.length - 1] = lastRange.with(undefined, lastRange.end.translate(numLines)) + } else { + const endLine = startIndex + numLines - 1 + this.ranges.push(new vscode.Range(startIndex, 0, endLine, Number.MAX_SAFE_INTEGER)) + } + + this.editor.setDecorations(this.getDecoration(), this.ranges) + } + + clear() { + this.ranges = [] + this.editor.setDecorations(this.getDecoration(), this.ranges) + } + + updateOverlayAfterLine(line: number, totalLines: number) { + this.ranges = this.ranges.filter((range) => range.end.line < line) + if (line < totalLines - 1) { + this.ranges.push( + new vscode.Range( + new vscode.Position(line + 1, 0), + new vscode.Position(totalLines - 1, Number.MAX_SAFE_INTEGER) + ) + ) + } + this.editor.setDecorations(this.getDecoration(), this.ranges) + } + + setActiveLine(line: number) { + this.ranges = [new vscode.Range(line, 0, line, Number.MAX_SAFE_INTEGER)] + this.editor.setDecorations(this.getDecoration(), this.ranges) + } +} diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index 16965e2f41f..50529acea75 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -104,6 +104,18 @@ import { getCursorState } from '../utils' import { focusAmazonQPanel } from './commands' import { ChatMessage } from '@aws/language-server-runtimes/server-interface' import { CommentUtils } from 'aws-core-vscode/utils' +import { DiffAnimationHandler } from './diffAnimation/diffAnimationHandler' +import { getLogger } from 'aws-core-vscode/shared' + +// Create a singleton instance of DiffAnimationHandler for streaming functionality +let diffAnimationHandler: DiffAnimationHandler | undefined + +function getDiffAnimationHandler(): DiffAnimationHandler { + if (!diffAnimationHandler) { + diffAnimationHandler = new DiffAnimationHandler() + } + return diffAnimationHandler +} export function registerActiveEditorChangeListener(languageClient: LanguageClient) { let debounceTimer: NodeJS.Timeout | undefined @@ -154,6 +166,17 @@ export function registerLanguageServerEventListener(languageClient: LanguageClie } }) } +async function cleanupTempFiles(): Promise { + try { + const animationHandler = getDiffAnimationHandler() + const streamingController = (animationHandler as any).streamingDiffController + if (streamingController && streamingController.cleanupChatSession) { + await streamingController.cleanupChatSession() + } + } catch (error) { + getLogger().warn(`Failed to cleanup temp files: ${error}`) + } +} export function registerMessageListeners( languageClient: LanguageClient, @@ -162,6 +185,12 @@ export function registerMessageListeners( ) { const chatStreamTokens = new Map() // tab id -> token + const initializingStreamsByFile = new Map>() // filePath -> Set of toolUseIds + + const processedChunks = new Map>() // toolUseId -> Set of content hashes + + // Initialize DiffAnimationHandler + // Keep track of pending chat options to send when webview UI is ready const pendingChatOptions = languageClient.initializeResult?.awsServerCapabilities?.chatOptions @@ -281,6 +310,9 @@ export function registerMessageListeners( token?.cancel() token?.dispose() chatStreamTokens.delete(tabId) + + initializingStreamsByFile.clear() + await cleanupTempFiles() break } case chatRequestType.method: { @@ -691,7 +723,17 @@ export function registerMessageListeners( await ecc.viewDiff(viewDiffMessage, amazonQDiffScheme) }) - languageClient.onNotification(chatUpdateNotificationType.method, (params: ChatUpdateParams) => { + languageClient.onNotification(chatUpdateNotificationType.method, async (params: ChatUpdateParams) => { + if ((params.data as any)?.streamingChunk) { + const animationHandler = getDiffAnimationHandler() + await animationHandler.handleStreamingChunk( + (params.data as any).streamingChunk, + initializingStreamsByFile, + processedChunks + ) + return + } + void provider.webview?.postMessage({ command: chatUpdateNotificationType.method, params: params, @@ -706,6 +748,14 @@ export function registerMessageListeners( }) } +// Clean up on extension deactivation +export function dispose() { + if (diffAnimationHandler) { + void diffAnimationHandler.dispose() + diffAnimationHandler = undefined + } +} + function isServerEvent(command: string) { return command.startsWith('aws/chat/') || command === 'telemetry/event' } diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index bc065c8f620..c01bf308783 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -181,6 +181,7 @@ export async function startLanguageServer( shortcut: true, reroute: true, modelSelection: true, + diffAnimation: true, workspaceFilePath: vscode.workspace.workspaceFile?.fsPath, codeReviewInChat: codeReviewInChat, // feature flag for displaying findings found not through CodeReview in the Code Issues Panel