clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-05-02 02:17:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Code mode on v8 (#15276)

Browse Source 
Moves Code Mode to a new crate with no dependencies on codex. This
create encodes the code mode semantics that we want for lifetime,
mounting, tool calling.

The model-facing surface is mostly unchanged. `exec` still runs raw
JavaScript, `wait` still resumes or terminates a `cell_id`, nested tools
are still available through `tools.*`, and helpers like `text`, `image`,
`store`, `load`, `notify`, `yield_control`, and `exit` still exist.

The major change is underneath that surface:

- Old code mode was an external Node runtime.
- New code mode is an in-process V8 runtime embedded directly in Rust.
- Old code mode managed cells inside a long-lived Node runner process.
- New code mode manages cells in Rust, with one V8 runtime thread per
active `exec`.
- Old code mode used JSON protocol messages over child stdin/stdout plus
Node worker-thread messages.
- New code mode uses Rust channels and direct V8 callbacks/events.

This PR also fixes the two migration regressions that fell out of that
substrate change:

- `wait { terminate: true }` now waits for the V8 runtime to actually
stop before reporting termination.
- synchronous top-level `exit()` now succeeds again instead of surfacing
as a script error.

---

- `core/src/tools/code_mode/*` is now mostly an adapter layer for the
public `exec` / `wait` tools.
- `code-mode/src/service.rs` owns cell sessions and async control flow
in Rust.
- `code-mode/src/runtime/*.rs` owns the embedded V8 isolate and
JavaScript execution.
- each `exec` spawns a dedicated runtime thread plus a Rust
session-control task.
- helper globals are installed directly into the V8 context instead of
being injected through a source prelude.
- helper modules like `tools.js` and `@openai/code_mode` are synthesized
through V8 module resolution callbacks in Rust.

---

Also added a benchmark for showing the speed of init and use of a code
mode env:
```
$ cargo bench -p codex-code-mode --bench exec_overhead -- --samples 30 --warm-iterations 25 --tool-counts 0,32,128
Finished [`bench` profile [optimized]](https://doc.rust-lang.org/cargo/reference/profiles.html#default-profiles) target(s) in 0.18s
     Running benches/exec_overhead.rs (target/release/deps/exec_overhead-008c440d800545ae)
exec_overhead: samples=30, warm_iterations=25, tool_counts=[0, 32, 128]
scenario       tools samples    warmups      iters      mean/exec       p95/exec       rssΔ p50       rssΔ max
cold_exec          0      30          0          1         1.13ms         1.20ms        8.05MiB        8.06MiB
warm_exec          0      30          1         25       473.43us       512.49us      912.00KiB        1.33MiB
cold_exec         32      30          0          1         1.03ms         1.15ms        8.08MiB        8.11MiB
warm_exec         32      30          1         25       509.73us       545.76us      960.00KiB        1.30MiB
cold_exec        128      30          0          1         1.14ms         1.19ms        8.30MiB        8.34MiB
warm_exec        128      30          1         25       575.08us       591.03us      736.00KiB      864.00KiB
memory uses a fresh-process max RSS delta for each scenario
```

---------

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Channing Conger
2026-03-20 23:36:58 -07:00
committed by GitHub
parent ec32866c37
commit e4eedd6170
31 changed files with 2730 additions and 2265 deletions
Download Patch File Download Diff File Expand all files Collapse all files

555
codex-rs/code-mode/src/description.rs Normal file
View File

@@ -0,0 +1,555 @@
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value as JsonValue;
use crate::PUBLIC_TOOL_NAME;
const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1;
const CODE_MODE_ONLY_PREFACE: &str =
"Use `exec/wait` tool to run all other tools, do not attempt to use any other tools directly";
const EXEC_DESCRIPTION_TEMPLATE: &str = r#"## exec
- Runs raw JavaScript in an isolated context (no Node, no file system, or network access, no console).
- Send raw JavaScript source text, not JSON, quoted strings, or markdown code fences.
- You may optionally start the tool input with a first-line pragma like `// @exec: {"yield_time_ms": 10000, "max_output_tokens": 1000}`.
- `yield_time_ms` asks `exec` to yield early after that many milliseconds if the script is still running.
- `max_output_tokens` sets the token budget for direct `exec` results. By default the result is truncated to 10000 tokens.
- All nested tools are available on the global `tools` object, for example `await tools.exec_command(...)`. Tool names are exposed as normalized JavaScript identifiers, for example `await tools.mcp__ologs__get_profile(...)`.
- Tool methods take either string or object as parameter.
- They return either a structured value or a string based on the description above.
- Global helpers:
- `exit()`: Immediately ends the current script successfully (like an early return from the top level).
- `text(value: string | number | boolean | undefined | null)`: Appends a text item. Non-string values are stringified with `JSON.stringify(...)` when possible.
- `image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null })`: Appends an image item. `image_url` can be an HTTPS URL or a base64-encoded `data:` URL.
- `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session.
- `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing.
- `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`.
- `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries.
- `yield_control()`: yields the accumulated output to the model immediately while the script keeps running."#;
const WAIT_DESCRIPTION_TEMPLATE: &str = r#"- Use `wait` only after `exec` returns `Script running with cell ID ...`.
- `cell_id` identifies the running `exec` cell to resume.
- `yield_time_ms` controls how long to wait for more output before yielding again. If omitted, `wait` uses its default wait timeout.
- `max_tokens` limits how much new output this wait call returns.
- `terminate: true` stops the running cell instead of waiting for more output.
- `wait` returns only the new output since the last yield, or the final completion or termination result for that cell.
- If the cell is still running, `wait` may yield again with the same `cell_id`.
- If the cell has already finished, `wait` returns the completed result and closes the cell."#;
pub const CODE_MODE_PRAGMA_PREFIX: &str = "// @exec:";
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CodeModeToolKind {
Function,
Freeform,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub kind: CodeModeToolKind,
pub input_schema: Option<JsonValue>,
pub output_schema: Option<JsonValue>,
}
#[derive(Debug, Default, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct CodeModeExecPragma {
#[serde(default)]
yield_time_ms: Option<u64>,
#[serde(default)]
max_output_tokens: Option<usize>,
}
#[derive(Debug, PartialEq, Eq)]
pub struct ParsedExecSource {
pub code: String,
pub yield_time_ms: Option<u64>,
pub max_output_tokens: Option<usize>,
}
pub fn parse_exec_source(input: &str) -> Result<ParsedExecSource, String> {
if input.trim().is_empty() {
return Err(
"exec expects raw JavaScript source text (non-empty). Provide JS only, optionally with first-line `// @exec: {\"yield_time_ms\": 10000, \"max_output_tokens\": 1000}`.".to_string(),
);
}
let mut args = ParsedExecSource {
code: input.to_string(),
yield_time_ms: None,
max_output_tokens: None,
};
let mut lines = input.splitn(2, '\n');
let first_line = lines.next().unwrap_or_default();
let rest = lines.next().unwrap_or_default();
let trimmed = first_line.trim_start();
let Some(pragma) = trimmed.strip_prefix(CODE_MODE_PRAGMA_PREFIX) else {
return Ok(args);
};
if rest.trim().is_empty() {
return Err(
"exec pragma must be followed by JavaScript source on subsequent lines".to_string(),
);
}
let directive = pragma.trim();
if directive.is_empty() {
return Err(
"exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
.to_string(),
);
}
let value: serde_json::Value = serde_json::from_str(directive).map_err(|err| {
format!(
"exec pragma must be valid JSON with supported fields `yield_time_ms` and `max_output_tokens`: {err}"
)
})?;
let object = value.as_object().ok_or_else(|| {
"exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
.to_string()
})?;
for key in object.keys() {
match key.as_str() {
"yield_time_ms" | "max_output_tokens" => {}
_ => {
return Err(format!(
"exec pragma only supports `yield_time_ms` and `max_output_tokens`; got `{key}`"
));
}
}
}
let pragma: CodeModeExecPragma = serde_json::from_value(value).map_err(|err| {
format!(
"exec pragma fields `yield_time_ms` and `max_output_tokens` must be non-negative safe integers: {err}"
)
})?;
if pragma
.yield_time_ms
.is_some_and(|yield_time_ms| yield_time_ms > MAX_JS_SAFE_INTEGER)
{
return Err(
"exec pragma field `yield_time_ms` must be a non-negative safe integer".to_string(),
);
}
if pragma.max_output_tokens.is_some_and(|max_output_tokens| {
u64::try_from(max_output_tokens)
.map(|max_output_tokens| max_output_tokens > MAX_JS_SAFE_INTEGER)
.unwrap_or(true)
}) {
return Err(
"exec pragma field `max_output_tokens` must be a non-negative safe integer".to_string(),
);
}
args.code = rest.to_string();
args.yield_time_ms = pragma.yield_time_ms;
args.max_output_tokens = pragma.max_output_tokens;
Ok(args)
}
pub fn is_code_mode_nested_tool(tool_name: &str) -> bool {
tool_name != crate::PUBLIC_TOOL_NAME && tool_name != crate::WAIT_TOOL_NAME
}
pub fn build_exec_tool_description(
enabled_tools: &[(String, String)],
code_mode_only: bool,
) -> String {
if !code_mode_only {
return EXEC_DESCRIPTION_TEMPLATE.to_string();
}
let mut sections = vec![
CODE_MODE_ONLY_PREFACE.to_string(),
EXEC_DESCRIPTION_TEMPLATE.to_string(),
];
if !enabled_tools.is_empty() {
let nested_tool_reference = enabled_tools
.iter()
.map(|(name, nested_description)| {
let global_name = normalize_code_mode_identifier(name);
format!(
"### `{global_name}` (`{name}`)\n{}",
nested_description.trim()
)
})
.collect::<Vec<_>>()
.join("\n\n");
sections.push(nested_tool_reference);
}
sections.join("\n\n")
}
pub fn build_wait_tool_description() -> &'static str {
WAIT_DESCRIPTION_TEMPLATE
}
pub fn normalize_code_mode_identifier(tool_key: &str) -> String {
let mut identifier = String::new();
for (index, ch) in tool_key.chars().enumerate() {
let is_valid = if index == 0 {
ch == '_' || ch == '$' || ch.is_ascii_alphabetic()
} else {
ch == '_' || ch == '$' || ch.is_ascii_alphanumeric()
};
if is_valid {
identifier.push(ch);
} else {
identifier.push('_');
}
}
if identifier.is_empty() {
"_".to_string()
} else {
identifier
}
}
pub fn augment_tool_definition(mut definition: ToolDefinition) -> ToolDefinition {
if definition.name != PUBLIC_TOOL_NAME {
definition.description = append_code_mode_sample_for_definition(&definition);
}
definition
}
pub fn enabled_tool_metadata(definition: &ToolDefinition) -> EnabledToolMetadata {
EnabledToolMetadata {
tool_name: definition.name.clone(),
global_name: normalize_code_mode_identifier(&definition.name),
description: definition.description.clone(),
kind: definition.kind,
}
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct EnabledToolMetadata {
pub tool_name: String,
pub global_name: String,
pub description: String,
pub kind: CodeModeToolKind,
}
pub fn append_code_mode_sample(
description: &str,
tool_name: &str,
input_name: &str,
input_type: String,
output_type: String,
) -> String {
let declaration = format!(
"declare const tools: {{ {} }};",
render_code_mode_tool_declaration(tool_name, input_name, input_type, output_type)
);
format!("{description}\n\nexec tool declaration:\n```ts\n{declaration}\n```")
}
fn append_code_mode_sample_for_definition(definition: &ToolDefinition) -> String {
let input_name = match definition.kind {
CodeModeToolKind::Function => "args",
CodeModeToolKind::Freeform => "input",
};
let input_type = match definition.kind {
CodeModeToolKind::Function => definition
.input_schema
.as_ref()
.map(render_json_schema_to_typescript)
.unwrap_or_else(|| "unknown".to_string()),
CodeModeToolKind::Freeform => "string".to_string(),
};
let output_type = definition
.output_schema
.as_ref()
.map(render_json_schema_to_typescript)
.unwrap_or_else(|| "unknown".to_string());
append_code_mode_sample(
&definition.description,
&definition.name,
input_name,
input_type,
output_type,
)
}
fn render_code_mode_tool_declaration(
tool_name: &str,
input_name: &str,
input_type: String,
output_type: String,
) -> String {
let tool_name = normalize_code_mode_identifier(tool_name);
format!("{tool_name}({input_name}: {input_type}): Promise<{output_type}>;")
}
pub fn render_json_schema_to_typescript(schema: &JsonValue) -> String {
render_json_schema_to_typescript_inner(schema)
}
fn render_json_schema_to_typescript_inner(schema: &JsonValue) -> String {
match schema {
JsonValue::Bool(true) => "unknown".to_string(),
JsonValue::Bool(false) => "never".to_string(),
JsonValue::Object(map) => {
if let Some(value) = map.get("const") {
return render_json_schema_literal(value);
}
if let Some(values) = map.get("enum").and_then(JsonValue::as_array) {
let rendered = values
.iter()
.map(render_json_schema_literal)
.collect::<Vec<_>>();
if !rendered.is_empty() {
return rendered.join(" | ");
}
}
for key in ["anyOf", "oneOf"] {
if let Some(variants) = map.get(key).and_then(JsonValue::as_array) {
let rendered = variants
.iter()
.map(render_json_schema_to_typescript_inner)
.collect::<Vec<_>>();
if !rendered.is_empty() {
return rendered.join(" | ");
}
}
}
if let Some(variants) = map.get("allOf").and_then(JsonValue::as_array) {
let rendered = variants
.iter()
.map(render_json_schema_to_typescript_inner)
.collect::<Vec<_>>();
if !rendered.is_empty() {
return rendered.join(" & ");
}
}
if let Some(schema_type) = map.get("type") {
if let Some(types) = schema_type.as_array() {
let rendered = types
.iter()
.filter_map(JsonValue::as_str)
.map(|schema_type| render_json_schema_type_keyword(map, schema_type))
.collect::<Vec<_>>();
if !rendered.is_empty() {
return rendered.join(" | ");
}
}
if let Some(schema_type) = schema_type.as_str() {
return render_json_schema_type_keyword(map, schema_type);
}
}
if map.contains_key("properties")
|| map.contains_key("additionalProperties")
|| map.contains_key("required")
{
return render_json_schema_object(map);
}
if map.contains_key("items") || map.contains_key("prefixItems") {
return render_json_schema_array(map);
}
"unknown".to_string()
}
_ => "unknown".to_string(),
}
}
fn render_json_schema_type_keyword(
map: &serde_json::Map<String, JsonValue>,
schema_type: &str,
) -> String {
match schema_type {
"string" => "string".to_string(),
"number" | "integer" => "number".to_string(),
"boolean" => "boolean".to_string(),
"null" => "null".to_string(),
"array" => render_json_schema_array(map),
"object" => render_json_schema_object(map),
_ => "unknown".to_string(),
}
}
fn render_json_schema_array(map: &serde_json::Map<String, JsonValue>) -> String {
if let Some(items) = map.get("items") {
let item_type = render_json_schema_to_typescript_inner(items);
return format!("Array<{item_type}>");
}
if let Some(items) = map.get("prefixItems").and_then(JsonValue::as_array) {
let item_types = items
.iter()
.map(render_json_schema_to_typescript_inner)
.collect::<Vec<_>>();
if !item_types.is_empty() {
return format!("[{}]", item_types.join(", "));
}
}
"unknown[]".to_string()
}
fn render_json_schema_object(map: &serde_json::Map<String, JsonValue>) -> String {
let required = map
.get("required")
.and_then(JsonValue::as_array)
.map(|items| {
items
.iter()
.filter_map(JsonValue::as_str)
.collect::<Vec<_>>()
})
.unwrap_or_default();
let properties = map
.get("properties")
.and_then(JsonValue::as_object)
.cloned()
.unwrap_or_default();
let mut sorted_properties = properties.iter().collect::<Vec<_>>();
sorted_properties.sort_unstable_by(|(name_a, _), (name_b, _)| name_a.cmp(name_b));
let mut lines = sorted_properties
.into_iter()
.map(|(name, value)| {
let optional = if required.iter().any(|required_name| required_name == name) {
""
} else {
"?"
};
let property_name = render_json_schema_property_name(name);
let property_type = render_json_schema_to_typescript_inner(value);
format!("{property_name}{optional}: {property_type};")
})
.collect::<Vec<_>>();
if let Some(additional_properties) = map.get("additionalProperties") {
let property_type = match additional_properties {
JsonValue::Bool(true) => Some("unknown".to_string()),
JsonValue::Bool(false) => None,
value => Some(render_json_schema_to_typescript_inner(value)),
};
if let Some(property_type) = property_type {
lines.push(format!("[key: string]: {property_type};"));
}
} else if properties.is_empty() {
lines.push("[key: string]: unknown;".to_string());
}
if lines.is_empty() {
return "{}".to_string();
}
format!("{{ {} }}", lines.join(" "))
}
fn render_json_schema_property_name(name: &str) -> String {
if normalize_code_mode_identifier(name) == name {
name.to_string()
} else {
serde_json::to_string(name).unwrap_or_else(|_| format!("\"{}\"", name.replace('"', "\\\"")))
}
}
fn render_json_schema_literal(value: &JsonValue) -> String {
serde_json::to_string(value).unwrap_or_else(|_| "unknown".to_string())
}
#[cfg(test)]
mod tests {
use super::CodeModeToolKind;
use super::ParsedExecSource;
use super::ToolDefinition;
use super::augment_tool_definition;
use super::build_exec_tool_description;
use super::normalize_code_mode_identifier;
use super::parse_exec_source;
use pretty_assertions::assert_eq;
use serde_json::json;
#[test]
fn parse_exec_source_without_pragma() {
assert_eq!(
parse_exec_source("text('hi')").unwrap(),
ParsedExecSource {
code: "text('hi')".to_string(),
yield_time_ms: None,
max_output_tokens: None,
}
);
}
#[test]
fn parse_exec_source_with_pragma() {
assert_eq!(
parse_exec_source("// @exec: {\"yield_time_ms\": 10}\ntext('hi')").unwrap(),
ParsedExecSource {
code: "text('hi')".to_string(),
yield_time_ms: Some(10),
max_output_tokens: None,
}
);
}
#[test]
fn normalize_identifier_rewrites_invalid_characters() {
assert_eq!(
"mcp__ologs__get_profile",
normalize_code_mode_identifier("mcp__ologs__get_profile")
);
assert_eq!(
"hidden_dynamic_tool",
normalize_code_mode_identifier("hidden-dynamic-tool")
);
}
#[test]
fn augment_tool_definition_appends_typed_declaration() {
let definition = ToolDefinition {
name: "hidden_dynamic_tool".to_string(),
description: "Test tool".to_string(),
kind: CodeModeToolKind::Function,
input_schema: Some(json!({
"type": "object",
"properties": { "city": { "type": "string" } },
"required": ["city"],
"additionalProperties": false
})),
output_schema: Some(json!({
"type": "object",
"properties": { "ok": { "type": "boolean" } },
"required": ["ok"]
})),
};
let description = augment_tool_definition(definition).description;
assert!(description.contains("declare const tools"));
assert!(
description.contains(
"hidden_dynamic_tool(args: { city: string; }): Promise<{ ok: boolean; }>;"
)
);
}
#[test]
fn code_mode_only_description_includes_nested_tools() {
let description =
build_exec_tool_description(&[("foo".to_string(), "bar".to_string())], true);
assert!(description.contains("### `foo` (`foo`)"));
}
}

30
codex-rs/code-mode/src/lib.rs Normal file
View File

@@ -0,0 +1,30 @@
mod description;
mod response;
mod runtime;
mod service;
pub use description::CODE_MODE_PRAGMA_PREFIX;
pub use description::CodeModeToolKind;
pub use description::ToolDefinition;
pub use description::append_code_mode_sample;
pub use description::augment_tool_definition;
pub use description::build_exec_tool_description;
pub use description::build_wait_tool_description;
pub use description::is_code_mode_nested_tool;
pub use description::normalize_code_mode_identifier;
pub use description::parse_exec_source;
pub use description::render_json_schema_to_typescript;
pub use response::FunctionCallOutputContentItem;
pub use response::ImageDetail;
pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS;
pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL;
pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS;
pub use runtime::ExecuteRequest;
pub use runtime::RuntimeResponse;
pub use runtime::WaitRequest;
pub use service::CodeModeService;
pub use service::CodeModeTurnHost;
pub use service::CodeModeTurnWorker;
pub const PUBLIC_TOOL_NAME: &str = "exec";
pub const WAIT_TOOL_NAME: &str = "wait";

24
codex-rs/code-mode/src/response.rs Normal file
View File

@@ -0,0 +1,24 @@
use serde::Deserialize;
use serde::Serialize;
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ImageDetail {
Auto,
Low,
High,
Original,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum FunctionCallOutputContentItem {
InputText {
text: String,
},
InputImage {
image_url: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
detail: Option<ImageDetail>,
},
}

209
codex-rs/code-mode/src/runtime/callbacks.rs Normal file
View File

@@ -0,0 +1,209 @@
use crate::response::FunctionCallOutputContentItem;
use super::EXIT_SENTINEL;
use super::RuntimeEvent;
use super::RuntimeState;
use super::value::json_to_v8;
use super::value::normalize_output_image;
use super::value::serialize_output_text;
use super::value::throw_type_error;
use super::value::v8_value_to_json;
pub(super) fn tool_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
mut retval: v8::ReturnValue<v8::Value>,
) {
let tool_name = args.data().to_rust_string_lossy(scope);
let input = if args.length() == 0 {
Ok(None)
} else {
v8_value_to_json(scope, args.get(0))
};
let input = match input {
Ok(input) => input,
Err(error_text) => {
throw_type_error(scope, &error_text);
return;
}
};
let Some(resolver) = v8::PromiseResolver::new(scope) else {
throw_type_error(scope, "failed to create tool promise");
return;
};
let promise = resolver.get_promise(scope);
let resolver = v8::Global::new(scope, resolver);
let Some(state) = scope.get_slot_mut::<RuntimeState>() else {
throw_type_error(scope, "runtime state unavailable");
return;
};
let id = format!("tool-{}", state.next_tool_call_id);
state.next_tool_call_id = state.next_tool_call_id.saturating_add(1);
let event_tx = state.event_tx.clone();
state.pending_tool_calls.insert(id.clone(), resolver);
let _ = event_tx.send(RuntimeEvent::ToolCall {
id,
name: tool_name,
input,
});
retval.set(promise.into());
}
pub(super) fn text_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
mut retval: v8::ReturnValue<v8::Value>,
) {
let value = if args.length() == 0 {
v8::undefined(scope).into()
} else {
args.get(0)
};
let text = match serialize_output_text(scope, value) {
Ok(text) => text,
Err(error_text) => {
throw_type_error(scope, &error_text);
return;
}
};
if let Some(state) = scope.get_slot::<RuntimeState>() {
let _ = state.event_tx.send(RuntimeEvent::ContentItem(
FunctionCallOutputContentItem::InputText { text },
));
}
retval.set(v8::undefined(scope).into());
}
pub(super) fn image_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
mut retval: v8::ReturnValue<v8::Value>,
) {
let value = if args.length() == 0 {
v8::undefined(scope).into()
} else {
args.get(0)
};
let image_item = match normalize_output_image(scope, value) {
Ok(image_item) => image_item,
Err(()) => return,
};
if let Some(state) = scope.get_slot::<RuntimeState>() {
let _ = state.event_tx.send(RuntimeEvent::ContentItem(image_item));
}
retval.set(v8::undefined(scope).into());
}
pub(super) fn store_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
_retval: v8::ReturnValue<v8::Value>,
) {
let key = match args.get(0).to_string(scope) {
Some(key) => key.to_rust_string_lossy(scope),
None => {
throw_type_error(scope, "store key must be a string");
return;
}
};
let value = args.get(1);
let serialized = match v8_value_to_json(scope, value) {
Ok(Some(value)) => value,
Ok(None) => {
throw_type_error(
scope,
&format!("Unable to store {key:?}. Only plain serializable objects can be stored."),
);
return;
}
Err(error_text) => {
throw_type_error(scope, &error_text);
return;
}
};
if let Some(state) = scope.get_slot_mut::<RuntimeState>() {
state.stored_values.insert(key, serialized);
}
}
pub(super) fn load_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
mut retval: v8::ReturnValue<v8::Value>,
) {
let key = match args.get(0).to_string(scope) {
Some(key) => key.to_rust_string_lossy(scope),
None => {
throw_type_error(scope, "load key must be a string");
return;
}
};
let value = scope
.get_slot::<RuntimeState>()
.and_then(|state| state.stored_values.get(&key))
.cloned();
let Some(value) = value else {
retval.set(v8::undefined(scope).into());
return;
};
let Some(value) = json_to_v8(scope, &value) else {
throw_type_error(scope, "failed to load stored value");
return;
};
retval.set(value);
}
pub(super) fn notify_callback(
scope: &mut v8::PinScope<'_, '_>,
args: v8::FunctionCallbackArguments,
mut retval: v8::ReturnValue<v8::Value>,
) {
let value = if args.length() == 0 {
v8::undefined(scope).into()
} else {
args.get(0)
};
let text = match serialize_output_text(scope, value) {
Ok(text) => text,
Err(error_text) => {
throw_type_error(scope, &error_text);
return;
}
};
if text.trim().is_empty() {
throw_type_error(scope, "notify expects non-empty text");
return;
}
if let Some(state) = scope.get_slot::<RuntimeState>() {
let _ = state.event_tx.send(RuntimeEvent::Notify {
call_id: state.tool_call_id.clone(),
text,
});
}
retval.set(v8::undefined(scope).into());
}
pub(super) fn yield_control_callback(
scope: &mut v8::PinScope<'_, '_>,
_args: v8::FunctionCallbackArguments,
_retval: v8::ReturnValue<v8::Value>,
) {
if let Some(state) = scope.get_slot::<RuntimeState>() {
let _ = state.event_tx.send(RuntimeEvent::YieldRequested);
}
}
pub(super) fn exit_callback(
scope: &mut v8::PinScope<'_, '_>,
_args: v8::FunctionCallbackArguments,
_retval: v8::ReturnValue<v8::Value>,
) {
if let Some(state) = scope.get_slot_mut::<RuntimeState>() {
state.exit_requested = true;
}
if let Some(error) = v8::String::new(scope, EXIT_SENTINEL) {
scope.throw_exception(error.into());
}
}

