diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 8379f372d..24eca2251 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -173,7 +173,7 @@ async fn run_cell( 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( .commit_completion( HashMap::new(), event, + /*pending_initial_yield_items*/ None, Arc::clone(&cell_state), ) .await @@ -429,6 +430,7 @@ async fn run_cell( .commit_completion( stored_value_writes, event, + /*pending_initial_yield_items*/ None, Arc::clone(&cell_state), ) .await diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index e217122b2..a4d93fc3e 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -45,9 +45,10 @@ impl CellHost for TestHost { &self, _stored_value_writes: HashMap, event: CellEvent, + pending_initial_yield_items: Option>, cell_state: Arc, ) -> 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, event: CellEvent, + pending_initial_yield_items: Option>, cell_state: Arc, ) -> 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))); +} diff --git a/codex-rs/code-mode/src/cell_actor/types.rs b/codex-rs/code-mode/src/cell_actor/types.rs index ee345d6ab..672820b93 100644 --- a/codex-rs/code-mode/src/cell_actor/types.rs +++ b/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, event: CellEvent, + pending_initial_yield_items: Option>, cell_state: Arc, ) -> impl Future + Send; @@ -113,7 +115,11 @@ enum CellPhase { Terminating { response_tx: oneshot::Sender>, }, - Completed(CellEvent), + Completed { + // Set only when `yield_control()` races the create-to-first-observe handoff. + pending_initial_yield_items: Option>, + 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>, 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>, ) -> 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>, +) -> 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, diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 71cf3d613..0f7557a75 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -256,6 +256,7 @@ impl CellHost for RuntimeCellHost { &self, stored_value_writes: HashMap, event: CellEvent, + pending_initial_yield_items: Option>, cell_state: Arc, ) -> CompletionCommit { let cancellation_token = cell_state.cancellation_token(); @@ -266,7 +267,9 @@ impl CellHost for RuntimeCellHost { } 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) { diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index a1dc18c0e..e0838d89d 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/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);