Batch hydrate normalized messages for session load

packages/ui/src/stores/instances.ts

@@ -13,7 +13,6 @@ import {
   fetchSessions,
   fetchAgents,
   fetchProviders,
-  removeSessionIndexes,
   clearInstanceDraftPrompts,
 } from "./sessions"
 import { fetchCommands, clearCommands } from "./commands"
@@ -21,6 +20,7 @@ import { preferences } from "./preferences"
 import { setSessionPendingPermission } from "./session-state"
 import { setHasInstances } from "./ui"
 import { messageStoreBus } from "./message-v2/bus"
+import { clearScrollCacheForInstance } from "../lib/scroll-cache"
 import type { MessageRecord } from "./message-v2/types"
 
 
@@ -296,7 +296,8 @@ function removeInstance(id: string) {
   }
 
   // Clean up session indexes and drafts for removed instance
-  removeSessionIndexes(id)
+  clearScrollCacheForInstance(id)
+  messageStoreBus.unregisterInstance(id)
   clearInstanceDraftPrompts(id)
 }
 

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

@@ -41,37 +41,27 @@ export function seedSessionMessagesV2(
   if (!session || !Array.isArray(messages)) return
   const store = messageStoreBus.getOrCreate(instanceId)
   const metadata: SessionMetadata = "id" in session ? { id: session.id, title: session.title, parentId: session.parentId ?? null } : session
-  const messageIds = messages.map((message) => message.id)
 
   store.addOrUpdateSession({
     id: metadata.id,
     title: metadata.title,
     parentId: metadata.parentId ?? null,
-    messageIds,
     revert: (session as Session)?.revert ?? undefined,
   })
 
-  messages.forEach((message) => {
-    store.upsertMessage({
-      id: message.id,
-      sessionId: message.sessionId,
-      role: message.type,
-      status: normalizeStatus(message.status),
-      createdAt: message.timestamp,
-      updatedAt: message.timestamp,
-      parts: message.parts,
-      isEphemeral: message.status === "sending" || message.status === "streaming",
-      bumpRevision: false,
-    })
-    const info = messageInfos?.get(message.id)
-    if (info) {
-      store.setMessageInfo(message.id, info)
-    }
-  })
+  const normalizedMessages = messages.map((message) => ({
+    id: message.id,
+    sessionId: message.sessionId,
+    role: message.type,
+    status: normalizeStatus(message.status),
+    createdAt: message.timestamp,
+    updatedAt: message.timestamp,
+    parts: message.parts,
+    isEphemeral: message.status === "sending" || message.status === "streaming",
+    bumpRevision: false,
+  }))
 
-  if (messageInfos) {
-    store.rebuildUsage(metadata.id, messageInfos.values())
-  }
+  store.hydrateMessages(metadata.id, normalizedMessages, messageInfos?.values())
 }
 
 interface MessageInfoOptions {

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

@@ -141,6 +141,7 @@ export interface InstanceMessageStore {
   state: InstanceMessageState
   setState: SetStoreFunction<InstanceMessageState>
   addOrUpdateSession: (input: SessionUpsertInput) => void
+  hydrateMessages: (sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) => void
   upsertMessage: (input: MessageUpsertInput) => void
   applyPartUpdate: (input: PartUpdateInput) => void
   bufferPendingPart: (entry: PendingPartEntry) => void
@@ -234,6 +235,81 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
     })
   }
 
