mlrs_core/
env.rs

1use anyhow::Result;
2use rhai::{Engine, Scope};
3
4use crate::protocol::StepResult;
5
6/// Build a Rhai engine with all RLM context functions registered.
7///
8/// The `ctx` string is stored as a Rhai variable in the returned `Scope`.
9/// Registered functions operate on a thread-local copy of the context so
10/// Rhai closures can call them without borrowing issues.
11pub fn build_engine(ctx: String) -> (Engine, Scope<'static>) {
12    use std::sync::{Arc, Mutex};
13
14    let ctx_arc: Arc<Mutex<String>> = Arc::new(Mutex::new(ctx.clone()));
15
16    let mut engine = Engine::new();
17
18    // ctx_len() -> int
19    {
20        let c = Arc::clone(&ctx_arc);
21        engine.register_fn("ctx_len", move || c.lock().unwrap().len() as i64);
22    }
23
24    // ctx_slice(start: int, end: int) -> String
25    {
26        let c = Arc::clone(&ctx_arc);
27        engine.register_fn("ctx_slice", move |start: i64, end: i64| {
28            let s = c.lock().unwrap();
29            let start = (start as usize).min(s.len());
30            let end = (end as usize).min(s.len());
31            s[start..end].to_string()
32        });
33    }
34
35    // ctx_grep(pattern: String) -> String
36    {
37        let c = Arc::clone(&ctx_arc);
38        engine.register_fn("ctx_grep", move |pattern: String| {
39            let s = c.lock().unwrap();
40            match regex::Regex::new(&pattern) {
41                Ok(re) => s
42                    .lines()
43                    .filter(|l| re.is_match(l))
44                    .collect::<Vec<_>>()
45                    .join("\n"),
46                Err(e) => format!("regex error: {e}"),
47            }
48        });
49    }
50
51    // print_cell(msg: String) — appended to cell output via a shared buffer.
52    // The buffer is drained by execute_script after the script runs.
53    let print_buf: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
54    {
55        let buf = Arc::clone(&print_buf);
56        engine.register_fn("print_cell", move |msg: String| {
57            buf.lock().unwrap().push(msg);
58        });
59    }
60
61    // `final` is a reserved keyword in Rhai — register as `done` as the exit signal.
62    // The system prompt instructs the model to call `done(answer)`.
63    let final_buf: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
64    {
65        let buf = Arc::clone(&final_buf);
66        engine.register_fn("done", move |answer: String| {
67            *buf.lock().unwrap() = Some(answer);
68        });
69    }
70
71    // Store the Arc refs on the engine's user data is not supported in rhai 1.x without
72    // a wrapper struct. Instead we return a scope with the ctx variable set, and expose
73    // print_buf / final_buf handles through thread-locals accessed by execute_script.
74    //
75    // SAFETY: We use a thread-local registry keyed by the engine pointer address as a
76    // workaround for rhai's lack of arbitrary user-data storage.
77    PRINT_BUF.with(|pb| *pb.borrow_mut() = Some(Arc::clone(&print_buf)));
78    FINAL_BUF.with(|fb| *fb.borrow_mut() = Some(Arc::clone(&final_buf)));
79
80    let mut scope = Scope::new();
81    scope.push("ctx", ctx);
82
83    (engine, scope)
84}
85
86thread_local! {
87    static PRINT_BUF: std::cell::RefCell<Option<std::sync::Arc<std::sync::Mutex<Vec<String>>>>> =
88        std::cell::RefCell::new(None);
89    static FINAL_BUF: std::cell::RefCell<Option<std::sync::Arc<std::sync::Mutex<Option<String>>>>> =
90        std::cell::RefCell::new(None);
91}
92
93/// Execute a Rhai script and return the `StepResult`.
94///
95/// On compile/runtime error the error message is returned as `Continue` so the
96/// model can self-correct (the caller counts retries).
97pub fn execute_script(engine: &Engine, scope: &mut Scope, script: &str) -> Result<StepResult> {
98    let exec_result = engine.eval_with_scope::<rhai::Dynamic>(scope, script);
99
100    // Drain print buffer regardless of success/failure.
101    let printed = PRINT_BUF.with(|pb| {
102        pb.borrow()
103            .as_ref()
104            .map(|arc| arc.lock().unwrap().drain(..).collect::<Vec<_>>().join("\n"))
105            .unwrap_or_default()
106    });
107
108    // Check if done() was called.
109    let final_answer = FINAL_BUF.with(|fb| {
110        fb.borrow()
111            .as_ref()
112            .and_then(|arc| arc.lock().unwrap().take())
113    });
114
115    if let Some(answer) = final_answer {
116        return Ok(StepResult::Final(answer));
117    }
118
119    match exec_result {
120        Ok(val) => {
121            let mut output = printed;
122            let val_str = val.to_string();
123            if !val_str.is_empty() && val_str != "()" {
124                if !output.is_empty() {
125                    output.push('\n');
126                }
127                output.push_str(&val_str);
128            }
129            Ok(StepResult::Continue(output))
130        }
131        Err(e) => Ok(StepResult::Continue(format!("script error: {e}"))),
132    }
133}
mlrs_core/env.rs

mlrs_core/
env.rs
