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:
parent
14549dd8a9
commit
5c6c607e25
9 changed files with 267 additions and 101 deletions
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue