diff --git a/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts b/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts index a982db829c6..e96fa3166e9 100644 --- a/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts +++ b/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts @@ -46,6 +46,7 @@ import { INesConfigs } from './nesConfigs'; import { CachedOrRebasedEdit, NextEditCache } from './nextEditCache'; import { LlmNESTelemetryBuilder, ReusedRequestKind } from './nextEditProviderTelemetry'; import { INextEditResult, NextEditResult } from './nextEditResult'; +import { SpeculativeCancelReason, SpeculativeRequestManager } from './speculativeRequestManager'; /** * Computes a reduced window range that encompasses both the original window (shrunk by one line @@ -167,28 +168,7 @@ export class NextEditProvider extends Disposable implements INextEditProvider | null = null; - /** - * Tracks a speculative request for the post-edit document state. - * When a suggestion is shown, we speculatively fetch the next edit as if the user had already accepted. - * This allows reusing the in-flight request when the user actually accepts the suggestion. - */ - private _speculativePendingRequest: { - request: StatelessNextEditRequest; - docId: DocumentId; - postEditContent: string; - } | null = null; - - /** - * A speculative request that is deferred until the originating stream completes. - * When a suggestion is shown while its stream is still running, we schedule the - * speculative request here instead of firing immediately. If more edits arrive - * from the stream, the schedule is cleared (the shown edit wasn't the last one). - * When the stream ends, if the schedule is still present, the speculative fires. - */ - private _scheduledSpeculativeRequest: { - suggestion: NextEditResult; - headerRequestId: string; - } | null = null; + private readonly _specManager: SpeculativeRequestManager; private _lastShownTime = 0; /** The requestId of the last shown suggestion. We store only the requestId (not the object) to avoid preventing garbage collection. */ @@ -230,35 +210,48 @@ export class NextEditProvider extends Disposable implements INextEditProvider { store.add(runOnChange(doc.value, (value) => { this._cancelPendingRequestDueToDocChange(doc.id, value); + // FIXME: don't invoke before fixing false positive cancellations + // this._specManager.onActiveDocumentChanged(doc.id, value.value); })); + // When the per-doc store is disposed, the document was removed from + // openDocuments. Cancel any speculative targeting it — its cached result + // would never be hit again. + store.add(toDisposable(() => this._specManager.onDocumentClosed(doc.id))); }).recomputeInitiallyAndOnChange(this._store); } - private _cancelSpeculativeRequest(): void { - this._scheduledSpeculativeRequest = null; - if (this._speculativePendingRequest) { - this._speculativePendingRequest.request.cancellationTokenSource.cancel(); - this._speculativePendingRequest = null; - } - } - + /** + * Cancels the in-flight stateless next-edit request when the document it + * was issued for has diverged from the request's expected post-edit state. + * + * Invoked from the per-document `runOnChange` autorun in the constructor + * whenever an open document's value changes. The pending request was built + * against a specific snapshot (`documentAfterEdits`); if the user has since + * typed something that makes the current value differ from that snapshot, + * the result would no longer be applicable and is cancelled eagerly. + * + * Skipped when: + * - the `InlineEditsAsyncCompletions` experiment is enabled (that path + * tolerates divergence and rebases later), or + * - there is no pending request, or + * - the changed document is not the one the pending request targets. + * + * Note: this only handles the regular pending stateless request. Speculative + * requests have their own divergence handling via + * `SpeculativeRequestManager.onActiveDocumentChanged` (trajectory check). + */ private _cancelPendingRequestDueToDocChange(docId: DocumentId, docValue: StringText) { - // Note: we intentionally do NOT cancel the speculative request here. - // The speculative request's postEditContent represents a *future* document state - // (after the user would accept the suggestion), so it will almost never match the - // current document value while the user is still typing. Cancelling here would - // wastefully kill and recreate the speculative request on every keystroke. - // Instead, speculative requests are cancelled by the appropriate lifecycle handlers: - // handleRejection, handleIgnored, _triggerSpeculativeRequest, and _executeNewNextEditRequest. - const isAsyncCompletions = this._configService.getExperimentBasedConfig(ConfigKey.TeamInternal.InlineEditsAsyncCompletions, this._expService); + if (isAsyncCompletions || this._pendingStatelessNextEditRequest === null) { return; } + const activeDoc = this._pendingStatelessNextEditRequest.getActiveDocument(); if (activeDoc.id === docId && activeDoc.documentAfterEdits.value !== docValue.value) { this._pendingStatelessNextEditRequest.cancellationTokenSource.cancel(); @@ -547,11 +540,12 @@ export class NextEditProvider extends Disposable implements INextEditProvider 1000 && suggestion.result.edit) { @@ -1471,9 +1474,15 @@ export class NextEditProvider extends Disposable implements INextEditProvider; + readonly docId: DocumentId; + readonly postEditContent: string; + /** preEditDocument[0..editStart] — the doc text before the edit window. */ + readonly trajectoryPrefix: string; + /** preEditDocument[editEnd..] — the doc text after the edit window. */ + readonly trajectorySuffix: string; + /** The replacement text the user would type to reach `postEditContent`. */ + readonly trajectoryNewText: string; +} + +export interface ScheduledSpeculativeRequest { + readonly suggestion: NextEditResult; + readonly headerRequestId: string; +} + +/** + * Owns the lifecycle of NES speculative requests: + * + * - the in-flight `pending` speculative (the bet on a specific post-accept document state) + * - the `scheduled` speculative deferred until its originating stream completes + * + * Centralizes cancellation with typed reasons so every triggered cancellation + * (reject, supersede, doc-close, trajectory divergence, dispose, ...) goes through + * one path and is logged on the request's log context. + */ +export class SpeculativeRequestManager extends Disposable { + + private _pending: SpeculativePendingRequest | null = null; + private _scheduled: ScheduledSpeculativeRequest | null = null; + + constructor(private readonly _logger: ILogger) { + super(); + } + + get pending(): SpeculativePendingRequest | null { + return this._pending; + } + + /** Replaces the current pending speculative; cancels the prior one as `Replaced`. */ + setPending(req: SpeculativePendingRequest): void { + if (this._pending && this._pending.request !== req.request) { + this._cancelPending(SpeculativeCancelReason.Replaced); + } + this._pending = req; + } + + /** Detaches the pending speculative without cancelling — caller is consuming it. */ + consumePending(): void { + this._pending = null; + } + + schedule(s: ScheduledSpeculativeRequest): void { + this._scheduled = s; + } + + clearScheduled(): void { + this._scheduled = null; + } + + /** + * Removes and returns the scheduled entry iff its `headerRequestId` matches. + * Used by the streaming path so that each stream only ever consumes its own + * schedule, never another stream's. + */ + consumeScheduled(headerRequestId: string): ScheduledSpeculativeRequest | null { + if (this._scheduled?.headerRequestId !== headerRequestId) { + return null; + } + const s = this._scheduled; + this._scheduled = null; + return s; + } + + cancelAll(reason: SpeculativeCancelReason): void { + this._scheduled = null; + this._cancelPending(reason); + } + + /** Cancels the pending speculative iff `(docId, postEditContent)` doesn't match. */ + cancelIfMismatch(docId: DocumentId, postEditContent: string, reason: SpeculativeCancelReason): void { + if (this._pending && (this._pending.docId !== docId || this._pending.postEditContent !== postEditContent)) { + this._cancelPending(reason); + } + } + + /** Cancels the pending and clears any scheduled targeting this document. */ + onDocumentClosed(docId: DocumentId): void { + if (this._scheduled?.suggestion.result?.targetDocumentId === docId) { + this._scheduled = null; + } + if (this._pending?.docId === docId) { + this._cancelPending(SpeculativeCancelReason.DocumentClosed); + } + } + + /** + * Trajectory check. The pending speculative is alive iff the current document + * value is a *type-through prefix* toward the speculative's `postEditContent`: + * + * cur === trajectoryPrefix + middle + trajectorySuffix + * where middle is some prefix of trajectoryNewText + * + * If not, the user's edits cannot reach `postEditContent` via continued typing + * and the speculative will never be consumed — cancel now. + */ + onActiveDocumentChanged(docId: DocumentId, currentDocValue: string): void { + const p = this._pending; + if (!p || p.docId !== docId) { + return; + } + // Cheap structural failure: doc shorter than the unedited frame. + if (currentDocValue.length < p.trajectoryPrefix.length + p.trajectorySuffix.length) { + this._cancelPending(SpeculativeCancelReason.DivergedFromTrajectoryForm); + return; + } + if (!currentDocValue.startsWith(p.trajectoryPrefix)) { + this._cancelPending(SpeculativeCancelReason.DivergedFromTrajectoryPrefix); + return; + } + if (!currentDocValue.endsWith(p.trajectorySuffix)) { + this._cancelPending(SpeculativeCancelReason.DivergedFromTrajectorySuffix); + return; + } + const middle = currentDocValue.slice(p.trajectoryPrefix.length, currentDocValue.length - p.trajectorySuffix.length); + if (!p.trajectoryNewText.startsWith(middle)) { + this._cancelPending(SpeculativeCancelReason.DivergedFromTrajectoryMiddle); + } + } + + private _cancelPending(reason: SpeculativeCancelReason): void { + const p = this._pending; + if (!p) { + return; + } + this._pending = null; + const headerRequestId = p.request.headerRequestId; + this._logger.trace(`cancelling speculative request: ${reason} (headerRequestId=${headerRequestId})`); + p.request.logContext.addLog(`speculative request cancelled: ${reason}`); + const cts = p.request.cancellationTokenSource; + cts.cancel(); + // Dispose to release the cancel-event listeners that the in-flight + // provider call hooked onto the token. Safe even though the runner may + // observe cancellation asynchronously — `cancel()` already fired the event. + cts.dispose(); + } + + override dispose(): void { + this.cancelAll(SpeculativeCancelReason.Disposed); + super.dispose(); + } +} diff --git a/extensions/copilot/src/extension/inlineEdits/test/node/nextEditProviderSpeculative.spec.ts b/extensions/copilot/src/extension/inlineEdits/test/node/nextEditProviderSpeculative.spec.ts index 9e48ebc88ab..73a7bd9104a 100644 --- a/extensions/copilot/src/extension/inlineEdits/test/node/nextEditProviderSpeculative.spec.ts +++ b/extensions/copilot/src/extension/inlineEdits/test/node/nextEditProviderSpeculative.spec.ts @@ -18,8 +18,8 @@ import { EditStreamingWithTelemetry, IStatelessNextEditProvider, NoNextEditReaso import { NesHistoryContextProvider } from '../../../../platform/inlineEdits/common/workspaceEditTracker/nesHistoryContextProvider'; import { NesXtabHistoryTracker } from '../../../../platform/inlineEdits/common/workspaceEditTracker/nesXtabHistoryTracker'; import { ILogger, ILogService, LogServiceImpl } from '../../../../platform/log/common/logService'; -import { NullRequestLogger } from '../../../../platform/requestLogger/node/nullRequestLogger'; import { IRequestLogger } from '../../../../platform/requestLogger/common/requestLogger'; +import { NullRequestLogger } from '../../../../platform/requestLogger/node/nullRequestLogger'; import { ISnippyService, NullSnippyService } from '../../../../platform/snippy/common/snippyService'; import { IExperimentationService, NullExperimentationService } from '../../../../platform/telemetry/common/nullExperimentationService'; import { mockNotebookService } from '../../../../platform/test/common/testNotebookService'; @@ -417,7 +417,7 @@ describe('NextEditProvider speculative requests', () => { await statelessProvider.calls[1].completed.p; }); - it('does not cancel speculative request when active document diverges from expected post-edit state', async () => { + it('cancels speculative request when active document edit moves off the type-through trajectory', async () => { await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); const statelessProvider = new TestStatelessNextEditProvider(); @@ -436,30 +436,28 @@ describe('NextEditProvider speculative requests', () => { nextEditProvider.handleShown(suggestion); await statelessProvider.waitForCall(2); - // Editing the active document should NOT cancel the speculative request. - // The speculative request targets a future post-edit state, not the current - // document value, so keystroke-level changes should not invalidate it. + // Inserting at the start of the document breaks the trajectory's prefix + // (the doc no longer starts with `pre[0..editStart]`). The speculative + // can no longer be reached via type-through-then-accept — cancel. doc.applyEdit(StringEdit.insert(0, '/* diverged */\n')); - await flushMicrotasks(); + await statelessProvider.calls[1].cancellationRequested.p; - expect(statelessProvider.calls[1].wasCancelled).toBe(false); - - // Clean up: reject so the speculative request gets cancelled properly - nextEditProvider.handleRejection(doc.id, suggestion); - await statelessProvider.calls[1].completed.p; + expect(statelessProvider.calls[1].wasCancelled).toBe(true); }); - it('keeps speculative request alive when user types in the active document', async () => { + it('keeps speculative alive while user types characters of the suggestion (type-through)', async () => { await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); const statelessProvider = new TestStatelessNextEditProvider(); - statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'const value = 2;') }); + // Suggestion inserts `'barbaz'` between `'foo'` and `'();'`. + // Resulting precise edit: replace [3, 3) with 'barbaz' (a pure insertion). + statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'foobarbaz();') }); statelessProvider.enqueueBehavior({ kind: 'waitForCancellation' }); const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); const doc = workspace.addDocument({ id: DocumentId.create(URI.file('/test/spec-typing.ts').toString()), - initialValue: 'const value = 1;\nconsole.log(value);', + initialValue: 'foo();\nconsole.log();', }); doc.setSelection([new OffsetRange(0, 0)], undefined); @@ -468,23 +466,28 @@ describe('NextEditProvider speculative requests', () => { nextEditProvider.handleShown(suggestion); await statelessProvider.waitForCall(2); - // Simulate multiple keystrokes in the active document while the speculative - // request is in flight — none of them should cancel it. - doc.applyEdit(StringEdit.insert(0, 'a')); + // User types characters of the suggestion at the edit position — each + // keystroke keeps the document on a type-through trajectory toward + // `postEditContent`, so the speculative must NOT be cancelled. + doc.applyEdit(StringEdit.insert(3, 'b')); await flushMicrotasks(); expect(statelessProvider.calls[1].wasCancelled).toBe(false); - doc.applyEdit(StringEdit.insert(1, 'b')); + doc.applyEdit(StringEdit.insert(4, 'a')); await flushMicrotasks(); expect(statelessProvider.calls[1].wasCancelled).toBe(false); - doc.applyEdit(StringEdit.insert(2, 'c')); + doc.applyEdit(StringEdit.insert(5, 'r')); await flushMicrotasks(); expect(statelessProvider.calls[1].wasCancelled).toBe(false); - // Clean up via rejection - nextEditProvider.handleRejection(doc.id, suggestion); - await statelessProvider.calls[1].completed.p; + // Now the user types a character that doesn't match the suggestion's + // next character (`'b'` would be expected; they typed `'X'`). The + // trajectory is broken — cancel. + doc.applyEdit(StringEdit.insert(6, 'X')); + await statelessProvider.calls[1].cancellationRequested.p; + + expect(statelessProvider.calls[1].wasCancelled).toBe(true); }); it('cancels mismatched speculative request when starting a request for another document', async () => { @@ -1370,4 +1373,83 @@ describe('NextEditProvider speculative requests', () => { } }); }); + + describe('lifecycle cancellation', () => { + it('cancels in-flight speculative when clearCache() is called', async () => { + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); + + const statelessProvider = new TestStatelessNextEditProvider(); + statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'const value = 2;') }); + statelessProvider.enqueueBehavior({ kind: 'waitForCancellation' }); + const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); + + const doc = workspace.addDocument({ + id: DocumentId.create(URI.file('/test/spec-clear-cache.ts').toString()), + initialValue: 'const value = 1;\nconsole.log(value);', + }); + doc.setSelection([new OffsetRange(0, 0)], undefined); + + const suggestion = await getNextEdit(nextEditProvider, doc.id); + assert(suggestion.result?.edit); + nextEditProvider.handleShown(suggestion); + await statelessProvider.waitForCall(2); + + nextEditProvider.clearCache(); + await statelessProvider.calls[1].cancellationRequested.p; + + expect(statelessProvider.calls[1].wasCancelled).toBe(true); + }); + + it('cancels in-flight speculative when its target document is closed', async () => { + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); + + const statelessProvider = new TestStatelessNextEditProvider(); + statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'const value = 2;') }); + statelessProvider.enqueueBehavior({ kind: 'waitForCancellation' }); + const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); + + const doc = workspace.addDocument({ + id: DocumentId.create(URI.file('/test/spec-doc-close.ts').toString()), + initialValue: 'const value = 1;\nconsole.log(value);', + }); + doc.setSelection([new OffsetRange(0, 0)], undefined); + + const suggestion = await getNextEdit(nextEditProvider, doc.id); + assert(suggestion.result?.edit); + nextEditProvider.handleShown(suggestion); + await statelessProvider.waitForCall(2); + + // Closing the document removes it from openDocuments — the speculative's + // cached result would never be hit again, so cancel it. + doc.dispose(); + await statelessProvider.calls[1].cancellationRequested.p; + + expect(statelessProvider.calls[1].wasCancelled).toBe(true); + }); + + it('cancels in-flight speculative when the provider is disposed', async () => { + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); + + const statelessProvider = new TestStatelessNextEditProvider(); + statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'const value = 2;') }); + statelessProvider.enqueueBehavior({ kind: 'waitForCancellation' }); + const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); + + const doc = workspace.addDocument({ + id: DocumentId.create(URI.file('/test/spec-provider-dispose.ts').toString()), + initialValue: 'const value = 1;\nconsole.log(value);', + }); + doc.setSelection([new OffsetRange(0, 0)], undefined); + + const suggestion = await getNextEdit(nextEditProvider, doc.id); + assert(suggestion.result?.edit); + nextEditProvider.handleShown(suggestion); + await statelessProvider.waitForCall(2); + + nextEditProvider.dispose(); + await statelessProvider.calls[1].cancellationRequested.p; + + expect(statelessProvider.calls[1].wasCancelled).toBe(true); + }); + }); });