a3s_flow/nodes/
code.rs

1//! Built-in `"code"` node — executes an inline script in a sandboxed Rhai
2//! engine and returns the result as JSON.
3//!
4//! Mirrors Dify's Code node. Rhai is a safe, embedded scripting language with
5//! Rust-like syntax. It has no file system, network, or OS access by default.
6//!
7//! # Config schema
8//!
9//! ```json
10//! {
11//!   "language": "rhai",
12//!   "code": "let total = inputs.items.len(); #{ total: total }"
13//! }
14//! ```
15//!
16//! | Field | Type | Description |
17//! |-------|------|-------------|
18//! | `language` | string | Must be `"rhai"` (the only supported language) |
19//! | `code` | string | Rhai script body |
20//!
21//! # Script context
22//!
23//! Two variables are injected into the script scope:
24//! - `inputs` — map keyed by upstream node ID (equivalent to `ctx.inputs`)
25//! - `variables` — global flow variables (equivalent to `ctx.variables`)
26//!
27//! # Output schema
28//!
29//! If the script returns a Rhai object map, it becomes the node output directly:
30//!
31//! ```rhai
32//! #{ status: inputs.fetch.status, ok: inputs.fetch.ok }
33//! // → { "status": 200, "ok": true }
34//! ```
35//!
36//! Any other return type is wrapped under `"output"`:
37//!
38//! ```rhai
39//! inputs.fetch.status == 200
40//! // → { "output": true }
41//! ```
42//!
43//! # Safety limits
44//!
45//! The engine enforces:
46//! - Max 100,000 operations (prevents infinite loops)
47//! - Max string size: 1 MB
48//! - Max array size: 10,000 elements
49//!
50//! # Rhai syntax reference
51//!
52//! - Variables: `let x = 42;`
53//! - Maps: `#{ key: value }`
54//! - Arrays: `[1, 2, 3]`
55//! - String ops: `s.len()`, `s.contains("x")`, `s.to_upper()`
56//! - Math: `+`, `-`, `*`, `/`, `%`
57//! - Conditionals: `if x > 0 { "pos" } else { "neg" }`
58//! - Loops: `for item in arr { ... }`
59//!
60//! Full docs: <https://rhai.rs/book/>
61
62use async_trait::async_trait;
63use serde_json::{json, Value};
64
65use crate::error::{FlowError, Result};
66use crate::node::{ExecContext, Node};
67
68/// Sandboxed script execution node (Dify-compatible, Rhai engine).
69pub struct CodeNode;
70
71#[async_trait]
72impl Node for CodeNode {
73    fn node_type(&self) -> &str {
74        "code"
75    }
76
77    async fn execute(&self, ctx: ExecContext) -> Result<Value> {
78        let language = ctx.data["language"].as_str().unwrap_or("rhai");
79        if language != "rhai" {
80            return Err(FlowError::InvalidDefinition(format!(
81                "code: unsupported language '{language}' — only 'rhai' is supported"
82            )));
83        }
84
85        let code = ctx.data["code"]
86            .as_str()
87            .ok_or_else(|| FlowError::InvalidDefinition("code: missing data.code".into()))?;
88
89        let mut engine = rhai::Engine::new();
90        engine.set_max_operations(100_000);
91        engine.set_max_string_size(1_000_000);
92        engine.set_max_array_size(10_000);
93
94        let mut scope = rhai::Scope::new();
95
96        let inputs_dyn = rhai::serde::to_dynamic(&ctx.inputs)
97            .map_err(|e| FlowError::Internal(format!("code: failed to serialize inputs: {e}")))?;
98        scope.push_dynamic("inputs", inputs_dyn);
99
100        let vars_dyn = rhai::serde::to_dynamic(&ctx.variables).map_err(|e| {
101            FlowError::Internal(format!("code: failed to serialize variables: {e}"))
102        })?;
103        scope.push_dynamic("variables", vars_dyn);
104
105        let result: rhai::Dynamic = engine
106            .eval_with_scope(&mut scope, code)
107            .map_err(|e| FlowError::Internal(format!("code: script error: {e}")))?;
108
109        // Object maps become the output directly; anything else wraps in {"output": ...}.
110        let output: Value = rhai::serde::from_dynamic(&result)
111            .map_err(|e| FlowError::Internal(format!("code: result serialization failed: {e}")))?;
112
113        if output.is_object() {
114            Ok(output)
115        } else {
116            Ok(json!({ "output": output }))
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use std::collections::HashMap;
125
126    fn ctx(inputs: HashMap<String, Value>, code: &str) -> ExecContext {
127        ExecContext {
128            data: json!({ "language": "rhai", "code": code }),
129            inputs,
130            variables: HashMap::new(),
131            ..Default::default()
132        }
133    }
134
135    #[tokio::test]
136    async fn scalar_wrapped_in_output() {
137        let out = CodeNode.execute(ctx(HashMap::new(), "42")).await.unwrap();
138        assert_eq!(out["output"], 42);
139    }
140
141    #[tokio::test]
142    async fn bool_result() {
143        let out = CodeNode
144            .execute(ctx(HashMap::new(), "1 + 1 == 2"))
145            .await
146            .unwrap();
147        assert_eq!(out["output"], true);
148    }
149
150    #[tokio::test]
151    async fn object_map_returned_directly() {
152        let out = CodeNode
153            .execute(ctx(HashMap::new(), "#{ x: 1, y: 2 }"))
154            .await
155            .unwrap();
156        assert_eq!(out["x"], 1);
157        assert_eq!(out["y"], 2);
158        assert!(out.get("output").is_none());
159    }
160
161    #[tokio::test]
162    async fn inputs_accessible_in_script() {
163        let out = CodeNode
164            .execute(ctx(
165                HashMap::from([("fetch".into(), json!({ "status": 200 }))]),
166                "inputs.fetch.status == 200",
167            ))
168            .await
169            .unwrap();
170        assert_eq!(out["output"], true);
171    }
172
173    #[tokio::test]
174    async fn rejects_unsupported_language() {
175        let err = CodeNode
176            .execute(ExecContext {
177                data: json!({ "language": "python", "code": "pass" }),
178                inputs: HashMap::new(),
179                variables: HashMap::new(),
180                ..Default::default()
181            })
182            .await
183            .unwrap_err();
184        assert!(matches!(err, FlowError::InvalidDefinition(_)));
185    }
186
187    #[tokio::test]
188    async fn rejects_missing_code() {
189        let err = CodeNode
190            .execute(ExecContext {
191                data: json!({ "language": "rhai" }),
192                inputs: HashMap::new(),
193                variables: HashMap::new(),
194                ..Default::default()
195            })
196            .await
197            .unwrap_err();
198        assert!(matches!(err, FlowError::InvalidDefinition(_)));
199    }
200
201    #[tokio::test]
202    async fn script_error_returns_internal() {
203        let err = CodeNode
204            .execute(ctx(HashMap::new(), "undefined_fn()"))
205            .await
206            .unwrap_err();
207        assert!(matches!(err, FlowError::Internal(_)));
208    }
209
210    #[tokio::test]
211    async fn string_manipulation() {
212        let out = CodeNode
213            .execute(ctx(
214                HashMap::from([("msg".into(), json!({ "text": "hello" }))]),
215                r#"let t = inputs.msg.text; t.to_upper()"#,
216            ))
217            .await
218            .unwrap();
219        assert_eq!(out["output"], "HELLO");
220    }
221}
a3s_flow/nodes/code.rs

a3s_flow/nodes/
code.rs
