diff --git a/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts b/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts index 97d0e6431..0cf406337 100644 --- a/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts +++ b/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts @@ -7,6 +7,7 @@ import { DocumentId } from '../../../platform/inlineEdits/common/dataTypes/docum import { NoNextEditReason, StreamedEdit } from '../../../platform/inlineEdits/common/statelessNextEditProvider'; import { ILogger } from '../../../platform/log/common/logService'; import { ErrorUtils } from '../../../util/common/errors'; +import { CancellationToken } from '../../../util/vs/base/common/cancellation'; import { isAbsolute } from '../../../util/vs/base/common/path'; import { URI } from '../../../util/vs/base/common/uri'; import { LineReplacement } from '../../../util/vs/editor/common/core/edits/lineEdit'; @@ -71,11 +72,15 @@ export class XtabCustomDiffPatchResponseHandler { window: OffsetRange | undefined, parentTracer: ILogger, getFetchFailure?: () => NoNextEditReason | undefined, + cancellationToken: CancellationToken = CancellationToken.None, ): AsyncGenerator { const tracer = parentTracer.createSubLogger(['XtabCustomDiffPatchResponseHandler', 'handleResponse']); const activeDocRelativePath = toUniquePath(activeDocumentId, workspaceRoot?.path); try { for await (const edit of XtabCustomDiffPatchResponseHandler.extractEdits(linesStream)) { + if (cancellationToken.isCancellationRequested) { + return new NoNextEditReason.GotCancelled('duringStreaming'); + } const fetchFailure = getFetchFailure?.(); if (fetchFailure) { return fetchFailure; diff --git a/src/extension/xtab/node/xtabProvider.ts b/src/extension/xtab/node/xtabProvider.ts index e1e8a89df..6aa380055 100644 --- a/src/extension/xtab/node/xtabProvider.ts +++ b/src/extension/xtab/node/xtabProvider.ts @@ -851,6 +851,7 @@ export class XtabProvider implements IStatelessNextEditProvider { pseudoEditWindow, tracer, () => chatResponseFailure ? mapChatFetcherErrorToNoNextEditReason(chatResponseFailure) : undefined, + cancellationToken, ); } else if (opts.responseFormat === xtabPromptOptions.ResponseFormat.UnifiedWithXml) { const linesIter = linesStream[Symbol.asyncIterator](); @@ -939,7 +940,7 @@ export class XtabProvider implements IStatelessNextEditProvider { let i = 0; let hasBeenDelayed = false; try { - for await (const edit of ResponseProcessor.diff(editWindowLines, cleanedLinesStream, cursorOriginalLinesOffset, diffOptions)) { + for await (const edit of ResponseProcessor.diff(editWindowLines, cleanedLinesStream, cursorOriginalLinesOffset, diffOptions, cancellationToken)) { tracer.trace(`ResponseProcessor streamed edit #${i} with latency ${fetchRequestStopWatch.elapsed()} ms`); @@ -1003,6 +1004,10 @@ export class XtabProvider implements IStatelessNextEditProvider { return mapChatFetcherErrorToNoNextEditReason(chatResponseFailure); } + if (cancellationToken.isCancellationRequested) { + return new NoNextEditReason.GotCancelled('duringStreaming'); + } + return new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow); } catch (err) { diff --git a/src/extension/xtab/test/common/responseProcessor.spec.ts b/src/extension/xtab/test/common/responseProcessor.spec.ts index 0a9eecbb5..1851d6cbb 100644 --- a/src/extension/xtab/test/common/responseProcessor.spec.ts +++ b/src/extension/xtab/test/common/responseProcessor.spec.ts @@ -6,6 +6,7 @@ import { describe, expect, it, suite, test } from 'vitest'; import { ResponseProcessor } from '../../../../platform/inlineEdits/common/responseProcessor'; import { AsyncIterUtils } from '../../../../util/common/asyncIterableUtils'; +import { CancellationTokenSource } from '../../../../util/vs/base/common/cancellation'; import { LineEdit, LineReplacement } from '../../../../util/vs/editor/common/core/edits/lineEdit'; import { LineRange } from '../../../../util/vs/editor/common/core/ranges/lineRange'; @@ -577,3 +578,49 @@ describe('isAdditiveEdit', () => { expect(ResponseProcessor.isAdditiveEdit('aaaa', 'aaa')).toMatchInlineSnapshot(`false`); }); }); + +describe('ResponseProcessor.diff cancellation', () => { + + it('stops yielding edits when the cancellation token is cancelled', async () => { + const original = ['line1', 'line2', 'line3', 'line4', 'line5']; + const modified = ['line1', 'CHANGED2', 'CHANGED3', 'CHANGED4', 'line5']; + + const cts = new CancellationTokenSource(); + + const edits: LineReplacement[] = []; + // Cancel before any iteration starts + cts.cancel(); + + for await (const edit of ResponseProcessor.diff(original, AsyncIterUtils.fromArray(modified), 0, ResponseProcessor.DEFAULT_DIFF_PARAMS, cts.token)) { + edits.push(edit); + } + + // No edits should have been yielded because the token was already cancelled + expect(edits).toHaveLength(0); + }); + + it('stops mid-stream when the cancellation token is cancelled during iteration', async () => { + const original = ['a', 'b', 'c', 'd', 'e']; + const modified = ['X', 'Y', 'Z', 'W', 'V']; + + const cts = new CancellationTokenSource(); + + // Cancel mid-stream using an async iterable that cancels after the first item + async function* cancelMidStream() { + yield modified[0]; + cts.cancel(); + yield modified[1]; + yield modified[2]; + yield modified[3]; + yield modified[4]; + } + + const edits: LineReplacement[] = []; + for await (const edit of ResponseProcessor.diff(original, cancelMidStream(), 0, ResponseProcessor.DEFAULT_DIFF_PARAMS, cts.token)) { + edits.push(edit); + } + + // Cancellation stops the generator — no final edit should be emitted + expect(edits).toHaveLength(0); + }); +}); diff --git a/src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts b/src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts index 5a35e89de..a31da4651 100644 --- a/src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts +++ b/src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts @@ -8,6 +8,7 @@ import { DocumentId } from '../../../../platform/inlineEdits/common/dataTypes/do import { NoNextEditReason, StreamedEdit } from '../../../../platform/inlineEdits/common/statelessNextEditProvider'; import { TestLogService } from '../../../../platform/testing/common/testLogService'; import { AsyncIterUtils } from '../../../../util/common/asyncIterableUtils'; +import { CancellationTokenSource } from '../../../../util/vs/base/common/cancellation'; import { Position } from '../../../../util/vs/editor/common/core/position'; import { StringText } from '../../../../util/vs/editor/common/core/text/abstractText'; import { ensureDependenciesAreSet } from '../../../../util/vs/editor/common/core/text/positionToOffset'; @@ -205,4 +206,79 @@ another_file.js: expect(edits).toHaveLength(0); expect(returnValue).toBe(cancellationReason); }); + + it('returns GotCancelled when the cancellation token is already cancelled', async () => { + const patchText = `/file.ts:0 +-old ++new +/file.ts:5 +-another old ++another new`; + const linesStream = AsyncIterUtils.fromArray(patchText.split('\n')); + const docId = DocumentId.create('file:///file.ts'); + const documentBeforeEdits = new CurrentDocument(new StringText('old\n'), new Position(1, 1)); + + const cts = new CancellationTokenSource(); + cts.cancel(); + + const { edits, returnValue } = await consumeHandleResponse( + linesStream, + documentBeforeEdits, + docId, + undefined, + undefined, + new TestLogService(), + undefined, + cts.token, + ); + + expect(edits).toHaveLength(0); + expect(returnValue).toBeInstanceOf(NoNextEditReason.GotCancelled); + expect((returnValue as NoNextEditReason.GotCancelled).message).toBe('duringStreaming'); + }); + + it('stops yielding edits when the cancellation token is cancelled mid-stream', async () => { + const patchText = `/file.ts:0 +-old ++new +/file.ts:5 +-another old ++another new`; + const linesStream = AsyncIterUtils.fromArray(patchText.split('\n')); + const docId = DocumentId.create('file:///file.ts'); + const documentBeforeEdits = new CurrentDocument(new StringText('old\n'), new Position(1, 1)); + + const cts = new CancellationTokenSource(); + let yieldCount = 0; + + const gen = XtabCustomDiffPatchResponseHandler.handleResponse( + linesStream, + documentBeforeEdits, + docId, + undefined, + undefined, + new TestLogService(), + undefined, + cts.token, + ); + + const edits: StreamedEdit[] = []; + for (; ;) { + const result = await gen.next(); + if (result.done) { + // Verify cancellation is returned + expect(result.value).toBeInstanceOf(NoNextEditReason.GotCancelled); + expect((result.value as NoNextEditReason.GotCancelled).message).toBe('duringStreaming'); + break; + } + edits.push(result.value); + yieldCount++; + if (yieldCount === 1) { + // Cancel after first edit is yielded + cts.cancel(); + } + } + + expect(edits).toHaveLength(1); + }); }); diff --git a/src/platform/inlineEdits/common/responseProcessor.ts b/src/platform/inlineEdits/common/responseProcessor.ts index cee7acd72..cbfe25dfb 100644 --- a/src/platform/inlineEdits/common/responseProcessor.ts +++ b/src/platform/inlineEdits/common/responseProcessor.ts @@ -3,6 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ import { illegalArgument } from '../../../util/vs/base/common/errors'; +import { CancellationToken } from '../../../util/vs/base/common/cancellation'; import { LineReplacement } from '../../../util/vs/editor/common/core/edits/lineEdit'; import { LineRange } from '../../../util/vs/editor/common/core/ranges/lineRange'; @@ -63,7 +64,7 @@ export namespace ResponseProcessor { * @param modifiedLines * @param cursorOriginalLinesOffset offset of cursor within original lines */ - export async function* diff(originalLines: string[], modifiedLines: AsyncIterable, cursorOriginalLinesOffset: number, params: DiffParams): AsyncIterable { + export async function* diff(originalLines: string[], modifiedLines: AsyncIterable, cursorOriginalLinesOffset: number, params: DiffParams, cancellationToken: CancellationToken = CancellationToken.None): AsyncIterable { const lineToIdxs = new ArrayMap(); for (const [i, line] of originalLines.entries()) { @@ -76,6 +77,9 @@ export namespace ResponseProcessor { let state: DivergenceState = { k: 'aligned' }; for await (const line of modifiedLines) { + if (cancellationToken.isCancellationRequested) { + return; + } ++updatedEditWindowIdx; // handle modifiedLines.length > originalLines.length