From 927cf77e41ecc2321ba35152e62725c7bebf9a17 Mon Sep 17 00:00:00 2001 From: Mark Erikson Date: Sun, 20 Sep 2026 16:54:41 -0400 Subject: [PATCH 1/5] fix(ui): refresh session usage info on step and usage events The V2 migration dropped the live updateSessionInfo() calls from the SSE handlers, so the Status tab's cost/token chips were only written at message load (and stayed at 0 for sessions that were empty when opened). Recompute them after projecting messages for session.step.ended, session.step.failed, session.usage.updated, and forced resyncs. --- packages/ui/src/stores/instances.ts | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/packages/ui/src/stores/instances.ts b/packages/ui/src/stores/instances.ts index 1ed49b576..eb79c68c5 100644 --- a/packages/ui/src/stores/instances.ts +++ b/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 { 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" @@ -1995,6 +1996,13 @@ async function sendFormCancel(instanceId: string, formId: string): Promise } } +// Events after which assistant message cost/token totals may have changed. +const USAGE_EVENT_TYPES = new Set([ + "session.step.ended", + "session.step.failed", + "session.usage.updated", +]) + function handleInstanceInvalidation(instanceId: string, event: Parameters>[1]): void { const instance = instances().get(instanceId) if (!instance?.client) return @@ -2021,13 +2029,16 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters, preserveOmitted = true) => { projectMessages(data, preserveOmitted) From 1085812537dec2f523cc7927c1c6de0fee183872 Mon Sep 17 00:00:00 2001 From: Mark Erikson Date: Sun, 20 Sep 2026 19:35:38 -0400 Subject: [PATCH 2/5] fix(ui): refresh session usage from usage events, hydration, and reverts - fall back to the session record's cost/tokens when no message usage is loaded - refresh on session.usage.updated for any session, after the record is updated - refresh after session.revert.committed removes messages - initialize info for already-loaded sessions on open - add session-usage-contract.test.ts to the force-exit CI list --- .github/workflows/pr-build.yml | 1 + packages/ui/src/stores/instances.ts | 2 +- .../ui/src/stores/message-v2/session-info.ts | 11 +- packages/ui/src/stores/session-api.ts | 6 +- packages/ui/src/stores/session-events.ts | 2 + .../src/stores/session-usage-contract.test.ts | 110 ++++++++++++++++++ 6 files changed, 129 insertions(+), 3 deletions(-) create mode 100644 packages/ui/src/stores/session-usage-contract.test.ts diff --git a/.github/workflows/pr-build.yml b/.github/workflows/pr-build.yml index 05f8ba084..292c3a3a3 100644 --- a/.github/workflows/pr-build.yml +++ b/.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 diff --git a/packages/ui/src/stores/instances.ts b/packages/ui/src/stores/instances.ts index eb79c68c5..2e785cb76 100644 --- a/packages/ui/src/stores/instances.ts +++ b/packages/ui/src/stores/instances.ts @@ -2000,7 +2000,6 @@ async function sendFormCancel(instanceId: string, formId: string): Promise const USAGE_EVENT_TYPES = new Set([ "session.step.ended", "session.step.failed", - "session.usage.updated", ]) function handleInstanceInvalidation(instanceId: string, event: Parameters>[1]): void { @@ -2049,6 +2048,7 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters= event.data.to) removeMessageV2(instanceId, messageId, sessionId) } + updateSessionInfo(instanceId, sessionId) } } const data = applyOpenCodeDataEvent(instanceId, instance.folder, event, project, (next) => { diff --git a/packages/ui/src/stores/message-v2/session-info.ts b/packages/ui/src/stores/message-v2/session-info.ts index fc1575102..cb0e6a6b1 100644 --- a/packages/ui/src/stores/message-v2/session-info.ts +++ b/packages/ui/src/stores/message-v2/session-info.ts @@ -40,7 +40,16 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void { let contextAvailableFromPrevious = false let isSubscriptionModel = false - if (!hasUsageEntries && previousInfo) { + // Message-derived totals win when the transcript is loaded, since they drop + // after a revert while the server's session counters only ever grow. The + // session record fills in for usage-only events and unloaded transcripts. + if (!hasUsageEntries && session.tokens) { + totalInputTokens = session.tokens.input + totalOutputTokens = session.tokens.output + totalReasoningTokens = session.tokens.reasoning + totalCost = session.cost ?? 0 + actualUsageTokens = previousInfo?.actualUsageTokens ?? 0 + } else if (!hasUsageEntries && previousInfo) { totalInputTokens = previousInfo.inputTokens totalOutputTokens = previousInfo.outputTokens totalReasoningTokens = previousInfo.reasoningTokens diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index 3e2a803fa..28ead5a3c 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/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 diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts index b02a388fd..a91b8b999 100644 --- a/packages/ui/src/stores/session-events.ts +++ b/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) diff --git a/packages/ui/src/stores/session-usage-contract.test.ts b/packages/ui/src/stores/session-usage-contract.test.ts new file mode 100644 index 000000000..a8c042f61 --- /dev/null +++ b/packages/ui/src/stores/session-usage-contract.test.ts @@ -0,0 +1,110 @@ +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 } = {}) { + 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, 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("a committed revert that removes a message lowers the 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) + } + 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() } +}) From 489e8a9c6ac67a264e3d8a1f98ef242841fe4e64 Mon Sep 17 00:00:00 2001 From: Mark Erikson Date: Sun, 20 Sep 2026 20:00:44 -0400 Subject: [PATCH 3/5] fix(ui): make server session usage authoritative and track reverts separately The loaded transcript is a bounded window, so message sums cannot stand in for the session totals when the server reports them. Reverted messages are recorded as a separate adjustment because the server's usage counters do not decrement on session.revert.committed. --- packages/ui/src/stores/instances.ts | 11 ++-- .../ui/src/stores/message-v2/session-info.ts | 60 +++++++++++++++---- .../src/stores/session-usage-contract.test.ts | 48 ++++++++++++++- 3 files changed, 104 insertions(+), 15 deletions(-) diff --git a/packages/ui/src/stores/instances.ts b/packages/ui/src/stores/instances.ts index 2e785cb76..209958b7c 100644 --- a/packages/ui/src/stores/instances.ts +++ b/packages/ui/src/stores/instances.ts @@ -48,7 +48,7 @@ import { } from "./session-state" import { setHasInstances } from "./ui" import { messageStoreBus } from "./message-v2/bus" -import { updateSessionInfo } from "./message-v2/session-info" +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" @@ -1359,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) @@ -2045,9 +2046,11 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters= 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) } } diff --git a/packages/ui/src/stores/message-v2/session-info.ts b/packages/ui/src/stores/message-v2/session-info.ts index cb0e6a6b1..885555c75 100644 --- a/packages/ui/src/stores/message-v2/session-info.ts +++ b/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() + +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,15 +79,16 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void { let contextAvailableFromPrevious = false let isSubscriptionModel = false - // Message-derived totals win when the transcript is loaded, since they drop - // after a revert while the server's session counters only ever grow. The - // session record fills in for usage-only events and unloaded transcripts. - if (!hasUsageEntries && session.tokens) { - totalInputTokens = session.tokens.input - totalOutputTokens = session.tokens.output - totalReasoningTokens = session.tokens.reasoning - totalCost = session.cost ?? 0 - actualUsageTokens = previousInfo?.actualUsageTokens ?? 0 + // 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 diff --git a/packages/ui/src/stores/session-usage-contract.test.ts b/packages/ui/src/stores/session-usage-contract.test.ts index a8c042f61..653a79660 100644 --- a/packages/ui/src/stores/session-usage-contract.test.ts +++ b/packages/ui/src/stores/session-usage-contract.test.ts @@ -94,9 +94,55 @@ test("opening an already-loaded session initializes totals from the hydrated ses } finally { cleanup() } }) -test("a committed revert that removes a message lowers the totals", () => { +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 } From 761968a2070b1f95d9488f7ab3b94317f7484a30 Mon Sep 17 00:00:00 2001 From: Mark Erikson Date: Sun, 20 Sep 2026 20:19:38 -0400 Subject: [PATCH 4/5] fix(ui): show lifetime session usage regardless of transcript residency - session totals come from the server's cumulative usage counters and are not adjusted on revert, so the display matches across windows and reloads - tests cover revert from an anchored window and post-reload hydration --- packages/ui/src/stores/instances.ts | 12 ++-- .../ui/src/stores/message-v2/session-info.ts | 57 +++---------------- .../src/stores/session-usage-contract.test.ts | 43 ++++++++++---- 3 files changed, 44 insertions(+), 68 deletions(-) diff --git a/packages/ui/src/stores/instances.ts b/packages/ui/src/stores/instances.ts index 209958b7c..6ee75b410 100644 --- a/packages/ui/src/stores/instances.ts +++ b/packages/ui/src/stores/instances.ts @@ -48,7 +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 { 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" @@ -1359,7 +1359,6 @@ 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) @@ -2046,12 +2045,9 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters 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) + for (const messageId of messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId)) { + if (messageId >= event.data.to) removeMessageV2(instanceId, messageId, sessionId) + } } } const data = applyOpenCodeDataEvent(instanceId, instance.folder, event, project, (next) => { diff --git a/packages/ui/src/stores/message-v2/session-info.ts b/packages/ui/src/stores/message-v2/session-info.ts index 885555c75..c2b290ab3 100644 --- a/packages/ui/src/stores/message-v2/session-info.ts +++ b/packages/ui/src/stores/message-v2/session-info.ts @@ -2,46 +2,7 @@ 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, 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() - -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) - } -} +import type { SessionUsageState } from "./types" function getLatestUsageEntry(usage?: SessionUsageState) { if (!usage?.latestMessageId) return undefined @@ -79,15 +40,15 @@ export function updateSessionInfo(instanceId: string, sessionId: string): void { let contextAvailableFromPrevious = false let isSubscriptionModel = false - // 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. + // The session record carries the server's lifetime usage, which covers the + // whole session regardless of which messages are currently loaded and is not + // reduced by reverts. 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)) + totalInputTokens = session.tokens.input + totalOutputTokens = session.tokens.output + totalReasoningTokens = session.tokens.reasoning + totalCost = session.cost ?? 0 if (!hasUsageEntries) actualUsageTokens = previousInfo?.actualUsageTokens ?? 0 } else if (!hasUsageEntries && previousInfo) { totalInputTokens = previousInfo.inputTokens diff --git a/packages/ui/src/stores/session-usage-contract.test.ts b/packages/ui/src/stores/session-usage-contract.test.ts index 653a79660..26b6ad78d 100644 --- a/packages/ui/src/stores/session-usage-contract.test.ts +++ b/packages/ui/src/stores/session-usage-contract.test.ts @@ -106,7 +106,7 @@ test("authoritative session totals win over a partially loaded transcript", () = } finally { cleanup() } }) -test("a committed revert subtracts the removed messages from the session totals", () => { +test("a committed revert keeps the lifetime session totals", () => { const instanceId = "usage-revert", sessionId = "s" const { emit, cleanup } = setup(instanceId, sessionId) try { @@ -118,28 +118,49 @@ test("a committed revert subtracts the removed messages from the session totals" 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. + assertUsage(instanceId, sessionId, { cost: 5, input: 50, output: 25 }) emit("session.usage.updated", { sessionID: sessionId, cost: 5, tokens: tokens(50, 25) }, 7) - assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 }) + assertUsage(instanceId, sessionId, { cost: 5, input: 50, output: 25 }) } finally { cleanup() } }) -test("reverting the only message shows zero usage", () => { - const instanceId = "usage-revert-all", sessionId = "s" +test("a revert from an anchored historical window shows the same lifetime totals", () => { + const instanceId = "usage-revert-anchored", sessionId = "s" const { emit, cleanup } = setup(instanceId, sessionId) try { + // Only the boundary message is resident; the newer $4 message that the + // native revert deletes was never loaded into this window. 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.usage.updated", { sessionID: sessionId, cost: 5, tokens: tokens(50, 25) }, 3) emit("session.revert.committed", { sessionID: sessionId, to: "m1" }, 4) - assertUsage(instanceId, sessionId, { cost: 0, input: 0, output: 0 }) + assert.equal(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId).length, 0) + assertUsage(instanceId, sessionId, { cost: 5, input: 50, output: 25 }) } finally { cleanup() } }) +test("hydrating a session after a revert shows the same lifetime totals as before reload", async () => { + const instanceId = "usage-revert-reload", sessionId = "s" + const first = setup(instanceId, sessionId) + try { + const base = { sessionID: sessionId, assistantMessageID: "m1" } + first.emit("session.step.started", { ...base, agent: "build", model, started: 1 }, 1) + first.emit("session.step.ended", { ...base, finish: "stop", cost: 1, tokens: tokens(10, 5) }, 2) + first.emit("session.usage.updated", { sessionID: sessionId, cost: 1, tokens: tokens(10, 5) }, 3) + first.emit("session.revert.committed", { sessionID: sessionId, to: "m1" }, 4) + assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 }) + } finally { first.cleanup() } + + // A fresh renderer only has the persisted session record and an empty transcript. + const reloaded = setup(instanceId, sessionId, { session: { cost: 1, tokens: tokens(10, 5) } }) + setMessagesLoaded(previous => new Map(previous).set(instanceId, new Set([sessionId]))) + try { + await loadMessages(instanceId, sessionId) + assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 }) + } finally { reloaded.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) @@ -150,7 +171,5 @@ test("message sums stand in when the server reports no session usage", () => { 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() } }) From 16d70b9fd968d7af8bba5a2b6edf9580c6aaac48 Mon Sep 17 00:00:00 2001 From: Mark Erikson Date: Sun, 20 Sep 2026 20:50:14 -0400 Subject: [PATCH 5/5] test(ui): hydrate the post-revert usage case under a fresh instance identity --- .../ui/src/stores/session-usage-contract.test.ts | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/packages/ui/src/stores/session-usage-contract.test.ts b/packages/ui/src/stores/session-usage-contract.test.ts index 26b6ad78d..719715ce4 100644 --- a/packages/ui/src/stores/session-usage-contract.test.ts +++ b/packages/ui/src/stores/session-usage-contract.test.ts @@ -152,12 +152,17 @@ test("hydrating a session after a revert shows the same lifetime totals as befor assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 }) } finally { first.cleanup() } - // A fresh renderer only has the persisted session record and an empty transcript. - const reloaded = setup(instanceId, sessionId, { session: { cost: 1, tokens: tokens(10, 5) } }) - setMessagesLoaded(previous => new Map(previous).set(instanceId, new Set([sessionId]))) + // A fresh renderer only has the persisted session record and an empty + // transcript. It gets its own instance identity because removeInstance() + // leaves the previous session info in place. + const reloadedInstanceId = `${instanceId}-fresh` + const reloaded = setup(reloadedInstanceId, sessionId, { session: { cost: 1, tokens: tokens(10, 5) } }) + setMessagesLoaded(previous => new Map(previous).set(reloadedInstanceId, new Set([sessionId]))) try { - await loadMessages(instanceId, sessionId) - assertUsage(instanceId, sessionId, { cost: 1, input: 10, output: 5 }) + assert.equal(getSessionInfo(reloadedInstanceId, sessionId), undefined) + assert.equal(messageStoreBus.getOrCreate(reloadedInstanceId).getSessionMessageIds(sessionId).length, 0) + await loadMessages(reloadedInstanceId, sessionId) + assertUsage(reloadedInstanceId, sessionId, { cost: 1, input: 10, output: 5 }) } finally { reloaded.cleanup() } })