diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs index cafaad95..4885364b 100644 --- a/src/agent/cost_guard.rs +++ b/src/agent/cost_guard.rs @@ -227,6 +227,11 @@ impl CostGuard { daily.reset_date = today; self.budget_exceeded.store(false, Ordering::Relaxed); tracing::info!("Cost guard: daily counter reset for {}", today); + + // Prune per-user entries from previous days to prevent + // unbounded HashMap growth in long-lived deployments. + let mut per_user = self.per_user_daily_cost.lock().await; + per_user.retain(|_, entry| entry.reset_date == today); } daily.total += cost; diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index d6e05c9b..c8c21b88 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -396,8 +396,12 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { Err(e) => return Err(e.into()), }; - // Record cost and track token usage (global + per-user) - let model_name = self.agent.llm().active_model_name(); + // Record cost and track token usage (global + per-user). + // Use the override model name if set so cost attribution is accurate. + let model_name = reason_ctx + .model_override + .clone() + .unwrap_or_else(|| self.agent.llm().active_model_name()); let read_discount = self.agent.llm().cache_read_discount(); let write_multiplier = self.agent.llm().cache_write_multiplier(); let call_cost = self diff --git a/src/agent/heartbeat.rs b/src/agent/heartbeat.rs index d58ed4db..3f3dfab5 100644 --- a/src/agent/heartbeat.rs +++ b/src/agent/heartbeat.rs @@ -512,10 +512,10 @@ pub fn spawn_heartbeat( }) } -/// Spawn a multi-user heartbeat runner that cycles through all users with -/// active routines. Each tick, it queries the DB for distinct user_ids that -/// own routines, creates a per-user workspace, and runs a heartbeat check -/// for each user. +/// Spawn a multi-user heartbeat runner that cycles through all users that +/// own routines (enabled or not). Each tick, it queries the DB for distinct +/// user_ids, creates a per-user workspace, and runs a heartbeat check for +/// each user concurrently. Per-user failure counts are tracked independently. pub fn spawn_multi_user_heartbeat( config: HeartbeatConfig, hygiene_config: HygieneConfig, @@ -537,6 +537,11 @@ pub fn spawn_multi_user_heartbeat( None }; + // Track consecutive failures per user so we can disable heartbeat + // for persistently-failing users (same semantics as single-user mode). + let mut user_failures: std::collections::HashMap = + std::collections::HashMap::new(); + tracing::info!("Starting multi-user heartbeat loop"); loop { @@ -569,7 +574,17 @@ pub fn spawn_multi_user_heartbeat( } }; + // Run all user heartbeats concurrently so one slow LLM call + // doesn't block others. + let mut join_set = tokio::task::JoinSet::new(); + for user_id in &user_ids { + // Skip users that have exceeded max_failures + let failures = user_failures.get(user_id).copied().unwrap_or(0); + if failures >= config.max_failures { + continue; + } + let workspace = Arc::new(Workspace::new_with_db(user_id, store.clone())); // Run memory hygiene per user (same as single-user heartbeat). @@ -589,28 +604,56 @@ pub fn spawn_multi_user_heartbeat( } }); - let mut runner = HeartbeatRunner::new( - config.clone(), - hygiene_config.clone(), - workspace, - llm.clone(), - ); - if let Some(ref tx) = response_tx { - runner = runner.with_response_channel(tx.clone()); - } - runner = runner.with_store(store.clone()); + let uid = user_id.clone(); + let cfg = config.clone(); + let hyg = hygiene_config.clone(); + let llm_clone = llm.clone(); + let tx = response_tx.clone(); + let st = store.clone(); - match runner.check_heartbeat().await { - HeartbeatResult::Ok => { - tracing::trace!(user_id, "Multi-user heartbeat OK"); + join_set.spawn(async move { + let mut runner = HeartbeatRunner::new(cfg, hyg, workspace, llm_clone); + if let Some(tx) = tx { + runner = runner.with_response_channel(tx); } - HeartbeatResult::NeedsAttention(msg) => { - tracing::info!(user_id, "Multi-user heartbeat needs attention"); - runner.send_notification(&msg).await; + runner = runner.with_store(st); + + let result = runner.check_heartbeat().await; + if let HeartbeatResult::NeedsAttention(msg) = &result { + runner.send_notification(msg).await; + } + (uid, result) + }); + } + + // Collect results and update failure counts + while let Some(Ok((uid, result))) = join_set.join_next().await { + match result { + HeartbeatResult::Ok => { + tracing::trace!(user_id = uid, "Multi-user heartbeat OK"); + user_failures.remove(&uid); + } + HeartbeatResult::NeedsAttention(_) => { + tracing::info!(user_id = uid, "Multi-user heartbeat needs attention"); + user_failures.remove(&uid); } HeartbeatResult::Skipped => {} HeartbeatResult::Failed(err) => { - tracing::error!(user_id, "Multi-user heartbeat failed: {}", err); + let count = user_failures.entry(uid.clone()).or_insert(0); + *count += 1; + tracing::error!( + user_id = uid, + consecutive_failures = *count, + "Multi-user heartbeat failed: {}", + err + ); + if *count >= config.max_failures { + tracing::error!( + user_id = uid, + "Multi-user heartbeat disabled for user after {} consecutive failures", + count + ); + } } } }