fix: address PR review — cost model attribution, heartbeat concurrency, pruning

Fixes from review comments on #1614:

- Cost tracking now uses the override model name (not active_model_name)
  when a per-user model override is active, for accurate attribution.
- Multi-user heartbeat runs per-user checks concurrently via JoinSet
  instead of sequentially, preventing one slow user from blocking others.
- Per-user failure counts tracked independently; users exceeding
  max_failures are skipped (matching single-user semantics).
- per_user_daily_cost HashMap pruned on day rollover to prevent
  unbounded growth in long-lived deployments.
- Doc comment fixed: says "routines" not "active routines".

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
2026-03-24 00:06:16 -07:00
co-authored by Claude Opus 4.6
parent 9ff4af5734
commit 6dfe246288
3 changed files with 75 additions and 23 deletions
+5
View File
@@ -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;
+6 -2
View File
@@ -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
+64 -21
View File
@@ -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<String, u32> =
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
);
}
}
}
}