fix(ui): refresh session usage info on step and usage events

.github/workflows/pr-build.yml

@@ -313,6 +313,7 @@ jobs:
           packages/ui/src/stores/session-pruning-pagination.test.ts
           packages/ui/src/stores/session-native-events.test.ts
           packages/ui/src/stores/runtime-contract.test.ts
+          packages/ui/src/stores/session-usage-contract.test.ts
           packages/ui/src/stores/session-move-restore.test.ts
           packages/ui/src/stores/session-request-authority.test.ts
           packages/ui/src/stores/session-send-lifecycle.test.ts

packages/ui/src/stores/instances.ts

@@ -48,6 +48,7 @@ import {
 } from "./session-state"
 import { setHasInstances } from "./ui"
 import { messageStoreBus } from "./message-v2/bus"
+import { clearRevertedUsage, recordRevertedUsage, updateSessionInfo } from "./message-v2/session-info"
 import { applyOpenCodeDataEvent, destroyOpenCodeData, projectOpenCodeMessages, syncOpenCodeSessionInbox } from "./opencode-data"
 import { isLatestWindow } from "./message-v2/message-window"
 import { upsertPermissionV2, removePermissionV2, removeMessageV2 } from "./message-v2/bridge"
@@ -1358,6 +1359,7 @@ function removeInstance(id: string, options: { authoritative?: boolean } = {}) {
   // Clean up session indexes and drafts for removed instance
   clearCacheForInstance(id)
   messageStoreBus.unregisterInstance(id)
+  clearRevertedUsage(id)
   clearInstanceDraftPrompts(id)
   clearSessionListRequestState(id)
   clearSessionCatalogState(id)
@@ -1995,6 +1997,12 @@ async function sendFormCancel(instanceId: string, formId: string): Promise<void>
   }
 }
 
