From f24e360d78fa6c3c7fedd7110795d72d76f7a906 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Tue, 6 Jan 2026 09:58:55 +0000 Subject: [PATCH] Optimize session status updates Reduce per-token store churn by updating status on transitions, caching instance-level indicators, and avoiding O(n) session-map cloning. --- packages/ui/src/stores/session-api.ts | 40 +++-- packages/ui/src/stores/session-events.ts | 79 +++++++--- packages/ui/src/stores/session-state.ts | 191 ++++++++++++++++++++--- packages/ui/src/stores/session-status.ts | 39 +---- 4 files changed, 261 insertions(+), 88 deletions(-) diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index 39a1853c..913d6b91 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/packages/ui/src/stores/session-api.ts @@ -24,6 +24,7 @@ import { loading, setLoading, cleanupBlankSessions, + syncInstanceSessionIndicator, } from "./session-state" import { DEFAULT_MODEL_OUTPUT_LIMIT, getDefaultModel, isModelValid } from "./session-models" import { normalizeMessagePart } from "./message-v2/normalizers" @@ -120,6 +121,8 @@ async function fetchSessions(instanceId: string): Promise { return next }) + syncInstanceSessionIndicator(instanceId, sessionMap) + setMessagesLoaded((prev) => { const next = new Map(prev) const loadedSet = next.get(instanceId) @@ -214,6 +217,8 @@ async function createSession(instanceId: string, agent?: string): Promise p.id === session.model.providerId) const initialModel = initialProvider?.models.find((m) => m.id === session.model.modelId) @@ -311,6 +316,8 @@ async function forkSession( return next }) + syncInstanceSessionIndicator(instanceId) + const instanceProviders = providers().get(instanceId) || [] const forkProvider = instanceProviders.find((p) => p.id === forkedSession.model.providerId) const forkModel = forkProvider?.models.find((m) => m.id === forkedSession.model.modelId) @@ -364,10 +371,15 @@ async function deleteSession(instanceId: string, sessionId: string): Promise { const next = new Map(prev) const nextInstanceSessions = next.get(instanceId) - if (nextInstanceSessions) { - const existingSession = nextInstanceSessions.get(sessionId) - if (existingSession) { - const updatedSession = { - ...existingSession, - agent: agentName || existingSession.agent, - model: providerID && modelID ? { providerId: providerID, modelId: modelID } : existingSession.model, - } - const updatedInstanceSessions = new Map(nextInstanceSessions) - updatedInstanceSessions.set(sessionId, updatedSession) - next.set(instanceId, updatedInstanceSessions) - } + if (!nextInstanceSessions) { + return next } + + const existingSession = nextInstanceSessions.get(sessionId) + if (!existingSession) { + return next + } + + const updatedSession = { + ...existingSession, + agent: agentName || existingSession.agent, + model: providerID && modelID ? { providerId: providerID, modelId: modelID } : existingSession.model, + } + + nextInstanceSessions.set(sessionId, updatedSession) + next.set(instanceId, nextInstanceSessions) return next }) diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts index f5427cfe..bfb49e3c 100644 --- a/packages/ui/src/stores/session-events.ts +++ b/packages/ui/src/stores/session-events.ts @@ -22,7 +22,7 @@ import { showToastNotification, ToastVariant } from "../lib/notifications" import { instances, addPermissionToQueue, removePermissionFromQueue } from "./instances" import { showAlertDialog } from "./alerts" import { createClientSession, Session, SessionStatus } from "../types/session" -import { sessions, setSessions, withSession } from "./session-state" +import { sessions, setSessions, syncInstanceSessionIndicator, withSession } from "./session-state" import { normalizeMessagePart } from "./message-v2/normalizers" import { updateSessionInfo } from "./message-v2/session-info" @@ -66,10 +66,19 @@ const mapSdkSessionStatus = (status: EventSessionStatus["properties"]["status"]) return "working" } -function applySessionStatus(instanceId: string, sessionId: string, status: SessionStatus, bumpUpdated = false) { +function applySessionStatus( + instanceId: string, + sessionId: string, + status: SessionStatus, + bumpUpdatedOnTransition = false, +) { withSession(instanceId, sessionId, (session) => { + const current = session.status ?? "idle" + if (current === status) return false + session.status = status - if (bumpUpdated) { + + if (bumpUpdatedOnTransition) { session.time = { ...(session.time ?? {}), updated: Date.now() } } }) @@ -87,21 +96,27 @@ async function fetchSessionInfo(instanceId: string, sessionId: string): Promise< const fetched = createClientSession(info, instanceId) + let updatedInstanceSessions: Map | undefined + setSessions((prev) => { const next = new Map(prev) - const instanceSessions = new Map(next.get(instanceId) ?? []) + const instanceSessions = next.get(instanceId) ?? new Map() const existing = instanceSessions.get(sessionId) - instanceSessions.set(sessionId, { + const merged: Session = { ...fetched, agent: existing?.agent ?? fetched.agent, model: existing?.model ?? fetched.model, status: existing?.status ?? fetched.status, pendingPermission: existing?.pendingPermission ?? fetched.pendingPermission, - }) + } + instanceSessions.set(sessionId, merged) next.set(instanceId, instanceSessions) + updatedInstanceSessions = instanceSessions return next }) + syncInstanceSessionIndicator(instanceId, updatedInstanceSessions) + return fetched } catch (error) { log.error("Failed to fetch session info", error) @@ -109,11 +124,19 @@ async function fetchSessionInfo(instanceId: string, sessionId: string): Promise< } } -function ensureSessionStatus(instanceId: string, sessionId: string, status: SessionStatus, bumpUpdated = false) { +function ensureSessionStatus( + instanceId: string, + sessionId: string, + status: SessionStatus, + bumpUpdatedOnTransition = false, +) { const instanceSessions = sessions().get(instanceId) const existing = instanceSessions?.get(sessionId) if (existing) { - applySessionStatus(instanceId, sessionId, status, bumpUpdated) + if ((existing.status ?? "idle") === status) { + return + } + applySessionStatus(instanceId, sessionId, status, bumpUpdatedOnTransition) return } @@ -125,7 +148,7 @@ function ensureSessionStatus(instanceId: string, sessionId: string, status: Sess const pending = (async () => { const fetched = await fetchSessionInfo(instanceId, sessionId) if (!fetched) return - applySessionStatus(instanceId, sessionId, status, bumpUpdated) + applySessionStatus(instanceId, sessionId, status, bumpUpdatedOnTransition) })() pendingSessionFetches.set(key, pending) @@ -170,14 +193,16 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes const sessionId = typeof part.sessionID === "string" ? part.sessionID : fallbackSessionId const messageId = typeof part.messageID === "string" ? part.messageID : fallbackMessageId if (!sessionId || !messageId) return - const session = instanceSessions?.get(sessionId) if (!session) { ensureSessionStatus(instanceId, sessionId, "working", true) return } - applySessionStatus(instanceId, sessionId, "working", true) + if (session.status !== "working" && session.status !== "compacting") { + applySessionStatus(instanceId, sessionId, "working", true) + } + const store = messageStoreBus.getOrCreate(instanceId) const role: MessageRole = resolveMessageRole(messageInfo) @@ -227,7 +252,9 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes return } - applySessionStatus(instanceId, sessionId, "working", true) + if (session.status !== "working" && session.status !== "compacting") { + applySessionStatus(instanceId, sessionId, "working", true) + } const store = messageStoreBus.getOrCreate(instanceId) @@ -297,13 +324,18 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo }, } as Session + let updatedInstanceSessions: Map | undefined + setSessions((prev) => { const next = new Map(prev) - const updated = new Map(prev.get(instanceId)) - updated.set(newSession.id, newSession) - next.set(instanceId, updated) + const instanceSessions = next.get(instanceId) ?? new Map() + instanceSessions.set(newSession.id, newSession) + next.set(instanceId, instanceSessions) + updatedInstanceSessions = instanceSessions return next }) + + syncInstanceSessionIndicator(instanceId, updatedInstanceSessions) setSessionRevertV2(instanceId, info.id, info.revert ?? null) log.info(`[SSE] New session created: ${info.id}`, newSession) @@ -331,13 +363,18 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo : existingSession.revert, } + let updatedInstanceSessions: Map | undefined + setSessions((prev) => { const next = new Map(prev) - const updated = new Map(prev.get(instanceId)) - updated.set(existingSession.id, updatedSession) - next.set(instanceId, updated) + const instanceSessions = next.get(instanceId) ?? new Map() + instanceSessions.set(existingSession.id, updatedSession) + next.set(instanceId, instanceSessions) + updatedInstanceSessions = instanceSessions return next }) + + syncInstanceSessionIndicator(instanceId, updatedInstanceSessions) setSessionRevertV2(instanceId, info.id, info.revert ?? null) } } @@ -346,7 +383,7 @@ function handleSessionIdle(instanceId: string, event: EventSessionIdle): void { const sessionId = event.properties?.sessionID if (!sessionId) return - ensureSessionStatus(instanceId, sessionId, "idle") + ensureSessionStatus(instanceId, sessionId, "idle", true) log.info(`[SSE] Session idle: ${sessionId}`) } @@ -355,7 +392,7 @@ function handleSessionStatus(instanceId: string, event: EventSessionStatus): voi if (!sessionId) return const status = mapSdkSessionStatus(event.properties.status) - ensureSessionStatus(instanceId, sessionId, status, status === "working") + ensureSessionStatus(instanceId, sessionId, status, true) log.info(`[SSE] Session status updated: ${sessionId}`, { status }) } @@ -366,7 +403,7 @@ function handleSessionCompacted(instanceId: string, event: EventSessionCompacted log.info(`[SSE] Session compacted: ${sessionID}`) setSessionCompactionState(instanceId, sessionID, false) - ensureSessionStatus(instanceId, sessionID, "idle") + ensureSessionStatus(instanceId, sessionID, "idle", true) withSession(instanceId, sessionID, (session) => { const time = { ...(session.time ?? {}) } diff --git a/packages/ui/src/stores/session-state.ts b/packages/ui/src/stores/session-state.ts index 0150e58b..be89a6f2 100644 --- a/packages/ui/src/stores/session-state.ts +++ b/packages/ui/src/stores/session-state.ts @@ -40,6 +40,130 @@ const [loading, setLoading] = createSignal({ const [messagesLoaded, setMessagesLoaded] = createSignal>>(new Map()) const [sessionInfoByInstance, setSessionInfoByInstance] = createSignal>>(new Map()) +export type InstanceSessionIndicatorStatus = "permission" | SessionStatus + +type InstanceIndicatorCounts = { + permission: number + working: number + compacting: number +} + +const [instanceIndicatorCounts, setInstanceIndicatorCounts] = createSignal>(new Map()) + +function getIndicatorBucket(session: Pick): InstanceSessionIndicatorStatus | "idle" { + if (session.pendingPermission) { + return "permission" + } + const status = session.status ?? "idle" + return status +} + +function adjustIndicatorCounts( + instanceId: string, + previous: InstanceSessionIndicatorStatus | "idle", + next: InstanceSessionIndicatorStatus | "idle", +): void { + if (previous === next) return + + const decKey = previous === "idle" ? null : previous + const incKey = next === "idle" ? null : next + + setInstanceIndicatorCounts((prev) => { + const current = prev.get(instanceId) ?? { permission: 0, working: 0, compacting: 0 } + const updated: InstanceIndicatorCounts = { ...current } + + if (decKey) { + updated[decKey] = Math.max(0, updated[decKey] - 1) + } + + if (incKey) { + updated[incKey] = updated[incKey] + 1 + } + + const hasAny = updated.permission > 0 || updated.working > 0 || updated.compacting > 0 + if (!hasAny) { + if (!prev.has(instanceId)) return prev + const nextMap = new Map(prev) + nextMap.delete(instanceId) + return nextMap + } + + const same = + current.permission === updated.permission && + current.working === updated.working && + current.compacting === updated.compacting + if (same && prev.has(instanceId)) { + return prev + } + + const nextMap = new Map(prev) + nextMap.set(instanceId, updated) + return nextMap + }) +} + +function recomputeIndicatorCounts(instanceId: string, instanceSessions: Map | undefined): void { + if (!instanceSessions || instanceSessions.size === 0) { + setInstanceIndicatorCounts((prev) => { + if (!prev.has(instanceId)) return prev + const next = new Map(prev) + next.delete(instanceId) + return next + }) + return + } + + let permission = 0 + let working = 0 + let compacting = 0 + + for (const session of instanceSessions.values()) { + if (session.pendingPermission) { + permission += 1 + continue + } + const status = session.status ?? "idle" + if (status === "compacting") { + compacting += 1 + } else if (status === "working") { + working += 1 + } + } + + if (permission === 0 && working === 0 && compacting === 0) { + setInstanceIndicatorCounts((prev) => { + if (!prev.has(instanceId)) return prev + const next = new Map(prev) + next.delete(instanceId) + return next + }) + return + } + + setInstanceIndicatorCounts((prev) => { + const current = prev.get(instanceId) + if (current && current.permission === permission && current.working === working && current.compacting === compacting) { + return prev + } + const next = new Map(prev) + next.set(instanceId, { permission, working, compacting }) + return next + }) +} + +export function getInstanceSessionIndicatorStatusCached(instanceId: string): InstanceSessionIndicatorStatus { + const counts = instanceIndicatorCounts().get(instanceId) + if (!counts) return "idle" + if (counts.permission > 0) return "permission" + if (counts.compacting > 0) return "compacting" + if (counts.working > 0) return "working" + return "idle" +} + +export function syncInstanceSessionIndicator(instanceId: string, instanceSessions?: Map): void { + recomputeIndicatorCounts(instanceId, instanceSessions ?? sessions().get(instanceId)) +} + function clearLoadedFlag(instanceId: string, sessionId: string) { if (!instanceId || !sessionId) return setMessagesLoaded((prev) => { @@ -131,33 +255,63 @@ function pruneDraftPrompts(instanceId: string, validSessionIds: Set) { }) } -function withSession(instanceId: string, sessionId: string, updater: (session: Session) => void) { - const instanceSessions = sessions().get(instanceId) - if (!instanceSessions) return - - const session = instanceSessions.get(sessionId) - if (!session) return - - updater(session) - - const updatedSession = { - ...session, - } +function withSession(instanceId: string, sessionId: string, updater: (session: Session) => void | boolean) { + let previousBucket: InstanceSessionIndicatorStatus | "idle" | null = null + let nextBucket: InstanceSessionIndicatorStatus | "idle" | null = null + let didUpdate = false setSessions((prev) => { + const instanceSessions = prev.get(instanceId) + if (!instanceSessions) return prev + + const current = instanceSessions.get(sessionId) + if (!current) return prev + + previousBucket = getIndicatorBucket(current) + + const updatedSession: Session = { ...current } + const result = updater(updatedSession) + if (result === false) { + return prev + } + + nextBucket = getIndicatorBucket(updatedSession) + + instanceSessions.set(sessionId, updatedSession) + didUpdate = true + const next = new Map(prev) - const newInstanceSessions = new Map(instanceSessions) - newInstanceSessions.set(sessionId, updatedSession) - next.set(instanceId, newInstanceSessions) + next.set(instanceId, instanceSessions) return next }) + + if (didUpdate && previousBucket && nextBucket) { + adjustIndicatorCounts(instanceId, previousBucket, nextBucket) + } } function setSessionCompactionState(instanceId: string, sessionId: string, isCompacting: boolean): void { withSession(instanceId, sessionId, (session) => { - const time = { ...(session.time ?? {}) } - time.compacting = isCompacting ? Date.now() : 0 + const time = { ...(session.time ?? {}) } as Session["time"] & { compacting?: number | boolean; updated?: number } + const compactingFlag = time.compacting + const wasCompacting = typeof compactingFlag === "number" ? compactingFlag > 0 : Boolean(compactingFlag) + + const shouldAlreadyBeCompacting = isCompacting + const isAlreadyCorrect = + wasCompacting === isCompacting && + (shouldAlreadyBeCompacting ? session.status === "compacting" : session.status !== "compacting") + + if (isAlreadyCorrect) { + return false + } + + if (wasCompacting !== isCompacting) { + time.compacting = isCompacting ? Date.now() : 0 + time.updated = Date.now() + } + session.time = time + if (isCompacting) { session.status = "compacting" } else if (session.status === "compacting") { @@ -168,7 +322,7 @@ function setSessionCompactionState(instanceId: string, sessionId: string, isComp function setSessionPendingPermission(instanceId: string, sessionId: string, pending: boolean): void { withSession(instanceId, sessionId, (session) => { - if (session.pendingPermission === pending) return + if (session.pendingPermission === pending) return false session.pendingPermission = pending }) } @@ -207,6 +361,7 @@ function clearActiveParentSession(instanceId: string): void { function setSessionStatus(instanceId: string, sessionId: string, status: SessionStatus): void { withSession(instanceId, sessionId, (session) => { + if (session.status === status) return false session.status = status }) } diff --git a/packages/ui/src/stores/session-status.ts b/packages/ui/src/stores/session-status.ts index f895441b..f5b30121 100644 --- a/packages/ui/src/stores/session-status.ts +++ b/packages/ui/src/stores/session-status.ts @@ -1,7 +1,7 @@ import type { Session, SessionStatus } from "../types/session" import type { MessageInfo } from "../types/message" import type { MessageRecord } from "./message-v2/types" -import { sessions } from "./sessions" +import { getInstanceSessionIndicatorStatusCached, sessions } from "./session-state" import { isSessionCompactionActive } from "./session-compaction" import { messageStoreBus } from "./message-v2/bus" @@ -173,42 +173,7 @@ export function getSessionStatus(instanceId: string, sessionId: string): Session export type InstanceSessionIndicatorStatus = "permission" | SessionStatus export function getInstanceSessionIndicatorStatus(instanceId: string): InstanceSessionIndicatorStatus { - const instanceSessions = sessions().get(instanceId) - if (!instanceSessions || instanceSessions.size === 0) { - return "idle" - } - - let bestRank = 0 - let best: InstanceSessionIndicatorStatus = "idle" - - for (const session of instanceSessions.values()) { - let rank = 0 - let status: InstanceSessionIndicatorStatus = "idle" - - if (session.pendingPermission) { - status = "permission" - rank = 3 - } else { - const sessionStatus = getSessionStatus(instanceId, session.id) - if (sessionStatus === "compacting") { - status = "compacting" - rank = 2 - } else if (sessionStatus === "working") { - status = "working" - rank = 1 - } - } - - if (rank > bestRank) { - bestRank = rank - best = status - if (bestRank === 3) { - break - } - } - } - - return best + return getInstanceSessionIndicatorStatusCached(instanceId) } export function isSessionBusy(instanceId: string, sessionId: string): boolean {