agent badges: split into ctx (last-inference) + cost (cumulative)

the existing ctx badge was misnamed: it summed `result.usage`, which is
the cumulative tokens billed across every inference in the turn. for
tool-heavy turns that easily exceeds the model's context window (a 600k
cached prefix × 15 sub-calls = 9M cache_read), making it useless as a
"should i compact?" signal.

now two separate badges:

  ctx · N    last inference's prompt size = actual context window in
             use right now. parsed from each `assistant` event's
             `.message.usage`; the harness tracks the most recent one
             across the stream and snapshots it when the `result`
             event lands.

  cost · M   cumulative tokens billed across the whole turn (the
             previous behaviour, now correctly labelled).

both update via a single `TokenUsageChanged { ctx, cost }` SSE event at
turn-end. turn_stats grows four columns (`last_input_tokens`,
`last_output_tokens`, `last_cache_read_input_tokens`,
`last_cache_creation_input_tokens`) so the cold-load seed can paint both
badges on page load. migrations run try-and-ignore ALTERs so existing
agent dbs catch up; pre-migration rows have last-inference zeros and
yield no `ctx` seed (badge stays empty until next turn) rather than a
misleading 0.
This commit is contained in:
müde 2026-05-18 18:48:35 +02:00
commit 5c6c607e25
9 changed files with 267 additions and 101 deletions

View file

