1use crate::program::ProgramCatalog;
4use crate::text::truncate_utf8;
5use crate::tools::types::{Tool, ToolContext, ToolOutput};
6use crate::tools::{registry_bound_tool_invoker, registry_tool_invoker, 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 PROGRAM_CANCELLATION_SETTLE_GRACE: Duration = Duration::from_millis(500);
25pub const MAX_PROGRAM_SCRIPT_SOURCE_BYTES: usize = 192 * 1024;
29
30pub struct ProgramTool {
31 fallback_invoker: Arc<dyn ToolInvoker>,
32}
33
34impl ProgramTool {
35 pub fn new(registry: Arc<ToolRegistry>) -> Self {
36 Self {
37 fallback_invoker: registry_tool_invoker(registry),
38 }
39 }
40
41 pub fn with_catalog(registry: Arc<ToolRegistry>, _catalog: ProgramCatalog) -> Self {
42 Self::new(registry)
43 }
44
45 pub(crate) fn with_catalog_registry_bound(
46 registry: Arc<ToolRegistry>,
47 _catalog: ProgramCatalog,
48 ) -> Self {
49 Self {
50 fallback_invoker: registry_bound_tool_invoker(registry),
51 }
52 }
53}
54
55#[async_trait]
56impl Tool for ProgramTool {
57 fn name(&self) -> &str {
58 "program"
59 }
60
61 fn description(&self) -> &str {
62 "Run a sandboxed JavaScript PTC script. The script defines async function run(ctx, inputs) and may call only allowed ctx tools."
63 }
64
65 fn parameters(&self) -> serde_json::Value {
66 serde_json::json!({
67 "type": "object",
68 "additionalProperties": false,
69 "properties": {
70 "type": {
71 "type": "string",
72 "description": "Required. Program kind. Only \"script\" is supported.",
73 "enum": ["script"]
74 },
75 "inputs": {
76 "type": "object",
77 "description": "Optional. JSON inputs passed to the script as the second argument."
78 },
79 "language": {
80 "type": "string",
81 "description": "Script language. Only JavaScript is supported.",
82 "enum": ["javascript"]
83 },
84 "source": {
85 "type": "string",
86 "description": "Inline JavaScript source defining async function run(ctx, inputs)."
87 },
88 "path": {
89 "type": "string",
90 "description": "Workspace-relative path to a .js or .mjs script defining async function run(ctx, inputs). Used when source is omitted."
91 },
92 "allowed_tools": {
93 "type": "array",
94 "description": "Tool names the script may call through ctx. Defaults to all registered tools except program, dynamic_workflow, and the legacy parallel_task alias.",
95 "items": { "type": "string" }
96 },
97 "limits": {
98 "type": "object",
99 "description": "Optional timeoutMs, maxToolCalls, and maxOutputBytes.",
100 "additionalProperties": false,
101 "properties": {
102 "timeoutMs": { "type": "integer", "minimum": 1 },
103 "maxToolCalls": { "type": "integer", "minimum": 1 },
104 "maxOutputBytes": { "type": "integer", "minimum": 1 }
105 }
106 }
107 },
108 "required": ["type"]
109 })
110 }
111
112 async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
113 let Some(kind) = args.get("type").and_then(|value| value.as_str()) else {
114 return Ok(ToolOutput::error("type parameter is required"));
115 };
116 if kind != "script" {
117 return Ok(ToolOutput::error(format!(
118 "Unsupported program type: {kind}. Only \"script\" is supported."
119 )));
120 }
121 let inputs = args
122 .get("inputs")
123 .cloned()
124 .unwrap_or_else(|| serde_json::json!({}));
125
126 let invoker = ctx
127 .tool_invoker()
128 .unwrap_or_else(|| Arc::clone(&self.fallback_invoker));
129 execute_script_program(args, inputs, invoker, ctx).await
130 }
131}
132
133#[derive(Debug, Deserialize)]
134#[serde(rename_all = "camelCase")]
135struct ScriptLimits {
136 timeout_ms: Option<u64>,
137 max_tool_calls: Option<usize>,
138 max_output_bytes: Option<usize>,
139}
140
141#[derive(Debug, Clone)]
142struct ScriptCallRecord {
143 tool_name: String,
144 success: bool,
145 exit_code: i32,
146 output_bytes: usize,
147 metadata: Option<serde_json::Value>,
148}
149
150async fn execute_script_program(
151 args: &serde_json::Value,
152 inputs: serde_json::Value,
153 invoker: Arc<dyn ToolInvoker>,
154 ctx: &ToolContext,
155) -> Result<ToolOutput> {
156 let language = args
157 .get("language")
158 .and_then(|value| value.as_str())
159 .unwrap_or("javascript");
160 if language != "javascript" {
161 return Ok(ToolOutput::error(format!(
162 "Unsupported script language: {language}"
163 )));
164 }
165
166 let source = match load_script_source(args, ctx).await {
167 Ok(source) => source,
168 Err(message) => return Ok(ToolOutput::error(message)),
169 };
170 if source.len() > MAX_PROGRAM_SCRIPT_SOURCE_BYTES {
171 return Ok(ToolOutput::error(format!(
172 "script source is too large: {} bytes exceeds {} bytes",
173 source.len(),
174 MAX_PROGRAM_SCRIPT_SOURCE_BYTES
175 )));
176 }
177 if let Err(message) = validate_script_source(&source) {
178 return Ok(ToolOutput::error(message));
179 }
180
181 let allowed_tools = script_allowed_tools(args, invoker.available_tools());
182 let limits = script_limits(args);
183 match run_quickjs_script(&source, inputs, invoker, ctx.clone(), allowed_tools, limits).await {
184 Ok(output) => Ok(output),
185 Err(err) => Ok(ToolOutput::error(format!("program script failed: {err}"))),
186 }
187}
188
189async fn load_script_source(
190 args: &serde_json::Value,
191 ctx: &ToolContext,
192) -> std::result::Result<String, String> {
193 if let Some(source) = args.get("source").and_then(|value| value.as_str()) {
194 return Ok(source.to_string());
195 }
196
197 let Some(path) = args.get("path").and_then(|value| value.as_str()) else {
198 return Err("program script requires either source or path".to_string());
199 };
200 if !(path.ends_with(".js") || path.ends_with(".mjs")) {
201 return Err("program script path must point to a .js or .mjs file".to_string());
202 }
203
204 let workspace_path = ctx
205 .resolve_workspace_path(path)
206 .map_err(|err| format!("failed to resolve script path: {err}"))?;
207 ctx.workspace_services
208 .fs()
209 .read_text(&workspace_path)
210 .await
211 .map_err(|err| format!("failed to read script path '{}': {err}", path))
212}
213
214fn script_allowed_tools(args: &serde_json::Value, available_tools: Vec<String>) -> HashSet<String> {
215 let mut allowed = args
216 .get("allowed_tools")
217 .and_then(|value| value.as_array())
218 .map(|items| {
219 items
220 .iter()
221 .filter_map(|item| item.as_str())
222 .map(ToString::to_string)
223 .collect::<HashSet<_>>()
224 })
225 .unwrap_or_else(|| available_tools.into_iter().collect());
226
227 allowed.remove("program");
228 allowed.remove("dynamic_workflow");
233 allowed.remove("parallel_task");
234 allowed
235}
236
237fn script_limits(args: &serde_json::Value) -> ScriptLimits {
238 args.get("limits")
239 .cloned()
240 .and_then(|value| serde_json::from_value(value).ok())
241 .unwrap_or(ScriptLimits {
242 timeout_ms: None,
243 max_tool_calls: None,
244 max_output_bytes: None,
245 })
246}
247
248fn validate_script_source(source: &str) -> std::result::Result<(), String> {
249 let forbidden = [
250 ("import ", "imports are not allowed inside PTC scripts"),
251 (
252 "import(",
253 "dynamic imports are not allowed inside PTC scripts",
254 ),
255 ("eval(", "eval is not allowed inside PTC scripts"),
256 (
257 "Function(",
258 "Function constructor is not allowed inside PTC scripts",
259 ),
260 ("Worker(", "Worker is not allowed inside PTC scripts"),
261 ("WebSocket", "WebSocket is not allowed inside PTC scripts"),
262 (
263 "fetch(",
264 "fetch is not allowed inside PTC scripts; use ctx tools instead",
265 ),
266 ];
267
268 for (needle, message) in forbidden {
269 if source.contains(needle) {
270 return Err(message.to_string());
271 }
272 }
273 Ok(())
274}
275
276async fn run_quickjs_script(
277 source: &str,
278 inputs: serde_json::Value,
279 invoker: Arc<dyn ToolInvoker>,
280 ctx: ToolContext,
281 allowed_tools: HashSet<String>,
282 limits: ScriptLimits,
283) -> Result<ToolOutput> {
284 let delegating = allowed_tools.contains("task");
290 let timeout_ms = limits.timeout_ms.unwrap_or(if delegating {
291 DELEGATION_SCRIPT_TIMEOUT_MS
292 } else {
293 DEFAULT_SCRIPT_TIMEOUT_MS
294 });
295 let max_tool_calls = limits
296 .max_tool_calls
297 .unwrap_or(DEFAULT_SCRIPT_MAX_TOOL_CALLS);
298 let max_output_bytes = limits
299 .max_output_bytes
300 .unwrap_or(DEFAULT_SCRIPT_MAX_OUTPUT_BYTES);
301 let executable_source = script_source_with_host_entrypoint(source)?;
302 let parent_cancellation = ctx.cancellation_token();
303 let program_cancellation = parent_cancellation.child_token();
304 let outer = tokio::runtime::Handle::current();
308 let state = Arc::new(Mutex::new(ScriptVmState {
309 invoker,
310 ctx: ctx.with_cancellation(program_cancellation.clone()),
311 allowed_tools,
312 max_tool_calls,
313 max_output_bytes,
314 tool_calls: 0,
315 records: Vec::new(),
316 outer,
317 }));
318
319 let vm_state = Arc::clone(&state);
320 let mut vm = tokio::task::spawn_blocking(move || {
321 let runtime = tokio::runtime::Builder::new_current_thread()
322 .enable_all()
323 .build()
324 .map_err(|err| anyhow!("failed to create program VM runtime: {err}"))?;
325 runtime.block_on(run_embedded_script(
326 executable_source,
327 inputs,
328 vm_state,
329 timeout_ms,
330 program_cancellation,
331 ))
332 });
333
334 enum Stop {
335 Cancelled,
336 TimedOut,
337 }
338 let result = tokio::select! {
339 biased;
340 _ = parent_cancellation.cancelled() => None,
341 result = &mut vm => Some(result),
342 _ = tokio::time::sleep(Duration::from_millis(timeout_ms)) => None,
343 };
344 let stop = if result.is_none() {
345 if parent_cancellation.is_cancelled() {
346 Some(Stop::Cancelled)
347 } else {
348 Some(Stop::TimedOut)
349 }
350 } else {
351 None
352 };
353
354 if let Some(stop) = stop {
355 state.lock().await.ctx.cancellation_token().cancel();
359 if timeout(PROGRAM_CANCELLATION_SETTLE_GRACE, &mut vm)
360 .await
361 .is_err()
362 {
363 vm.abort();
364 let _ = vm.await;
365 }
366 return Ok(ToolOutput::error(match stop {
367 Stop::Cancelled => "program script cancelled by caller".to_string(),
368 Stop::TimedOut => format!("program script timed out after {timeout_ms} ms"),
369 }));
370 }
371
372 let result = result.expect("completed VM result is present");
373
374 match result {
375 Ok(Ok(result)) => {
376 let records = state.lock().await.records.clone();
377 let output = render_script_output(&result, &records, "");
378 Ok(ToolOutput::success(output).with_metadata(serde_json::json!({
379 "program": {
380 "name": "script",
381 "language": "javascript",
382 "runtime": "embedded-quickjs",
383 "success": true,
384 "tool_calls": records.iter().map(script_record_to_value).collect::<Vec<_>>(),
385 },
386 "script_result": result,
387 })))
388 }
389 Ok(Err(err)) if is_quickjs_timeout(&err) => Ok(ToolOutput::error(format!(
390 "program script timed out after {timeout_ms} ms"
391 ))),
392 Ok(Err(err)) => Ok(ToolOutput::error(format!("program script error:\n{err}"))),
393 Err(err) => Ok(ToolOutput::error(format!(
394 "program VM thread failed: {err}"
395 ))),
396 }
397}
398
399fn script_source_with_host_entrypoint(source: &str) -> Result<String> {
400 let rewritten = if source.contains("export default async function run") {
401 source.replacen("export default async function run", "async function run", 1)
402 } else if source.contains("export default function run") {
403 source.replacen("export default function run", "function run", 1)
404 } else if source.contains("async function run") || source.contains("function run") {
405 source.to_string()
406 } else {
407 return Err(anyhow!(
408 "PTC script must define async function run(ctx, inputs)"
409 ));
410 };
411
412 Ok(format!(
413 r#"{rewritten}
414
415globalThis.__a3sResultJson = (async () => JSON.stringify(await run(globalThis.__a3sCtx, globalThis.__a3sInputs)))();
416"#
417 ))
418}
419
420async fn run_embedded_script(
421 source: String,
422 inputs: serde_json::Value,
423 state: Arc<Mutex<ScriptVmState>>,
424 timeout_ms: u64,
425 cancellation: tokio_util::sync::CancellationToken,
426) -> Result<serde_json::Value> {
427 let runtime = AsyncRuntime::new()?;
428 let started = Instant::now();
429 runtime
430 .set_interrupt_handler(Some(Box::new(move || {
431 cancellation.is_cancelled() || started.elapsed() >= Duration::from_millis(timeout_ms)
432 })))
433 .await;
434 runtime.set_memory_limit(64 * 1024 * 1024).await;
435 runtime.set_max_stack_size(512 * 1024).await;
436
437 let context = AsyncContext::full(&runtime).await?;
438 let inputs_json = serde_json::to_string(&inputs)?;
439 let script = format!("{}\n{}", embedded_script_bootstrap(&inputs_json), source);
440 let result_json = async_with!(context => |ctx| {
441 let state = Arc::clone(&state);
442 let host_tool = move |tool: String, args_json: String| {
443 let state = Arc::clone(&state);
444 async move { execute_host_tool_json(state, tool, args_json).await }
445 };
446 if let Err(err) = ctx.globals().set("__a3sHostTool", Func::from(Async(host_tool))) {
447 return Err(format!("failed to install program host tool: {err}"));
448 }
449 let promise: Promise = match ctx.eval(script) {
450 Ok(promise) => promise,
451 Err(err) => return Err(format!("failed to evaluate program script: {err}")),
452 };
453 promise
454 .into_future::<String>()
455 .await
456 .catch(&ctx)
457 .map_err(|err| err.to_string())
458 })
459 .await
460 .map_err(anyhow::Error::msg)?;
461
462 serde_json::from_str(&result_json)
463 .map_err(|err| anyhow!("program script returned invalid JSON: {err}"))
464}
465
466struct ScriptVmState {
467 invoker: Arc<dyn ToolInvoker>,
468 ctx: ToolContext,
469 allowed_tools: HashSet<String>,
470 max_tool_calls: usize,
471 max_output_bytes: usize,
472 tool_calls: usize,
473 records: Vec<ScriptCallRecord>,
474 outer: tokio::runtime::Handle,
478}
479
480fn embedded_script_bootstrap(inputs_json: &str) -> String {
481 format!(
482 r#"
483const __a3sCallTool = async (tool, args = {{}}) => {{
484 const response = await globalThis.__a3sHostTool(String(tool), JSON.stringify(args ?? {{}}));
485 return JSON.parse(response);
486}};
487
488const __a3sTools = Object.freeze(new Proxy({{}}, {{
489 get(_target, prop) {{
490 if (typeof prop !== "string" || prop === "then") return undefined;
491 return (args = {{}}) => __a3sCallTool(prop, args);
492 }},
493 has(_target, prop) {{
494 return typeof prop === "string";
495 }},
496}}));
497
498const __a3sReadArgs = (path, options = {{}}) => ({{ ...(options ?? {{}}), file_path: path }});
499const __a3sLegacySearchArgs = (mode, query, options = {{}}) => {{
500 const args = {{ ...(options ?? {{}}), mode, query }};
501 if (Object.prototype.hasOwnProperty.call(args, "glob")) {{
502 args.include = args.glob;
503 delete args.glob;
504 }}
505 if (Object.prototype.hasOwnProperty.call(args, "-i")) {{
506 args.case_sensitive = !args["-i"];
507 delete args["-i"];
508 }}
509 return args;
510}};
511const __a3sCtx = Object.freeze({{
512 tool: __a3sCallTool,
513 tools: __a3sTools,
514 readFile: (path, options = {{}}) => __a3sCallTool("read", __a3sReadArgs(path, options)).then((r) => r.text ?? r.output),
515 read: (path, options = {{}}) => __a3sCallTool("read", __a3sReadArgs(path, options)),
516 search: (query, options = {{}}) => __a3sCallTool("search", {{ ...options, query }}).then((r) => r.output),
517 grep: (query, options = {{}}) => __a3sCallTool("search", __a3sLegacySearchArgs("grep", query, options)).then((r) => r.output),
518 bm25: (query, options = {{}}) => __a3sCallTool("search", __a3sLegacySearchArgs("bm25", query, options)).then((r) => r.output),
519 glob: (query, options = {{}}) => __a3sCallTool("search", {{ ...options, mode: "glob", query }}).then((r) => r.output),
520 ls: (path = ".") => __a3sCallTool("ls", {{ path }}).then((r) => r.output),
521 bash: (command) => __a3sCallTool("bash", {{ command }}).then((r) => r.output),
522 git: (args = {{}}) => __a3sCallTool("git", args),
523 webSearch: (params) => __a3sCallTool("web_search", params),
524 verify: (args) => __a3sCallTool("bash", args),
525}});
526
527Object.defineProperty(globalThis, "__a3sCtx", {{ value: __a3sCtx, configurable: false }});
528Object.defineProperty(globalThis, "__a3sInputs", {{ value: {inputs_json}, configurable: false }});
529Object.defineProperty(globalThis, "fetch", {{ value: undefined, configurable: false, writable: false }});
530Object.defineProperty(globalThis, "WebSocket", {{ value: undefined, configurable: false, writable: false }});
531Object.defineProperty(globalThis, "Worker", {{ value: undefined, configurable: false, writable: false }});
532"#
533 )
534}
535
536async fn execute_host_tool_json(
537 state: Arc<Mutex<ScriptVmState>>,
538 tool: String,
539 args_json: String,
540) -> rquickjs::Result<String> {
541 let args = serde_json::from_str(&args_json).map_err(|err| {
542 JsError::new_from_js_message("string", "object", format!("invalid tool args JSON: {err}"))
543 })?;
544 let (invoker, ctx, max_output_bytes, outer) = {
545 let mut script = state.lock().await;
546 if !script_tool_is_allowed(&script.allowed_tools, &tool, &args) {
547 return Err(JsError::new_from_js_message(
548 "tool",
549 "allowed tool",
550 format!("tool '{tool}' is not allowed for this PTC script"),
551 ));
552 }
553 script.tool_calls += 1;
554 if script.tool_calls > script.max_tool_calls {
555 return Err(JsError::new_from_js_message(
556 "tool call",
557 "limited tool call",
558 format!("PTC script exceeded maxToolCalls={}", script.max_tool_calls),
559 ));
560 }
561 (
562 Arc::clone(&script.invoker),
563 script.ctx.clone(),
564 script.max_output_bytes,
565 script.outer.clone(),
566 )
567 };
568
569 let tool_for_spawn = tool.clone();
573 let result = outer
574 .spawn(async move {
575 invoker
576 .invoke(ctx.nested_tool_invocation(tool_for_spawn, args), &ctx)
577 .await
578 })
579 .await
580 .map_err(|err| JsError::new_from_js_message("tool", "spawn", err.to_string()))?;
581 let mut output = result.output;
582 if output.len() > max_output_bytes {
583 output = truncate_utf8(&output, max_output_bytes).to_string();
584 }
585 let success = result.exit_code == 0;
586 let metadata = result.metadata.clone();
587 let exit_code = result.exit_code;
588 let name = result.name;
589 let text = (tool == "read" && success)
590 .then(|| program_read_text(&output, metadata.as_ref()))
591 .flatten();
592
593 {
594 let mut script = state.lock().await;
595 script.records.push(ScriptCallRecord {
596 tool_name: tool,
597 success,
598 exit_code,
599 output_bytes: output.len(),
600 metadata: metadata.clone(),
601 });
602 }
603
604 serde_json::to_string(&serde_json::json!({
605 "name": name,
606 "output": output,
607 "text": text,
608 "exitCode": exit_code,
609 "metadata": metadata,
610 }))
611 .map_err(|err| JsError::new_from_js_message("tool result", "json", err.to_string()))
612}
613
614fn program_read_text(output: &str, metadata: Option<&serde_json::Value>) -> Option<String> {
620 let returned_lines = metadata?
621 .pointer("/range/returned_lines")?
622 .as_u64()
623 .and_then(|value| usize::try_from(value).ok())?;
624 if returned_lines == 0 {
625 return Some(String::new());
626 }
627
628 let mut text = String::new();
629 let mut rendered_lines = output.split_inclusive('\n');
630 for _ in 0..returned_lines {
631 let rendered = rendered_lines.next()?;
632 let (anchor, content) = rendered.split_once('\t')?;
633 if anchor.len() != 6 || anchor.trim().parse::<usize>().is_err() {
634 return None;
635 }
636 text.push_str(content);
637 }
638 Some(text)
639}
640
641fn script_tool_is_allowed(
642 allowed_tools: &HashSet<String>,
643 tool: &str,
644 args: &serde_json::Value,
645) -> bool {
646 if tool == "task"
647 && args
648 .get("tasks")
649 .and_then(serde_json::Value::as_array)
650 .is_some_and(|tasks| tasks.len() > 1)
651 {
652 return false;
653 }
654 if allowed_tools.contains(tool) {
655 return true;
656 }
657 tool == "search"
658 && args
659 .get("mode")
660 .and_then(serde_json::Value::as_str)
661 .is_some_and(|mode| allowed_tools.contains(mode))
662}
663
664fn is_quickjs_timeout(err: &anyhow::Error) -> bool {
665 let text = err.to_string();
666 text.contains("interrupted") || text.contains("InternalError")
667}
668
669fn script_record_to_value(record: &ScriptCallRecord) -> serde_json::Value {
670 serde_json::json!({
671 "tool_name": record.tool_name,
672 "success": record.success,
673 "exit_code": record.exit_code,
674 "output_bytes": record.output_bytes,
675 "metadata": record.metadata,
676 })
677}
678
679fn render_script_output(
680 result: &serde_json::Value,
681 records: &[ScriptCallRecord],
682 stderr: &str,
683) -> String {
684 let mut output = String::from("Program script completed.");
685 if let Some(summary) = result.get("summary").and_then(|value| value.as_str()) {
686 output.push('\n');
687 output.push_str(summary);
688 }
689
690 output.push_str(&format!("\n\nTool calls: {}", records.len()));
691 for (index, record) in records.iter().enumerate() {
692 output.push_str(&format!(
693 "\n{}. {} ({}, exit_code={}, output_bytes={})",
694 index + 1,
695 record.tool_name,
696 if record.success { "ok" } else { "failed" },
697 record.exit_code,
698 record.output_bytes
699 ));
700 }
701
702 output.push_str("\n\nResult:\n");
703 output.push_str(&serde_json::to_string_pretty(result).unwrap_or_else(|_| result.to_string()));
704
705 if !stderr.is_empty() {
706 output.push_str("\n\nstderr:\n");
707 output.push_str(stderr);
708 }
709
710 output
711}
712
713#[cfg(test)]
714#[path = "program_tool/tests.rs"]
715mod tests;