diff --git a/crates/goose/src/agents/code_execution_extension.rs b/crates/goose/src/agents/code_execution_extension.rs index 2d66599201ff..51c5e2658b0f 100644 --- a/crates/goose/src/agents/code_execution_extension.rs +++ b/crates/goose/src/agents/code_execution_extension.rs @@ -5,7 +5,6 @@ use anyhow::Result; use async_trait::async_trait; use boa_engine::builtins::promise::PromiseState; use boa_engine::module::{MapModuleLoader, Module, SyntheticModuleInitializer}; -use boa_engine::property::Attribute; use boa_engine::{js_string, Context, JsNativeError, JsString, JsValue, NativeFunction, Source}; use indoc::indoc; use regex::Regex; @@ -266,6 +265,8 @@ impl ToolInfo { thread_local! { static CALL_TX: std::cell::RefCell>> = const { std::cell::RefCell::new(None) }; + static RESULT_CELL: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; } fn create_server_module( @@ -360,6 +361,7 @@ fn run_js_module( call_tx: mpsc::UnboundedSender, ) -> Result { CALL_TX.with(|tx| *tx.borrow_mut() = Some(call_tx)); + RESULT_CELL.with(|cell| *cell.borrow_mut() = None); let loader = Rc::new(MapModuleLoader::new()); let mut ctx = Context::builder() @@ -367,12 +369,14 @@ fn run_js_module( .build() .map_err(|e| format!("Failed to create JS context: {e}"))?; - ctx.register_global_property( - js_string!("__result__"), - JsValue::undefined(), - Attribute::WRITABLE, - ) - .map_err(|e| format!("Failed to register __result__: {e}"))?; + let record_result = NativeFunction::from_copy_closure(|_this, args, _ctx| { + let value = args.first().cloned().unwrap_or(JsValue::undefined()); + RESULT_CELL.with(|cell| *cell.borrow_mut() = Some(value.display().to_string())); + Ok(value) + }); + + ctx.register_global_callable(js_string!("record_result"), 1, record_result) + .map_err(|e| format!("Failed to register record_result: {e}"))?; let mut by_server: BTreeMap<&str, Vec<&ToolInfo>> = BTreeMap::new(); for tool in tools { @@ -384,35 +388,7 @@ fn run_js_module( loader.insert(*server_name, module); } - let wrapped = { - let lines: Vec<&str> = code.trim().lines().collect(); - let last_idx = lines - .iter() - .rposition(|l| !l.trim().is_empty() && !l.trim().starts_with("//")) - .unwrap_or(0); - let last = lines.get(last_idx).map(|s| s.trim()).unwrap_or(""); - - const NO_WRAP: &[&str] = &["import ", "export ", "function ", "class "]; - if last.contains("__result__") || NO_WRAP.iter().any(|p| last.starts_with(p)) { - code.to_string() - } else { - let before = lines[..last_idx].join("\n"); - let mut result = None; - for decl in ["const ", "let ", "var "] { - if let Some(rest) = last.strip_prefix(decl) { - if let Some(name) = rest.split('=').next().map(str::trim) { - result = Some(format!("{before}\n{last}\n__result__ = {name};")); - } - break; - } - } - result.unwrap_or_else(|| { - format!("{before}\n__result__ = {};", last.trim_end_matches(';')) - }) - } - }; - - let user_module = Module::parse(Source::from_bytes(&wrapped), None, &mut ctx) + let user_module = Module::parse(Source::from_bytes(code), None, &mut ctx) .map_err(|e| format!("Parse error: {e}"))?; loader.insert("__main__", user_module.clone()); @@ -422,11 +398,8 @@ fn run_js_module( match promise.state() { PromiseState::Fulfilled(_) => { - let result = ctx - .global_object() - .get(js_string!("__result__"), &mut ctx) - .map_err(|e| format!("Failed to get result: {e}"))?; - Ok(result.display().to_string()) + let result = RESULT_CELL.with(|cell| cell.borrow().clone()); + Ok(result.unwrap_or_else(|| "undefined".to_string())) } PromiseState::Rejected(err) => Err(format!("Module error: {}", err.display())), PromiseState::Pending => Err("Module evaluation did not complete".to_string()), @@ -758,6 +731,7 @@ impl McpClientTrait for CodeExecutionClient { import { text_editor } from "developer"; const content = text_editor({ path: "/path/to/source.md", command: "view" }); text_editor({ path: "/path/to/dest.md", command: "write", file_text: content }); + record_result({ copied: true }); ``` EXAMPLE - Multiple operations chained: @@ -766,15 +740,14 @@ impl McpClientTrait for CodeExecutionClient { const files = shell({ command: "ls -la" }); const readme = text_editor({ path: "./README.md", command: "view" }); const status = shell({ command: "git status" }); - { files, readme, status } + record_result({ files, readme, status }); ``` SYNTAX: - Import: import { tool1, tool2 } from "serverName"; - Call: toolName({ param1: value, param2: value }) + - Result: record_result(value) - call this to return a value from the script - All calls are synchronous, return strings - - Last expression is the result - - No comments in code TOOL_GRAPH: Always provide tool_graph to describe the execution flow for the UI. Each node has: tool (server/name), description (what it does), depends_on (indices of dependencies). @@ -942,7 +915,10 @@ mod tests { let client = CodeExecutionClient::new(context).unwrap(); let mut args = JsonObject::new(); - args.insert("code".to_string(), Value::String("2 + 2".to_string())); + args.insert( + "code".to_string(), + Value::String("record_result(2 + 2)".to_string()), + ); let result = client .call_tool("execute_code", Some(args), CancellationToken::new()) diff --git a/crates/goose/tests/test_data/openai_builtin_execute.txt b/crates/goose/tests/test_data/openai_builtin_execute.txt index 56d27e2ca7d3..49d8b4992cb1 100644 --- a/crates/goose/tests/test_data/openai_builtin_execute.txt +++ b/crates/goose/tests/test_data/openai_builtin_execute.txt @@ -128,9 +128,9 @@ data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.c data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"n"}}]},"finish_reason":null}],"usage":null,"obfuscation":"FcZzRbewbWzs"} -data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"write"}}]},"finish_reason":null}],"usage":null,"obfuscation":"yJ0heSJY"} +data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"record_result"}}]},"finish_reason":null}],"usage":null,"obfuscation":"yJ0heSJY"} -data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"Res"}}]},"finish_reason":null}],"usage":null,"obfuscation":"VCJyV87lb5"} +data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"(writeRes)"}}]},"finish_reason":null}],"usage":null,"obfuscation":"VCJyV87lb5"} data: {"id":"chatcmpl-Cqnxk0p34ZkYNestA1fcahFGuLiRR","object":"chat.completion.chunk","created":1766701148,"model":"gpt-5-nano-2025-08-07","service_tier":"default","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\n"}}]},"finish_reason":null}],"usage":null,"obfuscation":"LniNHGfSN9"}