138
codex-rs/code-mode/src/runtime/globals.rs Normal file
View File

@@ -0,0 +1,138 @@
use super::RuntimeState;
use super::callbacks::exit_callback;
use super::callbacks::image_callback;
use super::callbacks::load_callback;
use super::callbacks::notify_callback;
use super::callbacks::store_callback;
use super::callbacks::text_callback;
use super::callbacks::tool_callback;
use super::callbacks::yield_control_callback;
pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), String> {
let global = scope.get_current_context().global(scope);
let console = v8::String::new(scope, "console")
.ok_or_else(|| "failed to allocate global `console`".to_string())?;
if global.delete(scope, console.into()) != Some(true) {
return Err("failed to remove global `console`".to_string());
}
let tools = build_tools_object(scope)?;
let all_tools = build_all_tools_value(scope)?;
let text = helper_function(scope, "text", text_callback)?;
let image = helper_function(scope, "image", image_callback)?;
let store = helper_function(scope, "store", store_callback)?;
let load = helper_function(scope, "load", load_callback)?;
let notify = helper_function(scope, "notify", notify_callback)?;
let yield_control = helper_function(scope, "yield_control", yield_control_callback)?;
let exit = helper_function(scope, "exit", exit_callback)?;
set_global(scope, global, "tools", tools.into())?;
set_global(scope, global, "ALL_TOOLS", all_tools)?;
set_global(scope, global, "text", text.into())?;
set_global(scope, global, "image", image.into())?;
set_global(scope, global, "store", store.into())?;
set_global(scope, global, "load", load.into())?;
set_global(scope, global, "notify", notify.into())?;
set_global(scope, global, "yield_control", yield_control.into())?;
set_global(scope, global, "exit", exit.into())?;
Ok(())
}
fn build_tools_object<'s>(
scope: &mut v8::PinScope<'s, '_>,
) -> Result<v8::Local<'s, v8::Object>, String> {
let tools = v8::Object::new(scope);
let enabled_tools = scope
.get_slot::<RuntimeState>()
.map(|state| state.enabled_tools.clone())
.unwrap_or_default();
for tool in enabled_tools {
let name = v8::String::new(scope, &tool.global_name)
.ok_or_else(|| "failed to allocate tool name".to_string())?;
let function = tool_function(scope, &tool.tool_name)?;
tools.set(scope, name.into(), function.into());
}
Ok(tools)
}
fn build_all_tools_value<'s>(
scope: &mut v8::PinScope<'s, '_>,
) -> Result<v8::Local<'s, v8::Value>, String> {
let enabled_tools = scope
.get_slot::<RuntimeState>()
.map(|state| state.enabled_tools.clone())
.unwrap_or_default();
let array = v8::Array::new(scope, enabled_tools.len() as i32);
let name_key = v8::String::new(scope, "name")
.ok_or_else(|| "failed to allocate ALL_TOOLS name key".to_string())?;
let description_key = v8::String::new(scope, "description")
.ok_or_else(|| "failed to allocate ALL_TOOLS description key".to_string())?;
for (index, tool) in enabled_tools.iter().enumerate() {
let item = v8::Object::new(scope);
let name = v8::String::new(scope, &tool.global_name)
.ok_or_else(|| "failed to allocate ALL_TOOLS name".to_string())?;
let description = v8::String::new(scope, &tool.description)
.ok_or_else(|| "failed to allocate ALL_TOOLS description".to_string())?;
if item.set(scope, name_key.into(), name.into()) != Some(true) {
return Err("failed to set ALL_TOOLS name".to_string());
}
if item.set(scope, description_key.into(), description.into()) != Some(true) {
return Err("failed to set ALL_TOOLS description".to_string());
}
if array.set_index(scope, index as u32, item.into()) != Some(true) {
return Err("failed to append ALL_TOOLS metadata".to_string());
}
}
Ok(array.into())
}
fn helper_function<'s, F>(
scope: &mut v8::PinScope<'s, '_>,
name: &str,
callback: F,
) -> Result<v8::Local<'s, v8::Function>, String>
where
F: v8::MapFnTo<v8::FunctionCallback>,
{
let name =
v8::String::new(scope, name).ok_or_else(|| "failed to allocate helper name".to_string())?;
let template = v8::FunctionTemplate::builder(callback)
.data(name.into())
.build(scope);
template
.get_function(scope)
.ok_or_else(|| "failed to create helper function".to_string())
}
fn tool_function<'s>(
scope: &mut v8::PinScope<'s, '_>,
tool_name: &str,
) -> Result<v8::Local<'s, v8::Function>, String> {
let data = v8::String::new(scope, tool_name)
.ok_or_else(|| "failed to allocate tool callback data".to_string())?;
let template = v8::FunctionTemplate::builder(tool_callback)
.data(data.into())
.build(scope);
template
.get_function(scope)
.ok_or_else(|| "failed to create tool function".to_string())
}
fn set_global<'s>(
scope: &mut v8::PinScope<'s, '_>,
global: v8::Local<'s, v8::Object>,
name: &str,
value: v8::Local<'s, v8::Value>,
) -> Result<(), String> {
let key = v8::String::new(scope, name)
.ok_or_else(|| format!("failed to allocate global `{name}`"))?;
if global.set(scope, key.into(), value) == Some(true) {
Ok(())
} else {
Err(format!("failed to set global `{name}`"))
}
}