@ -22,8 +22,9 @@ use anyhow::{Context, Result};
use rusqlite::{Connection, params};
/// SQL bootstrap. CREATE TABLE IF NOT EXISTS so first-boot agents
/// and existing ones converge on the same shape; ALTER-style
/// migrations land here as additional statements once we have any.
/// and existing ones converge on the same shape. The base table is
/// fresh-install only; additive migrations land via `MIGRATIONS`
/// below as try-and-ignore ALTERs so existing dbs catch up.
const SCHEMA: &str = "
CREATE TABLE IF NOT EXISTS turn_stats (
id INTEGER PRIMARY KEY AUTOINCREMENT,
@ -36,6 +37,10 @@ CREATE TABLE IF NOT EXISTS turn_stats (
output_tokens INTEGER NOT NULL DEFAULT 0,
cache_read_input_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0,
last_input_tokens INTEGER NOT NULL DEFAULT 0,
last_output_tokens INTEGER NOT NULL DEFAULT 0,
last_cache_read_input_tokens INTEGER NOT NULL DEFAULT 0,
last_cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0,
tool_call_count INTEGER NOT NULL DEFAULT 0,
tool_call_breakdown_json TEXT,
open_threads_count INTEGER,
@ -47,6 +52,17 @@ CREATE INDEX IF NOT EXISTS idx_turn_stats_started
ON turn_stats (started_at DESC);
";
/// Additive column migrations. Each runs unconditionally and ignores
/// `duplicate column name` errors — sqlite < 3.35 lacks
/// `ADD COLUMN IF NOT EXISTS`, so try-and-ignore is the portable path.
/// New columns MUST carry a default so existing rows decode.
const MIGRATIONS: &[&str] = &[
"ALTER TABLE turn_stats ADD COLUMN last_input_tokens INTEGER NOT NULL DEFAULT 0",
"ALTER TABLE turn_stats ADD COLUMN last_output_tokens INTEGER NOT NULL DEFAULT 0",
"ALTER TABLE turn_stats ADD COLUMN last_cache_read_input_tokens INTEGER NOT NULL DEFAULT 0",
"ALTER TABLE turn_stats ADD COLUMN last_cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0",
];
/// One row to be inserted. `Option`-wrapped fields default to NULL
/// when the harness couldn't gather them (e.g. socket roundtrip for
/// open_threads failed) so a partial row beats no row.
@ -57,10 +73,16 @@ pub struct TurnStatRow {
pub duration_ms: i64,
pub model: String,
pub wake_from: String,
/// Cumulative across every inference in the turn (cost signal).
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_input_tokens: u64,
pub cache_creation_input_tokens: u64,
/// Last inference's usage — the actual context size at turn end.
pub last_input_tokens: u64,
pub last_output_tokens: u64,
pub last_cache_read_input_tokens: u64,
pub last_cache_creation_input_tokens: u64,
pub tool_call_count: u64,
/// Per-tool breakdown as JSON: `{"Read":12,"Bash":3,...}`. None
/// when no tools were called (saves a sqlite write of `"{}"`).
@ -107,6 +129,18 @@ impl TurnStats {
.with_context(|| format!("open turn_stats db {}", path.display()))?;
conn.execute_batch(SCHEMA)
.context("apply turn_stats schema")?;
for stmt in MIGRATIONS {
// Ignore "duplicate column name" — the migration already ran.
// Any other error is logged but doesn't fail open() because the
// base schema works and we'd rather keep the harness alive than
// crash on an upgrade hiccup.
if let Err(e) = conn.execute(stmt, []) {
let msg = e.to_string();
if !msg.contains("duplicate column name") {
tracing::warn!(error = %msg, stmt, "turn_stats migration failed");
}
}
}
Ok(Self {
inner: std::sync::Arc::new(Mutex::new(conn)),
})
@ -121,6 +155,8 @@ impl TurnStats {
started_at, ended_at, duration_ms, model, wake_from,
input_tokens, output_tokens,
cache_read_input_tokens, cache_creation_input_tokens,
last_input_tokens, last_output_tokens,
last_cache_read_input_tokens, last_cache_creation_input_tokens,
tool_call_count, tool_call_breakdown_json,
open_threads_count, open_reminders_count,
result_kind, note
@ -130,7 +166,9 @@ impl TurnStats {
?8, ?9,
?10, ?11,
?12, ?13,
?14, ?15
?14, ?15,
?16, ?17,
?18, ?19
)",
params![
row.started_at,
@ -142,6 +180,10 @@ impl TurnStats {
i64::try_from(row.output_tokens).unwrap_or(i64::MAX),
i64::try_from(row.cache_read_input_tokens).unwrap_or(i64::MAX),
i64::try_from(row.cache_creation_input_tokens).unwrap_or(i64::MAX),
i64::try_from(row.last_input_tokens).unwrap_or(i64::MAX),
i64::try_from(row.last_output_tokens).unwrap_or(i64::MAX),
i64::try_from(row.last_cache_read_input_tokens).unwrap_or(i64::MAX),
i64::try_from(row.last_cache_creation_input_tokens).unwrap_or(i64::MAX),
i64::try_from(row.tool_call_count).unwrap_or(i64::MAX),
row.tool_call_breakdown_json,
row.open_threads_count
@ -157,32 +199,58 @@ impl TurnStats {
}
}
/// Token counts from the most recently inserted row, if any. Lets
/// the harness seed `Bus::last_usage` on startup so the per-agent
/// web UI's `ctx-badge` paints with real numbers on cold load
/// instead of waiting for the next `TokenUsageChanged` SSE event.
/// Best-effort: any sqlite error returns `None` and the caller
/// falls back to the empty state.
/// Token counts from the most recently inserted row, if any.
/// Returns `(ctx, cost)` — both backfill `Bus` on startup so the
/// per-agent web UI's ctx + cost badges paint with real numbers on
/// cold load instead of waiting for the next `TokenUsageChanged`
/// SSE event. Best-effort: any sqlite error returns `(None, None)`.
///
/// Pre-migration rows (before the `last_*_tokens` columns existed)
/// have last-inference zeros — those rows yield `ctx = None` so the
/// badge stays empty until the next real turn rather than showing a
/// misleading 0.
#[must_use]
pub fn last_usage(&self) -> Option<crate::events::TokenUsage> {
pub fn last_usage(
&self,
) -> (
Option<crate::events::TokenUsage>,
Option<crate::events::TokenUsage>,
) {
let conn = self.inner.lock().unwrap();
conn.query_row(
"SELECT input_tokens, output_tokens,
cache_read_input_tokens, cache_creation_input_tokens
cache_read_input_tokens, cache_creation_input_tokens,
last_input_tokens, last_output_tokens,
last_cache_read_input_tokens, last_cache_creation_input_tokens
FROM turn_stats
ORDER BY started_at DESC
LIMIT 1",
[],
|row| {
Ok(crate::events::TokenUsage {
input_tokens: u64::try_from(row.get::<_, i64>(0)?).unwrap_or(0),
output_tokens: u64::try_from(row.get::<_, i64>(1)?).unwrap_or(0),
cache_read_input_tokens: u64::try_from(row.get::<_, i64>(2)?).unwrap_or(0),
cache_creation_input_tokens: u64::try_from(row.get::<_, i64>(3)?).unwrap_or(0),
})
let g = |i: usize| -> rusqlite::Result<u64> {
Ok(u64::try_from(row.get::<_, i64>(i)?).unwrap_or(0))
};
let cost = crate::events::TokenUsage {
input_tokens: g(0)?,
output_tokens: g(1)?,
cache_read_input_tokens: g(2)?,
cache_creation_input_tokens: g(3)?,
};
let last = crate::events::TokenUsage {
input_tokens: g(4)?,
output_tokens: g(5)?,
cache_read_input_tokens: g(6)?,
cache_creation_input_tokens: g(7)?,
};
let ctx = if last == crate::events::TokenUsage::default() {
None
} else {
Some(last)
};
Ok((ctx, Some(cost)))
},
)
.ok()
.unwrap_or((None, None))
}
}