From 3b605b9c63f9123b39493b24d95d9a8428813e9f Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Sun, 21 Jun 2026 13:53:37 -0700 Subject: [PATCH] code-mode: preserve dropped observation output (#29288) ## Summary - Restore yielded output when an observation receiver disappears before delivery. - Preserve pending-frontier output and tool IDs across failed delivery. - Add dropped-observer coverage for yield and pending observations. ## Why Canceling a wait must not consume output or a pending frontier that the caller never received. ## Impact A later observation can recover undelivered incremental output without duplication. ## Validation - Stack-tip validation: `just test -p codex-code-mode -p codex-code-mode-protocol` (70 passed). - Parent branch: `cconger/code-mode-runtime-compact-03e-shutdown-hierarchy`. --- codex-rs/code-mode/src/cell_actor/mod.rs | 112 ++++++++-- codex-rs/code-mode/src/cell_actor/tests.rs | 246 ++++++++++++++++++++- 2 files changed, 337 insertions(+), 21 deletions(-) diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index d288bd2db..8379f372d 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -118,6 +118,7 @@ async fn run_cell( let callback_cancellation_token = cancellation_token.child_token(); let mut content_items = Vec::new(); let mut pending_tool_call_ids = Vec::new(); + let mut pending_frontier_ready = false; let mut observer = Some(initial_observer); let mut termination = false; let mut runtime_closed = false; @@ -169,6 +170,9 @@ async fn run_cell( cancellation_token.cancel(); continue; }; + if response_tx.is_closed() { + continue; + } let response_tx = match cell_state.route_observation(response_tx) { ObservationDelivery::Running(response_tx) => response_tx, ObservationDelivery::Delivered => break, @@ -185,8 +189,36 @@ async fn run_cell( let _ = response_tx.send(Err(CellError::Busy)); continue; } + if matches!(mode, ObserveMode::PendingFrontier) && pending_frontier_ready { + pending_frontier_ready = false; + match send_cell_event( + response_tx, + CellEvent::Pending { + content_items: std::mem::take(&mut content_items), + pending_tool_call_ids: std::mem::take(&mut pending_tool_call_ids), + }, + ) { + Ok(()) => {} + Err(CellEvent::Pending { + content_items: undelivered_items, + pending_tool_call_ids: undelivered_tool_call_ids, + }) => { + content_items = undelivered_items; + pending_tool_call_ids = undelivered_tool_call_ids; + pending_frontier_ready = true; + } + Err(event) => { + panic!("pending delivery returned an unexpected event: {event:?}") + } + } + continue; + } observer = Some(Observer { mode, response_tx }); yield_timer = observer.as_ref().and_then(observer_timer); + if runtime_paused && matches!(mode, ObserveMode::YieldAfter(_)) { + pending_frontier_ready = false; + pending_tool_call_ids.clear(); + } resume_for_observation( mode, &mut runtime_paused, @@ -202,11 +234,14 @@ async fn run_cell( } } => { yield_timer = None; - send_observer_event( - observer.take(), - CellEvent::Yielded { - content_items: std::mem::take(&mut content_items), - }, + restore_undelivered_yield( + send_observer_event( + observer.take(), + CellEvent::Yielded { + content_items: std::mem::take(&mut content_items), + }, + ), + &mut content_items, ); } maybe_event = async { @@ -285,7 +320,8 @@ async fn run_cell( Some(ObserveMode::PendingFrontier) ) { yield_timer = None; - send_observer_event( + pending_frontier_ready = false; + match send_observer_event( observer.take(), CellEvent::Pending { content_items: std::mem::take(&mut content_items), @@ -293,7 +329,20 @@ async fn run_cell( &mut pending_tool_call_ids, ), }, - ); + ) { + Ok(()) => {} + Err(CellEvent::Pending { + content_items: undelivered_items, + pending_tool_call_ids: undelivered_tool_call_ids, + }) => { + content_items = undelivered_items; + pending_tool_call_ids = undelivered_tool_call_ids; + pending_frontier_ready = true; + } + Err(event) => { + panic!("pending delivery returned an unexpected event: {event:?}") + } + } } else { pending_tool_call_ids.clear(); let _ = runtime_control_tx.send(RuntimeControlCommand::Continue); @@ -302,16 +351,20 @@ async fn run_cell( } RuntimeEvent::ContentItem(item) => content_items.push(output_item(item)), RuntimeEvent::YieldRequested => { - if matches!( + let yield_observer = matches!( observer.as_ref().map(|observer| observer.mode), Some(ObserveMode::YieldAfter(_)) - ) { + ); + if yield_observer { yield_timer = None; - send_observer_event( - observer.take(), - CellEvent::Yielded { - content_items: std::mem::take(&mut content_items), - }, + restore_undelivered_yield( + send_observer_event( + observer.take(), + CellEvent::Yielded { + content_items: std::mem::take(&mut content_items), + }, + ), + &mut content_items, ); } } @@ -429,9 +482,34 @@ async fn run_cell( host.closed().await; } -fn send_observer_event(observer: Option, event: CellEvent) { - if let Some(observer) = observer { - let _ = observer.response_tx.send(Ok(event)); +fn send_observer_event(observer: Option, event: CellEvent) -> Result<(), CellEvent> { + let Some(observer) = observer else { + return Err(event); + }; + send_cell_event(observer.response_tx, event) +} + +fn send_cell_event( + response_tx: oneshot::Sender>, + event: CellEvent, +) -> Result<(), CellEvent> { + match response_tx.send(Ok(event)) { + Ok(()) => Ok(()), + Err(Ok(event)) => Err(event), + Err(Err(error)) => panic!("cell event delivery returned an actor error: {error:?}"), + } +} + +fn restore_undelivered_yield(delivery: Result<(), CellEvent>, content_items: &mut Vec) { + match delivery { + Ok(()) => {} + Err(CellEvent::Yielded { + content_items: mut undelivered_items, + }) => { + undelivered_items.append(content_items); + *content_items = undelivered_items; + } + Err(event) => panic!("yield delivery returned an unexpected event: {event:?}"), } } diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index 7404b9fe1..e217122b2 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -1,5 +1,8 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc as std_mpsc; use std::time::Duration; use codex_code_mode_protocol::ExecuteRequest; @@ -11,10 +14,15 @@ use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; use super::*; -use crate::session_runtime::OutputItem as CellOutputItem; +use crate::session_runtime::OutputItem; struct TestHost; +#[derive(Default)] +struct RecordingHost { + notified: AtomicBool, +} + impl CellHost for TestHost { async fn invoke_tool( &self, @@ -45,20 +53,59 @@ impl CellHost for TestHost { async fn closed(&self) {} } +impl CellHost for RecordingHost { + async fn invoke_tool( + &self, + _invocation: CellToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + Err("unexpected tool call".to_string()) + } + + async fn notify( + &self, + _call_id: String, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + self.notified.store(true, Ordering::Release); + Ok(()) + } + + async fn commit_completion( + &self, + _stored_value_writes: HashMap, + event: CellEvent, + cell_state: Arc, + ) -> CompletionCommit { + cell_state.commit_completion(event, || {}) + } + + async fn closed(&self) {} +} + struct CellActorHarness { event_tx: mpsc::UnboundedSender, handle: CellHandle, initial_event_rx: oneshot::Receiver>, task: tokio::task::JoinHandle<()>, + runtime_control_rx: std_mpsc::Receiver, _runtime_event_rx: mpsc::UnboundedReceiver, } fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness { + spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost)) +} + +fn spawn_cell_actor_harness_with_host( + initial_observe_mode: ObserveMode, + host: Arc, +) -> CellActorHarness { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); let (initial_event_tx, initial_event_rx) = oneshot::channel(); let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel(); - let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime( + let (runtime_tx, _runtime_control_tx, runtime_terminate_handle) = spawn_runtime( HashMap::new(), ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -71,10 +118,11 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne PendingRuntimeMode::PauseUntilResumed, ) .unwrap(); + let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel(); let cell_state = Arc::new(CellState::new(CancellationToken::new())); let handle = CellHandle::new(command_tx, Arc::clone(&cell_state)); let task = tokio::spawn(run_cell( - Arc::new(TestHost), + host, CellContext { runtime_tx, runtime_control_tx, @@ -94,10 +142,21 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne handle, initial_event_rx, task, + runtime_control_rx, _runtime_event_rx: runtime_event_rx, } } +async fn wait_for_notification(host: &RecordingHost) { + tokio::time::timeout(Duration::from_secs(1), async { + while !host.notified.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("notification barrier timed out"); +} + #[tokio::test] async fn yield_timer_preempts_buffered_runtime_output() { let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO)); @@ -123,7 +182,7 @@ async fn yield_timer_preempts_buffered_runtime_output() { assert_eq!( termination.await, Ok(CellEvent::Terminated { - content_items: vec![CellOutputItem::Text { + content_items: vec![OutputItem::Text { text: "queued output".to_string(), }], }) @@ -151,6 +210,185 @@ async fn queued_termination_preempts_unobserved_runtime_completion() { harness.task.await.unwrap(); } +#[tokio::test] +async fn observation_dropped_before_dequeue_does_not_consume_output() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(60)), + Arc::clone(&host), + ); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert!(harness.initial_event_rx.await.unwrap().is_ok()); + + drop( + harness + .handle + .observe(ObserveMode::YieldAfter(Duration::from_secs(60))), + ); + harness + .event_tx + .send(RuntimeEvent::ContentItem( + FunctionCallOutputContentItem::InputText { + text: "survives pre-dequeue cancellation".to_string(), + }, + )) + .unwrap(); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + harness + .event_tx + .send(RuntimeEvent::Notify { + call_id: "after-dropped-command".to_string(), + text: "barrier".to_string(), + }) + .unwrap(); + wait_for_notification(&host).await; + + assert_eq!( + harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)) + .await, + Ok(CellEvent::Yielded { + content_items: vec![OutputItem::Text { + text: "survives pre-dequeue cancellation".to_string(), + }], + }) + ); + + let termination = harness.handle.terminate(); + drop(harness.event_tx); + assert_eq!( + termination.await, + Ok(CellEvent::Terminated { + content_items: Vec::new(), + }) + ); + harness.task.await.unwrap(); +} + +#[tokio::test] +async fn dropped_yield_observer_preserves_output_for_the_next_observation() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(60)), + Arc::clone(&host), + ); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert!(harness.initial_event_rx.await.unwrap().is_ok()); + + let dropped_observation = harness + .handle + .observe(ObserveMode::YieldAfter(Duration::from_secs(60))); + assert_eq!( + harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)) + .await, + Err(CellError::Busy) + ); + drop(dropped_observation); + harness + .event_tx + .send(RuntimeEvent::ContentItem( + FunctionCallOutputContentItem::InputText { + text: "survives active cancellation".to_string(), + }, + )) + .unwrap(); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + harness + .event_tx + .send(RuntimeEvent::Notify { + call_id: "after-dropped-observer".to_string(), + text: "barrier".to_string(), + }) + .unwrap(); + wait_for_notification(&host).await; + + assert_eq!( + harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)) + .await, + Ok(CellEvent::Yielded { + content_items: vec![OutputItem::Text { + text: "survives active cancellation".to_string(), + }], + }) + ); + + let termination = harness.handle.terminate(); + drop(harness.event_tx); + assert_eq!( + termination.await, + Ok(CellEvent::Terminated { + content_items: Vec::new(), + }) + ); + harness.task.await.unwrap(); +} + +#[tokio::test] +async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observation() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(60)), + Arc::clone(&host), + ); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert!(harness.initial_event_rx.await.unwrap().is_ok()); + + let dropped_observation = harness.handle.observe(ObserveMode::PendingFrontier); + assert_eq!( + harness.handle.observe(ObserveMode::PendingFrontier).await, + Err(CellError::Busy) + ); + drop(dropped_observation); + harness + .event_tx + .send(RuntimeEvent::ToolCall { + id: "tool-1".to_string(), + name: codex_protocol::ToolName { + name: "echo".to_string(), + namespace: None, + }, + kind: codex_code_mode_protocol::CodeModeToolKind::Function, + input: Some(serde_json::json!({})), + }) + .unwrap(); + harness.event_tx.send(RuntimeEvent::Pending).unwrap(); + harness + .event_tx + .send(RuntimeEvent::Notify { + call_id: "after-dropped-pending".to_string(), + text: "barrier".to_string(), + }) + .unwrap(); + wait_for_notification(&host).await; + + assert_eq!( + harness.handle.observe(ObserveMode::PendingFrontier).await, + Ok(CellEvent::Pending { + content_items: Vec::new(), + pending_tool_call_ids: vec!["tool-1".to_string()], + }) + ); + assert!(matches!( + harness.runtime_control_rx.try_recv(), + Err(std_mpsc::TryRecvError::Empty) + )); + + let termination = harness.handle.terminate(); + drop(harness.event_tx); + assert_eq!( + termination.await, + Ok(CellEvent::Terminated { + content_items: Vec::new(), + }) + ); + harness.task.await.unwrap(); +} + #[tokio::test] async fn only_the_first_termination_claims_a_buffered_completion() { let cell_state = CellState::new(CancellationToken::new());