Skip to content

Optimistic Next Edit Suggestions #222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+166 −9
Open

Optimistic Next Edit Suggestions #222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 165 additions & 8 deletions src/extension/inlineEdits/node/nextEditProvider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import type * as vscode from 'vscode';
import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
import { IGitExtensionService } from '../../../platform/git/common/gitExtensionService';
import { DocumentId } from '../../../platform/inlineEdits/common/dataTypes/documentId';
import { RootedEdit } from '../../../platform/inlineEdits/common/dataTypes/edit';
import { RootedLineEdit } from '../../../platform/inlineEdits/common/dataTypes/rootedLineEdit';
Expand All @@ -20,9 +21,9 @@ import { IExperimentationService } from '../../../platform/telemetry/common/null
import { Result } from '../../../util/common/result';
import { createTracer, ITracer } from '../../../util/common/tracing';
import { assert } from '../../../util/vs/base/common/assert';
import { DeferredPromise, timeout, TimeoutTimer } from '../../../util/vs/base/common/async';
import { DeferredPromise, TimeoutTimer } from '../../../util/vs/base/common/async';
import { CachedFunction } from '../../../util/vs/base/common/cache';
import { CancellationToken } from '../../../util/vs/base/common/cancellation';
import { CancellationToken, CancellationTokenSource } from '../../../util/vs/base/common/cancellation';
import { BugIndicatingError } from '../../../util/vs/base/common/errors';
import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../util/vs/base/common/lifecycle';
import { LRUCache } from '../../../util/vs/base/common/map';
Expand All @@ -38,7 +39,7 @@ import { RejectionCollector } from '../common/rejectionCollector';
import { DebugRecorder } from './debugRecorder';
import { INesConfigs } from './nesConfigs';
import { CachedOrRebasedEdit, NextEditCache } from './nextEditCache';
import { LlmNESTelemetryBuilder } from './nextEditProviderTelemetry';
import { LlmNESTelemetryBuilder, NextEditProviderTelemetryBuilder } from './nextEditProviderTelemetry';
import { INextEditResult, NextEditResult } from './nextEditResult';

export interface INextEditProvider<T extends INextEditResult, TTelemetry, TData = void> extends IDisposable {
Expand Down Expand Up @@ -66,6 +67,10 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
private readonly _nextEditCache: NextEditCache;
private readonly _recentlyShownCache = new RecentlyShownCache();

// Track active prefetch requests to avoid duplicates
private readonly _activePrefetches = new Map<string, CancellationTokenSource>();
private readonly _prefetchChainDepth = 3; // Maximum depth for prediction chain

private _pendingStatelessNextEditRequest: StatelessNextEditRequest<CachedOrRebasedEdit> | null = null;

private _lastShownTime = 0;
Expand All @@ -92,6 +97,7 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
@ISnippyService private readonly _snippyService: ISnippyService,
@ILogService private readonly _logService: ILogService,
@IExperimentationService private readonly _expService: IExperimentationService,
@IGitExtensionService private readonly _gitExtensionService: IGitExtensionService,
) {
super();

Expand All @@ -117,6 +123,7 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
}

public async getNextEdit(docId: DocumentId, context: vscode.InlineCompletionContext, logContext: InlineEditRequestLogContext, cancellationToken: CancellationToken, telemetryBuilder: LlmNESTelemetryBuilder): Promise<NextEditResult> {
const startTime = Date.now();
const tracer = this._tracer.sub('getNextEdit');

this._lastTriggerTime = Date.now();
Expand All @@ -141,6 +148,9 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
telemetryBuilder.setNESConfigs({ ...nesConfigs });
logContext.addCodeblockToLog(JSON.stringify(nesConfigs, null, '\t'));

// The existing cache lookup will now find our prefetched results!


const recentlyShownCachedEdit = this._recentlyShownCache.get(docId, documentAtInvocationTime);
const cachedEdit = this._nextEditCache.lookupNextEdit(docId, documentAtInvocationTime, doc.selection.get(), nesConfigs);

Expand All @@ -158,8 +168,6 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
let req: NextEditFetchRequest;
let targetDocumentId = docId;

const cacheDelay = this._configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsCacheDelay, this._expService);

