feat(nep): Data Instrumentation #7109

tomcat323 · 2025-04-18T20:23:24Z

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.
License: I confirm that my contribution is made under the terms of the Apache 2.0 license.

github-actions · 2025-04-18T20:23:38Z

This pull request modifies code in src/* but no tests were added/updated.
- Confirm whether tests should be added or ensure the PR description explains why tests are not required.
This pull request implements a feat or fix, so it must include a changelog entry (unless the fix is for an unreleased feature). Review the changelog guidelines.
- Note: beta or "experiment" features that have active users should announce fixes in the changelog.
- If this is not a feature or fix, use an appropriate type from the title guidelines. For example, telemetry-only changes should use the telemetry type.

justinmk3 · 2025-04-22T18:30:37Z

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

@@ -0,0 +1,118 @@
+/*!


should this these NEP modules live in the same dir as the other "trackers"?

https://github.com/aws/aws-toolkit-vscode/tree/master/packages/core/src/codewhisperer/tracker

are they really not sharing any concepts at all? this PR is all new code.

They are not shared IMO, the existing tracker collects statistical metrics for telemetry, the NEP tracker tracks content changes. Their hyper parameters are also very different.

Since the new tracker is populating the supplementalContext field, maybe we can also consider https://github.com/aws/aws-toolkit-vscode/tree/master/packages/core/src/codewhisperer/util/supplementalContext?

opieter-aws

Is clean up of snapshots not needed when the extension is closed?

packages/core/src/codewhisperer/activation.ts

packages/core/src/codewhisperer/service/recommendationHandler.ts

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

nkomonen-amazon · 2025-04-24T20:41:19Z

packages/core/src/codewhisperer/activation.ts

 let localize: nls.LocalizeFunc

 export async function activate(context: ExtContext): Promise<void> {
+    activateNextEditPrediction(context)


I would move this to the bottom of the function since the beginning is for more general activation (eg auth).

Also I think it is better to have this in activateAmazonQCommon(). We want to move as much Q specific code in to their so that we can reduce our dependency on core. You may need to export this in one of the index.ts, see amazonQDiffScheme and how it is exported

nit : Lets scope the namining to EditTracking

I'm not sure if this needs to have it's own index.ts and export it to QCommon, since it's not a stand-alone feature, but rather an addition on existing codeWhisper API flow. Will move it to the bottom of codeWhisper activation.

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

nkomonen-amazon · 2025-04-24T20:50:12Z

packages/core/src/codewhisperer/util/editorContext.ts

+            predictionSupplementalContext = await predictionTracker.generatePredictionSupplementalContext()
+        }
+    } catch (error) {
+        getLogger().error(`Error getting prediction supplemental context: ${error}`)


You have the final call, but is it dangerous to swallow the error here? If a user is having issues we will never have any telemetry for this.

It's not dangerous, we are collecting initial data, and from the bug-bash no issues has been found. Telemetry will come later in this project.
The main thing is if anything goes wrong it doesn't crash the inline completions flow, hence we swallow it here.

dk19y

Given this codebase will move to Flare eventually. Can we do a minimal refactor this to make the migration straight forward ?

APIs which may not have a direct mapping

VS Code API	LSP Equivalent	Notes
`vscode.window.onDidChangeVisibleTextEditors`	No direct equivalent in LSP	LSP focuses on opened/closed documents (`didOpen`, `didClose`) not "visibility". You may need to track `didOpen`/`didClose` and assume editors are visible.
`vscode.workspace.onDidChangeTextDocument`	`textDocument/didChange` (notification)	LSP server receives `textDocument/didChange` to reflect edits in documents.
`vscode.window.visibleTextEditors`	No direct equivalent in LSP	LSP does not track visible editors. It only tracks open documents via `didOpen`, `didClose`, `didChange`.
`vscode.window.activeTextEditor`	No direct equivalent in LSP	LSP does not track which document is active. It only manages documents' content and lifecycles. (The client — e.g., VS Code — knows which one is active.)

Also lets not use fs APIs and rely on message passing as-much as we can.

packages/core/src/codewhisperer/util/editorContext.ts

dk19y · 2025-04-28T13:51:05Z

packages/core/src/codewhisperer/activation.ts

 let localize: nls.LocalizeFunc

 export async function activate(context: ExtContext): Promise<void> {
+    activateNextEditPrediction(context)


nit : Lets scope the namining to EditTracking

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

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

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

tomcat323 · 2025-04-30T17:41:04Z

/runIntegrationTests

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

## 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.

Added Prediction folder, integrated with API, status bar visualier

808d95e

tomcat323 added 5 commits April 18, 2025 16:24

refactor supplemental context generation into diffGenerator

ced1ec6

clean up pakcage.json

475c5dd

update supplemental context logic according to science feedback

440c44e

not include debug webview in release

ea601c9

refactor config to constants.ts

2bfa873

tomcat323 changed the title ~~feat(NEP): Data Instrumentation~~ feat(next edit prediction): Data Instrumentation Apr 21, 2025

tomcat323 changed the title ~~feat(next edit prediction): Data Instrumentation~~ feat(nep): Data Instrumentation Apr 21, 2025

added tests for prediction tracker

5aa4c51

tomcat323 marked this pull request as ready for review April 22, 2025 16:31

tomcat323 requested review from a team as code owners April 22, 2025 16:31

tomcat323 added 3 commits April 22, 2025 13:11

clean up

9cdb6f5

remove tracker dispose in activation

a371128

remove unused imports

dabf670

justinmk3 reviewed Apr 22, 2025

View reviewed changes

opieter-aws reviewed Apr 22, 2025

View reviewed changes

tomcat323 added 4 commits April 22, 2025 16:14

refactored filePath to getter functions

9476765

refactor load from storage, use reduce in getTotalCount

d233c0c

remove predictionTypes from request

fa6b40b

enforce supplemental context char limit

b935964

nkomonen-amazon reviewed Apr 24, 2025

View reviewed changes

packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts Outdated Show resolved Hide resolved

nkomonen-amazon reviewed Apr 24, 2025

View reviewed changes

packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts Outdated Show resolved Hide resolved

nkomonen-amazon reviewed Apr 24, 2025

View reviewed changes

packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts Outdated Show resolved Hide resolved

nkomonen-amazon reviewed Apr 24, 2025

View reviewed changes

dk19y suggested changes Apr 28, 2025

View reviewed changes

tomcat323 added 2 commits April 28, 2025 16:31

address comments, deprecate fs usage, use in memory content tracking

25cc9f9

clean up logging, fix predictionTracker tests

18603fa

tomcat323 added 3 commits April 29, 2025 01:27

Merge branch 'master' into NEP/DataInstrumentationLaunch

6c3d6ab

refactor paths for windows test fail

c683cb8

fix uri path leading seperator

7d25d3f

dk19y suggested changes Apr 30, 2025

View reviewed changes

add more try-catch for error handling

0446ccb

remove unused maxfiles constant

1a9b305

dk19y approved these changes Apr 30, 2025

View reviewed changes

tomcat323 added 3 commits April 30, 2025 14:13

add tests for diff context generation

4250b65

Merge branch 'master' into NEP/DataInstrumentationLaunch

ed30502

resolve merge conflict in settings

6261398

jpinkney-aws reviewed May 3, 2025

View reviewed changes

packages/core/src/codewhisperer/nextEditPrediction/PredictionKeyStrokeHandler.ts Show resolved Hide resolved

packages/core/src/codewhisperer/nextEditPrediction/PredictionKeyStrokeHandler.ts Show resolved Hide resolved

tomcat323 added 5 commits May 5, 2025 11:08

renamed files, removed redundent checks

adfe12e

Rename PredictionTracker.ts to predictionTracker.ts

092d0b7

Rename PredictionKeyStrokeHandler.ts to predictionKeyStrokeHandler.ts

71384d6

Rename PredictionTracker.test.ts to predictionTracker.test.ts

399d67d

clean up

d94891d

jpinkney-aws approved these changes May 5, 2025

View reviewed changes

jpinkney-aws merged commit f1f5a36 into aws:master May 5, 2025
42 of 46 checks passed

feat(nep): Data Instrumentation #7109

feat(nep): Data Instrumentation #7109

Uh oh!

Conversation

tomcat323 commented Apr 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Solution

Key Components

How it Works

Memory management

Changes

Uh oh!

github-actions bot commented Apr 18, 2025

Uh oh!

justinmk3 Apr 22, 2025

Choose a reason for hiding this comment

Uh oh!

tomcat323 Apr 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

opieter-aws left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nkomonen-amazon Apr 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dk19y Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

tomcat323 Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nkomonen-amazon Apr 24, 2025

Choose a reason for hiding this comment

Uh oh!

tomcat323 Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

dk19y left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dk19y Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomcat323 commented Apr 30, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

tomcat323 commented Apr 18, 2025 •

edited

Loading

tomcat323 Apr 22, 2025 •

edited

Loading

nkomonen-amazon Apr 24, 2025 •

edited

Loading