diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 6d9f9c47..4bc81569 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -544,6 +544,8 @@ impl Agent { // In multi-tenant mode, extract the owning user_id from // the response metadata so notifications reach the // correct user rather than the agent's owner. + // This intentionally overrides the configured notify_target + // because each user's heartbeat should notify that user. let effective_user = if is_multi_tenant { response .metadata diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index c8c21b88..eabf5dc4 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -397,11 +397,17 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { }; // 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()); + // When a model override is active, use the override name for attribution + // and let CostGuard look up pricing via costs::model_cost() instead of + // using the default provider's cost_per_token (which reflects the wrong model). + let (model_name, cost_per_token) = if let Some(ref ovr) = reason_ctx.model_override { + (ovr.clone(), None) + } else { + ( + self.agent.llm().active_model_name(), + Some(self.agent.llm().cost_per_token()), + ) + }; let read_discount = self.agent.llm().cache_read_discount(); let write_multiplier = self.agent.llm().cache_write_multiplier(); let call_cost = self @@ -416,7 +422,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { output.usage.cache_creation_input_tokens, read_discount, write_multiplier, - Some(self.agent.llm().cost_per_token()), + cost_per_token, ) .await; tracing::debug!( diff --git a/src/agent/heartbeat.rs b/src/agent/heartbeat.rs index 32f79444..5c2c9b1f 100644 --- a/src/agent/heartbeat.rs +++ b/src/agent/heartbeat.rs @@ -574,8 +574,9 @@ pub fn spawn_multi_user_heartbeat( } }; - // Run all user heartbeats concurrently so one slow LLM call - // doesn't block others. + // Run user heartbeats concurrently so one slow LLM call doesn't + // block others. Cap concurrency to avoid flooding the LLM provider. + const MAX_CONCURRENT_HEARTBEATS: usize = 8; let mut join_set = tokio::task::JoinSet::new(); for user_id in &user_ids { @@ -604,6 +605,13 @@ pub fn spawn_multi_user_heartbeat( } }); + // Drain completed tasks to stay within the concurrency cap. + while join_set.len() >= MAX_CONCURRENT_HEARTBEATS { + if let Some(join_result) = join_set.join_next().await { + collect_heartbeat_result(join_result, &mut user_failures, &config); + } + } + let uid = user_id.clone(); let cfg = config.clone(); let hyg = hygiene_config.clone(); @@ -626,48 +634,57 @@ pub fn spawn_multi_user_heartbeat( }); } - // Collect results and update failure counts + // Collect remaining results and update failure counts while let Some(join_result) = join_set.join_next().await { - let (uid, result) = match join_result { - Ok(pair) => pair, - Err(e) => { - tracing::error!("Multi-user heartbeat task panicked: {}", e); - continue; - } - }; - 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) => { - 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 - ); - } - } - } + collect_heartbeat_result(join_result, &mut user_failures, &config); } } }) } +/// Process a single JoinSet result from the multi-user heartbeat loop. +fn collect_heartbeat_result( + join_result: Result<(String, HeartbeatResult), tokio::task::JoinError>, + user_failures: &mut std::collections::HashMap, + config: &HeartbeatConfig, +) { + let (uid, result) = match join_result { + Ok(pair) => pair, + Err(e) => { + tracing::error!("Multi-user heartbeat task panicked: {}", e); + return; + } + }; + 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) => { + 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 + ); + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index f9980283..eaa179c1 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -737,12 +737,19 @@ impl RoutineEngine { }); } + // Per-user workspace (same pattern as spawn_fire). + let routine_workspace = if routine.user_id == self.workspace.user_id() { + self.workspace.clone() + } else { + Arc::new(Workspace::new_with_db(&routine.user_id, self.store.clone())) + }; + // Execute inline for manual triggers (caller wants to wait) let engine = EngineContext { config: self.config.clone(), store: self.store.clone(), llm: self.llm.clone(), - workspace: self.workspace.clone(), + workspace: routine_workspace, notify_tx: self.notify_tx.clone(), running_count: self.running_count.clone(), scheduler: self.scheduler.clone(), diff --git a/src/config/agent.rs b/src/config/agent.rs index e91b9140..81a82c60 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -120,10 +120,9 @@ impl AgentConfig { "AGENT_MAX_TOKENS_PER_JOB", settings.agent.max_tokens_per_job, )?, - multi_tenant: parse_bool_env( - "MULTI_TENANT", - optional_env("GATEWAY_USER_TOKENS")?.is_some(), - )?, + // Auto-detected from GATEWAY_USER_TOKENS presence. Not a separate + // knob — multi-tenant mode is always implied by configuring user tokens. + multi_tenant: optional_env("GATEWAY_USER_TOKENS")?.is_some(), }) } } diff --git a/src/llm/rig_adapter.rs b/src/llm/rig_adapter.rs index 07aec57f..ab304edf 100644 --- a/src/llm/rig_adapter.rs +++ b/src/llm/rig_adapter.rs @@ -599,11 +599,12 @@ fn build_rig_request( /// Inject a per-request model override into the rig request's `additional_params`. /// -/// Rig-core bakes the model name at construction time. For OpenAI, Anthropic, and -/// Ollama, the `model` field in the request body determines which model serves the -/// request. Rig-core's `#[serde(flatten)]` on `additional_params` emits these fields -/// AFTER the struct's own `model` field. Most API servers (Python, Go) use -/// last-key-wins when deserializing duplicate JSON keys, so the override takes effect. +/// Rig-core bakes the model name at construction time inside each provider's +/// `CompletionModel` implementation. The actual HTTP request body includes a +/// `model` field set by the provider. Rig-core's `#[serde(flatten)]` on +/// `additional_params` emits these fields AFTER the provider's own fields. +/// Most API servers (Python, Go) use last-key-wins when deserializing +/// duplicate JSON keys, so the injected `model` value takes effect. fn inject_model_override(rig_req: &mut RigRequest, model_override: Option<&str>) { let Some(model) = model_override else { return;