chore: add message store v2 baseline

2025-11-26 09:42:10 +00:00
parent 9313b2bd6c
commit 16b76385e2
13 changed files with 1665 additions and 57 deletions
--- a/packages/ui/src/stores/message-v2/bridge.ts
+++ b/packages/ui/src/stores/message-v2/bridge.ts
@@ -0,0 +1,175 @@
+import type { Permission } from "@opencode-ai/sdk"
+import type { Message, MessageInfo, ClientPart } from "../../types/message"
+import type { Session } from "../../types/session"
+import { messageStoreBus } from "./bus"
+import type { MessageStatus, SessionRevertState } from "./types"
+
+interface SessionMetadata {
+  id: string
+  title?: string
+  parentId?: string | null
+}
+
+function resolveSessionMetadata(session?: Session | null): SessionMetadata | undefined {
+  if (!session) return undefined
+  return {
+    id: session.id,
+    title: session.title,
+    parentId: session.parentId ?? null,
+  }
+}
+
+function normalizeStatus(status: Message["status"]): MessageStatus {
+  switch (status) {
+    case "sending":
+    case "sent":
+    case "streaming":
+    case "complete":
+    case "error":
+      return status
+    default:
+      return "complete"
+  }
+}
+
+export function seedSessionMessagesV2(
+  instanceId: string,
+  session: Session | SessionMetadata,
+  messages: Message[],
+  messageInfos?: Map<string, MessageInfo>,
+): void {
+  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)
+    }
+  })
+
+  if (messageInfos) {
+    store.rebuildUsage(metadata.id, messageInfos.values())
+  }
+}
+
+interface MessageInfoOptions {
+  status?: MessageStatus
+  bumpRevision?: boolean
+}
+
+export function upsertMessageInfoV2(instanceId: string, info: MessageInfo | null | undefined, options?: MessageInfoOptions): void {
+  if (!info || typeof info.id !== "string" || typeof info.sessionID !== "string") {
+    return
+  }
+  const store = messageStoreBus.getOrCreate(instanceId)
+  const timeInfo = (info.time ?? {}) as { created?: number; completed?: number }
+  const createdAt = typeof timeInfo.created === "number" ? timeInfo.created : Date.now()
+  const completedAt = typeof timeInfo.completed === "number" ? timeInfo.completed : undefined
+
+  store.upsertMessage({
+    id: info.id,
+    sessionId: info.sessionID,
+    role: info.role === "user" ? "user" : "assistant",
+    status: options?.status ?? "complete",
+    createdAt,
+    updatedAt: completedAt ?? createdAt,
+    bumpRevision: Boolean(options?.bumpRevision),
+  })
+  store.setMessageInfo(info.id, info)
+}
+
+export function applyPartUpdateV2(instanceId: string, part: ClientPart | null | undefined): void {
+  if (!part || typeof part.messageID !== "string") {
+    return
+  }
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.applyPartUpdate({
+    messageId: part.messageID,
+    part,
+  })
+}
+
+export function replaceMessageIdV2(instanceId: string, oldId: string, newId: string): void {
+  if (!oldId || !newId || oldId === newId) return
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.replaceMessageId({ oldId, newId })
+}
+
+function extractPermissionMessageId(permission: Permission): string | undefined {
+  return (permission as any).messageID || (permission as any).messageId
+}
+
+function extractPermissionPartId(permission: Permission): string | undefined {
+  const metadata = (permission as any).metadata || {}
+  return (
+    (permission as any).callID ||
+    (permission as any).callId ||
+    (permission as any).toolCallID ||
+    (permission as any).toolCallId ||
+    metadata.partId ||
+    metadata.partID ||
+    metadata.callID ||
+    metadata.callId ||
+    undefined
+  )
+}
+
+export function upsertPermissionV2(instanceId: string, permission: Permission): void {
+  if (!permission) return
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.upsertPermission({
+    permission,
+    messageId: extractPermissionMessageId(permission),
+    partId: extractPermissionPartId(permission),
+    enqueuedAt: (permission as any).time?.created ?? Date.now(),
+  })
+}
+
+export function removePermissionV2(instanceId: string, permissionId: string): void {
+  if (!permissionId) return
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.removePermission(permissionId)
+}
+
+export function ensureSessionMetadataV2(instanceId: string, session: Session | null | undefined): void {
+  if (!session) return
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.addOrUpdateSession({
+    id: session.id,
+    title: session.title,
+    parentId: session.parentId ?? null,
+    messageIds: session.messages.map((message: Message) => message.id),
+  })
+}
+
+export function getSessionMetadataFromStore(session?: Session | null): SessionMetadata | undefined {
+  return resolveSessionMetadata(session ?? undefined)
+}
+
+export function setSessionRevertV2(instanceId: string, sessionId: string, revert?: SessionRevertState | null): void {
+  if (!sessionId) return
+  const store = messageStoreBus.getOrCreate(instanceId)
+  store.setSessionRevert(sessionId, revert ?? null)
+}
--- a/packages/ui/src/stores/message-v2/bus.ts
+++ b/packages/ui/src/stores/message-v2/bus.ts
@@ -0,0 +1,41 @@
+import { createInstanceMessageStore } from "./instance-store"
+import type { InstanceMessageStore } from "./instance-store"
+
+class MessageStoreBus {
+  private stores = new Map<string, InstanceMessageStore>()
+
+  registerInstance(instanceId: string, store?: InstanceMessageStore): InstanceMessageStore {
+    if (this.stores.has(instanceId)) {
+      return this.stores.get(instanceId) as InstanceMessageStore
+    }
+
+    const resolved = store ?? createInstanceMessageStore(instanceId)
+    this.stores.set(instanceId, resolved)
+    return resolved
+  }
+
+  getInstance(instanceId: string): InstanceMessageStore | undefined {
+    return this.stores.get(instanceId)
+  }
+
+  getOrCreate(instanceId: string): InstanceMessageStore {
+    return this.registerInstance(instanceId)
+  }
+
+  unregisterInstance(instanceId: string) {
+    const store = this.stores.get(instanceId)
+    if (store) {
+      store.clearInstance()
+    }
+    this.stores.delete(instanceId)
+  }
+
+  clearAll() {
+    for (const [instanceId, store] of this.stores.entries()) {
+      store.clearInstance()
+      this.stores.delete(instanceId)
+    }
+  }
+}
+
+export const messageStoreBus = new MessageStoreBus()
--- a/packages/ui/src/stores/message-v2/instance-store.ts
+++ b/packages/ui/src/stores/message-v2/instance-store.ts
@@ -0,0 +1,523 @@
+import { createStore, produce, reconcile } from "solid-js/store"
+import type { SetStoreFunction } from "solid-js/store"
+import type { ClientPart, MessageInfo } from "../../types/message"
+import type {
+  InstanceMessageState,
+  MessageRecord,
+  MessageUpsertInput,
+  PartUpdateInput,
+  PendingPartEntry,
+  PermissionEntry,
+  ReplaceMessageIdOptions,
+  ScrollSnapshot,
+  SessionRecord,
+  SessionUpsertInput,
+  SessionUsageState,
+  UsageEntry,
+} from "./types"
+
+function createInitialState(instanceId: string): InstanceMessageState {
+  return {
+    instanceId,
+    sessions: {},
+    sessionOrder: [],
+    messages: {},
+    messageInfo: {},
+    pendingParts: {},
+    permissions: {
+      queue: [],
+      active: null,
+      byMessage: {},
+    },
+    usage: {},
+    scrollState: {},
+  }
+}
+
+function ensurePartId(messageId: string, part: ClientPart, index: number): string {
+  if (typeof part.id === "string" && part.id.length > 0) {
+    return part.id
+  }
+  return `${messageId}-part-${index}`
+}
+
+function clonePart(part: ClientPart): ClientPart {
+  return JSON.parse(JSON.stringify(part)) as ClientPart
+}
+
+function createEmptyUsageState(): SessionUsageState {
+  return {
+    entries: {},
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    totalReasoningTokens: 0,
+    totalCost: 0,
+    actualUsageTokens: 0,
+    latestMessageId: undefined,
+  }
+}
+
+function extractUsageEntry(info: MessageInfo | undefined): UsageEntry | null {
+  if (!info || info.role !== "assistant") return null
+  const messageId = typeof info.id === "string" ? info.id : undefined
+  if (!messageId) return null
+  const tokens = info.tokens
+  if (!tokens) return null
+  const inputTokens = tokens.input ?? 0
+  const outputTokens = tokens.output ?? 0
+  const reasoningTokens = tokens.reasoning ?? 0
+  const cacheReadTokens = tokens.cache?.read ?? 0
+  const cacheWriteTokens = tokens.cache?.write ?? 0
+  if (inputTokens === 0 && outputTokens === 0 && reasoningTokens === 0 && cacheReadTokens === 0 && cacheWriteTokens === 0) {
+    return null
+  }
+  const combinedTokens = info.summary ? outputTokens : inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens + reasoningTokens
+  return {
+    messageId,
+    inputTokens,
+    outputTokens,
+    reasoningTokens,
+    cacheReadTokens,
+    cacheWriteTokens,
+    combinedTokens,
+    cost: info.cost ?? 0,
+    timestamp: info.time?.created ?? 0,
+    hasContextUsage: inputTokens + cacheReadTokens + cacheWriteTokens > 0,
+  }
+}
+
+function applyUsageState(state: SessionUsageState, entry: UsageEntry | null) {
+  if (!entry) return
+  state.entries[entry.messageId] = entry
+  state.totalInputTokens += entry.inputTokens
+  state.totalOutputTokens += entry.outputTokens
+  state.totalReasoningTokens += entry.reasoningTokens
+  state.totalCost += entry.cost
+  if (!state.latestMessageId || entry.timestamp >= (state.entries[state.latestMessageId]?.timestamp ?? 0)) {
+    state.latestMessageId = entry.messageId
+    state.actualUsageTokens = entry.combinedTokens
+  }
+}
+
+function removeUsageEntry(state: SessionUsageState, messageId: string | undefined) {
+  if (!messageId) return
+  const existing = state.entries[messageId]
+  if (!existing) return
+  state.totalInputTokens -= existing.inputTokens
+  state.totalOutputTokens -= existing.outputTokens
+  state.totalReasoningTokens -= existing.reasoningTokens
+  state.totalCost -= existing.cost
+  delete state.entries[messageId]
+  if (state.latestMessageId === messageId) {
+    state.latestMessageId = undefined
+    state.actualUsageTokens = 0
+    let latest: UsageEntry | null = null
+    for (const candidate of Object.values(state.entries) as UsageEntry[]) {
+      if (!latest || candidate.timestamp >= latest.timestamp) {
+        latest = candidate
+      }
+    }
+    if (latest) {
+      state.latestMessageId = latest.messageId
+      state.actualUsageTokens = latest.combinedTokens
+    }
+  }
+}
+
+function rebuildUsageStateFromInfos(infos: Iterable<MessageInfo>): SessionUsageState {
+  const usageState = createEmptyUsageState()
+  for (const info of infos) {
+    const entry = extractUsageEntry(info)
+    if (entry) {
+      applyUsageState(usageState, entry)
+    }
+  }
+  return usageState
+}
+
+export interface InstanceMessageStore {
+
+  instanceId: string
+  state: InstanceMessageState
+  setState: SetStoreFunction<InstanceMessageState>
+  addOrUpdateSession: (input: SessionUpsertInput) => void
+  upsertMessage: (input: MessageUpsertInput) => void
+  applyPartUpdate: (input: PartUpdateInput) => void
+  bufferPendingPart: (entry: PendingPartEntry) => void
+  flushPendingParts: (messageId: string) => void
+  replaceMessageId: (options: ReplaceMessageIdOptions) => void
+  setMessageInfo: (messageId: string, info: MessageInfo) => void
+  getMessageInfo: (messageId: string) => MessageInfo | undefined
+  upsertPermission: (entry: PermissionEntry) => void
+  removePermission: (permissionId: string) => void
+  getPermissionState: (messageId?: string, partId?: string) => { entry: PermissionEntry; active: boolean } | null
+  setSessionRevert: (sessionId: string, revert?: SessionRecord["revert"] | null) => void
+  getSessionRevert: (sessionId: string) => SessionRecord["revert"] | undefined | null
+  rebuildUsage: (sessionId: string, infos: Iterable<MessageInfo>) => void
+  getSessionUsage: (sessionId: string) => SessionUsageState | undefined
+  setScrollSnapshot: (sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) => void
+  getScrollSnapshot: (sessionId: string, scope: string) => ScrollSnapshot | undefined
+  getSessionMessageIds: (sessionId: string) => string[]
+  getMessage: (messageId: string) => MessageRecord | undefined
+  clearInstance: () => void
+}
+
+export function createInstanceMessageStore(instanceId: string): InstanceMessageStore {
+  const [state, setState] = createStore<InstanceMessageState>(createInitialState(instanceId))
+
+  function withUsageState(sessionId: string, updater: (draft: SessionUsageState) => void) {
+    setState("usage", sessionId, (current) => {
+      const draft = current
+        ? {
+            ...current,
+            entries: { ...current.entries },
+          }
+        : createEmptyUsageState()
+      updater(draft)
+      return draft
+    })
+  }
+
+  function updateUsageWithInfo(info: MessageInfo | undefined) {
+    if (!info || typeof info.sessionID !== "string") return
+    const messageId = typeof info.id === "string" ? info.id : undefined
+    if (!messageId) return
+    withUsageState(info.sessionID, (draft) => {
+      removeUsageEntry(draft, messageId)
+      const entry = extractUsageEntry(info)
+      if (entry) {
+        applyUsageState(draft, entry)
+      }
+    })
+  }
+
+  function rebuildUsage(sessionId: string, infos: Iterable<MessageInfo>) {
+    const usageState = rebuildUsageStateFromInfos(infos)
+    setState("usage", sessionId, usageState)
+  }
+
+  function getSessionUsage(sessionId: string) {
+    return state.usage[sessionId]
+  }
+
+  function ensureSessionEntry(sessionId: string): SessionRecord {
+    const existing = state.sessions[sessionId]
+    if (existing) {
+      return existing
+    }
+
+    const now = Date.now()
+    const session: SessionRecord = {
+      id: sessionId,
+      createdAt: now,
+      updatedAt: now,
+      messageIds: [],
+    }
+
+    setState("sessions", sessionId, session)
+    setState("sessionOrder", (order) => (order.includes(sessionId) ? order : [...order, sessionId]))
+    return session
+  }
+
+  function addOrUpdateSession(input: SessionUpsertInput) {
+    const session = ensureSessionEntry(input.id)
+    const nextMessageIds = Array.isArray(input.messageIds) ? input.messageIds : session.messageIds
+
+    setState("sessions", input.id, {
+      ...session,
+      title: input.title ?? session.title,
+      parentId: input.parentId ?? session.parentId ?? null,
+      updatedAt: Date.now(),
+      messageIds: nextMessageIds,
+      revert: input.revert ?? session.revert ?? null,
+    })
+  }
+
+  function insertMessageIntoSession(sessionId: string, messageId: string) {
+    ensureSessionEntry(sessionId)
+    setState("sessions", sessionId, "messageIds", (ids = []) => {
+      if (ids.includes(messageId)) {
+        return ids
+      }
+      return [...ids, messageId]
+    })
+  }
+
+  function normalizeParts(messageId: string, parts: ClientPart[] | undefined) {
+    if (!parts || parts.length === 0) {
+      return null
+    }
+    const map: MessageRecord["parts"] = {}
+    const ids: string[] = []
+
+    parts.forEach((part, index) => {
+      const id = ensurePartId(messageId, part, index)
+      const cloned = clonePart(part)
+      if (typeof cloned.version !== "number") {
+        cloned.version = 0
+      }
+      map[id] = {
+        id,
+        data: cloned,
+        revision: cloned.version,
+      }
+      ids.push(id)
+    })
+
+    return { map, ids }
+  }
+
+  function upsertMessage(input: MessageUpsertInput) {
+    const normalizedParts = normalizeParts(input.id, input.parts)
+    const shouldBump = Boolean(input.bumpRevision || normalizedParts)
+    const now = Date.now()
+
+    setState("messages", input.id, (previous) => {
+      const revision = previous ? previous.revision + (shouldBump ? 1 : 0) : 0
+      return {
+        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,
+        partIds: normalizedParts ? normalizedParts.ids : previous?.partIds ?? [],
+        parts: normalizedParts ? normalizedParts.map : previous?.parts ?? {},
+      }
+    })
+
+    insertMessageIntoSession(input.sessionId, input.id)
+    flushPendingParts(input.id)
+  }
+
+  function bufferPendingPart(entry: PendingPartEntry) {
+    setState("pendingParts", entry.messageId, (list = []) => [...list, entry])
+  }
+
+  function applyPartUpdate(input: PartUpdateInput) {
+    const message = state.messages[input.messageId]
+    if (!message) {
+      bufferPendingPart({ messageId: input.messageId, part: input.part, receivedAt: Date.now() })
+      return
+    }
+
+    const partId = ensurePartId(input.messageId, input.part, message.partIds.length)
+    const cloned = clonePart(input.part)
+
+    setState(
+      "messages",
+      input.messageId,
+      produce((draft: MessageRecord) => {
+        if (!draft.partIds.includes(partId)) {
+          draft.partIds = [...draft.partIds, partId]
+        }
+        const existing = draft.parts[partId]
+        const nextRevision = existing ? existing.revision + 1 : cloned.version ?? 0
+        draft.parts[partId] = {
+          id: partId,
+          data: cloned,
+          revision: nextRevision,
+        }
+        draft.updatedAt = Date.now()
+        if (input.bumpRevision ?? true) {
+          draft.revision += 1
+        }
+      }),
+    )
+  }
+
+  function flushPendingParts(messageId: string) {
+    const pending = state.pendingParts[messageId]
+    if (!pending || pending.length === 0) {
+      return
+    }
+    pending.forEach((entry) => applyPartUpdate({ messageId, part: entry.part }))
+    setState("pendingParts", (prev) => {
+      const next = { ...prev }
+      delete next[messageId]
+      return next
+    })
+  }
+
+  function replaceMessageId(options: ReplaceMessageIdOptions) {
+    if (options.oldId === options.newId) return
+    const existing = state.messages[options.oldId]
+    if (!existing) return
+
+    const cloned: MessageRecord = {
+      ...existing,
+      id: options.newId,
+      isEphemeral: false,
+      updatedAt: Date.now(),
+    }
+
+    setState("messages", options.newId, cloned)
+    setState("messages", (prev) => {
+      const next = { ...prev }
+      delete next[options.oldId]
+      return next
+    })
+
+    Object.values(state.sessions).forEach((session) => {
+      const index = session.messageIds.indexOf(options.oldId)
+      if (index === -1) return
+      setState("sessions", session.id, "messageIds", (ids) => {
+        const next = [...ids]
+        next[index] = options.newId
+        return next
+      })
+    })
+
+    const infoEntry = state.messageInfo[options.oldId]
+    if (infoEntry) {
+      setState("messageInfo", options.newId, infoEntry)
+      setState("messageInfo", (prev) => {
+        const next = { ...prev }
+        delete next[options.oldId]
+        return next
+      })
+    }
+
+    const permissionMap = state.permissions.byMessage[options.oldId]
+    if (permissionMap) {
+      setState("permissions", "byMessage", options.newId, permissionMap)
+      setState("permissions", (prev) => {
+        const next = { ...prev }
+        const nextByMessage = { ...next.byMessage }
+        delete nextByMessage[options.oldId]
+        next.byMessage = nextByMessage
+        return next
+      })
+    }
+
+    const pending = state.pendingParts[options.oldId]
+    if (pending) {
+      setState("pendingParts", options.newId, pending)
+      setState("pendingParts", (prev) => {
+        const next = { ...prev }
+        delete next[options.oldId]
+        return next
+      })
+    }
+  }
+
+  function setMessageInfo(messageId: string, info: MessageInfo) {
+    if (!messageId) return
+    setState("messageInfo", messageId, info)
+    updateUsageWithInfo(info)
+  }
+
+  function getMessageInfo(messageId: string) {
+    return state.messageInfo[messageId]
+  }
+
+  function upsertPermission(entry: PermissionEntry) {
+    const messageKey = entry.messageId ?? "__global__"
+    const partKey = entry.partId ?? "__global__"
+
+    setState(
+      "permissions",
+      produce((draft) => {
+        draft.byMessage[messageKey] = draft.byMessage[messageKey] ?? {}
+        draft.byMessage[messageKey][partKey] = entry
+        const existingIndex = draft.queue.findIndex((item) => item.permission.id === entry.permission.id)
+        if (existingIndex === -1) {
+          draft.queue.push(entry)
+        } else {
+          draft.queue[existingIndex] = entry
+        }
+        if (!draft.active || draft.active.permission.id === entry.permission.id) {
+          draft.active = entry
+        }
+      }),
+    )
+  }
+
+  function removePermission(permissionId: string) {
+    setState(
+      "permissions",
+      produce((draft) => {
+        draft.queue = draft.queue.filter((item) => item.permission.id !== permissionId)
+        if (draft.active?.permission.id === permissionId) {
+          draft.active = draft.queue[0] ?? null
+        }
+        Object.keys(draft.byMessage).forEach((messageKey) => {
+          const partEntries = draft.byMessage[messageKey]
+          Object.keys(partEntries).forEach((partKey) => {
+            if (partEntries[partKey].permission.id === permissionId) {
+              delete partEntries[partKey]
+            }
+          })
+          if (Object.keys(partEntries).length === 0) {
+            delete draft.byMessage[messageKey]
+          }
+        })
+      }),
+    )
+  }
+
+  function getPermissionState(messageId?: string, partId?: string) {
+    const messageKey = messageId ?? "__global__"
+    const partKey = partId ?? "__global__"
+    const entry = state.permissions.byMessage[messageKey]?.[partKey]
+    if (!entry) return null
+    const active = state.permissions.active?.permission.id === entry.permission.id
+    return { entry, active }
+  }
+
+  function setSessionRevert(sessionId: string, revert?: SessionRecord["revert"] | null) {
+    if (!sessionId) return
+    ensureSessionEntry(sessionId)
+    setState("sessions", sessionId, "revert", revert ?? null)
+  }
+
+  function getSessionRevert(sessionId: string) {
+    return state.sessions[sessionId]?.revert ?? null
+  }
+
+  function makeScrollKey(sessionId: string, scope: string) {
+    return `${sessionId}:${scope}`
+  }
+
+  function setScrollSnapshot(sessionId: string, scope: string, snapshot: Omit<ScrollSnapshot, "updatedAt">) {
+    const key = makeScrollKey(sessionId, scope)
+    setState("scrollState", key, { ...snapshot, updatedAt: Date.now() })
+  }
+
+  function getScrollSnapshot(sessionId: string, scope: string) {
+    const key = makeScrollKey(sessionId, scope)
+    return state.scrollState[key]
+  }
+
+  function clearInstance() {
+    setState(reconcile(createInitialState(instanceId)))
+  }
+
+  return {
+    instanceId,
+    state,
+    setState,
+    addOrUpdateSession,
+    upsertMessage,
+    applyPartUpdate,
+    bufferPendingPart,
+    flushPendingParts,
+    replaceMessageId,
+    setMessageInfo,
+    getMessageInfo,
+    upsertPermission,
+    removePermission,
+    getPermissionState,
+    setSessionRevert,
+    getSessionRevert,
+    rebuildUsage,
+    getSessionUsage,
+    setScrollSnapshot,
+    getScrollSnapshot,
+    getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
+    getMessage: (messageId: string) => state.messages[messageId],
+    clearInstance,
+  }
+}
--- a/packages/ui/src/stores/message-v2/types.ts
+++ b/packages/ui/src/stores/message-v2/types.ts
@@ -0,0 +1,138 @@
+import type { ClientPart, MessageInfo } from "../../types/message"
+import type { Permission } from "@opencode-ai/sdk"
+
+export type MessageStatus = "sending" | "sent" | "streaming" | "complete" | "error"
+export type MessageRole = "user" | "assistant"
+
+export interface NormalizedPartRecord {
+  id: string
+  data: ClientPart
+  revision: number
+}
+
+export interface MessageRecord {
+  id: string
+  sessionId: string
+  role: MessageRole
+  status: MessageStatus
+  createdAt: number
+  updatedAt: number
+  revision: number
+  isEphemeral?: boolean
+  partIds: string[]
+  parts: Record<string, NormalizedPartRecord>
+}
+
+export interface SessionRevertState {
+  messageID?: string
+  partID?: string
+  snapshot?: string
+  diff?: string
+}
+
+export interface SessionRecord {
+  id: string
+  title?: string
+  parentId?: string | null
+  createdAt: number
+  updatedAt: number
+  messageIds: string[]
+  revert?: SessionRevertState | null
+}
+
+export interface PendingPartEntry {
+  messageId: string
+  part: ClientPart
+  receivedAt: number
+}
+
+export interface PermissionEntry {
+  permission: Permission
+  messageId?: string
+  partId?: string
+  enqueuedAt: number
+}
+
+export interface InstancePermissionState {
+  queue: PermissionEntry[]
+  active: PermissionEntry | null
+  byMessage: Record<string, Record<string, PermissionEntry>>
+}
+
+export interface ScrollSnapshot {
+  scrollTop: number
+  atBottom: boolean
+  updatedAt: number
+}
+
+export interface UsageEntry {
+  messageId: string
+  inputTokens: number
+  outputTokens: number
+  reasoningTokens: number
+  cacheReadTokens: number
+  cacheWriteTokens: number
+  combinedTokens: number
+  cost: number
+  timestamp: number
+  hasContextUsage: boolean
+}
+
+export interface SessionUsageState {
+  entries: Record<string, UsageEntry>
+  totalInputTokens: number
+  totalOutputTokens: number
+  totalReasoningTokens: number
+  totalCost: number
+  actualUsageTokens: number
+  latestMessageId?: string
+}
+
+export interface InstanceMessageState {
+  instanceId: string
+  sessions: Record<string, SessionRecord>
+  sessionOrder: string[]
+  messages: Record<string, MessageRecord>
+  messageInfo: Record<string, MessageInfo>
+  pendingParts: Record<string, PendingPartEntry[]>
+  permissions: InstancePermissionState
+  usage: Record<string, SessionUsageState>
+  scrollState: Record<string, ScrollSnapshot>
+}
+
+export interface SessionUpsertInput {
+  id: string
+  title?: string
+  parentId?: string | null
+  messageIds?: string[]
+  revert?: SessionRevertState | null
+}
+
+export interface MessageUpsertInput {
+  id: string
+  sessionId: string
+  role: MessageRole
+  status: MessageStatus
+  parts?: ClientPart[]
+  createdAt?: number
+  updatedAt?: number
+  isEphemeral?: boolean
+  bumpRevision?: boolean
+}
+
+export interface PartUpdateInput {
+  messageId: string
+  part: ClientPart
+  bumpRevision?: boolean
+}
+
+export interface ReplaceMessageIdOptions {
+  oldId: string
+  newId: string
+}
+
+export interface ScrollCacheKey {
+  instanceId: string
+  sessionId: string
+  scope: string
+}
--- a/packages/ui/src/stores/session-api.ts
+++ b/packages/ui/src/stores/session-api.ts
@@ -32,6 +32,7 @@ import {
  rebuildSessionUsage,
  updateSessionInfo,
 } from "./session-messages"
