feat(nep): Data Instrumentation (aws#7109)

## Problem
The science & service team in Next Edit Prediction (NEP) project needs
initial data for model training and tuning, to provide more contextually
relevant suggestions.

This requires us to track user edits in the IDE, and send changes (of
current active file in unified diff format) to the codeWhisper API.
**No impact on user experience.**
**This won't live forever, will be migrated to flare by end of May,
needed now for science data collection.**

## Solution

### Key Components
`PredictionKeyStrokeHandler`: Listens for document changes, maintains
shadow copies of visible documents, and processes edits.
`PredictionTracker`: Manages file snapshots, implementing a policy for
storing, retrieving, and pruning snapshots based on age and memory
constraints.
`DiffGenerator`: Creates unified diffs between file snapshots, produces
`supplementalContext` sent to the API.

### How it Works
- The system track shadow copies of editor visible files' content
- Once an edit is made to a tracked file, it takes a snapshot of the
file content before the edit
- When the Inline API fires, the snapshots of the current editing files
are used to generate diff context

### Memory management
- maxTotalSizeKb (default: 5000): Caps total size of snapshots storage
at ~5MB, purging oldest snapshots when exceeded.
- debounceIntervalMs (default: 2000): Prevents excessive snapshots by
requiring 2 seconds between captures for the same file.
- maxAgeMs (default: 30000): Auto-deletes snapshots after 30 seconds to
maintain recent-only history.
- maxSupplementalContext (default: 15): Limits `supplementalContext`
sent to API to 15 entries maximum.

### Changes
- Added new NextEditPrediction module in the CodeWhisperer package
- Updated activation code to initialize the NEP system
- Updated codeWhisper inline API requests to fit new format
---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: tomcat323

packages/core/src/codewhisperer/activation.ts

@@ -90,6 +90,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
@@ -505,6 +506,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 = []
+    }
+}