diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs index 2ddeae58..59d4ed85 100644 --- a/src/agent/cost_guard.rs +++ b/src/agent/cost_guard.rs @@ -4,7 +4,7 @@ //! to prevent runaway agents from burning through API credits. Especially //! important for daemon/heartbeat modes where the agent acts autonomously. -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Instant; @@ -53,6 +53,14 @@ impl std::fmt::Display for CostLimitExceeded { } } +/// Per-model token usage counters. +#[derive(Debug, Clone, Default)] +pub struct ModelTokens { + pub input_tokens: u64, + pub output_tokens: u64, + pub cost: Decimal, +} + /// Tracks costs and action rates, enforcing configurable limits. /// /// Thread-safe; designed to be shared via `Arc`. @@ -67,6 +75,9 @@ pub struct CostGuard { /// Flag set when daily budget is exceeded to short-circuit checks. budget_exceeded: AtomicBool, + + /// Per-model token usage since startup. + model_tokens: Mutex>, } struct DailyCost { @@ -85,6 +96,7 @@ impl CostGuard { }), action_window: Mutex::new(VecDeque::new()), budget_exceeded: AtomicBool::new(false), + model_tokens: Mutex::new(HashMap::new()), } } @@ -192,6 +204,15 @@ impl CostGuard { window.push_back(Instant::now()); } + // Track per-model token usage + { + let mut tokens = self.model_tokens.lock().await; + let entry = tokens.entry(model.to_string()).or_default(); + entry.input_tokens += u64::from(input_tokens); + entry.output_tokens += u64::from(output_tokens); + entry.cost += cost; + } + cost } @@ -215,6 +236,11 @@ impl CostGuard { } window.len() as u64 } + + /// Per-model token usage since startup. + pub async fn model_usage(&self) -> HashMap { + self.model_tokens.lock().await.clone() + } } /// Convert a Decimal USD amount to whole cents (truncated). @@ -336,4 +362,37 @@ mod tests { assert!(rate.to_string().contains("101 actions")); assert!(rate.to_string().contains("100 allowed")); } + + #[tokio::test] + async fn test_model_usage_per_model_tracking() { + let guard = CostGuard::new(CostGuardConfig::default()); + + // Initially empty + assert!(guard.model_usage().await.is_empty()); + + // Record calls for two different models + guard.record_llm_call("gpt-4o", 1000, 500).await; + guard.record_llm_call("gpt-4o", 2000, 1000).await; + guard + .record_llm_call("claude-3-5-sonnet-20241022", 500, 200) + .await; + + let usage = guard.model_usage().await; + assert_eq!(usage.len(), 2); + + let gpt = usage.get("gpt-4o").expect("gpt-4o should be tracked"); + assert_eq!(gpt.input_tokens, 3000); + assert_eq!(gpt.output_tokens, 1500); + assert!(gpt.cost > Decimal::ZERO); + + let claude = usage + .get("claude-3-5-sonnet-20241022") + .expect("claude should be tracked"); + assert_eq!(claude.input_tokens, 500); + assert_eq!(claude.output_tokens, 200); + assert!(claude.cost > Decimal::ZERO); + + // Costs should differ since models have different pricing + assert_ne!(gpt.cost, claude.cost); + } } diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index 90a4cd04..a212978e 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -89,6 +89,8 @@ impl GatewayChannel { skill_registry: None, skill_catalog: None, chat_rate_limiter: server::RateLimiter::new(30, 60), + cost_guard: None, + startup_time: std::time::Instant::now(), }); Self { @@ -119,6 +121,8 @@ impl GatewayChannel { skill_registry: self.state.skill_registry.clone(), skill_catalog: self.state.skill_catalog.clone(), chat_rate_limiter: server::RateLimiter::new(30, 60), + cost_guard: self.state.cost_guard.clone(), + startup_time: self.state.startup_time, }; mutate(&mut new_state); self.state = Arc::new(new_state); @@ -206,6 +210,12 @@ impl GatewayChannel { self } + /// Inject the cost guard for token/cost tracking in the status popover. + pub fn with_cost_guard(mut self, cg: Arc) -> Self { + self.rebuild_state(|s| s.cost_guard = Some(cg)); + self + } + /// Get the auth token (for printing to console on startup). pub fn auth_token(&self) -> &str { &self.auth_token diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 45a22686..a4ee00a2 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -148,6 +148,10 @@ pub struct GatewayState { pub skill_catalog: Option>, /// Rate limiter for chat endpoints (30 messages per 60 seconds). pub chat_rate_limiter: RateLimiter, + /// Cost guard for token/cost tracking. + pub cost_guard: Option>, + /// Server startup time for uptime calculation. + pub startup_time: std::time::Instant, } /// Start the gateway HTTP server. @@ -2569,18 +2573,57 @@ async fn gateway_status_handler( .map(|t| t.connection_count()) .unwrap_or(0); + let uptime_secs = state.startup_time.elapsed().as_secs(); + + let (daily_cost, actions_this_hour, model_usage) = if let Some(ref cg) = state.cost_guard { + let cost = cg.daily_spend().await; + let actions = cg.actions_this_hour().await; + let usage = cg.model_usage().await; + let models: Vec = usage + .into_iter() + .map(|(model, tokens)| ModelUsageEntry { + model, + input_tokens: tokens.input_tokens, + output_tokens: tokens.output_tokens, + cost: format!("{:.6}", tokens.cost), + }) + .collect(); + (Some(format!("{:.4}", cost)), Some(actions), Some(models)) + } else { + (None, None, None) + }; + Json(GatewayStatusResponse { sse_connections, ws_connections, total_connections: sse_connections + ws_connections, + uptime_secs, + daily_cost, + actions_this_hour, + model_usage, }) } +#[derive(serde::Serialize)] +struct ModelUsageEntry { + model: String, + input_tokens: u64, + output_tokens: u64, + cost: String, +} + #[derive(serde::Serialize)] struct GatewayStatusResponse { sse_connections: u64, ws_connections: u64, total_connections: u64, + uptime_secs: u64, + #[serde(skip_serializing_if = "Option::is_none")] + daily_cost: Option, + #[serde(skip_serializing_if = "Option::is_none")] + actions_this_hour: Option, + #[serde(skip_serializing_if = "Option::is_none")] + model_usage: Option>, } #[cfg(test)] diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 57ba3947..d24cf96b 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -2252,13 +2252,72 @@ function startGatewayStatusPolling() { gatewayStatusInterval = setInterval(fetchGatewayStatus, 30000); } +function formatTokenCount(n) { + if (n == null || n === 0) return '0'; + if (n >= 1000000) return (n / 1000000).toFixed(1) + 'M'; + if (n >= 1000) return (n / 1000).toFixed(1) + 'k'; + return '' + n; +} + +function formatCost(costStr) { + if (!costStr) return '$0.00'; + var n = parseFloat(costStr); + if (n < 0.01) return '$' + n.toFixed(4); + return '$' + n.toFixed(2); +} + +function shortModelName(model) { + // Strip provider prefix and shorten common model names + var m = model.indexOf('/') >= 0 ? model.split('/').pop() : model; + // Shorten dated suffixes + m = m.replace(/-20\d{6}$/, ''); + return m; +} + function fetchGatewayStatus() { - apiFetch('/api/gateway/status').then((data) => { - const popover = document.getElementById('gateway-popover'); - popover.innerHTML = '
SSE clients' + (data.sse_clients || 0) + '
' - + '
Log clients' + (data.log_clients || 0) + '
' - + '
Uptime' + formatDuration(data.uptime_secs) + '
'; - }).catch(() => {}); + apiFetch('/api/gateway/status').then(function(data) { + var popover = document.getElementById('gateway-popover'); + var html = ''; + + // Connection info + html += ''; + html += '
SSE' + (data.sse_connections || 0) + '
'; + html += '
WebSocket' + (data.ws_connections || 0) + '
'; + html += '
Uptime' + formatDuration(data.uptime_secs) + '
'; + + // Cost tracker + if (data.daily_cost != null) { + html += '
'; + html += ''; + html += '
Spent' + formatCost(data.daily_cost) + '
'; + if (data.actions_this_hour != null) { + html += '
Actions/hr' + data.actions_this_hour + '
'; + } + } + + // Per-model token usage + if (data.model_usage && data.model_usage.length > 0) { + html += '
'; + html += ''; + data.model_usage.sort(function(a, b) { + return (b.input_tokens + b.output_tokens) - (a.input_tokens + a.output_tokens); + }); + for (var i = 0; i < data.model_usage.length; i++) { + var m = data.model_usage[i]; + var name = escapeHtml(shortModelName(m.model)); + html += '
' + + '' + name + '' + + '' + escapeHtml(formatCost(m.cost)) + '' + + '
'; + html += '
' + + 'in: ' + formatTokenCount(m.input_tokens) + '' + + 'out: ' + formatTokenCount(m.output_tokens) + '' + + '
'; + } + } + + popover.innerHTML = html; + }).catch(function() {}); } // Show/hide popover on hover diff --git a/src/channels/web/static/style.css b/src/channels/web/static/style.css index f62c25b1..626834c3 100644 --- a/src/channels/web/static/style.css +++ b/src/channels/web/static/style.css @@ -2549,7 +2549,7 @@ mark { border: 1px solid var(--border); border-radius: var(--radius-lg); padding: 12px; - min-width: 180px; + min-width: 220px; box-shadow: var(--shadow); z-index: 100; } @@ -2558,6 +2558,20 @@ mark { display: block; } +.gw-section-label { + font-size: 10px; + text-transform: uppercase; + letter-spacing: 0.05em; + color: var(--text-muted, var(--text-secondary)); + margin-bottom: 4px; + font-weight: 600; +} + +.gw-divider { + border-top: 1px solid var(--border); + margin: 8px 0; +} + .gw-stat { display: flex; justify-content: space-between; @@ -2571,6 +2585,37 @@ mark { font-weight: 500; } +.gw-model-row { + display: flex; + justify-content: space-between; + font-size: 12px; + padding: 3px 0 0 0; +} + +.gw-model-name { + color: var(--text); + font-weight: 500; + font-size: 11px; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + max-width: 140px; +} + +.gw-model-cost { + color: var(--accent, var(--text)); + font-weight: 500; + font-size: 11px; +} + +.gw-token-detail { + display: flex; + gap: 12px; + font-size: 10px; + color: var(--text-secondary); + padding: 1px 0 4px 0; +} + /* --- Extension install form --- */ .ext-install-form { diff --git a/src/channels/web/ws.rs b/src/channels/web/ws.rs index 2ada57e1..6e91717d 100644 --- a/src/channels/web/ws.rs +++ b/src/channels/web/ws.rs @@ -490,6 +490,8 @@ mod tests { skill_registry: None, skill_catalog: None, chat_rate_limiter: crate::channels::web::server::RateLimiter::new(30, 60), + cost_guard: None, + startup_time: std::time::Instant::now(), } } } diff --git a/src/main.rs b/src/main.rs index b4a5831e..f02da188 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1319,6 +1319,14 @@ async fn main() -> anyhow::Result<()> { (None, None) }; + // Create cost guard early so gateway can reference it. + let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new( + ironclaw::agent::cost_guard::CostGuardConfig { + max_cost_per_day_cents: config.agent.max_cost_per_day_cents, + max_actions_per_hour: config.agent.max_actions_per_hour, + }, + )); + // Add web gateway channel if configured let mut gateway_url: Option = None; if let Some(ref gw_config) = config.channels.gateway { @@ -1345,6 +1353,7 @@ async fn main() -> anyhow::Result<()> { if let Some(ref sc) = skill_catalog { gw = gw.with_skill_catalog(Arc::clone(sc)); } + gw = gw.with_cost_guard(Arc::clone(&cost_guard)); if config.sandbox.enabled { gw = gw.with_prompt_queue(Arc::clone(&prompt_queue)); @@ -1379,12 +1388,6 @@ async fn main() -> anyhow::Result<()> { let boot_cheap_model = cheap_llm.as_ref().map(|c| c.model_name().to_string()); // Create and run the agent - let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new( - ironclaw::agent::cost_guard::CostGuardConfig { - max_cost_per_day_cents: config.agent.max_cost_per_day_cents, - max_actions_per_hour: config.agent.max_actions_per_hour, - }, - )); let deps = AgentDeps { store: db, llm, diff --git a/tests/openai_compat_integration.rs b/tests/openai_compat_integration.rs index cf5be23a..32623150 100644 --- a/tests/openai_compat_integration.rs +++ b/tests/openai_compat_integration.rs @@ -198,6 +198,8 @@ async fn start_test_server_with_provider( skill_registry: None, skill_catalog: None, chat_rate_limiter: ironclaw::channels::web::server::RateLimiter::new(30, 60), + cost_guard: None, + startup_time: std::time::Instant::now(), }); let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); @@ -683,6 +685,8 @@ async fn test_no_llm_provider_returns_503() { skill_registry: None, skill_catalog: None, chat_rate_limiter: ironclaw::channels::web::server::RateLimiter::new(30, 60), + cost_guard: None, + startup_time: std::time::Instant::now(), }); let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); diff --git a/tests/ws_gateway_integration.rs b/tests/ws_gateway_integration.rs index 3ae713a2..cb002f2f 100644 --- a/tests/ws_gateway_integration.rs +++ b/tests/ws_gateway_integration.rs @@ -56,6 +56,8 @@ async fn start_test_server() -> ( skill_registry: None, skill_catalog: None, chat_rate_limiter: ironclaw::channels::web::server::RateLimiter::new(30, 60), + cost_guard: None, + startup_time: std::time::Instant::now(), }); let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();