code-mode: preserve initial yield at completion (#29289)

## Summary

- Retain the first pre-observation `yield_control()` boundary when a
cell completes before observation.
- Deliver the preserved yield before the buffered completion.
- Keep later unattached yields as no-ops.

## Why

Create followed by the initial wait must preserve the former execute
response boundary even when the script runs to completion first.

## Impact

The first wait observes the same initial yield boundary as before create
and observe were decoupled.

## Validation

- Focused initial-yield signature regression passed.
- Stack-tip validation: `just test -p codex-code-mode -p
codex-code-mode-protocol` (70 passed).
- Parent branch:
`cconger/code-mode-runtime-compact-03e2-observation-delivery`.

codex-rs/code-mode/src/cell_actor/mod.rs

@@ -173,7 +173,7 @@ async fn run_cell<H: CellHost>(
                 if response_tx.is_closed() {
                     continue;
                 }
-                let response_tx = match cell_state.route_observation(response_tx) {
+                let response_tx = match cell_state.route_observation(mode, response_tx) {
                     ObservationDelivery::Running(response_tx) => response_tx,
                     ObservationDelivery::Delivered => break,
                     ObservationDelivery::Buffered | ObservationDelivery::Closed => continue,
@@ -284,6 +284,7 @@ async fn run_cell<H: CellHost>(
                         .commit_completion(
                             HashMap::new(),
                             event,
+                            /*pending_initial_yield_items*/ None,
                             Arc::clone(&cell_state),
                         )
                         .await
@@ -429,6 +430,7 @@ async fn run_cell<H: CellHost>(
                             .commit_completion(
                                 stored_value_writes,
                                 event,
+                                /*pending_initial_yield_items*/ None,
                                 Arc::clone(&cell_state),
                             )
                             .await

codex-rs/code-mode/src/cell_actor/tests.rs

@@ -45,9 +45,10 @@ impl CellHost for TestHost {
         &self,
         _stored_value_writes: HashMap<String, JsonValue>,
         event: CellEvent,
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,
     ) -> CompletionCommit {
-        cell_state.commit_completion(event, || {})
+        cell_state.commit_completion(event, pending_initial_yield_items, || {})
     }
 
     async fn closed(&self) {}
@@ -76,9 +77,10 @@ impl CellHost for RecordingHost {
         &self,
         _stored_value_writes: HashMap<String, JsonValue>,
         event: CellEvent,
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,
     ) -> CompletionCommit {
-        cell_state.commit_completion(event, || {})
+        cell_state.commit_completion(event, pending_initial_yield_items, || {})
     }
 
     async fn closed(&self) {}
@@ -397,7 +399,11 @@ async fn only_the_first_termination_claims_a_buffered_completion() {
         error_text: None,
     };
     assert_eq!(
-        cell_state.commit_completion(completion.clone(), || {}),
+        cell_state.commit_completion(
+            completion.clone(),
+            /*pending_initial_yield_items*/ None,
+            || {}
+        ),
         CompletionCommit::Committed
     );
     assert!(matches!(
@@ -430,7 +436,11 @@ async fn termination_claim_prevents_stored_value_commit() {
     };
 
     assert_eq!(
-        cell_state.commit_completion(completion.clone(), || commit_ran = true),
+        cell_state.commit_completion(
+            completion.clone(),
+            /*pending_initial_yield_items*/ None,
+            || commit_ran = true
+        ),
         CompletionCommit::Rejected(completion)
     );
     assert!(!commit_ran);
@@ -453,7 +463,11 @@ fn failed_completion_delivery_rebuffers_the_event() {
         error_text: None,
     };
     assert_eq!(
-        cell_state.commit_completion(event.clone(), || {}),
+        cell_state.commit_completion(
+            event.clone(),
+            /*pending_initial_yield_items*/ None,
+            || {}
+        ),
         CompletionCommit::Committed
     );
     let (response_tx, response_rx) = oneshot::channel();
@@ -466,8 +480,151 @@ fn failed_completion_delivery_rebuffers_the_event() {
 
     let (response_tx, mut response_rx) = oneshot::channel();
     assert!(matches!(
-        cell_state.route_observation(response_tx),
+        cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
         ObservationDelivery::Delivered
     ));
     assert_eq!(response_rx.try_recv(), Ok(Ok(event)));
 }
+
+#[test]
+fn buffered_initial_yield_precedes_buffered_completion_for_yield_observer() {
+    let cell_state = CellState::new(CancellationToken::new());
+    let completion = CellEvent::Completed {
+        content_items: vec![OutputItem::Text {
+            text: "after".to_string(),
+        }],
+        error_text: None,
+    };
+    assert_eq!(
+        cell_state.commit_completion(
+            completion.clone(),
+            Some(vec![OutputItem::Text {
+                text: "before".to_string(),
+            }]),
+            || {}
+        ),
+        CompletionCommit::Committed
+    );
+    assert!(matches!(
+        cell_state.deliver_completion(/*response_tx*/ None),
+        CompletionDelivery::Buffered
+    ));
+
+    let (response_tx, mut response_rx) = oneshot::channel();
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
+        ObservationDelivery::Buffered
+    ));
+    assert_eq!(
+        response_rx.try_recv(),
+        Ok(Ok(CellEvent::Yielded {
+            content_items: vec![OutputItem::Text {
+                text: "before".to_string(),
+            }],
+        }))
+    );
+
+    let (response_tx, mut response_rx) = oneshot::channel();
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
+        ObservationDelivery::Delivered
+    ));
+    assert_eq!(response_rx.try_recv(), Ok(Ok(completion)));
+}
+
+#[test]
+fn pending_observer_merges_initial_yield_and_completion_output() {
+    let cell_state = CellState::new(CancellationToken::new());
+    assert_eq!(
+        cell_state.commit_completion(
+            CellEvent::Completed {
+                content_items: vec![OutputItem::Text {
+                    text: "after".to_string(),
+                }],
+                error_text: None,
+            },
+            Some(vec![OutputItem::Text {
+                text: "before".to_string(),
+            }]),
+            || {}
+        ),
+        CompletionCommit::Committed
+    );
+    assert!(matches!(
+        cell_state.deliver_completion(/*response_tx*/ None),
+        CompletionDelivery::Buffered
+    ));
+
+    let (response_tx, mut response_rx) = oneshot::channel();
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::PendingFrontier, response_tx),
+        ObservationDelivery::Delivered
+    ));
+    assert_eq!(
+        response_rx.try_recv(),
+        Ok(Ok(CellEvent::Completed {
+            content_items: vec![
+                OutputItem::Text {
+                    text: "before".to_string(),
+                },
+                OutputItem::Text {
+                    text: "after".to_string(),
+                },
+            ],
+            error_text: None,
+        }))
+    );
+}
+
+#[test]
+fn dropped_pending_observation_preserves_the_initial_yield_boundary() {
+    let cell_state = CellState::new(CancellationToken::new());
+    let completion = CellEvent::Completed {
+        content_items: vec![OutputItem::Text {
+            text: "after".to_string(),
+        }],
+        error_text: None,
+    };
+    assert_eq!(
+        cell_state.commit_completion(
+            completion.clone(),
+            Some(vec![OutputItem::Text {
+                text: "before".to_string(),
+            }]),
+            || {}
+        ),
+        CompletionCommit::Committed
+    );
+    assert!(matches!(
+        cell_state.deliver_completion(/*response_tx*/ None),
+        CompletionDelivery::Buffered
+    ));
+
+    let (response_tx, response_rx) = oneshot::channel();
+    drop(response_rx);
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::PendingFrontier, response_tx),
+        ObservationDelivery::Buffered
+    ));
+
+    let (response_tx, mut response_rx) = oneshot::channel();
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
+        ObservationDelivery::Buffered
+    ));
+    assert_eq!(
+        response_rx.try_recv(),
+        Ok(Ok(CellEvent::Yielded {
+            content_items: vec![OutputItem::Text {
+                text: "before".to_string(),
+            }],
+        }))
+    );
+
+    let (response_tx, mut response_rx) = oneshot::channel();
+    assert!(matches!(
+        cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
+        ObservationDelivery::Delivered
+    ));
+    assert_eq!(response_rx.try_recv(), Ok(Ok(completion)));
+}

