From dfc1048417e161e29fb53766e8ed6d88db2d6cc9 Mon Sep 17 00:00:00 2001 From: Paul Jankowski <8bittitan@gmail.com> Date: Wed, 17 Jun 2026 12:55:36 -0400 Subject: [PATCH] fix(askai): Fix new conversation causing thread depth errors (#2900) --- AGENTS.md | 2 +- .../src/DocSearchAskAiModal.tsx | 54 +- .../src/Sidepanel/Sidepanel.tsx | 20 +- .../src/__tests__/useAskAi.test.ts | 609 ++++++++++++++++++ .../src/__tests__/useAskAi.test.tsx | 305 --------- .../src/types/StoredDocSearchHit.ts | 1 + packages/docsearch-react/src/useAskAi.ts | 81 ++- packages/docsearch-react/src/utils/ai.ts | 3 +- 8 files changed, 730 insertions(+), 345 deletions(-) create mode 100644 packages/docsearch-react/src/__tests__/useAskAi.test.ts delete mode 100644 packages/docsearch-react/src/__tests__/useAskAi.test.tsx diff --git a/AGENTS.md b/AGENTS.md index 9b3bf35d..9a67c1af 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,7 +77,7 @@ bun run pw:run:webkit ### Imports -Imports must be ordered alphabetically with newlines between groups: +Imported modules must be ordered alphabetically with newlines between groups: 1. Built-in modules 2. External dependencies diff --git a/packages/docsearch-react/src/DocSearchAskAiModal.tsx b/packages/docsearch-react/src/DocSearchAskAiModal.tsx index 695f5d01..df050a1a 100644 --- a/packages/docsearch-react/src/DocSearchAskAiModal.tsx +++ b/packages/docsearch-react/src/DocSearchAskAiModal.tsx @@ -28,6 +28,7 @@ import type { InternalDocSearchHit, OnAskAiFeedback, StoredAskAiMessage, + StoredAskAiState, SuggestedQuestionHit, } from './types'; import { type AskAiState } from './types/AskiAi'; @@ -172,17 +173,28 @@ export function DocSearchAskAiModal({ const [stoppedStream, setStoppedStream] = React.useState(false); - const { messages, status, setMessages, sendMessage, stopAskAiStreaming, askAiError, sendFeedback, conversations } = - useAskAi({ - assistantId: askAiConfigurationId, - apiKey: askAiConfig?.apiKey || apiKey, - appId: askAiConfig?.appId || appId, - indexName: askAiConfig?.indexName || defaultIndexName, - searchParameters: askAiSearchParameters, - tools, - memory: props.memory, - indices: askAiConfig?.indices, - }); + const { + chatId, + messages, + status, + setMessages, + sendMessage, + stopAskAiStreaming, + askAiError, + sendFeedback, + conversations, + startNewConversation, + restoreConversation, + } = useAskAi({ + assistantId: askAiConfigurationId, + apiKey: askAiConfig?.apiKey || apiKey, + appId: askAiConfig?.appId || appId, + indexName: askAiConfig?.indexName || defaultIndexName, + searchParameters: askAiSearchParameters, + tools, + memory: props.memory, + indices: askAiConfig?.indices, + }); const prevStatus = React.useRef(status); React.useEffect(() => { @@ -200,12 +212,12 @@ export function DocSearchAskAiModal({ for (const part of messages[0].parts) { if (part.type === 'text') { - conversations.add(buildDummyAskAiHit(part.text, messages)); + conversations.add(buildDummyAskAiHit(part.text, messages, chatId)); } } } prevStatus.current = status; - }, [status, messages, conversations, disableUserPersonalization, stoppedStream]); + }, [status, messages, conversations, disableUserPersonalization, stoppedStream, chatId]); // Check if there's a thread depth error (AI-217) const hasThreadDepthError = React.useMemo(() => { @@ -392,14 +404,19 @@ export function DocSearchAskAiModal({ useRefreshOnInitialQuery({ initialQuery, inputRef, refresh }); + const hasCurrentMessages = messages.length > 0; + // Refresh the autocomplete results when ask ai is toggled off // helps return to the previous ac state and start screen React.useEffect(() => { if (!isAskAiActive) { autocomplete.refresh(); - setMessages([]); + + if (hasCurrentMessages) { + startNewConversation(); + } } - }, [isAskAiActive, autocomplete, setMessages]); + }, [isAskAiActive, autocomplete, startNewConversation, hasCurrentMessages]); // Track external state in order to manage internal askAiState React.useEffect(() => { @@ -413,7 +430,7 @@ export function DocSearchAskAiModal({ }; const handleNewConversation = (): void => { - setMessages([]); + startNewConversation(); setAskAiState('new-conversation'); }; @@ -510,10 +527,11 @@ export function DocSearchAskAiModal({ onItemClick={(item, event) => { if (item.type === 'askAI' && item.query) { if (item.anchor === 'stored' && 'messages' in item) { - setMessages(item.messages as any); + const hitMessages = item.messages as StoredAskAiMessage[]; + restoreConversation(hitMessages, (item as StoredAskAiState).chatId); const initialMessage: InitialAskAiMessage = { query: item.query, - messageId: (item.messages as StoredAskAiMessage[])[0].id, + messageId: hitMessages[0].id, }; if (interceptAskAiEvent?.(initialMessage)) { diff --git a/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx b/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx index eddbd086..66955445 100644 --- a/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx +++ b/packages/docsearch-react/src/Sidepanel/Sidepanel.tsx @@ -176,16 +176,18 @@ function SidepanelInner( const searchClient = useSearchClient(appId, apiKey, setSidepanelSearchClient); const { + chatId, status, sendMessage, stopAskAiStreaming, isStreaming, exchanges, - setMessages, conversations, messages, sendFeedback, askAiError, + startNewConversation, + restoreConversation, } = useAskAi({ appId, indexName, @@ -213,13 +215,13 @@ function SidepanelInner( }; const handleStartNewConversation = (): void => { - setMessages([]); + startNewConversation(); setSidepanelState('new-conversation'); }; const handleSelectQuestion = (question: SuggestedQuestionHit): void => { setStoppedStreaming(false); - setMessages([]); + startNewConversation(); sendMessage( { text: question.question }, { @@ -239,14 +241,14 @@ function SidepanelInner( const handleSelectConversation = React.useCallback( (conversation: StoredAskAiState): void => { if (conversation.messages) { - setMessages(conversation.messages); + restoreConversation(conversation.messages, conversation.chatId); } else if (conversation.query) { sendMessage({ text: conversation.query }); } setSidepanelState('conversation'); }, - [sendMessage, setMessages], + [sendMessage, restoreConversation], ); useManageSidepanelLayout({ @@ -290,13 +292,13 @@ function SidepanelInner( for (const part of messages[0].parts) { if (part.type === 'text') { - conversations.add(buildDummyAskAiHit(part.text, messages)); + conversations.add(buildDummyAskAiHit(part.text, messages, chatId)); } } } prevStatus.current = status; - }, [conversations, status, messages, stoppedStreaming]); + }, [conversations, status, messages, stoppedStreaming, chatId]); React.useEffect(() => { function setFullViewportHeight(): void { @@ -327,7 +329,7 @@ function SidepanelInner( if (selectedConversation) { handleSelectConversation(selectedConversation); } else { - setMessages([]); + startNewConversation(); sendMessage( { text: initialMessage.query, @@ -342,7 +344,7 @@ function SidepanelInner( ); setSidepanelState('conversation'); } - }, [initialMessage, sendMessage, conversations, handleSelectConversation, setMessages]); + }, [initialMessage, sendMessage, conversations, handleSelectConversation, startNewConversation]); // Autofocus the prompt input when the sidepanel opens and blur it when // it closes. Disabled on mobile because focusing the textarea triggers the diff --git a/packages/docsearch-react/src/__tests__/useAskAi.test.ts b/packages/docsearch-react/src/__tests__/useAskAi.test.ts new file mode 100644 index 00000000..fc40f5f4 --- /dev/null +++ b/packages/docsearch-react/src/__tests__/useAskAi.test.ts @@ -0,0 +1,609 @@ +/* eslint-disable max-classes-per-file */ +import { act, renderHook } from '@testing-library/react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { useAskAi } from '../useAskAi'; + +interface ToolCall { + input: unknown; + toolCallId: string; + toolName: string; +} + +interface ChatMessage { + id: string; + role: 'assistant' | 'user'; + parts: Array<{ type: 'text'; text: string }>; +} + +interface ChatOptions { + id?: string; + messages?: ChatMessage[]; + onToolCall: (params: { toolCall: ToolCall }) => unknown; + transport: { + options: { + headers?: Record; + body?: Record; + }; + }; +} + +interface ChatInstance { + id?: string; + messages?: ChatMessage[]; + options: ChatOptions; + sendMessage: (message: unknown) => void; +} + +interface UseChatOptions { + chat: ChatInstance; +} + +type CustomOnToolCallParams = ToolCall & { + addToolOutput: (props: { output: unknown }) => Promise; +}; + +const mocks = vi.hoisted(() => ({ + addToolOutput: vi.fn(), + setMessages: vi.fn(), + sendMessage: vi.fn(), + useChat: vi.fn(), + generateId: vi.fn(), +})); + +vi.mock('@ai-sdk/react', () => ({ + Chat: class Chat { + id?: string; + messages?: ChatMessage[]; + options: ChatOptions; + + constructor(options: ChatOptions) { + this.id = options.id; + this.messages = options.messages; + this.options = options; + } + + sendMessage(message: unknown): void { + mocks.sendMessage({ chatId: this.id, message }); + } + }, + useChat: mocks.useChat, +})); + +vi.mock('ai', () => ({ + DefaultChatTransport: class DefaultChatTransport { + options: unknown; + + constructor(options: unknown) { + this.options = options; + } + }, + lastAssistantMessageIsCompleteWithToolCalls: vi.fn(() => false), + generateId: mocks.generateId, +})); + +interface SendMessageCall { + chatId?: string; + message: unknown; +} + +describe('useAskAi', () => { + let chatOptions: ChatOptions | undefined; + let chatOptionsHistory: ChatOptions[] = []; + let sendMessageCalls: SendMessageCall[] = []; + + function getOnToolCall(): ChatOptions['onToolCall'] { + if (!chatOptions) { + throw new Error('useChat was not initialized'); + } + + return chatOptions.onToolCall; + } + + function getTransportHeaders(): Record { + if (!chatOptions) { + throw new Error('useChat was not initialized'); + } + + return chatOptions.transport.options.headers ?? {}; + } + + function getTransportBody(): Record { + if (!chatOptions) { + throw new Error('useChat was not initialized'); + } + + return chatOptions.transport.options.body ?? {}; + } + + beforeEach(() => { + vi.clearAllMocks(); + + let idCounter = 0; + mocks.generateId.mockImplementation(() => { + idCounter += 1; + return `generated-id-${idCounter}`; + }); + + chatOptions = undefined; + chatOptionsHistory = []; + sendMessageCalls = []; + mocks.sendMessage.mockImplementation((call: SendMessageCall) => { + sendMessageCalls.push(call); + }); + mocks.useChat.mockImplementation(({ chat }: UseChatOptions) => { + chatOptions = chat.options; + chatOptionsHistory.push(chat.options); + + return { + addToolOutput: mocks.addToolOutput, + error: undefined, + messages: chat.messages ?? [], + sendMessage: chat.sendMessage, + setMessages: mocks.setMessages, + status: 'ready', + stop: vi.fn(), + }; + }); + }); + + it('forwards custom tool output to useChat addToolOutput', async () => { + let addToolOutput: CustomOnToolCallParams['addToolOutput'] | undefined; + const onToolCall = vi.fn((params: CustomOnToolCallParams) => { + addToolOutput = params.addToolOutput; + }); + + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: { + customAction: { + onToolCall, + render: () => 'Custom action complete', + }, + }, + }), + ); + + getOnToolCall()({ + toolCall: { + input: { value: 'input value' }, + toolCallId: 'tool-call-id', + toolName: 'customAction', + }, + }); + + expect(onToolCall).toHaveBeenCalledWith( + expect.objectContaining({ + input: { value: 'input value' }, + toolCallId: 'tool-call-id', + toolName: 'customAction', + }), + ); + + if (!addToolOutput) { + throw new Error('addToolOutput was not provided to the custom tool'); + } + + await addToolOutput({ output: { result: 'output value' } }); + + expect(mocks.addToolOutput).toHaveBeenCalledWith({ + output: { result: 'output value' }, + tool: 'customAction', + toolCallId: 'tool-call-id', + }); + }); + + it('does not wait for custom onToolCall to finish', () => { + const pendingToolCall = new Promise(() => {}); + const onToolCall = vi.fn(() => pendingToolCall); + + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: { + customAction: { + onToolCall, + render: () => 'Custom action complete', + }, + }, + }), + ); + + const result = getOnToolCall()({ + toolCall: { + input: { value: 'input value' }, + toolCallId: 'tool-call-id', + toolName: 'customAction', + }, + }); + + expect(onToolCall).toHaveBeenCalledTimes(1); + expect(result).toBeUndefined(); + }); + + it('sends the secure user token header when memory.userToken is provided', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + memory: { userToken: 'secure-user-token' }, + }), + ); + + expect(getTransportHeaders()).toMatchObject({ + 'x-algolia-secure-user-token': 'secure-user-token', + }); + }); + + it('omits the secure user token header when no memory token is provided', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token'); + }); + + it('sends an empty transport body when no search parameters or indices are provided', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + expect(getTransportBody()).toEqual({ algolia: {} }); + }); + + it('includes searchParameters under the algolia body when provided', () => { + const searchParameters = { + 'index-name': { distinct: false }, + }; + + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + searchParameters, + }), + ); + + expect(getTransportBody()).toEqual({ + algolia: { searchParameters }, + }); + }); + + it('includes indices under the algolia body when provided', () => { + const indices = [ + { + index: 'docsearch-markdown', + description: 'Use this to gather specific results.', + }, + ]; + + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + indices, + }), + ); + + expect(getTransportBody()).toEqual({ + algolia: { indices }, + }); + }); + + it('includes both searchParameters and indices under the algolia body when both are provided', () => { + const searchParameters = { + 'index-name': { distinct: false }, + }; + const indices = [ + { + index: 'docsearch-markdown', + description: 'Use this to gather specific results.', + }, + ]; + + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + searchParameters, + indices, + }), + ); + + expect(getTransportBody()).toEqual({ + algolia: { searchParameters, indices }, + }); + }); + + it('omits indices from the body when an empty indices array is provided', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + indices: [], + }), + ); + + expect(getTransportBody()).toEqual({ algolia: {} }); + }); + + describe('conversation id rotation', () => { + it('passes a generated id to useChat on mount', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + expect(chatOptions?.id).toBe('generated-id-1'); + }); + + it('rotates the chat id and clears messages on startNewConversation', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + const initialId = chatOptions?.id; + expect(initialId).toBe('generated-id-1'); + + act(() => { + result.current.startNewConversation(); + }); + + // Rotating the id recreates the `Chat` instance with no seeded messages, + // which is what clears the previous conversation. + expect(chatOptions?.messages).toBeUndefined(); + expect(chatOptions?.id).toBe('generated-id-2'); + expect(chatOptions?.id).not.toBe(initialId); + }); + + it('generates a fresh id for each new conversation', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + act(() => { + result.current.startNewConversation(); + }); + const secondId = chatOptions?.id; + + act(() => { + result.current.startNewConversation(); + }); + const thirdId = chatOptions?.id; + + expect(secondId).toBe('generated-id-2'); + expect(thirdId).toBe('generated-id-3'); + expect(secondId).not.toBe(thirdId); + }); + + it('restores a stored conversation using its persisted chat id', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + const storedMessages: ChatMessage[] = [ + { + id: 'message-1', + role: 'user', + parts: [{ type: 'text', text: 'stored question' }], + }, + ]; + + act(() => { + result.current.restoreConversation(storedMessages, 'stored-chat-id'); + }); + + // Restored messages are seeded through the `messages` option so they + // survive the `Chat` instance being recreated when the id changes. + expect(chatOptions?.messages).toEqual(storedMessages); + expect(chatOptions?.id).toBe('stored-chat-id'); + }); + + it('restores a stored conversation with a fresh id when no chat id is persisted', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + const storedMessages: ChatMessage[] = [ + { + id: 'message-1', + role: 'user', + parts: [{ type: 'text', text: 'stored question' }], + }, + ]; + + act(() => { + result.current.restoreConversation(storedMessages); + }); + + // The messages must still be seeded even when a fresh id is generated, + // otherwise the recreated `Chat` instance would render an empty thread. + expect(chatOptions?.messages).toEqual(storedMessages); + expect(chatOptions?.id).toBe('generated-id-2'); + }); + + it('does not seed initial messages on mount', () => { + renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + expect(chatOptions?.messages).toBeUndefined(); + }); + + it('clears seeded messages and rotates the id when starting a new conversation after a restore', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + const storedMessages: ChatMessage[] = [ + { + id: 'message-1', + role: 'user', + parts: [{ type: 'text', text: 'stored question' }], + }, + ]; + + act(() => { + result.current.restoreConversation(storedMessages, 'stored-chat-id'); + }); + + expect(chatOptions?.messages).toEqual(storedMessages); + + act(() => { + result.current.startNewConversation(); + }); + + // Starting a new conversation must drop the seeded messages so a later + // recreation does not re-hydrate the previous conversation. + expect(chatOptions?.messages).toBeUndefined(); + expect(chatOptions?.id).toBe('generated-id-2'); + }); + + it('sends the next message with the new chat id after startNewConversation', () => { + const { result } = renderHook(() => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + tools: {}, + }), + ); + + const initialId = chatOptions?.id; + expect(initialId).toBe('generated-id-1'); + + // Capture `sendMessage` from the current render, mirroring how a consumer + // (e.g. the Sidepanel) grabs it before rotating the conversation. + const sendMessage = result.current.sendMessage; + + act(() => { + result.current.startNewConversation(); + sendMessage({ + role: 'user', + parts: [{ type: 'text', text: 'new question' }], + }); + }); + + // The rotated id is what every later request must target. + expect(chatOptions?.id).toBe('generated-id-2'); + expect(chatOptions?.id).not.toBe(initialId); + + // Regression: the request must NOT go out on the stale (pre-rotation) id. + expect(sendMessageCalls).toHaveLength(1); + expect(sendMessageCalls[0].chatId).toBe('generated-id-2'); + expect(sendMessageCalls[0].chatId).not.toBe(initialId); + }); + + it('keeps conversation lifecycle callbacks stable when transport inputs change by reference', () => { + const { result, rerender } = renderHook( + ({ indices, searchParameters }) => + useAskAi({ + apiKey: 'api-key', + appId: 'app-id', + assistantId: 'assistant-id', + indexName: 'index-name', + searchParameters, + indices, + tools: {}, + }), + { + initialProps: { + searchParameters: { + 'index-name': { distinct: false }, + }, + indices: [{ index: 'index-name', description: 'Test index' }], + }, + }, + ); + + const startNewConversation = result.current.startNewConversation; + const restoreConversation = result.current.restoreConversation; + + rerender({ + searchParameters: { + 'index-name': { distinct: false }, + }, + indices: [{ index: 'index-name', description: 'Test index' }], + }); + + expect(result.current.startNewConversation).toBe(startNewConversation); + expect(result.current.restoreConversation).toBe(restoreConversation); + }); + }); +}); diff --git a/packages/docsearch-react/src/__tests__/useAskAi.test.tsx b/packages/docsearch-react/src/__tests__/useAskAi.test.tsx deleted file mode 100644 index 64ebe196..00000000 --- a/packages/docsearch-react/src/__tests__/useAskAi.test.tsx +++ /dev/null @@ -1,305 +0,0 @@ -import { renderHook } from '@testing-library/react'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; - -import { useAskAi } from '../useAskAi'; - -type ToolCall = { - input: unknown; - toolCallId: string; - toolName: string; -}; - -type ChatOptions = { - onToolCall: (params: { toolCall: ToolCall }) => unknown; - transport: { - options: { - headers?: Record; - body?: Record; - }; - }; -}; - -type CustomOnToolCallParams = ToolCall & { - addToolOutput: (props: { output: unknown }) => Promise; -}; - -const mocks = vi.hoisted(() => ({ - addToolOutput: vi.fn(), - useChat: vi.fn(), -})); - -vi.mock('@ai-sdk/react', () => ({ - useChat: mocks.useChat, -})); - -vi.mock('ai', () => ({ - DefaultChatTransport: class DefaultChatTransport { - options: unknown; - - constructor(options: unknown) { - this.options = options; - } - }, - lastAssistantMessageIsCompleteWithToolCalls: vi.fn(() => false), -})); - -describe('useAskAi', () => { - let chatOptions: ChatOptions | undefined; - - function getOnToolCall(): ChatOptions['onToolCall'] { - if (!chatOptions) { - throw new Error('useChat was not initialized'); - } - - return chatOptions.onToolCall; - } - - function getTransportHeaders(): Record { - if (!chatOptions) { - throw new Error('useChat was not initialized'); - } - - return chatOptions.transport.options.headers ?? {}; - } - - function getTransportBody(): Record { - if (!chatOptions) { - throw new Error('useChat was not initialized'); - } - - return chatOptions.transport.options.body ?? {}; - } - - beforeEach(() => { - vi.clearAllMocks(); - - chatOptions = undefined; - mocks.useChat.mockImplementation((options: ChatOptions) => { - chatOptions = options; - - return { - addToolOutput: mocks.addToolOutput, - error: undefined, - messages: [], - sendMessage: vi.fn(), - setMessages: vi.fn(), - status: 'ready', - stop: vi.fn(), - }; - }); - }); - - it('forwards custom tool output to useChat addToolOutput', async () => { - let addToolOutput: CustomOnToolCallParams['addToolOutput'] | undefined; - const onToolCall = vi.fn((params: CustomOnToolCallParams) => { - addToolOutput = params.addToolOutput; - }); - - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: { - customAction: { - onToolCall, - render: () => 'Custom action complete', - }, - }, - }), - ); - - getOnToolCall()({ - toolCall: { - input: { value: 'input value' }, - toolCallId: 'tool-call-id', - toolName: 'customAction', - }, - }); - - expect(onToolCall).toHaveBeenCalledWith( - expect.objectContaining({ - input: { value: 'input value' }, - toolCallId: 'tool-call-id', - toolName: 'customAction', - }), - ); - - if (!addToolOutput) { - throw new Error('addToolOutput was not provided to the custom tool'); - } - - await addToolOutput({ output: { result: 'output value' } }); - - expect(mocks.addToolOutput).toHaveBeenCalledWith({ - output: { result: 'output value' }, - tool: 'customAction', - toolCallId: 'tool-call-id', - }); - }); - - it('does not wait for custom onToolCall to finish', () => { - const pendingToolCall = new Promise(() => {}); - const onToolCall = vi.fn(() => pendingToolCall); - - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: { - customAction: { - onToolCall, - render: () => 'Custom action complete', - }, - }, - }), - ); - - const result = getOnToolCall()({ - toolCall: { - input: { value: 'input value' }, - toolCallId: 'tool-call-id', - toolName: 'customAction', - }, - }); - - expect(onToolCall).toHaveBeenCalledTimes(1); - expect(result).toBeUndefined(); - }); - - it('sends the secure user token header when memory.userToken is provided', () => { - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - memory: { userToken: 'secure-user-token' }, - }), - ); - - expect(getTransportHeaders()).toMatchObject({ - 'x-algolia-secure-user-token': 'secure-user-token', - }); - }); - - it('omits the secure user token header when no memory token is provided', () => { - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - }), - ); - - expect(getTransportHeaders()).not.toHaveProperty('x-algolia-secure-user-token'); - }); - - it('sends an empty transport body when no search parameters or indices are provided', () => { - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - }), - ); - - expect(getTransportBody()).toEqual({ algolia: {} }); - }); - - it('includes searchParameters under the algolia body when provided', () => { - const searchParameters = { - 'index-name': { distinct: false }, - }; - - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - searchParameters, - }), - ); - - expect(getTransportBody()).toEqual({ - algolia: { searchParameters }, - }); - }); - - it('includes indices under the algolia body when provided', () => { - const indices = [ - { - index: 'docsearch-markdown', - description: 'Use this to gather specific results.', - }, - ]; - - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - indices, - }), - ); - - expect(getTransportBody()).toEqual({ - algolia: { indices }, - }); - }); - - it('includes both searchParameters and indices under the algolia body when both are provided', () => { - const searchParameters = { - 'index-name': { distinct: false }, - }; - const indices = [ - { - index: 'docsearch-markdown', - description: 'Use this to gather specific results.', - }, - ]; - - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - searchParameters, - indices, - }), - ); - - expect(getTransportBody()).toEqual({ - algolia: { searchParameters, indices }, - }); - }); - - it('omits indices from the body when an empty indices array is provided', () => { - renderHook(() => - useAskAi({ - apiKey: 'api-key', - appId: 'app-id', - assistantId: 'assistant-id', - indexName: 'index-name', - tools: {}, - indices: [], - }), - ); - - expect(getTransportBody()).toEqual({ algolia: {} }); - }); -}); diff --git a/packages/docsearch-react/src/types/StoredDocSearchHit.ts b/packages/docsearch-react/src/types/StoredDocSearchHit.ts index a56c6b33..20835a8c 100644 --- a/packages/docsearch-react/src/types/StoredDocSearchHit.ts +++ b/packages/docsearch-react/src/types/StoredDocSearchHit.ts @@ -14,6 +14,7 @@ export type StoredAskAiMessage = AIMessage & { }; export type StoredAskAiState = StoredDocSearchHit & { + chatId?: string; stopped?: boolean; messages?: StoredAskAiMessage[]; }; diff --git a/packages/docsearch-react/src/useAskAi.ts b/packages/docsearch-react/src/useAskAi.ts index 0bbb37f1..7b15c349 100644 --- a/packages/docsearch-react/src/useAskAi.ts +++ b/packages/docsearch-react/src/useAskAi.ts @@ -1,8 +1,8 @@ import type { UseChatHelpers } from '@ai-sdk/react'; -import { useChat } from '@ai-sdk/react'; +import { Chat, useChat } from '@ai-sdk/react'; import type { ChatOnToolCallCallback } from 'ai'; -import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai'; -import { useCallback, useEffect, useMemo, useRef } from 'react'; +import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, generateId } from 'ai'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { agentStudioBaseUrl, getAgentStudioErrorMessage, postAgentStudioFeedback } from './askai'; import type { Exchange } from './AskAiScreen'; @@ -28,6 +28,7 @@ type UseAskAiParams = { }; type UseAskAiReturn = { + chatId: string; messages: AIMessage[]; status: UseChat['status']; sendMessage: UseChat['sendMessage']; @@ -38,6 +39,14 @@ type UseAskAiReturn = { exchanges: Exchange[]; conversations: StoredSearchPlugin; sendFeedback: OnAskAiFeedback; + /** + * Create's a new chat instance, clearing existing messages and generating a new conversation ID. + */ + startNewConversation: () => void; + /** + * Create's a new chat instance, seeded with an existing conversation's ID and its messages. + */ + restoreConversation: (restored: AIMessage[], existingConversationId?: string) => void; }; type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn; @@ -105,6 +114,12 @@ export const useAskAi: UseAskAi = ({ [apiKey, appId, assistantId, searchParameters, memory?.userToken, indices], ); + // Store transport in a ref since it is dependent on unstable dependencies: + // - searchParameters, an object whose changed values trigger a new transport + // - indices, an array whose changed values trigger a new transport + const askAiTransportRef = useRef(askAiTransport); + askAiTransportRef.current = askAiTransport; + // Sync ref during render so the stable `handleToolCall` (registered once // by useChat) always sees the latest `tools` without re-creating itself. // Safe because tool calls only fire after a commit, and writes are idempotent. @@ -133,10 +148,25 @@ export const useAskAi: UseAskAi = ({ }); }, []); - const { messages, sendMessage, status, setMessages, error, stop, addToolOutput } = useChat({ - sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, - transport: askAiTransport, - onToolCall: handleToolCall, + const createChatInstance = useCallback( + (messages?: AIMessage[], id = generateId()): Chat => + new Chat({ + id, + messages, + sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, + transport: askAiTransportRef.current, + onToolCall: handleToolCall, + }), + [handleToolCall], + ); + + const [chatInstance, setChatInstance] = useState((): Chat => createChatInstance()); + // Keep a stable reference to the chat instance so reading messages are render safe + const chatInstanceRef = useRef>(chatInstance); + chatInstanceRef.current = chatInstance; + + const { messages, status, setMessages, error, stop, addToolOutput } = useChat({ + chat: chatInstance, }); useEffect(() => { @@ -171,12 +201,12 @@ export const useAskAi: UseAskAi = ({ [assistantId, appId, apiKey, conversations], ); - const onStopStreaming = async (): Promise => { + const onStopStreaming = useCallback(async (): Promise => { abortControllerRef.current.abort(); await stop(); - }; + }, [stop]); - const exchanges = useMemo(() => { + const exchanges = useMemo((): Exchange[] => { const grouped: Exchange[] = []; for (let i = 0; i < messages.length; i++) { @@ -201,9 +231,36 @@ export const useAskAi: UseAskAi = ({ return getAgentStudioErrorMessage(error); }, [error]); + const updateChatInstance = useCallback( + (restored?: AIMessage[], existingConversationId?: string): void => { + const newChatInstance = createChatInstance(restored, existingConversationId); + chatInstanceRef.current = newChatInstance; + setChatInstance(newChatInstance); + }, + [createChatInstance], + ); + + const startNewConversation = useCallback((): void => { + updateChatInstance(); + }, [updateChatInstance]); + + const restoreConversation = useCallback( + (restored: AIMessage[], existingConversationId?: string): void => { + updateChatInstance(restored, existingConversationId); + }, + [updateChatInstance], + ); + + // This is so that the public `sendMessage` is always pointed to a stable reference of the chat instance + const sendMessageSafe = useCallback( + (...args): Promise => chatInstanceRef.current!.sendMessage(...args), + [], + ); + return { + chatId: chatInstance.id, messages, - sendMessage, + sendMessage: sendMessageSafe, status, setMessages, askAiError, @@ -212,5 +269,7 @@ export const useAskAi: UseAskAi = ({ exchanges, conversations, sendFeedback, + startNewConversation, + restoreConversation, }; }; diff --git a/packages/docsearch-react/src/utils/ai.ts b/packages/docsearch-react/src/utils/ai.ts index 1ce4d2b5..7fd549ac 100644 --- a/packages/docsearch-react/src/utils/ai.ts +++ b/packages/docsearch-react/src/utils/ai.ts @@ -79,13 +79,14 @@ export function extractLinksFromMessage(message: AIMessage | null): ExtractedLin return links; } -export const buildDummyAskAiHit = (query: string, messages: AIMessage[]): StoredAskAiState => { +export const buildDummyAskAiHit = (query: string, messages: AIMessage[], chatId?: string): StoredAskAiState => { const textPart = messages[0].parts.find((part) => part.type === 'text'); const sanitizedText = textPart?.text ? sanitizeUserInput(textPart.text) : ''; return { query, objectID: sanitizedText, + chatId, messages, type: 'askAI', anchor: 'stored',