diff --git a/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs index 23c4bc9634..6d2444b8be 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs @@ -41,7 +41,6 @@ impl CellHost for PanickingCallbackHost { &self, _stored_value_writes: HashMap>, _event: CellEvent, - _pending_initial_yield_items: Option>, _cell_state: Arc, ) -> CompletionCommit { panic!("unexpected completion commit"); diff --git a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs index 1ccdb01c14..c44228aad7 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs @@ -184,7 +184,7 @@ async fn run_cell( if response_tx.is_closed() { continue; } - let response_tx = match cell_state.route_observation(mode, response_tx) { + let response_tx = match cell_state.route_observation(response_tx) { ObservationDelivery::Running(response_tx) => response_tx, ObservationDelivery::Delivered => break, ObservationDelivery::Buffered | ObservationDelivery::Closed => continue, @@ -322,7 +322,6 @@ async fn run_cell( .commit_completion( HashMap::new(), event, - /*pending_initial_yield_items*/ None, Arc::clone(&cell_state), ) .await @@ -479,7 +478,6 @@ 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-runtime/src/cell_actor/tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs index e85ea96eb3..fb25a7491a 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs @@ -45,10 +45,9 @@ 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, pending_initial_yield_items, || {}) + cell_state.commit_completion(event, || {}) } async fn closed(&self) {} @@ -77,10 +76,9 @@ 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, pending_initial_yield_items, || {}) + cell_state.commit_completion(event, || {}) } async fn closed(&self) {} @@ -479,11 +477,7 @@ async fn only_the_first_termination_claims_a_buffered_completion() { error_text: None, }; assert_eq!( - cell_state.commit_completion( - completion.clone(), - /*pending_initial_yield_items*/ None, - || {} - ), + cell_state.commit_completion(completion.clone(), || {}), CompletionCommit::Committed ); assert!(matches!( @@ -517,11 +511,7 @@ async fn termination_claim_prevents_stored_value_commit() { }; assert_eq!( - cell_state.commit_completion( - completion.clone(), - /*pending_initial_yield_items*/ None, - || commit_ran = true - ), + cell_state.commit_completion(completion.clone(), || commit_ran = true), CompletionCommit::Rejected(completion) ); assert!(!commit_ran); @@ -539,17 +529,15 @@ async fn termination_claim_prevents_stored_value_commit() { #[tokio::test(start_paused = true)] async fn failed_completion_delivery_rebuffers_the_event() { let cell_state = CellState::new(CancellationToken::new()); - let event = CellEvent::Completed { - content_items: Vec::new(), - result_buffered_for: Duration::ZERO, + let completion = |result_buffered_for| CellEvent::Completed { + content_items: vec![OutputItem::Text { + text: "done".to_string(), + }], + result_buffered_for, error_text: None, }; assert_eq!( - cell_state.commit_completion( - event.clone(), - /*pending_initial_yield_items*/ None, - || {} - ), + cell_state.commit_completion(completion(Duration::ZERO), || {}), CompletionCommit::Committed ); let (response_tx, response_rx) = oneshot::channel(); @@ -560,159 +548,20 @@ async fn failed_completion_delivery_rebuffers_the_event() { )); assert!(cell_state.accepting_observations()); - 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(event))); -} - -#[tokio::test(start_paused = true)] -async 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(), - }], - result_buffered_for: Duration::ZERO, - 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))); -} - -#[tokio::test(start_paused = true)] -async 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(), - }], - result_buffered_for: Duration::ZERO, - 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 buffered = Duration::from_secs(30); - tokio::time::advance(buffered).await; - 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(), - }, - ], - result_buffered_for: buffered, - error_text: None, - })) - ); -} - -#[tokio::test(start_paused = true)] -async 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(), - }], - result_buffered_for: Duration::ZERO, - 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 interval = Duration::from_secs(5); + tokio::time::advance(interval).await; let (response_tx, response_rx) = oneshot::channel(); drop(response_rx); assert!(matches!( - cell_state.route_observation(ObserveMode::PendingFrontier, response_tx), + cell_state.route_observation(response_tx), ObservationDelivery::Buffered )); + tokio::time::advance(interval).await; 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), + cell_state.route_observation(response_tx), ObservationDelivery::Delivered )); - assert_eq!(response_rx.try_recv(), Ok(Ok(completion))); + assert_eq!(response_rx.try_recv(), Ok(Ok(completion(2 * interval)))); } diff --git a/codex-rs/code-mode-runtime/src/cell_actor/types.rs b/codex-rs/code-mode-runtime/src/cell_actor/types.rs index 3f6be17aac..179022c3f0 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/types.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/types.rs @@ -10,10 +10,9 @@ use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; +use super::send_cell_event; use crate::session_runtime::CellEvent; use crate::session_runtime::Observation; -use crate::session_runtime::ObserveMode; -use crate::session_runtime::OutputItem; use crate::session_runtime::ToolKind; use crate::session_runtime::ToolName; @@ -57,7 +56,6 @@ 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; @@ -121,8 +119,6 @@ enum CellPhase { response_tx: oneshot::Sender>, }, Completed { - // Set only when `yield_control()` races the create-to-first-observe handoff. - pending_initial_yield_items: Option>, event: CellEvent, completed_at: Instant, }, @@ -186,11 +182,10 @@ impl CellState { Box::pin(async { Err(CellError::AlreadyTerminating) }) } CellPhase::Completed { - pending_initial_yield_items, event, completed_at, } => { - let event = buffered_completion(event, pending_initial_yield_items, completed_at); + let event = buffered_completion(event, completed_at); *phase = CellPhase::CompletionClaimed(event.clone()); self.cancellation_token.cancel(); ready_event(event) @@ -206,7 +201,6 @@ impl CellState { pub(crate) fn commit_completion( &self, event: CellEvent, - pending_initial_yield_items: Option>, commit: impl FnOnce(), ) -> CompletionCommit { let mut phase = self @@ -218,7 +212,6 @@ impl CellState { } commit(); *phase = CellPhase::Completed { - pending_initial_yield_items, event, completed_at: Instant::now(), }; @@ -233,48 +226,40 @@ impl CellState { .phase .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let (pending_initial_yield_items, event, completed_at) = - match std::mem::replace(&mut *phase, CellPhase::Tombstone) { - CellPhase::Completed { - pending_initial_yield_items, - event, - completed_at, - } => (pending_initial_yield_items, event, completed_at), - previous => { - *phase = previous; - return CompletionDelivery::Rejected(response_tx); - } - }; + let (event, completed_at) = match std::mem::replace(&mut *phase, CellPhase::Tombstone) { + CellPhase::Completed { + event, + completed_at, + } => (event, completed_at), + previous => { + *phase = previous; + return CompletionDelivery::Rejected(response_tx); + } + }; let Some(response_tx) = response_tx else { *phase = CellPhase::Completed { - pending_initial_yield_items, event, completed_at, }; return CompletionDelivery::Buffered; }; - match response_tx.send(Ok(event)) { + match send_cell_event(response_tx, event) { Ok(()) => { self.cancellation_token.cancel(); CompletionDelivery::Delivered } - Err(Ok(event)) => { + Err(event) => { *phase = CellPhase::Completed { - pending_initial_yield_items, event, completed_at, }; CompletionDelivery::Buffered } - Err(Err(error)) => { - panic!("completion delivery unexpectedly carried an actor error: {error:?}") - } } } pub(crate) fn route_observation( &self, - mode: ObserveMode, response_tx: oneshot::Sender>, ) -> ObservationDelivery { let mut phase = self @@ -287,63 +272,21 @@ impl CellState { ObservationDelivery::Running(response_tx) } CellPhase::Completed { - pending_initial_yield_items: Some(content_items), event, completed_at, - } if matches!(mode, ObserveMode::YieldAfter(_)) => { - match response_tx.send(Ok(CellEvent::Yielded { content_items })) { - Ok(()) => { - *phase = CellPhase::Completed { - pending_initial_yield_items: None, - event, - completed_at, - }; - ObservationDelivery::Buffered - } - Err(Ok(CellEvent::Yielded { content_items })) => { - *phase = CellPhase::Completed { - pending_initial_yield_items: Some(content_items), - event, - completed_at, - }; - 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:?}") - } + } => match send_cell_event(response_tx, buffered_completion(event, completed_at)) { + Ok(()) => { + self.cancellation_token.cancel(); + ObservationDelivery::Delivered } - } - CellPhase::Completed { - pending_initial_yield_items, - event, - completed_at, - } => { - let delivered_event = buffered_completion( - event.clone(), - pending_initial_yield_items.clone(), - completed_at, - ); - match response_tx.send(Ok(delivered_event)) { - Ok(()) => { - self.cancellation_token.cancel(); - ObservationDelivery::Delivered - } - Err(Ok(_)) => { - *phase = CellPhase::Completed { - pending_initial_yield_items, - event, - completed_at, - }; - ObservationDelivery::Buffered - } - Err(Err(error)) => { - panic!("completion delivery unexpectedly carried an actor error: {error:?}") - } + Err(event) => { + *phase = CellPhase::Completed { + event, + completed_at, + }; + ObservationDelivery::Buffered } - } + }, CellPhase::Terminating { response_tx: termination_tx, } => { @@ -377,14 +320,9 @@ impl CellState { Some(event) } CellPhase::Completed { - pending_initial_yield_items, event, completed_at, - } => Some(buffered_completion( - event, - pending_initial_yield_items, - completed_at, - )), + } => Some(buffered_completion(event, completed_at)), CellPhase::CompletionClaimed(completed_event) => Some(completed_event), CellPhase::Tombstone => None, }; @@ -405,27 +343,13 @@ impl CellState { } } -fn buffered_completion( - mut event: CellEvent, - pending_initial_yield_items: Option>, - completed_at: Instant, -) -> CellEvent { - let content_items = match &mut event { - CellEvent::Completed { - content_items, - result_buffered_for, - .. - } => { - *result_buffered_for = completed_at.elapsed(); - content_items - } - CellEvent::Yielded { content_items } - | CellEvent::Pending { content_items, .. } - | CellEvent::Terminated { content_items } => content_items, - }; - if let Some(mut initial) = pending_initial_yield_items { - initial.append(content_items); - *content_items = initial; +fn buffered_completion(mut event: CellEvent, completed_at: Instant) -> CellEvent { + if let CellEvent::Completed { + result_buffered_for, + .. + } = &mut event + { + *result_buffered_for = completed_at.elapsed(); } event } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index 9b27e72f3d..c500b816bd 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs @@ -287,7 +287,6 @@ 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(); @@ -298,7 +297,7 @@ impl CellHost for RuntimeCellHost { } stored_values = self.inner.stored_values.lock() => stored_values, }; - cell_state.commit_completion(event, pending_initial_yield_items, || { + cell_state.commit_completion(event, || { stored_values.extend(stored_value_writes); }) } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs index 9d8a3f28a1..af413b1916 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs @@ -119,7 +119,6 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa Arc::new(JsonValue::String("lost".to_string())), )]), completion.clone(), - /*pending_initial_yield_items*/ None, Arc::clone(&cell_state), ); tokio::pin!(commit);