From e210e470b061bb42660bb0a07e412fc7f0595032 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 5 Apr 2026 20:07:06 +0000 Subject: [PATCH 1/3] Initial plan From acf56c9d4982f765e843f69c41dd13fe12331107 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 5 Apr 2026 20:14:16 +0000 Subject: [PATCH 2/3] fix: make ResponseProcessor.diff respect cancellation token Add a CancellationToken parameter (defaulting to CancellationToken.None) to ResponseProcessor.diff. The generator now checks the token at the beginning of each loop iteration and returns early when cancellation is requested. Also update the streamEdits caller in xtabProvider to: - Pass the cancellation token through to ResponseProcessor.diff - Check the cancellation token after the for-await loop exits, returning GotCancelled('duringStreaming') instead of NoSuggestions when cancelled Co-authored-by: ulugbekna <16353531+ulugbekna@users.noreply.github.com> --- src/extension/xtab/node/xtabProvider.ts | 7 ++- .../test/common/responseProcessor.spec.ts | 47 +++++++++++++++++++ .../inlineEdits/common/responseProcessor.ts | 6 ++- 3 files changed, 58 insertions(+), 2 deletions(-) diff --git a/src/extension/xtab/node/xtabProvider.ts b/src/extension/xtab/node/xtabProvider.ts index e1e8a89dfe..6aa380055f 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 0a9eecbb5f..1851d6cbb7 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/platform/inlineEdits/common/responseProcessor.ts b/src/platform/inlineEdits/common/responseProcessor.ts index cee7acd721..cbfe25dfb7 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 From 26df85d7b58024f3c64dde54ed689dbe7529895d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 5 Apr 2026 20:14:21 +0000 Subject: [PATCH 3/3] fix: make XtabCustomDiffPatchResponseHandler respect cancellation token Add a CancellationToken parameter (defaulting to CancellationToken.None) to XtabCustomDiffPatchResponseHandler.handleResponse. The handler now checks the token before processing each extracted patch and returns NoNextEditReason.GotCancelled('duringStreaming') when cancellation is requested. Also update the caller in xtabProvider to pass the cancellation token through. Co-authored-by: ulugbekna <16353531+ulugbekna@users.noreply.github.com> --- .../xtabCustomDiffPatchResponseHandler.ts | 5 ++ ...xtabCustomDiffPatchResponseHandler.spec.ts | 76 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts b/src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts index 97d0e6431b..0cf4063375 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/test/node/xtabCustomDiffPatchResponseHandler.spec.ts b/src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts index 5a35e89dec..a31da4651a 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); + }); });