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    // Delegation tools can't run inside a PTC script: child agents need the
217    // multi-threaded session runtime, but the script executes on a nested
218    // single-thread runtime where they can't fan out. Force the model to call
219    // them directly instead of `ctx.tool("parallel_task", ...)`.
220    allowed.remove("task");
221    allowed.remove("parallel_task");
222    allowed
223}
224
225fn script_limits(args: &serde_json::Value) -> ScriptLimits {
226    args.get("limits")
227        .cloned()
228        .and_then(|value| serde_json::from_value(value).ok())
229        .unwrap_or(ScriptLimits {
230            timeout_ms: None,
231            max_tool_calls: None,
232            max_output_bytes: None,
233        })
234}
235
236fn validate_script_source(source: &str) -> std::result::Result<(), String> {
237    let forbidden = [
238        ("import ", "imports are not allowed inside PTC scripts"),
239        (
240            "import(",
241            "dynamic imports are not allowed inside PTC scripts",
242        ),
243        ("eval(", "eval is not allowed inside PTC scripts"),
244        (
245            "Function(",
246            "Function constructor is not allowed inside PTC scripts",
247        ),
248        ("Worker(", "Worker is not allowed inside PTC scripts"),
249        ("WebSocket", "WebSocket is not allowed inside PTC scripts"),
250        (
251            "fetch(",
252            "fetch is not allowed inside PTC scripts; use ctx tools instead",
253        ),
254    ];
255
256    for (needle, message) in forbidden {
257        if source.contains(needle) {
258            return Err(message.to_string());
259        }
260    }
261    Ok(())
262}
263
264async fn run_quickjs_script(
265    source: &str,
266    inputs: serde_json::Value,
267    registry: Arc<ToolRegistry>,
268    ctx: ToolContext,
269    allowed_tools: HashSet<String>,
270    limits: ScriptLimits,
271) -> Result<ToolOutput> {
272    let timeout_ms = limits.timeout_ms.unwrap_or(DEFAULT_SCRIPT_TIMEOUT_MS);
273    let max_tool_calls = limits
274        .max_tool_calls
275        .unwrap_or(DEFAULT_SCRIPT_MAX_TOOL_CALLS);
276    let max_output_bytes = limits
277        .max_output_bytes
278        .unwrap_or(DEFAULT_SCRIPT_MAX_OUTPUT_BYTES);
279    let executable_source = script_source_with_host_entrypoint(source)?;
280    let state = Arc::new(Mutex::new(ScriptVmState {
281        registry,
282        ctx,
283        allowed_tools,
284        max_tool_calls,
285        max_output_bytes,
286        tool_calls: 0,
287        records: Vec::new(),
288    }));
289
290    let vm_state = Arc::clone(&state);
291    let result = timeout(
292        Duration::from_millis(timeout_ms),
293        tokio::task::spawn_blocking(move || {
294            let runtime = tokio::runtime::Builder::new_current_thread()
295                .enable_all()
296                .build()
297                .map_err(|err| anyhow!("failed to create program VM runtime: {err}"))?;
298            runtime.block_on(run_embedded_script(
299                executable_source,
300                inputs,
301                vm_state,
302                timeout_ms,
303            ))
304        }),
305    )
306    .await;
307
308    match result {
309        Ok(Ok(Ok(result))) => {
310            let records = state.lock().await.records.clone();
311            let output = render_script_output(&result, &records, "");
312            Ok(ToolOutput::success(output).with_metadata(serde_json::json!({
313                "program": {
314                    "name": "script",
315                    "language": "javascript",
316                    "runtime": "embedded-quickjs",
317                    "success": true,
318                    "tool_calls": records.iter().map(script_record_to_value).collect::<Vec<_>>(),
319                },
320                "script_result": result,
321            })))
322        }
323        Ok(Ok(Err(err))) if is_quickjs_timeout(&err) => Ok(ToolOutput::error(format!(
324            "program script timed out after {timeout_ms} ms"
325        ))),
326        Ok(Ok(Err(err))) => Ok(ToolOutput::error(format!("program script error:\n{err}"))),
327        Ok(Err(err)) => Ok(ToolOutput::error(format!(
328            "program VM thread failed: {err}"
329        ))),
330        Err(_) => Ok(ToolOutput::error(format!(
331            "program script timed out after {timeout_ms} ms"
332        ))),
333    }
334}
335
336fn script_source_with_host_entrypoint(source: &str) -> Result<String> {
337    let rewritten = if source.contains("export default async function run") {
338        source.replacen("export default async function run", "async function run", 1)
339    } else if source.contains("export default function run") {
340        source.replacen("export default function run", "function run", 1)
341    } else if source.contains("async function run") || source.contains("function run") {
342        source.to_string()
343    } else {
344        return Err(anyhow!(
345            "PTC script must define async function run(ctx, inputs)"
346        ));
347    };
348
349    Ok(format!(
350        r#"{rewritten}
351
352globalThis.__a3sResultJson = (async () => JSON.stringify(await run(globalThis.__a3sCtx, globalThis.__a3sInputs)))();
353"#
354    ))
355}
356
357async fn run_embedded_script(
358    source: String,
359    inputs: serde_json::Value,
360    state: Arc<Mutex<ScriptVmState>>,
361    timeout_ms: u64,
362) -> Result<serde_json::Value> {
363    let runtime = AsyncRuntime::new()?;
364    let started = Instant::now();
365    runtime
366        .set_interrupt_handler(Some(Box::new(move || {
367            started.elapsed() >= Duration::from_millis(timeout_ms)
368        })))
369        .await;
370    runtime.set_memory_limit(64 * 1024 * 1024).await;
371    runtime.set_max_stack_size(512 * 1024).await;
372
373    let context = AsyncContext::full(&runtime).await?;
374    let inputs_json = serde_json::to_string(&inputs)?;
375    let script = format!("{}\n{}", embedded_script_bootstrap(&inputs_json), source);
376    let result_json = async_with!(context => |ctx| {
377        let state = Arc::clone(&state);
378        let host_tool = move |tool: String, args_json: String| {
379            let state = Arc::clone(&state);
380            async move { execute_host_tool_json(state, tool, args_json).await }
381        };
382        if let Err(err) = ctx.globals().set("__a3sHostTool", Func::from(Async(host_tool))) {
383            return Err(format!("failed to install program host tool: {err}"));
384        }
385        let promise: Promise = match ctx.eval(script) {
386            Ok(promise) => promise,
387            Err(err) => return Err(format!("failed to evaluate program script: {err}")),
388        };
389        promise
390            .into_future::<String>()
391            .await
392            .catch(&ctx)
393            .map_err(|err| err.to_string())
394    })
395    .await
396    .map_err(anyhow::Error::msg)?;
397
398    serde_json::from_str(&result_json)
399        .map_err(|err| anyhow!("program script returned invalid JSON: {err}"))
400}
401
402struct ScriptVmState {
403    registry: Arc<ToolRegistry>,
404    ctx: ToolContext,
405    allowed_tools: HashSet<String>,
406    max_tool_calls: usize,
407    max_output_bytes: usize,
408    tool_calls: usize,
409    records: Vec<ScriptCallRecord>,
410}
411
412fn embedded_script_bootstrap(inputs_json: &str) -> String {
413    format!(
414        r#"
415const __a3sCallTool = async (tool, args = {{}}) => {{
416  const response = await globalThis.__a3sHostTool(String(tool), JSON.stringify(args ?? {{}}));
417  return JSON.parse(response);
418}};
419
420const __a3sCtx = Object.freeze({{
421  tool: __a3sCallTool,
422  readFile: (path) => __a3sCallTool("read", {{ file_path: path }}).then((r) => r.output),
423  read: (path) => __a3sCallTool("read", {{ file_path: path }}),
424  grep: (pattern, options = {{}}) => __a3sCallTool("grep", {{ pattern, ...options }}).then((r) => r.output),
425  glob: (pattern, options = {{}}) => __a3sCallTool("glob", {{ pattern, ...options }}).then((r) => r.output),
426  ls: (path = ".") => __a3sCallTool("ls", {{ path }}).then((r) => r.output),
427  bash: (command) => __a3sCallTool("bash", {{ command }}).then((r) => r.output),
428  git: (args = {{}}) => __a3sCallTool("git", args),
429  webSearch: (params) => __a3sCallTool("web_search", params),
430  verify: (args) => __a3sCallTool("bash", args),
431}});
432
433Object.defineProperty(globalThis, "__a3sCtx", {{ value: __a3sCtx, configurable: false }});
434Object.defineProperty(globalThis, "__a3sInputs", {{ value: {inputs_json}, configurable: false }});
435Object.defineProperty(globalThis, "fetch", {{ value: undefined, configurable: false, writable: false }});
436Object.defineProperty(globalThis, "WebSocket", {{ value: undefined, configurable: false, writable: false }});
437Object.defineProperty(globalThis, "Worker", {{ value: undefined, configurable: false, writable: false }});
438"#
439    )
440}
441
442async fn execute_host_tool_json(
443    state: Arc<Mutex<ScriptVmState>>,
444    tool: String,
445    args_json: String,
446) -> rquickjs::Result<String> {
447    let args = serde_json::from_str(&args_json).map_err(|err| {
448        JsError::new_from_js_message("string", "object", format!("invalid tool args JSON: {err}"))
449    })?;
450    let (registry, ctx, max_output_bytes) = {
451        let mut script = state.lock().await;
452        if !script.allowed_tools.contains(&tool) {
453            return Err(JsError::new_from_js_message(
454                "tool",
455                "allowed tool",
456                format!("tool '{tool}' is not allowed for this PTC script"),
457            ));
458        }
459        script.tool_calls += 1;
460        if script.tool_calls > script.max_tool_calls {
461            return Err(JsError::new_from_js_message(
462                "tool call",
463                "limited tool call",
464                format!("PTC script exceeded maxToolCalls={}", script.max_tool_calls),
465            ));
466        }
467        (
468            Arc::clone(&script.registry),
469            script.ctx.clone(),
470            script.max_output_bytes,
471        )
472    };
473
474    let result = registry
475        .execute_with_context(&tool, &args, &ctx)
476        .await
477        .map_err(|err| JsError::new_from_js_message("tool", "result", err.to_string()))?;
478    let mut output = result.output;
479    if output.len() > max_output_bytes {
480        output = truncate_utf8(&output, max_output_bytes).to_string();
481    }
482    let success = result.exit_code == 0;
483    let metadata = result.metadata.clone();
484    let exit_code = result.exit_code;
485    let name = result.name;
486
487    {
488        let mut script = state.lock().await;
489        script.records.push(ScriptCallRecord {
490            tool_name: tool,
491            success,
492            exit_code,
493            output_bytes: output.len(),
494            metadata: metadata.clone(),
495        });
496    }
497
498    serde_json::to_string(&serde_json::json!({
499        "name": name,
500        "output": output,
501        "exitCode": exit_code,
502        "metadata": metadata,
503    }))
504    .map_err(|err| JsError::new_from_js_message("tool result", "json", err.to_string()))
505}
506
507fn is_quickjs_timeout(err: &anyhow::Error) -> bool {
508    let text = err.to_string();
509    text.contains("interrupted") || text.contains("InternalError")
510}
511
512fn script_record_to_value(record: &ScriptCallRecord) -> serde_json::Value {
513    serde_json::json!({
514        "tool_name": record.tool_name,
515        "success": record.success,
516        "exit_code": record.exit_code,
517        "output_bytes": record.output_bytes,
518        "metadata": record.metadata,
519    })
520}
521
522fn render_script_output(
523    result: &serde_json::Value,
524    records: &[ScriptCallRecord],
525    stderr: &str,
526) -> String {
527    let mut output = String::from("Program script completed.");
528    if let Some(summary) = result.get("summary").and_then(|value| value.as_str()) {
529        output.push('\n');
530        output.push_str(summary);
531    }
532
533    output.push_str(&format!("\n\nTool calls: {}", records.len()));
534    for (index, record) in records.iter().enumerate() {
535        output.push_str(&format!(
536            "\n{}. {} ({}, exit_code={}, output_bytes={})",
537            index + 1,
538            record.tool_name,
539            if record.success { "ok" } else { "failed" },
540            record.exit_code,
541            record.output_bytes
542        ));
543    }
544
545    output.push_str("\n\nResult:\n");
546    output.push_str(&serde_json::to_string_pretty(result).unwrap_or_else(|_| result.to_string()));
547
548    if !stderr.is_empty() {
549        output.push_str("\n\nstderr:\n");
550        output.push_str(stderr);
551    }
552
553    output
554}
555
556#[cfg(test)]
557mod tests {
558    use super::*;
559    use async_trait::async_trait;
560    use std::path::PathBuf;
561
562    struct EchoTool;
563
564    #[async_trait]
565    impl Tool for EchoTool {
566        fn name(&self) -> &str {
567            "echo"
568        }
569
570        fn description(&self) -> &str {
571            "Echo test tool"
572        }
573
574        fn parameters(&self) -> serde_json::Value {
575            serde_json::json!({
576                "type": "object",
577                "properties": {
578                    "message": { "type": "string" }
579                }
580            })
581        }
582
583        async fn execute(
584            &self,
585            args: &serde_json::Value,
586            _ctx: &ToolContext,
587        ) -> Result<ToolOutput> {
588            let message = args
589                .get("message")
590                .and_then(|value| value.as_str())
591                .unwrap_or("");
592            Ok(ToolOutput::success(format!("echo:{message}")))
593        }
594    }
595
596    #[tokio::test]
597    async fn program_tool_rejects_non_script_type() {
598        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
599        let output = tool
600            .execute(
601                &serde_json::json!({ "type": "program_code_search" }),
602                &ToolContext::new(PathBuf::from("/tmp")),
603            )
604            .await
605            .unwrap();
606
607        assert!(!output.success);
608        assert!(output.content.contains("Only \"script\" is supported"));
609    }
610
611    #[tokio::test]
612    async fn program_tool_rejects_missing_script_source_and_path() {
613        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
614        let output = tool
615            .execute(
616                &serde_json::json!({ "type": "script" }),
617                &ToolContext::new(PathBuf::from("/tmp")),
618            )
619            .await
620            .unwrap();
621
622        assert!(!output.success);
623        assert!(output.content.contains("requires either source or path"));
624    }
625
626    #[tokio::test]
627    async fn program_tool_rejects_unsupported_language() {
628        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
629        let output = tool
630            .execute(
631                &serde_json::json!({
632                    "type": "script",
633                    "language": "typescript",
634                    "source": "async function run() { return {}; }"
635                }),
636                &ToolContext::new(PathBuf::from("/tmp")),
637            )
638            .await
639            .unwrap();
640
641        assert!(!output.success);
642        assert!(output.content.contains("Unsupported script language"));
643    }
644
645    #[tokio::test]
646    async fn program_tool_rejects_unsupported_script_path() {
647        let dir = tempfile::tempdir().unwrap();
648        std::fs::write(dir.path().join("script.txt"), "async function run() {}").unwrap();
649        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(dir.path().to_path_buf())));
650        let output = tool
651            .execute(
652                &serde_json::json!({
653                    "type": "script",
654                    "path": "script.txt"
655                }),
656                &ToolContext::new(dir.path().to_path_buf()),
657            )
658            .await
659            .unwrap();
660
661        assert!(!output.success);
662        assert!(output.content.contains(".js or .mjs file"));
663    }
664
665    #[test]
666    fn program_tool_default_allowed_tools_include_registry_tools_except_program() {
667        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
668        registry.register(Arc::new(EchoTool));
669        registry.register_builtin(Arc::new(ProgramTool::new(Arc::new(ToolRegistry::new(
670            PathBuf::from("/tmp"),
671        )))));
672
673        let allowed = script_allowed_tools(&serde_json::json!({}), &registry);
674
675        assert!(allowed.contains("echo"));
676        assert!(!allowed.contains("program"));
677    }
678
679    #[tokio::test]
680    async fn program_tool_source_uses_default_all_registered_tools() {
681        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
682        registry.register(Arc::new(EchoTool));
683        let tool = ProgramTool::new(Arc::clone(&registry));
684        let output = tool
685            .execute(
686                &serde_json::json!({
687                    "type": "script",
688                    "source": r#"
689                        async function run(ctx, inputs) {
690                            const result = await ctx.tool("echo", { message: inputs.message });
691                            return { summary: result.output, result };
692                        }
693                    "#,
694                    "inputs": { "message": "hello" }
695                }),
696                &ToolContext::new(PathBuf::from("/tmp")),
697            )
698            .await
699            .unwrap();
700
701        assert!(output.success, "{}", output.content);
702        assert!(output.content.contains("echo:hello"));
703        let metadata = output.metadata.unwrap();
704        assert_eq!(metadata["program"]["runtime"], "embedded-quickjs");
705        assert_eq!(metadata["script_result"]["summary"], "echo:hello");
706    }
707
708    #[tokio::test]
709    async fn program_tool_explicit_allowed_tools_restrict_default_tools() {
710        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
711        registry.register(Arc::new(EchoTool));
712        let tool = ProgramTool::new(Arc::clone(&registry));
713        let output = tool
714            .execute(
715                &serde_json::json!({
716                    "type": "script",
717                    "source": r#"
718                        async function run(ctx) {
719                            await ctx.tool("echo", { message: "blocked" });
720                            return {};
721                        }
722                    "#,
723                    "allowed_tools": ["read"]
724                }),
725                &ToolContext::new(PathBuf::from("/tmp")),
726            )
727            .await
728            .unwrap();
729
730        assert!(!output.success);
731        assert!(output.content.contains("tool 'echo' is not allowed"));
732    }
733
734    #[tokio::test]
735    async fn program_tool_enforces_max_tool_calls() {
736        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
737        registry.register(Arc::new(EchoTool));
738        let tool = ProgramTool::new(Arc::clone(&registry));
739        let output = tool
740            .execute(
741                &serde_json::json!({
742                    "type": "script",
743                    "source": r#"
744                        async function run(ctx) {
745                            await ctx.tool("echo", { message: "one" });
746                            await ctx.tool("echo", { message: "two" });
747                            return {};
748                        }
749                    "#,
750                    "limits": { "maxToolCalls": 1 }
751                }),
752                &ToolContext::new(PathBuf::from("/tmp")),
753            )
754            .await
755            .unwrap();
756
757        assert!(!output.success);
758        assert!(output.content.contains("exceeded maxToolCalls=1"));
759    }
760
761    #[test]
762    fn program_tool_rejects_fetch_source_access() {
763        let err =
764            validate_script_source("export default async function run() { return fetch('/'); }")
765                .unwrap_err();
766        assert!(err.contains("fetch is not allowed"));
767    }
768
769    #[test]
770    fn program_tool_accepts_plain_function_run_entrypoint() {
771        let source = script_source_with_host_entrypoint(
772            "async function run(ctx, inputs) { return { summary: inputs.message }; }",
773        )
774        .unwrap();
775
776        assert!(source.contains("globalThis.__a3sResultJson"));
777        assert!(source.contains("async function run"));
778    }
779
780    #[test]
781    fn program_tool_renders_result_summary_and_tool_records() {
782        let output = render_script_output(
783            &serde_json::json!({ "summary": "done", "items": [1] }),
784            &[ScriptCallRecord {
785                tool_name: "echo".to_string(),
786                success: true,
787                exit_code: 0,
788                output_bytes: 8,
789                metadata: Some(serde_json::json!({ "kind": "test" })),
790            }],
791            "",
792        );
793
794        assert!(output.contains("Program script completed."));
795        assert!(output.contains("done"));
796        assert!(output.contains("echo (ok"));
797        assert!(output.contains("\"items\""));
798    }
799}
a3s_code_core/tools/program_tool.rs

a3s_code_core/tools/
program_tool.rs
