Batch hydrate normalized messages for session load

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"