Skip to content

Commit f403fc6

Browse files
committed
fix(storage): SqlStorageAdapter persists + aggregates cache tokens
Closes the cache-token propagation sweep by migrating the SQL-backed storage layer to match the InMemory / ITokenUsage / UsageLedger work. Before this, consumers on the SQL adapter (persistent conversations across process restarts) saw the same 0/undefined cache fields even though the API, strategies, and ledger surfaces all now forward them. Changes to SqlStorageAdapter: - messages schema gains cacheReadTokens + cacheCreationTokens INTEGER columns. Fresh DBs get them via CREATE TABLE IF NOT EXISTS; existing DBs via an idempotent ALTER TABLE ADD COLUMN wrapped in try/catch for the duplicate-column case (SQLite does not support IF NOT EXISTS on ADD COLUMN until 3.35; Postgres does but the catch is harmless on success). - Message INSERT path writes the cache counters when the caller supplies them, null when undefined — preserves the "not reported" vs "zero hits" signal. - getConversationTokenUsage now SUMs cacheReadTokens + cacheCreationTokens and uses COUNT(col) to detect whether any message actually reported each value, only setting the optional return fields when at least one did. - rowToMessage hydrates cache fields back onto the message's usage when the row carries them. Typecheck clean. 278/278 core tests still pass (no adapter tests exist to extend yet; storage adapter integration is covered by the broader conversation-manager flows).
1 parent a392103 commit f403fc6

1 file changed

Lines changed: 58 additions & 8 deletions

File tree

src/core/storage/SqlStorageAdapter.ts

Lines changed: 58 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,8 @@ export class SqlStorageAdapter implements IStorageAdapter {
186186
promptTokens INTEGER,
187187
completionTokens INTEGER,
188188
totalTokens INTEGER,
189+
cacheReadTokens INTEGER,
190+
cacheCreationTokens INTEGER,
189191
toolCalls TEXT,
190192
toolCallId TEXT,
191193
name TEXT,
@@ -197,6 +199,22 @@ export class SqlStorageAdapter implements IStorageAdapter {
197199
CREATE INDEX IF NOT EXISTS idx_messages_role ON messages(role);
198200
`);
199201

202+
// Back-compat migration: add cacheReadTokens / cacheCreationTokens
203+
// columns on pre-existing databases. CREATE TABLE IF NOT EXISTS
204+
// above only helps on fresh DBs; on DBs that predate this field
205+
// set the table already exists and the new columns never land
206+
// without an explicit ALTER. Swallow the duplicate-column error
207+
// because neither SQLite nor Postgres has portable IF NOT EXISTS
208+
// on ADD COLUMN (Postgres does; SQLite does not until 3.35).
209+
for (const col of ['cacheReadTokens', 'cacheCreationTokens']) {
210+
try {
211+
await this.adapter.exec(`ALTER TABLE messages ADD COLUMN ${col} INTEGER;`);
212+
} catch {
213+
// Column already exists — expected on every startup after the
214+
// first post-migration run. No-op.
215+
}
216+
}
217+
200218
this.initialized = true;
201219
}
202220

@@ -384,9 +402,9 @@ export class SqlStorageAdapter implements IStorageAdapter {
384402

385403
// Insert message
386404
await this.adapter.run(
387-
`INSERT INTO messages
388-
(id, conversationId, role, content, timestamp, model, promptTokens, completionTokens, totalTokens, toolCalls, toolCallId, name, metadata)
389-
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
405+
`INSERT INTO messages
406+
(id, conversationId, role, content, timestamp, model, promptTokens, completionTokens, totalTokens, cacheReadTokens, cacheCreationTokens, toolCalls, toolCallId, name, metadata)
407+
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
390408
[
391409
message.id,
392410
message.conversationId,
@@ -397,6 +415,12 @@ export class SqlStorageAdapter implements IStorageAdapter {
397415
message.usage?.promptTokens || null,
398416
message.usage?.completionTokens || null,
399417
message.usage?.totalTokens || null,
418+
// Persist cache-token counters when the caller reports them so
419+
// getConversationTokenUsage can SUM a full cache picture later.
420+
// Null when undefined — distinguishes "provider did not report"
421+
// from a genuine zero.
422+
message.usage?.cacheReadTokens ?? null,
423+
message.usage?.cacheCreationTokens ?? null,
400424
toolCallsJson,
401425
message.toolCallId || null,
402426
message.name || null,
@@ -538,20 +562,36 @@ export class SqlStorageAdapter implements IStorageAdapter {
538562
this.ensureInitialized();
539563

540564
const row = await this.adapter.get<any>(
541-
`SELECT
565+
`SELECT
542566
SUM(promptTokens) as promptTokens,
543567
SUM(completionTokens) as completionTokens,
544-
SUM(totalTokens) as totalTokens
545-
FROM messages
568+
SUM(totalTokens) as totalTokens,
569+
SUM(cacheReadTokens) as cacheReadTokens,
570+
SUM(cacheCreationTokens) as cacheCreationTokens,
571+
COUNT(cacheReadTokens) as cacheReadCount,
572+
COUNT(cacheCreationTokens) as cacheCreationCount
573+
FROM messages
546574
WHERE conversationId = ?`,
547575
[conversationId]
548576
);
549577

550-
return {
578+
const usage: ITokenUsage = {
551579
promptTokens: row?.promptTokens || 0,
552580
completionTokens: row?.completionTokens || 0,
553581
totalTokens: row?.totalTokens || 0,
554582
};
583+
// COUNT returns the number of non-NULL rows; when nothing recorded
584+
// cache metrics, SUM collapses to NULL / 0 but the count is also 0.
585+
// We only populate the optional cache fields when at least one
586+
// message carried a value, preserving the "not reported" signal
587+
// on cache-silent providers (e.g. OpenAI).
588+
if (row?.cacheReadCount && Number(row.cacheReadCount) > 0) {
589+
usage.cacheReadTokens = Number(row.cacheReadTokens) || 0;
590+
}
591+
if (row?.cacheCreationCount && Number(row.cacheCreationCount) > 0) {
592+
usage.cacheCreationTokens = Number(row.cacheCreationTokens) || 0;
593+
}
594+
return usage;
555595
}
556596

557597
// ==================== Private Helper Methods ====================
@@ -604,12 +644,22 @@ export class SqlStorageAdapter implements IStorageAdapter {
604644
};
605645

606646
if (row.model) message.model = row.model;
607-
if (row.promptTokens || row.completionTokens || row.totalTokens) {
647+
if (
648+
row.promptTokens ||
649+
row.completionTokens ||
650+
row.totalTokens ||
651+
row.cacheReadTokens != null ||
652+
row.cacheCreationTokens != null
653+
) {
608654
message.usage = {
609655
promptTokens: row.promptTokens || 0,
610656
completionTokens: row.completionTokens || 0,
611657
totalTokens: row.totalTokens || 0,
612658
};
659+
// Only set cache fields when the row actually stored a value so
660+
// callers retain "not reported" vs "zero hits" distinction.
661+
if (row.cacheReadTokens != null) message.usage.cacheReadTokens = row.cacheReadTokens;
662+
if (row.cacheCreationTokens != null) message.usage.cacheCreationTokens = row.cacheCreationTokens;
613663
}
614664
if (row.toolCalls) message.toolCalls = JSON.parse(row.toolCalls);
615665
if (row.toolCallId) message.toolCallId = row.toolCallId;

0 commit comments

Comments
 (0)