refactor supplemental context generation into diffGenerator

Author: tomcat323

packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts

@@ -7,16 +7,14 @@ import * as vscode from 'vscode'
 import * as path from 'path'
 import fs from '../../shared/fs/fs'
 import { getLogger } from '../../shared/logger/logger'
-import { DiffGenerator } from './diffGenerator'
+import * as diffGenerator from './diffGenerator'
 import * as codewhispererClient from '../client/codewhisperer'
 
 export interface FileTrackerConfig {
     /** Maximum number of files to track (default: 15) */
     maxFiles: number
     /** Maximum total size in kilobytes (default: 200) */
     maxTotalSizeKb: number
-    /** Maximum size per file in kilobytes */
-    maxFileSizeKb: number
     /** Debounce interval in milliseconds (default: 2000) */
     debounceIntervalMs: number
     /** Maximum age of snapshots in milliseconds (default: 30000) */
@@ -49,7 +47,6 @@ export class PredictionTracker {
         const defaultConfig = {
             maxFiles: 25,
             maxTotalSizeKb: 50000,
-            maxFileSizeKb: 100, // Default max size per file
             debounceIntervalMs: 2000,
             maxAgeMs: 30000, // 30 sec
             maxSupplementalContext: 15, // Default max supplemental contexts
@@ -90,12 +87,6 @@ export class PredictionTracker {
         const content = previousContent
         const size = Buffer.byteLength(content, 'utf8')
 
-        // Skip if the file is too large
-        if (size > this.config.maxFileSizeKb * 1024) {
-            getLogger().info(`File ${filePath} exceeds maximum size limit`)
-            return
-        }
-
         const timestamp = Date.now()
         const storageKey = this.generateStorageKey(filePath, timestamp)
 
@@ -291,7 +282,7 @@ export class PredictionTracker {
             throw new Error('Storage path not available')
         }
 
-        const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+        const snapshotsDir = path.join(this.storagePath, 'AmazonQ-file-snapshots')
         if (!(await fs.existsDir(snapshotsDir))) {
             await fs.mkdir(snapshotsDir)
         }
@@ -309,7 +300,7 @@ export class PredictionTracker {
             return
         }
 
-        const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+        const snapshotsDir = path.join(this.storagePath, 'AmazonQ-file-snapshots')
         const filePath = path.join(snapshotsDir, `${snapshot.storageKey}.content`)
 
         if (await fs.exists(filePath)) {
@@ -331,7 +322,7 @@ export class PredictionTracker {
             throw new Error('Storage path not available')
         }
 
-        const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+        const snapshotsDir = path.join(this.storagePath, 'AmazonQ-file-snapshots')
         const filePath = path.join(snapshotsDir, `${snapshot.storageKey}.content`)
 
         try {
@@ -346,8 +337,6 @@ export class PredictionTracker {
      * Generates unified diffs between adjacent snapshots of a file
      * and between the newest snapshot and the current file content
      *
-     * @param filePath Path to the file for which diffs should be generated
-     * @param currentContent Current content of the file to compare with the latest snapshot
      * @returns Array of SupplementalContext objects containing diffs between snapshots and current content
      */
     public async generatePredictionSupplementalContext(): Promise<codewhispererClient.SupplementalContext[]> {
@@ -357,97 +346,44 @@ export class PredictionTracker {
         }
         const filePath = activeEditor.document.uri.fsPath
         const currentContent = activeEditor.document.getText()
+
         // Get all snapshots for this file
         const snapshots = this.getFileSnapshots(filePath)
 
         if (snapshots.length === 0) {
             return []
         }
 
-        // Sort snapshots by timestamp (oldest first)
-        const sortedSnapshots = [...snapshots].sort((a, b) => a.timestamp - b.timestamp)
-        const supplementalContexts: codewhispererClient.SupplementalContext[] = []
-        const currentTimestamp = Date.now()
-
-        // Generate diffs between adjacent snapshots
-        for (let i = 0; i < sortedSnapshots.length - 1; i++) {
-            const oldSnapshot = sortedSnapshots[i]
-            const newSnapshot = sortedSnapshots[i + 1]
-
-            try {
-                const oldContent = await this.getSnapshotContent(oldSnapshot)
-                const newContent = await this.getSnapshotContent(newSnapshot)
-
-                const diff = await DiffGenerator.generateUnifiedDiffWithTimestamps(
-                    oldSnapshot.filePath,
-                    newSnapshot.filePath,
-                    oldContent,
-                    newContent,
-                    oldSnapshot.timestamp,
-                    newSnapshot.timestamp
-                )
-
-                supplementalContexts.push({
-                    filePath: oldSnapshot.filePath,
-                    content: diff,
-                    type: 'PreviousEditorState',
-                    metadata: {
-                        previousEditorStateMetadata: {
-                            timeOffset: currentTimestamp - oldSnapshot.timestamp,
-                        },
-                    },
-                })
-            } catch (err) {
-                getLogger().error(`Failed to generate diff: ${err}`)
-            }
-        }
-
-        // Generate diff between the newest snapshot and the current file content
-        if (sortedSnapshots.length > 0) {
-            const newestSnapshot = sortedSnapshots[sortedSnapshots.length - 1]
-
+        // Load all snapshot contents
+        const snapshotContents: diffGenerator.SnapshotContent[] = []
+        for (const snapshot of snapshots) {
             try {
-                // Need to temporarily save files to compare
-                const newestContent = await this.getSnapshotContent(newestSnapshot)
-
-                const diff = await DiffGenerator.generateUnifiedDiffWithTimestamps(
-                    newestSnapshot.filePath,
-                    newestSnapshot.filePath,
-                    newestContent,
-                    currentContent,
-                    newestSnapshot.timestamp,
-                    currentTimestamp
-                )
-
-                supplementalContexts.push({
-                    filePath: newestSnapshot.filePath,
-                    content: diff,
-                    type: 'PreviousEditorState',
-                    metadata: {
-                        previousEditorStateMetadata: {
-                            timeOffset: currentTimestamp - newestSnapshot.timestamp,
-                        },
-                    },
+                const content = await this.getSnapshotContent(snapshot)
+                snapshotContents.push({
+                    filePath: snapshot.filePath,
+                    content,
+                    timestamp: snapshot.timestamp,
                 })
             } catch (err) {
-                getLogger().error(`Failed to generate diff with current content: ${err}`)
+                getLogger().error(`Failed to load snapshot content: ${err}`)
             }
         }
 
-        // Limit the number of supplemental contexts based on config
-        if (supplementalContexts.length > this.config.maxSupplementalContext) {
-            return supplementalContexts.slice(-this.config.maxSupplementalContext)
-        }
-
-        return supplementalContexts
+        // Use the diffGenerator module to generate supplemental contexts
+        return diffGenerator.generateDiffContexts(
+            filePath,
+            currentContent,
+            snapshotContents,
+            this.config.maxSupplementalContext
+        )
     }
 
     private async ensureStorageDirectoryExists(): Promise<void> {
         if (!this.storagePath) {
             return
         }
 
-        const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+        const snapshotsDir = path.join(this.storagePath, 'AmazonQ-file-snapshots')
         if (!(await fs.existsDir(snapshotsDir))) {
             await fs.mkdir(snapshotsDir)
         }
@@ -458,7 +394,7 @@ export class PredictionTracker {
             return
         }
 
-        const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+        const snapshotsDir = path.join(this.storagePath, 'AmazonQ-file-snapshots')
         if (!(await fs.existsDir(snapshotsDir))) {
             return
         }

packages/core/src/codewhisperer/nextEditPrediction/SnapshotVisualizer.ts

@@ -54,7 +54,7 @@ export class SnapshotVisualizer {
                     case 'showSnapshot':
                         await this.showSnapshotContent(message.filePath, message.timestamp)
                         break
-                    case 'fireAPI':
+                    case 'generateUdiff':
                         await this.generateDiffs()
                         break
                 }
@@ -266,7 +266,7 @@ export class SnapshotVisualizer {
                 <div class="button-group">
                     <span id="total-count">0 snapshots</span>
                     <button class="btn" id="refresh-btn">Refresh</button>
-                    <button class="btn" id="fire-api-btn">FireAPI</button>
+                    <button class="btn" id="generate-udiff-btn">generateUdiff</button>
                 </div>
             </div>
             
@@ -404,9 +404,9 @@ export class SnapshotVisualizer {
                         vscode.postMessage({ command: 'refresh' });
                     });
                     
-                    // FireAPI button handler
-                    document.getElementById('fire-api-btn').addEventListener('click', () => {
-                        vscode.postMessage({ command: 'fireAPI' });
+                    // generateUdiff button handler
+                    document.getElementById('generate-udiff-btn').addEventListener('click', () => {
+                        vscode.postMessage({ command: 'generateUdiff' });
                     });
                     
                     // Request initial data

packages/core/src/codewhisperer/nextEditPrediction/diffGenerator.ts

@@ -4,47 +4,125 @@
  */
 
 import * as diff from 'diff'
+import { getLogger } from '../../shared/logger/logger'
+import * as codewhispererClient from '../client/codewhisperer'
 
 /**
- * Class to generate a unified diff format between old and new file contents
+ * Generates a unified diff format between old and new file contents
+ *
+ * @param oldFilePath - Path of the old file
+ * @param newFilePath - Path of the new file
+ * @param oldContent - Content of the old file
+ * @param newContent - Content of the new file
+ * @param oldTimestamp - Timestamp of the old file version
+ * @param newTimestamp - Timestamp of the new file version
+ * @param contextSize - Number of context lines to include (default: 3)
+ * @returns Unified diff as a string
  */
-export class DiffGenerator {
-    /**
-     * @param oldFilePath - Path of the old file
-     * @param newFilePath - Path of the new file
-     * @param oldContent - Content of the old file
-     * @param newContent - Content of the new file
-     * @param oldTimestamp - Timestamp of the old file version
-     * @param newTimestamp - Timestamp of the new file version
-     * @param contextSize - Number of context lines to include (default: 3)
-     * @returns Unified diff as a string
-     */
-    public static async generateUnifiedDiffWithTimestamps(
-        oldFilePath: string,
-        newFilePath: string,
-        oldContent: string,
-        newContent: string,
-        oldTimestamp: number,
-        newTimestamp: number,
-        contextSize: number = 3
-    ): Promise<string> {
-        const patchResult = diff.createTwoFilesPatch(
-            oldFilePath,
-            newFilePath,
-            oldContent,
-            newContent,
-            `${oldTimestamp}`, // Old file label with timestamp
-            `${newTimestamp}`, // New file label with timestamp
-            { context: contextSize }
-        )
-
-        // Remove the "Index:" line and the separator line that follows it
-        const lines = patchResult.split('\n')
-        if (lines.length >= 2 && lines[0].startsWith('Index:')) {
-            lines.splice(0, 2)
-            return lines.join('\n')
+export async function generateUnifiedDiffWithTimestamps(
+    oldFilePath: string,
+    newFilePath: string,
+    oldContent: string,
+    newContent: string,
+    oldTimestamp: number,
+    newTimestamp: number,
+    contextSize: number = 3
+): Promise<string> {
+    const patchResult = diff.createTwoFilesPatch(
+        oldFilePath,
+        newFilePath,
+        oldContent,
+        newContent,
+        `${oldTimestamp}`, // Old file label with timestamp
+        `${newTimestamp}`, // New file label with timestamp
+        { context: contextSize }
+    )
+
+    // Remove unused headers
+    const lines = patchResult.split('\n')
+    if (lines.length >= 2 && lines[0].startsWith('Index:')) {
+        lines.splice(0, 2)
+        return lines.join('\n')
+    }
+
+    return patchResult
+}
+
+/**
+ * Interface for snapshot content with timestamp
+ */
+export interface SnapshotContent {
+    filePath: string
+    content: string
+    timestamp: number
+}
+
+/**
+ * Generates supplemental contexts from snapshot contents and current content
+ *
+ * @param filePath - Path to the file
+ * @param currentContent - Current content of the file
+ * @param snapshotContents - List of snapshot contents sorted by timestamp (oldest first)
+ * @param maxContexts - Maximum number of supplemental contexts to return
+ * @returns Array of SupplementalContext objects
+ */
+export async function generateDiffContexts(
+    filePath: string,
+    currentContent: string,
+    snapshotContents: SnapshotContent[],
+    maxContexts: number
+): Promise<codewhispererClient.SupplementalContext[]> {
+    if (snapshotContents.length === 0) {
+        return []
+    }
+
+    const supplementalContexts: codewhispererClient.SupplementalContext[] = []
+    const currentTimestamp = Date.now()
+
+    // Treat current content as the last snapshot
+    const allContents = [
+        ...snapshotContents,
+        {
+            filePath,
+            content: currentContent,
+            timestamp: currentTimestamp,
+        },
+    ]
+
+    // Generate diffs between all adjacent snapshots (including current content)
+    for (let i = 0; i < allContents.length - 1; i++) {
+        const oldSnapshot = allContents[i]
+        const newSnapshot = allContents[i + 1]
+
+        try {
+            const diff = await generateUnifiedDiffWithTimestamps(
+                oldSnapshot.filePath,
+                newSnapshot.filePath,
+                oldSnapshot.content,
+                newSnapshot.content,
+                oldSnapshot.timestamp,
+                newSnapshot.timestamp
+            )
+
+            supplementalContexts.push({
+                filePath: oldSnapshot.filePath,
+                content: diff,
+                type: 'PreviousEditorState',
+                metadata: {
+                    previousEditorStateMetadata: {
+                        timeOffset: currentTimestamp - oldSnapshot.timestamp,
+                    },
+                },
+            })
+        } catch (err) {
+            getLogger().error(`Failed to generate diff: ${err}`)
         }
+    }
 
-        return patchResult
+    // Limit the number of supplemental contexts based on config
+    if (supplementalContexts.length > maxContexts) {
+        return supplementalContexts.slice(-maxContexts)
     }
+
+    return supplementalContexts
 }