summaryrefslogtreecommitdiffhomepage
path: root/packages/core
diff options
context:
space:
mode:
authorAdam Malczewski <[email protected]>2026-06-02 13:08:36 +0900
committerAdam Malczewski <[email protected]>2026-06-02 13:08:36 +0900
commitb734eb96bf0af267fdfbef85df51940ca0b4e8c7 (patch)
tree3815dc9e569bb57384f53d9042288ff831f02e74 /packages/core
parent3f629a8469fe483243671e1ca15582a111e96541 (diff)
downloaddispatch-b734eb96bf0af267fdfbef85df51940ca0b4e8c7.tar.gz
dispatch-b734eb96bf0af267fdfbef85df51940ca0b4e8c7.zip
feat: persist per-tab token/cache usage across reload
Persist usage as invisible type:"usage" chunk rows (side channel): - core: add "usage" ChunkType + UsageData; exclude usage rows from getChunksForTab/getTotalChunkCount; add getUsageStatsForTab aggregate (exported from barrel); defensive skip in groupRowsToMessages. - api: agent-manager accumulates per-attempt usageRows and flushes them in the same atomic appendChunks call as the turn's content (discarded on a superseded fallback attempt). GET /tabs enriches rows with usageStats. - frontend: hydrateFromBackend seeds cacheStats from usageStats (reload only; no re-seed on statuses reconnect, so no double-count with live events). Tests: core DB-backed usage persistence/aggregate; api usage-row-per-event + fallback discard; routes GET /tabs usageStats; frontend hydrate seed + no-double-count + live-accumulation-after-seed. 495 -> 509 passing.
Diffstat (limited to 'packages/core')
-rw-r--r--packages/core/src/chunks/transform.ts5
-rw-r--r--packages/core/src/db/chunks.ts90
-rw-r--r--packages/core/src/index.ts1
-rw-r--r--packages/core/src/types/index.ts28
-rw-r--r--packages/core/tests/db/chunks.db.test.ts336
5 files changed, 451 insertions, 9 deletions
diff --git a/packages/core/src/chunks/transform.ts b/packages/core/src/chunks/transform.ts
index a4c6fc8..e8f4a18 100644
--- a/packages/core/src/chunks/transform.ts
+++ b/packages/core/src/chunks/transform.ts
@@ -209,6 +209,11 @@ export function groupRowsToMessages(rows: ChunkRow[]): MessageRow[] {
continue;
}
+ // Usage rows are an invisible side channel (persisted for the backend
+ // aggregate only). They're already query-excluded from getChunksForTab,
+ // so this is defensive insurance: never let one leak into render grouping.
+ if (row.type === "usage") continue;
+
// assistant / tool rows → part of the current assistant message
const c = ensureAssistant(row);
switch (row.type) {
diff --git a/packages/core/src/db/chunks.ts b/packages/core/src/db/chunks.ts
index 077259d..509cbbd 100644
--- a/packages/core/src/db/chunks.ts
+++ b/packages/core/src/db/chunks.ts
@@ -5,7 +5,7 @@ import {
groupRowsToMessages,
type MessageRow,
} from "../chunks/transform.js";
-import type { ChunkData, ChunkRow, ChunkRowDraft, TextData } from "../types/index.js";
+import type { ChunkData, ChunkRow, ChunkRowDraft, TextData, UsageData } from "../types/index.js";
import { getDatabase } from "./index.js";
// Re-export the DB-free transforms so existing barrel consumers
@@ -101,7 +101,7 @@ export function getChunksForTab(
const db = getDatabase();
if (!options) {
const rows = db
- .query("SELECT * FROM chunks WHERE tab_id = $tabId ORDER BY seq ASC")
+ .query("SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' ORDER BY seq ASC")
.all({ $tabId: tabId }) as Array<Record<string, unknown>>;
return rows.map(mapRow);
}
@@ -110,24 +110,28 @@ export function getChunksForTab(
if (limit !== undefined) {
const rows = db
.query(
- "SELECT * FROM chunks WHERE tab_id = $tabId AND seq < $before ORDER BY seq DESC LIMIT $limit",
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' AND seq < $before ORDER BY seq DESC LIMIT $limit",
)
.all({ $tabId: tabId, $before: before, $limit: limit }) as Array<Record<string, unknown>>;
return rows.map(mapRow).reverse();
}
const rows = db
- .query("SELECT * FROM chunks WHERE tab_id = $tabId AND seq < $before ORDER BY seq DESC")
+ .query(
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' AND seq < $before ORDER BY seq DESC",
+ )
.all({ $tabId: tabId, $before: before }) as Array<Record<string, unknown>>;
return rows.map(mapRow).reverse();
}
if (limit !== undefined) {
const rows = db
- .query("SELECT * FROM chunks WHERE tab_id = $tabId ORDER BY seq DESC LIMIT $limit")
+ .query(
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' ORDER BY seq DESC LIMIT $limit",
+ )
.all({ $tabId: tabId, $limit: limit }) as Array<Record<string, unknown>>;
return rows.map(mapRow).reverse();
}
const rows = db
- .query("SELECT * FROM chunks WHERE tab_id = $tabId ORDER BY seq ASC")
+ .query("SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' ORDER BY seq ASC")
.all({ $tabId: tabId }) as Array<Record<string, unknown>>;
return rows.map(mapRow);
}
@@ -145,11 +149,83 @@ export function getMessagesForTab(tabId: string): MessageRow[] {
export function getTotalChunkCount(tabId: string): number {
const db = getDatabase();
const row = db
- .query("SELECT COUNT(*) as count FROM chunks WHERE tab_id = $tabId")
+ .query("SELECT COUNT(*) as count FROM chunks WHERE tab_id = $tabId AND type != 'usage'")
.get({ $tabId: tabId }) as { count: number } | null;
return row?.count ?? 0;
}
+/**
+ * Aggregate per-tab token/cache usage across ALL persisted `usage` chunk rows.
+ *
+ * Usage rows are written as an invisible side channel (one row per `usage`
+ * AgentEvent) and are query-excluded from `getChunksForTab`/`getTotalChunkCount`,
+ * so this aggregate is the read path. Because it sums server-side over every
+ * row, it stays complete even after the frontend evicts/pages out old turns
+ * (eviction is in-memory only). The return shape is structurally identical to
+ * the frontend `CacheStats`, so reload can seed it directly.
+ *
+ * - cumulative `inputTokens`/`outputTokens`/`cacheReadTokens`/`cacheWriteTokens`
+ * = SUM over all usage rows;
+ * - `requests` = COUNT of usage rows;
+ * - `last` = the highest-seq usage row's split (most recent request);
+ * - `null` when the tab has no usage rows.
+ *
+ * Sums in JS after selecting the rows (mirroring `mapRow`) to avoid relying on
+ * `json_extract` over the freeform `data_json`.
+ */
+export function getUsageStatsForTab(tabId: string): {
+ inputTokens: number;
+ outputTokens: number;
+ cacheReadTokens: number;
+ cacheWriteTokens: number;
+ requests: number;
+ last: {
+ inputTokens: number;
+ outputTokens: number;
+ cacheReadTokens: number;
+ cacheWriteTokens: number;
+ } | null;
+} | null {
+ const db = getDatabase();
+ const rows = db
+ .query("SELECT data_json FROM chunks WHERE tab_id = $tabId AND type = 'usage' ORDER BY seq ASC")
+ .all({ $tabId: tabId }) as Array<{ data_json: string }>;
+ if (rows.length === 0) return null;
+
+ let inputTokens = 0;
+ let outputTokens = 0;
+ let cacheReadTokens = 0;
+ let cacheWriteTokens = 0;
+ let last: UsageData | null = null;
+ for (const row of rows) {
+ let u: UsageData;
+ try {
+ u = JSON.parse(row.data_json) as UsageData;
+ } catch {
+ continue;
+ }
+ inputTokens += u.inputTokens ?? 0;
+ outputTokens += u.outputTokens ?? 0;
+ cacheReadTokens += u.cacheReadTokens ?? 0;
+ cacheWriteTokens += u.cacheWriteTokens ?? 0;
+ last = {
+ inputTokens: u.inputTokens ?? 0,
+ outputTokens: u.outputTokens ?? 0,
+ cacheReadTokens: u.cacheReadTokens ?? 0,
+ cacheWriteTokens: u.cacheWriteTokens ?? 0,
+ };
+ }
+
+ return {
+ inputTokens,
+ outputTokens,
+ cacheReadTokens,
+ cacheWriteTokens,
+ requests: rows.length,
+ last,
+ };
+}
+
export function clearChunksForTab(tabId: string): void {
const db = getDatabase();
db.query("DELETE FROM chunks WHERE tab_id = $tabId").run({ $tabId: tabId });
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 327b0a5..47f9218 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -37,6 +37,7 @@ export {
getChunksForTab,
getMessagesForTab,
getTotalChunkCount,
+ getUsageStatsForTab,
groupRowsToMessages,
type MessageRow,
} from "./db/chunks.js";
diff --git a/packages/core/src/types/index.ts b/packages/core/src/types/index.ts
index ced3dc2..abade27 100644
--- a/packages/core/src/types/index.ts
+++ b/packages/core/src/types/index.ts
@@ -90,7 +90,14 @@ export interface ChatMessage {
export type ChunkRole = "user" | "assistant" | "tool" | "system";
/** Discriminator for a persisted chunk row's payload. */
-export type ChunkType = "text" | "thinking" | "tool_call" | "tool_result" | "error" | "system";
+export type ChunkType =
+ | "text"
+ | "thinking"
+ | "tool_call"
+ | "tool_result"
+ | "error"
+ | "system"
+ | "usage";
export interface TextData {
text: string;
@@ -119,6 +126,22 @@ export interface SystemData {
kind: SystemChunkKind;
text: string;
}
+/**
+ * Per-request token usage persisted as a SIDE-CHANNEL chunk row (one row per
+ * `usage` AgentEvent, i.e. one per LLM round-trip). These rows are deliberately
+ * EXCLUDED from `getChunksForTab`/`getTotalChunkCount` so they never enter the
+ * render, pagination, eviction, or agent-history-rebuild paths — they exist
+ * only to feed the backend aggregate `getUsageStatsForTab`, which seeds the
+ * frontend's `cacheStats` on reload. `inputTokens` is the TOTAL prompt
+ * (cached + fresh); `cacheReadTokens`/`cacheWriteTokens` are Anthropic's
+ * prompt-cache split. Mirrors the `usage` AgentEvent payload.
+ */
+export interface UsageData {
+ inputTokens: number;
+ outputTokens: number;
+ cacheReadTokens: number;
+ cacheWriteTokens: number;
+}
export type ChunkData =
| TextData
@@ -126,7 +149,8 @@ export type ChunkData =
| ToolCallData
| ToolResultData
| ErrorData
- | SystemData;
+ | SystemData
+ | UsageData;
/**
* A persisted chunk row — the append-only unit of conversation storage and
diff --git a/packages/core/tests/db/chunks.db.test.ts b/packages/core/tests/db/chunks.db.test.ts
new file mode 100644
index 0000000..4f7d517
--- /dev/null
+++ b/packages/core/tests/db/chunks.db.test.ts
@@ -0,0 +1,336 @@
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import type { ChunkRowDraft, UsageData } from "../../src/types/index.js";
+
+/**
+ * Internal row shape — matches the production `chunks` table columns.
+ * Kept loose at the `query()` boundary to mirror bun:sqlite's dynamic
+ * return type.
+ */
+interface ChunkRecord {
+ id: string;
+ tab_id: string;
+ seq: number;
+ turn_id: string;
+ step: number;
+ role: string;
+ type: string;
+ data_json: string;
+ created_at: number;
+}
+
+/**
+ * In-memory fake of `bun:sqlite`'s Database implementing only the queries
+ * `chunks.ts` actually issues. Same approach as `tabs.test.ts`: match exact
+ * normalized query strings as fixed branches (no SQL parser), so a query-string
+ * change fails loudly as "unsupported" instead of silently returning wrong data.
+ *
+ * This lets the DB-backed `getChunksForTab` / `getTotalChunkCount` /
+ * `getUsageStatsForTab` logic run under vitest, where `bun:sqlite` can't load.
+ */
+class FakeDatabase {
+ rows: ChunkRecord[] = [];
+ private idCounter = 0;
+
+ query(sql: string): {
+ all: (params?: Record<string, unknown>) => unknown[];
+ get: (params?: Record<string, unknown>) => unknown;
+ run: (params?: Record<string, unknown>) => void;
+ } {
+ return {
+ all: (params) => this.execSelect(sql, params),
+ get: (params) => this.execSelect(sql, params)[0] ?? null,
+ run: (params) => {
+ this.execMutation(sql, params);
+ },
+ };
+ }
+
+ /** bun:sqlite's `db.transaction(fn)` returns a callable that runs `fn`. */
+ transaction(fn: () => void): () => void {
+ return () => {
+ fn();
+ };
+ }
+
+ private execSelect(sql: string, params?: Record<string, unknown>): unknown[] {
+ const norm = sql.replace(/\s+/g, " ").trim();
+ const tabId = params?.$tabId as string | undefined;
+ const forTab = this.rows.filter((r) => r.tab_id === tabId);
+ const visible = forTab.filter((r) => r.type !== "usage");
+
+ // appendChunks: next-seq lookup (counts ALL rows, incl. usage)
+ if (norm === "SELECT COALESCE(MAX(seq), -1) as max_seq FROM chunks WHERE tab_id = $tabId") {
+ const seqs = forTab.map((r) => r.seq);
+ return [{ max_seq: seqs.length > 0 ? Math.max(...seqs) : -1 }];
+ }
+
+ // getChunksForTab — no options (usage excluded)
+ if (
+ norm === "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' ORDER BY seq ASC"
+ ) {
+ return [...visible].sort((a, b) => a.seq - b.seq);
+ }
+
+ // getChunksForTab — before + limit (usage excluded)
+ if (
+ norm ===
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' AND seq < $before ORDER BY seq DESC LIMIT $limit"
+ ) {
+ const before = params?.$before as number;
+ const limit = params?.$limit as number;
+ return visible
+ .filter((r) => r.seq < before)
+ .sort((a, b) => b.seq - a.seq)
+ .slice(0, limit);
+ }
+
+ // getChunksForTab — before only (usage excluded)
+ if (
+ norm ===
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' AND seq < $before ORDER BY seq DESC"
+ ) {
+ const before = params?.$before as number;
+ return visible.filter((r) => r.seq < before).sort((a, b) => b.seq - a.seq);
+ }
+
+ // getChunksForTab — limit only (usage excluded)
+ if (
+ norm ===
+ "SELECT * FROM chunks WHERE tab_id = $tabId AND type != 'usage' ORDER BY seq DESC LIMIT $limit"
+ ) {
+ const limit = params?.$limit as number;
+ return [...visible].sort((a, b) => b.seq - a.seq).slice(0, limit);
+ }
+
+ // getTotalChunkCount (usage excluded)
+ if (norm === "SELECT COUNT(*) as count FROM chunks WHERE tab_id = $tabId AND type != 'usage'") {
+ return [{ count: visible.length }];
+ }
+
+ // getUsageStatsForTab: usage rows only, in seq order
+ if (
+ norm ===
+ "SELECT data_json FROM chunks WHERE tab_id = $tabId AND type = 'usage' ORDER BY seq ASC"
+ ) {
+ return forTab
+ .filter((r) => r.type === "usage")
+ .sort((a, b) => a.seq - b.seq)
+ .map((r) => ({ data_json: r.data_json }));
+ }
+
+ throw new Error(`FakeDatabase: unsupported SELECT: ${norm}`);
+ }
+
+ private execMutation(sql: string, params?: Record<string, unknown>): void {
+ const norm = sql.replace(/\s+/g, " ").trim();
+
+ // appendChunks: single-row insert
+ if (
+ norm ===
+ "INSERT INTO chunks (id, tab_id, seq, turn_id, step, role, type, data_json, created_at) VALUES ($id, $tabId, $seq, $turnId, $step, $role, $type, $dataJson, $now)"
+ ) {
+ this.rows.push({
+ id: (params?.$id as string) ?? `c${this.idCounter++}`,
+ tab_id: params?.$tabId as string,
+ seq: params?.$seq as number,
+ turn_id: params?.$turnId as string,
+ step: (params?.$step as number) ?? 0,
+ role: params?.$role as string,
+ type: params?.$type as string,
+ data_json: params?.$dataJson as string,
+ created_at: (params?.$now as number) ?? 0,
+ });
+ return;
+ }
+
+ throw new Error(`FakeDatabase: unsupported mutation: ${norm}`);
+ }
+}
+
+let fakeDb: FakeDatabase;
+
+vi.mock("../../src/db/index.js", () => ({
+ getDatabase: vi.fn(() => fakeDb),
+}));
+
+const { appendChunks, getChunksForTab, getTotalChunkCount, getUsageStatsForTab } = await import(
+ "../../src/db/chunks.js"
+);
+
+function usageDraft(turnId: string, u: UsageData): ChunkRowDraft {
+ return { turnId, step: 0, role: "assistant", type: "usage", data: u };
+}
+
+beforeAll(() => {
+ fakeDb = new FakeDatabase();
+});
+
+beforeEach(() => {
+ fakeDb.rows = [];
+});
+
+// ---------------------------------------------------------------------------
+// usage chunk persistence + side-channel invariants
+// ---------------------------------------------------------------------------
+describe("usage chunk rows (DB-backed)", () => {
+ const TAB = "tab-usage";
+
+ it("persists usage rows alongside content rows with contiguous seqs", () => {
+ appendChunks(TAB, [
+ { turnId: "t1", step: 0, role: "user", type: "text", data: { text: "hi" } },
+ { turnId: "t1", step: 0, role: "assistant", type: "text", data: { text: "yo" } },
+ usageDraft("t1", {
+ inputTokens: 100,
+ outputTokens: 10,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 90,
+ }),
+ ]);
+ // All three rows landed with contiguous seqs.
+ expect(fakeDb.rows.map((r) => r.seq)).toEqual([0, 1, 2]);
+ expect(fakeDb.rows.map((r) => r.type)).toEqual(["text", "text", "usage"]);
+ });
+
+ it("excludes usage rows from getChunksForTab (all variants)", () => {
+ appendChunks(TAB, [
+ { turnId: "t1", step: 0, role: "user", type: "text", data: { text: "q" } },
+ usageDraft("t1", {
+ inputTokens: 100,
+ outputTokens: 10,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 90,
+ }),
+ { turnId: "t1", step: 0, role: "assistant", type: "text", data: { text: "a" } },
+ usageDraft("t1", {
+ inputTokens: 200,
+ outputTokens: 20,
+ cacheReadTokens: 150,
+ cacheWriteTokens: 0,
+ }),
+ ]);
+
+ // no options
+ const all = getChunksForTab(TAB);
+ expect(all.every((r) => r.type !== "usage")).toBe(true);
+ expect(all.map((r) => r.type)).toEqual(["text", "text"]);
+
+ // limit only
+ const limited = getChunksForTab(TAB, { limit: 10 });
+ expect(limited.every((r) => r.type !== "usage")).toBe(true);
+ expect(limited).toHaveLength(2);
+
+ // before only — `before` is a seq cursor; usage seqs must never surface
+ const before = getChunksForTab(TAB, { before: 100 });
+ expect(before.every((r) => r.type !== "usage")).toBe(true);
+ expect(before).toHaveLength(2);
+
+ // before + limit
+ const bl = getChunksForTab(TAB, { before: 100, limit: 10 });
+ expect(bl.every((r) => r.type !== "usage")).toBe(true);
+ expect(bl).toHaveLength(2);
+ });
+
+ it("excludes usage rows from getTotalChunkCount", () => {
+ appendChunks(TAB, [
+ { turnId: "t1", step: 0, role: "user", type: "text", data: { text: "q" } },
+ { turnId: "t1", step: 0, role: "assistant", type: "text", data: { text: "a" } },
+ usageDraft("t1", {
+ inputTokens: 100,
+ outputTokens: 10,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 90,
+ }),
+ ]);
+ // 3 rows total, but only 2 visible.
+ expect(getTotalChunkCount(TAB)).toBe(2);
+ });
+});
+
+// ---------------------------------------------------------------------------
+// getUsageStatsForTab — backend aggregate
+// ---------------------------------------------------------------------------
+describe("getUsageStatsForTab", () => {
+ const TAB = "tab-agg";
+
+ it("returns null when the tab has no usage rows", () => {
+ appendChunks(TAB, [
+ { turnId: "t1", step: 0, role: "assistant", type: "text", data: { text: "a" } },
+ ]);
+ expect(getUsageStatsForTab(TAB)).toBeNull();
+ });
+
+ it("sums cumulative tokens, counts requests, and reports the last request's split", () => {
+ appendChunks(TAB, [
+ usageDraft("t1", {
+ inputTokens: 1000,
+ outputTokens: 40,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 900,
+ }),
+ usageDraft("t1", {
+ inputTokens: 1200,
+ outputTokens: 60,
+ cacheReadTokens: 1000,
+ cacheWriteTokens: 100,
+ }),
+ ]);
+
+ const stats = getUsageStatsForTab(TAB);
+ expect(stats).not.toBeNull();
+ expect(stats?.requests).toBe(2);
+ expect(stats?.inputTokens).toBe(2200);
+ expect(stats?.outputTokens).toBe(100);
+ expect(stats?.cacheReadTokens).toBe(1000);
+ expect(stats?.cacheWriteTokens).toBe(1000);
+ // `last` = the most recent (highest-seq) usage row.
+ expect(stats?.last).toEqual({
+ inputTokens: 1200,
+ outputTokens: 60,
+ cacheReadTokens: 1000,
+ cacheWriteTokens: 100,
+ });
+ });
+
+ it("is structurally identical to the frontend CacheStats shape (seeds directly)", () => {
+ appendChunks(TAB, [
+ usageDraft("t1", {
+ inputTokens: 5,
+ outputTokens: 1,
+ cacheReadTokens: 2,
+ cacheWriteTokens: 3,
+ }),
+ ]);
+ const stats = getUsageStatsForTab(TAB);
+ expect(Object.keys(stats ?? {}).sort()).toEqual(
+ [
+ "cacheReadTokens",
+ "cacheWriteTokens",
+ "inputTokens",
+ "last",
+ "outputTokens",
+ "requests",
+ ].sort(),
+ );
+ });
+
+ it("is scoped per tab", () => {
+ appendChunks("tab-a", [
+ usageDraft("t1", {
+ inputTokens: 10,
+ outputTokens: 1,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 0,
+ }),
+ ]);
+ appendChunks("tab-b", [
+ usageDraft("t2", {
+ inputTokens: 20,
+ outputTokens: 2,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 0,
+ }),
+ ]);
+ expect(getUsageStatsForTab("tab-a")?.inputTokens).toBe(10);
+ expect(getUsageStatsForTab("tab-b")?.inputTokens).toBe(20);
+ });
+});