From b24cd1ecb7356d2228a801086e08bfdd6e6c1b8f Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sun, 16 Feb 2025 05:01:57 +0000
Subject: [PATCH] Improve token streaming support for think tags

Co-Authored-By: Jack Hacksman <slack@hannis.io>
---
 .../workbench/contrib/chat/common/chatModel.ts  | 17 ++++++++---------
 .../contrib/chat/test/common/chatModel.test.ts  | 16 ++++++++++++++--
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/vs/workbench/contrib/chat/common/chatModel.ts b/src/vs/workbench/contrib/chat/common/chatModel.ts
index df982d6c..7fb12aac 100644
--- a/src/vs/workbench/contrib/chat/common/chatModel.ts
+++ b/src/vs/workbench/contrib/chat/common/chatModel.ts
@@ -409,16 +409,15 @@ class ThinkTagSurroundingsRemover extends SurroundingsRemover {
 	}
 
 	removeThinkTags() {
-		const foundTag = this.removePrefix('<think>');
-		if (!foundTag) {
-			// Handle partial tags during streaming
-			if (this.originalS.startsWith('<thi', this.i)) {
-				this.i += 4;
-				return true;
-			}
-			return false;
+		// Try to remove opening tag, handling partial tokens during streaming
+		const foundOpenTag = this.removePrefix('<think>');
+		if (!foundOpenTag) {
+			// Let removePrefix handle partial matches character by character
+			this.removePrefix('<');
+			this.removePrefix('think');
+			this.removePrefix('>');
 		}
-		return true;
+		return foundOpenTag;
 	}
 
 	deltaInfo(recentlyAddedTextLen: number) {
diff --git a/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts b/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts
index 8269d4fa..a145fba6 100644
--- a/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts
+++ b/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts
@@ -8,8 +8,20 @@ import { Response, stripThinkTags } from '../../../common/chatModel';
 import { MarkdownString } from '../../../../../../base/common/htmlContent';
 
 suite('ChatModel - Think Tags', () => {
-	test('handles partial tags during streaming', () => {
-		const response = new Response(new MarkdownString('<thi'));
+	test('handles partial tokens during streaming', () => {
+		const response = new Response(new MarkdownString('<'));
+		assert.strictEqual(response.toString(), '');
+
+		response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<t') });
+		assert.strictEqual(response.toString(), '');
+
+		response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<thi') });
+		assert.strictEqual(response.toString(), '');
+
+		response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think') });
+		assert.strictEqual(response.toString(), '');
+
+		response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think>') });
 		assert.strictEqual(response.toString(), '');
 
 		response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think>test') });