fix: address session usage review

This commit is contained in:
Aiden Cline
2026-07-05 19:42:40 -05:00
parent 59f0ac4774
commit 3401d57bdb
7 changed files with 216 additions and 27 deletions
+6 -2
View File
@@ -8,6 +8,7 @@ import {
Message,
SystemPart,
isContextOverflowFailure,
type ProviderMetadata,
type ProviderErrorEvent,
} from "@opencode-ai/llm"
import { Cause, DateTime, Effect, Exit, FiberSet, Layer, Option, Semaphore, Stream } from "effect"
@@ -52,7 +53,10 @@ type StepTokens = {
readonly cache: { readonly read: number; readonly write: number }
}
export function calculateCost(costs: ModelV2.Info["cost"], tokens: StepTokens) {
export function calculateCost(costs: ModelV2.Info["cost"], tokens: StepTokens, metadata?: ProviderMetadata) {
const totalNanoAiu = metadata?.copilot?.totalNanoAiu
if (typeof totalNanoAiu === "number" && Number.isFinite(totalNanoAiu) && totalNanoAiu >= 0)
return totalNanoAiu / 100_000_000_000
const context = tokens.input + tokens.cache.read + tokens.cache.write
const tier = costs
.filter((cost) => cost.tier?.type === "context" && context > cost.tier.size)
@@ -322,7 +326,7 @@ const layer = Layer.effect(
sessionID: session.id,
assistantMessageID: yield* publisher.startAssistant(),
finish: settlement.finish,
cost: calculateCost(resolved.cost, settlement.tokens),
cost: calculateCost(resolved.cost, settlement.tokens, settlement.providerMetadata),
tokens: settlement.tokens,
snapshot: endSnapshot,
files,
@@ -69,7 +69,13 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
let assistantActive = false
let assistantFailed = false
let providerFailed = false
let stepSettlement: { readonly finish: string; readonly tokens: ReturnType<typeof tokens> } | undefined
let stepSettlement:
| {
readonly finish: string
readonly tokens: ReturnType<typeof tokens>
readonly providerMetadata?: ProviderMetadata
}
| undefined
const startAssistant = Effect.fnUntraced(function* () {
if (assistantMessageID !== undefined) return assistantMessageID
@@ -381,7 +387,11 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
yield* flush()
assistantActive = false
if (stepSettlement) return yield* Effect.die(new Error("Duplicate step finish"))
stepSettlement = { finish: event.reason, tokens: tokens(event.usage) }
stepSettlement = {
finish: event.reason,
tokens: tokens(event.usage),
providerMetadata: event.usage?.providerMetadata ?? event.providerMetadata,
}
return
case "finish":
return
+12
View File
@@ -128,6 +128,16 @@ test("calculates step cost using the matching context tier", () => {
).toBeCloseTo(0.0002926)
})
test("prefers Copilot billed cost over token estimation", () => {
expect(
SessionRunnerLLM.calculateCost(
[{ input: 1, output: 2, cache: { read: 0.1, write: 0.5 } }],
{ input: 80, output: 10, reasoning: 2, cache: { read: 20, write: 1 } },
{ copilot: { totalNanoAiu: 4_473_525_000 } },
),
).toBeCloseTo(0.04473525)
})
const authorizations: Tool.Context[] = []
const executions: string[] = []
const permission = Layer.succeed(
@@ -1543,6 +1553,7 @@ describe("SessionRunnerLLM", () => {
LLMEvent.stepFinish({
index: 0,
reason: "tool-calls",
providerMetadata: { copilot: { totalNanoAiu: 4_473_525_000 } },
usage: {
inputTokens: 10,
nonCachedInputTokens: 8,
@@ -1563,6 +1574,7 @@ describe("SessionRunnerLLM", () => {
{
type: "assistant",
finish: "tool-calls",
cost: 0.04473525,
tokens: { input: 8, output: 3, reasoning: 1, cache: { read: 2, write: 0 } },
content: [
{ type: "reasoning", id: "reasoning-1", text: "Think" },
+16 -17
View File
@@ -100,6 +100,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
directory: process.cwd(),
})
const messageIndex = new Map<string, Map<string, number>>()
const sessionRefreshGeneration = new Map<string, number>()
let connectionGeneration = 0
let statusChanges: Set<string> | undefined
let bootstrapping: Promise<void> | undefined
@@ -358,22 +359,8 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
})
})
break
case "session.step.ended":
case "session.step.ended": {
setSessionStatus(event.data.sessionID, "running")
if (store.session.info[event.data.sessionID])
setStore(
"session",
"info",
event.data.sessionID,
produce((draft) => {
draft.cost += event.data.cost
draft.tokens.input += event.data.tokens.input
draft.tokens.output += event.data.tokens.output
draft.tokens.reasoning += event.data.tokens.reasoning
draft.tokens.cache.read += event.data.tokens.cache.read
draft.tokens.cache.write += event.data.tokens.cache.write
}),
)
message.update(event.data.sessionID, (draft, index) => {
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
if (!currentAssistant) return
@@ -384,7 +371,11 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
if (event.data.snapshot)
currentAssistant.snapshot = { ...currentAssistant.snapshot, end: event.data.snapshot }
})
void result.session
.refresh(event.data.sessionID)
.catch((error) => console.error("Failed to refresh session usage", error))
break
}
case "session.step.failed":
message.update(event.data.sessionID, (draft, index) => {
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
@@ -691,7 +682,11 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
return store.session.status[sessionID] ?? "idle"
},
async refresh(sessionID: string) {
setStore("session", "info", sessionID, mutable(await sdk.api.session.get({ sessionID })))
const generation = (sessionRefreshGeneration.get(sessionID) ?? 0) + 1
sessionRefreshGeneration.set(sessionID, generation)
const info = mutable(await sdk.api.session.get({ sessionID }))
if (sessionRefreshGeneration.get(sessionID) !== generation) return
setStore("session", "info", sessionID, info)
registerSession(sessionID)
},
message: {
@@ -875,6 +870,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
async function bootstrap() {
if (bootstrapping) return bootstrapping
const generation = new Map(sessionRefreshGeneration)
bootstrapping = Promise.allSettled([
sdk.api.session
.list({
@@ -888,7 +884,10 @@ 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 ((sessionRefreshGeneration.get(session.id) ?? 0) !== (generation.get(session.id) ?? 0)) continue
draft[session.id] = mutable(session)
}
}),
)
for (const session of response.data) registerSession(session.id)
+7 -2
View File
@@ -5,8 +5,13 @@ export function isDefaultTitle(title: string) {
}
export function lastAssistantWithUsage(messages: ReadonlyArray<SessionMessage>, boundary?: string) {
const boundaryIndex = boundary ? messages.findIndex((message) => message.id === boundary) : -1
if (boundary && boundaryIndex === -1) return
return messages.findLast(
(message): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } =>
message.type === "assistant" && message.tokens !== undefined && (!boundary || message.id < boundary),
(
message,
index,
): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } =>
message.type === "assistant" && message.tokens !== undefined && (boundaryIndex === -1 || index < boundaryIndex),
)
}
+160 -2
View File
@@ -108,6 +108,159 @@ test("refreshes resources into reactive getters", async () => {
}
})
test("refreshes usage without applying stale session snapshots", async () => {
const events = createEventStream()
const sessionID = "ses_usage_refresh"
let resolveSessions!: (response: Response) => void
const resolveSession: Array<(response: Response) => void> = []
let sessionsRequested = false
const calls = createFetch((url) => {
if (url.pathname === "/api/session") {
sessionsRequested = true
return new Promise<Response>((resolve) => {
resolveSessions = resolve
})
}
if (url.pathname === `/api/session/${sessionID}`) {
return new Promise<Response>((resolve) => {
resolveSession.push(resolve)
})
}
}, events)
let data!: ReturnType<typeof useData>
function Probe() {
data = useData()
return <box />
}
const app = await testRender(() => (
<TestTuiContexts>
<SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
<ProjectProvider>
<DataProvider>
<Probe />
</DataProvider>
</ProjectProvider>
</SDKProvider>
</TestTuiContexts>
))
try {
await wait(() => sessionsRequested)
emitEvent(events, {
id: "evt_usage_2",
created: 2,
type: "session.step.ended",
durable: durable(sessionID, 2),
data: {
sessionID,
assistantMessageID: "msg_usage_2",
finish: "stop",
cost: 0.5,
tokens: { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } },
},
})
await wait(() => resolveSession.length === 1)
resolveSessions(
json({
data: [
{
id: sessionID,
projectID: "proj_test",
cost: 0,
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
time: { created: 0, updated: 0 },
title: "Stale usage",
location: { directory },
},
],
watermarks: {},
cursor: {},
}),
)
resolveSession[0]!(
json({
data: {
id: sessionID,
projectID: "proj_test",
cost: 0.5,
tokens: { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } },
time: { created: 0, updated: 0 },
title: "Current usage",
location: { directory },
},
}),
)
await wait(() => data.session.get(sessionID)?.cost === 0.5)
expect(data.session.get(sessionID)?.tokens).toEqual({
input: 5,
output: 2,
reasoning: 1,
cache: { read: 1, write: 1 },
})
emitEvent(events, {
id: "evt_usage_3",
created: 3,
type: "session.step.ended",
durable: durable(sessionID, 3),
data: {
sessionID,
assistantMessageID: "msg_usage_3",
finish: "stop",
cost: 0.25,
tokens: { input: 3, output: 1, reasoning: 0, cache: { read: 0, write: 0 } },
},
})
emitEvent(events, {
id: "evt_usage_4",
created: 4,
type: "session.step.ended",
durable: durable(sessionID, 4),
data: {
sessionID,
assistantMessageID: "msg_usage_4",
finish: "stop",
cost: 0.25,
tokens: { input: 2, output: 1, reasoning: 0, cache: { read: 0, write: 0 } },
},
})
await wait(() => resolveSession.length === 3)
resolveSession[2]!(
json({
data: {
id: sessionID,
projectID: "proj_test",
cost: 1,
tokens: { input: 10, output: 4, reasoning: 1, cache: { read: 1, write: 1 } },
time: { created: 0, updated: 0 },
title: "Latest usage",
location: { directory },
},
}),
)
await wait(() => data.session.get(sessionID)?.cost === 1)
resolveSession[1]!(
json({
data: {
id: sessionID,
projectID: "proj_test",
cost: 0.75,
tokens: { input: 8, output: 3, reasoning: 1, cache: { read: 1, write: 1 } },
time: { created: 0, updated: 0 },
title: "Older usage",
location: { directory },
},
}),
)
await Bun.sleep(20)
expect(data.session.get(sessionID)?.cost).toBe(1)
} finally {
app.renderer.destroy()
}
})
test("reconnects the event stream and bootstraps fresh data", async () => {
const events = createEventStream()
const requests = { active: 0, event: 0, model: 0 }
@@ -348,6 +501,7 @@ test("connectedOnce is false until first connect and persists across disconnect"
test("tracks session status from active sessions and execution events", async () => {
const events = createEventStream()
let settled = false
const calls = createFetch((url) => {
if (url.pathname === "/api/session/active")
return json({ data: { "session-active": { type: "running" } }, watermarks: {} })
@@ -356,8 +510,10 @@ test("tracks session status from active sessions and execution events", async ()
data: {
id: "session-live",
projectID: "proj_test",
cost: 0,
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
cost: settled ? 0.75 : 0,
tokens: settled
? { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }
: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
time: { created: 0, updated: 0 },
title: "Live session",
location: { directory },
@@ -388,6 +544,7 @@ test("tracks session status from active sessions and execution events", async ()
expect(data.session.status("session-idle")).toBe("idle")
await data.session.refresh("session-live")
settled = true
emitEvent(events, {
id: "evt_step_started",
created: 0,
@@ -419,6 +576,7 @@ test("tracks session status from active sessions and execution events", async ()
const assistant = data.session.message.get("session-live", "message-live")
return assistant?.type === "assistant" && assistant.finish === "stop"
})
await wait(() => data.session.get("session-live")?.cost === 0.75)
expect(data.session.status("session-live")).toBe("running")
expect(data.session.get("session-live")).toMatchObject({
cost: 0.75,
+3 -2
View File
@@ -19,10 +19,11 @@ describe("util.session", () => {
tokens: { input, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
time: { created: 0 },
})
const messages = [assistant("msg_1", 10), assistant("msg_3", 30)]
const messages = [assistant("msg_z", 10), assistant("msg_a", 30)]
expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30)
expect(lastAssistantWithUsage(messages, "msg_2")?.tokens.input).toBe(10)
expect(lastAssistantWithUsage(messages, "msg_a")?.tokens.input).toBe(10)
expect(lastAssistantWithUsage(messages, "msg_missing")).toBeUndefined()
expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30)
})
})