From 9d7cdc0cf1637aa598e4299365b09e54e15dea33 Mon Sep 17 00:00:00 2001 From: "ilblackdragon@gmail.com" Date: Mon, 23 Mar 2026 22:46:59 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20complete=20multi-tenant=20isolation=20?= =?UTF-8?q?=E2=80=94=20per-user=20budgets,=20model=20selection,=20heartbea?= =?UTF-8?q?t=20cycling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finishes the remaining isolation work from phases 2–4 of #59: Phase 2 (DB scoping): Fix /status and /list commands to use _for_user DB variants instead of global queries that leaked cross-user job data. Phase 3 (Runtime isolation): Per-user workspace in routine engine's spawn_fire so lightweight routines run in the correct user context. Per-user daily cost tracking in CostGuard with configurable budget via MAX_COST_PER_USER_PER_DAY_CENTS. Multi-user heartbeat that cycles through all users with routines, auto-detected from GATEWAY_USER_TOKENS. Phase 4 (Provider/tools): Per-user model selection via preferred_model setting — looked up from SettingsStore on first iteration, threaded through ReasoningContext.model_override to CompletionRequest. Works with providers that support per-request model overrides (NearAI). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/agent_loop.rs | 34 ++++-- src/agent/commands.rs | 8 +- src/agent/cost_guard.rs | 234 +++++++++++++++++++++++++++++++++++- src/agent/dispatcher.rs | 32 ++++- src/agent/heartbeat.rs | 95 ++++++++++++++- src/agent/mod.rs | 4 +- src/agent/routine_engine.rs | 11 +- src/agent/scheduler.rs | 1 + src/app.rs | 1 + src/config/agent.rs | 4 + src/config/heartbeat.rs | 10 ++ src/llm/reasoning.rs | 11 ++ src/testing/mod.rs | 1 + 13 files changed, 423 insertions(+), 23 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index ee91ea9a..ca1e9ddd 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -13,7 +13,7 @@ use futures::StreamExt; use uuid::Uuid; use crate::agent::context_monitor::ContextMonitor; -use crate::agent::heartbeat::spawn_heartbeat; +use crate::agent::heartbeat::{spawn_heartbeat, spawn_multi_user_heartbeat}; use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker}; use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair}; use crate::agent::session_manager::SessionManager; @@ -505,6 +505,7 @@ impl Agent { .with_interval(std::time::Duration::from_secs(hb_config.interval_secs)); config.quiet_hours_start = hb_config.quiet_hours_start; config.quiet_hours_end = hb_config.quiet_hours_end; + config.multi_tenant = hb_config.multi_tenant; config.timezone = hb_config .timezone .clone() @@ -570,14 +571,29 @@ impl Agent { .map(|h| h.to_workspace_config()) .unwrap_or_default(); - Some(spawn_heartbeat( - config, - hygiene, - workspace.clone(), - self.cheap_llm().clone(), - Some(notify_tx), - self.store().map(Arc::clone), - )) + if config.multi_tenant { + if let Some(store) = self.store() { + Some(spawn_multi_user_heartbeat( + config, + hygiene, + self.cheap_llm().clone(), + Some(notify_tx), + Arc::clone(store), + )) + } else { + tracing::warn!("Multi-tenant heartbeat requires a database store"); + None + } + } else { + Some(spawn_heartbeat( + config, + hygiene, + workspace.clone(), + self.cheap_llm().clone(), + Some(notify_tx), + self.store().map(Arc::clone), + )) + } } else { tracing::warn!("Heartbeat enabled but no workspace available"); None diff --git a/src/agent/commands.rs b/src/agent/commands.rs index 75c99359..479d59d7 100644 --- a/src/agent/commands.rs +++ b/src/agent/commands.rs @@ -162,14 +162,14 @@ impl Agent { let mut failed = 0; let mut stuck = 0; - if let Ok(s) = store.agent_job_summary().await { + if let Ok(s) = store.agent_job_summary_for_user(user_id).await { total += s.total; in_progress += s.in_progress; completed += s.completed; failed += s.failed; stuck += s.stuck; } - if let Ok(s) = store.sandbox_job_summary().await { + if let Ok(s) = store.sandbox_job_summary_for_user(user_id).await { total += s.total; in_progress += s.running; completed += s.completed; @@ -226,14 +226,14 @@ impl Agent { ) -> Result { // List from DB for consistency with Jobs tab. if let Some(store) = self.store() { - let agent_jobs = match store.list_agent_jobs().await { + let agent_jobs = match store.list_agent_jobs_for_user(user_id).await { Ok(jobs) => jobs, Err(e) => { tracing::warn!("Failed to list agent jobs: {}", e); Vec::new() } }; - let sandbox_jobs = match store.list_sandbox_jobs().await { + let sandbox_jobs = match store.list_sandbox_jobs_for_user(user_id).await { Ok(jobs) => jobs, Err(e) => { tracing::warn!("Failed to list sandbox jobs: {}", e); diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs index 4563bbbe..cafaad95 100644 --- a/src/agent/cost_guard.rs +++ b/src/agent/cost_guard.rs @@ -21,6 +21,9 @@ pub struct CostGuardConfig { pub max_cost_per_day_cents: Option, /// Maximum LLM calls per hour. None = unlimited. pub max_actions_per_hour: Option, + /// Maximum spend per user per day in cents. None = unlimited. + /// Applied independently per user alongside the global budget. + pub max_cost_per_user_per_day_cents: Option, } /// Error returned when a cost limit is exceeded. @@ -30,6 +33,12 @@ pub enum CostLimitExceeded { DailyBudget { spent_cents: u64, limit_cents: u64 }, /// Hourly action rate limit reached. HourlyRate { actions: u64, limit: u64 }, + /// Per-user daily spending cap reached. + UserDailyBudget { + user_id: String, + spent_cents: u64, + limit_cents: u64, + }, } impl std::fmt::Display for CostLimitExceeded { @@ -49,6 +58,17 @@ impl std::fmt::Display for CostLimitExceeded { "Hourly action limit exceeded: {} actions of {} allowed per hour", actions, limit ), + Self::UserDailyBudget { + user_id, + spent_cents, + limit_cents, + } => write!( + f, + "User '{}' daily cost limit exceeded: spent ${:.2} of ${:.2} allowed", + user_id, + *spent_cents as f64 / 100.0, + *limit_cents as f64 / 100.0 + ), } } } @@ -78,6 +98,9 @@ pub struct CostGuard { /// Per-model token usage since startup. model_tokens: Mutex>, + + /// Per-user daily cost tracking. Each entry resets independently at midnight UTC. + per_user_daily_cost: Mutex>, } struct DailyCost { @@ -97,6 +120,7 @@ impl CostGuard { action_window: Mutex::new(VecDeque::new()), budget_exceeded: AtomicBool::new(false), model_tokens: Mutex::new(HashMap::new()), + per_user_daily_cost: Mutex::new(HashMap::new()), } } @@ -248,6 +272,85 @@ impl CostGuard { cost } + /// Record an LLM call with per-user attribution. + /// + /// Delegates to `record_llm_call` for global tracking, then additionally + /// records the cost against the user's daily budget. + #[allow(clippy::too_many_arguments)] + pub async fn record_llm_call_for_user( + &self, + user_id: &str, + model: &str, + input_tokens: u32, + output_tokens: u32, + cache_read_input_tokens: u32, + cache_creation_input_tokens: u32, + cache_read_discount: Decimal, + cache_write_multiplier: Decimal, + cost_per_token: Option<(Decimal, Decimal)>, + ) -> Decimal { + let cost = self + .record_llm_call( + model, + input_tokens, + output_tokens, + cache_read_input_tokens, + cache_creation_input_tokens, + cache_read_discount, + cache_write_multiplier, + cost_per_token, + ) + .await; + + // Track per-user daily cost + { + let today = chrono::Utc::now().date_naive(); + let mut per_user = self.per_user_daily_cost.lock().await; + let entry = per_user + .entry(user_id.to_string()) + .or_insert_with(|| DailyCost { + total: Decimal::ZERO, + reset_date: today, + }); + if today != entry.reset_date { + entry.total = Decimal::ZERO; + entry.reset_date = today; + } + entry.total += cost; + } + + cost + } + + /// Check whether the next action is allowed for a specific user. + /// + /// Checks the global limits first (via `check_allowed`), then additionally + /// checks the per-user daily budget if configured. + pub async fn check_allowed_for_user(&self, user_id: &str) -> Result<(), CostLimitExceeded> { + // Check global limits first + self.check_allowed().await?; + + // Check per-user daily budget + if let Some(limit_cents) = self.config.max_cost_per_user_per_day_cents { + let today = chrono::Utc::now().date_naive(); + let per_user = self.per_user_daily_cost.lock().await; + if let Some(entry) = per_user.get(user_id) + && entry.reset_date == today + { + let spent_cents = to_cents(entry.total); + if spent_cents >= limit_cents { + return Err(CostLimitExceeded::UserDailyBudget { + user_id: user_id.to_string(), + spent_cents, + limit_cents, + }); + } + } + } + + Ok(()) + } + /// Current daily spend in USD (as Decimal). pub async fn daily_spend(&self) -> Decimal { let daily = self.daily_cost.lock().await; @@ -259,6 +362,16 @@ impl CostGuard { } } + /// Current daily spend for a specific user in USD (as Decimal). + pub async fn daily_spend_for_user(&self, user_id: &str) -> Decimal { + let today = chrono::Utc::now().date_naive(); + let per_user = self.per_user_daily_cost.lock().await; + match per_user.get(user_id) { + Some(entry) if entry.reset_date == today => entry.total, + _ => Decimal::ZERO, + } + } + /// Number of actions in the current hourly window. pub async fn actions_this_hour(&self) -> u64 { let mut window = self.action_window.lock().await; @@ -314,7 +427,7 @@ mod tests { async fn test_daily_budget_enforcement() { let guard = CostGuard::new(CostGuardConfig { max_cost_per_day_cents: Some(1), // $0.01 limit - max_actions_per_hour: None, + ..CostGuardConfig::default() }); // First call allowed @@ -350,8 +463,8 @@ mod tests { #[tokio::test] async fn test_hourly_rate_enforcement() { let guard = CostGuard::new(CostGuardConfig { - max_cost_per_day_cents: None, max_actions_per_hour: Some(3), + ..CostGuardConfig::default() }); // First 3 actions allowed @@ -633,8 +746,8 @@ mod tests { // A fresh CostGuard with rate limits should not panic even if // checked_sub returns None (simulating short uptime). let guard = CostGuard::new(CostGuardConfig { - max_cost_per_day_cents: None, max_actions_per_hour: Some(100), + ..CostGuardConfig::default() }); // These must not panic regardless of system uptime @@ -656,4 +769,119 @@ mod tests { let result = Instant::now().checked_sub(std::time::Duration::MAX); assert!(result.is_none()); } + + #[tokio::test] + async fn test_per_user_daily_budget_enforcement() { + let guard = CostGuard::new(CostGuardConfig { + max_cost_per_day_cents: None, + max_actions_per_hour: None, + max_cost_per_user_per_day_cents: Some(1), // $0.01 per user + }); + + // Both users initially allowed + assert!(guard.check_allowed_for_user("alice").await.is_ok()); + assert!(guard.check_allowed_for_user("bob").await.is_ok()); + + // Alice makes an expensive call + guard + .record_llm_call_for_user( + "alice", + "gpt-4o", + 10_000, + 10_000, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + // Alice should be blocked, Bob should still be allowed + let result = guard.check_allowed_for_user("alice").await; + assert!(result.is_err()); + match result.unwrap_err() { + CostLimitExceeded::UserDailyBudget { + user_id, + limit_cents, + .. + } => { + assert_eq!(user_id, "alice"); + assert_eq!(limit_cents, 1); + } + other => panic!("Expected UserDailyBudget, got {:?}", other), + } + assert!(guard.check_allowed_for_user("bob").await.is_ok()); + } + + #[tokio::test] + async fn test_per_user_daily_spend_tracking() { + let guard = CostGuard::new(CostGuardConfig::default()); + + assert_eq!(guard.daily_spend_for_user("alice").await, Decimal::ZERO); + assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO); + + let cost = guard + .record_llm_call_for_user( + "alice", + "gpt-4o", + 1000, + 500, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + assert_eq!(guard.daily_spend_for_user("alice").await, cost); + assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO); + // Global spend should also be tracked + assert_eq!(guard.daily_spend().await, cost); + } + + #[tokio::test] + async fn test_per_user_budget_independent_of_global() { + let guard = CostGuard::new(CostGuardConfig { + max_cost_per_day_cents: Some(100_000), // $1000 global limit + max_actions_per_hour: None, + max_cost_per_user_per_day_cents: Some(1), // $0.01 per user + }); + + // User hits their personal limit + guard + .record_llm_call_for_user( + "alice", + "gpt-4o", + 10_000, + 10_000, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + // Alice blocked by per-user limit, not global + assert!(guard.check_allowed_for_user("alice").await.is_err()); + // Global limit is far from reached + assert!(guard.check_allowed().await.is_ok()); + // Bob is unaffected + assert!(guard.check_allowed_for_user("bob").await.is_ok()); + } + + #[test] + fn test_user_cost_limit_display() { + let limit = CostLimitExceeded::UserDailyBudget { + user_id: "alice".to_string(), + spent_cents: 150, + limit_cents: 100, + }; + let msg = limit.to_string(); + assert!(msg.contains("alice")); + assert!(msg.contains("$1.50")); + assert!(msg.contains("$1.00")); + } } diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 5d39866b..faa945d6 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -331,8 +331,13 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { reason_ctx: &mut ReasoningContext, iteration: usize, ) -> Result { - // Enforce cost guardrails before the LLM call - if let Err(limit) = self.agent.cost_guard().check_allowed().await { + // Enforce cost guardrails before the LLM call (global + per-user) + if let Err(limit) = self + .agent + .cost_guard() + .check_allowed_for_user(&self.message.user_id) + .await + { return Err(crate::error::LlmError::InvalidResponse { provider: "agent".to_string(), reason: limit.to_string(), @@ -340,6 +345,21 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { .into()); } + // Apply per-user model override from settings (first iteration only + // to avoid repeated DB lookups within the same agentic loop). + if iteration == 0 + && let Some(store) = self.agent.store() + && let Ok(Some(value)) = store + .get_setting(&self.message.user_id, "preferred_model") + .await + && let Some(model) = value.as_str() + { + let model = model.trim(); + if !model.is_empty() { + reason_ctx.model_override = Some(model.to_string()); + } + } + let output = match reasoning.respond_with_tools(reason_ctx).await { Ok(output) => output, Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => { @@ -374,14 +394,15 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { Err(e) => return Err(e.into()), }; - // Record cost and track token usage + // Record cost and track token usage (global + per-user) let model_name = 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 .agent .cost_guard() - .record_llm_call( + .record_llm_call_for_user( + &self.message.user_id, &model_name, output.usage.input_tokens, output.usage.output_tokens, @@ -1248,6 +1269,7 @@ mod tests { allow_local_tools: false, max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, max_tool_iterations: 50, auto_approve_tools: false, default_timezone: "UTC".to_string(), @@ -2115,6 +2137,7 @@ mod tests { allow_local_tools: false, max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, max_tool_iterations, auto_approve_tools: true, default_timezone: "UTC".to_string(), @@ -2235,6 +2258,7 @@ mod tests { allow_local_tools: false, max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, max_tool_iterations: max_iter, auto_approve_tools: true, default_timezone: "UTC".to_string(), diff --git a/src/agent/heartbeat.rs b/src/agent/heartbeat.rs index ec4cd5e9..443cb434 100644 --- a/src/agent/heartbeat.rs +++ b/src/agent/heartbeat.rs @@ -57,6 +57,9 @@ pub struct HeartbeatConfig { pub quiet_hours_end: Option, /// Timezone for fire_at and quiet hours evaluation (IANA name). pub timezone: Option, + /// When true, cycle through all users with routines instead of + /// running heartbeat for a single user. Requires a database store. + pub multi_tenant: bool, } impl Default for HeartbeatConfig { @@ -71,6 +74,7 @@ impl Default for HeartbeatConfig { quiet_hours_start: None, quiet_hours_end: None, timezone: None, + multi_tenant: false, } } } @@ -396,7 +400,7 @@ impl HeartbeatRunner { } /// Send a notification about heartbeat findings. - async fn send_notification(&self, message: &str) { + pub(crate) async fn send_notification(&self, message: &str) { let Some(ref tx) = self.response_tx else { tracing::debug!("No response channel configured for heartbeat notifications"); return; @@ -508,6 +512,95 @@ 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. +pub fn spawn_multi_user_heartbeat( + config: HeartbeatConfig, + hygiene_config: HygieneConfig, + llm: Arc, + response_tx: Option>, + store: Arc, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + if !config.enabled { + tracing::info!("Multi-user heartbeat is disabled"); + return; + } + + let mut tick_interval = if config.fire_at.is_none() { + let mut iv = tokio::time::interval(config.interval); + iv.tick().await; // skip immediate tick + Some(iv) + } else { + None + }; + + tracing::info!("Starting multi-user heartbeat loop"); + + loop { + if let Some(fire_at) = config.fire_at { + let sleep_dur = duration_until_next_fire(fire_at, config.resolved_tz()); + tokio::time::sleep(sleep_dur).await; + } else if let Some(ref mut iv) = tick_interval { + iv.tick().await; + } + + if config.is_quiet_hours() { + continue; + } + + // Get distinct user_ids from routines + let user_ids = match store.list_all_routines().await { + Ok(routines) => { + let mut ids: Vec = routines + .iter() + .map(|r| r.user_id.clone()) + .collect::>() + .into_iter() + .collect(); + ids.sort(); + ids + } + Err(e) => { + tracing::error!("Multi-user heartbeat: failed to list routines: {}", e); + continue; + } + }; + + for user_id in &user_ids { + let workspace = Arc::new(Workspace::new_with_db(user_id, store.clone())); + + 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()); + + match runner.check_heartbeat().await { + HeartbeatResult::Ok => { + tracing::trace!(user_id, "Multi-user heartbeat OK"); + } + HeartbeatResult::NeedsAttention(msg) => { + tracing::info!(user_id, "Multi-user heartbeat needs attention"); + runner.send_notification(&msg).await; + } + HeartbeatResult::Skipped => {} + HeartbeatResult::Failed(err) => { + tracing::error!(user_id, "Multi-user heartbeat failed: {}", err); + } + } + } + } + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 84155666..e7242845 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -36,7 +36,9 @@ pub(crate) use agent_loop::truncate_for_preview; pub use agent_loop::{Agent, AgentDeps}; pub use compaction::{CompactionResult, ContextCompactor}; pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor}; -pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat}; +pub use heartbeat::{ + HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat, spawn_multi_user_heartbeat, +}; pub use router::{MessageIntent, Router}; pub use routine::{Routine, RoutineAction, RoutineRun, Trigger}; pub use routine_engine::{RoutineEngine, SandboxReadiness}; diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index 7c7ef5f3..47eb0875 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -821,11 +821,20 @@ impl RoutineEngine { created_at: Utc::now(), }; + // Use per-user workspace so each routine executes in the correct + // user's context. Fall back to the engine-wide workspace when the + // routine belongs to the same user (avoids unnecessary allocation). + 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())) + }; + 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/agent/scheduler.rs b/src/agent/scheduler.rs index 02953a4b..db5e049a 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -780,6 +780,7 @@ mod tests { allow_local_tools: true, max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, max_tool_iterations: 10, auto_approve_tools: true, default_timezone: "UTC".to_string(), diff --git a/src/app.rs b/src/app.rs index edd547d3..de04f748 100644 --- a/src/app.rs +++ b/src/app.rs @@ -886,6 +886,7 @@ impl AppBuilder { crate::agent::cost_guard::CostGuardConfig { max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents, max_actions_per_hour: self.config.agent.max_actions_per_hour, + max_cost_per_user_per_day_cents: self.config.agent.max_cost_per_user_per_day_cents, }, )); diff --git a/src/config/agent.rs b/src/config/agent.rs index cb09707d..6a724c25 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -23,6 +23,8 @@ pub struct AgentConfig { pub max_cost_per_day_cents: Option, /// Maximum LLM/tool actions per hour. None = unlimited. pub max_actions_per_hour: Option, + /// Maximum daily LLM spend per user in cents. None = unlimited. + pub max_cost_per_user_per_day_cents: Option, /// Maximum tool-call iterations per agentic loop invocation. Default 50. pub max_tool_iterations: usize, /// When true, skip tool approval checks entirely. For benchmarks/CI. @@ -49,6 +51,7 @@ impl AgentConfig { allow_local_tools: true, max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, max_tool_iterations: 10, auto_approve_tools: true, default_timezone: "UTC".to_string(), @@ -87,6 +90,7 @@ impl AgentConfig { allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?, max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?, max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?, + max_cost_per_user_per_day_cents: parse_option_env("MAX_COST_PER_USER_PER_DAY_CENTS")?, max_tool_iterations: parse_optional_env( "AGENT_MAX_TOOL_ITERATIONS", settings.agent.max_tool_iterations, diff --git a/src/config/heartbeat.rs b/src/config/heartbeat.rs index 1dd456d7..09b8f0cd 100644 --- a/src/config/heartbeat.rs +++ b/src/config/heartbeat.rs @@ -21,6 +21,9 @@ pub struct HeartbeatConfig { pub quiet_hours_end: Option, /// Timezone for fire_at and quiet hours evaluation (IANA name). pub timezone: Option, + /// When true, cycle through all users with routines. Auto-detected from + /// GATEWAY_USER_TOKENS or set explicitly via HEARTBEAT_MULTI_TENANT. + pub multi_tenant: bool, } impl Default for HeartbeatConfig { @@ -34,6 +37,7 @@ impl Default for HeartbeatConfig { quiet_hours_start: None, quiet_hours_end: None, timezone: None, + multi_tenant: false, } } } @@ -101,6 +105,12 @@ impl HeartbeatConfig { } tz }, + // Auto-detect multi-tenant mode from GATEWAY_USER_TOKENS presence, + // or allow explicit override via HEARTBEAT_MULTI_TENANT. + multi_tenant: parse_bool_env( + "HEARTBEAT_MULTI_TENANT", + optional_env("GATEWAY_USER_TOKENS")?.is_some(), + )?, }) } } diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index cbec297b..3dbe9757 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -199,6 +199,10 @@ pub struct ReasoningContext { /// instead of calling `build_system_prompt_with_tools`. Allows callers to build /// the prompt once and reuse it across iterations. pub system_prompt: Option, + /// Per-user model override. When set, completion requests use this model + /// instead of the provider's default. Only effective with providers that + /// support per-request model overrides (e.g. NearAI). + pub model_override: Option, } impl ReasoningContext { @@ -212,6 +216,7 @@ impl ReasoningContext { metadata: std::collections::HashMap::new(), force_text: false, system_prompt: None, + model_override: None, } } @@ -653,6 +658,9 @@ Respond in JSON format: .with_temperature(0.7) .with_tool_choice("auto"); request.metadata = context.metadata.clone(); + if let Some(ref model) = context.model_override { + request.model = Some(model.clone()); + } let response = self.llm.complete_with_tools(request).await?; let usage = TokenUsage { @@ -732,6 +740,9 @@ Respond in JSON format: .with_max_tokens(4096) .with_temperature(0.7); request.metadata = context.metadata.clone(); + if let Some(ref model) = context.model_override { + request.model = Some(model.clone()); + } let response = self.llm.complete(request).await?; let pre_truncated = truncate_at_tool_tags(&response.content); diff --git a/src/testing/mod.rs b/src/testing/mod.rs index a633e91c..1a6d4264 100644 --- a/src/testing/mod.rs +++ b/src/testing/mod.rs @@ -532,6 +532,7 @@ impl TestHarnessBuilder { let cost_guard = Arc::new(CostGuard::new(CostGuardConfig { max_cost_per_day_cents: None, max_actions_per_hour: None, + max_cost_per_user_per_day_cents: None, })); let channel = if self.stub_channel {