+// Events after which assistant message cost/token totals may have changed.
+const USAGE_EVENT_TYPES = new Set<string>([
+  "session.step.ended",
+  "session.step.failed",
+])
+
 function handleInstanceInvalidation(instanceId: string, event: Parameters<NonNullable<typeof sseManager.onInvalidation>>[1]): void {
   const instance = instances().get(instanceId)
   if (!instance?.client) return
@@ -2021,23 +2029,29 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters<NonNul
     if (sessionId && (force || event.type.startsWith("session.")) && (
       activeSessionId().get(instanceId) === sessionId
       && isLatestWindow(messageStoreBus.getOrCreate(instanceId).getMessageWindow(sessionId))
-    )) projectOpenCodeMessages(
-      instanceId,
-      sessionId,
-      data,
-      preserveOmitted,
-      force || event.type !== "session.inbox.enqueued",
-    )
+    )) {
+      projectOpenCodeMessages(
+        instanceId,
+        sessionId,
+        data,
+        preserveOmitted,
+        force || event.type !== "session.inbox.enqueued",
+      )
+      if (force || USAGE_EVENT_TYPES.has(event.type)) updateSessionInfo(instanceId, sessionId)
+    }
   }
   const project = (data: ReturnType<typeof applyOpenCodeDataEvent>, preserveOmitted = true) => {
     projectMessages(data, preserveOmitted)
     if (sessionId && event.type === "session.inbox.cancelled") {
       removeMessageV2(instanceId, event.data.inboxID, sessionId)
     }
     if (sessionId && event.type === "session.revert.committed") {
-      for (const messageId of messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId)) {
-        if (messageId >= event.data.to) removeMessageV2(instanceId, messageId, sessionId)
-      }
+      const store = messageStoreBus.getOrCreate(instanceId)
+      const removed = store.getSessionMessageIds(sessionId).filter((messageId) => messageId >= event.data.to)
+      const entries = store.getSessionUsage(sessionId)?.entries ?? {}
+      recordRevertedUsage(instanceId, sessionId, removed.flatMap((messageId) => entries[messageId] ? [entries[messageId]] : []))
+      for (const messageId of removed) removeMessageV2(instanceId, messageId, sessionId)
+      updateSessionInfo(instanceId, sessionId)
     }
   }
   const data = applyOpenCodeDataEvent(instanceId, instance.folder, event, project, (next) => {

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

@@ -2,7 +2,46 @@ import type { Provider } from "../../types/session"
 import { DEFAULT_MODEL_OUTPUT_LIMIT } from "../session-models"
 import { providers, sessions, sessionInfoByInstance, setSessionInfoByInstance, updateThreadTotalsForSession } from "../session-state"
 import { messageStoreBus } from "./bus"
-import type { SessionUsageState } from "./types"
+import type { SessionUsageState, UsageEntry } from "./types"
+
+interface UsageTotals {
+  cost: number
+  inputTokens: number
+  outputTokens: number
+  reasoningTokens: number
+}
+
+// OpenCode's session usage counters only ever grow; a committed revert deletes
+// messages without decrementing them. Track the reverted usage separately so
+// the displayed totals can subtract it from the authoritative session totals.
+const revertedUsage = new Map<string, UsageTotals>()
+
+function revertedKey(instanceId: string, sessionId: string) {
+  return `${instanceId}:${sessionId}`
+}
+
+export function recordRevertedUsage(instanceId: string, sessionId: string, entries: UsageEntry[]): void {
+  if (!entries.length) return
+  const key = revertedKey(instanceId, sessionId)
+  const current = revertedUsage.get(key) ?? { cost: 0, inputTokens: 0, outputTokens: 0, reasoningTokens: 0 }
+  for (const entry of entries) {
+    current.cost += entry.cost
+    current.inputTokens += entry.inputTokens
+    current.outputTokens += entry.outputTokens
+    current.reasoningTokens += entry.reasoningTokens
+  }
+  revertedUsage.set(key, current)
+}
+
+export function clearRevertedUsage(instanceId: string, sessionId?: string): void {
+  if (sessionId) {
+    revertedUsage.delete(revertedKey(instanceId, sessionId))
+    return
+  }
+  for (const key of revertedUsage.keys()) {
+    if (key.startsWith(`${instanceId}:`)) revertedUsage.delete(key)
+  }
+}
 
 function getLatestUsageEntry(usage?: SessionUsageState) {
   if (!usage?.latestMessageId) return undefined
@@ -40,7 +79,17 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void {
   let contextAvailableFromPrevious = false
   let isSubscriptionModel = false
 
-  if (!hasUsageEntries && previousInfo) {
+  // The session record carries the server's cumulative usage, which covers the
+  // whole session regardless of which messages are currently loaded. Message
+  // sums only stand in when the server does not report session usage.
+  if (session.tokens) {
+    const reverted = revertedUsage.get(revertedKey(instanceId, sessionId))
+    totalInputTokens = Math.max(0, session.tokens.input - (reverted?.inputTokens ?? 0))
+    totalOutputTokens = Math.max(0, session.tokens.output - (reverted?.outputTokens ?? 0))
+    totalReasoningTokens = Math.max(0, session.tokens.reasoning - (reverted?.reasoningTokens ?? 0))
+    totalCost = Math.max(0, (session.cost ?? 0) - (reverted?.cost ?? 0))
+    if (!hasUsageEntries) actualUsageTokens = previousInfo?.actualUsageTokens ?? 0
+  } else if (!hasUsageEntries && previousInfo) {
     totalInputTokens = previousInfo.inputTokens
     totalOutputTokens = previousInfo.outputTokens
     totalReasoningTokens = previousInfo.reasoningTokens

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

@@ -57,6 +57,7 @@ import {
   getSessionHasMore,
   getSessionNextCursor,
   getSessionListIds,
+  sessionInfoByInstance,
 } from "./session-state"
 import { deleteSessionAttachments } from "./attachments"
 import { DEFAULT_MODEL_OUTPUT_LIMIT, getActiveCatalogLocation, getDefaultModel, isModelValid } from "./session-models"
@@ -1460,7 +1461,10 @@ async function loadMessages(
   if (!planned) return
 
   const alreadyLoaded = messagesLoaded().get(instanceId)?.has(sessionId)
-  if (alreadyLoaded && !force) return
+  if (alreadyLoaded && !force) {
+    if (!sessionInfoByInstance().get(instanceId)?.has(sessionId)) updateSessionInfo(instanceId, sessionId)
+    return
+  }
 
   const previousError = getSessionMessagesLoadError(instanceId, sessionId)
   if (previousError && !force) return

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

@@ -54,6 +54,7 @@ import {
   setSessionRevertV2,
 } from "./message-v2/bridge"
 import { messageStoreBus } from "./message-v2/bus"
+import { updateSessionInfo } from "./message-v2/session-info"
 import { handleConversationAssistantPartUpdated } from "./conversation-speech"
 
 const log = getLogger("sse")
@@ -113,6 +114,7 @@ function handleNativeSessionEvent(instanceId: string, event: NativeSessionEvent)
         session.cost = event.data.cost as unknown as number
         session.tokens = event.data.tokens as Session["tokens"]
       })
+      updateSessionInfo(instanceId, event.data.sessionID)
       return
     case "session.moved":
       handleSessionMoved(instanceId, event.data)

packages/ui/src/stores/session-usage-contract.test.ts

@@ -0,0 +1,156 @@
+import assert from "node:assert/strict"
+import { test } from "node:test"
+import type { OpenCodeEvent } from "@opencode/client"
+import { destroyOpenCodeData } from "./opencode-data.ts"
+import { sdkManager } from "../lib/sdk-manager.ts"
+import { sseManager } from "../lib/sse-manager.ts"
+import { addInstance, handleInstanceInvalidation, removeInstance } from "./instances.ts"
+import { handleNativeSessionEvent } from "./session-events.ts"
+import { messageStoreBus } from "./message-v2/bus.ts"
+import { loadMessages } from "./session-api.ts"
+import { getSessionInfo, setActiveSession, setMessagesLoaded, setSessions } from "./session-state.ts"
+
+const model = { providerID: "fixture", id: "fixture" }
+
+function tokens(input: number, output: number) {
+  return { input, output, reasoning: 0, cache: { read: 0, write: 0 } }
+}
+
+function setup(instanceId: string, sessionId: string, options: { active?: boolean; session?: Record<string, unknown> } = {}) {
+  const statuses = sseManager.getStatuses
+  sseManager.getStatuses = () => new Map([[instanceId, "connected"]])
+  const client = {
+    session: { active: async () => ({}), get: async () => ({ id: sessionId, location: { directory: "/fixture" }, time: { created: 1, updated: 5 } }) },
+    message: { list: async () => ({ data: [], cursor: {} }) },
+  } as any
+  ;(sdkManager as any).clients.set(`${instanceId}:/workspaces/${instanceId}/instance`, client)
+  addInstance({ id: instanceId, folder: "/fixture", port: 0, pid: 0, proxyPath: "", status: "ready", client })
+  setSessions(previous => new Map(previous).set(instanceId, new Map([[sessionId, {
+    id: sessionId, instanceId, title: "Fixture", parentId: null, location: { directory: "/fixture" },
+    status: "idle", agent: "build", model: { providerId: "fixture", modelId: "fixture" }, time: { created: 1, updated: 1 },
+    ...options.session,
+  } as any]])))
+  if (options.active !== false) setActiveSession(instanceId, sessionId)
+  const emit = (type: string, data: Record<string, unknown>, created: number) => {
+    const event = { id: `evt_${created}`, type, created, data, location: { directory: "/fixture" } } as OpenCodeEvent
+    handleInstanceInvalidation(instanceId, event)
+    handleNativeSessionEvent(instanceId, event)
+  }
+  const cleanup = () => {
+    sseManager.getStatuses = statuses
+    destroyOpenCodeData(instanceId)
+    removeInstance(instanceId, { authoritative: false })
+    sdkManager.destroyClientsForInstance(instanceId)
+    if (messageStoreBus.getInstance(instanceId)) messageStoreBus.unregisterInstance(instanceId)
+  }
+  return { emit, cleanup }
+}
+
+function assertUsage(instanceId: string, sessionId: string, expected: { cost: number; input: number; output: number }) {
+  const info = getSessionInfo(instanceId, sessionId)
+  assert.ok(info, "session info must exist")
+  assert.equal(info.cost, expected.cost)
+  assert.equal(info.inputTokens, expected.input)
+  assert.equal(info.outputTokens, expected.output)
+}
+
+test("a completed step on the active session refreshes the session usage totals", () => {
+  const instanceId = "usage-step-ended", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    const base = { sessionID: sessionId, assistantMessageID: "m" }
+    emit("session.step.started", { ...base, agent: "build", model, started: 1 }, 1)
+    emit("session.step.ended", { ...base, finish: "stop", cost: 1, tokens: tokens(10, 5) }, 2)
+    assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 })
+  } finally { cleanup() }
+})
+
+test("a usage-only event fills the totals when no messages are loaded", () => {
+  const instanceId = "usage-only", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    emit("session.usage.updated", { sessionID: sessionId, cost: 2.5, tokens: tokens(1200, 300) }, 1)
+    assertUsage(instanceId, sessionId, { cost: 2.5, input: 1200, output: 300 })
+  } finally { cleanup() }
+})
+
+test("usage events for an inactive session still update its totals", () => {
+  const instanceId = "usage-inactive", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId, { active: false })
+  try {
+    emit("session.usage.updated", { sessionID: sessionId, cost: 2.5, tokens: tokens(1200, 300) }, 1)
+    assertUsage(instanceId, sessionId, { cost: 2.5, input: 1200, output: 300 })
+  } finally { cleanup() }
+})
+
+test("opening an already-loaded session initializes totals from the hydrated session record", async () => {
+  const instanceId = "usage-hydrated", sessionId = "s"
+  const { cleanup } = setup(instanceId, sessionId, { session: { cost: 2.5, tokens: tokens(1200, 300) } })
+  setMessagesLoaded(previous => new Map(previous).set(instanceId, new Set([sessionId])))
+  try {
+    assert.equal(getSessionInfo(instanceId, sessionId), undefined)
+    await loadMessages(instanceId, sessionId)
+    assertUsage(instanceId, sessionId, { cost: 2.5, input: 1200, output: 300 })
+  } finally { cleanup() }
+})
+
+test("authoritative session totals win over a partially loaded transcript", () => {
+  const instanceId = "usage-partial-window", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    const base = { sessionID: sessionId, assistantMessageID: "m" }
+    emit("session.step.started", { ...base, agent: "build", model, started: 1 }, 1)
+    emit("session.step.ended", { ...base, finish: "stop", cost: 1, tokens: tokens(10, 5) }, 2)
+    emit("session.usage.updated", { sessionID: sessionId, cost: 20, tokens: tokens(2000, 500) }, 3)
+    assertUsage(instanceId, sessionId, { cost: 20, input: 2000, output: 500 })
+  } finally { cleanup() }
+})
+
+test("a committed revert subtracts the removed messages from the session totals", () => {
+  const instanceId = "usage-revert", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    for (const [id, cost, created] of [["m1", 1, 1], ["m2", 4, 3]] as const) {
+      const base = { sessionID: sessionId, assistantMessageID: id }
+      emit("session.step.started", { ...base, agent: "build", model, started: created }, created)
+      emit("session.step.ended", { ...base, finish: "stop", cost, tokens: tokens(10 * cost, 5 * cost) }, created + 1)
+    }
+    emit("session.usage.updated", { sessionID: sessionId, cost: 5, tokens: tokens(50, 25) }, 5)
+    assertUsage(instanceId, sessionId, { cost: 5, input: 50, output: 25 })
+    emit("session.revert.committed", { sessionID: sessionId, to: "m2" }, 6)
+    assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 })
+    // The server's counters do not decrement; a later usage event must not
+    // restore the reverted usage.
+    emit("session.usage.updated", { sessionID: sessionId, cost: 5, tokens: tokens(50, 25) }, 7)
+    assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 })
+  } finally { cleanup() }
+})
+
+test("reverting the only message shows zero usage", () => {
+  const instanceId = "usage-revert-all", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    const base = { sessionID: sessionId, assistantMessageID: "m1" }
+    emit("session.step.started", { ...base, agent: "build", model, started: 1 }, 1)
+    emit("session.step.ended", { ...base, finish: "stop", cost: 1, tokens: tokens(10, 5) }, 2)
+    emit("session.usage.updated", { sessionID: sessionId, cost: 1, tokens: tokens(10, 5) }, 3)
+    assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 })
+    emit("session.revert.committed", { sessionID: sessionId, to: "m1" }, 4)
+    assertUsage(instanceId, sessionId, { cost: 0, input: 0, output: 0 })
+  } finally { cleanup() }
+})
+
+test("message sums stand in when the server reports no session usage", () => {
+  const instanceId = "usage-no-session-totals", sessionId = "s"
+  const { emit, cleanup } = setup(instanceId, sessionId)
+  try {
+    for (const [id, cost, created] of [["m1", 1, 1], ["m2", 4, 3]] as const) {
+      const base = { sessionID: sessionId, assistantMessageID: id }
+      emit("session.step.started", { ...base, agent: "build", model, started: created }, created)
+      emit("session.step.ended", { ...base, finish: "stop", cost, tokens: tokens(10 * cost, 5 * cost) }, created + 1)
+    }
+    assertUsage(instanceId, sessionId, { cost: 5, input: 50, output: 25 })
+    emit("session.revert.committed", { sessionID: sessionId, to: "m2" }, 5)
+    assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 })
+  } finally { cleanup() }
+})