diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 85e0494d75..eb4a40de76 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -822,12 +822,19 @@ pub(crate) struct Session { mailbox_rx: Mutex, idle_pending_input: Mutex>, // TODO (jif) merge with mailbox! jobs: Mutex, + job_timers_cancellation_token: CancellationToken, pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, js_repl: Arc, next_internal_sub_id: AtomicU64, } +impl Drop for Session { + fn drop(&mut self) { + self.job_timers_cancellation_token.cancel(); + } +} + #[derive(Clone, Debug)] pub(crate) struct TurnSkillsContext { pub(crate) outcome: Arc, @@ -1940,6 +1947,7 @@ impl Session { mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), jobs: Mutex::new(JobsState::default()), + job_timers_cancellation_token: CancellationToken::new(), guardian_review_session: GuardianReviewSessionManager::default(), services, js_repl, @@ -2716,13 +2724,18 @@ impl Session { .await; self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) .await; - tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.start_task( + let session = Arc::clone(self); + if let Err(err) = tokio::task::spawn_blocking(move || { + futures::executor::block_on(session.start_task( turn_context, Vec::new(), crate::tasks::RegularTask::new(), )); - }); + }) + .await + { + warn!("failed to start job turn: {err}"); + } } fn spawn_job_timer( @@ -2732,9 +2745,11 @@ impl Session { cancellation_token: CancellationToken, ) { let weak = Arc::downgrade(self); + let session_cancel = self.job_timers_cancellation_token.clone(); tokio::spawn(async move { loop { tokio::select! { + _ = session_cancel.cancelled() => break, _ = cancellation_token.cancelled() => break, _ = tokio::time::sleep(std::time::Duration::from_secs(seconds)) => {} } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 1e8684edce..32e023afba 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -2753,6 +2753,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), jobs: Mutex::new(JobsState::default()), + job_timers_cancellation_token: CancellationToken::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, js_repl, @@ -3597,6 +3598,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), jobs: Mutex::new(JobsState::default()), + job_timers_cancellation_token: CancellationToken::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, js_repl, @@ -3616,6 +3618,16 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await } +#[tokio::test] +async fn dropping_session_cancels_job_timers() { + let (session, _, _) = make_session_and_context_with_rx().await; + let cancel_token = session.job_timers_cancellation_token.clone(); + + drop(session); + + assert!(cancel_token.is_cancelled()); +} + #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/tests/suite/jobs.rs b/codex-rs/core/tests/suite/jobs.rs index 27ca4c5bc6..d58e6ab6d9 100644 --- a/codex-rs/core/tests/suite/jobs.rs +++ b/codex-rs/core/tests/suite/jobs.rs @@ -17,6 +17,15 @@ use pretty_assertions::assert_eq; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn create_job_emits_fired_background_event_when_job_starts() -> Result<()> { + assert_after_turn_job_starts_and_emits_fired_event().await +} + +#[tokio::test(flavor = "current_thread")] +async fn create_job_starts_on_current_thread_runtime() -> Result<()> { + assert_after_turn_job_starts_and_emits_fired_event().await +} + +async fn assert_after_turn_job_starts_and_emits_fired_event() -> Result<()> { let server = start_mock_server().await; let _mock = mount_sse_once( &server,