diff options
| author | Adam Malczewski <[email protected]> | 2026-06-02 13:08:36 +0900 |
|---|---|---|
| committer | Adam Malczewski <[email protected]> | 2026-06-02 13:08:36 +0900 |
| commit | b734eb96bf0af267fdfbef85df51940ca0b4e8c7 (patch) | |
| tree | 3815dc9e569bb57384f53d9042288ff831f02e74 /packages/core/src | |
| parent | 3f629a8469fe483243671e1ca15582a111e96541 (diff) | |
| download | dispatch-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/src')
| -rw-r--r-- | packages/core/src/chunks/transform.ts | 5 | ||||
| -rw-r--r-- | packages/core/src/db/chunks.ts | 90 | ||||
| -rw-r--r-- | packages/core/src/index.ts | 1 | ||||
| -rw-r--r-- | packages/core/src/types/index.ts | 28 |
4 files changed, 115 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 |