349
codex-rs/code-mode/src/runtime/mod.rs Normal file
View File

@@ -0,0 +1,349 @@
mod callbacks;
mod globals;
mod module_loader;
mod value;
use std::collections::HashMap;
use std::sync::OnceLock;
use std::sync::mpsc as std_mpsc;
use std::thread;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use crate::description::EnabledToolMetadata;
use crate::description::ToolDefinition;
use crate::description::enabled_tool_metadata;
use crate::response::FunctionCallOutputContentItem;
pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000;
pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000;
pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000;
const EXIT_SENTINEL: &str = "__codex_code_mode_exit__";
#[derive(Clone, Debug)]
pub struct ExecuteRequest {
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
pub source: String,
pub stored_values: HashMap<String, JsonValue>,
pub yield_time_ms: Option<u64>,
pub max_output_tokens: Option<usize>,
}
#[derive(Clone, Debug)]
pub struct WaitRequest {
pub cell_id: String,
pub yield_time_ms: u64,
pub terminate: bool,
}
#[derive(Debug, PartialEq)]
pub enum RuntimeResponse {
Yielded {
cell_id: String,
content_items: Vec<FunctionCallOutputContentItem>,
},
Terminated {
cell_id: String,
content_items: Vec<FunctionCallOutputContentItem>,
},
Result {
cell_id: String,
content_items: Vec<FunctionCallOutputContentItem>,
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
#[derive(Debug)]
pub(crate) enum TurnMessage {
ToolCall {
cell_id: String,
id: String,
name: String,
input: Option<JsonValue>,
},
Notify {
cell_id: String,
call_id: String,
text: String,
},
}
#[derive(Debug)]
pub(crate) enum RuntimeCommand {
ToolResponse { id: String, result: JsonValue },
ToolError { id: String, error_text: String },
Terminate,
}
#[derive(Debug)]
pub(crate) enum RuntimeEvent {
Started,
ContentItem(FunctionCallOutputContentItem),
YieldRequested,
ToolCall {
id: String,
name: String,
input: Option<JsonValue>,
},
Notify {
call_id: String,
text: String,
},
Result {
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
pub(crate) fn spawn_runtime(
request: ExecuteRequest,
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
) -> Result<(std_mpsc::Sender<RuntimeCommand>, v8::IsolateHandle), String> {
let (command_tx, command_rx) = std_mpsc::channel();
let (isolate_handle_tx, isolate_handle_rx) = std_mpsc::sync_channel(1);
let enabled_tools = request
.enabled_tools
.iter()
.map(enabled_tool_metadata)
.collect::<Vec<_>>();
let config = RuntimeConfig {
tool_call_id: request.tool_call_id,
enabled_tools,
source: request.source,
stored_values: request.stored_values,
};
thread::spawn(move || {
run_runtime(config, event_tx, command_rx, isolate_handle_tx);
});
let isolate_handle = isolate_handle_rx
.recv()
.map_err(|_| "failed to initialize code mode runtime".to_string())?;
Ok((command_tx, isolate_handle))
}
#[derive(Clone)]
struct RuntimeConfig {
tool_call_id: String,
enabled_tools: Vec<EnabledToolMetadata>,
source: String,
stored_values: HashMap<String, JsonValue>,
}
pub(super) struct RuntimeState {
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
pending_tool_calls: HashMap<String, v8::Global<v8::PromiseResolver>>,
stored_values: HashMap<String, JsonValue>,
enabled_tools: Vec<EnabledToolMetadata>,
next_tool_call_id: u64,
tool_call_id: String,
exit_requested: bool,
}
pub(super) enum CompletionState {
Pending,
Completed {
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
fn initialize_v8() {
static PLATFORM: OnceLock<v8::SharedRef<v8::Platform>> = OnceLock::new();
let _ = PLATFORM.get_or_init(|| {
let platform = v8::new_default_platform(0, false).make_shared();
v8::V8::initialize_platform(platform.clone());
v8::V8::initialize();
platform
});
}
fn run_runtime(
config: RuntimeConfig,
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
command_rx: std_mpsc::Receiver<RuntimeCommand>,
isolate_handle_tx: std_mpsc::SyncSender<v8::IsolateHandle>,
) {
initialize_v8();
let isolate = &mut v8::Isolate::new(v8::CreateParams::default());
let isolate_handle = isolate.thread_safe_handle();
if isolate_handle_tx.send(isolate_handle).is_err() {
return;
}
isolate.set_host_import_module_dynamically_callback(module_loader::dynamic_import_callback);
v8::scope!(let scope, isolate);
let context = v8::Context::new(scope, Default::default());
let scope = &mut v8::ContextScope::new(scope, context);
scope.set_slot(RuntimeState {
event_tx: event_tx.clone(),
pending_tool_calls: HashMap::new(),
stored_values: config.stored_values,
enabled_tools: config.enabled_tools,
next_tool_call_id: 1,
tool_call_id: config.tool_call_id,
exit_requested: false,
});
if let Err(error_text) = globals::install_globals(scope) {
send_result(&event_tx, HashMap::new(), Some(error_text));
return;
}
let _ = event_tx.send(RuntimeEvent::Started);
let pending_promise = match module_loader::evaluate_main_module(scope, &config.source) {
Ok(pending_promise) => pending_promise,
Err(error_text) => {
capture_scope_send_error(scope, &event_tx, Some(error_text));
return;
}
};
match module_loader::completion_state(scope, pending_promise.as_ref()) {
CompletionState::Completed {
stored_values,
error_text,
} => {
send_result(&event_tx, stored_values, error_text);
return;
}
CompletionState::Pending => {}
}
let mut pending_promise = pending_promise;
loop {
let Ok(command) = command_rx.recv() else {
break;
};
match command {
RuntimeCommand::Terminate => break,
RuntimeCommand::ToolResponse { id, result } => {
if let Err(error_text) =
module_loader::resolve_tool_response(scope, &id, Ok(result))
{
capture_scope_send_error(scope, &event_tx, Some(error_text));
return;
}
}
RuntimeCommand::ToolError { id, error_text } => {
if let Err(runtime_error) =
module_loader::resolve_tool_response(scope, &id, Err(error_text))
{
capture_scope_send_error(scope, &event_tx, Some(runtime_error));
return;
}
}
}
scope.perform_microtask_checkpoint();
match module_loader::completion_state(scope, pending_promise.as_ref()) {
CompletionState::Completed {
stored_values,
error_text,
} => {
send_result(&event_tx, stored_values, error_text);
return;
}
CompletionState::Pending => {}
}
if let Some(promise) = pending_promise.as_ref() {
let promise = v8::Local::new(scope, promise);
if promise.state() != v8::PromiseState::Pending {
pending_promise = None;
}
}
}
}
fn capture_scope_send_error(
scope: &mut v8::PinScope<'_, '_>,
event_tx: &mpsc::UnboundedSender<RuntimeEvent>,
error_text: Option<String>,
) {
let stored_values = scope
.get_slot::<RuntimeState>()
.map(|state| state.stored_values.clone())
.unwrap_or_default();
send_result(event_tx, stored_values, error_text);
}
fn send_result(
event_tx: &mpsc::UnboundedSender<RuntimeEvent>,
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
) {
let _ = event_tx.send(RuntimeEvent::Result {
stored_values,
error_text,
});
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::time::Duration;
use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
use super::ExecuteRequest;
use super::RuntimeEvent;
use super::spawn_runtime;
fn execute_request(source: &str) -> ExecuteRequest {
ExecuteRequest {
tool_call_id: "call_1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
stored_values: HashMap::new(),
yield_time_ms: Some(1),
max_output_tokens: None,
}
}
#[tokio::test]
async fn terminate_execution_stops_cpu_bound_module() {
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
let (_runtime_tx, runtime_terminate_handle) =
spawn_runtime(execute_request("while (true) {}"), event_tx).unwrap();
let started_event = tokio::time::timeout(Duration::from_secs(1), event_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(started_event, RuntimeEvent::Started));
assert!(runtime_terminate_handle.terminate_execution());
let result_event = tokio::time::timeout(Duration::from_secs(1), event_rx.recv())
.await
.unwrap()
.unwrap();
let RuntimeEvent::Result {
stored_values,
error_text,
} = result_event
else {
panic!("expected runtime result after termination");
};
assert_eq!(stored_values, HashMap::new());
assert!(error_text.is_some());
assert!(
tokio::time::timeout(Duration::from_secs(1), event_rx.recv())
.await
.unwrap()
.is_none()
);
}
}

235
codex-rs/code-mode/src/runtime/module_loader.rs Normal file
View File

@@ -0,0 +1,235 @@
use serde_json::Value as JsonValue;
use super::CompletionState;
use super::EXIT_SENTINEL;
use super::RuntimeState;
use super::value::json_to_v8;
use super::value::value_to_error_text;
pub(super) fn evaluate_main_module(
scope: &mut v8::PinScope<'_, '_>,
source_text: &str,
) -> Result<Option<v8::Global<v8::Promise>>, String> {
let tc = std::pin::pin!(v8::TryCatch::new(scope));
let mut tc = tc.init();
let source = v8::String::new(&tc, source_text)
.ok_or_else(|| "failed to allocate exec source".to_string())?;
let origin = script_origin(&mut tc, "exec_main.mjs")?;
let mut source = v8::script_compiler::Source::new(source, Some(&origin));
let module = v8::script_compiler::compile_module(&tc, &mut source).ok_or_else(|| {
tc.exception()
.map(|exception| value_to_error_text(&mut tc, exception))
.unwrap_or_else(|| "unknown code mode exception".to_string())
})?;
module
.instantiate_module(&tc, resolve_module_callback)
.ok_or_else(|| {
tc.exception()
.map(|exception| value_to_error_text(&mut tc, exception))
.unwrap_or_else(|| "unknown code mode exception".to_string())
})?;
let result = match module.evaluate(&tc) {
Some(result) => result,
None => {
if let Some(exception) = tc.exception() {
if is_exit_exception(&mut tc, exception) {
return Ok(None);
}
return Err(value_to_error_text(&mut tc, exception));
}
return Err("unknown code mode exception".to_string());
}
};
tc.perform_microtask_checkpoint();
if result.is_promise() {
let promise = v8::Local::<v8::Promise>::try_from(result)
.map_err(|_| "failed to read exec promise".to_string())?;
return Ok(Some(v8::Global::new(&tc, promise)));
}
Ok(None)
}
fn is_exit_exception(
scope: &mut v8::PinScope<'_, '_>,
exception: v8::Local<'_, v8::Value>,
) -> bool {
scope
.get_slot::<RuntimeState>()
.map(|state| state.exit_requested)
.unwrap_or(false)
&& exception.is_string()
&& exception.to_rust_string_lossy(scope) == EXIT_SENTINEL
}
pub(super) fn resolve_tool_response(
scope: &mut v8::PinScope<'_, '_>,
id: &str,
response: Result<JsonValue, String>,
) -> Result<(), String> {
let resolver = {
let state = scope
.get_slot_mut::<RuntimeState>()
.ok_or_else(|| "runtime state unavailable".to_string())?;
state.pending_tool_calls.remove(id)
}
.ok_or_else(|| format!("unknown tool call `{id}`"))?;
let tc = std::pin::pin!(v8::TryCatch::new(scope));
let mut tc = tc.init();
let resolver = v8::Local::new(&tc, &resolver);
match response {
Ok(result) => {
let value = json_to_v8(&mut tc, &result)
.ok_or_else(|| "failed to serialize tool response".to_string())?;
resolver.resolve(&tc, value);
}
Err(error_text) => {
let value = v8::String::new(&tc, &error_text)
.ok_or_else(|| "failed to allocate tool error".to_string())?;
resolver.reject(&tc, value.into());
}
}
if tc.has_caught() {
return Err(tc
.exception()
.map(|exception| value_to_error_text(&mut tc, exception))
.unwrap_or_else(|| "unknown code mode exception".to_string()));
}
Ok(())
}
pub(super) fn completion_state(
scope: &mut v8::PinScope<'_, '_>,
pending_promise: Option<&v8::Global<v8::Promise>>,
) -> CompletionState {
let stored_values = scope
.get_slot::<RuntimeState>()
.map(|state| state.stored_values.clone())
.unwrap_or_default();
let Some(pending_promise) = pending_promise else {
return CompletionState::Completed {
stored_values,
error_text: None,
};
};
let promise = v8::Local::new(scope, pending_promise);
match promise.state() {
v8::PromiseState::Pending => CompletionState::Pending,
v8::PromiseState::Fulfilled => CompletionState::Completed {
stored_values,
error_text: None,
},
v8::PromiseState::Rejected => {
let result = promise.result(scope);
let error_text = if is_exit_exception(scope, result) {
None
} else {
Some(value_to_error_text(scope, result))
};
CompletionState::Completed {
stored_values,
error_text,
}
}
}
}
fn script_origin<'s>(
scope: &mut v8::PinScope<'s, '_>,
resource_name_: &str,
) -> Result<v8::ScriptOrigin<'s>, String> {
let resource_name = v8::String::new(scope, resource_name_)
.ok_or_else(|| "failed to allocate script origin".to_string())?;
let source_map_url = v8::String::new(scope, resource_name_)
.ok_or_else(|| "failed to allocate source map url".to_string())?;
Ok(v8::ScriptOrigin::new(
scope,
resource_name.into(),
0,
0,
true,
0,
Some(source_map_url.into()),
true,
false,
true,
None,
))
}
fn resolve_module_callback<'s>(
context: v8::Local<'s, v8::Context>,
specifier: v8::Local<'s, v8::String>,
_import_attributes: v8::Local<'s, v8::FixedArray>,
_referrer: v8::Local<'s, v8::Module>,
) -> Option<v8::Local<'s, v8::Module>> {
v8::callback_scope!(unsafe scope, context);
let specifier = specifier.to_rust_string_lossy(scope);
resolve_module(scope, &specifier)
}
pub(super) fn dynamic_import_callback<'s>(
scope: &mut v8::PinScope<'s, '_>,
_host_defined_options: v8::Local<'s, v8::Data>,
_resource_name: v8::Local<'s, v8::Value>,
specifier: v8::Local<'s, v8::String>,
_import_attributes: v8::Local<'s, v8::FixedArray>,
) -> Option<v8::Local<'s, v8::Promise>> {
let specifier = specifier.to_rust_string_lossy(scope);
let resolver = v8::PromiseResolver::new(scope)?;
match resolve_module(scope, &specifier) {
Some(module) => {
if module.get_status() == v8::ModuleStatus::Uninstantiated
&& module
.instantiate_module(scope, resolve_module_callback)
.is_none()
{
let error = v8::String::new(scope, "failed to instantiate module")
.map(Into::into)
.unwrap_or_else(|| v8::undefined(scope).into());
resolver.reject(scope, error);
return Some(resolver.get_promise(scope));
}
if matches!(
module.get_status(),
v8::ModuleStatus::Instantiated | v8::ModuleStatus::Evaluated
) && module.evaluate(scope).is_none()
{
let error = v8::String::new(scope, "failed to evaluate module")
.map(Into::into)
.unwrap_or_else(|| v8::undefined(scope).into());
resolver.reject(scope, error);
return Some(resolver.get_promise(scope));
}
let namespace = module.get_module_namespace();
resolver.resolve(scope, namespace);
Some(resolver.get_promise(scope))
}
None => {
let error = v8::String::new(scope, "unsupported import in exec")
.map(Into::into)
.unwrap_or_else(|| v8::undefined(scope).into());
resolver.reject(scope, error);
Some(resolver.get_promise(scope))
}
}
}
fn resolve_module<'s>(
scope: &mut v8::PinScope<'s, '_>,
specifier: &str,
) -> Option<v8::Local<'s, v8::Module>> {
if let Some(message) =
v8::String::new(scope, &format!("Unsupported import in exec: {specifier}"))
{
scope.throw_exception(message.into());
} else {
scope.throw_exception(v8::undefined(scope).into());
}
None
}

