From b2a9317ed3375d6958562abde194663d5eeaeee2 Mon Sep 17 00:00:00 2001 From: Radoslaw Kubiak Date: Thu, 2 Apr 2026 18:38:25 +0200 Subject: [PATCH 1/2] feat(memory): implement token usage tracking (DC-MEM-005) - #94 --- .../src/db/migrations/008_token_usage.ts | 35 +++ packages/memory/src/db/migrations/registry.ts | 2 + .../db/repositories/TokenUsageRepository.ts | 243 ++++++++++++++++++ packages/memory/src/db/schemas/token-usage.ts | 55 ++++ packages/memory/src/index.ts | 8 + 5 files changed, 343 insertions(+) create mode 100644 packages/memory/src/db/migrations/008_token_usage.ts create mode 100644 packages/memory/src/db/repositories/TokenUsageRepository.ts create mode 100644 packages/memory/src/db/schemas/token-usage.ts diff --git a/packages/memory/src/db/migrations/008_token_usage.ts b/packages/memory/src/db/migrations/008_token_usage.ts new file mode 100644 index 0000000..70a3d99 --- /dev/null +++ b/packages/memory/src/db/migrations/008_token_usage.ts @@ -0,0 +1,35 @@ +import type { Database } from "better-sqlite3"; +import type { Migration } from "./runner.js"; + +export const migration008: Migration = { + version: 8, + description: "token usage tracking", + up(db: Database): void { + db.exec(` + CREATE TABLE IF NOT EXISTS token_usage ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + turn_id TEXT, + agent_id TEXT, + model TEXT NOT NULL, + provider TEXT, + tokens_in INTEGER NOT NULL DEFAULT 0, + tokens_out INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0.0, + timestamp TEXT NOT NULL DEFAULT (datetime('now')) + ); + + CREATE INDEX IF NOT EXISTS idx_token_usage_session_time + ON token_usage(session_id, timestamp); + + CREATE INDEX IF NOT EXISTS idx_token_usage_turn + ON token_usage(turn_id); + + CREATE INDEX IF NOT EXISTS idx_token_usage_model + ON token_usage(model); + + CREATE INDEX IF NOT EXISTS idx_token_usage_agent + ON token_usage(agent_id); + `); + }, +}; diff --git a/packages/memory/src/db/migrations/registry.ts b/packages/memory/src/db/migrations/registry.ts index c53a742..ea7d9e9 100644 --- a/packages/memory/src/db/migrations/registry.ts +++ b/packages/memory/src/db/migrations/registry.ts @@ -6,6 +6,7 @@ import { migration004 } from "./004_background_tasks.js"; import { migration005 } from "./005_sessions_and_messages.js"; import { migration006 } from "./006_fts5_search.js"; import { migration007 } from "./007_turns.js"; +import { migration008 } from "./008_token_usage.js"; export const migrations: Migration[] = [ migration001, @@ -15,4 +16,5 @@ export const migrations: Migration[] = [ migration005, migration006, migration007, + migration008, ]; diff --git a/packages/memory/src/db/repositories/TokenUsageRepository.ts b/packages/memory/src/db/repositories/TokenUsageRepository.ts new file mode 100644 index 0000000..0ff0480 --- /dev/null +++ b/packages/memory/src/db/repositories/TokenUsageRepository.ts @@ -0,0 +1,243 @@ +import type { Database } from "better-sqlite3"; +import { + type TokenUsage, + type RecordTokenUsageInput, + type SessionUsageSummary, + type ModelUsageBreakdown, + type AgentUsageSummary, + RecordTokenUsageInputSchema, + TokenUsageSchema, + SessionUsageSummarySchema, + ModelUsageBreakdownSchema, + AgentUsageSummarySchema, +} from "../schemas/token-usage.js"; + +interface TokenUsageRow { + id: string; + session_id: string; + turn_id: string | null; + agent_id: string | null; + model: string; + provider: string | null; + tokens_in: number; + tokens_out: number; + cost_usd: number; + timestamp: string; +} + +interface SummaryRow { + session_id: string; + total_tokens_in: number; + total_tokens_out: number; + total_cost_usd: number; + record_count: number; +} + +interface ModelBreakdownRow { + model: string; + total_tokens_in: number; + total_tokens_out: number; + total_cost_usd: number; + record_count: number; +} + +interface AgentSummaryRow { + agent_id: string | null; + total_tokens_in: number; + total_tokens_out: number; + total_cost_usd: number; + record_count: number; +} + +function rowToTokenUsage(row: TokenUsageRow): TokenUsage { + return { + id: row.id, + sessionId: row.session_id, + turnId: row.turn_id ?? undefined, + agentId: row.agent_id ?? undefined, + model: row.model, + provider: row.provider ?? undefined, + tokensIn: row.tokens_in, + tokensOut: row.tokens_out, + costUsd: row.cost_usd, + timestamp: row.timestamp, + }; +} + +function rowToSessionSummary(row: SummaryRow): SessionUsageSummary { + return { + sessionId: row.session_id, + totalTokensIn: row.total_tokens_in, + totalTokensOut: row.total_tokens_out, + totalCostUsd: row.total_cost_usd, + recordCount: row.record_count, + }; +} + +function rowToModelBreakdown(row: ModelBreakdownRow): ModelUsageBreakdown { + return { + model: row.model, + totalTokensIn: row.total_tokens_in, + totalTokensOut: row.total_tokens_out, + totalCostUsd: row.total_cost_usd, + recordCount: row.record_count, + }; +} + +function rowToAgentSummary(row: AgentSummaryRow): AgentUsageSummary { + return { + agentId: row.agent_id, + totalTokensIn: row.total_tokens_in, + totalTokensOut: row.total_tokens_out, + totalCostUsd: row.total_cost_usd, + recordCount: row.record_count, + }; +} + +export class TokenUsageRepository { + private readonly stmtInsert; + private readonly stmtGetById; + private readonly stmtGetBySession; + private readonly stmtGetByTurn; + private readonly stmtSessionTotals; + private readonly stmtAgentTotals; + private readonly stmtModelBreakdown; + private readonly stmtByTimeRange; + private readonly stmtBudgetCheck; + + constructor(private readonly db: Database) { + this.stmtInsert = db.prepare( + `INSERT INTO token_usage (id, session_id, turn_id, agent_id, model, provider, tokens_in, tokens_out, cost_usd) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING *`, + ); + + this.stmtGetById = db.prepare("SELECT * FROM token_usage WHERE id = ?"); + + this.stmtGetBySession = db.prepare( + "SELECT * FROM token_usage WHERE session_id = ? ORDER BY timestamp ASC", + ); + + this.stmtGetByTurn = db.prepare( + "SELECT * FROM token_usage WHERE turn_id = ? ORDER BY timestamp ASC", + ); + + this.stmtSessionTotals = db.prepare( + `SELECT session_id, + COALESCE(SUM(tokens_in), 0) AS total_tokens_in, + COALESCE(SUM(tokens_out), 0) AS total_tokens_out, + COALESCE(SUM(cost_usd), 0) AS total_cost_usd, + COUNT(*) AS record_count + FROM token_usage + WHERE session_id = ? + GROUP BY session_id`, + ); + + this.stmtAgentTotals = db.prepare( + `SELECT agent_id, + COALESCE(SUM(tokens_in), 0) AS total_tokens_in, + COALESCE(SUM(tokens_out), 0) AS total_tokens_out, + COALESCE(SUM(cost_usd), 0) AS total_cost_usd, + COUNT(*) AS record_count + FROM token_usage + WHERE session_id = ? AND agent_id = ? + GROUP BY agent_id`, + ); + + this.stmtModelBreakdown = db.prepare( + `SELECT model, + COALESCE(SUM(tokens_in), 0) AS total_tokens_in, + COALESCE(SUM(tokens_out), 0) AS total_tokens_out, + COALESCE(SUM(cost_usd), 0) AS total_cost_usd, + COUNT(*) AS record_count + FROM token_usage + WHERE session_id = ? + GROUP BY model + ORDER BY total_cost_usd DESC`, + ); + + this.stmtByTimeRange = db.prepare( + `SELECT * FROM token_usage + WHERE timestamp >= ? AND timestamp <= ? + ORDER BY timestamp ASC`, + ); + + this.stmtBudgetCheck = db.prepare( + `SELECT COALESCE(SUM(cost_usd), 0) AS total_cost_usd + FROM token_usage + WHERE session_id = ?`, + ); + } + + record(input: RecordTokenUsageInput): TokenUsage { + const parsed = RecordTokenUsageInputSchema.parse(input); + const row = (this.stmtInsert as { get(...args: unknown[]): TokenUsageRow | undefined }).get( + parsed.id, + parsed.sessionId, + parsed.turnId ?? null, + parsed.agentId ?? null, + parsed.model, + parsed.provider ?? null, + parsed.tokensIn, + parsed.tokensOut, + parsed.costUsd, + ); + if (row === undefined) throw new Error(`TokenUsage ${parsed.id} not found after insert`); + return TokenUsageSchema.parse(rowToTokenUsage(row)); + } + + getById(id: string): TokenUsage | undefined { + const row = (this.stmtGetById as { get(id: string): TokenUsageRow | undefined }).get(id); + if (!row) return undefined; + return TokenUsageSchema.parse(rowToTokenUsage(row)); + } + + getBySessionId(sessionId: string): TokenUsage[] { + return (this.stmtGetBySession as { all(sessionId: string): TokenUsageRow[] }) + .all(sessionId) + .map((row) => TokenUsageSchema.parse(rowToTokenUsage(row))); + } + + getByTurnId(turnId: string): TokenUsage[] { + return (this.stmtGetByTurn as { all(turnId: string): TokenUsageRow[] }) + .all(turnId) + .map((row) => TokenUsageSchema.parse(rowToTokenUsage(row))); + } + + getSessionTotals(sessionId: string): SessionUsageSummary | undefined { + const row = (this.stmtSessionTotals as { get(sessionId: string): SummaryRow | undefined }).get( + sessionId, + ); + if (!row) return undefined; + return SessionUsageSummarySchema.parse(rowToSessionSummary(row)); + } + + getAgentTotals(sessionId: string, agentId: string): AgentUsageSummary | undefined { + const row = ( + this.stmtAgentTotals as { + get(sessionId: string, agentId: string): AgentSummaryRow | undefined; + } + ).get(sessionId, agentId); + if (!row) return undefined; + return AgentUsageSummarySchema.parse(rowToAgentSummary(row)); + } + + getModelBreakdown(sessionId: string): ModelUsageBreakdown[] { + return (this.stmtModelBreakdown as { all(sessionId: string): ModelBreakdownRow[] }) + .all(sessionId) + .map((row) => ModelUsageBreakdownSchema.parse(rowToModelBreakdown(row))); + } + + getByTimeRange(from: string, to: string): TokenUsage[] { + return (this.stmtByTimeRange as { all(from: string, to: string): TokenUsageRow[] }) + .all(from, to) + .map((row) => TokenUsageSchema.parse(rowToTokenUsage(row))); + } + + checkBudget(sessionId: string, threshold: number): boolean { + const row = ( + this.stmtBudgetCheck as { get(sessionId: string): { total_cost_usd: number } } + ).get(sessionId); + if (!row) return false; + return row.total_cost_usd >= threshold; + } +} diff --git a/packages/memory/src/db/schemas/token-usage.ts b/packages/memory/src/db/schemas/token-usage.ts new file mode 100644 index 0000000..9ab2fda --- /dev/null +++ b/packages/memory/src/db/schemas/token-usage.ts @@ -0,0 +1,55 @@ +import { z } from "zod"; + +export const TokenUsageSchema = z.object({ + id: z.string().min(1), + sessionId: z.string().min(1), + turnId: z.string().nullable().optional(), + agentId: z.string().nullable().optional(), + model: z.string().min(1), + provider: z.string().nullable().optional(), + tokensIn: z.number().int().nonnegative().default(0), + tokensOut: z.number().int().nonnegative().default(0), + costUsd: z.number().nonnegative().default(0), + timestamp: z.string(), +}); +export type TokenUsage = z.infer; + +export const RecordTokenUsageInputSchema = z.object({ + id: z.string().min(1), + sessionId: z.string().min(1), + turnId: z.string().optional(), + agentId: z.string().optional(), + model: z.string().min(1), + provider: z.string().optional(), + tokensIn: z.number().int().nonnegative().default(0), + tokensOut: z.number().int().nonnegative().default(0), + costUsd: z.number().nonnegative().default(0), +}); +export type RecordTokenUsageInput = z.infer; + +export const SessionUsageSummarySchema = z.object({ + sessionId: z.string(), + totalTokensIn: z.number().int().nonnegative(), + totalTokensOut: z.number().int().nonnegative(), + totalCostUsd: z.number().nonnegative(), + recordCount: z.number().int().nonnegative(), +}); +export type SessionUsageSummary = z.infer; + +export const ModelUsageBreakdownSchema = z.object({ + model: z.string(), + totalTokensIn: z.number().int().nonnegative(), + totalTokensOut: z.number().int().nonnegative(), + totalCostUsd: z.number().nonnegative(), + recordCount: z.number().int().nonnegative(), +}); +export type ModelUsageBreakdown = z.infer; + +export const AgentUsageSummarySchema = z.object({ + agentId: z.string().nullable(), + totalTokensIn: z.number().int().nonnegative(), + totalTokensOut: z.number().int().nonnegative(), + totalCostUsd: z.number().nonnegative(), + recordCount: z.number().int().nonnegative(), +}); +export type AgentUsageSummary = z.infer; diff --git a/packages/memory/src/index.ts b/packages/memory/src/index.ts index 1202d6e..da5e873 100644 --- a/packages/memory/src/index.ts +++ b/packages/memory/src/index.ts @@ -56,3 +56,11 @@ export type { TurnTelemetry, TurnPartialResult, } from "./db/schemas/turn.js"; +export { TokenUsageRepository } from "./db/repositories/TokenUsageRepository.js"; +export type { + TokenUsage, + RecordTokenUsageInput, + SessionUsageSummary, + ModelUsageBreakdown, + AgentUsageSummary, +} from "./db/schemas/token-usage.js"; From 8aa4c0615ba0b45b862174832c4b9df5900eae82 Mon Sep 17 00:00:00 2001 From: Radoslaw Kubiak Date: Thu, 2 Apr 2026 18:38:56 +0200 Subject: [PATCH 2/2] fix(memory): remove unnecessary null check in checkBudget - #94 --- packages/memory/src/db/repositories/TokenUsageRepository.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/memory/src/db/repositories/TokenUsageRepository.ts b/packages/memory/src/db/repositories/TokenUsageRepository.ts index 0ff0480..c560682 100644 --- a/packages/memory/src/db/repositories/TokenUsageRepository.ts +++ b/packages/memory/src/db/repositories/TokenUsageRepository.ts @@ -237,7 +237,6 @@ export class TokenUsageRepository { const row = ( this.stmtBudgetCheck as { get(sessionId: string): { total_cost_usd: number } } ).get(sessionId); - if (!row) return false; return row.total_cost_usd >= threshold; } }