Improve message stream caching and virtualization for large sessions

packages/ui/src/lib/global-cache.ts

@@ -0,0 +1,126 @@
+export interface CacheEntryBaseParams {
+  instanceId?: string
+  sessionId?: string
+  scope: string
+}
+
+export interface CacheEntryParams extends CacheEntryBaseParams {
+  key: string
+}
+
+type CacheValueMap = Map<string, unknown>
+type CacheScopeMap = Map<string, CacheValueMap>
+type CacheSessionMap = Map<string, CacheScopeMap>
+
+const GLOBAL_KEY = "GLOBAL"
+const cacheStore = new Map<string, CacheSessionMap>()
+
+function resolveKey(value?: string) {
+  return value && value.length > 0 ? value : GLOBAL_KEY
+}
+
+function getScopeValueMap(params: CacheEntryParams, create: boolean): CacheValueMap | undefined {
+  const instanceKey = resolveKey(params.instanceId)
+  const sessionKey = resolveKey(params.sessionId)
+
+  let sessionMap = cacheStore.get(instanceKey)
+  if (!sessionMap) {
+    if (!create) return undefined
+    sessionMap = new Map()
+    cacheStore.set(instanceKey, sessionMap)
+  }
+
+  let scopeMap = sessionMap.get(sessionKey)
+  if (!scopeMap) {
+    if (!create) return undefined
+    scopeMap = new Map()
+    sessionMap.set(sessionKey, scopeMap)
+  }
+
+  let valueMap = scopeMap.get(params.scope)
+  if (!valueMap) {
+    if (!create) return undefined
+    valueMap = new Map()
+    scopeMap.set(params.scope, valueMap)
+  }
+
+  return valueMap
+}
+
+function cleanupHierarchy(instanceKey: string, sessionKey: string, scopeKey?: string) {
+  const sessionMap = cacheStore.get(instanceKey)
+  if (!sessionMap) {
+    return
+  }
+
+  const scopeMap = sessionMap.get(sessionKey)
+  if (!scopeMap) {
+    if (sessionMap.size === 0) {
+      cacheStore.delete(instanceKey)
+    }
+    return
+  }
+
+  if (scopeKey) {
+    const valueMap = scopeMap.get(scopeKey)
+    if (valueMap && valueMap.size === 0) {
+      scopeMap.delete(scopeKey)
+    }
+  }
+
+  if (scopeMap.size === 0) {
+    sessionMap.delete(sessionKey)
+  }
+
+  if (sessionMap.size === 0) {
+    cacheStore.delete(instanceKey)
+  }
+}
+
+export function setCacheEntry<T>(params: CacheEntryParams, value: T | undefined): void {
+  const instanceKey = resolveKey(params.instanceId)
+  const sessionKey = resolveKey(params.sessionId)
+
+  if (value === undefined) {
+    const existingMap = getScopeValueMap(params, false)
+    existingMap?.delete(params.key)
+    cleanupHierarchy(instanceKey, sessionKey, params.scope)
+    return
+  }
+
+  const scopeEntries = getScopeValueMap(params, true)
+  scopeEntries?.set(params.key, value)
+}
+
+export function getCacheEntry<T>(params: CacheEntryParams): T | undefined {
+  const scopeEntries = getScopeValueMap(params, false)
+  return scopeEntries?.get(params.key) as T | undefined
+}
+
+export function clearCacheScope(params: CacheEntryBaseParams): void {
+  const instanceKey = resolveKey(params.instanceId)
+  const sessionKey = resolveKey(params.sessionId)
+  const sessionMap = cacheStore.get(instanceKey)
+  if (!sessionMap) return
+  const scopeMap = sessionMap.get(sessionKey)
+  if (!scopeMap) return
+  scopeMap.delete(params.scope)
+  cleanupHierarchy(instanceKey, sessionKey)
+}
+
+export function clearCacheForSession(instanceId?: string, sessionId?: string): void {
+  const instanceKey = resolveKey(instanceId)
+  const sessionKey = resolveKey(sessionId)
+  const sessionMap = cacheStore.get(instanceKey)
+  if (!sessionMap) return
+  sessionMap.delete(sessionKey)
+  if (sessionMap.size === 0) {
+    cacheStore.delete(instanceKey)
+  }
+}
+
+export function clearCacheForInstance(instanceId?: string): void {
+  const instanceKey = resolveKey(instanceId)
+  cacheStore.delete(instanceKey)
+}
+

packages/ui/src/lib/hooks/use-global-cache.ts