codex-rs/code-mode/src/cell_actor/types.rs

@@ -11,6 +11,7 @@ use tokio_util::sync::CancellationToken;
 
 use crate::session_runtime::CellEvent;
 use crate::session_runtime::ObserveMode;
+use crate::session_runtime::OutputItem;
 use crate::session_runtime::ToolKind;
 use crate::session_runtime::ToolName;
 
@@ -54,6 +55,7 @@ pub(crate) trait CellHost: Send + Sync + 'static {
         &self,
         stored_value_writes: HashMap<String, JsonValue>,
         event: CellEvent,
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,
     ) -> impl Future<Output = CompletionCommit> + Send;
 
@@ -113,7 +115,11 @@ enum CellPhase {
     Terminating {
         response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
     },
-    Completed(CellEvent),
+    Completed {
+        // Set only when `yield_control()` races the create-to-first-observe handoff.
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
+        event: CellEvent,
+    },
     CompletionClaimed(CellEvent),
     Tombstone,
 }
@@ -152,7 +158,7 @@ impl CellState {
                 .phase
                 .lock()
                 .unwrap_or_else(std::sync::PoisonError::into_inner),
-            CellPhase::Running | CellPhase::Completed(_)
+            CellPhase::Running | CellPhase::Completed { .. }
         );
         accepting_phase && !self.cancellation_token.is_cancelled()
     }
