feat(nep): Data Instrumentation

packages/core/src/codewhisperer/activation.ts

@@ -95,6 +95,7 @@ import { SecurityIssueTreeViewProvider } from './service/securityIssueTreeViewPr
 import { setContext } from '../shared/vscode/setContext'
 import { syncSecurityIssueWebview } from './views/securityIssue/securityIssueWebview'
 import { detectCommentAboveLine } from '../shared/utilities/commentUtils'
+import { activateEditTracking } from './nextEditPrediction/activation'
 import { notifySelectDeveloperProfile } from './region/utils'
 
 let localize: nls.LocalizeFunc
@@ -529,6 +530,8 @@ export async function activate(context: ExtContext): Promise<void> {
             })
         )
     }
+
+    activateEditTracking(context)
 }
 
 export async function shutdown() {

packages/core/src/codewhisperer/models/constants.ts

@@ -945,3 +945,10 @@ export const testGenExcludePatterns = [
     '**/*.deb',
     '**/*.model',
 ]
+
+export const predictionTrackerDefaultConfig = {
+    maxStorageSizeKb: 5000,
+    debounceIntervalMs: 2000,
+    maxAgeMs: 30000,
+    maxSupplementalContext: 15,
+}

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

@@ -0,0 +1,32 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { PredictionTracker } from './predictionTracker'
+import { PredictionKeyStrokeHandler } from './predictionKeyStrokeHandler'
+import { getLogger } from '../../shared/logger/logger'
+import { ExtContext } from '../../shared/extensions'
+
+export let predictionTracker: PredictionTracker | undefined
+let keyStrokeHandler: PredictionKeyStrokeHandler | undefined
+
+export function activateEditTracking(context: ExtContext): void {
+    try {
+        predictionTracker = new PredictionTracker(context.extensionContext)
+
+        keyStrokeHandler = new PredictionKeyStrokeHandler(predictionTracker)
+        context.extensionContext.subscriptions.push(
+            vscode.Disposable.from({
+                dispose: () => {
+                    keyStrokeHandler?.dispose()
+                },
+            })
+        )
+
+        getLogger('nextEditPrediction').debug('Next Edit Prediction activated')
+    } catch (error) {
+        getLogger('nextEditPrediction').error(`Error in activateEditTracking: ${error}`)
+    }
+}

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

@@ -0,0 +1,154 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as diff from 'diff'
+import { getLogger } from '../../shared/logger/logger'
+import * as codewhispererClient from '../client/codewhisperer'
+import { supplementalContextMaxTotalLength, charactersLimit } from '../models/constants'
+
+const logger = getLogger('nextEditPrediction')
+
+/**
+ * Generates a unified diff format between old and new file contents
+ */
+function generateUnifiedDiffWithTimestamps(
+    oldFilePath: string,
+    newFilePath: string,
+    oldContent: string,
+    newContent: string,
+    oldTimestamp: number,
+    newTimestamp: number,
+    contextSize: number = 3
+): string {
+    const patchResult = diff.createTwoFilesPatch(
+        oldFilePath,
+        newFilePath,
+        oldContent,
+        newContent,
+        String(oldTimestamp),
+        String(newTimestamp),
+        { 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
+}
+
+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, T_0 being the snapshot of current file content:
+ *  U0: udiff of T_0 and T_1
+ *  U1: udiff of T_0 and T_2
+ *  U2: udiff of T_0 and T_3
+ */
+export function generateDiffContexts(
+    filePath: string,
+    currentContent: string,
+    snapshotContents: SnapshotContent[],
+    maxContexts: number
+): codewhispererClient.SupplementalContext[] {
+    if (snapshotContents.length === 0) {
+        return []
+    }
+
+    const supplementalContexts: codewhispererClient.SupplementalContext[] = []
+    const currentTimestamp = Date.now()
+
+    for (let i = snapshotContents.length - 1; i >= 0; i--) {
+        const snapshot = snapshotContents[i]
+        try {
+            const unifiedDiff = generateUnifiedDiffWithTimestamps(
+                snapshot.filePath,
+                filePath,
+                snapshot.content,
+                currentContent,
+                snapshot.timestamp,
+                currentTimestamp
+            )
+
+            supplementalContexts.push({
+                filePath: snapshot.filePath,
+                content: unifiedDiff,
+                type: 'PreviousEditorState',
+                metadata: {
+                    previousEditorStateMetadata: {
+                        timeOffset: currentTimestamp - snapshot.timestamp,
+                    },
+                },
+            })
+        } catch (err) {
+            logger.error(`Failed to generate diff: ${err}`)
+        }
+    }
+
+    const trimmedContext = trimSupplementalContexts(supplementalContexts, maxContexts)
+    logger.debug(
+        `supplemental contexts: ${trimmedContext.length} contexts, total size: ${trimmedContext.reduce((sum, ctx) => sum + ctx.content.length, 0)} characters`
+    )
+    return trimmedContext
+}
+
+/**
+ * Trims the supplementalContexts array to ensure it doesn't exceed the max number
+ * of contexts or total character length limit
+ *
+ * @param supplementalContexts - Array of SupplementalContext objects (already sorted with newest first)
+ * @param maxContexts - Maximum number of supplemental contexts allowed
+ * @returns Trimmed array of SupplementalContext objects
+ */
+export function trimSupplementalContexts(
+    supplementalContexts: codewhispererClient.SupplementalContext[],
+    maxContexts: number
+): codewhispererClient.SupplementalContext[] {
+    if (supplementalContexts.length === 0) {
+        return supplementalContexts
+    }
+
+    // First filter out any individual context that exceeds the character limit
+    let result = supplementalContexts.filter((context) => {
+        return context.content.length <= charactersLimit
+    })
+
+    // Then limit by max number of contexts
+    if (result.length > maxContexts) {
+        result = result.slice(0, maxContexts)
+    }
+
+    // Lastly enforce total character limit
+    let totalLength = 0
+    let i = 0
+
+    while (i < result.length) {
+        totalLength += result[i].content.length
+        if (totalLength > supplementalContextMaxTotalLength) {
+            break
+        }
+        i++
+    }
+
+    if (i === result.length) {
+        return result
+    }
+
+    const trimmedContexts = result.slice(0, i)
+    return trimmedContexts
+}

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

@@ -0,0 +1,117 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { PredictionTracker } from './predictionTracker'
+
+/**
+ * Monitors document changes in the editor and track them for prediction.
+ */
+export class PredictionKeyStrokeHandler {
+    private disposables: vscode.Disposable[] = []
+    private tracker: PredictionTracker
+    private shadowCopies: Map<string, string> = new Map()
+
+    /**
+     * Creates a new PredictionKeyStrokeHandler
+     * @param context The extension context
+     * @param tracker The prediction tracker instance
+     * @param config Configuration options
+     */
+    constructor(tracker: PredictionTracker) {
+        this.tracker = tracker
+
+        // Initialize shadow copies for currently visible editors when extension starts
+        this.initializeVisibleDocuments()
+
+        // Register event handlers
+        this.registerVisibleDocumentListener()
+        this.registerTextDocumentChangeListener()
+    }
+
+    /**
+     * Initializes shadow copies for all currently visible text editors
+     */
+    private initializeVisibleDocuments(): void {
+        const editors = vscode.window.visibleTextEditors
+
+        for (const editor of editors) {
+            if (editor.document.uri.scheme === 'file') {
+                this.updateShadowCopy(editor.document)
+            }
+        }
+    }
+
+    /**
+     * Registers listeners for visibility events to maintain shadow copies of document content
+     * Only store and update shadow copies for currently visible editors
+     * And remove shadow copies for files that are no longer visible
+     * And edits are processed only if a shadow copy exists
+     * This avoids the memory problem if hidden files are bulk edited, i.e. with global find/replace
+     */
+    private registerVisibleDocumentListener(): void {
+        // Track when documents become visible (switched to)
+        const visibleDisposable = vscode.window.onDidChangeVisibleTextEditors((editors) => {
+            const currentVisibleFiles = new Set<string>()
+
+            for (const editor of editors) {
+                if (editor.document.uri.scheme === 'file') {
+                    const filePath = editor.document.uri.fsPath
+                    currentVisibleFiles.add(filePath)
+                    this.updateShadowCopy(editor.document)
+                }
+            }
+
+            for (const filePath of this.shadowCopies.keys()) {
+                if (!currentVisibleFiles.has(filePath)) {
+                    this.shadowCopies.delete(filePath)
+                }
+            }
+        })
+
+        this.disposables.push(visibleDisposable)
+    }
+
+    private updateShadowCopy(document: vscode.TextDocument): void {
+        if (document.uri.scheme === 'file') {
+            this.shadowCopies.set(document.uri.fsPath, document.getText())
+        }
+    }
+
+    /**
+     * Registers listener for text document changes to send to tracker
+     */
+    private registerTextDocumentChangeListener(): void {
+        // Listen for document changes
+        const changeDisposable = vscode.workspace.onDidChangeTextDocument(async (event) => {
+            const filePath = event.document.uri.fsPath
+            const prevContent = this.shadowCopies.get(filePath)
+
+            // Skip if there are no content changes or if the file is not visible
+            if (
+                event.contentChanges.length === 0 ||
+                event.document.uri.scheme !== 'file' ||
+                prevContent === undefined
+            ) {
+                return
+            }
+
+            await this.tracker.processEdit(event.document, prevContent)
+            this.updateShadowCopy(event.document)
+        })
+
+        this.disposables.push(changeDisposable)
+    }
+
+    /**
+     * Disposes of all resources used by this handler
+     */
+    public dispose(): void {
+        for (const disposable of this.disposables) {
+            disposable.dispose()
+        }
+        this.disposables = []
+    }
+}