@@ -0,0 +1,86 @@
+import { type Accessor, createMemo } from "solid-js"
+import {
+  type CacheEntryParams,
+  getCacheEntry,
+  setCacheEntry,
+  clearCacheScope,
+  clearCacheForSession,
+  clearCacheForInstance,
+} from "../global-cache"
+
+/**
+ * `useGlobalCache` exposes a tiny typed facade over the shared cache helpers.
+ * Callers can pass raw values or accessors for the cache keys; empty identifiers
+ * automatically fall back to the global buckets.
+ */
+export function useGlobalCache(params: UseGlobalCacheParams): GlobalCacheHandle {
+  const resolvedEntry = createMemo<CacheEntryParams>(() => {
+    const instanceId = normalizeId(resolveValue(params.instanceId))
+    const sessionId = normalizeId(resolveValue(params.sessionId))
+    const scope = resolveValue(params.scope)
+    const key = resolveValue(params.key)
+    return { instanceId, sessionId, scope, key }
+  })
+
+  const scopeParams = createMemo(() => {
+    const entry = resolvedEntry()
+    return { instanceId: entry.instanceId, sessionId: entry.sessionId, scope: entry.scope }
+  })
+
+  const sessionParams = createMemo(() => {
+    const entry = resolvedEntry()
+    return { instanceId: entry.instanceId, sessionId: entry.sessionId }
+  })
+
+  return {
+    get<T>() {
+      return getCacheEntry<T>(resolvedEntry())
+    },
+    set<T>(value: T | undefined) {
+      setCacheEntry(resolvedEntry(), value)
+    },
+    clearScope() {
+      clearCacheScope(scopeParams())
+    },
+    clearSession() {
+      const params = sessionParams()
+      clearCacheForSession(params.instanceId, params.sessionId)
+    },
+    clearInstance() {
+      const params = sessionParams()
+      clearCacheForInstance(params.instanceId)
+    },
+    params() {
+      return resolvedEntry()
+    },
+  }
+}
+
+function normalizeId(value?: string): string | undefined {
+  return value && value.length > 0 ? value : undefined
+}
+
+function resolveValue<T>(value: MaybeAccessor<T> | undefined): T {
+  if (typeof value === "function") {
+    return (value as Accessor<T>)()
+  }
+  return value as T
+}
+
+type MaybeAccessor<T> = T | Accessor<T>
+
+interface UseGlobalCacheParams {
+  instanceId?: MaybeAccessor<string | undefined>
+  sessionId?: MaybeAccessor<string | undefined>
+  scope: MaybeAccessor<string>
+  key: MaybeAccessor<string>
+}
+
+interface GlobalCacheHandle {
+  get<T>(): T | undefined
+  set<T>(value: T | undefined): void
+  clearScope(): void
+  clearSession(): void
+  clearInstance(): void
+  params(): CacheEntryParams
+}

packages/ui/src/lib/hooks/use-scroll-cache.ts

