ares_tools/
rhai_tool.rs

1//! Rhai-based Turing-complete tool engine for A.R.E.S.
2//!
3//! Provides [`RhaiTool`] as the default sandboxed scripting engine.
4//! Scripts are compiled once to an [`rhai::AST`] and reused via
5//! [`rhai::Engine::call_fn`] inside `spawn_blocking` + timeout.
6
7use crate::registry::Tool;
8use ares_types::types::{AppError, Result};
9use async_trait::async_trait;
10use rhai::{Dynamic, Engine, Scope, AST};
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13use std::sync::Arc;
14use std::time::Duration;
15
16// =============================================================================
17// Configuration
18// =============================================================================
19
20fn default_entry() -> String {
21    "execute".to_string()
22}
23fn default_max_ops() -> u64 {
24    50000
25}
26fn default_timeout_ms() -> u64 {
27    2000
28}
29fn default_max_string_size() -> usize {
30    8192
31}
32fn default_max_call_levels() -> usize {
33    64
34}
35
36/// Configuration parsed from `execution_config` JSONB for [`RhaiTool`].
37#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
38pub struct RhaiToolConfig {
39    /// Rhai script source. Must define `fn <entry>(args)` or be a bare expression
40    /// that can be evaluated with `args` in scope.
41    pub script: String,
42
43    /// Entry function name (default `"execute"`).
44    #[serde(default)]
45    pub entry: Option<String>,
46
47    /// Max operations (default 50000, 0 = unlimited — but we clamp to default).
48    #[serde(default)]
49    pub max_ops: Option<u64>,
50
51    /// Execution timeout in milliseconds (default 2000).
52    #[serde(default)]
53    pub timeout_ms: Option<u64>,
54
55    /// Max string size in bytes (default 8192).
56    #[serde(default)]
57    pub max_string_size: Option<usize>,
58
59    /// Max call stack depth (default 64).
60    #[serde(default)]
61    pub max_call_levels: Option<usize>,
62}
63
64impl RhaiToolConfig {
65    /// Effective entry name.
66    pub fn effective_entry(&self) -> String {
67        self.entry.clone().unwrap_or_else(default_entry)
68    }
69    /// Effective max operations.
70    pub fn effective_max_ops(&self) -> u64 {
71        self.max_ops.unwrap_or_else(default_max_ops)
72    }
73    /// Effective timeout.
74    pub fn effective_timeout(&self) -> Duration {
75        Duration::from_millis(self.timeout_ms.unwrap_or_else(default_timeout_ms))
76    }
77    /// Effective max string size.
78    pub fn effective_max_string_size(&self) -> usize {
79        self.max_string_size.unwrap_or_else(default_max_string_size)
80    }
81    /// Effective max call levels.
82    pub fn effective_max_call_levels(&self) -> usize {
83        self.max_call_levels.unwrap_or_else(default_max_call_levels)
84    }
85}
86
87// =============================================================================
88// Engine helpers
89// =============================================================================
90
91fn build_engine(config: &RhaiToolConfig) -> Engine {
92    let mut engine = Engine::new();
93    // bounded limits
94    engine.set_max_operations(config.effective_max_ops());
95    engine.set_max_string_size(config.effective_max_string_size());
96    engine.set_max_call_levels(config.effective_max_call_levels());
97    // max_expr_depth 128 for both expression depth counters
98    engine.set_max_expr_depths(128, 128);
99    // disable printing to stdout
100    engine.on_print(|_| {});
101    engine.on_debug(|_, _, _| {});
102    // disallow eval likely by not registering it; also disable symbol if present
103    engine.disable_symbol("eval");
104    engine
105}
106
107fn compile_with_config(config: &RhaiToolConfig, engine: &Engine) -> Result<AST> {
108    engine
109        .compile(config.script.clone())
110        .map_err(|e| AppError::Configuration(format!("Invalid Rhai script: {e}")))
111}
112
113// =============================================================================
114// RhaiTool
115// =============================================================================
116
117/// Rhai-based tool — Turing-complete default engine.
118///
119/// Compiles `script` once to an [`AST`] and reuses it for every `execute`
120/// via `Engine::call_fn` inside `spawn_blocking` to avoid blocking the Axum
121/// worker.
122pub struct RhaiTool {
123    name: String,
124    description: String,
125    parameters_schema: Value,
126    engine: Arc<Engine>,
127    ast: AST,
128    entry: String,
129    timeout: Duration,
130}
131
132impl std::fmt::Debug for RhaiTool {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("RhaiTool")
135            .field("name", &self.name)
136            .field("entry", &self.entry)
137            .field("timeout", &self.timeout)
138            .finish()
139    }
140}
141
142impl RhaiTool {
143    /// Parse `execution_config` JSON into [`RhaiToolConfig`].
144    pub fn parse_config(execution_config: &Value) -> Result<RhaiToolConfig> {
145        serde_json::from_value(execution_config.clone())
146            .map_err(|e| AppError::Configuration(format!("Invalid Rhai tool config: {e}")))
147    }
148
149    /// Validate script syntax with bounded engine limits.
150    pub fn validate(config: &RhaiToolConfig) -> Result<()> {
151        let engine = build_engine(config);
152        compile_with_config(config, &engine).map(|_| ())
153    }
154
155    /// Create a new [`RhaiTool`] from validated config.
156    ///
157    /// Compiles the script once; the resulting [`AST`] is reused for all
158    /// executions.
159    pub fn new(
160        name: impl Into<String>,
161        description: impl Into<String>,
162        parameters_schema: Value,
163        config: RhaiToolConfig,
164    ) -> Result<Self> {
165        if config.script.trim().is_empty() {
166            return Err(AppError::Configuration(
167                "Rhai script must not be empty".to_string(),
168            ));
169        }
170        let engine = build_engine(&config);
171        let ast = compile_with_config(&config, &engine)?;
172        let entry = config.effective_entry();
173        let timeout = config.effective_timeout();
174        Ok(Self {
175            name: name.into(),
176            description: description.into(),
177            parameters_schema,
178            engine: Arc::new(engine),
179            ast,
180            entry,
181            timeout,
182        })
183    }
184
185    /// Convenience: parse + construct in one step from raw `execution_config`.
186    pub fn from_config(
187        name: impl Into<String>,
188        description: impl Into<String>,
189        parameters_schema: Value,
190        execution_config: &Value,
191    ) -> Result<Self> {
192        let cfg = Self::parse_config(execution_config)?;
193        Self::new(name, description, parameters_schema, cfg)
194    }
195
196    /// Timeout for executions.
197    pub fn timeout(&self) -> Duration {
198        self.timeout
199    }
200
201    /// Entry function name.
202    pub fn entry(&self) -> &str {
203        &self.entry
204    }
205}
206
207/// Convert a Rhai [`Dynamic`] value to [`serde_json::Value`].
208///
209/// Uses `rhai::serde::from_dynamic` for faithful conversion; falls back to
210/// stringification for opaque types. Unit `()` maps to `Null`.
211pub fn rhai_value_to_json(dynamic: &Dynamic) -> Value {
212    if dynamic.is_unit() {
213        return Value::Null;
214    }
215    // Try serde conversion to Value
216    match rhai::serde::from_dynamic::<Value>(dynamic) {
217        Ok(v) => v,
218        Err(_) => {
219            // manual fallback for common primitives
220            if let Ok(i) = dynamic.as_int() {
221                return Value::Number(i.into());
222            }
223            if let Ok(b) = dynamic.as_bool() {
224                return Value::Bool(b);
225            }
226            if let Some(f) = dynamic.clone().try_cast::<f64>() {
227                if let Some(n) = serde_json::Number::from_f64(f) {
228                    return Value::Number(n);
229                }
230                return Value::String(f.to_string());
231            }
232            if let Some(s) = dynamic.clone().try_cast::<String>() {
233                return Value::String(s);
234            }
235            // last resort
236            Value::String(dynamic.to_string())
237        }
238    }
239}
240
241// ---------------------------------------------------------------------------
242// Helpers extracted to reduce `Tool::execute` complexity (DRY, one level)
243// ---------------------------------------------------------------------------
244
245fn is_not_found_error(msg: &str) -> bool {
246    msg.contains("Function not found")
247        || msg.contains("not found")
248        || msg.contains("unknown function")
249        || msg.contains("Unable to find function")
250}
251
252fn is_arity_error(msg: &str) -> bool {
253    msg.contains("parameter") || msg.contains("argument") || msg.contains("signature")
254}
255
256fn build_scope(args: &Value) -> std::result::Result<(Dynamic, Scope<'static>), String> {
257    let dynamic =
258        rhai::serde::to_dynamic(args).map_err(|e| format!("args conversion failed: {e}"))?;
259    let mut scope = Scope::new();
260    scope.push_dynamic("args", dynamic.clone());
261    if let Some(map) = dynamic.clone().try_cast::<rhai::Map>() {
262        for (k, v) in map {
263            let _ = scope.push_dynamic(k.to_string(), v);
264        }
265    } else if let Some(obj) = args.as_object() {
266        for (k, v) in obj {
267            if let Ok(d) = rhai::serde::to_dynamic(v.clone()) {
268                let _ = scope.push_dynamic(k.clone(), d);
269            }
270        }
271    }
272    Ok((dynamic, scope))
273}
274
275fn fallback_direct_eval(
276    engine: &Engine,
277    ast: &AST,
278    entry: &str,
279    scope: &mut Scope,
280) -> std::result::Result<Dynamic, String> {
281    match engine.eval_ast_with_scope::<Dynamic>(scope, ast) {
282        Ok(v) => Ok(v),
283        Err(e2) => {
284            let msg2 = e2.to_string();
285            if is_not_found_error(&msg2) {
286                match engine.call_fn::<Dynamic>(scope, ast, entry, ()) {
287                    Ok(v) => Ok(v),
288                    Err(e3) => Err(e3.to_string()),
289                }
290            } else {
291                Err(msg2)
292            }
293        }
294    }
295}
296
297fn fallback_zero_arg(
298    engine: &Engine,
299    ast: &AST,
300    entry: &str,
301    scope: &mut Scope,
302) -> std::result::Result<Dynamic, String> {
303    match engine.call_fn::<Dynamic>(scope, ast, entry, ()) {
304        Ok(v) => Ok(v),
305        Err(e) => Err(e.to_string()),
306    }
307}
308
309fn invoke_with_fallbacks(
310    engine: &Engine,
311    ast: &AST,
312    entry: &str,
313    scope: &mut Scope,
314    arg: Dynamic,
315) -> std::result::Result<Dynamic, String> {
316    match engine.call_fn::<Dynamic>(scope, ast, entry, (arg.clone(),)) {
317        Ok(v) => Ok(v),
318        Err(e) => {
319            let msg = e.to_string();
320            if is_not_found_error(&msg) {
321                return fallback_direct_eval(engine, ast, entry, scope);
322            }
323            if is_arity_error(&msg) {
324                return fallback_zero_arg(engine, ast, entry, scope).or(Err(msg));
325            }
326            Err(msg)
327        }
328    }
329}
330
331fn execute_blocking(
332    engine: &Engine,
333    ast: &AST,
334    entry: &str,
335    args: Value,
336) -> std::result::Result<Value, String> {
337    let (arg, mut scope) = build_scope(&args)?;
338    let dynamic_result = invoke_with_fallbacks(engine, ast, entry, &mut scope, arg)?;
339    Ok(rhai_value_to_json(&dynamic_result))
340}
341
342#[async_trait]
343impl Tool for RhaiTool {
344    fn name(&self) -> &str {
345        &self.name
346    }
347
348    fn description(&self) -> &str {
349        &self.description
350    }
351
352    fn parameters_schema(&self) -> Value {
353        self.parameters_schema.clone()
354    }
355
356    async fn execute(&self, args: Value) -> Result<Value> {
357        let engine = Arc::clone(&self.engine);
358        let ast = self.ast.clone();
359        let entry = self.entry.clone();
360        let timeout_dur = self.timeout;
361        let blocking =
362            tokio::task::spawn_blocking(move || execute_blocking(&engine, &ast, &entry, args));
363        let timed = tokio::time::timeout(timeout_dur, blocking).await;
364        match timed {
365            Ok(join_res) => match join_res {
366                Ok(inner) => match inner {
367                    Ok(v) => Ok(v),
368                    Err(e) => Err(AppError::External(format!("Rhai error: {e}"))),
369                },
370                Err(join_err) => Err(AppError::Internal(format!("Rhai join error: {join_err}"))),
371            },
372            Err(_) => Err(AppError::External("Rhai execution timed out".to_string())),
373        }
374    }
375}
376
377// =============================================================================
378// Tests
379// =============================================================================
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384    use serde_json::json;
385
386    fn mk_tool(
387        script: &str,
388        entry: Option<&str>,
389        max_ops: Option<u64>,
390        timeout_ms: Option<u64>,
391    ) -> RhaiTool {
392        let cfg = RhaiToolConfig {
393            script: script.to_string(),
394            entry: entry.map(|s| s.to_string()),
395            max_ops,
396            timeout_ms,
397            max_string_size: None,
398            max_call_levels: None,
399        };
400        RhaiTool::new("test", "test tool", json!({}), cfg).expect("tool creation")
401    }
402
403    /// Shared helper that executes a Rhai script and asserts the JSON result
404    /// equals `expected`. Extracted to eliminate the 84% near-duplicate bodies
405    /// between `test_execute_simple_add` and `test_json_result` reported by
406    /// rust-doctor. Plain prose: create tool, execute, compare.
407    async fn assert_execute_success(script: &str, input: Value, expected: Value) {
408        let tool = mk_tool(script, None, None, None);
409        let out = tool
410            .execute(input.clone())
411            .await
412            .expect("execute should succeed");
413        assert_eq!(out, expected, "script `{script}` input `{input:?}`");
414    }
415
416    /// Shared helper that runs an infinite-loop script and asserts it fails
417    /// with a message containing one of `needles`. Extracted to eliminate the
418    /// 68-node exact duplicate bodies between `test_max_ops_exceeded` and
419    /// `test_timeout` (rust-doctor duplicate_function_body).
420    async fn assert_execution_fails(
421        script: &str,
422        max_ops: Option<u64>,
423        timeout_ms: Option<u64>,
424        needles: &[&str],
425    ) {
426        let tool = mk_tool(script, None, max_ops, timeout_ms);
427        let res = tool.execute(json!({})).await;
428        assert!(
429            res.is_err(),
430            "expected error for script `{script}`, got {res:?}"
431        );
432        let msg = res.unwrap_err().to_string().to_lowercase();
433        assert!(
434            needles.iter().any(|n| msg.contains(&n.to_lowercase())),
435            "msg `{msg}` should contain one of {needles:?}"
436        );
437    }
438
439    #[test]
440    fn test_parse_config_valid() {
441        let v = json!({
442            "script": "fn execute(args){ 42 }",
443            "entry": "execute",
444            "max_ops": 1000,
445            "timeout_ms": 500,
446            "max_string_size": 1024,
447            "max_call_levels": 10
448        });
449        let cfg = RhaiTool::parse_config(&v).expect("parse");
450        assert_eq!(cfg.script, "fn execute(args){ 42 }");
451        assert_eq!(cfg.entry.unwrap(), "execute");
452        assert_eq!(cfg.max_ops.unwrap(), 1000);
453        assert_eq!(cfg.timeout_ms.unwrap(), 500);
454        assert_eq!(cfg.max_string_size.unwrap(), 1024);
455        assert_eq!(cfg.max_call_levels.unwrap(), 10);
456    }
457
458    #[test]
459    fn test_parse_config_defaults() {
460        let v = json!({ "script": "fn execute(args){ 1 }" });
461        let cfg = RhaiTool::parse_config(&v).expect("parse");
462        assert_eq!(cfg.effective_entry(), "execute");
463        assert_eq!(cfg.effective_max_ops(), 50000);
464        assert_eq!(cfg.effective_timeout(), Duration::from_millis(2000));
465        assert_eq!(cfg.effective_max_string_size(), 8192);
466        assert_eq!(cfg.effective_max_call_levels(), 64);
467    }
468
469    #[test]
470    fn test_invalid_script_syntax_returns_error() {
471        let cfg = RhaiToolConfig {
472            script: "fn execute( { broken syntax".to_string(),
473            entry: None,
474            max_ops: None,
475            timeout_ms: None,
476            max_string_size: None,
477            max_call_levels: None,
478        };
479        let res = RhaiTool::validate(&cfg);
480        assert!(res.is_err(), "expected syntax error");
481        let err = res.unwrap_err();
482        assert!(
483            err.to_string().contains("Invalid Rhai script")
484                || err.to_string().contains("Configuration")
485        );
486
487        // also via ::new
488        let new_res = RhaiTool::new("t", "d", json!({}), cfg);
489        assert!(new_res.is_err());
490    }
491
492    #[tokio::test]
493    async fn test_execute_simple_add() {
494        // Rhai ints map to JSON numbers
495        assert_execute_success(
496            r#"fn execute(args){ args["a"] + args["b"] }"#,
497            json!({"a": 2, "b": 3}),
498            json!(5),
499        )
500        .await;
501    }
502
503    #[tokio::test]
504    async fn test_max_ops_exceeded() {
505        // low max_ops should trigger quickly
506        assert_execution_fails(
507            "fn execute(args){ while true {} }",
508            Some(1000),
509            Some(2000),
510            &["operation", "exceed", "rhai error"],
511        )
512        .await;
513    }
514
515    #[tokio::test]
516    async fn test_timeout() {
517        // very low timeout + huge max_ops so timeout triggers before ops limit
518        assert_execution_fails(
519            "fn execute(args){ while true {} }",
520            Some(1_000_000_000),
521            Some(50),
522            &["timed out", "timeout", "rhai"],
523        )
524        .await;
525    }
526
527    #[tokio::test]
528    async fn test_tenant_isolation() {
529        let tool = mk_tool(r#"fn execute(args){ args["tenant"] }"#, None, None, None);
530        let out_a = tool.execute(json!({"tenant": "a"})).await.expect("a");
531        let out_b = tool.execute(json!({"tenant": "b"})).await.expect("b");
532        assert_eq!(out_a, json!("a"));
533        assert_eq!(out_b, json!("b"));
534        assert_ne!(out_a, out_b);
535        // second call with a again still returns a (no cross-call leakage)
536        let out_a2 = tool.execute(json!({"tenant": "a"})).await.expect("a2");
537        assert_eq!(out_a, out_a2);
538    }
539
540    #[tokio::test]
541    async fn test_json_result() {
542        assert_execute_success(
543            r#"fn execute(args){ #{"sum": args["x"] + args["y"], "greeting": "hello " + args["name"] } }"#,
544            json!({"x": 10, "y": 5, "name": "world"}),
545            json!({"sum": 15, "greeting": "hello world"}),
546        )
547        .await;
548    }
549
550    #[tokio::test]
551    async fn test_eval_ast_fallback() {
552        // script without function, just expression using args
553        let tool = mk_tool(r#"args["a"] * 3"#, Some("nonexistent"), None, None);
554        // entry not found -> fallback to eval_ast_with_scope should produce 6
555        // Actually our fallback tries call_fn first, then eval_ast; since entry "nonexistent" not found, eval_ast will evaluate "args[\"a\"] * 3"
556        let out = tool.execute(json!({"a": 2})).await.expect("fallback eval");
557        assert_eq!(out, json!(6));
558    }
559
560    #[tokio::test]
561    async fn test_string_limit() {
562        let cfg = RhaiToolConfig {
563            script: r#"fn execute(args){ "a" + "b" }"#.to_string(),
564            entry: None,
565            max_ops: None,
566            timeout_ms: None,
567            max_string_size: Some(2),
568            max_call_levels: None,
569        };
570        // "a"+"b" => "ab" length 2 ok, but if we exceed limit later? This test just ensures creation works
571        let tool = RhaiTool::new("t", "d", json!({}), cfg).expect("tool");
572        let out = tool.execute(json!({})).await.expect("execute");
573        assert_eq!(out, json!("ab"));
574    }
575}
ares_tools/rhai_tool.rs

ares_tools/
rhai_tool.rs
