fix: make ResponseProcessor.diff and XtabCustomDiffPatchResponseHandler respect cancellation token

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<StreamedEdit, NoNextEditReason, void> {
 		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;

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) {

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);
+	});
+});

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);
+	});
 });

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<string>, cursorOriginalLinesOffset: number, params: DiffParams): AsyncIterable<LineReplacement> {
+	export async function* diff(originalLines: string[], modifiedLines: AsyncIterable<string>, cursorOriginalLinesOffset: number, params: DiffParams, cancellationToken: CancellationToken = CancellationToken.None): AsyncIterable<LineReplacement> {
 
 		const lineToIdxs = new ArrayMap<string, number>();
 		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