fix: update v2 session usage metrics (#35468)
This commit is contained in:
parent
81f6e06681
commit
910e37f6d8
23 changed files with 1013 additions and 107 deletions
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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 }),
|
||||
)
|
||||
}),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}),
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue