Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion core/acp/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ export interface AcpPromptRequest {
}

export interface AcpUpdate {
kind: 'text' | 'tool_call' | 'tool_result' | 'done' | 'error'
kind: 'text' | 'tool_call' | 'tool_result' | 'usage_update' | 'done' | 'error'
data: unknown
}

Expand Down
5 changes: 5 additions & 0 deletions core/acp/stdio-client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ describe('mapSessionUpdate', () => {
expect(mapSessionUpdate(raw as never)).toEqual({ kind: 'tool_result', data: raw })
})

it('maps usage_update to a usage_update update', () => {
const raw = { sessionUpdate: 'usage_update', used: 1500, size: 200000 }
expect(mapSessionUpdate(raw as never)).toEqual({ kind: 'usage_update', data: raw })
})

it('falls back to tool_result for update kinds with no dedicated AcpUpdate kind', () => {
const raw = { sessionUpdate: 'plan', entries: [] }
expect(mapSessionUpdate(raw as never)).toEqual({ kind: 'tool_result', data: raw })
Expand Down
6 changes: 4 additions & 2 deletions core/acp/stdio-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,9 @@ export function mapSessionUpdate(update: SessionUpdate): AcpUpdate {
return { kind: 'tool_call', data: update }
case 'tool_call_update':
return { kind: 'tool_result', data: update }
// plan / mode / config / usage / command updates don't have a bespoke
case 'usage_update':
return { kind: 'usage_update', data: update }
// plan / mode / config / command updates don't have a bespoke
// AcpUpdate kind yet (contract only defines text/tool_call/tool_result/
// done/error). Surface them as tool_result so nothing is silently
// dropped; revisit the contract if a UI needs to key off these directly.
Expand Down Expand Up @@ -338,7 +340,7 @@ class StdioAcpSession implements AcpSession {
if (settled) return
settled = true
timeout.cancel()
queue.push({ kind: 'done', data: { stopReason: res.stopReason } })
queue.push({ kind: 'done', data: { stopReason: res.stopReason, ...(res.usage ? { usage: res.usage } : {}) } })
queue.close()
})
.catch((err: unknown) => {
Expand Down
56 changes: 56 additions & 0 deletions core/chat/manager.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { upsertPaper } from '../library/repo.js'
import { openDb } from '../store/db.js'
import { ChatManager } from './manager.js'
import { getChatMessages } from './repo.js'
import { listUsageRecords, getUsageSummary } from '../usage/repo.js'

class FakeSession implements AcpSession {
disposed = false
Expand Down Expand Up @@ -235,4 +236,59 @@ describe('ChatManager', () => {
const reopened = manager.openChat({ db, paperSlug: slug, chatSessionId: first.session.id })
expect(reopened.session.id).toBe(first.session.id)
})

it('records turn telemetry when adapter emits usage metadata [L2-05]', async () => {
const usageData = {
inputTokens: 1500,
outputTokens: 300,
totalTokens: 1800,
}
const client = new FakeClient([
textUpdate('response with usage'),
{ kind: 'usage_update', data: { used: 12000, size: 200000, cost: { amount: 0.02 } } },
{ kind: 'done', data: { stopReason: 'end_turn', usage: usageData } },
])
const manager = new ChatManager(client)
const { session } = manager.openChat({ db, paperSlug: slug, backend: 'claude' })

await manager.runTurn({ db, chatSessionId: session.id, paperSlug: slug, mdPath, text: 'hello' }, () => {})

const records = listUsageRecords(db, { sessionId: session.id })
expect(records).toHaveLength(1)
expect(records[0].hasMetrics).toBe(true)
expect(records[0].inputTokens).toBe(1500)
expect(records[0].outputTokens).toBe(300)
expect(records[0].totalTokens).toBe(1800)
expect(records[0].contextUsed).toBe(12000)
expect(records[0].costAmount).toBe(0.02)

const summary = getUsageSummary(db)
expect(summary.totalTurns).toBe(1)
expect(summary.turnsWithMetrics).toBe(1)
expect(summary.totalTokens).toBe(1800)
})

it('records turn with honest unavailable state when adapter lacks telemetry [L2-05]', async () => {
const client = new FakeClient([
textUpdate('response without usage'),
{ kind: 'done', data: { stopReason: 'end_turn' } },
])
const manager = new ChatManager(client)
const { session } = manager.openChat({ db, paperSlug: slug, backend: 'codex' })

await manager.runTurn({ db, chatSessionId: session.id, paperSlug: slug, mdPath, text: 'hello' }, () => {})

const records = listUsageRecords(db, { sessionId: session.id })
expect(records).toHaveLength(1)
expect(records[0].hasMetrics).toBe(false)
expect(records[0].inputTokens).toBeNull()
expect(records[0].outputTokens).toBeNull()
expect(records[0].totalTokens).toBeNull()

const summary = getUsageSummary(db)
expect(summary.totalTurns).toBe(1)
expect(summary.turnsWithMetrics).toBe(0)
expect(summary.totalTokens).toBeNull() // Honest null, never 0
expect(summary.backends.codex.totalTokens).toBeNull()
})
})
50 changes: 50 additions & 0 deletions core/chat/manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import type { Database } from 'better-sqlite3'
import type { AcpBackend, AcpClient, AcpSession, AcpUpdate } from '../acp/client.js'
import type { ChatMessage, ChatSession } from './repo.js'
import { addChatMessage, createChatSession, getChatMessages, getChatSession, getLatestChatSession } from './repo.js'
import { recordTurnUsage } from '../usage/repo.js'

const DEFAULT_BACKEND: AcpBackend = 'claude'

Expand Down Expand Up @@ -190,6 +191,16 @@ export class ChatManager {
const promptText = buildPromptText(params.mdPath, params.text, freshlySpawned ? priorMessages : [])

let accumulated = ''
let latestUsageUpdate: { used?: number; size?: number; costAmount?: number; costCurrency?: string } | undefined
let promptUsage: {
inputTokens?: number | null
outputTokens?: number | null
thoughtTokens?: number | null
cachedReadTokens?: number | null
cachedWriteTokens?: number | null
totalTokens?: number | null
} | null = null

try {
for await (const update of acpSession.prompt({ text: promptText, contextFiles: [params.mdPath] })) {
if (update.kind === 'text') {
Expand All @@ -200,13 +211,27 @@ export class ChatManager {
}
} else if (update.kind === 'tool_call' || update.kind === 'tool_result') {
onEvent({ kind: 'tool_activity' })
} else if (update.kind === 'usage_update') {
const u = update.data as { used?: number; size?: number; cost?: { amount?: number; currency?: string } } | undefined
if (u) {
latestUsageUpdate = {
used: typeof u.used === 'number' ? u.used : undefined,
size: typeof u.size === 'number' ? u.size : undefined,
costAmount: typeof u.cost?.amount === 'number' ? u.cost.amount : undefined,
costCurrency: u.cost?.currency ?? 'USD',
}
}
} else if (update.kind === 'error') {
// A broken turn likely means a broken session — evict the cache
// so the next turn spawns fresh rather than reusing a dead one.
await this.disposeSession(params.paperSlug, backend)
onEvent({ kind: 'error', message: errorMessageFromUpdate(update) })
return
} else if (update.kind === 'done') {
const doneData = update.data as { usage?: { inputTokens?: number; outputTokens?: number; thoughtTokens?: number; cachedReadTokens?: number; cachedWriteTokens?: number; totalTokens?: number } | null } | undefined
if (doneData?.usage) {
promptUsage = doneData.usage
}
break
}
}
Expand All @@ -216,6 +241,31 @@ export class ChatManager {
return
}

const turnIndex = priorMessages.filter((m) => m.role === 'user').length
const hasMetrics = !!(promptUsage || latestUsageUpdate)
try {
recordTurnUsage(params.db, {
sessionId: params.chatSessionId,
backend,
turnIndex,
inputTokens: promptUsage?.inputTokens ?? null,
outputTokens: promptUsage?.outputTokens ?? null,
thoughtTokens: promptUsage?.thoughtTokens ?? null,
cachedReadTokens: promptUsage?.cachedReadTokens ?? null,
cachedWriteTokens: promptUsage?.cachedWriteTokens ?? null,
totalTokens: promptUsage?.totalTokens ?? (promptUsage?.inputTokens && promptUsage?.outputTokens ? promptUsage.inputTokens + promptUsage.outputTokens : null),
contextUsed: latestUsageUpdate?.used ?? null,
contextSize: latestUsageUpdate?.size ?? null,
costAmount: latestUsageUpdate?.costAmount ?? null,
costCurrency: latestUsageUpdate?.costCurrency ?? 'USD',
hasMetrics,
recordedAt: new Date().toISOString(),
})
} catch {
// Usage recording is non-blocking telemetry; database recording failure
// must never fail the user chat turn.
}

const message = addChatMessage(params.db, {
sessionId: params.chatSessionId,
role: 'assistant',
Expand Down
2 changes: 1 addition & 1 deletion core/store/db.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ describe('openDb', () => {
const row = db2.prepare('SELECT * FROM papers WHERE slug = ?').get('p1')
expect(row).toBeTruthy()
const version = db2.pragma('user_version', { simple: true })
expect(version).toBe(7)
expect(version).toBe(8)
} finally {
db2.close()
}
Expand Down
23 changes: 21 additions & 2 deletions core/store/migrate.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ describe('runMigrations', () => {
runMigrations(db)

const version = db.pragma('user_version', { simple: true })
expect(version).toBe(7)
expect(version).toBe(8)
})

it('adds the author_orcids column to papers [P2-04]', () => {
Expand Down Expand Up @@ -77,7 +77,7 @@ describe('runMigrations', () => {
runMigrations(db)

const version = db.pragma('user_version', { simple: true })
expect(version).toBe(7)
expect(version).toBe(8)

const row = db.prepare('SELECT * FROM papers WHERE slug = ?').get('a')
expect(row).toBeTruthy()
Expand Down Expand Up @@ -114,6 +114,25 @@ describe('runMigrations', () => {
expect(columns).toContain('trashed_at')
})

it('creates the usage_records table and indexes [L2-05]', () => {
db = new Database(':memory:')
runMigrations(db)

const tables = db
.prepare("SELECT name FROM sqlite_master WHERE type = 'table'")
.all()
.map((row) => (row as { name: string }).name)
expect(tables).toContain('usage_records')

const indexes = db
.prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'usage_records'")
.all()
.map((row) => (row as { name: string }).name)
expect(indexes).toContain('idx_usage_records_backend')
expect(indexes).toContain('idx_usage_records_recorded_at')
expect(indexes).toContain('idx_usage_records_session')
})

it('applies only migrations newer than the current version, in order', () => {
db = new Database(':memory:')

Expand Down
3 changes: 2 additions & 1 deletion core/store/migrate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

import type { Database } from 'better-sqlite3'

import { SCHEMA_V1, SCHEMA_V3_NOTES, SCHEMA_V4_HIGHLIGHTS, SCHEMA_V5_AUTHOR_ORCIDS, SCHEMA_V6_SUGGESTED_QUESTIONS, SCHEMA_V7_TRASH } from './schema.js'
import { SCHEMA_V1, SCHEMA_V3_NOTES, SCHEMA_V4_HIGHLIGHTS, SCHEMA_V5_AUTHOR_ORCIDS, SCHEMA_V6_SUGGESTED_QUESTIONS, SCHEMA_V7_TRASH, SCHEMA_V8_USAGE } from './schema.js'

export interface Migration {
version: number
Expand All @@ -28,6 +28,7 @@ export const MIGRATIONS: Migration[] = [
{ version: 5, sql: SCHEMA_V5_AUTHOR_ORCIDS }, // [P2-04] author_orcids column
{ version: 6, sql: SCHEMA_V6_SUGGESTED_QUESTIONS }, // [L2-03] suggested_questions table
{ version: 7, sql: SCHEMA_V7_TRASH }, // [L2-04] trashed_at column
{ version: 8, sql: SCHEMA_V8_USAGE }, // [L2-05] usage_records table
]

/**
Expand Down
29 changes: 29 additions & 0 deletions core/store/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,3 +129,32 @@ CREATE INDEX IF NOT EXISTS idx_suggested_questions_paper ON suggested_questions(
export const SCHEMA_V7_TRASH = `
ALTER TABLE papers ADD COLUMN trashed_at TEXT;
`

// [L2-05] Honest ACP usage metrics: per-turn and session telemetry recorded
// without storing privacy-sensitive prompt text. UNIQUE(session_id, turn_index)
// ensures retries/resumed turns update rather than double-counting.
export const SCHEMA_V8_USAGE = `
CREATE TABLE IF NOT EXISTS usage_records (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER REFERENCES chat_sessions(id) ON DELETE SET NULL,
backend TEXT NOT NULL,
model TEXT,
turn_index INTEGER NOT NULL DEFAULT 0,
input_tokens INTEGER,
output_tokens INTEGER,
thought_tokens INTEGER,
cached_read_tokens INTEGER,
cached_write_tokens INTEGER,
total_tokens INTEGER,
context_used INTEGER,
context_size INTEGER,
cost_amount REAL,
cost_currency TEXT,
has_metrics INTEGER NOT NULL DEFAULT 0,
recorded_at TEXT NOT NULL,
UNIQUE(session_id, turn_index)
);
CREATE INDEX IF NOT EXISTS idx_usage_records_backend ON usage_records(backend);
CREATE INDEX IF NOT EXISTS idx_usage_records_recorded_at ON usage_records(recorded_at);
CREATE INDEX IF NOT EXISTS idx_usage_records_session ON usage_records(session_id);
`
Loading
Loading