Improve session cache eviction

packages/ui/src/stores/message-v2/bus.ts

@@ -8,17 +8,39 @@ const log = getLogger("session")
 class MessageStoreBus {
   private stores = new Map<string, InstanceMessageStore>()
   private teardownHandlers = new Set<(instanceId: string) => void>()
+  private sessionClearHandlers = new Set<(instanceId: string, sessionId: string) => void>()
 
   registerInstance(instanceId: string, store?: InstanceMessageStore): InstanceMessageStore {
     if (this.stores.has(instanceId)) {
       return this.stores.get(instanceId) as InstanceMessageStore
     }
 
-    const resolved = store ?? createInstanceMessageStore(instanceId)
+    const resolved =
+      store ??
+      createInstanceMessageStore(instanceId, {
+        onSessionCleared: (id, sessionId) => this.notifySessionCleared(id, sessionId),
+      })
     this.stores.set(instanceId, resolved)
     return resolved
   }
 
+  onSessionCleared(handler: (instanceId: string, sessionId: string) => void): () => void {
+    this.sessionClearHandlers.add(handler)
+    return () => {
+      this.sessionClearHandlers.delete(handler)
+    }
+  }
+
+  private notifySessionCleared(instanceId: string, sessionId: string) {
+    for (const handler of this.sessionClearHandlers) {
+      try {
+        handler(instanceId, sessionId)
+      } catch (error) {
+        log.error("Failed to run session clear handler", error)
+      }
+    }
+  }
+
   getInstance(instanceId: string): InstanceMessageStore | undefined {
     return this.stores.get(instanceId)
   }

packages/ui/src/stores/message-v2/instance-store.ts

@@ -1,7 +1,9 @@
 import { batch } from "solid-js"
 import { createStore, produce, reconcile } from "solid-js/store"
 import type { SetStoreFunction } from "solid-js/store"
+import { getLogger } from "../../lib/logger"
 import type { ClientPart, MessageInfo } from "../../types/message"
+import { clearRecordDisplayCacheForMessages } from "./record-display-cache"
 import type {
   InstanceMessageState,
   MessageRecord,
@@ -17,6 +19,12 @@ import type {
   UsageEntry,
 } from "./types"
 
+const storeLog = getLogger("session")
+
+interface MessageStoreHooks {
+  onSessionCleared?: (instanceId: string, sessionId: string) => void
+}
+
 function createInitialState(instanceId: string): InstanceMessageState {
   return {
     instanceId,
@@ -202,7 +210,7 @@ export interface InstanceMessageStore {
   clearInstance: () => void
 }
 
-export function createInstanceMessageStore(instanceId: string): InstanceMessageStore {
+export function createInstanceMessageStore(instanceId: string, hooks?: MessageStoreHooks): InstanceMessageStore {
   const [state, setState] = createStore<InstanceMessageState>(createInitialState(instanceId))
 
 
@@ -696,80 +704,92 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
    function clearSession(sessionId: string) {
      if (!sessionId) return
 
-     const messageIds = Object.values(state.messages)
-       .filter((record) => record.sessionId === sessionId)
-       .map((record) => record.id)
+    const messageIds = Object.values(state.messages)
+      .filter((record) => record.sessionId === sessionId)
+      .map((record) => record.id)
+ 
+    storeLog.info("Clearing session data", { instanceId, sessionId, messageCount: messageIds.length })
+    clearRecordDisplayCacheForMessages(instanceId, messageIds)
+ 
+    batch(() => {
+      setState("messages", (prev) => {
+        const next = { ...prev }
+        messageIds.forEach((id) => delete next[id])
+        return next
+      })
 
-     // Remove message-level data
-     setState("messages", (prev) => {
-       const next = { ...prev }
-       messageIds.forEach((id) => delete next[id])
-       return next
-     })
+      setState("messageInfoVersion", (prev) => {
+        const next = { ...prev }
+        messageIds.forEach((id) => delete next[id])
+        return next
+      })
 
-     setState("messageInfoVersion", (prev) => {
-       const next = { ...prev }
-       messageIds.forEach((id) => delete next[id])
-       return next
-     })
+      messageIds.forEach((id) => messageInfoCache.delete(id))
 
-     messageIds.forEach((id) => messageInfoCache.delete(id))
+      setState("pendingParts", (prev) => {
+        const next = { ...prev }
+        messageIds.forEach((id) => {
+          if (next[id]) delete next[id]
+        })
+        return next
+      })
 
-     setState("pendingParts", (prev) => {
-       const next = { ...prev }
-       messageIds.forEach((id) => {
-         if (next[id]) delete next[id]
-       })
-       return next
-     })
+      setState("permissions", "byMessage", (prev) => {
+        const next = { ...prev }
+        messageIds.forEach((id) => {
+          if (next[id]) delete next[id]
+        })
+        return next
+      })
 
-     setState("permissions", "byMessage", (prev) => {
-       const next = { ...prev }
-       messageIds.forEach((id) => {
-         if (next[id]) delete next[id]
-       })
-       return next
-     })
+      setState("usage", (prev) => {
+        const next = { ...prev }
+        delete next[sessionId]
+        return next
+      })
 
-     // Remove session-level data
-     setState("usage", (prev) => {
-       const next = { ...prev }
-       delete next[sessionId]
-       return next
-     })
+      setState("sessionRevisions", (prev) => {
+        const next = { ...prev }
+        delete next[sessionId]
+        return next
+      })
 
-     setState("sessionRevisions", (prev) => {
-       const next = { ...prev }
-       delete next[sessionId]
-       return next
-     })
+      setState("scrollState", (prev) => {
+        const next = { ...prev }
+        const prefix = `${sessionId}:`
+        Object.keys(next).forEach((key) => {
+          if (key.startsWith(prefix)) {
+            delete next[key]
+          }
+        })
+        return next
+      })
 
-     setState("scrollState", (prev) => {
-       const next = { ...prev }
-       const prefix = `${sessionId}:`
-       Object.keys(next).forEach((key) => {
-         if (key.startsWith(prefix)) {
-           delete next[key]
-         }
-       })
-       return next
-     })
+      setState("sessions", sessionId, (current) => {
+        if (!current) return current
+        return { ...current, messageIds: [] }
+      })
 
-     setState("sessions", (prev) => {
-       const next = { ...prev }
-       delete next[sessionId]
-       return next
-     })
+      setState("sessions", (prev) => {
+        const next = { ...prev }
+        delete next[sessionId]
+        return next
+      })
 
-     setState("sessionOrder", (ids) => ids.filter((id) => id !== sessionId))
-   }
+      setState("sessionOrder", (ids) => ids.filter((id) => id !== sessionId))
+    })
+ 
+    hooks?.onSessionCleared?.(instanceId, sessionId)
+  }
 
+ 
    function clearInstance() {
      messageInfoCache.clear()
      setState(reconcile(createInitialState(instanceId)))
    }
  
    return {
+
      instanceId,
      state,
      setState,

packages/ui/src/stores/message-v2/record-display-cache.ts

@@ -44,3 +44,10 @@ export function clearRecordDisplayCacheForInstance(instanceId: string) {
     }
   }
 }
+
+export function clearRecordDisplayCacheForMessages(instanceId: string, messageIds: Iterable<string>) {
+  for (const messageId of messageIds) {
+    if (typeof messageId !== "string" || messageId.length === 0) continue
+    recordDisplayCache.delete(makeCacheKey(instanceId, messageId))
+  }
+}

packages/ui/src/stores/session-state.ts

@@ -39,7 +39,31 @@ const [loading, setLoading] = createSignal({
 const [messagesLoaded, setMessagesLoaded] = createSignal<Map<string, Set<string>>>(new Map())
 const [sessionInfoByInstance, setSessionInfoByInstance] = createSignal<Map<string, Map<string, SessionInfo>>>(new Map())
 
+function clearLoadedFlag(instanceId: string, sessionId: string) {
+  if (!instanceId || !sessionId) return
+  setMessagesLoaded((prev) => {
+    const existing = prev.get(instanceId)
+    if (!existing || !existing.has(sessionId)) {
+      return prev
+    }
+    const next = new Map(prev)
+    const updated = new Set(existing)
+    updated.delete(sessionId)
+    if (updated.size === 0) {
+      next.delete(instanceId)
+    } else {
+      next.set(instanceId, updated)
+    }
+    return next
+  })
+}
+
+messageStoreBus.onSessionCleared((instanceId, sessionId) => {
+  clearLoadedFlag(instanceId, sessionId)
+})
+
 function getDraftKey(instanceId: string, sessionId: string): string {
+
   return `${instanceId}:${sessionId}`
 }
 
@@ -357,8 +381,9 @@ export {
   setSessionCompactionState,
   setSessionPendingPermission,
   setActiveSession,
-
+ 
   setActiveParentSession,
+
   clearActiveParentSession,
   getActiveSession,
   getActiveParentSession,

packages/ui/src/stores/sessions.ts

@@ -26,7 +26,8 @@ import {
   setActiveParentSession,
   setActiveSession,
   setSessionDraftPrompt,
-} from "./session-state"
+ } from "./session-state"
+
 import { getDefaultModel } from "./session-models"
 import {
   createSession,