openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify code-mode cell completion delivery (#57229)

Browse Source 
## What changed

Remove unused pending initial-yield state and its plumbing from cell completion handling. Route completion delivery through `send_cell_event`, retaining the completion event and original timestamp when an observer has disconnected.

## Testing

Extend `failed_completion_delivery_rebuffers_the_event` to verify that failed delivery attempts preserve output and report the full elapsed buffering duration on a later successful observation.

GitOrigin-RevId: eebe990a4f8258f1ccb12d0694945678f0cc3a8f
This commit is contained in:
Channing Conger
2026-10-26 20:18:45 +00:00
committed by copyberry
parent 7948c07531
commit 0b021a5fc9
6 changed files with 51 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs
-1
View File
@@ -41,7 +41,6 @@ impl CellHost for PanickingCallbackHost {
&self,
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
_event: CellEvent,
_pending_initial_yield_items: Option<Vec<crate::session_runtime::OutputItem>>,
_cell_state: Arc<CellState>,
) -> CompletionCommit {
panic!("unexpected completion commit");
codex-rs/code-mode-runtime/src/cell_actor/mod.rs
+1 -3
View File
@@ -184,7 +184,7 @@ async fn run_cell<H: CellHost>(
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<H: CellHost>(
.commit_completion(
HashMap::new(),
event,
/*pending_initial_yield_items*/ None,
Arc::clone(&cell_state),
)
.await
@@ -479,7 +478,6 @@ 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-runtime/src/cell_actor/tests.rs
+16 -167
View File
@@ -45,10 +45,9 @@ impl CellHost for TestHost {
&self,
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
event: CellEvent,
pending_initial_yield_items: Option<Vec<OutputItem>>,
cell_state: Arc<CellState>,
) -> 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<String, Arc<JsonValue>>,
event: CellEvent,
pending_initial_yield_items: Option<Vec<OutputItem>>,
cell_state: Arc<CellState>,
) -> 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))));
}
codex-rs/code-mode-runtime/src/cell_actor/types.rs
+33 -109
View File
@@ -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<String, Arc<JsonValue>>,
event: CellEvent,
pending_initial_yield_items: Option<Vec<OutputItem>>,
cell_state: Arc<CellState>,
) -> impl Future<Output = CompletionCommit> + Send;
@@ -121,8 +119,6 @@ enum CellPhase {
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
},
Completed {
// Set only when `yield_control()` races the create-to-first-observe handoff.
pending_initial_yield_items: Option<Vec<OutputItem>>,
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<Vec<OutputItem>>,
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<Result<CellEvent, CellError>>,
) -> 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<Vec<OutputItem>>,
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
}
codex-rs/code-mode-runtime/src/session_runtime/mod.rs
+1 -2
View File
@@ -287,7 +287,6 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
&self,
stored_value_writes: HashMap<String, Arc<JsonValue>>,
event: CellEvent,
pending_initial_yield_items: Option<Vec<OutputItem>>,
cell_state: Arc<CellState>,
) -> CompletionCommit {
let cancellation_token = cell_state.cancellation_token();
@@ -298,7 +297,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
}
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);
})
}
codex-rs/code-mode-runtime/src/session_runtime/tests.rs
-1
View File
@@ -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);