From f1815041e44ae14c1dbecea2fa09a560b649b896 Mon Sep 17 00:00:00 2001 From: Liang-Ting Jiang Date: Tue, 12 May 2026 08:14:29 +0000 Subject: [PATCH] feat(code-mode): add io.write helper Expose a host-mediated io.write primitive in Code Mode that writes serializable text or bytes to env://current paths via the core harness. Co-authored-by: Codex --- codex-rs/code-mode/src/description.rs | 1 + codex-rs/code-mode/src/lib.rs | 1 + codex-rs/code-mode/src/runtime/callbacks.rs | 54 +++++++++++ codex-rs/code-mode/src/runtime/globals.rs | 17 ++++ codex-rs/code-mode/src/runtime/mod.rs | 14 +++ codex-rs/code-mode/src/service.rs | 49 ++++++++++ codex-rs/core/src/tools/code_mode/mod.rs | 102 ++++++++++++++++++++ codex-rs/core/tests/suite/code_mode.rs | 37 +++++++ 8 files changed, 275 insertions(+) diff --git a/codex-rs/code-mode/src/description.rs b/codex-rs/code-mode/src/description.rs index 4c5eb6fbdc..7bf7d40fdf 100644 --- a/codex-rs/code-mode/src/description.rs +++ b/codex-rs/code-mode/src/description.rs @@ -30,6 +30,7 @@ const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/co - `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(...)`. +- `io.write(value: string | number[] | TextContent | BlobResourceContents, destination: "env://current/")`: writes text or bytes to the current environment filesystem and returns `{ uri, path, bytes_written }`. - `setTimeout(callback: () => void, delayMs?: number)`: schedules a callback to run later and returns a timeout id. Pending timeouts do not keep `exec` alive by themselves; await an explicit promise if you need to wait for one. - `clearTimeout(timeoutId?: number)`: cancels a timeout created by `setTimeout`. - `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries. diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index bf0ce699cc..e23f6435ce 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -18,6 +18,7 @@ pub use description::render_json_schema_to_typescript; pub use response::DEFAULT_IMAGE_DETAIL; pub use response::FunctionCallOutputContentItem; pub use response::ImageDetail; +pub use runtime::CodeModeIoWrite; pub use runtime::CodeModeNestedToolCall; pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS; pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; diff --git a/codex-rs/code-mode/src/runtime/callbacks.rs b/codex-rs/code-mode/src/runtime/callbacks.rs index a9755f6eb0..4d2e13cca4 100644 --- a/codex-rs/code-mode/src/runtime/callbacks.rs +++ b/codex-rs/code-mode/src/runtime/callbacks.rs @@ -9,6 +9,7 @@ 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; +use serde_json::Value as JsonValue; pub(super) fn tool_callback( scope: &mut v8::PinScope<'_, '_>, @@ -74,6 +75,59 @@ pub(super) fn tool_callback( retval.set(promise.into()); } +pub(super) fn io_write_callback( + scope: &mut v8::PinScope<'_, '_>, + args: v8::FunctionCallbackArguments, + mut retval: v8::ReturnValue, +) { + if args.length() < 2 { + throw_type_error(scope, "io.write expects a value and destination"); + return; + } + + let value = match v8_value_to_json(scope, args.get(0)) { + Ok(Some(value)) => value, + Ok(None) => JsonValue::Null, + Err(error_text) => { + throw_type_error(scope, &error_text); + return; + } + }; + let destination = match args.get(1).to_string(scope) { + Some(destination) => destination.to_rust_string_lossy(scope), + None => { + throw_type_error(scope, "io.write destination must be a string"); + return; + } + }; + if destination.trim().is_empty() { + throw_type_error(scope, "io.write destination must be non-empty"); + return; + } + + let Some(resolver) = v8::PromiseResolver::new(scope) else { + throw_type_error(scope, "failed to create io.write promise"); + return; + }; + let promise = resolver.get_promise(scope); + let resolver = v8::Global::new(scope, resolver); + + let Some(state) = scope.get_slot_mut::() else { + throw_type_error(scope, "runtime state unavailable"); + return; + }; + let id = format!("io-write-{}", 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::IoWrite { + id, + value, + destination, + }); + retval.set(promise.into()); +} + pub(super) fn text_callback( scope: &mut v8::PinScope<'_, '_>, args: v8::FunctionCallbackArguments, diff --git a/codex-rs/code-mode/src/runtime/globals.rs b/codex-rs/code-mode/src/runtime/globals.rs index 2ec6953f09..266e68e7ca 100644 --- a/codex-rs/code-mode/src/runtime/globals.rs +++ b/codex-rs/code-mode/src/runtime/globals.rs @@ -2,6 +2,7 @@ use super::RuntimeState; use super::callbacks::clear_timeout_callback; use super::callbacks::exit_callback; use super::callbacks::image_callback; +use super::callbacks::io_write_callback; use super::callbacks::load_callback; use super::callbacks::notify_callback; use super::callbacks::set_timeout_callback; @@ -19,6 +20,7 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St let tools = build_tools_object(scope)?; let all_tools = build_all_tools_value(scope)?; + let io = build_io_object(scope)?; let clear_timeout = helper_function(scope, "clearTimeout", clear_timeout_callback)?; let set_timeout = helper_function(scope, "setTimeout", set_timeout_callback)?; let text = helper_function(scope, "text", text_callback)?; @@ -30,6 +32,7 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St let exit = helper_function(scope, "exit", exit_callback)?; set_global(scope, global, "tools", tools.into())?; + set_global(scope, global, "io", io.into())?; set_global(scope, global, "ALL_TOOLS", all_tools)?; set_global(scope, global, "clearTimeout", clear_timeout.into())?; set_global(scope, global, "setTimeout", set_timeout.into())?; @@ -43,6 +46,20 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St Ok(()) } +fn build_io_object<'s>( + scope: &mut v8::PinScope<'s, '_>, +) -> Result, String> { + let io = v8::Object::new(scope); + let write_key = v8::String::new(scope, "write") + .ok_or_else(|| "failed to allocate io.write key".to_string())?; + let write = helper_function(scope, "io.write", io_write_callback)?; + if io.set(scope, write_key.into(), write.into()) == Some(true) { + Ok(io) + } else { + Err("failed to set io.write".to_string()) + } +} + fn build_tools_object<'s>( scope: &mut v8::PinScope<'s, '_>, ) -> Result, String> { diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 200a47c989..e85d9ef8ba 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -100,9 +100,18 @@ pub struct CodeModeNestedToolCall { pub input: Option, } +#[derive(Debug)] +pub struct CodeModeIoWrite { + pub cell_id: String, + pub runtime_write_id: String, + pub value: JsonValue, + pub destination: String, +} + #[derive(Debug)] pub(crate) enum TurnMessage { ToolCall(CodeModeNestedToolCall), + IoWrite(CodeModeIoWrite), Notify { cell_id: String, call_id: String, @@ -128,6 +137,11 @@ pub(crate) enum RuntimeEvent { name: ToolName, input: Option, }, + IoWrite { + id: String, + value: JsonValue, + destination: String, + }, Notify { call_id: String, text: String, diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 7326c834e2..ec12e78db0 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -13,6 +13,7 @@ use tokio_util::sync::CancellationToken; use tracing::warn; use crate::FunctionCallOutputContentItem; +use crate::runtime::CodeModeIoWrite; use crate::runtime::CodeModeNestedToolCall; use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS; use crate::runtime::ExecuteRequest; @@ -33,6 +34,12 @@ pub trait CodeModeTurnHost: Send + Sync { ) -> Result; async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>; + + async fn write_file( + &self, + request: CodeModeIoWrite, + cancellation_token: CancellationToken, + ) -> Result; } #[derive(Clone)] @@ -222,6 +229,35 @@ impl CodeModeService { let _ = runtime_tx.send(command); }); } + TurnMessage::IoWrite(request) => { + let host = Arc::clone(&host); + let inner = Arc::clone(&inner); + tokio::spawn(async move { + let cell_id = request.cell_id.clone(); + let runtime_write_id = request.runtime_write_id.clone(); + let response = host.write_file(request, 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: runtime_write_id, + result, + }, + Err(error_text) => RuntimeCommand::ToolError { + id: runtime_write_id, + error_text, + }, + }; + let _ = runtime_tx.send(command); + }); + } } } }); @@ -397,6 +433,19 @@ async fn run_session_control( .send(TurnMessage::ToolCall(tool_call)) .await; } + RuntimeEvent::IoWrite { + id, + value, + destination, + } => { + let request = CodeModeIoWrite { + cell_id: cell_id.clone(), + runtime_write_id: id, + value, + destination, + }; + let _ = inner.turn_message_tx.send(TurnMessage::IoWrite(request)).await; + } RuntimeEvent::Result { stored_values, error_text, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 9ee8e0352a..430dcd1dcc 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -4,9 +4,13 @@ mod response_adapter; mod wait_handler; pub(crate) mod wait_spec; +use std::path::Path; use std::sync::Arc; use std::time::Duration; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_code_mode::CodeModeIoWrite; use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeTurnHost; use codex_code_mode::RuntimeResponse; @@ -158,6 +162,104 @@ impl CodeModeTurnHost for CoreTurnHost { format!("failed to inject exec notify message for cell {cell_id}: no active turn") }) } + + async fn write_file( + &self, + request: CodeModeIoWrite, + _cancellation_token: CancellationToken, + ) -> Result { + write_code_mode_io_value(self.exec.turn.as_ref(), request).await + } +} + +const CURRENT_ENV_URI_PREFIX: &str = "env://current/"; + +async fn write_code_mode_io_value( + turn: &TurnContext, + request: CodeModeIoWrite, +) -> Result { + let destination = resolve_code_mode_io_destination(turn, &request.destination)?; + let bytes = code_mode_io_value_bytes(&request.value)?; + if let Some(parent) = destination.parent() { + tokio::fs::create_dir_all(parent.as_path()) + .await + .map_err(|error| format!("failed to create parent directory for io.write: {error}"))?; + } + tokio::fs::write(destination.as_path(), &bytes) + .await + .map_err(|error| format!("failed to write `{}`: {error}", request.destination))?; + Ok(serde_json::json!({ + "uri": request.destination, + "path": destination.to_string_lossy(), + "bytes_written": bytes.len(), + })) +} + +fn resolve_code_mode_io_destination( + turn: &TurnContext, + destination: &str, +) -> Result { + let relative_path = destination + .strip_prefix(CURRENT_ENV_URI_PREFIX) + .ok_or_else(|| { + format!("io.write destination must start with `{CURRENT_ENV_URI_PREFIX}`") + })?; + if relative_path.trim().is_empty() { + return Err("io.write destination path must be non-empty".to_string()); + } + if Path::new(relative_path).is_absolute() { + return Err("io.write destination path must be relative".to_string()); + } + let resolved = turn.cwd.join(relative_path); + if !resolved.as_path().starts_with(turn.cwd.as_path()) { + return Err(format!( + "io.write destination `{destination}` escapes the current environment" + )); + } + Ok(resolved) +} + +fn code_mode_io_value_bytes(value: &JsonValue) -> Result, String> { + match value { + JsonValue::String(text) => Ok(text.as_bytes().to_vec()), + JsonValue::Array(items) if items.iter().all(JsonValue::is_number) => items + .iter() + .map(|item| { + let Some(byte) = item.as_u64().filter(|byte| *byte <= u8::MAX as u64) else { + return Err( + "io.write byte arrays must contain integers from 0 to 255".to_string() + ); + }; + Ok(byte as u8) + }) + .collect(), + JsonValue::Object(object) => { + if let Some(text) = object.get("text").and_then(JsonValue::as_str) { + return Ok(text.as_bytes().to_vec()); + } + if let Some(resource) = object.get("resource") { + return code_mode_io_value_bytes(resource); + } + if let Some(encoded) = object + .get("blob") + .or_else(|| object.get("data_base64")) + .or_else(|| object.get("base64")) + .and_then(JsonValue::as_str) + { + return BASE64_STANDARD + .decode(encoded) + .map_err(|error| format!("io.write received invalid base64 data: {error}")); + } + Err( + "io.write value must be a string, byte array, MCP text block, or base64 blob block" + .to_string(), + ) + } + _ => Err( + "io.write value must be a string, byte array, MCP text block, or base64 blob block" + .to_string(), + ), + } } pub(super) async fn handle_runtime_response( diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 3bcb37e7b2..ac49afbda3 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -2170,6 +2170,43 @@ async fn code_mode_can_apply_patch_via_nested_tool() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_io_write_writes_text_to_current_env() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let code = r#" +const result = await io.write("hello from io", "env://current/nested/io-output.txt"); +text(JSON.stringify(result)); +"#; + + let (test, second_mock) = + run_code_mode_turn(&server, "use io.write to write a file", code, false).await?; + + let req = second_mock.single_request(); + let (output, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "exec io.write failed unexpectedly: {output}" + ); + let parsed: Value = serde_json::from_str(&output)?; + assert_eq!( + parsed.get("uri").and_then(Value::as_str), + Some("env://current/nested/io-output.txt") + ); + assert_eq!( + parsed.get("bytes_written").and_then(Value::as_u64), + Some(13) + ); + assert_eq!( + fs::read_to_string(test.cwd_path().join("nested/io-output.txt"))?, + "hello from io" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_print_structured_mcp_tool_result_fields() -> Result<()> { skip_if_no_network!(Ok(()));