codex: address PR review feedback (#16282)

codex-rs/core/src/codex.rs

@@ -822,12 +822,19 @@ pub(crate) struct Session {
     mailbox_rx: Mutex<MailboxReceiver>,
     idle_pending_input: Mutex<Vec<ResponseInputItem>>, // TODO (jif) merge with mailbox!
     jobs: Mutex<JobsState>,
+    job_timers_cancellation_token: CancellationToken,
     pub(crate) guardian_review_session: GuardianReviewSessionManager,
     pub(crate) services: SessionServices,
     js_repl: Arc<JsReplHandle>,
     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<SkillLoadOutcome>,
@@ -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)) => {}
                 }

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;

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,