fix: update v2 session usage metrics (#35468)

This commit is contained in:
Aiden Cline 2026-07-07 14:31:54 -05:00 committed by GitHub
commit 910e37f6d8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1013 additions and 107 deletions

View file

@ -110,6 +110,9 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
directory: process.cwd(),
})
const messageIndex = new Map<string, Map<string, number>>()
const sessionRefreshGeneration = new Map<string, number>()
const sessionRefreshApplied = new Map<string, number>()
const sessionUsage = new Map<string, { generation: number; cost: number; tokens: SessionV2Info["tokens"] }>()
let connectionGeneration = 0
let statusChanges: Set<string> | undefined
let bootstrapping: Promise<void> | undefined
@ -119,6 +122,24 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
setStore("session", "status", sessionID, status)
}
function nextSessionRefresh(sessionID: string) {
const generation = (sessionRefreshGeneration.get(sessionID) ?? 0) + 1
sessionRefreshGeneration.set(sessionID, generation)
return generation
}
function applySessionRefresh(sessionID: string, generation: number) {
if ((sessionRefreshApplied.get(sessionID) ?? 0) > generation) return false
sessionRefreshApplied.set(sessionID, generation)
return true
}
function updateSessionUsage(sessionID: string, cost: number, tokens: SessionV2Info["tokens"]) {
sessionUsage.set(sessionID, { generation: (sessionUsage.get(sessionID)?.generation ?? 0) + 1, cost, tokens })
if (!store.session.info[sessionID]) return
setStore("session", "info", sessionID, { cost, tokens })
}
const message = {
update(sessionID: string, fn: (messages: SessionMessage[], index: Map<string, number>) => void) {
setStore(
@ -222,6 +243,8 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
}
function removeSession(sessionID: string) {
sessionRefreshApplied.set(sessionID, nextSessionRefresh(sessionID))
sessionUsage.delete(sessionID)
messageIndex.delete(sessionID)
setStore(
"session",
@ -250,6 +273,9 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
case "session.deleted":
removeSession(event.data.sessionID)
break
case "session.usage.updated":
updateSessionUsage(event.data.sessionID, event.data.cost, event.data.tokens)
break
case "catalog.updated":
void Promise.all([
result.location.model.refresh(event.location),
@ -420,7 +446,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
})
})
break
case "session.step.ended":
case "session.step.ended": {
message.update(event.data.sessionID, (draft, index) => {
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
if (!currentAssistant) return
@ -432,6 +458,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
currentAssistant.snapshot = { ...currentAssistant.snapshot, end: event.data.snapshot }
})
break
}
case "session.step.failed":
message.update(event.data.sessionID, (draft, index) => {
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
@ -440,6 +467,10 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
currentAssistant.finish = "error"
currentAssistant.error = event.data.error
currentAssistant.retry = undefined
if (event.data.cost !== undefined && event.data.tokens !== undefined) {
currentAssistant.cost = event.data.cost
currentAssistant.tokens = event.data.tokens
}
})
break
case "session.text.started":
@ -639,8 +670,9 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
setStore("session", "info", event.data.sessionID, "revert", undefined)
break
case "session.revert.committed":
if (store.session.info[event.data.sessionID])
if (store.session.info[event.data.sessionID]) {
setStore("session", "info", event.data.sessionID, "revert", undefined)
}
setStore(
"session",
"input",
@ -811,7 +843,17 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
return store.session.compaction[sessionID]
},
async refresh(sessionID: string) {
setStore("session", "info", sessionID, mutable(await sdk.api.session.get({ sessionID })))
const generation = nextSessionRefresh(sessionID)
const usageGeneration = sessionUsage.get(sessionID)?.generation ?? 0
const info = mutable(await sdk.api.session.get({ sessionID }))
if (!applySessionRefresh(sessionID, generation)) return
const usage = sessionUsage.get(sessionID)
setStore(
"session",
"info",
sessionID,
usage && usage.generation !== usageGeneration ? { ...info, cost: usage.cost, tokens: usage.tokens } : info,
)
registerSession(sessionID)
},
message: {
@ -994,6 +1036,8 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
async function bootstrap() {
if (bootstrapping) return bootstrapping
const generation = new Map(sessionRefreshApplied)
const usageGeneration = new Map(Array.from(sessionUsage, ([id, usage]) => [id, usage.generation]))
bootstrapping = Promise.allSettled([
sdk.api.session
.list({
@ -1007,7 +1051,15 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
"session",
"info",
produce((draft) => {
for (const session of response.data) draft[session.id] = mutable(session)
for (const session of response.data) {
if ((sessionRefreshApplied.get(session.id) ?? 0) !== (generation.get(session.id) ?? 0)) continue
const usage = sessionUsage.get(session.id)
draft[session.id] = mutable(
usage && usage.generation !== (usageGeneration.get(session.id) ?? 0)
? { ...session, cost: usage.cost, tokens: usage.tokens }
: session,
)
}
}),
)
for (const session of response.data) registerSession(session.id)

View file

@ -1,7 +1,8 @@
import type { AssistantMessage } from "@opencode-ai/sdk/v2"
import type { TuiPlugin, TuiPluginApi } from "@opencode-ai/plugin/tui"
import type { BuiltinTuiPlugin } from "../builtins"
import { createMemo } from "solid-js"
import { useData } from "../../context/data"
import { lastAssistantWithUsage } from "../../util/session"
const id = "internal:sidebar-context"
@ -11,13 +12,14 @@ const money = new Intl.NumberFormat("en-US", {
})
function View(props: { api: TuiPluginApi; session_id: string }) {
const data = useData()
const theme = () => props.api.theme.current
const msg = createMemo(() => props.api.state.session.messages(props.session_id))
const session = createMemo(() => props.api.state.session.get(props.session_id))
const msg = createMemo(() => data.session.message.list(props.session_id))
const session = createMemo(() => data.session.get(props.session_id))
const cost = createMemo(() => session()?.cost ?? 0)
const state = createMemo(() => {
const last = msg().findLast((item): item is AssistantMessage => item.role === "assistant" && item.tokens.output > 0)
const last = lastAssistantWithUsage(msg(), session()?.revert?.messageID)
if (!last) {
return {
tokens: 0,
@ -27,7 +29,9 @@ function View(props: { api: TuiPluginApi; session_id: string }) {
const tokens =
last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
const model = props.api.state.provider.find((item) => item.id === last.providerID)?.models[last.modelID]
const model = data.location
.model.list(session()?.location)
?.find((model) => model.providerID === last.model.providerID && model.id === last.model.id)
return {
tokens,
percent: model?.limit.context ? Math.round((tokens / model.limit.context) * 100) : null,

View file

@ -6,6 +6,7 @@ import { SplitBorder } from "../../ui/border"
import { Locale } from "../../util/locale"
import { useTerminalDimensions } from "@opentui/solid"
import { useCommandShortcut, useOpencodeKeymap } from "../../keymap"
import { lastAssistantWithUsage } from "../../util/session"
export function SubagentFooter() {
const route = useRouteData("session")
@ -22,17 +23,15 @@ export function SubagentFooter() {
const usage = createMemo(() => {
const current = session()
if (!current) return
const last = lastAssistantWithUsage(data.session.message.list(route.sessionID), current.revert?.messageID)
if (!last) return
const tokens =
current.tokens.input +
current.tokens.output +
current.tokens.reasoning +
current.tokens.cache.read +
current.tokens.cache.write
last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
if (tokens <= 0) return
const model = data.location
.model.list(current.location)
?.find((model) => model.providerID === current.model?.providerID && model.id === current.model.id)
?.find((model) => model.providerID === last.model.providerID && model.id === last.model.id)
const pct = model?.limit.context ? `${Math.round((tokens / model.limit.context) * 100)}%` : undefined
const cost = current.cost
@ -83,10 +82,10 @@ export function SubagentFooter() {
</box>
<box flexDirection="row" gap={2}>
<box
onMouseOver={() => setHover("parent")}
onMouseOut={() => setHover(null)}
onMouseOver={() => setHover("parent")}
onMouseOut={() => setHover(null)}
onMouseUp={() => keymap.dispatchCommand("session.parent")}
backgroundColor={hover() === "parent" ? theme.backgroundElement : theme.backgroundPanel}
backgroundColor={hover() === "parent" ? theme.backgroundElement : theme.backgroundPanel}
>
<text fg={theme.text}>
Parent <span style={{ fg: theme.textMuted }}>{parentShortcut()}</span>

View file

@ -1,3 +1,17 @@
import type { SessionMessage, SessionMessageAssistant } from "@opencode-ai/sdk/v2"
export function isDefaultTitle(title: string) {
return /^(New session - |Child session - )\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/.test(title)
}
export function lastAssistantWithUsage(messages: ReadonlyArray<SessionMessage>, boundary?: string) {
const boundaryIndex = boundary ? messages.findIndex((message) => message.id === boundary) : -1
if (boundary && boundaryIndex === -1) return undefined
return messages.findLast(
(
message,
index,
): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } =>
message.type === "assistant" && message.tokens !== undefined && (boundaryIndex === -1 || index < boundaryIndex),
)
}