diff --git a/packages/client/src/promise/generated/types.ts b/packages/client/src/promise/generated/types.ts index 3bb87c1111..07964f950e 100644 --- a/packages/client/src/promise/generated/types.ts +++ b/packages/client/src/promise/generated/types.ts @@ -1439,6 +1439,13 @@ export type SessionLogOutput = readonly sessionID: string readonly assistantMessageID: string readonly error: { readonly type: string; readonly message: string } + readonly cost?: number + readonly tokens?: { + readonly input: number + readonly output: number + readonly reasoning: number + readonly cache: { readonly read: number; readonly write: number } + } } } | { @@ -4542,6 +4549,23 @@ export type EventSubscribeOutput = readonly location?: { readonly directory: string; readonly workspaceID?: string } readonly data: { readonly sessionID: string; readonly title: string } } + | { + readonly id: string + readonly created: number + readonly metadata?: { readonly [x: string]: unknown } + readonly type: "session.usage.updated" + readonly location?: { readonly directory: string; readonly workspaceID?: string } + readonly data: { + readonly sessionID: string + readonly cost: number + readonly tokens: { + readonly input: number + readonly output: number + readonly reasoning: number + readonly cache: { readonly read: number; readonly write: number } + } + } + } | { readonly id: string readonly created: number @@ -4765,6 +4789,13 @@ export type EventSubscribeOutput = readonly sessionID: string readonly assistantMessageID: string readonly error: { readonly type: string; readonly message: string } + readonly cost?: number + readonly tokens?: { + readonly input: number + readonly output: number + readonly reasoning: number + readonly cache: { readonly read: number; readonly write: number } + } } } | { diff --git a/packages/client/test/promise.test.ts b/packages/client/test/promise.test.ts index b8433a0d1d..9b26efe773 100644 --- a/packages/client/test/promise.test.ts +++ b/packages/client/test/promise.test.ts @@ -32,7 +32,7 @@ test("exposes every standard HTTP API group", () => { "vcs", "debug", ]) - expect(Object.keys(client.debug)).toEqual(["location"]) + expect(Object.keys(client.debug)).toEqual(["location", "evictLocation"]) expect(Object.keys(client.message)).toEqual(["list"]) expect(Object.keys(client.integration)).toEqual([ "list", diff --git a/packages/core/src/aisdk.ts b/packages/core/src/aisdk.ts index f7474b244a..aecf7a6482 100644 --- a/packages/core/src/aisdk.ts +++ b/packages/core/src/aisdk.ts @@ -6,6 +6,7 @@ import type { JSONValue, LanguageModelV3, LanguageModelV3CallOptions, + LanguageModelV3FinishReason, LanguageModelV3FunctionTool, LanguageModelV3Message, LanguageModelV3Prompt, @@ -624,8 +625,8 @@ function usage(input: Extract["us return Object.values(output).some((value) => value !== undefined) ? output : undefined } -function finishReason(value: unknown): FinishReason { - return Schema.is(FinishReason)(value) ? value : "unknown" +function finishReason(value: LanguageModelV3FinishReason): FinishReason { + return value.unified === "other" ? "unknown" : value.unified } function providerMetadata(value: unknown) { diff --git a/packages/core/src/session/message-updater.ts b/packages/core/src/session/message-updater.ts index 96881e7c69..dc61f96dd4 100644 --- a/packages/core/src/session/message-updater.ts +++ b/packages/core/src/session/message-updater.ts @@ -143,6 +143,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) { return Effect.gen(function* () { yield* SessionEvent.All.match(event, { + "session.usage.updated": () => Effect.void, "session.agent.selected": (event) => { return adapter.appendMessage( SessionMessage.AgentSelected.make({ @@ -296,6 +297,10 @@ export function update(adapter: Adapter, event: SessionEvent.Event) { draft.finish = "error" draft.error = castDraft(event.data.error) draft.retry = undefined + if (event.data.cost !== undefined && event.data.tokens !== undefined) { + draft.cost = event.data.cost + draft.tokens = castDraft(event.data.tokens) + } }) }, "session.text.started": (event) => { diff --git a/packages/core/src/session/projector.ts b/packages/core/src/session/projector.ts index 625b6f1553..f0a2f83d3f 100644 --- a/packages/core/src/session/projector.ts +++ b/packages/core/src/session/projector.ts @@ -1,7 +1,7 @@ export * as SessionProjector from "./projector" import { and, asc, desc, eq, gt, gte, inArray, lt, or, sql } from "drizzle-orm" -import { DateTime, Effect, Layer, Schema } from "effect" +import { DateTime, Effect, Layer, Schema, Stream } from "effect" import { Database } from "../database/database" import { EventV2 } from "../event" import { makeGlobalNode } from "../effect/app-node" @@ -48,11 +48,6 @@ type Usage = { const ForkBatchSize = 500 -const emptyUsage = (): Usage => ({ - cost: 0, - tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, -}) - const forkTitle = (value: string) => { const match = value.match(/^(.+) \(fork #(\d+)\)$/) if (match) return `${match[1]} (fork #${Number.parseInt(match[2], 10) + 1})` @@ -67,22 +62,6 @@ function usage(part: (typeof SessionV1.Event.PartUpdated.Type)["data"]["part"] | return { cost: value.cost as Usage["cost"], tokens: value.tokens as Usage["tokens"] } } -function addUsage(target: Usage, value: Usage) { - target.cost += value.cost - target.tokens.input += value.tokens.input - target.tokens.output += value.tokens.output - target.tokens.reasoning += value.tokens.reasoning - target.tokens.cache.read += value.tokens.cache.read - target.tokens.cache.write += value.tokens.cache.write -} - -function messageUsage(row: typeof SessionMessageTable.$inferSelect): Usage | undefined { - if (row.type !== "assistant") return undefined - const message = decodeMessage({ ...row.data, id: row.id, type: row.type }) - if (message.type !== "assistant" || message.cost === undefined || message.tokens === undefined) return undefined - return { cost: message.cost, tokens: message.tokens } -} - function sessionRow(info: SessionV1.SessionInfo): typeof SessionTable.$inferInsert { return { id: info.id, @@ -151,6 +130,37 @@ function applyUsage( .pipe(Effect.orDie) } +const publishSessionUsage = Effect.fn("SessionProjector.publishUsage")(function* ( + db: DatabaseService, + events: EventV2.Interface, + sessionID: (typeof SessionEvent.Step.Ended.Type)["data"]["sessionID"], +) { + const row = yield* db + .select({ + cost: SessionTable.cost, + input: SessionTable.tokens_input, + output: SessionTable.tokens_output, + reasoning: SessionTable.tokens_reasoning, + cacheRead: SessionTable.tokens_cache_read, + cacheWrite: SessionTable.tokens_cache_write, + }) + .from(SessionTable) + .where(eq(SessionTable.id, sessionID)) + .get() + .pipe(Effect.orDie) + if (!row) return + yield* events.publish(SessionEvent.UsageUpdated, { + sessionID, + cost: row.cost, + tokens: { + input: row.input, + output: row.output, + reasoning: row.reasoning, + cache: { read: row.cacheRead, write: row.cacheWrite }, + }, + }) +}) + const projectFork = Effect.fn("SessionProjector.projectFork")(function* ( db: DatabaseService, event: typeof SessionEvent.Forked.Type, @@ -187,7 +197,7 @@ const projectFork = Effect.fn("SessionProjector.projectFork")(function* ( .limit(1) .get() .pipe(Effect.orDie) - const copiedSeq = copied?.seq ?? 0 + const copiedSeq = copied?.seq const stored = yield* db .insert(SessionTable) @@ -237,9 +247,8 @@ const projectFork = Effect.fn("SessionProjector.projectFork")(function* ( .pipe(Effect.orDie) } - const usage = emptyUsage() let cursor = -1 - while (true) { + while (copiedSeq !== undefined) { const rows = yield* db .select() .from(SessionMessageTable) @@ -247,7 +256,7 @@ const projectFork = Effect.fn("SessionProjector.projectFork")(function* ( and( eq(SessionMessageTable.session_id, event.data.parentID), gt(SessionMessageTable.seq, cursor), - copiedSeq === 0 ? undefined : lt(SessionMessageTable.seq, copiedSeq + 1), + lt(SessionMessageTable.seq, copiedSeq + 1), sql`${SessionMessageTable.type} != 'compaction' or json_extract(${SessionMessageTable.data}, '$.status') not in ('queued', 'running')`, ), ) @@ -318,27 +327,9 @@ const projectFork = Effect.fn("SessionProjector.projectFork")(function* ( .pipe(Effect.orDie) } - for (const row of rows) { - const value = messageUsage(row) - if (value) addUsage(usage, value) - } cursor = rows.at(-1)!.seq } - - yield* db - .update(SessionTable) - .set({ - cost: usage.cost, - tokens_input: usage.tokens.input, - tokens_output: usage.tokens.output, - tokens_reasoning: usage.tokens.reasoning, - tokens_cache_read: usage.tokens.cache.read, - tokens_cache_write: usage.tokens.cache.write, - }) - .where(eq(SessionTable.id, event.data.sessionID)) - .run() - .pipe(Effect.orDie) - if (copiedSeq > 0) yield* EventV2.reserveSequence(db, event.data.sessionID, copiedSeq) + if (copiedSeq !== undefined) yield* EventV2.reserveSequence(db, event.data.sessionID, copiedSeq) }) function run(db: DatabaseService, event: MessageEvent) { @@ -697,8 +688,19 @@ const layer = Layer.effectDiscard( yield* events.project(SessionEvent.Shell.Started, (event) => run(db, event)) yield* events.project(SessionEvent.Shell.Ended, (event) => run(db, event)) yield* events.project(SessionEvent.Step.Started, (event) => run(db, event)) - yield* events.project(SessionEvent.Step.Ended, (event) => run(db, event)) - yield* events.project(SessionEvent.Step.Failed, (event) => run(db, event)) + yield* events.project(SessionEvent.Step.Ended, (event) => + Effect.gen(function* () { + yield* run(db, event) + yield* applyUsage(db, event.data.sessionID, event.data) + }), + ) + yield* events.project(SessionEvent.Step.Failed, (event) => + Effect.gen(function* () { + yield* run(db, event) + if (event.data.cost !== undefined && event.data.tokens !== undefined) + yield* applyUsage(db, event.data.sessionID, { cost: event.data.cost, tokens: event.data.tokens }) + }), + ) yield* events.project(SessionEvent.Text.Started, (event) => run(db, event)) yield* events.project(SessionEvent.Text.Ended, (event) => run(db, event)) yield* events.project(SessionEvent.Tool.Input.Started, (event) => run(db, event)) @@ -790,6 +792,17 @@ const layer = Layer.effectDiscard( yield* InstructionCheckpoint.reset(db, event.data.sessionID) }), ) + yield* events.subscribe([SessionEvent.Step.Ended, SessionEvent.Step.Failed]).pipe( + Stream.runForEach((event) => { + if ( + event.type === SessionEvent.Step.Failed.type && + (event.data.cost === undefined || event.data.tokens === undefined) + ) + return Effect.void + return publishSessionUsage(db, events, event.data.sessionID) + }), + Effect.forkScoped({ startImmediately: true }), + ) }), ) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 2e05d1edcf..2d5220310c 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -17,6 +17,7 @@ import { Config } from "../../config" import { Database } from "../../database/database" import { EventV2 } from "../../event" import { Location } from "../../location" +import { ModelV2 } from "../../model" import { PermissionV2 } from "../../permission" import { Instructions } from "../../instructions/index" import { InstructionBuiltIns } from "../../instructions/builtins" @@ -50,6 +51,30 @@ import { StepFailedError, UserInterruptedError } from "../error" import { toSessionError } from "../to-session-error" import { SessionRunnerRetry } from "./retry" +type StepTokens = { + readonly input: number + readonly output: number + readonly reasoning: number + readonly cache: { readonly read: number; readonly write: number } +} + +// TODO(#35765): Use Copilot's reported billed amount once billing has a dedicated typed runtime contract. +export function calculateCost(costs: ModelV2.Info["cost"], tokens: StepTokens) { + const context = tokens.input + tokens.cache.read + tokens.cache.write + const tier = costs + .filter((cost) => cost.tier?.type === "context" && context > cost.tier.size) + .toSorted((a, b) => (b.tier?.size ?? 0) - (a.tier?.size ?? 0))[0] + const cost = tier ?? costs.find((cost) => cost.tier === undefined) + if (!cost) return 0 + return ( + (tokens.input * cost.input + + (tokens.output + tokens.reasoning) * cost.output + + tokens.cache.read * cost.cache.read + + tokens.cache.write * cost.cache.write) / + 1_000_000 + ) +} + /** * Runs one durable coding-agent Session until it settles. * @@ -312,6 +337,11 @@ const layer = Layer.effect( Effect.ensuring(serialized(publisher.flush())), ) + const stepUsage = (settlement: NonNullable>) => ({ + cost: calculateCost(resolved.cost, settlement.tokens), + tokens: settlement.tokens, + }) + // Captures the end snapshot, diffs it against the step's start, and durably ends the // assistant step. const publishStepEnd = (settlement: NonNullable>) => @@ -328,8 +358,7 @@ const layer = Layer.effect( sessionID: session.id, assistantMessageID: yield* publisher.startAssistant(), finish: settlement.finish, - cost: 0, - tokens: settlement.tokens, + ...stepUsage(settlement), snapshot: endSnapshot, files, }), @@ -452,7 +481,8 @@ const layer = Layer.effect( const stepEndedCleanly = !streamInterrupted && !toolsInterrupted && infraError === undefined && !providerFailed && !stepFailure if (stepSettlement && stepEndedCleanly) yield* publishStepEnd(stepSettlement) - if (stepFailure) yield* serialized(publisher.publishStepFailure()) + if (stepFailure) + yield* serialized(publisher.publishStepFailure(stepSettlement ? stepUsage(stepSettlement) : undefined)) if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) if (userDeclined) return yield* Effect.interrupt diff --git a/packages/core/src/session/runner/model.ts b/packages/core/src/session/runner/model.ts index 6cc0d0845d..defd95f9c1 100644 --- a/packages/core/src/session/runner/model.ts +++ b/packages/core/src/session/runner/model.ts @@ -82,6 +82,8 @@ export interface Resolved { readonly model: Model /** Selected catalog identity. Durable records and displays must use this, never the API model id. */ readonly ref: ModelV2.Ref + /** Catalog pricing in dollars per million tokens. */ + readonly cost: ModelV2.Info["cost"] } export interface Interface { @@ -94,13 +96,14 @@ export class Service extends Context.Service()("@opencode/v2 export const layerWith = (resolve: Interface["resolve"]) => Layer.succeed(Service, Service.of({ resolve })) /** Builds a Resolved whose catalog identity mirrors the route model. Test or embedding seam. */ -export const resolved = (model: Model, variant?: ModelV2.VariantID): Resolved => ({ +export const resolved = (model: Model, variant?: ModelV2.VariantID, cost: ModelV2.Info["cost"] = []): Resolved => ({ model, ref: ModelV2.Ref.make({ id: ModelV2.ID.make(model.id), providerID: ProviderV2.ID.make(model.provider), ...(variant === undefined ? {} : { variant }), }), + cost, }) const apiKey = (model: ModelV2.Info, credential?: Credential.Value) => { @@ -341,6 +344,7 @@ const layer = Layer.effect( providerID: selected.providerID, ...(session.model?.variant === undefined ? {} : { variant: session.model.variant }), }), + cost: selected.cost, } }), }) diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index 5f4eeebc84..7b4d30df87 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -216,7 +216,10 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input) if (replace || stepFailure === undefined) stepFailure = error }) - const publishStepFailure = Effect.fnUntraced(function* () { + const publishStepFailure = Effect.fnUntraced(function* (usage?: { + readonly cost: number + readonly tokens: ReturnType + }) { if (stepFailed || stepFailure === undefined) return const assistantMessageID = yield* startAssistant() stepFailed = true @@ -224,6 +227,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input) sessionID: input.sessionID, assistantMessageID, error: stepFailure, + ...usage, }) }) @@ -409,12 +413,12 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input) case "step-finish": yield* flush() if (stepSettlement) return yield* Effect.die(new Error("Duplicate step finish")) + stepSettlement = { finish: event.reason, tokens: tokens(event.usage) } if (event.reason === "content-filter") { providerFailed = true yield* failAssistant({ type: "provider.content-filter", message: "Provider blocked the response" }, true) return } - stepSettlement = { finish: event.reason, tokens: tokens(event.usage) } return case "finish": return diff --git a/packages/core/test/aisdk.test.ts b/packages/core/test/aisdk.test.ts index 495c125a1e..4037949427 100644 --- a/packages/core/test/aisdk.test.ts +++ b/packages/core/test/aisdk.test.ts @@ -51,13 +51,11 @@ it.effect("projects request settings, headers, and body overlays", () => apiKey: "secret", thinkingConfig: { thinkingBudget: 1024 }, }) - const resolved = yield* aisdk.model( - { - ...input, - headers: { "x-test": "header" }, - body: { safety_setting: "strict" }, - }, - ) + const resolved = yield* aisdk.model({ + ...input, + headers: { "x-test": "header" }, + body: { safety_setting: "strict" }, + }) const prepared = yield* LLMClient.prepare( LLM.request({ model: resolved, prompt: "Hello" }), ) diff --git a/packages/core/test/session-create.test.ts b/packages/core/test/session-create.test.ts index 0a2e57edf5..86e593c722 100644 --- a/packages/core/test/session-create.test.ts +++ b/packages/core/test/session-create.test.ts @@ -225,9 +225,17 @@ describe("SessionV2.create", () => { promotedSeq: 2, }) - yield* session.prompt({ sessionID: parent.id, prompt: PromptInput.Prompt.make({ text: "Parent changed" }), resume: false }) + yield* session.prompt({ + sessionID: parent.id, + prompt: PromptInput.Prompt.make({ text: "Parent changed" }), + resume: false, + }) yield* SessionInput.promoteSteers(db, events, parent.id) - yield* session.prompt({ sessionID: forked.id, prompt: PromptInput.Prompt.make({ text: "Child continues" }), resume: false }) + yield* session.prompt({ + sessionID: forked.id, + prompt: PromptInput.Prompt.make({ text: "Child continues" }), + resume: false, + }) yield* SessionInput.promoteSteers(db, events, forked.id) expect((yield* session.context(parent.id)).map((message) => message.type)).toEqual(["user", "synthetic", "user"]) @@ -260,8 +268,25 @@ describe("SessionV2.create", () => { resume: false, }) yield* SessionInput.promoteSteers(db, events, parent.id) + const assistantMessageID = SessionMessage.ID.create() + const model = ModelV2.Ref.make({ id: ModelV2.ID.make("model"), providerID: ProviderV2.ID.make("provider") }) + yield* events.publish(SessionEvent.Step.Started, { + sessionID: parent.id, + assistantMessageID, + agent: "build", + model, + }) + yield* events.publish(SessionEvent.Step.Ended, { + sessionID: parent.id, + assistantMessageID, + finish: "stop", + cost: 0.75, + tokens: { input: 6, output: 3, reasoning: 1, cache: { read: 2, write: 1 } }, + }) const forked = yield* session.fork({ sessionID: parent.id, messageID: second.id }) + const beforeFirst = yield* session.fork({ sessionID: parent.id, messageID: first.id }) + const complete = yield* session.fork({ sessionID: parent.id }) const context = yield* session.context(forked.id) const history = Array.from(yield* Stream.runCollect(logEvents(session, forked.id))) @@ -269,6 +294,13 @@ describe("SessionV2.create", () => { expect(context).toMatchObject([{ text: "First" }]) expect(context[0]?.id).not.toBe(first.id) expect(history[0]).toMatchObject({ data: { from: second.id } }) + expect(forked).toMatchObject({ cost: 0, tokens: { input: 0, output: 0, reasoning: 0 } }) + expect(yield* session.context(beforeFirst.id)).toEqual([]) + expect(beforeFirst).toMatchObject({ cost: 0, tokens: { input: 0, output: 0, reasoning: 0 } }) + expect(complete).toMatchObject({ + cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + }) }), ) @@ -375,7 +407,11 @@ describe("SessionV2.create", () => { const events = yield* EventV2.Service const { db } = yield* Database.Service const created = yield* session.create({ location }) - yield* session.prompt({ sessionID: created.id, prompt: PromptInput.Prompt.make({ text: "Hello" }), resume: false }) + yield* session.prompt({ + sessionID: created.id, + prompt: PromptInput.Prompt.make({ text: "Hello" }), + resume: false, + }) yield* SessionInput.promoteSteers(db, events, created.id) expect( diff --git a/packages/core/test/session-projector.test.ts b/packages/core/test/session-projector.test.ts index ce30170362..cc0655c55f 100644 --- a/packages/core/test/session-projector.test.ts +++ b/packages/core/test/session-projector.test.ts @@ -1,5 +1,5 @@ import { describe, expect } from "bun:test" -import { DateTime, Effect, Schema } from "effect" +import { DateTime, Effect, Fiber, Option, Schema, Stream } from "effect" import { asc, eq } from "drizzle-orm" import { Database } from "@opencode-ai/core/database/database" import { LayerNode } from "@opencode-ai/core/effect/layer-node" @@ -41,12 +41,15 @@ const assistantRow = ( id: SessionMessage.ID, seq: number, time: { created: DateTime.Utc; completed?: DateTime.Utc } = { created }, + usage?: Pick, ) => { const { id: _, type, ...data - } = encodeMessage(SessionMessage.Assistant.make({ id, type: "assistant", agent: "build", model, content: [], time })) + } = encodeMessage( + SessionMessage.Assistant.make({ id, type: "assistant", agent: "build", model, content: [], time, ...usage }), + ) return { id, session_id: sessionID, type, seq, time_created: DateTime.toEpochMillis(time.created), data } } @@ -67,6 +70,12 @@ describe("SessionProjector", () => { directory: "/project", title: "test", version: "test", + cost: 1.25, + tokens_input: 10, + tokens_output: 4, + tokens_reasoning: 2, + tokens_cache_read: 3, + tokens_cache_write: 1, }) .run() const boundary = SessionMessage.ID.make("msg_boundary") @@ -75,8 +84,24 @@ describe("SessionProjector", () => { .insert(SessionMessageTable) .values([ assistantRow(earlier, 0), - assistantRow(boundary, 1), - assistantRow(SessionMessage.ID.make("msg_later"), 2), + assistantRow( + boundary, + 1, + { created }, + { + cost: 0.5, + tokens: { input: 4, output: 1, reasoning: 1, cache: { read: 1, write: 0 } }, + }, + ), + assistantRow( + SessionMessage.ID.make("msg_later"), + 2, + { created }, + { + cost: 0.75, + tokens: { input: 6, output: 3, reasoning: 1, cache: { read: 2, write: 1 } }, + }, + ), ]) .run() yield* db @@ -106,6 +131,14 @@ describe("SessionProjector", () => { expect( (yield* db.select({ id: SessionMessageTable.id }).from(SessionMessageTable).all()).map((row) => row.id), ).toEqual([earlier]) + expect(yield* db.select().from(SessionTable).where(eq(SessionTable.id, sessionID)).get()).toMatchObject({ + cost: 1.25, + tokens_input: 10, + tokens_output: 4, + tokens_reasoning: 2, + tokens_cache_read: 3, + tokens_cache_write: 1, + }) // A committed revert resets the context checkpoint so the next turn re-initializes. expect(yield* db.select().from(InstructionCheckpointTable).get().pipe(Effect.orDie)).toBeUndefined() }), @@ -534,12 +567,15 @@ describe("SessionProjector", () => { .pipe(Effect.orDie) const service = yield* EventV2.Service + const usageUpdated = yield* service + .subscribe(SessionEvent.UsageUpdated) + .pipe(Stream.runHead, Effect.forkScoped({ startImmediately: true })) yield* service.publish(SessionEvent.Step.Ended, { sessionID, assistantMessageID: SessionMessage.ID.make("msg_assistant_2"), finish: "stop", - cost: 0, - tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + cost: 1.25, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, }) const rows = yield* db @@ -556,8 +592,25 @@ describe("SessionProjector", () => { expect(messages[1]).toMatchObject({ type: "assistant", finish: "stop", + cost: 1.25, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, time: { completed: DateTime.makeUnsafe(0) }, }) + expect( + yield* db.select().from(SessionTable).where(eq(SessionTable.id, sessionID)).get().pipe(Effect.orDie), + ).toMatchObject({ + cost: 1.25, + tokens_input: 10, + tokens_output: 4, + tokens_reasoning: 2, + tokens_cache_read: 3, + tokens_cache_write: 1, + }) + expect(Option.getOrThrow(yield* Fiber.join(usageUpdated)).data).toEqual({ + sessionID, + cost: 1.25, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, + }) }), ) diff --git a/packages/core/test/session-runner-tool-events.test.ts b/packages/core/test/session-runner-tool-events.test.ts index 4fdbb49e8b..0e54513e8f 100644 --- a/packages/core/test/session-runner-tool-events.test.ts +++ b/packages/core/test/session-runner-tool-events.test.ts @@ -151,15 +151,39 @@ test("step finish records settlement without publishing step ended", async () => test("content-filter finish retains failure evidence until step closeout", async () => { const { published, publisher } = capture() await Effect.runPromise(publisher.publish(LLMEvent.stepStart({ index: 0 }))) - await Effect.runPromise(publisher.publish(LLMEvent.stepFinish({ index: 0, reason: "content-filter" }))) + await Effect.runPromise( + publisher.publish( + LLMEvent.stepFinish({ + index: 0, + reason: "content-filter", + usage: { + nonCachedInputTokens: 8, + outputTokens: 3, + reasoningTokens: 1, + }, + }), + ), + ) expect(published.map((event) => event.type)).toEqual(["session.step.started.1"]) - await Effect.runPromise(publisher.publishStepFailure()) + const settlement = publisher.stepSettlement() + expect(settlement).toMatchObject({ + finish: "content-filter", + tokens: { input: 8, output: 2, reasoning: 1 }, + }) + if (!settlement) throw new Error("Expected content-filter settlement") + await Effect.runPromise( + publisher.publishStepFailure({ + cost: 1.25, + tokens: settlement.tokens, + }), + ) expect(published.map((event) => event.type)).toEqual(["session.step.started.1", "session.step.failed.1"]) expect(published.at(-1)?.data).toMatchObject({ error: { type: "provider.content-filter", message: "Provider blocked the response" }, + cost: 1.25, + tokens: { input: 8, output: 2, reasoning: 1 }, }) - expect(publisher.stepSettlement()).toBeUndefined() }) test("content-filter finish preserves partial streamed text and never ends the step successfully", async () => { diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index ec893a5527..b98c5ef106 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -1,4 +1,4 @@ -import { describe, expect } from "bun:test" +import { describe, expect, test } from "bun:test" import { LLMClient, LLMError, @@ -116,6 +116,28 @@ const recoveryModel = Model.make({ provider: "fake", route: OpenAIChat.route.with({ limits: { context: 20_000, output: 1_000 } }), }) + +test("calculates step cost using the matching context tier", () => { + expect( + SessionRunnerLLM.calculateCost( + [ + { input: 1, output: 2, cache: { read: 0.1, write: 0.5 } }, + { tier: { type: "context", size: 100 }, input: 3, output: 4, cache: { read: 0.2, write: 0.6 } }, + ], + { input: 80, output: 10, reasoning: 2, cache: { read: 20, write: 1 } }, + ), + ).toBeCloseTo(0.0002926) +}) + +test("does not apply an ineligible tier without base pricing", () => { + expect( + SessionRunnerLLM.calculateCost( + [{ tier: { type: "context", size: 100 }, input: 3, output: 4, cache: { read: 0.2, write: 0.6 } }], + { input: 80, output: 10, reasoning: 2, cache: { read: 20, write: 0 } }, + ), + ).toBe(0) +}) + const authorizations: Tool.Context[] = [] const executions: string[] = [] const permission = Layer.succeed( @@ -1704,6 +1726,7 @@ describe("SessionRunnerLLM", () => { { type: "assistant", finish: "tool-calls", + cost: 0, tokens: { input: 8, output: 3, reasoning: 1, cache: { read: 2, write: 0 } }, content: [ { type: "reasoning", text: "Think" }, @@ -3635,7 +3658,11 @@ describe("SessionRunnerLLM", () => { LLMEvent.stepStart({ index: 0 }), LLMEvent.textStart({ id: "partial" }), LLMEvent.textDelta({ id: "partial", text: "Partial" }), - LLMEvent.stepFinish({ index: 0, reason: "content-filter" }), + LLMEvent.stepFinish({ + index: 0, + reason: "content-filter", + usage: { nonCachedInputTokens: 8, outputTokens: 3, reasoningTokens: 1 }, + }), LLMEvent.finish({ reason: "content-filter" }), ] @@ -3646,9 +3673,15 @@ describe("SessionRunnerLLM", () => { type: "assistant", finish: "error", error: { type: "provider.content-filter" }, + cost: 0, + tokens: { input: 8, output: 2, reasoning: 1, cache: { read: 0, write: 0 } }, content: [{ type: "text", text: "Partial" }], }, ]) + expect(yield* session.get(sessionID)).toMatchObject({ + cost: 0, + tokens: { input: 8, output: 2, reasoning: 1, cache: { read: 0, write: 0 } }, + }) expect(yield* recordedEventTypes(sessionID)).not.toContain("session.step.ended.1") }), ) diff --git a/packages/schema/src/session-event.ts b/packages/schema/src/session-event.ts index 8e96321ce2..9f98b35226 100644 --- a/packages/schema/src/session-event.ts +++ b/packages/schema/src/session-event.ts @@ -30,6 +30,15 @@ export interface Source extends Schema.Schema.Type {} const Base = { sessionID: SessionID, } +const Tokens = Schema.Struct({ + input: Schema.Finite, + output: Schema.Finite, + reasoning: Schema.Finite, + cache: Schema.Struct({ + read: Schema.Finite, + write: Schema.Finite, + }), +}) const PromptFields = { ...Base, inputID: SessionMessage.ID, @@ -84,6 +93,16 @@ export const Renamed = Event.durable({ }) export type Renamed = typeof Renamed.Type +export const UsageUpdated = Event.ephemeral({ + type: "session.usage.updated", + schema: { + ...Base, + cost: Schema.Finite, + tokens: Tokens, + }, +}) +export type UsageUpdated = typeof UsageUpdated.Type + export const Deleted = Event.durable({ type: "session.deleted", durable: { @@ -224,15 +243,7 @@ export namespace Step { assistantMessageID: SessionMessage.ID, finish: FinishReason, cost: Schema.Finite, - tokens: Schema.Struct({ - input: Schema.Finite, - output: Schema.Finite, - reasoning: Schema.Finite, - cache: Schema.Struct({ - read: Schema.Finite, - write: Schema.Finite, - }), - }), + tokens: Tokens, snapshot: Schema.String.pipe(optional), files: Schema.Array(RelativePath).pipe(optional), }, @@ -246,6 +257,8 @@ export namespace Step { ...Base, assistantMessageID: SessionMessage.ID, error: SessionError.Error, + cost: Schema.Finite.pipe(optional), + tokens: Tokens.pipe(optional), }, }) export type Failed = typeof Failed.Type @@ -504,6 +517,7 @@ export const Definitions = Event.inventory( ModelSelected, Moved, Renamed, + UsageUpdated, Deleted, Forked, PromptPromoted, diff --git a/packages/schema/test/event-manifest.test.ts b/packages/schema/test/event-manifest.test.ts index 9ff3a05519..8f8f112176 100644 --- a/packages/schema/test/event-manifest.test.ts +++ b/packages/schema/test/event-manifest.test.ts @@ -144,6 +144,8 @@ describe("public event manifest", () => { expect(SessionEvent.DurableDefinitions).toEqual( SessionEvent.Definitions.filter((definition) => definition.durability === "durable"), ) + expect(SessionEvent.UsageUpdated.durability).toBe("ephemeral") + expect(EventManifest.ServerDefinitions).toContain(SessionEvent.UsageUpdated) expect(EventManifest.Definitions.every((definition) => definition.durability !== undefined)).toBe(true) }) diff --git a/packages/sdk/js/src/v2/gen/sdk.gen.ts b/packages/sdk/js/src/v2/gen/sdk.gen.ts index 458f9c2d2e..239783dc87 100644 --- a/packages/sdk/js/src/v2/gen/sdk.gen.ts +++ b/packages/sdk/js/src/v2/gen/sdk.gen.ts @@ -277,6 +277,8 @@ import type { V2CredentialUpdateErrors, V2CredentialUpdateResponses, V2DebugLocationErrors, + V2DebugLocationEvictErrors, + V2DebugLocationEvictResponses, V2DebugLocationResponses, V2EventSubscribeErrors, V2EventSubscribeResponses, @@ -8117,6 +8119,34 @@ export class Vcs2 extends HeyApiClient { } } +export class Location2 extends HeyApiClient { + /** + * Evict a loaded location + * + * Dispose the requested location's cached services so its next use boots them fresh. + */ + public evict( + parameters?: { + location?: { + directory?: string | null + workspace?: string | null + } | null + }, + options?: Options, + ) { + const params = buildClientParams([parameters], [{ args: [{ in: "query", key: "location" }] }]) + return (options?.client ?? this.client).delete< + V2DebugLocationEvictResponses, + V2DebugLocationEvictErrors, + ThrowOnError + >({ + url: "/api/debug/location", + ...options, + ...params, + }) + } +} + export class Debug extends HeyApiClient { /** * List loaded locations @@ -8129,6 +8159,11 @@ export class Debug extends HeyApiClient { ...options, }) } + + private _location?: Location2 + get location2(): Location2 { + return (this._location ??= new Location2({ client: this.client })) + } } export class V2 extends HeyApiClient { diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index 87aa5ddbe1..68ad7c2a0e 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -21,6 +21,7 @@ export type Event = | EventSessionModelSelected | EventSessionMoved | EventSessionRenamed + | EventSessionUsageUpdated | EventSessionForked | EventSessionPromptPromoted | EventSessionPromptAdmitted @@ -895,6 +896,23 @@ export type GlobalEvent = { title: string } } + | { + id: string + type: "session.usage.updated" + properties: { + sessionID: string + cost: number + tokens: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } + } + } | { id: string type: "session.forked" @@ -1042,6 +1060,16 @@ export type GlobalEvent = { sessionID: string assistantMessageID: string error: SessionStructuredError + cost?: number + tokens?: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } } } | { @@ -3079,6 +3107,7 @@ export type V2Event = | SessionModelSelected | SessionMoved | SessionRenamed + | SessionUsageUpdated | SessionForked | SessionPromptPromoted | SessionPromptAdmitted @@ -3982,6 +4011,16 @@ export type SyncEventSessionStepFailed = { sessionID: string assistantMessageID: string error: SessionStructuredError + cost?: number + tokens?: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } } } } @@ -5088,6 +5127,16 @@ export type SessionStepFailed = { sessionID: string assistantMessageID: string error: SessionStructuredError + cost?: number + tokens?: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } } } @@ -5955,6 +6004,29 @@ export type MessagePartRemoved = { } } +export type SessionUsageUpdated = { + id: string + created: number + metadata?: { + [key: string]: unknown + } + type: "session.usage.updated" + location?: LocationRef + data: { + sessionID: string + cost: number + tokens: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } + } +} + export type SessionTextDelta = { id: string created: number @@ -7010,6 +7082,24 @@ export type EventSessionRenamed = { } } +export type EventSessionUsageUpdated = { + id: string + type: "session.usage.updated" + properties: { + sessionID: string + cost: number + tokens: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } + } +} + export type EventSessionForked = { id: string type: "session.forked" @@ -7171,6 +7261,16 @@ export type EventSessionStepFailed = { sessionID: string assistantMessageID: string error: SessionStructuredError + cost?: number + tokens?: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } } } @@ -8525,6 +8625,7 @@ export type V2EventV2 = | SessionModelSelectedV2 | SessionMovedV2 | SessionRenamedV2 + | SessionUsageUpdatedV2 | SessionDeletedV2 | SessionForkedV2 | SessionPromptPromotedV2 @@ -9272,6 +9373,16 @@ export type SessionStepFailedV2 = { sessionID: string assistantMessageID: string error: SessionStructuredError + cost?: number + tokens?: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } } } @@ -9996,6 +10107,29 @@ export type MessagePartRemovedV2 = { } } +export type SessionUsageUpdatedV2 = { + id: string + created: number + metadata?: { + [key: string]: unknown + } + type: "session.usage.updated" + location?: LocationRefV2 + data: { + sessionID: string + cost: number + tokens: { + input: number + output: number + reasoning: number + cache: { + read: number + write: number + } + } + } +} + export type SessionTextDeltaV2 = { id: string created: number @@ -18549,6 +18683,40 @@ export type V2VcsDiffResponses = { export type V2VcsDiffResponse = V2VcsDiffResponses[keyof V2VcsDiffResponses] +export type V2DebugLocationEvictData = { + body?: never + path?: never + query?: { + location?: { + directory?: string | null + workspace?: string | null + } | null + } + url: "/api/debug/location" +} + +export type V2DebugLocationEvictErrors = { + /** + * InvalidRequestError + */ + 400: InvalidRequestErrorV2 + /** + * UnauthorizedError + */ + 401: UnauthorizedError +} + +export type V2DebugLocationEvictError = V2DebugLocationEvictErrors[keyof V2DebugLocationEvictErrors] + +export type V2DebugLocationEvictResponses = { + /** + * + */ + 204: void +} + +export type V2DebugLocationEvictResponse = V2DebugLocationEvictResponses[keyof V2DebugLocationEvictResponses] + export type V2DebugLocationData = { body?: never path?: never diff --git a/packages/tui/src/context/data.tsx b/packages/tui/src/context/data.tsx index 4136b5dc36..04a81cb0b4 100644 --- a/packages/tui/src/context/data.tsx +++ b/packages/tui/src/context/data.tsx @@ -110,6 +110,9 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({ directory: process.cwd(), }) const messageIndex = new Map>() + const sessionRefreshGeneration = new Map() + const sessionRefreshApplied = new Map() + const sessionUsage = new Map() let connectionGeneration = 0 let statusChanges: Set | undefined let bootstrapping: Promise | 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) => 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) diff --git a/packages/tui/src/feature-plugins/sidebar/context.tsx b/packages/tui/src/feature-plugins/sidebar/context.tsx index f1c99d9679..ae41cec303 100644 --- a/packages/tui/src/feature-plugins/sidebar/context.tsx +++ b/packages/tui/src/feature-plugins/sidebar/context.tsx @@ -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, diff --git a/packages/tui/src/routes/session/subagent-footer.tsx b/packages/tui/src/routes/session/subagent-footer.tsx index 71b3a161e5..f02ef3f20e 100644 --- a/packages/tui/src/routes/session/subagent-footer.tsx +++ b/packages/tui/src/routes/session/subagent-footer.tsx @@ -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() { 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} > Parent {parentShortcut()} diff --git a/packages/tui/src/util/session.ts b/packages/tui/src/util/session.ts index 94ccad22d0..3a49ea1a9c 100644 --- a/packages/tui/src/util/session.ts +++ b/packages/tui/src/util/session.ts @@ -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, 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 } => + 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 ef2aa57a9b..3d87dce647 100644 --- a/packages/tui/test/cli/tui/data.test.tsx +++ b/packages/tui/test/cli/tui/data.test.tsx @@ -108,6 +108,309 @@ test("refreshes resources into reactive getters", async () => { } }) +test("applies absolute usage events without losing full session updates", 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.usage.updated", + data: { + sessionID, + cost: 0.5, + tokens: { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } }, + }, + }) + const initialRefresh = data.session.refresh(sessionID) + 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 }, + }, + ], + 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 initialRefresh + 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 }, + }) + + const fullRefresh = data.session.refresh(sessionID) + emitEvent(events, { + id: "evt_usage_3", + created: 3, + type: "session.usage.updated", + data: { + sessionID, + cost: 1, + tokens: { input: 10, output: 4, reasoning: 1, cache: { read: 1, write: 1 } }, + }, + }) + 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 fullRefresh + await Bun.sleep(20) + expect(data.session.get(sessionID)?.cost).toBe(1) + expect(data.session.get(sessionID)?.title).toBe("Older usage") + + emitEvent(events, { + id: "evt_usage_6", + created: 6, + type: "session.usage.updated", + data: { + sessionID, + cost: 1.25, + tokens: { input: 12, output: 5, reasoning: 1, cache: { read: 1, write: 1 } }, + }, + }) + emitEvent(events, { + id: "evt_usage_7", + created: 7, + type: "session.usage.updated", + data: { + sessionID, + cost: 1.25, + tokens: { input: 12, output: 5, reasoning: 1, cache: { read: 1, write: 1 } }, + }, + }) + await wait(() => data.session.get(sessionID)?.cost === 1.25) + expect(data.session.get(sessionID)?.title).toBe("Older usage") + + emitEvent(events, { + id: "evt_usage_8", + created: 8, + type: "session.usage.updated", + data: { + sessionID, + cost: 1.5, + tokens: { input: 14, output: 6, reasoning: 1, cache: { read: 1, write: 1 } }, + }, + }) + emitEvent(events, { + id: "evt_usage_deleted", + created: 9, + type: "session.deleted", + durable: durable(sessionID, 9), + data: { sessionID }, + }) + await Bun.sleep(20) + expect(data.session.get(sessionID)).toBeUndefined() + } finally { + app.renderer.destroy() + } +}) + +test("truncates committed revert messages without changing lifetime usage", async () => { + const events = createEventStream() + const sessionID = "ses_revert_usage" + let cost = 0 + let tokens = { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } + const calls = createFetch((url) => { + if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: [], cursor: {} }) + if (url.pathname !== `/api/session/${sessionID}`) return + return json({ + data: { + id: sessionID, + projectID: "proj_test", + cost, + tokens, + time: { created: 0, updated: 0 }, + title: "Revert usage", + location: { directory }, + }, + }) + }, events) + let data!: ReturnType + + function Probe() { + data = useData() + return + } + + const app = await testRender(() => ( + + + + + + + + + + )) + + try { + await data.session.refresh(sessionID) + emitEvent(events, { + id: "evt_revert_boundary_started", + created: 1, + type: "session.step.started", + durable: durable(sessionID, 1), + data: { + sessionID, + assistantMessageID: "msg_revert_boundary", + agent: "build", + model: { providerID: "provider", id: "model" }, + }, + }) + cost = 0.5 + tokens = { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } } + emitEvent(events, { + id: "evt_revert_boundary_ended", + created: 2, + type: "session.step.ended", + durable: durable(sessionID, 2), + data: { + sessionID, + assistantMessageID: "msg_revert_boundary", + finish: "stop", + cost: 0.5, + tokens, + }, + }) + emitEvent(events, { + id: "evt_revert_boundary_usage", + created: 2, + type: "session.usage.updated", + data: { sessionID, cost, tokens }, + }) + await wait(() => data.session.get(sessionID)?.cost === 0.5) + + emitEvent(events, { + id: "evt_revert_later_started", + created: 3, + type: "session.step.started", + durable: durable(sessionID, 3), + data: { + sessionID, + assistantMessageID: "msg_revert_later", + agent: "build", + model: { providerID: "provider", id: "model" }, + }, + }) + cost = 0.75 + tokens = { input: 8, output: 3, reasoning: 1, cache: { read: 1, write: 1 } } + emitEvent(events, { + id: "evt_revert_later_ended", + created: 4, + type: "session.step.ended", + durable: durable(sessionID, 4), + data: { + sessionID, + assistantMessageID: "msg_revert_later", + finish: "stop", + cost: 0.25, + tokens: { input: 3, output: 1, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + }) + emitEvent(events, { + id: "evt_revert_later_usage", + created: 4, + type: "session.usage.updated", + data: { sessionID, cost, tokens }, + }) + await wait(() => data.session.get(sessionID)?.cost === 0.75) + emitEvent(events, { + id: "evt_revert_staged", + created: 5, + type: "session.revert.staged", + durable: durable(sessionID, 5), + data: { sessionID, revert: { messageID: "msg_revert_later" } }, + }) + await wait(() => data.session.get(sessionID)?.revert?.messageID === "msg_revert_later") + + emitEvent(events, { + id: "evt_revert_committed", + created: 6, + type: "session.revert.committed", + durable: durable(sessionID, 6), + data: { sessionID, to: "msg_revert_later" }, + }) + await wait(() => data.session.message.ids(sessionID).length === 1) + expect(data.session.get(sessionID)?.cost).toBe(0.75) + expect(data.session.message.ids(sessionID)).toEqual(["msg_revert_boundary"]) + expect(data.session.get(sessionID)?.revert).toBeUndefined() + expect(data.session.get(sessionID)?.tokens).toEqual(tokens) + } finally { + app.renderer.destroy() + } +}) + test("updates session location when moved", async () => { const events = createEventStream() const destination = "/tmp/opencode-moved" @@ -517,8 +820,35 @@ 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" } } }) + if (url.pathname === "/api/session/session-live") + return json({ + data: { + id: "session-live", + projectID: "proj_test", + 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 }, + }, + }) + if (url.pathname === "/api/session/session-failed") + return json({ + data: { + id: "session-failed", + projectID: "proj_test", + cost: 0.25, + tokens: { input: 5, output: 1, reasoning: 1, cache: { read: 1, write: 0 } }, + time: { created: 0, updated: 0 }, + title: "Failed session", + location: { directory }, + }, + }) }, events) let data!: ReturnType let rows!: SessionRow[] @@ -546,7 +876,9 @@ test("tracks session status from active sessions and execution events", async () try { await wait(() => data.session.status("session-active") === "running") expect(data.session.status("session-idle")).toBe("idle") + await data.session.refresh("session-live") + settled = true emitEvent(events, { id: "evt_execution_started", created: 0, @@ -577,15 +909,30 @@ test("tracks session status from active sessions and execution events", async () sessionID: "session-live", assistantMessageID: "message-live", finish: "stop", - cost: 0, - tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + cost: 0.75, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, + }, + }) + emitEvent(events, { + id: "evt_step_usage", + created: 0, + type: "session.usage.updated", + data: { + sessionID: "session-live", + cost: 0.75, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, }, }) await wait(() => { 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, + tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }, + }) emitEvent(events, { id: "evt_execution_succeeded", @@ -596,6 +943,7 @@ test("tracks session status from active sessions and execution events", async () }) await wait(() => data.session.status("session-live") === "idle") + await data.session.refresh("session-failed") emitEvent(events, { id: "evt_failed_execution_started", created: 0, @@ -626,6 +974,18 @@ test("tracks session status from active sessions and execution events", async () sessionID: "session-failed", assistantMessageID: "message-failed", error: { type: "provider.content-filter", message: "Provider blocked the response" }, + cost: 0.25, + tokens: { input: 5, output: 1, reasoning: 1, cache: { read: 1, write: 0 } }, + }, + }) + emitEvent(events, { + id: "evt_failed_step_usage", + created: 0, + type: "session.usage.updated", + data: { + sessionID: "session-failed", + cost: 0.25, + tokens: { input: 5, output: 1, reasoning: 1, cache: { read: 1, write: 0 } }, }, }) await wait(() => { @@ -636,6 +996,13 @@ test("tracks session status from active sessions and execution events", async () assistant.error?.type === "provider.content-filter" ) }) + await wait(() => data.session.get("session-failed")?.cost === 0.25) + expect(data.session.get("session-failed")?.tokens).toEqual({ + input: 5, + output: 1, + reasoning: 1, + cache: { read: 1, write: 0 }, + }) expect(data.session.status("session-failed")).toBe("running") emitEvent(events, { diff --git a/packages/tui/test/util/session.test.ts b/packages/tui/test/util/session.test.ts index faf889781f..752f490247 100644 --- a/packages/tui/test/util/session.test.ts +++ b/packages/tui/test/util/session.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test" -import { isDefaultTitle } from "../../src/util/session" +import type { SessionMessage } from "@opencode-ai/sdk/v2" +import { isDefaultTitle, lastAssistantWithUsage } from "../../src/util/session" describe("util.session", () => { test("recognizes generated parent and child titles", () => { @@ -7,4 +8,22 @@ describe("util.session", () => { expect(isDefaultTitle("Child session - 2026-06-06T12:34:56.789Z")).toBeTrue() expect(isDefaultTitle("New session - custom")).toBeFalse() }) + + test("tracks usage across undo and redo boundaries", () => { + const assistant = (id: string, input: number): SessionMessage => ({ + id, + type: "assistant", + agent: "build", + model: { id: "model", providerID: "provider" }, + content: [], + tokens: { input, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + time: { created: 0 }, + }) + const messages = [assistant("msg_z", 10), assistant("msg_a", 30)] + + expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30) + expect(lastAssistantWithUsage(messages, "msg_a")?.tokens.input).toBe(10) + expect(lastAssistantWithUsage(messages, "msg_missing")).toBeUndefined() + expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30) + }) })