diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 75eb368e33..74b995be19 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -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, diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index 9a334a0a6d..1e5dea2795 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -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 } | undefined + let stepSettlement: + | { + readonly finish: string + readonly tokens: ReturnType + 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 diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 80ad249f3a..1b75ac0a98 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -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" }, diff --git a/packages/tui/src/context/data.tsx b/packages/tui/src/context/data.tsx index e98b8c5733..6ac68371e4 100644 --- a/packages/tui/src/context/data.tsx +++ b/packages/tui/src/context/data.tsx @@ -100,6 +100,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({ directory: process.cwd(), }) const messageIndex = new Map>() + const sessionRefreshGeneration = new Map() let connectionGeneration = 0 let statusChanges: Set | undefined let bootstrapping: Promise | 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) diff --git a/packages/tui/src/util/session.ts b/packages/tui/src/util/session.ts index 14d5a78f20..2b6d883cca 100644 --- a/packages/tui/src/util/session.ts +++ b/packages/tui/src/util/session.ts @@ -5,8 +5,13 @@ export function isDefaultTitle(title: string) { } export function lastAssistantWithUsage(messages: ReadonlyArray, 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 } => - message.type === "assistant" && message.tokens !== undefined && (!boundary || message.id < boundary), + ( + message, + index, + ): message is SessionMessageAssistant & { tokens: NonNullable } => + message.type === "assistant" && message.tokens !== undefined && (boundaryIndex === -1 || index < boundaryIndex), ) } diff --git a/packages/tui/test/cli/tui/data.test.tsx b/packages/tui/test/cli/tui/data.test.tsx index 872e386bf2..84724173da 100644 --- a/packages/tui/test/cli/tui/data.test.tsx +++ b/packages/tui/test/cli/tui/data.test.tsx @@ -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((resolve) => { + resolveSessions = resolve + }) + } + if (url.pathname === `/api/session/${sessionID}`) { + return new Promise((resolve) => { + resolveSession.push(resolve) + }) + } + }, events) + let data!: ReturnType + + function Probe() { + data = useData() + return + } + + const app = await testRender(() => ( + + + + + + + + + + )) + + 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, diff --git a/packages/tui/test/util/session.test.ts b/packages/tui/test/util/session.test.ts index 98562623ae..752f490247 100644 --- a/packages/tui/test/util/session.test.ts +++ b/packages/tui/test/util/session.test.ts @@ -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) }) })