From 6a0b4769d1d514812fb384f22d8e90bf638764cc Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Fri, 10 Apr 2026 09:24:47 +0700
Subject: [PATCH] fix(ui): add LRU eviction to chat messages store

Keep at most 25 cached sessions in memory. When the limit is exceeded,
evict the oldest idle sessions (never evict sessions with an active run).
---
 ui/web/src/stores/use-chat-messages-store.ts | 150 ++++++++-----------
 1 file changed, 66 insertions(+), 84 deletions(-)

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<string, SessionMessages>;
-
-  // 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<string, SessionMessages>): Record<string, SessionMessages> {
+  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>): 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<ChatMessagesState>((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 }),
+      },
+    }));
   },
 }));