diff --git a/packages/core/schema.json b/packages/core/schema.json index 0f52ae1bed12..7ebbebef0859 100644 --- a/packages/core/schema.json +++ b/packages/core/schema.json @@ -1,8 +1,8 @@ { "version": "7", "dialect": "sqlite", - "id": "5067fae0-90ff-4601-82d7-dd050ea9dc8f", - "prevIds": ["e98619ef-c334-45e3-8c48-9932ed1ef290"], + "id": "9b8d5c43-f59d-4cca-9c90-f66c5bc96060", + "prevIds": ["5067fae0-90ff-4601-82d7-dd050ea9dc8f"], "ddl": [ { "name": "account_state", @@ -68,6 +68,10 @@ "name": "session_pending", "entityType": "tables" }, + { + "name": "session_revert_usage", + "entityType": "tables" + }, { "name": "session_v2", "entityType": "tables" @@ -1080,6 +1084,136 @@ "entityType": "columns", "table": "session_pending" }, + { + "type": "text", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "id", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "text", + "notNull": true, + "autoincrement": false, + "default": null, + "generated": null, + "name": "session_id", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "text", + "notNull": true, + "autoincrement": false, + "default": null, + "generated": null, + "name": "type", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": null, + "generated": null, + "name": "time_created", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "text", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "model_provider_id", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "text", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "model_id", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "text", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "model_variant", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "real", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "cost", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "tokens_input", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "tokens_output", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "tokens_reasoning", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "tokens_cache_read", + "entityType": "columns", + "table": "session_revert_usage" + }, + { + "type": "integer", + "notNull": true, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "tokens_cache_write", + "entityType": "columns", + "table": "session_revert_usage" + }, { "type": "text", "notNull": false, @@ -1629,6 +1763,17 @@ "entityType": "fks", "table": "session_pending" }, + { + "columns": ["session_id"], + "tableTo": "session_v2", + "columnsTo": ["id"], + "onUpdate": "NO ACTION", + "onDelete": "CASCADE", + "nameExplicit": false, + "name": "fk_session_revert_usage_session_id_session_v2_id_fk", + "entityType": "fks", + "table": "session_revert_usage" + }, { "columns": ["project_id"], "tableTo": "project", @@ -1770,6 +1915,13 @@ "table": "session_pending", "entityType": "pks" }, + { + "columns": ["id"], + "nameExplicit": false, + "name": "session_revert_usage_pk", + "table": "session_revert_usage", + "entityType": "pks" + }, { "columns": ["id"], "nameExplicit": false, @@ -2016,6 +2168,38 @@ "entityType": "indexes", "table": "session_pending" }, + { + "columns": [ + { + "value": "session_id", + "isExpression": false + }, + { + "value": "time_created", + "isExpression": false + } + ], + "isUnique": false, + "where": null, + "origin": "manual", + "name": "session_revert_usage_session_time_created_idx", + "entityType": "indexes", + "table": "session_revert_usage" + }, + { + "columns": [ + { + "value": "time_created", + "isExpression": false + } + ], + "isUnique": false, + "where": null, + "origin": "manual", + "name": "session_revert_usage_time_created_idx", + "entityType": "indexes", + "table": "session_revert_usage" + }, { "columns": [ { diff --git a/packages/core/src/database/migration.gen.ts b/packages/core/src/database/migration.gen.ts index 2752a953cf30..55a29b1dc261 100644 --- a/packages/core/src/database/migration.gen.ts +++ b/packages/core/src/database/migration.gen.ts @@ -49,6 +49,7 @@ import m46 from "./migration/20260910120000_clear_v1_session_permission.js" import m47 from "./migration/20260923013825_project_time_active.js" import m48 from "./migration/20261007190000_azure_cli_external_credential.js" import m49 from "./migration/20261010180456_project_time_archived.js" +import m50 from "./migration/20261022034137_session_revert_usage.js" export const migrations = [ m00, @@ -101,4 +102,5 @@ export const migrations = [ m47, m48, m49, + m50, ] satisfies DatabaseMigration.Migration[] diff --git a/packages/core/src/database/migration/20261022034137_session_revert_usage.ts b/packages/core/src/database/migration/20261022034137_session_revert_usage.ts new file mode 100644 index 000000000000..e8a825e98160 --- /dev/null +++ b/packages/core/src/database/migration/20261022034137_session_revert_usage.ts @@ -0,0 +1,36 @@ +import { Effect } from "effect" +import type { DatabaseMigration } from "../migration.js" + +const migration: DatabaseMigration.Migration = { + id: "20261022034137_session_revert_usage", + up(tx) { + return Effect.gen(function* () { + yield* tx.run(` + CREATE TABLE \`session_revert_usage\` ( + \`id\` text PRIMARY KEY, + \`session_id\` text NOT NULL, + \`type\` text NOT NULL, + \`time_created\` integer NOT NULL, + \`model_provider_id\` text, + \`model_id\` text, + \`model_variant\` text, + \`cost\` real DEFAULT 0 NOT NULL, + \`tokens_input\` integer DEFAULT 0 NOT NULL, + \`tokens_output\` integer DEFAULT 0 NOT NULL, + \`tokens_reasoning\` integer DEFAULT 0 NOT NULL, + \`tokens_cache_read\` integer DEFAULT 0 NOT NULL, + \`tokens_cache_write\` integer DEFAULT 0 NOT NULL, + CONSTRAINT \`fk_session_revert_usage_session_id_session_v2_id_fk\` FOREIGN KEY (\`session_id\`) REFERENCES \`session_v2\`(\`id\`) ON DELETE CASCADE + ); + `) + yield* tx.run( + `CREATE INDEX \`session_revert_usage_session_time_created_idx\` ON \`session_revert_usage\` (\`session_id\`,\`time_created\`);`, + ) + yield* tx.run( + `CREATE INDEX \`session_revert_usage_time_created_idx\` ON \`session_revert_usage\` (\`time_created\`);`, + ) + }) + }, +} + +export default migration diff --git a/packages/core/src/database/schema.gen.ts b/packages/core/src/database/schema.gen.ts index 3ce8c6fdab94..7501d86aeb63 100644 --- a/packages/core/src/database/schema.gen.ts +++ b/packages/core/src/database/schema.gen.ts @@ -180,6 +180,24 @@ const schema: Omit = { CONSTRAINT \`fk_session_pending_session_id_session_v2_id_fk\` FOREIGN KEY (\`session_id\`) REFERENCES \`session_v2\`(\`id\`) ON DELETE CASCADE ); `) + yield* tx.run(` + CREATE TABLE \`session_revert_usage\` ( + \`id\` text PRIMARY KEY, + \`session_id\` text NOT NULL, + \`type\` text NOT NULL, + \`time_created\` integer NOT NULL, + \`model_provider_id\` text, + \`model_id\` text, + \`model_variant\` text, + \`cost\` real DEFAULT 0 NOT NULL, + \`tokens_input\` integer DEFAULT 0 NOT NULL, + \`tokens_output\` integer DEFAULT 0 NOT NULL, + \`tokens_reasoning\` integer DEFAULT 0 NOT NULL, + \`tokens_cache_read\` integer DEFAULT 0 NOT NULL, + \`tokens_cache_write\` integer DEFAULT 0 NOT NULL, + CONSTRAINT \`fk_session_revert_usage_session_id_session_v2_id_fk\` FOREIGN KEY (\`session_id\`) REFERENCES \`session_v2\`(\`id\`) ON DELETE CASCADE + ); + `) yield* tx.run(` CREATE TABLE \`session_v2\` ( \`id\` text PRIMARY KEY, @@ -270,6 +288,12 @@ const schema: Omit = { yield* tx.run( `CREATE UNIQUE INDEX \`session_pending_session_admitted_seq_idx\` ON \`session_pending\` (\`session_id\`,\`admitted_seq\`);`, ) + yield* tx.run( + `CREATE INDEX \`session_revert_usage_session_time_created_idx\` ON \`session_revert_usage\` (\`session_id\`,\`time_created\`);`, + ) + yield* tx.run( + `CREATE INDEX \`session_revert_usage_time_created_idx\` ON \`session_revert_usage\` (\`time_created\`);`, + ) yield* tx.run(`CREATE INDEX \`session_v2_project_idx\` ON \`session_v2\` (\`project_id\`);`) yield* tx.run(`CREATE INDEX \`session_v2_workspace_idx\` ON \`session_v2\` (\`workspace_id\`);`) yield* tx.run(`CREATE INDEX \`session_v2_parent_idx\` ON \`session_v2\` (\`parent_id\`);`) diff --git a/packages/core/src/session/projector.ts b/packages/core/src/session/projector.ts index 041bfe86bde7..99cd875e3f99 100644 --- a/packages/core/src/session/projector.ts +++ b/packages/core/src/session/projector.ts @@ -793,6 +793,36 @@ const layer = Layer.effectDiscard( .get() .pipe(Effect.orDie) if (!boundary) return yield* Effect.die(new Error(`Revert boundary message not found: ${event.data.to}`)) + // Preserve the reverted rows' prompts, steps and usage for SessionStats before + // deleting them. The fork filter matches SessionStats: rows a fork copied from its + // parent belong to the parent's stats, not this session's. + yield* db + .run( + sql` + INSERT OR IGNORE INTO session_revert_usage (id, session_id, type, time_created, model_provider_id, model_id, model_variant, cost, tokens_input, tokens_output, tokens_reasoning, tokens_cache_read, tokens_cache_write) + SELECT + message.id, + message.session_id, + message.type, + message.time_created, + json_extract(message.data, '$.model.providerID'), + json_extract(message.data, '$.model.id'), + json_extract(message.data, '$.model.variant'), + coalesce(json_extract(message.data, '$.cost'), 0), + coalesce(json_extract(message.data, '$.tokens.input'), 0), + coalesce(json_extract(message.data, '$.tokens.output'), 0), + coalesce(json_extract(message.data, '$.tokens.reasoning'), 0), + coalesce(json_extract(message.data, '$.tokens.cache.read'), 0), + coalesce(json_extract(message.data, '$.tokens.cache.write'), 0) + FROM ${SessionMessageTable} AS message + JOIN ${SessionTable} AS session ON session.id = message.session_id + WHERE message.type IN ('user', 'assistant') + AND message.session_id = ${event.data.sessionID} + AND message.seq >= ${boundary.seq} + AND (session.fork_session_id IS NULL OR message.time_created >= session.time_created) + `, + ) + .pipe(Effect.orDie) yield* db .delete(SessionMessageTable) .where( diff --git a/packages/core/src/session/sql.ts b/packages/core/src/session/sql.ts index 88e9f925600b..a4f24f58f508 100644 --- a/packages/core/src/session/sql.ts +++ b/packages/core/src/session/sql.ts @@ -98,6 +98,35 @@ export const SessionMessageTable = sqliteTable( ], ) +// Messages removed by a revert commit, copied out before the rows are deleted so SessionStats +// can keep counting the reverted prompts, steps and usage. Written in the projector, which +// runs regardless of event persistence (unlike durable events, which only workerd stores). +export const SessionRevertUsageTable = sqliteTable( + "session_revert_usage", + { + id: text().$type().primaryKey(), + session_id: text() + .$type() + .notNull() + .references(() => SessionTable.id, { onDelete: "cascade" }), + type: text().$type().notNull(), + time_created: integer().notNull(), + model_provider_id: text(), + model_id: text(), + model_variant: text(), + cost: real().notNull().default(0), + tokens_input: integer().notNull().default(0), + tokens_output: integer().notNull().default(0), + tokens_reasoning: integer().notNull().default(0), + tokens_cache_read: integer().notNull().default(0), + tokens_cache_write: integer().notNull().default(0), + }, + (table) => [ + index("session_revert_usage_session_time_created_idx").on(table.session_id, table.time_created), + index("session_revert_usage_time_created_idx").on(table.time_created), + ], +) + export const SessionPendingTable = sqliteTable( "session_pending", { diff --git a/packages/core/src/session/stats.ts b/packages/core/src/session/stats.ts index 17a51576ec8e..7bcfccf85c80 100644 --- a/packages/core/src/session/stats.ts +++ b/packages/core/src/session/stats.ts @@ -12,7 +12,7 @@ import { ToolMode } from "@opencode/schema/session-stats" import { TokenUsage } from "@opencode/schema/token-usage" import { Database } from "../database/database.js" import { EventTable } from "../event/sql.js" -import { SessionMessageTable, SessionTable } from "./sql.js" +import { SessionMessageTable, SessionRevertUsageTable, SessionTable } from "./sql.js" type Input = { readonly from?: number @@ -92,13 +92,22 @@ export const get = Effect.fn("SessionStats.get")(function* (input: Input = {}) { (yield* db .get<{ time: number | null }>( sql` - SELECT min(message.time_created) AS time - FROM ${SessionMessageTable} AS message - JOIN ${SessionTable} AS session ON session.id = message.session_id - WHERE message.type IN ('user', 'assistant') - AND message.time_created < ${to} - AND (session.fork_session_id IS NULL OR message.time_created >= session.time_created) - ${project} + SELECT min(time_created) AS time + FROM ( + SELECT message.time_created AS time_created + FROM ${SessionMessageTable} AS message + JOIN ${SessionTable} AS session ON session.id = message.session_id + WHERE message.type IN ('user', 'assistant') + AND message.time_created < ${to} + AND (session.fork_session_id IS NULL OR message.time_created >= session.time_created) + ${project} + UNION ALL + SELECT message.time_created AS time_created + FROM ${SessionRevertUsageTable} AS message + JOIN ${SessionTable} AS session ON session.id = message.session_id + WHERE message.time_created < ${to} + ${project} + ) `, ) .pipe(Effect.orDie))?.time ?? @@ -146,6 +155,26 @@ export const get = Effect.fn("SessionStats.get")(function* (input: Input = {}) { AND message.time_created < ${range.to} AND (session.fork_session_id IS NULL OR message.time_created >= session.time_created) ${project} + UNION ALL + SELECT + message.session_id AS sessionID, + session.parent_id AS parentID, + message.type AS type, + message.time_created AS timeCreated, + message.model_provider_id AS providerID, + message.model_id AS modelID, + message.model_variant AS variant, + message.tokens_input AS input, + message.tokens_output AS output, + message.tokens_reasoning AS reasoning, + message.tokens_cache_read AS cacheRead, + message.tokens_cache_write AS cacheWrite, + message.cost AS cost + FROM ${SessionRevertUsageTable} AS message + JOIN ${SessionTable} AS session ON session.id = message.session_id + WHERE message.time_created >= ${range.from} + AND message.time_created < ${range.to} + ${project} `, ) .pipe( diff --git a/packages/core/test/session-revert.test.ts b/packages/core/test/session-revert.test.ts index bc3906857760..efbee77a581c 100644 --- a/packages/core/test/session-revert.test.ts +++ b/packages/core/test/session-revert.test.ts @@ -18,6 +18,9 @@ import { SessionExecution } from "@opencode/core/session/execution" import { SessionInbox } from "@opencode/core/session/inbox" import { SessionMessage } from "@opencode/core/session/message" import { SessionProjector } from "@opencode/core/session/projector" +import { SessionStats } from "@opencode/core/session/stats" +import { SessionMessageTable, SessionRevertUsageTable, SessionTable } from "@opencode/core/session/sql" +import { and, asc, eq } from "drizzle-orm" import { Snapshot } from "@opencode/core/snapshot" import { Money } from "@opencode/schema/money" import { LayerNode } from "@opencode/util/effect/layer-node" @@ -119,3 +122,187 @@ describe("Session.revert files", () => { { timeout: 15_000 }, ) }) + +// The production bus only persists events on workerd; local CLI and server runs deliver +// events in memory with an empty event table. These tests use the default (non-persisting) +// bus to prove stats keep reverted usage where the CLI and server actually run. +const itEphemeral = testEffect( + AppNodeBuilder.build( + LayerNode.group([Database.node, Bus.node, SessionProjector.node, Session.node, LocationServiceMap.node]), + [ + Global.node.replace(tempGlobalLayer), + SessionExecution.node.replace(SessionExecution.noopLayer), + offlineModels, + ], + ), +) + +describe("Session.revert stats", () => { + const initProject = (directory: string) => + Effect.promise(async () => { + await fs.mkdir(directory, { recursive: true }) + await $`git init -q`.cwd(directory).quiet() + }) + + const prompt = Effect.fnUntraced(function* (sessionID: Session.ID, text: string) { + const session = yield* Session.Service + const database = yield* Database.Service + const bus = yield* Bus.Service + const admitted = yield* session.prompt({ sessionID, text, resume: false }) + yield* SessionInbox.promote(database.db, bus, sessionID, "steer") + return admitted + }) + + // One assistant step with fixed usage; returns the assistant message id. + const step = Effect.fnUntraced(function* (sessionID: Session.ID) { + const bus = yield* Bus.Service + const snapshot = yield* Snapshot.Service + const before = yield* snapshot.capture() + if (!before) throw new Error("Initial snapshot missing") + const assistantMessageID = SessionMessage.ID.create() + yield* bus.publish(SessionEvent.Step.Started, { + sessionID, + assistantMessageID, + agent: Agent.defaultID, + model: { id: Model.ID.make("test-model"), providerID: Provider.ID.make("test-provider") }, + snapshot: before, + started: 0, + }) + const after = yield* snapshot.capture() + if (!after) throw new Error("Snapshot missing") + yield* bus.publish(SessionEvent.Step.Ended, { + sessionID, + assistantMessageID, + finish: "stop", + cost: Money.USD.make(0.5), + tokens: { input: 1000, output: 100, reasoning: 10, cache: { read: 5, write: 3 } }, + snapshot: after, + files: yield* snapshot.files({ from: before, to: after }), + }) + return assistantMessageID + }) + + itEphemeral.live( + "keeps reverted step usage in stats without double counting session totals", + () => + Effect.gen(function* () { + const tmp = yield* tmpdirScoped() + const directory = path.join(tmp.path, "project") + yield* initProject(directory) + + const session = yield* Session.Service + const database = yield* Database.Service + const created = yield* session.create({ location: { directory: AbsolutePath.make(directory) } }) + + yield* Effect.gen(function* () { + const plugins = yield* Plugin.Service + yield* plugins.awaitActivation + + yield* prompt(created.id, "First") + yield* step(created.id) + const second = yield* prompt(created.id, "Second") + yield* step(created.id) + + const before = yield* SessionStats.get({ timezone: "UTC", tools: "none" }) + expect(before.tokens).toEqual({ + input: 2000, + output: 200, + reasoning: 20, + cache: { read: 10, write: 6 }, + }) + + // Prompting with a staged revert commits it, deleting the second exchange's message rows. + yield* session.revert.stage({ sessionID: created.id, messageID: second.id }) + yield* prompt(created.id, "Third") + + const stats = yield* SessionStats.get({ timezone: "UTC", tools: "none" }) + expect(stats.tokens).toEqual({ + input: 2000, + output: 200, + reasoning: 20, + cache: { read: 10, write: 6 }, + }) + expect(stats.cost).toBe(Money.USD.make(1)) + + const row = yield* database.db + .select() + .from(SessionTable) + .where(eq(SessionTable.id, created.id)) + .get() + .pipe(Effect.orDie) + expect(row?.tokens_input).toBe(2000) + expect(row?.tokens_output).toBe(200) + expect(row?.tokens_reasoning).toBe(20) + expect(row?.tokens_cache_read).toBe(10) + expect(row?.tokens_cache_write).toBe(6) + expect(row?.cost).toBe(1) + }).pipe(Effect.provide(LocationServiceMap.Service.get(created.location))) + }), + { timeout: 15_000 }, + ) + + itEphemeral.live( + "does not recount parent rows a fork copied when the fork reverts them", + () => + Effect.gen(function* () { + const tmp = yield* tmpdirScoped() + const directory = path.join(tmp.path, "project") + yield* initProject(directory) + + const session = yield* Session.Service + const database = yield* Database.Service + const bus = yield* Bus.Service + const created = yield* session.create({ location: { directory: AbsolutePath.make(directory) } }) + + yield* Effect.gen(function* () { + const plugins = yield* Plugin.Service + yield* plugins.awaitActivation + + yield* prompt(created.id, "First") + yield* step(created.id) + yield* prompt(created.id, "Second") + const secondStep = yield* step(created.id) + + const forkID = Session.ID.make("ses_revert_stats_fork") + yield* bus.publish(SessionEvent.Forked, { + sessionID: forkID, + parentID: created.id, + boundary: { type: "through", messageID: secondStep }, + }) + + // The revert boundary is the fork's copy of the second prompt, not the original. + const copied = yield* database.db + .select({ id: SessionMessageTable.id }) + .from(SessionMessageTable) + .where(and(eq(SessionMessageTable.session_id, forkID), eq(SessionMessageTable.type, "user"))) + .orderBy(asc(SessionMessageTable.seq)) + .all() + .pipe(Effect.orDie) + expect(copied.length).toBe(2) + + yield* session.revert.stage({ sessionID: forkID, messageID: copied[1].id }) + yield* prompt(forkID, "Third") + + // The fork's copied second exchange must not be counted again (the filter matches + // SessionStats: copied rows belong to the parent's stats). No new usage was produced, + // so totals stay exactly where they were. + const stats = yield* SessionStats.get({ timezone: "UTC", tools: "none" }) + expect(stats.tokens).toEqual({ + input: 2000, + output: 200, + reasoning: 20, + cache: { read: 10, write: 6 }, + }) + expect(stats.cost).toBe(Money.USD.make(1)) + + const preserved = yield* database.db + .select({ id: SessionRevertUsageTable.id }) + .from(SessionRevertUsageTable) + .all() + .pipe(Effect.orDie) + expect(preserved.length).toBe(0) + }).pipe(Effect.provide(LocationServiceMap.Service.get(created.location))) + }), + { timeout: 15_000 }, + ) +})