mod allocator;
mod conversion;
mod memory;
use super::limits::{CodeModeLimits, budget_snapshot};
use crate::cancellation::AgentCancellation;
pub(crate) use conversion::{serialized_size, value_reservation};
use crossbeam_channel::{Receiver, Sender, bounded};
pub(crate) use memory::MemoryBudget;
use rquickjs::{Context, Ctx, Exception, Function, Runtime, Value as JsValue, function::MutFn};
use serde_json::Value;
use std::{
cell::{Cell, RefCell},
rc::Rc,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::{Duration, Instant},
};
#[derive(Debug, thiserror::Error)]
#[error("Code Mode {0} limit exceeded")]
pub(crate) struct ResourceLimitExceeded(pub(crate) &'static str);
impl ResourceLimitExceeded {
pub(crate) fn recovery(&self) -> &'static str {
match self.0 {
"session persisted history bytes" => {
"Compact the session before another Code Mode call. A smaller batch cannot restore session history space. Inspect execution evidence and re-read affected files; completed effects remain applied."
}
"batch persisted history bytes" | "inner calls" => {
"Continue remaining work in a new, smaller Code Mode call. Use tools.budget() to stop at a safe boundary and return progress before exhaustion. Inspect execution evidence and re-read affected files; do not replay completed actions."
}
"source" => "Shorten the script or split it into separate Code Mode calls.",
"return bytes" => {
"Return a smaller summary instead of full tool outputs. Inspect execution evidence before retrying; completed effects remain applied."
}
"JavaScript execution time" => {
"Split computation into smaller Code Mode calls. Host-tool waiting does not spend JavaScript time. Inspect execution evidence before continuing; completed effects remain applied."
}
_ => {
"Read narrower file ranges, retain less intermediate data, or split work into smaller Code Mode calls. Use tools.budget() before exhaustion. Inspect execution evidence and re-read affected files; completed effects remain applied."
}
}
}
}
pub(crate) struct SdkRequest {
pub(crate) operation: String,
pub(crate) name: String,
pub(crate) arguments: Value,
pub(crate) memory: Arc<MemoryBudget>,
}
pub(crate) trait CodeModeEngine {
fn execute(
&self,
source: &str,
limits: &CodeModeLimits,
cancellation: &AgentCancellation,
host: &mut dyn FnMut(SdkRequest) -> anyhow::Result<Value>,
) -> anyhow::Result<Value>;
}
pub(crate) struct QuickJsEngine;
impl CodeModeEngine for QuickJsEngine {
fn execute(
&self,
source: &str,
limits: &CodeModeLimits,
cancellation: &AgentCancellation,
host: &mut dyn FnMut(SdkRequest) -> anyhow::Result<Value>,
) -> anyhow::Result<Value> {
limits.validate()?;
if source.len() > limits.max_source_bytes {
return Err(
anyhow::Error::from(ResourceLimitExceeded("source")).context(format!(
"Script has {} bytes; maximum {}",
source.len(),
limits.max_source_bytes
)),
);
}
cancellation.check()?;
std::thread::scope(|scope| {
let (requests, incoming) = bounded(0);
let (responses, outgoing) = bounded(0);
let worker = std::thread::Builder::new()
.name("code-mode-js".into())
.stack_size(2 * 1024 * 1024)
.spawn_scoped(scope, || {
run_javascript(source, limits, cancellation.clone(), requests, outgoing)
})?;
let mut fatal = None;
while let Ok(request) = incoming.recv() {
let response = if fatal.is_some() {
Err("Code Mode host execution terminated".to_owned())
} else {
host(request).map_err(|error| {
fatal = Some(error);
"Code Mode host execution terminated".to_owned()
})
};
if responses.send(response).is_err() {
break;
}
}
let result = worker
.join()
.map_err(|_| anyhow::anyhow!("Code Mode engine terminated"))?;
if let Some(error) = fatal {
return Err(error);
}
cancellation.check()?;
result
})
}
}
fn run_javascript(
source: &str,
limits: &CodeModeLimits,
cancellation: AgentCancellation,
requests: Sender<SdkRequest>,
responses: Receiver<Result<Value, String>>,
) -> anyhow::Result<Value> {
let execution_deadline = Rc::new(Cell::new(
Instant::now() + Duration::from_millis(limits.max_runtime_ms),
));
let stopped = Arc::new(AtomicBool::new(false));
let termination_error = Rc::new(RefCell::new(None));
let memory = memory::MemoryBudget::new(limits.max_memory_bytes, stopped.clone());
let _source_buffer = memory.reserve(source.len().saturating_mul(2).saturating_add(4096))?;
let runtime = Runtime::new_with_alloc(allocator::BoundedAllocator(memory.clone()))?;
runtime.set_max_stack_size(512 * 1024);
let interrupt_stopped = stopped.clone();
let interrupt_cancellation = cancellation.clone();
let interrupt_deadline = execution_deadline.clone();
runtime.set_interrupt_handler(Some(Box::new(move || {
if interrupt_cancellation.is_canceled() || Instant::now() >= interrupt_deadline.get() {
interrupt_stopped.store(true, Ordering::SeqCst);
}
interrupt_stopped.load(Ordering::SeqCst)
})));
let context = Context::full(&runtime)?;
let limits = limits.clone();
let callback_stopped = stopped.clone();
let result = context.with(|ctx| -> anyhow::Result<Value> {
let own_property: Function = ctx.eval("Object.getOwnPropertyDescriptor")?;
let result = (|| -> rquickjs::Result<Value> {
let mut count = 0usize;
let mut total_bytes = 0usize;
let mut response_reservation = None;
let callback_memory = memory.clone();
let callback_cancellation = cancellation.clone();
let callback_limits = limits.clone();
let callback_deadline = execution_deadline.clone();
let callback_error = termination_error.clone();
let bridge = Function::new(
ctx.clone(),
MutFn::new(
move |ctx: Ctx<'_>,
operation: rquickjs::String<'_>,
name: rquickjs::String<'_>,
args: rquickjs::String<'_>|
-> rquickjs::Result<String> {
let execute = || -> anyhow::Result<String> {
let cancellation = &callback_cancellation;
cancellation.check()?;
if callback_stopped.load(Ordering::SeqCst) {
return Err(ResourceLimitExceeded(
if Instant::now() >= callback_deadline.get() {
"JavaScript execution time"
} else {
"memory"
},
)
.into());
}
response_reservation.take();
let operation = operation.to_cstring()?;
let name = name.to_cstring()?;
let args = args.to_cstring()?;
let operation = checked_utf8(&operation)?;
let name = checked_utf8(&name)?;
let args = checked_utf8(&args)?;
let argument_charge =
conversion::parse_reservation(args, cancellation)?
.saturating_add(name.len())
.saturating_add(operation.len());
let _arguments = callback_memory.reserve(argument_charge)?;
if operation == "call" {
if count >= callback_limits.max_inner_calls {
return Err(anyhow::Error::from(ResourceLimitExceeded(
"inner calls",
))
.context(format!(
"Rejected tool call: {count} calls used; maximum {}",
callback_limits.max_inner_calls
)));
}
count += 1;
}
let request = SdkRequest {
operation: operation.to_owned(),
name: name.to_owned(),
arguments: serde_json::from_str(args)?,
memory: callback_memory.clone(),
};
if Instant::now() >= callback_deadline.get() {
return Err(
ResourceLimitExceeded("JavaScript execution time").into()
);
}
let waiting = Instant::now();
requests.send(request)?;
let response = responses.recv();
callback_deadline.set(callback_deadline.get() + waiting.elapsed());
let mut value = response?.map_err(anyhow::Error::msg)?;
cancellation.check()?;
if operation == "budget" {
let remaining_ms = callback_deadline
.get()
.saturating_duration_since(Instant::now())
.as_millis()
as u64;
value["inner_calls"] = budget_snapshot(
count as u64,
callback_limits.max_inner_calls as u64,
);
value["result_bytes"] = budget_snapshot(
total_bytes as u64,
callback_limits.max_total_result_bytes as u64,
);
value["runtime_ms"] = budget_snapshot(
callback_limits.max_runtime_ms.saturating_sub(remaining_ms),
callback_limits.max_runtime_ms,
);
value["memory_bytes"] = callback_memory.snapshot();
value["inner_result_bytes_limit"] =
(callback_limits.max_inner_result_bytes as u64).into();
value["return_bytes_limit"] =
(callback_limits.max_return_bytes as u64).into();
}
let _native_result = callback_memory
.reserve(conversion::value_reservation(&value, cancellation)?)?;
let remaining_bytes = callback_limits
.max_total_result_bytes
.saturating_sub(total_bytes);
let limit = if operation == "budget" {
4096
} else {
callback_limits.max_inner_result_bytes.min(remaining_bytes)
};
let size = conversion::serialized_size(&value, limit, cancellation)
.map_err(|error| {
if error.downcast_ref::<ResourceLimitExceeded>().is_some() {
anyhow::Error::from(ResourceLimitExceeded(
if remaining_bytes
< callback_limits.max_inner_result_bytes
{
"total result bytes"
} else {
"inner result bytes"
},
))
.context(format!("SDK result exceeded {limit} bytes"))
} else {
error
}
})?;
response_reservation = Some(callback_memory.reserve(size)?);
let mut bytes = Vec::with_capacity(size);
serde_json::to_writer(&mut bytes, &value)?;
let text = String::from_utf8(bytes)?;
if operation != "budget" {
total_bytes = total_bytes.saturating_add(size);
}
Ok(text)
};
match execute() {
Ok(text) => Ok(text),
Err(error) => {
let message = format!("{error:#}");
if callback_error.borrow().is_none() {
*callback_error.borrow_mut() = Some(error);
}
callback_stopped.store(true, Ordering::SeqCst);
Err(Exception::throw_message(&ctx, &message))
}
}
},
),
)?;
ctx.globals().set("__bridge", bridge)?;
ctx.eval::<(), _>(
r#"
(() => {
const nativeBridge = globalThis.__bridge;
const stringify = JSON.stringify, parse = JSON.parse;
const bridge = (...args) => parse(nativeBridge(...args));
delete globalThis.__bridge;
Object.defineProperty(globalThis, 'tools', {value: Object.freeze({
list: () => bridge('list', '', '{}'),
search: query => bridge('search', String(query), '{}'),
describe: name => bridge('describe', String(name), '{}'),
budget: () => bridge('budget', '', '{}'),
call: (name, args) => bridge('call', String(name), stringify(args))
}), writable: false, configurable: false});
})();
"#,
)?;
let value: JsValue =
ctx.eval(format!("(function() {{ 'use strict';\n{source}\n}})()"))?;
if value.is_promise() {
return Err(Exception::throw_message(
&ctx,
"Code Mode requires synchronous JavaScript",
));
}
let text = ctx
.json_stringify(value)?
.map(|text| text.to_cstring())
.transpose()?;
let Some(text) = text else {
return Ok(Value::Null);
};
if text.len() > limits.max_return_bytes {
if termination_error.borrow().is_none() {
*termination_error.borrow_mut() = Some(
anyhow::Error::from(ResourceLimitExceeded("return bytes")).context(
format!(
"Return has {} bytes; maximum {}",
text.len(),
limits.max_return_bytes
),
),
);
}
stopped.store(true, Ordering::SeqCst);
return Err(Exception::throw_message(
&ctx,
"Code Mode return limit exceeded",
));
}
let text = checked_utf8(&text).map_err(|_| {
Exception::throw_message(&ctx, "Code Mode return must be UTF-8 JSON")
})?;
let charge = conversion::parse_reservation(text, &cancellation)
.and_then(|bytes| memory.reserve(bytes))
.map_err(|_| {
Exception::throw_message(&ctx, "Code Mode return conversion limit exceeded")
})?;
let value = serde_json::from_str(text)
.map_err(|_| Exception::throw_message(&ctx, "Code Mode return must be JSON"));
drop(charge);
value
})();
result.map_err(|error| {
let detail = (error.is_exception() && !stopped.load(Ordering::SeqCst))
.then(|| javascript_exception_message(&ctx, &own_property, &memory))
.flatten()
.unwrap_or_else(|| error.to_string());
anyhow::anyhow!(
"Code Mode JavaScript failed: {detail}; completed effects remain applied"
)
})
});
cancellation.check()?;
if let Some(error) = termination_error.borrow_mut().take() {
return Err(error);
}
if Instant::now() >= execution_deadline.get() {
return Err(ResourceLimitExceeded("JavaScript execution time").into());
}
if stopped.load(Ordering::SeqCst) {
return Err(ResourceLimitExceeded("memory").into());
}
result
}
fn javascript_exception_message<'js>(
ctx: &Ctx<'js>,
own_property: &Function<'js>,
memory: &Arc<MemoryBudget>,
) -> Option<String> {
let exception = ctx.catch();
let location = exception
.is_error()
.then(|| user_script_error_location(own_property, &exception))
.flatten();
let message = if exception.is_error() {
own_data_property_string(own_property, &exception, "message")?
} else {
exception.into_string()?
};
let text = message.to_cstring().ok()?;
let _redaction = memory.reserve(text.len().saturating_mul(4)).ok()?;
let text = crate::output::redact_sensitive_text(checked_utf8(&text).ok()?);
let text: String = text.chars().take(2048).collect();
Some(match location {
Some((line, column)) => format!("{text} (line {line}, column {column})"),
None => text,
})
}
fn own_data_property_string<'js>(
own_property: &Function<'js>,
object: &JsValue<'js>,
name: &str,
) -> Option<rquickjs::String<'js>> {
let descriptor: JsValue = own_property.call((object.clone(), name)).ok()?;
if !descriptor.is_object() {
return None;
}
let value: JsValue = own_property.call((descriptor, "value")).ok()?;
value.as_object()?.get::<_, rquickjs::String>("value").ok()
}
fn user_script_error_location<'js>(
own_property: &Function<'js>,
exception: &JsValue<'js>,
) -> Option<(usize, usize)> {
let stack = own_data_property_string(own_property, exception, "stack")?
.to_string()
.ok()?;
let (_, frame) = stack.split_once("eval_script:")?;
let mut parts = frame.splitn(3, ':');
let line: usize = parts.next()?.parse().ok()?;
let column: usize = parts
.next()?
.trim_end_matches(|character: char| !character.is_ascii_digit())
.parse()
.ok()?;
Some((line.checked_sub(1).filter(|line| *line > 0)?, column))
}
fn checked_utf8<'a>(text: &'a rquickjs::CString<'_>) -> anyhow::Result<&'a str> {
let bytes = unsafe { std::slice::from_raw_parts(text.as_ptr().cast::<u8>(), text.len()) };
Ok(std::str::from_utf8(bytes)?)
}