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 __a3sReadArgs = (path, options = {{}}) => ({{ ...(options ?? {{}}), file_path: path }});
451const __a3sCtx = Object.freeze({{
452  tool: __a3sCallTool,
453  tools: __a3sTools,
454  readFile: (path, options = {{}}) => __a3sCallTool("read", __a3sReadArgs(path, options)).then((r) => r.output),
455  read: (path, options = {{}}) => __a3sCallTool("read", __a3sReadArgs(path, options)),
456  grep: (pattern, options = {{}}) => __a3sCallTool("grep", {{ pattern, ...options }}).then((r) => r.output),
457  glob: (pattern, options = {{}}) => __a3sCallTool("glob", {{ pattern, ...options }}).then((r) => r.output),
458  ls: (path = ".") => __a3sCallTool("ls", {{ path }}).then((r) => r.output),
459  bash: (command) => __a3sCallTool("bash", {{ command }}).then((r) => r.output),
460  git: (args = {{}}) => __a3sCallTool("git", args),
461  webSearch: (params) => __a3sCallTool("web_search", params),
462  verify: (args) => __a3sCallTool("bash", args),
463}});
464
465Object.defineProperty(globalThis, "__a3sCtx", {{ value: __a3sCtx, configurable: false }});
466Object.defineProperty(globalThis, "__a3sInputs", {{ value: {inputs_json}, configurable: false }});
467Object.defineProperty(globalThis, "fetch", {{ value: undefined, configurable: false, writable: false }});
468Object.defineProperty(globalThis, "WebSocket", {{ value: undefined, configurable: false, writable: false }});
469Object.defineProperty(globalThis, "Worker", {{ value: undefined, configurable: false, writable: false }});
470"#
471    )
472}
473
474async fn execute_host_tool_json(
475    state: Arc<Mutex<ScriptVmState>>,
476    tool: String,
477    args_json: String,
478) -> rquickjs::Result<String> {
479    let args = serde_json::from_str(&args_json).map_err(|err| {
480        JsError::new_from_js_message("string", "object", format!("invalid tool args JSON: {err}"))
481    })?;
482    let (registry, ctx, max_output_bytes, outer) = {
483        let mut script = state.lock().await;
484        if !script.allowed_tools.contains(&tool) {
485            return Err(JsError::new_from_js_message(
486                "tool",
487                "allowed tool",
488                format!("tool '{tool}' is not allowed for this PTC script"),
489            ));
490        }
491        script.tool_calls += 1;
492        if script.tool_calls > script.max_tool_calls {
493            return Err(JsError::new_from_js_message(
494                "tool call",
495                "limited tool call",
496                format!("PTC script exceeded maxToolCalls={}", script.max_tool_calls),
497            ));
498        }
499        (
500            Arc::clone(&script.registry),
501            script.ctx.clone(),
502            script.max_output_bytes,
503            script.outer.clone(),
504        )
505    };
506
507    // Run the tool on the OUTER multi-threaded runtime (not this nested
508    // single-thread VM runtime) so host tools can use the session runtime
509    // normally. `parallel_task` is intentionally filtered before this point.
510    let tool_for_spawn = tool.clone();
511    let result = outer
512        .spawn(async move {
513            registry
514                .execute_with_context(&tool_for_spawn, &args, &ctx)
515                .await
516        })
517        .await
518        .map_err(|err| JsError::new_from_js_message("tool", "spawn", err.to_string()))?
519        .map_err(|err| JsError::new_from_js_message("tool", "result", err.to_string()))?;
520    let mut output = result.output;
521    if output.len() > max_output_bytes {
522        output = truncate_utf8(&output, max_output_bytes).to_string();
523    }
524    let success = result.exit_code == 0;
525    let metadata = result.metadata.clone();
526    let exit_code = result.exit_code;
527    let name = result.name;
528
529    {
530        let mut script = state.lock().await;
531        script.records.push(ScriptCallRecord {
532            tool_name: tool,
533            success,
534            exit_code,
535            output_bytes: output.len(),
536            metadata: metadata.clone(),
537        });
538    }
539
540    serde_json::to_string(&serde_json::json!({
541        "name": name,
542        "output": output,
543        "exitCode": exit_code,
544        "metadata": metadata,
545    }))
546    .map_err(|err| JsError::new_from_js_message("tool result", "json", err.to_string()))
547}
548
549fn is_quickjs_timeout(err: &anyhow::Error) -> bool {
550    let text = err.to_string();
551    text.contains("interrupted") || text.contains("InternalError")
552}
553
554fn script_record_to_value(record: &ScriptCallRecord) -> serde_json::Value {
555    serde_json::json!({
556        "tool_name": record.tool_name,
557        "success": record.success,
558        "exit_code": record.exit_code,
559        "output_bytes": record.output_bytes,
560        "metadata": record.metadata,
561    })
562}
563
564fn render_script_output(
565    result: &serde_json::Value,
566    records: &[ScriptCallRecord],
567    stderr: &str,
568) -> String {
569    let mut output = String::from("Program script completed.");
570    if let Some(summary) = result.get("summary").and_then(|value| value.as_str()) {
571        output.push('\n');
572        output.push_str(summary);
573    }
574
575    output.push_str(&format!("\n\nTool calls: {}", records.len()));
576    for (index, record) in records.iter().enumerate() {
577        output.push_str(&format!(
578            "\n{}. {} ({}, exit_code={}, output_bytes={})",
579            index + 1,
580            record.tool_name,
581            if record.success { "ok" } else { "failed" },
582            record.exit_code,
583            record.output_bytes
584        ));
585    }
586
587    output.push_str("\n\nResult:\n");
588    output.push_str(&serde_json::to_string_pretty(result).unwrap_or_else(|_| result.to_string()));
589
590    if !stderr.is_empty() {
591        output.push_str("\n\nstderr:\n");
592        output.push_str(stderr);
593    }
594
595    output
596}
597
598#[cfg(test)]
599mod tests {
600    use super::*;
601    use async_trait::async_trait;
602    use std::path::PathBuf;
603
604    struct EchoTool;
605
606    #[async_trait]
607    impl Tool for EchoTool {
608        fn name(&self) -> &str {
609            "echo"
610        }
611
612        fn description(&self) -> &str {
613            "Echo test tool"
614        }
615
616        fn parameters(&self) -> serde_json::Value {
617            serde_json::json!({
618                "type": "object",
619                "properties": {
620                    "message": { "type": "string" }
621                }
622            })
623        }
624
625        async fn execute(
626            &self,
627            args: &serde_json::Value,
628            _ctx: &ToolContext,
629        ) -> Result<ToolOutput> {
630            let message = args
631                .get("message")
632                .and_then(|value| value.as_str())
633                .unwrap_or("");
634            Ok(ToolOutput::success(format!("echo:{message}")))
635        }
636    }
637
638    struct ParallelTaskProbeTool;
639
640    #[async_trait]
641    impl Tool for ParallelTaskProbeTool {
642        fn name(&self) -> &str {
643            "parallel_task"
644        }
645
646        fn description(&self) -> &str {
647            "Probe tool used to verify PTC sandbox filtering."
648        }
649
650        fn parameters(&self) -> serde_json::Value {
651            serde_json::json!({ "type": "object" })
652        }
653
654        async fn execute(
655            &self,
656            _args: &serde_json::Value,
657            _ctx: &ToolContext,
658        ) -> Result<ToolOutput> {
659            Ok(ToolOutput::success("should-not-run"))
660        }
661    }
662
663    #[tokio::test]
664    async fn program_tool_rejects_non_script_type() {
665        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
666        let output = tool
667            .execute(
668                &serde_json::json!({ "type": "program_code_search" }),
669                &ToolContext::new(PathBuf::from("/tmp")),
670            )
671            .await
672            .unwrap();
673
674        assert!(!output.success);
675        assert!(output.content.contains("Only \"script\" is supported"));
676    }
677
678    #[tokio::test]
679    async fn program_tool_rejects_missing_script_source_and_path() {
680        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
681        let output = tool
682            .execute(
683                &serde_json::json!({ "type": "script" }),
684                &ToolContext::new(PathBuf::from("/tmp")),
685            )
686            .await
687            .unwrap();
688
689        assert!(!output.success);
690        assert!(output.content.contains("requires either source or path"));
691    }
692
693    #[tokio::test]
694    async fn program_tool_rejects_unsupported_language() {
695        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(PathBuf::from("/tmp"))));
696        let output = tool
697            .execute(
698                &serde_json::json!({
699                    "type": "script",
700                    "language": "typescript",
701                    "source": "async function run() { return {}; }"
702                }),
703                &ToolContext::new(PathBuf::from("/tmp")),
704            )
705            .await
706            .unwrap();
707
708        assert!(!output.success);
709        assert!(output.content.contains("Unsupported script language"));
710    }
711
712    #[tokio::test]
713    async fn program_tool_rejects_unsupported_script_path() {
714        let dir = tempfile::tempdir().unwrap();
715        std::fs::write(dir.path().join("script.txt"), "async function run() {}").unwrap();
716        let tool = ProgramTool::new(Arc::new(ToolRegistry::new(dir.path().to_path_buf())));
717        let output = tool
718            .execute(
719                &serde_json::json!({
720                    "type": "script",
721                    "path": "script.txt"
722                }),
723                &ToolContext::new(dir.path().to_path_buf()),
724            )
725            .await
726            .unwrap();
727
728        assert!(!output.success);
729        assert!(output.content.contains(".js or .mjs file"));
730    }
731
732    #[test]
733    fn program_tool_default_allowed_tools_include_registry_tools_except_program() {
734        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
735        registry.register(Arc::new(EchoTool));
736        registry.register_builtin(Arc::new(ProgramTool::new(Arc::new(ToolRegistry::new(
737            PathBuf::from("/tmp"),
738        )))));
739
740        let allowed = script_allowed_tools(&serde_json::json!({}), &registry);
741
742        assert!(allowed.contains("echo"));
743        assert!(!allowed.contains("program"));
744    }
745
746    #[test]
747    fn program_tool_forbids_parallel_task_in_scripts() {
748        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
749        let args = serde_json::json!({
750            "allowed_tools": ["parallel_task", "task", "program", "echo"]
751        });
752        let allowed = script_allowed_tools(&args, &registry);
753        assert!(!allowed.contains("parallel_task"));
754        assert!(allowed.contains("task"));
755        assert!(allowed.contains("echo"));
756        assert!(!allowed.contains("program"));
757    }
758
759    #[tokio::test]
760    async fn program_tool_rejects_parallel_task_even_when_explicitly_allowed() {
761        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
762        registry.register(Arc::new(ParallelTaskProbeTool));
763        let tool = ProgramTool::new(Arc::clone(&registry));
764
765        let output = tool
766            .execute(
767                &serde_json::json!({
768                    "type": "script",
769                    "source": r#"
770                        async function run(ctx) {
771                            return await ctx.tool("parallel_task", { tasks: [] });
772                        }
773                    "#,
774                    "allowed_tools": ["parallel_task"]
775                }),
776                &ToolContext::new(PathBuf::from("/tmp")),
777            )
778            .await
779            .unwrap();
780
781        assert!(!output.success);
782        assert!(output
783            .content
784            .contains("tool 'parallel_task' is not allowed"));
785        assert!(!output.content.contains("should-not-run"));
786    }
787
788    #[tokio::test]
789    async fn program_tool_source_uses_default_all_registered_tools() {
790        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
791        registry.register(Arc::new(EchoTool));
792        let tool = ProgramTool::new(Arc::clone(&registry));
793        let output = tool
794            .execute(
795                &serde_json::json!({
796                    "type": "script",
797                    "source": r#"
798                        async function run(ctx, inputs) {
799                            const result = await ctx.tool("echo", { message: inputs.message });
800                            return { summary: result.output, result };
801                        }
802                    "#,
803                    "inputs": { "message": "hello" }
804                }),
805                &ToolContext::new(PathBuf::from("/tmp")),
806            )
807            .await
808            .unwrap();
809
810        assert!(output.success, "{}", output.content);
811        assert!(output.content.contains("echo:hello"));
812        let metadata = output.metadata.unwrap();
813        assert_eq!(metadata["program"]["runtime"], "embedded-quickjs");
814        assert_eq!(metadata["script_result"]["summary"], "echo:hello");
815    }
816
817    #[tokio::test]
818    async fn program_tool_ctx_read_file_passes_offset_and_limit() {
819        let dir = tempfile::tempdir().unwrap();
820        std::fs::write(dir.path().join("notes.txt"), "one\ntwo\nthree\n").unwrap();
821        let executor = crate::tools::ToolExecutor::new(dir.path().to_string_lossy().to_string());
822        let tool = ProgramTool::new(Arc::clone(executor.registry()));
823
824        let output = tool
825            .execute(
826                &serde_json::json!({
827                    "type": "script",
828                    "source": r#"
829                        async function run(ctx) {
830                            const text = await ctx.readFile("notes.txt", { offset: 1, limit: 1 });
831                            const result = await ctx.read("notes.txt", { offset: 2, limit: 1 });
832                            return { text, raw: result.output };
833                        }
834                    "#,
835                    "allowed_tools": ["read"]
836                }),
837                &ToolContext::new(dir.path().to_path_buf()),
838            )
839            .await
840            .unwrap();
841
842        assert!(output.success, "{}", output.content);
843        let metadata = output.metadata.unwrap();
844        let text = metadata["script_result"]["text"].as_str().unwrap();
845        assert!(text.contains("two"));
846        assert!(!text.contains("one"));
847        let raw = metadata["script_result"]["raw"].as_str().unwrap();
848        assert!(raw.contains("three"));
849        assert!(!raw.contains("two"));
850    }
851
852    #[tokio::test]
853    async fn program_tool_exposes_ctx_tools_proxy_for_named_tools() {
854        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
855        registry.register(Arc::new(EchoTool));
856        let tool = ProgramTool::new(Arc::clone(&registry));
857        let output = tool
858            .execute(
859                &serde_json::json!({
860                    "type": "script",
861                    "source": r#"
862                        async function run(ctx, inputs) {
863                            const result = await ctx.tools.echo({ message: inputs.message });
864                            return { summary: result.output, result };
865                        }
866                    "#,
867                    "inputs": { "message": "proxy" }
868                }),
869                &ToolContext::new(PathBuf::from("/tmp")),
870            )
871            .await
872            .unwrap();
873
874        assert!(output.success, "{}", output.content);
875        let metadata = output.metadata.unwrap();
876        assert_eq!(metadata["script_result"]["summary"], "echo:proxy");
877        assert_eq!(metadata["program"]["tool_calls"][0]["tool_name"], "echo");
878    }
879
880    #[tokio::test]
881    async fn program_tool_ctx_tools_proxy_respects_allowed_tools() {
882        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
883        registry.register(Arc::new(EchoTool));
884        let tool = ProgramTool::new(Arc::clone(&registry));
885        let output = tool
886            .execute(
887                &serde_json::json!({
888                    "type": "script",
889                    "source": r#"
890                        async function run(ctx) {
891                            await ctx.tools.echo({ message: "blocked" });
892                            return {};
893                        }
894                    "#,
895                    "allowed_tools": ["read"]
896                }),
897                &ToolContext::new(PathBuf::from("/tmp")),
898            )
899            .await
900            .unwrap();
901
902        assert!(!output.success);
903        assert!(output.content.contains("tool 'echo' is not allowed"));
904    }
905
906    #[tokio::test]
907    async fn program_tool_explicit_allowed_tools_restrict_default_tools() {
908        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
909        registry.register(Arc::new(EchoTool));
910        let tool = ProgramTool::new(Arc::clone(&registry));
911        let output = tool
912            .execute(
913                &serde_json::json!({
914                    "type": "script",
915                    "source": r#"
916                        async function run(ctx) {
917                            await ctx.tool("echo", { message: "blocked" });
918                            return {};
919                        }
920                    "#,
921                    "allowed_tools": ["read"]
922                }),
923                &ToolContext::new(PathBuf::from("/tmp")),
924            )
925            .await
926            .unwrap();
927
928        assert!(!output.success);
929        assert!(output.content.contains("tool 'echo' is not allowed"));
930    }
931
932    #[tokio::test]
933    async fn program_tool_enforces_max_tool_calls() {
934        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
935        registry.register(Arc::new(EchoTool));
936        let tool = ProgramTool::new(Arc::clone(&registry));
937        let output = tool
938            .execute(
939                &serde_json::json!({
940                    "type": "script",
941                    "source": r#"
942                        async function run(ctx) {
943                            await ctx.tool("echo", { message: "one" });
944                            await ctx.tool("echo", { message: "two" });
945                            return {};
946                        }
947                    "#,
948                    "limits": { "maxToolCalls": 1 }
949                }),
950                &ToolContext::new(PathBuf::from("/tmp")),
951            )
952            .await
953            .unwrap();
954
955        assert!(!output.success);
956        assert!(output.content.contains("exceeded maxToolCalls=1"));
957    }
958
959    #[test]
960    fn program_tool_rejects_fetch_source_access() {
961        let err =
962            validate_script_source("export default async function run() { return fetch('/'); }")
963                .unwrap_err();
964        assert!(err.contains("fetch is not allowed"));
965    }
966
967    #[test]
968    fn program_tool_accepts_plain_function_run_entrypoint() {
969        let source = script_source_with_host_entrypoint(
970            "async function run(ctx, inputs) { return { summary: inputs.message }; }",
971        )
972        .unwrap();
973
974        assert!(source.contains("globalThis.__a3sResultJson"));
975        assert!(source.contains("async function run"));
976    }
977
978    #[test]
979    fn program_tool_renders_result_summary_and_tool_records() {
980        let output = render_script_output(
981            &serde_json::json!({ "summary": "done", "items": [1] }),
982            &[ScriptCallRecord {
983                tool_name: "echo".to_string(),
984                success: true,
985                exit_code: 0,
986                output_bytes: 8,
987                metadata: Some(serde_json::json!({ "kind": "test" })),
988            }],
989            "",
990        );
991
992        assert!(output.contains("Program script completed."));
993        assert!(output.contains("done"));
994        assert!(output.contains("echo (ok"));
995        assert!(output.contains("\"items\""));
996    }
997}