From b5cefd25ef376f2140a1b42d5ca1c276310db68e Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Mon, 13 Jul 2026 22:17:38 +0200 Subject: [PATCH] feat(chat): add module-level chat stream store for resumable streaming --- surfsense_web/lib/chat/stream-engine/store.ts | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 surfsense_web/lib/chat/stream-engine/store.ts diff --git a/surfsense_web/lib/chat/stream-engine/store.ts b/surfsense_web/lib/chat/stream-engine/store.ts new file mode 100644 index 000000000..44627663f --- /dev/null +++ b/surfsense_web/lib/chat/stream-engine/store.ts @@ -0,0 +1,177 @@ +import type { ThreadMessageLike } from "@assistant-ui/react"; +import { createTokenUsageStore } from "@/components/assistant-ui/token-usage-context"; +import type { PendingInterruptState } from "@/features/chat-messages/hitl"; + +/** + * Durable, per-thread streaming state for a single in-flight chat turn. + * + * Lives at module scope (see {@link chatStreamStore}) so a running turn's + * messages / interrupts survive the chat page unmounting during in-app + * navigation. The React tree consumes it through ``useChatStream`` via + * ``useSyncExternalStore`` (state lives outside the render cycle). + */ +export interface ThreadStreamState { + threadId: number; + messages: ThreadMessageLike[]; + isRunning: boolean; + pendingInterrupts: PendingInterruptState[]; +} + +type Listener = () => void; + +/** + * Module-level singleton external store for chat streaming. + * + * Single-active-stream invariant: at most one turn streams at a time, + * tracked by {@link active}. Per-thread state is still keyed by threadId so + * the currently-streaming thread and the currently-viewed thread can differ + * (e.g. a stream finishes while the user is on another chat). + * + * ponytail: unbounded ``states`` map is bounded in practice by + * ``clearInactive`` (called on navigation) + ``clear`` (called after the DB + * re-hydrates a finished turn). Upgrade path if that ever leaks: LRU cap. + */ +class ChatStreamStore { + private states = new Map(); + private listeners = new Set(); + + /** Shared, cross-navigation token-usage store (one instance app-wide). */ + readonly tokenUsage = createTokenUsageStore(); + + /** The one in-flight turn's abort handle, or null when idle. */ + private active: { threadId: number; controller: AbortController } | null = null; + + /** Timestamp of the last explicit cancel, for the THREAD_BUSY retry window. */ + recentCancelRequestedAt = 0; + + subscribe = (listener: Listener): (() => void) => { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + }; + + private notify(): void { + for (const l of this.listeners) l(); + } + + /** Snapshot for ``useSyncExternalStore``; stable ref between mutations. */ + getSnapshot = (threadId: number | null): ThreadStreamState | null => { + if (threadId == null) return null; + return this.states.get(threadId) ?? null; + }; + + isRunning(threadId: number | null): boolean { + if (threadId == null) return false; + return this.states.get(threadId)?.isRunning ?? false; + } + + getMessages(threadId: number): ThreadMessageLike[] { + return this.states.get(threadId)?.messages ?? []; + } + + getPendingInterrupts(threadId: number): PendingInterruptState[] { + return this.states.get(threadId)?.pendingInterrupts ?? []; + } + + private ensure(threadId: number): ThreadStreamState { + let s = this.states.get(threadId); + if (!s) { + s = { threadId, messages: [], isRunning: false, pendingInterrupts: [] }; + this.states.set(threadId, s); + } + return s; + } + + private commit(threadId: number, next: ThreadStreamState): void { + this.states.set(threadId, next); + this.notify(); + } + + /** Seed a thread's state at the start of a fresh turn (running=true). */ + begin(threadId: number, messages: ThreadMessageLike[]): void { + this.commit(threadId, { threadId, messages, isRunning: true, pendingInterrupts: [] }); + } + + setMessages(threadId: number, updater: (prev: ThreadMessageLike[]) => ThreadMessageLike[]): void { + const prev = this.ensure(threadId); + const messages = updater(prev.messages); + if (messages === prev.messages) return; + this.commit(threadId, { ...prev, messages }); + } + + setRunning(threadId: number, running: boolean): void { + const prev = this.ensure(threadId); + if (prev.isRunning === running) return; + this.commit(threadId, { ...prev, isRunning: running }); + } + + setPendingInterrupts( + threadId: number, + updater: (prev: PendingInterruptState[]) => PendingInterruptState[] + ): void { + const prev = this.ensure(threadId); + const pendingInterrupts = updater(prev.pendingInterrupts); + if (pendingInterrupts === prev.pendingInterrupts) return; + this.commit(threadId, { ...prev, pendingInterrupts }); + } + + /** + * A thread whose overlay must survive DB re-hydration / navigation: it is + * either streaming or paused awaiting a HITL decision (the pending + * interrupts + interrupt cards live only in the overlay). + */ + private isPinned(s: ThreadStreamState): boolean { + return s.isRunning || s.pendingInterrupts.length > 0; + } + + /** Drop a thread's overlay once the DB is authoritative. No-op while pinned. */ + clear(threadId: number): void { + const s = this.states.get(threadId); + if (!s || this.isPinned(s)) return; + this.states.delete(threadId); + this.notify(); + } + + /** Evict every non-pinned thread except ``exceptThreadId`` (memory bound). */ + clearInactive(exceptThreadId: number | null): void { + let changed = false; + for (const [id, s] of this.states) { + if (id === exceptThreadId || this.isPinned(s)) continue; + this.states.delete(id); + changed = true; + } + if (changed) this.notify(); + } + + // ---- active-stream lifecycle ------------------------------------------- + + /** Register a new in-flight turn, aborting any previous one first. */ + beginActive(threadId: number, controller: AbortController): void { + this.abortActive(); + this.active = { threadId, controller }; + } + + get activeThreadId(): number | null { + return this.active?.threadId ?? null; + } + + /** Clear the active handle iff it still points at ``controller``. */ + clearActive(controller: AbortController): void { + if (this.active?.controller === controller) this.active = null; + } + + /** Abort the in-flight turn's fetch (client disconnect, not a server stop). */ + abortActive(): void { + if (this.active) { + this.active.controller.abort(); + this.active = null; + } + } + + markRecentCancel(): void { + this.recentCancelRequestedAt = Date.now(); + } +} + +export const chatStreamStore = new ChatStreamStore();