diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 493c638da44..4a19a7c5e00 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -48,6 +48,7 @@ impl CodeModeExecuteHandler { let message = HostToNodeMessage::Start { request_id: request_id.clone(), session_id, + default_yield_time_ms: super::DEFAULT_EXEC_YIELD_TIME_MS, enabled_tools, stored_values, source, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index f6561c518e5..1e20bd11f81 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -35,6 +35,7 @@ const CODE_MODE_WAIT_DESCRIPTION_TEMPLATE: &str = include_str!("wait_description pub(crate) const PUBLIC_TOOL_NAME: &str = "exec"; pub(crate) const WAIT_TOOL_NAME: &str = "exec_wait"; +pub(crate) const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; #[derive(Clone)] diff --git a/codex-rs/core/src/tools/code_mode/protocol.rs b/codex-rs/core/src/tools/code_mode/protocol.rs index fe0ab861f3e..ee522098292 100644 --- a/codex-rs/core/src/tools/code_mode/protocol.rs +++ b/codex-rs/core/src/tools/code_mode/protocol.rs @@ -41,6 +41,7 @@ pub(super) enum HostToNodeMessage { Start { request_id: String, session_id: i32, + default_yield_time_ms: u64, enabled_tools: Vec, stored_values: HashMap, source: String, diff --git a/codex-rs/core/src/tools/code_mode/runner.cjs b/codex-rs/core/src/tools/code_mode/runner.cjs index 3f6cedd53f4..bc6afe561ce 100644 --- a/codex-rs/core/src/tools/code_mode/runner.cjs +++ b/codex-rs/core/src/tools/code_mode/runner.cjs @@ -572,6 +572,7 @@ function startSession(protocol, sessions, start) { const session = { completed: false, content_items: [], + default_yield_time_ms: normalizeYieldTime(start.default_yield_time_ms), id: start.session_id, initial_yield_timer: null, initial_yield_triggered: false, @@ -585,6 +586,7 @@ function startSession(protocol, sessions, start) { }), }; sessions.set(session.id, session); + scheduleInitialYield(protocol, session, session.default_yield_time_ms); session.worker.on('message', (message) => { void handleWorkerMessage(protocol, sessions, session, message).catch((error) => { @@ -697,6 +699,9 @@ async function sendYielded(protocol, session) { if (session.completed || session.request_id === null) { return; } + session.initial_yield_timer = clearTimer(session.initial_yield_timer); + session.initial_yield_triggered = true; + session.poll_yield_timer = clearTimer(session.poll_yield_timer); const contentItems = takeContentItems(session); const requestId = session.request_id; try { diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 02c4987dc3a..699dc2011fd 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -68,7 +68,7 @@ struct WriteStdinArgs { } fn default_exec_yield_time_ms() -> u64 { - 10000 + 10_000 } fn default_write_stdin_yield_time_ms() -> u64 { diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 4f17d0d6c6a..66e1fa7c641 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -466,6 +466,103 @@ output_text("phase 3"); Ok(()) } +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_yield_timeout_works_for_busy_loop() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mut builder = test_codex().with_config(move |config| { + let _ = config.features.enable(Feature::CodeMode); + }); + let test = builder.build(&server).await?; + + let code = r#" +import { output_text, set_yield_time } from "@openai/code_mode"; + +output_text("phase 1"); +set_yield_time(10); +while (true) {} +"#; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call("call-1", "exec", code), + ev_completed("resp-1"), + ]), + ) + .await; + let first_completion = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "waiting"), + ev_completed("resp-2"), + ]), + ) + .await; + + tokio::time::timeout( + Duration::from_secs(5), + test.submit_turn("start the busy loop"), + ) + .await??; + + let first_request = first_completion.single_request(); + let first_items = custom_tool_output_items(&first_request, "call-1"); + assert_eq!(first_items.len(), 2); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script running with session ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&first_items, 0), + ); + assert_eq!(text_item(&first_items, 1), "phase 1"); + let session_id = extract_running_session_id(text_item(&first_items, 0)); + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-3"), + responses::ev_function_call( + "call-2", + "exec_wait", + &serde_json::to_string(&serde_json::json!({ + "session_id": session_id, + "terminate": true, + }))?, + ), + ev_completed("resp-3"), + ]), + ) + .await; + let second_completion = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-2", "terminated"), + ev_completed("resp-4"), + ]), + ) + .await; + + test.submit_turn("terminate it").await?; + + let second_request = second_completion.single_request(); + let second_items = function_tool_output_items(&second_request, "call-2"); + assert_eq!(second_items.len(), 1); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script terminated\nWall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&second_items, 0), + ); + + Ok(()) +} + #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_run_multiple_yielded_sessions() -> Result<()> {