+import { seedSessionMessagesV2 } from "./message-v2/bridge"

 interface SessionForkResponse {
  id: string
@@ -610,6 +611,14 @@ async function loadMessages(instanceId: string, sessionId: string, force = false
      return next
    })

+    const sessionForV2 = sessions().get(instanceId)?.get(sessionId) ?? {
+      id: sessionId,
+      title: session?.title,
+      parentId: session?.parentId ?? null,
+      revert: session?.revert,
+    }
+    seedSessionMessagesV2(instanceId, sessionForV2, messages, messagesInfo)
+
  } catch (error) {
    console.error("Failed to load messages:", error)
    throw error
--- a/packages/ui/src/stores/session-events.ts
+++ b/packages/ui/src/stores/session-events.ts
@@ -34,6 +34,14 @@ import {
 } from "./session-messages"
 import { loadMessages } from "./session-api"
 import { setSessionCompactionState } from "./session-compaction"
+import {
+  applyPartUpdateV2,
+  replaceMessageIdV2,
+  upsertMessageInfoV2,
+  upsertPermissionV2,
+  removePermissionV2,
+  setSessionRevertV2,
+} from "./message-v2/bridge"

 interface TuiToastEvent {
  type: "tui.toast.show"
@@ -195,6 +203,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
          index.partIndex.delete(oldId)
          index.partIndex.set(message.id, existingPartMap)
        }
+        replaceMessageIdV2(instanceId, oldId, message.id)
      }

      if (filteredSynthetics || replacedTemp) {
@@ -212,6 +221,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
      /* mutations already applied above */
    })

