improv_server/
scheduler.rs

1//! Hosted refresh scheduler: a background tokio task mirroring the CLI's
2//! standalone `improv serve-refresh` daemon (`crates/cli/src/main.rs`,
3//! `cmd_serve_refresh`), but running in-process alongside the API instead of
4//! as a separate process.
5//!
6//! # Scope decision: CALL measures only, not SQL
7//!
8//! The CLI daemon refreshes two measure kinds: external-function (`CALL`)
9//! measures via [`improv_engine::external::refresh_external_measure`], and
10//! SQL-sourced measures via `improv_storage_sql::refresh_sql_measure`, the
11//! latter needing a *source connection* (a `rusqlite::Connection` opened from
12//! an optional `source.sqlite` CLI argument). The hosted server has no
13//! per-request or per-server notion of a SQL source connection string —
14//! `improv_core_model::SqlSource` (the model's SQL-measure metadata) carries
15//! only the query/column mapping, not a connection string, so inventing one
16//! here would be scope creep the model doesn't support. The hosted scheduler
17//! therefore refreshes only `CALL` (external-function) measures
18//! automatically; SQL-sourced measures still need `improv serve-refresh
19//! <db> <source.sqlite>` or a manual refresh path.
20//!
21//! # Consent: off by default, allowlisted functions only
22//!
23//! Refreshing a `CALL` measure executes code stored *in the model*, so serving
24//! a shared model must not run it implicitly. The scheduler only starts when
25//! `IMPROV_SCHEDULER=1`, and then only runs functions whose
26//! [`fn_digest`] (`sha256(name + "\0" + body)`, lowercase hex) is listed, one
27//! per line, in the file named by `IMPROV_EXTFN_ALLOW`. Anything else is
28//! skipped with a logged warning.
29use std::collections::{HashMap, HashSet};
30use std::sync::{Arc, Mutex};
31use std::time::{Duration, Instant};
32
33use improv_core_model::{MeasureId, RefreshPolicy};
34use improv_store::Store;
35
36use crate::AppState;
37
38/// Snapshot of the scheduler's last tick, for the `/scheduler/status` route.
39#[derive(Clone, Default)]
40pub struct SchedulerStatus {
41    pub tick_secs: u64,
42    /// Ids of measures refreshed on the most recent tick (empty if none were
43    /// due, or if the tick errored before refreshing anything).
44    pub last_tick_refreshed: Vec<u32>,
45    /// Error from the most recent tick, if any (cleared on a successful tick).
46    pub last_error: Option<String>,
47}
48
49pub type SharedStatus = Arc<Mutex<SchedulerStatus>>;
50
51/// Operator allowlist of external-function digests (see [`fn_digest`]).
52pub type Allowlist = Arc<HashSet<String>>;
53
54/// `IMPROV_SCHEDULER` value -> enabled? Only an explicit `1` turns it on.
55pub fn scheduler_enabled(env: Option<&str>) -> bool {
56    env.map(str::trim) == Some("1")
57}
58
59/// `sha256(name + "\0" + body)` as lowercase hex: the allowlist key.
60pub fn fn_digest(name: &str, body: &str) -> String {
61    use sha2::{Digest, Sha256};
62    let mut h = Sha256::new();
63    h.update(name.as_bytes());
64    h.update([0u8]);
65    h.update(body.as_bytes());
66    h.finalize().iter().map(|b| format!("{b:02x}")).collect()
67}
68
69/// Parse an allowlist file: one hex digest per line; blank lines and `#`
70/// comments ignored; case-insensitive.
71pub fn parse_allowlist(text: &str) -> HashSet<String> {
72    text.lines()
73        .map(str::trim)
74        .filter(|l| !l.is_empty() && !l.starts_with('#'))
75        .map(str::to_ascii_lowercase)
76        .collect()
77}
78
79/// Spawn the background refresh loop: every `tick`, reload the model, refresh
80/// any due `CALL` measures, and save if anything changed. Never panics or
81/// exits on a tick error — it logs (`eprintln!`, matching the CLI daemon's
82/// style) and tries again next tick.
83pub fn spawn_scheduler(
84    state: Arc<AppState>,
85    tick: Duration,
86    status: SharedStatus,
87    allow: Allowlist,
88) -> tokio::task::JoinHandle<()> {
89    tokio::spawn(async move {
90        let start = Instant::now();
91        // last_run per measure, in whole seconds since `start` — lives for the
92        // task's lifetime, tracked across iterations like the CLI daemon's
93        // `HashMap<MeasureId, u64>`.
94        let mut last_run: HashMap<MeasureId, u64> = HashMap::new();
95        eprintln!(
96            "scheduler: refreshing due CALL measures every {}s",
97            tick.as_secs()
98        );
99        loop {
100            tokio::time::sleep(tick).await;
101            match tick_once(&state, &allow, &mut last_run, start).await {
102                Ok(refreshed) => {
103                    if let Ok(mut s) = status.lock() {
104                        s.last_tick_refreshed = refreshed.iter().map(|m| m.0).collect();
105                        s.last_error = None;
106                    }
107                }
108                Err(e) => {
109                    eprintln!("scheduler: tick error (continuing): {e}");
110                    if let Ok(mut s) = status.lock() {
111                        s.last_tick_refreshed.clear();
112                        s.last_error = Some(e);
113                    }
114                }
115            }
116        }
117    })
118}
119
120/// One scheduler tick: load the model, refresh due CALL measures, save if any
121/// were refreshed. Returns the ids refreshed. All I/O runs on a blocking
122/// thread, matching the request handlers' `spawn_blocking` pattern.
123async fn tick_once(
124    state: &Arc<AppState>,
125    allow: &Allowlist,
126    last_run: &mut HashMap<MeasureId, u64>,
127    start: Instant,
128) -> Result<Vec<MeasureId>, String> {
129    let now_secs = start.elapsed().as_secs();
130    let path = state.db_path.clone();
131    // Same lock as the API writers: this is a load-modify-save too.
132    let _write = state.write_lock.lock().await;
133    let allow = allow.clone();
134
135    // Load the model, compute due CALL measures against `last_run`, and
136    // refresh them, all on a blocking thread (Mentat's store and the external
137    // runtime's subprocess calls are blocking).
138    let refreshed = tokio::task::spawn_blocking({
139        let last_run_snapshot = last_run.clone();
140        move || -> Result<Vec<MeasureId>, String> {
141            let mut store = Store::open(&path).map_err(|e| e.to_string())?;
142            let mut model = store.load_model().map_err(|e| e.to_string())?;
143
144            let policies: HashMap<MeasureId, RefreshPolicy> = model
145                .external_calls
146                .iter()
147                .map(|(id, c)| (*id, c.refresh_policy))
148                .collect();
149            let due =
150                improv_core_model::schedule::due_measures(&policies, now_secs, &last_run_snapshot);
151
152            // Belt and braces: extfn no longer derives the policy from the
153            // model-controlled `pure` flag (S2), but the allowlist digest
154            // covers only name + body, so pin it anyway. Restored before save.
155            let authored: Vec<(String, bool)> = model
156                .external_fns
157                .iter_mut()
158                .map(|(k, f)| (k.clone(), std::mem::replace(&mut f.pure, true)))
159                .collect();
160
161            let mut refreshed = Vec::new();
162            for mid in &due {
163                let allowed = model
164                    .external_calls
165                    .get(mid)
166                    .and_then(|c| model.external_fns.get(&c.func))
167                    .is_some_and(|f| allow.contains(&fn_digest(&f.name, &f.body)));
168                if !allowed {
169                    eprintln!(
170                        "scheduler: WARNING skipping measure {}: its function is not in \
171                         the IMPROV_EXTFN_ALLOW allowlist",
172                        mid.0
173                    );
174                    continue;
175                }
176                match improv_engine::external::refresh_external_measure(
177                    &mut model,
178                    improv_engine::external::DEFAULT_TIMEOUT,
179                    *mid,
180                ) {
181                    Ok(n) => {
182                        eprintln!(
183                            "scheduler: [{now_secs}s] refreshed measure {} ({n} cells)",
184                            mid.0
185                        );
186                        refreshed.push(*mid);
187                    }
188                    Err(e) => {
189                        // One measure's refresh failing must not stop the
190                        // others, nor crash the tick: log and continue.
191                        eprintln!("scheduler: measure {} refresh failed: {e}", mid.0);
192                    }
193                }
194            }
195            for (k, pure) in authored {
196                if let Some(f) = model.external_fns.get_mut(&k) {
197                    f.pure = pure;
198                }
199            }
200            if !refreshed.is_empty() {
201                store.save_model(&model).map_err(|e| e.to_string())?;
202            }
203            Ok(refreshed)
204        }
205    })
206    .await
207    .map_err(|e| format!("task join: {e}"))??;
208
209    if !refreshed.is_empty() {
210        // The store changed under the API's cached model: drop it (P4).
211        *state.cache.write().await = None;
212    }
213    for mid in &refreshed {
214        last_run.insert(*mid, now_secs);
215    }
216    Ok(refreshed)
217}
218
219#[cfg(test)]
220#[cfg_attr(coverage_nightly, coverage(off))]
221mod tests {
222    use super::*;
223    use improv_core_model::{
224        CategoryId, ExternalFn, ItemId, Language, Measure, MeasureKind, Model, Name, Value,
225        ValueType,
226    };
227    use improv_store::Store;
228
229    fn python_available() -> bool {
230        // python3 must run through the host sandbox: Restricted fails closed
231        // without a working bwrap, and the test then skips.
232        let probe = ExternalFn {
233            name: "probe".into(),
234            language: Language::Python,
235            body: "result = 1".into(),
236            arg_types: vec![],
237            return_type: ValueType::Number,
238            pure: true,
239        };
240        improv_extfn::eval(&probe, &[], std::time::Duration::from_secs(5)).is_ok()
241    }
242
243    fn temp_db_path(tag: &str) -> std::path::PathBuf {
244        let dir = std::env::temp_dir().join(format!("improv-sched-test-{}", std::process::id()));
245        std::fs::create_dir_all(&dir).unwrap();
246        let path = dir.join(format!(
247            "{tag}-{}.db",
248            std::time::SystemTime::now()
249                .duration_since(std::time::UNIX_EPOCH)
250                .unwrap()
251                .as_nanos()
252        ));
253        let _ = std::fs::remove_file(&path);
254        path
255    }
256
257    // A model with one input measure `X` and a `CALL(double, X)` measure `H`
258    // with the given refresh policy.
259    fn call_model(policy: RefreshPolicy) -> Model {
260        let mut m = Model::new();
261        let product = CategoryId(2);
262        m.add_category(product, "Product");
263        m.add_item(ItemId(20), product, "A");
264        m.add_measure(Measure {
265            id: MeasureId(100),
266            name: Name("X".into()),
267            value_type: ValueType::Number,
268            categories: vec![product],
269            kind: MeasureKind::Input,
270            description: None,
271        });
272        m.set_input(
273            MeasureId(100),
274            improv_core_model::Coordinate::from_pairs([(product, ItemId(20))]),
275            Value::Number(3.0),
276        );
277        m.add_measure(Measure {
278            id: MeasureId(200),
279            name: Name("H".into()),
280            value_type: ValueType::Number,
281            categories: vec![product],
282            kind: MeasureKind::Input, // populated by refresh, like a CALL measure
283            description: None,
284        });
285        m.external_fns.insert(
286            "double".into(),
287            ExternalFn {
288                name: "double".into(),
289                language: Language::Python,
290                body: "result = args[0] * 2".into(),
291                arg_types: vec![ValueType::Number],
292                return_type: ValueType::Number,
293                pure: true,
294            },
295        );
296        m.external_calls.insert(
297            MeasureId(200),
298            improv_core_model::ExternalCall {
299                func: "double".into(),
300                arg_measures: vec![MeasureId(100)],
301                refresh_policy: policy,
302            },
303        );
304        m
305    }
306
307    fn app_state(db_path: &str) -> Arc<AppState> {
308        Arc::new(AppState::new(
309            db_path.to_string(),
310            crate::Auth::Disabled,
311            None,
312        ))
313    }
314
315    // The scheduler actually runs: a due CALL measure gets refreshed and the
316    // written cells land on disk.
317    #[tokio::test]
318    async fn scheduler_refreshes_due_call_measure() {
319        if !python_available() {
320            println!("skipped: python3 not found");
321            return;
322        }
323        let path = temp_db_path("due");
324        let mut store = Store::open(path.to_str().unwrap()).unwrap();
325        store
326            .save_model(&call_model(RefreshPolicy::Interval { secs: 0 }))
327            .unwrap();
328
329        let state = app_state(path.to_str().unwrap());
330        let status: SharedStatus = Arc::new(Mutex::new(SchedulerStatus {
331            tick_secs: 0,
332            ..Default::default()
333        }));
334        let f = &call_model(RefreshPolicy::Manual).external_fns["double"];
335        let allow: Allowlist = Arc::new(HashSet::from([fn_digest(&f.name, &f.body)]));
336        let handle = spawn_scheduler(state, Duration::from_millis(50), status.clone(), allow);
337
338        // Give it a few ticks to run; python3 subprocess startup can take a
339        // while under load, so poll instead of a single fixed sleep.
340        let mut refreshed_val = None;
341        for _ in 0..50 {
342            tokio::time::sleep(Duration::from_millis(100)).await;
343            let mut probe = Store::open(path.to_str().unwrap()).unwrap();
344            let model = probe.load_model().unwrap();
345            let coord = improv_core_model::Coordinate::from_pairs([(CategoryId(2), ItemId(20))]);
346            if let Some(v) = model.input(MeasureId(200), &coord) {
347                refreshed_val = Some(v.clone());
348                break;
349            }
350        }
351        // The status update happens right after the disk write, in the same
352        // task but a separate poll; give it a moment to land.
353        let mut saw_status = false;
354        for _ in 0..20 {
355            if status.lock().unwrap().last_tick_refreshed.contains(&200) {
356                saw_status = true;
357                break;
358            }
359            tokio::time::sleep(Duration::from_millis(20)).await;
360        }
361        handle.abort();
362
363        assert_eq!(
364            refreshed_val,
365            Some(Value::Number(6.0)),
366            "scheduler should have refreshed H = double(X) = 6"
367        );
368        assert!(saw_status, "status should report the refreshed measure");
369    }
370
371    // A tick with no due measures (Manual policy) is a no-op: it must not
372    // error, and must not commit a transaction (unchanged tx high-water).
373    #[tokio::test]
374    async fn no_due_measures_is_a_noop() {
375        let path = temp_db_path("noop");
376        let mut store = Store::open(path.to_str().unwrap()).unwrap();
377        store
378            .save_model(&call_model(RefreshPolicy::Manual))
379            .unwrap();
380        // The store's transaction high-water, not the file mtime: opening a
381        // Mentat >= 1.11 store touches the file (WAL checkpoint/restart) even
382        // for a read, so the mtime no longer says whether anything was written.
383        let tx_before = store.tx_high_water().unwrap();
384        drop(store);
385
386        let state = app_state(path.to_str().unwrap());
387        let mut last_run = HashMap::new();
388        let refreshed = tick_once(&state, &Allowlist::default(), &mut last_run, Instant::now())
389            .await
390            .expect("a no-op tick must not error");
391        assert!(refreshed.is_empty());
392
393        let tx_after = Store::open(path.to_str().unwrap())
394            .unwrap()
395            .tx_high_water()
396            .unwrap();
397        assert_eq!(
398            tx_before, tx_after,
399            "a tick that refreshes nothing must not write to the store"
400        );
401    }
402
403    // A refresh failure (unregistered external function) on one measure must
404    // not panic the tick, and the scheduler loop must still be alive and able
405    // to run subsequent ticks.
406    #[tokio::test]
407    async fn refresh_error_does_not_kill_the_loop() {
408        let path = temp_db_path("err");
409        let mut model = call_model(RefreshPolicy::Interval { secs: 0 });
410        // Point the CALL measure at a function that isn't registered.
411        model.external_calls.get_mut(&MeasureId(200)).unwrap().func = "nope".into();
412        let mut store = Store::open(path.to_str().unwrap()).unwrap();
413        store.save_model(&model).unwrap();
414
415        let state = app_state(path.to_str().unwrap());
416        let status: SharedStatus = Arc::new(Mutex::new(SchedulerStatus::default()));
417        let handle = spawn_scheduler(
418            state,
419            Duration::from_millis(50),
420            status,
421            Allowlist::default(),
422        );
423
424        tokio::time::sleep(Duration::from_millis(200)).await;
425        assert!(
426            !handle.is_finished(),
427            "a per-measure refresh error must not end the scheduler task"
428        );
429        handle.abort();
430    }
431
432    // S1: an unlisted function never runs, even when it is due.
433    #[tokio::test]
434    async fn unlisted_function_is_skipped() {
435        let path = temp_db_path("unlisted");
436        let mut store = Store::open(path.to_str().unwrap()).unwrap();
437        store
438            .save_model(&call_model(RefreshPolicy::Interval { secs: 0 }))
439            .unwrap();
440        let state = app_state(path.to_str().unwrap());
441        let other: Allowlist = Arc::new(HashSet::from([fn_digest("double", "evil")]));
442        let refreshed = tick_once(&state, &other, &mut HashMap::new(), Instant::now())
443            .await
444            .unwrap();
445        assert!(refreshed.is_empty(), "unlisted function must be skipped");
446        let model = Store::open(path.to_str().unwrap())
447            .unwrap()
448            .load_model()
449            .unwrap();
450        let coord = improv_core_model::Coordinate::from_pairs([(CategoryId(2), ItemId(20))]);
451        assert!(model.input(MeasureId(200), &coord).is_none());
452    }
453
454    #[test]
455    fn scheduler_is_off_by_default() {
456        assert!(!scheduler_enabled(None));
457        assert!(!scheduler_enabled(Some("0")));
458        assert!(!scheduler_enabled(Some("true")));
459        assert!(scheduler_enabled(Some("1")));
460    }
461
462    #[test]
463    fn digest_and_allowlist_parse() {
464        let d = fn_digest("double", "result = args[0] * 2");
465        assert_eq!(d.len(), 64);
466        assert_ne!(d, fn_digest("doubl", "eresult = args[0] * 2"));
467        let set = parse_allowlist(&format!("# ops\n\n  {}  \n", d.to_uppercase()));
468        assert_eq!(set, HashSet::from([d]));
469    }
470}
improv_server/scheduler.rs

improv_server/
scheduler.rs
