Skip to main content

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