diff --git a/codex-rs/code-mode-runtime/src/runtime/mod.rs b/codex-rs/code-mode-runtime/src/runtime/mod.rs index 3077b23330..f2e75c6f50 100644 --- a/codex-rs/code-mode-runtime/src/runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/runtime/mod.rs @@ -103,7 +103,9 @@ pub(crate) fn spawn_runtime( stored_values, }; + let runtime_handle = tokio::runtime::Handle::current(); spawn_supervised_runtime_thread(event_tx.clone(), task_failure_handler, move || { + let _runtime_guard = runtime_handle.enter(); run_runtime( config, event_tx, @@ -512,4 +514,70 @@ await new Promise(() => {}); .send(RuntimeControlCommand::Terminate) .unwrap(); } + + #[tokio::test] + async fn timers_release_tasks_when_cleared_or_the_cell_finishes() { + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let (_runtime_tx, _runtime_control_tx, _runtime_terminate_handle) = spawn_runtime( + HashMap::new(), + execute_request( + r#" +clearTimeout(setTimeout(() => text("cancelled"), 3_600_000)); +await new Promise((resolve) => setTimeout(resolve, 3_600_000)); +text("done"); +setTimeout(() => text("late"), 3_600_000); +"#, + ), + event_tx, + PendingRuntimeMode::Continue, + /*task_failure_handler*/ None, + ) + .unwrap(); + + loop { + match tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .unwrap() + { + Some(RuntimeEvent::Pending) => break, + Some(_) => {} + None => panic!("runtime closed before the timer was pending"), + } + } + tokio::task::yield_now().await; + assert_eq!( + tokio::runtime::Handle::current() + .metrics() + .num_alive_tasks(), + 1 + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(3_600)).await; + tokio::time::resume(); + + let mut output = Vec::new(); + while let Some(event) = tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("timer runtime should finish") + { + match event { + RuntimeEvent::ContentItem(item) => output.push(item), + RuntimeEvent::Result { error_text, .. } => assert_eq!(error_text, None), + _ => {} + } + } + assert_eq!( + output, + vec![FunctionCallOutputContentItem::InputText { + text: "done".to_string() + }] + ); + tokio::task::yield_now().await; + assert_eq!( + tokio::runtime::Handle::current() + .metrics() + .num_alive_tasks(), + 0 + ); + } } diff --git a/codex-rs/code-mode-runtime/src/runtime/timers.rs b/codex-rs/code-mode-runtime/src/runtime/timers.rs index 01c414cefe..3c2841c6cd 100644 --- a/codex-rs/code-mode-runtime/src/runtime/timers.rs +++ b/codex-rs/code-mode-runtime/src/runtime/timers.rs @@ -1,12 +1,15 @@ -use std::thread; use std::time::Duration; +use tokio_util::task::AbortOnDropHandle; + use super::RuntimeCommand; use super::RuntimeState; use super::value::value_to_error_text; pub(super) struct ScheduledTimeout { callback: v8::Global, + // Clearing the timeout or dropping the isolate also cancels its sleep. + _task: AbortOnDropHandle<()>, } pub(super) fn schedule_timeout( @@ -33,13 +36,18 @@ pub(super) fn schedule_timeout( 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(); - state - .pending_timeouts - .insert(timeout_id, ScheduledTimeout { callback }); - thread::spawn(move || { - thread::sleep(Duration::from_millis(delay_ms)); + let sleep = tokio::time::sleep(Duration::from_millis(delay_ms)); + let task = tokio::spawn(async move { + sleep.await; let _ = runtime_command_tx.send(RuntimeCommand::TimeoutFired { id: timeout_id }); }); + state.pending_timeouts.insert( + timeout_id, + ScheduledTimeout { + callback, + _task: AbortOnDropHandle::new(task), + }, + ); Ok(timeout_id) }