if (recentlyShownCachedEdit) {
tracer.trace('using recently shown cached edit');
edit = recentlyShownCachedEdit[0];
Expand All @@ -173,10 +181,12 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
// back-date the recording bookmark of the cached edit to the bookmark of the original request.
logContext.recordingBookmark = req.log.recordingBookmark;

await timeout(cacheDelay);
// No artificial delay for cache hits

} else if (cachedEdit) {
const timeElapsed = Date.now() - startTime;
tracer.trace('using cached edit');
this._logService.logger.info(`[NextEditProvider] CACHE HIT (prefetched) - ${timeElapsed}ms, subsequentN=${cachedEdit.subsequentN}`);
edit = cachedEdit.rebasedEdit || cachedEdit.edit;
req = cachedEdit.source;
logContext.setIsCachedResult(cachedEdit.source.log);
Expand All @@ -187,7 +197,7 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
// back-date the recording bookmark of the cached edit to the bookmark of the original request.
logContext.recordingBookmark = req.log.recordingBookmark;

await timeout(cacheDelay);
// No artificial delay for cache hits

} else {
tracer.trace('fetching next edit');
Expand Down Expand Up @@ -305,6 +315,7 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
private async fetchNextEdit(req: NextEditFetchRequest, doc: IObservableDocument, nesConfigs: INesConfigs, telemetryBuilder: LlmNESTelemetryBuilder, cancellationToken: CancellationToken): Promise<Result<CachedOrRebasedEdit, NoNextEditReason>> {
const curDocId = doc.id;
const tracer = this._tracer.sub('fetchNextEdit');

const historyContext = this._historyContextProvider.getHistoryContext(curDocId);

if (!historyContext) {
Expand Down Expand Up @@ -630,11 +641,144 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne

public handleShown(suggestion: NextEditResult) {
this._lastShownTime = Date.now();

// Trigger optimistic prefetching as soon as the edit is shown
// This assumes the user will accept it and pre-fetches the next edits
if (suggestion.result?.edit) {
const docId = suggestion.result.targetDocumentId ||
DocumentId.create(suggestion.source.log.filePath);
const doc = this._workspace.getDocument(docId);
if (!doc) {
return;
}

// Calculate what the document state would be AFTER accepting this edit
const baseDocState = suggestion.result.documentBeforeEdits;
const stringEdit = StringEdit.create([suggestion.result.edit]);
const futureDocState = stringEdit.applyOnText(baseDocState);

// Start the prefetch chain
this._startPrefetchChain(docId, futureDocState, 0);
}
}

private async _startPrefetchChain(docId: DocumentId, documentState: StringText, depth: number): Promise<void> {
if (depth >= this._prefetchChainDepth) {
return;
}

// Create a unique key for this prefetch based on document state
const prefetchKey = `${docId}:${documentState.value.length}:${depth}`;

// Cancel any existing prefetch for this state
const existingPrefetch = this._activePrefetches.get(prefetchKey);
if (existingPrefetch) {
existingPrefetch.cancel();
this._activePrefetches.delete(prefetchKey);
}

const cancellationSource = new CancellationTokenSource();
this._activePrefetches.set(prefetchKey, cancellationSource);

try {

// Create a synthetic document that represents the future state
const syntheticDoc = this._workspace.getDocument(docId);
if (!syntheticDoc || cancellationSource.token.isCancellationRequested) {
return;
}

// Perform the prefetch - this will populate the cache
const result = await this._performPrefetch(docId, documentState, cancellationSource.token);

if (result && result.result?.edit && !cancellationSource.token.isCancellationRequested) {
// Apply the edit to get the next document state
const nextEdit = StringEdit.create([result.result.edit]);
const nextDocumentState = nextEdit.applyOnText(documentState);

// Check if the document actually changed
if (nextDocumentState.value === documentState.value) {
// Edit made no changes, no point in continuing the chain
return;
}

// Continue the chain only if the document changed
await this._startPrefetchChain(docId, nextDocumentState, depth + 1);
}
} catch (err) {
this._logService.logger.trace(`[Prefetch] Error at depth ${depth}: ${err}`);
} finally {
this._activePrefetches.delete(prefetchKey);
}
}

private async _performPrefetch(docId: DocumentId, futureDocumentState: StringText, cancellationToken: CancellationToken): Promise<NextEditResult | undefined> {
const doc = this._workspace.getDocument(docId);
if (!doc) {
return undefined;
}

// Create synthetic context and telemetry for the prefetch
const syntheticContext: vscode.InlineCompletionContext = {
triggerKind: 0 as vscode.InlineCompletionTriggerKind,
selectedCompletionInfo: undefined,
requestUuid: generateUuid()
};

const logContext = new InlineEditRequestLogContext(docId.toUri().toString(), 0, syntheticContext);
const req = new NextEditFetchRequest(logContext, Date.now());

// Create a proper telemetry builder for prefetch
const telemetryBuilder = new NextEditProviderTelemetryBuilder(this._gitExtensionService, this.ID, doc);
telemetryBuilder.nesBuilder.setHeaderRequestId(req.headerRequestId);
telemetryBuilder.nesBuilder.setIsFromCache(); // Mark as prefetched

const nesConfigs: INesConfigs = {
isAsyncCompletions: this._configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsAsyncCompletions, this._expService),
isRevisedCacheStrategy: this._configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsRevisedCacheStrategy, this._expService),
isCacheTracksRejections: this._configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsCacheTracksRejections, this._expService),
isRecentlyShownCacheEnabled: this._configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsRecentlyShownCacheEnabled, this._expService),
};

