mirror of
https://github.com/openai/codex.git
synced 2026-10-23 18:08:45 +00:00
code-mode: clean up terminal cell dispatch gates
This commit is contained in:
@@ -341,7 +341,7 @@ async fn run_cell<H: CellHost>(
|
||||
),
|
||||
error_text: Some("exec runtime ended unexpectedly".to_string()),
|
||||
};
|
||||
let rejected_event = match host
|
||||
let (completion_committed, rejected_event) = match host
|
||||
.commit_completion(
|
||||
HashMap::new(),
|
||||
event,
|
||||
@@ -350,14 +350,18 @@ async fn run_cell<H: CellHost>(
|
||||
)
|
||||
.await
|
||||
{
|
||||
CompletionCommit::Committed => None,
|
||||
CompletionCommit::Rejected(event) => Some(event),
|
||||
CompletionCommit::Committed => (true, None),
|
||||
CompletionCommit::Rejected(event) => (false, Some(event)),
|
||||
};
|
||||
match cell_state.deliver_completion(
|
||||
let completion_delivery = cell_state.deliver_completion(
|
||||
observer
|
||||
.take()
|
||||
.map(|observer| (observer.mode, observer.response_tx)),
|
||||
) {
|
||||
);
|
||||
if completion_committed {
|
||||
host.terminal();
|
||||
}
|
||||
match completion_delivery {
|
||||
CompletionDelivery::Delivered => break,
|
||||
CompletionDelivery::Buffered => {}
|
||||
CompletionDelivery::Rejected(response_tx) => {
|
||||
@@ -510,7 +514,7 @@ async fn run_cell<H: CellHost>(
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
error_text,
|
||||
};
|
||||
let rejected_event = match host
|
||||
let (completion_committed, rejected_event) = match host
|
||||
.commit_completion(
|
||||
stored_value_writes,
|
||||
event,
|
||||
@@ -519,14 +523,18 @@ async fn run_cell<H: CellHost>(
|
||||
)
|
||||
.await
|
||||
{
|
||||
CompletionCommit::Committed => None,
|
||||
CompletionCommit::Rejected(event) => Some(event),
|
||||
CompletionCommit::Committed => (true, None),
|
||||
CompletionCommit::Rejected(event) => (false, Some(event)),
|
||||
};
|
||||
match cell_state.deliver_completion(
|
||||
let completion_delivery = cell_state.deliver_completion(
|
||||
observer
|
||||
.take()
|
||||
.map(|observer| (observer.mode, observer.response_tx)),
|
||||
) {
|
||||
);
|
||||
if completion_committed {
|
||||
host.terminal();
|
||||
}
|
||||
match completion_delivery {
|
||||
CompletionDelivery::Delivered => break,
|
||||
CompletionDelivery::Buffered => {}
|
||||
CompletionDelivery::Rejected(response_tx) => {
|
||||
|
||||
@@ -52,6 +52,8 @@ impl CellHost for TestHost {
|
||||
cell_state.commit_completion(event, pending_initial_yield_items, || {})
|
||||
}
|
||||
|
||||
fn terminal(&self) {}
|
||||
|
||||
async fn closed(&self) {}
|
||||
}
|
||||
|
||||
@@ -86,6 +88,8 @@ impl CellHost for RecordingHost {
|
||||
})
|
||||
}
|
||||
|
||||
fn terminal(&self) {}
|
||||
|
||||
async fn closed(&self) {}
|
||||
}
|
||||
|
||||
|
||||
@@ -89,6 +89,9 @@ pub(crate) trait CellHost: Send + Sync + 'static {
|
||||
cell_state: Arc<CellState>,
|
||||
) -> impl Future<Output = CompletionCommit> + Send;
|
||||
|
||||
/// Reports that execution reached a terminal state, even when its result remains buffered.
|
||||
fn terminal(&self);
|
||||
|
||||
fn closed(&self) -> impl Future<Output = ()> + Send;
|
||||
}
|
||||
|
||||
|
||||
@@ -74,6 +74,8 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
|
||||
) -> NotificationFuture<'a> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -402,6 +404,10 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn cell_closed(&self, cell_id: &runtime::CellId) {
|
||||
self.delegate.cell_closed(&protocol_cell_id(cell_id));
|
||||
}
|
||||
}
|
||||
|
||||
fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest {
|
||||
|
||||
@@ -21,6 +21,11 @@ enum DelegateEvent {
|
||||
NotificationFinished,
|
||||
ToolStarted,
|
||||
ToolCancelled,
|
||||
CellClosed(CellId),
|
||||
}
|
||||
|
||||
fn record_cell_closed(events_tx: &mpsc::UnboundedSender<DelegateEvent>, cell_id: &CellId) {
|
||||
let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone()));
|
||||
}
|
||||
|
||||
struct BlockingDelegate {
|
||||
@@ -101,6 +106,10 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate {
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn cell_closed(&self, cell_id: &CellId) {
|
||||
record_cell_closed(&self.events_tx, cell_id);
|
||||
}
|
||||
}
|
||||
|
||||
impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
|
||||
@@ -135,6 +144,10 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn cell_closed(&self, cell_id: &CellId) {
|
||||
record_cell_closed(&self.events_tx, cell_id);
|
||||
}
|
||||
}
|
||||
|
||||
impl BlockingDelegate {
|
||||
@@ -193,6 +206,10 @@ impl CodeModeSessionDelegate for BlockingDelegate {
|
||||
Err("cancelled".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn cell_closed(&self, cell_id: &CellId) {
|
||||
record_cell_closed(&self.events_tx, cell_id);
|
||||
}
|
||||
}
|
||||
|
||||
fn cell_id(value: &str) -> CellId {
|
||||
@@ -427,6 +444,40 @@ await tools.block({});
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn background_completion_notifies_the_delegate_without_another_observation() {
|
||||
let (delegate, mut events_rx) = BlockingDelegate::new();
|
||||
let service = CodeModeService::with_delegate(delegate);
|
||||
let created_cell_id = service
|
||||
.create_cell(execute_request(
|
||||
r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
service.observe(ObserveRequest {
|
||||
idempotency_key: "background-completion".to_string(),
|
||||
cell_id: created_cell_id.clone(),
|
||||
yield_time_ms: 1,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.expect("initial observation should yield while the cell is still running")
|
||||
.unwrap(),
|
||||
ObserveOutcome::Yielded {
|
||||
cell_id: created_cell_id.clone(),
|
||||
content_items: Vec::new(),
|
||||
}
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(created_cell_id)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_and_resumes_from_the_pending_frontier() {
|
||||
let (delegate, mut events_rx) = BlockingDelegate::new();
|
||||
@@ -566,6 +617,10 @@ async fn termination_cancels_pending_callbacks_before_responding() {
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::NotificationCancelled
|
||||
);
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -599,6 +654,10 @@ await tools.block({});
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::ToolCancelled
|
||||
);
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
execute_with_yield_time(
|
||||
@@ -641,6 +700,10 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
|
||||
delegate.release_notification();
|
||||
|
||||
assert_eq!(shutdown.await.unwrap(), Ok(()));
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -687,6 +750,10 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
|
||||
content_items: Vec::new(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -730,6 +797,10 @@ async fn create_cell_returns_before_natural_completion() {
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::NotificationFinished
|
||||
);
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -840,4 +911,8 @@ async fn natural_completion_cleans_up_callbacks_before_responding() {
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::ToolCancelled
|
||||
);
|
||||
assert_eq!(
|
||||
next_event(&mut events_rx).await,
|
||||
DelegateEvent::CellClosed(cell_id("1"))
|
||||
);
|
||||
}
|
||||
|
||||
@@ -69,6 +69,8 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
|
||||
) -> NotificationFuture<'a> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
@@ -83,6 +85,9 @@ enum RecordedDelegateCall {
|
||||
text: String,
|
||||
cancellation_requested: bool,
|
||||
},
|
||||
CellClosed {
|
||||
cell_id: CellId,
|
||||
},
|
||||
}
|
||||
|
||||
struct RecordingDelegate {
|
||||
@@ -151,6 +156,15 @@ impl CodeModeSessionDelegate for RecordingDelegate {
|
||||
.expect("test must provide one result per notification")
|
||||
})
|
||||
}
|
||||
|
||||
fn cell_closed(&self, cell_id: &CellId) {
|
||||
self.calls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(RecordedDelegateCall::CellClosed {
|
||||
cell_id: cell_id.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn execute_request(source: &str) -> CreateCellRequest {
|
||||
@@ -1215,6 +1229,7 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
|
||||
.await,
|
||||
Err("notification failed".to_string())
|
||||
);
|
||||
runtime::SessionRuntimeDelegate::cell_closed(&adapter, &runtime::CellId::new("cell-d4"));
|
||||
|
||||
assert_eq!(
|
||||
delegate.take_calls(),
|
||||
@@ -1245,6 +1260,9 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
|
||||
text: "progress".to_string(),
|
||||
cancellation_requested: true,
|
||||
},
|
||||
RecordedDelegateCall::CellClosed {
|
||||
cell_id: cell_id("cell-d4"),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
@@ -273,6 +274,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
|
||||
let host = Arc::new(RuntimeCellHost {
|
||||
cell_id: cell_id.clone(),
|
||||
inner: Arc::clone(&self.inner),
|
||||
terminal_notified: AtomicBool::new(false),
|
||||
});
|
||||
let mut cells = self.inner.cells.lock().await;
|
||||
if self.inner.shutdown_token.is_cancelled() {
|
||||
@@ -380,6 +382,7 @@ impl PendingEvent {
|
||||
struct RuntimeCellHost<D: SessionRuntimeDelegate> {
|
||||
cell_id: CellId,
|
||||
inner: Arc<Inner<D>>,
|
||||
terminal_notified: AtomicBool,
|
||||
}
|
||||
|
||||
impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
|
||||
@@ -435,8 +438,15 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
|
||||
})
|
||||
}
|
||||
|
||||
fn terminal(&self) {
|
||||
if !self.terminal_notified.swap(true, Ordering::AcqRel) {
|
||||
self.inner.delegate.cell_closed(&self.cell_id);
|
||||
}
|
||||
}
|
||||
|
||||
async fn closed(&self) {
|
||||
self.inner.cells.lock().await.remove(&self.cell_id);
|
||||
self.terminal();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
use std::task::Waker;
|
||||
@@ -48,6 +49,8 @@ impl SessionRuntimeDelegate for RecordingDelegate {
|
||||
) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
impl SessionRuntimeDelegate for ImmediateToolDelegate {
|
||||
@@ -69,6 +72,8 @@ impl SessionRuntimeDelegate for ImmediateToolDelegate {
|
||||
) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
impl SessionRuntimeDelegate for BlockingToolDelegate {
|
||||
@@ -95,6 +100,8 @@ impl SessionRuntimeDelegate for BlockingToolDelegate {
|
||||
) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
impl SessionRuntimeDelegate for NonCooperativeToolDelegate {
|
||||
@@ -116,6 +123,8 @@ impl SessionRuntimeDelegate for NonCooperativeToolDelegate {
|
||||
) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cell_closed(&self, _cell_id: &CellId) {}
|
||||
}
|
||||
|
||||
fn tool_definition(name: &str) -> ToolDefinition {
|
||||
@@ -568,6 +577,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
|
||||
let host = RuntimeCellHost {
|
||||
cell_id: CellId::new("terminating-writer"),
|
||||
inner: Arc::clone(&runtime.inner),
|
||||
terminal_notified: AtomicBool::new(false),
|
||||
};
|
||||
let completion = ActorEvent::Completed {
|
||||
content_items: vec![OutputItem::Text {
|
||||
|
||||
@@ -206,7 +206,8 @@ pub struct NestedToolCall {
|
||||
///
|
||||
/// Implementations should forward callback cancellation tokens to downstream
|
||||
/// work. After cancellation begins, the runtime allows callbacks a bounded
|
||||
/// grace period to finish, then aborts their local tasks.
|
||||
/// grace period to finish, then aborts their local tasks. Terminal notifications
|
||||
/// must be non-blocking and are not acknowledged or retried.
|
||||
pub trait SessionRuntimeDelegate: Send + Sync + 'static {
|
||||
fn invoke_tool(
|
||||
&self,
|
||||
@@ -221,6 +222,11 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static {
|
||||
text: String,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> impl Future<Output = Result<(), String>> + Send;
|
||||
|
||||
/// Reports that execution is terminal and will issue no more callbacks.
|
||||
///
|
||||
/// The terminal result may remain buffered for a later observation.
|
||||
fn cell_closed(&self, cell_id: &CellId);
|
||||
}
|
||||
|
||||
/// A failure reported by a session runtime operation.
|
||||
|
||||
Reference in New Issue
Block a user