fix: address session usage review
This commit is contained in:
@@ -8,6 +8,7 @@ import {
|
||||
Message,
|
||||
SystemPart,
|
||||
isContextOverflowFailure,
|
||||
type ProviderMetadata,
|
||||
type ProviderErrorEvent,
|
||||
} from "@opencode-ai/llm"
|
||||
import { Cause, DateTime, Effect, Exit, FiberSet, Layer, Option, Semaphore, Stream } from "effect"
|
||||
@@ -52,7 +53,10 @@ type StepTokens = {
|
||||
readonly cache: { readonly read: number; readonly write: number }
|
||||
}
|
||||
|
||||
export function calculateCost(costs: ModelV2.Info["cost"], tokens: StepTokens) {
|
||||
export function calculateCost(costs: ModelV2.Info["cost"], tokens: StepTokens, metadata?: ProviderMetadata) {
|
||||
const totalNanoAiu = metadata?.copilot?.totalNanoAiu
|
||||
if (typeof totalNanoAiu === "number" && Number.isFinite(totalNanoAiu) && totalNanoAiu >= 0)
|
||||
return totalNanoAiu / 100_000_000_000
|
||||
const context = tokens.input + tokens.cache.read + tokens.cache.write
|
||||
const tier = costs
|
||||
.filter((cost) => cost.tier?.type === "context" && context > cost.tier.size)
|
||||
@@ -322,7 +326,7 @@ const layer = Layer.effect(
|
||||
sessionID: session.id,
|
||||
assistantMessageID: yield* publisher.startAssistant(),
|
||||
finish: settlement.finish,
|
||||
cost: calculateCost(resolved.cost, settlement.tokens),
|
||||
cost: calculateCost(resolved.cost, settlement.tokens, settlement.providerMetadata),
|
||||
tokens: settlement.tokens,
|
||||
snapshot: endSnapshot,
|
||||
files,
|
||||
|
||||
@@ -69,7 +69,13 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
||||
let assistantActive = false
|
||||
let assistantFailed = false
|
||||
let providerFailed = false
|
||||
let stepSettlement: { readonly finish: string; readonly tokens: ReturnType<typeof tokens> } | undefined
|
||||
let stepSettlement:
|
||||
| {
|
||||
readonly finish: string
|
||||
readonly tokens: ReturnType<typeof tokens>
|
||||
readonly providerMetadata?: ProviderMetadata
|
||||
}
|
||||
| undefined
|
||||
|
||||
const startAssistant = Effect.fnUntraced(function* () {
|
||||
if (assistantMessageID !== undefined) return assistantMessageID
|
||||
@@ -381,7 +387,11 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
||||
yield* flush()
|
||||
assistantActive = false
|
||||
if (stepSettlement) return yield* Effect.die(new Error("Duplicate step finish"))
|
||||
stepSettlement = { finish: event.reason, tokens: tokens(event.usage) }
|
||||
stepSettlement = {
|
||||
finish: event.reason,
|
||||
tokens: tokens(event.usage),
|
||||
providerMetadata: event.usage?.providerMetadata ?? event.providerMetadata,
|
||||
}
|
||||
return
|
||||
case "finish":
|
||||
return
|
||||
|
||||
@@ -128,6 +128,16 @@ test("calculates step cost using the matching context tier", () => {
|
||||
).toBeCloseTo(0.0002926)
|
||||
})
|
||||
|
||||
test("prefers Copilot billed cost over token estimation", () => {
|
||||
expect(
|
||||
SessionRunnerLLM.calculateCost(
|
||||
[{ input: 1, output: 2, cache: { read: 0.1, write: 0.5 } }],
|
||||
{ input: 80, output: 10, reasoning: 2, cache: { read: 20, write: 1 } },
|
||||
{ copilot: { totalNanoAiu: 4_473_525_000 } },
|
||||
),
|
||||
).toBeCloseTo(0.04473525)
|
||||
})
|
||||
|
||||
const authorizations: Tool.Context[] = []
|
||||
const executions: string[] = []
|
||||
const permission = Layer.succeed(
|
||||
@@ -1543,6 +1553,7 @@ describe("SessionRunnerLLM", () => {
|
||||
LLMEvent.stepFinish({
|
||||
index: 0,
|
||||
reason: "tool-calls",
|
||||
providerMetadata: { copilot: { totalNanoAiu: 4_473_525_000 } },
|
||||
usage: {
|
||||
inputTokens: 10,
|
||||
nonCachedInputTokens: 8,
|
||||
@@ -1563,6 +1574,7 @@ describe("SessionRunnerLLM", () => {
|
||||
{
|
||||
type: "assistant",
|
||||
finish: "tool-calls",
|
||||
cost: 0.04473525,
|
||||
tokens: { input: 8, output: 3, reasoning: 1, cache: { read: 2, write: 0 } },
|
||||
content: [
|
||||
{ type: "reasoning", id: "reasoning-1", text: "Think" },
|
||||
|
||||
@@ -100,6 +100,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
directory: process.cwd(),
|
||||
})
|
||||
const messageIndex = new Map<string, Map<string, number>>()
|
||||
const sessionRefreshGeneration = new Map<string, number>()
|
||||
let connectionGeneration = 0
|
||||
let statusChanges: Set<string> | undefined
|
||||
let bootstrapping: Promise<void> | undefined
|
||||
@@ -358,22 +359,8 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
})
|
||||
})
|
||||
break
|
||||
case "session.step.ended":
|
||||
case "session.step.ended": {
|
||||
setSessionStatus(event.data.sessionID, "running")
|
||||
if (store.session.info[event.data.sessionID])
|
||||
setStore(
|
||||
"session",
|
||||
"info",
|
||||
event.data.sessionID,
|
||||
produce((draft) => {
|
||||
draft.cost += event.data.cost
|
||||
draft.tokens.input += event.data.tokens.input
|
||||
draft.tokens.output += event.data.tokens.output
|
||||
draft.tokens.reasoning += event.data.tokens.reasoning
|
||||
draft.tokens.cache.read += event.data.tokens.cache.read
|
||||
draft.tokens.cache.write += event.data.tokens.cache.write
|
||||
}),
|
||||
)
|
||||
message.update(event.data.sessionID, (draft, index) => {
|
||||
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
|
||||
if (!currentAssistant) return
|
||||
@@ -384,7 +371,11 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
if (event.data.snapshot)
|
||||
currentAssistant.snapshot = { ...currentAssistant.snapshot, end: event.data.snapshot }
|
||||
})
|
||||
void result.session
|
||||
.refresh(event.data.sessionID)
|
||||
.catch((error) => console.error("Failed to refresh session usage", error))
|
||||
break
|
||||
}
|
||||
case "session.step.failed":
|
||||
message.update(event.data.sessionID, (draft, index) => {
|
||||
const currentAssistant = message.assistant(draft, index, event.data.assistantMessageID)
|
||||
@@ -691,7 +682,11 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
return store.session.status[sessionID] ?? "idle"
|
||||
},
|
||||
async refresh(sessionID: string) {
|
||||
setStore("session", "info", sessionID, mutable(await sdk.api.session.get({ sessionID })))
|
||||
const generation = (sessionRefreshGeneration.get(sessionID) ?? 0) + 1
|
||||
sessionRefreshGeneration.set(sessionID, generation)
|
||||
const info = mutable(await sdk.api.session.get({ sessionID }))
|
||||
if (sessionRefreshGeneration.get(sessionID) !== generation) return
|
||||
setStore("session", "info", sessionID, info)
|
||||
registerSession(sessionID)
|
||||
},
|
||||
message: {
|
||||
@@ -875,6 +870,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
|
||||
async function bootstrap() {
|
||||
if (bootstrapping) return bootstrapping
|
||||
const generation = new Map(sessionRefreshGeneration)
|
||||
bootstrapping = Promise.allSettled([
|
||||
sdk.api.session
|
||||
.list({
|
||||
@@ -888,7 +884,10 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
|
||||
"session",
|
||||
"info",
|
||||
produce((draft) => {
|
||||
for (const session of response.data) draft[session.id] = mutable(session)
|
||||
for (const session of response.data) {
|
||||
if ((sessionRefreshGeneration.get(session.id) ?? 0) !== (generation.get(session.id) ?? 0)) continue
|
||||
draft[session.id] = mutable(session)
|
||||
}
|
||||
}),
|
||||
)
|
||||
for (const session of response.data) registerSession(session.id)
|
||||
|
||||
@@ -5,8 +5,13 @@ export function isDefaultTitle(title: string) {
|
||||
}
|
||||
|
||||
export function lastAssistantWithUsage(messages: ReadonlyArray<SessionMessage>, boundary?: string) {
|
||||
const boundaryIndex = boundary ? messages.findIndex((message) => message.id === boundary) : -1
|
||||
if (boundary && boundaryIndex === -1) return
|
||||
return messages.findLast(
|
||||
(message): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } =>
|
||||
message.type === "assistant" && message.tokens !== undefined && (!boundary || message.id < boundary),
|
||||
(
|
||||
message,
|
||||
index,
|
||||
): message is SessionMessageAssistant & { tokens: NonNullable<SessionMessageAssistant["tokens"]> } =>
|
||||
message.type === "assistant" && message.tokens !== undefined && (boundaryIndex === -1 || index < boundaryIndex),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -108,6 +108,159 @@ test("refreshes resources into reactive getters", async () => {
|
||||
}
|
||||
})
|
||||
|
||||
test("refreshes usage without applying stale session snapshots", async () => {
|
||||
const events = createEventStream()
|
||||
const sessionID = "ses_usage_refresh"
|
||||
let resolveSessions!: (response: Response) => void
|
||||
const resolveSession: Array<(response: Response) => void> = []
|
||||
let sessionsRequested = false
|
||||
const calls = createFetch((url) => {
|
||||
if (url.pathname === "/api/session") {
|
||||
sessionsRequested = true
|
||||
return new Promise<Response>((resolve) => {
|
||||
resolveSessions = resolve
|
||||
})
|
||||
}
|
||||
if (url.pathname === `/api/session/${sessionID}`) {
|
||||
return new Promise<Response>((resolve) => {
|
||||
resolveSession.push(resolve)
|
||||
})
|
||||
}
|
||||
}, events)
|
||||
let data!: ReturnType<typeof useData>
|
||||
|
||||
function Probe() {
|
||||
data = useData()
|
||||
return <box />
|
||||
}
|
||||
|
||||
const app = await testRender(() => (
|
||||
<TestTuiContexts>
|
||||
<SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
|
||||
<ProjectProvider>
|
||||
<DataProvider>
|
||||
<Probe />
|
||||
</DataProvider>
|
||||
</ProjectProvider>
|
||||
</SDKProvider>
|
||||
</TestTuiContexts>
|
||||
))
|
||||
|
||||
try {
|
||||
await wait(() => sessionsRequested)
|
||||
emitEvent(events, {
|
||||
id: "evt_usage_2",
|
||||
created: 2,
|
||||
type: "session.step.ended",
|
||||
durable: durable(sessionID, 2),
|
||||
data: {
|
||||
sessionID,
|
||||
assistantMessageID: "msg_usage_2",
|
||||
finish: "stop",
|
||||
cost: 0.5,
|
||||
tokens: { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } },
|
||||
},
|
||||
})
|
||||
await wait(() => resolveSession.length === 1)
|
||||
resolveSessions(
|
||||
json({
|
||||
data: [
|
||||
{
|
||||
id: sessionID,
|
||||
projectID: "proj_test",
|
||||
cost: 0,
|
||||
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
time: { created: 0, updated: 0 },
|
||||
title: "Stale usage",
|
||||
location: { directory },
|
||||
},
|
||||
],
|
||||
watermarks: {},
|
||||
cursor: {},
|
||||
}),
|
||||
)
|
||||
resolveSession[0]!(
|
||||
json({
|
||||
data: {
|
||||
id: sessionID,
|
||||
projectID: "proj_test",
|
||||
cost: 0.5,
|
||||
tokens: { input: 5, output: 2, reasoning: 1, cache: { read: 1, write: 1 } },
|
||||
time: { created: 0, updated: 0 },
|
||||
title: "Current usage",
|
||||
location: { directory },
|
||||
},
|
||||
}),
|
||||
)
|
||||
await wait(() => data.session.get(sessionID)?.cost === 0.5)
|
||||
expect(data.session.get(sessionID)?.tokens).toEqual({
|
||||
input: 5,
|
||||
output: 2,
|
||||
reasoning: 1,
|
||||
cache: { read: 1, write: 1 },
|
||||
})
|
||||
|
||||
emitEvent(events, {
|
||||
id: "evt_usage_3",
|
||||
created: 3,
|
||||
type: "session.step.ended",
|
||||
durable: durable(sessionID, 3),
|
||||
data: {
|
||||
sessionID,
|
||||
assistantMessageID: "msg_usage_3",
|
||||
finish: "stop",
|
||||
cost: 0.25,
|
||||
tokens: { input: 3, output: 1, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
},
|
||||
})
|
||||
emitEvent(events, {
|
||||
id: "evt_usage_4",
|
||||
created: 4,
|
||||
type: "session.step.ended",
|
||||
durable: durable(sessionID, 4),
|
||||
data: {
|
||||
sessionID,
|
||||
assistantMessageID: "msg_usage_4",
|
||||
finish: "stop",
|
||||
cost: 0.25,
|
||||
tokens: { input: 2, output: 1, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
},
|
||||
})
|
||||
await wait(() => resolveSession.length === 3)
|
||||
resolveSession[2]!(
|
||||
json({
|
||||
data: {
|
||||
id: sessionID,
|
||||
projectID: "proj_test",
|
||||
cost: 1,
|
||||
tokens: { input: 10, output: 4, reasoning: 1, cache: { read: 1, write: 1 } },
|
||||
time: { created: 0, updated: 0 },
|
||||
title: "Latest usage",
|
||||
location: { directory },
|
||||
},
|
||||
}),
|
||||
)
|
||||
await wait(() => data.session.get(sessionID)?.cost === 1)
|
||||
resolveSession[1]!(
|
||||
json({
|
||||
data: {
|
||||
id: sessionID,
|
||||
projectID: "proj_test",
|
||||
cost: 0.75,
|
||||
tokens: { input: 8, output: 3, reasoning: 1, cache: { read: 1, write: 1 } },
|
||||
time: { created: 0, updated: 0 },
|
||||
title: "Older usage",
|
||||
location: { directory },
|
||||
},
|
||||
}),
|
||||
)
|
||||
await Bun.sleep(20)
|
||||
expect(data.session.get(sessionID)?.cost).toBe(1)
|
||||
} finally {
|
||||
app.renderer.destroy()
|
||||
}
|
||||
})
|
||||
|
||||
test("reconnects the event stream and bootstraps fresh data", async () => {
|
||||
const events = createEventStream()
|
||||
const requests = { active: 0, event: 0, model: 0 }
|
||||
@@ -348,6 +501,7 @@ test("connectedOnce is false until first connect and persists across disconnect"
|
||||
|
||||
test("tracks session status from active sessions and execution events", async () => {
|
||||
const events = createEventStream()
|
||||
let settled = false
|
||||
const calls = createFetch((url) => {
|
||||
if (url.pathname === "/api/session/active")
|
||||
return json({ data: { "session-active": { type: "running" } }, watermarks: {} })
|
||||
@@ -356,8 +510,10 @@ test("tracks session status from active sessions and execution events", async ()
|
||||
data: {
|
||||
id: "session-live",
|
||||
projectID: "proj_test",
|
||||
cost: 0,
|
||||
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
cost: settled ? 0.75 : 0,
|
||||
tokens: settled
|
||||
? { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 1 } }
|
||||
: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
time: { created: 0, updated: 0 },
|
||||
title: "Live session",
|
||||
location: { directory },
|
||||
@@ -388,6 +544,7 @@ test("tracks session status from active sessions and execution events", async ()
|
||||
expect(data.session.status("session-idle")).toBe("idle")
|
||||
await data.session.refresh("session-live")
|
||||
|
||||
settled = true
|
||||
emitEvent(events, {
|
||||
id: "evt_step_started",
|
||||
created: 0,
|
||||
@@ -419,6 +576,7 @@ test("tracks session status from active sessions and execution events", async ()
|
||||
const assistant = data.session.message.get("session-live", "message-live")
|
||||
return assistant?.type === "assistant" && assistant.finish === "stop"
|
||||
})
|
||||
await wait(() => data.session.get("session-live")?.cost === 0.75)
|
||||
expect(data.session.status("session-live")).toBe("running")
|
||||
expect(data.session.get("session-live")).toMatchObject({
|
||||
cost: 0.75,
|
||||
|
||||
@@ -19,10 +19,11 @@ describe("util.session", () => {
|
||||
tokens: { input, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
||||
time: { created: 0 },
|
||||
})
|
||||
const messages = [assistant("msg_1", 10), assistant("msg_3", 30)]
|
||||
const messages = [assistant("msg_z", 10), assistant("msg_a", 30)]
|
||||
|
||||
expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30)
|
||||
expect(lastAssistantWithUsage(messages, "msg_2")?.tokens.input).toBe(10)
|
||||
expect(lastAssistantWithUsage(messages, "msg_a")?.tokens.input).toBe(10)
|
||||
expect(lastAssistantWithUsage(messages, "msg_missing")).toBeUndefined()
|
||||
expect(lastAssistantWithUsage(messages)?.tokens.input).toBe(30)
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user