diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts index b0aa54e17a0..7e48c7549bc 100644 --- a/packages/core/src/codewhisperer/activation.ts +++ b/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 { }) ) } + + activateEditTracking(context) } export async function shutdown() { diff --git a/packages/core/src/codewhisperer/models/constants.ts b/packages/core/src/codewhisperer/models/constants.ts index 2fb3dd10069..dc0426376ce 100644 --- a/packages/core/src/codewhisperer/models/constants.ts +++ b/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, +} diff --git a/packages/core/src/codewhisperer/nextEditPrediction/activation.ts b/packages/core/src/codewhisperer/nextEditPrediction/activation.ts new file mode 100644 index 00000000000..f302cec2ad5 --- /dev/null +++ b/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}`) + } +} diff --git a/packages/core/src/codewhisperer/nextEditPrediction/diffContextGenerator.ts b/packages/core/src/codewhisperer/nextEditPrediction/diffContextGenerator.ts new file mode 100644 index 00000000000..9f379b82a4e --- /dev/null +++ b/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 +} diff --git a/packages/core/src/codewhisperer/nextEditPrediction/predictionKeyStrokeHandler.ts b/packages/core/src/codewhisperer/nextEditPrediction/predictionKeyStrokeHandler.ts new file mode 100644 index 00000000000..b09272f0d8b --- /dev/null +++ b/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 = 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() + + 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 = [] + } +} diff --git a/packages/core/src/codewhisperer/nextEditPrediction/predictionTracker.ts b/packages/core/src/codewhisperer/nextEditPrediction/predictionTracker.ts new file mode 100644 index 00000000000..1bfafb53114 --- /dev/null +++ b/packages/core/src/codewhisperer/nextEditPrediction/predictionTracker.ts @@ -0,0 +1,236 @@ +/*! + * 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 * as diffGenerator from './diffContextGenerator' +import * as codewhispererClient from '../client/codewhisperer' +import { predictionTrackerDefaultConfig } from '../models/constants' +import globals from '../../shared/extensionGlobals' + +// defaul values are stored in codewhisperer/model/constants +export interface FileTrackerConfig { + maxStorageSizeKb: number + debounceIntervalMs: number + maxAgeMs: number + maxSupplementalContext: number +} + +/** + * Represents a snapshot of a file at a specific point in time + */ +export interface FileSnapshot { + filePath: string + size: number + timestamp: number + content: string +} + +export class PredictionTracker { + private snapshots: Map = new Map() + private logger = getLogger('nextEditPrediction') + readonly config: FileTrackerConfig + private storageSize: number = 0 + + constructor(extensionContext: vscode.ExtensionContext, config?: Partial) { + this.config = { + ...predictionTrackerDefaultConfig, + ...config, + } + } + + /** + * Processes an edit to a document and takes a snapshot if needed + * @param document The document being edited + * @param previousContent The content of the document before the edit + */ + public async processEdit(document: vscode.TextDocument, previousContent: string): Promise { + const filePath = document.uri.fsPath + + try { + // Get existing snapshots for this file + const fileSnapshots = this.snapshots.get(filePath) || [] + const timestamp = globals.clock.Date.now() + + // Anti-throttling, only add snap shot after the debounce is cleared + const shouldAddSnapshot = + fileSnapshots.length === 0 || + timestamp - fileSnapshots[fileSnapshots.length - 1].timestamp > this.config.debounceIntervalMs + + if (!shouldAddSnapshot) { + return + } + + const content = previousContent + const size = Buffer.byteLength(content, 'utf8') + const snapshot: FileSnapshot = { + filePath, + size, + timestamp, + content, + } + + fileSnapshots.push(snapshot) + this.snapshots.set(filePath, fileSnapshots) + this.storageSize += size + this.logger.debug( + `Snapshot taken for file: ${filePath}, total snapshots: ${this.getTotalSnapshotCount()}, total size: ${Math.round(this.storageSize / 1024)} KB` + ) + + await this.enforceMemoryLimits() + this.enforceTimeLimits(snapshot) + } catch (err) { + this.logger.error(`Failed to save snapshot: ${err}`) + } + } + + /** + * Sets up a timeout to delete the given snapshot after it exceeds the max age + */ + private enforceTimeLimits(snapshot: FileSnapshot): void { + const fileSnapshots = this.snapshots.get(snapshot.filePath) + if (fileSnapshots === undefined) { + return + } + + setTimeout(() => { + // find the snapshot and remove it + const index = fileSnapshots.indexOf(snapshot) + if (index !== -1) { + fileSnapshots.splice(index, 1) + this.storageSize -= snapshot.size + if (fileSnapshots.length === 0) { + this.snapshots.delete(snapshot.filePath) + } + this.logger.debug( + `Snapshot deleted (aged out) for file: ${snapshot.filePath}, remaining snapshots: ${this.getTotalSnapshotCount()}, new size: ${Math.round(this.storageSize / 1024)} KB` + ) + } + }, this.config.maxAgeMs) + } + + /** + * Enforces memory limits by removing old snapshots if necessary + */ + private async enforceMemoryLimits(): Promise { + while (this.storageSize > this.config.maxStorageSizeKb * 1024) { + const oldestFile = this.findOldestFile() + if (!oldestFile) { + break + } + + const fileSnapshots = this.snapshots.get(oldestFile) + if (!fileSnapshots || fileSnapshots.length === 0) { + this.snapshots.delete(oldestFile) + continue + } + + const removedSnapshot = fileSnapshots.shift() + if (removedSnapshot) { + this.storageSize -= removedSnapshot.size + this.logger.debug( + `Snapshot deleted (memory limit) for file: ${removedSnapshot.filePath}, remaining snapshots: ${this.getTotalSnapshotCount()}, new size: ${Math.round(this.storageSize / 1024)} KB` + ) + } + + if (fileSnapshots.length === 0) { + 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 + } + + /** + * 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()) + } + + public getTotalSnapshotCount(): number { + return Array.from(this.snapshots.values()).reduce((count, snapshots) => count + snapshots.length, 0) + } + + public async getSnapshotContent(snapshot: FileSnapshot): Promise { + return snapshot.content + } + + /** + * Generates unified diffs between adjacent snapshots of a file + * and between the newest snapshot and the current file content + * + * @returns Array of SupplementalContext objects containing diffs between snapshots and current content + */ + public async generatePredictionSupplementalContext(): Promise { + try { + const activeEditor = vscode.window.activeTextEditor + if (activeEditor === undefined) { + return [] + } + const filePath = activeEditor.document.uri.fsPath + const currentContent = activeEditor.document.getText() + const snapshots = this.getFileSnapshots(filePath) + + if (snapshots.length === 0) { + return [] + } + + // Create SnapshotContent array from snapshots + const snapshotContents: diffGenerator.SnapshotContent[] = snapshots.map((snapshot) => ({ + filePath: snapshot.filePath, + content: snapshot.content, + timestamp: snapshot.timestamp, + })) + + // Use the diffGenerator module to generate supplemental contexts + return diffGenerator.generateDiffContexts( + filePath, + currentContent, + snapshotContents, + this.config.maxSupplementalContext + ) + } catch (err) { + // this ensures we are not breaking inline requests + this.logger.error(`Failed to generate prediction supplemental context: ${err}`) + return [] + } + } + + public getTotalSize() { + return this.storageSize + } +} diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/packages/core/src/codewhisperer/service/recommendationHandler.ts index 00d1f3254a5..8ab491b32e0 100644 --- a/packages/core/src/codewhisperer/service/recommendationHandler.ts +++ b/packages/core/src/codewhisperer/service/recommendationHandler.ts @@ -328,7 +328,7 @@ export class RecommendationHandler { msg += `\n ${index.toString().padStart(2, '0')}: ${indent(item.content, 8, true).trim()}` session.requestIdList.push(requestId) } - getLogger().debug(msg) + getLogger('nextEditPrediction').debug(`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 88c3d3847f1..0861b982d13 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -20,6 +20,7 @@ import { getOptOutPreference } from '../../shared/telemetry/util' import { indent } from '../../shared/utilities/textUtilities' import { isInDirectory } from '../../shared/filesystemUtilities' import { AuthUtil } from './authUtil' +import { predictionTracker } from '../nextEditPrediction/activation' let tabSize: number = getTabSizeSetting() @@ -119,8 +120,14 @@ export async function buildListRecommendationRequest( logSupplementalContext(supplementalContexts) + // Get predictionSupplementalContext from PredictionTracker + let predictionSupplementalContext: codewhispererClient.SupplementalContext[] = [] + if (predictionTracker) { + predictionSupplementalContext = await predictionTracker.generatePredictionSupplementalContext() + } + const selectedCustomization = getSelectedCustomization() - const supplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts + const completionSupplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts ? supplementalContexts.supplementalContextItems.map((v) => { return selectFrom(v, 'content', 'filePath') }) @@ -128,6 +135,10 @@ export async function buildListRecommendationRequest( const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile + const editorState = getEditorState(editor, fileContext) + + // Combine inline and prediction supplemental contexts + const finalSupplementalContext = completionSupplementalContext.concat(predictionSupplementalContext) return { request: { fileContext: fileContext, @@ -135,7 +146,9 @@ export async function buildListRecommendationRequest( referenceTrackerConfiguration: { recommendationsWithReferences: allowCodeWithReference ? 'ALLOW' : 'BLOCK', }, - supplementalContexts: supplementalContext, + supplementalContexts: finalSupplementalContext, + editorState: editorState, + maxResults: CodeWhispererConstants.maxRecommendations, customizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn, optOutPreference: getOptOutPreference(), workspaceId: await getWorkspaceId(editor), @@ -201,6 +214,29 @@ export function getTabSize(): number { return tabSize } +export function getEditorState(editor: vscode.TextEditor, fileContext: codewhispererClient.FileContext): any { + try { + return { + 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, + }, + }, + } + } catch (error) { + getLogger().error(`Error generating editor state: ${error}`) + return undefined + } +} + export function getLeftContext(editor: vscode.TextEditor, line: number): string { let lineText = '' try { diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts index 85df7b4e1f8..38ff40627c1 100644 --- a/packages/core/src/shared/logger/logger.ts +++ b/packages/core/src/shared/logger/logger.ts @@ -18,6 +18,7 @@ export type LogTopic = | 'chat' | 'stepfunctions' | 'unknown' + | 'nextEditPrediction' | 'resourceCache' class ErrorLog { diff --git a/packages/core/src/test/codewhisperer/nextEditPrediction/predictionTracker.test.ts b/packages/core/src/test/codewhisperer/nextEditPrediction/predictionTracker.test.ts new file mode 100644 index 00000000000..715044fdf73 --- /dev/null +++ b/packages/core/src/test/codewhisperer/nextEditPrediction/predictionTracker.test.ts @@ -0,0 +1,344 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as vscode from 'vscode' +import * as sinon from 'sinon' +import assert from 'assert' +import * as path from 'path' + +import { + FileSnapshot, + FileTrackerConfig, + PredictionTracker, +} from '../../../codewhisperer/nextEditPrediction/predictionTracker' +import { FakeExtensionContext } from '../../fakeExtensionContext' +import { createMockDocument } from '../testUtil' +import * as diffGenerator from '../../../codewhisperer/nextEditPrediction/diffContextGenerator' +import globals from '../../../shared/extensionGlobals' +import { charactersLimit, supplementalContextMaxTotalLength } from '../../../codewhisperer/models/constants' + +describe('PredictionTracker', function () { + let sandbox: sinon.SinonSandbox + let mockExtensionContext: vscode.ExtensionContext + let tracker: PredictionTracker + let clock: sinon.SinonFakeTimers + let dateNowStub: sinon.SinonStub + + beforeEach(async function () { + sandbox = sinon.createSandbox() + // Set a base time for tests + const startTime = new Date('2025-04-21T12:00:00Z').getTime() + + clock = sandbox.useFakeTimers({ + now: startTime, + shouldAdvanceTime: true, + }) + + // Set up a stub for globals.clock.Date.now() that we can control manually + dateNowStub = sandbox.stub(globals.clock.Date, 'now') + dateNowStub.returns(startTime) + + mockExtensionContext = await FakeExtensionContext.create() + }) + + afterEach(function () { + sandbox.restore() + clock.restore() + }) + + describe('processEdit', function () { + let filePath: string + let previousContent: string + let mockDocument: vscode.TextDocument + + beforeEach(function () { + filePath = testPath('path', 'to', 'file.js') + previousContent = 'previous content' + tracker = new PredictionTracker(mockExtensionContext) + + // Create a mock document + mockDocument = createMockDocument(previousContent, filePath) + }) + + it('should store snapshot in memory', async function () { + await tracker.processEdit(mockDocument, previousContent) + const snapshots = tracker.getFileSnapshots(filePath) + + assert.strictEqual(snapshots.length, 1) + assert.strictEqual(snapshots[0].content, previousContent) + assert.strictEqual(snapshots[0].size, Buffer.byteLength(previousContent, 'utf8')) + }) + + it('should not add new snapshot within debounce interval', async function () { + await tracker.processEdit(mockDocument, 'first edit') + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 1) + + // Another edit within debounce interval, should not add another snapshot + await tracker.processEdit(mockDocument, 'second edit') + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 1) + }) + + it('should add new snapshot after debounce interval', async function () { + const initialTime = globals.clock.Date.now() + await tracker.processEdit(mockDocument, 'first edit') + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 1) + + // Another edit after debounce interval, should add another snapshot + const laterTime = initialTime + tracker.config.debounceIntervalMs + 1000 + dateNowStub.returns(laterTime) + await tracker.processEdit(mockDocument, 'second edit') + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 2) + + // Verify the content of the second snapshot + const snapshots = tracker.getFileSnapshots(filePath) + assert.strictEqual(snapshots[1].content, 'second edit') + }) + + it('should delete snapshot after maxAgeMs', async function () { + const customConfig: Partial = { + maxAgeMs: 10000, + } + tracker = new PredictionTracker(mockExtensionContext, customConfig) + const initialTime = globals.clock.Date.now() + await tracker.processEdit(mockDocument, previousContent) + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 1) + + // Advance time just under the maxAgeMs, snapshot should still exist + dateNowStub.returns(initialTime + tracker.config.maxAgeMs - 1000) + await clock.tickAsync(tracker.config.maxAgeMs - 1000) + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 1) + + // Advance time past the maxAgeMs, snapshot should be removed + dateNowStub.returns(initialTime + tracker.config.maxAgeMs + 2000) + await clock.tickAsync(3000) + assert.strictEqual(tracker.getFileSnapshots(filePath).length, 0) + }) + }) + + describe('enforceMemoryLimits', function () { + beforeEach(function () { + tracker = new PredictionTracker(mockExtensionContext) + }) + + it('should remove oldest snapshots when storage size exceeds limit', async function () { + // Very small storage limit + const customConfig: Partial = { + maxStorageSizeKb: 0.1, + } + tracker = new PredictionTracker(mockExtensionContext, customConfig) + + const file1 = testPath('path', 'to', 'file1.js') + const file2 = testPath('path', 'to', 'file2.js') + + const initialTime = globals.clock.Date.now() + + // First snapshot for file1 (oldest) + const mockDocument1 = createMockDocument('content 1', file1) + await tracker.processEdit(mockDocument1, 'content 1') + dateNowStub.returns(initialTime + 1000) + await clock.tickAsync(1000) + + // Second snapshot for file1 + await tracker.processEdit(mockDocument1, 'content 2') + dateNowStub.returns(initialTime + 2000) + await clock.tickAsync(1000) + + // First snapshot for file2 + const mockDocument2 = createMockDocument('content 3', file2) + await tracker.processEdit(mockDocument2, 'content 3') + + await (tracker as any).enforceMemoryLimits() + + // Oldest snapshot should be removed + const file1Snapshots = tracker.getFileSnapshots(file1) + assert.strictEqual(file1Snapshots.length, 1) + }) + }) + + describe('getFileSnapshots', function () { + beforeEach(function () { + tracker = new PredictionTracker(mockExtensionContext) + }) + + it('should return empty array for non-existent file', function () { + const result = tracker.getFileSnapshots(testPath('non-existent', 'file.js')) + assert.deepStrictEqual(result, []) + }) + + it('should return snapshots for existing file', async function () { + const file = testPath('path', 'to', 'file.js') + const content = 'file content' + const mockDocument = createMockDocument(content, file) + await tracker.processEdit(mockDocument, content) + + const result = tracker.getFileSnapshots(file) + assert.strictEqual(result.length, 1) + assert.strictEqual(result[0].filePath, file) + assert.strictEqual(result[0].content, content) + }) + }) + + describe('getSnapshotContent', function () { + let file: string + let snapshotContent: string + let snapshot: FileSnapshot + + beforeEach(async function () { + tracker = new PredictionTracker(mockExtensionContext) + file = testPath('path', 'to', 'file.js') + snapshotContent = 'snapshot content' + const mockDocument = createMockDocument(snapshotContent, file) + await tracker.processEdit(mockDocument, snapshotContent) + + snapshot = tracker.getFileSnapshots(file)[0] + }) + + it('should retrieve snapshot content from memory', async function () { + const content = await tracker.getSnapshotContent(snapshot) + assert.strictEqual(content, snapshotContent) + }) + }) + + describe('generatePredictionSupplementalContext', function () { + let mockEditor: vscode.TextEditor + let diffGenerateStub: sinon.SinonStub + + beforeEach(function () { + tracker = new PredictionTracker(mockExtensionContext) + + // Mock active editor, we only care about document + mockEditor = { + document: createMockDocument('current content', testPath('path', 'to', 'active.js')), + selection: new vscode.Selection(0, 0, 0, 0), + selections: [new vscode.Selection(0, 0, 0, 0)], + options: {}, + visibleRanges: [], + edit: () => Promise.resolve(true), + insertSnippet: () => Promise.resolve(true), + setDecorations: () => {}, + revealRange: () => {}, + show: () => {}, + hide: () => {}, + viewColumn: vscode.ViewColumn.One, + } as vscode.TextEditor + + sandbox.stub(vscode.window, 'activeTextEditor').value(mockEditor) + + // Mock diffGenerator.generateDiffContexts + diffGenerateStub = sandbox.stub(diffGenerator, 'generateDiffContexts').resolves([]) + }) + + it('should return empty array if no snapshots', async function () { + const result = await tracker.generatePredictionSupplementalContext() + assert.deepStrictEqual(result, []) + }) + + it('should generate and return supplemental contexts', async function () { + const filePath = testPath('path', 'to', 'active.js') + const initialTime = globals.clock.Date.now() + + const mockDoc = createMockDocument('old content 1', filePath) + await tracker.processEdit(mockDoc, 'old content 1') + dateNowStub.returns(initialTime + tracker.config.debounceIntervalMs + 1000) + await clock.tickAsync(tracker.config.debounceIntervalMs + 1000) + await tracker.processEdit(mockDoc, 'old content 2') + + const mockContexts = [ + { filePath, content: 'diff1', type: 'PreviousEditorState' }, + { filePath, content: 'diff2', type: 'PreviousEditorState' }, + ] + diffGenerateStub.resolves(mockContexts) + + const result = await tracker.generatePredictionSupplementalContext() + + // Should have called generateDiffContexts with the right params + assert.ok(diffGenerateStub.called) + assert.strictEqual(diffGenerateStub.args[0][0], filePath) + assert.strictEqual(diffGenerateStub.args[0][1], 'current content') + assert.strictEqual(diffGenerateStub.args[0][2].length, 2) + assert.strictEqual(diffGenerateStub.args[0][3], tracker.config.maxSupplementalContext) + + // Should return the contexts from generateDiffContexts + assert.deepStrictEqual(result, mockContexts) + + // Check that the snapshot content is correctly passed to the diffContextGenerator + const snapshotContents = diffGenerateStub.args[0][2] + assert.strictEqual(snapshotContents[0].content, 'old content 1') + assert.strictEqual(snapshotContents[1].content, 'old content 2') + }) + }) + + function testPath(...segments: string[]): string { + // Mock the path from vscode uri + return path.sep + path.join(...segments) + } + + describe('trimSupplementalContexts', function () { + it('should filter out contexts that exceed individual character limit', function () { + const smallContext = { + filePath: 'file.js', + content: 'small content', + type: 'PreviousEditorState', + } + + // Create a context that exceeds the characters limit + const largeContent = 'a'.repeat(charactersLimit + 100) + const largeContext = { + filePath: 'file.js', + content: largeContent, + type: 'PreviousEditorState', + } + + const contexts = [smallContext, largeContext] + const result = diffGenerator.trimSupplementalContexts(contexts, 10) + + assert.strictEqual(result.length, 1) + assert.deepStrictEqual(result[0], smallContext) + }) + + it('should limit the number of contexts to maxContexts', function () { + const contexts = [ + { filePath: 'file1.js', content: 'content 1', type: 'PreviousEditorState' }, + { filePath: 'file2.js', content: 'content 2', type: 'PreviousEditorState' }, + { filePath: 'file3.js', content: 'content 3', type: 'PreviousEditorState' }, + { filePath: 'file4.js', content: 'content 4', type: 'PreviousEditorState' }, + { filePath: 'file5.js', content: 'content 5', type: 'PreviousEditorState' }, + ] + + const maxContexts = 3 + const result = diffGenerator.trimSupplementalContexts(contexts, maxContexts) + + assert.strictEqual(result.length, maxContexts) + }) + + it('should enforce total character length limit across all contexts', function () { + // Create contexts where total size exceeds the limit + const contentSize = Math.floor(supplementalContextMaxTotalLength / 2.5) + const contexts = [ + { + filePath: 'file1.js', + content: 'a'.repeat(contentSize), + type: 'PreviousEditorState', + }, + { + filePath: 'file2.js', + content: 'b'.repeat(contentSize), + type: 'PreviousEditorState', + }, + { + filePath: 'file3.js', + content: 'c'.repeat(contentSize), + type: 'PreviousEditorState', + }, + ] + + const result = diffGenerator.trimSupplementalContexts(contexts, 10) + + // Only the first two contexts should be included since the third would exceed the total limit + assert.strictEqual(result.length, 2) + assert.deepStrictEqual(result, contexts.slice(0, 2)) + }) + }) +})