core: make thread activity lifecycle atomic

codex-rs/core/src/session/mod.rs

@@ -792,14 +792,17 @@ impl Codex {
         }
         let is_shutdown = matches!(&sub.op, Op::Shutdown);
         let _send_guard = self.session.submission_send_lock.lock().await;
-        let Some(reservation) = self.session.submission_lifecycle.try_reserve(is_shutdown) else {
+        let Some(mut reservation) = self.session.try_reserve_submission(is_shutdown) else {
             if is_shutdown {
                 return Ok(());
             }
             return Err(CodexErr::InvalidRequest(
                 "thread is shutting down".to_string(),
             ));
         };
+        if !reservation.prepare_delivery() {
+            return Err(CodexErr::InternalAgentDied);
+        }
         if self.tx_sub.send(sub).await.is_err() {
             reservation.release_after_failed_delivery();
             return Err(CodexErr::InternalAgentDied);
@@ -825,9 +828,12 @@ impl Codex {
         if self.session.submission_lifecycle.is_closing() {
             return Ok(true);
         }
-        let Some(reservation) = self.session.try_claim_idle_shutdown().await else {
+        let Some(mut reservation) = self.session.try_claim_idle_shutdown().await else {
             return Ok(false);
         };
+        if !reservation.prepare_idle_shutdown_delivery() {
+            return Ok(false);
+        }
         let sub = Submission {
             id: new_submission_id(),
             op: Op::Shutdown,
@@ -1053,11 +1059,11 @@ impl Session {
 
     pub(crate) async fn try_claim_idle_shutdown(
         &self,
-    ) -> Option<session::SubmissionReservation<'_>> {
+    ) -> Option<session::SessionActivityReservation<'_>> {
         if self.has_idle_shutdown_blocker().await {
             return None;
         }
-        let reservation = self.submission_lifecycle.try_reserve_idle_shutdown()?;
+        let reservation = self.try_reserve_idle_shutdown()?;
         if self.has_idle_shutdown_blocker().await {
             return None;
         }

codex-rs/core/src/session/session.rs

@@ -8,6 +8,7 @@ use crate::shell_snapshot::ShellSnapshot;
 use crate::skills::SkillError;
 use crate::state::ActiveTurn;
 use crate::thread_activity::ThreadActivityHandle;
+use crate::thread_activity::ThreadActivityReservation;
 use codex_extension_api::ExtensionDataInit;
 use codex_login::auth::AgentIdentityAuthPolicy;
 use codex_protocol::SessionId;
@@ -108,12 +109,37 @@ pub(crate) struct SubmissionReservation<'a> {
     active: bool,
 }
 
+pub(crate) struct SessionActivityReservation<'a> {
+    submission: SubmissionReservation<'a>,
+    thread: ThreadActivityReservation,
+}
+
+impl SessionActivityReservation<'_> {
+    pub(crate) fn prepare_delivery(&mut self) -> bool {
+        self.thread.prepare_delivery()
+    }
+
+    pub(crate) fn prepare_idle_shutdown_delivery(&mut self) -> bool {
+        self.thread.prepare_idle_shutdown_delivery()
+    }
+
+    pub(crate) fn commit(self) {
+        self.submission.commit();
+        self.thread.commit();
+    }
+
+    pub(crate) fn release_after_failed_delivery(self) {
+        self.submission.release_after_failed_delivery();
+        self.thread.release_after_failed_delivery();
+    }
+}
+
 impl SubmissionReservation<'_> {
     pub(crate) fn commit(mut self) {
         self.active = false;
     }
 
-    pub(crate) fn release_after_failed_delivery(mut self) {
+    fn release_after_failed_delivery(mut self) {
         if self.active {
             self.lifecycle.release(/*clear_closing*/ false);
             self.active = false;
@@ -572,12 +598,28 @@ impl Session {
         self.thread_id
     }
 
-    pub(crate) fn try_reserve_activity(&self) -> Option<SubmissionReservation<'_>> {
-        self.submission_lifecycle.try_reserve(/*close*/ false)
+    pub(crate) fn try_reserve_submission(
+        &self,
+        close: bool,
+    ) -> Option<SessionActivityReservation<'_>> {
+        let thread = self.thread_activity.try_reserve(close)?;
+        let submission = self.submission_lifecycle.try_reserve(close)?;
+        Some(SessionActivityReservation { submission, thread })
+    }
+
+    pub(crate) fn try_reserve_activity(&self) -> Option<SessionActivityReservation<'_>> {
+        self.try_reserve_submission(/*close*/ false)
+    }
+
+    pub(crate) fn try_reserve_idle_shutdown(&self) -> Option<SessionActivityReservation<'_>> {
+        let thread = self.thread_activity.try_reserve_idle_shutdown()?;
+        let submission = self.submission_lifecycle.try_reserve_idle_shutdown()?;
+        Some(SessionActivityReservation { submission, thread })
     }
 
     pub(crate) fn finish_submission(&self) {
         self.submission_lifecycle.finish_submission();
+        self.thread_activity.finish_submission();
     }
 
     /// Returns the identity shared by the root thread and all descendant threads.