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