diff --git a/lib/shared/src/chat/transcript/index.ts b/lib/shared/src/chat/transcript/index.ts index 00a5d2ad7b84..ac95a6e7f240 100644 --- a/lib/shared/src/chat/transcript/index.ts +++ b/lib/shared/src/chat/transcript/index.ts @@ -25,7 +25,11 @@ export interface SerializedChatInteraction { export function serializeChatMessage(chatMessage: ChatMessage): SerializedChatMessage { return { - ...chatMessage, + speaker: chatMessage.speaker, + model: chatMessage.model, + contextFiles: chatMessage.contextFiles, + editorState: chatMessage.editorState, + error: chatMessage.error, text: chatMessage.text ? chatMessage.text.toString() : undefined, } } diff --git a/lib/shared/src/editor/editorState.ts b/lib/shared/src/editor/editorState.ts index 41c1489ab4a3..ea45dc8e1f4c 100644 --- a/lib/shared/src/editor/editorState.ts +++ b/lib/shared/src/editor/editorState.ts @@ -5,6 +5,9 @@ export function setEditorWindowIsFocused(editorWindowIsFocused: () => boolean): } export function editorWindowIsFocused(): boolean { + if (process.env.VITEST) { + return true + } if (!_editorWindowIsFocused) { throw new Error('must call setEditorWindowIsFocused first') } diff --git a/lib/shared/src/index.ts b/lib/shared/src/index.ts index ae764be00e9b..38ba7c64c9fe 100644 --- a/lib/shared/src/index.ts +++ b/lib/shared/src/index.ts @@ -47,6 +47,7 @@ export { modelTier, parseModelRef, toLegacyModel, + FIXTURE_MODEL, } from './models/model' export { type EditModel, @@ -328,7 +329,12 @@ export { type ContextMentionProviderID, type ContextMentionProviderMetadata, } from './mentions/api' -export { TokenCounter, getTokenCounterUtils, TokenCounterUtils } from './token/counter' +export { + TokenCounter, + getTokenCounterUtils, + TokenCounterUtils, + useFakeTokenCounterUtils, +} from './token/counter' export { CORPUS_CONTEXT_ALLOCATION as ENHANCED_CONTEXT_ALLOCATION } from './token/constants' export { tokensToChars, charsToTokens } from './token/utils' export * from './prompt/prompt-string' diff --git a/lib/shared/src/models/model.ts b/lib/shared/src/models/model.ts index efd9beb256b4..047ae354e666 100644 --- a/lib/shared/src/models/model.ts +++ b/lib/shared/src/models/model.ts @@ -242,3 +242,9 @@ export function getServerModelTags( } return tags } + +export const FIXTURE_MODEL = createModel({ + id: 'my-model', + usage: [ModelUsage.Chat], + tags: [ModelTag.Enterprise], +}) diff --git a/lib/shared/src/token/counter.ts b/lib/shared/src/token/counter.ts index 5dc6e07326de..309ebd0e8222 100644 --- a/lib/shared/src/token/counter.ts +++ b/lib/shared/src/token/counter.ts @@ -10,14 +10,47 @@ export interface TokenCounterUtils { decode(encoded: number[]): string countTokens(text: string): number countPromptString(text: PromptString): number - getMessagesTokenCount(messages: Message[]): number - getTokenCountForMessage(message: Message): number + getMessagesTokenCount(messages: (Message | undefined)[]): number + getTokenCountForMessage(message: Message | undefined): number +} + +let _useFakeTokenCounterUtils: TokenCounterUtils | undefined + +/** + * @internal For testing only. Importing the weights for the token counter is slow and is not + * necessary for most tests. + */ +export function useFakeTokenCounterUtils(): void { + _useFakeTokenCounterUtils = { + encode(text) { + return text.split(' ').map(word => word.length) + }, + decode(encoded) { + return encoded.map(n => ' '.repeat(n)).join('') + }, + countTokens(text) { + return text.split(' ').length + }, + countPromptString(text) { + return text.split(' ').length + }, + getMessagesTokenCount(messages) { + return messages.reduce((acc, m) => acc + (m?.text?.split(' ').length ?? 0), 0) + }, + getTokenCountForMessage(message) { + return message?.text?.split(' ').length ?? 0 + }, + } } /** * Get the tokenizer, which is lazily-loaded it because it requires reading ~1 MB of tokenizer data. */ export async function getTokenCounterUtils(): Promise { + if (_useFakeTokenCounterUtils) { + return _useFakeTokenCounterUtils + } + // This could have been implemented in a separate file that is wholly async-imported, but that // carries too much risk of accidental non-async importing. if (!_tokenCounterUtilsPromise) { diff --git a/vscode/src/chat/agentic/CodyTool.ts b/vscode/src/chat/agentic/CodyTool.ts index b35acbb8fafe..ebd46fc54ef3 100644 --- a/vscode/src/chat/agentic/CodyTool.ts +++ b/vscode/src/chat/agentic/CodyTool.ts @@ -97,7 +97,7 @@ class SearchTool extends CodyTool { } constructor( - private contextRetriever: ContextRetriever, + private contextRetriever: Pick, private span: Span ) { super() @@ -133,6 +133,9 @@ class SearchTool extends CodyTool { } } -export function getCodyTools(contextRetriever: ContextRetriever, span: Span): CodyTool[] { +export function getCodyTools( + contextRetriever: Pick, + span: Span +): CodyTool[] { return [new SearchTool(contextRetriever, span), new CliTool(), new FileTool()] } diff --git a/vscode/src/chat/agentic/DeepCody.ts b/vscode/src/chat/agentic/DeepCody.ts index bfaf7a724572..c37c9d2fc503 100644 --- a/vscode/src/chat/agentic/DeepCody.ts +++ b/vscode/src/chat/agentic/DeepCody.ts @@ -49,7 +49,7 @@ export class DeepCodyAgent { constructor( private readonly chatBuilder: ChatBuilder, - private readonly chatClient: ChatClient, + private readonly chatClient: Pick, private readonly tools: CodyTool[], mentions: ContextItem[] = [] ) { diff --git a/vscode/src/chat/chat-view/ChatController.test.ts b/vscode/src/chat/chat-view/ChatController.test.ts index 2af8948eeb1b..922359387be6 100644 --- a/vscode/src/chat/chat-view/ChatController.test.ts +++ b/vscode/src/chat/chat-view/ChatController.test.ts @@ -1,8 +1,334 @@ -import { describe, expect, it, vi } from 'vitest' +import { + AUTH_STATUS_FIXTURE_AUTHED, + CLIENT_CAPABILITIES_FIXTURE, + type CompletionGeneratorValue, + FIXTURE_MODEL, + type Guardrails, + PromptString, + errorToChatError, + mockAuthStatus, + mockClientCapabilities, + mockResolvedConfig, + modelsService, + ps, + useFakeTokenCounterUtils, +} from '@sourcegraph/cody-shared' +import { Observable } from 'observable-fns' +import { beforeAll, beforeEach, describe, expect, it, test, vi } from 'vitest' import { Uri } from 'vscode' +import { URI } from 'vscode-uri' +import * as featureFlagProviderModule from '../../../../lib/shared/src/experimentation/FeatureFlagProvider' +import type { VSCodeEditor } from '../../editor/vscode-editor' +import type { ExtensionClient } from '../../extension-client' +import * as githubRepoMetadataModule from '../../repository/githubRepoMetadata' +import { mockLocalStorage } from '../../services/LocalStorageProvider' +import type { ExtensionMessage } from '../protocol' +import { ChatController, type ChatControllerOptions } from './ChatController' import { manipulateWebviewHTML } from './ChatController' -vi.mock('../../services/AuthProvider', () => ({})) +describe('ChatController', () => { + beforeAll(() => { + useFakeTokenCounterUtils() + }) + + const mockChatClient = { + chat: vi.fn(), + } satisfies ChatControllerOptions['chatClient'] + + const mockContextRetriever = { + retrieveContext: vi.fn(), + } satisfies ChatControllerOptions['contextRetriever'] + + const mockEditor: VSCodeEditor = {} as any + const mockExtensionClient: Pick = { + capabilities: {}, + } + const mockGuardrails: Guardrails = {} as any + + vi.spyOn(featureFlagProviderModule.featureFlagProvider, 'evaluatedFeatureFlag').mockReturnValue( + Observable.of(true) + ) + + vi.spyOn( + githubRepoMetadataModule, + 'publicRepoMetadataIfAllWorkspaceReposArePublic', + 'get' + ).mockReturnValue(Observable.of({ isPublic: false, repoMetadata: undefined })) + + const mockNowDate = new Date(123456) + + let chatController: ChatController + beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers() + mockAuthStatus(AUTH_STATUS_FIXTURE_AUTHED) + mockResolvedConfig({ auth: { serverEndpoint: AUTH_STATUS_FIXTURE_AUTHED.endpoint } }) + mockClientCapabilities(CLIENT_CAPABILITIES_FIXTURE) + mockLocalStorage() + vi.setSystemTime(mockNowDate) + + vi.spyOn(modelsService, 'getDefaultModel').mockReturnValue(Observable.of(FIXTURE_MODEL)) + + chatController = new ChatController({ + extensionUri: URI.file('x'), + chatClient: mockChatClient, + editor: mockEditor, + extensionClient: mockExtensionClient, + guardrails: mockGuardrails, + contextRetriever: mockContextRetriever, + chatIntentAPIClient: null, + }) + }) + + test('send, followup, and edit', async () => { + const postMessageSpy = vi.spyOn(chatController as any, 'postMessage') + const addBotMessageSpy = vi.spyOn(chatController as any, 'addBotMessage') + + mockChatClient.chat.mockReturnValue( + (async function* () { + yield { type: 'change', text: 'Test reply 1' } + yield { type: 'complete', text: 'Test reply 1' } + })() satisfies AsyncGenerator + ) + mockContextRetriever.retrieveContext.mockResolvedValue([]) + + // Send the first message in a new chat. + await chatController.handleUserMessageSubmission({ + requestID: '1', + inputText: PromptString.unsafe_fromUserQuery('Test input'), + mentions: [], + editorState: null, + signal: new AbortController().signal, + source: 'chat', + }) + expect(postMessageSpy.mock.calls.at(0)?.at(0)).toStrictEqual< + Extract + >({ + type: 'transcript', + isMessageInProgress: true, + chatID: mockNowDate.toUTCString(), + messages: [ + { + speaker: 'human', + text: 'Test input', + model: undefined, + error: undefined, + editorState: null, + contextFiles: undefined, + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: undefined, + contextFiles: undefined, + }, + ], + }) + + // Make sure it was sent and the reply was received. + await vi.runOnlyPendingTimersAsync() + expect(mockChatClient.chat).toBeCalledTimes(1) + expect(addBotMessageSpy).toHaveBeenCalledWith('1', ps`Test reply 1`, 'my-model') + expect(postMessageSpy.mock.calls.at(3)?.at(0)).toStrictEqual< + Extract + >({ + type: 'transcript', + isMessageInProgress: true, + chatID: mockNowDate.toUTCString(), + messages: [ + { + speaker: 'human', + text: 'Test input', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: 'Test reply 1', + contextFiles: undefined, + }, + ], + }) + + // Send a followup. + vi.clearAllMocks() + mockChatClient.chat.mockReturnValue( + (async function* () { + yield { type: 'change', text: 'Test reply 2' } + yield { type: 'complete', text: 'Test reply 2' } + })() satisfies AsyncGenerator + ) + await chatController.handleUserMessageSubmission({ + requestID: '2', + inputText: PromptString.unsafe_fromUserQuery('Test followup'), + mentions: [], + editorState: null, + signal: new AbortController().signal, + source: 'chat', + }) + await vi.runOnlyPendingTimersAsync() + expect(mockChatClient.chat).toBeCalledTimes(1) + expect(addBotMessageSpy).toHaveBeenCalledWith('2', ps`Test reply 2`, 'my-model') + expect(postMessageSpy.mock.calls.at(3)?.at(0)).toStrictEqual< + Extract + >({ + type: 'transcript', + isMessageInProgress: false, + chatID: mockNowDate.toUTCString(), + messages: [ + { + speaker: 'human', + text: 'Test input', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: 'Test reply 1', + contextFiles: undefined, + }, + { + speaker: 'human', + text: 'Test followup', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: 'Test reply 2', + contextFiles: undefined, + }, + ], + }) + + // Now try editing the message. + vi.clearAllMocks() + mockChatClient.chat.mockReturnValue( + (async function* () { + yield { type: 'change', text: 'Test reply 3' } + yield { type: 'complete', text: 'Test reply 3' } + })() satisfies AsyncGenerator + ) + await chatController.handleEdit({ + requestID: '3', + index: 2, + text: PromptString.unsafe_fromUserQuery('Test edit'), + contextFiles: [], + editorState: null, + }) + await vi.runOnlyPendingTimersAsync() + expect(mockChatClient.chat).toBeCalledTimes(1) + expect(addBotMessageSpy).toHaveBeenCalledWith('3', ps`Test reply 3`, 'my-model') + expect(postMessageSpy.mock.calls.at(3)?.at(0)).toStrictEqual< + Extract + >({ + type: 'transcript', + isMessageInProgress: false, + chatID: mockNowDate.toUTCString(), + messages: [ + { + speaker: 'human', + text: 'Test input', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: 'Test reply 1', + contextFiles: undefined, + }, + { + speaker: 'human', + text: 'Test edit', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: 'my-model', + error: undefined, + editorState: undefined, + text: 'Test reply 3', + contextFiles: undefined, + }, + ], + }) + }) + + test('send error', async () => { + const postMessageSpy = vi.spyOn(chatController as any, 'postMessage') + const addBotMessageSpy = vi.spyOn(chatController as any, 'addBotMessage') + + mockChatClient.chat.mockReturnValue( + (async function* () { + yield { type: 'error', error: new Error('my-error') } + })() satisfies AsyncGenerator + ) + mockContextRetriever.retrieveContext.mockResolvedValue([]) + + // Send the first message in a new chat. + await chatController.handleUserMessageSubmission({ + requestID: '1', + inputText: PromptString.unsafe_fromUserQuery('Test input'), + mentions: [], + editorState: null, + signal: new AbortController().signal, + source: 'chat', + }) + await vi.runOnlyPendingTimersAsync() + expect(mockChatClient.chat).toBeCalledTimes(1) + expect(addBotMessageSpy).toHaveBeenCalledWith('1', ps``, 'my-model') + expect(postMessageSpy.mock.calls.at(3)?.at(0)).toStrictEqual< + Extract + >({ + type: 'transcript', + isMessageInProgress: false, + chatID: mockNowDate.toUTCString(), + messages: [ + { + speaker: 'human', + text: 'Test input', + model: undefined, + error: undefined, + editorState: null, + contextFiles: [], + }, + { + speaker: 'assistant', + model: undefined, + error: errorToChatError(new Error('my-error')), + editorState: undefined, + text: undefined, + contextFiles: undefined, + }, + ], + }) + }) +}) describe('manipulateWebviewHTML', () => { const options = { diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts index ed4456cc3811..f48817e50307 100644 --- a/vscode/src/chat/chat-view/ChatController.ts +++ b/vscode/src/chat/chat-view/ChatController.ts @@ -131,14 +131,14 @@ import { getChatPanelTitle } from './chat-helpers' import { type HumanInput, getPriorityContext } from './context' import { DefaultPrompter, type PromptInfo } from './prompt' -interface ChatControllerOptions { +export interface ChatControllerOptions { extensionUri: vscode.Uri - chatClient: ChatClient + chatClient: Pick - contextRetriever: ContextRetriever + contextRetriever: Pick chatIntentAPIClient: ChatIntentAPIClient | null - extensionClient: ExtensionClient + extensionClient: Pick editor: VSCodeEditor guardrails: Guardrails @@ -180,12 +180,12 @@ export interface ChatSession { export class ChatController implements vscode.Disposable, vscode.WebviewViewProvider, ChatSession { private chatBuilder: ChatBuilder - private readonly chatClient: ChatClient + private readonly chatClient: ChatControllerOptions['chatClient'] - private readonly contextRetriever: ContextRetriever + private readonly contextRetriever: ChatControllerOptions['contextRetriever'] - private readonly editor: VSCodeEditor - private readonly extensionClient: ExtensionClient + private readonly editor: ChatControllerOptions['editor'] + private readonly extensionClient: ChatControllerOptions['extensionClient'] private readonly guardrails: Guardrails private readonly startTokenReceiver: typeof startTokenReceiver | undefined @@ -650,161 +650,173 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv tokenCounterUtils ) - tracer.startActiveSpan('chat.submit.firstToken', async (firstTokenSpan): Promise => { - if (inputText.toString().match(/^\/reset$/)) { - span.addEvent('clearAndRestartSession') - span.end() - return this.clearAndRestartSession() - } - - this.chatBuilder.addHumanMessage({ - text: inputText, - editorState, - intent: detectedIntent, - }) - await this.saveSession() - signal.throwIfAborted() - - this.postEmptyMessageInProgress(model) + await tracer.startActiveSpan( + 'chat.submit.firstToken', + async (firstTokenSpan): Promise => { + if (inputText.toString().match(/^\/reset$/)) { + span.addEvent('clearAndRestartSession') + span.end() + return this.clearAndRestartSession() + } - // All mentions we receive are either source=initial or source=user. If the caller - // forgot to set the source, assume it's from the user. - mentions = mentions.map(m => (m.source ? m : { ...m, source: ContextItemSource.User })) + this.chatBuilder.addHumanMessage({ + text: inputText, + editorState, + intent: detectedIntent, + }) + await this.saveSession() + signal.throwIfAborted() - const contextAlternatives = await this.computeContext( - { text: inputText, mentions }, - requestID, - editorState, - span, - signal - ) - signal.throwIfAborted() - const corpusContext = contextAlternatives[0].items + this.postEmptyMessageInProgress(model) - const repositoryMentioned = mentions.find(contextItem => - ['repository', 'tree'].includes(contextItem.type) - ) + // All mentions we receive are either source=initial or source=user. If the caller + // forgot to set the source, assume it's from the user. + mentions = mentions.map(m => + m.source ? m : { ...m, source: ContextItemSource.User } + ) - // We are checking the feature flag here to log non-undefined intent only if the feature flag is on - let intent: ChatMessage['intent'] | undefined = this.featureCodyExperimentalOneBox - ? detectedIntent - : undefined + const contextAlternatives = await this.computeContext( + { text: inputText, mentions }, + requestID, + editorState, + span, + signal + ) + signal.throwIfAborted() + const corpusContext = contextAlternatives[0].items - let intentScores: { intent: string; score: number }[] | undefined | null = this - .featureCodyExperimentalOneBox - ? detectedIntentScores - : undefined + const repositoryMentioned = mentions.find(contextItem => + ['repository', 'tree'].includes(contextItem.type) + ) - const userSpecifiedIntent = this.featureCodyExperimentalOneBox - ? manuallySelectedIntent && detectedIntent + // We are checking the feature flag here to log non-undefined intent only if the feature flag is on + let intent: ChatMessage['intent'] | undefined = this.featureCodyExperimentalOneBox ? detectedIntent - : 'auto' - : undefined - - if (this.featureCodyExperimentalOneBox && repositoryMentioned) { - const inputTextWithoutContextChips = editorState - ? PromptString.unsafe_fromUserQuery( - inputTextWithoutContextChipsFromPromptEditorState(editorState) - ) - : inputText - - const finalIntentDetectionResponse = detectedIntent - ? { intent: detectedIntent, allScores: detectedIntentScores } - : await this.detectChatIntent({ - requestID, - text: inputTextWithoutContextChips.toString(), - }) - .then(async response => { - signal.throwIfAborted() - this.chatBuilder.setLastMessageIntent(response?.intent) - this.postViewTranscript() - return response + : undefined + + let intentScores: { intent: string; score: number }[] | undefined | null = this + .featureCodyExperimentalOneBox + ? detectedIntentScores + : undefined + + const userSpecifiedIntent = this.featureCodyExperimentalOneBox + ? manuallySelectedIntent && detectedIntent + ? detectedIntent + : 'auto' + : undefined + + if (this.featureCodyExperimentalOneBox && repositoryMentioned) { + const inputTextWithoutContextChips = editorState + ? PromptString.unsafe_fromUserQuery( + inputTextWithoutContextChipsFromPromptEditorState(editorState) + ) + : inputText + + const finalIntentDetectionResponse = detectedIntent + ? { intent: detectedIntent, allScores: detectedIntentScores } + : await this.detectChatIntent({ + requestID, + text: inputTextWithoutContextChips.toString(), }) - .catch(() => undefined) - intent = finalIntentDetectionResponse?.intent - intentScores = finalIntentDetectionResponse?.allScores - signal.throwIfAborted() - if (intent === 'search') { + .then(async response => { + signal.throwIfAborted() + this.chatBuilder.setLastMessageIntent(response?.intent) + this.postViewTranscript() + return response + }) + .catch(() => undefined) + intent = finalIntentDetectionResponse?.intent + intentScores = finalIntentDetectionResponse?.allScores + signal.throwIfAborted() + if (intent === 'search') { + telemetryEvents['cody.chat-question/executed'].record( + { + ...telemetryProperties, + context: corpusContext, + userSpecifiedIntent, + detectedIntent: intent, + detectedIntentScores: intentScores, + }, + { current: span, firstToken: firstTokenSpan, addMetadata: true } + ) + + return await this.handleSearchIntent({ + context: corpusContext, + signal, + contextAlternatives, + }) + } + } + + // Experimental Feature: Deep Cody + if (model === DeepCodyAgent.ModelRef) { + const agenticContext = await new DeepCodyAgent( + this.chatBuilder, + this.chatClient, + getCodyTools(this.contextRetriever, span), + corpusContext + ).getContext(signal) + corpusContext.push(...agenticContext) + } + + const { explicitMentions, implicitMentions } = getCategorizedMentions(corpusContext) + + const prompter = new DefaultPrompter( + explicitMentions, + implicitMentions, + command !== undefined + ) + + try { + const { prompt, context } = await this.buildPrompt( + prompter, + signal, + requestID, + authStatus.codyApiVersion, + contextAlternatives + ) + telemetryEvents['cody.chat-question/executed'].record( { ...telemetryProperties, - context: corpusContext, + context, userSpecifiedIntent, detectedIntent: intent, detectedIntentScores: intentScores, }, - { current: span, firstToken: firstTokenSpan, addMetadata: true } + { + addMetadata: true, + current: span, + firstToken: firstTokenSpan, + } ) - return await this.handleSearchIntent({ - context: corpusContext, - signal, - contextAlternatives, - }) - } - } - - // Experimental Feature: Deep Cody - if (model === DeepCodyAgent.ModelRef) { - const agenticContext = await new DeepCodyAgent( - this.chatBuilder, - this.chatClient, - getCodyTools(this.contextRetriever, span), - corpusContext - ).getContext(signal) - corpusContext.push(...agenticContext) - } - - const { explicitMentions, implicitMentions } = getCategorizedMentions(corpusContext) - - const prompter = new DefaultPrompter( - explicitMentions, - implicitMentions, - command !== undefined - ) - - try { - const { prompt, context } = await this.buildPrompt( - prompter, - signal, - requestID, - authStatus.codyApiVersion, - contextAlternatives - ) - - telemetryEvents['cody.chat-question/executed'].record( - { - ...telemetryProperties, - context, - userSpecifiedIntent, - detectedIntent: intent, - detectedIntentScores: intentScores, - }, - { - addMetadata: true, - current: span, - firstToken: firstTokenSpan, - } - ) - - signal.throwIfAborted() - this.streamAssistantResponse(requestID, prompt, model, span, firstTokenSpan, signal) - } catch (error) { - if (isAbortErrorOrSocketHangUp(error as Error)) { - return - } - if (isRateLimitError(error) || isContextWindowLimitError(error)) { - this.postError(error, 'transcript') - } else { - this.postError( - isError(error) - ? error - : new Error(`Error generating assistant response: ${error}`) + signal.throwIfAborted() + this.streamAssistantResponse( + requestID, + prompt, + model, + span, + firstTokenSpan, + signal ) + } catch (error) { + if (isAbortErrorOrSocketHangUp(error as Error)) { + return + } + if (isRateLimitError(error) || isContextWindowLimitError(error)) { + this.postError(error, 'transcript') + } else { + this.postError( + isError(error) + ? error + : new Error(`Error generating assistant response: ${error}`) + ) + } + recordErrorToSpan(span, error as Error) } - recordErrorToSpan(span, error as Error) } - }) + ) }) } @@ -944,8 +956,10 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv * Removes any existing messages from the provided index, * before submitting the replacement text as a new question. * When no index is provided, default to the last human message. + * + * @internal Public for testing only. */ - private async handleEdit({ + public async handleEdit({ requestID, text, index, @@ -1713,7 +1727,7 @@ export function revealWebviewViewOrPanel(viewOrPanel: vscode.WebviewView | vscod * Set HTML for webview (panel) & webview view (sidebar) */ async function addWebviewViewHTML( - extensionClient: ExtensionClient, + extensionClient: Pick, extensionUri: vscode.Uri, view: vscode.WebviewView | vscode.WebviewPanel ): Promise {