@@ -0,0 +1,102 @@
+import { type Accessor, createMemo } from "solid-js"
+import { messageStoreBus } from "../../stores/message-v2/bus"
+import type { ScrollSnapshot } from "../../stores/message-v2/types"
+
+interface UseScrollCacheParams {
+  instanceId: MaybeAccessor<string>
+  sessionId: MaybeAccessor<string>
+  scope: MaybeAccessor<string>
+}
+
+interface PersistScrollOptions {
+  atBottomOffset?: number
+}
+
+interface RestoreScrollOptions {
+  behavior?: ScrollBehavior
+  fallback?: () => void
+  onApplied?: (snapshot: ScrollSnapshot | undefined) => void
+}
+
+interface ScrollCacheHandle {
+  persist: (element: HTMLElement | null | undefined, options?: PersistScrollOptions) => ScrollSnapshot | undefined
+  restore: (element: HTMLElement | null | undefined, options?: RestoreScrollOptions) => void
+}
+
+const DEFAULT_BOTTOM_OFFSET = 48
+
+/**
+ * Wraps the message-store scroll snapshot helpers so components can
+ * persist/restore scroll positions without duplicating requestAnimationFrame
+ * boilerplate.
+ */
+export function useScrollCache(params: UseScrollCacheParams): ScrollCacheHandle {
+  const resolved = createMemo(() => ({
+    instanceId: resolveValue(params.instanceId),
+    sessionId: resolveValue(params.sessionId),
+    scope: resolveValue(params.scope),
+  }))
+
+  const store = createMemo(() => {
+    const { instanceId } = resolved()
+    return messageStoreBus.getOrCreate(instanceId)
+  })
+
+  function persist(element: HTMLElement | null | undefined, options?: PersistScrollOptions) {
+    if (!element) {
+      return undefined
+    }
+    const target = resolved()
+    if (!target.sessionId) {
+      return undefined
+    }
+    const snapshot: Omit<ScrollSnapshot, "updatedAt"> = {
+      scrollTop: element.scrollTop,
+      atBottom: isNearBottom(element, options?.atBottomOffset ?? DEFAULT_BOTTOM_OFFSET),
+    }
+    store().setScrollSnapshot(target.sessionId, target.scope, snapshot)
+    return { ...snapshot, updatedAt: Date.now() }
+  }
+
+  function restore(element: HTMLElement | null | undefined, options?: RestoreScrollOptions) {
+    const target = resolved()
+    if (!element || !target.sessionId) {
+      options?.fallback?.()
+      options?.onApplied?.(undefined)
+      return
+    }
+    const snapshot = store().getScrollSnapshot(target.sessionId, target.scope)
+    requestAnimationFrame(() => {
+      if (!element) {
+        options?.onApplied?.(snapshot)
+        return
+      }
+      if (!snapshot) {
+        options?.fallback?.()
+        options?.onApplied?.(undefined)
+        return
+      }
+      const maxScrollTop = Math.max(element.scrollHeight - element.clientHeight, 0)
+      const nextTop = snapshot.atBottom ? maxScrollTop : Math.min(snapshot.scrollTop, maxScrollTop)
+      const behavior = options?.behavior ?? "auto"
+      element.scrollTo({ top: nextTop, behavior })
+      options?.onApplied?.(snapshot)
+    })
+  }
+
+  return { persist, restore }
+}
+
+function isNearBottom(element: HTMLElement, offset: number) {
+  const { scrollTop, scrollHeight, clientHeight } = element
+  return scrollHeight - (scrollTop + clientHeight) <= offset
+}
+
+function resolveValue<T>(value: MaybeAccessor<T>): T {
+  if (typeof value === "function") {
+    return (value as Accessor<T>)()
+  }
+  return value
+}
+
+type MaybeAccessor<T> = T | Accessor<T>

packages/ui/src/lib/scroll-cache.ts

@@ -1,53 +0,0 @@
-import type { ScrollSnapshot } from "../stores/message-v2/types"
-
-interface ScrollCacheParams {
-  instanceId?: string
-  sessionId?: string
-  scope?: string
-}
-
-const scrollCache = new Map<string, ScrollSnapshot>()
-const DEFAULT_SCOPE = "session"
-
-function resolve(value?: string) {
-  return value && value.length > 0 ? value : "GLOBAL"
-}
-
-function makeKey(params: ScrollCacheParams) {
-  return `${resolve(params.instanceId)}:${resolve(params.sessionId)}:${params.scope ?? DEFAULT_SCOPE}`
-}
-
-export function setScrollCache(params: ScrollCacheParams, snapshot: Omit<ScrollSnapshot, "updatedAt">) {
-  scrollCache.set(makeKey(params), { ...snapshot, updatedAt: Date.now() })
-}
-
-export function getScrollCache(params: ScrollCacheParams): ScrollSnapshot | undefined {
-  return scrollCache.get(makeKey(params))
-}
-
-export function clearScrollCacheScope(params: ScrollCacheParams) {
-  const key = makeKey(params)
-  scrollCache.delete(key)
-}
-
-export function clearScrollCacheForSession(instanceId?: string, sessionId?: string) {
-  const match = `${resolve(instanceId)}:${resolve(sessionId)}:`
-  for (const key of scrollCache.keys()) {
-    if (key.startsWith(match)) {
-      scrollCache.delete(key)
-    }
-  }
-}
-
-export function clearScrollCacheForInstance(instanceId?: string) {
-  const match = `${resolve(instanceId)}:`
-  for (const key of scrollCache.keys()) {
-    if (key.startsWith(match)) {
-      scrollCache.delete(key)
-    }
-  }
-}
-
-export function clearAllScrollCache() {
-  scrollCache.clear()
-}

packages/ui/src/lib/tool-render-cache.ts

@@ -1,22 +0,0 @@
-import type { RenderCache } from "../types/message"
-
-const toolRenderCache = new Map<string, RenderCache>()
-
-export function getToolRenderCache(key?: string | null): RenderCache | undefined {
-  if (!key) return undefined
-  return toolRenderCache.get(key)
-}
-
-export function setToolRenderCache(key: string | undefined | null, cache?: RenderCache): void {
-  if (!key) return
-  if (cache) {
-    toolRenderCache.set(key, cache)
-  } else {
-    toolRenderCache.delete(key)
-  }
-}
-
-export function clearToolRenderCache(key?: string | null): void {
-  if (!key) return
-  toolRenderCache.delete(key)
-}