diff --git a/ui/web/src/stores/use-chat-messages-store.ts b/ui/web/src/stores/use-chat-messages-store.ts index 98ee7239..bc271a9d 100644 --- a/ui/web/src/stores/use-chat-messages-store.ts +++ b/ui/web/src/stores/use-chat-messages-store.ts @@ -1,119 +1,101 @@ import { create } from "zustand"; import type { ChatMessage } from "@/types/chat"; +const MAX_CACHED_SESSIONS = 25; + interface SessionMessages { messages: ChatMessage[]; streamText: string | null; thinkingText: string | null; isRunning: boolean; + lastAccessedAt: number; } interface ChatMessagesState { - // Map of sessionKey -> session message state sessions: Record; - - // Set messages for a specific session setSessionMessages: (sessionKey: string, messages: ChatMessage[]) => void; - - // Update session messages incrementally updateSessionMessages: (sessionKey: string, updater: (prev: ChatMessage[]) => ChatMessage[]) => void; - - // Set streaming text for a session setSessionStream: (sessionKey: string, streamText: string | null) => void; - - // Set thinking text for a session setSessionThinking: (sessionKey: string, thinkingText: string | null) => void; - - // Set running state for a session setSessionRunning: (sessionKey: string, isRunning: boolean) => void; } +// Evict oldest idle sessions when cache exceeds limit. +// Never evict sessions with an active run (isRunning). +function evictStale(sessions: Record): Record { + const keys = Object.keys(sessions); + if (keys.length <= MAX_CACHED_SESSIONS) return sessions; + + const evictable = keys + .filter((k) => !sessions[k]?.isRunning) + .sort((a, b) => (sessions[a]?.lastAccessedAt ?? 0) - (sessions[b]?.lastAccessedAt ?? 0)); + + const toRemove = evictable.slice(0, keys.length - MAX_CACHED_SESSIONS); + if (toRemove.length === 0) return sessions; + + const next = { ...sessions }; + for (const k of toRemove) delete next[k]; + return next; +} + +function touchSession(existing: SessionMessages | undefined, patch: Partial): SessionMessages { + return { + messages: patch.messages ?? existing?.messages ?? [], + streamText: patch.streamText !== undefined ? patch.streamText : (existing?.streamText ?? null), + thinkingText: patch.thinkingText !== undefined ? patch.thinkingText : (existing?.thinkingText ?? null), + isRunning: patch.isRunning !== undefined ? patch.isRunning : (existing?.isRunning ?? false), + lastAccessedAt: Date.now(), + }; +} + export const useChatMessagesStore = create((set) => ({ sessions: {}, - setSessionMessages: (sessionKey: string, messages: ChatMessage[]) => { + setSessionMessages: (sessionKey, messages) => { + set((state) => ({ + sessions: evictStale({ + ...state.sessions, + [sessionKey]: touchSession(state.sessions[sessionKey], { messages }), + }), + })); + }, + + updateSessionMessages: (sessionKey, updater) => { set((state) => { - const existing = state.sessions[sessionKey]; + const current = state.sessions[sessionKey]?.messages ?? []; return { - sessions: { + sessions: evictStale({ ...state.sessions, - [sessionKey]: { - messages, - streamText: existing?.streamText ?? null, - thinkingText: existing?.thinkingText ?? null, - isRunning: existing?.isRunning ?? false, - }, - }, + [sessionKey]: touchSession(state.sessions[sessionKey], { messages: updater(current) }), + }), }; }); }, - updateSessionMessages: (sessionKey: string, updater: (prev: ChatMessage[]) => ChatMessage[]) => { - set((state) => { - const existing = state.sessions[sessionKey]; - const current = existing?.messages ?? []; - return { - sessions: { - ...state.sessions, - [sessionKey]: { - messages: updater(current), - streamText: existing?.streamText ?? null, - thinkingText: existing?.thinkingText ?? null, - isRunning: existing?.isRunning ?? false, - }, - }, - }; - }); + setSessionStream: (sessionKey, streamText) => { + set((state) => ({ + sessions: { + ...state.sessions, + [sessionKey]: touchSession(state.sessions[sessionKey], { streamText }), + }, + })); }, - setSessionStream: (sessionKey: string, streamText: string | null) => { - set((state) => { - const existing = state.sessions[sessionKey]; - return { - sessions: { - ...state.sessions, - [sessionKey]: { - messages: existing?.messages ?? [], - streamText, - thinkingText: existing?.thinkingText ?? null, - isRunning: existing?.isRunning ?? false, - }, - }, - }; - }); + setSessionThinking: (sessionKey, thinkingText) => { + set((state) => ({ + sessions: { + ...state.sessions, + [sessionKey]: touchSession(state.sessions[sessionKey], { thinkingText }), + }, + })); }, - setSessionThinking: (sessionKey: string, thinkingText: string | null) => { - set((state) => { - const existing = state.sessions[sessionKey]; - return { - sessions: { - ...state.sessions, - [sessionKey]: { - messages: existing?.messages ?? [], - streamText: existing?.streamText ?? null, - thinkingText, - isRunning: existing?.isRunning ?? false, - }, - }, - }; - }); - }, - - setSessionRunning: (sessionKey: string, isRunning: boolean) => { - set((state) => { - const existing = state.sessions[sessionKey]; - return { - sessions: { - ...state.sessions, - [sessionKey]: { - messages: existing?.messages ?? [], - streamText: existing?.streamText ?? null, - thinkingText: existing?.thinkingText ?? null, - isRunning, - }, - }, - }; - }); + setSessionRunning: (sessionKey, isRunning) => { + set((state) => ({ + sessions: { + ...state.sessions, + [sessionKey]: touchSession(state.sessions[sessionKey], { isRunning }), + }, + })); }, }));