163
codex-rs/code-mode/src/runtime/value.rs Normal file
View File

@@ -0,0 +1,163 @@
use serde_json::Value as JsonValue;
use crate::response::FunctionCallOutputContentItem;
use crate::response::ImageDetail;
pub(super) fn serialize_output_text(
scope: &mut v8::PinScope<'_, '_>,
value: v8::Local<'_, v8::Value>,
) -> Result<String, String> {
if value.is_undefined()
|| value.is_null()
|| value.is_boolean()
|| value.is_number()
|| value.is_big_int()
|| value.is_string()
{
return Ok(value.to_rust_string_lossy(scope));
}
let tc = std::pin::pin!(v8::TryCatch::new(scope));
let mut tc = tc.init();
if let Some(stringified) = v8::json::stringify(&tc, value) {
return Ok(stringified.to_rust_string_lossy(&tc));
}
if tc.has_caught() {
return Err(tc
.exception()
.map(|exception| value_to_error_text(&mut tc, exception))
.unwrap_or_else(|| "unknown code mode exception".to_string()));
}
Ok(value.to_rust_string_lossy(&tc))
}
pub(super) fn normalize_output_image(
scope: &mut v8::PinScope<'_, '_>,
value: v8::Local<'_, v8::Value>,
) -> Result<FunctionCallOutputContentItem, ()> {
let result = (|| -> Result<FunctionCallOutputContentItem, String> {
let (image_url, detail) = if value.is_string() {
(value.to_rust_string_lossy(scope), None)
} else if value.is_object() && !value.is_array() {
let object = v8::Local::<v8::Object>::try_from(value).map_err(|_| {
"image expects a non-empty image URL string or an object with image_url and optional detail".to_string()
})?;
let image_url_key = v8::String::new(scope, "image_url")
.ok_or_else(|| "failed to allocate image helper keys".to_string())?;
let detail_key = v8::String::new(scope, "detail")
.ok_or_else(|| "failed to allocate image helper keys".to_string())?;
let image_url = object
.get(scope, image_url_key.into())
.filter(|value| value.is_string())
.map(|value| value.to_rust_string_lossy(scope))
.ok_or_else(|| {
"image expects a non-empty image URL string or an object with image_url and optional detail"
.to_string()
})?;
let detail = match object.get(scope, detail_key.into()) {
Some(value) if value.is_string() => Some(value.to_rust_string_lossy(scope)),
Some(value) if value.is_null() || value.is_undefined() => None,
Some(_) => return Err("image detail must be a string when provided".to_string()),
None => None,
};
(image_url, detail)
} else {
return Err(
"image expects a non-empty image URL string or an object with image_url and optional detail"
.to_string(),
);
};
if image_url.is_empty() {
return Err(
"image expects a non-empty image URL string or an object with image_url and optional detail"
.to_string(),
);
}
let lower = image_url.to_ascii_lowercase();
if !(lower.starts_with("http://")
|| lower.starts_with("https://")
|| lower.starts_with("data:"))
{
return Err("image expects an http(s) or data URL".to_string());
}
let detail = match detail {
Some(detail) => {
let normalized = detail.to_ascii_lowercase();
Some(match normalized.as_str() {
"auto" => ImageDetail::Auto,
"low" => ImageDetail::Low,
"high" => ImageDetail::High,
"original" => ImageDetail::Original,
_ => {
return Err(
"image detail must be one of: auto, low, high, original".to_string()
);
}
})
}
None => None,
};
Ok(FunctionCallOutputContentItem::InputImage { image_url, detail })
})();
match result {
Ok(item) => Ok(item),
Err(error_text) => {
throw_type_error(scope, &error_text);
Err(())
}
}
}
pub(super) fn v8_value_to_json(
scope: &mut v8::PinScope<'_, '_>,
value: v8::Local<'_, v8::Value>,
) -> Result<Option<JsonValue>, String> {
let tc = std::pin::pin!(v8::TryCatch::new(scope));
let mut tc = tc.init();
let Some(stringified) = v8::json::stringify(&tc, value) else {
if tc.has_caught() {
return Err(tc
.exception()
.map(|exception| value_to_error_text(&mut tc, exception))
.unwrap_or_else(|| "unknown code mode exception".to_string()));
}
return Ok(None);
};
serde_json::from_str(&stringified.to_rust_string_lossy(&tc))
.map(Some)
.map_err(|err| format!("failed to serialize JavaScript value: {err}"))
}
pub(super) fn json_to_v8<'s>(
scope: &mut v8::PinScope<'s, '_>,
value: &JsonValue,
) -> Option<v8::Local<'s, v8::Value>> {
let json = serde_json::to_string(value).ok()?;
let json = v8::String::new(scope, &json)?;
v8::json::parse(scope, json)
}
pub(super) fn value_to_error_text(
scope: &mut v8::PinScope<'_, '_>,
value: v8::Local<'_, v8::Value>,
) -> String {
if value.is_object()
&& let Ok(object) = v8::Local::<v8::Object>::try_from(value)
&& let Some(key) = v8::String::new(scope, "stack")
&& let Some(stack) = object.get(scope, key.into())
&& stack.is_string()
{
return stack.to_rust_string_lossy(scope);
}
value.to_rust_string_lossy(scope)
}
pub(super) fn throw_type_error(scope: &mut v8::PinScope<'_, '_>, message: &str) {
if let Some(message) = v8::String::new(scope, message) {
scope.throw_exception(message.into());
}
}

