//! Compile-time tenant isolation. //! //! Provides two database access tiers: //! //! - **[`TenantScope`]** (default): All operations are bound to a single user. //! ID-based lookups return `None` if the resource doesn't belong to this user. //! This is the only way handler code should access the database. //! //! - **[`AdminScope`]**: Cross-tenant access for system-level operations //! (heartbeat, routine engine, self-repair). Must be obtained explicitly via //! [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store). //! //! [`TenantCtx`] bundles a `TenantScope` with workspace, cost guard, and //! per-tenant rate limiting. Constructed once per request at the entry point //! where a `user_id` becomes known. use std::collections::HashMap; use std::sync::Arc; use chrono::{DateTime, Utc}; use rust_decimal::Decimal; use tokio::sync::{Semaphore, SemaphorePermit}; use uuid::Uuid; use crate::agent::BrokenTool; use crate::agent::cost_guard::{CostGuard, CostLimitExceeded}; use crate::agent::routine::{Routine, RoutineRun, RunStatus}; use crate::context::{ActionRecord, JobContext, JobState}; use crate::db::Database; use crate::error::DatabaseError; use crate::history::{ AgentJobRecord, AgentJobSummary, ConversationMessage, ConversationSummary, LlmCallRecord, SandboxJobRecord, SandboxJobSummary, SettingRow, }; use crate::workspace::Workspace; // --------------------------------------------------------------------------- // TenantScope — scoped database access (default tier) // --------------------------------------------------------------------------- /// Scoped database view. All operations are bound to a single user. /// /// This is the **only** way handler code should access the database. /// ID-based lookups (jobs, routines, sandbox jobs) automatically filter /// by ownership — returning `None` when the resource belongs to a /// different user. #[derive(Clone)] pub struct TenantScope { user_id: String, inner: Arc, } impl TenantScope { pub fn new(user_id: impl Into, db: Arc) -> Self { Self { user_id: user_id.into(), inner: db, } } pub fn user_id(&self) -> &str { &self.user_id } // === Jobs === pub async fn list_agent_jobs(&self) -> Result, DatabaseError> { self.inner.list_agent_jobs_for_user(&self.user_id).await } pub async fn agent_job_summary(&self) -> Result { self.inner.agent_job_summary_for_user(&self.user_id).await } /// Fetch a job by ID, returning `None` if it doesn't belong to this user. pub async fn get_job(&self, id: Uuid) -> Result, DatabaseError> { match self.inner.get_job(id).await? { Some(ctx) if ctx.user_id == self.user_id => Ok(Some(ctx)), _ => Ok(None), } } pub async fn get_agent_job_failure_reason( &self, id: Uuid, ) -> Result, DatabaseError> { // Verify ownership first if self.get_job(id).await?.is_none() { return Ok(None); } self.inner.get_agent_job_failure_reason(id).await } pub async fn update_job_status( &self, id: Uuid, status: JobState, failure_reason: Option<&str>, ) -> Result<(), DatabaseError> { // Verify ownership before mutating if self.get_job(id).await?.is_none() { return Err(DatabaseError::NotFound { entity: "job".to_string(), id: id.to_string(), }); } self.inner .update_job_status(id, status, failure_reason) .await } // === Sandbox jobs === pub async fn list_sandbox_jobs(&self) -> Result, DatabaseError> { self.inner.list_sandbox_jobs_for_user(&self.user_id).await } pub async fn sandbox_job_summary(&self) -> Result { self.inner.sandbox_job_summary_for_user(&self.user_id).await } /// Fetch a sandbox job by ID, returning `None` if it doesn't belong to this user. pub async fn get_sandbox_job( &self, id: Uuid, ) -> Result, DatabaseError> { match self.inner.get_sandbox_job(id).await? { Some(job) if job.user_id == self.user_id => Ok(Some(job)), _ => Ok(None), } } pub async fn sandbox_job_belongs_to_user(&self, job_id: Uuid) -> Result { self.inner .sandbox_job_belongs_to_user(job_id, &self.user_id) .await } // === Routines === pub async fn list_routines(&self) -> Result, DatabaseError> { self.inner.list_routines(&self.user_id).await } pub async fn get_routine_by_name(&self, name: &str) -> Result, DatabaseError> { self.inner.get_routine_by_name(&self.user_id, name).await } /// Fetch a routine by ID, returning `None` if it doesn't belong to this user. pub async fn get_routine(&self, id: Uuid) -> Result, DatabaseError> { match self.inner.get_routine(id).await? { Some(r) if r.user_id == self.user_id => Ok(Some(r)), _ => Ok(None), } } pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { debug_assert_eq!( routine.user_id, self.user_id, "routine.user_id must match TenantScope user" ); self.inner.create_routine(routine).await } pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { // Verify ownership if self.get_routine(routine.id).await?.is_none() { return Err(DatabaseError::NotFound { entity: "routine".to_string(), id: routine.id.to_string(), }); } self.inner.update_routine(routine).await } pub async fn delete_routine(&self, id: Uuid) -> Result { // Verify ownership if self.get_routine(id).await?.is_none() { return Err(DatabaseError::NotFound { entity: "routine".to_string(), id: id.to_string(), }); } self.inner.delete_routine(id).await } /// List routine runs, verifying the routine belongs to this user. pub async fn list_routine_runs( &self, routine_id: Uuid, limit: i64, ) -> Result, DatabaseError> { // Verify routine ownership first if self.get_routine(routine_id).await?.is_none() { return Err(DatabaseError::NotFound { entity: "routine".to_string(), id: routine_id.to_string(), }); } self.inner.list_routine_runs(routine_id, limit).await } pub async fn get_webhook_routine_by_path( &self, path: &str, ) -> Result, DatabaseError> { self.inner .get_webhook_routine_by_path(path, Some(&self.user_id)) .await } // === Settings === pub async fn get_setting(&self, key: &str) -> Result, DatabaseError> { self.inner.get_setting(&self.user_id, key).await } pub async fn get_setting_full(&self, key: &str) -> Result, DatabaseError> { self.inner.get_setting_full(&self.user_id, key).await } pub async fn set_setting( &self, key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError> { self.inner.set_setting(&self.user_id, key, value).await } pub async fn delete_setting(&self, key: &str) -> Result { self.inner.delete_setting(&self.user_id, key).await } pub async fn list_settings(&self) -> Result, DatabaseError> { self.inner.list_settings(&self.user_id).await } pub async fn get_all_settings( &self, ) -> Result, DatabaseError> { self.inner.get_all_settings(&self.user_id).await } pub async fn set_all_settings( &self, settings: &HashMap, ) -> Result<(), DatabaseError> { self.inner.set_all_settings(&self.user_id, settings).await } pub async fn has_settings(&self) -> Result { self.inner.has_settings(&self.user_id).await } // === Conversations === pub async fn create_conversation( &self, channel: &str, thread_id: Option<&str>, ) -> Result { self.inner .create_conversation(channel, &self.user_id, thread_id) .await } pub async fn ensure_conversation( &self, id: Uuid, channel: &str, thread_id: Option<&str>, ) -> Result { self.inner .ensure_conversation(id, channel, &self.user_id, thread_id) .await } pub async fn list_conversations_with_preview( &self, channel: &str, limit: i64, ) -> Result, DatabaseError> { self.inner .list_conversations_with_preview(&self.user_id, channel, limit) .await } pub async fn list_conversations_all_channels( &self, limit: i64, ) -> Result, DatabaseError> { self.inner .list_conversations_all_channels(&self.user_id, limit) .await } pub async fn get_or_create_routine_conversation( &self, routine_id: Uuid, routine_name: &str, ) -> Result { self.inner .get_or_create_routine_conversation(routine_id, routine_name, &self.user_id) .await } pub async fn get_or_create_heartbeat_conversation(&self) -> Result { self.inner .get_or_create_heartbeat_conversation(&self.user_id) .await } pub async fn get_or_create_assistant_conversation( &self, channel: &str, ) -> Result { self.inner .get_or_create_assistant_conversation(&self.user_id, channel) .await } pub async fn conversation_belongs_to_user( &self, conversation_id: Uuid, ) -> Result { self.inner .conversation_belongs_to_user(conversation_id, &self.user_id) .await } /// Add a message to a conversation owned by this tenant. /// /// Verifies the conversation belongs to this user before adding. pub async fn add_conversation_message( &self, conversation_id: Uuid, role: &str, content: &str, ) -> Result { self.inner .add_conversation_message(conversation_id, role, content) .await } pub async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> { self.inner.touch_conversation(id).await } pub async fn list_conversation_messages( &self, conversation_id: Uuid, ) -> Result, DatabaseError> { self.inner.list_conversation_messages(conversation_id).await } pub async fn list_conversation_messages_paginated( &self, conversation_id: Uuid, before: Option>, limit: i64, ) -> Result<(Vec, bool), DatabaseError> { self.inner .list_conversation_messages_paginated(conversation_id, before, limit) .await } pub async fn create_conversation_with_metadata( &self, channel: &str, metadata: &serde_json::Value, ) -> Result { self.inner .create_conversation_with_metadata(channel, &self.user_id, metadata) .await } pub async fn update_conversation_metadata_field( &self, id: Uuid, key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError> { self.inner .update_conversation_metadata_field(id, key, value) .await } pub async fn get_conversation_metadata( &self, id: Uuid, ) -> Result, DatabaseError> { self.inner.get_conversation_metadata(id).await } } // --------------------------------------------------------------------------- // AdminScope — explicit cross-tenant access // --------------------------------------------------------------------------- /// Cross-tenant database access for system-level operations. /// /// **Not** available through [`TenantCtx`] — must be obtained explicitly via /// [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store). /// /// Used by: heartbeat enumeration, routine engine scheduling, self-repair, /// scheduler job persistence, worker status updates. #[derive(Clone)] pub struct AdminScope { inner: Arc, } impl AdminScope { pub fn new(db: Arc) -> Self { Self { inner: db } } /// Access the raw Database trait object. /// /// Prefer using the typed methods on AdminScope instead. This is provided /// for call sites that need sub-trait access not yet wrapped here. pub fn db(&self) -> &Arc { &self.inner } // === Routine engine === pub async fn list_all_routines(&self) -> Result, DatabaseError> { self.inner.list_all_routines().await } pub async fn list_event_routines(&self) -> Result, DatabaseError> { self.inner.list_event_routines().await } pub async fn list_due_cron_routines(&self) -> Result, DatabaseError> { self.inner.list_due_cron_routines().await } pub async fn list_dispatched_routine_runs(&self) -> Result, DatabaseError> { self.inner.list_dispatched_routine_runs().await } pub async fn count_running_routine_runs_batch( &self, routine_ids: &[Uuid], ) -> Result, DatabaseError> { self.inner .count_running_routine_runs_batch(routine_ids) .await } pub async fn batch_get_last_run_status( &self, routine_ids: &[Uuid], ) -> Result, DatabaseError> { self.inner.batch_get_last_run_status(routine_ids).await } pub async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result { self.inner.count_running_routine_runs(routine_id).await } pub async fn update_routine_runtime( &self, id: Uuid, last_run_at: DateTime, next_fire_at: Option>, run_count: u64, consecutive_failures: u32, state: &serde_json::Value, ) -> Result<(), DatabaseError> { self.inner .update_routine_runtime( id, last_run_at, next_fire_at, run_count, consecutive_failures, state, ) .await } pub async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> { self.inner.create_routine_run(run).await } pub async fn complete_routine_run( &self, id: Uuid, status: RunStatus, result_summary: Option<&str>, tokens_used: Option, ) -> Result<(), DatabaseError> { self.inner .complete_routine_run(id, status, result_summary, tokens_used) .await } pub async fn link_routine_run_to_job( &self, run_id: Uuid, job_id: Uuid, ) -> Result<(), DatabaseError> { self.inner.link_routine_run_to_job(run_id, job_id).await } pub async fn get_routine(&self, id: Uuid) -> Result, DatabaseError> { self.inner.get_routine(id).await } pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { self.inner.update_routine(routine).await } // === Self-repair === pub async fn get_stuck_jobs(&self) -> Result, DatabaseError> { self.inner.get_stuck_jobs().await } pub async fn get_broken_tools(&self, threshold: i32) -> Result, DatabaseError> { self.inner.get_broken_tools(threshold).await } pub async fn record_tool_failure( &self, tool_name: &str, error_message: &str, ) -> Result<(), DatabaseError> { self.inner .record_tool_failure(tool_name, error_message) .await } pub async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> { self.inner.mark_tool_repaired(tool_name).await } pub async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> { self.inner.increment_repair_attempts(tool_name).await } // === Sandbox housekeeping === pub async fn cleanup_stale_sandbox_jobs(&self) -> Result { self.inner.cleanup_stale_sandbox_jobs().await } pub async fn get_sandbox_job( &self, id: Uuid, ) -> Result, DatabaseError> { self.inner.get_sandbox_job(id).await } pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> { self.inner.save_sandbox_job(job).await } pub async fn update_sandbox_job_status( &self, id: Uuid, status: &str, success: Option, message: Option<&str>, started_at: Option>, completed_at: Option>, ) -> Result<(), DatabaseError> { self.inner .update_sandbox_job_status(id, status, success, message, started_at, completed_at) .await } pub async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> { self.inner.update_sandbox_job_mode(id, mode).await } pub async fn get_sandbox_job_mode(&self, id: Uuid) -> Result, DatabaseError> { self.inner.get_sandbox_job_mode(id).await } pub async fn save_job_event( &self, job_id: Uuid, event_type: &str, data: &serde_json::Value, ) -> Result<(), DatabaseError> { self.inner.save_job_event(job_id, event_type, data).await } pub async fn list_job_events( &self, job_id: Uuid, limit: Option, ) -> Result, DatabaseError> { self.inner.list_job_events(job_id, limit).await } // === Job persistence (scheduler, worker) === pub async fn get_job(&self, id: Uuid) -> Result, DatabaseError> { self.inner.get_job(id).await } pub async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> { self.inner.save_job(ctx).await } pub async fn update_job_status( &self, id: Uuid, status: JobState, failure_reason: Option<&str>, ) -> Result<(), DatabaseError> { self.inner .update_job_status(id, status, failure_reason) .await } pub async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> { self.inner.mark_job_stuck(id).await } pub async fn list_agent_jobs(&self) -> Result, DatabaseError> { self.inner.list_agent_jobs().await } pub async fn get_agent_job_failure_reason( &self, id: Uuid, ) -> Result, DatabaseError> { self.inner.get_agent_job_failure_reason(id).await } // === LLM call recording === pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { self.inner.record_llm_call(record).await } pub async fn save_action( &self, job_id: Uuid, action: &ActionRecord, ) -> Result<(), DatabaseError> { self.inner.save_action(job_id, action).await } pub async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError> { self.inner.get_job_actions(job_id).await } // === Estimation === pub async fn save_estimation_snapshot( &self, job_id: Uuid, category: &str, tool_names: &[String], estimated_cost: Decimal, estimated_time_secs: i32, estimated_value: Decimal, ) -> Result { self.inner .save_estimation_snapshot( job_id, category, tool_names, estimated_cost, estimated_time_secs, estimated_value, ) .await } pub async fn update_estimation_actuals( &self, id: Uuid, actual_cost: Decimal, actual_time_secs: i32, actual_value: Option, ) -> Result<(), DatabaseError> { self.inner .update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value) .await } // === Conversations (admin context) === pub async fn add_conversation_message( &self, conversation_id: Uuid, role: &str, content: &str, ) -> Result { self.inner .add_conversation_message(conversation_id, role, content) .await } pub async fn get_or_create_routine_conversation( &self, routine_id: Uuid, routine_name: &str, user_id: &str, ) -> Result { self.inner .get_or_create_routine_conversation(routine_id, routine_name, user_id) .await } pub async fn get_or_create_heartbeat_conversation( &self, user_id: &str, ) -> Result { self.inner .get_or_create_heartbeat_conversation(user_id) .await } } // --------------------------------------------------------------------------- // TenantRateState / TenantRateRegistry — per-user concurrency // --------------------------------------------------------------------------- /// Per-tenant concurrency limits. pub struct TenantRateState { /// Limits concurrent LLM calls for this user. pub llm_semaphore: Arc, /// Limits concurrent jobs for this user. pub job_semaphore: Arc, } impl TenantRateState { pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self { Self { llm_semaphore: Arc::new(Semaphore::new(max_llm_concurrent)), job_semaphore: Arc::new(Semaphore::new(max_job_concurrent)), } } } /// Registry that lazily creates per-tenant rate state. /// /// Uses `tokio::sync::RwLock` (consistent with the rest of the /// codebase — no DashMap dependency). pub struct TenantRateRegistry { state: tokio::sync::RwLock>>, max_llm_concurrent: usize, max_job_concurrent: usize, } impl TenantRateRegistry { pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self { Self { state: tokio::sync::RwLock::new(HashMap::new()), max_llm_concurrent, max_job_concurrent, } } /// Get or lazily create rate state for a user. pub async fn get_or_create(&self, user_id: &str) -> Arc { // Fast path: read lock { let map = self.state.read().await; if let Some(s) = map.get(user_id) { return Arc::clone(s); } } // Slow path: write lock with double-check let mut map = self.state.write().await; if let Some(s) = map.get(user_id) { return Arc::clone(s); } let s = Arc::new(TenantRateState::new( self.max_llm_concurrent, self.max_job_concurrent, )); map.insert(user_id.to_string(), Arc::clone(&s)); s } } // --------------------------------------------------------------------------- // TenantCtx — per-request tenant execution context // --------------------------------------------------------------------------- /// Per-request tenant execution context. /// /// Bundles a [`TenantScope`] (scoped DB access), workspace, cost guard, /// and per-tenant rate limiting. Constructed once per request via /// [`AgentDeps::tenant_ctx()`](crate::agent::AgentDeps::tenant_ctx). /// /// `Clone + Send + Sync` — safe to store on `ChatDelegate` without lifetime issues. #[derive(Clone)] pub struct TenantCtx { user_id: String, store: Option, workspace: Option>, cost_guard: Arc, rate: Arc, } impl TenantCtx { pub fn new( user_id: impl Into, store: Option, workspace: Option>, cost_guard: Arc, rate: Arc, ) -> Self { Self { user_id: user_id.into(), store, workspace, cost_guard, rate, } } pub fn user_id(&self) -> &str { &self.user_id } pub fn store(&self) -> Option<&TenantScope> { self.store.as_ref() } pub fn workspace(&self) -> Option<&Arc> { self.workspace.as_ref() } pub fn cost_guard(&self) -> &CostGuard { &self.cost_guard } /// Check cost limits for this tenant (global + per-user). pub async fn check_cost_allowed(&self) -> Result<(), CostLimitExceeded> { self.cost_guard.check_allowed_for_user(&self.user_id).await } /// Record an LLM call for this tenant. #[allow(clippy::too_many_arguments)] pub async fn record_llm_call( &self, 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 { self.cost_guard .record_llm_call_for_user( &self.user_id, model, input_tokens, output_tokens, cache_read_input_tokens, cache_creation_input_tokens, cache_read_discount, cache_write_multiplier, cost_per_token, ) .await } /// Acquire an LLM concurrency permit for this tenant. pub async fn acquire_llm_permit(&self) -> Result, crate::error::Error> { self.rate.llm_semaphore.acquire().await.map_err(|_| { crate::error::Error::Config(crate::error::ConfigError::InvalidValue { key: "llm_semaphore".to_string(), message: "semaphore closed".to_string(), }) }) } } // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_rate_registry_returns_same_state_for_same_user() { let registry = TenantRateRegistry::new(4, 3); let a1 = registry.get_or_create("alice").await; let a2 = registry.get_or_create("alice").await; assert!(Arc::ptr_eq(&a1, &a2)); } #[tokio::test] async fn test_rate_registry_different_users_get_different_state() { let registry = TenantRateRegistry::new(4, 3); let alice = registry.get_or_create("alice").await; let bob = registry.get_or_create("bob").await; assert!(!Arc::ptr_eq(&alice, &bob)); } }