1use 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;
19const 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 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 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 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 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 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(®istry));
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(®istry));
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(®istry));
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(®istry));
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(®istry));
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(®istry));
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}