From 401fe38d71270e7559603ec18a147e94f3aef6da Mon Sep 17 00:00:00 2001 From: ulugbekna Date: Fri, 17 Apr 2026 23:38:28 +0200 Subject: [PATCH] other pr --- .../inlineEdits/node/nextEditProvider.ts | 152 ++++---- .../node/speculativeRequestManager.ts | 361 ++++++++++++++++++ .../node/nextEditProviderSpeculative.spec.ts | 216 +++++++++-- 3 files changed, 634 insertions(+), 95 deletions(-) create mode 100644 extensions/copilot/src/extension/inlineEdits/node/speculativeRequestManager.ts diff --git a/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts b/extensions/copilot/src/extension/inlineEdits/node/nextEditProvider.ts index a982db829c6..3d3d17bf63a 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,47 @@ export class NextEditProvider extends Disposable implements INextEditProvider { store.add(runOnChange(doc.value, (value) => { this._cancelPendingRequestDueToDocChange(doc.id, value); + 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 +539,12 @@ export class NextEditProvider extends Disposable implements INextEditProvider 1000 && suggestion.result.edit) { @@ -1471,9 +1478,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; + /** + * Length of the text replaced by the suggestion (preEdit[editStart..editEnd]). + * The type-through trajectory check only models pure insertions; for + * substitutions/removals (`trajectoryOldTextLen > 0`) the check is skipped + * and the speculative is kept alive until another lifecycle trigger fires. + */ + readonly trajectoryOldTextLen: number; +} + +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`. + * + * The slot stays populated until the request is explicitly cancelled or + * superseded — multiple `provideNextEdit` invocations targeting the same + * `(docId, postEditContent)` can all join the same in-flight stream, + * mirroring `_pendingStatelessNextEditRequest`'s dedupe behavior. + * Once the speculative settles, `_nextEditCache` covers further reuse; + * the now-stale slot is cleared by the next `setPending` / `cancelAll` + * / `cancelIfMismatch` / trajectory-divergence trigger. + */ + setPending(req: SpeculativePendingRequest): void { + if (this._pending && this._pending.request !== req.request) { + this._cancelPending(SpeculativeCancelReason.Replaced); + } + this._pending = req; + // Record the trajectory data on the request's own log context so any + // subsequent cancellation diagnostic can be cross-referenced against + // the state captured at speculative creation time. + req.request.logContext.addLog('speculative request installed (trajectory captured)'); + req.request.logContext.addCodeblockToLog(JSON.stringify({ + docId: req.docId.toString(), + postEditContentLen: req.postEditContent.length, + trajectoryPrefixLen: req.trajectoryPrefix.length, + trajectoryPrefixHead: truncate(req.trajectoryPrefix.slice(0, 80), 80), + trajectoryPrefixTail: truncate(req.trajectoryPrefix.slice(Math.max(0, req.trajectoryPrefix.length - 80)), 80), + trajectorySuffixLen: req.trajectorySuffix.length, + trajectorySuffixHead: truncate(req.trajectorySuffix.slice(0, 80), 80), + trajectorySuffixTail: truncate(req.trajectorySuffix.slice(Math.max(0, req.trajectorySuffix.length - 80)), 80), + trajectoryNewText: truncate(req.trajectoryNewText, 200), + trajectoryNewTextLen: req.trajectoryNewText.length, + trajectoryOldTextLen: req.trajectoryOldTextLen, + trajectoryTracked: req.trajectoryOldTextLen === 0, + }, null, 2), 'json'); + } + + 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) { return; } + if (this._pending.docId === docId && this._pending.postEditContent === postEditContent) { return; } + this._cancelPending(reason, { + pendingDocId: this._pending.docId.toString(), + incomingDocId: docId.toString(), + pendingPostEditContentLen: this._pending.postEditContent.length, + incomingPostEditContentLen: postEditContent.length, + docIdMatches: this._pending.docId === docId, + postEditContentMatches: this._pending.postEditContent === postEditContent, + ...(this._pending.postEditContent !== postEditContent + ? { mismatch: describeStringMismatch(this._pending.postEditContent, postEditContent) } + : {}), + }); + } + + /** 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; + } + const check = checkTrajectory(p, currentDocValue); + if (check.ok) { + return; + } + this._cancelPending(check.reason, check.details); + } + + private _cancelPending(reason: SpeculativeCancelReason, diagnostic?: Record): 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}\``); + // Always include the request's own trajectory metadata so the log entry + // is self-contained — the reader doesn't need to cross-reference the + // originating `setPending` call. + const payload: Record = { + reason, + docId: p.docId.toString(), + postEditContentLen: p.postEditContent.length, + trajectoryPrefixLen: p.trajectoryPrefix.length, + trajectorySuffixLen: p.trajectorySuffix.length, + trajectoryNewText: truncate(p.trajectoryNewText, 200), + trajectoryNewTextLen: p.trajectoryNewText.length, + ...(diagnostic ?? {}), + }; + p.request.logContext.addCodeblockToLog(JSON.stringify(payload, null, 2), 'json'); + 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(); + } +} + +type TrajectoryCheckResult = + | { ok: true } + | { + ok: false; + reason: SpeculativeCancelReason; + details: Record; + }; + +/** + * Checks whether `currentDocValue` is on the type-through trajectory toward + * `p.postEditContent`. Returns a structured result so the cancellation site + * can record exactly what mismatched. + */ +function checkTrajectory(p: SpeculativePendingRequest, currentDocValue: string): TrajectoryCheckResult { + // The type-through trajectory model only fits pure insertions: starting from + // preEdit (`prefix + "" + suffix`), the user types characters of `newText` + // into the gap. For substitutions/removals (oldText non-empty), the initial + // state already has `oldText` in the gap, which is generally not a prefix of + // `newText`, so the check would over-cancel on the first keystroke. Skip it. + if (p.trajectoryOldTextLen > 0) { + return { ok: true }; + } + + const minLen = p.trajectoryPrefix.length + p.trajectorySuffix.length; + + // Cheap structural failure: doc shorter than the unedited frame. + if (currentDocValue.length < minLen) { + return { + ok: false, + reason: SpeculativeCancelReason.DivergedFromTrajectoryPrefix, + details: { + explanation: 'currentDocValue is shorter than trajectoryPrefix + trajectorySuffix', + currentDocLen: currentDocValue.length, + minRequiredLen: minLen, + docHead: truncate(currentDocValue, 200), + }, + }; + } + + if (!currentDocValue.startsWith(p.trajectoryPrefix)) { + return { + ok: false, + reason: SpeculativeCancelReason.DivergedFromTrajectoryMiddle, + details: { + explanation: 'currentDocValue does not start with trajectoryPrefix', + currentDocLen: currentDocValue.length, + ...describeStringMismatch(p.trajectoryPrefix, currentDocValue.slice(0, p.trajectoryPrefix.length)), + }, + }; + } + + if (!currentDocValue.endsWith(p.trajectorySuffix)) { + return { + ok: false, + reason: SpeculativeCancelReason.DivergedFromTrajectoryMiddle, + details: { + explanation: 'currentDocValue does not end with trajectorySuffix', + currentDocLen: currentDocValue.length, + ...describeStringMismatch( + p.trajectorySuffix, + currentDocValue.slice(currentDocValue.length - p.trajectorySuffix.length), + ), + }, + }; + } + + const middle = currentDocValue.slice(p.trajectoryPrefix.length, currentDocValue.length - p.trajectorySuffix.length); + if (!p.trajectoryNewText.startsWith(middle)) { + return { + ok: false, + reason: SpeculativeCancelReason.DivergedFromTrajectorySuffix, + details: { + explanation: 'middle is not a prefix of trajectoryNewText', + middle: truncate(middle, 200), + middleLen: middle.length, + ...describeStringMismatch(p.trajectoryNewText.slice(0, middle.length), middle), + }, + }; + } + + return { ok: true }; +} + +/** + * Compares two strings character-by-character and returns the first index at + * which they differ, plus a small window of context around it. + */ +function describeStringMismatch(expected: string, actual: string): Record { + const minLen = Math.min(expected.length, actual.length); + let i = 0; + while (i < minLen && expected.charCodeAt(i) === actual.charCodeAt(i)) { + i++; + } + const ctxStart = Math.max(0, i - 20); + const ctxEnd = i + 20; + return { + firstMismatchAt: i, + expectedLen: expected.length, + actualLen: actual.length, + expectedAround: visualize(expected.slice(ctxStart, ctxEnd)), + actualAround: visualize(actual.slice(ctxStart, ctxEnd)), + expectedCharAtMismatch: i < expected.length ? toCharRepr(expected.charCodeAt(i)) : '', + actualCharAtMismatch: i < actual.length ? toCharRepr(actual.charCodeAt(i)) : '', + }; +} + +function toCharRepr(code: number): string { + return `U+${code.toString(16).toUpperCase().padStart(4, '0')} (${JSON.stringify(String.fromCharCode(code))})`; +} + +function visualize(s: string): string { + return s.replace(/\r/g, '␍').replace(/\n/g, '␊').replace(/\t/g, '␉'); +} + +function truncate(s: string, max: number): string { + if (s.length <= max) { + return visualize(s); + } + return visualize(s.slice(0, max)) + `…(+${s.length - max} chars)`; +} 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..a7f745a45b5 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'; @@ -333,6 +333,46 @@ describe('NextEditProvider speculative requests', () => { expect(secondSuggestion.result.edit.newText).toBe('console.log(value + 1);'); }); + it('lets multiple parallel provideNextEdit calls reuse the same speculative request', async () => { + // Mirrors the dedupe behavior of `_pendingStatelessNextEditRequest`: once a + // speculative request is in-flight for `(docId, postEditContent)`, multiple + // `provideNextEdit` invocations targeting that same state should join it + // rather than triggering additional stateless calls. + 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: 'yieldEditThenNoSuggestions', edit: lineReplacement(2, 'console.log(value + 1);') }); + const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); + + const doc = workspace.addDocument({ + id: DocumentId.create(URI.file('/test/spec-multi.ts').toString()), + initialValue: 'const value = 1;\nconsole.log(value);', + }); + doc.setSelection([new OffsetRange(0, 0)], undefined); + + const firstSuggestion = await getNextEdit(nextEditProvider, doc.id); + assert(firstSuggestion.result?.edit); + nextEditProvider.handleShown(firstSuggestion); + await statelessProvider.waitForCall(2); + await statelessProvider.calls[1].completed.p; + + nextEditProvider.handleAcceptance(doc.id, firstSuggestion); + doc.applyEdit(firstSuggestion.result.edit.toEdit()); + + // Two parallel reuses of the same speculative — neither should trigger a new stateless call. + const [a, b] = await Promise.all([ + getNextEdit(nextEditProvider, doc.id), + getNextEdit(nextEditProvider, doc.id), + ]); + + expect(statelessProvider.calls.length).toBe(2); + assert(a.result?.edit); + assert(b.result?.edit); + expect(a.result.edit.newText).toBe('console.log(value + 1);'); + expect(b.result.edit.newText).toBe('console.log(value + 1);'); + }); + it('cancels speculative request on rejection', async () => { await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); @@ -417,17 +457,20 @@ 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(); - statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'const value = 2;') }); + // Pure-insertion suggestion (replace `'foo();'` with `'foobar();'` is + // minimised to inserting `'bar'` at offset 3) — the trajectory check + // only applies to insertions. + statelessProvider.enqueueBehavior({ kind: 'yieldEditThenNoSuggestions', edit: lineReplacement(1, 'foobar();') }); statelessProvider.enqueueBehavior({ kind: 'waitForCancellation' }); const { nextEditProvider, workspace } = createProviderAndWorkspace(statelessProvider); const doc = workspace.addDocument({ id: DocumentId.create(URI.file('/test/spec-diverge.ts').toString()), - initialValue: 'const value = 1;\nconsole.log(value);', + initialValue: 'foo();\nconsole.log();', }); doc.setSelection([new OffsetRange(0, 0)], undefined); @@ -436,29 +479,76 @@ 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: 'foo();\nconsole.log();', + }); + 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); + + // 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(4, 'a')); + await flushMicrotasks(); + expect(statelessProvider.calls[1].wasCancelled).toBe(false); + + doc.applyEdit(StringEdit.insert(5, 'r')); + await flushMicrotasks(); + expect(statelessProvider.calls[1].wasCancelled).toBe(false); + + // 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('does not cancel speculative for substitution suggestions on user edits (trajectory check skipped)', async () => { + // The type-through trajectory check only models pure insertions. For + // substitutions the live document still contains the to-be-replaced + // `oldText` in the gap, which is generally not a prefix of `newText` + // so a naive check would over-cancel on the first keystroke. Verify + // substitution speculatives stay alive across user edits. + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsSpeculativeRequests, SpeculativeRequestsEnablement.On); + + const statelessProvider = new TestStatelessNextEditProvider(); + // Suggestion: replace `'const value = 1;'` with `'const value = 2;'` (substitution). + 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-substitution.ts').toString()), initialValue: 'const value = 1;\nconsole.log(value);', }); doc.setSelection([new OffsetRange(0, 0)], undefined); @@ -468,21 +558,14 @@ 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')); + // Make a small unrelated edit further down — for a pure-insertion + // trajectory this would diverge, but for a substitution the + // trajectory check is skipped and the speculative must stay alive. + doc.applyEdit(StringEdit.insert(doc.value.get().value.length, '\n// extra')); await flushMicrotasks(); + expect(statelessProvider.calls[1].wasCancelled).toBe(false); - doc.applyEdit(StringEdit.insert(1, 'b')); - await flushMicrotasks(); - expect(statelessProvider.calls[1].wasCancelled).toBe(false); - - doc.applyEdit(StringEdit.insert(2, 'c')); - await flushMicrotasks(); - expect(statelessProvider.calls[1].wasCancelled).toBe(false); - - // Clean up via rejection nextEditProvider.handleRejection(doc.id, suggestion); await statelessProvider.calls[1].completed.p; }); @@ -1370,4 +1453,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); + }); + }); });