Fix chat rendering issue (#230677)

* Fix chat rendering issue Sometimes not properly rendering non-markdown content after consuming all markdown words. Got stuck when the markdown part ends in whitespace. * Fix test * Fix test
microsoft · Oct 7, 2024 · 75d515e · 75d515e
1 parent e2e048d
commit 75d515e
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 30 deletions.
diff --git a/src/vs/workbench/contrib/chat/browser/chatListRenderer.ts b/src/vs/workbench/contrib/chat/browser/chatListRenderer.ts
@@ -590,30 +590,27 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 		}
 
 		templateData.rowContainer.classList.toggle('chat-response-loading', true);
-		let isFullyRendered = false;
 		this.traceLayout('doNextProgressiveRender', `START progressive render, index=${index}, renderData=${JSON.stringify(element.renderData)}`);
 		const contentForThisTurn = this.getNextProgressiveRenderContent(element);
 		const partsToRender = this.diff(templateData.renderedParts ?? [], contentForThisTurn.content, element);
-		isFullyRendered = partsToRender.every(part => part === null);
 
-		if (isFullyRendered) {
-			if (element.isComplete) {
-				// Response is done and content is rendered, so do a normal render
+		const contentIsAlreadyRendered = partsToRender.every(part => part === null);
+		if (contentIsAlreadyRendered) {
+			if (contentForThisTurn.moreContentAvailable) {
+				// The content that we want to render in this turn is already rendered, but there is more content to render on the next tick
+				this.traceLayout('doNextProgressiveRender', 'not rendering any new content this tick, but more available');
+				return false;
+			} else if (element.isComplete) {
+				// All content is rendered, and response is done, so do a normal render
 				this.traceLayout('doNextProgressiveRender', `END progressive render, index=${index} and clearing renderData, response is complete`);
 				element.renderData = undefined;
 				this.basicRenderElement(element, index, templateData);
 				return true;
-			}
-
-			if (!contentForThisTurn.moreContentAvailable) {
+			} else {
 				// Nothing new to render, stop rendering until next model update
 				this.traceLayout('doNextProgressiveRender', 'caught up with the stream- no new content to render');
 				return true;
 			}
-
-			// The content that we want to render in this turn is already rendered, but there is more content to render on the next tick
-			this.traceLayout('doNextProgressiveRender', 'not rendering any new content this tick, but more available');
-			return false;
 		}
 
 		// Do an actual progressive render
@@ -690,29 +687,26 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 			const part = renderableResponse[i];
 			if (part.kind === 'markdownContent') {
 				const wordCountResult = getNWords(part.content.value, numNeededWords);
-				if (wordCountResult.isFullString) {
-					partsToRender.push(part);
-				} else {
-					partsToRender.push({ kind: 'markdownContent', content: new MarkdownString(wordCountResult.value, part.content) });
-				}
-
 				this.traceLayout('getNextProgressiveRenderContent', `  Chunk ${i}: Want to render ${numNeededWords} words and found ${wordCountResult.returnedWordCount} words. Total words in chunk: ${wordCountResult.totalWordCount}`);
 				numNeededWords -= wordCountResult.returnedWordCount;
 
-				if (numNeededWords <= 0) {
-					// No more markdown needed, but need to ensure that all following non-markdown parts are rendered
-					i++;
-					while (i < renderableResponse.length) {
-						const nextPart = renderableResponse[i];
+				if (wordCountResult.isFullString) {
+					partsToRender.push(part);
+
+					// Consumed full markdown chunk- need to ensure that all following non-markdown parts are rendered
+					for (const nextPart of renderableResponse.slice(i + 1)) {
 						if (nextPart.kind !== 'markdownContent') {
+							i++;
 							partsToRender.push(nextPart);
 						} else {
 							break;
 						}
-
-						i++;
 					}
+				} else {
+					partsToRender.push({ kind: 'markdownContent', content: new MarkdownString(wordCountResult.value, part.content) });
+				}
 
+				if (numNeededWords <= 0) {
 					// Collected all words and following non-markdown parts if needed, done
 					break;
 				}

diff --git a/src/vs/workbench/contrib/chat/common/chatWordCounter.ts b/src/vs/workbench/contrib/chat/common/chatWordCounter.ts
@@ -50,7 +50,7 @@ export function getNWords(str: string, numWordsToCount: number): IWordCountResul
 
 	const targetWords = allWordMatches.slice(0, numWordsToCount);
 
-	const endIndex = numWordsToCount > allWordMatches.length
+	const endIndex = numWordsToCount >= allWordMatches.length
 		? str.length // Reached end of string
 		: targetWords.length ? targetWords.at(-1)!.index + targetWords.at(-1)![0].length : 0;
 

diff --git a/src/vs/workbench/contrib/chat/test/common/chatWordCounter.test.ts b/src/vs/workbench/contrib/chat/test/common/chatWordCounter.test.ts
@@ -5,7 +5,7 @@
 
 import assert from 'assert';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
-import { getNWords } from '../../common/chatWordCounter.js';
+import { getNWords, IWordCountResult } from '../../common/chatWordCounter.js';
 
 suite('ChatWordCounter', () => {
 	ensureNoDisposablesAreLeakedInTestSuite();
@@ -23,15 +23,42 @@ suite('ChatWordCounter', () => {
 				['hello', 1, 'hello'],
 				['hello world', 0, ''],
 				['here\'s, some.   punctuation?', 3, 'here\'s, some.   punctuation?'],
-				['| markdown | _table_ | header |', 3, '| markdown | _table_ | header'],
+				['| markdown | _table_ | header |', 3, '| markdown | _table_ | header |'],
 				['| --- | --- | --- |', 1, '| ---'],
-				['| --- | --- | --- |', 3, '| --- | --- | ---'],
-				[' \t some \n whitespace     \n\n\nhere   ', 3, ' \t some \n whitespace     \n\n\nhere'],
+				['| --- | --- | --- |', 3, '| --- | --- | --- |'],
+				[' \t some \n whitespace     \n\n\nhere   ', 3, ' \t some \n whitespace     \n\n\nhere   '],
 			];
 
 			cases.forEach(([str, nWords, result]) => doTest(str, nWords, result));
 		});
 
+		test('whitespace', () => {
+			assert.deepStrictEqual(
+				getNWords('hello ', 1),
+				{
+					value: 'hello ',
+					returnedWordCount: 1,
+					isFullString: true,
+					totalWordCount: 1,
+				} satisfies IWordCountResult);
+			assert.deepStrictEqual(
+				getNWords('hello\n\n', 1),
+				{
+					value: 'hello\n\n',
+					returnedWordCount: 1,
+					isFullString: true,
+					totalWordCount: 1,
+				} satisfies IWordCountResult);
+			assert.deepStrictEqual(
+				getNWords('\nhello', 1),
+				{
+					value: '\nhello',
+					returnedWordCount: 1,
+					isFullString: true,
+					totalWordCount: 1,
+				} satisfies IWordCountResult);
+		});
+
 		test('matching links', () => {
 			const cases: [string, number, string][] = [
 				['[hello](https://example.com) world', 1, '[hello](https://example.com)'],