From 28a264fbc766a59f2b550b8318f88e4a4b8dffe1 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Mon, 5 Oct 2026 16:49:17 +0000 Subject: [PATCH] Add promise settlement streaming helpers to code mode (#51126) ## What changed Expose two JavaScript globals for consuming results as inputs settle: - `as_settled(promises)` returns an async iterator of fulfilled and rejected records in settlement order, preserving queued results while the consumer is busy. - `stream_settled(promises, emit)` delivers those records to a callback, awaiting each invocation and propagating callback errors. Records include an `index`, a `status`, and either a `value` or a `reason`. Iterable inputs use positional indices; `Map` inputs retain their keys. Both helpers accept promises, thenables, and plain values. ## Testing Add runtime coverage for ordering, input types, object identity, callback behavior, early iteration exit, and termination. Add stdio and gRPC host tests for partial results across waits, plus a scenario snapshot verifying fulfilled and rejected records reach the model. GitOrigin-RevId: cddc9f3679ae04f2421e154caa0417afb15f550d --- codex-rs/code-mode-host/tests/settled.rs | 181 ++++++++++++ codex-rs/code-mode-runtime/BUILD.bazel | 1 + .../code-mode-runtime/src/runtime/globals.rs | 6 + .../code-mode-runtime/src/runtime/settled.js | 67 +++++ .../code-mode-runtime/src/service_tests.rs | 2 + codex-rs/code-mode-runtime/tests/settled.rs | 257 ++++++++++++++++++ codex-rs/core/tests/suite/scenarios.rs | 3 + ...enarios_code_mode_settled_helpers_tests.rs | 64 +++++ ...ed_helpers__code_mode_settled_helpers.snap | 43 +++ 9 files changed, 624 insertions(+) create mode 100644 codex-rs/code-mode-host/tests/settled.rs create mode 100644 codex-rs/code-mode-runtime/src/runtime/settled.js create mode 100644 codex-rs/code-mode-runtime/tests/settled.rs create mode 100644 codex-rs/core/tests/suite/scenarios_code_mode_settled_helpers_tests.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_settled_helpers__code_mode_settled_helpers.snap diff --git a/codex-rs/code-mode-host/tests/settled.rs b/codex-rs/code-mode-host/tests/settled.rs new file mode 100644 index 0000000000..fff96761ff --- /dev/null +++ b/codex-rs/code-mode-host/tests/settled.rs @@ -0,0 +1,181 @@ +//! Verifies partial settlement output survives waits on both host transports. + +use std::sync::Arc; +use std::time::Duration; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeToolKind; +use codex_code_mode::ExecuteRequest; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::GrpcCodeModeSessionProvider; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; +use codex_code_mode::RuntimeResponse; +use codex_code_mode::ToolDefinition; +use codex_code_mode::ToolInvocationFuture; +use codex_code_mode::WaitOutcome; +use codex_code_mode::WaitRequest; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[path = "support/host.rs"] +mod host; + +#[derive(Default)] +struct SlowDelegate { + release: Notify, +} + +impl CodeModeSessionDelegate for SlowDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + tokio::select! { + _ = self.release.notified() => Ok(json!("slow")), + _ = cancellation.cancelled() => Err("cancelled".to_string()), + } + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +#[allow(clippy::expect_used)] +async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) { + for consumer in [ + "for await (const result of as_settled(promises)) text(result);", + "await stream_settled(promises, text);", + ] { + let delegate = Arc::new(SlowDelegate::default()); + let session = provider.create_session().await.expect("create session"); + let started = session + .execute(ExecuteRequest { + tool_call_id: "settled".to_string(), + enabled_tools: vec![ToolDefinition { + name: "slow".to_string(), + tool_name: ToolName::plain("slow"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + input_schema_max_bytes: None, + output_schema: None, + }], + source: format!( + r#"const promises = [tools.slow({{}}), Promise.resolve("fast"), Promise.reject("failed")]; +{consumer}"# + ), + yield_time_ms: Some(/*value*/ 10), + max_output_tokens: None, + }, delegate.clone(), /*preempt*/ None) + .await + .expect("start cell"); + let cell_id = started.cell_id.clone(); + let mut response = started.initial_response().await.expect("initial response"); + let mut partial = Vec::new(); + loop { + let RuntimeResponse::Yielded { content_items, .. } = response else { + panic!("cell completed before the slow tool was released: {response:?}"); + }; + partial.extend(content_items.into_iter().map(|item| { + let FunctionCallOutputContentItem::InputText { text } = item else { + panic!("unexpected output: {item:?}"); + }; + serde_json::from_str::(&text).expect("JSON output") + })); + if partial.len() >= 2 { + break; + } + let WaitOutcome::LiveCell(next) = session + .wait( + WaitRequest { + cell_id: cell_id.clone(), + yield_time_ms: 10, + }, + /*preempt*/ None, + ) + .await + .expect("wait for partial results") + else { + panic!("cell disappeared while the slow tool was pending"); + }; + response = next; + } + assert_eq!( + partial, + [ + json!({"index": 1, "status": "fulfilled", "value": "fast"}), + json!({"index": 2, "status": "rejected", "reason": "failed"}), + ] + ); + + delegate.release.notify_one(); + let completed = session + .wait( + WaitRequest { + cell_id: cell_id.clone(), + yield_time_ms: 5_000, + }, + /*preempt*/ None, + ) + .await + .expect("wait for completion"); + assert_eq!( + completed, + WaitOutcome::LiveCell(RuntimeResponse::Result { + code_mode_host_duration: completed.code_mode_host_duration(), + cell_id, + content_items: vec![FunctionCallOutputContentItem::InputText { + text: json!({"index": 0, "status": "fulfilled", "value": "slow"}).to_string(), + }], + error_text: None, + }) + ); + session.shutdown().await.expect("shutdown session"); + } +} + +#[tokio::test] +async fn settled_results_stream_over_stdio() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + tokio::time::timeout( + Duration::from_secs(/*secs*/ 20), + assert_partial_results(&provider), + ) + .await + .expect("settlement streaming timed out"); +} + +#[tokio::test] +async fn settled_results_stream_over_grpc() { + let host = host::HostHarness::start("grpc://127.0.0.1:0") + .await + .expect("start gRPC host"); + let provider = GrpcCodeModeSessionProvider::new(host.endpoint); + tokio::time::timeout( + Duration::from_secs(/*secs*/ 20), + assert_partial_results(&provider), + ) + .await + .expect("settlement streaming timed out"); +} diff --git a/codex-rs/code-mode-runtime/BUILD.bazel b/codex-rs/code-mode-runtime/BUILD.bazel index b3903eb77d..22de889d3a 100644 --- a/codex-rs/code-mode-runtime/BUILD.bazel +++ b/codex-rs/code-mode-runtime/BUILD.bazel @@ -2,5 +2,6 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "code-mode-runtime", + compile_data = ["src/runtime/settled.js"], crate_name = "codex_code_mode_runtime", ) diff --git a/codex-rs/code-mode-runtime/src/runtime/globals.rs b/codex-rs/code-mode-runtime/src/runtime/globals.rs index 66601deb49..4ce226f68c 100644 --- a/codex-rs/code-mode-runtime/src/runtime/globals.rs +++ b/codex-rs/code-mode-runtime/src/runtime/globals.rs @@ -46,6 +46,12 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St set_global(scope, global, "notify", notify.into())?; set_global(scope, global, "yield_control", yield_control.into())?; set_global(scope, global, "exit", exit.into())?; + + let source = v8::String::new(scope, include_str!("settled.js")) + .ok_or_else(|| "failed to allocate settlement helpers".to_string())?; + v8::Script::compile(scope, source, /*origin*/ None) + .and_then(|script| script.run(scope)) + .ok_or_else(|| "failed to initialize settlement helpers".to_string())?; Ok(()) } diff --git a/codex-rs/code-mode-runtime/src/runtime/settled.js b/codex-rs/code-mode-runtime/src/runtime/settled.js new file mode 100644 index 0000000000..5c4a837d62 --- /dev/null +++ b/codex-rs/code-mode-runtime/src/runtime/settled.js @@ -0,0 +1,67 @@ +// Runtime-owned promise combinators; all state stays in the cell's V8 heap. +// Observe each input once and queue settlements while the consumer is busy. +// Repeated races can lose that order and attach extra reactions to pending inputs. +(() => { + function as_settled(promises) { + let ready = []; + let pending = 0; + let wake; + let closed = false; + + function settle(result) { + pending--; + if (closed) return; + ready.push(result); + if (wake) { + const resolve = wake; + wake = undefined; + resolve(); + } + } + + const keyed = promises instanceof Map; + let position = 0; + try { + for (const entry of promises) { + const [index, promise] = keyed ? entry : [position++, entry]; + pending++; + Promise.resolve(promise).then( + (value) => settle({ index, status: "fulfilled", value }), + (reason) => settle({ index, status: "rejected", reason }), + ); + } + } catch (error) { + closed = true; + throw error; + } + + return (async function* () { + try { + while (pending > 0 || ready.length > 0) { + if (ready.length === 0) { + await new Promise((resolve) => { + wake = resolve; + }); + } + const batch = ready; + ready = []; + yield* batch; + } + } finally { + closed = true; + ready = []; + } + })(); + } + + async function stream_settled(promises, emit) { + if (typeof emit !== "function") { + throw new TypeError("stream_settled requires a callback"); + } + for await (const result of as_settled(promises)) { + await emit(result); + } + } + + Object.assign(globalThis, { as_settled, stream_settled }); +})(); diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index 730b9ec707..62d15d6648 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -1127,6 +1127,7 @@ async fn global_scope_contains_only_allowed_items() { "WeakSet", "__codexContentItems", "add_content", + "as_settled", "audio", "decodeURI", "decodeURIComponent", @@ -1146,6 +1147,7 @@ async fn global_scope_contains_only_allowed_items() { "parseInt", "setTimeout", "store", + "stream_settled", "text", "tools", "undefined", diff --git a/codex-rs/code-mode-runtime/tests/settled.rs b/codex-rs/code-mode-runtime/tests/settled.rs new file mode 100644 index 0000000000..d9ed6a5c7c --- /dev/null +++ b/codex-rs/code-mode-runtime/tests/settled.rs @@ -0,0 +1,257 @@ +//! Exercises settlement ordering and callback behavior through V8 cell execution. + +use std::sync::Arc; +use std::time::Duration; + +use codex_code_mode_runtime::ExecuteRequest; +use codex_code_mode_runtime::ExecuteToPendingOutcome; +use codex_code_mode_runtime::FunctionCallOutputContentItem; +use codex_code_mode_runtime::InProcessCodeModeSession; +use codex_code_mode_runtime::NoopCodeModeSessionDelegate; +use codex_code_mode_runtime::RuntimeResponse; +use codex_code_mode_runtime::WaitOutcome; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; + +#[allow(clippy::expect_used)] +async fn assert_output(source: &str, expected: Value) { + let session = InProcessCodeModeSession::new(); + let started = session + .execute( + ExecuteRequest { + tool_call_id: "settled".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: Some(/*value*/ 5_000), + max_output_tokens: None, + }, + Arc::new(NoopCodeModeSessionDelegate), + /*preempt*/ None, + ) + .await + .expect("start cell"); + let response = started.initial_response().await.expect("execute cell"); + let RuntimeResponse::Result { + content_items, + error_text: None, + .. + } = response + else { + panic!("unexpected cell response: {response:?}"); + }; + let actual: Vec = content_items + .into_iter() + .map(|item| { + let FunctionCallOutputContentItem::InputText { text } = item else { + panic!("unexpected output: {item:?}"); + }; + serde_json::from_str(&text).expect("JSON output") + }) + .collect(); + assert_eq!(actual, vec![expected]); +} + +#[tokio::test] +async fn preserves_settlement_order_before_and_during_iteration() { + assert_output( + r#" +const resolve = []; +const results = as_settled(Array.from({length: 4}, (_, index) => + new Promise(done => { resolve[index] = done; }))); +resolve[2]("two"); +await Promise.resolve(); +resolve[0]("zero"); +await Promise.resolve(); +const seen = []; +for await (const result of results) { + seen.push(result); + if (seen.length === 1) { + resolve[3]("three"); + resolve[1]("one"); + await Promise.resolve(); + } +} +text(seen); +"#, + json!([ + {"index": 2, "status": "fulfilled", "value": "two"}, + {"index": 0, "status": "fulfilled", "value": "zero"}, + {"index": 3, "status": "fulfilled", "value": "three"}, + {"index": 1, "status": "fulfilled", "value": "one"}, + ]), + ) + .await; +} + +#[tokio::test] +async fn accepts_iterables_values_thenables_and_duplicate_promises() { + assert_output( + r#" +const shared = Promise.resolve("shared"); +function* inputs() { + yield 42; + yield {then(resolve) { resolve("thenable"); }}; + yield shared; + yield shared; + yield Promise.reject({message: "failed"}); + yield {get then() { throw "getter failed"; }}; +} +const seen = []; +for await (const result of as_settled(inputs())) seen.push(result); +text(seen.sort((a, b) => a.index - b.index)); +"#, + json!([ + {"index": 0, "status": "fulfilled", "value": 42}, + {"index": 1, "status": "fulfilled", "value": "thenable"}, + {"index": 2, "status": "fulfilled", "value": "shared"}, + {"index": 3, "status": "fulfilled", "value": "shared"}, + {"index": 4, "status": "rejected", "reason": {"message": "failed"}}, + {"index": 5, "status": "rejected", "reason": "getter failed"}, + ]), + ) + .await; +} + +#[tokio::test] +async fn preserves_map_keys_and_completes_empty_inputs() { + assert_output( + r#" +const seen = []; +await stream_settled([], result => seen.push(result)); +for await (const result of as_settled(new Map())) seen.push(result); +await stream_settled(new Map([ + ["success", 7], + ["failure", Promise.reject("failed")], +]), result => seen.push(result)); +text(seen); +"#, + json!([ + {"index": "success", "status": "fulfilled", "value": 7}, + {"index": "failure", "status": "rejected", "reason": "failed"}, + ]), + ) + .await; +} + +#[tokio::test] +async fn awaits_each_callback_including_rejected_inputs() { + assert_output( + r#" +const seen = []; +await stream_settled([Promise.reject("failed"), 2], async result => { + seen.push(result); + await new Promise(resolve => setTimeout(resolve, 1)); + seen.push(["callback finished", result.index]); +}); +seen.push("stream finished"); +text(seen); +"#, + json!([ + {"index": 0, "status": "rejected", "reason": "failed"}, + ["callback finished", 0], + {"index": 1, "status": "fulfilled", "value": 2}, + ["callback finished", 1], + "stream finished", + ]), + ) + .await; +} + +#[tokio::test] +async fn propagates_callback_errors_and_closes_early_iterations() { + assert_output( + r#" +const seen = []; +const failure = new Error("callback failed"); +try { + await stream_settled([1, 2], async result => { + seen.push(result); + await new Promise(resolve => setTimeout(resolve, 1)); + throw failure; + }); +} catch (error) { + seen.push(error === failure); +} +const results = as_settled([3, new Promise(() => {}), Promise.reject("unconsumed")]); +for await (const result of results) { + seen.push(result); + break; +} +seen.push(await results.next()); +text(seen); +"#, + json!([ + {"index": 0, "status": "fulfilled", "value": 1}, + true, + {"index": 0, "status": "fulfilled", "value": 3}, + {"done": true}, + ]), + ) + .await; +} + +#[tokio::test] +async fn preserves_identity_of_keys_values_and_rejections() { + assert_output( + r#" +const key = {}; +const value = {method() { return 7; }}; +value.self = value; +const failure = new Error("failed"); +const seen = []; +await stream_settled(new Map([ + [key, Promise.resolve(value)], + [failure, Promise.reject(failure)], +]), result => { + seen.push(result.status === "fulfilled" + ? result.index === key && result.value === value && result.value.self === value + && result.value.method() === 7 + : result.index === failure && result.reason === failure); +}); +text(seen); +"#, + json!([true, true]), + ) + .await; +} + +#[tokio::test] +async fn terminates_helpers_waiting_for_unsettled_inputs() { + let session = InProcessCodeModeSession::new(); + for source in [ + "for await (const result of as_settled([new Promise(() => {})])) text(result);", + "await stream_settled([new Promise(() => {})], text);", + ] { + let response = session + .execute_to_pending( + ExecuteRequest { + tool_call_id: "settled".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + Arc::new(NoopCodeModeSessionDelegate), + ) + .await + .expect("execute to pending"); + let ExecuteToPendingOutcome::Pending { cell_id, .. } = response else { + panic!("expected pending helper: {response:?}"); + }; + assert_eq!( + tokio::time::timeout( + Duration::from_secs(/*secs*/ 5), + session.terminate(cell_id.clone()) + ) + .await + .expect("termination timed out") + .expect("terminate cell"), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id, + content_items: Vec::new(), + code_mode_host_duration: None, + }), + ); + } +} diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 7b69faff01..765b75f83f 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -92,6 +92,9 @@ const ONE_PIXEL_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ #[path = "scenarios_incremental_tools.rs"] mod incremental_tools; +#[path = "scenarios_code_mode_settled_helpers_tests.rs"] +mod code_mode_settled_helpers; + #[path = "scenarios_strict_3p_cache.rs"] mod strict_3p_cache; diff --git a/codex-rs/core/tests/suite/scenarios_code_mode_settled_helpers_tests.rs b/codex-rs/core/tests/suite/scenarios_code_mode_settled_helpers_tests.rs new file mode 100644 index 0000000000..bedd3c8d96 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_code_mode_settled_helpers_tests.rs @@ -0,0 +1,64 @@ +//! Snapshots settlement records delivered to the model after code-mode execution. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_settled_helpers_report_successes_and_failures() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_model("gpt-6-astra") + .with_config(|config| { + configure_scenario_catalog(config); + config.workspace_roots = vec![config.cwd.clone()]; + for feature in [ + Feature::CodeMode, + Feature::CodeModeOnly, + Feature::CodeModeHost, + ] { + config.features.enable(feature).expect("enable code mode"); + } + }) + .build_with_auto_env(&server) + .await?; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("exec-response"), + ev_custom_tool_call( + "exec-call", + "exec", + r#"for await (const result of as_settled([Promise.resolve("ready")])) text(result); +await stream_settled(new Map([["failed", Promise.reject("unavailable")], ["completed", "done"]]), text);"#, + ), + ev_completed("exec-response"), + ]), + sse(vec![ + ev_assistant_message("final", "Two inputs succeeded; one failed with unavailable."), + ev_completed("final-response"), + ]), + ], + ) + .await; + test.submit_turn( + "Collect each completed result, including failures, using the settlement helpers.", + ) + .await?; + let requests = mock.requests(); + assert_eq!(requests.len(), 2); + let output = requests[1].custom_tool_call_output("exec-call"); + let output = output.to_string(); + assert!(output.contains("fulfilled"), "{output}"); + assert!(output.contains("rejected"), "{output}"); + insta::assert_snapshot!( + "code_mode_settled_helpers", + context_snapshot::format_request_history_snapshot( + "Settlement helpers deliver fulfilled and rejected records with input indices and map keys.", + &requests, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_settled_helpers__code_mode_settled_helpers.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_settled_helpers__code_mode_settled_helpers.snap new file mode 100644 index 0000000000..d33a16adea --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_settled_helpers__code_mode_settled_helpers.snap @@ -0,0 +1,43 @@ +--- +source: core/tests/suite/scenarios_code_mode_settled_helpers_tests.rs +expression: "context_snapshot::format_request_history_snapshot(\"Settlement helpers deliver fulfilled and rejected records with input indices and map keys.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Settlement helpers deliver fulfilled and rejected records with input indices and map keys. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (3; hash=2c5a5b48dbcfce9b): + - namespace/functions; hash=fada4a476a3754dd + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=b263e3c84ad07c37 + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent +01:message/developer: + +02:message/developer[2]: + [01] + [02] +03:message/developer: + +04:message/developer: + +05:message/user: + +06:message/user: + Collect each completed result, including failures, using the settlement helpers. +-- request 2 (turn) -- +07:custom_tool_call/exec:for await (const result of as_settled([Promise.resolve("ready")])) text(result); + await stream_settled(new Map([["failed", Promise.reject("unavailable")], ["completed", "done"]]), text); +08:custom_tool_call_output:Script completed + Wall time seconds + Output: + | {"index":0,"status":"fulfilled","value":"ready"} | {"index":"failed","status":"rejected","reason":"unavailable"} | {"index":"completed","status":"fulfilled","value":"done"}