diff --git a/src/openhuman/memory_sync/composio/periodic.rs b/src/openhuman/memory_sync/composio/periodic.rs index b7da0f811..34f66465b 100644 --- a/src/openhuman/memory_sync/composio/periodic.rs +++ b/src/openhuman/memory_sync/composio/periodic.rs @@ -52,7 +52,10 @@ use tokio::time::interval; use crate::openhuman::config::rpc as config_rpc; use crate::openhuman::config::DEFAULT_MEMORY_SYNC_INTERVAL_SECS; -use crate::openhuman::scheduler_gate::gate::current_policy; +use crate::openhuman::memory_sources::{ + memory_sync_defaults_for_toolkit, MemorySourceEntry, SourceKind, +}; +use crate::openhuman::scheduler_gate::gate::{current_policy, resume_notify}; use crate::openhuman::scheduler_gate::policy::PauseReason; use super::providers::{get_provider, ComposioUsage, ProviderContext, SyncReason}; @@ -192,6 +195,57 @@ fn persisted_since_last_sync( }) } +/// Outcome of consulting the per-source registry for one Composio connection +/// during a periodic tick (#2831). +#[derive(Debug, PartialEq, Eq)] +enum PeriodicSourceDecision { + /// The user toggled this source **off** — skip background sync entirely. + /// Manual `memory_sources_sync` still works (it has its own `enabled` + /// guard); only the automatic loop honours this here. + Skip, + /// Sync this connection with the given caps (`None` = uncapped for that + /// dimension). + Sync { + max_items: Option, + sync_depth_days: Option, + }, +} + +/// Decide whether — and with what caps — to periodically sync one connection, +/// honouring the per-source `enabled` toggle (#2831). Pure so the three +/// branches can be unit-tested without async/registry I/O. +/// +/// - **disabled row** → [`PeriodicSourceDecision::Skip`]: the background loop +/// must not sync a source the user switched off (this was the row-2 leak — +/// previously the loop only read the registry for caps and synced disabled +/// sources anyway, uncapped). +/// - **enabled row** → `Sync` with the row's caps. +/// - **no row yet** → `Sync` with conservative per-toolkit defaults +/// ([`memory_sync_defaults_for_toolkit`]). `reconcile` normally backfills an +/// enabled, capped row for every connection; this covers the brief +/// pre-reconcile window. The no-match path defaults to *sync-bounded*, never +/// *skip* and never *uncapped*, so a missing/mismatched row degrades safely +/// (data keeps flowing, just capped) instead of silently going dark. +fn decide_periodic_source( + source: Option<&MemorySourceEntry>, + toolkit: &str, +) -> PeriodicSourceDecision { + match source { + Some(s) if !s.enabled => PeriodicSourceDecision::Skip, + Some(s) => PeriodicSourceDecision::Sync { + max_items: s.max_items, + sync_depth_days: s.sync_depth_days, + }, + None => { + let (max_items, sync_depth_days) = memory_sync_defaults_for_toolkit(toolkit); + PeriodicSourceDecision::Sync { + max_items, + sync_depth_days, + } + } + } +} + /// Spawn the periodic sync background task. Idempotent: only the /// first call actually spawns the loop, every subsequent call is a /// cheap no-op (logged at `debug` so it's visible during startup @@ -221,14 +275,35 @@ pub fn start_periodic_sync() { /// Inner loop, broken out so it's easy to mock-replace in tests if we /// ever want to drive ticks deterministically. +/// +/// Each iteration waits on whichever comes first (#2831): +/// * the 20-min `ticker` — the steady-state cadence, or +/// * the scheduler-gate **resume** notify — fired when the user toggles +/// Memory Tree back on or signs back in. +/// +/// On a resume wake we run a tick **immediately** (so sync restarts within +/// seconds, not at the next ≤20-min boundary) and `reset()` the ticker so the +/// *next* scheduled tick is a full `TICK_SECONDS` out. The reset is what stops +/// rapid off-on toggling from bunch-firing: many wakes collapse into at most +/// one extra tick (the `Notify` stores a single permit), and the cadence +/// re-bases from the last actual tick. async fn run_loop() { let mut ticker = interval(Duration::from_secs(TICK_SECONDS)); + let resume = resume_notify(); // Skip the immediate-fire tick so startup isn't slammed before the // user even has time to sign in. ticker.tick().await; loop { - ticker.tick().await; + tokio::select! { + _ = ticker.tick() => {} + _ = resume.notified() => { + // Woke early on a resume transition. Re-base the cadence so the + // next scheduled tick is TICK_SECONDS from now, then fall + // through and run the tick immediately. + ticker.reset(); + } + } if let Err(e) = run_one_tick().await { tracing::warn!( error = %e, @@ -378,6 +453,28 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { let audit_index = index_last_success_by_connection(&read_audit_log(&config)); let now = Utc::now(); + // Per-source registry snapshot (#2831). The periodic loop gates on the + // per-source `enabled` toggle so a source the user switched off stops + // syncing in the background — matching the manual paths + // (`memory_sources::sync_source`, `memory_sources_sync_all`), which already + // early-return on `!enabled`. Index every Composio source (enabled and + // disabled) by connection id; the per-connection branch below resolves + // skip/caps via `decide_periodic_source`. + // + // Built from the **already-loaded** `config` snapshot (Step 1), not a second + // `list_sources()` read. A separate read whose error we swallowed to an + // empty map would make every disabled source fall through to the + // `decide_periodic_source(None, ..)` default-caps path — silently + // re-enabling background sync for sources the user switched off on a + // transient config-read failure. Reusing the tick's snapshot is fail-closed + // (a disabled row stays disabled) and avoids the extra read entirely. + let composio_sources: HashMap = config + .memory_sources + .iter() + .filter(|s| s.kind == SourceKind::Composio) + .filter_map(|s| s.connection_id.clone().map(|id| (id, s.clone()))) + .collect(); + let mut considered = 0usize; let mut fired = 0usize; for conn in resp.connections { @@ -425,20 +522,25 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { continue; } - // Look up per-source caps from the memory_sources registry. - // Non-fatal: if the lookup fails we proceed without caps. - let (src_max_items, src_sync_depth_days) = { - let registry_sources = crate::openhuman::memory_sources::list_enabled_by_kind( - crate::openhuman::memory_sources::SourceKind::Composio, - ) - .await - .unwrap_or_default(); - registry_sources - .iter() - .find(|s| s.connection_id.as_deref() == Some(conn.id.as_str())) - .map(|s| (s.max_items, s.sync_depth_days)) - .unwrap_or((None, None)) - }; + // Per-source gate + caps from the memory_sources registry (#2831). + // A disabled source is skipped here (the background-sync half of the + // toggle); enabled sources sync with their caps; a connection with no + // registry row yet syncs with conservative per-toolkit defaults. + let (src_max_items, src_sync_depth_days) = + match decide_periodic_source(composio_sources.get(&conn.id), &toolkit) { + PeriodicSourceDecision::Skip => { + tracing::debug!( + toolkit = %toolkit, + connection_id = %conn.id, + "[composio:periodic] source disabled — skipping periodic sync" + ); + continue; + } + PeriodicSourceDecision::Sync { + max_items, + sync_depth_days, + } => (max_items, sync_depth_days), + }; tracing::debug!( toolkit = %toolkit, @@ -633,6 +735,126 @@ mod tests { )); } + /// Build a minimal Composio `MemorySourceEntry` for the per-source gate + /// tests — only the fields `decide_periodic_source` reads are meaningful. + fn composio_source( + enabled: bool, + max_items: Option, + sync_depth_days: Option, + ) -> MemorySourceEntry { + MemorySourceEntry { + id: "src_test".to_string(), + kind: SourceKind::Composio, + label: "test".to_string(), + enabled, + toolkit: Some("gmail".to_string()), + connection_id: Some("cmp-1".to_string()), + path: None, + glob: None, + url: None, + branch: None, + paths: Vec::new(), + max_commits: None, + max_issues: None, + max_prs: None, + query: None, + since_days: None, + max_items, + selector: None, + max_tokens_per_sync: None, + max_cost_per_sync_usd: None, + sync_depth_days, + } + } + + /// #2831 row 2: a source explicitly toggled **off** must be skipped by the + /// background loop — this is the leak the gate closes. + #[test] + fn decide_periodic_source_skips_disabled_source() { + let src = composio_source(false, Some(100), Some(30)); + assert_eq!( + decide_periodic_source(Some(&src), "gmail"), + PeriodicSourceDecision::Skip + ); + } + + /// An enabled source syncs with exactly its configured caps (no defaulting). + #[test] + fn decide_periodic_source_uses_enabled_source_caps() { + let src = composio_source(true, Some(42), Some(7)); + assert_eq!( + decide_periodic_source(Some(&src), "gmail"), + PeriodicSourceDecision::Sync { + max_items: Some(42), + sync_depth_days: Some(7), + } + ); + } + + /// A connection with no registry row yet (pre-reconcile window) syncs with + /// the conservative per-toolkit defaults — **bounded**, never uncapped, and + /// never skipped. This is the safe-direction fallback for a missing match. + #[test] + fn decide_periodic_source_defaults_caps_when_no_row() { + let (want_items, want_depth) = memory_sync_defaults_for_toolkit("gmail"); + assert_eq!( + decide_periodic_source(None, "gmail"), + PeriodicSourceDecision::Sync { + max_items: want_items, + sync_depth_days: want_depth, + } + ); + // The defaults are bounded for a known toolkit (regression guard against + // an accidental return to uncapped background fetches). + assert!(want_items.is_some()); + } + + /// Multi-account regression (#3443 added multiple account connections per + /// toolkit): two live connections of the *same* toolkit must be gated + /// **independently** by their own per-`connection_id` source rows. This + /// pins the loop's connection-id keying — re-keying the lookup by toolkit + /// would collapse the two accounts and is the regression this guards. + #[test] + fn per_connection_gate_is_independent_across_accounts_of_same_toolkit() { + // gmail account A: enabled with caps; gmail account B: disabled. + let mut a = composio_source(true, Some(10), Some(5)); + a.connection_id = Some("conn-A".to_string()); + a.toolkit = Some("gmail".to_string()); + let mut b = composio_source(false, Some(99), Some(99)); + b.connection_id = Some("conn-B".to_string()); + b.toolkit = Some("gmail".to_string()); + + // Build the same connection_id → entry index the live tick builds. + let index: HashMap = [a, b] + .into_iter() + .filter_map(|s| s.connection_id.clone().map(|id| (id, s))) + .collect(); + + // Account A (enabled) syncs with its own caps... + assert_eq!( + decide_periodic_source(index.get("conn-A"), "gmail"), + PeriodicSourceDecision::Sync { + max_items: Some(10), + sync_depth_days: Some(5), + } + ); + // ...account B (disabled) is skipped, even though it shares the toolkit. + assert_eq!( + decide_periodic_source(index.get("conn-B"), "gmail"), + PeriodicSourceDecision::Skip + ); + // A third, not-yet-registered account of the same toolkit falls back to + // bounded defaults (never skipped, never uncapped). + let (def_items, def_depth) = memory_sync_defaults_for_toolkit("gmail"); + assert_eq!( + decide_periodic_source(index.get("conn-C"), "gmail"), + PeriodicSourceDecision::Sync { + max_items: def_items, + sync_depth_days: def_depth, + } + ); + } + /// End-to-end simulation of the scheduler's per-connection decision: prove /// that **changing the global setting changes when the next sync fires** /// (issue #3302 acceptance criterion). We drive the same two pure helpers diff --git a/src/openhuman/scheduler_gate/gate.rs b/src/openhuman/scheduler_gate/gate.rs index 96fd7a846..728893d1e 100644 --- a/src/openhuman/scheduler_gate/gate.rs +++ b/src/openhuman/scheduler_gate/gate.rs @@ -10,7 +10,7 @@ use std::sync::{Arc, OnceLock}; use std::time::Duration; use parking_lot::RwLock; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; use crate::openhuman::config::{Config, SchedulerGateConfig}; use crate::openhuman::scheduler_gate::policy::{decide, PauseReason, Policy}; @@ -240,12 +240,50 @@ pub fn init_global(config: &Config) { }); } +/// Process-wide resume signal (#2831). Fired whenever the gate transitions +/// **out of** a paused state — the user toggles Memory Tree back on +/// ([`update_config`]) or signs back in ([`set_signed_out`]). Background loops +/// (e.g. the Composio periodic scheduler) park on [`resume_notify`] so they can +/// resume work within seconds instead of waiting out their next tick boundary. +static RESUME_NOTIFY: OnceLock> = OnceLock::new(); + +/// Handle to the process-wide resume [`Notify`] (#2831). +/// +/// Both the firing side (`update_config` / `set_signed_out`) and the waiting +/// side (background loops) call this, so they share one instance. Use +/// `notify_one()` to fire: if a loop is parked it wakes immediately; if it's +/// mid-tick, a single permit is stored so the *next* `notified()` returns at +/// once — a resume that arrives during a tick is never lost. +/// +/// **Over-notifying is safe by design.** A spurious wake (e.g. Memory Tree +/// toggled on while still signed out, so the effective policy is still paused) +/// just causes one cheap gate-checked tick that re-reads [`current_policy`] and +/// no-ops. We therefore fire on each individual un-pause transition rather than +/// computing the precise combined (config × signed-out) edge. +pub fn resume_notify() -> Arc { + RESUME_NOTIFY + .get_or_init(|| Arc::new(Notify::new())) + .clone() +} + /// Update the gate's view of user config (e.g. after a settings change). +/// +/// Fires [`resume_notify`] when this update moves the policy out of a paused +/// state (e.g. Memory Tree toggled back on), so parked background loops resume +/// promptly (#2831). pub fn update_config(cfg: SchedulerGateConfig) { - if let Some(state) = STATE.get() { + let Some(state) = STATE.get() else { + return; + }; + let resumed = { let mut guard = state.write(); + let was_paused = matches!(guard.policy, Policy::Paused { .. }); guard.cfg = cfg; guard.policy = decide(&guard.signals, &guard.cfg); + was_paused && !matches!(guard.policy, Policy::Paused { .. }) + }; + if resumed { + resume_notify().notify_one(); } } @@ -311,6 +349,12 @@ pub fn set_signed_out(signed_out: bool) { let prev = SIGNED_OUT.swap(signed_out, Ordering::AcqRel); if prev != signed_out { log::info!("[scheduler_gate] signed_out {} -> {}", prev, signed_out); + // #2831: signing back in (true -> false) is a transition out of + // `Policy::Paused { SignedOut }`. Wake any periodic loop so background + // sync restarts immediately rather than at the next tick boundary. + if prev && !signed_out { + resume_notify().notify_one(); + } } } @@ -325,6 +369,11 @@ pub fn set_signed_out(signed_out: bool) { let prev = test_state::set_signed_out_for(id, signed_out); if prev != signed_out { log::info!("[scheduler_gate] signed_out {} -> {}", prev, signed_out); + // #2831: mirror the production sign-in wake so tests exercise the + // same resume-notify path (true -> false fires the loop wake). + if prev && !signed_out { + resume_notify().notify_one(); + } } } @@ -700,4 +749,89 @@ mod tests { let p2 = try_acquire_llm_permit().expect("permit free after drop"); drop(p2); } + + /// #2831: both the firing side (`update_config` / `set_signed_out`) and the + /// waiting side (the periodic loop) must observe the *same* `Notify`, so + /// `resume_notify` must hand back one process-wide instance. + #[tokio::test] + async fn resume_notify_is_a_stable_singleton() { + let _g = lock(); + assert!( + Arc::ptr_eq(&resume_notify(), &resume_notify()), + "resume_notify must return one shared instance" + ); + } + + /// A `notify_one()` wakes a task parked on `notified()` — the mechanism the + /// periodic loop relies on to resume early. Proves the singleton wiring + /// end-to-end (fire on one handle, wake on another). + #[tokio::test] + async fn resume_notify_wakes_a_parked_waiter() { + let _g = lock(); + let waiter = resume_notify(); + let parked = tokio::spawn(async move { waiter.notified().await }); + // Yield so the spawned task reaches `.notified()` before we fire. + tokio::task::yield_now().await; + resume_notify().notify_one(); + timeout(TokioDuration::from_secs(1), parked) + .await + .expect("parked waiter must wake promptly after notify_one") + .expect("waiter task must not panic"); + } + + /// #2831 wiring: a paused→running `update_config` transition and a + /// sign-in (`set_signed_out` true→false) each fire the resume notify. + /// + /// Seeds `STATE` directly (the test module can reach it) so `update_config` + /// / `set_signed_out` are live without spawning the real sampler. We drive + /// `Off → AlwaysOn`, which is deterministically `Paused → Aggressive` + /// regardless of any policy a prior test left behind, so the transition — + /// and thus the `notify_one()` — is guaranteed. + #[tokio::test] + async fn resume_transitions_fire_the_notify() { + use crate::openhuman::config::SchedulerGateMode; + let _g = lock(); + + // Ensure STATE is initialised. `set` is a no-op if an earlier test + // already promoted it — that's fine, we re-drive the transition below. + let cfg = SchedulerGateConfig { + mode: SchedulerGateMode::Off, + ..Default::default() + }; + let signals = Signals::sample(); + let policy = decide(&signals, &cfg); + let _ = STATE.set(Arc::new(RwLock::new(State { + cfg, + signals, + policy, + }))); + + // --- update_config: Paused -> running fires the notify --- + let waiter = resume_notify(); + let parked = tokio::spawn(async move { waiter.notified().await }); + tokio::task::yield_now().await; + update_config(SchedulerGateConfig { + mode: SchedulerGateMode::Off, + ..Default::default() + }); // -> Paused { UserDisabled } + update_config(SchedulerGateConfig { + mode: SchedulerGateMode::AlwaysOn, + ..Default::default() + }); // Paused -> Aggressive => resume fires + timeout(TokioDuration::from_secs(1), parked) + .await + .expect("update_config un-pause must wake the resume waiter") + .expect("waiter task must not panic"); + + // --- set_signed_out true -> false fires the notify --- + let waiter2 = resume_notify(); + let parked2 = tokio::spawn(async move { waiter2.notified().await }); + tokio::task::yield_now().await; + set_signed_out(true); + set_signed_out(false); // true -> false => resume fires + timeout(TokioDuration::from_secs(1), parked2) + .await + .expect("sign-in (signed_out true->false) must wake the resume waiter") + .expect("waiter task must not panic"); + } }