diff --git a/src/vs/workbench/contrib/chat/browser/actions/codeBlockOperations.ts b/src/vs/workbench/contrib/chat/browser/actions/codeBlockOperations.ts index 947199fc..3149f453 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/codeBlockOperations.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/codeBlockOperations.ts @@ -11,7 +11,6 @@ import { ResourceMap } from '../../../../../base/common/map.js'; import { isEqual } from '../../../../../base/common/resources.js'; import * as strings from '../../../../../base/common/strings.js'; import { URI } from '../../../../../base/common/uri.js'; -import { stripThinkTags } from '../../common/chatModel.js'; import { IActiveCodeEditor, isCodeEditor, isDiffEditor } from '../../../../../editor/browser/editorBrowser.js'; import { IBulkEditService, ResourceTextEdit } from '../../../../../editor/browser/services/bulkEditService.js'; import { ICodeEditorService } from '../../../../../editor/browser/services/codeEditorService.js'; @@ -383,11 +382,13 @@ function collectDocumentContextFromContext(context: ICodeBlockActionContext, res } function getChatConversation(context: ICodeBlockActionContext): (ConversationRequest | ConversationResponse)[] { + // TODO@aeschli for now create a conversation with just the current element + // this will be expanded in the future to include the request and any other responses + if (isResponseVM(context.element)) { return [{ type: 'response', - // Strip think tags before computing diffs - message: stripThinkTags(context.element.response.toMarkdown()), + message: context.element.response.toMarkdown(), references: getReferencesAsDocumentContext(context.element.contentReferences) }]; } else if (isRequestVM(context.element)) { diff --git a/src/vs/workbench/contrib/chat/common/chatModel.ts b/src/vs/workbench/contrib/chat/common/chatModel.ts index 81b0c970..f11a40b7 100644 --- a/src/vs/workbench/contrib/chat/common/chatModel.ts +++ b/src/vs/workbench/contrib/chat/common/chatModel.ts @@ -8,7 +8,6 @@ import { DeferredPromise } from '../../../../base/common/async.js'; import { Emitter, Event } from '../../../../base/common/event.js'; import { IMarkdownString, MarkdownString, isMarkdownString } from '../../../../base/common/htmlContent.js'; import { Disposable } from '../../../../base/common/lifecycle.js'; -import { SurroundingsRemover } from '../../../void/browser/helpers/extractCodeFromResult.js'; import { revive } from '../../../../base/common/marshalling.js'; import { equals } from '../../../../base/common/objects.js'; import { basename, isEqual } from '../../../../base/common/resources.js'; @@ -250,11 +249,6 @@ export class Response extends Disposable implements IResponse { updateContent(progress: IChatProgressResponseContent | IChatTextEdit | IChatTask, quiet?: boolean): void { if (progress.kind === 'markdownContent') { - // Handle streaming for think tags - const remover = new ThinkTagSurroundingsRemover(progress.content.value); - const [delta, ignoredSuffix] = remover.deltaInfo(progress.content.value.length); - progress.content.value = delta; - const responsePartLength = this._responseParts.length - 1; const lastResponsePart = this._responseParts[responsePartLength]; @@ -360,78 +354,6 @@ export class Response extends Disposable implements IResponse { } } -/** - * Strips tags and their content from a text string. - * Handles nested tags using a stack-based approach. - * @param text The text to strip tags from - * @returns The text with all tags and their content removed - */ -export function stripThinkTags(text: string): string { - // Handle nested tags with a stack-based approach - let result = ''; - let depth = 0; - let i = 0; - - while (i < text.length) { - if (text.startsWith('', i)) { - depth++; - i += 7; // length of '' - } else if (text.startsWith('', i)) { - if (depth > 0) depth--; - i += 8; // length of '' - } else if (depth === 0) { - result += text[i]; - i++; - } else { - i++; - } - } - - return result; -} - -class ThinkTagSurroundingsRemover extends SurroundingsRemover { - constructor(s: string) { - super(s); - } - - removeThinkTags() { - // Handle token streaming at a more granular level - let foundTag = false; - - // Try to remove opening tag, handling partial tokens - foundTag = this.removePrefix('<'); - if (foundTag) { - foundTag = this.removePrefix('t'); - if (foundTag) { - foundTag = this.removePrefix('h'); - if (foundTag) { - foundTag = this.removePrefix('i'); - if (foundTag) { - foundTag = this.removePrefix('n'); - if (foundTag) { - foundTag = this.removePrefix('k'); - if (foundTag) { - foundTag = this.removePrefix('>'); - } - } - } - } - } - } - - return foundTag; - } - - deltaInfo(recentlyAddedTextLen: number) { - // Get the delta and suffix from parent class - const [delta, ignoredSuffix] = super.deltaInfo(recentlyAddedTextLen); - - // Strip any think tags from the delta before returning - return [stripThinkTags(delta), ignoredSuffix] as const; - } -} - export class ChatResponseModel extends Disposable implements IChatResponseModel { private readonly _onDidChange = this._register(new Emitter()); readonly onDidChange = this._onDidChange.event; 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 24ff8129..0a49d0bf 100644 --- a/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts +++ b/src/vs/workbench/contrib/chat/test/common/chatModel.test.ts @@ -3,57 +3,276 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import * as assert from 'assert'; -import { Response, stripThinkTags } from '../../../common/chatModel'; -import { MarkdownString } from '../../../../../../base/common/htmlContent'; +import assert from 'assert'; +import { timeout } from '../../../../../base/common/async.js'; +import { MarkdownString } from '../../../../../base/common/htmlContent.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { assertSnapshot } from '../../../../../base/test/common/snapshot.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { OffsetRange } from '../../../../../editor/common/core/offsetRange.js'; +import { Range } from '../../../../../editor/common/core/range.js'; +import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js'; +import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; +import { MockContextKeyService } from '../../../../../platform/keybinding/test/common/mockKeybindingService.js'; +import { ILogService, NullLogService } from '../../../../../platform/log/common/log.js'; +import { IStorageService } from '../../../../../platform/storage/common/storage.js'; +import { ChatAgentLocation, ChatAgentService, IChatAgentService } from '../../common/chatAgents.js'; +import { ChatModel, ISerializableChatData1, ISerializableChatData2, ISerializableChatData3, normalizeSerializableChatData, Response } from '../../common/chatModel.js'; +import { ChatRequestTextPart } from '../../common/chatParserTypes.js'; +import { IExtensionService } from '../../../../services/extensions/common/extensions.js'; +import { TestExtensionService, TestStorageService } from '../../../../test/common/workbenchTestServices.js'; -suite('ChatModel - Think Tags', () => { - test('handles partial tokens during streaming', () => { - const response = new Response(new MarkdownString('<')); - assert.strictEqual(response.toString(), ''); +suite('ChatModel', () => { + const testDisposables = ensureNoDisposablesAreLeakedInTestSuite(); - response.updateContent({ kind: 'markdownContent', content: new MarkdownString('') }); - assert.strictEqual(response.toString(), ''); - - response.updateContent({ kind: 'markdownContent', content: new MarkdownString('test') }); - assert.strictEqual(response.toString(), ''); - - response.updateContent({ kind: 'markdownContent', content: new MarkdownString('test') }); - assert.strictEqual(response.toString(), 'test'); + setup(async () => { + instantiationService = testDisposables.add(new TestInstantiationService()); + instantiationService.stub(IStorageService, testDisposables.add(new TestStorageService())); + instantiationService.stub(ILogService, new NullLogService()); + instantiationService.stub(IExtensionService, new TestExtensionService()); + instantiationService.stub(IContextKeyService, new MockContextKeyService()); + instantiationService.stub(IChatAgentService, instantiationService.createInstance(ChatAgentService)); }); - test('handles malformed tags', () => { - assert.strictEqual(stripThinkTags('unclosed'), ''); - assert.strictEqual(stripThinkTags('unopened'), 'unopened'); - assert.strictEqual(stripThinkTags('nestedtags'), ''); + test('Waits for initialization', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + + let hasInitialized = false; + model.waitForInitialization().then(() => { + hasInitialized = true; + }); + + await timeout(0); + assert.strictEqual(hasInitialized, false); + + model.startInitialize(); + model.initialize(undefined); + await timeout(0); + assert.strictEqual(hasInitialized, true); }); - test('handles half-nested tags from conversation', () => { - const input = `Okay, the user wants me to create nested tags in my thought process. Let me start by recalling what they asked for. They initially mentioned using nested tags with multiple layers. In my first attempt, I probably just used a single pair of tags without nesting. + test('must call startInitialize before initialize', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); -So, I need to make sure each opening tag is properly closed with a tag, and that these tags are nested. For example, starting with one tag, then another inside it, and so on. + let hasInitialized = false; + model.waitForInitialization().then(() => { + hasInitialized = true; + }); -i tried`; - assert.strictEqual(stripThinkTags(input), 'i tried'); + await timeout(0); + assert.strictEqual(hasInitialized, false); + + assert.throws(() => model.initialize(undefined)); + assert.strictEqual(hasInitialized, false); }); - test('preserves text mentions of tags', () => { - const input = 'Let me try with tags and see how it works'; - assert.strictEqual(stripThinkTags(input), 'Let me try with tags and see how it works'); + test('deinitialize/reinitialize', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + + let hasInitialized = false; + model.waitForInitialization().then(() => { + hasInitialized = true; + }); + + model.startInitialize(); + model.initialize(undefined); + await timeout(0); + assert.strictEqual(hasInitialized, true); + + model.deinitialize(); + let hasInitialized2 = false; + model.waitForInitialization().then(() => { + hasInitialized2 = true; + }); + + model.startInitialize(); + model.initialize(undefined); + await timeout(0); + assert.strictEqual(hasInitialized2, true); + }); + + test('cannot initialize twice', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + + model.startInitialize(); + model.initialize(undefined); + assert.throws(() => model.initialize(undefined)); + }); + + test('Initialization fails when model is disposed', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + model.dispose(); + + assert.throws(() => model.initialize(undefined)); + }); + + test('removeRequest', async () => { + const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + + model.startInitialize(); + model.initialize(undefined); + const text = 'hello'; + model.addRequest({ text, parts: [new ChatRequestTextPart(new OffsetRange(0, text.length), new Range(1, text.length, 1, text.length), text)] }, { variables: [] }, 0); + const requests = model.getRequests(); + assert.strictEqual(requests.length, 1); + + model.removeRequest(requests[0].id); + assert.strictEqual(model.getRequests().length, 0); + }); + + test('adoptRequest', async function () { + const model1 = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Editor)); + const model2 = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel)); + + model1.startInitialize(); + model1.initialize(undefined); + + model2.startInitialize(); + model2.initialize(undefined); + + const text = 'hello'; + const request1 = model1.addRequest({ text, parts: [new ChatRequestTextPart(new OffsetRange(0, text.length), new Range(1, text.length, 1, text.length), text)] }, { variables: [] }, 0); + + assert.strictEqual(model1.getRequests().length, 1); + assert.strictEqual(model2.getRequests().length, 0); + assert.ok(request1.session === model1); + assert.ok(request1.response?.session === model1); + + model2.adoptRequest(request1); + + assert.strictEqual(model1.getRequests().length, 0); + assert.strictEqual(model2.getRequests().length, 1); + assert.ok(request1.session === model2); + assert.ok(request1.response?.session === model2); + + model2.acceptResponseProgress(request1, { content: new MarkdownString('Hello'), kind: 'markdownContent' }); + + assert.strictEqual(request1.response.response.toString(), 'Hello'); + }); +}); + +suite('Response', () => { + const store = ensureNoDisposablesAreLeakedInTestSuite(); + + test('mergeable markdown', async () => { + const response = store.add(new Response([])); + response.updateContent({ content: new MarkdownString('markdown1'), kind: 'markdownContent' }); + response.updateContent({ content: new MarkdownString('markdown2'), kind: 'markdownContent' }); + await assertSnapshot(response.value); + + assert.strictEqual(response.toString(), 'markdown1markdown2'); + }); + + test('not mergeable markdown', async () => { + const response = store.add(new Response([])); + const md1 = new MarkdownString('markdown1'); + md1.supportHtml = true; + response.updateContent({ content: md1, kind: 'markdownContent' }); + response.updateContent({ content: new MarkdownString('markdown2'), kind: 'markdownContent' }); + await assertSnapshot(response.value); + }); + + test('inline reference', async () => { + const response = store.add(new Response([])); + response.updateContent({ content: new MarkdownString('text before'), kind: 'markdownContent' }); + response.updateContent({ inlineReference: URI.parse('https://microsoft.com'), kind: 'inlineReference' }); + response.updateContent({ content: new MarkdownString('text after'), kind: 'markdownContent' }); + await assertSnapshot(response.value); + }); +}); + +suite('normalizeSerializableChatData', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('v1', () => { + const v1Data: ISerializableChatData1 = { + creationDate: Date.now(), + initialLocation: undefined, + isImported: false, + requesterAvatarIconUri: undefined, + requesterUsername: 'me', + requests: [], + responderAvatarIconUri: undefined, + responderUsername: 'bot', + sessionId: 'session1', + }; + + const newData = normalizeSerializableChatData(v1Data); + assert.strictEqual(newData.creationDate, v1Data.creationDate); + assert.strictEqual(newData.lastMessageDate, v1Data.creationDate); + assert.strictEqual(newData.version, 3); + assert.ok('customTitle' in newData); + }); + + test('v2', () => { + const v2Data: ISerializableChatData2 = { + version: 2, + creationDate: 100, + lastMessageDate: Date.now(), + initialLocation: undefined, + isImported: false, + requesterAvatarIconUri: undefined, + requesterUsername: 'me', + requests: [], + responderAvatarIconUri: undefined, + responderUsername: 'bot', + sessionId: 'session1', + computedTitle: 'computed title' + }; + + const newData = normalizeSerializableChatData(v2Data); + assert.strictEqual(newData.version, 3); + assert.strictEqual(newData.creationDate, v2Data.creationDate); + assert.strictEqual(newData.lastMessageDate, v2Data.lastMessageDate); + assert.strictEqual(newData.customTitle, v2Data.computedTitle); + }); + + test('old bad data', () => { + const v1Data: ISerializableChatData1 = { + // Testing the scenario where these are missing + sessionId: undefined!, + creationDate: undefined!, + + initialLocation: undefined, + isImported: false, + requesterAvatarIconUri: undefined, + requesterUsername: 'me', + requests: [], + responderAvatarIconUri: undefined, + responderUsername: 'bot', + }; + + const newData = normalizeSerializableChatData(v1Data); + assert.strictEqual(newData.version, 3); + assert.ok(newData.creationDate > 0); + assert.ok(newData.lastMessageDate > 0); + assert.ok(newData.sessionId); + }); + + test('v3 with bug', () => { + const v3Data: ISerializableChatData3 = { + // Test case where old data was wrongly normalized and these fields were missing + creationDate: undefined!, + lastMessageDate: undefined!, + + version: 3, + initialLocation: undefined, + isImported: false, + requesterAvatarIconUri: undefined, + requesterUsername: 'me', + requests: [], + responderAvatarIconUri: undefined, + responderUsername: 'bot', + sessionId: 'session1', + customTitle: 'computed title' + }; + + const newData = normalizeSerializableChatData(v3Data); + assert.strictEqual(newData.version, 3); + assert.ok(newData.creationDate > 0); + assert.ok(newData.lastMessageDate > 0); + assert.ok(newData.sessionId); }); });