ares_tools/
script_tool.rs

1//! Runtime script tool executor for A.R.E.S.
2//!
3//! This tool executes JavaScript or Python scripts configured via `execution_config`.
4//! Scripts receive tool arguments via parameter substitution (`{{param}}`) and as an
5//! `args` global variable.
6//!
7//! # Execution environments
8//!
9//! * **JavaScript** — executed in-process via `boa_engine` (pure-Rust, sandboxed:
10//!   no filesystem or network access).
11//! * **Python** — executed in a separate OS process with a restricted builtin set
12//!   and configurable timeout / memory limits.
13
14use crate::registry::Tool;
15use ares_types::Result;
16use async_trait::async_trait;
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19use std::time::Duration;
20use tokio::process::Command;
21use tokio::time::timeout;
22
23// =============================================================================
24// Configuration
25// =============================================================================
26
27/// Script-specific configuration parsed from `execution_config` JSONB.
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
29pub struct ScriptToolConfig {
30    /// Scripting language: `"javascript"` or `"python"`.
31    pub language: String,
32    /// Script source code. Supports `{{param}}` placeholders which are replaced
33    /// with JSON-encoded argument values before execution.
34    pub script: String,
35    /// Execution timeout in seconds (default: 30).
36    #[serde(default)]
37    pub timeout_secs: Option<u64>,
38    /// Optional memory limit in megabytes (Python only; default: 256).
39    #[serde(default)]
40    pub memory_limit_mb: Option<u64>,
41}
42
43impl Default for ScriptToolConfig {
44    fn default() -> Self {
45        Self {
46            language: "javascript".to_string(),
47            script: String::new(),
48            timeout_secs: Some(30),
49            memory_limit_mb: Some(256),
50        }
51    }
52}
53
54// =============================================================================
55// Tool implementation
56// =============================================================================
57
58/// Runtime script tool that executes JavaScript or Python code.
59pub struct ScriptTool {
60    name: String,
61    description: String,
62    parameters_schema: Value,
63    config: ScriptToolConfig,
64}
65
66impl ScriptTool {
67    /// Create a script tool from its runtime configuration.
68    pub fn new(
69        name: impl Into<String>,
70        description: impl Into<String>,
71        parameters_schema: Value,
72        config: ScriptToolConfig,
73    ) -> Self {
74        Self {
75            name: name.into(),
76            description: description.into(),
77            parameters_schema,
78            config,
79        }
80    }
81
82    /// Parse `execution_config` JSONB into [`ScriptToolConfig`].
83    pub fn parse_config(execution_config: &Value) -> Result<ScriptToolConfig> {
84        serde_json::from_value(execution_config.clone()).map_err(|e| {
85            ares_types::AppError::Configuration(format!("Invalid script tool config: {e}"))
86        })
87    }
88
89    /// Resolve the effective timeout (default 30 s).
90    fn timeout(&self) -> Duration {
91        Duration::from_secs(self.config.timeout_secs.unwrap_or(30))
92    }
93
94    /// Resolve the effective memory limit in MB (default 256).
95    fn memory_limit_mb(&self) -> u64 {
96        self.config.memory_limit_mb.unwrap_or(256)
97    }
98}
99
100#[async_trait]
101impl Tool for ScriptTool {
102    fn name(&self) -> &str {
103        &self.name
104    }
105
106    fn description(&self) -> &str {
107        &self.description
108    }
109
110    fn parameters_schema(&self) -> Value {
111        self.parameters_schema.clone()
112    }
113
114    async fn execute(&self, args: Value) -> Result<Value> {
115        let args_map = args.as_object().ok_or_else(|| {
116            ares_types::AppError::InvalidInput("args must be a JSON object".to_string())
117        })?;
118
119        // Replace {{param}} placeholders with JSON-encoded values.
120        let script = substitute_script_params(&self.config.script, args_map);
121
122        match self.config.language.to_ascii_lowercase().as_str() {
123            "javascript" | "js" => execute_javascript(&script, args, self.timeout()).await,
124            "python" | "py" => {
125                execute_python(&script, args, self.timeout(), self.memory_limit_mb()).await
126            }
127            other => Err(ares_types::AppError::InvalidInput(format!(
128                "Unsupported script language: {other}"
129            ))),
130        }
131    }
132}
133
134// =============================================================================
135// JavaScript execution (boa_engine)
136// =============================================================================
137
138/// Execute JavaScript in a sandboxed `boa_engine` context.
139async fn execute_javascript(script: &str, args: Value, ttl: Duration) -> Result<Value> {
140    let script = script.to_string();
141
142    let handle = tokio::task::spawn_blocking(move || {
143        let mut context = boa_engine::Context::default();
144
145        // Defensive resource limit: cap loop iterations so infinite loops
146        // terminate quickly instead of consuming a blocking thread forever.
147        context
148            .runtime_limits_mut()
149            .set_loop_iteration_limit(10_000_000);
150
151        // Inject `args` as a global JavaScript variable.
152        let args_json = serde_json::to_string(&args).map_err(|e| {
153            ares_types::AppError::Internal(format!("JSON serialization error: {e}"))
154        })?;
155        let init = format!("let args = {args_json};");
156        context
157            .eval(boa_engine::Source::from_bytes(init.as_bytes()))
158            .map_err(|e| ares_types::AppError::External(format!("JS init error: {e}")))?;
159
160        // Run the user script.
161        let result = context
162            .eval(boa_engine::Source::from_bytes(script.as_bytes()))
163            .map_err(|e| ares_types::AppError::External(format!("JS execution error: {e}")))?;
164
165        // boa_engine 0.20 panics on undefined → JSON; handle it explicitly.
166        if result.is_undefined() {
167            return Ok(Value::Null);
168        }
169
170        // Convert back to serde_json::Value.
171        let json_value = result.to_json(&mut context).map_err(|e| {
172            ares_types::AppError::External(format!("JS result conversion error: {e}"))
173        })?;
174
175        Ok(json_value)
176    });
177
178    match timeout(ttl, handle).await {
179        Ok(Ok(result)) => result,
180        Ok(Err(join_err)) => Err(ares_types::AppError::External(format!(
181            "JS task panicked: {join_err}"
182        ))),
183        Err(_) => Err(ares_types::AppError::External(
184            "JS execution timed out".to_string(),
185        )),
186    }
187}
188
189// =============================================================================
190// Python execution (subprocess)
191// =============================================================================
192
193/// Execute Python in a restricted subprocess.
194async fn execute_python(
195    script: &str,
196    args: Value,
197    ttl: Duration,
198    _memory_limit_mb: u64,
199) -> Result<Value> {
200    let args_json = serde_json::to_string(&args)
201        .map_err(|e| ares_types::AppError::Internal(format!("JSON serialization error: {e}")))?;
202
203    // Spawn Python with the wrapper script.
204    let mut cmd = Command::new("python3");
205    cmd.arg("-c")
206        .arg(PYTHON_WRAPPER)
207        .stdin(std::process::Stdio::piped())
208        .stdout(std::process::Stdio::piped())
209        .stderr(std::process::Stdio::piped())
210        .env("SCRIPT_TOOL_ARGS_JSON", args_json)
211        .env("SCRIPT_TOOL_USER_CODE", script);
212
213    let mut child = cmd
214        .spawn()
215        .map_err(|e| ares_types::AppError::External(format!("Failed to spawn python3: {e}")))?;
216
217    // Close stdin immediately — all data is passed via env vars.
218    drop(child.stdin.take());
219
220    let mut stdout_pipe = child.stdout.take().unwrap();
221    let mut stderr_pipe = child.stderr.take().unwrap();
222
223    // Read stdout/stderr concurrently so the pipe buffer never blocks the child.
224    let stdout_task = tokio::spawn(async move {
225        let mut buf = Vec::new();
226        tokio::io::AsyncReadExt::read_to_end(&mut stdout_pipe, &mut buf)
227            .await
228            .ok();
229        buf
230    });
231    let stderr_task = tokio::spawn(async move {
232        let mut buf = Vec::new();
233        tokio::io::AsyncReadExt::read_to_end(&mut stderr_pipe, &mut buf)
234            .await
235            .ok();
236        buf
237    });
238
239    let status = match timeout(ttl, child.wait()).await {
240        Ok(Ok(s)) => s,
241        Ok(Err(e)) => {
242            return Err(ares_types::AppError::External(format!(
243                "Python subprocess error: {e}"
244            )))
245        }
246        Err(_) => {
247            let _ = child.start_kill();
248            return Err(ares_types::AppError::External(
249                "Python execution timed out".to_string(),
250            ));
251        }
252    };
253
254    let stdout = stdout_task.await.unwrap_or_default();
255    let stderr = stderr_task.await.unwrap_or_default();
256
257    if !status.success() {
258        let stderr_str = String::from_utf8_lossy(&stderr);
259        return Err(ares_types::AppError::External(format!(
260            "Python script exited with code {:?}: {stderr_str}",
261            status.code()
262        )));
263    }
264
265    let stdout_str = String::from_utf8_lossy(&stdout);
266    let trimmed = stdout_str.trim();
267
268    if trimmed.is_empty() {
269        return Ok(Value::Null);
270    }
271
272    // Try to parse stdout as JSON.
273    match serde_json::from_str::<Value>(trimmed) {
274        Ok(v) => Ok(v),
275        Err(_) => Ok(Value::String(trimmed.to_string())),
276    }
277}
278
279/// Python wrapper script executed via `python3 -c`.
280///
281/// Reads the user script from `SCRIPT_TOOL_USER_CODE` and arguments from
282/// `SCRIPT_TOOL_ARGS_JSON`.  Executes the user code with a restricted builtin
283/// set and emits `{"result": ..., "stdout": ...}` on stdout.
284const PYTHON_WRAPPER: &str = r#"
285import json, sys, io, os
286
287args = json.loads(os.environ.get("SCRIPT_TOOL_ARGS_JSON", "{}"))
288user_script = os.environ.get("SCRIPT_TOOL_USER_CODE", "")
289
290_SAFE_BUILTINS = {
291    'abs': abs, 'all': all, 'any': any, 'bool': bool,
292    'dict': dict, 'enumerate': enumerate, 'filter': filter,
293    'float': float, 'frozenset': frozenset, 'int': int,
294    'isinstance': isinstance, 'issubclass': issubclass,
295    'len': len, 'list': list, 'map': map, 'max': max,
296    'min': min, 'next': next, 'pow': pow, 'print': print,
297    'range': range, 'reversed': reversed, 'round': round,
298    'set': set, 'slice': slice, 'sorted': sorted,
299    'str': str, 'sum': sum, 'tuple': tuple, 'type': type,
300    'zip': zip, 'json': json, 'io': io,
301    'Exception': Exception, 'TypeError': TypeError,
302    'ValueError': ValueError, 'KeyError': KeyError,
303    'IndexError': IndexError, 'AttributeError': AttributeError,
304    'ArithmeticError': ArithmeticError, 'RuntimeError': RuntimeError,
305}
306
307_old_stdout = sys.stdout
308_sys_stdout = io.StringIO()
309sys.stdout = _sys_stdout
310
311_locals = {'args': args}
312exec(user_script, {"__builtins__": _SAFE_BUILTINS}, _locals)
313
314sys.stdout = _old_stdout
315_stdout_text = _sys_stdout.getvalue()
316
317_result = _locals.get('result')
318print(json.dumps({"result": _result, "stdout": _stdout_text}))
319"#;
320
321// =============================================================================
322// Helpers
323// =============================================================================
324
325/// Replace `{{key}}` placeholders in `script` with JSON-encoded values from `args`.
326fn substitute_script_params(script: &str, args: &serde_json::Map<String, Value>) -> String {
327    let mut result = script.to_string();
328    for (key, value) in args {
329        let placeholder = format!("{{{{{}}}}}", key);
330        let replacement = serde_json::to_string(value).unwrap_or_default();
331        result = result.replace(&placeholder, &replacement);
332    }
333    result
334}
335
336// =============================================================================
337// Tests
338// =============================================================================
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use serde_json::json;
344
345    fn make_tool(config: ScriptToolConfig) -> ScriptTool {
346        ScriptTool::new(
347            "test_script",
348            "A test script tool",
349            json!({
350                "type": "object",
351                "properties": {
352                    "name": { "type": "string" },
353                    "count": { "type": "number" }
354                }
355            }),
356            config,
357        )
358    }
359
360    // -------------------------------------------------------------------------
361    // Config parsing
362    // -------------------------------------------------------------------------
363
364    #[test]
365    fn test_parse_config_javascript() {
366        let raw = json!({
367            "language": "javascript",
368            "script": "args.name",
369            "timeout_secs": 10,
370            "memory_limit_mb": 128
371        });
372        let cfg = ScriptTool::parse_config(&raw).unwrap();
373        assert_eq!(cfg.language, "javascript");
374        assert_eq!(cfg.script, "args.name");
375        assert_eq!(cfg.timeout_secs, Some(10));
376        assert_eq!(cfg.memory_limit_mb, Some(128));
377    }
378
379    #[test]
380    fn test_parse_config_python() {
381        let raw = json!({
382            "language": "python",
383            "script": "result = args['name']"
384        });
385        let cfg = ScriptTool::parse_config(&raw).unwrap();
386        assert_eq!(cfg.language, "python");
387        assert_eq!(cfg.script, "result = args['name']");
388        assert_eq!(cfg.timeout_secs, None);
389        assert_eq!(cfg.memory_limit_mb, None);
390    }
391
392    #[test]
393    fn test_parse_config_invalid() {
394        let raw = json!({ "language": 42 });
395        assert!(ScriptTool::parse_config(&raw).is_err());
396    }
397
398    // -------------------------------------------------------------------------
399    // Parameter substitution
400    // -------------------------------------------------------------------------
401
402    #[test]
403    fn test_substitute_params_strings_and_numbers() {
404        let mut map = serde_json::Map::new();
405        map.insert("name".into(), json!("Alice"));
406        map.insert("count".into(), json!(42));
407
408        let script = r#"let greeting = "Hello, {{name}}"; let n = {{count}};"#;
409        let out = substitute_script_params(script, &map);
410        assert_eq!(out, r#"let greeting = "Hello, "Alice""; let n = 42;"#);
411    }
412
413    #[test]
414    fn test_substitute_params_escapes_quotes() {
415        let mut map = serde_json::Map::new();
416        map.insert("unsafe".into(), json!(r#""; DROP TABLE users; --"#));
417
418        let script = r#"let x = {{unsafe}};"#;
419        let out = substitute_script_params(script, &map);
420        assert_eq!(out, "let x = \"\\\"; DROP TABLE users; --\";");
421    }
422
423    #[test]
424    fn test_substitute_params_missing_placeholder_unchanged() {
425        let mut map = serde_json::Map::new();
426        map.insert("a".into(), json!(1));
427
428        let script = "{{a}} {{b}}";
429        let out = substitute_script_params(script, &map);
430        assert_eq!(out, "1 {{b}}");
431    }
432
433    // -------------------------------------------------------------------------
434    // JavaScript execution
435    // -------------------------------------------------------------------------
436
437    #[tokio::test]
438    async fn test_js_returns_last_expression() {
439        let tool = make_tool(ScriptToolConfig {
440            language: "javascript".into(),
441            script: "args.count * 2".into(),
442            ..Default::default()
443        });
444        let out = tool.execute(json!({"count": 21})).await.unwrap();
445        assert_eq!(out, json!(42));
446    }
447
448    #[tokio::test]
449    async fn test_js_string_concatenation() {
450        let tool = make_tool(ScriptToolConfig {
451            language: "javascript".into(),
452            script: r#""Hello, " + args.name"#.into(),
453            ..Default::default()
454        });
455        let out = tool.execute(json!({"name": "World"})).await.unwrap();
456        assert_eq!(out, json!("Hello, World"));
457    }
458
459    #[tokio::test]
460    async fn test_js_object_manipulation() {
461        let tool = make_tool(ScriptToolConfig {
462            language: "javascript".into(),
463            script: "args.count + 1".into(),
464            ..Default::default()
465        });
466        let out = tool.execute(json!({"count": 5})).await.unwrap();
467        assert_eq!(out, json!(6));
468    }
469
470    #[tokio::test]
471    async fn test_js_template_substitution() {
472        let tool = make_tool(ScriptToolConfig {
473            language: "javascript".into(),
474            script: r#""value is " + {{count}}"#.into(),
475            ..Default::default()
476        });
477        let out = tool.execute(json!({"count": 99})).await.unwrap();
478        assert_eq!(out, json!("value is 99"));
479    }
480
481    #[tokio::test]
482    async fn test_js_infinite_loop_hits_limit() {
483        let tool = make_tool(ScriptToolConfig {
484            language: "javascript".into(),
485            script: "while(true) {{ }}".into(),
486            ..Default::default()
487        });
488        let err = tool.execute(json!({})).await.unwrap_err();
489        let msg = format!("{err}");
490        assert!(
491            msg.contains("execution error"),
492            "Expected execution error, got: {msg}"
493        );
494    }
495
496    #[tokio::test]
497    async fn test_js_syntax_error() {
498        let tool = make_tool(ScriptToolConfig {
499            language: "javascript".into(),
500            script: "this is not valid js !!!".into(),
501            ..Default::default()
502        });
503        let err = tool.execute(json!({})).await.unwrap_err();
504        let msg = format!("{err}");
505        assert!(
506            msg.contains("execution error") || msg.contains("SyntaxError"),
507            "Expected JS error, got: {msg}"
508        );
509    }
510
511    #[tokio::test]
512    async fn test_js_undefined_returns_null() {
513        let tool = make_tool(ScriptToolConfig {
514            language: "javascript".into(),
515            script: "let x = undefined; x".into(),
516            ..Default::default()
517        });
518        let out = tool.execute(json!({})).await.unwrap();
519        assert!(out.is_null());
520    }
521
522    // -------------------------------------------------------------------------
523    // Python execution
524    // -------------------------------------------------------------------------
525
526    fn is_python_available() -> bool {
527        std::process::Command::new("python3")
528            .arg("--version")
529            .output()
530            .is_ok()
531    }
532
533    #[tokio::test]
534    async fn test_python_math() {
535        if !is_python_available() {
536            eprintln!("Skipping Python test: python3 not found");
537            return;
538        }
539        let tool = make_tool(ScriptToolConfig {
540            language: "python".into(),
541            script: "result = args['a'] + args['b']".into(),
542            ..Default::default()
543        });
544        let out = tool.execute(json!({"a": 10, "b": 32})).await.unwrap();
545        // stdout wrapper emits {"result": 42, "stdout": ""}
546        assert_eq!(out["result"], json!(42));
547    }
548
549    #[tokio::test]
550    async fn test_python_string_result() {
551        if !is_python_available() {
552            eprintln!("Skipping Python test: python3 not found");
553            return;
554        }
555        let tool = make_tool(ScriptToolConfig {
556            language: "python".into(),
557            script: "result = 'Hello, ' + args['name']".into(),
558            ..Default::default()
559        });
560        let out = tool.execute(json!({"name": "World"})).await.unwrap();
561        assert_eq!(out["result"], json!("Hello, World"));
562    }
563
564    #[tokio::test]
565    async fn test_python_template_substitution() {
566        if !is_python_available() {
567            eprintln!("Skipping Python test: python3 not found");
568            return;
569        }
570        let tool = make_tool(ScriptToolConfig {
571            language: "python".into(),
572            script: "result = 'count = ' + str({{count}})".into(),
573            ..Default::default()
574        });
575        let out = tool.execute(json!({"count": 7})).await.unwrap();
576        assert_eq!(out["result"], json!("count = 7"));
577    }
578
579    #[tokio::test]
580    async fn test_python_captured_stdout() {
581        if !is_python_available() {
582            eprintln!("Skipping Python test: python3 not found");
583            return;
584        }
585        let tool = make_tool(ScriptToolConfig {
586            language: "python".into(),
587            script: "print('debug line')".into(),
588            ..Default::default()
589        });
590        let out = tool.execute(json!({})).await.unwrap();
591        assert_eq!(out["stdout"], json!("debug line\n"));
592    }
593
594    #[tokio::test]
595    async fn test_python_timeout() {
596        if !is_python_available() {
597            eprintln!("Skipping Python test: python3 not found");
598            return;
599        }
600        let tool = make_tool(ScriptToolConfig {
601            language: "python".into(),
602            script: "while True:\n    pass".into(),
603            timeout_secs: Some(1),
604            ..Default::default()
605        });
606        let err = tool.execute(json!({})).await.unwrap_err();
607        let msg = format!("{err}");
608        assert!(
609            msg.contains("timed out"),
610            "Expected timeout error, got: {msg}"
611        );
612    }
613
614    #[tokio::test]
615    async fn test_python_syntax_error() {
616        if !is_python_available() {
617            eprintln!("Skipping Python test: python3 not found");
618            return;
619        }
620        let tool = make_tool(ScriptToolConfig {
621            language: "python".into(),
622            script: "this is not valid python !!!".into(),
623            ..Default::default()
624        });
625        let err = tool.execute(json!({})).await.unwrap_err();
626        let msg = format!("{err}");
627        assert!(
628            msg.contains("exited with code") || msg.contains("SyntaxError"),
629            "Expected Python error, got: {msg}"
630        );
631    }
632
633    #[tokio::test]
634    async fn test_python_restricted_builtin_blocks_file_open() {
635        if !is_python_available() {
636            eprintln!("Skipping Python test: python3 not found");
637            return;
638        }
639        let tool = make_tool(ScriptToolConfig {
640            language: "python".into(),
641            script: "open('/etc/passwd')".into(),
642            ..Default::default()
643        });
644        let err = tool.execute(json!({})).await.unwrap_err();
645        let msg = format!("{err}");
646        assert!(
647            msg.contains("NameError") || msg.contains("exited with code"),
648            "Expected blocked open(), got: {msg}"
649        );
650    }
651
652    // -------------------------------------------------------------------------
653    // Tool trait plumbing
654    // -------------------------------------------------------------------------
655
656    #[test]
657    fn test_name_and_description() {
658        let tool = make_tool(ScriptToolConfig::default());
659        assert_eq!(tool.name(), "test_script");
660        assert_eq!(tool.description(), "A test script tool");
661    }
662
663    #[test]
664    fn test_parameters_schema() {
665        let tool = make_tool(ScriptToolConfig::default());
666        let schema = tool.parameters_schema();
667        assert_eq!(schema["type"], "object");
668        assert!(schema["properties"].get("name").is_some());
669    }
670
671    #[tokio::test]
672    async fn test_unsupported_language() {
673        let tool = make_tool(ScriptToolConfig {
674            language: "ruby".into(),
675            script: "1+1".into(),
676            ..Default::default()
677        });
678        let err = tool.execute(json!({})).await.unwrap_err();
679        let msg = format!("{err}");
680        assert!(msg.contains("Unsupported"));
681    }
682
683    #[tokio::test]
684    async fn test_invalid_args_type() {
685        let tool = make_tool(ScriptToolConfig::default());
686        let err = tool.execute(json!("not an object")).await.unwrap_err();
687        let msg = format!("{err}");
688        assert!(msg.contains("must be a JSON object"));
689    }
690}
ares_tools/script_tool.rs

ares_tools/
script_tool.rs
