mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-30 08:17:53 +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]> * fix: use selected_model setting key to match /model command persistence The dispatcher was reading "preferred_model" but the /model command (merged from staging) persists to "selected_model". Since set_setting is already per-user scoped, using the same key makes /model work as the per-user model override in multi-tenant mode. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: heartbeat hygiene, /model multi-tenant guard, RigAdapter model override Three follow-up fixes for multi-tenant isolation: 1. Multi-user heartbeat now runs memory hygiene per user before each heartbeat check, matching single-user heartbeat behavior. 2. /model command in multi-tenant mode only persists to per-user settings (selected_model) without calling set_model() on the shared LlmProvider. The per-request model_override in the dispatcher reads from the same setting. Added multi_tenant flag to AgentConfig (auto-detected from GATEWAY_USER_TOKENS). 3. RigAdapter now supports per-request model overrides by injecting the model name into rig-core's additional_params. OpenAI/Anthropic/Ollama API servers use last-key-wins for duplicate JSON keys, so the override takes effect via serde's flatten serialization order. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * 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]> * fix: /status ownership, model persistence scoping, heartbeat robustness Addresses second round of PR review on #1614: - /status <job_id> DB path now validates job.user_id == requesting user before returning data (was missing ownership check, security fix). - persist_selected_model takes user_id param instead of owner_id, and skips .env/TOML writes in multi-tenant mode (these are shared global files). handle_system_command now receives user_id from caller. - JoinSet collection handles Err(JoinError) explicitly instead of silently dropping panicked tasks. - Notification forwarder extracts owner_id from response metadata in multi-tenant mode for per-user routing instead of broadcasting to the agent owner. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: cost pricing, fire_manual workspace, heartbeat concurrency cap Round 3 review fixes: - Cost tracking passes None for cost_per_token when model override is active, letting CostGuard look up pricing by model name instead of using the default provider's rates (serrrfirat). - fire_manual() now uses per-user workspace, matching spawn_fire() pattern (serrrfirat). - Removed MULTI_TENANT env var — multi-tenant mode is auto-detected solely from GATEWAY_USER_TOKENS presence (serrrfirat + Copilot). - Multi-user heartbeat capped at 8 concurrent tasks to avoid flooding the LLM provider (serrrfirat + Copilot). - Fixed inject_model_override doc comment accuracy (Copilot). - Added comment explaining multi-tenant notification routing priority (Copilot). Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: user-scoped webhook endpoint for multi-tenant isolation Adds POST /api/webhooks/u/{user_id}/{path} — a user-scoped webhook endpoint that filters the routine lookup by user_id, preventing cross-user webhook triggering when paths collide. The existing /api/webhooks/{path} endpoint remains unchanged for backward compatibility in single-user deployments. Changes: - get_webhook_routine_by_path gains user_id: Option<&str> param - Both postgres and libsql implementations add AND user_id = ? filter when user_id is provided - New webhook_trigger_user_scoped_handler extracts (user_id, path) from URL and passes to shared fire_webhook_inner logic - Route registered on public router (webhooks are called by external services that can't send bearer tokens) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: add TenantCtx for compile-time tenant isolation Implements zmanian's architectural proposal from #1614 review: two-tier scoped database access (TenantScope/AdminScope) so handler code cannot accidentally bypass tenant scoping. TenantScope (default): wraps user_id + Arc<dyn Database>, auto-binds user_id on every operation. ID-based lookups return None for cross- tenant resources. No escape hatch — forgetting to scope is a compile error. AdminScope (explicit opt-in): cross-tenant access for system-level components (heartbeat, routine engine, self-repair, scheduler, worker). TenantCtx bundles TenantScope + workspace + cost guard + per-user rate limiting. Constructed once per request in handle_message, threaded through all command handlers and ChatDelegate. Key changes: - New src/tenant.rs (~920 lines): TenantScope, AdminScope, TenantCtx, TenantRateState, TenantRateRegistry - All command handlers: user_id: &str → ctx: &TenantCtx - ChatDelegate: cost check/record/settings via self.tenant - System components: store field changed to AdminScope - Config: TENANT_MAX_LLM_CONCURRENT, TENANT_MAX_JOBS_CONCURRENT env vars - Fixes bug: /status <job_id> cross-tenant leak (now auto-filtered) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
970 lines
33 KiB
Rust
970 lines
33 KiB
Rust
//! Proactive heartbeat system for periodic execution.
|
|
//!
|
|
//! The heartbeat runner executes periodically (default: every 30 minutes) and:
|
|
//! 1. Reads the HEARTBEAT.md checklist
|
|
//! 2. Runs an agent turn to process the checklist
|
|
//! 3. Reports any findings to the configured channel
|
|
//!
|
|
//! If nothing needs attention, the agent replies "HEARTBEAT_OK" and no
|
|
//! message is sent to the user.
|
|
//!
|
|
//! # Usage
|
|
//!
|
|
//! Create a HEARTBEAT.md in the workspace with a checklist of things to monitor:
|
|
//!
|
|
//! ```markdown
|
|
//! # Heartbeat Checklist
|
|
//!
|
|
//! - [ ] Check for unread emails
|
|
//! - [ ] Review calendar for upcoming events
|
|
//! - [ ] Check project build status
|
|
//! ```
|
|
//!
|
|
//! The agent will process this checklist on each heartbeat and only notify
|
|
//! if action is needed.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use chrono::TimeZone as _;
|
|
use chrono_tz::Tz;
|
|
use tokio::sync::mpsc;
|
|
|
|
use crate::channels::OutgoingResponse;
|
|
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
|
use crate::tenant::AdminScope;
|
|
use crate::workspace::Workspace;
|
|
use crate::workspace::hygiene::HygieneConfig;
|
|
|
|
/// Configuration for the heartbeat runner.
|
|
#[derive(Debug, Clone)]
|
|
pub struct HeartbeatConfig {
|
|
/// Interval between heartbeat checks (used when fire_at is not set).
|
|
pub interval: Duration,
|
|
/// Whether heartbeat is enabled.
|
|
pub enabled: bool,
|
|
/// Maximum consecutive failures before disabling.
|
|
pub max_failures: u32,
|
|
/// User ID to notify on heartbeat findings.
|
|
pub notify_user_id: Option<String>,
|
|
/// Channel to notify on heartbeat findings.
|
|
pub notify_channel: Option<String>,
|
|
/// Fixed time-of-day to fire (24h). When set, interval is ignored.
|
|
pub fire_at: Option<chrono::NaiveTime>,
|
|
/// Hour (0-23) when quiet hours start.
|
|
pub quiet_hours_start: Option<u32>,
|
|
/// Hour (0-23) when quiet hours end.
|
|
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 {
|
|
fn default() -> Self {
|
|
Self {
|
|
interval: Duration::from_secs(30 * 60), // 30 minutes
|
|
enabled: true,
|
|
max_failures: 3,
|
|
notify_user_id: None,
|
|
notify_channel: None,
|
|
fire_at: None,
|
|
quiet_hours_start: None,
|
|
quiet_hours_end: None,
|
|
timezone: None,
|
|
multi_tenant: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl HeartbeatConfig {
|
|
/// Create a config with a specific interval.
|
|
pub fn with_interval(mut self, interval: Duration) -> Self {
|
|
self.interval = interval;
|
|
self
|
|
}
|
|
|
|
/// Disable heartbeat.
|
|
pub fn disabled(mut self) -> Self {
|
|
self.enabled = false;
|
|
self
|
|
}
|
|
|
|
/// Check whether the current time falls within configured quiet hours.
|
|
pub fn is_quiet_hours(&self) -> bool {
|
|
use chrono::Timelike;
|
|
let (Some(start), Some(end)) = (self.quiet_hours_start, self.quiet_hours_end) else {
|
|
return false;
|
|
};
|
|
let tz = self
|
|
.timezone
|
|
.as_deref()
|
|
.and_then(crate::timezone::parse_timezone)
|
|
.unwrap_or(chrono_tz::UTC);
|
|
let now_hour = crate::timezone::now_in_tz(tz).hour();
|
|
if start <= end {
|
|
now_hour >= start && now_hour < end
|
|
} else {
|
|
// Wraps midnight, e.g. 22..06
|
|
now_hour >= start || now_hour < end
|
|
}
|
|
}
|
|
|
|
/// Set the notification target.
|
|
pub fn with_notify(mut self, user_id: impl Into<String>, channel: impl Into<String>) -> Self {
|
|
self.notify_user_id = Some(user_id.into());
|
|
self.notify_channel = Some(channel.into());
|
|
self
|
|
}
|
|
|
|
/// Set a fixed time-of-day to fire (overrides interval).
|
|
pub fn with_fire_at(mut self, time: chrono::NaiveTime, tz: Option<String>) -> Self {
|
|
self.fire_at = Some(time);
|
|
self.timezone = tz;
|
|
self
|
|
}
|
|
|
|
/// Resolve timezone string to chrono_tz::Tz (defaults to UTC).
|
|
fn resolved_tz(&self) -> Tz {
|
|
self.timezone
|
|
.as_deref()
|
|
.and_then(crate::timezone::parse_timezone)
|
|
.unwrap_or(chrono_tz::UTC)
|
|
}
|
|
}
|
|
|
|
/// Result of a heartbeat check.
|
|
#[derive(Debug)]
|
|
pub enum HeartbeatResult {
|
|
/// Nothing needs attention.
|
|
Ok,
|
|
/// Something needs attention, with the message to send.
|
|
NeedsAttention(String),
|
|
/// Heartbeat was skipped (no checklist or disabled).
|
|
Skipped,
|
|
/// Heartbeat failed.
|
|
Failed(String),
|
|
}
|
|
|
|
/// Compute how long to sleep until the next occurrence of `fire_at` in `tz`.
|
|
///
|
|
/// If the target time today is still in the future, sleep until then.
|
|
/// Otherwise sleep until the same time tomorrow.
|
|
fn duration_until_next_fire(fire_at: chrono::NaiveTime, tz: Tz) -> Duration {
|
|
let now = chrono::Utc::now().with_timezone(&tz);
|
|
let today = now.date_naive();
|
|
|
|
// Try to build today's target datetime in the given timezone.
|
|
// `.earliest()` picks the first occurrence if DST creates ambiguity.
|
|
let candidate = tz.from_local_datetime(&today.and_time(fire_at)).earliest();
|
|
|
|
let target = match candidate {
|
|
Some(t) if t > now => t,
|
|
_ => {
|
|
// Already past (or ambiguous) — schedule for tomorrow
|
|
let tomorrow = today + chrono::Duration::days(1);
|
|
tz.from_local_datetime(&tomorrow.and_time(fire_at))
|
|
.earliest()
|
|
.unwrap_or_else(|| now + chrono::Duration::days(1))
|
|
}
|
|
};
|
|
|
|
let secs = (target - now).num_seconds().max(1) as u64;
|
|
Duration::from_secs(secs)
|
|
}
|
|
|
|
/// Heartbeat runner for proactive periodic execution.
|
|
pub struct HeartbeatRunner {
|
|
config: HeartbeatConfig,
|
|
hygiene_config: HygieneConfig,
|
|
workspace: Arc<Workspace>,
|
|
llm: Arc<dyn LlmProvider>,
|
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
|
store: Option<AdminScope>,
|
|
consecutive_failures: u32,
|
|
}
|
|
|
|
impl HeartbeatRunner {
|
|
/// Create a new heartbeat runner.
|
|
pub fn new(
|
|
config: HeartbeatConfig,
|
|
hygiene_config: HygieneConfig,
|
|
workspace: Arc<Workspace>,
|
|
llm: Arc<dyn LlmProvider>,
|
|
) -> Self {
|
|
Self {
|
|
config,
|
|
hygiene_config,
|
|
workspace,
|
|
llm,
|
|
response_tx: None,
|
|
store: None,
|
|
consecutive_failures: 0,
|
|
}
|
|
}
|
|
|
|
/// Set the response channel for notifications.
|
|
pub fn with_response_channel(mut self, tx: mpsc::Sender<OutgoingResponse>) -> Self {
|
|
self.response_tx = Some(tx);
|
|
self
|
|
}
|
|
|
|
/// Set the admin-scoped database store for persistent heartbeat conversations.
|
|
pub fn with_store(mut self, store: AdminScope) -> Self {
|
|
self.store = Some(store);
|
|
self
|
|
}
|
|
|
|
/// Run the heartbeat loop.
|
|
///
|
|
/// This runs forever, checking periodically based on the configured interval.
|
|
pub async fn run(&mut self) {
|
|
if !self.config.enabled {
|
|
tracing::info!("Heartbeat is disabled, not starting loop");
|
|
return;
|
|
}
|
|
|
|
// Two scheduling modes:
|
|
// fire_at → sleep until the next occurrence (recalculated each iteration)
|
|
// interval → tokio::time::interval (drift-free, accounts for loop body time)
|
|
let mut tick_interval = if self.config.fire_at.is_none() {
|
|
let mut iv = tokio::time::interval(self.config.interval);
|
|
// Don't fire immediately on startup.
|
|
iv.tick().await;
|
|
Some(iv)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
if let Some(fire_at) = self.config.fire_at {
|
|
tracing::info!(
|
|
"Starting heartbeat loop: fire daily at {:?} {:?}",
|
|
fire_at,
|
|
self.config.timezone
|
|
);
|
|
} else {
|
|
tracing::info!(
|
|
"Starting heartbeat loop with interval {:?}",
|
|
self.config.interval
|
|
);
|
|
}
|
|
|
|
loop {
|
|
if let Some(fire_at) = self.config.fire_at {
|
|
let sleep_dur = duration_until_next_fire(fire_at, self.config.resolved_tz());
|
|
tracing::info!("Next heartbeat in {:.1}h", sleep_dur.as_secs_f64() / 3600.0);
|
|
tokio::time::sleep(sleep_dur).await;
|
|
} else if let Some(ref mut iv) = tick_interval {
|
|
iv.tick().await;
|
|
}
|
|
|
|
// Skip during quiet hours
|
|
if self.config.is_quiet_hours() {
|
|
tracing::trace!("Heartbeat skipped: quiet hours");
|
|
continue;
|
|
}
|
|
|
|
// Run memory hygiene in the background so it never delays the
|
|
// heartbeat checklist. Failures are logged inside run_if_due.
|
|
let hygiene_workspace = Arc::clone(&self.workspace);
|
|
let hygiene_config = self.hygiene_config.clone();
|
|
tokio::spawn(async move {
|
|
let report =
|
|
crate::workspace::hygiene::run_if_due(&hygiene_workspace, &hygiene_config)
|
|
.await;
|
|
if report.had_work() {
|
|
tracing::info!(
|
|
daily_logs_deleted = report.daily_logs_deleted,
|
|
conversation_docs_deleted = report.conversation_docs_deleted,
|
|
"heartbeat: memory hygiene deleted stale documents"
|
|
);
|
|
}
|
|
});
|
|
|
|
match self.check_heartbeat().await {
|
|
HeartbeatResult::Ok => {
|
|
tracing::trace!("Heartbeat OK");
|
|
self.consecutive_failures = 0;
|
|
}
|
|
HeartbeatResult::NeedsAttention(message) => {
|
|
tracing::info!("Heartbeat needs attention: {}", message);
|
|
self.consecutive_failures = 0;
|
|
self.send_notification(&message).await;
|
|
}
|
|
HeartbeatResult::Skipped => {
|
|
tracing::trace!("Heartbeat skipped");
|
|
}
|
|
HeartbeatResult::Failed(error) => {
|
|
tracing::error!("Heartbeat failed: {}", error);
|
|
self.consecutive_failures += 1;
|
|
|
|
if self.consecutive_failures >= self.config.max_failures {
|
|
tracing::error!(
|
|
"Heartbeat disabled after {} consecutive failures",
|
|
self.consecutive_failures
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Run a single heartbeat check.
|
|
pub async fn check_heartbeat(&self) -> HeartbeatResult {
|
|
// Get the heartbeat checklist
|
|
let checklist = match self.workspace.heartbeat_checklist().await {
|
|
Ok(Some(content)) if !is_effectively_empty(&content) => content,
|
|
Ok(_) => return HeartbeatResult::Skipped,
|
|
Err(e) => return HeartbeatResult::Failed(format!("Failed to read checklist: {}", e)),
|
|
};
|
|
|
|
// Build the heartbeat prompt
|
|
let prompt = format!(
|
|
"Read the HEARTBEAT.md checklist below and follow it strictly. \
|
|
Do not infer or repeat old tasks. Check each item and report findings.\n\
|
|
\n\
|
|
If nothing needs attention, reply EXACTLY with: HEARTBEAT_OK\n\
|
|
\n\
|
|
If something needs attention, provide a concise summary of what needs action.\n\
|
|
\n\
|
|
## HEARTBEAT.md\n\
|
|
\n\
|
|
{}",
|
|
checklist
|
|
);
|
|
|
|
// Get the system prompt for context
|
|
let system_prompt = match self.workspace.system_prompt().await {
|
|
Ok(p) => p,
|
|
Err(e) => {
|
|
tracing::warn!("Failed to get system prompt for heartbeat: {}", e);
|
|
String::new()
|
|
}
|
|
};
|
|
|
|
// Run the agent turn
|
|
let messages = if system_prompt.is_empty() {
|
|
vec![ChatMessage::user(&prompt)]
|
|
} else {
|
|
vec![
|
|
ChatMessage::system(&system_prompt),
|
|
ChatMessage::user(&prompt),
|
|
]
|
|
};
|
|
|
|
// Use the model's context_length to set max_tokens. The API returns
|
|
// the total context window; we cap output at half of that (the rest is
|
|
// the prompt) with a floor of 4096.
|
|
let max_tokens = match self.llm.model_metadata().await {
|
|
Ok(meta) => {
|
|
let from_api = meta.context_length.map(|ctx| ctx / 2).unwrap_or(4096);
|
|
from_api.max(4096)
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Could not fetch model metadata, using default max_tokens: {}",
|
|
e
|
|
);
|
|
4096
|
|
}
|
|
};
|
|
|
|
let request = CompletionRequest::new(messages)
|
|
.with_max_tokens(max_tokens)
|
|
.with_temperature(0.3);
|
|
|
|
let reasoning =
|
|
Reasoning::new(self.llm.clone()).with_model_name(self.llm.active_model_name());
|
|
let (content, _usage) = match reasoning.complete(request).await {
|
|
Ok(r) => r,
|
|
Err(e) => return HeartbeatResult::Failed(format!("LLM call failed: {}", e)),
|
|
};
|
|
|
|
let content = content.trim();
|
|
|
|
// Guard against empty content. Reasoning models (e.g. GLM-4.7) may
|
|
// burn all output tokens on chain-of-thought and return content: null.
|
|
if content.is_empty() {
|
|
return HeartbeatResult::Failed("LLM returned empty content.".to_string());
|
|
}
|
|
|
|
// Check if nothing needs attention
|
|
if content == "HEARTBEAT_OK" || content.contains("HEARTBEAT_OK") {
|
|
return HeartbeatResult::Ok;
|
|
}
|
|
|
|
HeartbeatResult::NeedsAttention(content.to_string())
|
|
}
|
|
|
|
/// Send a notification about heartbeat findings.
|
|
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;
|
|
};
|
|
|
|
let user_id = self
|
|
.config
|
|
.notify_user_id
|
|
.as_deref()
|
|
.unwrap_or_else(|| self.workspace.user_id());
|
|
|
|
// Persist to heartbeat conversation and get thread_id
|
|
let thread_id = if let Some(ref store) = self.store {
|
|
match store.get_or_create_heartbeat_conversation(user_id).await {
|
|
Ok(conv_id) => {
|
|
if let Err(e) = store
|
|
.add_conversation_message(conv_id, "assistant", message)
|
|
.await
|
|
{
|
|
tracing::error!("Failed to persist heartbeat message: {}", e);
|
|
}
|
|
Some(conv_id.to_string())
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("Failed to get heartbeat conversation: {}", e);
|
|
None
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let response = OutgoingResponse {
|
|
content: format!("🔔 *Heartbeat Alert*\n\n{}", message),
|
|
thread_id,
|
|
attachments: Vec::new(),
|
|
metadata: serde_json::json!({
|
|
"source": "heartbeat",
|
|
"owner_id": self.workspace.user_id(),
|
|
}),
|
|
};
|
|
|
|
if let Err(e) = tx.send(response).await {
|
|
tracing::error!("Failed to send heartbeat notification: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Check if heartbeat content is effectively empty.
|
|
///
|
|
/// Returns true if the content contains only:
|
|
/// - Whitespace
|
|
/// - Markdown headers (lines starting with #)
|
|
/// - HTML comments (`<!-- ... -->`)
|
|
/// - Empty list items (`- [ ]`, `- [x]`, `-`, `*`)
|
|
///
|
|
/// This skips the LLM call when the user hasn't added real tasks yet,
|
|
/// saving API costs.
|
|
fn is_effectively_empty(content: &str) -> bool {
|
|
let without_comments = strip_html_comments(content);
|
|
|
|
without_comments.lines().all(|line| {
|
|
let trimmed = line.trim();
|
|
trimmed.is_empty()
|
|
|| trimmed.starts_with('#')
|
|
|| trimmed == "- [ ]"
|
|
|| trimmed == "- [x]"
|
|
|| trimmed == "-"
|
|
|| trimmed == "*"
|
|
})
|
|
}
|
|
|
|
/// Remove HTML comments from content.
|
|
fn strip_html_comments(content: &str) -> String {
|
|
let mut result = String::with_capacity(content.len());
|
|
let mut rest = content;
|
|
while let Some(start) = rest.find("<!--") {
|
|
result.push_str(&rest[..start]);
|
|
match rest[start..].find("-->") {
|
|
Some(end) => rest = &rest[start + end + 3..],
|
|
None => return result, // unclosed comment, treat rest as comment
|
|
}
|
|
}
|
|
result.push_str(rest);
|
|
result
|
|
}
|
|
|
|
/// Spawn the heartbeat runner as a background task.
|
|
///
|
|
/// Returns a handle that can be used to stop the runner.
|
|
pub fn spawn_heartbeat(
|
|
config: HeartbeatConfig,
|
|
hygiene_config: HygieneConfig,
|
|
workspace: Arc<Workspace>,
|
|
llm: Arc<dyn LlmProvider>,
|
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
|
store: Option<AdminScope>,
|
|
) -> tokio::task::JoinHandle<()> {
|
|
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
|
if let Some(tx) = response_tx {
|
|
runner = runner.with_response_channel(tx);
|
|
}
|
|
if let Some(s) = store {
|
|
runner = runner.with_store(s);
|
|
}
|
|
|
|
tokio::spawn(async move {
|
|
runner.run().await;
|
|
})
|
|
}
|
|
|
|
/// 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,
|
|
llm: Arc<dyn LlmProvider>,
|
|
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
|
store: AdminScope,
|
|
) -> 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
|
|
};
|
|
|
|
// 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 {
|
|
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;
|
|
}
|
|
};
|
|
|
|
// 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 {
|
|
// 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, Arc::clone(store.db())));
|
|
|
|
// Run memory hygiene per user (same as single-user heartbeat).
|
|
let hygiene_ws = Arc::clone(&workspace);
|
|
let hygiene_cfg = hygiene_config.clone();
|
|
let hygiene_user = user_id.clone();
|
|
tokio::spawn(async move {
|
|
let report =
|
|
crate::workspace::hygiene::run_if_due(&hygiene_ws, &hygiene_cfg).await;
|
|
if report.had_work() {
|
|
tracing::info!(
|
|
user_id = hygiene_user,
|
|
daily_logs_deleted = report.daily_logs_deleted,
|
|
conversation_docs_deleted = report.conversation_docs_deleted,
|
|
"multi-user heartbeat: memory hygiene deleted stale documents"
|
|
);
|
|
}
|
|
});
|
|
|
|
// 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();
|
|
let llm_clone = llm.clone();
|
|
let tx = response_tx.clone();
|
|
let admin = store.clone();
|
|
|
|
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);
|
|
}
|
|
runner = runner.with_store(admin);
|
|
|
|
let result = runner.check_heartbeat().await;
|
|
if let HeartbeatResult::NeedsAttention(msg) = &result {
|
|
runner.send_notification(msg).await;
|
|
}
|
|
(uid, result)
|
|
});
|
|
}
|
|
|
|
// Collect remaining results and update failure counts
|
|
while let Some(join_result) = join_set.join_next().await {
|
|
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<String, u32>,
|
|
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::*;
|
|
|
|
#[test]
|
|
fn test_heartbeat_config_defaults() {
|
|
let config = HeartbeatConfig::default();
|
|
assert!(config.enabled);
|
|
assert_eq!(config.interval, Duration::from_secs(30 * 60));
|
|
assert_eq!(config.max_failures, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_heartbeat_config_builders() {
|
|
let config = HeartbeatConfig::default()
|
|
.with_interval(Duration::from_secs(60))
|
|
.with_notify("user1", "telegram");
|
|
|
|
assert_eq!(config.interval, Duration::from_secs(60));
|
|
assert_eq!(config.notify_user_id, Some("user1".to_string()));
|
|
assert_eq!(config.notify_channel, Some("telegram".to_string()));
|
|
|
|
let disabled = HeartbeatConfig::default().disabled();
|
|
assert!(!disabled.enabled);
|
|
}
|
|
|
|
// ==================== strip_html_comments ====================
|
|
|
|
#[test]
|
|
fn test_strip_html_comments_no_comments() {
|
|
assert_eq!(strip_html_comments("hello world"), "hello world");
|
|
}
|
|
|
|
#[test]
|
|
fn test_strip_html_comments_single() {
|
|
assert_eq!(
|
|
strip_html_comments("before<!-- gone -->after"),
|
|
"beforeafter"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_strip_html_comments_multiple() {
|
|
let input = "a<!-- 1 -->b<!-- 2 -->c";
|
|
assert_eq!(strip_html_comments(input), "abc");
|
|
}
|
|
|
|
#[test]
|
|
fn test_strip_html_comments_multiline() {
|
|
let input = "# Title\n<!-- multi\nline\ncomment -->\nreal content";
|
|
assert_eq!(strip_html_comments(input), "# Title\n\nreal content");
|
|
}
|
|
|
|
#[test]
|
|
fn test_strip_html_comments_unclosed() {
|
|
let input = "before<!-- never closed";
|
|
assert_eq!(strip_html_comments(input), "before");
|
|
}
|
|
|
|
// ==================== is_effectively_empty ====================
|
|
|
|
#[test]
|
|
fn test_effectively_empty_empty_string() {
|
|
assert!(is_effectively_empty(""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_whitespace() {
|
|
assert!(is_effectively_empty(" \n\n \n "));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_headers_only() {
|
|
assert!(is_effectively_empty("# Title\n## Subtitle\n### Section"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_html_comments_only() {
|
|
assert!(is_effectively_empty("<!-- this is a comment -->"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_empty_checkboxes() {
|
|
assert!(is_effectively_empty("# Checklist\n- [ ]\n- [x]"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_bare_list_markers() {
|
|
assert!(is_effectively_empty("-\n*\n-"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_seeded_template() {
|
|
let template = "\
|
|
# Heartbeat Checklist
|
|
|
|
<!-- Keep this file empty to skip heartbeat API calls.
|
|
Add tasks below when you want the agent to check something periodically.
|
|
|
|
Example:
|
|
- [ ] Check for unread emails needing a reply
|
|
- [ ] Review today's calendar for upcoming meetings
|
|
- [ ] Check CI build status for main branch
|
|
-->";
|
|
assert!(is_effectively_empty(template));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_real_checklist() {
|
|
let content = "\
|
|
# Heartbeat Checklist
|
|
|
|
- [ ] Check for unread emails needing a reply
|
|
- [ ] Review today's calendar for upcoming meetings";
|
|
assert!(!is_effectively_empty(content));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_mixed_real_and_headers() {
|
|
let content = "# Title\n\nDo something important";
|
|
assert!(!is_effectively_empty(content));
|
|
}
|
|
|
|
#[test]
|
|
fn test_effectively_empty_comment_plus_real_content() {
|
|
let content = "<!-- comment -->\nActual task here";
|
|
assert!(!is_effectively_empty(content));
|
|
}
|
|
|
|
// ==================== quiet hours ====================
|
|
|
|
#[test]
|
|
fn test_quiet_hours_inside() {
|
|
use chrono::{Timelike, Utc};
|
|
|
|
let now_utc = Utc::now();
|
|
let hour = now_utc.hour();
|
|
let start = hour;
|
|
let end = (hour + 1) % 24;
|
|
|
|
let config = HeartbeatConfig {
|
|
quiet_hours_start: Some(start),
|
|
quiet_hours_end: Some(end),
|
|
timezone: Some("UTC".to_string()),
|
|
..HeartbeatConfig::default()
|
|
};
|
|
// Current UTC hour is inside [start, end) by construction
|
|
assert!(config.is_quiet_hours());
|
|
}
|
|
|
|
#[test]
|
|
fn test_quiet_hours_outside() {
|
|
use chrono::{Timelike, Utc};
|
|
|
|
let now_utc = Utc::now();
|
|
let hour = now_utc.hour();
|
|
let start = (hour + 1) % 24;
|
|
let end = (hour + 2) % 24;
|
|
|
|
let config = HeartbeatConfig {
|
|
quiet_hours_start: Some(start),
|
|
quiet_hours_end: Some(end),
|
|
timezone: Some("UTC".to_string()),
|
|
..HeartbeatConfig::default()
|
|
};
|
|
// Current UTC hour is outside [start, end) by construction
|
|
assert!(!config.is_quiet_hours());
|
|
}
|
|
|
|
#[test]
|
|
fn test_quiet_hours_wraparound_excludes_now() {
|
|
use chrono::{Timelike, Utc};
|
|
|
|
let now_utc = Utc::now();
|
|
let hour = now_utc.hour();
|
|
// Window covers all hours except the current one
|
|
let start = (hour + 1) % 24;
|
|
let end = hour;
|
|
|
|
let config = HeartbeatConfig {
|
|
quiet_hours_start: Some(start),
|
|
quiet_hours_end: Some(end),
|
|
timezone: Some("UTC".to_string()),
|
|
..HeartbeatConfig::default()
|
|
};
|
|
assert!(!config.is_quiet_hours());
|
|
}
|
|
|
|
#[test]
|
|
fn test_quiet_hours_none_configured() {
|
|
let config = HeartbeatConfig::default();
|
|
assert!(!config.is_quiet_hours());
|
|
}
|
|
|
|
#[test]
|
|
fn test_quiet_hours_same_start_end() {
|
|
let config = HeartbeatConfig {
|
|
quiet_hours_start: Some(10),
|
|
quiet_hours_end: Some(10),
|
|
timezone: Some("UTC".to_string()),
|
|
..HeartbeatConfig::default()
|
|
};
|
|
// start == end means zero-width window, should be false
|
|
assert!(!config.is_quiet_hours());
|
|
}
|
|
|
|
#[test]
|
|
fn test_spawn_heartbeat_accepts_store_param() {
|
|
// Regression: spawn_heartbeat must accept an optional Database store
|
|
// for persisting heartbeat notifications to a dedicated conversation.
|
|
// Compile-time check: the 7th parameter is `Option<Arc<dyn Database>>`.
|
|
#[allow(clippy::type_complexity)]
|
|
let _fn_ptr: fn(
|
|
HeartbeatConfig,
|
|
HygieneConfig,
|
|
Arc<crate::workspace::Workspace>,
|
|
Arc<dyn crate::llm::LlmProvider>,
|
|
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
|
|
Option<AdminScope>,
|
|
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
|
|
let _ = _fn_ptr;
|
|
}
|
|
|
|
// ==================== fire_at scheduling ====================
|
|
|
|
#[test]
|
|
fn test_default_config_has_no_fire_at() {
|
|
let config = HeartbeatConfig::default();
|
|
assert!(config.fire_at.is_none());
|
|
// Interval-based scheduling should be the default
|
|
assert_eq!(config.interval, Duration::from_secs(30 * 60));
|
|
}
|
|
|
|
#[test]
|
|
fn test_with_fire_at_builder() {
|
|
let time = chrono::NaiveTime::from_hms_opt(9, 0, 0).unwrap();
|
|
let config =
|
|
HeartbeatConfig::default().with_fire_at(time, Some("Pacific/Auckland".to_string()));
|
|
assert_eq!(config.fire_at, Some(time));
|
|
assert_eq!(config.timezone, Some("Pacific/Auckland".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_duration_until_next_fire_is_bounded() {
|
|
// Result must always be between 1 second and ~24 hours
|
|
let time = chrono::NaiveTime::from_hms_opt(14, 0, 0).unwrap();
|
|
let dur = duration_until_next_fire(time, chrono_tz::UTC);
|
|
assert!(dur.as_secs() >= 1, "duration must be at least 1 second");
|
|
assert!(
|
|
dur.as_secs() <= 86_401,
|
|
"duration must be at most ~24 hours, got {}s",
|
|
dur.as_secs()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_duration_until_next_fire_dst_timezone_no_panic() {
|
|
// Use a timezone with DST (US Eastern) — should never panic
|
|
let tz: Tz = "America/New_York".parse().unwrap();
|
|
// Test a range of times including midnight boundaries
|
|
for hour in [0, 2, 3, 12, 23] {
|
|
let time = chrono::NaiveTime::from_hms_opt(hour, 30, 0).unwrap();
|
|
let dur = duration_until_next_fire(time, tz);
|
|
assert!(dur.as_secs() >= 1);
|
|
assert!(dur.as_secs() <= 86_401);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_resolved_tz_defaults_to_utc() {
|
|
let config = HeartbeatConfig::default();
|
|
assert_eq!(config.resolved_tz(), chrono_tz::UTC);
|
|
}
|
|
|
|
#[test]
|
|
fn test_resolved_tz_parses_iana() {
|
|
let time = chrono::NaiveTime::from_hms_opt(9, 0, 0).unwrap();
|
|
let config =
|
|
HeartbeatConfig::default().with_fire_at(time, Some("Europe/London".to_string()));
|
|
assert_eq!(config.resolved_tz(), chrono_tz::Europe::London);
|
|
}
|
|
}
|