feat(nep): Data Instrumentation

packages/core/src/codewhisperer/activation.ts

@@ -95,10 +95,12 @@ import { SecurityIssueTreeViewProvider } from './service/securityIssueTreeViewPr
 import { setContext } from '../shared/vscode/setContext'
 import { syncSecurityIssueWebview } from './views/securityIssue/securityIssueWebview'
 import { detectCommentAboveLine } from '../shared/utilities/commentUtils'
+import { activateNextEditPrediction } from './nextEditPrediction/activation'
 
 let localize: nls.LocalizeFunc
 
 export async function activate(context: ExtContext): Promise<void> {
+    activateNextEditPrediction(context)
     localize = nls.loadMessageBundle()
 
     // Import old CodeWhisperer settings into Amazon Q

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

@@ -931,3 +931,11 @@ export const testGenExcludePatterns = [
     '**/*.deb',
     '**/*.model',
 ]
+
+export const predictionTrackerDefaultConfig = {
+    maxFiles: 25,
+    maxStorageSizeKb: 10000,
+    debounceIntervalMs: 2000,
+    maxAgeMs: 30000,
+    maxSupplementalContext: 15,
+}

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

@@ -0,0 +1,118 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { getLogger } from '../../shared/logger/logger'
+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
+     */
+    private registerVisibleDocumentListener(): void {
+        // Track when documents become visible (switched to)
+        const visibleDisposable = vscode.window.onDidChangeVisibleTextEditors((editors) => {
+            const currentVisibleFiles = new Set<string>()
+
+            // Update shadow copies for currently visible editors
+            for (const editor of editors) {
+                if (editor.document.uri.scheme === 'file') {
+                    const filePath = editor.document.uri.fsPath
+                    currentVisibleFiles.add(filePath)
+                    this.updateShadowCopy(editor.document)
+                }
+            }
+
+            // Remove shadow copies for files that are no longer visible
+            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())
+            getLogger().debug(`Updated shadow copy for ${document.uri.fsPath}`)
+        }
+    }
+
+    /**
+     * Registers listener for text document changes to send to tracker
+     */
+    private registerTextDocumentChangeListener(): void {
+        // Listen for document changes
+        const changeDisposable = vscode.workspace.onDidChangeTextDocument((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
+            // This avoids tracking bulk edits on non-visible files
+            if (
+                event.contentChanges.length === 0 ||
+                event.document.uri.scheme !== 'file' ||
+                prevContent === undefined
+            ) {
+                return
+            }
+
+            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 = []
+    }
+}