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

@ -6,6 +6,7 @@ import type {
JSONValue,
LanguageModelV3,
LanguageModelV3CallOptions,
LanguageModelV3FinishReason,
LanguageModelV3FunctionTool,
LanguageModelV3Message,
LanguageModelV3Prompt,
@ -624,8 +625,8 @@ function usage(input: Extract<LanguageModelV3StreamPart, { type: "finish" }>["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) {

View file

@ -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) => {

View file

@ -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 }),
)
}),
)

View file

@ -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<ReturnType<typeof publisher.stepSettlement>>) => ({
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<ReturnType<typeof publisher.stepSettlement>>) =>
@ -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

View file

@ -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<Service, Interface>()("@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,
}
}),
})

View file

@ -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<typeof tokens>
}) {
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