nes: patch: fix: don't yield malformed edits on fetch fail/cancellation (#4904)

Author: ulugbekna

src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts

@@ -70,11 +70,16 @@ export class XtabCustomDiffPatchResponseHandler {
 		workspaceRoot: URI | undefined,
 		window: OffsetRange | undefined,
 		parentTracer: ILogger,
+		getFetchFailure?: () => NoNextEditReason | undefined,
 	): 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)) {
+				const fetchFailure = getFetchFailure?.();
+				if (fetchFailure) {
+					return fetchFailure;
+				}
 				const targetDocument = edit.filePath === activeDocRelativePath
 					? activeDocumentId
 					: XtabCustomDiffPatchResponseHandler.resolveTargetDocument(edit.filePath, workspaceRoot);

src/extension/xtab/node/xtabProvider.ts

@@ -850,6 +850,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 				activeDoc.workspaceRoot,
 				pseudoEditWindow,
 				tracer,
+				() => chatResponseFailure ? mapChatFetcherErrorToNoNextEditReason(chatResponseFailure) : undefined,
 			);
 		} else if (opts.responseFormat === xtabPromptOptions.ResponseFormat.UnifiedWithXml) {
 			const linesIter = linesStream[Symbol.asyncIterator]();

src/extension/xtab/test/node/xtabCustomDiffPatchResponseHandler.spec.ts

@@ -3,12 +3,37 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { describe, expect, it } from 'vitest';
+import { beforeEach, describe, expect, it } from 'vitest';
+import { DocumentId } from '../../../../platform/inlineEdits/common/dataTypes/documentId';
+import { NoNextEditReason, StreamedEdit } from '../../../../platform/inlineEdits/common/statelessNextEditProvider';
+import { TestLogService } from '../../../../platform/testing/common/testLogService';
 import { AsyncIterUtils } from '../../../../util/common/asyncIterableUtils';
+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';
+import { CurrentDocument } from '../../common/xtabCurrentDocument';
 import { XtabCustomDiffPatchResponseHandler } from '../../node/xtabCustomDiffPatchResponseHandler';
 
+async function consumeHandleResponse(
+	...args: Parameters<typeof XtabCustomDiffPatchResponseHandler.handleResponse>
+): Promise<{ edits: StreamedEdit[]; returnValue: NoNextEditReason }> {
+	const gen = XtabCustomDiffPatchResponseHandler.handleResponse(...args);
+	const edits: StreamedEdit[] = [];
+	for (; ;) {
+		const result = await gen.next();
+		if (result.done) {
+			return { edits, returnValue: result.value };
+		}
+		edits.push(result.value);
+	}
+}
+
 describe('XtabCustomDiffPatchResponseHandler', () => {
 
+	beforeEach(() => {
+		ensureDependenciesAreSet();
+	});
+
 	async function collectPatches(patchText: string): Promise<string> {
 		const linesStream = AsyncIterUtils.fromArray(patchText.split('\n'));
 		const patches = await AsyncIterUtils.toArray(XtabCustomDiffPatchResponseHandler.extractEdits(linesStream));
@@ -113,4 +138,71 @@ another_file.js:
 			+Old line 1"
 		`);
 	});
+
+	it('stops yielding edits when getFetchFailure returns a failure', 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));
+
+		let yieldCount = 0;
+		const cancellationReason = new NoNextEditReason.GotCancelled('afterFetchCall');
+
+		const { edits, returnValue } = await consumeHandleResponse(
+			linesStream,
+			documentBeforeEdits,
+			docId,
+			undefined,
+			undefined,
+			new TestLogService(),
+			() => {
+				// Signal failure after the first edit has been yielded
+				if (yieldCount++ > 0) {
+					return cancellationReason;
+				}
+				return undefined;
+			},
+		);
+
+		expect(edits).toHaveLength(1);
+		expect(returnValue).toBe(cancellationReason);
+	});
+
+	it('does not yield incomplete patch when fetch is cancelled mid-stream', async () => {
+		// Stream is truncated before the last line "+one more new" is emitted,
+		// simulating a cancellation that cuts off the response mid-patch.
+		const truncatedPatchText = `/file.ts:0
+-old
++new
+/file.ts:5
+-another old
++another new`;
+		const linesStream = AsyncIterUtils.fromArray(truncatedPatchText.split('\n'));
+		const docId = DocumentId.create('file:///file.ts');
+		const documentBeforeEdits = new CurrentDocument(new StringText('old\n'), new Position(1, 1));
+
+		const cancellationReason = new NoNextEditReason.GotCancelled('afterFetchCall');
+
+		const { edits, returnValue } = await consumeHandleResponse(
+			linesStream,
+			documentBeforeEdits,
+			docId,
+			undefined,
+			undefined,
+			new TestLogService(),
+			() => {
+				// Fetch was cancelled — the full patch had "+one more new" but the
+				// stream was resolved early so it's missing from the second patch.
+				return cancellationReason;
+			},
+		);
+
+		expect(edits).toHaveLength(0);
+		expect(returnValue).toBe(cancellationReason);
+	});
 });