fix: clean up result recording for code mode (#6343)

This commit is contained in:
Alex Hancock
2026-01-13 14:35:00 -05:00
committed by GitHub
parent d9a30a8f99
commit c222bdeab3
2 changed files with 23 additions and 47 deletions
@@ -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<Option<mpsc::UnboundedSender<ToolCallRequest>>> =
const { std::cell::RefCell::new(None) };
static RESULT_CELL: std::cell::RefCell<Option<String>> =
const { std::cell::RefCell::new(None) };
}
fn create_server_module(
@@ -360,6 +361,7 @@ fn run_js_module(
call_tx: mpsc::UnboundedSender<ToolCallRequest>,
) -> Result<String, String> {
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()),
@@ -764,6 +737,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:
@@ -772,15 +746,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).
@@ -950,7 +923,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())
@@ -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"}