a3s_code_core/tools/
program_tool.rs

1//! Tool wrapper for programmatic tool calling.
2
3use crate::program::ProgramCatalog;
4use crate::text::truncate_utf8;
5use crate::tools::types::{Tool, ToolContext, ToolOutput};
6use crate::tools::ToolRegistry;
7use anyhow::{anyhow, Result};
8use async_trait::async_trait;
9use rquickjs::function::{Async, Func};
10use rquickjs::{async_with, AsyncContext, AsyncRuntime, CatchResultExt, Error as JsError, Promise};
11use serde::Deserialize;
12use std::collections::HashSet;
13use std::sync::Arc;
14use std::time::Instant;
15use tokio::sync::Mutex;
16use tokio::time::{timeout, Duration};
17
18const DEFAULT_SCRIPT_TIMEOUT_MS: u64 = 30_000;
19const DEFAULT_SCRIPT_MAX_TOOL_CALLS: usize = 20;
20const DEFAULT_SCRIPT_MAX_OUTPUT_BYTES: usize = 64 * 1024;
21const MAX_SCRIPT_SOURCE_BYTES: usize = 64 * 1024;
22
23pub struct ProgramTool {
24    registry: Arc<ToolRegistry>,
25}
26
27impl ProgramTool {
28    pub fn new(registry: Arc<ToolRegistry>) -> Self {
29        Self { registry }
30    }
31
32    pub fn with_catalog(registry: Arc<ToolRegistry>, _catalog: ProgramCatalog) -> Self {
33        Self { registry }
34    }
35}
36
37#[async_trait]
38impl Tool for ProgramTool {
39    fn name(&self) -> &str {
40        "program"
41    }
42
43    fn description(&self) -> &str {
44        "Run a sandboxed JavaScript PTC script. The script defines async function run(ctx, inputs) and may call only allowed ctx tools."
45    }
46
47    fn parameters(&self) -> serde_json::Value {
48        serde_json::json!({
49            "type": "object",
50            "additionalProperties": false,
51            "properties": {
52                "type": {
53                    "type": "string",
54                    "description": "Required. Program kind. Only \"script\" is supported.",
55                    "enum": ["script"]
56                },
57                "inputs": {
58                    "type": "object",
59                    "description": "Optional. JSON inputs passed to the script as the second argument."
60                },
61                "language": {
62                    "type": "string",
63                    "description": "Script language. Only JavaScript is supported.",
64                    "enum": ["javascript"]
65                },
66                "source": {
67                    "type": "string",
68                    "description": "Inline JavaScript source defining async function run(ctx, inputs)."
69                },
70                "path": {
71                    "type": "string",
72                    "description": "Workspace-relative path to a .js or .mjs script defining async function run(ctx, inputs). Used when source is omitted."
73                },
74                "allowed_tools": {
75                    "type": "array",
76                    "description": "Tool names the script may call through ctx. Defaults to all registered tools except program.",
77                    "items": { "type": "string" }
78                },
79                "limits": {
80                    "type": "object",
81                    "description": "Optional timeoutMs, maxToolCalls, and maxOutputBytes.",
82                    "additionalProperties": false,
83                    "properties": {
84                        "timeoutMs": { "type": "integer", "minimum": 1 },
85                        "maxToolCalls": { "type": "integer", "minimum": 1 },
86                        "maxOutputBytes": { "type": "integer", "minimum": 1 }
87                    }
88                }
89            },
90            "required": ["type"]
91        })
92    }
93
94    async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
95        let Some(kind) = args.get("type").and_then(|value| value.as_str()) else {
96            return Ok(ToolOutput::error("type parameter is required"));
97        };
98        if kind != "script" {
99            return Ok(ToolOutput::error(format!(
100                "Unsupported program type: {kind}. Only \"script\" is supported."
101            )));
102        }
103        let inputs = args
104            .get("inputs")
105            .cloned()
106            .unwrap_or_else(|| serde_json::json!({}));
107
108        execute_script_program(args, inputs, Arc::clone(&self.registry), ctx).await
109    }
110}
111
112#[derive(Debug, Deserialize)]
113#[serde(rename_all = "camelCase")]
114struct ScriptLimits {
115    timeout_ms: Option<u64>,
116    max_tool_calls: Option<usize>,
117    max_output_bytes: Option<usize>,
118}
119
120#[derive(Debug, Clone)]
121struct ScriptCallRecord {
122    tool_name: String,
123    success: bool,
124    exit_code: i32,
125    output_bytes: usize,
126    metadata: Option<serde_json::Value>,
127}
128
129async fn execute_script_program(
130    args: &serde_json::Value,
131    inputs: serde_json::Value,
132    registry: Arc<ToolRegistry>,
133    ctx: &ToolContext,
134) -> Result<ToolOutput> {
135    let language = args
136        .get("language")
137        .and_then(|value| value.as_str())
138        .unwrap_or("javascript");
139    if language != "javascript" {
140        return Ok(ToolOutput::error(format!(
141            "Unsupported script language: {language}"
142        )));
143    }
144
145    let source = match load_script_source(args, ctx).await {
146        Ok(source) => source,
147        Err(message) => return Ok(ToolOutput::error(message)),
148    };
149    if source.len() > MAX_SCRIPT_SOURCE_BYTES {
150        return Ok(ToolOutput::error(format!(
151            "script source is too large: {} bytes exceeds {} bytes",
152            source.len(),
153            MAX_SCRIPT_SOURCE_BYTES
154        )));
155    }
156    if let Err(message) = validate_script_source(&source) {
157        return Ok(ToolOutput::error(message));
158    }
159
160    let allowed_tools = script_allowed_tools(args, &registry);
161    let limits = script_limits(args);
162    match run_quickjs_script(
163        &source,
164        inputs,
165        registry,
166        ctx.clone(),
167        allowed_tools,
168        limits,
169    )
170    .await
171    {
172        Ok(output) => Ok(output),
173        Err(err) => Ok(ToolOutput::error(format!("program script failed: {err}"))),
174    }
175}
176
177async fn load_script_source(
178    args: &serde_json::Value,
179    ctx: &ToolContext,
180) -> std::result::Result<String, String> {
181    if let Some(source) = args.get("source").and_then(|value| value.as_str()) {
182        return Ok(source.to_string());
183    }
184
185    let Some(path) = args.get("path").and_then(|value| value.as_str()) else {
186        return Err("program script requires either source or path".to_string());
187    };
188    if !(path.ends_with(".js") || path.ends_with(".mjs")) {
189        return Err("program script path must point to a .js or .mjs file".to_string());
190    }
191
192    let workspace_path = ctx
193        .resolve_workspace_path(path)
194        .map_err(|err| format!("failed to resolve script path: {err}"))?;
195    ctx.workspace_services
196        .fs()
197        .read_text(&workspace_path)
198        .await
199        .map_err(|err| format!("failed to read script path '{}': {err}", path))
200}
201
202fn script_allowed_tools(args: &serde_json::Value, registry: &ToolRegistry) -> HashSet<String> {
203    let mut allowed = args
204        .get("allowed_tools")
205        .and_then(|value| value.as_array())
206        .map(|items| {
207            items
208                .iter()
209                .filter_map(|item| item.as_str())
210                .map(ToString::to_string)
211                .collect::<HashSet<_>>()
212        })
213        .unwrap_or_else(|| registry.list().into_iter().collect());
214
215    allowed.remove("program");
216    allowed
217}
218
219fn script_limits(args: &serde_json::Value) -> ScriptLimits {
220    args.get("limits")
221        .cloned()
222        .and_then(|value| serde_json::from_value(value).ok())
223        .unwrap_or(ScriptLimits {
224            timeout_ms: None,
225            max_tool_calls: None,
226            max_output_bytes: None,
227        })
228}
229
230fn validate_script_source(source: &str) -> std::result::Result<(), String> {
231    let forbidden = [
232        ("import ", "imports are not allowed inside PTC scripts"),
233        (
234            "import(",
235            "dynamic imports are not allowed inside PTC scripts",
236        ),
237        ("eval(", "eval is not allowed inside PTC scripts"),
238        (
239            "Function(",
240            "Function constructor is not allowed inside PTC scripts",
241        ),
242        ("Worker(", "Worker is not allowed inside PTC scripts"),
243        ("WebSocket", "WebSocket is not allowed inside PTC scripts"),
244        (
245            "fetch(",
246            "fetch is not allowed inside PTC scripts; use ctx tools instead",
247        ),
248    ];
249
250    for (needle, message) in forbidden {
251        if source.contains(needle) {
252            return Err(message.to_string());
253        }
254    }
255    Ok(())
256}
257
258async fn run_quickjs_script(
259    source: &str,
260    inputs: serde_json::Value,
261    registry: Arc<ToolRegistry>,
262    ctx: ToolContext,
263    allowed_tools: HashSet<String>,
264    limits: ScriptLimits,
265) -> Result<ToolOutput> {
266    let timeout_ms = limits.timeout_ms.unwrap_or(DEFAULT_SCRIPT_TIMEOUT_MS);
267    let max_tool_calls = limits
268        .max_tool_calls
269        .unwrap_or(DEFAULT_SCRIPT_MAX_TOOL_CALLS);
270    let max_output_bytes = limits
271        .max_output_bytes
272        .unwrap_or(DEFAULT_SCRIPT_MAX_OUTPUT_BYTES);
273    let executable_source = script_source_with_host_entrypoint(source)?;
274    let state = Arc::new(Mutex::new(ScriptVmState {
275        registry,
276        ctx,
277        allowed_tools,
278        max_tool_calls,
279        max_output_bytes,
280        tool_calls: 0,
281        records: Vec::new(),
282    }));
283
284    let vm_state = Arc::clone(&state);
285    let result = timeout(
286        Duration::from_millis(timeout_ms),
287        tokio::task::spawn_blocking(move || {
288            let runtime = tokio::runtime::Builder::new_current_thread()
289                .enable_all()
290                .build()
291                .map_err(|err| anyhow!("failed to create program VM runtime: {err}"))?;
292            runtime.block_on(run_embedded_script(
293                executable_source,
294                inputs,
295                vm_state,
296                timeout_ms,
297            ))
298        }),
299    )
300    .await;
301
302    match result {
303        Ok(Ok(Ok(result))) => {
304            let records = state.lock().await.records.clone();
305            let output = render_script_output(&result, &records, "");
306            Ok(ToolOutput::success(output).with_metadata(serde_json::json!({
307                "program": {
308                    "name": "script",
309                    "language": "javascript",
310                    "runtime": "embedded-quickjs",
311                    "success": true,
312                    "tool_calls": records.iter().map(script_record_to_value).collect::<Vec<_>>(),
313                },
314                "script_result": result,
315            })))
316        }
317        Ok(Ok(Err(err))) if is_quickjs_timeout(&err) => Ok(ToolOutput::error(format!(
318            "program script timed out after {timeout_ms} ms"
319        ))),
320        Ok(Ok(Err(err))) => Ok(ToolOutput::error(format!("program script error:\n{err}"))),
321        Ok(Err(err)) => Ok(ToolOutput::error(format!(
322            "program VM thread failed: {err}"
323        ))),
324        Err(_) => Ok(ToolOutput::error(format!(
325            "program script timed out after {timeout_ms} ms"
326        ))),
327    }
328}
329
330fn script_source_with_host_entrypoint(source: &str) -> Result<String> {
331    let rewritten = if source.contains("export default async function run") {
332        source.replacen("export default async function run", "async function run", 1)
333    } else if source.contains("export default function run") {
334        source.replacen("export default function run", "function run", 1)
335    } else if source.contains("async function run") || source.contains("function run") {
336        source.to_string()
337    } else {
338        return Err(anyhow!(
339            "PTC script must define async function run(ctx, inputs)"
340        ));
341    };
342
343    Ok(format!(
344        r#"{rewritten}
345
346globalThis.__a3sResultJson = (async () => JSON.stringify(await run(globalThis.__a3sCtx, globalThis.__a3sInputs)))();
347"#
348    ))
349}
350
351async fn run_embedded_script(
352    source: String,
353    inputs: serde_json::Value,
354    state: Arc<Mutex<ScriptVmState>>,
355    timeout_ms: u64,
356) -> Result<serde_json::Value> {
357    let runtime = AsyncRuntime::new()?;
358    let started = Instant::now();
359    runtime
360        .set_interrupt_handler(Some(Box::new(move || {
361            started.elapsed() >= Duration::from_millis(timeout_ms)
362        })))
363        .await;
364    runtime.set_memory_limit(64 * 1024 * 1024).await;
365    runtime.set_max_stack_size(512 * 1024).await;
366
367    let context = AsyncContext::full(&runtime).await?;
368    let inputs_json = serde_json::to_string(&inputs)?;
369    let script = format!("{}\n{}", embedded_script_bootstrap(&inputs_json), source);
370    let result_json = async_with!(context => |ctx| {
371        let state = Arc::clone(&state);
372        let host_tool = move |tool: String, args_json: String| {
373            let state = Arc::clone(&state);
374            async move { execute_host_tool_json(state, tool, args_json).await }
375        };
376        if let Err(err) = ctx.globals().set("__a3sHostTool", Func::from(Async(host_tool))) {
377            return Err(format!("failed to install program host tool: {err}"));
378        }
379        let promise: Promise = match ctx.eval(script) {
380            Ok(promise) => promise,
381            Err(err) => return Err(format!("failed to evaluate program script: {err}")),
382        };
383        promise
384            .into_future::<String>()
385            .await
386            .catch(&ctx)
387            .map_err(|err| err.to_string())
388    })
389    .await
390    .map_err(anyhow::Error::msg)?;
391
392    serde_json::from_str(&result_json)
393        .map_err(|err| anyhow!("program script returned invalid JSON: {err}"))
394}
395
396struct ScriptVmState {
397    registry: Arc<ToolRegistry>,
398    ctx: ToolContext,
399    allowed_tools: HashSet<String>,
400    max_tool_calls: usize,
401    max_output_bytes: usize,
402    tool_calls: usize,
403    records: Vec<ScriptCallRecord>,
404}
405
406fn embedded_script_bootstrap(inputs_json: &str) -> String {
407    format!(
408        r#"
409const __a3sCallTool = async (tool, args = {{}}) => {{
410  const response = await globalThis.__a3sHostTool(String(tool), JSON.stringify(args ?? {{}}));
411  return JSON.parse(response);
412}};
413
414const __a3sCtx = Object.freeze({{
415  tool: __a3sCallTool,
416  readFile: (path) => __a3sCallTool("read", {{ file_path: path }}).then((r) => r.output),
417  read: (path) => __a3sCallTool("read", {{ file_path: path }}),
418  grep: (pattern, options = {{}}) => __a3sCallTool("grep", {{ pattern, ...options }}).then((r) => r.output),
419  glob: (pattern, options = {{}}) => __a3sCallTool("glob", {{ pattern, ...options }}).then((r) => r.output),
420  ls: (path = ".") => __a3sCallTool("ls", {{ path }}).then((r) => r.output),
421  bash: (command) => __a3sCallTool("bash", {{ command }}).then((r) => r.output),
422  git: (args = {{}}) => __a3sCallTool("git", args),
423  webSearch: (params) => __a3sCallTool("web_search", params),
424  verify: (args) => __a3sCallTool("bash", args),
425}});
426
427Object.defineProperty(globalThis, "__a3sCtx", {{ value: __a3sCtx, configurable: false }});
428Object.defineProperty(globalThis, "__a3sInputs", {{ value: {inputs_json}, configurable: false }});
429Object.defineProperty(globalThis, "fetch", {{ value: undefined, configurable: false, writable: false }});
430Object.defineProperty(globalThis, "WebSocket", {{ value: undefined, configurable: false, writable: false }});
431Object.defineProperty(globalThis, "Worker", {{ value: undefined, configurable: false, writable: false }});
432"#
433    )
434}
435
436async fn execute_host_tool_json(
437    state: Arc<Mutex<ScriptVmState>>,
438    tool: String,
439    args_json: String,
440) -> rquickjs::Result<String> {
441    let args = serde_json::from_str(&args_json).map_err(|err| {
442        JsError::new_from_js_message("string", "object", format!("invalid tool args JSON: {err}"))
443    })?;
444    let (registry, ctx, max_output_bytes) = {
445        let mut script = state.lock().await;
446        if !script.allowed_tools.contains(&tool) {
447            return Err(JsError::new_from_js_message(
448                "tool",
449                "allowed tool",
450                format!("tool '{tool}' is not allowed for this PTC script"),
451            ));
452        }
453        script.tool_calls += 1;
454        if script.tool_calls > script.max_tool_calls {
455            return Err(JsError::new_from_js_message(
456                "tool call",
457                "limited tool call",
458                format!("PTC script exceeded maxToolCalls={}", script.max_tool_calls),
459            ));
460        }
461        (
462            Arc::clone(&script.registry),
463            script.ctx.clone(),
464            script.max_output_bytes,
465        )
466    };
467
468    let result = registry
469        .execute_with_context(&tool, &args, &ctx)
470        .await
471        .map_err(|err| JsError::new_from_js_message("tool", "result", err.to_string()))?;
472    let mut output = result.output;
473    if output.len() > max_output_bytes {
474        output = truncate_utf8(&output, max_output_bytes).to_string();
475    }
476    let success = result.exit_code == 0;
477    let metadata = result.metadata.clone();
478    let exit_code = result.exit_code;
479    let name = result.name;
480
481    {
482        let mut script = state.lock().await;
483        script.records.push(ScriptCallRecord {
484            tool_name: tool,
485            success,
486            exit_code,
487            output_bytes: output.len(),
488            metadata: metadata.clone(),
489        });
490    }
491
492    serde_json::to_string(&serde_json::json!({
493        "name": name,
494        "output": output,
495        "exitCode": exit_code,
496        "metadata": metadata,
497    }))
498    .map_err(|err| JsError::new_from_js_message("tool result", "json", err.to_string()))
499}
500
501fn is_quickjs_timeout(err: &anyhow::Error) -> bool {
502    let text = err.to_string();
503    text.contains("interrupted") || text.contains("InternalError")
504}
505
506fn script_record_to_value(record: &ScriptCallRecord) -> serde_json::Value {
507    serde_json::json!({
508        "tool_name": record.tool_name,
509        "success": record.success,
510        "exit_code": record.exit_code,
511        "output_bytes": record.output_bytes,
512        "metadata": record.metadata,
513    })
514}
515
516fn render_script_output(
517    result: &serde_json::Value,
518    records: &[ScriptCallRecord],
519    stderr: &str,
520) -> String {
521    let mut output = String::from("Program script completed.");
522    if let Some(summary) = result.get("summary").and_then(|value| value.as_str()) {
523        output.push('\n');
524        output.push_str(summary);
525    }
526
527    output.push_str(&format!("\n\nTool calls: {}", records.len()));
528    for (index, record) in records.iter().enumerate() {
529        output.push_str(&format!(
530            "\n{}. {} ({}, exit_code={}, output_bytes={})",
531            index + 1,
532            record.tool_name,
533            if record.success { "ok" } else { "failed" },
534            record.exit_code,
535            record.output_bytes
536        ));
537    }
538
539    output.push_str("\n\nResult:\n");
540    output.push_str(&serde_json::to_string_pretty(result).unwrap_or_else(|_| result.to_string()));
541
542    if !stderr.is_empty() {
543        output.push_str("\n\nstderr:\n");
544        output.push_str(stderr);
545    }
546
547    output
548}
549
550#[cfg(test)]
551mod tests {
552    use super::*;
553    use async_trait::async_trait;
554    use std::path::PathBuf;
555
556    struct EchoTool;
557
558    #[async_trait]
559    impl Tool for EchoTool {
560        fn name(&self) -> &str {
561            "echo"
562        }
563
564        fn description(&self) -> &str {
565            "Echo test tool"
566        }
567
568        fn parameters(&self) -> serde_json::Value {
569            serde_json::json!({
570                "type": "object",
571                "properties": {
572                    "message": { "type": "string" }
573                }
574            })
575        }
576
577        async fn execute(
578            &self,
579            args: &serde_json::Value,
580            _ctx: &ToolContext,
581        ) -> Result<ToolOutput> {
582            let message = args
583                .get("message")
584                .and_then(|value| value.as_str())
585                .unwrap_or("");
586            Ok(ToolOutput::success(format!("echo:{message}")))
587        }
588    }
589
590    #[tokio::test]
591    async fn program_tool_rejects_non_script_type() {
592        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
593        let output = tool
594            .execute(
595                &serde_json::json!({ "type": "program_code_search" }),
596                &ToolContext::new(PathBuf::from("/tmp")),
597            )
598            .await
599            .unwrap();
600
601        assert!(!output.success);
602        assert!(output.content.contains("Only \"script\" is supported"));
603    }
604
605    #[tokio::test]
606    async fn program_tool_rejects_missing_script_source_and_path() {
607        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
608        let output = tool
609            .execute(
610                &serde_json::json!({ "type": "script" }),
611                &ToolContext::new(PathBuf::from("/tmp")),
612            )
613            .await
614            .unwrap();
615
616        assert!(!output.success);
617        assert!(output.content.contains("requires either source or path"));
618    }
619
620    #[tokio::test]
621    async fn program_tool_rejects_unsupported_language() {
622        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
623        let output = tool
624            .execute(
625                &serde_json::json!({
626                    "type": "script",
627                    "language": "typescript",
628                    "source": "async function run() { return {}; }"
629                }),
630                &ToolContext::new(PathBuf::from("/tmp")),
631            )
632            .await
633            .unwrap();
634
635        assert!(!output.success);
636        assert!(output.content.contains("Unsupported script language"));
637    }
638
639    #[tokio::test]
640    async fn program_tool_rejects_unsupported_script_path() {
641        let dir = tempfile::tempdir().unwrap();
642        std::fs::write(dir.path().join("script.txt"), "async function run() {}").unwrap();
643        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(dir.path().to_path_buf())));
644        let output = tool
645            .execute(
646                &serde_json::json!({
647                    "type": "script",
648                    "path": "script.txt"
649                }),
650                &ToolContext::new(dir.path().to_path_buf()),
651            )
652            .await
653            .unwrap();
654
655        assert!(!output.success);
656        assert!(output.content.contains(".js or .mjs file"));
657    }
658
659    #[test]
660    fn program_tool_default_allowed_tools_include_registry_tools_except_program() {
661        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
662        registry.register(Arc::new(EchoTool));
663        registry.register_builtin(Arc::new(ProgramTool::new(Arc::new(ToolRegistry::new(
664            PathBuf::from("/tmp"),
665        )))));
666
667        let allowed = script_allowed_tools(&serde_json::json!({}), &registry);
668
669        assert!(allowed.contains("echo"));
670        assert!(!allowed.contains("program"));
671    }
672
673    #[tokio::test]
674    async fn program_tool_source_uses_default_all_registered_tools() {
675        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
676        registry.register(Arc::new(EchoTool));
677        let tool = ProgramTool::new(Arc::clone(&registry));
678        let output = tool
679            .execute(
680                &serde_json::json!({
681                    "type": "script",
682                    "source": r#"
683                        async function run(ctx, inputs) {
684                            const result = await ctx.tool("echo", { message: inputs.message });
685                            return { summary: result.output, result };
686                        }
687                    "#,
688                    "inputs": { "message": "hello" }
689                }),
690                &ToolContext::new(PathBuf::from("/tmp")),
691            )
692            .await
693            .unwrap();
694
695        assert!(output.success, "{}", output.content);
696        assert!(output.content.contains("echo:hello"));
697        let metadata = output.metadata.unwrap();
698        assert_eq!(metadata["program"]["runtime"], "embedded-quickjs");
699        assert_eq!(metadata["script_result"]["summary"], "echo:hello");
700    }
701
702    #[tokio::test]
703    async fn program_tool_explicit_allowed_tools_restrict_default_tools() {
704        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
705        registry.register(Arc::new(EchoTool));
706        let tool = ProgramTool::new(Arc::clone(&registry));
707        let output = tool
708            .execute(
709                &serde_json::json!({
710                    "type": "script",
711                    "source": r#"
712                        async function run(ctx) {
713                            await ctx.tool("echo", { message: "blocked" });
714                            return {};
715                        }
716                    "#,
717                    "allowed_tools": ["read"]
718                }),
719                &ToolContext::new(PathBuf::from("/tmp")),
720            )
721            .await
722            .unwrap();
723
724        assert!(!output.success);
725        assert!(output.content.contains("tool 'echo' is not allowed"));
726    }
727
728    #[tokio::test]
729    async fn program_tool_enforces_max_tool_calls() {
730        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
731        registry.register(Arc::new(EchoTool));
732        let tool = ProgramTool::new(Arc::clone(&registry));
733        let output = tool
734            .execute(
735                &serde_json::json!({
736                    "type": "script",
737                    "source": r#"
738                        async function run(ctx) {
739                            await ctx.tool("echo", { message: "one" });
740                            await ctx.tool("echo", { message: "two" });
741                            return {};
742                        }
743                    "#,
744                    "limits": { "maxToolCalls": 1 }
745                }),
746                &ToolContext::new(PathBuf::from("/tmp")),
747            )
748            .await
749            .unwrap();
750
751        assert!(!output.success);
752        assert!(output.content.contains("exceeded maxToolCalls=1"));
753    }
754
755    #[test]
756    fn program_tool_rejects_fetch_source_access() {
757        let err =
758            validate_script_source("export default async function run() { return fetch('/'); }")
759                .unwrap_err();
760        assert!(err.contains("fetch is not allowed"));
761    }
762
763    #[test]
764    fn program_tool_accepts_plain_function_run_entrypoint() {
765        let source = script_source_with_host_entrypoint(
766            "async function run(ctx, inputs) { return { summary: inputs.message }; }",
767        )
768        .unwrap();
769
770        assert!(source.contains("globalThis.__a3sResultJson"));
771        assert!(source.contains("async function run"));
772    }
773
774    #[test]
775    fn program_tool_renders_result_summary_and_tool_records() {
776        let output = render_script_output(
777            &serde_json::json!({ "summary": "done", "items": [1] }),
778            &[ScriptCallRecord {
779                tool_name: "echo".to_string(),
780                success: true,
781                exit_code: 0,
782                output_bytes: 8,
783                metadata: Some(serde_json::json!({ "kind": "test" })),
784            }],
785            "",
786        );
787
788        assert!(output.contains("Program script completed."));
789        assert!(output.contains("done"));
790        assert!(output.contains("echo (ok"));
791        assert!(output.contains("\"items\""));
792    }
793}
a3s_code_core/tools/program_tool.rs

a3s_code_core/tools/
program_tool.rs
