Merge pull request #4869 from microsoft/ulugbekna/nes-dont-yield-fetch-with-error-candidate

nes: fix: don't yield edits when fetch errors or is cancelled

Author: eleanorjboyd

src/extension/xtab/node/xtabCustomDiffPatchResponseHandler.ts

@@ -69,10 +69,15 @@ export class XtabCustomDiffPatchResponseHandler {
 		workspaceRoot: URI | undefined,
 		window: OffsetRange | undefined,
 		originalWindow?: OffsetRange,
+		getFetchFailure?: () => NoNextEditReason | undefined,
 	): AsyncGenerator<StreamedEdit, NoNextEditReason, void> {
 		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) ?? activeDocumentId; // FIXME@ulugbekna: it's wrong to fallback to active document but just ignoring edits is also bad

src/extension/xtab/node/xtabProvider.ts

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

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

@@ -4,9 +4,26 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { describe, expect, it } from 'vitest';
+import { DocumentId } from '../../../../platform/inlineEdits/common/dataTypes/documentId';
+import { NoNextEditReason, StreamedEdit } from '../../../../platform/inlineEdits/common/statelessNextEditProvider';
 import { AsyncIterUtils } from '../../../../util/common/asyncIterableUtils';
+import { StringText } from '../../../../util/vs/editor/common/core/text/abstractText';
 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', () => {
 
 	async function collectPatches(patchText: string): Promise<string> {
@@ -101,4 +118,77 @@ another_file.js:
 			-Old line 2"
 		`);
 	});
+
+	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 StringText('old\n');
+
+		let yieldCount = 0;
+		const cancellationReason = new NoNextEditReason.GotCancelled('afterFetchCall');
+
+		const { edits, returnValue } = await consumeHandleResponse(
+			linesStream,
+			documentBeforeEdits,
+			docId,
+			undefined,
+			undefined,
+			undefined,
+			() => {
+				// 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 truncated patch when fetch is cancelled mid-stream', async () => {
+		// Full response would be:
+		//   file.ts:0 / -old / +new / file.ts:5 / -another old / +another new / +one more new
+		// But the fetch is cancelled right before "+one more new" is emitted,
+		// so the stream only delivers lines up to "+another new".
+		// The second patch is incomplete and must not be yielded.
+		const truncatedPatchText = `file.ts:0
+-old
++new
+file.ts:5
+-another old
++another new`;
+		// "+one more new" is missing — fetch was cancelled before it arrived
+		const linesStream = AsyncIterUtils.fromArray(truncatedPatchText.split('\n'));
+		const docId = DocumentId.create('file:///file.ts');
+		const documentBeforeEdits = new StringText('old\n');
+
+		let checkCount = 0;
+		const cancellationReason = new NoNextEditReason.GotCancelled('afterFetchCall');
+
+		const { edits, returnValue } = await consumeHandleResponse(
+			linesStream,
+			documentBeforeEdits,
+			docId,
+			undefined,
+			undefined,
+			undefined,
+			() => {
+				// Fetch was still running (not yet cancelled) when the first edit was checked,
+				// but was cancelled before the second (incomplete) edit is about to be yielded.
+				return checkCount++ > 0 ? cancellationReason : undefined;
+			},
+		);
+
+		// The first (complete) edit is yielded; the second (truncated) edit is not
+		expect(edits).toHaveLength(1);
+		expect(returnValue).toBe(cancellationReason);
+	});
 });