@@ -173,7 +179,11 @@ impl CellState {
                 *phase = CellPhase::Terminating { response_tx };
                 Box::pin(async { Err(CellError::AlreadyTerminating) })
             }
-            CellPhase::Completed(event) => {
+            CellPhase::Completed {
+                pending_initial_yield_items,
+                event,
+            } => {
+                let event = prepend_initial_yield(event, pending_initial_yield_items);
                 *phase = CellPhase::CompletionClaimed(event.clone());
                 self.cancellation_token.cancel();
                 ready_event(event)
@@ -189,6 +199,7 @@ impl CellState {
     pub(crate) fn commit_completion(
         &self,
         event: CellEvent,
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
         commit: impl FnOnce(),
     ) -> CompletionCommit {
         let mut phase = self
@@ -199,7 +210,10 @@ impl CellState {
             return CompletionCommit::Rejected(event);
         }
         commit();
-        *phase = CellPhase::Completed(event);
+        *phase = CellPhase::Completed {
+            pending_initial_yield_items,
+            event,
+        };
         CompletionCommit::Committed
     }
 
@@ -211,15 +225,22 @@ impl CellState {
             .phase
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner);
-        let event = match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
-            CellPhase::Completed(event) => event,
-            previous => {
-                *phase = previous;
-                return CompletionDelivery::Rejected(response_tx);
-            }
-        };
+        let (pending_initial_yield_items, event) =
+            match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
+                CellPhase::Completed {
+                    pending_initial_yield_items,
+                    event,
+                } => (pending_initial_yield_items, event),
+                previous => {
+                    *phase = previous;
+                    return CompletionDelivery::Rejected(response_tx);
+                }
+            };
         let Some(response_tx) = response_tx else {
-            *phase = CellPhase::Completed(event);
+            *phase = CellPhase::Completed {
+                pending_initial_yield_items,
+                event,
+            };
             return CompletionDelivery::Buffered;
         };
         match response_tx.send(Ok(event)) {
@@ -228,7 +249,10 @@ impl CellState {
                 CompletionDelivery::Delivered
             }
             Err(Ok(event)) => {
-                *phase = CellPhase::Completed(event);
+                *phase = CellPhase::Completed {
+                    pending_initial_yield_items,
+                    event,
+                };
                 CompletionDelivery::Buffered
             }
             Err(Err(error)) => {
@@ -239,6 +263,7 @@ impl CellState {
 
     pub(crate) fn route_observation(
         &self,
+        mode: ObserveMode,
         response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
     ) -> ObservationDelivery {
         let mut phase = self
@@ -250,19 +275,56 @@ impl CellState {
                 *phase = CellPhase::Running;
                 ObservationDelivery::Running(response_tx)
             }
-            CellPhase::Completed(event) => match response_tx.send(Ok(event)) {
-                Ok(()) => {
-                    self.cancellation_token.cancel();
-                    ObservationDelivery::Delivered
+            CellPhase::Completed {
+                pending_initial_yield_items: Some(content_items),
+                event,
+            } if matches!(mode, ObserveMode::YieldAfter(_)) => {
+                match response_tx.send(Ok(CellEvent::Yielded { content_items })) {
+                    Ok(()) => {
+                        *phase = CellPhase::Completed {
+                            pending_initial_yield_items: None,
+                            event,
+                        };
+                        ObservationDelivery::Buffered
+                    }
+                    Err(Ok(CellEvent::Yielded { content_items })) => {
+                        *phase = CellPhase::Completed {
+                            pending_initial_yield_items: Some(content_items),
+                            event,
+                        };
+                        ObservationDelivery::Buffered
+                    }
+                    Err(Ok(event)) => {
+                        panic!("initial yield delivery returned an unexpected event: {event:?}")
+                    }
+                    Err(Err(error)) => {
+                        panic!("initial yield delivery returned an actor error: {error:?}")
+                    }
                 }
-                Err(Ok(event)) => {
-                    *phase = CellPhase::Completed(event);
-                    ObservationDelivery::Buffered
+            }
+            CellPhase::Completed {
+                pending_initial_yield_items,
+                event,
+            } => {
+                let delivered_event =
+                    prepend_initial_yield(event.clone(), pending_initial_yield_items.clone());
+                match response_tx.send(Ok(delivered_event)) {
+                    Ok(()) => {
+                        self.cancellation_token.cancel();
+                        ObservationDelivery::Delivered
+                    }
+                    Err(Ok(_)) => {
+                        *phase = CellPhase::Completed {
+                            pending_initial_yield_items,
+                            event,
+                        };
+                        ObservationDelivery::Buffered
+                    }
+                    Err(Err(error)) => {
+                        panic!("completion delivery unexpectedly carried an actor error: {error:?}")
+                    }
                 }
-                Err(Err(error)) => {
-                    panic!("completion delivery unexpectedly carried an actor error: {error:?}")
-                }
-            },
+            }
             CellPhase::Terminating {
                 response_tx: termination_tx,
             } => {
@@ -295,7 +357,10 @@ impl CellState {
                 let _ = response_tx.send(Ok(event.clone()));
                 Some(event)
             }
-            CellPhase::Completed(completed_event) => Some(completed_event),
+            CellPhase::Completed {
+                pending_initial_yield_items,
+                event,
+            } => Some(prepend_initial_yield(event, pending_initial_yield_items)),
             CellPhase::CompletionClaimed(completed_event) => Some(completed_event),
             CellPhase::Tombstone => None,
         };
@@ -316,6 +381,49 @@ impl CellState {
     }
 }
 
+fn prepend_initial_yield(
+    event: CellEvent,
+    pending_initial_yield_items: Option<Vec<OutputItem>>,
+) -> CellEvent {
+    let Some(mut pending_initial_yield_items) = pending_initial_yield_items else {
+        return event;
+    };
+    match event {
+        CellEvent::Yielded { mut content_items } => {
+            pending_initial_yield_items.append(&mut content_items);
+            CellEvent::Yielded {
+                content_items: pending_initial_yield_items,
+            }
+        }
+        CellEvent::Pending {
+            mut content_items,
+            pending_tool_call_ids,
+        } => {
+            pending_initial_yield_items.append(&mut content_items);
+            CellEvent::Pending {
+                content_items: pending_initial_yield_items,
+                pending_tool_call_ids,
+            }
+        }
+        CellEvent::Completed {
+            mut content_items,
+            error_text,
+        } => {
+            pending_initial_yield_items.append(&mut content_items);
+            CellEvent::Completed {
+                content_items: pending_initial_yield_items,
+                error_text,
+            }
+        }
+        CellEvent::Terminated { mut content_items } => {
+            pending_initial_yield_items.append(&mut content_items);
+            CellEvent::Terminated {
+                content_items: pending_initial_yield_items,
+            }
+        }
+    }
+}
+
 pub(super) enum CellCommand {
     Observe {
         mode: ObserveMode,

codex-rs/code-mode/src/session_runtime/mod.rs

@@ -256,6 +256,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
         &self,
         stored_value_writes: HashMap<String, JsonValue>,
         event: CellEvent,
+        pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,
     ) -> CompletionCommit {
         let cancellation_token = cell_state.cancellation_token();
@@ -266,7 +267,9 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
             }
             stored_values = self.inner.stored_values.lock() => stored_values,
         };
-        cell_state.commit_completion(event, || stored_values.extend(stored_value_writes))
+        cell_state.commit_completion(event, pending_initial_yield_items, || {
+            stored_values.extend(stored_value_writes);
+        })
     }
 
     async fn closed(&self) {

codex-rs/code-mode/src/session_runtime/tests.rs

@@ -59,6 +59,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
             JsonValue::String("lost".to_string()),
         )]),
         completion.clone(),
+        /*pending_initial_yield_items*/ None,
         Arc::clone(&cell_state),
     );
     tokio::pin!(commit);