673
codex-rs/code-mode/src/service.rs Normal file
View File

@@ -0,0 +1,673 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;
use async_trait::async_trait;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use crate::FunctionCallOutputContentItem;
use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS;
use crate::runtime::ExecuteRequest;
use crate::runtime::RuntimeCommand;
use crate::runtime::RuntimeEvent;
use crate::runtime::RuntimeResponse;
use crate::runtime::TurnMessage;
use crate::runtime::WaitRequest;
use crate::runtime::spawn_runtime;
#[async_trait]
pub trait CodeModeTurnHost: Send + Sync {
async fn invoke_tool(
&self,
tool_name: String,
input: Option<JsonValue>,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String>;
async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>;
}
#[derive(Clone)]
struct SessionHandle {
control_tx: mpsc::UnboundedSender<SessionControlCommand>,
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
}
struct Inner {
stored_values: Mutex<HashMap<String, JsonValue>>,
sessions: Mutex<HashMap<String, SessionHandle>>,
turn_message_tx: mpsc::UnboundedSender<TurnMessage>,
turn_message_rx: Arc<Mutex<mpsc::UnboundedReceiver<TurnMessage>>>,
next_cell_id: AtomicU64,
}
pub struct CodeModeService {
inner: Arc<Inner>,
}
impl CodeModeService {
pub fn new() -> Self {
let (turn_message_tx, turn_message_rx) = mpsc::unbounded_channel();
Self {
inner: Arc::new(Inner {
stored_values: Mutex::new(HashMap::new()),
sessions: Mutex::new(HashMap::new()),
turn_message_tx,
turn_message_rx: Arc::new(Mutex::new(turn_message_rx)),
next_cell_id: AtomicU64::new(1),
}),
}
}
pub async fn stored_values(&self) -> HashMap<String, JsonValue> {
self.inner.stored_values.lock().await.clone()
}
pub async fn replace_stored_values(&self, values: HashMap<String, JsonValue>) {
*self.inner.stored_values.lock().await = values;
}
pub async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String> {
let cell_id = self
.inner
.next_cell_id
.fetch_add(1, Ordering::Relaxed)
.to_string();
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (runtime_tx, runtime_terminate_handle) = spawn_runtime(request.clone(), event_tx)?;
let (control_tx, control_rx) = mpsc::unbounded_channel();
let (response_tx, response_rx) = oneshot::channel();
self.inner.sessions.lock().await.insert(
cell_id.clone(),
SessionHandle {
control_tx: control_tx.clone(),
runtime_tx: runtime_tx.clone(),
},
);
tokio::spawn(run_session_control(
Arc::clone(&self.inner),
SessionControlContext {
cell_id: cell_id.clone(),
runtime_tx,
runtime_terminate_handle,
},
event_rx,
control_rx,
response_tx,
request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS),
));
response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())
}
pub async fn wait(&self, request: WaitRequest) -> Result<RuntimeResponse, String> {
let cell_id = request.cell_id.clone();
let handle = self
.inner
.sessions
.lock()
.await
.get(&request.cell_id)
.cloned();
let Some(handle) = handle else {
return Ok(missing_cell_response(cell_id));
};
let (response_tx, response_rx) = oneshot::channel();
let control_message = if request.terminate {
SessionControlCommand::Terminate { response_tx }
} else {
SessionControlCommand::Poll {
yield_time_ms: request.yield_time_ms,
response_tx,
}
};
if handle.control_tx.send(control_message).is_err() {
return Ok(missing_cell_response(cell_id));
}
match response_rx.await {
Ok(response) => Ok(response),
Err(_) => Ok(missing_cell_response(request.cell_id)),
}
}
pub fn start_turn_worker(&self, host: Arc<dyn CodeModeTurnHost>) -> CodeModeTurnWorker {
let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
let inner = Arc::clone(&self.inner);
let turn_message_rx = Arc::clone(&self.inner.turn_message_rx);
tokio::spawn(async move {
loop {
let next_message = tokio::select! {
_ = &mut shutdown_rx => break,
message = async {
let mut turn_message_rx = turn_message_rx.lock().await;
turn_message_rx.recv().await
} => message,
};
let Some(next_message) = next_message else {
break;
};
match next_message {
TurnMessage::Notify {
cell_id,
call_id,
text,
} => {
if let Err(err) = host.notify(call_id, cell_id.clone(), text).await {
warn!(
"failed to deliver code mode notification for cell {cell_id}: {err}"
);
}
}
TurnMessage::ToolCall {
cell_id,
id,
name,
input,
} => {
let host = Arc::clone(&host);
let inner = Arc::clone(&inner);
tokio::spawn(async move {
let response = host
.invoke_tool(name, input, CancellationToken::new())
.await;
let runtime_tx = inner
.sessions
.lock()
.await
.get(&cell_id)
.map(|handle| handle.runtime_tx.clone());
let Some(runtime_tx) = runtime_tx else {
return;
};
let command = match response {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
};
let _ = runtime_tx.send(command);
});
}
}
}
});
CodeModeTurnWorker {
shutdown_tx: Some(shutdown_tx),
}
}
}
impl Default for CodeModeService {
fn default() -> Self {
Self::new()
}
}
pub struct CodeModeTurnWorker {
shutdown_tx: Option<oneshot::Sender<()>>,
}
impl Drop for CodeModeTurnWorker {
fn drop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
}
}
enum SessionControlCommand {
Poll {
yield_time_ms: u64,
response_tx: oneshot::Sender<RuntimeResponse>,
},
Terminate {
response_tx: oneshot::Sender<RuntimeResponse>,
},
}
struct PendingResult {
content_items: Vec<FunctionCallOutputContentItem>,
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
}
struct SessionControlContext {
cell_id: String,
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
runtime_terminate_handle: v8::IsolateHandle,
}
fn missing_cell_response(cell_id: String) -> RuntimeResponse {
RuntimeResponse::Result {
error_text: Some(format!("exec cell {cell_id} not found")),
cell_id,
content_items: Vec::new(),
stored_values: HashMap::new(),
}
}
fn pending_result_response(cell_id: &str, result: PendingResult) -> RuntimeResponse {
RuntimeResponse::Result {
cell_id: cell_id.to_string(),
content_items: result.content_items,
stored_values: result.stored_values,
error_text: result.error_text,
}
}
fn send_or_buffer_result(
cell_id: &str,
result: PendingResult,
response_tx: &mut Option<oneshot::Sender<RuntimeResponse>>,
pending_result: &mut Option<PendingResult>,
) -> bool {
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(pending_result_response(cell_id, result));
return true;
}
*pending_result = Some(result);
false
}
async fn run_session_control(
inner: Arc<Inner>,
context: SessionControlContext,
mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
mut control_rx: mpsc::UnboundedReceiver<SessionControlCommand>,
initial_response_tx: oneshot::Sender<RuntimeResponse>,
initial_yield_time_ms: u64,
) {
let SessionControlContext {
cell_id,
runtime_tx,
runtime_terminate_handle,
} = context;
let mut content_items = Vec::new();
let mut pending_result: Option<PendingResult> = None;
let mut response_tx = Some(initial_response_tx);
let mut termination_requested = false;
let mut runtime_closed = false;
let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
loop {
tokio::select! {
maybe_event = async {
if runtime_closed {
std::future::pending::<Option<RuntimeEvent>>().await
} else {
event_rx.recv().await
}
} => {
let Some(event) = maybe_event else {
runtime_closed = true;
if termination_requested {
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(RuntimeResponse::Terminated {
cell_id: cell_id.clone(),
content_items: std::mem::take(&mut content_items),
});
}
break;
}
if pending_result.is_none() {
let result = PendingResult {
content_items: std::mem::take(&mut content_items),
stored_values: HashMap::new(),
error_text: Some("exec runtime ended unexpectedly".to_string()),
};
if send_or_buffer_result(
&cell_id,
result,
&mut response_tx,
&mut pending_result,
) {
break;
}
}
continue;
};
match event {
RuntimeEvent::Started => {
yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(initial_yield_time_ms))));
}
RuntimeEvent::ContentItem(item) => {
content_items.push(item);
}
RuntimeEvent::YieldRequested => {
yield_timer = None;
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: std::mem::take(&mut content_items),
});
}
}
RuntimeEvent::Notify { call_id, text } => {
let _ = inner.turn_message_tx.send(TurnMessage::Notify {
cell_id: cell_id.clone(),
call_id,
text,
});
}
RuntimeEvent::ToolCall { id, name, input } => {
let _ = inner.turn_message_tx.send(TurnMessage::ToolCall {
cell_id: cell_id.clone(),
id,
name,
input,
});
}
RuntimeEvent::Result {
stored_values,
error_text,
} => {
yield_timer = None;
if termination_requested {
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(RuntimeResponse::Terminated {
cell_id: cell_id.clone(),
content_items: std::mem::take(&mut content_items),
});
}
break;
}
let result = PendingResult {
content_items: std::mem::take(&mut content_items),
stored_values,
error_text,
};
if send_or_buffer_result(
&cell_id,
result,
&mut response_tx,
&mut pending_result,
) {
break;
}
}
}
}
maybe_command = control_rx.recv() => {
let Some(command) = maybe_command else {
break;
};
match command {
SessionControlCommand::Poll {
yield_time_ms,
response_tx: next_response_tx,
} => {
if let Some(result) = pending_result.take() {
let _ = next_response_tx.send(pending_result_response(&cell_id, result));
break;
}
response_tx = Some(next_response_tx);
yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms))));
}
SessionControlCommand::Terminate { response_tx: next_response_tx } => {
if let Some(result) = pending_result.take() {
let _ = next_response_tx.send(pending_result_response(&cell_id, result));
break;
}
response_tx = Some(next_response_tx);
termination_requested = true;
yield_timer = None;
let _ = runtime_tx.send(RuntimeCommand::Terminate);
let _ = runtime_terminate_handle.terminate_execution();
if runtime_closed {
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(RuntimeResponse::Terminated {
cell_id: cell_id.clone(),
content_items: std::mem::take(&mut content_items),
});
}
break;
} else {
continue;
}
}
}
}
_ = async {
if let Some(yield_timer) = yield_timer.as_mut() {
yield_timer.await;
} else {
std::future::pending::<()>().await;
}
} => {
yield_timer = None;
if let Some(response_tx) = response_tx.take() {
let _ = response_tx.send(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: std::mem::take(&mut content_items),
});
}
}
}
}
let _ = runtime_tx.send(RuntimeCommand::Terminate);
inner.sessions.lock().await.remove(&cell_id);
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::time::Duration;
use pretty_assertions::assert_eq;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use super::CodeModeService;
use super::Inner;
use super::RuntimeCommand;
use super::RuntimeResponse;
use super::SessionControlCommand;
use super::SessionControlContext;
use super::run_session_control;
use crate::FunctionCallOutputContentItem;
use crate::runtime::ExecuteRequest;
use crate::runtime::RuntimeEvent;
use crate::runtime::spawn_runtime;
fn execute_request(source: &str) -> ExecuteRequest {
ExecuteRequest {
tool_call_id: "call_1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
stored_values: HashMap::new(),
yield_time_ms: Some(1),
max_output_tokens: None,
}
}
fn test_inner() -> Arc<Inner> {
let (turn_message_tx, turn_message_rx) = mpsc::unbounded_channel();
Arc::new(Inner {
stored_values: Mutex::new(HashMap::new()),
sessions: Mutex::new(HashMap::new()),
turn_message_tx,
turn_message_rx: Arc::new(Mutex::new(turn_message_rx)),
next_cell_id: AtomicU64::new(1),
})
}
#[tokio::test]
async fn synchronous_exit_returns_successfully() {
let service = CodeModeService::new();
let response = service
.execute(ExecuteRequest {
source: r#"text("before"); exit(); text("after");"#.to_string(),
yield_time_ms: None,
..execute_request("")
})
.await
.unwrap();
assert_eq!(
response,
RuntimeResponse::Result {
cell_id: "1".to_string(),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "before".to_string(),
}],
stored_values: HashMap::new(),
error_text: None,
}
);
}
#[tokio::test]
async fn v8_console_is_not_exposed_on_global_this() {
let service = CodeModeService::new();
let response = service
.execute(ExecuteRequest {
source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(),
yield_time_ms: None,
..execute_request("")
})
.await
.unwrap();
assert_eq!(
response,
RuntimeResponse::Result {
cell_id: "1".to_string(),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "false".to_string(),
}],
stored_values: HashMap::new(),
error_text: None,
}
);
}
#[tokio::test]
async fn output_helpers_return_undefined() {
let service = CodeModeService::new();
let response = service
.execute(ExecuteRequest {
source: r#"
const returnsUndefined = [
text("first"),
image("https://example.com/image.jpg"),
notify("ping"),
].map((value) => value === undefined);
text(JSON.stringify(returnsUndefined));
"#
.to_string(),
yield_time_ms: None,
..execute_request("")
})
.await
.unwrap();
assert_eq!(
response,
RuntimeResponse::Result {
cell_id: "1".to_string(),
content_items: vec![
FunctionCallOutputContentItem::InputText {
text: "first".to_string(),
},
FunctionCallOutputContentItem::InputImage {
image_url: "https://example.com/image.jpg".to_string(),
detail: None,
},
FunctionCallOutputContentItem::InputText {
text: "[true,true,true]".to_string(),
},
],
stored_values: HashMap::new(),
error_text: None,
}
);
}
#[tokio::test]
async fn terminate_waits_for_runtime_shutdown_before_responding() {
let inner = test_inner();
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (control_tx, control_rx) = mpsc::unbounded_channel();
let (initial_response_tx, initial_response_rx) = oneshot::channel();
let (runtime_event_tx, _runtime_event_rx) = mpsc::unbounded_channel();
let (runtime_tx, runtime_terminate_handle) = spawn_runtime(
ExecuteRequest {
source: "await new Promise(() => {})".to_string(),
yield_time_ms: None,
..execute_request("")
},
runtime_event_tx,
)
.unwrap();
tokio::spawn(run_session_control(
inner,
SessionControlContext {
cell_id: "cell-1".to_string(),
runtime_tx: runtime_tx.clone(),
runtime_terminate_handle,
},
event_rx,
control_rx,
initial_response_tx,
60_000,
));
event_tx.send(RuntimeEvent::Started).unwrap();
event_tx.send(RuntimeEvent::YieldRequested).unwrap();
assert_eq!(
initial_response_rx.await.unwrap(),
RuntimeResponse::Yielded {
cell_id: "cell-1".to_string(),
content_items: Vec::new(),
}
);
let (terminate_response_tx, terminate_response_rx) = oneshot::channel();
control_tx
.send(SessionControlCommand::Terminate {
response_tx: terminate_response_tx,
})
.unwrap();
let terminate_response = async { terminate_response_rx.await.unwrap() };
tokio::pin!(terminate_response);
assert!(
tokio::time::timeout(Duration::from_millis(100), terminate_response.as_mut())
.await
.is_err()
);
drop(event_tx);
assert_eq!(
terminate_response.await,
RuntimeResponse::Terminated {
cell_id: "cell-1".to_string(),
content_items: Vec::new(),
}
);
let _ = runtime_tx.send(RuntimeCommand::Terminate);
}
}