agent ctx-badge: seed Bus::last_usage from latest turn_stats row on startup
This commit is contained in:
parent
8a3e8bfb7f
commit
f827187341
4 changed files with 47 additions and 0 deletions
|
|
@ -74,6 +74,11 @@ async fn main() -> Result<()> {
|
||||||
let login_state = Arc::new(Mutex::new(initial));
|
let login_state = Arc::new(Mutex::new(initial));
|
||||||
let bus = Bus::new();
|
let bus = Bus::new();
|
||||||
let stats = TurnStats::open_default();
|
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 files = turn::TurnFiles::prepare(&cli.socket, &label, mcp::Flavor::Agent).await?;
|
||||||
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
|
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
|
||||||
plugins::install_configured(&cli.socket, Some("manager")).await;
|
plugins::install_configured(&cli.socket, Some("manager")).await;
|
||||||
|
|
|
||||||
|
|
@ -64,6 +64,11 @@ async fn main() -> Result<()> {
|
||||||
let login_state = Arc::new(Mutex::new(initial));
|
let login_state = Arc::new(Mutex::new(initial));
|
||||||
let bus = Bus::new();
|
let bus = Bus::new();
|
||||||
let stats = TurnStats::open_default();
|
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 files = turn::TurnFiles::prepare(&cli.socket, &label, mcp::Flavor::Manager).await?;
|
||||||
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
|
let turn_lock: TurnLock = Arc::new(tokio::sync::Mutex::new(()));
|
||||||
plugins::install_configured(&cli.socket, None).await;
|
plugins::install_configured(&cli.socket, None).await;
|
||||||
|
|
|
||||||
|
|
@ -378,6 +378,15 @@ impl Bus {
|
||||||
self.emit(LiveEvent::ModelChanged { model: value });
|
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.
|
/// Record the latest token usage from a completed turn.
|
||||||
pub fn record_usage(&self, usage: TokenUsage) {
|
pub fn record_usage(&self, usage: TokenUsage) {
|
||||||
*self.last_usage.lock().unwrap() = Some(usage);
|
*self.last_usage.lock().unwrap() = Some(usage);
|
||||||
|
|
|
||||||
|
|
@ -156,6 +156,34 @@ impl TurnStats {
|
||||||
tracing::warn!(error = ?e, "turn_stats: insert failed");
|
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 {
|
fn default_path() -> PathBuf {
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue