mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
feat: complete multi-tenant isolation — per-user budgets, model selection, heartbeat cycling
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) <[email protected]>
This commit is contained in:
+25
-9
@@ -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
|
||||
|
||||
@@ -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<String, Error> {
|
||||
// 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);
|
||||
|
||||
+231
-3
@@ -21,6 +21,9 @@ pub struct CostGuardConfig {
|
||||
pub max_cost_per_day_cents: Option<u64>,
|
||||
/// Maximum LLM calls per hour. None = unlimited.
|
||||
pub max_actions_per_hour: Option<u64>,
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
/// 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<HashMap<String, ModelTokens>>,
|
||||
|
||||
/// Per-user daily cost tracking. Each entry resets independently at midnight UTC.
|
||||
per_user_daily_cost: Mutex<HashMap<String, DailyCost>>,
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
||||
+28
-4
@@ -331,8 +331,13 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
|
||||
reason_ctx: &mut ReasoningContext,
|
||||
iteration: usize,
|
||||
) -> Result<crate::llm::RespondOutput, Error> {
|
||||
// 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(),
|
||||
|
||||
+94
-1
@@ -57,6 +57,9 @@ pub struct HeartbeatConfig {
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
/// 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<dyn LlmProvider>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: Arc<dyn Database>,
|
||||
) -> 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<String> = routines
|
||||
.iter()
|
||||
.map(|r| r.user_id.clone())
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
.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::*;
|
||||
|
||||
+3
-1
@@ -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};
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
));
|
||||
|
||||
|
||||
@@ -23,6 +23,8 @@ pub struct AgentConfig {
|
||||
pub max_cost_per_day_cents: Option<u64>,
|
||||
/// Maximum LLM/tool actions per hour. None = unlimited.
|
||||
pub max_actions_per_hour: Option<u64>,
|
||||
/// Maximum daily LLM spend per user in cents. None = unlimited.
|
||||
pub max_cost_per_user_per_day_cents: Option<u64>,
|
||||
/// 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,
|
||||
|
||||
@@ -21,6 +21,9 @@ pub struct HeartbeatConfig {
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for fire_at and quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
/// 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(),
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String>,
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user