openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

code-mode: preserve dropped observation output (#29288)

Browse Source 
## 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`.
This commit is contained in:
Channing Conger
2026-06-21 13:53:37 -07:00
committed by GitHub
parent 6bfc58a688
commit 3b605b9c63
2 changed files with 337 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
codex-rs/code-mode/src/cell_actor/mod.rs
View File

@@ -118,6 +118,7 @@ async fn run_cell<H: CellHost>(
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<H: CellHost>(
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<H: CellHost>(
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<H: CellHost>(
}
} => {
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<H: CellHost>(
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<H: CellHost>(
&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<H: CellHost>(
}
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<H: CellHost>(
host.closed().await;
}
fn send_observer_event(observer: Option<Observer>, event: CellEvent) {
if let Some(observer) = observer {
let _ = observer.response_tx.send(Ok(event));
fn send_observer_event(observer: Option<Observer>, 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<Result<CellEvent, CellError>>,
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<OutputItem>) {
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:?}"),
}
}

246
codex-rs/code-mode/src/cell_actor/tests.rs
View File

@@ -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<JsonValue, String> {
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<String, JsonValue>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> CompletionCommit {
cell_state.commit_completion(event, || {})
}
async fn closed(&self) {}
}
struct CellActorHarness {
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
handle: CellHandle,
initial_event_rx: oneshot::Receiver<Result<CellEvent, CellError>>,
task: tokio::task::JoinHandle<()>,
runtime_control_rx: std_mpsc::Receiver<RuntimeControlCommand>,
_runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
}
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<H: CellHost>(
initial_observe_mode: ObserveMode,
host: Arc<H>,
) -> 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());