From 1bba183e8d21869642cc4da9d8915df4738515dc Mon Sep 17 00:00:00 2001 From: Jan Calanog Date: Mon, 26 Jan 2026 22:57:16 +0100 Subject: [PATCH] Fix AskAI chat persistence and use well-known `idb-keyval` library --- .../Assets/web-components/AskAi/Chat.test.tsx | 13 +- .../web-components/AskAi/ChatMessage.test.tsx | 13 +- .../web-components/AskAi/chat.store.test.ts | 21 +- .../Assets/web-components/AskAi/chat.store.ts | 248 +++++++++++------- .../package-lock.json | 19 +- src/Elastic.Documentation.Site/package.json | 4 +- 6 files changed, 185 insertions(+), 133 deletions(-) diff --git a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/Chat.test.tsx b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/Chat.test.tsx index 5a9898084..676248f45 100644 --- a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/Chat.test.tsx +++ b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/Chat.test.tsx @@ -6,13 +6,12 @@ import { render, screen, waitFor } from '@testing-library/react' import userEvent from '@testing-library/user-event' import * as React from 'react' -// Mock zustand-indexeddb (IndexedDB not available in Node.js test environment) -jest.mock('zustand-indexeddb', () => ({ - createIndexedDBStorage: () => ({ - getItem: jest.fn().mockResolvedValue(null), - setItem: jest.fn().mockResolvedValue(undefined), - removeItem: jest.fn().mockResolvedValue(undefined), - }), +// Mock idb-keyval (IndexedDB not available in Node.js test environment) +jest.mock('idb-keyval', () => ({ + get: jest.fn().mockResolvedValue(null), + set: jest.fn().mockResolvedValue(undefined), + del: jest.fn().mockResolvedValue(undefined), + createStore: jest.fn().mockReturnValue({}), })) // Create a fresh QueryClient for each test diff --git a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/ChatMessage.test.tsx b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/ChatMessage.test.tsx index 8d51c0335..c440f28a5 100644 --- a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/ChatMessage.test.tsx +++ b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/ChatMessage.test.tsx @@ -6,13 +6,12 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { render, screen } from '@testing-library/react' import * as React from 'react' -// Mock zustand-indexeddb (IndexedDB not available in Node.js test environment) -jest.mock('zustand-indexeddb', () => ({ - createIndexedDBStorage: () => ({ - getItem: jest.fn().mockResolvedValue(null), - setItem: jest.fn().mockResolvedValue(undefined), - removeItem: jest.fn().mockResolvedValue(undefined), - }), +// Mock idb-keyval (IndexedDB not available in Node.js test environment) +jest.mock('idb-keyval', () => ({ + get: jest.fn().mockResolvedValue(null), + set: jest.fn().mockResolvedValue(undefined), + del: jest.fn().mockResolvedValue(undefined), + createStore: jest.fn().mockReturnValue({}), })) // Create a fresh QueryClient for each test diff --git a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.test.ts b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.test.ts index 1e38556c1..02b306680 100644 --- a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.test.ts +++ b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.test.ts @@ -2,13 +2,12 @@ import { chatStore } from './chat.store' import { act } from 'react' import { v4 as uuidv4 } from 'uuid' -// Mock zustand-indexeddb (IndexedDB not available in Node.js test environment) -jest.mock('zustand-indexeddb', () => ({ - createIndexedDBStorage: () => ({ - getItem: jest.fn().mockResolvedValue(null), - setItem: jest.fn().mockResolvedValue(undefined), - removeItem: jest.fn().mockResolvedValue(undefined), - }), +// Mock idb-keyval (IndexedDB not available in Node.js test environment) +jest.mock('idb-keyval', () => ({ + get: jest.fn().mockResolvedValue(null), + set: jest.fn().mockResolvedValue(undefined), + del: jest.fn().mockResolvedValue(undefined), + createStore: jest.fn().mockReturnValue({}), })) // Mock uuid @@ -182,7 +181,7 @@ describe('chat.store', () => { // The persist middleware adds a persist property to the store expect(chatStore.persist).toBeDefined() expect(chatStore.persist.getOptions().name).toBe( - 'elastic-docs-conversations-index' + 'ask-ai/conversations-index' ) expect(chatStore.persist.getOptions().version).toBe(1) }) @@ -243,9 +242,7 @@ describe('chat.store', () => { expect(state.conversations['conv-123'].messageCount).toBe(2) expect(state.conversations['conv-123'].createdAt).toBeDefined() expect(state.conversations['conv-123'].updatedAt).toBeDefined() - expect(state.conversations['conv-123'].aiProvider).toBe( - 'LlmGateway' - ) + // aiProvider is stored per-conversation, not in the index }) it('should update existing ConversationMeta when setConversationId is called with existing ID', () => { @@ -297,7 +294,7 @@ describe('chat.store', () => { }) it('should clear all conversations when clearAllConversations is called', () => { - // Create multiple conversations + // Create multiple conversations (MAX_CONVERSATIONS = 2) act(() => { chatStore.getState().actions.submitQuestion('Question 1') chatStore.getState().actions.setConversationId('conv-1') diff --git a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.ts b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.ts index dd5bd9e49..a7acbc5e6 100644 --- a/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.ts +++ b/src/Elastic.Documentation.Site/Assets/web-components/AskAi/chat.store.ts @@ -4,77 +4,117 @@ import { cooldownStore } from '../shared/cooldown.store' import { ApiError, isApiError, isRateLimitError } from '../shared/errorHandling' import { AskAiEvent } from './AskAiEvent' import { MessageThrottler } from './MessageThrottler' +import { + get as idbGet, + set as idbSet, + del as idbDel, + createStore as createIdbStore, +} from 'idb-keyval' import { v4 as uuidv4 } from 'uuid' import { createStore } from 'zustand' -import { createIndexedDBStorage } from 'zustand-indexeddb' -import { persist } from 'zustand/middleware' +import { persist, PersistStorage, StorageValue } from 'zustand/middleware' import { useStore } from 'zustand/react' -// IndexedDB storage for conversations index (lightweight, always loaded) -const conversationsIndexStorage = createIndexedDBStorage( - 'elastic-docs', - 'conversations-index' +// ============================================================================= +// IndexedDB Storage via idb-keyval +// ============================================================================= +// We use idb-keyval for IndexedDB access. It provides a simple key-value API +// without the complexity of managing object stores and schema versions. +// +// IndexedDB supports structured cloning, so we store objects directly without +// JSON serialization (unlike localStorage which only stores strings). +// +// Key naming convention: +// - "ask-ai/conversations-index" - lightweight index for listing (zustand persist) +// - "ask-ai/conversation-{id}" - conversation data (messages, feedback, etc.) +// +// To add new IndexedDB storage for other features, simply use new key prefixes +// with the get/set/del helpers below. +// ============================================================================= + +// Custom IndexedDB store for elastic-docs data +const elasticDocsStore = createIdbStore( + 'elastic-docs-keyval-store', + 'elastic-docs-keyval-store' ) -// IndexedDB storage for per-conversation messages (lazy loaded) -// Each conversation gets its own storage key -const messageStorageCache = new Map< - ConversationId, - ReturnType ->() +// Typed wrappers around idb-keyval that use our custom store +const get = (key: string) => idbGet(key, elasticDocsStore) +const set = (key: string, value: T) => idbSet(key, value, elasticDocsStore) +const del = (key: string) => idbDel(key, elasticDocsStore) + +// PersistStorage adapter for zustand's persist middleware +// Uses IndexedDB's native object storage (no JSON serialization needed) +function createIdbStorage(): PersistStorage { + return { + getItem: async (name: string): Promise | null> => { + const value = await get>(name) + return value ?? null + }, + setItem: async ( + name: string, + value: StorageValue + ): Promise => { + await set(name, value) + }, + removeItem: async (name: string): Promise => { + await del(name) + }, + } +} -function getMessageStorage(conversationId: ConversationId) { - if (!messageStorageCache.has(conversationId)) { - messageStorageCache.set( - conversationId, - createIndexedDBStorage('elastic-docs', `messages-${conversationId}`) - ) +// Zustand persist storage for conversations index +const conversationsIndexStorage = createIdbStorage>() + +// Conversation storage key helper +const getConversationKey = (conversationId: ConversationId) => + `ask-ai/conversation-${conversationId}` + +// Type for persisted conversation data (messages, feedback, and settings) +interface PersistedConversation { + state: { + chatMessages: ChatMessage[] + totalMessageCount: number + messageFeedback: Record + aiProvider: AiProvider } - return messageStorageCache.get(conversationId)! + version: number } -// Helper to save messages for a conversation -async function saveConversationMessages( +// Helper to save a conversation's data +async function saveConversationData( conversationId: ConversationId, messages: ChatMessage[], totalMessageCount: number, - messageFeedback: Record + messageFeedback: Record, + aiProvider: AiProvider ) { - const storage = getMessageStorage(conversationId) - await storage.setItem(`messages-${conversationId}`, { + const data: PersistedConversation = { state: { - chatMessages: messages.slice(-MAX_PERSISTED_MESSAGES), + chatMessages: messages, totalMessageCount, messageFeedback, + aiProvider, }, version: 1, - }) + } + await set(getConversationKey(conversationId), data) } -// Helper to load messages for a conversation -async function loadConversationMessages( - conversationId: ConversationId -): Promise<{ +// Helper to load a conversation's data +async function loadConversationData(conversationId: ConversationId): Promise<{ chatMessages: ChatMessage[] totalMessageCount: number messageFeedback: Record + aiProvider: AiProvider } | null> { - const storage = getMessageStorage(conversationId) try { - const stored = await storage.getItem(`messages-${conversationId}`) - if (stored && typeof stored === 'object' && 'state' in stored) { - const state = ( - stored as { - state: { - chatMessages?: ChatMessage[] - totalMessageCount?: number - messageFeedback?: Record - } - } - ).state - + const stored = await get( + getConversationKey(conversationId) + ) + if (stored?.state) { // Mark any streaming messages as interrupted - const messages = (state.chatMessages ?? []).map( + const messages = (stored.state.chatMessages ?? []).map( (msg: ChatMessage) => msg.status === 'streaming' ? { ...msg, status: 'interrupted' as const } @@ -83,8 +123,10 @@ async function loadConversationMessages( return { chatMessages: messages, - totalMessageCount: state.totalMessageCount ?? messages.length, - messageFeedback: state.messageFeedback ?? {}, + totalMessageCount: + stored.state.totalMessageCount ?? messages.length, + messageFeedback: stored.state.messageFeedback ?? {}, + aiProvider: stored.state.aiProvider ?? 'LlmGateway', } } } catch { @@ -93,11 +135,9 @@ async function loadConversationMessages( return null } -// Helper to delete messages for a conversation -async function deleteConversationMessages(conversationId: ConversationId) { - const storage = getMessageStorage(conversationId) - await storage.removeItem(`messages-${conversationId}`) - messageStorageCache.delete(conversationId) +// Helper to delete a conversation's data +async function deleteConversationData(conversationId: ConversationId) { + await del(getConversationKey(conversationId)) } export type { AiProvider } @@ -116,13 +156,13 @@ export interface ChatMessage { reasoning?: AskAiEvent[] // Reasoning steps for status display (search, tool calls, etc.) } +// Conversation metadata stored in the index (for listing and sorting) export interface ConversationMeta { id: ConversationId // Backend conversation ID (set on conversation_start) title: string // First user message, truncated createdAt: number // When conversation started updatedAt: number // When last message was added messageCount: number - aiProvider: AiProvider // Which AI provider this conversation uses } interface ActiveStream { @@ -136,8 +176,10 @@ const activeStreams = new Map() const sentAiMessageIds = new Set() -// Maximum number of messages to persist in localStorage -const MAX_PERSISTED_MESSAGES = 50 +// Maximum number of conversations to keep (oldest are deleted when exceeded) +// This is a temporary limit to prevent the IndexedDB from growing too large +// As soon as we support multiple conversation in the UI, we will set a more reasonable limit +const MAX_CONVERSATIONS = 2 interface ChatState { // Conversation index (lightweight, always loaded) - keyed by conversation ID for O(1) lookup @@ -290,14 +332,37 @@ export const chatStore = createStore()( createdAt: Date.now(), updatedAt: Date.now(), messageCount: state.chatMessages.length, - aiProvider: state.aiProvider, } + + // Enforce MAX_CONVERSATIONS limit - delete oldest if exceeded + let updatedConversations = { + ...state.conversations, + [conversationId]: newConv, + } + + const convList = Object.values(updatedConversations) + if (convList.length > MAX_CONVERSATIONS) { + // Sort by updatedAt descending, keep only the newest + const sorted = convList.sort( + (a, b) => b.updatedAt - a.updatedAt + ) + const toKeep = sorted.slice(0, MAX_CONVERSATIONS) + const toDelete = sorted.slice(MAX_CONVERSATIONS) + + // Delete old conversation data from IndexedDB + toDelete.forEach((conv) => + deleteConversationData(conv.id) + ) + + // Rebuild conversations map with only kept ones + updatedConversations = Object.fromEntries( + toKeep.map((c) => [c.id, c]) + ) + } + set({ activeConversationId: conversationId, - conversations: { - ...state.conversations, - [conversationId]: newConv, - }, + conversations: updatedConversations, }) } }, @@ -404,31 +469,29 @@ export const chatStore = createStore()( switchConversation: async (id) => { const state = get() - // Save current conversation's messages before switching + // Save current conversation's data before switching if ( state.activeConversationId && state.chatMessages.length > 0 ) { - await saveConversationMessages( + await saveConversationData( state.activeConversationId, state.chatMessages, state.totalMessageCount, - state.messageFeedback + state.messageFeedback, + state.aiProvider ) } - // Find the conversation to get its aiProvider - const targetConv = state.conversations[id] - - // Load the new conversation's messages - const loaded = await loadConversationMessages(id) + // Load the new conversation's data + const loaded = await loadConversationData(id) set({ activeConversationId: id, chatMessages: loaded?.chatMessages ?? [], totalMessageCount: loaded?.totalMessageCount ?? 0, messageFeedback: loaded?.messageFeedback ?? {}, - aiProvider: targetConv?.aiProvider ?? state.aiProvider, + aiProvider: loaded?.aiProvider ?? state.aiProvider, scrollPosition: 0, inputValue: '', }) @@ -442,11 +505,12 @@ export const chatStore = createStore()( state.activeConversationId && state.chatMessages.length > 0 ) { - saveConversationMessages( + saveConversationData( state.activeConversationId, state.chatMessages, state.totalMessageCount, - state.messageFeedback + state.messageFeedback, + state.aiProvider ) } @@ -465,7 +529,7 @@ export const chatStore = createStore()( const state = get() // Delete from IndexedDB - await deleteConversationMessages(id) + await deleteConversationData(id) // Remove from conversations map // eslint-disable-next-line @typescript-eslint/no-unused-vars @@ -481,7 +545,7 @@ export const chatStore = createStore()( )[0] if (mostRecent) { - const loaded = await loadConversationMessages( + const loaded = await loadConversationData( mostRecent.id ) set({ @@ -491,7 +555,8 @@ export const chatStore = createStore()( totalMessageCount: loaded?.totalMessageCount ?? 0, messageFeedback: loaded?.messageFeedback ?? {}, - aiProvider: mostRecent.aiProvider, + aiProvider: + loaded?.aiProvider ?? state.aiProvider, scrollPosition: 0, }) } else { @@ -522,7 +587,7 @@ export const chatStore = createStore()( // Delete all conversation messages from IndexedDB for (const conv of Object.values(state.conversations)) { - deleteConversationMessages(conv.id) + deleteConversationData(conv.id) } set({ @@ -538,7 +603,7 @@ export const chatStore = createStore()( }, }), { - name: 'elastic-docs-conversations-index', + name: 'ask-ai/conversations-index', version: 1, storage: conversationsIndexStorage, skipHydration: true, // Manual hydration to properly sequence loading @@ -563,7 +628,7 @@ if (typeof window !== 'undefined') { try { // Phase 1: Load conversations index const stored = await conversationsIndexStorage.getItem( - 'elastic-docs-conversations-index' + 'ask-ai/conversations-index' ) let conversations: Record = {} @@ -582,29 +647,22 @@ if (typeof window !== 'undefined') { persistedState.activeConversationId ?? null scrollPosition = persistedState.scrollPosition ?? 0 inputValue = persistedState.inputValue ?? '' - - // Get aiProvider from the active conversation, or fall back to stored/default - const activeConv = activeConversationId - ? conversations[activeConversationId] - : undefined - aiProvider = - activeConv?.aiProvider ?? - persistedState.aiProvider ?? - 'LlmGateway' + // Use stored aiProvider as fallback (will be overridden by conversation data if available) + aiProvider = persistedState.aiProvider ?? 'LlmGateway' } - // Phase 2: Load active conversation's messages (if any) + // Phase 2: Load active conversation's data (if any) let chatMessages: ChatMessage[] = [] let totalMessageCount = 0 let messageFeedback: Record = {} if (activeConversationId) { - const loaded = - await loadConversationMessages(activeConversationId) + const loaded = await loadConversationData(activeConversationId) if (loaded) { chatMessages = loaded.chatMessages totalMessageCount = loaded.totalMessageCount messageFeedback = loaded.messageFeedback + aiProvider = loaded.aiProvider } } @@ -648,13 +706,14 @@ if (typeof window !== 'undefined') { return msg }) - // Save current conversation's messages to IndexedDB + // Save current conversation's data to IndexedDB if (state.activeConversationId && messages.length > 0) { - saveConversationMessages( + saveConversationData( state.activeConversationId, messages, state.totalMessageCount, - state.messageFeedback + state.messageFeedback, + state.aiProvider ) } }) @@ -744,14 +803,15 @@ function handleStreamEvent(messageId: string, event: AskAiEvent): void { stream.throttler.clear() activeStreams.delete(messageId) - // Save messages to IndexedDB after streaming completes + // Save conversation data to IndexedDB after streaming completes const state = chatStore.getState() if (state.activeConversationId) { - saveConversationMessages( + saveConversationData( state.activeConversationId, state.chatMessages, state.totalMessageCount, - state.messageFeedback + state.messageFeedback, + state.aiProvider ) } } diff --git a/src/Elastic.Documentation.Site/package-lock.json b/src/Elastic.Documentation.Site/package-lock.json index 0b96646f4..cca06c21f 100644 --- a/src/Elastic.Documentation.Site/package-lock.json +++ b/src/Elastic.Documentation.Site/package-lock.json @@ -34,6 +34,7 @@ "htmx-ext-head-support": "2.0.4", "htmx-ext-preload": "2.1.2", "htmx.org": "2.0.8", + "idb-keyval": "^6.2.2", "katex": "^0.16.25", "lodash": "^4.17.21", "marked": "17.0.1", @@ -43,8 +44,7 @@ "ua-parser-js": "2.0.6", "uuid": "11.1.0", "zod": "4.1.12", - "zustand": "5.0.8", - "zustand-indexeddb": "^0.1.1" + "zustand": "5.0.8" }, "devDependencies": { "@babel/core": "7.28.4", @@ -27335,6 +27335,12 @@ "node": ">=0.10.0" } }, + "node_modules/idb-keyval": { + "version": "6.2.2", + "resolved": "https://registry.npmjs.org/idb-keyval/-/idb-keyval-6.2.2.tgz", + "integrity": "sha512-yjD9nARJ/jb1g+CvD0tlhUHOrJ9Sy0P8T9MF3YaLlHnSRpwPfpTX0XIvpmw3gAJUmEu3FiICLBDPXVwyEvrleg==", + "license": "Apache-2.0" + }, "node_modules/identity-obj-proxy": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/identity-obj-proxy/-/identity-obj-proxy-3.0.0.tgz", @@ -33415,15 +33421,6 @@ } } }, - "node_modules/zustand-indexeddb": { - "version": "0.1.1", - "resolved": "https://registry.npmjs.org/zustand-indexeddb/-/zustand-indexeddb-0.1.1.tgz", - "integrity": "sha512-2IErbvNdzxHkPLRep53e2dRcX0TtDmc940aPyEjzLxAYWVxNY9crrtkufrE/k/jkvfxybuu8xN+emSdWFMVP6w==", - "license": "MIT", - "peerDependencies": { - "zustand": "^5.0.0" - } - }, "node_modules/zwitch": { "version": "1.0.5", "resolved": "https://registry.npmjs.org/zwitch/-/zwitch-1.0.5.tgz", diff --git a/src/Elastic.Documentation.Site/package.json b/src/Elastic.Documentation.Site/package.json index 61cb2de0f..d5a56a204 100644 --- a/src/Elastic.Documentation.Site/package.json +++ b/src/Elastic.Documentation.Site/package.json @@ -111,6 +111,7 @@ "htmx-ext-head-support": "2.0.4", "htmx-ext-preload": "2.1.2", "htmx.org": "2.0.8", + "idb-keyval": "6.2.2", "katex": "^0.16.25", "lodash": "^4.17.21", "marked": "17.0.1", @@ -120,7 +121,6 @@ "ua-parser-js": "2.0.6", "uuid": "11.1.0", "zod": "4.1.12", - "zustand": "5.0.8", - "zustand-indexeddb": "^0.1.1" + "zustand": "5.0.8" } }