+    applyPartUpdateV2(instanceId, part)
    updateSessionInfo(instanceId, part.sessionID)
    refreshPermissionsForSession(instanceId, part.sessionID)
  } else if (event.type === "message.updated") {
@@ -270,6 +280,7 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
            index.partIndex.delete(oldId)
            index.partIndex.set(message.id, existingPartMap)
          }
+          replaceMessageIdV2(instanceId, oldId, message.id)
        }
      } else {
        const newMessage: any = {
@@ -305,8 +316,11 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
      message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks)
    }

+    upsertMessageInfoV2(instanceId, info, { status: "complete" })
+
    session.messagesInfo.set(info.id, info)
    updateUsageFromMessageInfo(instanceId, info.sessionID, info)
+
    withSession(instanceId, info.sessionID, () => {
      /* ensure reactivity */
    })
@@ -359,6 +373,7 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo
      next.set(instanceId, updated)
      return next
    })
+    setSessionRevertV2(instanceId, info.id, info.revert ?? null)

    console.log(`[SSE] New session created: ${info.id}`, newSession)
  } else {
@@ -391,6 +406,7 @@ function handleSessionUpdate(instanceId: string, event: EventSessionUpdated): vo
      next.set(instanceId, updated)
      return next
    })
+    setSessionRevertV2(instanceId, info.id, info.revert ?? null)
  }
 }