+  function hydrateMessages(sessionId: string, inputs: MessageUpsertInput[], infos?: Iterable<MessageInfo>) {
+    if (!Array.isArray(inputs) || inputs.length === 0) return
+
+    ensureSessionEntry(sessionId)
+
+    const incomingIds = inputs.map((item) => item.id)
+    const incomingIdSet = new Set(incomingIds)
+    const existingIds = state.sessions[sessionId]?.messageIds ?? []
+    const removedIds = existingIds.filter((id) => !incomingIdSet.has(id))
+
+    const normalizedRecords: Record<string, MessageRecord> = {}
+    const now = Date.now()
+
+    inputs.forEach((input) => {
+      const normalizedParts = normalizeParts(input.id, input.parts)
+      const shouldBump = Boolean(input.bumpRevision || normalizedParts)
+      const previous = state.messages[input.id]
+      normalizedRecords[input.id] = {
+        id: input.id,
+        sessionId: input.sessionId,
+        role: input.role,
+        status: input.status,
+        createdAt: input.createdAt ?? previous?.createdAt ?? now,
+        updatedAt: input.updatedAt ?? now,
+        isEphemeral: input.isEphemeral ?? previous?.isEphemeral ?? false,
+        revision: previous ? previous.revision + (shouldBump ? 1 : 0) : 0,
+        partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
+        parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
+      }
+    })
+
+    const infoList = infos ? Array.from(infos) : undefined
+    const usageState = infoList ? rebuildUsageStateFromInfos(infoList) : state.usage[sessionId]
+
+    setState(
+      produce((draft) => {
+        removedIds.forEach((id) => {
+          if (draft.messages[id]?.sessionId === sessionId) {
+            delete draft.messages[id]
+            delete draft.messageInfoVersion[id]
+            delete draft.pendingParts[id]
+            if (draft.permissions.byMessage[id]) {
+              delete draft.permissions.byMessage[id]
+            }
+          }
+        })
+
+        Object.entries(normalizedRecords).forEach(([id, record]) => {
+          draft.messages[id] = record as MessageRecord
+        })
+
+        const session = draft.sessions[sessionId]!
+        session.messageIds = incomingIds
+        session.updatedAt = Date.now()
+
+        if (usageState) {
+          draft.usage[sessionId] = usageState
+        }
+
+        if (infoList) {
+          for (const info of infoList) {
+            const messageId = info.id as string
+            messageInfoCache.set(messageId, info)
+            const currentVersion = draft.messageInfoVersion[messageId] ?? 0
+            draft.messageInfoVersion[messageId] = currentVersion + 1
+          }
+        }
+      }),
+    )
+
+    removedIds.forEach((id) => {
+      messageInfoCache.delete(id)
+    })
+  }
+
   function insertMessageIntoSession(sessionId: string, messageId: string) {
     ensureSessionEntry(sessionId)
     setState("sessions", sessionId, "messageIds", (ids = []) => {
@@ -508,6 +584,7 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
     state,
     setState,
     addOrUpdateSession,
+    hydrateMessages,
     upsertMessage,
     applyPartUpdate,
     bufferPendingPart,

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

@@ -0,0 +1,96 @@
+import { decodeHtmlEntities } from "../../lib/markdown"
+import { partHasRenderableText } from "../../types/message"
+import type { MessageDisplayParts, Message } from "../../types/message"
+
+function decodeTextSegment(segment: any): any {
+  if (typeof segment === "string") {
+    return decodeHtmlEntities(segment)
+  }
+
+  if (segment && typeof segment === "object") {
+    const updated: Record<string, any> = { ...segment }
+
+    if (typeof updated.text === "string") {
+      updated.text = decodeHtmlEntities(updated.text)
+    }
+
+    if (typeof updated.value === "string") {
+      updated.value = decodeHtmlEntities(updated.value)
+    }
+
+    if (Array.isArray(updated.content)) {
+      updated.content = updated.content.map((item: any) => decodeTextSegment(item))
+    }
+
+    return updated
+  }
+
+  return segment
+}
+
+export function normalizeMessagePart(part: any): any {
+  if (!part || typeof part !== "object") {
+    return part
+  }
+
+  if (part.type !== "text") {
+    return part
+  }
+
+  const normalized: Record<string, any> = { ...part, renderCache: undefined }
+
+  if (typeof normalized.text === "string") {
+    normalized.text = decodeHtmlEntities(normalized.text)
+  } else if (normalized.text && typeof normalized.text === "object") {
+    const textObject: Record<string, any> = { ...normalized.text }
+
+    if (typeof textObject.value === "string") {
+      textObject.value = decodeHtmlEntities(textObject.value)
+    }
+
+    if (Array.isArray(textObject.content)) {
+      textObject.content = textObject.content.map((item: any) => decodeTextSegment(item))
+    }
+
+    if (typeof textObject.text === "string") {
+      textObject.text = decodeHtmlEntities(textObject.text)
+    }
+
+    normalized.text = textObject
+  }
+
+  if (Array.isArray(normalized.content)) {
+    normalized.content = normalized.content.map((item: any) => decodeTextSegment(item))
+  }
+
+  if (normalized.thinking && typeof normalized.thinking === "object") {
+    const thinking: Record<string, any> = { ...normalized.thinking }
+    if (Array.isArray(thinking.content)) {
+      thinking.content = thinking.content.map((item: any) => decodeTextSegment(item))
+    }
+    normalized.thinking = thinking
+  }
+
+  return normalized
+}
+
+export function computeDisplayParts(message: Message, showThinking: boolean): MessageDisplayParts {
+  const text: any[] = []
+  const tool: any[] = []
+  const reasoning: any[] = []
+
+  for (const part of message.parts) {
+    if (part.type === "text" && !part.synthetic && partHasRenderableText(part)) {
+      text.push(part)
+    } else if (part.type === "tool") {
+      tool.push(part)
+    } else if (part.type === "reasoning" && showThinking && partHasRenderableText(part)) {
+      reasoning.push(part)
+    }
+  }
+
+  const combined = reasoning.length > 0 ? [...text, ...reasoning] : [...text]
+  const version = typeof message.version === "number" ? message.version : 0
+
+  return { text, tool, reasoning, combined, showThinking, version }
+}

packages/ui/src/stores/message-v2/session-info.ts

@@ -0,0 +1,139 @@
+import type { Provider } from "../../types/session"
+import { DEFAULT_MODEL_OUTPUT_LIMIT } from "../session-models"
+import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance } from "../session-state"
+import { messageStoreBus } from "./bus"
+import type { SessionUsageState } from "./types"
+
+function getLatestUsageEntry(usage?: SessionUsageState) {
+  if (!usage?.latestMessageId) return undefined
+  return usage.entries[usage.latestMessageId]
+}
+
+function resolveSelectedModel(instanceProviders: Provider[], providerId?: string, modelId?: string) {
+  if (!providerId || !modelId) return undefined
+  const provider = instanceProviders.find((p) => p.id === providerId)
+  return provider?.models.find((m) => m.id === modelId)
+}
+
+export function updateSessionInfo(instanceId: string, sessionId: string): void {
+  const instanceSessions = sessions().get(instanceId)
+  if (!instanceSessions) return
+  const session = instanceSessions.get(sessionId)
+  if (!session) return
+
+  const store = messageStoreBus.getOrCreate(instanceId)
+  const usage = store.getSessionUsage(sessionId)
+  const hasUsageEntries = Boolean(usage && Object.keys(usage.entries).length > 0)
+
+  let totalInputTokens = usage?.totalInputTokens ?? 0
+  let totalOutputTokens = usage?.totalOutputTokens ?? 0
+  let totalReasoningTokens = usage?.totalReasoningTokens ?? 0
+  let totalCost = usage?.totalCost ?? 0
+  let actualUsageTokens = usage?.actualUsageTokens ?? 0
+
+  const latestEntry = getLatestUsageEntry(usage)
+  let latestHasContextUsage = latestEntry?.hasContextUsage ?? false
+
+  const previousInfo = sessionInfoByInstance().get(instanceId)?.get(sessionId)
+  let contextWindow = 0
+  let contextAvailableTokens: number | null = null
+  let contextAvailableFromPrevious = false
+  let isSubscriptionModel = false
+
+  if (!hasUsageEntries && previousInfo) {
+    totalInputTokens = previousInfo.inputTokens
+    totalOutputTokens = previousInfo.outputTokens
+    totalReasoningTokens = previousInfo.reasoningTokens
+    totalCost = previousInfo.cost
+    actualUsageTokens = previousInfo.actualUsageTokens
+  }
+
+  const instanceProviders = providers().get(instanceId) || []
+
+  const sessionModel = session.model
+  const sessionProviderId = sessionModel?.providerId
+  const sessionModelId = sessionModel?.modelId
+
+  const latestInfo = latestEntry?.messageId ? store.getMessageInfo(latestEntry.messageId) : undefined
+  const latestProviderId = (latestInfo as any)?.providerID || (latestInfo as any)?.providerId || ""
+  const latestModelId = (latestInfo as any)?.modelID || (latestInfo as any)?.modelId || ""
+
+  const selectedModel =
+    resolveSelectedModel(instanceProviders, sessionProviderId, sessionModelId) ??
+    resolveSelectedModel(instanceProviders, latestProviderId, latestModelId)
+
+  let modelOutputLimit = DEFAULT_MODEL_OUTPUT_LIMIT
+
+  if (selectedModel) {
+    contextWindow = selectedModel.limit?.context ?? 0
+    const outputLimit = selectedModel.limit?.output
+    if (typeof outputLimit === "number" && outputLimit > 0) {
+      modelOutputLimit = Math.min(outputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+    }
+    if ((selectedModel.cost?.input ?? 0) === 0 && (selectedModel.cost?.output ?? 0) === 0) {
+      isSubscriptionModel = true
+    }
+  }
+
+  if (contextWindow === 0 && previousInfo) {
+    contextWindow = previousInfo.contextWindow
+  }
+
+  modelOutputLimit = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+
+  if (previousInfo) {
+    const previousContextWindow = previousInfo.contextWindow
+    const previousContextAvailable = previousInfo.contextAvailableTokens ?? null
+    const previousHasContextUsage = previousContextAvailable !== null && previousContextWindow > 0
+      ? previousContextAvailable < previousContextWindow
+      : false
+
+    if (contextWindow !== previousContextWindow) {
+      contextAvailableTokens = null
+      contextAvailableFromPrevious = false
+      latestHasContextUsage = previousHasContextUsage
+    } else {
+      contextAvailableTokens = previousContextAvailable
+      contextAvailableFromPrevious = true
+      latestHasContextUsage = previousHasContextUsage
+    }
+
+    if (!hasUsageEntries) {
+      isSubscriptionModel = previousInfo.isSubscriptionModel
+    } else if (!isSubscriptionModel) {
+      isSubscriptionModel = previousInfo.isSubscriptionModel
+    }
+  }
+
+  const outputBudget = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
+
+  if (!contextAvailableFromPrevious) {
+    if (contextWindow > 0) {
+      if (latestHasContextUsage && actualUsageTokens > 0) {
+        contextAvailableTokens = Math.max(contextWindow - (actualUsageTokens + outputBudget), 0)
+      } else {
+        contextAvailableTokens = contextWindow
+      }
+    } else {
+      contextAvailableTokens = null
+    }
+  }
+
+  setSessionInfoByInstance((prev) => {
+    const next = new Map(prev)
+    const instanceInfo = new Map(prev.get(instanceId))
+    instanceInfo.set(sessionId, {
+      cost: totalCost,
+      contextWindow,
+      isSubscriptionModel,
+      inputTokens: totalInputTokens,
+      outputTokens: totalOutputTokens,
+      reasoningTokens: totalReasoningTokens,
+      actualUsageTokens,
+      modelOutputLimit,
+      contextAvailableTokens,
+    })
+    next.set(instanceId, instanceInfo)
+    return next
+  })
+}

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

@@ -1,4 +1,4 @@
-import type { ClientPart, MessageInfo } from "../../types/message"
+import type { ClientPart } from "../../types/message"
 import type { Permission } from "@opencode-ai/sdk"
 
 export type MessageStatus = "sending" | "sent" | "streaming" | "complete" | "error"

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

@@ -4,7 +4,7 @@ import { instances } from "./instances"
 import { addRecentModelPreference, setAgentModelPreference } from "./preferences"
 import { sessions, withSession } from "./session-state"
 import { getDefaultModel, isModelValid } from "./session-models"
-import { updateSessionInfo } from "./session-messages"
+import { updateSessionInfo } from "./message-v2/session-info"
 import { messageStoreBus } from "./message-v2/bus"
 
 const ID_LENGTH = 26

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

@@ -22,7 +22,8 @@ import {
   setLoading,
 } from "./session-state"
 import { DEFAULT_MODEL_OUTPUT_LIMIT, getDefaultModel, isModelValid } from "./session-models"
-import { normalizeMessagePart, updateSessionInfo } from "./session-messages"
+import { normalizeMessagePart } from "./message-v2/normalizers"
+import { updateSessionInfo } from "./message-v2/session-info"
 import { seedSessionMessagesV2 } from "./message-v2/bridge"
 
 interface SessionForkResponse {
@@ -92,8 +93,6 @@ async function fetchSessions(instanceId: string): Promise<void> {
               diff: apiSession.revert.diff,
             }
           : undefined,
-        messages: [],
-        messagesInfo: new Map(),
       })
     }
 
@@ -188,8 +187,6 @@ async function createSession(instanceId: string, agent?: string): Promise<Sessio
             diff: response.data.revert.diff,
           }
         : undefined,
