agent ctx-badge: seed Bus::last_usage from latest turn_stats row on startup

This commit is contained in:
müde 2026-05-18 18:00:48 +02:00
parent 8a3e8bfb7f
commit f827187341
4 changed files with 47 additions and 0 deletions

View file

@ -74,6 +74,11 @@ async fn main() -> Result<()> {
let login_state = Arc::new(Mutex::new(initial));
let bus = Bus::new();
let stats = TurnStats::open_default();
if let Some(s) = &stats
&& let Some(u) = s.last_usage()
{
bus.seed_usage(u);
}
let files = turn::TurnFiles::prepare(&cli.socket, &label, mcp::Flavor::Agent).await?;
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
plugins::install_configured(&cli.socket, Some("manager")).await;

View file

@ -64,6 +64,11 @@ async fn main() -> Result<()> {
let login_state = Arc::new(Mutex::new(initial));
let bus = Bus::new();
let stats = TurnStats::open_default();
if let Some(s) = &stats
&& let Some(u) = s.last_usage()
{
bus.seed_usage(u);
}
let files = turn::TurnFiles::prepare(&cli.socket, &label, mcp::Flavor::Manager).await?;
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
plugins::install_configured(&cli.socket, None).await;

View file

@ -378,6 +378,15 @@ impl Bus {
self.emit(LiveEvent::ModelChanged { model: value });
}
/// Seed `last_usage` at startup without emitting a SSE event.
/// Used by the bin entrypoints to backfill from the most recent
/// `turn_stats` row so the per-agent web UI's `ctx-badge` paints
/// real numbers on cold load instead of staying empty until the
/// next turn finishes.
pub fn seed_usage(&self, usage: TokenUsage) {
*self.last_usage.lock().unwrap() = Some(usage);
}
/// Record the latest token usage from a completed turn.
pub fn record_usage(&self, usage: TokenUsage) {
*self.last_usage.lock().unwrap() = Some(usage);

View file

@ -156,6 +156,34 @@ impl TurnStats {
tracing::warn!(error = ?e, "turn_stats: insert failed");
}
}
/// 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.
#[must_use]
pub fn last_usage(&self) -> 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
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),
})
},
)
.ok()
}
}
fn default_path() -> PathBuf {