Close input admission while a turn finishes recording (#53452)

## Why

Inputs arriving during turn finalization could be queued after the final batch had already been drained, or recorded ahead of that batch.

## What changed

Mark active turns as finishing before draining pending input. Reject injection into finishing turns and make late client responses and notifications wait until the final batch is recorded before appending to history. Continue accepting input into turns reserved but not yet started.

## Testing

Extend the task-completion test to block final recording, verify that input admission is closed, and confirm that a late notification is recorded after queued input while preserving the notification truncation budget. Strengthen the retained-tool-output scenario to assert tool output names and ordering.

GitOrigin-RevId: a57e44bd974e60e134183dab989795e2c3538596

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

@@ -19,7 +19,7 @@ impl Session {
         input: Vec<T>,
     ) -> Result<(), Vec<T>> {
         let mut active = self.active_turn.lock().await;
-        match active.as_mut() {
+        match active.as_mut().filter(|turn| turn.finishing.is_none()) {
             Some(active_turn) => {
                 self.input_queue
                     .extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
@@ -81,7 +81,7 @@ impl Session {
             .map(|item| self.annotate_client_response_item(item))
             .collect::<Vec<_>>();
         let mut active = self.active_turn.lock().await;
-        if let Some(active_turn) = active.as_mut() {
+        if let Some(active_turn) = active.as_mut().filter(|turn| turn.finishing.is_none()) {
             self.input_queue
                 .extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
                     active_turn.turn_state.as_ref(),
@@ -93,7 +93,11 @@ impl Session {
                 .await;
             return;
         }
+        let finishing = active.as_ref().and_then(|turn| turn.finishing.clone());
         drop(active);
+        if let Some(finishing) = finishing {
+            finishing.cancelled().await;
+        }
         self.record_annotated_conversation_items(turn_context, turn_context.model_info(), items)
             .await;
     }
@@ -175,6 +179,15 @@ impl Session {
         let Err(items) = self.inject_if_running(items).await else {
             return;
         };
+        let finishing = self
+            .active_turn
+            .lock()
+            .await
+            .as_ref()
+            .and_then(|turn| turn.finishing.clone());
+        if let Some(finishing) = finishing {
+            finishing.cancelled().await;
+        }
         let default_turn_context;
         let turn_context = match current_turn_context {
             Some(turn_context) => turn_context,

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

@@ -47,9 +47,31 @@ async fn task_finish_preserves_notification_budget_and_queued_user_input() {
         codex_protocol::openai_models::TruncationPolicyConfig::tokens(/*limit*/ 400);
     turn.next_step_settings.store(Arc::new(current));
 
-    session
-        .on_task_finished(Arc::clone(&turn), /*task_result*/ Ok(None))
-        .await;
+    // Stall recording the first batch, then inject while finalization is in progress.
+    let state = session.state.lock().await;
+    let finish = session.on_task_finished(Arc::clone(&turn), /*task_result*/ Ok(None));
+    let late = vec![notification("late notification".to_string())];
+    tokio::join!(finish, async {
+        while session
+            .active_turn
+            .lock()
+            .await
+            .as_ref()
+            .expect("turn remains active while final recording is blocked")
+            .finishing
+            .is_none()
+        {
+            tokio::task::yield_now().await;
+        }
+        // Admission is closed even if recording has not yet reached the history lock.
+        let notice = vec![notification("producer still running".to_string())];
+        assert_eq!(session.inject_if_running(notice.clone()).await, Err(notice));
+        let inject = session.inject_client_response_items(late.clone(), &turn);
+        tokio::pin!(inject);
+        assert!(futures::poll!(inject.as_mut()).is_pending());
+        drop(state);
+        inject.await;
+    });
 
     let expected = vec![
         notification(codex_utils_output_truncation::truncate_text(
@@ -57,6 +79,7 @@ async fn task_finish_preserves_notification_budget_and_queued_user_input() {
             codex_utils_output_truncation::TruncationPolicy::Tokens(120),
         )),
         session.response_item_from_user_input(user_input),
+        late[0].clone(),
     ];
     let history = session.clone_history().await;
     assert_eq!(

codex-rs/core/src/state/turn.rs

@@ -29,6 +29,9 @@ use codex_protocol::protocol::TokenUsage;
 /// Metadata about the currently running turn.
 pub(crate) struct ActiveTurn {
     pub(crate) task: Option<RunningTask>,
+    /// Present when input admission is closed; cancelled when the final batch is recorded.
+    /// Reservations that have not started still accept inputs.
+    pub(crate) finishing: Option<CancellationToken>,
     pub(crate) turn_state: Arc<Mutex<TurnState>>,
 }
 
@@ -57,6 +60,7 @@ impl Default for ActiveTurn {
     fn default() -> Self {
         Self {
             task: None,
+            finishing: None,
             turn_state: Arc::new(Mutex::new(TurnState::default())),
         }
     }

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

@@ -322,6 +322,7 @@ impl Session {
             .await;
             let turn = active.get_or_insert_with(ActiveTurn::default);
             debug_assert!(turn.task.is_none());
+            turn.finishing = None;
             Arc::clone(&turn.turn_state)
         };
         turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone();
@@ -659,11 +660,14 @@ impl Session {
             .turn_metadata_state
             .cancel_git_enrichment_task();
 
+        let finishing = CancellationToken::new();
+        let finish_recording = finishing.clone().drop_guard();
         let turn_state = {
             let mut active = self.active_turn.lock().await;
             active.as_mut().and_then(|active_turn| {
                 let task = active_turn.task.take()?;
                 task.handle.detach();
+                active_turn.finishing = Some(finishing);
                 Some(Arc::clone(&active_turn.turn_state))
             })
         };
@@ -696,6 +700,7 @@ impl Session {
             PersistContext::Standard,
         )
         .await;
+        drop(finish_recording);
         let turn_telemetry = &turn_context.session_telemetry;
         // Emit token usage metrics.
         {

codex-rs/core/tests/suite/scenarios_retained_tool_outputs.rs

@@ -98,7 +98,18 @@ async fn retained_tool_output_survives_compaction_and_resume() -> Result<()> {
             assert!(request.inputs_of_type("compaction_trigger").is_empty());
         }
         let compact_outputs = requests[3].inputs_of_type("function_call_output");
-        assert_eq!(compact_outputs.len(), 4);
+        assert_eq!(
+            compact_outputs
+                .iter()
+                .map(|output| output["name"].as_str())
+                .collect::<Vec<_>>(),
+            [
+                Some("send_message_to_thread"),
+                Some("task_details"),
+                Some("send_message_to_thread"),
+                Some("progress"),
+            ]
+        );
         // Emergency trimming removes the trailing log, then stops at the marked instruction.
         assert_eq!(compact_outputs[1]["output"], details);
         if enabled {