From 165299214855c2e48facf9bb3bc3c0fc277ee639 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Sat, 10 Oct 2026 02:04:15 +0000 Subject: [PATCH] Make code-mode `exit()` stop the entire cell (#52748) ## Why A catchable exception lets JavaScript continue after `exit()`, including in `catch`, `finally`, and queued promise callbacks. Exiting should finish the cell successfully without further output or state changes. ## What changed - Terminate V8 execution when `exit()` is called and complete the cell even if its module promise remains pending. - Guard host callbacks against effects after exit, including when getters or value conversion request exit. - Preserve stored writes made before exit and suppress errors when exit occurs during exception formatting. ## Testing Add regression tests for synchronous and asynchronous exit, timers, parallel promises, value conversion, exception getters, and resuming after a yield and tool response. Verify that only prior output and stored writes survive, and separately exercise host guards while JavaScript continues running. GitOrigin-RevId: e9487ba65ef0358612c5009bd26953b786a12798 --- .../src/runtime/callbacks.rs | 58 +++-- .../src/runtime/callbacks_tests.rs | 194 +++++++++++++++++ codex-rs/code-mode-runtime/src/runtime/mod.rs | 39 ++-- .../src/runtime/module_loader.rs | 44 +--- .../code-mode-runtime/src/runtime/timers.rs | 12 +- .../code-mode-runtime/src/runtime/value.rs | 10 + .../code-mode-runtime/src/service_tests.rs | 198 ++++++++++++++++++ 7 files changed, 475 insertions(+), 80 deletions(-) create mode 100644 codex-rs/code-mode-runtime/src/runtime/callbacks_tests.rs diff --git a/codex-rs/code-mode-runtime/src/runtime/callbacks.rs b/codex-rs/code-mode-runtime/src/runtime/callbacks.rs index 7d18ece900..36bd5f3dec 100644 --- a/codex-rs/code-mode-runtime/src/runtime/callbacks.rs +++ b/codex-rs/code-mode-runtime/src/runtime/callbacks.rs @@ -1,7 +1,6 @@ use codex_code_mode_protocol::FunctionCallOutputContentItem; use std::sync::Arc; -use super::EXIT_SENTINEL; use super::RuntimeEvent; use super::RuntimeState; use super::timers; @@ -37,18 +36,15 @@ pub(super) fn tool_callback( } }; - 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 (tool_name, tool_kind) = { let Some(state) = scope.get_slot::() else { throw_type_error(scope, "runtime state unavailable"); return; }; + // Input serialization can call exit() from a getter or toJSON. + if state.exit_requested { + return; + } let Some(tool) = state.enabled_tools.get(tool_index) else { throw_type_error(scope, "tool callback data is out of range"); return; @@ -56,6 +52,13 @@ pub(super) fn tool_callback( (tool.tool_name.clone(), tool.kind) }; + 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::() else { throw_type_error(scope, "runtime state unavailable"); return; @@ -90,7 +93,9 @@ pub(super) fn text_callback( return; } }; - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::ContentItem( FunctionCallOutputContentItem::InputText { text }, )); @@ -112,7 +117,9 @@ pub(super) fn audio_callback( Ok(audio_item) => audio_item, Err(()) => return, }; - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::ContentItem(audio_item)); } retval.set(v8::undefined(scope).into()); @@ -145,7 +152,9 @@ pub(super) fn image_callback( Ok(image_item) => image_item, Err(()) => return, }; - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::ContentItem(image_item)); } retval.set(v8::undefined(scope).into()); @@ -172,7 +181,9 @@ pub(super) fn generated_image_callback( Ok(image_item) => image_item, Err(()) => return, }; - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::ContentItem(image_item)); if let Some(text) = output_hint { let _ = state.event_tx.send(RuntimeEvent::ContentItem( @@ -230,7 +241,9 @@ pub(super) fn store_callback( return; } }; - if let Some(state) = scope.get_slot_mut::() { + if let Some(state) = scope.get_slot_mut::() + && !state.exit_requested + { let serialized = Arc::new(serialized); state .stored_values @@ -287,7 +300,9 @@ pub(super) fn notify_callback( throw_type_error(scope, "notify expects non-empty text"); return; } - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::Notify { call_id: state.tool_call_id.clone(), text, @@ -302,7 +317,8 @@ pub(super) fn set_timeout_callback( mut retval: v8::ReturnValue, ) { let timeout_id = match timers::schedule_timeout(scope, args) { - Ok(timeout_id) => timeout_id, + Ok(Some(timeout_id)) => timeout_id, + Ok(None) => return, Err(error_text) => { throw_type_error(scope, &error_text); return; @@ -330,7 +346,9 @@ pub(super) fn yield_control_callback( _args: v8::FunctionCallbackArguments, _retval: v8::ReturnValue, ) { - if let Some(state) = scope.get_slot::() { + if let Some(state) = scope.get_slot::() + && !state.exit_requested + { let _ = state.event_tx.send(RuntimeEvent::YieldRequested); } } @@ -343,7 +361,9 @@ pub(super) fn exit_callback( if let Some(state) = scope.get_slot_mut::() { state.exit_requested = true; } - if let Some(error) = v8::String::new(scope, EXIT_SENTINEL) { - scope.throw_exception(error.into()); - } + scope.terminate_execution(); } + +#[cfg(test)] +#[path = "callbacks_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-runtime/src/runtime/callbacks_tests.rs b/codex-rs/code-mode-runtime/src/runtime/callbacks_tests.rs new file mode 100644 index 0000000000..aae92e04ed --- /dev/null +++ b/codex-rs/code-mode-runtime/src/runtime/callbacks_tests.rs @@ -0,0 +1,194 @@ +//! Exercises the host guard separately from V8's interrupt mechanism. Simulate +//! V8 continuing to execute after exit was requested, including during getters. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::mpsc as std_mpsc; + +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::EnabledToolMetadata; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::sync::mpsc; + +use super::super::RuntimeEvent; +use super::super::RuntimeState; +use super::super::ToolCallbackMetadata; +use super::super::globals; +use crate::FunctionCallOutputContentItem; +use crate::v8_init::ensure_v8_initialized; + +// Deliberately leave JS running: the test must verify host-side admission even +// when V8 has not processed its interrupt. Returning a value exercises coercion. +fn mark_exit_requested( + scope: &mut v8::PinScope<'_, '_>, + args: v8::FunctionCallbackArguments, + mut retval: v8::ReturnValue, +) { + scope + .get_slot_mut::() + .expect("runtime state") + .exit_requested = true; + retval.set(args.get(0)); +} + +#[tokio::test] +async fn host_callbacks_admit_no_new_effects_after_exit_is_requested() { + for (label, source) in [ + ( + "ordinary calls", + r#" +markExit(); +text("still working"); +image("data:image/png;base64,AAA"); +audio("data:audio/mpeg;base64,YXVkaW8="); +generatedImage({image_url: "data:image/png;base64,AAA", output_hint: "late"}); +notify("late"); store("phase", "late"); +tools.echo("late"); yield_control(); +clearTimeout(earlyTimer); setTimeout(() => {}, 60000); +"#, + ), + ( + "text conversion", + r#"text({toJSON() { return markExit("late"); }});"#, + ), + ( + "notify conversion", + r#"notify({toJSON() { return markExit("late"); }});"#, + ), + ( + "store conversion", + r#"store("phase", {toJSON() { return markExit("late"); }});"#, + ), + ( + "store key", + r#"store({toString() { return markExit("phase"); }}, "late");"#, + ), + ( + "tool conversion", + r#"tools.echo({toJSON() { return markExit("late"); }});"#, + ), + ( + "image getter", + r#"image({get image_url() { return markExit("data:image/png;base64,AAA"); }});"#, + ), + ( + "audio getter", + r#"audio({get audio_url() { return markExit("data:audio/mpeg;base64,YXVkaW8="); }});"#, + ), + ( + "generated image hint", + r#"generatedImage({image_url: "data:image/png;base64,AAA", get output_hint() {return markExit("late"); }});"#, + ), + ( + "timer delay", + r#"setTimeout(() => {}, {valueOf() {return markExit(60000); }});"#, + ), + ( + "clear timer id", + r#"clearTimeout({valueOf() {return markExit(earlyTimer); }});"#, + ), + ] { + ensure_v8_initialized().expect("initialize V8"); + let isolate = &mut v8::Isolate::new(v8::CreateParams::default()); + v8::scope!(let scope, isolate); + let context = v8::Context::new(scope, Default::default()); + let scope = &mut v8::ContextScope::new(scope, context); + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let (runtime_command_tx, _runtime_command_rx) = std_mpsc::channel(); + let tool_name = ToolName::plain("echo"); + globals::install_globals( + scope, + &[EnabledToolMetadata { + global_name: "echo".to_string(), + tool_name: tool_name.clone(), + description: String::new(), + kind: CodeModeToolKind::Function, + }], + ) + .expect("install globals"); + scope.set_slot(RuntimeState { + event_tx, + pending_tool_calls: HashMap::new(), + pending_timeouts: HashMap::new(), + stored_values: HashMap::new(), + stored_value_writes: HashMap::new(), + enabled_tools: vec![ToolCallbackMetadata { + tool_name, + kind: CodeModeToolKind::Function, + }], + next_tool_call_id: 1, + next_timeout_id: 1, + tool_call_id: "cell".to_string(), + runtime_command_tx, + exit_requested: false, + }); + + let marker = v8::Function::new(scope, mark_exit_requested).expect("exit marker"); + let key = v8::String::new(scope, "markExit").expect("marker name"); + context.global(scope).set(scope, key.into(), marker.into()); + let program = v8::String::new( + scope, + &format!( + r#" +text("before"); store("phase", "before"); notify("before"); +tools.echo("before"); +const earlyTimer = setTimeout(() => text("early timer"), 60000); +{source} +"# + ), + ) + .expect("test program"); + let tc = std::pin::pin!(v8::TryCatch::new(scope)); + let tc = tc.init(); + let result = + v8::Script::compile(&tc, program, /*origin*/ None).and_then(|script| script.run(&tc)); + assert!( + result.is_some(), + "{label}: {:?}", + tc.exception().map(|e| e.to_rust_string_lossy(&tc)) + ); + + assert!( + matches!(event_rx.try_recv(), Ok(RuntimeEvent::ContentItem( + FunctionCallOutputContentItem::InputText { text } + )) if text == "before"), + "{label}: pre-exit output" + ); + assert!( + matches!(event_rx.try_recv(), Ok(RuntimeEvent::Notify { text, .. }) if text == "before"), + "{label}: pre-exit notification" + ); + assert!( + matches!(event_rx.try_recv(), Ok(RuntimeEvent::ToolCall { id, input, .. }) + if id == "tool-1" && input == Some(json!("before"))), + "{label}: pre-exit tool" + ); + assert!( + matches!(event_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)), + "{label}: emitted an effect after exit" + ); + + let state = tc.get_slot::().expect("runtime state"); + assert_eq!( + state.stored_value_writes, + HashMap::from([("phase".to_string(), Arc::new(json!("before")))]), + "{label}: stored writes" + ); + assert_eq!( + state.stored_values, state.stored_value_writes, + "{label}: cell storage" + ); + assert_eq!(state.next_tool_call_id, 2, "{label}: allocated tool ID"); + assert_eq!(state.next_timeout_id, 2, "{label}: allocated timeout ID"); + assert!( + state.pending_tool_calls.contains_key("tool-1"), + "{label}: lost prior tool" + ); + assert!( + state.pending_timeouts.contains_key(&1), + "{label}: lost prior timer" + ); + } +} diff --git a/codex-rs/code-mode-runtime/src/runtime/mod.rs b/codex-rs/code-mode-runtime/src/runtime/mod.rs index 7a3904fbce..6e00316bbd 100644 --- a/codex-rs/code-mode-runtime/src/runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/runtime/mod.rs @@ -24,8 +24,6 @@ use tokio::sync::mpsc; use crate::TaskFailureHandler; use crate::v8_init::ensure_v8_initialized; -const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; - #[derive(Debug)] pub(crate) enum RuntimeCommand { ToolResponse { id: String, result: JsonValue }, @@ -174,10 +172,7 @@ pub(super) struct RuntimeState { pub(super) enum CompletionState { Pending, - Completed { - stored_value_writes: HashMap>, - error_text: Option, - }, + Completed { error_text: Option }, } fn run_runtime( @@ -231,17 +226,14 @@ fn run_runtime( 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)); + send_scope_result(scope, &event_tx, Some(error_text)); return; } }; match module_loader::completion_state(scope, pending_promise.as_ref()) { - CompletionState::Completed { - stored_value_writes, - error_text, - } => { - send_result(&event_tx, stored_value_writes, error_text); + CompletionState::Completed { error_text } => { + send_scope_result(scope, &event_tx, error_text); return; } CompletionState::Pending => {} @@ -257,7 +249,7 @@ fn run_runtime( if let Err(error_text) = module_loader::resolve_tool_response(scope, &id, Ok(result)) { - capture_scope_send_error(scope, &event_tx, Some(error_text)); + send_scope_result(scope, &event_tx, Some(error_text)); return; } } @@ -265,13 +257,13 @@ fn run_runtime( 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)); + send_scope_result(scope, &event_tx, Some(runtime_error)); return; } } RuntimeCommand::TimeoutFired { id } => { if let Err(runtime_error) = timers::invoke_timeout_callback(scope, id) { - capture_scope_send_error(scope, &event_tx, Some(runtime_error)); + send_scope_result(scope, &event_tx, Some(runtime_error)); return; } } @@ -280,11 +272,8 @@ fn run_runtime( scope.perform_microtask_checkpoint(); match module_loader::completion_state(scope, pending_promise.as_ref()) { - CompletionState::Completed { - stored_value_writes, - error_text, - } => { - send_result(&event_tx, stored_value_writes, error_text); + CompletionState::Completed { error_text } => { + send_scope_result(scope, &event_tx, error_text); return; } CompletionState::Pending => {} @@ -325,16 +314,18 @@ fn next_runtime_command( } } -fn capture_scope_send_error( +fn send_scope_result( scope: &mut v8::PinScope<'_, '_>, event_tx: &mpsc::UnboundedSender, error_text: Option, ) { - let stored_value_writes = scope - .get_slot::() + // Error rendering can itself run JS getters, including exit(). Capture + // the result only after every call into the exiting isolate has returned. + let state = scope.get_slot::(); + let stored_value_writes = state .map(|state| state.stored_value_writes.clone()) .unwrap_or_default(); - + let error_text = error_text.filter(|_| !state.is_some_and(|state| state.exit_requested)); send_result(event_tx, stored_value_writes, error_text); } diff --git a/codex-rs/code-mode-runtime/src/runtime/module_loader.rs b/codex-rs/code-mode-runtime/src/runtime/module_loader.rs index a3af147365..67c32b1ebb 100644 --- a/codex-rs/code-mode-runtime/src/runtime/module_loader.rs +++ b/codex-rs/code-mode-runtime/src/runtime/module_loader.rs @@ -1,7 +1,6 @@ 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; @@ -32,9 +31,6 @@ pub(super) fn evaluate_main_module( 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()); @@ -51,18 +47,6 @@ pub(super) fn evaluate_main_module( Ok(None) } -fn is_exit_exception( - scope: &mut v8::PinScope<'_, '_>, - exception: v8::Local<'_, v8::Value>, -) -> bool { - scope - .get_slot::() - .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, @@ -106,35 +90,27 @@ pub(super) fn completion_state( scope: &mut v8::PinScope<'_, '_>, pending_promise: Option<&v8::Global>, ) -> CompletionState { - let stored_value_writes = scope + // Termination during a microtask can leave the module promise pending. + // This isolate belongs to the exiting cell; never pump or await it again. + if scope .get_slot::() - .map(|state| state.stored_value_writes.clone()) - .unwrap_or_default(); + .is_some_and(|state| state.exit_requested) + { + return CompletionState::Completed { error_text: None }; + } let Some(pending_promise) = pending_promise else { - return CompletionState::Completed { - stored_value_writes, - error_text: None, - }; + return CompletionState::Completed { 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_value_writes, - error_text: None, - }, + v8::PromiseState::Fulfilled => CompletionState::Completed { 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_value_writes, - error_text, + error_text: Some(value_to_error_text(scope, result)), } } } diff --git a/codex-rs/code-mode-runtime/src/runtime/timers.rs b/codex-rs/code-mode-runtime/src/runtime/timers.rs index 3c2841c6cd..4fa8d17ae3 100644 --- a/codex-rs/code-mode-runtime/src/runtime/timers.rs +++ b/codex-rs/code-mode-runtime/src/runtime/timers.rs @@ -15,7 +15,7 @@ pub(super) struct ScheduledTimeout { pub(super) fn schedule_timeout( scope: &mut v8::PinScope<'_, '_>, args: v8::FunctionCallbackArguments, -) -> Result { +) -> Result, String> { let callback = args.get(0); if !callback.is_function() { return Err("setTimeout expects a function callback".to_string()); @@ -33,6 +33,10 @@ pub(super) fn schedule_timeout( let state = scope .get_slot_mut::() .ok_or_else(|| "runtime state unavailable".to_string())?; + // Converting the delay can execute JavaScript that calls exit(). + if state.exit_requested { + return Ok(None); + } let timeout_id = state.next_timeout_id; state.next_timeout_id = state.next_timeout_id.saturating_add(1); let runtime_command_tx = state.runtime_command_tx.clone(); @@ -49,7 +53,7 @@ pub(super) fn schedule_timeout( }, ); - Ok(timeout_id) + Ok(Some(timeout_id)) } pub(super) fn clear_timeout( @@ -63,7 +67,9 @@ pub(super) fn clear_timeout( let Some(state) = scope.get_slot_mut::() else { return Err("runtime state unavailable".to_string()); }; - state.pending_timeouts.remove(&timeout_id); + if !state.exit_requested { + state.pending_timeouts.remove(&timeout_id); + } Ok(()) } diff --git a/codex-rs/code-mode-runtime/src/runtime/value.rs b/codex-rs/code-mode-runtime/src/runtime/value.rs index 13e49159e4..2616354e8b 100644 --- a/codex-rs/code-mode-runtime/src/runtime/value.rs +++ b/codex-rs/code-mode-runtime/src/runtime/value.rs @@ -4,6 +4,7 @@ use codex_code_mode_protocol::DEFAULT_IMAGE_DETAIL; use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::ImageDetail; +use super::RuntimeState; use super::audio::wav_duration_seconds; const IMAGE_HELPER_EXPECTS_MESSAGE: &str = "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block"; @@ -362,6 +363,15 @@ pub(super) fn value_to_error_text( { return stack.to_rust_string_lossy(scope); } + // The stack getter can call exit() and V8 can clear its termination before + // returning here. Never enter a user-defined toString() after that. + if scope + .get_slot::() + .is_some_and(|state| state.exit_requested) + { + // send_scope_result reports this as successful exit and preserves prior stores. + return String::new(); + } value.to_rust_string_lossy(scope) } diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index 138d6234d1..9fb77e4fe6 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -577,6 +577,204 @@ async fn synchronous_exit_returns_successfully() { ); } +#[tokio::test] +async fn exit_stops_the_whole_cell_and_preserves_only_prior_writes() { + for (label, source) in [ + ("exit inside infinite loop", r#"for (;;) { exit(); }"#), + ("infinite loop after exit", r#"exit(); for (;;) {}"#), + ( + "invalid helper after exit", + r#"try { exit(); image(null); } catch { for (;;) {} }"#, + ), + ( + "catch and finally", + r#"try { exit(); } catch { text("caught"); } finally { store("phase", "finally"); text("finally"); }"#, + ), + ( + "queued microtask", + r#"Promise.resolve().then(() => { store("phase", "microtask"); text("microtask"); }); exit();"#, + ), + ( + "unawaited promise", + r#"Promise.resolve().then(() => exit()).catch(() => text("caught")); await new Promise(() => {});"#, + ), + ( + "async module", + r#"await Promise.resolve(); try { exit(); } finally { text("finally"); }"#, + ), + ( + "parallel promises", + r#"await Promise.all([Promise.resolve().then(() => exit()), Promise.resolve().then(() => { store("phase", "sibling"); text("sibling"); })]);"#, + ), + ( + "timer callback", + r#"setTimeout(() => { try { exit(); } catch { text("caught"); } finally { text("finally"); } }, 0); await new Promise(() => {});"#, + ), + ( + "timer microtask", + r#"setTimeout(() => { Promise.resolve().then(() => exit()); }, 0); await new Promise(() => {});"#, + ), + ( + "after awaited timer", + r#"await new Promise((resolve) => setTimeout(resolve, 0)); try { exit(); } catch { text("caught"); }"#, + ), + ( + "text toJSON", + r#"try { text({ toJSON() { exit(); } }); } catch { text("caught"); } finally { text("finally"); }"#, + ), + ( + "store toJSON", + r#"try { store("phase", { toJSON() { exit(); } }); } catch { text("caught"); } finally { text("finally"); }"#, + ), + ( + "store key coercion", + r#"try { store({ toString() { exit(); } }, "late"); } catch { text("caught"); } finally { text("finally"); }"#, + ), + ( + "image getter", + r#"try { image({ get image_url() { exit(); } }); } catch { text("caught"); } finally { text("finally"); }"#, + ), + ( + "exception stack getter", + r#"store("phase", "before formatting"); +throw { + get stack() { store("phase", "before"); exit(); for (;;) {} }, + toString() { for (;;) {} } +};"#, + ), + ] { + let service = InProcessCodeModeSession::new(); + let response = tokio::time::timeout( + Duration::from_secs(5), + execute( + &service, + ExecuteRequest { + source: format!( + r#"store("phase", "before"); text("before"); {source} store("phase", "after"); text("after");"# + ), + yield_time_ms: Some(30_000), + ..execute_request("") + }, + ), + ) + .await + .unwrap_or_else(|_| panic!("exit did not stop {label}")); + assert_eq!( + response, + RuntimeResponse::Result { + code_mode_host_duration: None, + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }], + error_text: None, + }, + "{label}" + ); + + let response = execute( + &service, + ExecuteRequest { + yield_time_ms: None, + ..execute_request(r#"text(load("phase"));"#) + }, + ) + .await; + assert_eq!( + response, + RuntimeResponse::Result { + code_mode_host_duration: None, + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }], + error_text: None, + }, + "{label}: next cell" + ); + } +} + +#[tokio::test] +async fn exit_after_yield_and_tool_response_finishes_successfully() { + let delegate = Arc::new(ReleasableToolDelegate::default()); + let service = InProcessCodeModeSession::new(); + let yield_signal = CancellationToken::new(); + let started = service + .execute( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#" +text("before yield"); +await tools.echo({}); +store("phase", "after tool"); +text("after tool"); +try { exit(); } catch { text("caught"); } finally { text("finally"); } +text("after exit"); +"# + .to_string(), + yield_time_ms: Some(30_000), + ..execute_request("") + }, + delegate.clone(), + Some(yield_signal.clone()), + ) + .await + .unwrap(); + let response = tokio::spawn(started.initial_response()); + wait_until_tool_started(&delegate).await.unwrap(); + yield_signal.cancel(); + assert_eq!( + response.await.unwrap().unwrap(), + RuntimeResponse::Yielded { + code_mode_host_duration: None, + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before yield".to_string(), + }], + } + ); + delegate.release_tool(); + assert_eq!( + service + .wait( + WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 30_000, + }, + /*preempt*/ None, + ) + .await + .unwrap(), + WaitOutcome::LiveCell(RuntimeResponse::Result { + code_mode_host_duration: None, + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "after tool".to_string(), + }], + error_text: None, + }) + ); + assert_eq!( + execute( + &service, + ExecuteRequest { + yield_time_ms: None, + ..execute_request(r#"text(load("phase"));"#) + }, + ) + .await, + RuntimeResponse::Result { + code_mode_host_duration: None, + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "after tool".to_string(), + }], + error_text: None, + } + ); +} + #[tokio::test] async fn stored_values_are_shared_between_cells_but_not_sessions() { let first_session = InProcessCodeModeSession::new();