From c8948c7a331157082696b60adc75fc4b47ec3933 Mon Sep 17 00:00:00 2001
From: Tom Zu <138054255+tomcat323@users.noreply.github.com>
Date: Mon, 14 Apr 2025 10:23:03 -0400
Subject: [PATCH 1/2] update codeWhisper API for new model inputs
---
.../codewhisperer/client/user-service-2.json | 404 +++++++++++++++++-
1 file changed, 399 insertions(+), 5 deletions(-)
diff --git a/packages/core/src/codewhisperer/client/user-service-2.json b/packages/core/src/codewhisperer/client/user-service-2.json
index 93b857f6ac0..969abf41f1a 100644
--- a/packages/core/src/codewhisperer/client/user-service-2.json
+++ b/packages/core/src/codewhisperer/client/user-service-2.json
@@ -66,6 +66,25 @@
"documentation": "
Creates a pre-signed, S3 write URL for uploading a repository zip archive.
",
"idempotent": true
},
+ "CreateUserMemoryEntry": {
+ "name": "CreateUserMemoryEntry",
+ "http": {
+ "method": "POST",
+ "requestUri": "/"
+ },
+ "input": { "shape": "CreateUserMemoryEntryInput" },
+ "output": { "shape": "CreateUserMemoryEntryOutput" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "InternalServerException" },
+ { "shape": "ValidationException" },
+ { "shape": "AccessDeniedException" }
+ ],
+ "documentation": "API to create a single user memory entry
",
+ "idempotent": true
+ },
"CreateWorkspace": {
"name": "CreateWorkspace",
"http": {
@@ -100,6 +119,24 @@
],
"documentation": "API to delete task assist conversation.
"
},
+ "DeleteUserMemoryEntry": {
+ "name": "DeleteUserMemoryEntry",
+ "http": {
+ "method": "POST",
+ "requestUri": "/"
+ },
+ "input": { "shape": "DeleteUserMemoryEntryInput" },
+ "output": { "shape": "DeleteUserMemoryEntryOutput" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "InternalServerException" },
+ { "shape": "ValidationException" },
+ { "shape": "AccessDeniedException" }
+ ],
+ "documentation": "API to delete a single user memory entry
",
+ "idempotent": true
+ },
"DeleteWorkspace": {
"name": "DeleteWorkspace",
"http": {
@@ -312,6 +349,23 @@
],
"documentation": "Return configruations for each feature that has been setup for A/B testing.
"
},
+ "ListUserMemoryEntries": {
+ "name": "ListUserMemoryEntries",
+ "http": {
+ "method": "POST",
+ "requestUri": "/"
+ },
+ "input": { "shape": "ListUserMemoryEntriesInput" },
+ "output": { "shape": "ListUserMemoryEntriesOutput" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "InternalServerException" },
+ { "shape": "ValidationException" },
+ { "shape": "AccessDeniedException" }
+ ],
+ "documentation": "API to list user memories
"
+ },
"ListWorkspaceMetadata": {
"name": "ListWorkspaceMetadata",
"http": {
@@ -647,6 +701,17 @@
"min": 0,
"sensitive": true
},
+ "AttributesMap": {
+ "type": "map",
+ "key": { "shape": "AttributesMapKeyString" },
+ "value": { "shape": "StringList" },
+ "documentation": "Attributes is a map of key-value pairs
"
+ },
+ "AttributesMapKeyString": {
+ "type": "string",
+ "max": 128,
+ "min": 1
+ },
"Base64EncodedPaginationToken": {
"type": "string",
"max": 2048,
@@ -665,6 +730,17 @@
"toggle": { "shape": "OptInFeatureToggle" }
}
},
+ "ChangeLogGranularityType": {
+ "type": "string",
+ "enum": ["STANDARD", "BUSINESS"]
+ },
+ "ChangeLogOptions": {
+ "type": "structure",
+ "required": ["granularity"],
+ "members": {
+ "granularity": { "shape": "ChangeLogGranularityType" }
+ }
+ },
"ChatAddMessageEvent": {
"type": "structure",
"required": ["conversationId", "messageId"],
@@ -689,7 +765,7 @@
"type": "list",
"member": { "shape": "ChatMessage" },
"documentation": "Indicates Participant in Chat conversation
",
- "max": 100,
+ "max": 250,
"min": 0
},
"ChatInteractWithMessageEvent": {
@@ -803,6 +879,23 @@
"type": "integer",
"min": 0
},
+ "CodeDescription": {
+ "type": "structure",
+ "required": ["href"],
+ "members": {
+ "href": {
+ "shape": "CodeDescriptionHrefString",
+ "documentation": "An URI to open with more information about the diagnostic error.
"
+ }
+ },
+ "documentation": "Structure to capture a description for an error code.
"
+ },
+ "CodeDescriptionHrefString": {
+ "type": "string",
+ "max": 1024,
+ "min": 1,
+ "sensitive": true
+ },
"CodeFixAcceptanceEvent": {
"type": "structure",
"required": ["jobId"],
@@ -1086,6 +1179,40 @@
"requestHeaders": { "shape": "RequestHeaders" }
}
},
+ "CreateUserMemoryEntryInput": {
+ "type": "structure",
+ "required": ["memoryEntryString", "origin"],
+ "members": {
+ "memoryEntryString": { "shape": "CreateUserMemoryEntryInputMemoryEntryStringString" },
+ "origin": { "shape": "Origin" },
+ "profileArn": {
+ "shape": "CreateUserMemoryEntryInputProfileArnString",
+ "documentation": "ProfileArn for the managing Q Profile
"
+ },
+ "clientToken": {
+ "shape": "IdempotencyToken",
+ "idempotencyToken": true
+ }
+ }
+ },
+ "CreateUserMemoryEntryInputMemoryEntryStringString": {
+ "type": "string",
+ "max": 500,
+ "min": 1,
+ "sensitive": true
+ },
+ "CreateUserMemoryEntryInputProfileArnString": {
+ "type": "string",
+ "min": 1,
+ "pattern": "arn:aws:codewhisperer:[-.a-z0-9]{1,63}:\\d{12}:profile/([a-zA-Z0-9]){12}"
+ },
+ "CreateUserMemoryEntryOutput": {
+ "type": "structure",
+ "required": ["memoryEntry"],
+ "members": {
+ "memoryEntry": { "shape": "MemoryEntry" }
+ }
+ },
"CreateWorkspaceRequest": {
"type": "structure",
"required": ["workspaceRoot"],
@@ -1171,6 +1298,32 @@
},
"documentation": "Structure to represent bootstrap conversation response.
"
},
+ "DeleteUserMemoryEntryInput": {
+ "type": "structure",
+ "required": ["id"],
+ "members": {
+ "id": { "shape": "DeleteUserMemoryEntryInputIdString" },
+ "profileArn": {
+ "shape": "DeleteUserMemoryEntryInputProfileArnString",
+ "documentation": "ProfileArn for the managing Q Profile
"
+ }
+ }
+ },
+ "DeleteUserMemoryEntryInputIdString": {
+ "type": "string",
+ "max": 36,
+ "min": 36,
+ "pattern": "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
+ },
+ "DeleteUserMemoryEntryInputProfileArnString": {
+ "type": "string",
+ "min": 1,
+ "pattern": "arn:aws:codewhisperer:[-.a-z0-9]{1,63}:\\d{12}:profile/([a-zA-Z0-9]){12}"
+ },
+ "DeleteUserMemoryEntryOutput": {
+ "type": "structure",
+ "members": {}
+ },
"DeleteWorkspaceRequest": {
"type": "structure",
"required": ["workspaceId"],
@@ -1204,11 +1357,66 @@
"documentation": "Represents a Diagnostic message
",
"union": true
},
+ "DiagnosticLocation": {
+ "type": "structure",
+ "required": ["uri", "range"],
+ "members": {
+ "uri": { "shape": "DiagnosticLocationUriString" },
+ "range": { "shape": "Range" }
+ },
+ "documentation": "Represents a location inside a resource, such as a line inside a text file.
"
+ },
+ "DiagnosticLocationUriString": {
+ "type": "string",
+ "max": 1024,
+ "min": 1,
+ "sensitive": true
+ },
+ "DiagnosticRelatedInformation": {
+ "type": "structure",
+ "required": ["location", "message"],
+ "members": {
+ "location": {
+ "shape": "DiagnosticLocation",
+ "documentation": "The location of this related diagnostic information.
"
+ },
+ "message": {
+ "shape": "DiagnosticRelatedInformationMessageString",
+ "documentation": "The message of this related diagnostic information.
"
+ }
+ },
+ "documentation": "Represents a related message and source code location for a diagnostic.
"
+ },
+ "DiagnosticRelatedInformationList": {
+ "type": "list",
+ "member": { "shape": "DiagnosticRelatedInformation" },
+ "documentation": "List of DiagnosticRelatedInformation
",
+ "max": 1024,
+ "min": 0
+ },
+ "DiagnosticRelatedInformationMessageString": {
+ "type": "string",
+ "max": 1024,
+ "min": 0,
+ "sensitive": true
+ },
"DiagnosticSeverity": {
"type": "string",
"documentation": "Diagnostic Error types
",
"enum": ["ERROR", "WARNING", "INFORMATION", "HINT"]
},
+ "DiagnosticTag": {
+ "type": "string",
+ "documentation": "The diagnostic tags.
",
+ "enum": ["UNNECESSARY", "DEPRECATED"]
+ },
+ "DiagnosticTagList": {
+ "type": "list",
+ "member": { "shape": "DiagnosticTag" },
+ "documentation": "List of DiagnosticTag
",
+ "max": 1024,
+ "min": 0
+ },
"Dimension": {
"type": "structure",
"members": {
@@ -1378,7 +1586,8 @@
"required": ["type"],
"members": {
"scope": { "shape": "DocumentationIntentContextScopeString" },
- "type": { "shape": "DocumentationType" }
+ "type": { "shape": "DocumentationType" },
+ "changeLogOptions": { "shape": "ChangeLogOptions" }
}
},
"DocumentationIntentContextScopeString": {
@@ -1389,12 +1598,26 @@
},
"DocumentationType": {
"type": "string",
- "enum": ["README"]
+ "enum": ["README", "CHANGE_LOG"]
},
"Double": {
"type": "double",
"box": true
},
+ "Edit": {
+ "type": "structure",
+ "required": ["content"],
+ "members": {
+ "content": { "shape": "EditContentString" },
+ "references": { "shape": "References" }
+ }
+ },
+ "EditContentString": {
+ "type": "string",
+ "max": 5120,
+ "min": 1,
+ "sensitive": true
+ },
"EditorState": {
"type": "structure",
"members": {
@@ -1691,7 +1914,9 @@
"required": ["fileContext"],
"members": {
"fileContext": { "shape": "FileContext" },
+ "editorState": { "shape": "EditorState" },
"maxResults": { "shape": "GenerateCompletionsRequestMaxResultsInteger" },
+ "predictionTypes": { "shape": "PredictionTypes" },
"nextToken": { "shape": "GenerateCompletionsRequestNextTokenString" },
"referenceTrackerConfiguration": { "shape": "ReferenceTrackerConfiguration" },
"supplementalContexts": { "shape": "SupplementalContextList" },
@@ -1718,6 +1943,7 @@
"GenerateCompletionsResponse": {
"type": "structure",
"members": {
+ "predictions": { "shape": "Predictions" },
"completions": { "shape": "Completions" },
"nextToken": { "shape": "SensitiveString" }
}
@@ -2126,9 +2352,48 @@
"featureEvaluations": { "shape": "FeatureEvaluationsList" }
}
},
+ "ListUserMemoryEntriesInput": {
+ "type": "structure",
+ "members": {
+ "maxResults": { "shape": "ListUserMemoryEntriesInputMaxResultsInteger" },
+ "profileArn": {
+ "shape": "ListUserMemoryEntriesInputProfileArnString",
+ "documentation": "ProfileArn for the managing Q Profile
"
+ },
+ "nextToken": { "shape": "ListUserMemoryEntriesInputNextTokenString" }
+ }
+ },
+ "ListUserMemoryEntriesInputMaxResultsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 100,
+ "min": 1
+ },
+ "ListUserMemoryEntriesInputNextTokenString": {
+ "type": "string",
+ "min": 1,
+ "pattern": "\\S+"
+ },
+ "ListUserMemoryEntriesInputProfileArnString": {
+ "type": "string",
+ "min": 1,
+ "pattern": "arn:aws:codewhisperer:[-.a-z0-9]{1,63}:\\d{12}:profile/([a-zA-Z0-9]){12}"
+ },
+ "ListUserMemoryEntriesOutput": {
+ "type": "structure",
+ "required": ["memoryEntries"],
+ "members": {
+ "memoryEntries": { "shape": "MemoryEntryList" },
+ "nextToken": { "shape": "ListUserMemoryEntriesOutputNextTokenString" }
+ }
+ },
+ "ListUserMemoryEntriesOutputNextTokenString": {
+ "type": "string",
+ "min": 1,
+ "pattern": "\\S+"
+ },
"ListWorkspaceMetadataRequest": {
"type": "structure",
- "required": ["workspaceRoot"],
"members": {
"workspaceRoot": { "shape": "ListWorkspaceMetadataRequestWorkspaceRootString" },
"nextToken": { "shape": "String" },
@@ -2154,6 +2419,47 @@
"type": "long",
"box": true
},
+ "MemoryEntry": {
+ "type": "structure",
+ "required": ["id", "memoryEntryString", "metadata"],
+ "members": {
+ "id": {
+ "shape": "MemoryEntryIdString",
+ "documentation": "A unique identifier for a single memory entry
"
+ },
+ "memoryEntryString": { "shape": "MemoryEntryMemoryEntryStringString" },
+ "metadata": { "shape": "MemoryEntryMetadata" }
+ },
+ "documentation": "MemoryEntry corresponds to a single user memory
"
+ },
+ "MemoryEntryIdString": {
+ "type": "string",
+ "max": 36,
+ "min": 36,
+ "pattern": "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
+ },
+ "MemoryEntryList": {
+ "type": "list",
+ "member": { "shape": "MemoryEntry" },
+ "documentation": "List of user memories
"
+ },
+ "MemoryEntryMemoryEntryStringString": {
+ "type": "string",
+ "max": 500,
+ "min": 1,
+ "sensitive": true
+ },
+ "MemoryEntryMetadata": {
+ "type": "structure",
+ "required": ["origin", "createdAt", "updatedAt"],
+ "members": {
+ "origin": { "shape": "Origin" },
+ "attributes": { "shape": "AttributesMap" },
+ "createdAt": { "shape": "Timestamp" },
+ "updatedAt": { "shape": "Timestamp" }
+ },
+ "documentation": "Metadata for a single memory entry
"
+ },
"MessageId": {
"type": "string",
"documentation": "Unique identifier for the chat message
",
@@ -2307,6 +2613,35 @@
"min": 1,
"sensitive": true
},
+ "Prediction": {
+ "type": "structure",
+ "members": {
+ "completion": { "shape": "Completion" },
+ "edit": { "shape": "Edit" }
+ },
+ "union": true
+ },
+ "PredictionType": {
+ "type": "string",
+ "enum": ["Completions", "Edits"]
+ },
+ "PredictionTypes": {
+ "type": "list",
+ "member": { "shape": "PredictionType" }
+ },
+ "Predictions": {
+ "type": "list",
+ "member": { "shape": "Prediction" },
+ "max": 10,
+ "min": 0
+ },
+ "PreviousEditorStateMetadata": {
+ "type": "structure",
+ "required": ["timeOffset"],
+ "members": {
+ "timeOffset": { "shape": "Integer" }
+ }
+ },
"PrimitiveInteger": { "type": "integer" },
"Profile": {
"type": "structure",
@@ -2918,6 +3253,19 @@
"documentation": "Structure to represent stop code transformation response.
"
},
"String": { "type": "string" },
+ "StringList": {
+ "type": "list",
+ "member": { "shape": "StringListMemberString" },
+ "documentation": "A list of strings
",
+ "max": 50,
+ "min": 0
+ },
+ "StringListMemberString": {
+ "type": "string",
+ "max": 256,
+ "min": 1,
+ "sensitive": true
+ },
"SuggestedFix": {
"type": "structure",
"members": {
@@ -2947,7 +3295,9 @@
"required": ["filePath", "content"],
"members": {
"filePath": { "shape": "SupplementalContextFilePathString" },
- "content": { "shape": "SupplementalContextContentString" }
+ "content": { "shape": "SupplementalContextContentString" },
+ "type": { "shape": "SupplementalContextType" },
+ "metadata": { "shape": "SupplementalContextMetadata" }
}
},
"SupplementalContextContentString": {
@@ -2968,6 +3318,17 @@
"max": 5,
"min": 0
},
+ "SupplementalContextMetadata": {
+ "type": "structure",
+ "members": {
+ "previousEditorStateMetadata": { "shape": "PreviousEditorStateMetadata" }
+ },
+ "union": true
+ },
+ "SupplementalContextType": {
+ "type": "string",
+ "enum": ["PreviousEditorState", "WorkspaceContext"]
+ },
"SupplementaryWebLink": {
"type": "structure",
"required": ["url", "title"],
@@ -3306,10 +3667,42 @@
"message": {
"shape": "TextDocumentDiagnosticMessageString",
"documentation": "The diagnostic's message.
"
+ },
+ "code": {
+ "shape": "TextDocumentDiagnosticCodeString",
+ "documentation": "The diagnostic's code, which might appear in the user interface.
"
+ },
+ "codeDescription": {
+ "shape": "CodeDescription",
+ "documentation": "An optional property to describe the error code.
"
+ },
+ "tags": {
+ "shape": "DiagnosticTagList",
+ "documentation": "Additional metadata about the diagnostic.
"
+ },
+ "relatedInformation": {
+ "shape": "DiagnosticRelatedInformationList",
+ "documentation": "an array of related diagnostic information, e.g. when symbol-names within a scope collide all definitions can be marked via this property.
"
+ },
+ "data": {
+ "shape": "TextDocumentDiagnosticDataString",
+ "documentation": "A data entry field that is preserved between a textDocument/publishDiagnostics notification and textDocument/codeAction request.
"
}
},
"documentation": "Structure to represent metadata about a TextDocument Diagnostic
"
},
+ "TextDocumentDiagnosticCodeString": {
+ "type": "string",
+ "max": 1024,
+ "min": 0,
+ "sensitive": true
+ },
+ "TextDocumentDiagnosticDataString": {
+ "type": "string",
+ "max": 4096,
+ "min": 0,
+ "sensitive": true
+ },
"TextDocumentDiagnosticMessageString": {
"type": "string",
"max": 1024,
@@ -3936,6 +4329,7 @@
"members": {
"workspaceId": { "shape": "UUID" },
"workspaceStatus": { "shape": "WorkspaceStatus" },
+ "environmentAddress": { "shape": "SensitiveString" },
"environmentId": { "shape": "SensitiveString" }
}
},
From 8488cbe3344c561fd08572fe7900225465d02e10 Mon Sep 17 00:00:00 2001
From: Tom Zu <138054255+tomcat323@users.noreply.github.com>
Date: Fri, 18 Apr 2025 14:23:23 -0400
Subject: [PATCH 2/2] data instrumentation
---
packages/amazonq/package.json | 35 ++
packages/core/package.json | 3 +-
packages/core/src/codewhisperer/activation.ts | 3 +
.../PredictionKeyStrokeHandler.ts | 118 ++++
.../nextEditPrediction/PredictionTracker.ts | 546 ++++++++++++++++++
.../nextEditPrediction/SnapshotVisualizer.ts | 419 ++++++++++++++
.../nextEditPrediction/activation.ts | 65 +++
.../nextEditPrediction/diffGenerator.ts | 50 ++
.../service/recommendationHandler.ts | 3 +
.../src/codewhisperer/util/editorContext.ts | 40 +-
.../util/supplementalContext/utgUtils.ts | 2 +-
.../core/src/shared/settings-amazonq.gen.ts | 7 +-
12 files changed, 1286 insertions(+), 5 deletions(-)
create mode 100644 packages/core/src/codewhisperer/nextEditPrediction/PredictionKeyStrokeHandler.ts
create mode 100644 packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts
create mode 100644 packages/core/src/codewhisperer/nextEditPrediction/SnapshotVisualizer.ts
create mode 100644 packages/core/src/codewhisperer/nextEditPrediction/activation.ts
create mode 100644 packages/core/src/codewhisperer/nextEditPrediction/diffGenerator.ts
diff --git a/packages/amazonq/package.json b/packages/amazonq/package.json
index 9d90694d93f..9dc60f8604e 100644
--- a/packages/amazonq/package.json
+++ b/packages/amazonq/package.json
@@ -187,6 +187,41 @@
"items": {
"type": "string"
}
+ },
+ "amazonQ.nextEditPrediction.maxFiles": {
+ "type": "number",
+ "default": 15,
+ "description": "Maximum number of files to track for edit predictions",
+ "minimum": 1,
+ "scope": "application"
+ },
+ "amazonQ.nextEditPrediction.maxTotalSizeKb": {
+ "type": "number",
+ "default": 50000,
+ "description": "Maximum total size of snapshots in kilobytes",
+ "minimum": 10,
+ "scope": "application"
+ },
+ "amazonQ.nextEditPrediction.maxFileSizeKb": {
+ "type": "number",
+ "default": 100,
+ "description": "Maximum size per file in kilobytes",
+ "minimum": 5,
+ "scope": "application"
+ },
+ "amazonQ.nextEditPrediction.debounceIntervalMs": {
+ "type": "number",
+ "default": 1000,
+ "description": "Debounce interval in milliseconds between taking snapshots",
+ "minimum": 500,
+ "scope": "application"
+ },
+ "amazonQ.nextEditPrediction.maxAgeMs": {
+ "type": "number",
+ "default": 30000,
+ "description": "Maximum age of snapshots in milliseconds",
+ "minimum": 5000,
+ "scope": "application"
}
}
},
diff --git a/packages/core/package.json b/packages/core/package.json
index 3da76eb1465..86e9092e060 100644
--- a/packages/core/package.json
+++ b/packages/core/package.json
@@ -31,7 +31,8 @@
"./feedback": "./dist/src/feedback/index.js",
"./telemetry": "./dist/src/shared/telemetry/index.js",
"./dev": "./dist/src/dev/index.js",
- "./notifications": "./dist/src/notifications/index.js"
+ "./notifications": "./dist/src/notifications/index.js",
+ "./nep": "./dist/src/nep/index.js"
},
"contributes": {
"icons": {
diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts
index 73f65ca4ada..16cb9750bdf 100644
--- a/packages/core/src/codewhisperer/activation.ts
+++ b/packages/core/src/codewhisperer/activation.ts
@@ -89,10 +89,13 @@ 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 {
+ // Activate the Next Edit Prediction system
+ activateNextEditPrediction(context)
localize = nls.loadMessageBundle()
// Import old CodeWhisperer settings into Amazon Q
diff --git a/packages/core/src/codewhisperer/nextEditPrediction/PredictionKeyStrokeHandler.ts b/packages/core/src/codewhisperer/nextEditPrediction/PredictionKeyStrokeHandler.ts
new file mode 100644
index 00000000000..3d2fbc7e2b3
--- /dev/null
+++ b/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 = 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()
+
+ // 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 = []
+ }
+}
diff --git a/packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts b/packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts
new file mode 100644
index 00000000000..2fa88f4fc6c
--- /dev/null
+++ b/packages/core/src/codewhisperer/nextEditPrediction/PredictionTracker.ts
@@ -0,0 +1,546 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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 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) */
+ maxAgeMs: number
+ /** Maximum number of supplemental contexts to return (default: 15) */
+ maxSupplementalContext: number
+}
+
+/**
+ * Represents a snapshot of a file at a specific point in time
+ */
+export interface FileSnapshot {
+ filePath: string
+ size: number
+ timestamp: number
+ storageKey: string
+}
+
+export class PredictionTracker {
+ private snapshots: Map = new Map()
+ readonly config: FileTrackerConfig
+ private totalSize: number = 0
+ private storagePath?: string
+ private debounceTracker: Set = new Set()
+
+ constructor(extensionContext: vscode.ExtensionContext, config?: Partial) {
+ getLogger().debug('Initializing PredictionTracker')
+
+ // Default configuration values
+ 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
+ }
+
+ this.config = {
+ ...defaultConfig,
+ ...config,
+ }
+
+ // Use workspace storage
+ this.storagePath = extensionContext.storageUri?.fsPath
+
+ void this.ensureStorageDirectoryExists()
+ void this.loadSnapshotsFromStorage()
+
+ // Schedule periodic cleanup
+ // setInterval(() => this.cleanupOldSnapshots(), this.config.maxAgeMs / 2)
+ }
+
+ public processEdit(document: vscode.TextDocument, previousContent: string): void {
+ const filePath = document.uri.fsPath
+ getLogger().debug(`Processing edit for file: ${filePath}`)
+
+ if (!this.storagePath || filePath.startsWith('untitled:') || !document.uri.scheme.startsWith('file')) {
+ return
+ }
+
+ void this.takeSnapshot(filePath, previousContent)
+ }
+
+ /**
+ * Takes a snapshot with provided previous content
+ * @param filePath The path with of document
+ * @param previousContent It's content before the edit
+ */
+ private async takeSnapshot(filePath: string, previousContent: string): Promise {
+ 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)
+
+ const snapshot: FileSnapshot = {
+ filePath,
+ size,
+ timestamp,
+ storageKey,
+ }
+
+ // Get existing snapshots for this file
+ const fileSnapshots = this.snapshots.get(filePath) || []
+
+ // Check if we should add a new snapshot given the debounce
+ const shouldAddSnapshot =
+ fileSnapshots.length === 0 ||
+ timestamp - fileSnapshots[fileSnapshots.length - 1].timestamp > this.config.debounceIntervalMs
+
+ if (shouldAddSnapshot) {
+ try {
+ // Save to workspace storage
+ await this.saveSnapshotContentToStorage(storageKey, content)
+
+ fileSnapshots.push(snapshot)
+ this.snapshots.set(filePath, fileSnapshots)
+
+ // Update total size
+ this.totalSize += size
+
+ // Enforce memory limits
+ await this.enforceMemoryLimits()
+
+ // Set a timeout to delete the snapshot after maxAgeMs
+ setTimeout(async () => {
+ const index = fileSnapshots.indexOf(snapshot)
+ if (index !== -1) {
+ fileSnapshots.splice(index, 1)
+ this.totalSize -= size
+ await this.deleteSnapshotFromStorage(snapshot)
+ if (fileSnapshots.length === 0) {
+ this.snapshots.delete(filePath)
+ }
+ }
+ }, this.config.maxAgeMs)
+ } catch (err) {
+ getLogger().error(`Failed to save snapshot to Storage: ${err}`)
+ }
+ }
+ }
+
+ /**
+ * Generates a unique storage key for a snapshot
+ */
+ private generateStorageKey(filePath: string, timestamp: number): string {
+ const fileName = path.basename(filePath)
+ return `${fileName}-${timestamp}`
+ }
+
+ /**
+ * Enforces memory limits by removing old snapshots if necessary
+ */
+ private async enforceMemoryLimits(): Promise {
+ // Enforce total size limit
+ while (this.totalSize > this.config.maxTotalSizeKb * 1024) {
+ const oldestFile = this.findOldestFile()
+ if (!oldestFile) {
+ break
+ }
+
+ const fileSnapshots = this.snapshots.get(oldestFile)
+ if (!fileSnapshots || fileSnapshots.length === 0) {
+ this.snapshots.delete(oldestFile)
+ continue
+ }
+
+ // Remove the oldest snapshot
+ const removedSnapshot = fileSnapshots.shift()
+ if (removedSnapshot) {
+ this.totalSize -= removedSnapshot.size
+ await this.deleteSnapshotFromStorage(removedSnapshot)
+ }
+
+ // If no snapshots left for this file, remove the file entry
+ if (fileSnapshots.length === 0) {
+ this.snapshots.delete(oldestFile)
+ }
+ }
+
+ // Enforce max files limit
+ while (this.snapshots.size > this.config.maxFiles) {
+ const oldestFile = this.findOldestFile()
+ if (!oldestFile) {
+ break
+ }
+
+ const fileSnapshots = this.snapshots.get(oldestFile)
+ if (fileSnapshots) {
+ // Subtract all snapshot sizes from the total
+ for (const snapshot of fileSnapshots) {
+ this.totalSize -= snapshot.size
+ await this.deleteSnapshotFromStorage(snapshot)
+ }
+ }
+
+ this.snapshots.delete(oldestFile)
+ }
+ }
+
+ /**
+ * Finds the file with the oldest snapshot
+ * @returns The file path of the oldest snapshot
+ */
+ private findOldestFile(): string | undefined {
+ let oldestTime = Number.MAX_SAFE_INTEGER
+ let oldestFile: string | undefined
+
+ for (const [filePath, snapshots] of this.snapshots.entries()) {
+ if (snapshots.length === 0) {
+ continue
+ }
+
+ const oldestSnapshot = snapshots[0]
+ if (oldestSnapshot.timestamp < oldestTime) {
+ oldestTime = oldestSnapshot.timestamp
+ oldestFile = filePath
+ }
+ }
+
+ return oldestFile
+ }
+
+ /**
+ * Removes snapshots that are older than the maximum age
+ */
+ private cleanupOldSnapshots(): void {
+ const now = Date.now()
+ const maxAge = this.config.maxAgeMs
+
+ for (const [filePath, snapshots] of this.snapshots.entries()) {
+ const validSnapshots = snapshots.filter((snapshot) => {
+ const isValid = now - snapshot.timestamp <= maxAge
+ if (!isValid) {
+ this.totalSize -= snapshot.size
+ void this.deleteSnapshotFromStorage(snapshot)
+ }
+ return isValid
+ })
+
+ if (validSnapshots.length === 0) {
+ this.snapshots.delete(filePath)
+ } else {
+ this.snapshots.set(filePath, validSnapshots)
+ }
+ }
+ }
+
+ /**
+ * Gets all snapshots for a specific file
+ * @param filePath The path to the file
+ * @returns Array of snapshots for the file
+ */
+ public getFileSnapshots(filePath: string): FileSnapshot[] {
+ return this.snapshots.get(filePath) || []
+ }
+
+ /**
+ * Gets all tracked files
+ * @returns Array of file paths
+ */
+ public getTrackedFiles(): string[] {
+ return Array.from(this.snapshots.keys())
+ }
+
+ /**
+ * Gets the total number of snapshots across all files
+ * @returns Total snapshot count
+ */
+ public getTotalSnapshotCount(): number {
+ let count = 0
+ for (const snapshots of this.snapshots.values()) {
+ count += snapshots.length
+ }
+ return count
+ }
+
+ /**
+ * Saves snapshot content to Storage
+ * @param storageKey The storage key for the snapshot
+ * @param content The content to save
+ */
+ private async saveSnapshotContentToStorage(storageKey: string, content: string): Promise {
+ if (!this.storagePath) {
+ throw new Error('Storage path not available')
+ }
+
+ const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+ if (!(await fs.existsDir(snapshotsDir))) {
+ await fs.mkdir(snapshotsDir)
+ }
+
+ const filePath = path.join(snapshotsDir, `${storageKey}.content`)
+ await fs.writeFile(filePath, content)
+ }
+
+ /**
+ * Deletes a snapshot content from Storage
+ * @param snapshot The snapshot to delete
+ */
+ private async deleteSnapshotFromStorage(snapshot: FileSnapshot): Promise {
+ if (!this.storagePath) {
+ return
+ }
+
+ const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+ const filePath = path.join(snapshotsDir, `${snapshot.storageKey}.content`)
+
+ if (await fs.exists(filePath)) {
+ try {
+ await fs.delete(filePath)
+ } catch (err) {
+ getLogger().error(`Failed to delete snapshot from Storage: ${err}`)
+ }
+ }
+ }
+
+ /**
+ * Loads snapshot content from Storage
+ * @param snapshot The snapshot metadata
+ * @returns The string content of the snapshot
+ */
+ public async getSnapshotContent(snapshot: FileSnapshot): Promise {
+ if (!this.storagePath) {
+ throw new Error('Storage path not available')
+ }
+
+ const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+ const filePath = path.join(snapshotsDir, `${snapshot.storageKey}.content`)
+
+ try {
+ return await fs.readFileText(filePath)
+ } catch (err) {
+ getLogger().error(`Failed to read snapshot content from Storage: ${err}`)
+ throw new Error(`Failed to read snapshot content: ${err}`)
+ }
+ }
+
+ /**
+ * 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 {
+ const activeEditor = vscode.window.activeTextEditor
+ if (activeEditor === undefined) {
+ return []
+ }
+ 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]
+
+ 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,
+ },
+ },
+ })
+ } catch (err) {
+ getLogger().error(`Failed to generate diff with current 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
+ }
+
+ private async ensureStorageDirectoryExists(): Promise {
+ if (!this.storagePath) {
+ return
+ }
+
+ const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+ if (!(await fs.existsDir(snapshotsDir))) {
+ await fs.mkdir(snapshotsDir)
+ }
+ }
+
+ private async loadSnapshotsFromStorage(): Promise {
+ if (!this.storagePath) {
+ return
+ }
+
+ const snapshotsDir = path.join(this.storagePath, 'file-snapshots')
+ if (!(await fs.existsDir(snapshotsDir))) {
+ return
+ }
+
+ try {
+ const files = await fs.readdir(snapshotsDir)
+ const metadataFiles = new Map()
+
+ // First, collect all the metadata files
+ for (const [filename, fileType] of files) {
+ if (!filename.endsWith('.content') || fileType !== vscode.FileType.File) {
+ continue
+ }
+
+ const storageKey = filename.substring(0, filename.length - '.content'.length)
+ const parts = storageKey.split('-')
+ const timestamp = parseInt(parts[parts.length - 1], 10)
+ const originalFilename = parts.slice(0, parts.length - 1).join('-')
+
+ // This helps us match the files back to their original source
+ metadataFiles.set(storageKey, {
+ timestamp,
+ filePath: originalFilename,
+ })
+ }
+
+ // Now process each file that we found
+ for (const [storageKey, metadata] of metadataFiles.entries()) {
+ const contentPath = path.join(snapshotsDir, `${storageKey}.content`)
+
+ try {
+ if (!(await fs.exists(metadata.filePath))) {
+ await fs.delete(contentPath)
+ continue
+ }
+
+ // Calculate size from the content file
+ const stats = await fs.stat(contentPath)
+ const size = stats.size
+
+ // Create a metadata-only snapshot
+ const snapshot: FileSnapshot = {
+ filePath: metadata.filePath,
+ timestamp: metadata.timestamp,
+ size,
+ storageKey,
+ }
+
+ // Add to memory tracking
+ const fileSnapshots = this.snapshots.get(metadata.filePath) || []
+ fileSnapshots.push(snapshot)
+ this.snapshots.set(metadata.filePath, fileSnapshots)
+ this.totalSize += size
+ } catch (err) {
+ // Remove invalid files
+ getLogger().error(`Error processing snapshot file ${storageKey}: ${err}`)
+ await fs.delete(contentPath)
+ }
+ }
+
+ // Sort snapshots by timestamp
+ for (const [filePath, snapshots] of this.snapshots.entries()) {
+ this.snapshots.set(
+ filePath,
+ snapshots.sort((a, b) => a.timestamp - b.timestamp)
+ )
+ }
+
+ // Apply memory limits after loading
+ await this.enforceMemoryLimits()
+ this.cleanupOldSnapshots()
+
+ getLogger().info(`Loaded ${this.getTotalSnapshotCount()} snapshots for ${this.snapshots.size} files`)
+ } catch (err) {
+ getLogger().error(`Failed to load snapshots from Storage: ${err}`)
+ }
+ }
+
+ /**
+ * Disposes of resources used by the tracker
+ */
+ public dispose(): void {
+ this.debounceTracker.clear()
+ }
+}
diff --git a/packages/core/src/codewhisperer/nextEditPrediction/SnapshotVisualizer.ts b/packages/core/src/codewhisperer/nextEditPrediction/SnapshotVisualizer.ts
new file mode 100644
index 00000000000..f597354048b
--- /dev/null
+++ b/packages/core/src/codewhisperer/nextEditPrediction/SnapshotVisualizer.ts
@@ -0,0 +1,419 @@
+/*!
+ * 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 { FileSnapshot, PredictionTracker } from './PredictionTracker'
+
+/**
+ * Snapshot visualizer for developers to debug file tracking
+ */
+export class SnapshotVisualizer {
+ public static readonly viewType = 'amazonQ.nextEditPrediction.SnapshotVisualizer'
+ private panel: vscode.WebviewPanel | undefined
+ private readonly predictionTracker: PredictionTracker
+
+ constructor(
+ private readonly extensionContext: vscode.ExtensionContext,
+ predictionTracker: PredictionTracker
+ ) {
+ this.predictionTracker = predictionTracker
+ }
+
+ /**
+ * Shows the snapshot visualizer panel
+ */
+ public show(): void {
+ if (this.panel) {
+ this.panel.reveal(vscode.ViewColumn.Beside)
+ return
+ }
+
+ this.panel = vscode.window.createWebviewPanel(
+ SnapshotVisualizer.viewType,
+ 'File Snapshot Visualizer',
+ vscode.ViewColumn.Beside,
+ {
+ enableScripts: true,
+ retainContextWhenHidden: true,
+ localResourceRoots: [vscode.Uri.file(this.extensionContext.extensionPath)],
+ }
+ )
+
+ this.panel.webview.html = this.getWebviewContent()
+
+ // Message handling
+ this.panel.webview.onDidReceiveMessage(
+ async (message) => {
+ switch (message.command) {
+ case 'refresh':
+ this.updateContent()
+ break
+ case 'showSnapshot':
+ await this.showSnapshotContent(message.filePath, message.timestamp)
+ break
+ case 'fireAPI':
+ await this.generateDiffs()
+ break
+ }
+ },
+ undefined,
+ this.extensionContext.subscriptions
+ )
+
+ this.panel.onDidDispose(
+ () => {
+ this.panel = undefined
+ },
+ undefined,
+ this.extensionContext.subscriptions
+ )
+
+ // Update content every 0.5 seconds
+ const interval = setInterval(() => {
+ if (this.panel) {
+ this.updateContent()
+ } else {
+ clearInterval(interval)
+ }
+ }, 500)
+ this.extensionContext.subscriptions.push({ dispose: () => clearInterval(interval) })
+ }
+
+ /**
+ * Updates the content of the webview
+ */
+ private updateContent(): void {
+ if (!this.panel) {
+ return
+ }
+
+ const trackedFiles = this.predictionTracker.getTrackedFiles()
+ const fileData: { [key: string]: FileSnapshot[] } = {}
+
+ for (const filePath of trackedFiles) {
+ fileData[filePath] = this.predictionTracker.getFileSnapshots(filePath)
+ }
+
+ void this.panel.webview.postMessage({
+ command: 'updateFiles',
+ files: fileData,
+ totalCount: this.predictionTracker.getTotalSnapshotCount(),
+ })
+ }
+
+ /**
+ * Generates diffs between adjacent snapshots of the currently opened file
+ * and between the newest snapshot and the current file content
+ */
+ private async generateDiffs(): Promise {
+ if (!this.panel) {
+ return
+ }
+
+ // Get the currently active text editor
+ const editor = vscode.window.activeTextEditor
+ if (!editor) {
+ void vscode.window.showErrorMessage('No active text editor found')
+ return
+ }
+
+ const filePath = editor.document.uri.fsPath
+
+ // Generate diffs using the PredictionTracker's method
+ try {
+ const supplementalContexts = await this.predictionTracker.generatePredictionSupplementalContext()
+
+ if (supplementalContexts.length === 0) {
+ void vscode.window.showInformationMessage('No snapshots found for the current file')
+ return
+ }
+
+ // Send the full supplemental contexts to webview
+ void this.panel.webview.postMessage({
+ command: 'showDiffs',
+ filePath,
+ diffs: supplementalContexts,
+ })
+ } catch (err) {
+ getLogger().error(`Failed to generate diffs: ${err}`)
+ void vscode.window.showErrorMessage('Failed to generate diffs')
+ }
+ }
+
+ /**
+ * Shows the content of a specific snapshot
+ */
+ private async showSnapshotContent(filePath: string, timestamp: number): Promise {
+ if (!this.panel) {
+ return
+ }
+
+ const snapshots = this.predictionTracker.getFileSnapshots(filePath)
+ const snapshot = snapshots.find((s: FileSnapshot) => s.timestamp === timestamp)
+
+ if (snapshot) {
+ try {
+ // Load content from storage on demand
+ const content = await this.predictionTracker.getSnapshotContent(snapshot)
+
+ void this.panel.webview.postMessage({
+ command: 'showSnapshotContent',
+ filePath: snapshot.filePath,
+ timestamp: snapshot.timestamp,
+ content: content,
+ })
+ } catch (err) {
+ getLogger().error(`Failed to load snapshot content: ${err}`)
+
+ void this.panel.webview.postMessage({
+ command: 'showSnapshotContent',
+ filePath: snapshot.filePath,
+ timestamp: snapshot.timestamp,
+ content: '(Error loading content)',
+ })
+ }
+ }
+ }
+
+ /**
+ * Gets the HTML content for the webview
+ */
+ private getWebviewContent(): string {
+ return `
+
+
+
+
+ File Snapshot Visualizer
+
+
+
+
+
+
+
No files tracked yet. Make some edits to see snapshots appear.
+
+
+
+
+
+
+ `
+ }
+}
diff --git a/packages/core/src/codewhisperer/nextEditPrediction/activation.ts b/packages/core/src/codewhisperer/nextEditPrediction/activation.ts
new file mode 100644
index 00000000000..e3dae72efc7
--- /dev/null
+++ b/packages/core/src/codewhisperer/nextEditPrediction/activation.ts
@@ -0,0 +1,65 @@
+/*!
+ * 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'
+import { SnapshotVisualizer } from './SnapshotVisualizer'
+
+export let predictionTracker: PredictionTracker | undefined
+let keyStrokeHandler: PredictionKeyStrokeHandler | undefined
+
+/**
+ * Activates the Next Edit Prediction system
+ */
+export function activateNextEditPrediction(context: ExtContext): void {
+ // Initialize the tracker
+ predictionTracker = new PredictionTracker(context.extensionContext)
+
+ // Initialize the keystroke handler
+ keyStrokeHandler = new PredictionKeyStrokeHandler(predictionTracker)
+
+ // Register for disposal
+ context.extensionContext.subscriptions.push(
+ vscode.Disposable.from({
+ dispose: () => {
+ getLogger().info('Disposing Next Edit Prediction resources')
+ keyStrokeHandler?.dispose()
+ predictionTracker?.dispose()
+ },
+ })
+ )
+
+ // Register snapshot visualizer
+ registerSnapshotVisualizer(context, predictionTracker)
+
+ getLogger().info('Next Edit Prediction activated')
+}
+
+/**
+ * Registers the snapshot visualizer command and status bar item
+ */
+function registerSnapshotVisualizer(context: ExtContext, tracker: PredictionTracker): void {
+ // Create the visualizer
+ const visualizer = new SnapshotVisualizer(context.extensionContext, tracker)
+
+ // Register command
+ context.extensionContext.subscriptions.push(
+ vscode.commands.registerCommand('amazonQ.nextEditPrediction.showSnapshotVisualizer', () => {
+ visualizer.show()
+ })
+ )
+
+ // Add a status bar item to open the visualizer
+ const statusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 100)
+ statusBarItem.text = '$(history) NEP'
+ statusBarItem.tooltip = 'Show Next Edit Prediction snapshot visualizer'
+ statusBarItem.command = 'amazonQ.nextEditPrediction.showSnapshotVisualizer'
+ statusBarItem.show()
+ statusBarItem.show()
+ context.extensionContext.subscriptions.push(statusBarItem)
+}
diff --git a/packages/core/src/codewhisperer/nextEditPrediction/diffGenerator.ts b/packages/core/src/codewhisperer/nextEditPrediction/diffGenerator.ts
new file mode 100644
index 00000000000..fa903a836e6
--- /dev/null
+++ b/packages/core/src/codewhisperer/nextEditPrediction/diffGenerator.ts
@@ -0,0 +1,50 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as diff from 'diff'
+
+/**
+ * Class to generate a unified diff format between old and new file contents
+ */
+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 {
+ 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')
+ }
+
+ return patchResult
+ }
+}
diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/packages/core/src/codewhisperer/service/recommendationHandler.ts
index 00d1f3254a5..b1999dd65d1 100644
--- a/packages/core/src/codewhisperer/service/recommendationHandler.ts
+++ b/packages/core/src/codewhisperer/service/recommendationHandler.ts
@@ -220,6 +220,8 @@ export class RecommendationHandler {
session.requestContext = await EditorContext.buildGenerateRecommendationRequest(editor as vscode.TextEditor)
}
const request = session.requestContext.request
+ // eslint-disable-next-line aws-toolkits/no-json-stringify-in-log
+ getLogger().info(JSON.stringify(request))
// record preprocessing end time
TelemetryHelper.instance.setPreprocessEndTime()
@@ -329,6 +331,7 @@ export class RecommendationHandler {
session.requestIdList.push(requestId)
}
getLogger().debug(msg)
+ getLogger().info(`NEP testing: codeWhisper request ${requestId}`)
if (invocationResult === 'Succeeded') {
CodeWhispererCodeCoverageTracker.getTracker(session.language)?.incrementServiceInvocationCount()
UserWrittenCodeTracker.instance.onQFeatureInvoked()
diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts
index 11598cfe20c..2ab84dece75 100644
--- a/packages/core/src/codewhisperer/util/editorContext.ts
+++ b/packages/core/src/codewhisperer/util/editorContext.ts
@@ -18,6 +18,7 @@ import { checkLeftContextKeywordsForJson } from './commonUtil'
import { CodeWhispererSupplementalContext } from '../models/model'
import { getOptOutPreference } from '../../shared/telemetry/util'
import { indent } from '../../shared/utilities/textUtilities'
+import { predictionTracker } from '../nextEditPrediction/activation'
let tabSize: number = getTabSizeSetting()
@@ -101,13 +102,45 @@ export async function buildListRecommendationRequest(
logSupplementalContext(supplementalContexts)
+ // Get predictionSupplementalContext from PredictionTracker
+ let predictionSupplementalContext: codewhispererClient.SupplementalContext[] = []
+ try {
+ if (predictionTracker) {
+ predictionSupplementalContext = await predictionTracker.generatePredictionSupplementalContext()
+ }
+ } catch (error) {
+ getLogger().error(`Error getting prediction supplemental context: ${error}`)
+ }
+
const selectedCustomization = getSelectedCustomization()
- const supplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts
+ const inlineSupplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts
? supplementalContexts.supplementalContextItems.map((v) => {
return selectFrom(v, 'content', 'filePath')
})
: []
+ // Dynamically create editorState from current editor context
+ const editorState = {
+ document: {
+ programmingLanguage: {
+ languageName: fileContext.programmingLanguage.languageName,
+ },
+ relativeFilePath: fileContext.filename,
+ text: editor.document.getText(),
+ },
+ cursorState: {
+ position: {
+ line: editor.selection.active.line,
+ character: editor.selection.active.character,
+ },
+ },
+ }
+
+ const predictionTypes = ['EDITS']
+
+ // Combine inline and prediction supplemental contexts
+ const finalSupplementalContext = inlineSupplementalContext.concat(predictionSupplementalContext)
+
return {
request: {
fileContext: fileContext,
@@ -115,7 +148,10 @@ export async function buildListRecommendationRequest(
referenceTrackerConfiguration: {
recommendationsWithReferences: allowCodeWithReference ? 'ALLOW' : 'BLOCK',
},
- supplementalContexts: supplementalContext,
+ supplementalContexts: finalSupplementalContext,
+ editorState: editorState,
+ predictionTypes: predictionTypes,
+ maxResults: CodeWhispererConstants.maxRecommendations,
customizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn,
optOutPreference: getOptOutPreference(),
},
diff --git a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts b/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts
index 0d33969773e..afd12499304 100644
--- a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts
+++ b/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts
@@ -102,7 +102,7 @@ export async function fetchSupplementalContextForTest(
async function generateSupplementalContextFromFocalFile(
filePath: string,
strategy: UtgStrategy,
- cancellationToken: vscode.CancellationToken
+ cancellationToken?: vscode.CancellationToken
): Promise {
const fileContent = await fs.readFileText(vscode.Uri.parse(filePath!).fsPath)
diff --git a/packages/core/src/shared/settings-amazonq.gen.ts b/packages/core/src/shared/settings-amazonq.gen.ts
index f0a3d47f989..551234806fe 100644
--- a/packages/core/src/shared/settings-amazonq.gen.ts
+++ b/packages/core/src/shared/settings-amazonq.gen.ts
@@ -31,7 +31,12 @@ export const amazonqSettings = {
"amazonQ.workspaceIndexWorkerThreads": {},
"amazonQ.workspaceIndexUseGPU": {},
"amazonQ.workspaceIndexMaxSize": {},
- "amazonQ.ignoredSecurityIssues": {}
+ "amazonQ.ignoredSecurityIssues": {},
+ "amazonQ.nextEditPrediction.maxFiles": {},
+ "amazonQ.nextEditPrediction.maxTotalSizeKb": {},
+ "amazonQ.nextEditPrediction.maxFileSizeKb": {},
+ "amazonQ.nextEditPrediction.debounceIntervalMs": {},
+ "amazonQ.nextEditPrediction.maxAgeMs": {}
}
export default amazonqSettings