salvacybersec/CodeNomad
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize session status updates

Browse Source 
Reduce per-token store churn by updating status on transitions, caching instance-level indicators, and avoiding O(n) session-map cloning.
This commit is contained in:
Shantur Rathore
2026-01-06 09:58:55 +00:00
parent 1a6f1fdbae
commit f24e360d78
4 changed files with 261 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
packages/ui/src/stores/session-api.ts
View File

@@ -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<void> {
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<Sessio
return next
})
syncInstanceSessionIndicator(instanceId)
const instanceProviders = providers().get(instanceId) || []
const initialProvider = instanceProviders.find((p) => 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<voi
const instanceSessions = next.get(instanceId)
if (instanceSessions) {
instanceSessions.delete(sessionId)
if (instanceSessions.size === 0) {
next.delete(instanceId)
}
}
return next
})
syncInstanceSessionIndicator(instanceId)
setSessionCompactionState(instanceId, sessionId, false)
clearSessionDraftPrompt(instanceId, sessionId)
@@ -578,19 +590,23 @@ async function loadMessages(instanceId: string, sessionId: string, force = false
setSessions((prev) => {
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
})

79
packages/ui/src/stores/session-events.ts
View File

@@ -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<string, Session> | undefined
setSessions((prev) => {
const next = new Map(prev)
const instanceSessions = new Map(next.get(instanceId) ?? [])
const instanceSessions = next.get(instanceId) ?? new Map<string, Session>()
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<string, Session> | 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<string, Session>()
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<string, Session> | 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<string, Session>()
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 ?? {}) }

191
packages/ui/src/stores/session-state.ts
View File

@@ -40,6 +40,130 @@ 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())
export type InstanceSessionIndicatorStatus = "permission" | SessionStatus
type InstanceIndicatorCounts = {
permission: number
working: number
compacting: number
}
const [instanceIndicatorCounts, setInstanceIndicatorCounts] = createSignal<Map<string, InstanceIndicatorCounts>>(new Map())
function getIndicatorBucket(session: Pick<Session, "status" | "pendingPermission">): 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<string, Session> | 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<string, Session>): 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<string>) {
})
}
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
})
}

39
packages/ui/src/stores/session-status.ts
View File

@@ -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 {