// Create a synthetic document wrapper that returns our future state
const syntheticDocWrapper = {
...doc,
value: {
get: () => futureDocumentState
}
};

try {
// This will fetch the edit and populate the cache via pushEdit
const result = await this.fetchNextEdit(req, syntheticDocWrapper as IObservableDocument, nesConfigs, telemetryBuilder.nesBuilder, cancellationToken);

if (result.isOk() && result.val.edit) {
const nextEditResult = new NextEditResult(
logContext.requestId,
req,
{
edit: result.val.edit,
documentBeforeEdits: futureDocumentState,
showRangePreference: this._statelessNextEditProvider.showNextEditPreference ?? ShowNextEditPreference.AroundEdit,
targetDocumentId: docId
}
);
return nextEditResult;
}
} catch (err) {
this._logService.logger.trace(`[Prefetch] Error during fetch: ${err}`);
} finally {
telemetryBuilder.dispose();
}

return undefined;
}

public handleAcceptance(docId: DocumentId, suggestion: NextEditResult) {
this.runSnippy(docId, suggestion);
this._statelessNextEditProvider.handleAcceptance?.();

// Optimistic fetching already triggered in handleShown
}

public handleRejection(docId: DocumentId, suggestion: NextEditResult) {
Expand All @@ -656,7 +800,19 @@ export class NextEditProvider extends Disposable implements INextEditProvider<Ne
this._statelessNextEditProvider.handleRejection?.();
}

public handleIgnored(docId: DocumentId, suggestion: NextEditResult, supersededBy: INextEditResult | undefined): void { }
public handleIgnored(_docId: DocumentId, _suggestion: NextEditResult, _supersededBy: INextEditResult | undefined): void {
// No-op for now
}

public override dispose(): void {
// Cancel all active prefetches
for (const [, cancellationSource] of this._activePrefetches) {
cancellationSource.cancel();
}
this._activePrefetches.clear();

super.dispose();
}

private async runSnippy(docId: DocumentId, suggestion: NextEditResult) {
if (suggestion.result === undefined) {
Expand Down Expand Up @@ -718,3 +874,4 @@ class RecentlyShownCache {
return docId.uri + ';' + documentContent.value;
}
}

2 changes: 1 addition & 1 deletion src/extension/inlineEdits/test/node/nextEditProviderCaching.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ describe('NextEditProvider Caching', () => {
}
};

const nextEditProvider: NextEditProvider = new NextEditProvider(obsWorkspace, statelessNextEditProvider, new NesHistoryContextProvider(obsWorkspace, obsGit), new NesXtabHistoryTracker(obsWorkspace), undefined, configService, snippyService, logService, expService);
const nextEditProvider: NextEditProvider = new NextEditProvider(obsWorkspace, statelessNextEditProvider, new NesHistoryContextProvider(obsWorkspace, obsGit), new NesXtabHistoryTracker(obsWorkspace), undefined, configService, snippyService, logService, expService, gitExtensionService);

const doc = obsWorkspace.addDocument({
id: DocumentId.create(URI.file('/test/test.ts').toString()),
Expand Down