-      messages: [],
-      messagesInfo: new Map(),
     }
 
     setSessions((prev) => {
@@ -291,8 +288,6 @@ async function forkSession(
           diff: info.revert.diff,
         }
       : undefined,
-    messages: [],
-    messagesInfo: new Map(),
   } as unknown as Session
 
   setSessions((prev) => {

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

@@ -22,7 +22,8 @@ import {
   setSessions,
   withSession,
 } from "./session-state"
-import { normalizeMessagePart, updateSessionInfo } from "./session-messages"
+import { normalizeMessagePart } from "./message-v2/normalizers"
+import { updateSessionInfo } from "./message-v2/session-info"
 import { loadMessages } from "./session-api"
 import { setSessionCompactionState } from "./session-compaction"
 import {
@@ -89,7 +90,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
     if (!session) return
 
     const store = messageStoreBus.getOrCreate(instanceId)
-    const messageInfo = event.properties?.message as MessageInfo | undefined
+    const messageInfo = (event as any)?.properties?.message as MessageInfo | undefined
     const role: MessageRole = resolveMessageRole(messageInfo)
     const createdAt = typeof messageInfo?.time?.created === "number" ? messageInfo.time.created : Date.now()
 
@@ -204,8 +205,6 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo
             created: Date.now(),
             updated: Date.now(),
           },
-      messages: [],
-      messagesInfo: new Map(),
     } as any
 
     setSessions((prev) => {

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

@@ -1,462 +0,0 @@
-import type { Message, MessageDisplayParts } from "../types/message"
-import { partHasRenderableText, type MessageInfo } from "../types/message"
-import type { Provider } from "../types/session"
-
-import { decodeHtmlEntities } from "../lib/markdown"
-import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance } from "./session-state"
-import { DEFAULT_MODEL_OUTPUT_LIMIT } from "./session-models"
-
-interface SessionIndexCache {
-  messageIndex: Map<string, number>
-  partIndex: Map<string, Map<string, number>>
-}
-
-interface AssistantUsageEntry {
-  info: MessageInfo
-  inputTokens: number
-  outputTokens: number
-  reasoningTokens: number
-  combinedTokens: number
-  cost: number
-  hasContextUsage: boolean
-  timestamp: number
-}
-
-interface SessionUsageState {
-  entries: Map<string, AssistantUsageEntry>
-  totalInputTokens: number
-  totalOutputTokens: number
-  totalReasoningTokens: number
-  totalCost: number
-  latestEntry: AssistantUsageEntry | null
-}
-
-const sessionIndexes = new Map<string, Map<string, SessionIndexCache>>()
-const sessionUsageStates = new Map<string, Map<string, SessionUsageState>>()
-
-function createEmptyUsageState(): SessionUsageState {
-  return {
-    entries: new Map(),
-    totalInputTokens: 0,
-    totalOutputTokens: 0,
-    totalReasoningTokens: 0,
-    totalCost: 0,
-    latestEntry: null,
-  }
-}
-
-function getUsageInstance(instanceId: string): Map<string, SessionUsageState> {
-  let usageMap = sessionUsageStates.get(instanceId)
-  if (!usageMap) {
-    usageMap = new Map()
-    sessionUsageStates.set(instanceId, usageMap)
-  }
-  return usageMap
-}
-
-function getSessionUsageState(instanceId: string, sessionId: string): SessionUsageState {
-  const usageMap = getUsageInstance(instanceId)
-  let state = usageMap.get(sessionId)
-  if (!state) {
-    state = createEmptyUsageState()
-    usageMap.set(sessionId, state)
-  }
-  return state
-}
-
-function recomputeLatestEntry(state: SessionUsageState) {
-  state.latestEntry = null
-  for (const entry of state.entries.values()) {
-    if (!state.latestEntry || entry.timestamp >= state.latestEntry.timestamp) {
-      state.latestEntry = entry
-    }
-  }
-}
-
-function extractAssistantUsage(info: MessageInfo): AssistantUsageEntry | null {
-  if (!info || info.role !== "assistant") return null
-  if (!info.tokens) return null
-  const tokens = info.tokens
-  const inputTokens = tokens.input ?? 0
-  const outputTokens = tokens.output ?? 0
-  const reasoningTokens = tokens.reasoning ?? 0
-  if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0) {
-    return null
-  }
-  const cacheReadTokens = tokens.cache?.read ?? 0
-  const cacheWriteTokens = tokens.cache?.write ?? 0
-  const combinedTokens = info.summary
-    ? outputTokens
-    : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
-  const cost = info.cost ?? 0
-  const hasContextUsage = inputTokens + cacheReadTokens + cacheWriteTokens > 0
-  return {
-    info,
-    inputTokens,
-    outputTokens,
-    reasoningTokens,
-    combinedTokens,
-    cost,
-    hasContextUsage,
-    timestamp: info.time?.created ?? 0,
-  }
-}
-
-function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) {
-  if (!messageId) return
-  const existing = state.entries.get(messageId)
-  if (!existing) return
-  state.entries.delete(messageId)
-  state.totalInputTokens -= existing.inputTokens
-  state.totalOutputTokens -= existing.outputTokens
-  state.totalReasoningTokens -= existing.reasoningTokens
-  state.totalCost -= existing.cost
-  if (state.latestEntry?.info.id === messageId) {
-    recomputeLatestEntry(state)
-  }
-}
-
-function addUsageEntry(state: SessionUsageState, entry: AssistantUsageEntry) {
-  state.entries.set(entry.info.id, entry)
-  state.totalInputTokens += entry.inputTokens
-  state.totalOutputTokens += entry.outputTokens
-  state.totalReasoningTokens += entry.reasoningTokens
-  state.totalCost += entry.cost
-  if (!state.latestEntry || entry.timestamp >= state.latestEntry.timestamp) {
-    state.latestEntry = entry
-  }
-}
-
-function updateUsageFromMessageInfo(instanceId: string, sessionId: string, info: MessageInfo) {
-  const messageId = typeof info.id === "string" ? info.id : undefined
-  if (!messageId) return
-  const state = getSessionUsageState(instanceId, sessionId)
-  removeUsageEntry(state, messageId)
-  const entry = extractAssistantUsage(info)
-  if (entry) {
-    addUsageEntry(state, entry)
-  }
-}
-
-function rebuildSessionUsage(instanceId: string, sessionId: string, messagesInfo: Map<string, MessageInfo>) {
-  const usageMap = getUsageInstance(instanceId)
-  const nextState = createEmptyUsageState()
-  for (const info of messagesInfo.values()) {
-    const entry = extractAssistantUsage(info)
-    if (entry) {
-      addUsageEntry(nextState, entry)
-    }
-  }
-  usageMap.set(sessionId, nextState)
-}
-
-function clearSessionUsage(instanceId: string, sessionId: string) {
-  const usageMap = sessionUsageStates.get(instanceId)
-  if (!usageMap) return
-  usageMap.delete(sessionId)
-  if (usageMap.size === 0) {
-    sessionUsageStates.delete(instanceId)
-  }
-}
-
-function decodeTextSegment(segment: any): any {
-  if (typeof segment === "string") {
-    return decodeHtmlEntities(segment)
-  }
-
-  if (segment && typeof segment === "object") {
-    const updated: Record<string, any> = { ...segment }
-
-    if (typeof updated.text === "string") {
-      updated.text = decodeHtmlEntities(updated.text)
-    }
-
-    if (typeof updated.value === "string") {
-      updated.value = decodeHtmlEntities(updated.value)
-    }
-
-    if (Array.isArray(updated.content)) {
-      updated.content = updated.content.map((item: any) => decodeTextSegment(item))
-    }
-
-    return updated
-  }
-
-  return segment
-}
-
-function normalizeMessagePart(part: any): any {
-  if (!part || typeof part !== "object") {
-    return part
-  }
-
-  if (part.type !== "text") {
-    return part
-  }
-
-  const normalized: Record<string, any> = { ...part, renderCache: undefined }
-
-  if (typeof normalized.text === "string") {
-    normalized.text = decodeHtmlEntities(normalized.text)
-  } else if (normalized.text && typeof normalized.text === "object") {
-    const textObject: Record<string, any> = { ...normalized.text }
-
-    if (typeof textObject.value === "string") {
-      textObject.value = decodeHtmlEntities(textObject.value)
-    }
-
-    if (Array.isArray(textObject.content)) {
-      textObject.content = textObject.content.map((item: any) => decodeTextSegment(item))
-    }
-
-    if (typeof textObject.text === "string") {
-      textObject.text = decodeHtmlEntities(textObject.text)
-    }
-
-    normalized.text = textObject
-  }
-
-  if (Array.isArray(normalized.content)) {
-    normalized.content = normalized.content.map((item: any) => decodeTextSegment(item))
-  }
-
-  if (normalized.thinking && typeof normalized.thinking === "object") {
-    const thinking: Record<string, any> = { ...normalized.thinking }
-    if (Array.isArray(thinking.content)) {
-      thinking.content = thinking.content.map((item: any) => decodeTextSegment(item))
-    }
-    normalized.thinking = thinking
-  }
-
-  return normalized
-}
-
-function computeDisplayParts(message: Message, showThinking: boolean): MessageDisplayParts {
-  const text: any[] = []
-  const tool: any[] = []
-  const reasoning: any[] = []
-
-  for (const part of message.parts) {
-    if (part.type === "text" && !part.synthetic && partHasRenderableText(part)) {
-      text.push(part)
-    } else if (part.type === "tool") {
-      tool.push(part)
-    } else if (part.type === "reasoning" && showThinking && partHasRenderableText(part)) {
-      reasoning.push(part)
-    }
-  }
-
-  const combined = reasoning.length > 0 ? [...text, ...reasoning] : [...text]
-  const version = typeof message.version === "number" ? message.version : 0
-
-  return { text, tool, reasoning, combined, showThinking, version }
-}
-
-function initializePartVersion(part: any, version = 0) {
-  if (!part || typeof part !== "object") return
-  const partAny = part as any
-  if (typeof partAny.version !== "number") {
-    partAny.version = version
-  }
-}
-
-function bumpPartVersion(previousPart: any, nextPart: any): number {
-  const prevVersion = typeof previousPart?.version === "number" ? previousPart.version : -1
-  const nextVersion = prevVersion + 1
-  nextPart.version = nextVersion
-  return nextVersion
-}
-
-function getSessionIndex(instanceId: string, sessionId: string) {
-  let instanceMap = sessionIndexes.get(instanceId)
-  if (!instanceMap) {
-    instanceMap = new Map()
-    sessionIndexes.set(instanceId, instanceMap)
-  }
-
-  let sessionMap = instanceMap.get(sessionId)
-  if (!sessionMap) {
-    sessionMap = { messageIndex: new Map(), partIndex: new Map() }
-    instanceMap.set(sessionId, sessionMap)
-  }
-
-  return sessionMap
-}
-
-function rebuildSessionIndex(instanceId: string, sessionId: string, messages: Message[]) {
-  const index = getSessionIndex(instanceId, sessionId)
-  index.messageIndex.clear()
-  index.partIndex.clear()
-
-  messages.forEach((message, messageIdx) => {
-    index.messageIndex.set(message.id, messageIdx)
-
-    const partMap = new Map<string, number>()
-    message.parts.forEach((part, partIdx) => {
-      if (part.id && typeof part.id === "string") {
-        partMap.set(part.id, partIdx)
-      }
-    })
-    index.partIndex.set(message.id, partMap)
-  })
-}
-
-function clearSessionIndex(instanceId: string, sessionId: string) {
-  const instanceMap = sessionIndexes.get(instanceId)
-  if (instanceMap) {
-    instanceMap.delete(sessionId)
-    if (instanceMap.size === 0) {
-      sessionIndexes.delete(instanceId)
-    }
-  }
-  clearSessionUsage(instanceId, sessionId)
-}
-
-function removeSessionIndexes(instanceId: string) {
-  sessionIndexes.delete(instanceId)
-  sessionUsageStates.delete(instanceId)
-}
-
-function updateSessionInfo(instanceId: string, sessionId: string) {
-  const instanceSessions = sessions().get(instanceId)
-  if (!instanceSessions) return
-
-  const session = instanceSessions.get(sessionId)
-  if (!session) return
-
-  let contextWindow = 0
-  let isSubscriptionModel = false
-  let modelID = ""
-  let providerID = ""
-  let actualUsageTokens = 0
-
-  const usageState = getSessionUsageState(instanceId, sessionId)
-  const hasUsageEntries = usageState.entries.size > 0
-
-  let totalInputTokens = hasUsageEntries ? usageState.totalInputTokens : 0
-  let totalOutputTokens = hasUsageEntries ? usageState.totalOutputTokens : 0
-  let totalReasoningTokens = hasUsageEntries ? usageState.totalReasoningTokens : 0
-  let totalCost = hasUsageEntries ? usageState.totalCost : 0
-
-  let latestAssistantInfo: MessageInfo | null = usageState.latestEntry?.info ?? null
-  let latestHasContextUsage = usageState.latestEntry?.hasContextUsage ?? false
-  const previousInfo = sessionInfoByInstance().get(instanceId)?.get(sessionId)
-  let contextAvailableTokens: number | null = null
-  let contextAvailableFromPrevious = false
-
-  if (latestAssistantInfo) {
-    const infoAny = latestAssistantInfo as any
-    actualUsageTokens = usageState.latestEntry?.combinedTokens ?? 0
-    modelID = infoAny.modelID || ""
-    providerID = infoAny.providerID || ""
-  } else if (previousInfo) {
-    totalInputTokens = previousInfo.inputTokens
-    totalOutputTokens = previousInfo.outputTokens
-    totalReasoningTokens = previousInfo.reasoningTokens
-    totalCost = previousInfo.cost
-    actualUsageTokens = previousInfo.actualUsageTokens
-
-    const previousContextWindow = previousInfo.contextWindow
-    const previousContextAvailable = previousInfo.contextAvailableTokens ?? null
-    const previousHasContextUsage =
-      previousContextAvailable !== null && previousContextWindow > 0
-        ? previousContextAvailable < previousContextWindow
-        : false
-
-    if (contextWindow === 0) {
-      contextWindow = previousContextWindow
-    }
-
-    if (contextWindow !== previousContextWindow) {
-      contextAvailableTokens = null
-      contextAvailableFromPrevious = false
-      latestHasContextUsage = previousHasContextUsage
-    } else {
-      contextAvailableTokens = previousContextAvailable
-      contextAvailableFromPrevious = true
-      latestHasContextUsage = previousHasContextUsage
-    }
-
-    isSubscriptionModel = previousInfo.isSubscriptionModel
-  }
-
-  const instanceProviders = providers().get(instanceId) || []
-
-
-
-
-  const sessionModel = session.model
-  let selectedModel: Provider["models"][number] | undefined
-
-  if (sessionModel?.providerId && sessionModel?.modelId) {
-    const provider = instanceProviders.find((p) => p.id === sessionModel.providerId)
-    selectedModel = provider?.models.find((m) => m.id === sessionModel.modelId)
-  }
-
-  if (!selectedModel && modelID && providerID) {
-    const provider = instanceProviders.find((p) => p.id === providerID)
-    selectedModel = provider?.models.find((m) => m.id === modelID)
-  }
-
-  let modelOutputLimit = DEFAULT_MODEL_OUTPUT_LIMIT
-
-  if (selectedModel) {
-    if (selectedModel.limit?.context) {
-      contextWindow = selectedModel.limit.context
-    }
-
-    if (selectedModel.limit?.output && selectedModel.limit.output > 0) {
-      modelOutputLimit = selectedModel.limit.output
-    }
-
-    if (selectedModel.cost?.input === 0 && selectedModel.cost?.output === 0) {
-      isSubscriptionModel = true
-    }
-  }
-
-  const outputBudget = Math.min(modelOutputLimit, DEFAULT_MODEL_OUTPUT_LIMIT)
-
-  if (!contextAvailableFromPrevious) {
-    if (contextWindow > 0) {
-      if (latestHasContextUsage && actualUsageTokens > 0) {
-        contextAvailableTokens = Math.max(contextWindow - (actualUsageTokens + outputBudget), 0)
-      } else {
-        contextAvailableTokens = contextWindow
-      }
-    } else {
-      contextAvailableTokens = null
-    }
-  }
-
-  setSessionInfoByInstance((prev) => {
-    const next = new Map(prev)
-    const instanceInfo = new Map(prev.get(instanceId))
-    instanceInfo.set(sessionId, {
-      cost: totalCost,
-      contextWindow,
-      isSubscriptionModel,
-      inputTokens: totalInputTokens,
-      outputTokens: totalOutputTokens,
-      reasoningTokens: totalReasoningTokens,
-      actualUsageTokens,
-      modelOutputLimit,
-      contextAvailableTokens,
-    })
-    next.set(instanceId, instanceInfo)
-    return next
-  })
-}
-
-export {
-  bumpPartVersion,
-  clearSessionIndex,
-  computeDisplayParts,
-  getSessionIndex,
-  initializePartVersion,
-  normalizeMessagePart,
-  rebuildSessionIndex,
-  rebuildSessionUsage,
-  removeSessionIndexes,
-  updateSessionInfo,
-  updateUsageFromMessageInfo,
-}

packages/ui/src/stores/sessions.ts

@@ -28,7 +28,6 @@ import {
   setSessionDraftPrompt,
 } from "./session-state"
 import { getDefaultModel } from "./session-models"
-import { computeDisplayParts, removeSessionIndexes } from "./session-messages"
 import {
   createSession,
   deleteSession,
@@ -79,7 +78,6 @@ export {
   clearActiveParentSession,
   clearInstanceDraftPrompts,
   clearSessionDraftPrompt,
-  computeDisplayParts,
   createSession,
   deleteSession,
   executeCustomCommand,
@@ -102,7 +100,6 @@ export {
   loadMessages,
   loading,
   providers,
-  removeSessionIndexes,
   sendMessage,
   sessionInfoByInstance,
   sessions,