@@ -490,6 +506,7 @@ function handlePermissionUpdated(instanceId: string, event: EventPermissionUpdat

  console.log(`[SSE] Permission updated: ${permission.id} (${permission.type})`)
  addPermissionToQueue(instanceId, permission)
+  upsertPermissionV2(instanceId, permission)
 }

 function handlePermissionReplied(instanceId: string, event: EventPermissionReplied): void {
@@ -498,6 +515,7 @@ function handlePermissionReplied(instanceId: string, event: EventPermissionRepli

  console.log(`[SSE] Permission replied: ${permissionID}`)
  removePermissionFromQueue(instanceId, permissionID)
+  removePermissionV2(instanceId, permissionID)
 }

 export {
--- a/packages/ui/src/stores/session-status.ts
+++ b/packages/ui/src/stores/session-status.ts
@@ -1,7 +1,9 @@
 import type { Session, SessionStatus } from "../types/session"
 import type { Message, MessageInfo } from "../types/message"
+import type { MessageRecord } from "./message-v2/types"
 import { sessions } from "./sessions"
 import { isSessionCompactionActive } from "./session-compaction"
+import { messageStoreBus } from "./message-v2/bus"

 function getSession(instanceId: string, sessionId: string): Session | null {
  const instanceSessions = sessions().get(instanceId)
@@ -17,21 +19,49 @@ function isSessionCompacting(session: Session): boolean {
  return Boolean(compactingFlag)
 }

-function getMessageTimestamp(session: Session, message?: Message): number {
-  if (!message) return Number.NEGATIVE_INFINITY
-  if (typeof message.timestamp === "number" && Number.isFinite(message.timestamp)) {
-    return message.timestamp
+function getLatestInfoFromStore(instanceId: string, sessionId: string, role?: MessageInfo["role"]): MessageInfo | undefined {
+  const store = messageStoreBus.getOrCreate(instanceId)
+  const messageIds = store.getSessionMessageIds(sessionId)
+  let latest: MessageInfo | undefined
+  let latestTimestamp = Number.NEGATIVE_INFINITY
+  for (const id of messageIds) {
+    const info = store.getMessageInfo(id)
+    if (!info) continue
+    if (role && info.role !== role) continue
+    const timestamp = info.time?.created ?? 0
+    if (timestamp >= latestTimestamp) {
+      latest = info
+      latestTimestamp = timestamp
+    }
  }
-  const info = session.messagesInfo.get(message.id)
-  return info?.time?.created ?? Number.NEGATIVE_INFINITY
+  return latest
 }

-function getLastMessage(session: Session): Message | undefined {
+function getLastMessageFromStore(instanceId: string, sessionId: string): MessageRecord | undefined {
+  const store = messageStoreBus.getOrCreate(instanceId)
+  const messageIds = store.getSessionMessageIds(sessionId)
+  let latest: MessageRecord | undefined
+  let latestTimestamp = Number.NEGATIVE_INFINITY
+  for (const id of messageIds) {
+    const record = store.getMessage(id)
+    if (!record) continue
+    const info = store.getMessageInfo(id)
+    const timestamp = info?.time?.created ?? record.createdAt ?? Number.NEGATIVE_INFINITY
+    if (timestamp >= latestTimestamp) {
+      latest = record
+      latestTimestamp = timestamp
+    }
+  }
+  return latest
+}
+
+function getLegacyLastMessage(session: Session): Message | undefined {
  let latest: Message | undefined
  let latestTimestamp = Number.NEGATIVE_INFINITY
  for (const message of session.messages) {
    if (!message) continue
-    const timestamp = getMessageTimestamp(session, message)
+    const info = session.messagesInfo.get(message.id)
+    const timestamp = info?.time?.created ?? message.timestamp ?? Number.NEGATIVE_INFINITY
    if (timestamp >= latestTimestamp) {
      latest = message
      latestTimestamp = timestamp
@@ -40,7 +70,7 @@ function getLastMessage(session: Session): Message | undefined {
  return latest
 }

-function getLastMessageInfo(session: Session, role?: MessageInfo["role"]): MessageInfo | undefined {
+function getLegacyLastMessageInfo(session: Session, role?: MessageInfo["role"]): MessageInfo | undefined {
  if (session.messagesInfo.size === 0) {
    return undefined
  }
@@ -92,7 +122,28 @@ function isAssistantInfoPending(info?: MessageInfo): boolean {
  return completed < created
 }

-function isAssistantStillGenerating(message: Message, info?: MessageInfo): boolean {
+function isAssistantStillGeneratingRecord(record: MessageRecord, info?: MessageInfo): boolean {
+  if (record.role !== "assistant") {
+    return false
+  }
+
+  if (record.status === "error") {
+    return false
+  }
+
+  if (record.status === "streaming" || record.status === "sending") {
+    return true
+  }
+
+  const completedAt = (info?.time as { completed?: number } | undefined)?.completed
+  if (completedAt !== undefined && completedAt !== null) {
+    return false
+  }
+
+  return !(record.status === "complete" || record.status === "sent")
+}
+
+function isAssistantStillGeneratingLegacy(message: Message, info?: MessageInfo): boolean {
  if (message.type !== "assistant") {
    return false
  }
@@ -119,15 +170,20 @@ export function getSessionStatus(instanceId: string, sessionId: string): Session
    return "idle"
  }

+  const store = messageStoreBus.getOrCreate(instanceId)
+
  if (isSessionCompactionActive(instanceId, sessionId) || isSessionCompacting(session)) {
    return "compacting"
  }

-  const latestUserInfo = getLastMessageInfo(session, "user")
-  const latestAssistantInfo = getLastMessageInfo(session, "assistant")
-  const lastMessage = getLastMessage(session)
-  if (!lastMessage) {
-    const latestInfo = getLastMessageInfo(session)
+  const latestUserInfo = getLatestInfoFromStore(instanceId, sessionId, "user") ?? getLegacyLastMessageInfo(session, "user")
+  const latestAssistantInfo = getLatestInfoFromStore(instanceId, sessionId, "assistant") ?? getLegacyLastMessageInfo(session, "assistant")
+
+  const lastRecord = getLastMessageFromStore(instanceId, sessionId)
+  const legacyFallbackMessage = lastRecord ? undefined : getLegacyLastMessage(session)
+
+  if (!lastRecord && !legacyFallbackMessage) {
+    const latestInfo = latestUserInfo ?? latestAssistantInfo ?? getLegacyLastMessageInfo(session)
    if (!latestInfo) {
      return "idle"
    }
@@ -138,13 +194,22 @@ export function getSessionStatus(instanceId: string, sessionId: string): Session
    return infoCompleted ? "idle" : "working"
  }

-  if (lastMessage.type === "user") {
-    return "working"
-  }
-
-  const infoForMessage = session.messagesInfo.get(lastMessage.id) ?? latestAssistantInfo
-  if (isAssistantStillGenerating(lastMessage, infoForMessage)) {
-    return "working"
+  if (lastRecord) {
+    if (lastRecord.role === "user") {
+      return "working"
+    }
+    const infoForRecord = store.getMessageInfo(lastRecord.id) ?? latestAssistantInfo
+    if (infoForRecord && isAssistantStillGeneratingRecord(lastRecord, infoForRecord)) {
+      return "working"
+    }
+  } else if (legacyFallbackMessage) {
+    if (legacyFallbackMessage.type === "user") {
+      return "working"
+    }
+    const infoForLegacy = session.messagesInfo.get(legacyFallbackMessage.id) ?? latestAssistantInfo
+    if (isAssistantStillGeneratingLegacy(legacyFallbackMessage, infoForLegacy)) {
+      return "working"
+    }
  }

  if (isAssistantInfoPending(latestAssistantInfo)) {