Skip to main content

leviath_scripting/
types.rs

1//! Leviath types registered in Rhai.
2
3use rhai::Engine;
4
5/// Register Leviath types in the Rhai engine.
6pub fn register_types(engine: &mut Engine) {
7    // Region kind constructors
8    engine.register_fn("region_pinned", || -> String { "pinned".to_string() });
9
10    engine.register_fn("region_temporary", || -> String { "temporary".to_string() });
11
12    engine.register_fn("region_clearable", || -> String { "clearable".to_string() });
13
14    engine.register_fn("region_sliding_window", |max_items: i64| -> rhai::Map {
15        let mut map = rhai::Map::new();
16        map.insert(
17            "kind".into(),
18            rhai::Dynamic::from("sliding_window".to_string()),
19        );
20        map.insert("max_items".into(), rhai::Dynamic::from(max_items));
21        map
22    });
23
24    engine.register_fn("region_compacting", |threshold: i64| -> rhai::Map {
25        let mut map = rhai::Map::new();
26        map.insert("kind".into(), rhai::Dynamic::from("compacting".to_string()));
27        map.insert("threshold_tokens".into(), rhai::Dynamic::from(threshold));
28        map
29    });
30
31    engine.register_fn(
32        "region_custom",
33        |script: String, persistent: bool| -> rhai::Map {
34            let mut map = rhai::Map::new();
35            map.insert("kind".into(), rhai::Dynamic::from("custom".to_string()));
36            map.insert("script".into(), rhai::Dynamic::from(script));
37            map.insert("persistent".into(), rhai::Dynamic::from(persistent));
38            map
39        },
40    );
41
42    // Region entry constructor
43    engine.register_fn(
44        "region_entry",
45        |content: String, tokens: i64| -> rhai::Map {
46            let mut map = rhai::Map::new();
47            map.insert("content".into(), rhai::Dynamic::from(content));
48            map.insert("tokens".into(), rhai::Dynamic::from(tokens));
49            map
50        },
51    );
52
53    // Content format validator
54    engine.register_fn("content_format", |format: &str| -> String {
55        match format {
56            "text" | "json" | "mermaid" | "markdown" | "code" => format.to_string(),
57            _ => "text".to_string(),
58        }
59    });
60
61    // Token budget helpers. The operands come from a script (ultimately from
62    // model output), so plain `-` would panic on overflow in a debug build -
63    // and a panic inside a Rhai native fn aborts the process (#109).
64    engine.register_fn(
65        "tokens_remaining",
66        |max_tokens: i64, current_tokens: i64| -> i64 { max_tokens.saturating_sub(current_tokens) },
67    );
68
69    engine.register_fn(
70        "usage_ratio",
71        |max_tokens: i64, current_tokens: i64| -> f64 {
72            if max_tokens == 0 {
73                return 1.0;
74            }
75            current_tokens as f64 / max_tokens as f64
76        },
77    );
78
79    engine.register_fn(
80        "needs_eviction",
81        |max_tokens: i64, current_tokens: i64, threshold: f64| -> bool {
82            if max_tokens == 0 {
83                return true;
84            }
85            (current_tokens as f64 / max_tokens as f64) >= threshold
86        },
87    );
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use rhai::Engine;
94
95    fn engine() -> Engine {
96        let mut e = Engine::new();
97        register_types(&mut e);
98        e
99    }
100
101    // --- region constructors ---
102
103    #[test]
104    fn region_pinned_returns_pinned() {
105        let e = engine();
106        let result: String = e.eval("region_pinned()").unwrap();
107        assert_eq!(result, "pinned");
108    }
109
110    #[test]
111    fn region_temporary_returns_temporary() {
112        let e = engine();
113        let result: String = e.eval("region_temporary()").unwrap();
114        assert_eq!(result, "temporary");
115    }
116
117    #[test]
118    fn region_clearable_returns_clearable() {
119        let e = engine();
120        let result: String = e.eval("region_clearable()").unwrap();
121        assert_eq!(result, "clearable");
122    }
123
124    // --- region_sliding_window ---
125
126    #[test]
127    fn region_sliding_window_returns_map_with_kind_and_max_items() {
128        let e = engine();
129        let result: rhai::Map = e.eval("region_sliding_window(10)").unwrap();
130        assert_eq!(
131            result.get("kind").unwrap().clone_cast::<String>(),
132            "sliding_window"
133        );
134        assert_eq!(result.get("max_items").unwrap().clone_cast::<i64>(), 10);
135    }
136
137    #[test]
138    fn region_sliding_window_zero() {
139        let e = engine();
140        let result: rhai::Map = e.eval("region_sliding_window(0)").unwrap();
141        assert_eq!(result.get("max_items").unwrap().clone_cast::<i64>(), 0);
142    }
143
144    // --- region_compacting ---
145
146    #[test]
147    fn region_compacting_returns_map_with_kind_and_threshold() {
148        let e = engine();
149        let result: rhai::Map = e.eval("region_compacting(5000)").unwrap();
150        assert_eq!(
151            result.get("kind").unwrap().clone_cast::<String>(),
152            "compacting"
153        );
154        assert_eq!(
155            result.get("threshold_tokens").unwrap().clone_cast::<i64>(),
156            5000
157        );
158    }
159
160    // --- region_custom ---
161
162    #[test]
163    fn region_custom_returns_map_with_script_and_persistent() {
164        let e = engine();
165        let result: rhai::Map = e.eval(r#"region_custom("hooks/conv.rhai", true)"#).unwrap();
166        assert_eq!(result.get("kind").unwrap().clone_cast::<String>(), "custom");
167        assert_eq!(
168            result.get("script").unwrap().clone_cast::<String>(),
169            "hooks/conv.rhai"
170        );
171        assert!(result.get("persistent").unwrap().clone_cast::<bool>());
172    }
173
174    #[test]
175    fn region_custom_non_persistent() {
176        let e = engine();
177        let result: rhai::Map = e.eval(r#"region_custom("r.rhai", false)"#).unwrap();
178        assert!(!result.get("persistent").unwrap().clone_cast::<bool>());
179    }
180
181    // --- region_entry ---
182
183    #[test]
184    fn region_entry_returns_map_with_content_and_tokens() {
185        let e = engine();
186        let result: rhai::Map = e.eval(r#"region_entry("content", 42)"#).unwrap();
187        assert_eq!(
188            result.get("content").unwrap().clone_cast::<String>(),
189            "content"
190        );
191        assert_eq!(result.get("tokens").unwrap().clone_cast::<i64>(), 42);
192    }
193
194    #[test]
195    fn region_entry_empty_content() {
196        let e = engine();
197        let result: rhai::Map = e.eval(r#"region_entry("", 0)"#).unwrap();
198        assert_eq!(result.get("content").unwrap().clone_cast::<String>(), "");
199        assert_eq!(result.get("tokens").unwrap().clone_cast::<i64>(), 0);
200    }
201
202    // --- content_format ---
203
204    #[test]
205    fn content_format_valid_formats() {
206        let e = engine();
207        for fmt in &["text", "json", "mermaid", "markdown", "code"] {
208            let script = format!(r#"content_format("{fmt}")"#);
209            let result: String = e.eval(&script).unwrap();
210            assert_eq!(result, *fmt);
211        }
212    }
213
214    #[test]
215    fn content_format_invalid_falls_back_to_text() {
216        let e = engine();
217        let result: String = e.eval(r#"content_format("invalid")"#).unwrap();
218        assert_eq!(result, "text");
219    }
220
221    #[test]
222    fn content_format_empty_falls_back_to_text() {
223        let e = engine();
224        let result: String = e.eval(r#"content_format("")"#).unwrap();
225        assert_eq!(result, "text");
226    }
227
228    // --- tokens_remaining ---
229
230    #[test]
231    fn tokens_remaining_basic() {
232        let e = engine();
233        let result: i64 = e.eval("tokens_remaining(100, 30)").unwrap();
234        assert_eq!(result, 70);
235    }
236
237    #[test]
238    fn tokens_remaining_zero_used() {
239        let e = engine();
240        let result: i64 = e.eval("tokens_remaining(100, 0)").unwrap();
241        assert_eq!(result, 100);
242    }
243
244    #[test]
245    fn tokens_remaining_all_used() {
246        let e = engine();
247        let result: i64 = e.eval("tokens_remaining(100, 100)").unwrap();
248        assert_eq!(result, 0);
249    }
250
251    #[test]
252    fn tokens_remaining_saturates_instead_of_overflowing() {
253        // The operands come from a script, so extreme values must not panic -
254        // a panic in a Rhai native fn aborts the daemon (issue #109).
255        let e = engine();
256        let low: i64 = e.eval("tokens_remaining(-9223372036854775808, 1)").unwrap();
257        assert_eq!(low, i64::MIN);
258        let high: i64 = e.eval("tokens_remaining(9223372036854775807, -1)").unwrap();
259        assert_eq!(high, i64::MAX);
260    }
261
262    // --- usage_ratio ---
263
264    #[test]
265    fn usage_ratio_half() {
266        let e = engine();
267        let result: f64 = e.eval("usage_ratio(100, 50)").unwrap();
268        assert!((result - 0.5).abs() < f64::EPSILON);
269    }
270
271    #[test]
272    fn usage_ratio_zero_max_returns_one() {
273        let e = engine();
274        let result: f64 = e.eval("usage_ratio(0, 50)").unwrap();
275        assert!((result - 1.0).abs() < f64::EPSILON);
276    }
277
278    #[test]
279    fn usage_ratio_none_used() {
280        let e = engine();
281        let result: f64 = e.eval("usage_ratio(100, 0)").unwrap();
282        assert!((result - 0.0).abs() < f64::EPSILON);
283    }
284
285    // --- needs_eviction ---
286
287    #[test]
288    fn needs_eviction_above_threshold() {
289        let e = engine();
290        let result: bool = e.eval("needs_eviction(100, 90, 0.8)").unwrap();
291        assert!(result);
292    }
293
294    #[test]
295    fn needs_eviction_below_threshold() {
296        let e = engine();
297        let result: bool = e.eval("needs_eviction(100, 50, 0.8)").unwrap();
298        assert!(!result);
299    }
300
301    #[test]
302    fn needs_eviction_at_exact_threshold() {
303        let e = engine();
304        let result: bool = e.eval("needs_eviction(100, 80, 0.8)").unwrap();
305        assert!(result);
306    }
307
308    #[test]
309    fn needs_eviction_zero_max_returns_true() {
310        let e = engine();
311        let result: bool = e.eval("needs_eviction(0, 0, 0.8)").unwrap();
312        assert!(result);
313    }
314}