From 1819fe2711cdcbcfd8f077f87bb669f29f218308 Mon Sep 17 00:00:00 2001 From: Dax Raad Date: Sun, 10 May 2026 20:23:40 -0400 Subject: [PATCH] Move session usage backfill to data migration --- .../migration.sql | 22 +---- packages/opencode/src/data-migration.ts | 83 ++++++++++++++++++- 2 files changed, 82 insertions(+), 23 deletions(-) diff --git a/packages/opencode/migration/20260510033149_session_usage/migration.sql b/packages/opencode/migration/20260510033149_session_usage/migration.sql index 40a79258f3..68e12aad09 100644 --- a/packages/opencode/migration/20260510033149_session_usage/migration.sql +++ b/packages/opencode/migration/20260510033149_session_usage/migration.sql @@ -3,24 +3,4 @@ ALTER TABLE `session` ADD `tokens_input` integer DEFAULT 0 NOT NULL;--> statemen ALTER TABLE `session` ADD `tokens_output` integer DEFAULT 0 NOT NULL;--> statement-breakpoint ALTER TABLE `session` ADD `tokens_reasoning` integer DEFAULT 0 NOT NULL;--> statement-breakpoint ALTER TABLE `session` ADD `tokens_cache_read` integer DEFAULT 0 NOT NULL;--> statement-breakpoint -ALTER TABLE `session` ADD `tokens_cache_write` integer DEFAULT 0 NOT NULL;--> statement-breakpoint -CREATE TEMP TABLE `session_usage` AS -SELECT - `part`.`session_id`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.cost'), 0)), 0) AS `cost`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.input'), 0)), 0) AS `tokens_input`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.output'), 0)), 0) AS `tokens_output`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.reasoning'), 0)), 0) AS `tokens_reasoning`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.cache.read'), 0)), 0) AS `tokens_cache_read`, - COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.cache.write'), 0)), 0) AS `tokens_cache_write` -FROM `part` -WHERE json_extract(`part`.`data`, '$.type') = 'step-finish' -GROUP BY `part`.`session_id`;--> statement-breakpoint -UPDATE `session` SET - `cost` = COALESCE((SELECT `cost` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0), - `tokens_input` = COALESCE((SELECT `tokens_input` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0), - `tokens_output` = COALESCE((SELECT `tokens_output` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0), - `tokens_reasoning` = COALESCE((SELECT `tokens_reasoning` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0), - `tokens_cache_read` = COALESCE((SELECT `tokens_cache_read` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0), - `tokens_cache_write` = COALESCE((SELECT `tokens_cache_write` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0);--> statement-breakpoint -DROP TABLE `session_usage`; +ALTER TABLE `session` ADD `tokens_cache_write` integer DEFAULT 0 NOT NULL; diff --git a/packages/opencode/src/data-migration.ts b/packages/opencode/src/data-migration.ts index c3e5a9d2b0..7a2393e120 100644 --- a/packages/opencode/src/data-migration.ts +++ b/packages/opencode/src/data-migration.ts @@ -2,7 +2,9 @@ import { Context, Effect, Layer } from "effect" import { Database } from "./storage/db" import { DataMigrationTable } from "./data-migration.sql" import * as Log from "@opencode-ai/core/util/log" -import { eq } from "drizzle-orm" +import { asc, eq, gt, inArray } from "drizzle-orm" +import { PartTable, SessionTable } from "./session/session.sql" +import type { SessionID } from "./session/schema" export type Migration = { name: string @@ -18,7 +20,84 @@ export class Service extends Context.Service()("@opencode/Da export const layer = Layer.effect( Service, Effect.gen(function* () { - const migrations: Migration[] = [] + const migrations: Migration[] = [ + { + name: "session_usage_from_parts", + run: Effect.sync(() => + Database.transaction((db) => { + type Usage = { + cost: number + tokens: { + input: number + output: number + reasoning: number + cache: { read: number; write: number } + } + } + + for (let cursor: SessionID | undefined; ; ) { + const sessions = db + .select({ id: SessionTable.id }) + .from(SessionTable) + .where(cursor ? gt(SessionTable.id, cursor) : undefined) + .orderBy(asc(SessionTable.id)) + .limit(500) + .all() + if (sessions.length === 0) return + + const usageBySession = new Map( + sessions.map((session) => [ + session.id, + { cost: 0, tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } }, + ]), + ) + + for (const row of db + .select({ session_id: PartTable.session_id, data: PartTable.data }) + .from(PartTable) + .where(inArray(PartTable.session_id, sessions.map((session) => session.id))) + .all()) { + const part = row.data + if (part.type !== "step-finish") continue + if (!("cost" in part) || typeof part.cost !== "number") continue + if (!("tokens" in part) || typeof part.tokens !== "object" || part.tokens === null) continue + if (!("input" in part.tokens) || typeof part.tokens.input !== "number") continue + if (!("output" in part.tokens) || typeof part.tokens.output !== "number") continue + if (!("reasoning" in part.tokens) || typeof part.tokens.reasoning !== "number") continue + if (!("cache" in part.tokens) || typeof part.tokens.cache !== "object" || part.tokens.cache === null) continue + if (!("read" in part.tokens.cache) || typeof part.tokens.cache.read !== "number") continue + if (!("write" in part.tokens.cache) || typeof part.tokens.cache.write !== "number") continue + + const current = usageBySession.get(row.session_id) + if (!current) continue + current.cost += part.cost + current.tokens.input += part.tokens.input + current.tokens.output += part.tokens.output + current.tokens.reasoning += part.tokens.reasoning + current.tokens.cache.read += part.tokens.cache.read + current.tokens.cache.write += part.tokens.cache.write + } + + for (const [sessionID, value] of usageBySession) { + db.update(SessionTable) + .set({ + cost: value.cost, + tokens_input: value.tokens.input, + tokens_output: value.tokens.output, + tokens_reasoning: value.tokens.reasoning, + tokens_cache_read: value.tokens.cache.read, + tokens_cache_write: value.tokens.cache.write, + }) + .where(eq(SessionTable.id, sessionID)) + .run() + } + + cursor = sessions.at(-1)?.id + } + }), + ), + }, + ] yield* Effect.gen(function* () { if (migrations.length === 0) return