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

@ -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<LanguageModelV3CallOptions>(
LLM.request({ model: resolved, prompt: "Hello" }),
)

View file

@ -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(

View file

@ -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<SessionMessage.Assistant, "cost" | "tokens">,
) => {
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 } },
})
}),
)

View file

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

View file

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